Skip to main content

hetorch.compiler.compiler

Main HETorch compiler

Classes​

HETorchCompiler(context: <class 'hetorch.compiler.context.CompilationContext'>, pipeline: <class 'hetorch.passes.pipeline.PassPipeline'>)​

Main compiler for transforming PyTorch models to HE operations

The compiler orchestrates the compilation pipeline:

  1. Capture model graph using torch.fx
  2. Apply transformation passes
  3. Return compiled graph module

Args: context: Compilation context with scheme, params, and backend pipeline: Pass pipeline to apply

Methods:

__init__(self, context: hetorch.compiler.context.CompilationContext, pipeline: hetorch.passes.pipeline.PassPipeline)​

Initialize self. See help(type(self)) for accurate signature.

compile(self, model: torch.nn.modules.module.Module, example_inputs: Union[torch.Tensor, Tuple[torch.Tensor, ...]], concrete_args: Optional[dict] = None) -\> torch.fx.graph_module.GraphModule​

Compile a PyTorch model to HE operations

Args: model: PyTorch model to compile example_inputs: Example inputs for tracing concrete_args: Optional concrete arguments for tracing

Returns: Compiled graph module with HE operations

Raises: RuntimeError: If compilation fails

compile_function(self, func: \<built-in function callable\>, example_inputs: Union[torch.Tensor, Tuple[torch.Tensor, ...]]) -\> torch.fx.graph_module.GraphModule​

Compile a function to HE operations

Args: func: Function to compile example_inputs: Example inputs for tracing

Returns: Compiled graph module with HE operations