Index

_

_data (ejkernel.ops.config.cache.ConfigCache attribute)

_logger (ejkernel.loggings.ProgressLogger attribute) _registry (ejkernel.kernels._registry.KernelRegistry attribute)

A

abstractify() (in module ejkernel.ops.utils.fingerprint) algorithm (ejkernel.kernels._registry.KernelSpec attribute), [1] allow_autotune (ejkernel.ops.config.selection.AutotunePolicy attribute), [1] allow_fallback_timing (ejkernel.ops.execution.tuning.FNAutotuner attribute) allow_heuristics (ejkernel.ops.config.selection.AutotunePolicy attribute), [1] ANY (ejkernel.kernels._registry.Backend attribute), [1] args (ejkernel.ops.core.kernel.Invocation attribute), [1] as_overlay() (ejkernel.ops.execution.tuning.AutotuningResult method) autotune() (ejkernel.ops.config.selection.Tuner method) (ejkernel.ops.execution.tuning.Autotuner method) (in module ejkernel.ops.execution.tuning)

autotune_lowered() (in module ejkernel.ops.execution.offline) autotune_recorded() (in module ejkernel.ops.execution.tuning) AutotuneData (class in ejkernel.ops.execution.tuning) AutotunePolicy (class in ejkernel.ops.config.selection) Autotuner (class in ejkernel.ops.execution.tuning) AutotuningResult (class in ejkernel.ops.execution.tuning) aval_size_bytes() (in module ejkernel.callib._triton_call) avals_to_layouts() (in module ejkernel.callib._triton_call)

B

Backend (class in ejkernel.kernels._registry) backend (ejkernel.kernels._registry.KernelSpec attribute), [1] batch_axes (ejkernel.ops.core.kernel.Invocation attribute), [1] benchmark() (in module ejkernel.ops.execution.tuning) binary (ejkernel.callib._triton_call.CompilationResult attribute), [1] blocksize_heads (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] blocksize_k (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1]

blocksize_keys (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] blocksize_m (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] blocksize_n (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] buffered_pallas_call() (in module ejkernel.callib._pallas_call) build_cute_ffi_call() (in module ejkernel.callib._cute_ffi) BwdParams (class in ejkernel.ops.utils.datacarrier)

C

cache (ejkernel.ops.config.selection.ConfigSelectorChain attribute) cache_miss_fallback (ejkernel.ops.config.selection.AutotunePolicy attribute), [1] cache_size_limit (ejkernel.ops.execution.tuning.FNAutotuner attribute) call_key (ejkernel.ops.core.kernel.Invocation property) (ejkernel.ops.execution.tuning.Entry attribute), [1] calls_per_round (ejkernel.ops.execution.tuning.FNAutotuner attribute) candidate_cfgs() (ejkernel.ops.core.kernel.Kernel method) cdiv() (in module ejkernel.callib._utils) cfg (ejkernel.ops.execution.tuning.Entry attribute), [1] (ejkernel.ops.execution.tuning.Measurement attribute), [1] check_bool_flag() (in module ejkernel.callib._utils) check_vma (ejkernel.ops.core.kernel.Invocation attribute), [1] choose() (ejkernel.ops.config.selection.ConfigSelectorChain method) (ejkernel.ops.execution.executor.ConfigChooser method), [1] choose_config() (ejkernel.ops.execution.executor.Executor method) chooser (ejkernel.ops.execution.executor.Executor attribute) clear() (ejkernel.ops.config.cache.ConfigCache method)

cluster_dims (ejkernel.callib._triton_call.CompilationResult attribute), [1] ColorFormatter (class in ejkernel.loggings) CompilationResult (class in ejkernel.callib._triton_call) compile() (ejkernel.ops.execution.executor.Executor method) compile_ttir_inplace() (in module ejkernel.callib._triton_call) compile_ttir_to_hsaco_inplace() (in module ejkernel.callib._triton_call) compile_ttir_to_ptx_inplace() (in module ejkernel.callib._triton_call) complete() (ejkernel.loggings.ProgressLogger method) ConfigCache (class in ejkernel.ops.config.cache) ConfigChooser (class in ejkernel.ops.execution.executor) ConfigSelectorChain (class in ejkernel.ops.config.selection) CPU (ejkernel.kernels._registry.Backend attribute), [1] create_shard_map_wrapper() (ejkernel.ops.core.kernel.Kernel method) create_step_profiler() (in module ejkernel.loggings) CUDA (ejkernel.kernels._registry.Platform attribute), [1] CUTE (ejkernel.kernels._registry.Platform attribute), [1] cute_call() (in module ejkernel.callib._cute_call)

D

