Basic Concepts

This guide introduces the core concepts and abstractions in HETorch. Understanding these concepts will help you effectively use and extend the framework.

Overview: The HETorch Stack

HETorch is organized into layers, each building on the previous one:

┌─────────────────────────────────────────┐
│     PyTorch Model (nn.Module)           │  ← Your neural network
└─────────────────────────────────────────┘
                  │
                  ▼
┌─────────────────────────────────────────┐
│     Frontend (torch.fx)                 │  ← Graph capture
└─────────────────────────────────────────┘
                  │
                  ▼
┌─────────────────────────────────────────┐
│     HE-Aware IR (fx.GraphModule)        │  ← Intermediate representation
└─────────────────────────────────────────┘
                  │
                  ▼
┌─────────────────────────────────────────┐
│     Pass Pipeline                       │  ← Transformations
└─────────────────────────────────────────┘
                  │
                  ▼
┌─────────────────────────────────────────┐
│     Backend (Fake or Real)              │  ← Execution
└─────────────────────────────────────────┘

Let's explore each layer and the key abstractions.

1. Homomorphic Encryption Basics

What is Homomorphic Encryption?

Homomorphic Encryption (HE) allows computations on encrypted data without decryption:

Encrypt(x) + Encrypt(y) = Encrypt(x + y)
Encrypt(x) × Encrypt(y) = Encrypt(x × y)

This enables privacy-preserving machine learning: run inference on encrypted data without seeing the actual values.

Why is HE Compilation Challenging?

1. Limited Operations: HE supports addition and multiplication, but not arbitrary operations
2. Noise Accumulation: Each operation adds noise; too much noise corrupts results
3. Performance: HE operations are 1000-10000× slower than plaintext
4. Packing: Data must be packed into ciphertext slots efficiently

HETorch addresses these challenges through compilation and optimization.

2. Core Abstractions

2.1 HE Schemes

HETorch supports three main HE schemes:

CKKS (Cheon-Kim-Kim-Song)

• Arithmetic: Approximate (floating-point-like)
• Use Case: Neural networks, signal processing
• Key Feature: Supports rescaling to manage scale
• Example:

  from hetorch import HEScheme, CKKSParameters
  
  scheme = HEScheme.CKKS
  params = CKKSParameters(
      poly_modulus_degree=8192,
      coeff_modulus=[60, 40, 40, 60],
      scale=2**40
  )

BFV (Brakerski-Fan-Vercauteren)

• Arithmetic: Exact (integer)
• Use Case: Integer computations, decision trees
• Key Feature: Exact results, no approximation
• Example:

  from hetorch import HEScheme, BFVParameters
  
  scheme = HEScheme.BFV
  params = BFVParameters(
      poly_modulus_degree=8192,
      coeff_modulus=[60, 60, 60],
      plain_modulus=1024
  )

BGV (Brakerski-Gentry-Vaikuntanathan)

• Arithmetic: Exact (integer)
• Use Case: Similar to BFV
• Key Feature: Different noise management than BFV
• Example:

  from hetorch import HEScheme, BGVParameters
  
  scheme = HEScheme.BGV
  params = BGVParameters(
      poly_modulus_degree=8192,
      coeff_modulus=[60, 60, 60],
      plain_modulus=1024
  )

Choosing a Scheme: For neural networks, use CKKS. For exact integer computations, use BFV or BGV.

2.2 Compilation Context

The CompilationContext is the central object that maintains state throughout compilation:

from hetorch import CompilationContext, HEScheme, CKKSParameters, FakeBackend

context = CompilationContext(
    scheme=HEScheme.CKKS,           # Which HE scheme to use
    params=CKKSParameters(...),     # Encryption parameters
    backend=FakeBackend(),          # Backend implementation
    metadata={}                     # Optional custom metadata
)

Key Properties:

• Immutable during pass execution: Ensures consistency
• Shared across all passes: Provides global state
• Contains backend: Provides HE operations and cost models

Relationship to Other Components:

CompilationContext
    ├── scheme: HEScheme (CKKS/BFV/BGV)
    ├── params: EncryptionParameters
    ├── backend: HEBackend
    └── metadata: Dict[str, Any]

2.3 Ciphertext Metadata

CiphertextInfo describes the state of encrypted tensors:

from hetorch.core import CiphertextInfo, PackingInfo

