KeyMgmt

‘key_mgmt’ Dialect

A dialect for managing the lifetime of FHE key material

KeyMgmt is the control dialect for an abstracted key management runtime. It individually manages key lifetimes to lower memory utilization in FHE programs that utilize a large amount of evaluation keys. It treats key material as dataflow elements and provides automatic, fine-grained key management.

The key_mgmt dialect provides:

• New types for keys (e.g., !key_mgmt.rot_key).
• Operations for loading and clearing keys (key_mgmt.load_key, key_mgmt.clear_key).
• Prefetching hints to overlap computation with key movement (key_mgmt.prefetch_key).

For more details, see the paper KeyMemRT Compiler and Runtime: Unlocking Memory-Scalable FHE

KeyMgmt types

RotKeyType

A rotation key type with optional static index parameter

Syntax:

!key_mgmt.rot_key<
  std::optional<int64_t>   # rotation_index
>

The RotKeyType represents rotation keys used in the KeyMgmt system.

• If rotation_index is present: Static constant rotation key
• If rotation_index is absent: Dynamic rotation key (determined at runtime, e.g., from affine loop induction variables)

For dynamic rotation keys, analysis passes examine the LoadKeyOp’s index operand to determine the range of possible rotation values by tracing through SSA def-use chains and analyzing affine contexts.

Examples: !key_mgmt.rot_key<rotation_index = 5> // Static rotation key for index 5 !key_mgmt.rot_key // Dynamic rotation key

Parameters:

KeyMgmt ops

key_mgmt.assume_loaded (heir::key_mgmt::AssumeLoadedOp)

Assume a rotation key is already loaded in memory

Syntax:

operation ::= `key_mgmt.assume_loaded` $index attr-dict `:` type($index) `->` type($rot_key)

The assume_loaded operation creates a reference to a rotation key that is assumed to already be in memory. This is used for post-loop optimization where a key loaded inside a loop has its clear operation skipped, leaving it in memory for use after the loop.

Unlike load_key, this operation does NOT perform any I/O or deserialization - it simply creates an SSA value representing the key that is already in memory. The compiler/optimizer is responsible for ensuring the key is actually loaded.

This operation is typically emitted as a no-op or debug assertion at runtime.

Example: affine.for %iv = 1 to 16 { %rk = key_mgmt.load_key %iv : i64 -> !key_mgmt.rot_key<> // use %rk… affine.if (%iv != 5) { key_mgmt.clear_key %rk : <> // Skip clearing when iv == 5 } } // Key 5 is still in memory (wasn’t cleared) %c5 = arith.constant 5 : i64 %rk5 = key_mgmt.assume_loaded %c5 : i64 -> !key_mgmt.rot_key<rotation_index = 5> // use %rk5… key_mgmt.clear_key %rk5 : <>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

key_mgmt.clear_key (heir::key_mgmt::ClearKeyOp)

Clear a rotation key from memory

Syntax:

operation ::= `key_mgmt.clear_key` $rot_key attr-dict `:` type($rot_key)

The clear_key operation removes a rotation key from memory or cache. This operation consumes the rotation key and has side effects.

This operation should typically be preceded by a load_key operation that produced the rotation key being cleared.

Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{MemoryEffects::Write on ::mlir::SideEffects::DefaultResource}

Operands:

key_mgmt.load_key (heir::key_mgmt::LoadKeyOp)

Load a rotation key by index

Syntax:

operation ::= `key_mgmt.load_key` $index attr-dict `:` type($index) `->` type($rot_key)

The load_key operation loads a rotation key identified by the given index and returns it as a rot_key type.

The index operand can be:

• A constant value (e.g., from arith.constant) → produces static rot_key type
• An affine induction variable or computed index → produces dynamic rot_key type

The result type indicates whether this is a static or dynamic rotation:

• !key_mgmt.rot_key<rotation_index = N> for static constant N
• !key_mgmt.rot_key for dynamic (runtime-determined) indices

For dynamic rotation keys, analysis passes examine the index operand’s def-use chain to determine possible values (e.g., from affine loop bounds).

This operation should typically be preceded by a prefetch_key operation with the same index, and followed by a clear_key operation on the result.

Examples: // Static constant index %c5 = arith.constant 5 : i64 %rk = key_mgmt.load_key %c5 : !key_mgmt.rot_key<rotation_index = 5>

// Dynamic index from affine loop affine.for %iv = 1 to 16 { %iv_i64 = arith.index_cast %iv : index to i64 %rk = key_mgmt.load_key %iv_i64 : !key_mgmt.rot_key }

// Can also accept index type directly affine.for %iv = 1 to 16 { %rk = key_mgmt.load_key %iv : !key_mgmt.rot_key }

Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{MemoryEffects::Read on ::mlir::SideEffects::DefaultResource}

Operands:

Results:

key_mgmt.prefetch_key (heir::key_mgmt::PrefetchKeyOp)

Prefetch a rotation key by index

Syntax:

operation ::= `key_mgmt.prefetch_key` $index attr-dict

The prefetch_key operation initiates prefetching of a rotation key identified by the given index. This operation has side effects as it may trigger memory operations or cache warming.

This operation should typically be followed by a load_key operation with the same index.

Interfaces: MemoryEffectOpInterface (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{MemoryEffects::Write on ::mlir::SideEffects::DefaultResource}

Operands:

key_mgmt.use_key (heir::key_mgmt::UseKeyOp)

Reference an already-loaded rotation key

Syntax:

operation ::= `key_mgmt.use_key` $input_key attr-dict `:` type($input_key) `->` type($rot_key)

The use_key operation converts a static rotation key to a dynamic rotation key for use in contexts where the type system requires dynamic keys (e.g., inside loops with dynamic indices). This is used for key reuse optimization where a key loaded outside a loop can be referenced inside the loop.

Unlike load_key, this operation does NOT perform any memory operations - it simply provides a type conversion from static to dynamic rot_key. The key must already be loaded and live.

The input is typically a static rotation key (!key_mgmt.rot_key<rotation_index = N>) and the output is always a dynamic rotation key (!key_mgmt.rot_key<>).

Example: // Load key outside loop %c4 = arith.constant 4 : i64 %rk_static = key_mgmt.load_key %c4 : i64 -> !key_mgmt.rot_key<rotation_index = 4>

affine.for %iv = 1 to 16 { %iv_i64 = arith.index_cast %iv : index to i64 // When iv == 4, reuse the pre-loaded key instead of loading again %rk = scf.if (%iv == 4) { %reused = key_mgmt.use_key %rk_static : !key_mgmt.rot_key<rotation_index = 4> -> !key_mgmt.rot_key scf.yield %reused } else { %loaded = key_mgmt.load_key %iv_i64 : i64 -> !key_mgmt.rot_key scf.yield %loaded } … }

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results: