Crate torsh_jit 

ToRSh JIT compilation and kernel fusion module

This module provides Just-In-Time (JIT) compilation capabilities for ToRSh, enabling automatic kernel fusion and optimization of computational graphs.

Features

• Kernel Fusion: Automatically fuses compatible operations to reduce memory bandwidth
• Graph Optimization: Applies various optimization passes to the computation graph
• Multiple Backends: Supports Cranelift and (future) MLIR code generation
• TorchScript-like API: Compatible with PyTorch’s JIT compilation model

Example

use torsh_jit::{jit_compile, FusionStrategy};

// Define a model
let model = MyModel::new();

// JIT compile with fusion enabled
let jit_model = jit_compile(model, FusionStrategy::Aggressive)?;

// Use the JIT-compiled model
let output = jit_model.forward(input);

Re-exports

pub use abstract_interpretation::AbstractAnalysisResult;

pub use abstract_interpretation::AbstractDomain;

pub use abstract_interpretation::AbstractInterpretationConfig;

pub use abstract_interpretation::AbstractInterpreter;

pub use abstract_interpretation::AbstractValue;

pub use abstract_interpretation::ConstantDomain;

pub use abstract_interpretation::IntervalDomain;

pub use abstract_interpretation::SignDomain;

pub use adaptive_compilation::AdaptiveCompiler;

pub use adaptive_compilation::AdaptiveConfig;

pub use adaptive_compilation::CompilationStrategy;

pub use adaptive_compilation::PerformanceMetrics;

pub use codegen::CodeGenerator;

pub use const_eval::ConstEvalConfig;

pub use const_eval::ConstantEvaluator;

pub use const_eval::ConstantValue;

pub use const_eval::EvaluationResult;

pub use custom_ops::get_custom_op;

pub use custom_ops::list_custom_ops;

pub use custom_ops::register_custom_op;

pub use custom_ops::CustomOpBuilder;

pub use debug_symbols::DebugSymbolConfig;

pub use debug_symbols::DebugSymbolManager;

pub use debug_symbols::SourceLocation;

pub use debug_symbols::SymbolTable;

pub use debugger::BreakpointLocation;

pub use debugger::DebugCommand;

pub use debugger::DebugSession;

pub use debugger::DebugState;

pub use debugger::DebugValue;

pub use debugger::DebuggerConfig;

pub use debugger::ExecutionLocation;

pub use debugger::InspectionTarget;

pub use debugger::JitDebugger;

pub use differentiable_compilation::CompilationParams;

pub use differentiable_compilation::CompilationTrainer;

pub use differentiable_compilation::DiffCompilationResult;

pub use differentiable_compilation::DifferentiableCompiler;

pub use differentiable_compilation::GumbelSoftmax;

pub use differentiable_compilation::PerformanceMetrics as DiffPerformanceMetrics;

pub use differentiable_compilation::SoftDecision;

pub use error_diagnostics::DiagnosticError;

pub use error_diagnostics::ErrorCategory;

pub use error_diagnostics::ErrorDiagnosticsManager;

pub use error_diagnostics::ErrorSeverity;

pub use fusion::FusionStrategy;

pub use fusion::KernelFusion;

pub use generics::create_type_param;

pub use generics::shape_constraint;

pub use generics::trait_constraint;

pub use generics::GenericFunctionManager;

pub use generics::GenericFunctionTemplate;

pub use generics::ParameterKind;

pub use generics::TypeConstraint;

pub use generics::TypeParameter;

pub use graph::ComputationGraph;

pub use graph::Edge;

pub use graph::Node;

pub use graph::NodeId;

pub use hardware_tuning::Architecture;

pub use hardware_tuning::HardwareInfo;

pub use hardware_tuning::HardwareTuner;

pub use hardware_tuning::HardwareTuningConfig;

pub use hardware_tuning::TuningRecommendation;

pub use llvm_backend::LlvmBackend;

pub use llvm_backend::LlvmOptimizer;

pub use metaprogramming::CodeTemplate;

pub use metaprogramming::DynamicCodeGenerator;

pub use metaprogramming::GeneratedCode;

pub use metaprogramming::GraphReflection;

pub use metaprogramming::MacroDefinition;