decorate() (ejkernel.ops.execution.tuning.FNAutotuner method) default_key_builder_with_sharding() (in module ejkernel.ops.utils.fingerprint) device (ejkernel.ops.execution.tuning.AutotuningResult attribute), [1] device_fingerprint() (in module ejkernel.ops.utils.fingerprint)

device_kind() (in module ejkernel.ops.utils.fingerprint) dtype (ejkernel.callib._utils.ShapeDtype property) DummyStream (class in ejkernel.callib._utils) dumper (ejkernel.ops.config.persistent.PersistentCache attribute)

E

ejit() (in module ejkernel.callib._ejit) ejkernel.callib._cute_call module ejkernel.callib._cute_ffi module ejkernel.callib._ejit module ejkernel.callib._pallas_call module ejkernel.callib._triton_call module ejkernel.callib._utils module ejkernel.errors module ejkernel.kernels._registry module ejkernel.loggings module ejkernel.ops.config.cache module ejkernel.ops.config.persistent module ejkernel.ops.config.selection module ejkernel.ops.core.kernel module ejkernel.ops.core.types module

ejkernel.ops.execution.batch module ejkernel.ops.execution.executor module ejkernel.ops.execution.offline module ejkernel.ops.execution.profiler module ejkernel.ops.execution.tuning module ejkernel.ops.registry module ejkernel.ops.utils.datacarrier module ejkernel.ops.utils.fingerprint module ejkernel.ops.utils.meta module ejkernel.ops.utils.serialize module EjkernelRuntimeError enable_detailed_logging (ejkernel.ops.execution.tuning.FNAutotuner attribute) entries (ejkernel.ops.execution.tuning.AutotuningResult attribute), [1] Entry (class in ejkernel.ops.execution.tuning) event_filter_regex (ejkernel.ops.execution.profiler.Profiler attribute) Executor (class in ejkernel.ops.execution.executor) extinguish_profiler() (in module ejkernel.loggings) extract_labels_from_hlo_text() (in module ejkernel.ops.utils.meta)

F

F (in module ejkernel.kernels._registry) fastest_config (ejkernel.ops.execution.tuning.AutotuneData property) find_device_plane_ids() (ejkernel.ops.execution.profiler.Profiler static method) find_labels_in_lowered() (in module ejkernel.ops.utils.meta) find_optimal_layouts_automatically (ejkernel.ops.execution.tuning.FNAutotuner attribute) flush() (ejkernel.callib._utils.DummyStream method) FNAutotuner (class in ejkernel.ops.execution.tuning)

forbid_reautotune (ejkernel.ops.config.selection.ConfigSelectorChain attribute) format() (ejkernel.loggings.ColorFormatter method), [1] forward_autotune_only (class in ejkernel.ops.config.selection) from_json() (in module ejkernel.ops.utils.serialize) fwd_with_residuals() (ejkernel.ops.core.kernel.Kernel method) fwd_with_residuals_shard_map() (ejkernel.ops.core.kernel.Kernel method) fwd_with_residuals_shard_map_gpu() (ejkernel.ops.core.kernel.Kernel method) FwdParams (class in ejkernel.ops.utils.datacarrier)

G

get() (ejkernel.kernels._registry.KernelRegistry method) (ejkernel.ops.config.cache.ConfigCache method) (ejkernel.ops.config.persistent.PersistentCache method) get_cache_dir() (in module ejkernel.callib._utils) get_cuda_backend() (in module ejkernel.callib._triton_call) get_device_platform() (in module ejkernel.ops.utils.fingerprint) get_events_from_plane() (ejkernel.ops.execution.profiler.Profiler method) get_hash_of_lowering() (in module ejkernel.callib._ejit)

get_hip_backend() (in module ejkernel.callib._triton_call) get_invocations() (in module ejkernel.ops.registry) get_logger() (in module ejkernel.loggings) get_or_create_triton_kernel() (in module ejkernel.callib._triton_call) get_safe_hash_int() (in module ejkernel.callib._ejit) (in module ejkernel.ops.utils.datacarrier) get_triton_type() (in module ejkernel.callib._triton_call) GPU (ejkernel.kernels._registry.Backend attribute), [1]

H

has_cute_ffi_support() (in module ejkernel.callib._cute_ffi) hash_fn() (in module ejkernel.callib._ejit) (in module ejkernel.ops.utils.datacarrier)

heuristic_cfg() (ejkernel.ops.core.kernel.Kernel method) hyperparams (ejkernel.ops.execution.tuning.TimingResult attribute), [1]

I

ignite_profiler() (in module ejkernel.loggings) implementation (ejkernel.kernels._registry.KernelSpec attribute), [1] in_specs (ejkernel.ops.core.kernel.Invocation attribute), [1]