ciphertext_info = CiphertextInfo(
    shape=(1, 10),              # Tensor shape
    dtype=torch.float32,        # Data type
    level=3,                    # Remaining multiplication depth
    scale=2**40,                # Current scale (CKKS)
    packing=PackingInfo(...),   # How data is packed
    noise_budget=85.0           # Estimated noise budget (optional)
)

Why Metadata?

• Passes need to reason about ciphertext properties without executing HE operations
• Metadata is attached to graph nodes: node.meta['ciphertext_info']
• Enables static analysis and optimization

Key Fields:

• level: Tracks multiplication depth (decreases with each multiplication)
• scale: Tracks scaling factor (CKKS only)
• noise_budget: Estimates remaining noise capacity
• packing: Describes slot layout

2.4 Packing Strategies

PackingInfo describes how data is packed into ciphertext slots:

from hetorch.core import PackingInfo

packing = PackingInfo(
    strategy="row_major",       # Packing strategy name
    slot_count=8192,            # Number of slots
    dimensions={"rows": 32, "cols": 256},  # Layout
    metadata={}                 # Strategy-specific data
)

Why Packing Matters:

• Ciphertexts have fixed slot count (e.g., 8192 slots)
• Efficient packing reduces ciphertext count
• Different operations require different packing

Common Strategies:

• row_major: Pack rows sequentially (good for matrix-vector)
• column_major: Pack columns sequentially
• diagonal: Pack diagonals (good for matrix-matrix)
• custom: User-defined packing

Example:

Tensor: [1, 2, 3, 4]
Slots:  [1, 2, 3, 4, 0, 0, 0, 0, ...]  (row_major)

3. Transformation Passes

Passes are the core of HETorch's compilation pipeline. Each pass transforms the computation graph.

Pass Lifecycle

Input Graph → Validate → Transform → Output Graph

1. Validate: Check preconditions (e.g., required passes ran)
2. Transform: Modify the graph
3. Return: New graph with transformations applied

Pass Categories

3.1 Graph Transformation Passes

Modify the computation graph:

• InputPackingPass: Annotate inputs with packing info
• NonlinearToPolynomialPass: Replace activations with polynomials
• LinearLayerBSGSPass: Optimize matrix operations

3.2 Scheme-Specific Passes

Handle scheme requirements:

• RescalingInsertionPass: Insert rescaling (CKKS)
• RelinearizationInsertionPass: Insert relinearization (all schemes)
• BootstrappingInsertionPass: Insert bootstrapping

3.3 Optimization Passes

Improve performance:

• DeadCodeEliminationPass: Remove unused nodes
• Lazy strategies: Reduce unnecessary operations

3.4 Analysis Passes

Extract information without modification:

• CostAnalysisPass: Estimate latency, memory, operations
• PrintGraphPass: Debug visualization

Pass Dependencies

Passes declare dependencies:

class MyPass(TransformationPass):
    name = "my_pass"
    requires = ["input_packed"]  # Must run after InputPackingPass
    provides = ["my_property"]   # Guarantees this property

Dependency Graph Example:

InputPackingPass
    ↓
NonlinearToPolynomialPass
    ↓
LinearLayerBSGSPass (requires: input_packed)
    ↓
RescalingInsertionPass
    ↓
BootstrappingInsertionPass (requires: rescaling_inserted)

Pass Pipeline

PassPipeline executes passes in sequence:

from hetorch.passes import PassPipeline
from hetorch.passes.builtin import (
    InputPackingPass,
    NonlinearToPolynomialPass,
    RescalingInsertionPass,
)

pipeline = PassPipeline([
    InputPackingPass(strategy="row_major"),
    NonlinearToPolynomialPass(degree=8),
    RescalingInsertionPass(strategy="lazy"),
])

# Run pipeline
transformed_graph = pipeline.run(graph, context)

Pipeline Execution:

1. Validate each pass before running
2. Execute passes in order
3. Pass context to each pass
4. Return final transformed graph

4. Intermediate Representation (IR)

HETorch uses torch.fx's graph representation with HE-specific operations.

Graph Structure

# torch.fx graph
graph = fx.Graph()

# Nodes represent operations
node = graph.call_function(
    torch.ops.hetorch.cadd,  # HE operation
    args=(ct1, ct2),         # Arguments
)

# Metadata attached to nodes
node.meta['ciphertext_info'] = CiphertextInfo(...)