pub use metaprogramming::MetaprogrammingEngine;

pub use metaprogramming::RuntimeReflector;

pub use metaprogramming::TemplateParameters;

pub use mlir_backend::MlirBackend;

pub use mlir_backend::MlirOptimizer;

pub use mlir_backend::MlirPass;

pub use neural_compilation::CompilationStrategy as NeuralCompilationStrategy;

pub use neural_compilation::GraphFeatures;

pub use neural_compilation::NeuralCompiler;

pub use neural_compilation::NeuralCompilerConfig;

pub use neural_compilation::OptimizationDecision;

pub use optimizer::GraphOptimizer;

pub use partial_evaluation::ConstantFolder;

pub use partial_evaluation::EvaluationStatistics;

pub use partial_evaluation::FunctionSpecializer;

pub use partial_evaluation::OptimizedGraph;

pub use partial_evaluation::OptimizedIrModule;

pub use partial_evaluation::PartialEvalConfig;

pub use partial_evaluation::PartialEvaluator;

pub use pgo::OptimizationRecommendation as PgoRecommendation;

pub use pgo::OptimizationType as PgoOptimizationType;

pub use pgo::PgoConfig;

pub use pgo::ProfileGuidedOptimizer;

pub use plugin_system::load_all_plugins;

pub use plugin_system::load_plugin;

pub use plugin_system::Plugin;

pub use plugin_system::PluginCapability;

pub use plugin_system::PluginManager;

pub use plugin_system::PluginMetadata;

pub use plugin_system::PluginRegistry;

pub use polyhedral_optimization::AffineExpr;

pub use polyhedral_optimization::AffineSchedule;

pub use polyhedral_optimization::LoopNest;

pub use polyhedral_optimization::PolyhedralConfig;

pub use polyhedral_optimization::PolyhedralOptimizer;

pub use polyhedral_optimization::Polyhedron;

pub use polyhedral_optimization::TransformationMatrix;

pub use polyhedral_optimization::TransformationType;

pub use probabilistic_compilation::BetaDistribution;

pub use probabilistic_compilation::MonteCarloResult;

pub use probabilistic_compilation::NormalDistribution;

pub use probabilistic_compilation::ProbabilisticCompilationResult;

pub use probabilistic_compilation::ProbabilisticCompiler;

pub use probabilistic_compilation::ProbabilisticConfig;

pub use probabilistic_compilation::ProbabilisticPerformance;

pub use probabilistic_compilation::UncertainDecision;

pub use profiler::PerformanceEvent;

pub use profiler::ProfilerConfig;

pub use profiler::ProfilerManager;

pub use profiler::ProfilingSession;

pub use program_synthesis::ExampleBuilder;

pub use program_synthesis::ProgramSynthesizer;

pub use program_synthesis::SynthesisExample;

pub use program_synthesis::SynthesisResult;

pub use program_synthesis::SynthesisStrategy;

pub use program_synthesis::SynthesisTemplate;

pub use program_synthesis::SynthesisValue;

pub use runtime::JitRuntime;

pub use script::export_torchscript;

pub use script::import_torchscript;

pub use script::ScriptCompiler;

pub use specialization::create_specialized_type;

pub use specialization::SpecializationConfig;

pub use specialization::SpecializedType;

pub use specialization::TypeSpecializer;

pub use speculative_optimization::DeoptimizationEvent;

pub use speculative_optimization::SpeculationResult;

pub use speculative_optimization::SpeculativeConfig;

pub use speculative_optimization::SpeculativeOptimizer;

pub use symbolic_execution::Constraint;

pub use symbolic_execution::ConstraintSet;

pub use symbolic_execution::ExecutionState;

pub use symbolic_execution::SymbolicExecutionConfig;

pub use symbolic_execution::SymbolicExecutionEngine;

pub use symbolic_execution::SymbolicExecutionResult;

pub use symbolic_execution::SymbolicGraph;

pub use symbolic_execution::SymbolicValue;

pub use trace_viz::TraceEvent;

pub use trace_viz::TraceVisualizationManager;

pub use trace_viz::VisualizationConfig;

pub use trace_viz::VisualizationSession;

pub use optimization_advisor::analyze_computation_graph;

pub use optimization_advisor::analyze_with_benchmarks;