Invocation (class in ejkernel.ops.core.kernel) iters (ejkernel.ops.config.selection.Tuner attribute) (ejkernel.ops.execution.tuning.Autotuner attribute)

K

Kernel (class in ejkernel.ops.core.kernel) kernel_registry (in module ejkernel.kernels._registry) KernelRegistry (class in ejkernel.kernels._registry) KernelSpec (class in ejkernel.kernels._registry)

key_builder (ejkernel.ops.core.kernel.Kernel attribute), [1] keys() (ejkernel.ops.config.cache.ConfigCache method) kv_blocksize (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] kwargs (ejkernel.ops.core.kernel.Invocation attribute), [1]

L

label() (in module ejkernel.ops.utils.meta) labels_to_configs() (in module ejkernel.ops.utils.meta) LazyLogger (class in ejkernel.loggings) level (ejkernel.loggings.LazyLogger attribute) (ejkernel.loggings.LazyLogger property) list_algorithms() (ejkernel.kernels._registry.KernelRegistry method)

list_implementations() (ejkernel.kernels._registry.KernelRegistry method) llir (ejkernel.callib._triton_call.CompilationResult attribute), [1] load_cached_functions() (in module ejkernel.callib._ejit) load_compiled_fn() (in module ejkernel.callib._ejit) loader (ejkernel.ops.config.persistent.PersistentCache attribute) logger (in module ejkernel.loggings)

M

make_key() (ejkernel.ops.core.kernel.Invocation method) max_compilation_time_seconds (ejkernel.ops.execution.tuning.FNAutotuner attribute) max_events_per_profile (ejkernel.ops.execution.profiler.Profiler attribute) measure() (ejkernel.ops.config.selection.Tuner method) Measurement (class in ejkernel.ops.execution.tuning) measurements (ejkernel.ops.execution.tuning.AutotuneData attribute), [1] mesh (ejkernel.ops.core.kernel.Invocation attribute), [1] method (ejkernel.ops.core.kernel.Invocation attribute), [1] min_duration_ns (ejkernel.ops.execution.profiler.Profiler attribute) module ejkernel.callib._cute_call ejkernel.callib._cute_ffi ejkernel.callib._ejit ejkernel.callib._pallas_call ejkernel.callib._triton_call ejkernel.callib._utils ejkernel.errors ejkernel.kernels._registry ejkernel.loggings ejkernel.ops.config.cache ejkernel.ops.config.persistent ejkernel.ops.config.selection ejkernel.ops.core.kernel ejkernel.ops.core.types ejkernel.ops.execution.batch ejkernel.ops.execution.executor ejkernel.ops.execution.offline ejkernel.ops.execution.profiler ejkernel.ops.execution.tuning ejkernel.ops.registry ejkernel.ops.utils.datacarrier ejkernel.ops.utils.fingerprint ejkernel.ops.utils.meta ejkernel.ops.utils.serialize

must_find_profiler_fraction (ejkernel.ops.execution.tuning.FNAutotuner attribute)

N

name (ejkernel.callib._triton_call.CompilationResult attribute), [1] (ejkernel.loggings.LazyLogger attribute) (ejkernel.loggings.LazyLogger property) (ejkernel.loggings.ProgressLogger attribute) next_power_of_2() (in module ejkernel.callib._utils)

normalize_grid() (in module ejkernel.callib._triton_call) num_key_splits (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] num_stages (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1] num_warps (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1]

O

on_event (ejkernel.ops.config.selection.ConfigSelectorChain attribute) op_id (ejkernel.ops.core.kernel.Invocation attribute), [1] (ejkernel.ops.core.kernel.Kernel attribute), [1]

op_id_v (ejkernel.ops.execution.tuning.Entry attribute), [1] out_specs (ejkernel.ops.core.kernel.Invocation attribute), [1] overlay_cache (class in ejkernel.ops.config.cache) override_cfg (ejkernel.ops.core.kernel.Invocation attribute), [1]

P

PALLAS (ejkernel.kernels._registry.Platform attribute), [1] parse_profile_from_bytes() (ejkernel.ops.execution.profiler.Profiler static method) path (ejkernel.ops.config.persistent.PersistentCache attribute) persist_autotune (ejkernel.ops.config.selection.ConfigSelectorChain attribute) persistent (ejkernel.ops.config.selection.ConfigSelectorChain attribute) PersistentCache (class in ejkernel.ops.config.persistent) Platform (class in ejkernel.kernels._registry) platform (ejkernel.kernels._registry.KernelSpec attribute), [1] pmap_with_config() (in module ejkernel.ops.execution.batch) policy (ejkernel.ops.config.selection.ConfigSelectorChain attribute) policy_override (class in ejkernel.ops.config.selection)

