Lattigo: Add Negate op and e2e CMUX example #1627

ZenithalHourlyRate · 2025-03-30T03:35:03Z

Fixes #1623

Depends on #1620.

This PR construct a Lattigo Negate Op via ring operation in Lattigo (they seem not to offer a high level one).

To better demonstrate its use a CMUX e2e example from #1620 is also added.

ZenithalHourlyRate · 2025-04-02T16:44:54Z

Rebased after #1620 is in. Ready for review now.

asraa

Thank you! Very small nits

asraa · 2025-04-02T17:45:07Z

lib/Target/Lattigo/LattigoEmitter.cpp

-              valueString = "[]int64{";
-              for (auto value : denseAttr.getValues<APInt>()) {
-                valueString += std::to_string(value.getSExtValue()) + ", ";
+              if (denseAttr.getType().getElementType().isInteger(1)) {


nit: i was noticing the i1 / bool case isn't handled for the IntegerAttr case above, but then in convertType it is handlded below.

Maybe you can remove this typeswitch and instead rely on auto res = convertType(valueAttr.getType())?

Now arith.constant is honest on type: what is in IR is converted to corresponding type using convertType. The reason int64/float64 is deliberately here was that Encode/Decode only accepts int64[]/float64[] and we have the following common pattern

cst = arith.constant <4xi16> Decode(pt, cst)

so for convenience I directly convert i16 to i64. Now the extra step is done in DecodeOp.

lib/Target/Lattigo/LattigoEmitter.cpp

tests/Dialect/Lattigo/Emitters/emit_lattigo.mlir

asraa

Thanks as always!

j2kun · 2025-04-06T16:00:54Z

It looks like this is failing the cmux test at head

llvm::SmallVectorTemplateCommon<long>::operator[](size_type) [T = long]: idx < size()
*** Check failure stack trace: ***
    @     0x55af1f1a3e09  absl::log_internal::LogMessage::SendToLog()
    @     0x55af1f1a3d7e  absl::log_internal::LogMessage::Flush()
    @     0x55af1f177af4  __assert_fail
    @     0x55af1a658830  mlir::heir::ScaleAnalysisBackward<>::visitOperation()::{lambda()#1}::operator()<>()
    @     0x55af1a654e53  mlir::heir::ScaleAnalysisBackward<>::visitOperation()
    @     0x55af1b81c52d  mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::visitOperation()
    @     0x55af1b81b3ac  mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::initializeRecursively()
    @     0x55af1b81b559  mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::initializeRecursively()
    @     0x55af1b81b559  mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::initializeRecursively()
    @     0x55af1b81b559  mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::initializeRecursively()
    @     0x55af1b81d82e  mlir::DataFlowSolver::initializeAndRun()
    @     0x55af1a63a46f  mlir::heir::PopulateScaleBGV::runOnOperation()

This was internal so I will try to reproduce externally

heir/tools/heir-opt "--annotate-module=backend=lattigo scheme=bgv" --mlir-to-bgv --scheme-to-lattigo -o blaze-out/k8-fastbuild/bin/tests/Examples/lattigo/bgv/cmux/cmux_heir_opt.mlir tests/Examples/lattigo/bgv/cmux/cmux.mlir

ZenithalHourlyRate · 2025-04-06T16:04:19Z

Is this the same problem as #1671

j2kun · 2025-04-06T16:04:29Z

Confirmed. Here's the trace:

 bazel run //tools:heir-opt -- "--annotate-module=backend=lattigo scheme=bgv" --mlir-to-bgv --scheme-to-lattigo $PWD/tests/Examples/lattigo/bgv/cmux/cmux.mlir
...
heir-opt: external/llvm-project/llvm/include/llvm/ADT/SmallVector.h:291: reference llvm::SmallVectorTemplateCommon<long>::operator[](size_type) [T = long]: Assertion `idx < size()' failed.
PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
0.      Program arguments: /home/j2kun/.cache/bazel/_bazel_j2kun/b07a48cf9da8c73a462d3608e5c508dd/execroot/_main/bazel-out/k8-dbg/bin/tools/heir-opt "--annotate-module=backend=lattigo scheme=bgv" --mlir-to-bgv --scheme-to-lattigo /home/j2kun/fhe/heir/tests/Examples/lattigo/bgv/cmux/cmux.mlir
<snip>
#12 0x000055ec18841c82 llvm::SmallVectorTemplateCommon<long, void>::operator[](unsigned long) /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/SmallVector.h:0:5
#13 0x000055ec1a60f1df auto mlir::heir::ScaleAnalysisBackward<mlir::heir::BGVScaleModel>::visitOperation(mlir::Operation*, llvm::ArrayRef<mlir::heir::ScaleLattice*>, llvm::ArrayRef<mlir::heir::ScaleLattice const*>)::'lambda'(auto)::operator()<mlir::arith::MulIOp>(auto) const /proc/self/cwd/lib/Analysis/ScaleAnalysis/ScaleAnalysis.cpp:334:37
#14 0x000055ec1a60ee1c llvm::TypeSwitch<mlir::Operation&, void>& llvm::TypeSwitch<mlir::Operation&, void>::Case<mlir::arith::MulIOp, mlir::heir::ScaleAnalysisBackward<mlir::heir::BGVScaleModel>::visitOperation(mlir::Operation*, llvm::ArrayRef<mlir::heir::ScaleLattice*>, llvm::ArrayRef<mlir::heir::ScaleLattice const*>)::'lambda'(auto)&>(mlir::heir::ScaleAnalysisBackward<mlir::heir::BGVScaleModel>::visitOperation(mlir::Operation*, llvm::ArrayRef<mlir::heir::ScaleLattice*>, llvm::ArrayRef<mlir::heir::ScaleLattice const*>)::'lambda'(auto)&) /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/TypeSwitch.h:149:7
#15 0x000055ec1a60a3bf mlir::heir::ScaleAnalysisBackward<mlir::heir::BGVScaleModel>::visitOperation(mlir::Operation*, llvm::ArrayRef<mlir::heir::ScaleLattice*>, llvm::ArrayRef<mlir::heir::ScaleLattice const*>) /proc/self/cwd/lib/Analysis/ScaleAnalysis/ScaleAnalysis.cpp:310:17
#16 0x000055ec1a5ce66d mlir::dataflow::SparseBackwardDataFlowAnalysis<mlir::heir::ScaleLattice>::visitOperationImpl(mlir::Operation*, llvm::ArrayRef<mlir::dataflow::AbstractSparseLattice*>, llvm::ArrayRef<mlir::dataflow::AbstractSparseLattice const*>) /proc/self/cwd/external/llvm-project/mlir/include/mlir/Analysis/DataFlow/SparseAnalysis.h:532:12
#17 0x000055ec1e9fbba1 mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::visitOperation(mlir::Operation*) /proc/self/cwd/external/llvm-project/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp:537:10
#18 0x000055ec1e9fa4a7 mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::initializeRecursively(mlir::Operation*) /proc/self/cwd/external/llvm-project/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp:343:14
#19 0x000055ec1e9fa5cc mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::initializeRecursively(mlir::Operation*) /proc/self/cwd/external/llvm-project/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp:354:20
#20 0x000055ec1e9fa5cc mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::initializeRecursively(mlir::Operation*) /proc/self/cwd/external/llvm-project/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp:354:20
#21 0x000055ec1e9fa5cc mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::initializeRecursively(mlir::Operation*) /proc/self/cwd/external/llvm-project/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp:354:20
#22 0x000055ec1e9fa46d mlir::dataflow::AbstractSparseBackwardDataFlowAnalysis::initialize(mlir::Operation*) /proc/self/cwd/external/llvm-project/mlir/lib/Analysis/DataFlow/SparseAnalysis.cpp:338:10
#23 0x000055ec1ea018bd mlir::DataFlowSolver::initializeAndRun(mlir::Operation*) /proc/self/cwd/external/llvm-project/mlir/lib/Analysis/DataFlowFramework.cpp:116:25
#24 0x000055ec1a5cc3e2 mlir::heir::PopulateScaleBGV::runOnOperation() /proc/self/cwd/lib/Transforms/PopulateScale/PopulateScaleBGV.cpp:81:23
<snip>

j2kun · 2025-04-06T16:06:06Z

ok I will remove my smoke test from that PR so it can go in on its own

j2kun · 2025-04-06T16:11:24Z

ok I will remove my smoke test from that PR so it can go in on its own

Cherry-picking #1671 indeed fixes the issue.

ZenithalHourlyRate · 2025-04-07T15:53:47Z

Do I need to rebase this (how is the internal patching status for it) (changes in #1661 also apply here so there will be some new change if I rebase)

j2kun · 2025-04-07T17:49:44Z

I am syncing the internal patch now, will let you know.

j2kun · 2025-04-07T18:02:56Z

Looks like the patch is working, and should land shortly

ZenithalHourlyRate force-pushed the lattigo-negate branch 2 times, most recently from 9fd8287 to e64986e Compare April 2, 2025 16:44

ZenithalHourlyRate marked this pull request as ready for review April 2, 2025 16:44

asraa approved these changes Apr 2, 2025

View reviewed changes

ZenithalHourlyRate force-pushed the lattigo-negate branch 2 times, most recently from bfdefac to 7398502 Compare April 3, 2025 08:40

Lattigo: Add Negate op and e2e CMUX example

f7ca8da

ZenithalHourlyRate force-pushed the lattigo-negate branch from 7398502 to f7ca8da Compare April 4, 2025 06:58

asraa approved these changes Apr 4, 2025

View reviewed changes

asraa added the pull_ready Indicates whether a PR is ready to pull. The copybara worker will import for internal testing label Apr 4, 2025

ZenithalHourlyRate mentioned this pull request Apr 6, 2025

LattigoEmitter: add tensor.insert, improve tensor.extract, improve constant emitter #1674

Closed

j2kun added a commit to j2kun/heir that referenced this pull request Apr 6, 2025

cherry-pick google#1627, google#1671, google#1674

633877c

copybara-service bot merged commit d1ddd5e into google:main Apr 7, 2025
9 checks passed

ZenithalHourlyRate mentioned this pull request Apr 8, 2025

tests: use export files for common src #1683

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Lattigo: Add Negate op and e2e CMUX example #1627

Lattigo: Add Negate op and e2e CMUX example #1627

ZenithalHourlyRate commented Mar 30, 2025

ZenithalHourlyRate commented Apr 2, 2025

asraa left a comment

asraa Apr 2, 2025

ZenithalHourlyRate Apr 3, 2025 •

edited

Loading

asraa left a comment

j2kun commented Apr 6, 2025 •

edited

Loading

ZenithalHourlyRate commented Apr 6, 2025

j2kun commented Apr 6, 2025

j2kun commented Apr 6, 2025

j2kun commented Apr 6, 2025

ZenithalHourlyRate commented Apr 7, 2025

j2kun commented Apr 7, 2025

j2kun commented Apr 7, 2025

Lattigo: Add Negate op and e2e CMUX example #1627

Lattigo: Add Negate op and e2e CMUX example #1627

Conversation

ZenithalHourlyRate commented Mar 30, 2025

ZenithalHourlyRate commented Apr 2, 2025

asraa left a comment

Choose a reason for hiding this comment

asraa Apr 2, 2025

Choose a reason for hiding this comment

ZenithalHourlyRate Apr 3, 2025 • edited Loading

Choose a reason for hiding this comment

asraa left a comment

Choose a reason for hiding this comment

j2kun commented Apr 6, 2025 • edited Loading

ZenithalHourlyRate commented Apr 6, 2025

j2kun commented Apr 6, 2025

j2kun commented Apr 6, 2025

j2kun commented Apr 6, 2025

ZenithalHourlyRate commented Apr 7, 2025

j2kun commented Apr 7, 2025

j2kun commented Apr 7, 2025

ZenithalHourlyRate Apr 3, 2025 •

edited

Loading

j2kun commented Apr 6, 2025 •

edited

Loading