openfhe-alloc-to-inplace overwrites value which has uses.
- Dominant language
- MLIR
- Stars
- 906
- Forks
- 171
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 32
Description
The `--openfhe-alloc-to-inplace` (and `--lattigo-alloc-to-inplace`) passes
promote `openfhe.add` / `lattigo.add` to their in-place variants using a
liveness check that only queries MLIR's `Liveness::isDeadAfter`
(`lib/Utils/AllocToInPlaceUtils.h:355-360`, `isSafeToMutateInPlace`). This
misses the case where the first operand is captured by a `func.return`,
`secret.yield`, or `tensor.insert{,_slice}` op earlier in program order:
MLIR liveness reports the operand as dead after the add because its next
use is the mutation itself, but the outlasting capture still needs the
pre-mutation value. Since `Ciphertext` (OpenFHE) and its Lattigo
analogue are reference-counted / shared-pointer types in the generated C++,
the in-place mutation corrupts the value observed by every other consumer
of the same SSA def.
We hit this in a multi-ciphertext output kernel where two
`secret.yield` results (`%sum` and `%sum + %rot`) end up both equal to
`%sum + %rot` at runtime, silently producing "decrypted output disagrees
with reference" errors — symbolic verification passes because SSA semantics
model each value as immutable.
Suggested fix: extend `isSafeToMutateInPlace` to reuse the existing
`CallerProvidedStorageInfo::isStored` helper
(`AllocToInPlaceUtils.h:59-93`), which already knows how to detect capture
by `tensor.insert{,_slice}` / `scf.yield` / `func.return`, and reject
promotion when any such user exists downstream of the mutation site.
Minimal reproducer (also usable as a regression test — the `CHECK-NOT`
directives fail on current HEAD):
// RUN: heir-opt --openfhe-alloc-to-inplace %s | FileCheck %s
!cc = !openfhe.crypto_context
!ct = !openfhe.ciphertext
// CHECK-LABEL: func.func @bug_returned_operand_still_live
func.func @bug_returned_operand_still_live(
%cc: !cc, %ct_a: !ct, %ct_b: !ct, %ct_c: !ct) -> (!ct, !ct) {
%sum = openfhe.add %cc, %ct_a, %ct_b : (!cc, !ct, !ct) -> !ct
// CHECK-NOT: openfhe.add_inplace %{{.*}}, %sum
%sum2 = openfhe.add %cc, %sum, %ct_c : (!cc, !ct, !ct) -> !ct
return %sum, %sum2 : !ct, !ct
}
Affected files:
- lib/Dialect/Openfhe/Transforms/AllocToInPlace.cpp:33-53
- lib/Utils/AllocToInPlaceUtils.h:355-360 (root cause)
- lib/Dialect/Lattigo/Transforms/AllocToInPlace.cpp (same helper, same bug)
```
// Reproducer: `openfhe.add` is promoted to `openfhe.add_inplace` even when
// the first operand is still live via a `func.return` (or a
// `tensor.insert{,_slice}` / `secret.yield` that outlasts the mutation site).
// The lattigo `--lattigo-alloc-to-inplace` and openfhe
// `--openfhe-alloc-to-inplace` passes share the buggy liveness check in
// `lib/Utils/AllocToInPlaceUtils.h:355-360` (`isSafeToMutateInPlace`) — it
// only queries MLIR `Liveness::isDeadAfter`, which treats the second use of
// the operand as "dead after" the first use in program order even if that
// first use is a mutation of the same shared storage.
//
// Impact: `Ciphertext` (OpenFHE) is a `std::shared_ptr` in the
// C++ runtime. `EvalAddInPlace(a, b)` mutates the pointed-to object. If any
// other Op captured the shared_ptr (return, tensor.insert, secret.yield),
// that consumer sees the mutated value instead of the pre-add snapshot.
//
// RUN: heir-opt --openfhe-alloc-to-inplace %s | FileCheck %s
!cc = !openfhe.crypto_context
!ct = !openfhe.ciphertext
// Minimal shape: compute %sum, then compute %sum2 = %sum + %ct_c, return
// both. Symbolic (SSA) semantics: result 0 = %sum, result 1 = %sum + %ct_c.
// Buggy runtime (after add_inplace promotion): result 0 == result 1.
//
// CHECK-LABEL: func.func @bug_returned_operand_still_live
func.func @bug_returned_operand_still_live(
%cc: !cc, %ct_a: !ct, %ct_b: !ct, %ct_c: !ct) -> (!ct, !ct) {
%sum = openfhe.add %cc, %ct_a, %ct_b : (!cc, !ct, !ct) -> !ct
// BUG: the pass promotes this to `openfhe.add_inplace %cc, %sum, %ct_c`,
// mutating %sum. But %sum is captured by the func.return below — its
// pre-mutation value is required, so the promotion is unsafe.
//
// A correct pass output would keep this as `openfhe.add` (out-of-place),
// so %sum stays intact for the return.
//
// CHECK-NOT: openfhe.add_inplace %{{.*}}, %sum
%sum2 = openfhe.add %cc, %sum, %ct_c : (!cc, !ct, !ct) -> !ct
return %sum, %sum2 : !ct, !ct
}
// Same bug shape but with %sum captured by a `tensor.insert` first (as
// produced by HEIR's multi-ct output-packing lowering) rather than
// directly by a return. Same root cause: MLIR liveness treats %sum as dead
// after the `openfhe.add` because its next use is the mutation itself, but
// the `%packed` tensor built from %sum before the add still needs %sum's
// pre-mutation value.
//
// CHECK-LABEL: func.func @bug_captured_by_tensor_insert
func.func @bug_captured_by_tensor_insert(
%cc: !cc, %ct_a: !ct, %ct_b: !ct, %ct_c: !ct) -> (tensor<2x!ct>) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
%sum = openfhe.add %cc, %ct_a, %ct_b : (!cc, !ct, !ct) -> !ct
%empty = tensor.empty() : tensor<2x!ct>
%with_sum = tensor.insert %sum into %empty[%c0] : tensor<2x!ct>
// BUG: promoted to add_inplace, mutating %sum. The `%with_sum` tensor's
// slot 0 shares the shared_ptr and now sees the mutated value.
//
// CHECK-NOT: openfhe.add_inplace %{{.*}}, %sum
%sum2 = openfhe.add %cc, %sum, %ct_c : (!cc, !ct, !ct) -> !ct
%packed = tensor.insert %sum2 into %with_sum[%c1] : tensor<2x!ct>
return %packed : tensor<2x!ct>
}
```
PS: Triaged using claude opus
Contributor guide
Research direction
Start in lib/Utils/AllocToInPlaceUtils.h at isSafeToMutateInPlace and CallerProvidedStorageInfo::isStored, then compare the OpenFHE and Lattigo AllocToInPlace.cpp callers. Run the supplied heir-opt --openfhe-alloc-to-inplace reproducer with FileCheck. Done means captured operands are not promoted to in-place operations and the regression cases pass while safe promotions remain available.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers, cryptography
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 74/100