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

[VedantParanjape]

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<DCRTPoly> (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<DCRTPoly>` (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