rlx_ir/
logical_kernel.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4//! One logical kernel, many backends — dispatch policy and registry.
5//!
6//! A **logical kernel** is a single [`OpKind`] (e.g. [`OpKind::GaussianSplatRender`]) with a
7//! documented semantic contract. Backends may provide a **native** implementation (fast path:
8//! custom thunk, MSL, MPS, etc.). When native is unavailable or [`KernelDispatchPolicy::ForceCommon`]
9//! is set, the compiler lowers to a **common** subgraph built only from primitive MIR ops so each
10//! backend schedules the same math through its usual fusion/GEMM/elementwise paths.
11//!
12//! Native kernels are never removed from backends; common lowering is additive.
13
14use crate::env;
15use crate::op::OpKind;
16
17pub mod splat_common;
18
19/// When to use native backend kernels vs the shared IR common body.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
21pub enum KernelDispatchPolicy {
22    /// Native thunk when `OpKind` is in the backend `supported_ops`; else common IR lower.
23    #[default]
24    PreferNative,
25    /// Always lower registered logical kernels to common IR (parity / minimal backends).
26    ForceCommon,
27    /// Never common-lower; legalization must succeed with native ops only.
28    ForceNative,
29}
30
31impl KernelDispatchPolicy {
32    pub fn from_env() -> Self {
33        let v = env::var("KERNEL_DISPATCH").or_else(|| env::var("RLX_KERNEL_DISPATCH"));
34        match v.as_deref() {
35            Some("common") | Some("force_common") | Some("ForceCommon") => Self::ForceCommon,
36            Some("native") | Some("force_native") | Some("ForceNative") => Self::ForceNative,
37            _ => Self::PreferNative,
38        }
39    }
40}
41
42/// Registered logical kernel: native [`OpKind`] plus optional common lower pass name.
43#[derive(Debug, Clone, Copy)]
44pub struct LogicalKernelEntry {
45    pub kind: OpKind,
46    /// Human-readable id (logging / docs).
47    pub name: &'static str,
48}
49
50/// Logical kernels that have a registered common IR body in `rlx-fusion`.
51pub fn registered_logical_kernels() -> &'static [LogicalKernelEntry] {
52    &[
53        LogicalKernelEntry {
54            kind: OpKind::GroupNorm,
55            name: "group_norm",
56        },
57        LogicalKernelEntry {
58            kind: OpKind::BatchNormInference,
59            name: "batch_norm_inference",
60        },
61        LogicalKernelEntry {
62            kind: OpKind::ResizeNearest2x,
63            name: "resize_nearest_2x",
64        },
65        LogicalKernelEntry {
66            kind: OpKind::GaussianSplatRender,
67            name: "gaussian_splat_render",
68        },
69        LogicalKernelEntry {
70            kind: OpKind::GaussianSplatRenderBackward,
71            name: "gaussian_splat_render_backward",
72        },
73    ]
74}
75
76/// Per-compile overrides on top of [`KernelDispatchPolicy`].
77#[derive(Debug, Clone, Copy, Default)]
78pub struct KernelDispatchConfig {
79    pub policy: KernelDispatchPolicy,
80    /// Always common-lower these kinds (e.g. splat on CPU while keeping native matmul).
81    pub force_common_kinds: &'static [OpKind],
82    /// Never common-lower these kinds (overrides `ForceCommon` for listed kinds).
83    pub force_native_kinds: &'static [OpKind],
84}
85
86impl KernelDispatchConfig {
87    pub fn new(policy: KernelDispatchPolicy) -> Self {
88        Self {
89            policy,
90            ..Self::default()
91        }
92    }
93
94    pub fn from_env() -> Self {
95        Self::new(KernelDispatchPolicy::from_env())
96    }
97}
98
99/// Whether `kind` should be common-lowered for this backend claim set and config.
100pub fn should_lower_to_common(
101    kind: OpKind,
102    supported: &[OpKind],
103    config: KernelDispatchConfig,
104) -> bool {
105    if !registered_logical_kernels().iter().any(|e| e.kind == kind) {
106        return false;
107    }
108    if config.force_native_kinds.contains(&kind) {
109        return false;
110    }
111    if config.force_common_kinds.contains(&kind) {
112        return true;
113    }
114    match config.policy {
115        KernelDispatchPolicy::ForceCommon => true,
116        KernelDispatchPolicy::ForceNative => false,
117        KernelDispatchPolicy::PreferNative => !supported.is_empty() && !supported.contains(&kind),
118    }
119}
120
121/// Op kinds that appear in the graph and may need common lowering.
122pub fn logical_kinds_in_graph(
123    graph: &crate::Graph,
124    supported: &[OpKind],
125    config: KernelDispatchConfig,
126) -> Vec<OpKind> {
127    let mut kinds = Vec::new();
128    for node in graph.nodes() {
129        let k = node.op.kind();
130        if should_lower_to_common(k, supported, config) && !kinds.contains(&k) {
131            kinds.push(k);
132        }
133    }
134    kinds
135}
rlx_ir/logical_kernel.rs

rlx_ir/
logical_kernel.rs
