google / google/heir

[Build][Enhancement] Enable Automatic Kernel Fusion & Integrate FIDESlib GPU Backend into HEIR

Open
#2,652 6 comments 0 reactions 0 assignees View on GitHub
Dominant language
MLIR
Stars
906
Forks
171
Avg merge
4d 12h
Merged PRs (30d)
32

Description

## Background: FIDESlib

**FIDESlib** is an open-source, server-side CKKS GPU library designed for high-performance homomorphic encryption on CUDA GPUs, with interoperability to OpenFHE workflows.

- Paper: **"FIDESlib: A Fully-Fledged Open-Source FHE Library for Efficient CKKS on GPUs"**
https://arxiv.org/abs/2507.04775
- GitHub: **CAPS-UMU/FIDESlib**
https://github.com/CAPS-UMU/FIDESlib

---

## Problem Statement

HEIR currently provides mature scheme-level lowering and multiple **library exit paths** (e.g., OpenFHE/Lattigo/Jaxite/tfhe-rs related flows), but there is no GPU backend path for CKKS that directly targets a CUDA library such as FIDESlib.

Because of this gap:

1. It is hard to lower CKKS programs from HEIR into a GPU-native execution path.
2. There is no dedicated IR stage to make kernel-fusion decisions before backend codegen.
3. End-to-end fused kernel emission against a CKKS GPU library is not available yet.

---

## Why This Contribution

Based on the above, I would like to contribute a project to HEIR that adds:

1. A fusion-friendly intermediate representation for CKKS GPU lowering.
2. Conversion passes into a FIDESlib-oriented exit representation.
3. A kernel-fusion pass that decides what to fuse.
4. Code generation that can emit CUDA code (including fused kernels when needed) for FIDESlib execution.

---

## Summary of Proposed Extension

I propose a GPU backend extension path for HEIR targeting CKKS workloads with FIDESlib integration:

1. New `limb` dialect + `lwe-to-limb` conversion
2. A kernel fusion pass on `limb` IR (horizontal + vertical fusion)
3. New `fideslib` dialect + `limb-to-fideslib` conversion
4. `heir-translate` extension for `.cuh` / `.cu` emission and fused kernel codegen

---

## Proposed Design

### 1) `limb` dialect + `lwe-to-limb` conversion pass

#### Purpose
Bridge existing LWE-level representation to a fusion-friendly, low-level CKKS representation.

#### Dialect content (initial)
- Primitive ops (initial scope):
- modular add/sub
- modular mul/reduce
- NTT / INTT
- rotate
- Types/attributes needed for fusion planning
- `fusion_id` (or equivalent) attribute for downstream fused-kernel grouping

#### Pass
- `lwe-to-limb` conversion to map existing LWE ops/types into `limb` equivalents.

---

### 2) Fusion pass on `limb` dialect

#### Scope
Implement transform logic for:
- Horizontal fusion
- Vertical fusion

#### Initial algorithm
1. Build graph over `limb` ops
2. Assign / track memory buffers per node
3. Identify candidates sharing/reusing compatible buffers
4. Score candidates using expected memory savings
5. Apply highest-benefit legal fusion(s)

---

### 3) `fideslib` dialect + `limb-to-fideslib` conversion pass

#### Purpose
Provide an exit dialect representing FIDESlib-facing operations and API-level semantics.

#### Dialect content (initial)
- Ops mapped to FIDESlib-supported operations
- Types for FIDESlib context/config and API arguments
- `fusion_id` propagation for fused codegen

#### Pass
- `limb-to-fideslib` conversion to lower fusion-ready `limb` IR into FIDESlib-level IR.

---

### 4) Codegen (`heir-translate`)

#### Codegen
Extend `heir-translate` to emit:
- header files (`.cuh`)
- source files (`.cu`)

Behavior:
- Emit direct FIDESlib calls for supported kernels
- Emit generated fused kernels when fusion patterns are not natively provided by FIDESlib

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.