API ReferenceOther APIsAituneTorchKernel Forge

aitune.torch.kernel_forge.kernel_optimizer

View as Markdown

Kernel optimizer based on module function kernel profiling.

Module Contents

Classes

NameDescription
KernelOptimizerProfiles and replaces eligible torch.nn.functional kernels in a module.
_FunctionSearchSamples and accumulated candidates for one PyTorch function.
_GenerationTaskOne submitted generator task and its expected function.
_KernelCandidateA kernel candidate.

Data

KernelSource

KernelSourceType

_PROFILE_SUMMARY_LIMIT

logger

API

class aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer(
    kernel_providers: aitune.torch.kernel_forge.kernel_provider.KernelProvider | list[aitune.torch.kernel_forge.kernel_provider.KernelProvider] | None = None,
    kernel_generators: aitune.torch.kernel_forge.kernel_provider.KernelGenerator | list[aitune.torch.kernel_forge.kernel_provider.KernelGenerator] | None = None,
    kernel_profiler_factory: collections.abc.Callable[..., aitune.torch.kernel_forge.module_function_kernel_profiler.ModuleFunctionKernelProfiler] = ModuleFunctionKernelProfiler,
    kernel_utils: aitune.torch.kernel_forge.kernel_utils.KernelUtils | None = None,
    provider_min_time_share_percent: float = 0.0,
    generator_min_time_share_percent: float = 10.0,
    generation_timeout: float = AITUNE_KERNEL_GENERATION_TI...
)

Profiles and replaces eligible torch.nn.functional kernels in a module.

make_plan(function, data, module=...) profiles forwards, starts asynchronous generators, benchmarks static providers while generation runs, collects and benchmarks generated candidates, and returns the selected runtime plan.

kernel_utils
aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._evaluate_candidate(
    search: aitune.torch.kernel_forge.kernel_optimizer._FunctionSearch,
    provider: aitune.torch.kernel_forge.kernel_provider.KernelProvider,
    description: str
) -> None

Run the common correctness validation and benchmark for one candidate.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._evaluate_generation_result(
    searches: dict[str, aitune.torch.kernel_forge.kernel_optimizer._FunctionSearch],
    task: aitune.torch.kernel_forge.kernel_optimizer._GenerationTask,
    result: aitune.torch.kernel_forge.kernel_provider.KernelGenerationResult
) -> None

Evaluate one completed generation result.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._evaluate_generation_results(
    searches: dict[str, aitune.torch.kernel_forge.kernel_optimizer._FunctionSearch],
    tasks: list[aitune.torch.kernel_forge.kernel_optimizer._GenerationTask]
) -> None

Evaluate generated candidates sequentially as their futures finish.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._evaluate_provider_candidates(
    searches: dict[str, aitune.torch.kernel_forge.kernel_optimizer._FunctionSearch],
    provider_functions: list[str]
) -> None

Prepare and benchmark all provider candidates while generators run.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._generators_for_function(
    func_name: str
) -> list[aitune.torch.kernel_forge.kernel_provider.KernelGenerator]

Return dynamic generators supporting one function.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._kernel_sources() -> list[aitune.torch.kernel_forge.kernel_optimizer.KernelSource]

Return all available static and dynamic kernel sources.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._kernel_sources_for_function(
    sources: list[aitune.torch.kernel_forge.kernel_optimizer.KernelSourceType],
    func_name: str
) -> list[aitune.torch.kernel_forge.kernel_optimizer.KernelSourceType]
staticmethod

Return sources supporting one function.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._log_kernel_candidates_results(
    func_name: str,
    candidates: list[aitune.torch.kernel_forge.kernel_optimizer._KernelCandidate],
    baseline_latency: float
)

Log candidate latencies and their improvement over the baseline.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._log_summary(
    summary_df: pandas.DataFrame,
    provider_function_names: list[str],
    generator_function_names: list[str]
) -> None

Log profiler results with markers for the selected kernel sources.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._make_plan(
    function_data: dict[str, aitune.torch.kernel_forge.module_function_kernel_profiler.FunctionData],
    profiled_functions: list[str],
    provider_functions: list[str],
    generator_functions: list[str]
) -> aitune.torch.kernel_forge.kernel_optimization_plan.KernelOptimizationPlan

Evaluate kernel candidates and return the best providers.

All generator tasks are scheduled first. While asynchronous generation is running, candidates from static providers are prepared and benchmarked. Generated results are then collected and evaluated using the same validation and benchmarking path. Finally, the fastest candidate for each function is selected when it outperforms the original PyTorch function.

Parameters:

function_data
dict[str, FunctionData]

Collected input samples for each supported function.

profiled_functions
list[str]

Profiled functions in descending kernel-time order.

provider_functions
list[str]

Profiled functions selected for static providers.

generator_functions
list[str]

Profiled functions selected for asynchronous generators.

Returns: KernelOptimizationPlan

