Cornami MX2 systolic array backend by ai-mannamalai · Pull Request #2594 · google/heir

ai-mannamalai · 2026-01-31T09:23:42Z

Cornami MX2 systolic array backend for TFHE and CKKS schemes lowering into SCIFR Bool/Ckks dialects

map the said dialect into SCIFR Concrete and SCIFR Liberate topologies suitable for hardware

cornami backend: move tests into common area

move scifr-opt and scifr-translate into the parent HEIR tools to avoid rebuild etc.

update heir-translate to include HEIR_BACKEND_CORNAMI and Target SCIFRBoolEmitter

Cornami: HEIR dialect SCIFRBool emitter and optimizer merged into heir-opt or heir-translate

$ bazelisk test //tests/Dialect/SCIFRBool/... //tests/Dialect/SCIFRCkks/...

j2kun

I think it would be helpful to upstream this in smaller pieces. For example, the SCIFRBool dialect first, then the SCIFRBool Analysis, then the conversions, then the C++ code gen. The the other dialect, etc.

I say this mainly because scanning through this PR, my initial feedback is mainly structural (files in the wrong places, unexpected dependencies, etc.) and it would probably be best to go through this one piece at a time so that we can be on the same page when it comes to the other dialect and future passes.

In particular:

I don't see a reason to separate the code into lib/Backend/cornami; we should just put the dialects in lib/Dialects like the rest of the codebase. Moreover, the current location would not (I think) be compatible with the website doc process)
Conversions going from dialect A to dialect B should live in the lib/Dialect/A/Conversions, not lib/Dialect/B/Conversions
The dirs named Analysis should instead be Transform, and if there is a generic analysis that is used across multiple passes, it can live in lib/Analysis. It's not clear to me how similar the graph analysis in the two examples are that they could be abstracted into a single analysis that is instantiated with different data, but we can get there as the review progresses.
The codegen is currently in lib/Backend/cornami/Dialect/SCIFRBool/Transforms, but it should instead be in lib/Target/, also, I'd recommend using the emitC framework to generate the C++ code, as we're kind of sick of manually emitting C++ via strings and it would likely simplify the code significantly.
Most of the shell scripts should probably not be included, or else be wrapped in a bazel target (or, really, a python script in scripts/ if you want a helper)
Most of the docs should probably be put in docs/ so they show up on the website, and all the configuration for how to use properly use the Cornami backend consolidated into a single pipeline in lib/Pipeline that combines all the pieces. (Ofc, Cornami may ultimately want to build their own, slimmed-down binary for this purpose).

ai-mannamalai · 2026-02-04T17:34:38Z

Okay

ai-mannamalai · 2026-02-04T17:57:27Z

I will make the changes but I prefer to have one big commit.

ai-mannamalai · 2026-02-10T09:50:11Z

@j2kun - many of the items you have mentioned are re-worked. Please take a look

j2kun

Still working through the files, and in the meantime there's some more concrete feedback along the lines of testing.

BUILD.bazel

tests/Dialect/SCIFRBool/Transforms/and_gate.mlir

scripts/cornami/README_cornami.md

scripts/cornami/scifr.bats

lib/Dialect/SCIFRCkks/docs/SCIFRCkks.md

j2kun

Some more feedback on the SCIFRBool dialect definitions.

lib/Dialect/SCIFRBool/IR/SCIFRBoolDialect.td

lib/Dialect/SCIFRBool/IR/SCIFRBoolTypes.td

lib/Dialect/SCIFRBool/IR/SCIFRBoolOps.td

ai-mannamalai · 2026-02-25T20:18:22Z

@j2kun - I think most of the items are addressed for a good first pass commit IMHO; lmk if this is okay from your end

j2kun

As discussed in OH, let's give one more pass over this PR and then we can get it in.

lib/Dialect/SCIFRBool/IR/SCIFRBoolOps.td

lib/Dialect/SCIFRBool/IR/SCIFRBoolDialect.td

lib/Dialect/SCIFRCkks/IR/SCIFRCkksOps.td

lib/Dialect/SCIFRBool/IR/SCIFRBoolOps.td

j2kun

I still have to do:

CGGIToSCIFRBool.cpp
SCIFRBoolEmitter.cpp

Will get back to it later today.

scripts/cornami/README_cornami.md

lib/Backend/cornami/scifr.bats

lib/Backend/cornami/Dialect/SCIFRBool/docs/Cornami-HEIR.png

lib/Dialect/SCIFRBool/IR/SCIFRBoolOps.td

lib/Dialect/SCIFRCkks/IR/SCIFRCkksOps.cpp

lib/Dialect/SCIFRCkks/IR/SCIFRCkksDialect.h

lib/Dialect/SCIFRCkks/Transforms/CKKSEstimator.cpp

lib/Dialect/CGGI/Conversions/CGGIToSCIFRBool/CGGIToSCIFRBool.cpp

j2kun · 2026-03-04T17:59:47Z

lib/Dialect/CGGI/Conversions/CGGIToSCIFRBool/CGGIToSCIFRBool.cpp

+
+    target.addLegalDialect<scifrbool::SCIFRBoolDialect>();
+    target.addIllegalDialect<cggi::CGGIDialect>();
+    target.addIllegalDialect<lwe::LWEDialect>();


I'm a bit confused by this line. Wouldn't this make lwe types also illegal? But you are leaving the LWE types in the IR and not converting them. Maybe this line should be removed (I would expect it to cause the pass to fail...)

lib/Target/SCIFRBool/SCIFRBoolEmitter.cpp

j2kun · 2026-03-04T19:19:55Z

lib/Target/SCIFRBool/SCIFRBoolEmitter.cpp

+        });
+  });
+  m_os << "engineBoolean.m_engine->routeToHost("
+       << m_oprToStrm[sectionOpResultName].back() << ");\n";


The lack of a terminator on section, plus this code, makes the implicit assumption that only the results of the very last op in a section are the output of the entire section op. If you want to be able to have a section produce multiple results, you would probably need a terminator at that point.

If these are the intended semantics of section, this should be added to the tablegen description for the section op.

I propose to work on section op pass in a later PR

lib/Target/SCIFRBool/SCIFRBoolEmitter.cpp

j2kun · 2026-03-07T03:56:39Z

Please run pre-commit run --all-files on this PR to fix the linter errors.

@j2kun

… into SCIFR Bool/Ckks dialects - map the said dialect into SCIFR Concrete and SCIFR Liberate topologies suitable for hardware cornami backend: move tests into common area move scifr-opt and scifr-translate into the parent HEIR tools to avoid rebuild etc. - update heir-translate to include HEIR_BACKEND_CORNAMI and Target SCIFRBoolEmitter Cornami: HEIR dialect SCIFRBool emitter and optimizer merged into heir-opt or heir-translate Move Dialect from lib/Backend/cornami/Dialect to lib/Dialect - move analysis passes internal to SCIFRCkks -> SCIFRCkks/Transforms/ - tests passing - cornami folder update [cfair]: tests updated [SCIFRBool]: hide operators for PBS, KS and Linear which are lower level right now and not exposed from the SCIFRBool APIs; perhaps in future more lowers are possible with these ops and optimizations [SCIFRBool]: Dialect update and documentation update - remove some docs files per Jeremy Khun's review - Remove ASCII art - Cornami SCIFRBool type names updated per review - Remove SCIFRBool_Ciphertext - it is identical to LWECipherTextLike [Cornami MX2 Backend for HEIR]: remove build flag --enable_cornami_mx2 as per review comments; Cornami MX2 backend for HEIR is built always - $ bazelisk test //tests/Dialect/SCIFRBool/... //tests/Dialect/SCIFRCkks/... - move tests to basic FileCheck related w/o bats format etc. - [clang-fmt] fix - add filecheck test for conversion Document the SCIFRBool and SCIFRCkks dialect ops; - remove the unused operators Add description for SectionOp Update lib/Dialect/SCIFRBool/IR/SCIFRBoolDialect.td Update lib/Dialect/SCIFRBool/Conversions/ReplaceOpWithSection.cpp Cornami MX2 systolic array backend for TFHE and CKKS schemes lowering into SCIFR Bool/Ckks dialects - map the said dialect into SCIFR Concrete and SCIFR Liberate topologies suitable for hardware cornami backend: move tests into common area move scifr-opt and scifr-translate into the parent HEIR tools to avoid rebuild etc. - update heir-translate to include HEIR_BACKEND_CORNAMI and Target SCIFRBoolEmitter Cornami: HEIR dialect SCIFRBool emitter and optimizer merged into heir-opt or heir-translate [Cornami MX2 Backend for HEIR]: remove build flag --enable_cornami_mx2 as per review comments; Cornami MX2 backend for HEIR is built always Remove Cornami SCIFR BATS not useful for the project (outside Cornami) updates per review comments remove Cornami backend folder Update lib/Dialect/SCIFRBool/Conversions/ReplaceOpWithSection.cpp Repeal deprecate factory method with new factory method Apply suggestion from @j2kun Apply suggestions from code review part n remove comments and unused code add per review comment on operator identification from SCIFRBool dialect remove unused result SCIFRBoolEmitter: linked and used from lib/Target/ path remove third_party folder; spurious Move folder over to Transforms [scifrbool]: update namespace to mlir::cornami::scifrbool Code Cleanup : CGGIToSCIFRBool conversion Code Cleanup : CGGIToSCIFRBool conversion per review comments Apply suggestion from @j2kun Apply suggestions from code review Remove tests/Backend/cornami in favor of lit tests lit test for option --replace-op-with-section remove comments [scifrbool-visualization-dir option added to the code] pre-commit run --all-files