pub use optimization_advisor::analyze_with_profiling;

pub use optimization_advisor::create_advisor;

pub use optimization_advisor::create_advisor_with_config;

pub use optimization_advisor::create_fast_config;

pub use optimization_advisor::create_production_config;

pub use optimization_advisor::create_thorough_config;

pub use optimization_advisor::quick_analyze;

pub use optimization_advisor::AdvisorConfig;

pub use optimization_advisor::AnalysisInput;

pub use optimization_advisor::CostBenefitAnalysis;

pub use optimization_advisor::OptimizationAdvisor;

pub use optimization_advisor::OptimizationRecommendation;

pub use optimization_advisor::OptimizationReport;

pub use optimization_advisor::OptimizationType;

pub use optimization_advisor::PatternAnalysis;

pub use optimization_advisor::PerformanceAnalysis;

pub use optimization_advisor::SystemConstraints;

pub use optimization_advisor::TargetPlatform;

pub use optimization_advisor::UserPreferences;

Modules

abstract_interpretation

Comprehensive Abstract Interpretation Framework

adaptive_compilation

Adaptive Compilation for ToRSh JIT

advisor

Optimization advisor system for JIT compilation

analysis

Analysis utilities for JIT compilation

benchmarking

Comprehensive Benchmarking Suite for ToRSh JIT

codegen

Code generation backend for JIT compilation

const_eval

Compile-Time Evaluation for ToRSh JIT

cranelift_backend

Cranelift backend for JIT code generation

custom_ops

Custom operators for JIT compilation

debug_symbols

Debug symbol generation for JIT compilation

debugger

Interactive JIT debugger with breakpoints, stepping, and introspection

differentiable_compilation

Differentiable Compilation

error_diagnostics

Error diagnostics for JIT compilation

fusion

Kernel fusion optimization for JIT compilation

generics

Generic function support for JIT compilation

graph

Computation graph representation and analysis

hardware_tuning

Hardware-Specific Tuning for ToRSh JIT

ir

Intermediate Representation (IR) for JIT compilation

llvm_backend

LLVM Backend for ToRSh JIT Compilation

lowering

Graph to IR lowering pass

metaprogramming

Metaprogramming support for dynamic code generation and reflection

mlir_backend

MLIR Backend for ToRSh JIT Compilation

neural_compilation

Neural Compilation

optimization_advisor

Intelligent optimization advisor with automated recommendations and analysis

optimizer

Graph optimization passes for JIT compilation

partial_evaluation

Partial evaluation system for compile-time optimization and specialization

pgo

Profile-Guided Optimization for ToRSh JIT

plugin_system

Plugin System for ToRSh JIT Compilation

polyhedral_optimization

Polyhedral Optimization

prelude

Prelude module for convenient imports

probabilistic_compilation

Probabilistic Compilation

profiler

Profiler integration for JIT compilation

program_synthesis

Program Synthesis Module

runtime

Runtime execution engine for JIT-compiled code

script

Script mode for JIT compilation

specialization

Type specialization for JIT compilation

speculative_optimization

Speculative Optimization for ToRSh JIT

symbolic_execution

Symbolic execution engine for path analysis and constraint solving

trace_viz

Trace visualization for JIT compilation

tracing

Graph capture and tracing for JIT compilation

type_inference

Type inference and shape propagation for JIT compilation

utils

Utility functions for common JIT operations

Macros

register_op

Macro for easier custom operator registration

trace

Utility macro for easy tracing

Structs

CompiledKernel

Compiled kernel representation

CompiledModule

A compiled module ready for execution

ExecutionStats

Execution statistics

JitCompiler

Main JIT compiler interface

JitConfig

JIT compilation configuration

KernelMetadata

Kernel metadata for runtime execution

TensorDesc

Tensor description for kernel interface

TensorRef

Placeholder for tensor references (will be properly integrated with torsh-tensor)

Enums

JitError

JIT compilation errors

Constants

VERSION

VERSION_MAJOR

VERSION_MINOR

VERSION_PATCH

Traits

ScriptableModule

Trait for scriptable modules

Functions

script

JIT script a module

trace

JIT trace a function to capture its computation graph

Type Aliases

IrFunction

IrInstruction

JitResult