HE Operations

HETorch defines custom operations:

Ciphertext-Ciphertext Operations

torch.ops.hetorch.cadd(ct1, ct2)    # Ciphertext addition
torch.ops.hetorch.cmult(ct1, ct2)   # Ciphertext multiplication
torch.ops.hetorch.rotate(ct, steps) # Ciphertext rotation

Plaintext-Ciphertext Operations

torch.ops.hetorch.padd(ct, pt)      # Plaintext addition
torch.ops.hetorch.pmult(ct, pt)     # Plaintext multiplication

Scheme-Specific Operations

torch.ops.hetorch.rescale(ct)       # Rescale (CKKS)
torch.ops.hetorch.relinearize(ct)   # Relinearize (reduce size)
torch.ops.hetorch.bootstrap(ct)     # Bootstrap (refresh noise)

Why Custom Operations?

1. Explicit HE semantics: Clear distinction between HE operations
2. torch.fx compatibility: Seamless integration with PyTorch
3. Type checking: Enable shape inference and validation
4. Extensibility: Easy to add new operations

5. Backend Interface

The backend provides HE operation implementations.

Backend Abstraction

from hetorch.backend import HEBackend

class HEBackend(ABC):
    @abstractmethod
    def cadd(self, ct1, ct2):
        """Ciphertext addition"""
        pass

    @abstractmethod
    def cmult(self, ct1, ct2):
        """Ciphertext multiplication"""
        pass

    # ... more operations

Fake Backend

For rapid development and testing:

from hetorch import FakeBackend

backend = FakeBackend(
    simulate_noise=False  # No noise simulation (fast)
)

# Or with noise simulation
backend = FakeBackend(
    simulate_noise=True,
    initial_noise_budget=100.0,
    warn_on_low_noise=True,
)

How it Works:

• Simulates HE operations using PyTorch tensors
• No actual encryption (fast)
• Optional noise budget tracking
• Perfect for testing and development

Real Backend (Future)

Integration with HE libraries:

from hetorch.backend import SEALBackend

backend = SEALBackend(
    seal_context=...  # Microsoft SEAL context
)

6. Compilation Workflow

Putting it all together:

# 1. Define model
model = MyNeuralNetwork()

# 2. Create context
context = CompilationContext(
    scheme=HEScheme.CKKS,
    params=CKKSParameters(...),
    backend=FakeBackend()
)

# 3. Build pipeline
pipeline = PassPipeline([
    InputPackingPass(),
    NonlinearToPolynomialPass(),
    RescalingInsertionPass(),
])

# 4. Compile
compiler = HETorchCompiler(context, pipeline)
compiled_model = compiler.compile(model, example_input)

# 5. Execute
output = compiled_model(input_tensor)

What Happens:

1. torch.fx traces the model into a graph
2. Pipeline transforms the graph (passes run in order)
3. Backend executes the transformed graph
4. Output is returned

7. Key Relationships

Understanding how components relate:

CompilationContext
    ├── Contains: HEScheme, EncryptionParameters, HEBackend
    └── Used by: All passes, Compiler

TransformationPass
    ├── Receives: fx.GraphModule, CompilationContext
    ├── Modifies: Graph nodes, metadata
    └── Returns: Transformed fx.GraphModule

fx.GraphModule (IR)
    ├── Contains: Graph nodes (operations)
    ├── Metadata: CiphertextInfo, PackingInfo
    └── Operations: HE custom ops (cadd, cmult, etc.)

HEBackend
    ├── Implements: HE operations
    ├── Provides: CostModel
    └── Used by: Compiled model execution

PassPipeline
    ├── Contains: List of TransformationPass
    ├── Validates: Pass dependencies
    └── Executes: Passes in sequence

8. Design Principles

HETorch follows these principles:

1. Modularity: Passes are independent and composable
2. Extensibility: Easy to add custom passes and backends
3. Metadata-Driven: Ciphertext properties tracked as metadata
4. Backend-Agnostic: Same IR works with different backends
5. Performance-Aware: Cost models guide optimization

Next Steps

Now that you understand the core concepts:

1. Try Examples: See Examples for complete code
2. Learn Passes: Read Builtin Passes for available transformations
3. Understand Workflow: Study Compilation Workflow for details
4. Build Custom Passes: Follow Custom Passes tutorial