ai-mannamalai · 2026-03-08T15:12:48Z

Please run pre-commit run --all-files on this PR to fix the linter errors.

done; thanks

j2kun requested changes Feb 4, 2026

View reviewed changes

ai-mannamalai requested a review from j2kun February 10, 2026 16:03

j2kun requested changes Feb 13, 2026

View reviewed changes

ai-mannamalai force-pushed the public-cornami-backend branch 3 times, most recently from 9dc9005 to 5ebc690 Compare February 24, 2026 21:58

ai-mannamalai requested a review from j2kun February 25, 2026 20:18

ai-mannamalai force-pushed the public-cornami-backend branch 3 times, most recently from 6b95d2e to 42cd707 Compare February 26, 2026 18:32

j2kun reviewed Mar 2, 2026

View reviewed changes

ai-mannamalai force-pushed the public-cornami-backend branch from 0f25e24 to 21b2a44 Compare March 3, 2026 01:29

j2kun requested changes Mar 3, 2026

View reviewed changes

ai-mannamalai force-pushed the public-cornami-backend branch from 934a98c to 877854b Compare March 3, 2026 20:50

j2kun requested changes Mar 4, 2026

View reviewed changes

ai-mannamalai force-pushed the public-cornami-backend branch 6 times, most recently from 85cb8a9 to f63096a Compare March 6, 2026 21:54

ai-mannamalai mentioned this pull request Mar 6, 2026

SCIFRBool - Update SectionOp generation pass #2738

Open

j2kun approved these changes Mar 7, 2026

View reviewed changes

ai-mannamalai force-pushed the public-cornami-backend branch from f63096a to 245400c Compare March 8, 2026 15:12

j2kun approved these changes Mar 9, 2026

View reviewed changes

j2kun added the pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing label Mar 9, 2026

Conversation

ai-mannamalai commented Jan 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

j2kun left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ai-mannamalai commented Feb 4, 2026

Uh oh!

ai-mannamalai commented Feb 4, 2026

Uh oh!

ai-mannamalai commented Feb 10, 2026

Uh oh!

j2kun left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

j2kun left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ai-mannamalai commented Feb 25, 2026

Uh oh!

j2kun left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

j2kun left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

j2kun Mar 4, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

j2kun Mar 4, 2026

Choose a reason for hiding this comment

Uh oh!

ai-mannamalai Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

j2kun commented Mar 7, 2026

Uh oh!

ai-mannamalai commented Mar 8, 2026

Uh oh!

Reviewers

Assignees

ai-mannamalai commented Jan 31, 2026 •

edited

Loading

j2kun left a comment •

edited

Loading