diff --git a/.github/workflows/format.yaml b/.github/workflows/format.yaml old mode 100644 new mode 100755 index a7475748..7d6fc637 --- a/.github/workflows/format.yaml +++ b/.github/workflows/format.yaml @@ -43,12 +43,28 @@ jobs: - name: Install clang-format run: | sudo apt-get update - sudo apt-get install -y clang-format-17 + sudo apt-get install -y clang-format-15 - name: Check changed files run: | + set -euo pipefail + + if [ "${{ github.event_name }}" = "pull_request" ]; then + diff_range="${{ github.event.pull_request.base.sha }}...HEAD" + elif [ "${{ github.event_name }}" = "push" ] && \ + [ -n "${{ github.event.before }}" ] && \ + [ "${{ github.event.before }}" != "0000000000000000000000000000000000000000" ]; then + diff_range="${{ github.event.before }}..${{ github.sha }}" + else + default_branch="${{ github.event.repository.default_branch }}" + git fetch --no-tags --prune origin \ + "+refs/heads/${default_branch}:refs/remotes/origin/${default_branch}" + diff_range="origin/${default_branch}...HEAD" + fi + + echo "Checking files changed in ${diff_range}" git diff --diff-filter=d --name-only \ - origin/${{ github.base_ref }}...HEAD \ + "${diff_range}" -- \ > all_changed_files.txt awk '/\.(cc|cpp|cxx|h|hpp)$/ && $0 !~ /^third_party\// { print }' \ @@ -60,7 +76,7 @@ jobs: exit 0 fi - xargs -a changed_files.txt clang-format-17 --dry-run -Werror + xargs -r -a changed_files.txt clang-format-15 --dry-run -Werror python-black: needs: markdownlint diff --git a/backend/ascend_autotune_hooks.py b/backend/ascend_autotune_hooks.py index f5f5eade..6a62e2bb 100644 --- a/backend/ascend_autotune_hooks.py +++ b/backend/ascend_autotune_hooks.py @@ -48,7 +48,7 @@ def _ascend_autotune_fn(): """Lazy singleton accessor for ascend ``autotune``.""" global _ASCEND_AUTOTUNE if _ASCEND_AUTOTUNE is None: - from .ascend_autotune_runtime.autotuner import autotune as _fn + from .ascend_autotune_runtime.ascend_kernel_autotuner import autotune as _fn _ASCEND_AUTOTUNE = _fn return _ASCEND_AUTOTUNE @@ -58,7 +58,7 @@ def _ascend_max_autotune_fn(): """Lazy singleton accessor for ascend ``max_autotune``.""" global _ASCEND_MAX_AUTOTUNE if _ASCEND_MAX_AUTOTUNE is None: - from .ascend_autotune_runtime.autotuner import max_autotune as _fn + from .ascend_autotune_runtime.ascend_kernel_autotuner import max_autotune as _fn _ASCEND_MAX_AUTOTUNE = _fn return _ASCEND_MAX_AUTOTUNE @@ -68,7 +68,7 @@ def _ascend_max_autotune_fn(): # Proxies (installed once on import of this module) # ------------------------------------------------------------------ def _autotune_proxy(configs, key, **kwargs): - if _is_ascend_backend(): + if _is_ascend_backend() or kwargs.get("hints") is not None: return _ascend_autotune_fn()(configs=configs, key=key, **kwargs) from triton.runtime.autotuner import autotune as _stock diff --git a/backend/ascend_autotune_runtime/__init__.py b/backend/ascend_autotune_runtime/__init__.py index b5a99865..36345bba 100644 --- a/backend/ascend_autotune_runtime/__init__.py +++ b/backend/ascend_autotune_runtime/__init__.py @@ -4,7 +4,7 @@ Adapted for DLCompiler (triton.backends.dicp_triton). """ -from .autoparser import ( +from .kernel_ast_analyzer import ( AutoParser, AxesKeyParser, SplitAxesParser, @@ -13,13 +13,22 @@ LowDimsAxesParser, PtrNumsParser, ) -from .tile_generator import AxisInfo, BlockInfo, KernelMeta, TileGenerator -from .compile_options import ( +from .tile_candidate_generator import AxisInfo, BlockInfo, KernelMeta, TileGenerator +from .schedule_profiles import ( CompileOptionsSpec, + CompileFailureRegionSet, + classify_compile_failure, + compile_profile_to_config, + effective_compile_profile_key, expand_compile_option_configs, + generate_linked_compile_neighbors, + get_stage1_probe_configs, + get_stage1_probe_profiles, + make_stage2_seed_profiles, parse_compile_options_hint, + validate_compile_profile, ) -from .autotuner import ( +from .ascend_kernel_autotuner import ( AutoTilingTuner, autotune, max_autotune, @@ -46,8 +55,17 @@ "KernelMeta", "TileGenerator", "CompileOptionsSpec", + "CompileFailureRegionSet", + "classify_compile_failure", + "compile_profile_to_config", + "effective_compile_profile_key", "expand_compile_option_configs", + "generate_linked_compile_neighbors", + "get_stage1_probe_configs", + "get_stage1_probe_profiles", + "make_stage2_seed_profiles", "parse_compile_options_hint", + "validate_compile_profile", "AutoTilingTuner", "autotune", "max_autotune", diff --git a/backend/ascend_autotune_runtime/autotuner.py b/backend/ascend_autotune_runtime/ascend_kernel_autotuner.py similarity index 52% rename from backend/ascend_autotune_runtime/autotuner.py rename to backend/ascend_autotune_runtime/ascend_kernel_autotuner.py index 00a4f167..736179b1 100644 --- a/backend/ascend_autotune_runtime/autotuner.py +++ b/backend/ascend_autotune_runtime/ascend_kernel_autotuner.py @@ -27,6 +27,7 @@ import functools import ast import inspect +import math import os import time from concurrent.futures import ThreadPoolExecutor @@ -39,22 +40,55 @@ from triton.runtime.autotuner import Autotuner, Config from triton.backends.dicp_triton.utils import is_compile_on_910_95 -from .autoparser import ( +from .kernel_ast_analyzer import ( DotCallParser, LowDimsAxesParser, - PtrNumsParser, ReductionAxesParser, SplitAxesParser, TilingAxesParser, ) -from .compile_options import ( +from .schedule_profiles import ( + apply_fixed_compile_options_to_profile, + classify_compile_failure, + compile_profile_to_config, + effective_compile_profile_key, + get_stage1_probe_profiles, expand_compile_option_configs, format_compile_option_result, get_compile_option_param_names, + _hashable_value, parse_compile_options_hint, summarize_compile_option_configs, ) -from .benchmark import select_benchmark_strategy +from .kernel_archetype import ( + OperatorKind, + SearchPolicy, + analyze_operator_policy, + make_search_policy, +) +from .tile_search_policy import ( + apply_no_dot_search_defaults, + filter_shapes_by_ub_and_timing, + parse_search_params_hint, + propose_evolved_shape_configs, + select_initial_percentile_shapes, + select_top_shape_entries, + stage1_children_per_round, + stage1_initial_budget, + stage1_initial_percentiles, + stage1_total_budget, + ub_timing_weighted_key, +) +from .tile_acquisition_model import resource_overflow_ratio +from .measurement_cache import SearchMeasureCache +from .tile_shape_space import ( + effective_search_value_map, + expand_search_param_shapes, + extract_shape, + shape_key, +) +from .schedule_profile_search import Stage2CompileSearcher +from .measurement_strategy import select_benchmark_strategy, ub_bytes_of from .utils import get_byte_per_numel, is_valid_axis_name, valid_axis_names @@ -110,7 +144,6 @@ def __init__( rep=None, use_cuda_graph=False, do_bench=None, - cache_results=False, auto_profile_dir=None, hints=None, ): @@ -135,7 +168,7 @@ def __init__( rep, use_cuda_graph, do_bench, - cache_results, + False, ) self.user_defined_do_bench = do_bench is not None if not hints: @@ -156,10 +189,17 @@ def __init__( ) self.auto_gen_config = not configs or self.hints.get("auto_gen_config", False) + self.operator_features = None + self.operator_policy = self._infer_operator_policy() self._infer_compile_options_hint_if_needed() self.compile_options = parse_compile_options_hint( self.hints.get("compile_options", None) ) + self.search_params = parse_search_params_hint( + self.hints.get("search_params", None) + ) + if self.operator_policy.no_dot_shape_values: + apply_no_dot_search_defaults(self.search_params) self.gen_configs = [] # generated configs from TileGenerator self.auto_profile_dir = auto_profile_dir if not configs: @@ -175,20 +215,76 @@ def __init__( self.print_autotuning_timings = ( os.getenv("TRITON_PRINT_AUTOTUNING_TIMINGS", None) == "1" ) - # Compile kernels in parallel by default for triton.runtime.JITFunction, - # but not for others, e.g., LibEntry, since it's not compatible with AsyncCompileMode + self.fixed_compile_options = {} + self._search_params_runtime_enabled = False + self._search_param_anchor_shape_keys = set() + self._search_ub_cache = {} + self._search_measure_cache = SearchMeasureCache( + self.search_params.params, self._search_ub_cache + ) + self._search_stage1_measured_summary = [] + self._search_stage1_failed_summary = [] + self._search_bench_call_count = 0 + self._search_bench_config_count = 0 + self.last_search_stats = {} + # AsyncCompileMode is available for triton.runtime.JITFunction, but not + # for wrappers such as LibEntry. + self._can_async_compile = isinstance(self.fn, triton.runtime.JITFunction) + # General autotune keeps the historical env switch. Stage 1 search_params + # probe batches force async compile independently via _batch_bench(..., + # force_parallel=True). self.compile_parallel = ( - isinstance(self.fn, triton.runtime.JITFunction) + self._can_async_compile and os.getenv("TRITON_AUTOTUNE_PARALLEL_COMPILE", "1") == "1" ) def _infer_compile_options_hint_if_needed(self): if "compile_options" in self.hints: return - self.hints["compile_options"] = ( - "mixcv" if self._autoparse_has_dot() else "vector" + self.hints["compile_options"] = self.operator_policy.compile_options_kernel_type + + def _use_search_params_runtime(self) -> bool: + return ( + self.search_params.enabled + and self.compile_options.enabled + and self.compile_options.kernel_type + == self.operator_policy.compile_options_kernel_type + ) + + def _use_search_params_mixcv(self) -> bool: + return ( + self._use_search_params_runtime() + and self.compile_options.kernel_type == "mixcv" ) + def _apply_fixed_compile_profile(self, profile): + return apply_fixed_compile_options_to_profile( + profile, getattr(self, "fixed_compile_options", {}) + ) + + def _apply_fixed_compile_profiles(self, profiles): + deduped = [] + seen = set() + for profile in profiles: + fixed_profile = self._apply_fixed_compile_profile(profile) + key = effective_compile_profile_key(fixed_profile) + if key in seen: + continue + seen.add(key) + deduped.append(fixed_profile) + return deduped + + def _infer_operator_policy(self) -> SearchPolicy: + try: + features, policy = analyze_operator_policy( + self._parse_ast(), self._get_capture_scope(), {id(self.ast_fn)} + ) + self.operator_features = features + return policy + except Exception as exc: # noqa: BLE001 - autotune must fail closed. + self.operator_features = None + return make_search_policy(OperatorKind.UNKNOWN) + def _parse_ast(self): parse = getattr(self.ast_fn, "parse", None) if not callable(parse): @@ -338,6 +434,10 @@ def _init_axis_params( def _autoparse_axis_params(self, all_args): miss_params = [arg for arg in self.arg_names if arg not in all_args.keys()] + search_param_names = set() + search_params = getattr(self, "search_params", None) + if search_params is not None and getattr(search_params, "enabled", False): + search_param_names = set(search_params.params) # parse pointer params nums if self.num_buffers == -1: self.num_buffers = self._autoparse_ptr_nums(all_args) @@ -413,6 +513,7 @@ def _autoparse_axis_params(self, all_args): miss_params = [ arg for arg in miss_params if arg not in self.tiling_params.values() ] + miss_params = [arg for arg in miss_params if arg not in search_param_names] if miss_params: raise ValueError( f"Missing required arguments: {miss_params}. " @@ -423,7 +524,7 @@ def _autoparse_axis_params(self, all_args): def _gen_tile_configs( self, kv_dict: Dict[str, int], dtype: torch.dtype ) -> List[Config]: - from .tile_generator import KernelMeta, TileGenerator + from .tile_candidate_generator import KernelMeta, TileGenerator axis_sizes = {} for k, v in kv_dict.items(): @@ -496,18 +597,59 @@ def generate_key_and_configs(self, *args, **kwargs): if dtype is None: raise NotImplementedError("Not support for non-Tensor inputs") + fixed_compile_options = { + name: kwargs[name] + for name in get_compile_option_param_names(self.compile_options) + if name in kwargs and kwargs[name] is not None + } + if fixed_compile_options: + key.append( + tuple( + sorted( + (name, _hashable_value(value)) + for name, value in fixed_compile_options.items() + ) + ) + ) + self.fixed_compile_options = fixed_compile_options + self._search_params_runtime_enabled = self._use_search_params_runtime() key = tuple(key) if key not in self.cache: - if self.auto_gen_config: + if self.auto_gen_config and not self._search_params_runtime_enabled: self._autoparse_axis_params(all_args) _kv_dict = {k: _args[v] for k, v in self.keys.items() if v in _args} self._gen_tile_configs(_kv_dict, dtype) - fixed_compile_options = { - name: kwargs[name] - for name in get_compile_option_param_names(self.compile_options) - if name in kwargs and kwargs[name] is not None - } - self.fixed_compile_options = fixed_compile_options + if self._search_params_runtime_enabled: + self.configs = self.gen_configs + self.user_configs + if not self.configs: + self.configs = [ + _make_config_compat( + kwargs={}, + num_warps=4, + num_stages=2, + num_ctas=1, + num_buffers_warp_spec=0, + num_consumer_groups=0, + reg_dec_producer=0, + reg_inc_consumer=0, + ) + ] + if self.print_autotuning or self.search_params.debug: + print( + "Search params autotuning: enabled; " + f"operator_kind={self.operator_policy.operator_kind.value}, " + f"compile_options={self.compile_options.kernel_type}, " + f"params={self.search_params.params}, " + f"values={self.search_params.values}, " + f"fixed_runtime_params={sorted(fixed_compile_options)}, " + f"base_configs={len(self.configs)}, " + "skip old compile_options cartesian expansion", + flush=True, + ) + notice = self._search_params_correctness_notice() + if notice: + print(notice, flush=True) + return key gen_configs = expand_compile_option_configs( self.gen_configs, self.compile_options, @@ -574,41 +716,60 @@ def run(self, *args, **kwargs): kwargs["simt_stack_limit"] = self.simt_stack_limit used_cached_result = True if key not in self.cache: - # prune configs - pruned_configs = self.prune_configs(kwargs) - if len(pruned_configs) > 1: + if self._search_params_runtime_enabled: used_cached_result = False - - def benchmark(): - bench_start = time.time() - timings = self._batch_bench(*args, configs=pruned_configs, **kwargs) - bench_end = time.time() - self.bench_time = bench_end - bench_start - self.cache[key] = builtins.min(timings, key=timings.get) - full_nargs = { - **self.nargs, - **kwargs, - **self.cache[key].all_kwargs(), - } - self.pre_hook(full_nargs, reset_only=True) - self.configs_timings = timings - if self.print_autotuning_timings: - self._print_config_timings(timings) - - if self.cache_results: - used_cached_result = self.check_disk_cache( - key, pruned_configs, benchmark - ) - else: - benchmark() + bench_start = time.time() + try: + config, timings = self._run_search_params_autotune(*args, **kwargs) + finally: + self.bench_time = time.time() - bench_start + self.last_search_stats["bench_time"] = self.bench_time + self.cache[key] = config + full_nargs = { + **self.nargs, + **kwargs, + **self.cache[key].all_kwargs(), + } + self.pre_hook(full_nargs, reset_only=True) + self.configs_timings = timings + if self.print_autotuning_timings: + self._print_config_timings(timings) config = self.cache[key] else: - self.cache[key] = pruned_configs[0] - config = self.cache[key] + # prune configs + pruned_configs = self.prune_configs(kwargs) + if len(pruned_configs) > 1: + used_cached_result = False + + def benchmark(): + bench_start = time.time() + timings = self._batch_bench( + *args, configs=pruned_configs, **kwargs + ) + bench_end = time.time() + self.bench_time = bench_end - bench_start + self.cache[key] = builtins.min(timings, key=timings.get) + full_nargs = { + **self.nargs, + **kwargs, + **self.cache[key].all_kwargs(), + } + self.pre_hook(full_nargs, reset_only=True) + self.configs_timings = timings + if self.print_autotuning_timings: + self._print_config_timings(timings) + + benchmark() + config = self.cache[key] + else: + self.cache[key] = pruned_configs[0] + config = self.cache[key] else: config = self.cache[key] self.best_config = config + if self._search_params_runtime_enabled and used_cached_result: + self._record_search_stats(searched=False) if self.print_autotuning and not used_cached_result: print( f"Triton autotuning for function {self.base_fn.__name__} finished after " @@ -621,7 +782,7 @@ def benchmark(): if config.pre_hook is not None: full_nargs = {**self.nargs, **kwargs, **config.all_kwargs()} config.pre_hook(full_nargs) - final_kwargs = dict(config.all_kwargs(), **kwargs) + final_kwargs = dict(kwargs, **config.all_kwargs()) ret = self.fn.run( *args, **final_kwargs, @@ -635,6 +796,10 @@ def _timing_sort_key(cost): return cost[0] return cost + def _is_finite_timing(self, cost): + value = self._timing_sort_key(cost) + return isinstance(value, (int, float)) and math.isfinite(value) + def _print_config_timings(self, timings): sorted_timings = sorted( timings.items(), key=lambda item: self._timing_sort_key(item[1]) @@ -647,7 +812,9 @@ def _print_config_timings(self, timings): flush=True, ) - def _batch_bench(self, *args, configs, **kwargs): + def _batch_bench( + self, *args, configs, force_parallel=False, return_errors=False, **kwargs + ): from triton.compiler.errors import CompilationError, CompileTimeAssertionFailure from triton.runtime.errors import OutOfResources @@ -656,13 +823,18 @@ def _batch_bench(self, *args, configs, **kwargs): for config in configs } run_fns = {} + errors = {} exc = None exc_stack = "" - if self.compile_parallel: + use_parallel = self.compile_parallel or ( + force_parallel and self._can_async_compile + ) + if use_parallel: import psutil - max_workers = min(psutil.cpu_count(logical=False) // 2, len(kernels_call)) + cpu_count = psutil.cpu_count(logical=False) or 1 + max_workers = max(1, min(max(1, cpu_count // 2), len(kernels_call))) future_kernels = [] try: with ( @@ -690,6 +862,7 @@ def _batch_bench(self, *args, configs, **kwargs): exc_stack = traceback.format_exc() exc = e + errors[config] = e except Exception as e: # ignore exception from __exit__() of AsyncCompileMode triton.runtime._async_compile.active_mode.set(None) @@ -699,6 +872,9 @@ def _batch_bench(self, *args, configs, **kwargs): exc_stack = traceback.format_exc() exc = e + for config in configs: + if config not in run_fns: + errors.setdefault(config, e) else: for config, fn in kernels_call.items(): try: @@ -715,25 +891,116 @@ def _batch_bench(self, *args, configs, **kwargs): exc_stack = traceback.format_exc() exc = e + errors[config] = e if len(run_fns) == 0: + if return_errors: + return {}, errors raise RuntimeError( f"No valid triton configs. {type(exc).__name__}: {exc} \nStack trace: {exc_stack}" ) + bench_fn = self.do_bench + if ( + getattr(self, "search_params", None) is not None + and self.search_params.bench_warmup is not None + and self.search_params.bench_active is not None + ): + bench_fn = functools.partial( + self.do_bench, + warmup=self.search_params.bench_warmup, + active=self.search_params.bench_active, + ) strategy = select_benchmark_strategy( - self.do_bench, + bench_fn, self.user_defined_do_bench, len(run_fns), ) + timings = {} try: - return strategy.bench(run_fns) + timings = strategy.bench(run_fns) except Exception: _empty_npu_cache_after_failure() - return self._batch_bench_fallback(strategy, run_fns) + fallback_timings, fallback_errors = self._batch_bench_fallback( + strategy, run_fns + ) + timings.update(fallback_timings) + errors.update(fallback_errors) + if return_errors: + return timings, errors + return timings + + def _collect_search_ub_for_configs(self, *args, configs, **kwargs): + """Cache compile-time UB bytes for Stage 1 shape ranking. + + UB collection is intentionally separate from timing benchmark. It reads + ``CompiledKernel.metadata.required_ub_bits``, which the backend fills + from ``memory_info_{aic,aiv}.json`` when ``TRITON_MEMORY_DISPLAY=1`` enables + ``--enable-memory-display=true``. + """ + if not getattr(self, "search_params", None): + return + if not hasattr(self, "_search_ub_cache"): + self._search_ub_cache = {} + configs = [config for config in configs if config not in self._search_ub_cache] + if not configs: + return + + prev_memory_display = os.environ.get("TRITON_MEMORY_DISPLAY") + os.environ["TRITON_MEMORY_DISPLAY"] = "1" + try: + if self._can_async_compile: + import psutil + + cpu_count = psutil.cpu_count(logical=False) or 1 + max_workers = max(1, min(max(1, cpu_count // 2), len(configs))) + future_kernels = [] + try: + with ( + ThreadPoolExecutor(max_workers=max_workers) as executor, + triton.AsyncCompileMode(executor), + ): + for config in configs: + kernel_call = self._make_kernel_call( + *args, config=config, **kwargs + ) + future_kernels.append((config, kernel_call(warmup=True))) + + for config, fut in future_kernels: + try: + if hasattr(fut, "result"): + fut = fut.result() + ub = ub_bytes_of(fut) + if ub is not None: + self._search_ub_cache[config] = ub + except Exception: + _empty_npu_cache_after_failure() + except Exception: + triton.runtime._async_compile.active_mode.set(None) + _empty_npu_cache_after_failure() + return + + for config in configs: + try: + kernel_call = self._make_kernel_call(*args, config=config, **kwargs) + compiled = kernel_call(warmup=True) + if hasattr(compiled, "result"): + compiled = compiled.result() + except Exception: + _empty_npu_cache_after_failure() + continue + ub = ub_bytes_of(compiled) + if ub is not None: + self._search_ub_cache[config] = ub + finally: + if prev_memory_display is None: + os.environ.pop("TRITON_MEMORY_DISPLAY", None) + else: + os.environ["TRITON_MEMORY_DISPLAY"] = prev_memory_display def _batch_bench_fallback(self, strategy, run_fns): timings = {} + errors = {} for config, fn in run_fns.items(): try: cost = strategy.bench({config: fn}) @@ -745,10 +1012,904 @@ def _batch_bench_fallback(self, strategy, run_fns): timings[config] = cost else: timings[config] = float("inf") - except Exception: + errors[config] = RuntimeError( + "benchmark fallback returned non-scalar timing" + ) + except Exception as exc: _empty_npu_cache_after_failure() timings[config] = float("inf") - return timings + errors[config] = exc + return timings, errors + + def _bench_stage1_fast_configs(self, *args, configs, **kwargs): + """Compile Stage-1 probe configs in parallel, then time them cheaply. + + Stage 1 only ranks shape candidates, so it intentionally avoids the + CANN profiler based do_bench_npu path. Stage 2 uses the same fast + timing path; Stage 3 is the only precise do_bench_npu pass. + """ + from triton.compiler.errors import CompilationError, CompileTimeAssertionFailure + from triton.runtime.errors import OutOfResources + from triton.testing import do_bench + + self._search_bench_call_count += 1 + self._search_bench_config_count += len(configs) + + kernels_call = { + config: self._make_kernel_call(*args, config=config, **kwargs) + for config in configs + } + run_fns = {} + errors = {} + + prev_memory_display = os.environ.get("TRITON_MEMORY_DISPLAY") + os.environ["TRITON_MEMORY_DISPLAY"] = "1" + try: + if self._can_async_compile: + import psutil + + cpu_count = psutil.cpu_count(logical=False) or 1 + max_workers = max(1, min(max(1, cpu_count // 2), len(kernels_call))) + future_kernels = [] + try: + with ( + ThreadPoolExecutor(max_workers=max_workers) as executor, + triton.AsyncCompileMode(executor), + ): + for config, fn in kernels_call.items(): + future_kernels.append((config, fn(warmup=True))) + + for config, fut in future_kernels: + try: + if hasattr(fut, "result"): + fut = fut.result() + ub = ub_bytes_of(fut) + if ub is not None: + self._search_ub_cache[config] = ub + run_fns[config] = functools.partial( + kernels_call[config], warmup=False + ) + except ( + CompileTimeAssertionFailure, + CompilationError, + OutOfResources, + Exception, + ) as exc: + _empty_npu_cache_after_failure() + errors[config] = exc + except Exception as exc: + triton.runtime._async_compile.active_mode.set(None) + _empty_npu_cache_after_failure() + for config in configs: + errors.setdefault(config, exc) + else: + for config, fn in kernels_call.items(): + try: + compiled = fn(warmup=True) + if hasattr(compiled, "result"): + compiled = compiled.result() + ub = ub_bytes_of(compiled) + if ub is not None: + self._search_ub_cache[config] = ub + run_fns[config] = functools.partial(fn, warmup=False) + except ( + CompileTimeAssertionFailure, + CompilationError, + OutOfResources, + Exception, + ) as exc: + _empty_npu_cache_after_failure() + errors[config] = exc + finally: + if prev_memory_display is None: + os.environ.pop("TRITON_MEMORY_DISPLAY", None) + else: + os.environ["TRITON_MEMORY_DISPLAY"] = prev_memory_display + + timings = {} + for config, run_fn in run_fns.items(): + try: + cost = do_bench( + run_fn, + warmup=self.search_params.stage1_bench_warmup, + rep=self.search_params.stage1_bench_rep, + quantiles=(0.5, 0.2, 0.8), + ) + if isinstance(cost, (list, tuple)): + cost = cost[0] if cost else float("inf") + if self._is_finite_timing(cost): + timings[config] = cost + else: + errors[config] = RuntimeError("stage1 fast benchmark returned inf") + except Exception as exc: # noqa: BLE001 + _empty_npu_cache_after_failure() + errors[config] = exc + return timings, errors + + def _bench_stage2_fast_configs(self, *args, configs, **kwargs): + """Stage-2 fast timing for every compile-profile candidate.""" + from triton.compiler.errors import CompilationError, CompileTimeAssertionFailure + from triton.runtime.errors import OutOfResources + from triton.testing import do_bench + + self._search_bench_call_count += 1 + self._search_bench_config_count += len(configs) + + prev_memory_display = os.environ.pop("TRITON_MEMORY_DISPLAY", None) + prev_memory_display_debug = os.environ.pop("TRITON_MEMORY_DISPLAY_DEBUG", None) + try: + kernels_call = { + config: self._make_kernel_call(*args, config=config, **kwargs) + for config in configs + } + run_fns = {} + errors = {} + + if self._can_async_compile: + import psutil + + cpu_count = psutil.cpu_count(logical=False) or 1 + max_workers = max(1, min(max(1, cpu_count // 2), len(kernels_call))) + future_kernels = [] + try: + with ( + ThreadPoolExecutor(max_workers=max_workers) as executor, + triton.AsyncCompileMode(executor), + ): + for config, fn in kernels_call.items(): + future_kernels.append((config, fn(warmup=True))) + + for config, fut in future_kernels: + try: + if hasattr(fut, "result"): + fut = fut.result() + run_fns[config] = functools.partial( + kernels_call[config], warmup=False + ) + except ( + CompileTimeAssertionFailure, + CompilationError, + OutOfResources, + Exception, + ) as exc: + _empty_npu_cache_after_failure() + errors[config] = exc + except Exception as exc: + triton.runtime._async_compile.active_mode.set(None) + _empty_npu_cache_after_failure() + for config in configs: + if config not in run_fns: + errors.setdefault(config, exc) + else: + for config, fn in kernels_call.items(): + try: + compiled = fn(warmup=True) + if hasattr(compiled, "result"): + compiled = compiled.result() + run_fns[config] = functools.partial(fn, warmup=False) + except ( + CompileTimeAssertionFailure, + CompilationError, + OutOfResources, + Exception, + ) as exc: + _empty_npu_cache_after_failure() + errors[config] = exc + + timings = {} + for config, run_fn in run_fns.items(): + try: + cost = do_bench( + run_fn, + warmup=self.search_params.stage1_bench_warmup, + rep=self.search_params.stage1_bench_rep, + quantiles=(0.5, 0.2, 0.8), + ) + if isinstance(cost, (list, tuple)): + cost = cost[0] if cost else float("inf") + if self._is_finite_timing(cost): + timings[config] = cost + else: + errors[config] = RuntimeError( + "stage2 fast benchmark returned inf" + ) + except Exception as exc: # noqa: BLE001 + _empty_npu_cache_after_failure() + errors[config] = exc + return timings, errors + finally: + if prev_memory_display is not None: + os.environ["TRITON_MEMORY_DISPLAY"] = prev_memory_display + if prev_memory_display_debug is not None: + os.environ["TRITON_MEMORY_DISPLAY_DEBUG"] = prev_memory_display_debug + + def _search_debug(self, message: str): + if self.search_params.debug or self.print_autotuning: + print(f"Search params autotuning: {message}", flush=True) + + def _format_search_error_log(self, error) -> str: + if error is None: + return "" + text = str(error) + if not text: + text = repr(error) + return text.replace("\r\n", "\n").replace("\r", "\n") + + def _record_search_stats(self, *, searched: bool, stage2_candidates: int = 0): + self.last_search_stats = { + "searched": searched, + "bench_time": getattr(self, "bench_time", 0.0), + "bench_calls": self._search_bench_call_count if searched else 0, + "bench_configs": self._search_bench_config_count if searched else 0, + "measurements": len(self._search_measure_cache) if searched else 0, + "stage1_success": ( + len(self._search_stage1_measured_summary) if searched else 0 + ), + "stage1_fail": (len(self._search_stage1_failed_summary) if searched else 0), + "stage2_final_candidates": stage2_candidates if searched else 0, + } + + def _bench_search_config(self, *args, config, **kwargs): + try: + timings, errors = self._batch_bench( + *args, configs=[config], return_errors=True, **kwargs + ) + except Exception as exc: # noqa: BLE001 + return None, exc + cost = timings.get(config, float("inf")) + if not self._is_finite_timing(cost): + return None, errors.get(config) or RuntimeError("benchmark returned inf") + return cost, None + + def _bench_search_configs(self, *args, configs, force_parallel=False, **kwargs): + self._search_bench_call_count += 1 + self._search_bench_config_count += len(configs) + try: + return self._batch_bench( + *args, + configs=configs, + force_parallel=force_parallel, + return_errors=True, + **kwargs, + ) + except Exception as exc: # noqa: BLE001 + return {}, {config: exc for config in configs} + + def _shape_config_pool(self, kwargs): + base_config = self.configs[0] + self._search_param_anchor_shape_keys = { + shape_key( + extract_shape(config, self.search_params.params), + self.search_params.params, + ) + for config in self.configs + if all(name in config.kwargs for name in self.search_params.params) + } + shape_configs = expand_search_param_shapes(base_config, self.search_params) + old_configs = self.configs + try: + self.configs = shape_configs + pruned_configs = self.prune_configs(kwargs) + finally: + self.configs = old_configs + self._search_debug( + f"Stage 0 shape pool: total={len(shape_configs)}, after_prune={len(pruned_configs)}" + ) + return pruned_configs + + def _run_stage1_probe_for_shape(self, *args, shape_config, **kwargs): + results, _ = self._run_stage1_probes_for_shapes( + *args, shape_configs=[shape_config], **kwargs + ) + return results[0] if results else None + + def _run_stage1_probes_for_shapes(self, *args, shape_configs, **kwargs): + operator_policy = getattr(self, "operator_policy", None) + profile_family = ( + operator_policy.stage1_profile_family + if operator_policy is not None + else "mixcv" + ) + profiles = self._apply_fixed_compile_profiles( + get_stage1_probe_profiles(profile_family) + ) + if not profiles: + return [], list(shape_configs) + + pending = [ + { + "shape": extract_shape(config, self.search_params.params), + "shape_config": config, + } + for config in shape_configs + ] + if not pending: + return [], [] + + results = [] + failed_configs = [] + + def _bench_profile(items, profile, label): + configs = [ + compile_profile_to_config( + profile, + shape_kwargs=item["shape"], + base_config=item["shape_config"], + ) + for item in items + ] + timings = {} + errors = {} + batch_configs = [] + for shape_item, config in zip(items, configs): + self._search_debug( + f"Stage 1 {label} shape={shape_item['shape']}, profile={profile}" + ) + cached = self._search_measure_cache.get(shape_item["shape"], profile) + if cached is None: + batch_configs.append(config) + continue + if cached["ok"] and cached.get("source") == "stage1_fast": + timings[config] = cached["time"] + if cached.get("ub") is not None: + self._search_ub_cache[config] = cached["ub"] + self._search_debug( + f"Stage 1 {label} cache hit " + f"shape={shape_item['shape']}, cost={cached['time']}" + ) + continue + if not cached["ok"]: + errors[config] = cached.get("error") or RuntimeError( + f"cached failure classified_as={cached.get('failure')}" + ) + self._search_debug( + f"Stage 1 {label} cache failure " + f"shape={shape_item['shape']}, " + f"classified_as={cached.get('failure')}" + ) + continue + batch_configs.append(config) + if batch_configs: + batch_timings, batch_errors = self._bench_stage1_fast_configs( + *args, configs=batch_configs, **kwargs + ) + timings.update(batch_timings) + errors.update(batch_errors) + return configs, timings, errors + + remaining = pending + last_failed = {} + for probe_index, profile in enumerate(profiles, 1): + if not remaining: + break + probe_configs, probe_timings, probe_errors = _bench_profile( + remaining, profile, f"probe[{probe_index}]" + ) + next_remaining = [] + for item, config in zip(remaining, probe_configs): + shape = item["shape"] + cost = probe_timings.get(config, float("inf")) + if self._is_finite_timing(cost): + self._search_measure_cache.put( + shape, profile, config, cost=cost, source="stage1_fast" + ) + self._search_debug( + f"Stage 1 probe[{probe_index}] success shape={shape}, " + f"cost={cost}, " + f"ub={getattr(self, '_search_ub_cache', {}).get(config)}" + ) + results.append( + { + "shape": shape, + "shape_config": item["shape_config"], + "profile": profile, + "config": config, + "time": cost, + "ub": getattr(self, "_search_ub_cache", {}).get(config), + } + ) + continue + + error = probe_errors.get(config) + failure = ( + classify_compile_failure(error) + if error is not None + else "EXACT_ONLY" + ) + self._search_measure_cache.put( + shape, + profile, + config, + error=error, + failure=failure, + source="stage1_fast", + ) + last_failed[shape_key(shape, self.search_params.params)] = ( + item, + failure, + error, + ) + next_remaining.append(item) + self._search_debug( + f"Stage 1 probe[{probe_index}] fail shape={shape}, " + f"classified_as={failure}, " + f"resource_ratio={resource_overflow_ratio(error)}, " + f"try_next={probe_index < len(profiles)}, " + f"error_log:\n{self._format_search_error_log(error)}" + ) + remaining = next_remaining + + for item in remaining: + shape = item["shape"] + failure_info = last_failed.get(shape_key(shape, self.search_params.params)) + failure = failure_info[1] if failure_info is not None else "EXACT_ONLY" + error = failure_info[2] if failure_info is not None else None + failed_configs.append(item["shape_config"]) + self._search_debug( + "Stage 1 all probes failed " + f"shape={shape}, classified_as={failure}, " + f"resource_ratio={resource_overflow_ratio(error)}, " + f"error_log:\n{self._format_search_error_log(error)}" + ) + + return results, failed_configs + + def _screen_search_param_shapes(self, *args, **kwargs): + shape_configs = self._shape_config_pool(kwargs) + if not shape_configs: + raise RuntimeError("No valid search_param shapes after early_config_prune") + + measured = [] + failed_configs = [] + observed_configs = [] + seen_shape_keys = set() + value_map = effective_search_value_map(shape_configs, self.search_params) + initial_percentiles = stage1_initial_percentiles(self.search_params, value_map) + total_stage1_budget = stage1_total_budget(self.search_params, shape_configs) + initial_budget = stage1_initial_budget(self.search_params, shape_configs) + initial = select_initial_percentile_shapes( + shape_configs, + self.search_params, + limit=initial_budget, + ) + initial_keys = { + shape_key( + extract_shape(config, self.search_params.params), + self.search_params.params, + ) + for config in initial + } + anchor_configs = [] + for config in shape_configs: + key = shape_key( + extract_shape(config, self.search_params.params), + self.search_params.params, + ) + if key in self._search_param_anchor_shape_keys and key not in initial_keys: + anchor_configs.append(config) + initial.extend(anchor_configs) + deduped_initial = [] + deduped_initial_keys = set() + for config in initial: + key = shape_key( + extract_shape(config, self.search_params.params), + self.search_params.params, + ) + if key in deduped_initial_keys: + continue + deduped_initial_keys.add(key) + deduped_initial.append(config) + initial = deduped_initial + if len(initial) > total_stage1_budget: + initial = initial[:total_stage1_budget] + self._search_debug( + f"Stage 1 initial cases={len(initial)} " + f"(percentile + anchors={len(anchor_configs)}), " + f"budget={total_stage1_budget}, " + f"effective_value_counts=" + f"{ {name: len(values) for name, values in value_map.items()} }, " + f"initial_percentiles={initial_percentiles}" + ) + + for config in initial: + shape = extract_shape(config, self.search_params.params) + seen_shape_keys.add(shape_key(shape, self.search_params.params)) + observed_configs.append(config) + initial_results, initial_failures = self._run_stage1_probes_for_shapes( + *args, shape_configs=initial, **kwargs + ) + measured.extend(initial_results) + failed_configs.extend(initial_failures) + + for round_id in range(self.search_params.shape_refine_rounds): + candidates_for_parents = filter_shapes_by_ub_and_timing( + measured, timing_sort_key=self._timing_sort_key + ) + if len(candidates_for_parents) < len(measured): + self._search_debug( + "Stage 1 evolve round=" + f"{round_id + 1}, UB/timing filter dropped " + f"{len(measured) - len(candidates_for_parents)} entries" + ) + parents = select_top_shape_entries( + candidates_for_parents, + k=self.search_params.shape_final_top_k, + key_fn=lambda item: ub_timing_weighted_key( + item, timing_sort_key=self._timing_sort_key + ), + ) + if ( + getattr(self.operator_policy, "operator_kind", None) + == OperatorKind.DOT_STATEFUL + and len(self.search_params.params) == 1 + and len(parents) >= 2 + ): + self._search_debug( + "Stage 1 evolve stop: DOT_STATEFUL single-param search " + f"already has {len(parents)} successful parent shapes" + ) + break + unseen_shapes = sum( + 1 + for config in shape_configs + if shape_key( + extract_shape(config, self.search_params.params), + self.search_params.params, + ) + not in seen_shape_keys + ) + remaining_budget = total_stage1_budget - len(seen_shape_keys) + remaining_rounds = self.search_params.shape_refine_rounds - round_id + children_per_round = stage1_children_per_round( + self.search_params, + shape_configs, + remaining_unseen=unseen_shapes, + remaining_budget=remaining_budget, + remaining_rounds=remaining_rounds, + ) + proposals = propose_evolved_shape_configs( + parents=parents, + successes=measured, + failures=failed_configs, + failure_observations=( + self._search_measure_cache.stage1_failure_observations( + failed_configs + ) + ), + observed=observed_configs, + all_configs=shape_configs, + spec=self.search_params, + seen_keys=seen_shape_keys, + limit=children_per_round, + ) + self._search_debug( + "Stage 1 evolve round=" + f"{round_id + 1}, proposals={len(proposals)}, " + f"parents={len(parents)}, successes={len(measured)}, " + f"failures={len(failed_configs)}, target_children={children_per_round}, " + f"unseen_shapes={unseen_shapes}, " + "proposal_shapes=" + + ", ".join( + str(extract_shape(config, self.search_params.params)) + for config in proposals + ) + ) + if not proposals: + self._search_debug( + "Stage 1 evolve stop: no new proposal " + f"(unseen_shapes={unseen_shapes}, " + "all candidates may already be measured or rejected by " + "failure-cone/acquisition filters)" + ) + break + for config in proposals: + shape = extract_shape(config, self.search_params.params) + seen_shape_keys.add(shape_key(shape, self.search_params.params)) + observed_configs.append(config) + proposal_results, proposal_failures = self._run_stage1_probes_for_shapes( + *args, shape_configs=proposals, **kwargs + ) + measured.extend(proposal_results) + failed_configs.extend(proposal_failures) + + if not measured: + raise RuntimeError("No valid search_param shapes after Stage 1 screening") + + selected = select_top_shape_entries( + measured, + k=self.search_params.shape_final_top_k, + key_fn=lambda item: ub_timing_weighted_key( + item, timing_sort_key=self._timing_sort_key + ), + ) + self._search_debug( + "Stage 1 selected shapes: " + + ", ".join( + f"{item['shape']}@{self._timing_sort_key(item['time'])}" + for item in selected + ) + ) + self._search_stage1_measured_summary = select_top_shape_entries( + measured, + k=len(measured), + key_fn=lambda item: self._timing_sort_key(item["time"]), + ) + self._search_stage1_failed_summary = [ + extract_shape(config, self.search_params.params) + for config in failed_configs + ] + return selected + + def _search_param_baseline_candidate(self): + for config in list(self.user_configs) + list(self.gen_configs): + if all(name in config.kwargs for name in self.search_params.params): + profile = { + "mode": "BASELINE", + "num_stages": getattr(config, "num_stages", None), + } + self._search_debug(f"Stage 3 baseline guard candidate: config={config}") + return config, profile, float("inf") + self._search_debug( + "Stage 3 baseline guard unavailable: no original config contains " + f"all search params {self.search_params.params}" + ) + return None + + def _make_stage2_searcher(self, *args, **kwargs): + def bench_configs(configs): + return self._bench_stage2_fast_configs(*args, configs=configs, **kwargs) + + return Stage2CompileSearcher( + search_params=self.search_params, + operator_policy=self.operator_policy, + cache=self._search_measure_cache, + apply_fixed_profile=self._apply_fixed_compile_profile, + bench_configs=bench_configs, + timing_sort_key=self._timing_sort_key, + is_finite_timing=self._is_finite_timing, + debug=self._search_debug, + format_error_log=self._format_search_error_log, + ) + + def _run_stage3_precise_pick( + self, + *args, + candidates, + baseline_candidate=None, + **kwargs, + ): + reference_fn = getattr(self.search_params, "reference_fn", None) + + ranked_candidates = [] + seen_configs = set() + for config, profile, cost in sorted( + candidates, key=lambda item: self._timing_sort_key(item[2]) + ): + if config in seen_configs: + continue + seen_configs.add(config) + ranked_candidates.append((config, profile, cost)) + + check_queue = list(ranked_candidates[:3]) + if baseline_candidate is not None: + baseline_config, baseline_profile, baseline_cost = baseline_candidate + if baseline_config not in seen_configs: + check_queue.append((baseline_config, baseline_profile, baseline_cost)) + seen_configs.add(baseline_config) + self._search_debug( + f"Stage 3 precise check: top-{len(check_queue)} of " + f"{len(candidates)} candidates" + ) + failures = [] + passing = [] + checked_configs = set() + next_ranked_index = 3 + + while True: + while check_queue: + config, profile, cost = check_queue.pop(0) + if config in checked_configs: + continue + checked_configs.add(config) + rank = len(checked_configs) + precise_timings, precise_errors = self._bench_search_configs( + *args, configs=[config], **kwargs + ) + precise_cost = precise_timings.get(config, float("inf")) + if not self._is_finite_timing(precise_cost): + error = precise_errors.get(config) + message = self._format_search_error_log(error) + failures.append(f"rank={rank} precise failed: {message}") + self._search_debug( + "Stage 3 precise fail " + f"rank={rank}, config={config}, error={message}" + ) + continue + + try: + passed = ( + True + if reference_fn is None + else reference_fn(self, config, *args, **kwargs) + ) + except Exception as exc: + failures.append(f"rank={rank} accuracy error: {exc}") + self._search_debug( + f"Stage 3 accuracy error rank={rank}, " + f"config={config}, error={exc}" + ) + continue + if not passed: + failures.append(f"rank={rank} accuracy failed") + self._search_debug( + f"Stage 3 accuracy fail rank={rank}, " + f"config={config}, cost={cost}" + ) + continue + self._search_debug( + "Stage 3 precise pass " + f"rank={rank}, config={config}, fast_cost={cost}, " + f"precise_cost={precise_cost}" + ) + passing.append((config, profile, precise_cost, rank, cost)) + + if passing or reference_fn is None: + break + while ( + next_ranked_index < len(ranked_candidates) + and ranked_candidates[next_ranked_index][0] in checked_configs + ): + next_ranked_index += 1 + if next_ranked_index >= len(ranked_candidates): + break + extra = ranked_candidates[next_ranked_index] + next_ranked_index += 1 + check_queue.append(extra) + self._search_debug( + "Stage 3 extend precise check after all checked candidates " + f"failed correctness; next_config={extra[0]}" + ) + + if passing: + baseline = next( + ( + item + for item in passing + if isinstance(item[1], dict) and item[1].get("mode") == "BASELINE" + ), + None, + ) + tuned = [ + item + for item in passing + if not (isinstance(item[1], dict) and item[1].get("mode") == "BASELINE") + ] + if baseline is not None and tuned: + best_tuned = min(tuned, key=lambda item: self._timing_sort_key(item[2])) + baseline_time = self._timing_sort_key(baseline[2]) + tuned_time = self._timing_sort_key(best_tuned[2]) + if tuned_time < baseline_time: + config, profile, precise_cost, rank, fast_cost = best_tuned + self._search_debug( + "Stage 3 baseline guard accept tuned " + f"tuned_precise={precise_cost}, " + f"baseline_precise={baseline[2]}" + ) + else: + config, profile, precise_cost, rank, fast_cost = baseline + self._search_debug( + "Stage 3 baseline guard fallback " + f"best_tuned_precise={best_tuned[2]}, " + f"baseline_precise={precise_cost}" + ) + else: + config, profile, precise_cost, rank, fast_cost = min( + passing, key=lambda item: self._timing_sort_key(item[2]) + ) + self._search_debug( + "Stage 3 final pick " + f"rank={rank}, config={config}, fast_cost={fast_cost}, " + f"precise_cost={precise_cost}" + ) + return config, profile, precise_cost + + raise RuntimeError( + "Stage 3 top-3 candidates failed precise benchmark or accuracy check: " + + "; ".join(failures) + ) + + def _run_search_params_autotune(self, *args, **kwargs): + self._search_ub_cache = {} + self._search_measure_cache = SearchMeasureCache( + self.search_params.params, + self._search_ub_cache, + ) + self._search_bench_call_count = 0 + self._search_bench_config_count = 0 + all_candidates = [] + try: + stage1_entries = self._screen_search_param_shapes(*args, **kwargs) + if not stage1_entries: + raise RuntimeError("No valid search_param shapes after Stage 1") + stage1_top1_time = min(entry["time"] for entry in stage1_entries) + stage2_searcher = self._make_stage2_searcher(*args, **kwargs) + timings = {} + for stage1_rank, entry in enumerate(stage1_entries, 1): + if not getattr(self.operator_policy, "stage2_enabled", True): + all_candidates.append( + (entry["config"], entry["profile"], entry["time"]) + ) + timings[entry["config"]] = entry["time"] + continue + shape_candidates = stage2_searcher.search( + stage1_entry=entry, + stage1_rank=stage1_rank, + stage1_top1_time=stage1_top1_time, + ) + for config, profile, cost in shape_candidates: + all_candidates.append((config, profile, cost)) + timings[config] = cost + if not all_candidates: + raise RuntimeError( + "No valid compile profiles after search_params compile search" + ) + baseline_candidate = self._search_param_baseline_candidate() + stage3_pick = self._run_stage3_precise_pick( + *args, + candidates=all_candidates, + baseline_candidate=baseline_candidate, + **kwargs, + ) + best_config, best_profile, best_time = stage3_pick + timings[best_config] = best_time + self._search_debug( + f"Stage 3 best time={best_time}, profile={best_profile}, " + f"config={best_config} (precise-picked)" + ) + self._search_debug( + "Stage 1 measured summary: " + + ", ".join( + f"{item['shape']}@{self._timing_sort_key(item['time'])}" + for item in self._search_stage1_measured_summary + ) + ) + if self._search_stage1_failed_summary: + self._search_debug( + "Stage 1 failed summary: " + + ", ".join( + str(shape) for shape in self._search_stage1_failed_summary + ) + ) + self._record_search_stats( + searched=True, stage2_candidates=len(all_candidates) + ) + return best_config, timings + except Exception: + self._record_search_stats( + searched=True, stage2_candidates=len(all_candidates) + ) + raise + + def _search_params_correctness_notice(self): + search_params = getattr(self, "search_params", None) + if search_params is None or getattr(search_params, "reference_fn", None): + return None + kind = getattr(self.operator_policy, "operator_kind", None) + risky_kinds = { + OperatorKind.DOT_STATEFUL, + OperatorKind.VECTOR_DISCRETE_OR_STATEFUL, + OperatorKind.UNKNOWN, + } + if kind not in risky_kinds: + return None + return ( + "Search params autotuning: correctness notice; " + f"operator_kind={kind.value} has no reference_fn. " + "Only use search_params for true performance tile knobs. " + "Keep coverage, atomic granularity, state partition, and " + "hand-derived UB-safe constants fixed unless a correctness gate " + "is provided." + ) def _make_kernel_call(self, *args, config, **meta): # check for conflicts, i.e. meta-parameters both provided @@ -1050,7 +2211,6 @@ def autotune( rep=None, use_cuda_graph=False, do_bench=None, - cache_results=False, *, auto_prof_dir=None, hints=None, @@ -1108,8 +2268,6 @@ def kernel(x_ptr, x_size, **META): :type rep: int :param do_bench: a benchmark function to measure the time of each run. :type do_bench: lambda fn, quantiles - :param cache_results: whether to cache autotune timings to disk. - :type cache_results: bool :param auto_prof_dir: the specified directory to store the profiling results of the best config. If this parameter is None or the best config is retrieved from cache, the profiling process will be ignored. :type auto_prof_dir: str @@ -1131,7 +2289,6 @@ def decorator(fn): rep=rep, use_cuda_graph=use_cuda_graph, do_bench=do_bench, - cache_results=cache_results, auto_profile_dir=auto_prof_dir, hints=hints, ) @@ -1510,7 +2667,6 @@ def max_autotune( rep=None, use_cuda_graph=False, do_bench=None, - cache_results=False, **tuning_params, ): """ @@ -1532,7 +2688,6 @@ def max_autotune( :param rep: Deprecated. :param use_cuda_graph: Deprecated. :param do_bench: Same as in autotune. - :param cache_results: Same as in autotune. :param tuning_params: Additional tuning parameters as keyword arguments. Each value must be a list; the Cartesian product of these lists will be combined with each base config. @@ -1563,7 +2718,6 @@ def decorator(fn): rep=rep, use_cuda_graph=use_cuda_graph, do_bench=do_bench, - cache_results=cache_results, )(fn) return decorator diff --git a/backend/ascend_autotune_runtime/compile_options.py b/backend/ascend_autotune_runtime/compile_options.py deleted file mode 100644 index 1f4957a0..00000000 --- a/backend/ascend_autotune_runtime/compile_options.py +++ /dev/null @@ -1,612 +0,0 @@ -from __future__ import annotations - -import inspect -import itertools -from dataclasses import dataclass, field -from typing import Any, Dict, List, Optional - -from triton.runtime.autotuner import Config -from triton.backends.dicp_triton.utils import is_compile_on_910_95 - - -DEFAULT_MAX_CONFIGS = None - -_VALID_VALUES = { - "num_stages": [1, 2], - "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], - "set_workspace_multibuffer": [2, 4], - "tile_mix_vector_loop": [1, 2, 4, 8], - "tile_mix_cube_loop": [1, 2, 4, 8], -} - -_BOOLEAN_PARAMS = { - "enable_tuning_mode", - "multibuffer", - "unit_flag", - "limit_auto_multi_buffer_only_for_local_buffer", - "enable_hivm_auto_cv_balance", - "enable_ubuf_saving", - # "enable_preload", - "enable_auto_bind_sub_block", -} - -_SUPPORTED_PARAMS = { - "cube": { - "enable_tuning_mode", - "num_stages", - "unit_flag", - "limit_auto_multi_buffer_of_local_buffer", - }, - "mixcv": { - "enable_tuning_mode", - "num_stages", - "multibuffer", - "unit_flag", - "limit_auto_multi_buffer_only_for_local_buffer", - "limit_auto_multi_buffer_of_local_buffer", - "set_workspace_multibuffer", - "enable_hivm_auto_cv_balance", - "tile_mix_vector_loop", - "tile_mix_cube_loop", - "enable_ubuf_saving", - # "enable_preload", - "enable_auto_bind_sub_block", - }, - "vector": { - "enable_tuning_mode", - "num_stages", - "enable_ubuf_saving", - }, -} - -_ALL_PARAMS = set().union(*_SUPPORTED_PARAMS.values()) - -_MIXCV_910_95_UNSUPPORTED_PARAMS = { - "tile_mix_vector_loop", - "tile_mix_cube_loop", -} - -_AUTO_SEARCH_PRESETS = { - "cube": { - "enable_tuning_mode": [True], - "num_stages": [1, 2], - "unit_flag": [False, True], - "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], - }, - "mixcv": { - "enable_tuning_mode": [True], - "num_stages": [1, 2], - "unit_flag": [False, True], - "limit_auto_multi_buffer_only_for_local_buffer": [True, False], - "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], - "set_workspace_multibuffer": [2, 4], - "enable_hivm_auto_cv_balance": [True], - "tile_mix_vector_loop": [1, 2, 4], - "tile_mix_cube_loop": [1, 2, 4], - "enable_ubuf_saving": [False, True], - # "enable_preload": [False, True], - "enable_auto_bind_sub_block": [True], - }, - "vector": { - "enable_tuning_mode": [True], - "num_stages": [1, 2], - "enable_ubuf_saving": [True, False], - }, -} - - -def _is_mixcv_multi_buffer_auto_enabled( - num_stages: int, - combo: Dict[str, Any], - config: Config, - fixed_options: Dict[str, Any], -) -> bool: - """Whether `enable_auto_multi_buffer` takes effect for this combination.""" - if num_stages == 1: - return False - multibuffer = _resolve_compile_option( - "multibuffer", combo, config, fixed_options, default=None - ) - return multibuffer is not False - - -def _is_mixcv_limit_to_local_only_active( - num_stages: int, - combo: Dict[str, Any], - config: Config, - fixed_options: Dict[str, Any], -) -> bool: - return _is_mixcv_multi_buffer_auto_enabled(num_stages, combo, config, fixed_options) - - -def _is_mixcv_workspace_multibuffer_active( - num_stages: int, - combo: Dict[str, Any], - config: Config, - fixed_options: Dict[str, Any], -) -> bool: - if not _is_mixcv_multi_buffer_auto_enabled( - num_stages, combo, config, fixed_options - ): - return False - limit_to_local_only = _resolve_compile_option( - "limit_auto_multi_buffer_only_for_local_buffer", - combo, - config, - fixed_options, - default=True, - ) - return limit_to_local_only is False - - -_MIXCV_OPTION_ACTIVITY_RULES = { - "limit_auto_multi_buffer_only_for_local_buffer": _is_mixcv_limit_to_local_only_active, - "limit_auto_multi_buffer_of_local_buffer": _is_mixcv_limit_to_local_only_active, - "set_workspace_multibuffer": _is_mixcv_workspace_multibuffer_active, - "tile_mix_vector_loop": _is_mixcv_workspace_multibuffer_active, - "tile_mix_cube_loop": _is_mixcv_workspace_multibuffer_active, -} - - -@dataclass -class CompileOptionsSpec: - enabled: bool = False - kernel_type: str = "mixcv" - params: Dict[str, List[Any]] = field(default_factory=dict) - max_configs: Optional[int] = DEFAULT_MAX_CONFIGS - - -def _normalize_kernel_type(kernel_type: str) -> str: - if kernel_type == "mix": - return "mixcv" - if kernel_type not in _SUPPORTED_PARAMS: - raise ValueError( - "compile_options kernel_type must be one of: cube, mix, mixcv, vector" - ) - return kernel_type - - -def _as_value_list(name: str, value: Any) -> List[Any]: - values = list(value) if isinstance(value, (list, tuple)) else [value] - if not values: - raise ValueError(f"compile_options parameter '{name}' must not be empty") - return values - - -def validate_compile_option_values(name: str, values: List[Any]) -> None: - if name in _BOOLEAN_PARAMS and not all(isinstance(v, bool) for v in values): - raise ValueError(f"compile_options parameter '{name}' expects boolean values") - - if name in _VALID_VALUES and not all(v in _VALID_VALUES[name] for v in values): - raise ValueError( - f"compile_options parameter '{name}' expects values in {_VALID_VALUES[name]}" - ) - - -def parse_compile_options_hint(hint: Any) -> CompileOptionsSpec: - if hint is None or hint is False: - return CompileOptionsSpec(enabled=False) - - if hint is True: - return CompileOptionsSpec(enabled=True) - - if isinstance(hint, str): - return CompileOptionsSpec( - enabled=True, - kernel_type=_normalize_kernel_type(hint), - ) - - if not isinstance(hint, dict): - raise TypeError("hints['compile_options'] must be bool, str, or dict") - - raw = dict(hint) - kernel_type = _normalize_kernel_type( - raw.pop("kernel_type", raw.pop("type", "mixcv")) - ) - max_configs = raw.pop("max_configs", DEFAULT_MAX_CONFIGS) - if max_configs is not None and ( - not isinstance(max_configs, int) or max_configs <= 0 - ): - raise ValueError( - "compile_options max_configs must be a positive integer or None" - ) - - nested_options = raw.pop("options", {}) - if nested_options: - if not isinstance(nested_options, dict): - raise TypeError("compile_options options must be a dict") - raw.update(nested_options) - - supported = _SUPPORTED_PARAMS[kernel_type] - params: Dict[str, List[Any]] = {} - for name, value in raw.items(): - if name not in _ALL_PARAMS: - raise ValueError(f"Unknown compile_options parameter: {name}") - if name not in supported: - print( - f"[WARNING] compile_options parameter '{name}' is not supported " - f"for kernel_type '{kernel_type}' and will be ignored." - ) - continue - values = _as_value_list(name, value) - validate_compile_option_values(name, values) - params[name] = values - - return CompileOptionsSpec( - enabled=True, - kernel_type=kernel_type, - params=params, - max_configs=max_configs, - ) - - -def _make_config_compat(**kwargs): - supported_config_args = inspect.signature(Config).parameters - return Config( - **{key: value for key, value in kwargs.items() if key in supported_config_args} - ) - - -def _value_space_for_config( - config: Config, spec: CompileOptionsSpec, *, generated_tiling: bool -) -> Dict[str, List[Any]]: - supported = _SUPPORTED_PARAMS[spec.kernel_type] - preset = _AUTO_SEARCH_PRESETS[spec.kernel_type] - - value_space = {} - for name in sorted(supported): - if name in spec.params: - values = spec.params[name] - elif name not in preset: - continue - else: - values = preset[name] - validate_compile_option_values(name, values) - value_space[name] = values - if spec.kernel_type == "mixcv" and is_compile_on_910_95: - for name in _MIXCV_910_95_UNSUPPORTED_PARAMS: - value_space.pop(name, None) - return value_space - - -def _is_inactive_reason( - name: str, - num_stages: int, - combo: Dict[str, Any], - config: Config, - fixed_options: Dict[str, Any], -) -> Optional[str]: - if not _is_param_effective(name, num_stages, combo, config, fixed_options): - if name in { - "limit_auto_multi_buffer_only_for_local_buffer", - "limit_auto_multi_buffer_of_local_buffer", - "set_workspace_multibuffer", - "tile_mix_vector_loop", - "tile_mix_cube_loop", - }: - return "depends on auto multi-buffer" - return None - - -def _is_param_effective( - name: str, - num_stages: int, - combo: Dict[str, Any], - config: Config, - fixed_options: Dict[str, Any], -) -> bool: - if name not in _MIXCV_OPTION_ACTIVITY_RULES: - return True - return _MIXCV_OPTION_ACTIVITY_RULES[name](num_stages, combo, config, fixed_options) - - -def _resolve_compile_option( - name: str, - combo: Dict[str, Any], - config: Config, - fixed_options: Dict[str, Any], - default: Any, -) -> Any: - if name in combo: - return combo[name] - if name in fixed_options: - return fixed_options[name] - return default - - -def _effective_values( - name: str, - value_space: Dict[str, List[Any]], - fixed_options: Dict[str, Any], - default: Any, -) -> List[Any]: - if name in fixed_options: - return [fixed_options[name]] - if name in value_space: - return value_space[name] - return [default] - - -def _emit_values( - names: List[str], - value_space: Dict[str, List[Any]], - fixed_options: Dict[str, Any], -) -> List[tuple[str, List[Any]]]: - return [ - (name, value_space[name]) - for name in names - if name in value_space and name not in fixed_options - ] - - -def _product_dict(items: List[tuple[str, List[Any]]]): - if not items: - yield {} - return - names = [name for name, _ in items] - values = [values for _, values in items] - for combo in itertools.product(*values): - yield dict(zip(names, combo)) - - -def _mixcv_branch_items( - *, - num_stages: int, - value_space: Dict[str, List[Any]], - fixed_options: Dict[str, Any], -) -> List[List[tuple[str, List[Any]]]]: - independent_names = [ - "enable_tuning_mode", - "unit_flag", - "enable_ubuf_saving", - "enable_hivm_auto_cv_balance", - "enable_auto_bind_sub_block", - ] - independent_items = _emit_values(independent_names, value_space, fixed_options) - - multibuffer_values = _effective_values( - "multibuffer", value_space, fixed_options, None - ) - branches = [] - for multibuffer in multibuffer_values: - multibuffer_items = [] - if "multibuffer" in value_space and "multibuffer" not in fixed_options: - multibuffer_items = [("multibuffer", [multibuffer])] - branch_base = independent_items + multibuffer_items - - if num_stages == 1 or multibuffer is False: - branches.append(branch_base) - continue - - limit_only_values = _effective_values( - "limit_auto_multi_buffer_only_for_local_buffer", - value_space, - fixed_options, - True, - ) - for limit_only in limit_only_values: - limit_only_items = [] - if ( - "limit_auto_multi_buffer_only_for_local_buffer" in value_space - and "limit_auto_multi_buffer_only_for_local_buffer" not in fixed_options - ): - limit_only_items = [ - ("limit_auto_multi_buffer_only_for_local_buffer", [limit_only]) - ] - if limit_only is True: - branch_names = ["limit_auto_multi_buffer_of_local_buffer"] - else: - branch_names = [ - "limit_auto_multi_buffer_of_local_buffer", - "set_workspace_multibuffer", - "tile_mix_vector_loop", - "tile_mix_cube_loop", - ] - branches.append( - branch_base - + limit_only_items - + _emit_values(branch_names, value_space, fixed_options) - ) - - return branches - - -def _make_expanded_config( - config: Config, - spec: CompileOptionsSpec, - combo_value: Dict[str, Any], - num_stages: int, -) -> Config: - new_kwargs = dict(config.kwargs) - for name in _SUPPORTED_PARAMS[spec.kernel_type]: - new_kwargs.pop(name, None) - for name, value in combo_value.items(): - new_kwargs[name] = value - - return _make_config_compat( - kwargs=new_kwargs, - num_warps=getattr(config, "num_warps", 4), - num_stages=num_stages, - num_ctas=getattr(config, "num_ctas", 1), - maxnreg=getattr(config, "maxnreg", None), - pre_hook=getattr(config, "pre_hook", None), - ir_override=getattr(config, "ir_override", None), - num_buffers_warp_spec=getattr(config, "num_buffers_warp_spec", None), - num_consumer_groups=getattr(config, "num_consumer_groups", None), - reg_dec_producer=getattr(config, "reg_dec_producer", None), - reg_inc_consumer=getattr(config, "reg_inc_consumer", None), - ) - - -def _hashable_value(value: Any): - if isinstance(value, dict): - return tuple(sorted((key, _hashable_value(val)) for key, val in value.items())) - if isinstance(value, (list, tuple)): - return tuple(_hashable_value(item) for item in value) - if isinstance(value, set): - return tuple(sorted(_hashable_value(item) for item in value)) - try: - hash(value) - except TypeError: - return repr(value) - return value - - -def _config_key(config: Config) -> tuple: - return ( - tuple( - sorted( - (key, _hashable_value(value)) for key, value in config.kwargs.items() - ) - ), - getattr(config, "num_warps", 4), - getattr(config, "num_stages", None), - getattr(config, "num_ctas", 1), - getattr(config, "maxnreg", None), - id(getattr(config, "pre_hook", None)), - _hashable_value(getattr(config, "ir_override", None)), - getattr(config, "num_buffers_warp_spec", None), - getattr(config, "num_consumer_groups", None), - getattr(config, "reg_dec_producer", None), - getattr(config, "reg_inc_consumer", None), - ) - - -def expand_compile_option_configs( - configs: List[Config], - spec: CompileOptionsSpec, - *, - generated_tiling: bool, - fixed_options: Optional[Dict[str, Any]] = None, -) -> List[Config]: - if not spec.enabled or not configs: - return configs - - fixed_options = fixed_options or {} - expanded_configs = [] - emitted_config_keys = set() - for config in configs: - value_space = _value_space_for_config( - config, spec, generated_tiling=generated_tiling - ) - if "num_stages" in fixed_options: - num_stage_values = [fixed_options["num_stages"]] - value_space.pop("num_stages", None) - else: - num_stage_values = value_space.pop("num_stages") - - for name in fixed_options: - if name != "num_stages": - value_space.pop(name, None) - - for num_stages in num_stage_values: - if spec.kernel_type == "mixcv": - branch_items = _mixcv_branch_items( - num_stages=num_stages, - value_space=value_space, - fixed_options=fixed_options, - ) - combo_iter = itertools.chain.from_iterable( - _product_dict(items) for items in branch_items - ) - else: - combo_iter = _product_dict(list(value_space.items())) - - for combo_value in combo_iter: - new_config = _make_expanded_config( - config, spec, combo_value, num_stages - ) - config_key = _config_key(new_config) - if config_key in emitted_config_keys: - continue - emitted_config_keys.add(config_key) - - if ( - spec.max_configs is not None - and len(expanded_configs) >= spec.max_configs - ): - raise ValueError( - "compile_options generated more than " - f"{spec.max_configs} configs. Narrow the search space or raise max_configs." - ) - expanded_configs.append(new_config) - - return expanded_configs - - -def get_compile_option_param_names(spec: CompileOptionsSpec) -> set[str]: - if not spec.enabled: - return set() - return set(_SUPPORTED_PARAMS[spec.kernel_type]) - - -def format_compile_option_result( - config: Config, - spec: CompileOptionsSpec, - fixed_options: Optional[Dict[str, Any]] = None, -) -> str: - if not spec.enabled: - return str(config) - - fixed_options = fixed_options or {} - compile_param_names = _SUPPORTED_PARAMS[spec.kernel_type] - {"num_stages"} - selected_meta = { - key: value - for key, value in sorted(config.kwargs.items()) - if key not in compile_param_names - } - selected_meta["num_stages"] = getattr(config, "num_stages", None) - - effective = { - key: value - for key, value in sorted(config.kwargs.items()) - if key in compile_param_names - } - effective.update( - { - key: value - for key, value in sorted(fixed_options.items()) - if key in compile_param_names - } - ) - - if spec.kernel_type == "mixcv": - num_stages = selected_meta["num_stages"] - multibuffer = effective.get("multibuffer", None) - effective["enable_auto_multi_buffer"] = ( - False if multibuffer is False or num_stages == 1 else True - ) - for name in _SUPPORTED_PARAMS[spec.kernel_type]: - if name == "num_stages": - continue - reason = _is_inactive_reason( - name, num_stages, config.kwargs, config, fixed_options - ) - if reason is not None and name not in effective: - effective[name] = f"" - if reason is not None and name in effective: - effective[name] = f"" - if is_compile_on_910_95: - for name in _MIXCV_910_95_UNSUPPORTED_PARAMS: - effective[name] = "" - - selected_items = [f"{key}={value}" for key, value in selected_meta.items()] - effective_items = [f"{key}={value}" for key, value in sorted(effective.items())] - return ( - "selected_meta: " - + ", ".join(selected_items) - + "; effective_compile_options: " - + ", ".join(effective_items) - ) - - -def summarize_compile_option_configs( - configs: List[Config], limit: Optional[int] = None -) -> List[str]: - summary = [] - selected_configs = configs if limit is None else configs[:limit] - for config in selected_configs: - items = [f"{key}={value}" for key, value in sorted(config.kwargs.items())] - items.append(f"num_stages={getattr(config, 'num_stages', None)}") - summary.append(", ".join(items)) - return summary diff --git a/backend/ascend_autotune_runtime/kernel_archetype.py b/backend/ascend_autotune_runtime/kernel_archetype.py new file mode 100644 index 00000000..af1c464e --- /dev/null +++ b/backend/ascend_autotune_runtime/kernel_archetype.py @@ -0,0 +1,451 @@ +from __future__ import annotations + +import ast +from dataclasses import dataclass +from enum import Enum +from typing import Any, Dict, Optional, Set + + +class OperatorKind(str, Enum): + DOT_REGULAR = "DOT_REGULAR" + DOT_STATEFUL = "DOT_STATEFUL" + VECTOR_REDUCTION = "VECTOR_REDUCTION" + VECTOR_AFFINE = "VECTOR_AFFINE" + VECTOR_DISCRETE_OR_STATEFUL = "VECTOR_DISCRETE_OR_STATEFUL" + UNKNOWN = "UNKNOWN" + + +class MemoryAccessKind(str, Enum): + AFFINE = "AFFINE" + INDIRECT = "INDIRECT" + UNKNOWN = "UNKNOWN" + + +class ControlKind(str, Enum): + STATIC_OR_RUNTIME = "STATIC_OR_RUNTIME" + DATA_DEPENDENT = "DATA_DEPENDENT" + UNKNOWN = "UNKNOWN" + + +class MemoryEffectKind(str, Enum): + NORMAL = "NORMAL" + SCATTER_OR_UNKNOWN = "SCATTER_OR_UNKNOWN" + UNKNOWN = "UNKNOWN" + + +class _Taint(str, Enum): + INDEX = "INDEX" + MEMORY = "MEMORY" + UNKNOWN = "UNKNOWN" + + +@dataclass(frozen=True) +class OperatorFeatures: + has_dot: bool = False + has_explicit_reduction: bool = False + memory_access_kind: MemoryAccessKind = MemoryAccessKind.AFFINE + control_kind: ControlKind = ControlKind.STATIC_OR_RUNTIME + memory_effect_kind: MemoryEffectKind = MemoryEffectKind.NORMAL + analysis_failed: bool = False + reason: str = "" + + +@dataclass(frozen=True) +class SearchPolicy: + operator_kind: OperatorKind + compile_options_kernel_type: str + stage1_profile_family: str + stage2_enabled: bool + allow_final_unit_flag: bool = False + mixcv_quick_gate: bool = False + no_dot_shape_values: bool = False + + +def analyze_operator_ast( + func_ast: ast.AST, + scope: Optional[Dict[str, Any]] = None, + seen: Optional[Set[int]] = None, +) -> OperatorFeatures: + analyzer = _OperatorFeatureAnalyzer(scope=scope or {}, seen=seen or set()) + try: + analyzer.visit(func_ast) + except Exception as exc: # noqa: BLE001 - classifier must fail closed. + return OperatorFeatures( + memory_access_kind=MemoryAccessKind.UNKNOWN, + control_kind=ControlKind.UNKNOWN, + memory_effect_kind=MemoryEffectKind.UNKNOWN, + analysis_failed=True, + reason=str(exc), + ) + return analyzer.features() + + +def classify_operator(features: OperatorFeatures) -> OperatorKind: + if features.analysis_failed: + return OperatorKind.UNKNOWN + + if features.has_dot: + if features.memory_access_kind != MemoryAccessKind.AFFINE: + return OperatorKind.DOT_STATEFUL + if features.control_kind != ControlKind.STATIC_OR_RUNTIME: + return OperatorKind.DOT_STATEFUL + if features.memory_effect_kind != MemoryEffectKind.NORMAL: + return OperatorKind.DOT_STATEFUL + return OperatorKind.DOT_REGULAR + + if features.memory_access_kind != MemoryAccessKind.AFFINE: + return OperatorKind.VECTOR_DISCRETE_OR_STATEFUL + if features.control_kind == ControlKind.DATA_DEPENDENT: + return OperatorKind.VECTOR_DISCRETE_OR_STATEFUL + if features.memory_effect_kind != MemoryEffectKind.NORMAL: + return OperatorKind.VECTOR_DISCRETE_OR_STATEFUL + if features.has_explicit_reduction: + return OperatorKind.VECTOR_REDUCTION + return OperatorKind.VECTOR_AFFINE + + +def make_search_policy(kind: OperatorKind) -> SearchPolicy: + if kind == OperatorKind.DOT_REGULAR: + return SearchPolicy( + operator_kind=kind, + compile_options_kernel_type="mixcv", + stage1_profile_family="mixcv", + stage2_enabled=True, + allow_final_unit_flag=True, + mixcv_quick_gate=True, + ) + if kind == OperatorKind.DOT_STATEFUL: + return SearchPolicy( + operator_kind=kind, + compile_options_kernel_type="mixcv", + stage1_profile_family="conservative_mixcv", + stage2_enabled=True, + allow_final_unit_flag=False, + mixcv_quick_gate=True, + ) + if kind == OperatorKind.VECTOR_REDUCTION: + return SearchPolicy( + operator_kind=kind, + compile_options_kernel_type="vector", + stage1_profile_family="vector", + stage2_enabled=True, + no_dot_shape_values=True, + ) + if kind == OperatorKind.VECTOR_AFFINE: + return SearchPolicy( + operator_kind=kind, + compile_options_kernel_type="vector", + stage1_profile_family="vector", + stage2_enabled=True, + no_dot_shape_values=True, + ) + if kind == OperatorKind.VECTOR_DISCRETE_OR_STATEFUL: + return SearchPolicy( + operator_kind=kind, + compile_options_kernel_type="vector", + stage1_profile_family="vector", + stage2_enabled=False, + no_dot_shape_values=True, + ) + return SearchPolicy( + operator_kind=OperatorKind.UNKNOWN, + compile_options_kernel_type="vector", + stage1_profile_family="vector", + stage2_enabled=False, + no_dot_shape_values=True, + ) + + +def analyze_operator_policy( + func_ast: ast.AST, + scope: Optional[Dict[str, Any]] = None, + seen: Optional[Set[int]] = None, +) -> tuple[OperatorFeatures, SearchPolicy]: + features = analyze_operator_ast(func_ast, scope=scope, seen=seen) + kind = classify_operator(features) + return features, make_search_policy(kind) + + +class _OperatorFeatureAnalyzer(ast.NodeVisitor): + _REDUCTION_CALLS = {"sum", "max", "min", "prod"} + _DOT_CALLS = {"dot", "dot_scaled"} + _AFFINE_TL_CALLS = { + "arange", + "program_id", + "cdiv", + "minimum", + "maximum", + "where", + "zeros", + "full", + "broadcast_to", + "expand_dims", + } + + def __init__(self, *, scope: Dict[str, Any], seen: Set[int]): + self.scope = scope + self.seen = seen + self.env: Dict[str, _Taint] = {} + self.has_dot = False + self.has_explicit_reduction = False + self.memory_access_kind = MemoryAccessKind.AFFINE + self.control_kind = ControlKind.STATIC_OR_RUNTIME + self.memory_effect_kind = MemoryEffectKind.NORMAL + + def features(self) -> OperatorFeatures: + return OperatorFeatures( + has_dot=self.has_dot, + has_explicit_reduction=self.has_explicit_reduction, + memory_access_kind=self.memory_access_kind, + control_kind=self.control_kind, + memory_effect_kind=self.memory_effect_kind, + ) + + def visit_FunctionDef(self, node: ast.FunctionDef): + for stmt in node.body: + self.visit(stmt) + + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef): + self.visit_FunctionDef(node) # type: ignore[arg-type] + + def visit_Assign(self, node: ast.Assign): + taint = self._expr_taint(node.value) + for target in node.targets: + self._assign_target(target, taint) + self.visit(node.value) + + def visit_AnnAssign(self, node: ast.AnnAssign): + taint = ( + self._expr_taint(node.value) if node.value is not None else _Taint.UNKNOWN + ) + self._assign_target(node.target, taint) + if node.value is not None: + self.visit(node.value) + + def visit_AugAssign(self, node: ast.AugAssign): + taint = self._combine_taint( + self._expr_taint(node.target), self._expr_taint(node.value) + ) + self._assign_target(node.target, taint) + self.visit(node.value) + + def visit_For(self, node: ast.For): + self._update_control_kind(self._expr_taint(node.iter)) + self._assign_target(node.target, _Taint.INDEX) + for stmt in node.body: + self.visit(stmt) + for stmt in node.orelse: + self.visit(stmt) + + def visit_If(self, node: ast.If): + self._update_control_kind(self._expr_taint(node.test)) + for stmt in node.body: + self.visit(stmt) + for stmt in node.orelse: + self.visit(stmt) + + def visit_While(self, node: ast.While): + self._update_control_kind(self._expr_taint(node.test)) + for stmt in node.body: + self.visit(stmt) + for stmt in node.orelse: + self.visit(stmt) + + def visit_Call(self, node: ast.Call): + if self._is_tl_call(node.func, self._DOT_CALLS): + self.has_dot = True + if self._is_tl_call(node.func, self._REDUCTION_CALLS): + self.has_explicit_reduction = True + if self._is_tl_call(node.func, {"load"}): + self._record_load_address(node) + elif self._is_tl_call(node.func, {"store"}): + self._record_store_address(node) + else: + self._merge_called_jit_features(node.func) + self.generic_visit(node) + + def _record_load_address(self, node: ast.Call): + if not node.args: + self.memory_access_kind = MemoryAccessKind.UNKNOWN + return + taint = self._expr_taint(node.args[0]) + if taint == _Taint.MEMORY: + self.memory_access_kind = MemoryAccessKind.INDIRECT + elif ( + taint == _Taint.UNKNOWN + and self.memory_access_kind != MemoryAccessKind.INDIRECT + ): + self.memory_access_kind = MemoryAccessKind.UNKNOWN + + def _record_store_address(self, node: ast.Call): + if not node.args: + self.memory_effect_kind = MemoryEffectKind.UNKNOWN + return + taint = self._expr_taint(node.args[0]) + if taint in {_Taint.MEMORY, _Taint.UNKNOWN}: + self.memory_effect_kind = MemoryEffectKind.SCATTER_OR_UNKNOWN + + def _merge_called_jit_features(self, func: ast.AST): + if not isinstance(func, ast.Name): + return + callee = self.scope.get(func.id) + if callee is None or not callable(getattr(callee, "parse", None)): + return + callee_id = id(callee) + if callee_id in self.seen: + return + self.seen.add(callee_id) + callee_scope = ( + callee.get_capture_scope() + if callable(getattr(callee, "get_capture_scope", None)) + else self.scope + ) + features = analyze_operator_ast( + callee.parse(), scope=callee_scope, seen=self.seen + ) + self.has_dot = self.has_dot or features.has_dot + self.has_explicit_reduction = ( + self.has_explicit_reduction or features.has_explicit_reduction + ) + self.memory_access_kind = _merge_memory_access( + self.memory_access_kind, features.memory_access_kind + ) + self.control_kind = _merge_control(self.control_kind, features.control_kind) + self.memory_effect_kind = _merge_memory_effect( + self.memory_effect_kind, features.memory_effect_kind + ) + + def _assign_target(self, target: ast.AST, taint: _Taint): + if isinstance(target, ast.Name): + self.env[target.id] = taint + return + if isinstance(target, (ast.Tuple, ast.List)): + for elt in target.elts: + self._assign_target(elt, taint) + + def _update_control_kind(self, taint: _Taint): + if taint == _Taint.MEMORY: + self.control_kind = ControlKind.DATA_DEPENDENT + elif ( + taint == _Taint.UNKNOWN and self.control_kind != ControlKind.DATA_DEPENDENT + ): + self.control_kind = ControlKind.UNKNOWN + + def _expr_taint(self, node: Optional[ast.AST]) -> _Taint: + if node is None: + return _Taint.INDEX + if isinstance(node, ast.Name): + return self.env.get(node.id, _Taint.INDEX) + if isinstance(node, ast.Constant): + return _Taint.INDEX + if isinstance(node, ast.Attribute): + return self._expr_taint(node.value) + if isinstance(node, ast.Subscript): + return self._combine_taint( + self._expr_taint(node.value), self._expr_taint(node.slice) + ) + if isinstance(node, ast.Slice): + return self._combine_all( + self._expr_taint(node.lower), + self._expr_taint(node.upper), + self._expr_taint(node.step), + ) + if isinstance(node, ast.UnaryOp): + return self._expr_taint(node.operand) + if isinstance(node, ast.BinOp): + return self._combine_taint( + self._expr_taint(node.left), self._expr_taint(node.right) + ) + if isinstance(node, ast.BoolOp): + return self._combine_all( + *(self._expr_taint(value) for value in node.values) + ) + if isinstance(node, ast.Compare): + return self._combine_all( + self._expr_taint(node.left), + *(self._expr_taint(comp) for comp in node.comparators), + ) + if isinstance(node, ast.IfExp): + return self._combine_all( + self._expr_taint(node.test), + self._expr_taint(node.body), + self._expr_taint(node.orelse), + ) + if isinstance(node, ast.Call): + return self._call_taint(node) + if isinstance(node, (ast.Tuple, ast.List)): + return self._combine_all(*(self._expr_taint(elt) for elt in node.elts)) + return _Taint.UNKNOWN + + def _call_taint(self, node: ast.Call) -> _Taint: + arg_taint = self._combine_all( + *(self._expr_taint(arg) for arg in node.args), + *(self._expr_taint(keyword.value) for keyword in node.keywords), + ) + if self._is_tl_call(node.func, {"load"}): + return _Taint.MEMORY + if self._is_tl_call(node.func, self._DOT_CALLS | self._REDUCTION_CALLS): + return arg_taint + if self._is_tl_call(node.func, self._AFFINE_TL_CALLS): + return arg_taint + if isinstance(node.func, ast.Name) and node.func.id in { + "range", + "min", + "max", + "int", + }: + return arg_taint + if arg_taint == _Taint.MEMORY: + return _Taint.MEMORY + return _Taint.UNKNOWN + + @staticmethod + def _is_tl_call(func: ast.AST, names: Set[str]) -> bool: + return ( + isinstance(func, ast.Attribute) + and func.attr in names + and isinstance(func.value, ast.Name) + and func.value.id == "tl" + ) + + @staticmethod + def _combine_taint(left: _Taint, right: _Taint) -> _Taint: + if _Taint.MEMORY in (left, right): + return _Taint.MEMORY + if _Taint.UNKNOWN in (left, right): + return _Taint.UNKNOWN + return _Taint.INDEX + + def _combine_all(self, *values: _Taint) -> _Taint: + result = _Taint.INDEX + for value in values: + result = self._combine_taint(result, value) + return result + + +def _merge_memory_access( + left: MemoryAccessKind, right: MemoryAccessKind +) -> MemoryAccessKind: + if MemoryAccessKind.INDIRECT in (left, right): + return MemoryAccessKind.INDIRECT + if MemoryAccessKind.UNKNOWN in (left, right): + return MemoryAccessKind.UNKNOWN + return MemoryAccessKind.AFFINE + + +def _merge_control(left: ControlKind, right: ControlKind) -> ControlKind: + if ControlKind.DATA_DEPENDENT in (left, right): + return ControlKind.DATA_DEPENDENT + if ControlKind.UNKNOWN in (left, right): + return ControlKind.UNKNOWN + return ControlKind.STATIC_OR_RUNTIME + + +def _merge_memory_effect( + left: MemoryEffectKind, right: MemoryEffectKind +) -> MemoryEffectKind: + if MemoryEffectKind.SCATTER_OR_UNKNOWN in (left, right): + return MemoryEffectKind.SCATTER_OR_UNKNOWN + if MemoryEffectKind.UNKNOWN in (left, right): + return MemoryEffectKind.UNKNOWN + return MemoryEffectKind.NORMAL diff --git a/backend/ascend_autotune_runtime/autoparser.py b/backend/ascend_autotune_runtime/kernel_ast_analyzer.py similarity index 100% rename from backend/ascend_autotune_runtime/autoparser.py rename to backend/ascend_autotune_runtime/kernel_ast_analyzer.py diff --git a/backend/ascend_autotune_runtime/measurement_cache.py b/backend/ascend_autotune_runtime/measurement_cache.py new file mode 100644 index 00000000..d6ea098c --- /dev/null +++ b/backend/ascend_autotune_runtime/measurement_cache.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +from typing import Any, Mapping, Sequence + +from triton.runtime.autotuner import Config + +from .schedule_profiles import classify_compile_failure, effective_compile_profile_key +from .tile_acquisition_model import resource_overflow_ratio +from .tile_shape_space import extract_shape, shape_key + + +class SearchMeasureCache: + """Exact shape/profile measurement cache for search_params autotune.""" + + def __init__( + self, params: Sequence[str], ub_cache: Mapping[Config, Any] | None = None + ): + self.params = list(params) + self.ub_cache = ub_cache if ub_cache is not None else {} + self._records = {} + + def __len__(self) -> int: + return len(self._records) + + def key(self, shape: Mapping[str, Any], profile: Mapping[str, Any]): + return ( + shape_key(dict(shape), self.params), + effective_compile_profile_key(profile), + ) + + def get(self, shape: Mapping[str, Any], profile: Mapping[str, Any]): + return self._records.get(self.key(shape, profile)) + + def put( + self, + shape: Mapping[str, Any], + profile: Mapping[str, Any], + config: Config | None, + *, + cost=None, + error=None, + failure=None, + source="stage2", + ): + ok = cost is not None and cost != float("inf") and error is None + failure_kind = None if ok else failure or classify_compile_failure(error) + ub = None + if ok and config is not None: + ub = self.ub_cache.get(config) + self._records[self.key(shape, profile)] = { + "ok": ok, + "shape": dict(shape), + "profile": dict(profile), + "config": config, + "time": cost if ok else None, + "failure": failure_kind, + "error": error, + "resource_ratio": None if ok else resource_overflow_ratio(error), + "ub": ub, + "source": source, + } + + def attempt_count(self, shape: Mapping[str, Any]) -> int: + shape_id = shape_key(dict(shape), self.params) + return sum(1 for cached_shape, _ in self._records if cached_shape == shape_id) + + def stage1_failure_observations(self, failed_configs: Sequence[Config]): + failed_shape_keys = { + shape_key(extract_shape(config, self.params), self.params) + for config in failed_configs + } + observations = [] + for cached_shape, profile_key in self._records: + if cached_shape not in failed_shape_keys: + continue + cached = self._records[(cached_shape, profile_key)] + if cached.get("ok") or cached.get("source") != "stage1_fast": + continue + observations.append(cached) + return observations diff --git a/backend/ascend_autotune_runtime/benchmark.py b/backend/ascend_autotune_runtime/measurement_strategy.py similarity index 54% rename from backend/ascend_autotune_runtime/benchmark.py rename to backend/ascend_autotune_runtime/measurement_strategy.py index 25c28b37..9bc2b3c7 100644 --- a/backend/ascend_autotune_runtime/benchmark.py +++ b/backend/ascend_autotune_runtime/measurement_strategy.py @@ -2,7 +2,21 @@ from __future__ import annotations -from typing import Mapping +from typing import Mapping, Optional + + +def ub_bytes_of(compiled_kernel) -> Optional[int]: + """UB allocation in bytes from compile-time metadata. + + Returns None when the kernel was compiled without ``TRITON_MEMORY_DISPLAY=1`` + / ``--enable-memory-display=true``, because the backend then leaves + ``required_ub_bits`` at its 0 default. The Ascend profiler does not expose + static UB size at runtime, so compile metadata is the source of truth. + """ + bits = ( + getattr(getattr(compiled_kernel, "metadata", None), "required_ub_bits", 0) or 0 + ) + return bits // 8 if bits else None class NpuProfilerBenchStrategy: @@ -10,6 +24,9 @@ def bench(self, run_fns: Mapping): from ..testing import do_bench_npu costs = do_bench_npu(list(run_fns.values()), clear_l2_cache=False) + if len(run_fns) == 1 and isinstance(costs, (int, float)): + config = next(iter(run_fns)) + return {config: costs} if not isinstance(costs, (list, tuple)): raise RuntimeError( "do_bench_npu must return one timing per autotune config " diff --git a/backend/ascend_autotune_runtime/resource_memory_parser.py b/backend/ascend_autotune_runtime/resource_memory_parser.py new file mode 100644 index 00000000..4c060df5 --- /dev/null +++ b/backend/ascend_autotune_runtime/resource_memory_parser.py @@ -0,0 +1,119 @@ +"""Parse ``memory_info_{aic,aiv}.json`` emitted by bishengir-compile when +invoked with ``--enable-memory-display=true``. + +Each ``Record`` covers one memory scope and lists every ``memref.alloc()`` +the compiler placed there. The buffer's ``extent`` field is in BITS +(verified empirically: ``extent == element_count * dtype_size_bytes * 8``). +The buffer's ``offset`` field is in BYTES. + +AIC kernels expose UB as ``scope == "cbuf"``; AIV kernels use ``scope == "ub"``. +When offsets are present, the required allocation is the highest planned +end address: ``max(offset + ceil(extent_bits / 8)) * 8``. If a toolchain does +not emit offsets, fall back to a live-buffer overlap over the closed +``life_time_in_ir`` intervals. +""" + +from __future__ import annotations + +import json +from typing import Optional + +import numpy as np + +UB_SCOPE_NAMES = ("cbuf", "ub") + + +def _tuple_or_scalar(value): + if isinstance(value, list): + return tuple(value) + return (value,) if value is not None else () + + +def _peak_overlap(buffers): + """Max sum(extent) over closed live intervals via vectorized sweep. + + Emits ``(start, +extent)`` and ``(end, -extent)`` events, then + sorts start events before end events at the same ``t``. That preserves + closed interval semantics where ``[0, 10]`` overlaps ``[10, 20]`` at 10. + A cumsum over deltas then yields the live allocation at each event; + the max is the peak. + """ + events = [] + for b in buffers: + life = b.get("life_time_in_ir") + extent = b.get("extent") + if ( + isinstance(life, list) + and len(life) == 2 + and isinstance(extent, int) + and extent > 0 + ): + events.append((int(life[0]), +extent)) + events.append((int(life[1]), -extent)) + if not events: + return 0 + times, deltas = zip(*events) + deltas = np.asarray(deltas) + order = np.lexsort((-deltas, np.asarray(times))) + return int(np.cumsum(deltas[order]).max()) + + +def _required_bits_from_offsets(buffers) -> Optional[int]: + """Return planner footprint in bits using byte offsets and bit extents.""" + max_end_bytes = 0 + found = False + for b in buffers: + extent = b.get("extent") + if not isinstance(extent, int) or extent <= 0: + continue + extent_bytes = (extent + 7) // 8 + offsets = b.get("offset") or [] + if not isinstance(offsets, list): + offsets = [offsets] + for offset in offsets: + if isinstance(offset, int) and offset >= 0: + found = True + max_end_bytes = max(max_end_bytes, offset + extent_bytes) + return max_end_bytes * 8 if found else None + + +def peak_ub_bits(json_path: str) -> Optional[int]: + """Required/peak UB allocation in bits, or ``None`` if no UB scope is present. + + Buffers seen in the JSON are deduplicated because some toolchains emit each + allocation twice. ``error_info`` is intentionally ignored; the value is + derived from structured ``memory_info_array`` fields only. + """ + with open(json_path, "r", encoding="utf-8") as f: + records = json.load(f).get("Record") or [] + peaks = [] + for rec in records: + if (rec.get("scope") or "") not in UB_SCOPE_NAMES: + continue + buffers = rec.get("memory_info_array") or [] + seen = set() + dedup = [] + for b in buffers: + key = ( + b.get("buffer"), + b.get("extent"), + _tuple_or_scalar(b.get("life_time_in_ir")), + _tuple_or_scalar(b.get("offset")), + ) + if key in seen: + continue + seen.add(key) + dedup.append(b) + required = _required_bits_from_offsets(dedup) + peaks.append(required if required is not None else _peak_overlap(dedup)) + return max(peaks) if peaks else None + + +if __name__ == "__main__": + import sys + + bits = peak_ub_bits(sys.argv[1]) + if bits is None: + print("no UB scope found") + sys.exit(1) + print(f"peak UB: {bits} bits ({bits / 8 / 1024:.1f} KiB)") diff --git a/backend/ascend_autotune_runtime/schedule_profile_search.py b/backend/ascend_autotune_runtime/schedule_profile_search.py new file mode 100644 index 00000000..f88d0a98 --- /dev/null +++ b/backend/ascend_autotune_runtime/schedule_profile_search.py @@ -0,0 +1,534 @@ +from __future__ import annotations + +import math +import random +from dataclasses import dataclass, field +from typing import Any, Callable, Mapping, Sequence + +from triton.backends.dicp_triton.utils import is_compile_on_910_95 + +from .schedule_profiles import ( + CompileFailureRegionSet, + classify_compile_failure, + compile_profile_resource_not_less, + compile_profile_to_config, + effective_compile_profile_key, + generate_linked_compile_neighbors, + make_stage2_seed_profiles, +) +from .measurement_cache import SearchMeasureCache + + +Stage2Candidate = tuple[Any, Mapping[str, Any], Any] + + +@dataclass +class _Stage2State: + base_profile: Mapping[str, Any] + rng: random.Random + timing_sort_key: Callable[[Any], float] + candidates: list[Stage2Candidate] = field(default_factory=list) + failed_regions: CompileFailureRegionSet = field( + default_factory=CompileFailureRegionSet + ) + seen_profiles: set = field(default_factory=set) + current_profile: Mapping[str, Any] | None = None + current_time: Any = float("inf") + best_profile: Mapping[str, Any] | None = None + best_time: Any = float("inf") + + def __post_init__(self): + self.current_profile = self.base_profile + self.best_profile = self.base_profile + self.seen_profiles.add(effective_compile_profile_key(self.base_profile)) + + def accept_annealed(self, new_cost, temperature: float) -> bool: + old_value = self.timing_sort_key(self.current_time) + new_value = self.timing_sort_key(new_cost) + if new_value <= old_value: + return True + if temperature <= 0: + return False + return self.rng.random() < math.exp((old_value - new_value) / temperature) + + def add_success(self, config, profile, cost, temperature: float): + self.candidates.append((config, profile, cost)) + if self.timing_sort_key(cost) < self.timing_sort_key(self.best_time): + self.best_profile = profile + self.best_time = cost + if self.accept_annealed(cost, temperature): + self.current_profile = profile + self.current_time = cost + + +class Stage2CompileSearcher: + """Branch-aware compile option search for one selected Stage-1 shape.""" + + def __init__( + self, + *, + search_params, + operator_policy, + cache: SearchMeasureCache, + apply_fixed_profile: Callable[[Mapping[str, Any]], dict], + bench_configs: Callable[[Sequence[Any]], tuple[dict, dict]], + timing_sort_key: Callable[[Any], float], + is_finite_timing: Callable[[Any], bool], + debug: Callable[[str], None], + format_error_log: Callable[[Any], str], + ): + self.search_params = search_params + self.operator_policy = operator_policy + self.cache = cache + self.apply_fixed_profile = apply_fixed_profile + self.bench_configs = bench_configs + self.timing_sort_key = timing_sort_key + self.is_finite_timing = is_finite_timing + self.debug = debug + self.format_error_log = format_error_log + + def search(self, *, stage1_entry, stage1_rank: int, stage1_top1_time): + shape = stage1_entry["shape"] + shape_config = stage1_entry["shape_config"] + stage1_ub = stage1_entry.get("ub") + rng = random.Random(self.search_params.shape_final_top_k + len(shape)) + state = _Stage2State( + base_profile=stage1_entry["profile"], + rng=rng, + timing_sort_key=self.timing_sort_key, + ) + trials = max(1, self.cache.attempt_count(shape)) + + self.debug( + "Stage 2 start " + f"shape={shape}, stage1_base_time={stage1_entry['time']}, " + f"stage1_ub={stage1_ub}, base={stage1_entry['profile']}" + ) + if not getattr(self.operator_policy, "stage2_enabled", True): + self.debug( + "Stage 2 skipped " + f"shape={shape}, operator_kind={self.operator_policy.operator_kind.value}" + ) + return [] + + batch_size = max(1, self.search_params.neighbors_per_step) + main_trial_budget = max(1, self.search_params.max_compile_trials_per_shape - 1) + prefer_resource_relax = stage1_ub is not None and stage1_ub >= 160 * 1024 + + base_result = self._cached_or_bench_profile( + shape=shape, + shape_config=shape_config, + profile=stage1_entry["profile"], + item_index=0, + ) + self._apply_result( + state, + base_result, + "base cache" if base_result["source"] == "cache" else "base fast", + self.search_params.compile_initial_temperature, + shape, + ) + base_cost = ( + base_result["cost"] if base_result["error"] is None else float("inf") + ) + if not self.is_finite_timing(base_cost): + self.debug( + f"Stage 2 gate stop shape={shape}, rank={stage1_rank}, " + "reason=base_failed" + ) + return state.candidates + + stage1_top1_value = self.timing_sort_key(stage1_top1_time) + base_ratio = self.timing_sort_key(base_cost) / stage1_top1_value + if base_ratio > 3.0: + self.debug( + f"Stage 2 gate stop shape={shape}, rank={stage1_rank}, " + f"base_ratio={base_ratio:.3f} > 3.000" + ) + return state.candidates + + if base_ratio >= 1.5 and self.operator_policy.mixcv_quick_gate: + trials += self._run_quick_gate( + state, + shape=shape, + shape_config=shape_config, + base_profile=stage1_entry["profile"], + base_ratio=base_ratio, + stage1_rank=stage1_rank, + ) + quick_ratio = self.timing_sort_key(state.best_time) / stage1_top1_value + if quick_ratio > 1.5: + self.debug( + f"Stage 2 gate stop shape={shape}, rank={stage1_rank}, " + f"quick_ratio={quick_ratio:.3f} > 1.500" + ) + return self._top_candidates(state) + self.debug( + f"Stage 2 gate pass shape={shape}, rank={stage1_rank}, " + f"quick_ratio={quick_ratio:.3f} <= 1.500" + ) + + trials = self._run_seed_phase( + state, + shape=shape, + shape_config=shape_config, + stage1_ub=stage1_ub, + trials=trials, + main_trial_budget=main_trial_budget, + batch_size=batch_size, + ) + trials = self._run_anneal_phase( + state, + shape=shape, + shape_config=shape_config, + trials=trials, + main_trial_budget=main_trial_budget, + batch_size=batch_size, + prefer_resource_relax=prefer_resource_relax, + ) + self._run_final_unit_flag_trial( + state, + shape=shape, + shape_config=shape_config, + trials=trials, + ) + self.debug( + f"Stage 2 done shape={shape}, trials={trials}, " + f"best_time={state.best_time}, best_profile={state.best_profile}" + ) + return self._top_candidates(state) + + def _run_quick_gate( + self, + state: _Stage2State, + *, + shape, + shape_config, + base_profile, + base_ratio, + stage1_rank, + ) -> int: + quick_items = [] + for quick_index, (workspace, cube, vector) in enumerate( + ((4, 2, 2), (2, 1, 1)), 1 + ): + profile = dict( + base_profile, set_workspace_multibuffer=workspace, unit_flag=False + ) + if not is_compile_on_910_95: + profile["tile_mix_cube_loop"] = cube + profile["tile_mix_vector_loop"] = vector + profile = self.apply_fixed_profile(profile) + if not self._mark_runnable(state, profile, shape, f"quick[{quick_index}]"): + continue + quick_items.append((quick_index, profile)) + + if not quick_items: + return 0 + self.debug( + f"Stage 2 quick gate shape={shape}, rank={stage1_rank}, " + f"base_ratio={base_ratio:.3f}, size={len(quick_items)}" + ) + for result in self._bench_profile_batch( + shape=shape, + shape_config=shape_config, + profile_items=quick_items, + ): + self._apply_result(state, result, f"quick[{result['index']}]", 0.20, shape) + return len(quick_items) + + def _run_seed_phase( + self, + state: _Stage2State, + *, + shape, + shape_config, + stage1_ub, + trials, + main_trial_budget, + batch_size, + ) -> int: + seeds = make_stage2_seed_profiles( + state.base_profile, + shape_kwargs=shape, + seed_budget=self.search_params.seed_budget, + allow_unit_flag=False, + stage1_ub_bytes=stage1_ub, + ) + seed_batch = [] + + def flush_seed_batch(): + nonlocal trials, seed_batch + if not seed_batch: + return + self.debug(f"Stage 2 seed batch shape={shape}, size={len(seed_batch)}") + results = self._bench_profile_batch( + shape=shape, + shape_config=shape_config, + profile_items=seed_batch, + ) + trials += len(seed_batch) + for result in results: + self._apply_result( + state, result, f"seed[{result['index']}]", 0.20, shape + ) + seed_batch = [] + + for seed_index, profile in enumerate(seeds, 1): + profile = self.apply_fixed_profile(profile) + if not self._mark_runnable(state, profile, shape, f"seed[{seed_index}]"): + continue + cached = self.cache.get(shape, profile) + if cached is not None: + self._apply_cached_result( + state, cached, profile, f"seed[{seed_index}]", shape, 0.20 + ) + continue + if trials >= main_trial_budget: + break + if seed_batch and self._has_resource_dependency(profile, seed_batch): + flush_seed_batch() + if state.failed_regions.is_forbidden(profile): + self.debug(f"Stage 2 seed[{seed_index}] forbidden shape={shape}") + continue + self.debug( + f"Stage 2 seed[{seed_index}] queued shape={shape}, profile={profile}" + ) + seed_batch.append((seed_index, profile)) + if ( + len(seed_batch) >= batch_size + or trials + len(seed_batch) >= main_trial_budget + ): + flush_seed_batch() + + flush_seed_batch() + return trials + + def _run_anneal_phase( + self, + state: _Stage2State, + *, + shape, + shape_config, + trials, + main_trial_budget, + batch_size, + prefer_resource_relax, + ) -> int: + temperature = self.search_params.compile_initial_temperature + while trials < main_trial_budget: + neighbors = [ + self.apply_fixed_profile(profile) + for profile in generate_linked_compile_neighbors( + state.current_profile, + shape_kwargs=shape, + limit=self.search_params.neighbors_per_step, + allow_unit_flag=False, + prefer_resource_relax=prefer_resource_relax, + ) + ] + runnable = [ + profile for profile in neighbors if self._is_runnable(state, profile) + ] + if not runnable: + break + + anneal_batch = [] + while ( + runnable + and len(anneal_batch) < batch_size + and trials + len(anneal_batch) < main_trial_budget + ): + profile = state.rng.choice(runnable) + runnable.remove(profile) + if anneal_batch and self._has_resource_dependency( + profile, anneal_batch + ): + runnable.append(profile) + break + state.seen_profiles.add(effective_compile_profile_key(profile)) + cached = self.cache.get(shape, profile) + if cached is None: + anneal_batch.append((trials + len(anneal_batch) + 1, profile)) + continue + self._apply_cached_result( + state, cached, profile, "anneal", shape, temperature + ) + temperature *= self.search_params.compile_cooling + + if not anneal_batch: + continue + self.debug(f"Stage 2 anneal batch shape={shape}, size={len(anneal_batch)}") + results = self._bench_profile_batch( + shape=shape, + shape_config=shape_config, + profile_items=anneal_batch, + ) + trials += len(anneal_batch) + for result in results: + self.debug( + "Stage 2 anneal trial=" + f"{result['index']}, shape={shape}, profile={result['profile']}" + ) + self._apply_result(state, result, "anneal", temperature, shape) + temperature *= self.search_params.compile_cooling + return trials + + def _run_final_unit_flag_trial(self, state, *, shape, shape_config, trials): + if ( + not self.operator_policy.allow_final_unit_flag + or state.best_profile.get("unit_flag", False) + or trials >= self.search_params.max_compile_trials_per_shape + ): + return + trial_profile = self.apply_fixed_profile( + dict(state.best_profile, unit_flag=True) + ) + trial_key = effective_compile_profile_key(trial_profile) + if trial_key in state.seen_profiles: + return + state.seen_profiles.add(trial_key) + self.debug(f"Stage 2 final unit_flag trial shape={shape}") + for result in self._bench_profile_batch( + shape=shape, + shape_config=shape_config, + profile_items=[(trials + 1, trial_profile)], + ): + if result["error"] is None: + state.candidates.append( + (result["config"], trial_profile, result["cost"]) + ) + + def _cached_or_bench_profile(self, *, shape, shape_config, profile, item_index): + cached = self.cache.get(shape, profile) + if cached is not None: + return self._cache_result(profile, cached, item_index) + return self._bench_profile_batch( + shape=shape, + shape_config=shape_config, + profile_items=[(item_index, profile)], + )[0] + + def _bench_profile_batch(self, *, shape, shape_config, profile_items): + profile_items = [ + (item_index, self.apply_fixed_profile(profile)) + for item_index, profile in profile_items + ] + configs = [ + compile_profile_to_config( + profile, shape_kwargs=shape, base_config=shape_config + ) + for _, profile in profile_items + ] + timings, errors = self.bench_configs(configs) + results = [] + for (item_index, profile), config in zip(profile_items, configs): + cost = timings.get(config, float("inf")) + if self.is_finite_timing(cost): + self.cache.put(shape, profile, config, cost=cost, source="stage2_fast") + results.append( + self._result(item_index, profile, config, cost, None, None, "bench") + ) + continue + error = errors.get(config) or RuntimeError("benchmark returned inf") + failure = classify_compile_failure(error) + self.cache.put( + shape, + profile, + config, + error=error, + failure=failure, + source="stage2_fast", + ) + results.append( + self._result(item_index, profile, config, None, error, failure, "bench") + ) + return results + + def _apply_result(self, state, result, stage_name, temperature, shape): + profile = result["profile"] + if result["error"] is not None: + failure = state.failed_regions.add(profile, result["failure"]) + self.debug( + f"Stage 2 {stage_name} fail shape={shape}, " + f"classified_as={failure}, " + f"error_log:\n{self.format_error_log(result['error'])}" + ) + return + state.add_success(result["config"], profile, result["cost"], temperature) + self.debug(f"Stage 2 {stage_name} success shape={shape}, cost={result['cost']}") + + def _apply_cached_result( + self, state, cached, profile, stage_name, shape, temperature + ): + result = self._cache_result(profile, cached, 0) + self._apply_result(state, result, f"{stage_name} cache", temperature, shape) + + def _cache_result(self, profile, cached, item_index): + if cached["ok"]: + return self._result( + item_index, + profile, + cached["config"], + cached["time"], + None, + None, + "cache", + ) + return self._result( + item_index, + profile, + cached.get("config"), + None, + cached.get("error"), + cached.get("failure"), + "cache", + ) + + @staticmethod + def _result(index, profile, config, cost, error, failure, source): + return { + "index": index, + "profile": profile, + "config": config, + "cost": cost, + "error": error, + "failure": failure, + "source": source, + } + + def _mark_runnable(self, state, profile, shape, label, *, quiet=False): + profile_key = effective_compile_profile_key(profile) + if profile_key in state.seen_profiles: + if label and not quiet: + self.debug(f"Stage 2 {label} duplicate shape={shape}") + return False + if state.failed_regions.is_forbidden(profile): + if label and not quiet: + self.debug(f"Stage 2 {label} forbidden shape={shape}") + return False + state.seen_profiles.add(profile_key) + return True + + @staticmethod + def _is_runnable(state, profile): + profile_key = effective_compile_profile_key(profile) + if profile_key in state.seen_profiles: + return False + if state.failed_regions.is_forbidden(profile): + return False + return True + + @staticmethod + def _has_resource_dependency(profile, profile_batch): + return any( + compile_profile_resource_not_less(profile, pending_profile) + or compile_profile_resource_not_less(pending_profile, profile) + for _, pending_profile in profile_batch + ) + + def _top_candidates(self, state): + return sorted( + state.candidates, + key=lambda item: self.timing_sort_key(item[2]), + )[: self.search_params.candidate_pool_per_shape] diff --git a/backend/ascend_autotune_runtime/schedule_profiles.py b/backend/ascend_autotune_runtime/schedule_profiles.py new file mode 100644 index 00000000..c822e844 --- /dev/null +++ b/backend/ascend_autotune_runtime/schedule_profiles.py @@ -0,0 +1,1662 @@ +from __future__ import annotations + +import inspect +import itertools +import re +from dataclasses import dataclass, field +from typing import Any, Dict, List, Mapping, Optional + +from triton.runtime.autotuner import Config +from triton.backends.dicp_triton.utils import is_compile_on_910_95 + + +DEFAULT_MAX_CONFIGS = None + +_VALID_VALUES = { + "num_stages": [1, 2], + "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], + "set_workspace_multibuffer": [2, 4], + "tile_mix_vector_loop": [1, 2, 4, 8], + "tile_mix_cube_loop": [1, 2, 4, 8], +} + +_BOOLEAN_PARAMS = { + "enable_tuning_mode", + "multibuffer", + "unit_flag", + "limit_auto_multi_buffer_only_for_local_buffer", + "enable_hivm_auto_cv_balance", + "enable_ubuf_saving", + # "enable_preload", + "enable_auto_bind_sub_block", +} + +_SUPPORTED_PARAMS = { + "cube": { + "enable_tuning_mode", + "num_stages", + "unit_flag", + "limit_auto_multi_buffer_of_local_buffer", + }, + "mixcv": { + "enable_tuning_mode", + "num_stages", + "multibuffer", + "unit_flag", + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "enable_hivm_auto_cv_balance", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + "enable_ubuf_saving", + # "enable_preload", + "enable_auto_bind_sub_block", + }, + "vector": { + "num_stages", + "enable_ubuf_saving", + }, +} + +_ALL_PARAMS = set().union(*_SUPPORTED_PARAMS.values()) + +_MIXCV_910_95_UNSUPPORTED_PARAMS = { + "tile_mix_vector_loop", + "tile_mix_cube_loop", +} + +COMPILE_MODE_KEY = "mode" +COMPILE_MODE_VECTOR = "VECTOR" +COMPILE_MODE_MB_OFF = "MB_OFF" +COMPILE_MODE_LOCAL_MB = "LOCAL_MB" +COMPILE_MODE_WORKSPACE_CV = "WORKSPACE_CV" + +FAILURE_EXACT_ONLY = "EXACT_ONLY" +FAILURE_RESOURCE_UB = "RESOURCE_UB" +FAILURE_RESOURCE_L0C = "RESOURCE_L0C" +FAILURE_RESOURCE_L1 = "RESOURCE_L1" +FAILURE_RESOURCE_WORKSPACE = "RESOURCE_WORKSPACE" +FAILURE_SYNC_OR_CORRECTNESS = "SYNC_OR_CORRECTNESS" +FAILURE_COMPILER_INTERNAL = "COMPILER_INTERNAL" + +_RESOURCE_FAILURES = { + FAILURE_RESOURCE_UB, + FAILURE_RESOURCE_L0C, + FAILURE_RESOURCE_L1, + FAILURE_RESOURCE_WORKSPACE, +} + +_ALL_FAILURE_KINDS = { + FAILURE_EXACT_ONLY, + FAILURE_RESOURCE_UB, + FAILURE_RESOURCE_L0C, + FAILURE_RESOURCE_L1, + FAILURE_RESOURCE_WORKSPACE, + FAILURE_SYNC_OR_CORRECTNESS, + FAILURE_COMPILER_INTERNAL, +} + +_MULTI_BUFFER_CHILD_PARAMS = { + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "tile_mix_vector_loop", + "tile_mix_cube_loop", +} + +_WORKSPACE_CV_PARAMS = { + "set_workspace_multibuffer", + "tile_mix_vector_loop", + "tile_mix_cube_loop", +} + +WORKSPACE_CV_AGGRESSIVE_PROBE = { + COMPILE_MODE_KEY: COMPILE_MODE_WORKSPACE_CV, + "enable_tuning_mode": True, + "num_stages": 2, + "multibuffer": True, + "enable_auto_bind_sub_block": True, + "enable_hivm_auto_cv_balance": True, + "enable_ubuf_saving": True, + "limit_auto_multi_buffer_only_for_local_buffer": False, + "limit_auto_multi_buffer_of_local_buffer": "no-limit", + "set_workspace_multibuffer": 4, + "tile_mix_cube_loop": 4, + "tile_mix_vector_loop": 4, + "unit_flag": False, +} + +WORKSPACE_CV_MIX1_PROBE = { + COMPILE_MODE_KEY: COMPILE_MODE_WORKSPACE_CV, + "enable_tuning_mode": True, + "num_stages": 2, + "multibuffer": True, + "enable_auto_bind_sub_block": True, + "enable_hivm_auto_cv_balance": True, + "enable_ubuf_saving": True, + "limit_auto_multi_buffer_only_for_local_buffer": False, + "limit_auto_multi_buffer_of_local_buffer": "no-l0c", + "set_workspace_multibuffer": 2, + "tile_mix_cube_loop": 1, + "tile_mix_vector_loop": 1, + "unit_flag": False, +} + +WORKSPACE_CV_LOW_RESOURCE_PROBE = { + COMPILE_MODE_KEY: COMPILE_MODE_WORKSPACE_CV, + "enable_tuning_mode": True, + "num_stages": 2, + "multibuffer": True, + "enable_auto_bind_sub_block": True, + "enable_hivm_auto_cv_balance": True, + "enable_ubuf_saving": True, + "limit_auto_multi_buffer_only_for_local_buffer": False, + "limit_auto_multi_buffer_of_local_buffer": "no-l0c", + "set_workspace_multibuffer": 2, + "tile_mix_cube_loop": 4, + "tile_mix_vector_loop": 4, + "unit_flag": False, +} + +DEFAULT_STAGE1_PROBE_PROFILES = [ + WORKSPACE_CV_AGGRESSIVE_PROBE, + WORKSPACE_CV_LOW_RESOURCE_PROBE, + WORKSPACE_CV_MIX1_PROBE, +] + +VECTOR_STAGE1_PROBE = { + COMPILE_MODE_KEY: COMPILE_MODE_VECTOR, + "num_stages": 2, + "enable_ubuf_saving": True, +} + +DEFAULT_VECTOR_STAGE1_PROBE_PROFILES = [ + VECTOR_STAGE1_PROBE, +] + +CONSERVATIVE_MIXCV_STAGE1_PROBE_PROFILES = [ + WORKSPACE_CV_LOW_RESOURCE_PROBE, + WORKSPACE_CV_MIX1_PROBE, +] + +DEFAULT_COMPILE_ANNEALING_OPTIONS = { + "seed_budget": 8, + "max_compile_trials_per_shape": 16, + "neighbors_per_step": 2, + "compile_initial_temperature": 0.20, + "compile_cooling": 0.85, + "candidate_pool_per_shape": 3, + "random_seed": 0, +} + +DEFAULT_UB_LIMIT_BYTES = 192 * 1024 +LOCAL_ONLY_PROBE_UB_THRESHOLD_BYTES = 100 * 1024 +TIGHT_UB_MARGIN_BYTES = 32 * 1024 + +_AUTO_SEARCH_PRESETS = { + "cube": { + "enable_tuning_mode": [True], + "num_stages": [1, 2], + "unit_flag": [False, True], + "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], + }, + "mixcv": { + "enable_tuning_mode": [True], + "num_stages": [1, 2], + "unit_flag": [False, True], + "limit_auto_multi_buffer_only_for_local_buffer": [True, False], + "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], + "set_workspace_multibuffer": [2, 4], + "enable_hivm_auto_cv_balance": [True], + "tile_mix_vector_loop": [1, 2, 4], + "tile_mix_cube_loop": [1, 2, 4], + "enable_ubuf_saving": [False, True], + # "enable_preload": [False, True], + "enable_auto_bind_sub_block": [True], + }, + "vector": { + "num_stages": [1, 2], + "enable_ubuf_saving": [True, False], + }, +} + + +def _is_mixcv_multi_buffer_auto_enabled( + num_stages: int, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], +) -> bool: + """Whether `enable_auto_multi_buffer` takes effect for this combination.""" + if num_stages == 1: + return False + multibuffer = _resolve_compile_option( + "multibuffer", combo, config, fixed_options, default=None + ) + return multibuffer is not False + + +def _is_mixcv_limit_to_local_only_active( + num_stages: int, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], +) -> bool: + return _is_mixcv_multi_buffer_auto_enabled(num_stages, combo, config, fixed_options) + + +def _is_mixcv_workspace_multibuffer_active( + num_stages: int, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], +) -> bool: + if not _is_mixcv_multi_buffer_auto_enabled( + num_stages, combo, config, fixed_options + ): + return False + limit_to_local_only = _resolve_compile_option( + "limit_auto_multi_buffer_only_for_local_buffer", + combo, + config, + fixed_options, + default=True, + ) + return limit_to_local_only is False + + +_MIXCV_OPTION_ACTIVITY_RULES = { + "limit_auto_multi_buffer_only_for_local_buffer": _is_mixcv_limit_to_local_only_active, + "limit_auto_multi_buffer_of_local_buffer": _is_mixcv_limit_to_local_only_active, + "set_workspace_multibuffer": _is_mixcv_workspace_multibuffer_active, + "tile_mix_vector_loop": _is_mixcv_workspace_multibuffer_active, + "tile_mix_cube_loop": _is_mixcv_workspace_multibuffer_active, +} + + +@dataclass +class CompileOptionsSpec: + enabled: bool = False + kernel_type: str = "mixcv" + params: Dict[str, List[Any]] = field(default_factory=dict) + max_configs: Optional[int] = DEFAULT_MAX_CONFIGS + + +def _normalize_kernel_type(kernel_type: str) -> str: + if kernel_type == "mix": + return "mixcv" + if kernel_type not in _SUPPORTED_PARAMS: + raise ValueError( + "compile_options kernel_type must be one of: cube, mix, mixcv, vector" + ) + return kernel_type + + +def _as_value_list(name: str, value: Any) -> List[Any]: + values = list(value) if isinstance(value, (list, tuple)) else [value] + if not values: + raise ValueError(f"compile_options parameter '{name}' must not be empty") + return values + + +def validate_compile_option_values(name: str, values: List[Any]) -> None: + if name in _BOOLEAN_PARAMS and not all(isinstance(v, bool) for v in values): + raise ValueError(f"compile_options parameter '{name}' expects boolean values") + + if name in _VALID_VALUES and not all(v in _VALID_VALUES[name] for v in values): + raise ValueError( + f"compile_options parameter '{name}' expects values in {_VALID_VALUES[name]}" + ) + + +def parse_compile_options_hint(hint: Any) -> CompileOptionsSpec: + if hint is None or hint is False: + return CompileOptionsSpec(enabled=False) + + if hint is True: + return CompileOptionsSpec(enabled=True) + + if isinstance(hint, str): + return CompileOptionsSpec( + enabled=True, + kernel_type=_normalize_kernel_type(hint), + ) + + if not isinstance(hint, dict): + raise TypeError("hints['compile_options'] must be bool, str, or dict") + + raw = dict(hint) + kernel_type = _normalize_kernel_type( + raw.pop("kernel_type", raw.pop("type", "mixcv")) + ) + max_configs = raw.pop("max_configs", DEFAULT_MAX_CONFIGS) + if max_configs is not None and ( + not isinstance(max_configs, int) or max_configs <= 0 + ): + raise ValueError( + "compile_options max_configs must be a positive integer or None" + ) + + nested_options = raw.pop("options", {}) + if nested_options: + if not isinstance(nested_options, dict): + raise TypeError("compile_options options must be a dict") + raw.update(nested_options) + + supported = _SUPPORTED_PARAMS[kernel_type] + params: Dict[str, List[Any]] = {} + for name, value in raw.items(): + if name not in _ALL_PARAMS: + raise ValueError(f"Unknown compile_options parameter: {name}") + if name not in supported: + raise ValueError( + f"compile_options parameter '{name}' is not supported " + f"for kernel_type '{kernel_type}'" + ) + values = _as_value_list(name, value) + validate_compile_option_values(name, values) + params[name] = values + + return CompileOptionsSpec( + enabled=True, + kernel_type=kernel_type, + params=params, + max_configs=max_configs, + ) + + +def _make_config_compat(**kwargs): + supported_config_args = inspect.signature(Config).parameters + return Config( + **{key: value for key, value in kwargs.items() if key in supported_config_args} + ) + + +def _drop_unsupported_profile_options(profile: Mapping[str, Any]) -> Dict[str, Any]: + copied = dict(profile) + if is_compile_on_910_95: + for name in _MIXCV_910_95_UNSUPPORTED_PARAMS: + copied.pop(name, None) + return copied + + +def _profile_mode(profile: Mapping[str, Any]) -> str: + if COMPILE_MODE_KEY in profile: + return profile[COMPILE_MODE_KEY] + num_stages = profile.get("num_stages") + multibuffer = profile.get("multibuffer", None) + if num_stages == 1 or multibuffer is False: + return COMPILE_MODE_MB_OFF + if profile.get("limit_auto_multi_buffer_only_for_local_buffer", True) is True: + return COMPILE_MODE_LOCAL_MB + return COMPILE_MODE_WORKSPACE_CV + + +def _compile_profile_error(profile: Mapping[str, Any]) -> Optional[str]: + mode = _profile_mode(profile) + if mode not in { + COMPILE_MODE_VECTOR, + COMPILE_MODE_MB_OFF, + COMPILE_MODE_LOCAL_MB, + COMPILE_MODE_WORKSPACE_CV, + }: + return f"unknown compile profile mode: {mode}" + + supported = ( + _SUPPORTED_PARAMS["vector"] + if mode == COMPILE_MODE_VECTOR + else _SUPPORTED_PARAMS["mixcv"] + ) | {COMPILE_MODE_KEY} + unknown = sorted(name for name in profile if name not in supported) + if unknown: + return f"unknown compile profile option(s): {unknown}" + + if "num_stages" not in profile: + return "compile profile must include num_stages" + + for name, value in profile.items(): + if name == COMPILE_MODE_KEY: + continue + if name == "num_stages": + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + return "compile profile num_stages must be a positive integer" + continue + try: + validate_compile_option_values(name, [value]) + except ValueError as exc: + return str(exc) + + if is_compile_on_910_95: + unsupported = sorted(_MIXCV_910_95_UNSUPPORTED_PARAMS & set(profile)) + if unsupported: + return f"compile profile uses unsupported 910_95 option(s): {unsupported}" + + num_stages = profile["num_stages"] + if mode == COMPILE_MODE_VECTOR: + return None + + multibuffer = profile.get("multibuffer", None) + auto_mb_disabled = num_stages == 1 or multibuffer is False + child_params = _MULTI_BUFFER_CHILD_PARAMS & set(profile) + + if mode == COMPILE_MODE_MB_OFF: + if num_stages != 1: + return "MB_OFF profile must use num_stages=1" + if child_params: + return f"MB_OFF profile must not include multi-buffer child option(s): {sorted(child_params)}" + return None + + if auto_mb_disabled: + return ( + "num_stages=1 or multibuffer=False disables auto multi-buffer; " + "use MB_OFF without child options" + ) + + if num_stages != 2: + return f"{mode} profile must use num_stages=2" + if multibuffer is not True: + return f"{mode} profile must include multibuffer=True" + + limit_local = profile.get("limit_auto_multi_buffer_only_for_local_buffer", None) + if mode == COMPILE_MODE_LOCAL_MB: + if limit_local is not True: + return "LOCAL_MB profile must set limit_auto_multi_buffer_only_for_local_buffer=True" + workspace_params = _WORKSPACE_CV_PARAMS & set(profile) + if workspace_params: + return f"LOCAL_MB profile must not include workspace option(s): {sorted(workspace_params)}" + if "limit_auto_multi_buffer_of_local_buffer" not in profile: + return ( + "LOCAL_MB profile must include limit_auto_multi_buffer_of_local_buffer" + ) + return None + + if limit_local is not False: + return "WORKSPACE_CV profile must set limit_auto_multi_buffer_only_for_local_buffer=False" + if "set_workspace_multibuffer" not in profile: + return "WORKSPACE_CV profile must include set_workspace_multibuffer" + if not is_compile_on_910_95: + for name in ("tile_mix_cube_loop", "tile_mix_vector_loop"): + if name not in profile: + return f"WORKSPACE_CV profile must include {name}" + return None + + +def validate_compile_profile( + profile: Mapping[str, Any], *, raise_on_error: bool = False +) -> bool: + """Validate a generated search profile without repairing inactive fields.""" + error = _compile_profile_error(profile) + if error is None: + return True + if raise_on_error: + raise ValueError(error) + return False + + +def apply_fixed_compile_options_to_profile( + profile: Mapping[str, Any], + fixed_options: Optional[Mapping[str, Any]] = None, +) -> Dict[str, Any]: + """Overlay runtime-fixed compile options onto a generated profile.""" + fixed_options = fixed_options or {} + copied = dict(profile) + if not fixed_options: + return copied + + mode = _profile_mode(copied) + supported = ( + _SUPPORTED_PARAMS["vector"] + if mode == COMPILE_MODE_VECTOR + else _SUPPORTED_PARAMS["mixcv"] + ) + unsupported = sorted(name for name in fixed_options if name not in supported) + if unsupported: + raise ValueError( + f"fixed compile option(s) {unsupported} are not valid for profile mode {mode}" + ) + copied.update(fixed_options) + return copied + + +def effective_compile_profile_key(profile: Mapping[str, Any]) -> tuple: + validate_compile_profile(profile, raise_on_error=True) + effective = { + key: value for key, value in profile.items() if key != COMPILE_MODE_KEY + } + return ( + _profile_mode(profile), + tuple( + sorted((key, _hashable_value(value)) for key, value in effective.items()) + ), + ) + + +def compile_profile_to_config( + profile: Mapping[str, Any], + *, + shape_kwargs: Optional[Mapping[str, Any]] = None, + base_config: Optional[Config] = None, +) -> Config: + """Convert a validated profile to Triton Config and keep mode internal.""" + validate_compile_profile(profile, raise_on_error=True) + kwargs = dict(base_config.kwargs) if base_config is not None else {} + if shape_kwargs: + kwargs.update(shape_kwargs) + for name in _ALL_PARAMS: + kwargs.pop(name, None) + for name, value in profile.items(): + if name in {COMPILE_MODE_KEY, "num_stages"}: + continue + kwargs[name] = value + + return _make_config_compat( + kwargs=kwargs, + num_warps=getattr(base_config, "num_warps", 4), + num_stages=profile["num_stages"], + num_ctas=getattr(base_config, "num_ctas", 1), + maxnreg=getattr(base_config, "maxnreg", None), + pre_hook=getattr(base_config, "pre_hook", None), + ir_override=getattr(base_config, "ir_override", None), + num_buffers_warp_spec=getattr(base_config, "num_buffers_warp_spec", None), + num_consumer_groups=getattr(base_config, "num_consumer_groups", None), + reg_dec_producer=getattr(base_config, "reg_dec_producer", None), + reg_inc_consumer=getattr(base_config, "reg_inc_consumer", None), + ) + + +def get_stage1_probe_profiles(profile_family: str = "mixcv") -> List[Dict[str, Any]]: + if profile_family == "vector": + profiles = DEFAULT_VECTOR_STAGE1_PROBE_PROFILES + elif profile_family == "conservative_mixcv": + profiles = CONSERVATIVE_MIXCV_STAGE1_PROBE_PROFILES + else: + profiles = DEFAULT_STAGE1_PROBE_PROFILES + return [_drop_unsupported_profile_options(profile) for profile in profiles] + + +def get_stage1_probe_configs( + shape_kwargs: Optional[Mapping[str, Any]] = None, + *, + base_config: Optional[Config] = None, + profile_family: str = "mixcv", +) -> List[Config]: + return [ + compile_profile_to_config( + profile, shape_kwargs=shape_kwargs, base_config=base_config + ) + for profile in get_stage1_probe_profiles(profile_family) + ] + + +def _shape_size(shape_kwargs: Optional[Mapping[str, Any]]) -> int: + size = 1 + if not shape_kwargs: + return size + for value in shape_kwargs.values(): + if isinstance(value, bool) or not isinstance(value, int): + continue + size *= max(1, value) + return size + + +def _is_large_shape( + shape_kwargs: Optional[Mapping[str, Any]], large_tile: Optional[bool] +) -> bool: + if large_tile is not None: + return large_tile + return _shape_size(shape_kwargs) >= 128 * 128 + + +def _workspace_cv_profile( + *, + set_workspace_multibuffer: int, + tile_mix_cube_loop: int, + tile_mix_vector_loop: int, + local_buffer_strategy: str = "no-l0c", + enable_ubuf_saving: bool = True, + unit_flag: bool = False, +) -> Dict[str, Any]: + profile = { + COMPILE_MODE_KEY: COMPILE_MODE_WORKSPACE_CV, + "num_stages": 2, + "multibuffer": True, + "enable_tuning_mode": True, + "enable_auto_bind_sub_block": True, + "enable_hivm_auto_cv_balance": True, + "enable_ubuf_saving": enable_ubuf_saving, + "limit_auto_multi_buffer_only_for_local_buffer": False, + "limit_auto_multi_buffer_of_local_buffer": local_buffer_strategy, + "set_workspace_multibuffer": set_workspace_multibuffer, + "tile_mix_cube_loop": tile_mix_cube_loop, + "tile_mix_vector_loop": tile_mix_vector_loop, + "unit_flag": unit_flag, + } + return _drop_unsupported_profile_options(profile) + + +def _local_mb_profile( + *, + local_buffer_strategy: str, + enable_ubuf_saving: bool = True, + unit_flag: bool = False, +) -> Dict[str, Any]: + return { + COMPILE_MODE_KEY: COMPILE_MODE_LOCAL_MB, + "num_stages": 2, + "multibuffer": True, + "enable_tuning_mode": True, + "enable_auto_bind_sub_block": True, + "enable_hivm_auto_cv_balance": True, + "enable_ubuf_saving": enable_ubuf_saving, + "limit_auto_multi_buffer_only_for_local_buffer": True, + "limit_auto_multi_buffer_of_local_buffer": local_buffer_strategy, + "unit_flag": unit_flag, + } + + +def _mb_off_profile( + *, enable_ubuf_saving: bool = True, unit_flag: bool = False +) -> Dict[str, Any]: + return { + COMPILE_MODE_KEY: COMPILE_MODE_MB_OFF, + "num_stages": 1, + "enable_tuning_mode": True, + "enable_auto_bind_sub_block": True, + "enable_hivm_auto_cv_balance": True, + "enable_ubuf_saving": enable_ubuf_saving, + "unit_flag": unit_flag, + } + + +def _append_valid_unique_profile( + profiles: List[Dict[str, Any]], + profile: Mapping[str, Any], + seen: set, + *, + allow_unit_flag: bool, +) -> None: + candidate = _drop_unsupported_profile_options(profile) + if not allow_unit_flag and candidate.get("unit_flag") is True: + return + if not validate_compile_profile(candidate): + return + key = effective_compile_profile_key(candidate) + if key in seen: + return + seen.add(key) + profiles.append(candidate) + + +def _workspace_profile_from_probe( + probe: Mapping[str, Any], + *, + set_workspace_multibuffer: int, + tile_mix_cube_loop: int, + tile_mix_vector_loop: int, + local_buffer_strategy: Optional[str] = None, +) -> Dict[str, Any]: + if _profile_mode(probe) == COMPILE_MODE_WORKSPACE_CV: + return _replace_profile( + probe, + set_workspace_multibuffer=set_workspace_multibuffer, + tile_mix_cube_loop=tile_mix_cube_loop, + tile_mix_vector_loop=tile_mix_vector_loop, + limit_auto_multi_buffer_of_local_buffer=( + local_buffer_strategy + if local_buffer_strategy is not None + else probe.get("limit_auto_multi_buffer_of_local_buffer", "no-l0c") + ), + ) + return _workspace_cv_profile( + set_workspace_multibuffer=set_workspace_multibuffer, + tile_mix_cube_loop=tile_mix_cube_loop, + tile_mix_vector_loop=tile_mix_vector_loop, + local_buffer_strategy=local_buffer_strategy or "no-l0c", + enable_ubuf_saving=probe.get("enable_ubuf_saving", True), + unit_flag=probe.get("unit_flag", False), + ) + + +def _tile_mix_probe_order( + base_pair: tuple[int, int], *, stage1_ub_bytes: Optional[int] +) -> List[tuple[int, int]]: + if stage1_ub_bytes is None: + order = [(4, 4), (4, 2), (2, 2), (2, 4), (1, 1)] + elif stage1_ub_bytes < LOCAL_ONLY_PROBE_UB_THRESHOLD_BYTES: + order = [base_pair, (4, 2), (2, 2), (2, 4), (1, 1), (4, 4)] + elif stage1_ub_bytes >= DEFAULT_UB_LIMIT_BYTES - TIGHT_UB_MARGIN_BYTES: + order = [base_pair, (4, 4), (4, 2), (2, 4), (2, 2)] + else: + order = [base_pair, (4, 4), (4, 2), (2, 2), (2, 4), (1, 1)] + + result: List[tuple[int, int]] = [] + for pair in order: + if pair not in result: + result.append(pair) + return result + + +def _workspace_probe_order( + base_workspace: int, *, stage1_ub_bytes: Optional[int] +) -> List[int]: + if stage1_ub_bytes is None: + order = [base_workspace, 4, 2] + elif stage1_ub_bytes < LOCAL_ONLY_PROBE_UB_THRESHOLD_BYTES: + order = [4, base_workspace, 2] + elif stage1_ub_bytes >= DEFAULT_UB_LIMIT_BYTES - TIGHT_UB_MARGIN_BYTES: + order = [base_workspace, 2, 4] + else: + order = [base_workspace, 4, 2] + + result: List[int] = [] + for value in order: + if value in (2, 4) and value not in result: + result.append(value) + return result + + +def _stage2_workspace_tile_profiles_from_probe( + stage1_profile: Mapping[str, Any], + *, + stage1_ub_bytes: Optional[int], +) -> List[Dict[str, Any]]: + mode = _profile_mode(stage1_profile) + if mode == COMPILE_MODE_WORKSPACE_CV: + base_workspace = stage1_profile.get("set_workspace_multibuffer", 2) + base_pair = ( + stage1_profile.get("tile_mix_cube_loop", 4), + stage1_profile.get("tile_mix_vector_loop", 4), + ) + base_local_strategy = stage1_profile.get( + "limit_auto_multi_buffer_of_local_buffer", "no-l0c" + ) + else: + base_workspace = 2 + base_pair = (4, 4) + base_local_strategy = stage1_profile.get( + "limit_auto_multi_buffer_of_local_buffer", "no-l0c" + ) + + workspaces = _workspace_probe_order(base_workspace, stage1_ub_bytes=stage1_ub_bytes) + tile_pairs = _tile_mix_probe_order(base_pair, stage1_ub_bytes=stage1_ub_bytes) + + profiles: List[Dict[str, Any]] = [] + for workspace in workspaces: + for cube, vector in tile_pairs: + profiles.append( + _workspace_profile_from_probe( + stage1_profile, + set_workspace_multibuffer=workspace, + tile_mix_cube_loop=cube, + tile_mix_vector_loop=vector, + local_buffer_strategy=base_local_strategy, + ) + ) + + # After workspace/tile_mix ranking, try the no-limit variant of the same + # probe-derived tile order. If the Stage 1 probe was already no-limit this + # only contributes non-duplicates. + for workspace in workspaces: + for cube, vector in tile_pairs: + profiles.append( + _workspace_profile_from_probe( + stage1_profile, + set_workspace_multibuffer=workspace, + tile_mix_cube_loop=cube, + tile_mix_vector_loop=vector, + local_buffer_strategy="no-limit", + ) + ) + return profiles + + +def _make_vector_stage2_seed_profiles( + stage1_profile: Mapping[str, Any], + *, + seed_budget: int, +) -> List[Dict[str, Any]]: + seeds: List[Dict[str, Any]] = [] + seen = set() + base = { + COMPILE_MODE_KEY: COMPILE_MODE_VECTOR, + "num_stages": stage1_profile.get("num_stages", 2), + "enable_ubuf_saving": stage1_profile.get("enable_ubuf_saving", True), + } + candidates = [ + base, + {**base, "num_stages": 1 if base["num_stages"] == 2 else 2}, + {**base, "enable_ubuf_saving": not bool(base["enable_ubuf_saving"])}, + { + **base, + "num_stages": 1 if base["num_stages"] == 2 else 2, + "enable_ubuf_saving": not bool(base["enable_ubuf_saving"]), + }, + ] + for profile in candidates: + _append_valid_unique_profile(seeds, profile, seen, allow_unit_flag=False) + if len(seeds) >= seed_budget: + break + return seeds + + +def make_stage2_seed_profiles( + stage1_profile: Mapping[str, Any], + *, + shape_kwargs: Optional[Mapping[str, Any]] = None, + large_tile: Optional[bool] = None, + seed_budget: int = 8, + allow_unit_flag: bool = False, + stage1_ub_bytes: Optional[int] = None, +) -> List[Dict[str, Any]]: + """Build Stage 2 seeds from the Stage 1 winning probe. + + The first seed is always the exact Stage 1 probe. Subsequent workspace CV + seeds change workspace/tile_mix around that probe, ordered by Stage 1 UB + margin. Fixed branch seeds are fallback only. + """ + if _profile_mode(stage1_profile) == COMPILE_MODE_VECTOR: + return _make_vector_stage2_seed_profiles( + stage1_profile, seed_budget=seed_budget + ) + + seeds: List[Dict[str, Any]] = [] + seen = set() + _append_valid_unique_profile( + seeds, stage1_profile, seen, allow_unit_flag=allow_unit_flag + ) + if allow_unit_flag: + _append_valid_unique_profile( + seeds, + _replace_profile(stage1_profile, unit_flag=True), + seen, + allow_unit_flag=allow_unit_flag, + ) + + workspace_tile_profiles = _stage2_workspace_tile_profiles_from_probe( + stage1_profile, stage1_ub_bytes=stage1_ub_bytes + ) + if ( + stage1_ub_bytes is not None + and stage1_ub_bytes < LOCAL_ONLY_PROBE_UB_THRESHOLD_BYTES + ): + early_workspace_profile_count = min(4, max(1, seed_budget - len(seeds) - 2)) + else: + early_workspace_profile_count = len(workspace_tile_profiles) + + for profile in workspace_tile_profiles[:early_workspace_profile_count]: + _append_valid_unique_profile( + seeds, profile, seen, allow_unit_flag=allow_unit_flag + ) + + local_strategies = [] + if ( + stage1_ub_bytes is not None + and stage1_ub_bytes < LOCAL_ONLY_PROBE_UB_THRESHOLD_BYTES + ): + local_strategies.append("no-l0c") + local_strategies.append("no-limit") + for local_strategy in local_strategies: + _append_valid_unique_profile( + seeds, + _local_mb_profile( + local_buffer_strategy=local_strategy, + enable_ubuf_saving=stage1_profile.get("enable_ubuf_saving", True), + ), + seen, + allow_unit_flag=allow_unit_flag, + ) + + for profile in workspace_tile_profiles[early_workspace_profile_count:]: + _append_valid_unique_profile( + seeds, profile, seen, allow_unit_flag=allow_unit_flag + ) + + _append_valid_unique_profile( + seeds, _mb_off_profile(), seen, allow_unit_flag=allow_unit_flag + ) + return seeds[:seed_budget] + + +def _replace_profile(profile: Mapping[str, Any], **updates: Any) -> Dict[str, Any]: + copied = dict(profile) + copied.update(updates) + return copied + + +def _next_tile_mix_up(value: int) -> Optional[int]: + order = [1, 2, 4] + if value not in order: + return None + index = order.index(value) + return order[index + 1] if index + 1 < len(order) else None + + +def _next_tile_mix_down(value: int) -> Optional[int]: + order = [1, 2, 4] + if value not in order: + return None + index = order.index(value) + return order[index - 1] if index > 0 else None + + +def _tile_pair_balance_candidates( + profile: Mapping[str, Any], *, large_shape: bool +) -> List[Dict[str, Any]]: + if _profile_mode(profile) != COMPILE_MODE_WORKSPACE_CV: + return [] + if "tile_mix_cube_loop" not in profile or "tile_mix_vector_loop" not in profile: + return [] + + order = ( + [(4, 4), (4, 2), (2, 4), (2, 2), (1, 1)] + if large_shape + else [(2, 2), (4, 2), (2, 4), (4, 4), (1, 1)] + ) + current = (profile["tile_mix_cube_loop"], profile["tile_mix_vector_loop"]) + if current not in order: + return [ + _replace_profile( + profile, + tile_mix_cube_loop=order[0][0], + tile_mix_vector_loop=order[0][1], + ) + ] + index = order.index(current) + candidates = [] + for next_index in (index + 1, index - 1): + if 0 <= next_index < len(order): + cube, vector = order[next_index] + candidates.append( + _replace_profile( + profile, + tile_mix_cube_loop=cube, + tile_mix_vector_loop=vector, + ) + ) + return candidates + + +def generate_linked_compile_neighbors( + profile: Mapping[str, Any], + *, + shape_kwargs: Optional[Mapping[str, Any]] = None, + large_tile: Optional[bool] = None, + limit: Optional[int] = None, + allow_unit_flag: bool = False, + prefer_resource_relax: bool = False, +) -> List[Dict[str, Any]]: + """Generate legal action-based neighbors for compile-option annealing.""" + validate_compile_profile(profile, raise_on_error=True) + large_shape = _is_large_shape(shape_kwargs, large_tile) + mode = _profile_mode(profile) + candidates: List[Dict[str, Any]] = [] + if allow_unit_flag and mode != COMPILE_MODE_VECTOR: + candidates.append( + _replace_profile( + profile, + unit_flag=not bool(profile.get("unit_flag", False)), + ) + ) + + if mode == COMPILE_MODE_VECTOR: + candidates.append( + _replace_profile( + profile, + num_stages=1 if profile.get("num_stages") == 2 else 2, + ) + ) + candidates.append( + _replace_profile( + profile, + enable_ubuf_saving=not bool(profile.get("enable_ubuf_saving", True)), + ) + ) + + elif mode == COMPILE_MODE_WORKSPACE_CV: + resource_relax = [] + if profile.get("set_workspace_multibuffer") == 4: + resource_relax.append( + _replace_profile(profile, set_workspace_multibuffer=2) + ) + for name in ("tile_mix_vector_loop", "tile_mix_cube_loop"): + if name in profile: + next_value = _next_tile_mix_up(profile[name]) + if next_value is not None: + resource_relax.append( + _replace_profile(profile, **{name: next_value}) + ) + if profile.get("limit_auto_multi_buffer_of_local_buffer") == "no-limit": + resource_relax.append( + _replace_profile( + profile, limit_auto_multi_buffer_of_local_buffer="no-l0c" + ) + ) + if profile.get("enable_ubuf_saving") is False: + resource_relax.append(_replace_profile(profile, enable_ubuf_saving=True)) + resource_relax.append( + _local_mb_profile( + local_buffer_strategy=profile.get( + "limit_auto_multi_buffer_of_local_buffer", "no-l0c" + ), + enable_ubuf_saving=profile.get("enable_ubuf_saving", True), + unit_flag=profile.get("unit_flag", False), + ) + ) + + perf_push = [] + if profile.get("set_workspace_multibuffer") == 2: + perf_push.append(_replace_profile(profile, set_workspace_multibuffer=4)) + if profile.get("limit_auto_multi_buffer_of_local_buffer") == "no-l0c": + perf_push.append( + _replace_profile( + profile, limit_auto_multi_buffer_of_local_buffer="no-limit" + ) + ) + if profile.get("enable_ubuf_saving") is True: + perf_push.append(_replace_profile(profile, enable_ubuf_saving=False)) + for name in ("tile_mix_vector_loop", "tile_mix_cube_loop"): + if name in profile: + next_value = _next_tile_mix_down(profile[name]) + if next_value is not None: + perf_push.append(_replace_profile(profile, **{name: next_value})) + + balance = _tile_pair_balance_candidates(profile, large_shape=large_shape) + candidates.extend(resource_relax if prefer_resource_relax else perf_push) + candidates.extend(balance) + candidates.extend(perf_push if prefer_resource_relax else resource_relax) + + elif mode == COMPILE_MODE_LOCAL_MB: + candidates.append( + _workspace_cv_profile( + set_workspace_multibuffer=2, + tile_mix_cube_loop=4 if large_shape else 2, + tile_mix_vector_loop=4 if large_shape else 2, + local_buffer_strategy=profile.get( + "limit_auto_multi_buffer_of_local_buffer", "no-l0c" + ), + enable_ubuf_saving=profile.get("enable_ubuf_saving", True), + unit_flag=profile.get("unit_flag", False), + ) + ) + candidates.append( + _mb_off_profile( + enable_ubuf_saving=profile.get("enable_ubuf_saving", True), + unit_flag=profile.get("unit_flag", False), + ) + ) + if profile.get("limit_auto_multi_buffer_of_local_buffer") == "no-l0c": + candidates.append( + _replace_profile( + profile, limit_auto_multi_buffer_of_local_buffer="no-limit" + ) + ) + else: + candidates.append( + _replace_profile( + profile, limit_auto_multi_buffer_of_local_buffer="no-l0c" + ) + ) + + else: + candidates.append( + _local_mb_profile( + local_buffer_strategy="no-l0c", + enable_ubuf_saving=profile.get("enable_ubuf_saving", True), + unit_flag=profile.get("unit_flag", False), + ) + ) + + selected: List[Dict[str, Any]] = [] + seen = set() + for candidate in candidates: + _append_valid_unique_profile( + selected, candidate, seen, allow_unit_flag=allow_unit_flag + ) + if limit is not None and len(selected) >= limit: + break + return selected + + +_SPACE_OVERFLOW_RE = re.compile(r"\b([a-z][a-z0-9_]*)\s+overflow\b") +_SYNC_FAILURE_TOKENS = ( + "injectsync", + "block-sync", + "syncblock", + "graphsyncsolver", + "syncsolver", + "barrier", + "event", + "set flag", + "wait flag", + "memory conflict", +) +_INTERNAL_FAILURE_TOKENS = ( + "internal error", + "report_fatal_error", + "llvm_unreachable", + "assertion failed", + "segmentation fault", + "core dumped", + "unhandled case", + "unexpected op", + "error in cv-pipelining", + "postprocesscubefunc failed", + "postprocessvectorfunc failed", +) +_WORKSPACE_FAILURE_TOKENS = ( + "workspace", + "alloc_workspace", + "allocworkspace", + "failed to multibuffer", + "multibuffer", + "unknown buffer size", + "alloc-like op", + "cv-pipelining", + "unable to pipeline", + "cannot pipeline", + "failed to pipelinine", +) + + +def _failure_text(error: Any) -> str: + if error is None: + return "" + if isinstance(error, (list, tuple)): + text = " ".join(_failure_text(item) for item in error) + elif isinstance(error, BaseException): + text = f"{type(error).__name__} {error}" + else: + text = str(error) + return re.sub(r"\s+", " ", text.lower()).strip() + + +def classify_compile_failure(error: Any) -> str: + text = _failure_text(error) + match = _SPACE_OVERFLOW_RE.search(text) + if match: + space = match.group(1) + if space in {"ub", "ubuf"}: + return FAILURE_RESOURCE_UB + if space in {"cbuf", "l1", "l1a", "l1b"}: + return FAILURE_RESOURCE_L1 + if space in {"l0c", "cc"}: + return FAILURE_RESOURCE_L0C + if space in {"workspace", "gm"}: + return FAILURE_RESOURCE_WORKSPACE + return FAILURE_EXACT_ONLY + if any(token in text for token in _SYNC_FAILURE_TOKENS): + return FAILURE_SYNC_OR_CORRECTNESS + if any(token in text for token in _INTERNAL_FAILURE_TOKENS): + return FAILURE_COMPILER_INTERNAL + if any(token in text for token in _WORKSPACE_FAILURE_TOKENS): + return FAILURE_RESOURCE_WORKSPACE + return FAILURE_EXACT_ONLY + + +def is_resource_failure(failure_kind: str) -> bool: + return failure_kind in _RESOURCE_FAILURES + + +def _is_workspace_safest_tile_point(profile: Mapping[str, Any]) -> bool: + return ( + _profile_mode(profile) == COMPILE_MODE_WORKSPACE_CV + and profile.get("tile_mix_vector_loop") == 4 + and profile.get("tile_mix_cube_loop") == 4 + and profile.get("enable_ubuf_saving", True) is True + ) + + +def should_prune_compile_direction( + profile: Mapping[str, Any], failure_kind: str +) -> bool: + """Whether a failed profile can prune higher-resource neighbors. + + Resource overflow failures are always directional. For workspace CV, a + failure at tile_mix=(4,4) with ubuf_saving=True is also treated as a + directional failure because it is the lowest-UB point for that workspace / + local-buffer strategy; smaller tile_mix values should not be tested first. + """ + if is_resource_failure(failure_kind): + return True + return _is_workspace_safest_tile_point(profile) + + +def _pressure_rank_local_buffer(value: Any) -> int: + return {"no-l0c": 0, "no-limit": 1}.get(value, 0) + + +def _pressure_rank_ubuf_saving(value: Any) -> int: + return 0 if value is True else 1 + + +def compile_profile_resource_not_less( + candidate: Mapping[str, Any], failed_profile: Mapping[str, Any] +) -> bool: + if _profile_mode(candidate) != _profile_mode(failed_profile): + return False + if _profile_mode(candidate) != COMPILE_MODE_WORKSPACE_CV: + return False + + for name in ("tile_mix_vector_loop", "tile_mix_cube_loop"): + if name in candidate and name in failed_profile: + if candidate[name] > failed_profile[name]: + return False + if ( + "set_workspace_multibuffer" in candidate + and "set_workspace_multibuffer" in failed_profile + and candidate["set_workspace_multibuffer"] + < failed_profile["set_workspace_multibuffer"] + ): + return False + if _pressure_rank_local_buffer( + candidate.get("limit_auto_multi_buffer_of_local_buffer", "no-l0c") + ) < _pressure_rank_local_buffer( + failed_profile.get("limit_auto_multi_buffer_of_local_buffer", "no-l0c") + ): + return False + if _pressure_rank_ubuf_saving( + candidate.get("enable_ubuf_saving", True) + ) < _pressure_rank_ubuf_saving(failed_profile.get("enable_ubuf_saving", True)): + return False + return True + + +@dataclass +class CompileFailureRegionSet: + exact_keys: set = field(default_factory=set) + resource_failures: List[Dict[str, Any]] = field(default_factory=list) + forbidden_modes: set = field(default_factory=set) + + def add(self, profile: Mapping[str, Any], failure: Any) -> str: + failure_kind = ( + failure + if isinstance(failure, str) and failure in _ALL_FAILURE_KINDS + else classify_compile_failure(failure) + ) + self.exact_keys.add(effective_compile_profile_key(profile)) + mode = _profile_mode(profile) + if should_prune_compile_direction(profile, failure_kind): + if mode == COMPILE_MODE_WORKSPACE_CV: + self.resource_failures.append(dict(profile)) + elif mode == COMPILE_MODE_LOCAL_MB: + self.forbidden_modes.add(COMPILE_MODE_LOCAL_MB) + elif mode == COMPILE_MODE_LOCAL_MB: + self.forbidden_modes.add(COMPILE_MODE_LOCAL_MB) + return failure_kind + + def is_forbidden(self, profile: Mapping[str, Any]) -> bool: + key = effective_compile_profile_key(profile) + if key in self.exact_keys: + return True + if _profile_mode(profile) in self.forbidden_modes: + return True + return any( + compile_profile_resource_not_less(profile, failed) + for failed in self.resource_failures + ) + + +def _value_space_for_config( + config: Config, spec: CompileOptionsSpec, *, generated_tiling: bool +) -> Dict[str, List[Any]]: + supported = _SUPPORTED_PARAMS[spec.kernel_type] + preset = _AUTO_SEARCH_PRESETS[spec.kernel_type] + + value_space = {} + for name in sorted(supported): + if name in spec.params: + values = spec.params[name] + elif name not in preset: + continue + else: + values = preset[name] + validate_compile_option_values(name, values) + value_space[name] = values + if spec.kernel_type == "mixcv" and is_compile_on_910_95: + for name in _MIXCV_910_95_UNSUPPORTED_PARAMS: + value_space.pop(name, None) + return value_space + + +def _is_inactive_reason( + name: str, + num_stages: int, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], +) -> Optional[str]: + if not _is_param_effective(name, num_stages, combo, config, fixed_options): + if name in { + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + }: + return "depends on auto multi-buffer" + return None + + +def _is_param_effective( + name: str, + num_stages: int, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], +) -> bool: + if name not in _MIXCV_OPTION_ACTIVITY_RULES: + return True + return _MIXCV_OPTION_ACTIVITY_RULES[name](num_stages, combo, config, fixed_options) + + +def _resolve_compile_option( + name: str, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], + default: Any, +) -> Any: + if name in combo: + return combo[name] + if name in fixed_options: + return fixed_options[name] + return default + + +def _effective_values( + name: str, + value_space: Dict[str, List[Any]], + fixed_options: Dict[str, Any], + default: Any, +) -> List[Any]: + if name in fixed_options: + return [fixed_options[name]] + if name in value_space: + return value_space[name] + return [default] + + +def _emit_values( + names: List[str], + value_space: Dict[str, List[Any]], + fixed_options: Dict[str, Any], +) -> List[tuple[str, List[Any]]]: + return [ + (name, value_space[name]) + for name in names + if name in value_space and name not in fixed_options + ] + + +def _product_dict(items: List[tuple[str, List[Any]]]): + if not items: + yield {} + return + names = [name for name, _ in items] + values = [values for _, values in items] + for combo in itertools.product(*values): + yield dict(zip(names, combo)) + + +def _mixcv_branch_items( + *, + num_stages: int, + value_space: Dict[str, List[Any]], + fixed_options: Dict[str, Any], +) -> List[List[tuple[str, List[Any]]]]: + independent_names = [ + "enable_tuning_mode", + "unit_flag", + "enable_ubuf_saving", + "enable_hivm_auto_cv_balance", + "enable_auto_bind_sub_block", + ] + independent_items = _emit_values(independent_names, value_space, fixed_options) + + multibuffer_values = _effective_values( + "multibuffer", value_space, fixed_options, None + ) + branches = [] + for multibuffer in multibuffer_values: + multibuffer_items = [] + if "multibuffer" in value_space and "multibuffer" not in fixed_options: + multibuffer_items = [("multibuffer", [multibuffer])] + branch_base = independent_items + multibuffer_items + + if num_stages == 1 or multibuffer is False: + branches.append(branch_base) + continue + + limit_only_values = _effective_values( + "limit_auto_multi_buffer_only_for_local_buffer", + value_space, + fixed_options, + True, + ) + for limit_only in limit_only_values: + limit_only_items = [] + if ( + "limit_auto_multi_buffer_only_for_local_buffer" in value_space + and "limit_auto_multi_buffer_only_for_local_buffer" not in fixed_options + ): + limit_only_items = [ + ("limit_auto_multi_buffer_only_for_local_buffer", [limit_only]) + ] + if limit_only is True: + branch_names = ["limit_auto_multi_buffer_of_local_buffer"] + else: + branch_names = [ + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + ] + branches.append( + branch_base + + limit_only_items + + _emit_values(branch_names, value_space, fixed_options) + ) + + return branches + + +def _make_expanded_config( + config: Config, + spec: CompileOptionsSpec, + combo_value: Dict[str, Any], + num_stages: int, +) -> Config: + new_kwargs = dict(config.kwargs) + for name in _ALL_PARAMS: + new_kwargs.pop(name, None) + for name, value in combo_value.items(): + new_kwargs[name] = value + + return _make_config_compat( + kwargs=new_kwargs, + num_warps=getattr(config, "num_warps", 4), + num_stages=num_stages, + num_ctas=getattr(config, "num_ctas", 1), + maxnreg=getattr(config, "maxnreg", None), + pre_hook=getattr(config, "pre_hook", None), + ir_override=getattr(config, "ir_override", None), + num_buffers_warp_spec=getattr(config, "num_buffers_warp_spec", None), + num_consumer_groups=getattr(config, "num_consumer_groups", None), + reg_dec_producer=getattr(config, "reg_dec_producer", None), + reg_inc_consumer=getattr(config, "reg_inc_consumer", None), + ) + + +def _hashable_value(value: Any): + if isinstance(value, dict): + return tuple(sorted((key, _hashable_value(val)) for key, val in value.items())) + if isinstance(value, (list, tuple)): + return tuple(_hashable_value(item) for item in value) + if isinstance(value, set): + return tuple(sorted(_hashable_value(item) for item in value)) + try: + hash(value) + except TypeError: + return repr(value) + return value + + +def _config_key(config: Config) -> tuple: + return ( + tuple( + sorted( + (key, _hashable_value(value)) for key, value in config.kwargs.items() + ) + ), + getattr(config, "num_warps", 4), + getattr(config, "num_stages", None), + getattr(config, "num_ctas", 1), + getattr(config, "maxnreg", None), + id(getattr(config, "pre_hook", None)), + _hashable_value(getattr(config, "ir_override", None)), + getattr(config, "num_buffers_warp_spec", None), + getattr(config, "num_consumer_groups", None), + getattr(config, "reg_dec_producer", None), + getattr(config, "reg_inc_consumer", None), + ) + + +def expand_compile_option_configs( + configs: List[Config], + spec: CompileOptionsSpec, + *, + generated_tiling: bool, + fixed_options: Optional[Dict[str, Any]] = None, +) -> List[Config]: + if not spec.enabled or not configs: + return configs + + fixed_options = fixed_options or {} + expanded_configs = [] + emitted_config_keys = set() + for config in configs: + value_space = _value_space_for_config( + config, spec, generated_tiling=generated_tiling + ) + if "num_stages" in fixed_options: + num_stage_values = [fixed_options["num_stages"]] + value_space.pop("num_stages", None) + else: + num_stage_values = value_space.pop("num_stages") + + for name in fixed_options: + if name != "num_stages": + value_space.pop(name, None) + + for num_stages in num_stage_values: + if spec.kernel_type == "mixcv": + branch_items = _mixcv_branch_items( + num_stages=num_stages, + value_space=value_space, + fixed_options=fixed_options, + ) + combo_iter = itertools.chain.from_iterable( + _product_dict(items) for items in branch_items + ) + else: + combo_iter = _product_dict(list(value_space.items())) + + for combo_value in combo_iter: + new_config = _make_expanded_config( + config, spec, combo_value, num_stages + ) + config_key = _config_key(new_config) + if config_key in emitted_config_keys: + continue + emitted_config_keys.add(config_key) + + if ( + spec.max_configs is not None + and len(expanded_configs) >= spec.max_configs + ): + raise ValueError( + "compile_options generated more than " + f"{spec.max_configs} configs. Narrow the search space or raise max_configs." + ) + expanded_configs.append(new_config) + + return expanded_configs + + +def get_compile_option_param_names(spec: CompileOptionsSpec) -> set[str]: + if not spec.enabled: + return set() + return set(_SUPPORTED_PARAMS[spec.kernel_type]) + + +def format_compile_option_result( + config: Config, + spec: CompileOptionsSpec, + fixed_options: Optional[Dict[str, Any]] = None, +) -> str: + if not spec.enabled: + return str(config) + + fixed_options = fixed_options or {} + compile_param_names = _SUPPORTED_PARAMS[spec.kernel_type] - {"num_stages"} + selected_meta = { + key: value + for key, value in sorted(config.kwargs.items()) + if key not in compile_param_names + } + selected_meta["num_stages"] = getattr(config, "num_stages", None) + + effective = { + key: value + for key, value in sorted(config.kwargs.items()) + if key in compile_param_names + } + effective.update( + { + key: value + for key, value in sorted(fixed_options.items()) + if key in compile_param_names + } + ) + + if spec.kernel_type == "mixcv": + num_stages = selected_meta["num_stages"] + multibuffer = effective.get("multibuffer", None) + effective["enable_auto_multi_buffer"] = ( + False if multibuffer is False or num_stages == 1 else True + ) + for name in _SUPPORTED_PARAMS[spec.kernel_type]: + if name == "num_stages": + continue + reason = _is_inactive_reason( + name, num_stages, config.kwargs, config, fixed_options + ) + if reason is not None and name not in effective: + effective[name] = f"" + if reason is not None and name in effective: + effective[name] = f"" + if is_compile_on_910_95: + for name in _MIXCV_910_95_UNSUPPORTED_PARAMS: + effective[name] = "" + + selected_items = [f"{key}={value}" for key, value in selected_meta.items()] + effective_items = [f"{key}={value}" for key, value in sorted(effective.items())] + return ( + "selected_meta: " + + ", ".join(selected_items) + + "; effective_compile_options: " + + ", ".join(effective_items) + ) + + +def summarize_compile_option_configs( + configs: List[Config], limit: Optional[int] = None +) -> List[str]: + summary = [] + selected_configs = configs if limit is None else configs[:limit] + for config in selected_configs: + items = [f"{key}={value}" for key, value in sorted(config.kwargs.items())] + items.append(f"num_stages={getattr(config, 'num_stages', None)}") + summary.append(", ".join(items)) + return summary diff --git a/backend/ascend_autotune_runtime/tile_acquisition_model.py b/backend/ascend_autotune_runtime/tile_acquisition_model.py new file mode 100644 index 00000000..e16f00a3 --- /dev/null +++ b/backend/ascend_autotune_runtime/tile_acquisition_model.py @@ -0,0 +1,484 @@ +from __future__ import annotations + +import math +import re +from dataclasses import dataclass +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple + +from triton.runtime.autotuner import Config + +from .tile_shape_space import DiscreteShapeSpace + + +DEFAULT_UB_LIMIT_BYTES = 192 * 1024 +ACQUISITION_FAILURE_WEIGHT = 0.25 +ACQUISITION_EXPLORATION_WEIGHT = 0.03 +ACQUISITION_UB_WEIGHT = 0.12 +ACQUISITION_UB_TIE_WEIGHT = 0.01 +ACQUISITION_TRUST_WEIGHT = 0.05 +ACQUISITION_BOUNDARY_WEIGHT = 0.25 +ACQUISITION_BATCH_DIVERSITY_WEIGHT = 0.03 +UB_BARRIER_SAFE_RATIO = 0.85 +_RESOURCE_OVERFLOW_RE = re.compile( + r"requires\s+(\d+)\s+bits\s+while\s+(\d+)\s+bits\s+available", + re.IGNORECASE, +) + + +ShapePoint = Tuple[int, ...] + + +def resource_penalty(rho: Optional[float]) -> float: + """Smooth UB overflow penalty. + + Legal UB usage is not penalized. Once the predicted or observed usage + exceeds the hardware limit, the first 30% overflow is soft and larger + overflow grows quadratically. + """ + if rho is None or rho <= 1.0: + return 0.0 + soft_margin = 0.3 + soft_cost = 0.25 + hard_cost = 2.0 + if rho <= 1.0 + soft_margin: + t = (rho - 1.0) / soft_margin + return soft_cost * (t * t * (3.0 - 2.0 * t)) + t = (rho - 1.0 - soft_margin) / soft_margin + return soft_cost + hard_cost * t * t + + +def resource_overflow_ratio(error: Any) -> Optional[float]: + """Parse compiler overflow severity as required / available UB ratio.""" + text = _error_text(error) + if not text: + return None + match = _RESOURCE_OVERFLOW_RE.search(text) + if not match: + return None + required = int(match.group(1)) + available = int(match.group(2)) + if available <= 0: + return None + return required / available + + +def timing_value(cost: Any) -> float: + if isinstance(cost, (list, tuple)) and cost: + return float(cost[0]) + return float(cost) + + +@dataclass +class ShapeAcquisition: + success_points: Sequence[Tuple[ShapePoint, float]] + failure_records: Sequence[Tuple[ShapePoint, Optional[float]]] + observed_points: Sequence[ShapePoint] + parent_points: Sequence[ShapePoint] + resource_model: Optional[Mapping[str, Any]] + latency_model: Optional[Mapping[str, Any]] + spec: object + + @classmethod + def from_history( + cls, + *, + space: DiscreteShapeSpace, + parents: Sequence[Mapping[str, Any]], + successes: Sequence[Mapping[str, Any]], + failures: Sequence[Config], + failure_observations: Sequence[Mapping[str, Any]], + observed: Sequence[Config], + spec: object, + ) -> "ShapeAcquisition": + success_points = [ + ( + space.point_from_shape(item["shape"]), + timing_value(item["time"]), + ) + for item in successes + ] + failure_records = [(space.point(config), None) for config in failures] + failure_records.extend( + (point, _failure_observation_resource_ratio(item)) + for point, item in ( + (_failure_observation_point(item, space), item) + for item in failure_observations + ) + if point is not None + ) + observed_points = [space.point(config) for config in observed] + parent_points = [ + space.point_from_shape(parent["shape"]) + for parent in parents + if isinstance(parent, Mapping) and "shape" in parent + ] + resource_model = _fit_resource_log_model( + _resource_observation_points(successes, failure_observations, space) + ) + latency_model = _fit_latency_quadratic_model(success_points, spec) + return cls( + success_points=success_points, + failure_records=failure_records, + observed_points=observed_points, + parent_points=parent_points, + resource_model=resource_model, + latency_model=latency_model, + spec=spec, + ) + + def cost(self, point: ShapePoint) -> float: + latency = _predict_latency_quadratic(point, self.latency_model, self.spec) + novelty = _min_distance(point, self.observed_points, self.spec) + predicted_rho = _predict_resource_rho(point, self.resource_model) + failure = _failure_potential(point, self.failure_records, self.spec) + ub = _resource_barrier(predicted_rho) + ub_tie = _resource_tie_breaker(predicted_rho) + trust = _trust_region_penalty( + point, + success_points=self.success_points, + observed_points=self.observed_points, + parent_points=self.parent_points, + spec=self.spec, + ) + boundary = _boundary_barrier(point, self.spec) + return ( + latency + + ACQUISITION_FAILURE_WEIGHT * failure + + ACQUISITION_UB_WEIGHT * ub + + ACQUISITION_UB_TIE_WEIGHT * ub_tie + + ACQUISITION_TRUST_WEIGHT * trust + + ACQUISITION_BOUNDARY_WEIGHT * boundary + - ACQUISITION_EXPLORATION_WEIGHT * novelty + ) + + def batch_adjusted_cost( + self, + point: ShapePoint, + selected_points: Sequence[ShapePoint], + ) -> float: + return self.cost(point) - ACQUISITION_BATCH_DIVERSITY_WEIGHT * _min_distance( + point, selected_points, self.spec + ) + + +def _error_text(error: Any) -> str: + if error is None: + return "" + if isinstance(error, (list, tuple)): + return " ".join(_error_text(item) for item in error) + if isinstance(error, BaseException): + return f"{type(error).__name__} {error}" + return str(error) + + +def _failure_observation_point( + observation: Mapping[str, Any], space: DiscreteShapeSpace +) -> Optional[ShapePoint]: + shape = observation.get("shape") + if shape is None and observation.get("config") is not None: + shape = space.shape(observation["config"]) + if not isinstance(shape, Mapping): + return None + try: + return space.point_from_shape(dict(shape)) + except (KeyError, ValueError): + return None + + +def _failure_observation_resource_ratio( + observation: Mapping[str, Any], +) -> Optional[float]: + ratio = observation.get("resource_ratio") + if ratio is None: + ratio = resource_overflow_ratio(observation.get("error")) + if not isinstance(ratio, (int, float)) or ratio <= 0: + return None + return float(ratio) + + +def _resource_observation_points( + successes: Sequence[Mapping[str, Any]], + failure_observations: Sequence[Mapping[str, Any]], + space: DiscreteShapeSpace, +) -> List[Tuple[ShapePoint, float]]: + points: List[Tuple[ShapePoint, float]] = [] + for item in successes: + ub = item.get("ub") + if not isinstance(ub, (int, float)) or ub <= 0: + continue + try: + point = space.point_from_shape(item["shape"]) + except (KeyError, ValueError): + continue + points.append((point, max(float(ub) / DEFAULT_UB_LIMIT_BYTES, 1e-9))) + + for item in failure_observations: + ratio = _failure_observation_resource_ratio(item) + if ratio is None: + continue + point = _failure_observation_point(item, space) + if point is None: + continue + points.append((point, ratio)) + return points + + +def _fit_resource_log_model( + observations: Sequence[Tuple[ShapePoint, float]], +) -> Optional[Dict[str, Any]]: + if not observations: + return None + + records = [ + (tuple(float(v) for v in point), math.log(max(ratio, 1e-9))) + for point, ratio in observations + ] + dims = len(records[0][0]) + center = tuple( + sum(point[dim] for point, _ in records) / len(records) for dim in range(dims) + ) + value_center = sum(value for _, value in records) / len(records) + + gradient = [0.0 for _ in range(dims)] + weight_sum = 0.0 + for i, (point_i, value_i) in enumerate(records): + for point_j, value_j in records[i + 1 :]: + delta = tuple(b - a for a, b in zip(point_i, point_j)) + norm2 = sum(value * value for value in delta) + if norm2 <= 0.0: + continue + scale = (value_j - value_i) / norm2 + weight = math.sqrt(norm2) + for dim, value in enumerate(delta): + gradient[dim] += weight * scale * value + weight_sum += weight + + if weight_sum > 0.0: + gradient = [value / weight_sum for value in gradient] + + intercept = value_center - sum(g * x for g, x in zip(gradient, center)) + return {"intercept": intercept, "gradient": tuple(gradient)} + + +def _predict_resource_rho( + point: ShapePoint, resource_model: Optional[Mapping[str, Any]] +) -> Optional[float]: + if resource_model is None: + return None + log_rho = resource_model["intercept"] + sum( + g * x for g, x in zip(resource_model["gradient"], point) + ) + log_rho = min(20.0, max(-20.0, log_rho)) + return math.exp(log_rho) + + +def _fit_latency_quadratic_model( + success_points: Sequence[Tuple[ShapePoint, float]], + spec: object, +) -> Optional[Dict[str, Any]]: + if not success_points: + return None + + best_point, best_cost = min(success_points, key=lambda item: item[1]) + best_cost = max(float(best_cost), 1e-12) + center = tuple(float(value) for value in best_point) + param_count = _quadratic_feature_count(len(best_point)) + ridge = 0.03 * param_count / max(1, len(success_points)) + + lhs = [[0.0 for _ in range(param_count)] for _ in range(param_count)] + rhs = [0.0 for _ in range(param_count)] + sigma = _kernel_sigma(spec) + for success_point, cost in success_points: + feature = _quadratic_features(success_point, center, spec) + y = math.log(max(float(cost), 1e-12) / best_cost) + distance = _index_distance(success_point, best_point, spec) + weight = math.exp(-(distance * distance) / (2.0 * sigma * sigma)) + weight = max(weight, 0.05) + for row in range(param_count): + rhs[row] += weight * feature[row] * y + for col in range(param_count): + lhs[row][col] += weight * feature[row] * feature[col] + + for index in range(param_count): + lhs[index][index] += ridge + theta = _solve_linear_system(lhs, rhs) + if theta is None: + return None + return {"theta": theta, "center": center} + + +def _predict_latency_quadratic( + point: ShapePoint, + model: Optional[Mapping[str, Any]], + spec: object, +) -> float: + if model is None: + return 0.0 + feature = _quadratic_features(point, model["center"], spec) + value = sum(coef * x for coef, x in zip(model["theta"], feature)) + return max(-1.0, min(2.0, value)) + + +def _failure_potential( + point: ShapePoint, + failure_records: Sequence[Tuple[ShapePoint, Optional[float]]], + spec: object, +) -> float: + if not failure_records: + return 0.0 + sigma = _kernel_sigma(spec) + potential = 0.0 + for failure_point, ratio in failure_records: + distance = _index_distance(point, failure_point, spec) + severity = 1.0 + if ratio is not None and ratio > 1.0: + severity += math.log(max(ratio, 1.0), 2.0) + potential += severity * math.exp(-(distance * distance) / (2.0 * sigma * sigma)) + return potential + + +def _resource_barrier(rho: Optional[float]) -> float: + if rho is None or rho <= UB_BARRIER_SAFE_RATIO: + return 0.0 + if rho < 1.0: + remaining = max((1.0 - rho) / (1.0 - UB_BARRIER_SAFE_RATIO), 1e-6) + return -math.log(remaining) + overflow = min(rho - 1.0, 16.0) + return -math.log(1e-6) + overflow * overflow + + +def _resource_tie_breaker(rho: Optional[float]) -> float: + if rho is None: + return 0.0 + capped = max(0.0, min(rho, UB_BARRIER_SAFE_RATIO)) + return (capped / UB_BARRIER_SAFE_RATIO) ** 2 + + +def _boundary_barrier(point: ShapePoint, spec: object) -> float: + """Softly discourage multi-dimensional edge/corner proposals.""" + if len(point) <= 1: + return 0.0 + denom = max(1, len(spec.values) - 1) + if denom <= 1: + return 0.0 + edge_values = [] + for index in point: + t = float(index) / denom + edge = max(0.0, (abs(2.0 * t - 1.0) - 0.70) / 0.30) + edge_values.append(edge * edge) + return sum(edge_values) / len(edge_values) + + +def _trust_region_penalty( + point: ShapePoint, + *, + success_points: Sequence[Tuple[ShapePoint, float]], + observed_points: Sequence[ShapePoint], + parent_points: Sequence[ShapePoint], + spec: object, +) -> float: + if success_points: + center = min(success_points, key=lambda item: item[1])[0] + spread = max( + _index_distance(center, success_point, spec) + for success_point, _ in success_points + ) + radius = max(0.25, min(1.0, 0.25 + spread)) + elif parent_points: + center = parent_points[0] + radius = 1.0 + elif observed_points: + center = _mean_point(observed_points) + radius = 1.0 + else: + return 0.0 + distance = _index_distance_float(point, center, spec) + return (distance / max(radius, 1e-6)) ** 4 + + +def _min_distance( + point: ShapePoint, + points: Sequence[ShapePoint], + spec: object, +) -> float: + if not points: + return 1.0 + return min(_index_distance(point, other, spec) for other in points) + + +def _index_distance(a: ShapePoint, b: ShapePoint, spec: object) -> float: + return _index_distance_float(a, b, spec) + + +def _index_distance_float( + a: Sequence[float], + b: Sequence[float], + spec: object, +) -> float: + denom = max(1, len(spec.values) - 1) + squared = 0.0 + for ai, bi in zip(a, b): + squared += ((ai - bi) / denom) ** 2 + return math.sqrt(squared) + + +def _kernel_sigma(spec: object) -> float: + return max(0.18, 1.5 / max(1, len(spec.values) - 1)) + + +def _quadratic_feature_count(dims: int) -> int: + return 1 + dims + dims + (dims * (dims - 1)) // 2 + + +def _quadratic_features( + point: Sequence[float], + center: Sequence[float], + spec: object, +) -> List[float]: + denom = max(1, len(spec.values) - 1) + z = [(float(value) - float(base)) / denom for value, base in zip(point, center)] + features = [1.0] + features.extend(z) + features.extend(value * value for value in z) + for left in range(len(z)): + for right in range(left + 1, len(z)): + features.append(z[left] * z[right]) + return features + + +def _solve_linear_system( + matrix: Sequence[Sequence[float]], + vector: Sequence[float], +) -> Optional[List[float]]: + size = len(vector) + if size == 0: + return [] + aug = [ + [float(matrix[row][col]) for col in range(size)] + [float(vector[row])] + for row in range(size) + ] + for col in range(size): + pivot = max(range(col, size), key=lambda row: abs(aug[row][col])) + if abs(aug[pivot][col]) <= 1e-12: + return None + if pivot != col: + aug[col], aug[pivot] = aug[pivot], aug[col] + pivot_value = aug[col][col] + for idx in range(col, size + 1): + aug[col][idx] /= pivot_value + for row in range(size): + if row == col: + continue + factor = aug[row][col] + if abs(factor) <= 1e-18: + continue + for idx in range(col, size + 1): + aug[row][idx] -= factor * aug[col][idx] + return [aug[row][size] for row in range(size)] + + +def _mean_point(points: Sequence[ShapePoint]) -> Tuple[float, ...]: + dims = len(points[0]) + return tuple( + sum(float(point[dim]) for point in points) / len(points) for dim in range(dims) + ) diff --git a/backend/ascend_autotune_runtime/tile_generator.py b/backend/ascend_autotune_runtime/tile_candidate_generator.py similarity index 100% rename from backend/ascend_autotune_runtime/tile_generator.py rename to backend/ascend_autotune_runtime/tile_candidate_generator.py diff --git a/backend/ascend_autotune_runtime/tile_search_policy.py b/backend/ascend_autotune_runtime/tile_search_policy.py new file mode 100644 index 00000000..6b9d1738 --- /dev/null +++ b/backend/ascend_autotune_runtime/tile_search_policy.py @@ -0,0 +1,657 @@ +from __future__ import annotations + +import itertools +import math +from dataclasses import dataclass, field +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple + +from triton.runtime.autotuner import Config + +from .tile_acquisition_model import ShapeAcquisition +from .tile_shape_space import ( + DiscreteShapeSpace, + effective_search_value_map, + extract_shape, + shape_key, +) + + +DEFAULT_SEARCH_VALUES = [16, 32, 64, 128, 256, 512, 1024, 2048] +DEFAULT_NO_DOT_SEARCH_VALUES = [ + 1, + 2, + 4, + 8, + 16, + 32, + 64, + 128, + 256, + 512, + 1024, + 2048, + 4096, + 8192, +] +DEFAULT_INITIAL_PERCENTILES = [0.20, 0.50, 0.80] +SMALL_VALUE_PERCENTILES = [0.30, 0.70] +SMALL_VALUE_COUNT_THRESHOLD = 6 +DEFAULT_MAX_STAGE1_CHILDREN_PER_ROUND = 16 +DEFAULT_STAGE1_TOTAL_BUDGET_CAP = 18 + +UB_FLOOR_RATIO = 0.5 +TIMING_CEILING_RATIO = 1.4 + + +@dataclass +class SearchParamsSpec: + enabled: bool = False + params: List[str] = field(default_factory=list) + values: List[int] = field(default_factory=lambda: list(DEFAULT_SEARCH_VALUES)) + shape_initial_percentiles: List[float] = field( + default_factory=lambda: list(DEFAULT_INITIAL_PERCENTILES) + ) + shape_refine_rounds: int = 2 + shape_final_top_k: int = 3 + seed_budget: int = 8 + max_compile_trials_per_shape: int = 16 + neighbors_per_step: int = 2 + compile_initial_temperature: float = 0.20 + compile_cooling: float = 0.85 + candidate_pool_per_shape: int = 3 + debug: bool = True + reference_fn: Any = None + bench_warmup: Optional[int] = None + bench_active: Optional[int] = None + stage1_bench_warmup: int = 5 + stage1_bench_rep: int = 10 + max_stage1_children_per_round: int = DEFAULT_MAX_STAGE1_CHILDREN_PER_ROUND + values_were_default: bool = True + + +def _as_list(name: str, value: Any) -> List[Any]: + values = list(value) if isinstance(value, (list, tuple)) else [value] + if not values: + raise ValueError(f"search_params '{name}' must not be empty") + return values + + +def parse_search_params_hint(hint: Any) -> SearchParamsSpec: + if hint is None or hint is False: + return SearchParamsSpec(enabled=False) + if not isinstance(hint, dict): + raise TypeError("hints['search_params'] must be a dict") + + raw = dict(hint) + params = _as_list("params", raw.pop("params", [])) + if not all(isinstance(name, str) and name for name in params): + raise ValueError("search_params params must be non-empty strings") + + values_were_default = "values" not in raw + values = _as_list("values", raw.pop("values", DEFAULT_SEARCH_VALUES)) + if not all(isinstance(value, int) and value > 0 for value in values): + raise ValueError("search_params values must be positive integers") + + percentiles = _as_list( + "shape_initial_percentiles", + raw.pop("shape_initial_percentiles", DEFAULT_INITIAL_PERCENTILES), + ) + if not all( + isinstance(value, (int, float)) and 0 <= value <= 1 for value in percentiles + ): + raise ValueError("shape_initial_percentiles must be in [0, 1]") + + spec = SearchParamsSpec( + enabled=True, + params=params, + values=sorted(set(values)), + shape_initial_percentiles=[float(value) for value in percentiles], + shape_refine_rounds=int(raw.pop("shape_refine_rounds", 2)), + shape_final_top_k=int(raw.pop("shape_final_top_k", 3)), + seed_budget=int(raw.pop("seed_budget", 8)), + max_compile_trials_per_shape=int(raw.pop("max_compile_trials_per_shape", 16)), + neighbors_per_step=int(raw.pop("neighbors_per_step", 2)), + compile_initial_temperature=float(raw.pop("compile_initial_temperature", 0.20)), + compile_cooling=float(raw.pop("compile_cooling", 0.85)), + candidate_pool_per_shape=int(raw.pop("candidate_pool_per_shape", 3)), + debug=bool(raw.pop("debug", True)), + reference_fn=raw.pop("reference_fn", None), + bench_warmup=raw.pop("bench_warmup", None), + bench_active=raw.pop("bench_active", None), + stage1_bench_warmup=int(raw.pop("stage1_bench_warmup", 5)), + stage1_bench_rep=int(raw.pop("stage1_bench_rep", 10)), + max_stage1_children_per_round=int( + raw.pop( + "max_stage1_children_per_round", + DEFAULT_MAX_STAGE1_CHILDREN_PER_ROUND, + ) + ), + values_were_default=values_were_default, + ) + if raw: + raise ValueError(f"Unknown search_params option(s): {sorted(raw)}") + if spec.shape_refine_rounds < 0: + raise ValueError("shape_refine_rounds must be >= 0") + if spec.shape_final_top_k <= 0: + raise ValueError("shape_final_top_k must be positive") + if spec.stage1_bench_warmup < 0 or spec.stage1_bench_rep <= 0: + raise ValueError( + "stage1_bench_warmup must be >= 0 and stage1_bench_rep must be > 0" + ) + if spec.max_stage1_children_per_round <= 0: + raise ValueError("max_stage1_children_per_round must be positive") + return spec + + +def apply_no_dot_search_defaults(spec: SearchParamsSpec) -> None: + """Expand default search values for no-dot operators. + + Explicit user-provided values are left untouched. The autotuner calls this + only after operator classification decides the kernel should use the vector + search policy. + """ + if not spec.enabled or not spec.values_were_default: + return + spec.values = sorted(set(DEFAULT_NO_DOT_SEARCH_VALUES)) + + +def select_initial_percentile_shapes( + configs: Sequence[Config], + spec: SearchParamsSpec, + *, + limit: Optional[int] = None, +) -> List[Config]: + if not configs: + return [] + + space = DiscreteShapeSpace(configs, spec) + selected = [] + seen = set() + value_map = space.value_map + percentiles = stage1_initial_percentiles(spec, value_map) + target_values_by_param = [ + [values[index] for index in _percentile_indices(len(values), percentiles)] + for values in (value_map[name] for name in spec.params) + ] + + for value_combo in itertools.product(*target_values_by_param): + target_shape = dict(zip(spec.params, value_combo)) + config = _nearest_unseen_config(target_shape, configs, spec, seen) + if config is None: + continue + key = shape_key(extract_shape(config, spec.params), spec.params) + if space.contains_key(key) and key not in seen: + seen.add(key) + selected.append(config) + + if limit is None or len(selected) == limit: + return selected + if len(selected) > limit: + return _select_center_axis_configs(configs, spec, limit) + + return _select_diverse_configs( + list(selected) + [config for config in configs if config not in selected], + spec, + limit, + seed_configs=selected, + ) + + +def propose_evolved_shape_configs( + *, + parents: Sequence[Dict[str, Any]], + successes: Sequence[Dict[str, Any]], + failures: Sequence[Config], + failure_observations: Optional[Sequence[Mapping[str, Any]]] = None, + observed: Sequence[Config], + all_configs: Sequence[Config], + spec: SearchParamsSpec, + seen_keys: set, + limit: Optional[int] = None, +) -> List[Config]: + """Generate Stage-1 children with a bounded acquisition model. + + The search space is discrete. Every unseen point gets a score in log-time + units from nearby successful timings, nearby failures, predicted UB risk and + distance from already measured points. Lower score is better. Batch diversity + is applied greedily so one round does not spend all children around the same + local basin. + """ + + child_limit = ( + limit if limit is not None else stage1_children_per_round(spec, all_configs) + ) + if child_limit <= 0: + return [] + + space = DiscreteShapeSpace(all_configs, spec) + failure_observations = list(failure_observations or []) + acquisition = ShapeAcquisition.from_history( + space=space, + parents=parents, + successes=successes, + failures=failures, + failure_observations=failure_observations, + observed=observed, + spec=spec, + ) + + candidates = [] + for key, config in space.by_key.items(): + if key in seen_keys: + continue + shape = extract_shape(config, spec.params) + point = space.point_from_shape(shape) + score = acquisition.cost(point) + candidates.append([score, config, point]) + + selected = [] + selected_keys = set() + selected_points = [] + while candidates and len(selected) < child_limit: + best_index = min( + range(len(candidates)), + key=lambda index: ( + acquisition.batch_adjusted_cost(candidates[index][2], selected_points), + candidates[index][2], + ), + ) + _, config, point = candidates.pop(best_index) + key = shape_key(extract_shape(config, spec.params), spec.params) + if key in selected_keys: + continue + selected.append(config) + selected_keys.add(key) + selected_points.append(point) + return selected + + +def stage1_initial_percentiles( + spec: SearchParamsSpec, + value_map: Optional[Dict[str, List[int]]] = None, +) -> List[float]: + if not value_map: + return list(spec.shape_initial_percentiles) + counts = [len(values) for values in value_map.values()] + if counts and min(counts) < SMALL_VALUE_COUNT_THRESHOLD: + return list(SMALL_VALUE_PERCENTILES) + return list(spec.shape_initial_percentiles) + + +def stage1_total_budget( + spec: SearchParamsSpec, + configs: Optional[Sequence[Config]] = None, +) -> int: + total = ( + len(configs) if configs is not None else len(spec.values) ** len(spec.params) + ) + if total <= 0: + return 0 + budget = 8 + 3 * len(spec.params) + budget = min(budget, DEFAULT_STAGE1_TOTAL_BUDGET_CAP) + return min(total, max(spec.shape_final_top_k, budget)) + + +def stage1_initial_budget( + spec: SearchParamsSpec, + configs: Optional[Sequence[Config]] = None, +) -> int: + total_budget = stage1_total_budget(spec, configs) + if total_budget <= 0: + return 0 + initial = 2 * len(spec.params) + (4 if len(spec.params) >= 2 else 3) + return min(total_budget, max(1, initial)) + + +def stage1_children_per_round( + spec: SearchParamsSpec, + configs: Optional[Sequence[Config]] = None, + *, + remaining_unseen: Optional[int] = None, + remaining_budget: Optional[int] = None, + remaining_rounds: Optional[int] = None, +) -> int: + value_map = ( + effective_search_value_map(configs, spec) if configs is not None else None + ) + points_per_dim = len(stage1_initial_percentiles(spec, value_map)) + count = min( + points_per_dim ** len(spec.params), + spec.max_stage1_children_per_round, + ) + if remaining_unseen is not None: + count = min(count, max(0, remaining_unseen)) + if remaining_budget is not None: + budget = max(0, remaining_budget) + if remaining_rounds is not None and remaining_rounds > 0: + budget = int(math.ceil(budget / remaining_rounds)) + count = min(count, budget) + return count + + +def select_top_shape_entries(entries: Sequence[Any], *, k: int, key_fn) -> List[Any]: + return sorted(entries, key=key_fn)[:k] + + +def filter_shapes_by_ub_and_timing( + entries: Sequence[Any], + *, + ub_floor_ratio: float = UB_FLOOR_RATIO, + timing_ceiling_ratio: float = TIMING_CEILING_RATIO, + timing_sort_key, +) -> List[Any]: + """Drop entries whose UB is below ``ub_floor_ratio * max_ub`` AND whose timing + is above ``timing_ceiling_ratio * min_timing``. Both conditions must hold for a + drop; a small-UB shape that is also fast is kept. + + UB=None on any entry disables the filter (returns the input unchanged). This + keeps backward compatibility when compile-time memory info is unavailable. + """ + if not entries: + return list(entries) + ub_values = [entry.get("ub") for entry in entries] + if any(value is None for value in ub_values): + return list(entries) + max_ub = max(ub_values) + if max_ub <= 0: + return list(entries) + timings = [timing_sort_key(entry["time"]) for entry in entries] + min_timing = min(timings) + if min_timing <= 0: + return list(entries) + ub_floor = ub_floor_ratio * max_ub + timing_ceiling = timing_ceiling_ratio * min_timing + kept = [] + for entry, ub, timing in zip(entries, ub_values, timings): + if ub < ub_floor and timing > timing_ceiling: + continue + kept.append(entry) + return kept + + +def ub_timing_weighted_key(entry, *, timing_sort_key): + """UB-aware rank key after ``filter_shapes_by_ub_and_timing``. + + The hard policy is handled by ``filter_shapes_by_ub_and_timing``: + low-UB and clearly-slow shapes are removed. Among the remaining shapes, + timing stays primary and UB is only a tie-breaker. This keeps the search + performance-oriented while still preferring larger UB when latency is close. + + Returns a tuple that sorts ascending under ``sorted``. + """ + ub = entry.get("ub") + time_value = timing_sort_key(entry["time"]) + if not ub or time_value <= 0: + return (1, time_value, 0) + return (0, time_value, -ub) + + +def _percentile_indices(count: int, percentiles: Sequence[float]) -> List[int]: + if count <= 0: + return [] + if count == 1: + return [0] + + indices = [] + for percentile in percentiles: + raw = percentile * (count - 1) + if percentile < 0.5: + index = int(round(raw)) + else: + index = int(math.ceil(raw)) + if count > len(percentiles) + 1: + index = min(max(index, 1), count - 2) + else: + index = min(max(index, 0), count - 1) + if index not in indices: + indices.append(index) + + # Very short domains can collapse nearby percentiles. Fill from the closest + # remaining high-percentile side first, so [0.30, 0.70] over 3 values becomes + # middle/high instead of a single middle point. + if len(indices) < min(len(percentiles), count): + for index in range(count - 1, -1, -1): + if index in indices: + continue + indices.append(index) + if len(indices) >= min(len(percentiles), count): + break + return indices + + +def _nearest_unseen_config( + target_shape: Dict[str, int], + configs: Sequence[Config], + spec: SearchParamsSpec, + seen: set, +) -> Optional[Config]: + best = None + best_distance = None + for config in configs: + shape = extract_shape(config, spec.params) + key = shape_key(shape, spec.params) + if key in seen: + continue + distance = _shape_distance(target_shape, shape, spec) + if best is None or distance < best_distance: + best = config + best_distance = distance + return best + + +def _select_diverse_configs( + configs: Sequence[Config], + spec: SearchParamsSpec, + limit: int, + *, + seed_configs: Sequence[Config] = (), +) -> List[Config]: + if limit <= 0: + return [] + + space = DiscreteShapeSpace(configs, spec) + unique = list(space.unique_configs) + + selected = [] + selected_keys = set() + for config in seed_configs: + key = space.key(config) + if key in selected_keys or key not in space.by_key: + continue + selected.append(space.by_key[key]) + selected_keys.add(key) + if len(selected) >= limit: + return selected + + if not selected and unique: + center = tuple((len(spec.values) - 1) / 2.0 for _ in spec.params) + first = min( + unique, + key=lambda config: _raw_norm( + tuple( + float(index) - center_dim + for index, center_dim in zip( + space.point(config), + center, + ) + ) + ), + ) + selected.append(first) + selected_keys.add(space.key(first)) + + while len(selected) < limit: + remaining = [ + config for config in unique if space.key(config) not in selected_keys + ] + if not remaining: + break + selected_points = [space.point(config) for config in selected] + best = max( + remaining, + key=lambda config: ( + _min_distance( + space.point(config), + selected_points, + spec, + ), + _edge_distance( + space.point(config), + spec, + ), + ), + ) + selected.append(best) + selected_keys.add(space.key(best)) + return selected + + +def _select_center_axis_configs( + configs: Sequence[Config], + spec: SearchParamsSpec, + limit: int, +) -> List[Config]: + """Trim initial sampling without picking expensive corners first. + + The initial budget is small, so the first round should estimate a local + shape surface around the middle of the discrete domain instead of probing + high-high corners such as 1024x1024. This is a discrete central-composite + design: center hypercube first, then one-axis near-ring probes. + """ + if limit <= 0: + return [] + + space = DiscreteShapeSpace(configs, spec) + by_point = space.by_point + + if len(by_point) <= limit: + return [by_point[point] for point in sorted(by_point)] + + values_by_dim = space.values_by_dim + center_choices = [] + low_center = [] + for values in values_by_dim: + if not values: + center_choices.append([]) + low_center.append(0) + continue + upper = len(values) // 2 + lower = max(0, upper - 1) if len(values) % 2 == 0 else upper + choices = [values[lower]] + if values[upper] not in choices: + choices.append(values[upper]) + center_choices.append(choices) + low_center.append(values[lower]) + + selected_points: List[Tuple[int, ...]] = [] + + def add_point(point): + if point in by_point and point not in selected_points: + selected_points.append(point) + + for point in itertools.product(*center_choices): + add_point(tuple(point)) + if len(selected_points) >= limit: + return [by_point[point] for point in selected_points] + + # Add one-axis high-side probes before low-side probes. For the other + # dimensions we enumerate the center choices, so 2-D domains cover both + # (low_center, high_axis) and (high_center, high_axis) within an 8 point + # budget. That keeps useful middle-large tiles visible without jumping to + # high-high corners. + for dim in range(len(spec.params)): + values = values_by_dim[dim] + base_index = values.index(low_center[dim]) + if base_index + 2 < len(values): + for center_combo in itertools.product(*center_choices): + point = list(center_combo) + point[dim] = values[base_index + 2] + add_point(tuple(point)) + if len(selected_points) >= limit: + return [by_point[point] for point in selected_points] + + for dim in range(len(spec.params)): + values = values_by_dim[dim] + base_index = values.index(low_center[dim]) + if base_index - 1 >= 0: + for center_combo in itertools.product(*center_choices): + point = list(center_combo) + point[dim] = values[base_index - 1] + add_point(tuple(point)) + if len(selected_points) >= limit: + return [by_point[point] for point in selected_points] + + center = tuple((values[0] + values[-1]) / 2.0 for values in values_by_dim) + remaining = [point for point in by_point if point not in selected_points] + remaining.sort( + key=lambda point: ( + _index_distance_float(point, center, spec), + -sum(point), + max(point), + point, + ) + ) + for point in remaining: + selected_points.append(point) + if len(selected_points) >= limit: + break + return [by_point[point] for point in selected_points] + + +def _shape_distance( + a: Dict[str, int], b: Dict[str, int], spec: SearchParamsSpec +) -> float: + denom = max(1, len(spec.values) - 1) + distance = 0.0 + for name in spec.params: + ai = spec.values.index(a[name]) + bi = spec.values.index(b[name]) + distance += abs(ai - bi) / denom + return distance + + +def _shape_index_vector( + shape: Dict[str, int], spec: SearchParamsSpec +) -> Tuple[int, ...]: + return tuple(spec.values.index(shape[name]) for name in spec.params) + + +def _edge_distance(point: Tuple[int, ...], spec: SearchParamsSpec) -> float: + max_index = len(spec.values) - 1 + if max_index <= 0: + return 0.0 + return min(min(index, max_index - index) / max_index for index in point) + + +def _min_distance( + point: Tuple[int, ...], + points: Sequence[Tuple[int, ...]], + spec: SearchParamsSpec, +) -> float: + if not points: + return 1.0 + return min(_index_distance(point, other, spec) for other in points) + + +def _index_distance( + a: Tuple[int, ...], + b: Tuple[int, ...], + spec: SearchParamsSpec, +) -> float: + return _index_distance_float(a, b, spec) + + +def _index_distance_float( + a: Sequence[float], + b: Sequence[float], + spec: SearchParamsSpec, +) -> float: + denom = max(1, len(spec.values) - 1) + squared = 0.0 + for ai, bi in zip(a, b): + squared += ((ai - bi) / denom) ** 2 + return math.sqrt(squared) + + +def _raw_norm(values: Sequence[float]) -> float: + return math.sqrt(sum(value * value for value in values)) diff --git a/backend/ascend_autotune_runtime/tile_shape_space.py b/backend/ascend_autotune_runtime/tile_shape_space.py new file mode 100644 index 00000000..54c732f0 --- /dev/null +++ b/backend/ascend_autotune_runtime/tile_shape_space.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict, List, Sequence, Tuple + +from triton.runtime.autotuner import Config + + +Shape = Dict[str, int] +ShapeKey = Tuple[int, ...] +ShapePoint = Tuple[int, ...] + + +def shape_key(shape: Shape, params: Sequence[str]) -> ShapeKey: + return tuple(shape[name] for name in params) + + +def extract_shape(config: Config, params: Sequence[str]) -> Shape: + return {name: config.kwargs[name] for name in params} + + +def make_shape_config( + base_config: Config, shape: Shape, params: Sequence[str] +) -> Config: + kwargs = dict(base_config.kwargs) + for name in params: + kwargs.pop(name, None) + kwargs.update(shape) + return Config( + kwargs=kwargs, + num_warps=getattr(base_config, "num_warps", 4), + num_stages=getattr(base_config, "num_stages", 2), + num_ctas=getattr(base_config, "num_ctas", 1), + maxnreg=getattr(base_config, "maxnreg", None), + pre_hook=getattr(base_config, "pre_hook", None), + ir_override=getattr(base_config, "ir_override", None), + ) + + +def expand_search_param_shapes(base_config: Config, spec) -> List[Config]: + import itertools + + return [ + make_shape_config(base_config, dict(zip(spec.params, combo)), spec.params) + for combo in itertools.product(spec.values, repeat=len(spec.params)) + ] + + +def effective_search_value_map(configs: Sequence[Config], spec) -> Dict[str, List[int]]: + return { + name: sorted( + {config.kwargs[name] for config in configs if name in config.kwargs} + ) + for name in spec.params + } + + +@dataclass(frozen=True) +class DiscreteShapeSpace: + """Index-space view of search-param configs. + + Search algorithms operate on small integer index vectors, while Triton + kernels need concrete ``Config`` objects. This class keeps that translation + in one place. + """ + + configs: Sequence[Config] + spec: object + + def __post_init__(self): + by_key = {} + by_point = {} + unique = [] + for config in self.configs: + shape = self.shape(config) + key = shape_key(shape, self.spec.params) + if key not in by_key: + by_key[key] = config + unique.append(config) + point = self.point_from_shape(shape) + by_point.setdefault(point, config) + object.__setattr__(self, "by_key", by_key) + object.__setattr__(self, "by_point", by_point) + object.__setattr__(self, "unique_configs", tuple(unique)) + object.__setattr__(self, "points", tuple(by_point)) + object.__setattr__( + self, "value_map", effective_search_value_map(self.configs, self.spec) + ) + object.__setattr__( + self, + "values_by_dim", + tuple( + sorted({point[dim] for point in by_point}) + for dim in range(len(self.spec.params)) + ), + ) + + def shape(self, config: Config) -> Shape: + return extract_shape(config, self.spec.params) + + def key(self, config: Config) -> ShapeKey: + return shape_key(self.shape(config), self.spec.params) + + def point(self, config: Config) -> ShapePoint: + return self.point_from_shape(self.shape(config)) + + def point_from_shape(self, shape: Shape) -> ShapePoint: + return tuple(self.spec.values.index(shape[name]) for name in self.spec.params) + + def config_for_point(self, point: ShapePoint) -> Config | None: + return self.by_point.get(point) + + def contains_key(self, key: ShapeKey) -> bool: + return key in self.by_key diff --git a/backend/npu.py b/backend/npu.py index 758bc936..53e254fb 100644 --- a/backend/npu.py +++ b/backend/npu.py @@ -29,7 +29,6 @@ _check_bishengir_api_change, _check_bishengir_able_save_ir, _is_debug_line_info_disabled, - _enable_print_ub_bits, _enable_dump_memory_info, _enable_msdebug, _enable_unpublished_feature, @@ -515,6 +514,41 @@ def get_libdevice(): return os.path.join(current, "lib/libdevice.10.bc") +def _collect_required_ub_bits_from_memory_info(tmpdir: str, *, debug: bool = False): + """Parse compiler memory display JSONs produced in ``tmpdir``. + + This runs only after bishengir-compile succeeds and while ``tmpdir`` still + exists. ``memory_info`` is a compiler side artifact, so do not try to infer + it from runtime benchmark/profiler data. + """ + from .ascend_autotune_runtime.resource_memory_parser import peak_ub_bits + + peak = 0 + parsed = [] + for side in ("aic", "aiv"): + path = os.path.join(tmpdir, f"memory_info_{side}.json") + if not os.path.isfile(path): + parsed.append((side, "missing", None)) + continue + bits = peak_ub_bits(path) + parsed.append((side, "ok", bits)) + if bits is not None and bits > peak: + peak = bits + + if debug or os.getenv("TRITON_MEMORY_DISPLAY_DEBUG", "false").lower() in ( + "true", + "1", + ): + details = ", ".join(f"{side}:{status}:{bits}" for side, status, bits in parsed) + print( + f"[DEBUG] memory display parsed required_ub_bits={peak or 0}; " + f"files={details}", + flush=True, + ) + + return peak + + # --------------------------------------------------------------------------- # Shared NPU compilation orchestration # --------------------------------------------------------------------------- @@ -598,6 +632,7 @@ def _compile_linalg_to_npu_bin( ret = subprocess.run( cmd_list, env=env, + cwd=tmpdir, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=True, @@ -625,9 +660,13 @@ def _compile_linalg_to_npu_bin( stdout_bytes, stderr_bytes, tmpdir, metadata["hash"] ) - match = re.search(r"UB\s+size\s*=\s*(\d+)\s*bits", stdout_str) - if match: - metadata["required_ub_bits"] = int(match.group(1)) + # When `--enable-memory-display=true` is in the option list, parse the + # compiler memory plan while cwd tmpdir still contains + # memory_info_{aic,aiv}.json. + if "--enable-memory-display=true" in _compile_option_list: + peak = _collect_required_ub_bits_from_memory_info(tmpdir, debug=opt.debug) + if peak > 0: + metadata["required_ub_bits"] = peak if not Path(bin_path).exists(): error_msg = ret.stderr.decode("utf-8") if ret.stderr else "" @@ -688,8 +727,6 @@ def _build_options(m, o): opts.append("--enable-sanitizer=true") if not _is_debug_line_info_disabled(): opts.append("--enable-debug-info=true") - if _enable_print_ub_bits(): - opts.append("--enable-print-memory-allocated-size") enable_hivm_auto_cv_balance = m["enable_hivm_auto_cv_balance"] if enable_hivm_auto_cv_balance is not None: @@ -866,8 +903,6 @@ def _build_options(m, o): opts.append("--enable-sanitizer=true") if not _is_debug_line_info_disabled(): opts.append("--enable-debug-info=true") - if _enable_print_ub_bits(): - opts.append("--enable-print-memory-allocated-size") if _enable_dump_memory_info(): opts.append("--enable-memory-display=true") if _enable_msdebug(): diff --git a/backend/testing.py b/backend/testing.py index 6ba65d8b..0ac88e3b 100644 --- a/backend/testing.py +++ b/backend/testing.py @@ -31,8 +31,24 @@ def get_home_dir(): return os.getenv("TRITON_HOME", Path.home()) +_AIC_METRIC_NAMES = { + "pipe": "PipeUtilization", + "memory": "Memory", + "memory_l0": "MemoryL0", + "memory_ub": "MemoryUB", + "l2": "L2Cache", + "arith": "ArithmeticUtilization", +} + + def do_bench_npu( - funcs, warmup=5, active=30, clear_l2_cache=False, prof_dir=None, keep_res=False + funcs, + warmup=5, + active=30, + clear_l2_cache=False, + prof_dir=None, + keep_res=False, + aic_metrics="pipe", ): import torch import torch_npu @@ -40,13 +56,27 @@ def do_bench_npu( if not isinstance(funcs, list): funcs = [funcs] + if aic_metrics not in _AIC_METRIC_NAMES: + raise ValueError( + f"unknown aic_metrics={aic_metrics!r}, " + f"candidates={list(_AIC_METRIC_NAMES)}" + ) + metric_attr_name = _AIC_METRIC_NAMES[aic_metrics] + metric_attr = getattr(torch_npu.profiler.AiCMetrics, metric_attr_name, None) + if metric_attr is None: + raise AttributeError( + f"torch_npu.profiler.AiCMetrics has no member " + f"{metric_attr_name!r} (aic_metrics={aic_metrics!r}); " + f"installed torch_npu may be too old" + ) + # warmup kernel for fn in funcs: fn() torch.npu.synchronize() experimental_config = torch_npu.profiler._ExperimentalConfig( - aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + aic_metrics=metric_attr, profiler_level=torch_npu.profiler.ProfilerLevel.Level1, l2_cache=False, data_simplification=False, diff --git a/backend/utils.py b/backend/utils.py index 4e999516..189ad7c2 100644 --- a/backend/utils.py +++ b/backend/utils.py @@ -573,10 +573,6 @@ def _is_debug_line_info_disabled() -> bool: return os.getenv("TRITON_DISABLE_LINE_INFO", "true").lower() in ("true", "1") -def _enable_print_ub_bits() -> bool: - return os.getenv("ENABLE_PRINT_UB_BITS", "false").lower() in ("true", "1") - - def _enable_dump_memory_info() -> bool: return os.getenv("TRITON_MEMORY_DISPLAY", "false").lower() in ("true", "1") diff --git a/test/ascend/autotune/01-vector-add.py b/test/ascend/autotune/01-vector-add.py deleted file mode 100644 index 1486d0c7..00000000 --- a/test/ascend/autotune/01-vector-add.py +++ /dev/null @@ -1,48 +0,0 @@ -import os - -import torch -import torch_npu -import triton -import triton.language as tl -from backend.testing import do_bench_npu -import triton.backends.dicp_triton.ascend_autotune_hooks - - -@triton.autotune(configs=[], key=["n_elements"]) -@triton.jit -def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - tl.store(output_ptr + offsets, output, mask=mask) - - -def add_torch(x, y): - return x + y - - -def add_autotune(x, y): - output = torch.empty_like(x) - n_elements = output.numel() - add_kernel[lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)]( - x, y, output, n_elements - ) - return output - - -def test_add(size: int): - x = torch.rand(size, device="npu") - y = torch.rand(size, device="npu") - - output_torch = add_torch(x, y) - output_triton = add_autotune(x, y) - assert torch.allclose(output_triton, output_torch) - print(f"Vector Add {size} PASSED!") - - -if __name__ == "__main__": - test_add(98432) diff --git a/test/ascend/autotune/02-fused-softmax.py b/test/ascend/autotune/02-fused-softmax.py deleted file mode 100644 index 659fea80..00000000 --- a/test/ascend/autotune/02-fused-softmax.py +++ /dev/null @@ -1,73 +0,0 @@ -import os - -import torch -import torch_npu -import triton -import triton.language as tl -from backend.testing import do_bench_npu -import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 — install proxy before @triton.autotune - - -@triton.autotune( - configs=[], - key=["n_rows", "n_cols"], -) -@triton.jit -def softmax_kernel( - output_ptr, - input_ptr, - input_row_stride, - output_row_stride, - n_rows, - n_cols, - BLOCK_SIZE: tl.constexpr, - XBLOCK: tl.constexpr, - XBLOCK_SUB: tl.constexpr, -): - row_start = tl.program_id(0) * XBLOCK - for row_idx in tl.range(0, XBLOCK, XBLOCK_SUB): - row_offsets = row_start + row_idx + tl.arange(0, XBLOCK_SUB)[:, None] - col_offsets = tl.arange(0, BLOCK_SIZE)[None, :] - xmask = row_offsets < n_rows - ymask = col_offsets < n_cols - mask = xmask & ymask - input_ptrs = input_ptr + (row_offsets * input_row_stride + col_offsets) - row = tl.load(input_ptrs, mask=mask, other=-float("inf")) - row_minus_max = row - tl.max(row, axis=1).reshape(XBLOCK_SUB, 1).broadcast_to( - XBLOCK_SUB, BLOCK_SIZE - ) - numerator = tl.exp(row_minus_max) - denominator = ( - tl.sum(numerator, axis=1) - .reshape(XBLOCK_SUB, 1) - .broadcast_to(XBLOCK_SUB, BLOCK_SIZE) - ) - softmax_output = numerator / denominator - output_ptrs = output_ptr + (row_offsets * output_row_stride + col_offsets) - tl.store(output_ptrs, softmax_output, mask=mask) - - -def softmax_torch(x): - return torch.softmax(x, axis=-1) - - -def softmax_autotune(x): - n_rows, n_cols = x.shape - BLOCK_SIZE = n_cols - y = torch.empty_like(x) - softmax_kernel[lambda meta: (triton.cdiv(n_rows, meta["XBLOCK"]), 1, 1)]( - y, x, x.stride(0), y.stride(0), n_rows, n_cols, BLOCK_SIZE=BLOCK_SIZE - ) - return y - - -def test_softmax(shape, dtype): - x = torch.randn(shape, dtype=dtype, device="npu") - y_torch = softmax_torch(x) - y_triton = softmax_autotune(x) - assert torch.allclose(y_triton, y_torch) - print(f"Fused Softmax {shape} {dtype} PASSED!") - - -if __name__ == "__main__": - test_softmax((16896, 1024), torch.float32) diff --git a/test/ascend/autotune/03-layer-norm.py b/test/ascend/autotune/03-layer-norm.py deleted file mode 100644 index f5568afe..00000000 --- a/test/ascend/autotune/03-layer-norm.py +++ /dev/null @@ -1,105 +0,0 @@ -import os - -import torch -import torch_npu -import triton -import triton.language as tl -from backend.testing import do_bench_npu -import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 — install proxy before @triton.autotune - - -@triton.autotune( - configs=[], - key=["M", "N"], -) -@triton.jit -def _layer_norm_fwd_fused( - X, - Y, - W, - B, - Mean, - Rstd, - stride, - N, - M, - eps, - XBLOCK_SIZE: tl.constexpr, - RBLOCK_SIZE: tl.constexpr, -): - row_begin = tl.program_id(0) * XBLOCK_SIZE - row_idx = row_begin + tl.arange(0, XBLOCK_SIZE) - row_mask = row_idx < M - row_offsets = row_idx[:, None] * stride - _mean = tl.zeros((XBLOCK_SIZE, RBLOCK_SIZE), dtype=tl.float32) - for off in range(0, N, RBLOCK_SIZE): - col_idx = off + tl.arange(0, RBLOCK_SIZE) - col_mask = col_idx < N - mask = row_mask[:, None] & col_mask[None, :] - a = tl.load(X + row_offsets + col_idx[None, :], mask=mask, other=0.0).to( - tl.float32 - ) - _mean += a - mean = tl.sum(_mean, axis=1, keep_dims=True) / N - _var = tl.zeros((XBLOCK_SIZE, RBLOCK_SIZE), dtype=tl.float32) - for off in range(0, N, RBLOCK_SIZE): - col_idx = off + tl.arange(0, RBLOCK_SIZE) - col_mask = col_idx < N - mask = row_mask[:, None] & col_mask[None, :] - x = tl.load(X + row_offsets + col_idx[None, :], mask=mask, other=0.0).to( - tl.float32 - ) - x = tl.where(mask, x - mean, 0.0) - _var += x * x - var = tl.sum(_var, axis=1, keep_dims=True) / N - rstd = 1 / tl.sqrt(var + eps) - tl.store(Mean + row_idx[:, None], mean, mask=row_mask[:, None]) - tl.store(Rstd + row_idx[:, None], rstd, mask=row_mask[:, None]) - for off in range(0, N, RBLOCK_SIZE): - col_idx = off + tl.arange(0, RBLOCK_SIZE) - col_mask = col_idx < N - mask = row_mask[:, None] & col_mask[None, :] - w = tl.load(W + col_idx, mask=col_mask).reshape((1, RBLOCK_SIZE)) - b = tl.load(B + col_idx, mask=col_mask).reshape((1, RBLOCK_SIZE)) - x = tl.load(X + row_offsets + col_idx[None, :], mask=mask, other=0.0).to( - tl.float32 - ) - x_hat = (x - mean) * rstd - y = x_hat * w + b - tl.store(Y + row_offsets + col_idx[None, :], y, mask=mask) - - -def layer_norm_torch(args): - x, w_shape, weight, bias, eps, dtype = args - return torch.nn.functional.layer_norm(x, w_shape, weight, bias, eps).to(dtype) - - -def layer_norm_autotune(args): - x, weight, bias, eps = args - y = torch.empty_like(x) - x_arg = x.reshape(-1, x.shape[-1]) - M, N = x_arg.shape - mean = torch.empty((M,), dtype=torch.float32, device=x.device) - rstd = torch.empty((M,), dtype=torch.float32, device=x.device) - _layer_norm_fwd_fused[lambda meta: (triton.cdiv(M, meta["XBLOCK_SIZE"]), 1, 1)]( - x_arg, y, weight, bias, mean, rstd, x_arg.stride(0), N, M, eps - ) - return y - - -def test_layer_norm(shape, dtype, eps=1e-5): - M, N = shape - device = "npu" - x_shape = shape - w_shape = (x_shape[-1],) - weight = torch.rand(w_shape, dtype=dtype, device=device) - bias = torch.rand(w_shape, dtype=dtype, device=device) - x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device=device) - y_torch = layer_norm_torch((x, w_shape, weight, bias, eps, dtype)) - y_triton = layer_norm_autotune((x, weight, bias, eps)) - assert torch.allclose(y_triton, y_torch, atol=1e-2, rtol=0) - print(f"Layer Normalization {M},{N} {dtype} PASSED!") - - -if __name__ == "__main__": - test_layer_norm((128, 32), torch.float16) diff --git a/test/ascend/autotune/04-libentry.py b/test/ascend/autotune/04-libentry.py deleted file mode 100644 index a229bd0b..00000000 --- a/test/ascend/autotune/04-libentry.py +++ /dev/null @@ -1,59 +0,0 @@ -import os - -import torch -import torch_npu -import triton -import triton.language as tl - -from language.deeplink.runtime import libentry - -from backend.testing import do_bench_npu -import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 — install proxy before @triton.autotune - - -@triton.autotune( - configs=[ - triton.Config({"BLOCK_SIZE": 1 * 1024, "multibuffer": True}), - triton.Config({"BLOCK_SIZE": 12 * 1024, "multibuffer": True}), - triton.Config({"BLOCK_SIZE": 12 * 1024, "multibuffer": False}), - triton.Config({"BLOCK_SIZE": 8 * 1024, "multibuffer": True}), - ], - key=["n_elements"], -) -@libentry() -@triton.jit -def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - tl.store(output_ptr + offsets, output, mask=mask) - - -def add_torch(x, y): - return x + y - - -def add_autotune(x, y): - output = torch.empty_like(x) - n_elements = output.numel() - add_kernel[lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)]( - x, y, output, n_elements - ) - return output - - -def test_add(size: int): - x = torch.rand(size, device="npu") - y = torch.rand(size, device="npu") - output_torch = add_torch(x, y) - output_triton = add_autotune(x, y) - assert torch.allclose(output_triton, output_torch) - print(f"Vector Add {size} with libentry PASSED!") - - -if __name__ == "__main__": - test_add(98432) diff --git a/test/ascend/autotune/autotune_regression_goldens.json b/test/ascend/autotune/autotune_regression_goldens.json new file mode 100644 index 00000000..ce02efc9 --- /dev/null +++ b/test/ascend/autotune/autotune_regression_goldens.json @@ -0,0 +1,102 @@ +{ + "fa2": { + "script": "fa2_case.py", + "expected_operator_kind": "DOT_STATEFUL", + "expected_compile_options": "mixcv", + "golden_precise_cost": [ + 0.06627206666666666, + 13.949780766666668 + ], + "golden_triton_ms": [ + 66.048, + 13963.672666666665 + ], + "expected_autotune_runs": 2, + "max_search_time_s": [ + 165.0, + 175.0 + ], + "max_precise_cost_ratio": 1.2, + "max_triton_ms_ratio": 1.25, + "timeout_s": 900 + }, + "fla_cumsum": { + "script": "fla_cumsum_case.py", + "expected_operator_kind": "VECTOR_DISCRETE_OR_STATEFUL", + "expected_compile_options": "vector", + "golden_precise_cost": [ + 0.003098033333333334 + ], + "golden_triton_ms": [ + 0.285451, + 0.301398, + 0.287805 + ], + "expected_autotune_runs": 1, + "max_search_time_s": 50.0, + "max_precise_cost_ratio": 1.2, + "max_triton_ms_ratio": 1.25, + "timeout_s": 300 + }, + "fla_chunk_delta_hupdate": { + "script": "fla_chunk_delta_hupdate_case.py", + "expected_operator_kind": "DOT_STATEFUL", + "expected_compile_options": "mixcv", + "golden_precise_cost": [ + 0.013484333333333336 + ], + "golden_triton_ms": [ + 0.598872, + 0.559796, + 0.55266 + ], + "expected_autotune_runs": 1, + "max_search_time_s": 120.0, + "max_precise_cost_ratio": 1.2, + "max_triton_ms_ratio": 1.25, + "timeout_s": 600 + }, + "gdn_chunk_meta": { + "script": "gdn_chunk_meta_case.py", + "expected_operator_kind": "VECTOR_AFFINE", + "expected_compile_options": "vector", + "golden_precise_cost": [ + 0.0015280666666666668, + 0.04506756666666667, + 0.0012441 + ], + "golden_triton_ms": [ + 0.7816, + 1.05994, + 4.20512 + ], + "expected_autotune_runs": 3, + "max_search_time_s": [ + 70.0, + 55.0, + 55.0 + ], + "max_precise_cost_ratio": 1.2, + "max_triton_ms_ratio": 1.25, + "timeout_s": 600 + }, + "batch_invariant_mean": { + "script": "batch_invariant_mean_case.py", + "expected_operator_kind": "VECTOR_REDUCTION", + "expected_compile_options": "vector", + "golden_precise_cost": [ + 0.0218057 + ], + "golden_triton_ms": [ + 0.307106, + 0.28717, + 0.27884, + 0.277523 + ], + "expected_autotune_runs": 1, + "max_search_time_s": 75.0, + "max_precise_cost_ratio": 1.2, + "max_triton_ms_ratio": 1.25, + "timeout_s": 300 + } +} diff --git a/test/ascend/autotune/conftest.py b/test/ascend/autotune/conftest.py index a2e61e64..9b3ea98e 100644 --- a/test/ascend/autotune/conftest.py +++ b/test/ascend/autotune/conftest.py @@ -2,3 +2,7 @@ # We import this early so every test module sees the proxy (which auto-detects # ascend via triton.runtime.driver.active.target at call time). import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 — side-effect import + + +def pytest_configure(config): + config.addinivalue_line("markers", "autotune: tests that require Ascend autotune") diff --git a/test/ascend/autotune/demo_backend_runtime_autotune.py b/test/ascend/autotune/demo_backend_runtime_autotune.py deleted file mode 100644 index 401aab40..00000000 --- a/test/ascend/autotune/demo_backend_runtime_autotune.py +++ /dev/null @@ -1,48 +0,0 @@ -import torch -import torch_npu -import triton -import triton.language as tl - -from backend.ascend_autotune_runtime import autotune as ascend_autotune - - -@ascend_autotune( - configs=[ - triton.Config({"BLOCK_SIZE": 256}, num_warps=4), - triton.Config({"BLOCK_SIZE": 512}, num_warps=4), - triton.Config({"BLOCK_SIZE": 1024}, num_warps=8), - ], - key=["n_elements"], - hints={"compile_options": "vector"}, -) -@triton.jit -def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(0) - offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - tl.store(out_ptr + offsets, x + y, mask=mask) - - -def add(x, y): - out = torch.empty_like(x) - n_elements = out.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) - add_kernel[grid](x, y, out, n_elements) - return out - - -def main(): - n_elements = 98432 - x = torch.rand(n_elements, device="npu", dtype=torch.float32) - y = torch.rand(n_elements, device="npu", dtype=torch.float32) - - out = add(x, y) - torch.npu.synchronize() - torch.testing.assert_close(out, x + y, rtol=1e-3, atol=1e-3) - print("backend.ascend_autotune_runtime autotune demo PASSED") - - -if __name__ == "__main__": - main() diff --git a/test/ascend/autotune/regression_cases/batch_invariant_mean_case.py b/test/ascend/autotune/regression_cases/batch_invariant_mean_case.py new file mode 100644 index 00000000..ac6aed9f --- /dev/null +++ b/test/ascend/autotune/regression_cases/batch_invariant_mean_case.py @@ -0,0 +1,168 @@ +import torch +import triton +import triton.language as tl +from common import ( + assert_close, + bench, + case_values, + dtype_from_name, + init, + parse_args, + report, +) + +MODELS = "Runtime batch-invariant replacement; kernel coverage does not depend on model availability" + + +def mean_reference_check(tuner, config, *args, **kwargs): + input_3d = args[0] + output_2d = args[1] + M = args[7] + N = args[8] + K = args[9] + output_2d.zero_() + tuner._make_kernel_call(*args, config=config, **kwargs)(warmup=False) + ref = torch.mean(input_3d.reshape(M, N, K).to(torch.float16), dim=1) + return torch.allclose(output_2d, ref.reshape_as(output_2d), atol=1e-2, rtol=1e-2) + + +MEAN_SEARCH_HINTS = { + "search_params": { + "params": ["BLOCK_SIZE"], + "reference_fn": mean_reference_check, + }, +} + + +@triton.autotune( + configs=[], + key=[], + hints=MEAN_SEARCH_HINTS, +) +@triton.jit +def mean_kernel( + input_ptr, + output_ptr, + input_stride0, + input_stride1, + input_stride2, + output_stride0, + output_stride1, + M, # size before reduction dim + N, # size of reduction dim + K, # size after reduction dim + BLOCK_SIZE: tl.constexpr, +): + """ + Kernel for computing mean along a single dimension. + Input is viewed as (M, N, K) where N is the dimension being reduced. + """ + # Program ID gives us which output element we're computing + pid = tl.program_id(0) + + # Compute output indices + m_idx = pid // K + k_idx = pid % K + + # Bounds check + if m_idx >= M or k_idx >= K: + return + # Accumulate sum across reduction dimension + acc = 0.0 + for n_start in range(0, N, BLOCK_SIZE): + n_offsets = n_start + tl.arange(0, BLOCK_SIZE) + mask = n_offsets < N + + # Calculate input indices + input_idx = ( + m_idx * input_stride0 + n_offsets * input_stride1 + k_idx * input_stride2 + ) + # Load and accumulate + vals = tl.load(input_ptr + input_idx, mask=mask, other=0.0) + acc += tl.sum(vals) + + # Compute mean and store + mean_val = acc / N + output_idx = m_idx * output_stride0 + k_idx * output_stride1 + tl.store(output_ptr + output_idx, mean_val) + + +def mean_dim( + input_: torch.Tensor, + dim: int, + keepdim: bool = False, + dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + assert ( + -input_.ndim <= dim < input_.ndim + ), f"Invalid dimension {dim} for tensor with {input_.ndim} dimensions" + if dim < 0: + dim = dim + input_.ndim + if dtype is None: + if input_.dtype in [torch.int8, torch.int16, torch.int32, torch.int64]: + dtype = torch.float32 + else: + dtype = input_.dtype + if input_.dtype != dtype: + input_ = input_.to(dtype) + shape = list(input_.shape) + M = 1 + for i in range(dim): + M *= shape[i] + N = shape[dim] + K = 1 + for i in range(dim + 1, len(shape)): + K *= shape[i] + input_3d = input_.reshape(M, N, K) + if keepdim: + output_shape = shape.copy() + output_shape[dim] = 1 + else: + output_shape = shape[:dim] + shape[dim + 1 :] + output = torch.empty(output_shape, dtype=dtype, device=input_.device) + if keepdim: + output_2d = output.reshape(M, 1, K).squeeze(1) + else: + output_2d = output.reshape(M, K) + grid = (M * K,) + mean_kernel[grid]( + input_3d, + output_2d, + input_3d.stride(0), + input_3d.stride(1), + input_3d.stride(2), + output_2d.stride(0), + output_2d.stride(1) if output_2d.ndim > 1 else 0, + M, + N, + K, + ) + return output + + +def main(): + args = parse_args() + init(args.seed) + dtype = dtype_from_name(args.dtype) + for rows, hidden in case_values( + ((512, 4096), (1024, 5120), (2048, 7168), (4096, 8192)) + ): + x = torch.randn(rows, hidden, dtype=dtype, device=args.device) + tri = mean_dim(x, dim=-1, dtype=torch.float16) + ref = torch.mean(x.to(torch.float16), dim=-1) + assert_close("mean_kernel", tri, ref, args.rtol or 1e-2, args.atol or 1e-2) + torch_ms, _ = bench( + lambda: torch.mean(x.to(torch.float16), dim=-1), args.warmup, args.repeat + ) + triton_ms, _ = bench( + lambda: mean_dim(x, dim=-1, dtype=torch.float16), args.warmup, args.repeat + ) + report(f"mean_kernel shape=({rows}, {hidden})", torch_ms, triton_ms, MODELS) + + +def test_benchmark(): + main() + + +if __name__ == "__main__": + main() diff --git a/test/ascend/autotune/regression_cases/common.py b/test/ascend/autotune/regression_cases/common.py new file mode 100644 index 00000000..a9d7e7ce --- /dev/null +++ b/test/ascend/autotune/regression_cases/common.py @@ -0,0 +1,134 @@ +import argparse +import os +import sys +import time + +import torch + +import triton +import triton.language as tl + +try: + import backend.ascend_autotune_hooks # noqa: F401 +except ImportError: + pass + +try: + import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 +except ImportError: + pass + +try: + import triton.language.extra.deeplink.cann.extension as cann_ext +except ImportError: + try: + import triton.language.extra.cann.extension as cann_ext + except ImportError: + cann_ext = None + +if cann_ext is not None: + extract_slice = cann_ext.extract_slice + insert_slice = cann_ext.insert_slice + get_element = cann_ext.get_element +else: + extract_slice = getattr(tl, "extract_slice", None) + insert_slice = getattr(tl, "insert_slice", None) + get_element = getattr(tl, "get_element", None) + + +class _TLDeviceFallback: + fast_dividef = staticmethod(lambda x, y: x / y) + fast_expf = staticmethod(tl.exp) + fast_logf = staticmethod(tl.log) + fast_log2f = staticmethod(tl.log2) + + +tldevice = _TLDeviceFallback() + + +@triton.jit +def safe_exp(x): + return tl.exp(tl.where(x <= 0, x, float("-inf"))) + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--repeat", type=int, default=5) + parser.add_argument("--device", default="npu") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--dtype", default="float16", choices=["float16", "bfloat16", "float32"] + ) + parser.add_argument("--rtol", type=float, default=None) + parser.add_argument("--atol", type=float, default=None) + if "pytest" in sys.modules: + return parser.parse_args([]) + return parser.parse_args() + + +def case_values(cases): + cases = tuple(cases) + limit = int(os.getenv("DLC_AUTOTUNE_CASE_LIMIT", "0")) + if limit <= 0: + return cases + return cases[:limit] + + +def dtype_from_name(name): + return { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, + }[name] + + +def init(seed=0): + torch.manual_seed(seed) + + +def get_vectorcore_num(device=None): + if device is None: + if hasattr(torch, "npu") and torch.npu.is_available(): + device = torch.npu.current_device() + elif torch.cuda.is_available(): + device = torch.cuda.current_device() + else: + device = 0 + props = triton.runtime.driver.active.utils.get_device_properties(device) + return props.get("num_vectorcore", 1) + + +def sync(): + if hasattr(torch, "npu") and torch.npu.is_available(): + torch.npu.synchronize() + elif torch.cuda.is_available(): + torch.cuda.synchronize() + + +def bench(fn, warmup=10, repeat=50): + out = fn() + sync() + ms = triton.testing.do_bench(fn, warmup=warmup, rep=repeat) + return ms, out + + +def assert_close(name, actual, expected, rtol=None, atol=None): + if rtol is None: + rtol = 1e-2 if actual.dtype in (torch.float16, torch.bfloat16) else 1e-4 + if atol is None: + atol = 1e-2 if actual.dtype in (torch.float16, torch.bfloat16) else 1e-4 + if isinstance(actual, torch.Tensor): + actual = actual.detach().cpu() + if isinstance(expected, torch.Tensor): + expected = expected.detach().cpu() + torch.testing.assert_close(actual, expected, rtol=rtol, atol=atol, equal_nan=True) + + +def report(name, torch_ms, triton_ms, models): + speedup = torch_ms / triton_ms if triton_ms > 0 else float("inf") + print(f"op: {name}") + print(f"torch_ms: {torch_ms:.6f}") + print(f"triton_ms: {triton_ms:.6f}") + print(f"speedup: {speedup:.4f}x") + print(f"models: {models}") diff --git a/test/ascend/autotune/regression_cases/fa2_case.py b/test/ascend/autotune/regression_cases/fa2_case.py new file mode 100644 index 00000000..e4bf5cf5 --- /dev/null +++ b/test/ascend/autotune/regression_cases/fa2_case.py @@ -0,0 +1,482 @@ +from __future__ import annotations + +import argparse +import json +import os +from functools import cache +from typing import Any, Dict, Tuple + +import torch +import torch_npu +import triton +import triton.language as tl +from common import case_values +from triton.backends.dicp_triton.testing import do_bench_npu + + +torch_npu.npu.current_device() +import backend.ascend_autotune_hooks as ascend_autotune_hooks + +_ASCEND_AUTOTUNE_HOOKS = ascend_autotune_hooks + +os.environ.setdefault("TRITON_PRINT_AUTOTUNING", "1") +os.environ.setdefault("TRITON_PRINT_AUTOTUNING_TIMINGS", "1") + +DEVICE = "npu" +BENCH_WARMUP = 1 +BENCH_ACTIVE = 3 + + +ATTENTION_SEARCH_HINTS = { + "search_params": { + "params": ["BLOCK_M", "BLOCK_N"], + }, +} + + +FA_TEST_CASES = [ + (1, 32, 512, 128, False, torch.float16), + (1, 32, 1024, 128, False, torch.float16), + (1, 32, 2048, 128, False, torch.float16), + (1, 32, 10240, 128, False, torch.float16), +] + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--warmup", + type=int, + default=int(os.getenv("DLC_AUTOTUNE_BENCH_WARMUP", str(BENCH_WARMUP))), + ) + parser.add_argument( + "--repeat", + type=int, + default=int(os.getenv("DLC_AUTOTUNE_BENCH_ACTIVE", str(BENCH_ACTIVE))), + ) + return parser.parse_args() + + +def prune_attention_configs(configs, nargs, **kwargs): + n_ctx = kwargs.get("N_CTX", nargs.get("N_CTX")) + if n_ctx is None: + return configs + pruned = [ + config + for config in configs + if config.kwargs["BLOCK_M"] <= n_ctx and config.kwargs["BLOCK_N"] <= n_ctx + ] + return pruned or configs + + +@cache +def get_device_properties() -> Tuple[int, int]: + device = torch.npu.current_device() + device_properties: Dict[str, Any] = ( + triton.runtime.driver.active.utils.get_device_properties(device) + ) + + num_aicore = device_properties.get("num_aicore", -1) + num_vectorcore = device_properties.get("num_vectorcore", -1) + + assert num_aicore > 0 and num_vectorcore > 0, "Failed to detect device properties." + return num_aicore, num_vectorcore + + +@triton.jit +def _attn_fwd_inner2( + acc, + l_i, + m_i, + q, + K_block_ptr, + V_block_ptr, + start_m, + qk_scale: tl.constexpr, + BLOCK_M: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, + STAGE: tl.constexpr, + offs_m: tl.constexpr, + offs_n: tl.constexpr, + N_CTX: tl.constexpr, + fp8_v: tl.constexpr, +): + if STAGE == 1: + tl.static_assert(BLOCK_M >= BLOCK_N) + lo, hi = 0, start_m * BLOCK_M + elif STAGE == 2: + tl.static_assert(BLOCK_M >= BLOCK_N) + lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M + lo = tl.multiple_of(lo, BLOCK_M) + else: + lo, hi = 0, N_CTX + + K_block_ptr = tl.advance(K_block_ptr, (lo, 0)) + V_block_ptr = tl.advance(V_block_ptr, (lo, 0)) + for start_n in range(lo, hi, BLOCK_N): + start_n = tl.multiple_of(start_n, BLOCK_N) + k = tl.load(K_block_ptr) + + trans_k = tl.trans(k) + qk = tl.dot(q, trans_k) + + if STAGE == 2: + mask = offs_m[:, None] >= (start_n + offs_n[None, :]) + qk = qk * qk_scale + tl.where(mask, 0, -1.0e6) + m_ij = tl.maximum(m_i, tl.max(qk, 1), propagate_nan=tl.PropagateNan.ALL) + qk -= m_ij[:, None] + else: + qk = qk * qk_scale + m_ij = tl.maximum(m_i, tl.max(qk, 1), propagate_nan=tl.PropagateNan.ALL) + qk = qk - m_ij[:, None] + + p = tl.math.exp(qk) + p_cast = p.to(k.dtype) + v = tl.load(V_block_ptr) + pv = tl.dot(p_cast, v) + tl.extra.deeplink.cann.extension.compile_hint(pv, "hivm.tile_mix_cube_num", 2) + l_ij = tl.sum(p, 1) + + alpha = tl.math.exp(m_i - m_ij) + l_i = l_i * alpha + l_ij + acc = acc * alpha[:, None] + pv + + m_i = m_ij + V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) + K_block_ptr = tl.advance(K_block_ptr, (BLOCK_N, 0)) + return acc, l_i, m_i + + +@triton.autotune( + configs=[], + key=["N_CTX", "HEAD_DIM"], + prune_configs_by={"early_config_prune": prune_attention_configs}, + hints=ATTENTION_SEARCH_HINTS, +) +@triton.jit +def _attn_fwd2( + Q, + K, + V, + M, + Out, + sm_scale: tl.constexpr, + stride_qz: tl.constexpr, + stride_qh: tl.constexpr, + stride_qm: tl.constexpr, + stride_qk: tl.constexpr, + stride_kz: tl.constexpr, + stride_kh: tl.constexpr, + stride_kn: tl.constexpr, + stride_kk: tl.constexpr, + stride_vz: tl.constexpr, + stride_vh: tl.constexpr, + stride_vn: tl.constexpr, + stride_vk: tl.constexpr, + stride_oz: tl.constexpr, + stride_oh: tl.constexpr, + stride_om: tl.constexpr, + stride_on: tl.constexpr, + Z: tl.constexpr, + H: tl.constexpr, + N_CTX: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + STAGE: tl.constexpr, +): + pid = tl.program_id(0) + core_step = tl.num_programs(0) + NUM_BLOCKS_M = tl.cdiv(N_CTX, BLOCK_M) + NUM_BLOCKS = NUM_BLOCKS_M * Z * H + task_m_idx = 0 + task_hz_idx = 0 + + for block_idx in range(pid, NUM_BLOCKS, core_step): + task_hz_idx = block_idx // NUM_BLOCKS_M + task_m_idx = block_idx % NUM_BLOCKS_M + off_z = task_hz_idx // H + off_h = task_hz_idx % H + qvk_offset = off_z.to(tl.int64) * stride_qz + off_h.to(tl.int64) * stride_qh + Q_block_ptr = tl.make_block_ptr( + base=Q + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_qm, stride_qk), + offsets=(task_m_idx * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), + order=(1, 0), + ) + V_block_ptr = tl.make_block_ptr( + base=V + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_vn, stride_vk), + offsets=(0, 0), + block_shape=(BLOCK_N, HEAD_DIM), + order=(1, 0), + ) + K_block_ptr = tl.make_block_ptr( + base=K + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_kn, stride_kk), + offsets=(0, 0), + block_shape=(BLOCK_N, HEAD_DIM), + order=(1, 0), + ) + O_block_ptr = tl.make_block_ptr( + base=Out + qvk_offset, + shape=(N_CTX, HEAD_DIM), + strides=(stride_om, stride_on), + offsets=(task_m_idx * BLOCK_M, 0), + block_shape=(BLOCK_M, HEAD_DIM), + order=(1, 0), + ) + offs_m = task_m_idx * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = tl.arange(0, BLOCK_N) + + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 + acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + q = tl.load(Q_block_ptr) + + if STAGE & 1: + acc, l_i, m_i = _attn_fwd_inner2( + acc, + l_i, + m_i, + q, + K_block_ptr, + V_block_ptr, + task_m_idx, + sm_scale, + BLOCK_M, + HEAD_DIM, + BLOCK_N, + 4 - STAGE, + offs_m, + offs_n, + N_CTX, + V.dtype.element_ty == tl.float8e5, + ) + + if STAGE & 2: + acc, l_i, m_i = _attn_fwd_inner2( + acc, + l_i, + m_i, + q, + K_block_ptr, + V_block_ptr, + task_m_idx, + sm_scale, + BLOCK_M, + HEAD_DIM, + BLOCK_N, + 2, + offs_m, + offs_n, + N_CTX, + V.dtype.element_ty == tl.float8e5, + ) + + m_i += tl.math.log(l_i) + acc = acc / l_i[:, None] + tl.store(O_block_ptr, acc.to(Out.type.element_ty)) + + +class _attention(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k, v, causal, sm_scale, BM=128, BN=None): + head_dim_q, head_dim_k = q.shape[-1], k.shape[-1] + head_dim_v = v.shape[-1] + assert head_dim_q == head_dim_k and head_dim_k == head_dim_v + assert head_dim_k in {16, 32, 64, 128, 512} + + o = torch.empty_like(q) + stage = 3 if causal else 1 + num_cores, _ = get_device_properties() + M = torch.empty( + (q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32 + ) + launch_meta: Dict[str, int] = {} + if BN is not None: + launch_meta["BLOCK_M"] = BM + launch_meta["BLOCK_N"] = BN + _attn_fwd2[(num_cores,)]( + q, + k, + v, + M, + o, + sm_scale, + q.stride(0), + q.stride(1), + q.stride(2), + q.stride(3), + k.stride(0), + k.stride(1), + k.stride(2), + k.stride(3), + v.stride(0), + v.stride(1), + v.stride(2), + v.stride(3), + o.stride(0), + o.stride(1), + o.stride(2), + o.stride(3), + q.shape[0], + q.shape[1], + N_CTX=q.shape[2], + HEAD_DIM=head_dim_k, + STAGE=stage, + debug=os.environ.get("TRITON_DEBUG", "0") == "1", + **launch_meta, + ) + + ctx.save_for_backward(q, k, v, o, M) + ctx.sm_scale = sm_scale + ctx.HEAD_DIM = head_dim_k + ctx.causal = causal + return o + + +attention = _attention.apply + + +def bench_operator(name, fn): + print( + f"--------------------benchmark_{name} for {BENCH_ACTIVE} times--------------------" + ) + took = do_bench_npu( + fn, + warmup=BENCH_WARMUP, + active=BENCH_ACTIVE, + clear_l2_cache=False, + ) + print(f" [op time] {name}: {took:.6f} s ({took * 1000000:.3f} us)", flush=True) + return took + + +def snapshot_attention_search_stats(): + current = dict(getattr(_attn_fwd2, "last_search_stats", {})) + saved = dict(getattr(_attn_fwd2, "_last_case_search_stats", {})) + if saved.get("searched", False) and not current.get("searched", False): + return saved + return current or saved + + +def benchmark(Z, H, N_CTX, HEAD_DIM, causal, dtype, BM=128, BN=None): + torch.manual_seed(20) + q = ( + torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE) + .normal_(mean=0.0, std=0.5) + .requires_grad_() + ) + k = ( + torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE) + .normal_(mean=0.0, std=0.5) + .requires_grad_() + ) + v = ( + torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE) + .normal_(mean=0.0, std=0.5) + .requires_grad_() + ) + sm_scale = 0.5 + + ref_out = torch_npu.npu_fusion_attention( + q, + k, + v, + H, + padding_mask=None, + atten_mask=None, + scale=sm_scale, + keep_prob=1.0, + input_layout="BNSD", + pre_tockens=65535, + next_tockens=65535, + sparse_mode=0, + )[0] + npu_time = bench_operator( + "torch_npu.npu_fusion_attention", + lambda: torch_npu.npu_fusion_attention( + q, + k, + v, + H, + padding_mask=None, + atten_mask=None, + scale=sm_scale, + keep_prob=1.0, + input_layout="BNSD", + pre_tockens=65535, + next_tockens=65535, + sparse_mode=0, + )[0], + ) + + tri_out = attention(q, k, v, causal, sm_scale, BM, BN) + search_stats = snapshot_attention_search_stats() + _attn_fwd2._last_case_search_stats = dict(search_stats) + print(f" [selected config] _attn_fwd2: {_attn_fwd2.best_config}", flush=True) + triton_time = bench_operator( + "attention.forward", + lambda: attention(q, k, v, causal, sm_scale, BM, BN), + ) + + assert torch.allclose(ref_out, tri_out, atol=1e-2, rtol=0.0) + return { + "npu_time": npu_time, + "triton_time": triton_time, + "search_stats": search_stats, + "best_config": str(_attn_fwd2.best_config), + } + + +def load_fa_cases(): + cases = list(FA_TEST_CASES) + selected_cases = os.getenv("TEST_FA2_CASES", "1,4") + if selected_cases: + selected = {int(item) for item in selected_cases.split(",") if item.strip()} + cases = [ + case for case_index, case in enumerate(cases, 1) if case_index in selected + ] + return cases + + +def main(): + global BENCH_ACTIVE, BENCH_WARMUP + + args = parse_args() + BENCH_WARMUP = args.warmup + BENCH_ACTIVE = args.repeat + results = [] + for case in case_values(load_fa_cases()): + _attn_fwd2._last_case_search_stats = {} + results.append(benchmark(*case)) + + payload = { + "case": "fa2", + "ok": True, + "triton_ms": [item["triton_time"] * 1000.0 for item in results], + "search_time": [ + item.get("search_stats", {}).get("bench_time", 0.0) for item in results + ], + "bench_configs": [ + item.get("search_stats", {}).get("bench_configs", 0) for item in results + ], + "measurements": [ + item.get("search_stats", {}).get("measurements", 0) for item in results + ], + "best_config": [item.get("best_config", "") for item in results], + } + print( + "DLC_AUTOTUNE_REGRESSION_RESULT=" + json.dumps(payload, sort_keys=True), + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/test/ascend/autotune/regression_cases/fla_chunk_delta_hupdate_case.py b/test/ascend/autotune/regression_cases/fla_chunk_delta_hupdate_case.py new file mode 100644 index 00000000..86229fc4 --- /dev/null +++ b/test/ascend/autotune/regression_cases/fla_chunk_delta_hupdate_case.py @@ -0,0 +1,259 @@ +# SPDX-License-Identifier: Apache-2.0 +# Exact-copy Triton benchmark/smoke generated from vllm_ascend source kernels. + +import torch +from common import assert_close, bench, case_values, dtype_from_name, init, parse_args +from common import tl, triton + + +def report_smoke(name, triton_ms, models): + print(f"op: {name}") + print("torch_ms: smoke-only") + print(f"triton_ms: {triton_ms:.6f}") + print("speedup: smoke-only") + print(f"models: {models}") + + +from common import safe_exp + +MODELS = "FLA chunk gated delta hupdate exact-copy kernel" + + +@triton.autotune( + configs=[], + key=[], + hints={"search_params": {"params": ["BT"]}}, +) +@triton.heuristics( + { + "USE_G": lambda args: args["g"] is not None, + "IS_VARLEN": lambda args: args["cu_seqlens"] is not None, + } +) +@triton.jit(do_not_specialize=["T"]) +def chunk_gated_delta_rule_fwd_kernel_hupdate_blockdim64( + k, + w, + g, + cu_seqlens, + chunk_offsets, + h_update, + T, + H: tl.constexpr, + Hg: tl.constexpr, + K: tl.constexpr, + BT: tl.constexpr, + USE_G: tl.constexpr, + IS_VARLEN: tl.constexpr, +): + i_nh = tl.program_id(1) + i_n, i_h = i_nh // H, i_nh % H + T_max = 1 * T + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + if IS_VARLEN: + bos, eos = ( + tl.load(cu_seqlens + i_n).to(tl.int32), + tl.load(cu_seqlens + i_n + 1).to(tl.int32), + ) + T = eos - bos + NT = tl.cdiv(T, BT) + boh = tl.load(chunk_offsets + i_n).to(tl.int32) + else: + bos, eos = i_n * T, i_n * T + T + NT = tl.cdiv(T, BT) + boh = i_n * NT + + stride_k = Hg * K + stride_w = H * K + + # create b_hupd_bv1 and b_hupd_bv2 + off_hupd_1_top = tl.arange(0, 64)[:, None] + off_hupd_2_top = tl.arange(0, 64)[None, :] + + # main recurrence + for i_t in range(NT): + last_idx = min((i_t + 1) * BT, T) - 1 + b_g_last = tl.load(g + bos + i_h * T_max + last_idx) + + offs_t = i_t * BT + tl.arange(0, BT) + mask_t = offs_t < T + g_ptr = g + bos + i_h * T_max + b_g = tl.load(g_ptr + offs_t, mask=mask_t, other=0.0) + + b_g = safe_exp(b_g_last - b_g) + b_g_last = tl.exp(b_g_last) + + offs_t_wv = (i_t * BT + tl.arange(0, BT))[:, None] + w_base = w + bos * H * K + i_h * K + # get column-sliced w [BT, 64] + offs_w_upd1 = tl.arange(0, 64)[None, :] + mask_w_upd1 = (offs_t_wv < T) & (offs_w_upd1 < K) + ptr_w_upd1 = w_base + offs_t_wv * stride_w + offs_w_upd1 * 1 + b_w_upd1 = tl.load(ptr_w_upd1, mask=mask_w_upd1, other=0.0).to(tl.float32) + + offs_w_upd2 = 64 + tl.arange(0, 64)[None, :] + mask_w_upd2 = (offs_t_wv < T) & (offs_w_upd2 < K) + ptr_w_upd2 = w_base + offs_t_wv * stride_w + offs_w_upd2 * 1 + b_w_upd2 = tl.load(ptr_w_upd2, mask=mask_w_upd2, other=0.0).to(tl.float32) + + k_base = k + bos * Hg * K + (i_h // (H // Hg)) * K + # get row-sliced k [64, T] + p_k_upd1 = tl.make_block_ptr( + k_base, (K, T), (1, stride_k), (0, i_t * BT), (64, BT), (0, 1) + ) + b_k_upd1 = tl.load(p_k_upd1, boundary_check=(0, 1)) + p_k_upd2 = tl.make_block_ptr( + k_base, (K, T), (1, stride_k), (64, i_t * BT), (64, BT), (0, 1) + ) + b_k_upd2 = tl.load(p_k_upd2, boundary_check=(0, 1)) + + if USE_G: + b_w_upd1 = b_w_upd1 * b_g[:, None] + b_w_upd2 = b_w_upd2 * b_g[:, None] + + # compute [64, BT] @ [BT, 64] + b_hupd_local_11 = (off_hupd_1_top == off_hupd_2_top).to(tl.float32) + b_hupd_local_22 = (off_hupd_1_top == off_hupd_2_top).to(tl.float32) + + # fp32 + if USE_G: + b_hupd_local_11 = b_hupd_local_11 * b_g_last + b_hupd_local_22 = b_hupd_local_22 * b_g_last + + b_hupd_local_11 -= tl.dot(b_k_upd1, b_w_upd1.to(b_k_upd1.dtype)) + b_hupd_local_22 -= tl.dot(b_k_upd2, b_w_upd2.to(b_k_upd2.dtype)) + b_hupd_local_12 = -tl.dot(b_k_upd1, b_w_upd2.to(b_k_upd1.dtype)).to(tl.float32) + b_hupd_local_21 = -tl.dot(b_k_upd2, b_w_upd1.to(b_k_upd2.dtype)).to(tl.float32) + + hupd_base = h_update + (boh + i_t + i_n) * H * K * K + i_h * K * K + p_hupd_11 = tl.make_block_ptr( + hupd_base, (K, K), (K, 1), (0, 0), (64, 64), (1, 0) + ) + b_hupd_11 = tl.load(p_hupd_11, boundary_check=(1, 0)) + p_hupd_21 = tl.make_block_ptr( + hupd_base, (K, K), (K, 1), (64, 0), (64, 64), (1, 0) + ) + b_hupd_21 = tl.load(p_hupd_21, boundary_check=(1, 0)) + p_hupd_12 = tl.make_block_ptr( + hupd_base, (K, K), (K, 1), (0, 64), (64, 64), (1, 0) + ) + b_hupd_12 = tl.load(p_hupd_12, boundary_check=(1, 0)) + p_hupd_22 = tl.make_block_ptr( + hupd_base, (K, K), (K, 1), (64, 64), (64, 64), (1, 0) + ) + b_hupd_22 = tl.load(p_hupd_22, boundary_check=(1, 0)) + + b_hupd11_new = tl.dot(b_hupd_local_11.to(b_hupd_11.dtype), b_hupd_11).to( + tl.float32 + ) + b_hupd11_new += tl.dot(b_hupd_local_12.to(b_hupd_21.dtype), b_hupd_21) + + b_hupd21_new = tl.dot(b_hupd_local_21.to(b_hupd_11.dtype), b_hupd_11).to( + tl.float32 + ) + b_hupd21_new += tl.dot(b_hupd_local_22.to(b_hupd_21.dtype), b_hupd_21) + + b_hupd12_new = tl.dot(b_hupd_local_11.to(b_hupd_12.dtype), b_hupd_12).to( + tl.float32 + ) + b_hupd12_new += tl.dot(b_hupd_local_12.to(b_hupd_22.dtype), b_hupd_22) + + b_hupd22_new = tl.dot(b_hupd_local_21.to(b_hupd_12.dtype), b_hupd_12).to( + tl.float32 + ) + b_hupd22_new += tl.dot(b_hupd_local_22.to(b_hupd_22.dtype), b_hupd_22) + + hupd_next = h_update + (boh + i_t + i_n + 1) * H * K * K + i_h * K * K + p_hupd_11 = tl.make_block_ptr( + hupd_next, (K, K), (K, 1), (0, 0), (64, 64), (1, 0) + ) + tl.store( + p_hupd_11, + b_hupd11_new.to(p_hupd_11.dtype.element_ty), + boundary_check=(0, 1), + ) + + p_hupd_21 = tl.make_block_ptr( + hupd_next, (K, K), (K, 1), (64, 0), (64, 64), (1, 0) + ) + tl.store( + p_hupd_21, + b_hupd21_new.to(p_hupd_21.dtype.element_ty), + boundary_check=(0, 1), + ) + + p_hupd_12 = tl.make_block_ptr( + hupd_next, (K, K), (K, 1), (0, 64), (64, 64), (1, 0) + ) + tl.store( + p_hupd_12, + b_hupd12_new.to(p_hupd_12.dtype.element_ty), + boundary_check=(0, 1), + ) + + p_hupd_22 = tl.make_block_ptr( + hupd_next, (K, K), (K, 1), (64, 64), (64, 64), (1, 0) + ) + tl.store( + p_hupd_22, + b_hupd22_new.to(p_hupd_22.dtype.element_ty), + boundary_check=(0, 1), + ) + + +def launch(k, w, g): + b, t, hg, kk = k.shape + h = g.shape[1] + bt = 64 + nt = triton.cdiv(t, bt) + h_update = torch.zeros((b, nt + b, h, kk, kk), device=k.device, dtype=torch.float32) + eye = torch.eye(kk, device=k.device, dtype=torch.float32) + h_update[:, 0, :, :, :] = eye + cu_seqlens = torch.arange(0, (b + 1) * t, t, device=k.device, dtype=torch.int64) + chunk_offsets = torch.arange(0, b * nt, nt, device=k.device, dtype=torch.int32) + chunk_gated_delta_rule_fwd_kernel_hupdate_blockdim64[(nt, b * h)]( + k, + w, + g, + cu_seqlens, + chunk_offsets, + h_update, + t, + h, + hg, + kk, + ) + return h_update + + +def main(): + args = parse_args() + init(args.seed) + dtype = dtype_from_name(args.dtype) + for b, t, h, hg, kk in case_values( + ((1, 64, 1, 1, 128), (2, 128, 4, 2, 128), (4, 256, 8, 4, 128)) + ): + k = torch.randn(b, t, hg, kk, device=args.device, dtype=dtype) + w = torch.randn(b, t, hg, kk, device=args.device, dtype=dtype) + g = torch.randn(b, h, t, device=args.device, dtype=torch.float32) + launch(k, w, g) + triton_ms, _ = bench(lambda: launch(k, w, g), args.warmup, args.repeat) + report_smoke( + f"chunk_gated_delta_rule_fwd_kernel_hupdate_blockdim64 shape=({b}, {t}, {h}, {hg}, {kk})", + triton_ms, + MODELS, + ) + + +def test_benchmark(): + main() + + +if __name__ == "__main__": + main() diff --git a/test/ascend/autotune/regression_cases/fla_cumsum_case.py b/test/ascend/autotune/regression_cases/fla_cumsum_case.py new file mode 100644 index 00000000..51e7e599 --- /dev/null +++ b/test/ascend/autotune/regression_cases/fla_cumsum_case.py @@ -0,0 +1,185 @@ +import torch +import triton +import triton.language as tl +from common import ( + assert_close, + bench, + case_values, + dtype_from_name, + init, + parse_args, + report, +) + +MODELS = "FLA chunk local cumsum exact-copy Triton kernel benchmark" + + +@triton.autotune( + configs=[], + key=[], + hints={"search_params": {"params": ["NUM_CHUNKS"]}}, +) +@triton.jit(do_not_specialize=["T"]) +def chunk_local_cumsum_scalar_kernel( + s, + o, + scale, + cu_seqlens, + chunk_indices, + T, + H: tl.constexpr, + REVERSE: tl.constexpr, + HAS_SCALE: tl.constexpr, + IS_VARLEN: tl.constexpr, + HEAD_FIRST: tl.constexpr, + CHUNK_SIZE: tl.constexpr = 64, + NUM_CHUNKS: tl.constexpr = 1, +): + i_block, i_b = tl.program_id(0), tl.program_id(1) + BLOCK_T: tl.constexpr = NUM_CHUNKS * CHUNK_SIZE + + if IS_VARLEN: + i_s, i_block = ( + tl.load(chunk_indices + i_block * 2).to(tl.int32), + tl.load(chunk_indices + i_block * 2 + 1).to(tl.int32), + ) + bos, eos = tl.load(cu_seqlens + i_s).to(tl.int32), tl.load( + cu_seqlens + i_s + 1 + ).to(tl.int32) + T = eos - bos + else: + bos, eos = i_b * T, i_b * T + T + + if HEAD_FIRST: + ptr_s = tl.make_block_ptr( + s + bos * H, (H, T), (T, 1), (0, i_block * BLOCK_T), (H, BLOCK_T), (1, 0) + ) + ptr_o = tl.make_block_ptr( + o + bos * H, (H, T), (T, 1), (0, i_block * BLOCK_T), (H, BLOCK_T), (1, 0) + ) + b_s = tl.load(ptr_s, boundary_check=(0,)).to(tl.float32) + b_s = tl.reshape(b_s, (H, NUM_CHUNKS, CHUNK_SIZE)) + b_s = tl.trans(b_s, (2, 0, 1)) + b_o = tl.cumsum(b_s, axis=0, reverse=REVERSE) + if HAS_SCALE: + b_o *= scale + b_o = tl.trans(b_o, (2, 0, 1)) + b_o = tl.reshape(b_o, (H, BLOCK_T)) + else: + ptr_s = tl.make_block_ptr( + s + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0) + ) + ptr_o = tl.make_block_ptr( + o + bos * H, (T, H), (H, 1), (i_block * BLOCK_T, 0), (BLOCK_T, H), (1, 0) + ) + b_s = tl.load(ptr_s, boundary_check=(0,)).to(tl.float32) + b_s = tl.reshape(b_s, (NUM_CHUNKS, CHUNK_SIZE, H)) + b_s = tl.trans(b_s, (1, 0, 2)) + b_o = tl.cumsum(b_s, axis=0, reverse=REVERSE) + if HAS_SCALE: + b_o *= scale + b_o = tl.trans(b_o, (1, 0, 2)) + b_o = tl.reshape(b_o, (BLOCK_T, H)) + + tl.store(ptr_o, b_o.to(s.dtype.element_ty), boundary_check=(0,)) + return + + +def run_triton( + g, chunk_size, reverse=False, scale=None, head_first=False, output_dtype=None +): + if head_first: + B, H, T = g.shape + else: + B, T, H = g.shape + assert chunk_size == 2 ** ( + chunk_size.bit_length() - 1 + ), "chunk_size must be a power of 2" + g_org, g_out = g, torch.empty_like(g, dtype=output_dtype or g.dtype) + grid = lambda meta: (triton.cdiv(T, meta["NUM_CHUNKS"] * chunk_size), B) + chunk_local_cumsum_scalar_kernel[grid]( + s=g_org, + o=g_out, + scale=scale, + cu_seqlens=None, + chunk_indices=None, + T=T, + H=H, + CHUNK_SIZE=chunk_size, + HEAD_FIRST=head_first, + REVERSE=reverse, + HAS_SCALE=scale is not None, + IS_VARLEN=False, + num_warps=8, + num_stages=3, + ) + return g_out + + +def torch_ref(g, chunk_size, reverse=False, scale=None, head_first=False): + if head_first: + b, h, t = g.shape + y = g.float().reshape(b, h, t // chunk_size, chunk_size) + y = ( + torch.flip(torch.cumsum(torch.flip(y, dims=[-1]), dim=-1), dims=[-1]) + if reverse + else torch.cumsum(y, dim=-1) + ) + y = y.reshape_as(g) + else: + b, t, h = g.shape + y = g.float().reshape(b, t // chunk_size, chunk_size, h) + y = ( + torch.flip(torch.cumsum(torch.flip(y, dims=[2]), dim=2), dims=[2]) + if reverse + else torch.cumsum(y, dim=2) + ) + y = y.reshape_as(g) + if scale is not None: + y = y * scale + return y.to(g.dtype) + + +def main(): + args = parse_args() + init(args.seed) + dtype = dtype_from_name(args.dtype) + scale = 0.5 + for B, T, H, chunk_size in case_values( + ((2, 1024, 32, 64), (4, 2048, 64, 64), (8, 4096, 64, 64)) + ): + g = torch.randn(B, T, H, dtype=dtype, device=args.device) + tri = run_triton( + g, + chunk_size, + reverse=False, + scale=scale, + head_first=False, + output_dtype=dtype, + ) + ref = torch_ref(g, chunk_size, reverse=False, scale=scale, head_first=False) + assert_close("chunk_local_cumsum_scalar_kernel", tri, ref, args.rtol, args.atol) + torch_ms, _ = bench( + lambda: torch_ref(g, chunk_size, False, scale, False), + args.warmup, + args.repeat, + ) + triton_ms, _ = bench( + lambda: run_triton(g, chunk_size, False, scale, False, dtype), + args.warmup, + args.repeat, + ) + report( + f"chunk_local_cumsum_scalar_kernel shape=({B}, {T}, {H}, {chunk_size})", + torch_ms, + triton_ms, + MODELS, + ) + + +def test_benchmark(): + main() + + +if __name__ == "__main__": + main() diff --git a/test/ascend/autotune/regression_cases/gdn_chunk_meta_case.py b/test/ascend/autotune/regression_cases/gdn_chunk_meta_case.py new file mode 100644 index 00000000..9bd8c5a3 --- /dev/null +++ b/test/ascend/autotune/regression_cases/gdn_chunk_meta_case.py @@ -0,0 +1,173 @@ +import torch +import triton +import triton.language as tl +from common import assert_close, bench, case_values, init, parse_args, report + +MODELS = "Qwen3-Next and Qwen3.5 GDN metadata paths; kernel coverage does not depend on model availability" + + +@triton.autotune( + configs=[], + key=[], + hints={"search_params": {"params": ["BLOCK_SIZE"]}}, +) +@triton.jit +def _build_chunk_counts_kernel( + cu_seqlens_ptr, + chunk_counts_ptr, + num_seqs: tl.constexpr, + chunk_size, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_seqs + + bos = tl.load(cu_seqlens_ptr + offsets, mask=mask, other=0).to(tl.int32) + eos = tl.load(cu_seqlens_ptr + offsets + 1, mask=mask, other=0).to(tl.int32) + seq_lens = eos - bos + chunk_counts = (seq_lens + chunk_size - 1) // chunk_size + + tl.store(chunk_counts_ptr + offsets, chunk_counts, mask=mask) + + +@triton.autotune( + configs=[], + key=[], + hints={"search_params": {"params": ["BLOCK_SIZE"]}}, +) +@triton.jit +def _build_chunk_offsets_kernel( + chunk_counts_ptr, + out_offsets_ptr, + num_seqs: tl.constexpr, + ADD_ONE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets <= num_seqs + prefix = tl.zeros([BLOCK_SIZE], dtype=tl.int32) + + for seq_idx in range(0, num_seqs): + chunk_count = tl.load( + chunk_counts_ptr + seq_idx, mask=seq_idx < num_seqs, other=0 + ).to(tl.int32) + prefix += tl.where(mask & (offsets > seq_idx), chunk_count + ADD_ONE, 0) + + tl.store( + out_offsets_ptr + offsets, + prefix.to(out_offsets_ptr.dtype.element_ty), + mask=mask, + ) + + +@triton.autotune( + configs=[], + key=[], + hints={"search_params": {"params": ["BLOCK_SIZE"]}}, +) +@triton.jit +def _build_final_chunk_indices_kernel( + update_chunk_offsets_ptr, + out_final_chunk_indices_ptr, + num_seqs: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_seqs + final_indices = ( + tl.load(update_chunk_offsets_ptr + offsets + 1, mask=mask, other=0).to(tl.int32) + - 1 + ) + tl.store( + out_final_chunk_indices_ptr + offsets, + final_indices.to(out_final_chunk_indices_ptr.dtype.element_ty), + mask=mask, + ) + + +def run_triton(cu_seqlens, chunk_size): + num_seqs = cu_seqlens.numel() - 1 + chunk_counts = torch.empty(num_seqs, dtype=torch.int32, device=cu_seqlens.device) + chunk_offsets = torch.empty( + num_seqs + 1, dtype=torch.int32, device=cu_seqlens.device + ) + update_offsets = torch.empty( + num_seqs + 1, dtype=torch.int32, device=cu_seqlens.device + ) + final_indices = torch.empty(num_seqs, dtype=torch.int32, device=cu_seqlens.device) + + grid_counts = lambda meta: (triton.cdiv(num_seqs, meta["BLOCK_SIZE"]),) + _build_chunk_counts_kernel[grid_counts]( + cu_seqlens, chunk_counts, num_seqs, chunk_size + ) + + grid_offsets = lambda meta: (triton.cdiv(num_seqs + 1, meta["BLOCK_SIZE"]),) + _build_chunk_offsets_kernel[grid_offsets]( + chunk_counts, chunk_offsets, num_seqs, ADD_ONE=0 + ) + _build_chunk_offsets_kernel[grid_offsets]( + chunk_counts, update_offsets, num_seqs, ADD_ONE=1 + ) + + grid_final = lambda meta: (triton.cdiv(num_seqs, meta["BLOCK_SIZE"]),) + _build_final_chunk_indices_kernel[grid_final]( + update_offsets, final_indices, num_seqs + ) + return chunk_counts, chunk_offsets, update_offsets, final_indices + + +def torch_ref(cu_seqlens, chunk_size): + seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] + counts = torch.div(seq_lens + chunk_size - 1, chunk_size, rounding_mode="floor").to( + torch.int32 + ) + chunk_offsets = torch.empty( + cu_seqlens.numel(), dtype=torch.int32, device=cu_seqlens.device + ) + update_offsets = torch.empty_like(chunk_offsets) + chunk_offsets[0] = 0 + update_offsets[0] = 0 + chunk_offsets[1:] = torch.cumsum(counts, dim=0) + update_offsets[1:] = torch.cumsum(counts + 1, dim=0) + final_indices = update_offsets[1:] - 1 + return counts, chunk_offsets, update_offsets, final_indices + + +def main(): + args = parse_args() + init(args.seed) + for num_seqs, max_len, chunk_size in case_values( + ((256, 2048, 64), (1024, 4096, 64), (2048, 8192, 128)) + ): + lens = torch.randint( + 1, max_len, (num_seqs,), dtype=torch.int32, device=args.device + ) + cu_seqlens = torch.zeros(num_seqs + 1, dtype=torch.int32, device=args.device) + cu_seqlens[1:] = torch.cumsum(lens, dim=0) + tri = run_triton(cu_seqlens, chunk_size) + ref = torch_ref(cu_seqlens, chunk_size) + for actual, expected in zip(tri, ref): + assert_close("gdn_chunk_meta", actual, expected, 0, 0) + torch_ms, _ = bench( + lambda: torch_ref(cu_seqlens, chunk_size), args.warmup, args.repeat + ) + triton_ms, _ = bench( + lambda: run_triton(cu_seqlens, chunk_size), args.warmup, args.repeat + ) + report( + f"gdn_chunk_meta_kernels shape=({num_seqs}, {max_len}, {chunk_size})", + torch_ms, + triton_ms, + MODELS, + ) + + +def test_benchmark(): + main() + + +if __name__ == "__main__": + main() diff --git a/test/ascend/autotune/regression_config.py b/test/ascend/autotune/regression_config.py new file mode 100644 index 00000000..9a88689e --- /dev/null +++ b/test/ascend/autotune/regression_config.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import json +import os +import shutil +import sys +from pathlib import Path + + +THIS_DIR = Path(__file__).resolve().parent +REPO_ROOT = THIS_DIR.parents[2] +CASE_DIR = THIS_DIR / "regression_cases" +GOLDEN_PATH = THIS_DIR / "autotune_regression_goldens.json" + +DEFAULT_BENCH_WARMUP = 1 +DEFAULT_BENCH_REPEAT = 3 +FA2_SELECTED_CASES = "1,4" + +PARALLEL_DEVICES = { + "fa2": 1, + "fla_chunk_delta_hupdate": 2, + "fla_cumsum": 3, + "gdn_chunk_meta": 6, + "batch_invariant_mean": 4, +} + + +def load_regression_specs() -> dict: + return json.loads(GOLDEN_PATH.read_text()) + + +def case_env( + base_env: dict[str, str], + cache_root: Path, + name: str, + device: int | None = None, +) -> dict[str, str]: + env = dict(base_env) + env.update( + { + "PYTHONPATH": f"{REPO_ROOT}:{env.get('PYTHONPATH', '')}", + "PYTHONUNBUFFERED": "1", + "TRITON_PRINT_AUTOTUNING": "1", + "TRITON_PRINT_AUTOTUNING_TIMINGS": "1", + "TRITON_CACHE_DIR": str(cache_root), + "DLC_AUTOTUNE_BENCH_WARMUP": str(DEFAULT_BENCH_WARMUP), + "DLC_AUTOTUNE_BENCH_ACTIVE": str(DEFAULT_BENCH_REPEAT), + } + ) + if device is not None: + env["ASCEND_RT_VISIBLE_DEVICES"] = str(device) + if name == "fa2": + env["TEST_FA2_CASES"] = FA2_SELECTED_CASES + return env + + +def prepare_clean_cache(cache_root: Path) -> None: + shutil.rmtree(cache_root, ignore_errors=True) + cache_root.mkdir(parents=True, exist_ok=True) + + +def case_cmd(script: Path) -> list[str]: + return [ + sys.executable, + str(script), + "--warmup", + str(DEFAULT_BENCH_WARMUP), + "--repeat", + str(DEFAULT_BENCH_REPEAT), + ] diff --git a/test/ascend/autotune/run_autotune_regression_parallel.py b/test/ascend/autotune/run_autotune_regression_parallel.py new file mode 100644 index 00000000..7df48096 --- /dev/null +++ b/test/ascend/autotune/run_autotune_regression_parallel.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import argparse +import os +import subprocess +import time +from pathlib import Path + +from regression_config import ( + CASE_DIR, + PARALLEL_DEVICES, + case_cmd, + case_env, + load_regression_specs, + prepare_clean_cache, +) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument( + "--log-root", + default="/tmp/dlc_autotune_parallel_devices_20260629_trust_region_final", + ) + parser.add_argument("--poll-seconds", type=float, default=30.0) + args = parser.parse_args() + + root = Path(args.log_root) + root.mkdir(parents=True, exist_ok=True) + + specs = load_regression_specs() + procs = [] + start = time.time() + for name, device in PARALLEL_DEVICES.items(): + spec = specs[name] + case_root = root / name + case_root.mkdir(parents=True, exist_ok=True) + cache_root = case_root / "triton_cache" + prepare_clean_cache(cache_root) + log_path = case_root / "run.log" + log_file = log_path.open("w") + cmd = case_cmd(CASE_DIR / spec["script"]) + proc = subprocess.Popen( + cmd, + cwd=str(CASE_DIR), + env=case_env(os.environ, cache_root, name, device), + stdout=log_file, + stderr=subprocess.STDOUT, + ) + procs.append( + { + "name": name, + "proc": proc, + "log_file": log_file, + "log_path": log_path, + "timeout_s": int(spec.get("timeout_s", 600)), + "start": time.time(), + "device": device, + } + ) + print(f"started {name} device={device} pid={proc.pid}", flush=True) + + remaining = {item["name"] for item in procs} + worst_rc = 0 + while remaining: + for item in procs: + name = item["name"] + if name not in remaining: + continue + proc = item["proc"] + rc = proc.poll() + elapsed = time.time() - item["start"] + if rc is None and elapsed > item["timeout_s"]: + proc.kill() + rc = proc.wait() + item["log_file"].write(f"\nTIMEOUT after {item['timeout_s']}s\n") + if rc is not None: + item["log_file"].close() + remaining.remove(name) + worst_rc = worst_rc or rc + print( + f"finished {name} rc={rc} elapsed={elapsed:.1f}s " + f"log={item['log_path']}", + flush=True, + ) + if remaining: + print( + "running " + + ", ".join(sorted(remaining)) + + f" elapsed={time.time() - start:.1f}s", + flush=True, + ) + time.sleep(args.poll_seconds) + + print(f"all_done root={root}", flush=True) + return worst_rc + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/test/ascend/autotune/test_autotune_doc_e2e.py b/test/ascend/autotune/test_autotune_doc_e2e.py deleted file mode 100644 index 4150b82e..00000000 --- a/test/ascend/autotune/test_autotune_doc_e2e.py +++ /dev/null @@ -1,275 +0,0 @@ -import os - -import pytest -import torch -import torch_npu -import triton -import triton.language as tl - -import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 - install proxy before decorators - -os.environ.setdefault("TRITON_AUTOTUNE_PARALLEL_COMPILE", "0") - - -@triton.autotune( - configs=[ - triton.Config({"XS": 128, "multibuffer": True}), - triton.Config({"XS": 1024, "multibuffer": True}), - triton.Config({"XS": 1024, "multibuffer": False}), - ], - key=["numel"], -) -@triton.jit -def _explicit_config_exp_add_kernel(out_ptr, x_ptr, y_ptr, numel, XS: tl.constexpr): - offsets = tl.program_id(0) * XS + tl.arange(0, XS) - mask = offsets < numel - x = tl.load(x_ptr + offsets, mask=mask, other=0.0) - y = tl.load(y_ptr + offsets, mask=mask, other=0.0) - out = tl.full((XS,), 0.0, tl.float32) - for i in range(8): - out = tl.exp(x) + y + i - tl.store(out_ptr + offsets, out, mask=mask) - - -def _explicit_config_exp_add(x, y): - out = torch.empty_like(x) - numel = out.numel() - grid = lambda meta: (triton.cdiv(numel, meta["XS"]), 1, 1) - _explicit_config_exp_add_kernel[grid](out, x, y, numel) - return out - - -@triton.autotune(configs=[], key=["n_elements"]) -@triton.jit -def _auto_tiling_add_kernel( - x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr -): - offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask, other=0.0) - y = tl.load(y_ptr + offsets, mask=mask, other=0.0) - tl.store(out_ptr + offsets, x + y, mask=mask) - - -def _auto_tiling_add(x, y): - out = torch.empty_like(x) - n_elements = out.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) - _auto_tiling_add_kernel[grid](x, y, out, n_elements) - return out - - -@triton.autotune( - configs=[], - key=["n_elements"], - hints={"compile_options": "vector"}, -) -@triton.jit -def _auto_tiling_compile_options_add_kernel( - x_ptr, - y_ptr, - out_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, -): - offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask, other=0.0) - y = tl.load(y_ptr + offsets, mask=mask, other=0.0) - tl.store(out_ptr + offsets, x + y * 3.0, mask=mask) - - -def _auto_tiling_compile_options_add(x, y): - out = torch.empty_like(x) - n_elements = out.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) - _auto_tiling_compile_options_add_kernel[grid](x, y, out, n_elements) - return out - - -@triton.autotune( - configs=[], - key={"x": "n_elements"}, - hints={ - "split_params": {"x": "BLOCK_SIZE"}, - "tiling_params": {"x": "BLOCK_SIZE_SUB"}, - "low_dim_axes": ["x"], - "reduction_axes": [], - }, -) -@triton.jit -def _hinted_tiling_add_kernel( - x_ptr, - y_ptr, - out_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, - BLOCK_SIZE_SUB: tl.constexpr, -): - block_start = tl.program_id(0) * BLOCK_SIZE - sub_offsets = tl.arange(0, BLOCK_SIZE_SUB) - loops = (BLOCK_SIZE + BLOCK_SIZE_SUB - 1) // BLOCK_SIZE_SUB - for loop in range(loops): - offsets = block_start + loop * BLOCK_SIZE_SUB + sub_offsets - mask = offsets < min(block_start + BLOCK_SIZE, n_elements) - x = tl.load(x_ptr + offsets, mask=mask, other=0.0) - y = tl.load(y_ptr + offsets, mask=mask, other=0.0) - tl.store(out_ptr + offsets, x + y, mask=mask) - - -def _hinted_tiling_add(x, y): - out = torch.empty_like(x) - n_elements = out.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) - _hinted_tiling_add_kernel[grid](x, y, out, n_elements) - return out - - -@triton.autotune( - configs=[triton.Config({"BLOCK_SIZE": 128, "multibuffer": False})], - key=["n_elements"], - hints={"auto_gen_config": True}, -) -@triton.jit -def _auto_and_user_config_add_kernel( - x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr -): - offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask, other=0.0) - y = tl.load(y_ptr + offsets, mask=mask, other=0.0) - tl.store(out_ptr + offsets, x - y, mask=mask) - - -def _auto_and_user_config_add(x, y): - out = torch.empty_like(x) - n_elements = out.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) - _auto_and_user_config_add_kernel[grid](x, y, out, n_elements) - return out - - -@triton.max_autotune( - configs=[triton.Config({"BLOCK_SIZE": 256})], - key=["n_elements"], - kernel_type="vector", - num_stages=[1, 2], - enable_ubuf_saving=[True, False], -) -@triton.jit -def _max_autotune_vector_kernel( - x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr -): - offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask, other=0.0) - y = tl.load(y_ptr + offsets, mask=mask, other=0.0) - tl.store(out_ptr + offsets, x * 2.0 + y, mask=mask) - - -def _max_autotune_vector(x, y): - out = torch.empty_like(x) - n_elements = out.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) - _max_autotune_vector_kernel[grid](x, y, out, n_elements) - return out - - -@triton.max_autotune( - configs=[triton.Config({"BLOCK_SIZE": 128})], - key=["n_elements"], - kernel_type="vector", - enable_ubuf_saving=[True, False], -) -@triton.jit -def _max_autotune_default_stage_kernel( - x_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr -): - offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask, other=0.0) - tl.store(out_ptr + offsets, x + 1.0, mask=mask) - - -def _max_autotune_default_stage(x): - out = torch.empty_like(x) - n_elements = out.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) - _max_autotune_default_stage_kernel[grid](x, out, n_elements) - return out - - -@pytest.mark.autotune -def test_community_autotune_explicit_configs_e2e(): - x = torch.randn(4096, dtype=torch.float32, device="npu") - y = torch.randn(4096, dtype=torch.float32, device="npu") - - actual = _explicit_config_exp_add(x, y) - expected = torch.exp(x) + y + 7 - - torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) - - -@pytest.mark.autotune -def test_advanced_autotune_empty_configs_auto_tiling_e2e(): - x = torch.randn(4096, dtype=torch.float32, device="npu") - y = torch.randn(4096, dtype=torch.float32, device="npu") - - actual = _auto_tiling_add(x, y) - expected = x + y - - torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) - - -@pytest.mark.autotune -def test_advanced_autotune_auto_tiling_compile_options_e2e(): - x = torch.randn(4096, dtype=torch.float32, device="npu") - y = torch.randn(4096, dtype=torch.float32, device="npu") - - actual = _auto_tiling_compile_options_add(x, y) - expected = x + y * 3.0 - - torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) - - -@pytest.mark.autotune -def test_advanced_autotune_hints_dict_key_e2e(): - x = torch.randn(4096, dtype=torch.float32, device="npu") - y = torch.randn(4096, dtype=torch.float32, device="npu") - - actual = _hinted_tiling_add(x, y) - expected = x + y - - torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) - - -@pytest.mark.autotune -def test_advanced_autotune_user_configs_merge_auto_configs_e2e(): - x = torch.randn(4096, dtype=torch.float32, device="npu") - y = torch.randn(4096, dtype=torch.float32, device="npu") - - actual = _auto_and_user_config_add(x, y) - expected = x - y - - torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) - - -@pytest.mark.autotune -def test_max_autotune_vector_expanded_configs_e2e(): - x = torch.randn(4096, dtype=torch.float32, device="npu") - y = torch.randn(4096, dtype=torch.float32, device="npu") - - actual = _max_autotune_vector(x, y) - expected = x * 2.0 + y - - torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) - - -@pytest.mark.autotune -def test_max_autotune_uses_ascend_default_num_stages_e2e(): - x = torch.randn(4096, dtype=torch.float32, device="npu") - - actual = _max_autotune_default_stage(x) - expected = x + 1.0 - - torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) diff --git a/test/ascend/autotune/test_autotune_param_valid.py b/test/ascend/autotune/test_autotune_param_valid.py index e4fdca42..a4e023ee 100644 --- a/test/ascend/autotune/test_autotune_param_valid.py +++ b/test/ascend/autotune/test_autotune_param_valid.py @@ -1,5 +1,3 @@ -import os - import pytest import torch import torch_npu @@ -29,7 +27,7 @@ def add_kernel( offset = tl.program_id(0) * BLOCK_SIZE loops1 = (BLOCK_SIZE + BLOCK_SIZE_SUB - 1) // BLOCK_SIZE_SUB for loop in range(0, loops1): - x0 = offset + loop * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE_SUB) + x0 = offset + loop * BLOCK_SIZE_SUB + tl.arange(0, BLOCK_SIZE_SUB) mask = x0 < n_elements x = tl.load(x_ptr + x0, mask) y = tl.load(y_ptr + x0, mask) @@ -68,91 +66,25 @@ def test_add(size: int): @pytest.mark.autotune def test_add_no_reduction_axes(): - try: - - @triton.autotune( - configs=[], - key={"x": "n_elements"}, - hints={ - "split_params": {"x": "BLOCK_SIZE"}, - "tiling_params": {"x": "BLOCK_SIZE_SUB"}, - "low_dim_axes": ["x"], - }, - ) - @triton.jit - def add_kernel_exception(): - pass - - except ValueError as e: - assert "reduction_axes must be a list" in str(e) - - -@pytest.mark.autotune -def test_add_no_low_dim_axes(): - try: - - @triton.autotune( - configs=[], - key={"x": "n_elements"}, - hints={ - "split_params": {"x": "BLOCK_SIZE"}, - "tiling_params": {"x": "BLOCK_SIZE_SUB"}, - "reduction_axes": [], - }, - ) - @triton.jit - def add_kernel_exception(): - pass - - except ValueError as e: - assert "low_dim_axes must be a list" in str(e) - - -@pytest.mark.autotune -def test_add_no_tiling_params(): - try: + with pytest.raises(ValueError, match="reduction_axes must be a list"): @triton.autotune( configs=[], key={"x": "n_elements"}, hints={ "split_params": {"x": "BLOCK_SIZE"}, - "low_dim_axes": ["x"], - "reduction_axes": [], - }, - ) - @triton.jit - def add_kernel_exception(): - pass - - except ValueError as e: - assert "tiling_params must be a dict" in str(e) - - -@pytest.mark.autotune -def test_add_no_split_params(): - try: - - @triton.autotune( - configs=[], - key={"x": "n_elements"}, - hints={ "tiling_params": {"x": "BLOCK_SIZE_SUB"}, "low_dim_axes": ["x"], - "reduction_axes": [], }, ) @triton.jit def add_kernel_exception(): pass - except ValueError as e: - assert "split_params must be a dict" in str(e) - @pytest.mark.autotune def test_add_no_keyname(): - try: + with pytest.raises(ValueError, match="All keys in 'key' must be valid axis names"): @triton.autotune( configs=[], @@ -166,6 +98,3 @@ def test_add_no_keyname(): @triton.jit def add_kernel_exception(): pass - - except ValueError as e: - assert "All keys in 'key' must be valid axis names" in str(e) diff --git a/test/ascend/autotune/test_autotune_regression_ci.py b/test/ascend/autotune/test_autotune_regression_ci.py new file mode 100644 index 00000000..2f9afc70 --- /dev/null +++ b/test/ascend/autotune/test_autotune_regression_ci.py @@ -0,0 +1,349 @@ +from __future__ import annotations + +import json +import os +import re +import subprocess +from dataclasses import dataclass +from pathlib import Path + +import pytest + +from regression_config import ( + CASE_DIR, + case_cmd, + case_env, + load_regression_specs, + prepare_clean_cache, +) + + +MAX_ATTEMPTS = 3 + + +@dataclass +class CaseRun: + name: str + attempt: int + rc: int + log_path: Path + operator_kinds: list[str] + compile_options: list[str] + precise_costs: list[float] + triton_ms: list[float] + search_times: list[float] + selected_configs: list[str] + autotune_enabled_count: int + stage1_initial_count: int + stage1_selected_count: int + stage3_final_pick_count: int + fatal_log_markers: list[str] + + +def _as_float_list(value) -> list[float]: + if value is None: + return [] + if isinstance(value, list): + return [float(item) for item in value] + return [float(value)] + + +def _extend_from_payload(target: list, value): + if value is None: + return + if isinstance(value, list): + target.extend(value) + else: + target.insert(0, value) + + +def _has_npu_runtime() -> bool: + try: + import torch + import torch_npu # noqa: F401 + except Exception: + return False + return bool(hasattr(torch, "npu") and torch.npu.is_available()) + + +def _load_goldens(): + return load_regression_specs() + + +def _artifact_root(tmp_path: Path) -> Path: + configured = os.getenv("DLC_AUTOTUNE_REGRESSION_LOG_DIR") + root = Path(configured) if configured else tmp_path / "autotune_regression_logs" + root.mkdir(parents=True, exist_ok=True) + return root + + +def _parse_run(name: str, attempt: int, rc: int, log_path: Path) -> CaseRun: + text = log_path.read_text(errors="ignore") + payloads = [ + json.loads(match) + for match in re.findall(r"DLC_AUTOTUNE_REGRESSION_RESULT=(\{.*\})", text) + ] + precise_costs = [ + float(value) + for value in re.findall( + r"Stage 3 final pick .*?precise_cost=([0-9.eE+-]+)", text + ) + ] + triton_ms = [ + float(value) for value in re.findall(r"triton_ms: ([0-9.eE+-]+)", text) + ] + search_times = [ + float(value) + for value in re.findall( + r"finished after ([0-9.eE+-]+)s; best config selected", text + ) + ] + selected_configs = re.findall(r"best config selected: (.*)", text) + if payloads: + payload = payloads[-1] + _extend_from_payload(triton_ms, payload.get("triton_ms")) + if not search_times: + _extend_from_payload(search_times, payload.get("search_time")) + if not selected_configs: + _extend_from_payload(selected_configs, payload.get("best_config")) + fatal_markers = [ + marker + for marker in ( + "Traceback", + "TIMEOUT after", + "Segmentation fault", + "Aborted", + "core dumped", + ) + if marker in text + ] + return CaseRun( + name=name, + attempt=attempt, + rc=rc, + log_path=log_path, + operator_kinds=re.findall(r"operator_kind=([A-Z_]+)", text), + compile_options=re.findall(r"compile_options=([a-z0-9_]+)", text), + precise_costs=precise_costs, + triton_ms=triton_ms, + search_times=search_times, + selected_configs=selected_configs, + autotune_enabled_count=text.count("Search params autotuning: enabled"), + stage1_initial_count=text.count("Stage 1 initial"), + stage1_selected_count=text.count("Stage 1 selected shapes"), + stage3_final_pick_count=text.count("Stage 3 final pick"), + fatal_log_markers=fatal_markers, + ) + + +def _run_once(name: str, spec: dict, attempt: int, artifact_root: Path) -> CaseRun: + script = CASE_DIR / spec["script"] + assert script.is_file(), f"missing regression case script: {script}" + + case_root = artifact_root / name / f"attempt_{attempt}" + cache_root = case_root / "triton_cache" + case_root.mkdir(parents=True, exist_ok=True) + prepare_clean_cache(cache_root) + log_path = case_root / "run.log" + + env = case_env(os.environ, cache_root, name) + env["DLC_AUTOTUNE_CASE_LIMIT"] = env.get("DLC_AUTOTUNE_CASE_LIMIT", "0") + + cmd = case_cmd(script) + try: + proc = subprocess.run( + cmd, + cwd=str(CASE_DIR), + env=env, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + timeout=int(spec.get("timeout_s", 600)), + check=False, + ) + log_path.write_text(proc.stdout) + return _parse_run(name, attempt, proc.returncode, log_path) + except subprocess.TimeoutExpired as exc: + output = exc.stdout or "" + if isinstance(output, bytes): + output = output.decode(errors="ignore") + log_path.write_text(output + f"\nTIMEOUT after {spec.get('timeout_s', 600)}s\n") + return _parse_run(name, attempt, 124, log_path) + + +def _performance_failures(run: CaseRun, spec: dict) -> list[str]: + failures = [] + golden_precise = _as_float_list(spec.get("golden_precise_cost")) + if golden_precise: + if len(run.precise_costs) < len(golden_precise): + failures.append( + f"precise_cost count {len(run.precise_costs)} < golden count {len(golden_precise)}" + ) + for index, golden in enumerate(golden_precise): + if index >= len(run.precise_costs): + break + limit = golden * spec.get("max_precise_cost_ratio", 1.5) + if run.precise_costs[index] > limit: + failures.append( + f"precise_cost[{index}] {run.precise_costs[index]:.6g} > {limit:.6g}" + ) + golden_ms = _as_float_list(spec.get("golden_triton_ms")) + if golden_ms: + if len(run.triton_ms) < len(golden_ms): + failures.append( + f"triton_ms count {len(run.triton_ms)} < golden count {len(golden_ms)}" + ) + for index, golden in enumerate(golden_ms): + if index >= len(run.triton_ms): + break + limit = golden * spec.get("max_triton_ms_ratio", 2.0) + if run.triton_ms[index] > limit: + failures.append( + f"triton_ms[{index}] {run.triton_ms[index]:.6g} > {limit:.6g}" + ) + return failures + + +def _performance_score(run: CaseRun, spec: dict) -> float: + score = 0.0 + golden_precise = _as_float_list(spec.get("golden_precise_cost")) + for index, golden in enumerate(golden_precise): + if index < len(run.precise_costs) and golden > 0: + score = max(score, run.precise_costs[index] / golden) + golden_ms = _as_float_list(spec.get("golden_triton_ms")) + for index, golden in enumerate(golden_ms): + if index < len(run.triton_ms) and golden > 0: + score = max(score, run.triton_ms[index] / golden) + return score if score > 0.0 else float("inf") + + +def _max_search_time_failures(run: CaseRun, spec: dict) -> list[str]: + limit_hint = spec.get("max_search_time_s") + if limit_hint is None: + return [] + + failures = [] + limits = _as_float_list(limit_hint) + if len(limits) == 1: + limits = limits * max(1, len(run.search_times)) + if len(run.search_times) < len(limits): + failures.append( + f"search_time count {len(run.search_times)} < expected count {len(limits)}" + ) + for index, limit in enumerate(limits): + if index >= len(run.search_times): + break + if run.search_times[index] > limit: + failures.append( + f"search_time[{index}] {run.search_times[index]:.2f}s > {limit:.2f}s" + ) + return failures + + +def _log_failures(run: CaseRun, spec: dict) -> list[str]: + failures = [] + if run.fatal_log_markers: + failures.append( + f"fatal log markers found: {sorted(set(run.fatal_log_markers))}" + ) + + expected_runs = int(spec.get("expected_autotune_runs", 1)) + log_counts = { + "autotune enabled": run.autotune_enabled_count, + "Stage 1 initial": run.stage1_initial_count, + "Stage 1 selected shapes": run.stage1_selected_count, + "Stage 3 final pick": run.stage3_final_pick_count, + "best config selected": len(run.selected_configs), + } + for label, observed in log_counts.items(): + if observed < expected_runs: + failures.append(f"{label} count {observed} < expected {expected_runs}") + + failures.extend(_max_search_time_failures(run, spec)) + return failures + + +def _semantic_failures(run: CaseRun, spec: dict) -> list[str]: + failures = [] + expected_kind = spec.get("expected_operator_kind") + if expected_kind and expected_kind not in run.operator_kinds: + failures.append( + f"operator_kind {expected_kind!r} not found; observed={run.operator_kinds}" + ) + expected_compile = spec.get("expected_compile_options") + if expected_compile and expected_compile not in run.compile_options: + failures.append( + f"compile_options {expected_compile!r} not found; observed={run.compile_options}" + ) + return failures + + +def _format_attempts(attempts: list[CaseRun]) -> str: + lines = [] + for run in attempts: + lines.append( + "attempt={attempt}, rc={rc}, precise={precise}, triton_ms={triton}, " + "search_times={search}, selected_configs={selected}, log={log}".format( + attempt=run.attempt, + rc=run.rc, + precise=run.precise_costs, + triton=run.triton_ms, + search=run.search_times[:3], + selected=len(run.selected_configs), + log=run.log_path, + ) + ) + return "\n".join(lines) + + +def _case_names(): + return list(_load_goldens()) + + +@pytest.mark.parametrize("case_name", _case_names()) +def test_autotune_regression_case(case_name, tmp_path): + if not _has_npu_runtime(): + pytest.skip("Ascend NPU runtime is not available") + + goldens = _load_goldens() + spec = goldens[case_name] + artifact_root = _artifact_root(tmp_path) + attempts = [] + + for attempt in range(1, MAX_ATTEMPTS + 1): + run = _run_once(case_name, spec, attempt, artifact_root) + attempts.append(run) + if run.rc != 0: + break + if _semantic_failures(run, spec): + break + if _log_failures(run, spec): + break + if not _performance_failures(run, spec): + break + + successful_runs = [run for run in attempts if run.rc == 0] + passing_runs = [ + run + for run in successful_runs + if not _semantic_failures(run, spec) + and not _log_failures(run, spec) + and not _performance_failures(run, spec) + ] + best = ( + passing_runs[0] + if passing_runs + else min( + successful_runs, + key=lambda run: _performance_score(run, spec), + default=attempts[-1], + ) + ) + + failures = [] + if best.rc != 0: + failures.append(f"process failed with rc={best.rc}") + failures.extend(_semantic_failures(best, spec)) + failures.extend(_log_failures(best, spec)) + failures.extend(_performance_failures(best, spec)) + assert not failures, "\n".join(failures + ["", _format_attempts(attempts)]) diff --git a/test/ascend/autotune/test_benchmark_strategy.py b/test/ascend/autotune/test_benchmark_strategy.py new file mode 100644 index 00000000..547d6971 --- /dev/null +++ b/test/ascend/autotune/test_benchmark_strategy.py @@ -0,0 +1,38 @@ +from backend.ascend_autotune_runtime.measurement_strategy import ( + NpuProfilerBenchStrategy, +) + + +def test_npu_bench_strategy_accepts_single_scalar_timing(monkeypatch): + import backend.testing as testing + + config = object() + + def fake_do_bench_npu(funcs, clear_l2_cache=False): + assert len(funcs) == 1 + assert clear_l2_cache is False + return 0.123 + + monkeypatch.setattr(testing, "do_bench_npu", fake_do_bench_npu) + + assert NpuProfilerBenchStrategy().bench({config: lambda: None}) == {config: 0.123} + + +def test_npu_bench_strategy_accepts_multi_config_timing(monkeypatch): + import backend.testing as testing + + configs = [object(), object()] + + def fake_do_bench_npu(funcs, clear_l2_cache=False): + assert len(funcs) == 2 + assert clear_l2_cache is False + return [0.2, 0.1] + + monkeypatch.setattr(testing, "do_bench_npu", fake_do_bench_npu) + + assert NpuProfilerBenchStrategy().bench( + {configs[0]: lambda: None, configs[1]: lambda: None} + ) == { + configs[0]: 0.2, + configs[1]: 0.1, + } diff --git a/test/ascend/autotune/test_common.py b/test/ascend/autotune/test_common.py deleted file mode 100644 index 350ba874..00000000 --- a/test/ascend/autotune/test_common.py +++ /dev/null @@ -1,83 +0,0 @@ -import unittest.mock as mock -import pytest -import torch - - -def MockAutoTilingTunerRun(self, *args, **kwargs): - self.nargs = dict(zip(self.arg_names, args)) - - # generate key - all_args = {**self.nargs, **kwargs} - try: - self._autoparse_axis_params(all_args) - except ValueError as e: - if "Missing required arguments" in str(e): - pass - else: - raise - return { - "keys": self.keys, - "split_params": self.split_params, - "tiling_params": self.tiling_params, - "low_dim_axes": self.low_dim_axes, - "reduction_axes": self.reduction_axes, - "persistent_reduction": self.persistent_reduction, - } - - -def check_axes_parse_res(act: dict, ref: dict): - ref_keys = ref["keys"] - act_keys = act["keys"] - - assert set(ref_keys.values()) == set( - act_keys.values() - ), f"Semantic dimensions mismatch: ref={set(ref_keys.values())}, act={set(act_keys.values())}" - - def normalize_param_dict(param_dict: dict, sym_to_sem: dict) -> dict: - return {sym_to_sem[sym]: value for sym, value in param_dict.items()} - - ref_split = normalize_param_dict(ref["split_params"], ref_keys) - act_split = normalize_param_dict(act["split_params"], act_keys) - - ref_tiling = normalize_param_dict(ref["tiling_params"], ref_keys) - act_tiling = normalize_param_dict(act["tiling_params"], act_keys) - - def normalize_axis_list(axis_list: list, sym_to_sem: dict) -> list: - return sorted(sym_to_sem[sym] for sym in axis_list) - - ref_low = normalize_axis_list(ref["low_dim_axes"], ref_keys) - act_low = normalize_axis_list(act["low_dim_axes"], act_keys) - - ref_red = normalize_axis_list(ref["reduction_axes"], ref_keys) - act_red = normalize_axis_list(act["reduction_axes"], act_keys) - - assert ref_split == act_split, f"split_params mismatch: {ref_split} vs {act_split}" - assert ( - ref_tiling == act_tiling - ), f"tiling_params mismatch: {ref_tiling} vs {act_tiling}" - assert ref_low == act_low, f"low_dim_axes mismatch: {ref_low} vs {act_low}" - assert ref_red == act_red, f"reduction_axes mismatch: {ref_red} vs {act_red}" - - -@pytest.fixture -def mock_autotuner(): - with mock.patch( - "triton.backends.dicp_triton.ascend_autotune_runtime.autotuner.AutoTilingTuner.run", - new=MockAutoTilingTunerRun, - ): - yield - - -def generate_tensor(shape, dtype): - if dtype == "float32" or dtype == "float16" or dtype == "bfloat16": - return torch.randn(size=shape, dtype=eval("torch." + dtype)) - elif dtype == "int32" or dtype == "int64" or dtype == "int16": - return torch.randint(low=0, high=2000, size=shape, dtype=eval("torch." + dtype)) - elif dtype == "int8": - return torch.randint(low=0, high=127, size=shape, dtype=eval("torch." + dtype)) - elif dtype == "bool": - return torch.randint(low=0, high=2, size=shape).bool() - elif dtype == "uint8": - return torch.randint(low=0, high=255, size=shape, dtype=torch.uint8) - else: - raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) diff --git a/test/ascend/autotune/test_compile_options.py b/test/ascend/autotune/test_compile_options.py index 9a3d18eb..ab5e0ac4 100644 --- a/test/ascend/autotune/test_compile_options.py +++ b/test/ascend/autotune/test_compile_options.py @@ -1,406 +1,133 @@ import pytest import triton -from backend.ascend_autotune_runtime.compile_options import ( - expand_compile_option_configs, - format_compile_option_result, +from backend.ascend_autotune_runtime.schedule_profiles import ( + COMPILE_MODE_KEY, + COMPILE_MODE_VECTOR, + CompileFailureRegionSet, + WORKSPACE_CV_AGGRESSIVE_PROBE, + WORKSPACE_CV_LOW_RESOURCE_PROBE, + WORKSPACE_CV_MIX1_PROBE, + classify_compile_failure, + compile_profile_to_config, + effective_compile_profile_key, + get_stage1_probe_configs, + get_stage1_probe_profiles, + make_stage2_seed_profiles, parse_compile_options_hint, ) -def test_compile_options_string_hint_expands_vector_defaults(): - spec = parse_compile_options_hint("vector") - configs = expand_compile_option_configs( - [triton.Config({"BLOCK_SIZE": 1024})], - spec, - generated_tiling=True, - ) - - assert len(configs) == 4 - assert {cfg.num_stages for cfg in configs} == {1, 2} - assert {cfg.kwargs["enable_ubuf_saving"] for cfg in configs} == {True, False} - assert {cfg.kwargs["enable_tuning_mode"] for cfg in configs} == {True} - assert {cfg.kwargs["BLOCK_SIZE"] for cfg in configs} == {1024} - - -def test_compile_options_string_hint_expands_mixcv_auto_search_space_without_limit(): - spec = parse_compile_options_hint("mixcv") - configs = expand_compile_option_configs( - [triton.Config({"BLOCK_SIZE": 1024})], - spec, - generated_tiling=True, - ) +def test_stage1_probe_order_matches_search_params_design(): + profiles = get_stage1_probe_profiles() - assert spec.max_configs is None - assert len(configs) == 156 - assert {cfg.num_stages for cfg in configs} == {1, 2} - assert {cfg.kwargs["enable_tuning_mode"] for cfg in configs} == {True} - assert {cfg.kwargs["unit_flag"] for cfg in configs} == {False, True} - assert {cfg.kwargs["enable_hivm_auto_cv_balance"] for cfg in configs} == {True} - assert {cfg.kwargs["enable_ubuf_saving"] for cfg in configs} == {False, True} - - stage1_configs = [cfg for cfg in configs if cfg.num_stages == 1] - assert len(stage1_configs) == 4 - assert all("multibuffer" not in cfg.kwargs for cfg in stage1_configs) - assert all( - "limit_auto_multi_buffer_only_for_local_buffer" not in cfg.kwargs - for cfg in stage1_configs - ) - assert all( - "limit_auto_multi_buffer_of_local_buffer" not in cfg.kwargs - for cfg in stage1_configs - ) - assert all("set_workspace_multibuffer" not in cfg.kwargs for cfg in stage1_configs) - assert all("tile_mix_vector_loop" not in cfg.kwargs for cfg in stage1_configs) - assert all("tile_mix_cube_loop" not in cfg.kwargs for cfg in stage1_configs) - assert all("enable_preload" not in cfg.kwargs for cfg in stage1_configs) - assert {cfg.kwargs["enable_auto_bind_sub_block"] for cfg in stage1_configs} == { - True - } - - stage2_configs = [cfg for cfg in configs if cfg.num_stages == 2] - assert len(stage2_configs) == 152 - assert all("multibuffer" not in cfg.kwargs for cfg in stage2_configs) - assert { - cfg.kwargs["limit_auto_multi_buffer_only_for_local_buffer"] - for cfg in stage2_configs - } == {False, True} - assert { - cfg.kwargs["limit_auto_multi_buffer_of_local_buffer"] for cfg in stage2_configs - } == {"no-limit", "no-l0c"} - assert { - cfg.kwargs["set_workspace_multibuffer"] - for cfg in stage2_configs - if "set_workspace_multibuffer" in cfg.kwargs - } == {2, 4} - assert { - cfg.kwargs["tile_mix_vector_loop"] - for cfg in stage2_configs - if "tile_mix_vector_loop" in cfg.kwargs - } == {1, 2, 4} - assert { - cfg.kwargs["tile_mix_cube_loop"] - for cfg in stage2_configs - if "tile_mix_cube_loop" in cfg.kwargs - } == {1, 2, 4} - assert {cfg.kwargs["enable_auto_bind_sub_block"] for cfg in stage2_configs} == { - True - } - assert any( - "set_workspace_multibuffer" in cfg.kwargs - and "tile_mix_vector_loop" in cfg.kwargs - and "tile_mix_cube_loop" in cfg.kwargs - for cfg in stage2_configs - ) - assert any( - "set_workspace_multibuffer" not in cfg.kwargs - and "tile_mix_vector_loop" not in cfg.kwargs - and "tile_mix_cube_loop" not in cfg.kwargs - for cfg in stage2_configs - ) - workspace_configs = [ - cfg for cfg in stage2_configs if "set_workspace_multibuffer" in cfg.kwargs + assert profiles == [ + WORKSPACE_CV_AGGRESSIVE_PROBE, + WORKSPACE_CV_LOW_RESOURCE_PROBE, + WORKSPACE_CV_MIX1_PROBE, ] - assert len(workspace_configs) == 144 - assert {cfg.num_stages for cfg in workspace_configs} == {2} - assert { - cfg.kwargs["limit_auto_multi_buffer_only_for_local_buffer"] - for cfg in workspace_configs - } == {False} - - -def test_compile_options_workspace_pruned_when_auto_multibuffer_disabled(): - spec = parse_compile_options_hint("mixcv") - configs = expand_compile_option_configs( - [triton.Config({"BLOCK_SIZE": 1024})], - spec, - generated_tiling=True, - fixed_options={"multibuffer": False, "num_stages": 2}, - ) - - assert configs - assert all("set_workspace_multibuffer" not in cfg.kwargs for cfg in configs) - - -def test_compile_options_workspace_pruned_when_workspace_limit_enabled(): - spec = parse_compile_options_hint("mixcv") - configs = expand_compile_option_configs( - [triton.Config({"BLOCK_SIZE": 1024})], - spec, - generated_tiling=True, - fixed_options={ - "limit_auto_multi_buffer_only_for_local_buffer": True, - "num_stages": 2, - }, - ) - - assert configs - assert all("set_workspace_multibuffer" not in cfg.kwargs for cfg in configs) - - -def test_compile_options_explicit_mixcv_values_are_not_restricted_by_auto_search_space(): - spec = parse_compile_options_hint( - { - "kernel_type": "mixcv", - "num_stages": [2], - "multibuffer": [False], - "unit_flag": [True], - "limit_auto_multi_buffer_only_for_local_buffer": [True], - "limit_auto_multi_buffer_of_local_buffer": ["no-limit"], - "set_workspace_multibuffer": [4], - "enable_hivm_auto_cv_balance": [False], - "tile_mix_vector_loop": [8], - "tile_mix_cube_loop": [8], - "enable_ubuf_saving": [False], - "enable_auto_bind_sub_block": [False], - } - ) - configs = expand_compile_option_configs( - [triton.Config({"BLOCK_SIZE": 1024})], - spec, - generated_tiling=True, - ) - - assert len(configs) == 1 - assert configs[0].num_stages == 2 - assert configs[0].kwargs["enable_tuning_mode"] is True - assert configs[0].kwargs["multibuffer"] is False - assert configs[0].kwargs["unit_flag"] is True - assert configs[0].kwargs["enable_hivm_auto_cv_balance"] is False - assert configs[0].kwargs["enable_auto_bind_sub_block"] is False - assert "limit_auto_multi_buffer_only_for_local_buffer" not in configs[0].kwargs - assert "limit_auto_multi_buffer_of_local_buffer" not in configs[0].kwargs - assert "set_workspace_multibuffer" not in configs[0].kwargs - assert "tile_mix_vector_loop" not in configs[0].kwargs - assert "tile_mix_cube_loop" not in configs[0].kwargs + assert {profile["unit_flag"] for profile in profiles} == {False} + assert profiles[0]["set_workspace_multibuffer"] == 4 + assert profiles[0]["limit_auto_multi_buffer_of_local_buffer"] == "no-limit" + assert profiles[1]["set_workspace_multibuffer"] == 2 + assert profiles[1]["limit_auto_multi_buffer_of_local_buffer"] == "no-l0c" + assert ( + profiles[2]["tile_mix_cube_loop"], + profiles[2]["tile_mix_vector_loop"], + ) == (1, 1) + + configs = get_stage1_probe_configs({"BLOCK_M": 128, "BLOCK_N": 512}) + assert len(configs) == 3 + assert {cfg.num_stages for cfg in configs} == {2} + assert {cfg.kwargs["unit_flag"] for cfg in configs} == {False} + assert configs[0].kwargs["BLOCK_M"] == 128 + assert configs[0].kwargs["BLOCK_N"] == 512 + assert COMPILE_MODE_KEY not in configs[0].kwargs -def test_compile_options_explicit_mixcv_stage1_values_are_preserved(): - spec = parse_compile_options_hint( - { - "kernel_type": "mixcv", - "num_stages": [1], - "multibuffer": [True], - "limit_auto_multi_buffer_only_for_local_buffer": [False], - "limit_auto_multi_buffer_of_local_buffer": ["no-l0c"], - "set_workspace_multibuffer": [2], - "tile_mix_vector_loop": [4], - "tile_mix_cube_loop": [4], - "enable_auto_bind_sub_block": [True], - } - ) - configs = expand_compile_option_configs( - [triton.Config({"BLOCK_SIZE": 1024})], - spec, - generated_tiling=True, - ) +def test_vector_compile_options_reject_mixcv_only_params(): + spec = parse_compile_options_hint({"kernel_type": "vector", "num_stages": [1, 2]}) + assert spec.params == {"num_stages": [1, 2]} - assert len(configs) == 4 - assert {cfg.num_stages for cfg in configs} == {1} - assert {cfg.kwargs["enable_tuning_mode"] for cfg in configs} == {True} - assert {cfg.kwargs["unit_flag"] for cfg in configs} == {False, True} - assert {cfg.kwargs["enable_ubuf_saving"] for cfg in configs} == {False, True} - assert all(cfg.kwargs["multibuffer"] is True for cfg in configs) - assert all("set_workspace_multibuffer" not in cfg.kwargs for cfg in configs) - assert all("tile_mix_vector_loop" not in cfg.kwargs for cfg in configs) - assert all("tile_mix_cube_loop" not in cfg.kwargs for cfg in configs) + for name, value in ( + ("unit_flag", False), + ("tile_mix_vector_loop", 2), + ("set_workspace_multibuffer", 2), + ): + with pytest.raises(ValueError, match=f"'{name}' is not supported"): + parse_compile_options_hint({"kernel_type": "vector", name: value}) -def test_compile_options_auto_search_overrides_base_compile_options(): - spec = parse_compile_options_hint("mixcv") - configs = expand_compile_option_configs( - [ - triton.Config( - { - "BLOCK_SIZE": 1024, - "enable_tuning_mode": False, - "unit_flag": True, - "enable_ubuf_saving": True, - "set_workspace_multibuffer": 4, - "tile_mix_vector_loop": 4, - "tile_mix_cube_loop": 4, - }, - num_stages=1, - ) - ], - spec, - generated_tiling=False, +def test_stage2_seeds_keep_stage1_winner_first_and_unit_flag_disabled(): + seeds = make_stage2_seed_profiles( + WORKSPACE_CV_MIX1_PROBE, + shape_kwargs={"BLOCK_M": 128, "BLOCK_N": 512}, + seed_budget=8, + allow_unit_flag=False, ) - assert len(configs) == 156 - assert {cfg.kwargs["enable_tuning_mode"] for cfg in configs} == {True} - assert {cfg.kwargs["unit_flag"] for cfg in configs} == {False, True} - assert {cfg.kwargs["enable_ubuf_saving"] for cfg in configs} == {False, True} - assert all( - "multibuffer" not in cfg.kwargs for cfg in configs if cfg.num_stages == 1 - ) - assert all( - "limit_auto_multi_buffer_only_for_local_buffer" not in cfg.kwargs - for cfg in configs - if cfg.num_stages == 1 - ) + assert seeds[0] == WORKSPACE_CV_MIX1_PROBE + assert len(seeds) <= 8 + assert len({effective_compile_profile_key(seed) for seed in seeds}) == len(seeds) + assert {seed["unit_flag"] for seed in seeds} == {False} assert all( - "limit_auto_multi_buffer_of_local_buffer" not in cfg.kwargs - for cfg in configs - if cfg.num_stages == 1 + seed[COMPILE_MODE_KEY] == WORKSPACE_CV_MIX1_PROBE[COMPILE_MODE_KEY] + for seed in seeds ) - assert all( - "set_workspace_multibuffer" not in cfg.kwargs - for cfg in configs - if cfg.num_stages == 1 - ) - assert all( - "tile_mix_vector_loop" not in cfg.kwargs - for cfg in configs - if cfg.num_stages == 1 - ) - assert all( - "tile_mix_cube_loop" not in cfg.kwargs for cfg in configs if cfg.num_stages == 1 - ) - assert all( - "enable_preload" not in cfg.kwargs for cfg in configs if cfg.num_stages == 1 - ) - assert { - cfg.kwargs["enable_auto_bind_sub_block"] - for cfg in configs - if cfg.num_stages == 1 - } == {True} -def test_compile_options_runtime_fixed_options_are_not_redefined(): - spec = parse_compile_options_hint("mixcv") - configs = expand_compile_option_configs( - [ - triton.Config( - { - "BLOCK_SIZE": 1024, - "multibuffer": False, - }, - num_stages=2, - ) - ], - spec, - generated_tiling=False, - fixed_options={ - "multibuffer": True, - }, +def test_compile_profile_to_config_keeps_mode_internal(): + config = compile_profile_to_config( + WORKSPACE_CV_LOW_RESOURCE_PROBE, + shape_kwargs={"BLOCK_M": 64, "BLOCK_N": 64}, + base_config=triton.Config({"EXTRA": 1}, num_warps=8), ) - assert configs - assert all("multibuffer" not in cfg.kwargs for cfg in configs) - assert {cfg.num_stages for cfg in configs} == {1, 2} + assert config.num_stages == 2 + assert config.num_warps == 8 + assert config.kwargs["EXTRA"] == 1 + assert config.kwargs["BLOCK_M"] == 64 + assert config.kwargs["tile_mix_cube_loop"] == 4 + assert COMPILE_MODE_KEY not in config.kwargs -def test_compile_options_runtime_fixed_num_stages_limits_search(): - spec = parse_compile_options_hint("mixcv") - configs = expand_compile_option_configs( - [triton.Config({"BLOCK_SIZE": 1024})], - spec, - generated_tiling=False, - fixed_options={"num_stages": 2}, +def test_compile_failure_classification_and_resource_region_prune(): + assert ( + classify_compile_failure( + "ub overflow, requires 3260416 bits while 1572864 bits available" + ) + == "RESOURCE_UB" ) - - assert len(configs) == 152 - assert {cfg.num_stages for cfg in configs} == {2} - - -def test_compile_options_format_stage1_effective_options(): - spec = parse_compile_options_hint("mixcv") - config = triton.Config( - { - "BLOCK_M": 32, - "BLOCK_N": 32, - "enable_tuning_mode": True, - "enable_ubuf_saving": True, - "enable_hivm_auto_cv_balance": True, - "unit_flag": True, - }, - num_stages=1, + assert ( + classify_compile_failure("internal error: dummyOps size is not 1") + == "COMPILER_INTERNAL" ) - text = format_compile_option_result(config, spec) - - assert "selected_meta: BLOCK_M=32, BLOCK_N=32, num_stages=1" in text - assert "enable_auto_multi_buffer=False" in text - assert "set_workspace_multibuffer=" in text - assert "tile_mix_vector_loop=" in text - assert "tile_mix_cube_loop=" in text - assert "enable_preload" not in text - - -def test_compile_options_format_stage2_effective_options(): - spec = parse_compile_options_hint("mixcv") - config = triton.Config( + failed_regions = CompileFailureRegionSet() + failed_regions.add( { - "BLOCK_M": 64, - "BLOCK_N": 128, - "enable_tuning_mode": True, - "enable_ubuf_saving": False, - "enable_hivm_auto_cv_balance": True, - "limit_auto_multi_buffer_only_for_local_buffer": False, - "limit_auto_multi_buffer_of_local_buffer": "no-l0c", + **WORKSPACE_CV_LOW_RESOURCE_PROBE, "set_workspace_multibuffer": 4, + "tile_mix_cube_loop": 2, "tile_mix_vector_loop": 2, - "tile_mix_cube_loop": 4, - "unit_flag": False, }, - num_stages=2, + "RESOURCE_UB", ) - text = format_compile_option_result(config, spec) - - assert "selected_meta: BLOCK_M=64, BLOCK_N=128, num_stages=2" in text - assert "enable_auto_multi_buffer=True" in text - assert "set_workspace_multibuffer=4" in text - assert "tile_mix_vector_loop=2" in text - assert "tile_mix_cube_loop=4" in text - assert " 1 and point[1] > 1 for point in proposal_points) + + +def test_resource_failure_sampling_backoffs_when_initial_points_are_too_large(): + spec = SearchParamsSpec( + enabled=True, + params=["NUM_CHUNKS"], + values=[1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192], + ) + all_configs = expand_search_param_shapes( + triton.Config({"NUM_CHUNKS": 64}), + spec, + ) + failed_values = [64, 512, 4096] + failed_configs = [ + next(config for config in all_configs if config.kwargs["NUM_CHUNKS"] == value) + for value in failed_values + ] + failure_observations = [ + {"shape": {"NUM_CHUNKS": 64}, "resource_ratio": 14680064 / 1572864}, + {"shape": {"NUM_CHUNKS": 512}, "resource_ratio": 117440512 / 1572864}, + {"shape": {"NUM_CHUNKS": 4096}, "resource_ratio": 512.0}, + ] + seen = { + shape_key(extract_shape(config, spec.params), spec.params) + for config in failed_configs + } + + proposals = propose_evolved_shape_configs( + parents=[], + successes=[], + failures=failed_configs, + failure_observations=failure_observations, + observed=failed_configs, + all_configs=all_configs, + spec=spec, + seen_keys=seen, + limit=4, + ) + proposed_values = [config.kwargs["NUM_CHUNKS"] for config in proposals] + + assert proposed_values + assert len(proposed_values) == len(set(proposed_values)) + assert not (set(proposed_values) & set(failed_values)) + assert min(proposed_values) <= 4 + assert max(proposed_values) <= 32 + + +def test_stage1_probe_failure_falls_back_until_success(): + # First probe fails, second succeeds, third must not run. + tuner = _DummyStage1Tuner( + batch_timings=[ + [float("inf")], + [0.25], + ] + ) + shape_config = triton.Config({"BLOCK_M": 128, "BLOCK_N": 512}) + + result = AutoTilingTuner._run_stage1_probe_for_shape( + tuner, shape_config=shape_config + ) + + assert result is not None + assert result["time"] == 0.25 + assert result["profile"] == WORKSPACE_CV_LOW_RESOURCE_PROBE + assert tuner.batch_calls == [1, 1] + + +def test_stage1_probe_rounds_only_retry_failed_shapes(): + # 2 shapes: A passes first probe, B retries and passes second probe. + tuner = _DummyStage1Tuner( + batch_timings=[ + [0.30, float("inf")], + [0.25], + ] + ) + shape_configs = [ + triton.Config({"BLOCK_M": 128, "BLOCK_N": 256}), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 512}), + ] + + results, failures = AutoTilingTuner._run_stage1_probes_for_shapes( + tuner, shape_configs=shape_configs + ) + + assert not failures + assert len(results) == 2 + assert results[0]["shape"] == {"BLOCK_M": 128, "BLOCK_N": 256} + assert results[0]["time"] == 0.30 + assert results[0]["profile"] == WORKSPACE_CV_AGGRESSIVE_PROBE + assert results[1]["shape"] == {"BLOCK_M": 128, "BLOCK_N": 512} + assert results[1]["time"] == 0.25 + assert results[1]["profile"] == WORKSPACE_CV_LOW_RESOURCE_PROBE + assert tuner.batch_calls == [2, 1] + assert tuner.force_parallel_flags == [True, True] + + +def _timing_key(value): + return float(value) + + +def test_ub_filter_drops_low_ub_and_slow_entries(): + entries = [ + {"shape": {"BLOCK_M": 128, "BLOCK_N": 128}, "ub": 10000, "time": 0.10}, + {"shape": {"BLOCK_M": 64, "BLOCK_N": 64}, "ub": 4000, "time": 0.20}, + {"shape": {"BLOCK_M": 32, "BLOCK_N": 32}, "ub": 9000, "time": 0.15}, + ] + kept = filter_shapes_by_ub_and_timing(entries, timing_sort_key=_timing_key) + shapes = [entry["shape"] for entry in kept] + # max_ub=10000, ub_floor=5000; min_time=0.10, timing_ceiling=0.14. + # (64,64) has ub=4000<5000 AND time=0.20>0.14 -> dropped. + assert {"BLOCK_M": 32, "BLOCK_N": 32} in shapes + assert {"BLOCK_M": 128, "BLOCK_N": 128} in shapes + assert {"BLOCK_M": 64, "BLOCK_N": 64} not in shapes diff --git a/test/ascend/autotune/test_split_axis_parse.py b/test/ascend/autotune/test_split_axis_parse.py deleted file mode 100644 index 473ea095..00000000 --- a/test/ascend/autotune/test_split_axis_parse.py +++ /dev/null @@ -1,163 +0,0 @@ -import triton -import triton.language as tl - -from test_common import check_axes_parse_res, mock_autotuner - - -def test_split_axis_parse_base_case1(mock_autotuner): - @triton.autotune(configs=[], key=["n_elements"]) - @triton.jit - def triton_split_axis_parse_base_case1( - x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr - ): - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - - tl.store(output_ptr + offsets, output, mask=mask) - - ref_res = { - "keys": {"x": "n_elements"}, - "split_params": {"x": "BLOCK_SIZE"}, - "tiling_params": {}, - "low_dim_axes": ["x"], - "reduction_axes": [], - } - grid = lambda meta: (meta["BLOCK_SIZE"],) - act_res = triton_split_axis_parse_base_case1[grid]() - - check_axes_parse_res(act_res, ref_res) - - -def test_split_axis_parse_base_case2(mock_autotuner): - @triton.autotune(configs=[], key=["n_elements"]) - @triton.jit - def triton_split_axis_parse_base_case2( - x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr - ): - block_start = tl.program_id(axis=0) * BLOCK_SIZE - - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - - tl.store(output_ptr + offsets, output, mask=mask) - - ref_res = { - "keys": {"x": "n_elements"}, - "split_params": {"x": "BLOCK_SIZE"}, - "tiling_params": {}, - "low_dim_axes": ["x"], - "reduction_axes": [], - } - grid = lambda meta: (meta["BLOCK_SIZE"],) - act_res = triton_split_axis_parse_base_case2[grid]() - - check_axes_parse_res(act_res, ref_res) - - -def test_split_axis_parse_base_case3(mock_autotuner): - @triton.autotune(configs=[], key=["n_elements"]) - @triton.jit - def triton_split_axis_parse_base_case3( - x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr - ): - offsets = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - - tl.store(output_ptr + offsets, output, mask=mask) - - ref_res = { - "keys": {"x": "n_elements"}, - "split_params": {"x": "BLOCK_SIZE"}, - "tiling_params": {}, - "low_dim_axes": ["x"], - "reduction_axes": [], - } - grid = lambda meta: (meta["BLOCK_SIZE"],) - act_res = triton_split_axis_parse_base_case3[grid]() - - check_axes_parse_res(act_res, ref_res) - - -def test_split_axis_parse_base_case4(mock_autotuner): - @triton.autotune(configs=[], key=["n_elements"]) - @triton.jit - def triton_split_axis_parse_base_case4( - x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr - ): - offsets = tl.program_id(axis=0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output_x = x[:, None].to(tl.float32) * 1 - output_y = 1 * y[None, :].to(tl.float32) - - output_offsets = ( - tl.program_id(axis=0) * BLOCK_SIZE * BLOCK_SIZE - + tl.arange(0, BLOCK_SIZE)[:, None] * BLOCK_SIZE - + tl.arange(0, BLOCK_SIZE)[None, :] - ) - tl.store(output_ptr + output_offsets, output_x + output_y, mask=mask) - - ref_res = { - "keys": {"x": "n_elements"}, - "split_params": {"x": "BLOCK_SIZE"}, - "tiling_params": {}, - "low_dim_axes": ["x"], - "reduction_axes": [], - } - grid = lambda meta: (meta["BLOCK_SIZE"],) - act_res = triton_split_axis_parse_base_case4[grid]() - - check_axes_parse_res(act_res, ref_res) - - -def test_grid_stride_loop_block_only_tiling_semantics(mock_autotuner): - @triton.autotune(configs=[], key=["N", "index_len"]) - @triton.jit - def triton_grid_stride_loop_block_only_tiling_semantics( - input_ptr, - output_ptr, - index_ptr, - N: tl.constexpr, - index_len: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, - ): - pid_x = tl.program_id(axis=0) - pid_y = tl.program_id(axis=1) - grid_x = tl.num_programs(axis=0) - grid_y = tl.num_programs(axis=1) - for x in range(pid_x * BLOCK_M, index_len, grid_x * BLOCK_M): - row_offsets = x + tl.arange(0, BLOCK_M) - indices = tl.load( - index_ptr + row_offsets, mask=row_offsets < index_len, other=0 - ) - for y in range(pid_y * BLOCK_N, N, grid_y * BLOCK_N): - col_offsets = y + tl.arange(0, BLOCK_N) - col_mask = col_offsets < N - inp_offset = indices[:, None] * N + col_offsets[None, :] - out_offset = row_offsets[:, None] * N + col_offsets[None, :] - selected = tl.load( - input_ptr + inp_offset, mask=col_mask[None, :], other=0.0 - ) - tl.store(output_ptr + out_offset, selected, mask=col_mask[None, :]) - - act_res = triton_grid_stride_loop_block_only_tiling_semantics[(1, 1)]() - assert act_res["split_params"] == {} - assert act_res["tiling_params"] == {"y": "BLOCK_M", "x": "BLOCK_N"} diff --git a/test/ascend/autotune/test_tiling_axis_parse.py b/test/ascend/autotune/test_tiling_axis_parse.py deleted file mode 100644 index de9e5f65..00000000 --- a/test/ascend/autotune/test_tiling_axis_parse.py +++ /dev/null @@ -1,115 +0,0 @@ -import pytest -import triton -import triton.language as tl -from test_common import check_axes_parse_res, mock_autotuner - - -def test_tiling_axis_parse_base_case1(mock_autotuner): - @triton.autotune(configs=[], key=["n_elements"]) - @triton.jit - def triton_tiling_axis_parse_base_case1( - x_ptr, - y_ptr, - output_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, - BLOCK_SUB: tl.constexpr, - ): - offset = tl.program_id(axis=0) * BLOCK_SIZE - base = tl.arange(0, BLOCK_SUB) - loops = (BLOCK_SIZE + BLOCK_SUB - 1) // BLOCK_SUB - for loop in range(loops): - offsets = offset + (loop * BLOCK_SUB) + base - mask = offsets < min(BLOCK_SIZE + offset, n_elements) - - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - - tl.store(output_ptr + offsets, output, mask=mask) - - ref_res = { - "keys": {"x": "n_elements"}, - "split_params": {"x": "BLOCK_SIZE"}, - "tiling_params": {"x": "BLOCK_SUB"}, - "low_dim_axes": ["x"], - "reduction_axes": [], - } - grid = lambda meta: (meta["BLOCK_SIZE"],) - act_res = triton_tiling_axis_parse_base_case1[grid]() - - check_axes_parse_res(act_res, ref_res) - - -@pytest.mark.skip -def test_tiling_axis_parse_base_case2(mock_autotuner): - @triton.autotune(configs=[], key=["n_elements"]) - @triton.jit - def triton_tiling_axis_parse_base_case2( - x_ptr, - y_ptr, - output_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, - BLOCK_SUB: tl.constexpr, - ): - offset = tl.program_id(axis=0) * BLOCK_SIZE - base = tl.arange(0, BLOCK_SUB) - for offset_sub in range(0, BLOCK_SIZE, BLOCK_SUB): - offsets = offset + offset_sub + base[:] - mask = offsets < min(BLOCK_SIZE + offset, n_elements) - - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - - tl.store(output_ptr + offsets, output, mask=mask) - - ref_res = { - "keys": {"x": "n_elements"}, - "split_params": {"x": "BLOCK_SIZE"}, - "tiling_params": {"x": "BLOCK_SUB"}, - "low_dim_axes": ["x"], - "reduction_axes": [], - } - grid = lambda meta: (meta["BLOCK_SIZE"],) - act_res = triton_tiling_axis_parse_base_case2[grid]() - - check_axes_parse_res(act_res, ref_res) - - -@pytest.mark.skip -def test_tiling_axis_parse_base_case3(mock_autotuner): - @triton.autotune(configs=[], key=["n_elements"]) - @triton.jit - def triton_tiling_axis_parse_base_case3( - x_ptr, - y_ptr, - output_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, - BLOCK_SUB: tl.constexpr, - ): - offset = tl.program_id(axis=0) * BLOCK_SIZE - base = tl.arange(0, BLOCK_SUB)[:] - for offset_sub in range(0, BLOCK_SIZE, BLOCK_SUB): - offsets = offset + offset_sub + base - mask = offsets < min(BLOCK_SIZE + offset, n_elements) - - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - - tl.store(output_ptr + offsets, output, mask=mask) - - ref_res = { - "keys": {"x": "n_elements"}, - "split_params": {"x": "BLOCK_SIZE"}, - "tiling_params": {"x": "BLOCK_SUB"}, - "low_dim_axes": ["x"], - "reduction_axes": [], - } - grid = lambda meta: (meta["BLOCK_SIZE"],) - act_res = triton_tiling_axis_parse_base_case3[grid]() - - check_axes_parse_res(act_res, ref_res) diff --git a/tools/dicp_triton_opt/CMakeLists.txt b/tools/dicp_triton_opt/CMakeLists.txt index 1dd64186..6e849e14 100644 --- a/tools/dicp_triton_opt/CMakeLists.txt +++ b/tools/dicp_triton_opt/CMakeLists.txt @@ -9,6 +9,7 @@ llvm_update_compile_flags(dicp_opt) target_link_libraries(dicp_opt PRIVATE TritonAnalysis TritonTransforms + MLIRFunctionInterfaces ${dialect_libs} ${translation_libs} ${conversion_libs} diff --git a/tools/dicp_triton_opt/dicp_triton_opt.cpp b/tools/dicp_triton_opt/dicp_triton_opt.cpp index 676b7677..ee7c6524 100644 --- a/tools/dicp_triton_opt/dicp_triton_opt.cpp +++ b/tools/dicp_triton_opt/dicp_triton_opt.cpp @@ -48,6 +48,7 @@ #include "mlir/Dialect/GPU/TransformOps/GPUTransformOps.h" #include "mlir/Dialect/Index/IR/IndexDialect.h" #include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/IR/ValueBoundsOpInterfaceImpl.h" #include "mlir/Dialect/Linalg/Passes.h" @@ -89,6 +90,9 @@ inline void registerDICPDialects(mlir::DialectRegistry ®istry) { mlir::registerAllPasses(); mlir::registerLinalgPasses(); + mlir::func::registerAllExtensions(registry); + mlir::LLVM::registerInlinerInterface(registry); + // triton-dicp pass registrations triton::registerAutoBlockifyPass(); triton::registerAscendLegalizePass();