google / google/heir

openfhe-alloc-to-inplace overwrites value which has uses.

Open
#3,195 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.