kmrt: add KeyMemRuntime dialect

[j2kun]

Part of #2598

Ports https://github.com/eymay/KeyMemRT-Compiler/tree/main/lib/Dialect/KMRT/IR with minor naming/documentation changes.

[j2kun]

@eymay Would you be open to reviewing the pieces of my port? I also have one question about the clear_key op: is the semantics of this lazy? (e.g., marking a key for eviction without doing a memory op at that moment?)

[eymay]

kmrt.use_key and kmrt.assume_loaded are quick fix ops to make loop nesting work. While use_key is more harmless, assume_loaded can cause key management bugs and is not safe. The tension is between having the keys globally and the IR attempting to treat them locally. I would appreciate some improvement in how to ensure type safety while being able to make use of loop nestings.

[eymay]

This is strongly tied to the loop representation and the nesting model but this can be improved to only accept tighter ranges such as affine maps. Having a dynamic type breaks the type safety in the IR to a great extent.

[j2kun]

So we actually did some static analysis of this since this PR was first pushed (cf. RotationAnalysis), and while it currently only annotates the overall module with the total set of rotation indices required (dynamic or static), we could go further and annotate the individual rotation ops with the indices that it uses.

I would be interested to see a bit more details about the static analysis you did, since tracing the def-use chain and analyzing the affine context is what I wanted to do, and I only had time for "brute force simulation" of the loop in cleartext 🙈

I'll leave this as-is and leave this comment here for context :)

[eymay]

Hi @j2kun thanks for upstreaming KeyMemRT, I'm happy to review this. I think the naming KeyMemRuntime can cause confusion with the actual runtime. We named it KMRT because it was a core dialect of KeyMemRT compiler in the paper. I think for HEIR it can be a good idea to make it clear by calling it something like KeyManagement, KeyMgmt or KeyMemMgmt to show it is about the key memory management. But it is still true that for some features (like prefetching) there needs to be a runtime so might as well keep runtime in the name.

An important thing is to ensure the dialect is library agnostic as possible. We focused on OpenFHE so we modified some ops (e.g. openfhe.rot here) in its dialect but that might not be necessary.

[eymay]

@eymay Would you be open to reviewing the pieces of my port? I also have one question about the clear_key op: is the semantics of this lazy? (e.g., marking a key for eviction without doing a memory op at that moment?)

Sure! The compiler marks it for clearing but the runtime can react to it differently depending on its modes of execution. I figured it's best to keep clear_key eagerly executed to reduce the memory usage immediately regardless of whether the runtime is prefetching (KeyMemRT - Balanced) or eagerly following the ops (KeyMemRT - Low Memory). The compiler should do analysis to remove immature clears by merging with loads beforehand so the runtime does not have much prediction to do. In OpenFHE, that becomes finding the key in the key dictionary and removing its value.

[j2kun]

I think the naming KeyMemRuntime can cause confusion with the actual runtime. We named it KMRT because it was a core dialect of KeyMemRT compiler in the paper. I think for HEIR it can be a good idea to make it clear by calling it something like KeyManagement, KeyMgmt or KeyMemMgmt to show it is about the key memory management. But it is still true that for some features (like prefetching) there needs to be a runtime so might as well keep runtime in the name.

There is a prior art of MLIR dialects with "runtime" in the name, for example tensorflow runtime (tfrt) and AIR runtime (airrt). So I'm not too worried, but I'll think on it.

[asraa]

nit: also add to heir-lsp