ejkernel.ops.execution.executor

ejkernel.ops.execution.executor#

Main execution engine for kernels with configuration management.

This module provides the Executor class, which serves as the central orchestrator for running kernel operations with automatic configuration selection, custom gradient support, and comprehensive profiling capabilities.

Key Components:

Executor: Main execution engine coordinating the entire execution pipeline ConfigChooser: Protocol defining configuration selection interface

The Executor handles:

  • Argument preprocessing via kernel.prepare()

  • Configuration selection through ConfigChooser strategies

  • Custom VJP (Vector-Jacobian Product) implementation for gradients

  • Profiling metadata injection for performance analysis

  • Invocation recording for batch optimization

  • JAX compilation with pre-selected configurations

Execution Flow:

  1. Preprocess arguments using kernel.prepare()

  2. Create Invocation object with argument metadata

  3. Select configuration via ConfigChooser.choose()

  4. Set up custom VJP if kernel implements it

  5. Add profiling metadata based on environment settings

  6. Execute kernel with chosen configuration

  7. Record invocation for future optimization (if enabled)

Environment Variables:

EJKERNEL_OPS_RECORD: Set to “1” to enable invocation recording EJKERNEL_OPS_STAMP: Controls profiling metadata format:

  • “hash”: Use operation hash for labeling (default)

  • “json”: Use full JSON payload for labeling

  • “none”: Disable profiling metadata

Example Usage:
>>> cache = ConfigCache()
>>> selector = ConfigSelectorChain(cache)
>>> executor = Executor(selector)
>>>
>>>
>>> result = executor(my_kernel, input_data)
>>>
>>>
>>> compiled_fn = executor.compile(my_kernel, example_input)
>>> result = compiled_fn(actual_input)
class ejkernel.ops.execution.executor.ConfigChooser(*args, **kwargs)[source]#

Bases: Protocol

Protocol for configuration selection strategies.

Defines the interface that configuration selection strategies must implement. The primary implementer is ConfigSelectorChain, which provides a sophisticated multi-tier selection system with caching and autotuning.

choose()[source]#

Select optimal configuration for the given invocation and kernel

choose(inv: Invocation[Cfg, Out], kernel: Kernel[Cfg, Out]) → Cfg[source]#

Select optimal configuration for the given invocation.

Parameters

  • inv – Invocation object containing arguments and metadata

  • kernel – Kernel implementation requiring configuration

Returns

Configuration object suitable for the kernel and invocation

class ejkernel.ops.execution.executor.Executor(chooser: ConfigChooser, stamp_prefix: str = 'ejkernel_ops')[source]#

Bases: Generic[Cfg, Out]

Main execution engine for kernels with automatic configuration selection.

The Executor coordinates the entire execution process: 1. Preprocess arguments via kernel.prepare() 2. Select configuration via the ConfigChooser 3. Handle custom VJP if implemented by the kernel 4. Add profiling metadata if requested 5. Execute the kernel with the chosen configuration

Supports both regular operations and custom gradient implementations.

chooser#

Configuration selection strategy (typically ConfigSelectorChain)

stamp_prefix#

Prefix for profiling metadata labels

choose_config(kernel: Kernel[Cfg, Out], *args, cfg: Optional[Cfg] = None, **kwargs) → Cfg[source]#

Select configuration for kernel without executing it.

Useful for inspecting what configuration would be chosen for given arguments, or for pre-selecting configurations for compilation.

Parameters

  • kernel – Kernel implementation requiring configuration

  • *args – Positional arguments for the kernel

  • cfg – Optional configuration override

  • **kwargs – Keyword arguments for the kernel

Returns

Configuration that would be selected for this invocation

compile(kernel: Kernel[Cfg, Out], *example_args, cfg: Optional[Cfg] = None, **example_kwargs)[source]#

Compile kernel with pre-selected configuration.

Selects optimal configuration based on example arguments, then returns a JAX-compiled function that uses that configuration for all subsequent calls. This avoids configuration selection overhead during execution.

Parameters

  • kernel – Kernel implementation to compile

  • *example_args – Example positional arguments for configuration selection

  • cfg – Optional configuration override

  • **example_kwargs – Example keyword arguments for configuration selection

Returns

JAX-compiled function with pre-selected configuration

Example

>>> compiled_matmul = executor.compile(matmul_kernel, x_example, y_example)
>>>
>>> result = compiled_matmul(x_actual, y_actual)
Contents