Expand description
JIT Compilation for Axonml
This crate provides Just-In-Time compilation for tensor operations, enabling significant performance improvements through:
- Operation tracing and graph construction
- Graph optimization (fusion, constant folding, dead code elimination)
- Native code generation via Cranelift
- Compiled function caching
§Example
ⓘ
use axonml_jit::{JitCompiler, trace};
// Trace operations to build a computation graph
let graph = trace(|tracer| {
let a = tracer.input("a", &[2, 3]);
let b = tracer.input("b", &[2, 3]);
let c = a.add(&b);
let d = c.mul_scalar(2.0);
tracer.output("result", d)
});
// Compile the graph
let compiler = JitCompiler::new();
let compiled = compiler.compile(&graph)?;
// Execute with real tensors
let a = Tensor::randn(&[2, 3]);
let b = Tensor::randn(&[2, 3]);
let result = compiled.run(&[("a", &a), ("b", &b)])?;@version 0.1.0 @author AutomataNexus Development Team
Re-exports§
pub use ir::Graph;pub use ir::Node;pub use ir::NodeId;pub use ir::Op;pub use ir::DataType;pub use ir::Shape;pub use trace::Tracer;pub use trace::TracedValue;pub use trace::trace;pub use optimize::Optimizer;pub use optimize::OptimizationPass;pub use codegen::JitCompiler;pub use codegen::CompiledFunction;pub use cache::FunctionCache;pub use error::JitError;pub use error::JitResult;