1pub mod ad;
28pub mod apple_bnns;
29pub mod async_copy;
30pub mod attention_layout;
31pub mod audio;
32pub mod bytes;
33pub mod capability;
34pub mod const_check;
35pub mod dtype;
36pub mod dynamic;
37pub mod env;
38pub mod env_catalog;
39pub mod env_registry;
40pub mod equivalence;
41pub mod fft;
42pub mod gdn;
43pub mod graph;
44pub mod hir;
45pub mod infer;
46pub mod infer_shape;
47pub mod inspect;
48pub mod kernel_schedule;
49pub mod lanes;
50pub mod layout;
51pub mod lir;
52pub mod logical_kernel;
53pub mod lowp_codec;
54pub mod measure;
55pub mod mir;
56pub mod module;
57pub mod numeric_check;
58pub mod nvfp4;
59pub mod op;
60pub mod op_registry;
61pub mod opcodes;
62pub mod ops;
63pub mod perfetto;
64pub mod phase;
65pub mod pretty;
66pub mod provenance;
67pub mod quant;
68pub mod region_encode;
69pub mod repr_check;
70pub mod residual;
71pub mod rng;
72pub mod tensor_inspect;
73pub use nvfp4::{FP4_E2M1_LUT, NVFP4_GROUP_SIZE, fp4_e2m1_to_f32, fp8_e4m3_scale_to_f32};
74pub mod binding_manifest;
75pub mod component;
76pub mod hir_extension;
77pub mod reflect;
78pub mod rf;
79#[cfg(feature = "test-support")]
82pub mod sample_ops;
83#[cfg(feature = "serialize")]
84pub mod serialize;
85pub mod shape;
86pub mod target;
87pub mod text;
88pub mod variant;
89pub mod verify;
90
91pub use ad::AdPipelineStage;
92pub use async_copy::{AsyncCopy, BarrierToken, DoubleBuffer, SyncCopy};
93pub use attention_layout::{
94 ATTENTION_FLASH_MAX_HEAD_DIM, AttentionGeom, AttentionLaunchStrides,
95 attention_dispatch_use_row, attention_geom, attention_launch_strides, cpu_attention_bshd,
96 cpu_attention_packed_bshd_qkv, detect_packed_bshd_qkv_attention, mask_strides_bhsd,
97 mask_strides_for_shape, packed_bshd_narrow_elidable, packed_bshd_qkv_strides, strides_bhsd,
98 strides_bshd, strides_for_shape,
99};
100pub use capability::{N_KINDS, OpCaps, fused_kinds};
101pub use dtype::{DType, Element, ElementSubtype, scalar_constant_bytes};
102pub use dynamic::sym;
103pub use dynamic::{
104 DimEnv, bind_graph, collect_dynamic_symbols, has_dynamic_dims, infer_bindings_from_f32_inputs,
105 infer_bindings_from_inputs, same_binding, sync_concat_shapes, sync_expand_ops,
106 sync_graph_shapes, sync_narrow_ops, sync_reshape_ops,
107};
108pub use env::{RlxEnv, RuntimeOverrides, flag, is_unset, parse_or, set, unset, var, var_os};
109pub use env_catalog::{EnvVarDoc, format_catalog as format_env_catalog};
110pub use env_registry::{
111 EnvKind, EnvLayer, EnvStability, EnvVarEntry, REGISTRY as ENV_REGISTRY, canonical_name,
112 format_registry_markdown, is_registered, lookup as env_lookup, public_catalog_docs,
113};
114
115pub fn env_catalog() -> &'static [EnvVarDoc] {
117 public_catalog_docs()
118}
119
120pub use equivalence::{
121 IgnoreConfig, fingerprint, node_value_key, ops_deep_eq, structural_hash, structurally_eq,
122};
123pub use fft::{FftGpuPlan, FftMeta, FftNorm, fft_meta, fftn_axes_all, normalize_fftn_axes};
124pub use gdn::GdnBackwardLayout;
125pub use graph::{Graph, Node, NodeId};
126pub use hir::{
127 FusionPolicy, HirGraphExt, HirModule, HirMut, HirNode, HirNodeId, HirOp,
128 window_token_gather_bsn, window_token_scatter_bsn,
129};
130pub use infer::GraphExt;
131pub use inspect::{
132 inspect_buffer_plan, inspect_graph, inspect_graph_diff, inspect_hir, inspect_hir_stats,
133 inspect_lir, inspect_mir, inspect_mir_diff, inspect_mir_stats,
134};
135pub use layout::{Coord2, Ragged, ShapeTuple, Strides2, Strides3, Tile2, Tile3};
136pub use lir::{
137 LirBufferPlan, LirBufferSlot, LirFingerprint, LirIoManifest, LirModule, LirViewAlias,
138};
139pub use logical_kernel::{
140 KernelDispatchConfig, KernelDispatchPolicy, LogicalKernelEntry, logical_kinds_in_graph,
141 registered_logical_kernels, should_lower_to_common,
142};
143pub use measure::{CacheBuster, Tick, time_ns};
144pub use mir::{MirModule, MirNode, MirNodeId, MirOp};
145pub use module::{GraphModule, GraphStage};
146pub use numeric_check::{BadValue, DebugMode, DebugScanner, NanReport, check_node, first_bad};
147pub use op::{
148 AdaNormKind, ChainOperand, ChainStep, Op, OpKind, PadMode, RegionPrologue, RopeStyle,
149 ScatterNdReduction, SpdMatFn, SynthBwdWrt, SynthKind, TransformStep,
150};
151pub use op_registry::{
152 CustomOpError, JvpContext, LowerContext, OpExtension, OpRegistry, VjpContext, VmapContext,
153 global_registry, is_op_registered, lookup_op, register_op, register_op_strict,
154};
155pub use ops::attention::attention_kind_op;
156pub use ops::backward::{unpack_ada_layer_norm_backward, unpack_gated_residual_backward};
157pub use ops::dsp::{FirMode, iir_impulse_response};
158pub use phase::{Phase, PhaseSchedule, derive_phases};
159pub use provenance::{NodeOrigin, node_label, stamp_pass_origins};
160pub use public_catalog_docs as ENV_CATALOG;
162pub use quant::{QuantMap, QuantScheme, ScaleLayout, ScaledFormat};
163pub use region_encode::{
164 FK_BATCH_SINGLE_KERNEL_MAX, PrologueLaunchGrid, REGION_META_WORDS, REGION_PROLOGUE_NONE,
165 REGION_PROLOGUE_RESIZE_NEAREST_2X_NCHW, RegionNchwDims, batch_region_slice_dst_off_f32,
166 batch_region_slice_elems, batch_region_slice_shape, encode_chain_operand, encode_chain_steps,
167 encode_elementwise_region_meta, encode_prologue_tail, fk_batch_single_kernel_enabled,
168 fk_batch_use_single_launch,
169};
170pub use rng::{
171 BNNS_AES_CTR_STATE_SIZE, BnnsAesCtr128, GUMBEL_UNIFORM_EPS, NativeBnnsGumbelMax, Philox4x32,
172 RngBackend, RngOptions, argmax_f32, bnns_categorical_from_uniform, bnns_entropy_seed,
173 bnns_generator_state_is_fresh, bnns_gumbel_from_uniform, bnns_gumbel01_from_bytes,
174 bnns_seed_from_generator_state, bnns_uniform_f32_from_bytes, bnns_uniform_f32_from_u32,
175 bnns_uniform_int_from_u32, combine_seed, fill_normal_like, fill_uniform_like,
176 gumbel_max_argmax_with, gumbel01_from_uniform, ort_engine_seed,
177};
178#[cfg(feature = "serialize")]
179pub use serialize::{
180 LIR_BIN_VERSION, hir_from_json, hir_to_json, lir_from_bytes, lir_from_json, lir_to_bytes,
181 lir_to_json,
182};
183pub use verify::{VerifyError, verify, verify_all, verify_shapes, verify_unique_leaf_names};
184
185pub fn hir_to_graph(hir: HirModule) -> Result<Graph, hir::LowerError> {
187 Ok(hir.lower_to_mir()?.into_graph())
188}
189pub use binding_manifest::{BindingManifest, IoBindingEntry, WeightBlock};
190pub use component::{CompilationMode, ModelComponent};
191pub use hir_extension::{
192 HirExtensionFn, apply_hir_extensions, apply_hir_extensions_named, register_hir_extension,
193 registered_hir_extensions,
194};
195pub use reflect::{
196 BlockSpecialization, HirReflection, ManifestDiff, MirReflection, SpecializeBlockRecord,
197 layout_for_binding, layout_from_lir, probe_block_specialization, symbolic_layout_hint,
198};
199pub use rf::{
200 complex_div, const_f32, cs_degen_z_in, find_param_node, find_param_nodes, mag2, s11_from_z,
201 scalar_f32,
202};
203pub use shape::{Dim, DimBinding, Shape, ada_modulation_launch, ada_modulation_lead_pack};
204pub use variant::{ModelPhase, ModelVariant};