fix(Secretize): Selective output wrapping in WrapGeneric using SecretnessAnalysis

[pantha704]

Summary

Fixes #2553 - WrapGeneric was forcing all outputs to secret types, causing downstream crashes for functions that return values not depending on secrets.

Changes

Core Change: Operation Partitioning

Instead of wrapping ALL function body operations in secret.generic, this PR:

1. Classifies operations using SecretnessAnalysis - tracks which values depend on secrets
2. Partitions operations into secret-dependent vs plaintext
3. Only wraps secret-dependent ops in secret.generic
4. Keeps plaintext operations outside the generic block

Key Algorithm

Phase 1: Track secret values (starting from secret arguments)
Phase 2: Mark ops as secret if any operand is secret
Phase 3: Determine which outputs are secret vs plaintext
Phase 4: Build IR with plaintext ops outside, secret ops inside generic

Edge Cases Handled

• Direct secret return: return %secret_arg still creates a generic (no ops but secret output)
• Plaintext-only return: return constant does NOT create a generic
• Mixed inputs: Non-secret inputs passed to secret ops work correctly

Example

Before (Issue #2553 - causing crash):

func.func @foo(%x: i32 {secret.secret}) -> i8 {
  %c = arith.constant 42 : i8
  return %c : i8
}
// Output: !secret.secret<i8> ← WRONG, crashes downstream

After (This PR):

func.func @foo(%x: !secret.secret<i32>) -> i8 {
  %c = arith.constant 42 : i8
  return %c : i8  // No generic, plaintext preserved
}

Testing

• Updated existing test case for issue WrapGeneric forces output to secret type, leading to crashes later in the pipeline #2553
• All existing tests should continue to pass

[asraa]

Thank you! I think instead of cloning only secret ops in, I think this particular pattern should just be wrapping the entire block in a generic, but modifying just the selective output wrapping that you did with newOutputs.

There are other patterns that would also benefit from taking the solver in this file, but maybe we can go pattern by pattern to integrate the secretness analysis (for example, HoistPlaintextOps would benefit, and that would hoist plaintext ops outside of the secret body if they can be hoisted).

[j2kun]

I think the core issue is that if you secret.yield a plaintext SSA value, it becomes secret according to the secretness analysis and won't be automatically converted to a public value. We had this special behavior back in the CGGI pipeline because we wanted to make an empty memref secret as the initializer for a loop that put secret items inside it.

Maybe the relevant pattern could have an option to control its behavior, and we could have this pass specialize... Thoughts?

[asraa]

heir/lib/Dialect/Secret/IR/SecretPatterns.cpp

Line 107 in a6a0414

if (!op.getOps<memref::AllocOp>().empty()) {

Right - that's why I think this pattern should also take secretness analysis and we should test if that allocated value will be used for secret storing values later. (That being said, I think that would mean secretness analysis would need a backwards analysis as well). But I think that's why I'd prefer this pass stay minimal and tackle that problem in a later PR

[j2kun]

I think we agree, but that would make this issue unresolvable until the secretness analysis is improved.

[pantha704]

Thanks for the feedback @asraa @j2kun!

I've simplified the implementation based on your guidance:

Changes

• Removed op partitioning - now wraps entire block in generic (like original)
• Kept selective output types using SecretnessAnalysis
• Skip generic creation when no outputs depend on secrets (fixes WrapGeneric forces output to secret type, leading to crashes later in the pipeline #2553)

Key Logic

// Determine output types based on secretness analysis
for (auto [i, resultType] : llvm::enumerate(op.getResultTypes())) {
  Value returnVal = returnOp->getOperand(i);
  if (isSecret(returnVal, solver)) {
    newOutputs.push_back(secret::SecretType::get(resultType));
    hasSecretOutputs = true;
  } else {
    newOutputs.push_back(resultType);
  }
}

// Skip generic if no outputs depend on secrets
if (!hasSecretOutputs) {
  return success();
}