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