prefix_filter (ejkernel.ops.execution.profiler.Profiler attribute) PREFIX_FN (ejkernel.ops.execution.tuning.FNAutotuner attribute) prepare() (ejkernel.ops.core.kernel.Kernel method) priority (ejkernel.kernels._registry.KernelSpec attribute), [1] profile_time_by_function_id() (ejkernel.ops.execution.profiler.Profiler method) Profiler (class in ejkernel.ops.execution.profiler) profiler (ejkernel.ops.execution.tuning.FNAutotuner attribute) profiling_samples (ejkernel.ops.execution.tuning.FNAutotuner attribute) ProfilingError ProgressLogger (class in ejkernel.loggings) put() (ejkernel.ops.config.cache.ConfigCache method) (ejkernel.ops.config.persistent.PersistentCache method)

Q

q_blocksize (ejkernel.ops.utils.datacarrier.BwdParams attribute), [1] (ejkernel.ops.utils.datacarrier.FwdParams attribute), [1]

quiet() (in module ejkernel.callib._utils)

R

record_invocation() (in module ejkernel.ops.registry) register() (ejkernel.kernels._registry.KernelRegistry method) require_tf (ejkernel.ops.execution.profiler.Profiler attribute)

run() (ejkernel.ops.core.kernel.Kernel method) run_shard_map() (ejkernel.ops.core.kernel.Kernel method) run_shard_map_gpu() (ejkernel.ops.core.kernel.Kernel method)

S

save_compiled_fn() (in module ejkernel.callib._ejit) seconds (ejkernel.ops.execution.tuning.Measurement attribute), [1] shape (ejkernel.callib._utils.ShapeDtype property) ShapeDtype (class in ejkernel.callib._utils) sharding_fingerprint() (in module ejkernel.ops.utils.fingerprint) shared_mem_bytes (ejkernel.callib._triton_call.CompilationResult attribute), [1] short_hash() (in module ejkernel.ops.utils.fingerprint)

size() (ejkernel.ops.config.cache.ConfigCache method) smart_compile() (in module ejkernel.callib._ejit) stable_json() (in module ejkernel.ops.utils.fingerprint) stamp (ejkernel.ops.core.kernel.Invocation attribute), [1] stamp_prefix (ejkernel.ops.execution.executor.Executor attribute) start_time (ejkernel.loggings.ProgressLogger attribute) strides_from_shape() (in module ejkernel.callib._utils)

T

t_mean (ejkernel.ops.execution.tuning.TimingResult attribute), [1] t_std (ejkernel.ops.execution.tuning.TimingResult attribute), [1] timing_rounds (ejkernel.ops.execution.tuning.FNAutotuner attribute) timing_warmup_iterations (ejkernel.ops.execution.tuning.FNAutotuner attribute) TimingResult (class in ejkernel.ops.execution.tuning) to_json() (in module ejkernel.ops.utils.serialize) TPU (ejkernel.kernels._registry.Backend attribute), [1] TRITON (ejkernel.kernels._registry.Platform attribute), [1]

triton_call() (in module ejkernel.callib._triton_call) triton_kernel_call_abstract_eval() (in module ejkernel.callib._triton_call) triton_kernel_call_lowering() (in module ejkernel.callib._triton_call) triton_kernel_call_raise_on_jvp() (in module ejkernel.callib._triton_call) triton_kernel_call_raise_on_vmap() (in module ejkernel.callib._triton_call) ttgir (ejkernel.callib._triton_call.CompilationResult attribute), [1] tune() (ejkernel.ops.execution.tuning.FNAutotuner method) Tuner (class in ejkernel.ops.config.selection) tuner (ejkernel.ops.config.selection.ConfigSelectorChain attribute)

U

update() (ejkernel.loggings.ProgressLogger method)

update_simple() (ejkernel.loggings.ProgressLogger method) use_tty (ejkernel.loggings.ProgressLogger attribute)

V

validate_backward (ejkernel.ops.config.selection.AutotunePolicy attribute), [1] validate_signatures() (ejkernel.kernels._registry.KernelRegistry method) verbose (ejkernel.ops.execution.profiler.Profiler attribute) version (ejkernel.ops.core.kernel.Kernel attribute), [1]

vjp() (ejkernel.ops.core.kernel.Kernel method) vjp_shard_map() (ejkernel.ops.core.kernel.Kernel method) vjp_shard_map_gpu() (ejkernel.ops.core.kernel.Kernel method) vmap_with_config() (in module ejkernel.ops.execution.batch)

W

warmup (ejkernel.ops.config.selection.Tuner attribute) (ejkernel.ops.execution.tuning.Autotuner attribute)

write() (ejkernel.callib._utils.DummyStream method)

X

XLA (ejkernel.kernels._registry.Platform attribute), [1]