Skip to main content

hetorch.passes.builtin.rescaling

RescalingInsertionPass: Insert rescaling operations after multiplications (CKKS)

Classes​

RescalingInsertionPass(strategy: <class 'str'>, target_level: typing.Optional[int])​

Insert rescaling operations after multiplications for CKKS scheme

In CKKS, multiplications increase the scale of ciphertexts. Rescaling operations are needed to bring the scale back to the target level.

This pass supports two strategies:

  • "eager": Insert rescale after every multiplication
  • "lazy": Fuse consecutive multiplications before rescaling (future work)

Attributes: strategy: Rescaling strategy ("eager" or "lazy") target_level: Target multiplication depth (optional)

Methods:

__init__(self, strategy: str = 'eager', target_level: Optional[int] = None)​

Initialize RescalingInsertionPass

Args: strategy: "eager" (rescale after each mult) or "lazy" (fuse mults) target_level: Target multiplication depth (None = no limit)

analyze_cost(self, graph_module: torch.fx.graph_module.GraphModule, context: hetorch.compiler.context.CompilationContext) -\> hetorch.backend.cost_model.CostAnalysis​

Analyze cost impact of this pass

Args: graph_module: Graph module to analyze context: Compilation context

Returns: Cost analysis result

transform(self, graph_module: torch.fx.graph_module.GraphModule, context: hetorch.compiler.context.CompilationContext) -\> torch.fx.graph_module.GraphModule​

Apply rescaling insertion transformation

Args: graph_module: Input graph module context: Compilation context

Returns: Transformed graph module with rescaling operations

validate(self, graph_module: torch.fx.graph_module.GraphModule, context: hetorch.compiler.context.CompilationContext) -\> bool​

Validate that the pass can be applied

Args: graph_module: Graph module to validate context: Compilation context

Returns: True if validation passes

Raises: PassValidationError: If validation fails