A runtime plan containing the fastest providers that outperform their baselines.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._prepare_provider_samples(
    function_data: aitune.torch.kernel_forge.module_function_kernel_profiler.FunctionData
) -> tuple[list[aitune.torch.module.sample_store.Sample], list[aitune.torch.module.sample_store.Sample]]
staticmethod

Return unique validation samples and a GCD-reduced benchmark distribution.

By dividing counters by the GCD of the counters, we can reduce the number of samples to a minimum while maintaining the same distribution.

Parameters:

function_data
FunctionData

the function data to prepare the provider samples for

Returns: tuple[list[Sample], list[Sample]]

a tuple of unique validation samples and a GCD-reduced benchmark distribution

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._prepare_searches(
    function_data: dict[str, aitune.torch.kernel_forge.module_function_kernel_profiler.FunctionData],
    profiled_functions: list[str]
) -> dict[str, aitune.torch.kernel_forge.kernel_optimizer._FunctionSearch]

Prepare each selected function’s samples exactly once.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._providers_for_function(
    func_name: str
) -> list[aitune.torch.kernel_forge.kernel_provider.KernelProvider]

Return static providers supporting one function.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._select_best_candidate(
    func_name: str,
    search: aitune.torch.kernel_forge.kernel_optimizer._FunctionSearch
) -> aitune.torch.kernel_forge.kernel_optimizer._KernelCandidate | None

Select the fastest valid candidate when it beats the single baseline measurement.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._select_source_functions(
    summary_df: pandas.DataFrame,
    sources: list[aitune.torch.kernel_forge.kernel_optimizer.KernelSourceType],
    min_time_share_percent: float
) -> list[str]

Select functions supported by a source and meeting its time-share threshold.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._submit_generation_tasks(
    searches: dict[str, aitune.torch.kernel_forge.kernel_optimizer._FunctionSearch],
    generator_functions: list[str]
) -> list[aitune.torch.kernel_forge.kernel_optimizer._GenerationTask]

Submit all supported generator-function pairs before static work.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._supported_function_names() -> set[str]

Return function names supported by providers or generators.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer._supported_functions(
    source: aitune.torch.kernel_forge.kernel_optimizer.KernelSource
) -> list[str]
staticmethod

Return function names supported by a provider or generator.

aitune.torch.kernel_forge.kernel_optimizer.KernelOptimizer.make_plan(
    function: collections.abc.Callable,
    data: collections.abc.Sequence[aitune.torch.module.sample_store.Sample] | None = None,
    module: torch.nn.Module | None = None
) -> aitune.torch.kernel_forge.kernel_optimization_plan.KernelOptimizationPlan

Optimize a module by replacing eligible kernels with the best ones.

Parameters:

function
Callable

the function to run inference with

data
Sequence[Sample] | NoneDefaults to None

the data to use for optimization, if None the function will be called without arguments

module
nn.Module | NoneDefaults to None

the module whose hierarchy should be optimized. If omitted and function is an nn.Module, function is used as the module.

Returns: KernelOptimizationPlan

A plan containing the selected runtime kernel providers.

Raises:

  • ValueError: If the module scope cannot be resolved.
class aitune.torch.kernel_forge.kernel_optimizer._FunctionSearch(
    real_function: collections.abc.Callable,
    unique_samples: list[aitune.torch.module.sample_store.Sample],
    benchmark_samples: list[aitune.torch.module.sample_store.Sample],
    candidates: list[aitune.torch.kernel_forge.kernel_optimizer._KernelCandidate] = list()
)
Dataclass

Samples and accumulated candidates for one PyTorch function.

benchmark_samples
list[Sample]
candidates
list[_KernelCandidate] = field(default_factory=list)
real_function
Callable
unique_samples
list[Sample]
class aitune.torch.kernel_forge.kernel_optimizer._GenerationTask(
    function_name: str,
    generator_name: str,
    future: concurrent.futures.Future[aitune.torch.kernel_forge.kernel_provider.KernelGenerationResult]
)
Dataclass

One submitted generator task and its expected function.

function_name
str
future
Future[KernelGenerationResult]
generator_name
str
class aitune.torch.kernel_forge.kernel_optimizer._KernelCandidate(
    provider: aitune.torch.kernel_forge.kernel_provider.KernelProvider,
    latency: float,
    description: str
)
Dataclass

A kernel candidate.

description
str

The description of the candidate.

latency
float

The latency of the candidate.

provider
KernelProvider

The selected runtime kernel provider.

aitune.torch.kernel_forge.kernel_optimizer.KernelSource = KernelProvider | KernelGenerator
 
aitune.torch.kernel_forge.kernel_optimizer.KernelSourceType = TypeVar('KernelSourceType', KernelProvider, KernelGenerator)
 
aitune.torch.kernel_forge.kernel_optimizer._PROFILE_SUMMARY_LIMIT = 100
 
aitune.torch.kernel_forge.kernel_optimizer.logger = getLogger(__name__)