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