Add pattern to ensure mgmt init on plaintext branch terminators

In some cases, an if/else can be produced in which the else branch
consists of a single yield of a plaintext constant. This is not
directly used in a ct-pt op, which causes the normal
UseInitOpForPlaintextOperand pattern to fail to recognize it as requiring
intialization.

This change adds a new pattern, UseInitForPlaintextBranchTerminators,
that looks for RegionBranchTerminatorOpInterface, determines if the
yielded value is not secret, and then checks all other terminators
of the op that forward to the same successor region, if any of them
are secret, then this pattern will insert a mgmt.init op on the
non-secret operand so it can later be understood to be a plaintext.

PiperOrigin-RevId: 883691202

Author: j2kun

lib/Transforms/Halo/Patterns.cpp

@@ -201,7 +201,6 @@ FailureOr<SmallVector<Value>> isLoopStructuredForHaloUnroll(
 // Inject an scf::ForOp overload of this function, which exists upstream for
 // affine already.
 FailureOr<int64_t> getConstantTripCount(scf::ForOp forOp) {
-  LLVM_DEBUG(llvm::dbgs() << "Getting constant trip count\n");
   if (auto step = forOp.getConstantStep();
       !step.has_value() || !step->isOne()) {
     if (step.has_value()) {
@@ -287,11 +286,9 @@ LogicalResult doPartialUnroll(ForOp forOp, PatternRewriter& rewriter,
           "Start and end levels were not inferable to be concrete integers");
     }
 
-    LLVM_DEBUG(llvm::dbgs() << "doPartialUnroll: Getting levelEndVal\n");
     int levelEndVal = levelEnd->getValue().getInt();
     LLVM_DEBUG(llvm::dbgs()
                << "doPartialUnroll: levelEndVal=" << levelEndVal << "\n");
-    LLVM_DEBUG(llvm::dbgs() << "doPartialUnroll: Getting levelStartVal\n");
     int levelStartVal = levelStart->getValue().getInt();
     LLVM_DEBUG(llvm::dbgs()
                << "doPartialUnroll: levelStartVal=" << levelStartVal << "\n");

lib/Transforms/SecretInsertMgmt/BUILD

@@ -92,6 +92,7 @@ cc_library(
         "@llvm-project//llvm:Support",
         "@llvm-project//mlir:Analysis",
         "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:ControlFlowInterfaces",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:SideEffectInterfaces",
         "@llvm-project//mlir:Support",

lib/Transforms/SecretInsertMgmt/Pipeline.cpp

@@ -223,19 +223,23 @@ void makeLoopsTypeAndLevelInvariant(Operation* top) {
 
   DataFlowSolver solver;
   makeAndRunSecretnessSolver(top, solver);
-
   RewritePatternSet patterns(ctx);
   patterns.add<PeelPlaintextAffineForInit, PeelPlaintextScfForInit>(ctx,
                                                                     &solver);
   walkAndApplyPatterns(top, std::move(patterns));
 
   DataFlowSolver solver2;
   makeAndRunSecretnessSolver(top, solver2);
-
   patterns.clear();
   patterns.add<BootstrapIterArgsPattern<affine::AffineForOp>,
                BootstrapIterArgsPattern<scf::ForOp>>(ctx, &solver2);
   walkAndApplyPatterns(top, std::move(patterns));
+
+  DataFlowSolver solver3;
+  makeAndRunSecretnessSolver(top, solver3);
+  patterns.clear();
+  patterns.add<UseInitForPlaintextBranchTerminators>(ctx, &solver3);
+  walkAndApplyPatterns(top, std::move(patterns));
 }
 
 void unrollLoopsForLevelUtilization(Operation* top, int levelBudget) {

lib/Transforms/SecretInsertMgmt/SecretInsertMgmtPatterns.cpp

@@ -16,6 +16,7 @@
 #include "mlir/include/mlir/IR/BuiltinAttributes.h"      // from @llvm-project
 #include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
 #include "mlir/include/mlir/IR/Value.h"                  // from @llvm-project
+#include "mlir/include/mlir/Interfaces/ControlFlowInterfaces.h"  // from @llvm-project
 #include "mlir/include/mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
 #include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project
 
@@ -274,13 +275,60 @@ LogicalResult MatchCrossMulDepth<Op>::matchAndRewrite(
     }
   }
 
-  // FIXME: replace with updateResultMulDepthLattice(op, solver);
-  // propagateIfChanged only push workitem to the worklist queue
-  // actually execute the transfer for the new values
   solver->eraseAllStates();
   return solver->initializeAndRun(top);
 }
 
+LogicalResult UseInitForPlaintextBranchTerminators::matchAndRewrite(
+    RegionBranchTerminatorOpInterface op, PatternRewriter& rewriter) const {
+  auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op->getParentOp());
+  if (!regionBranchOp) {
+    return rewriter.notifyMatchFailure(
+        op, "parent does not implement RegionBranchOpInterface");
+  }
+
+  // If the op's operands are all already secret, return failure
+  if (!op->getOperands().empty() &&
+      llvm::all_of(op->getOperands(),
+                   [&](Value operand) { return isSecret(operand, solver); })) {
+    return rewriter.notifyMatchFailure(op, "all operands are already secret");
+  }
+
+  bool changed = false;
+  SmallVector<RegionSuccessor> successors;
+  SmallVector<Attribute> operands(op->getNumOperands(), nullptr);
+  op.getSuccessorRegions(operands, successors);
+
+  for (const auto& successor : successors) {
+    auto successorOperands = op.getSuccessorOperands(successor);
+    auto successorInputs = regionBranchOp.getSuccessorInputs(successor);
+
+    if (successorOperands.empty()) continue;
+
+    for (unsigned i = 0; i < successorOperands.size(); ++i) {
+      Value operand = successorOperands[i];
+      Value target = successorInputs[i];
+
+      if (isSecret(target, solver) && !isSecret(operand, solver)) {
+        // Check if already initted
+        if (auto definingOp = operand.getDefiningOp()) {
+          if (isa<mgmt::InitOp>(definingOp)) continue;
+        }
+
+        rewriter.setInsertionPoint(op);
+        auto initOp = mgmt::InitOp::create(rewriter, op.getLoc(),
+                                           operand.getType(), operand);
+
+        op->setOperand(successorOperands.getBeginOperandIndex() + i,
+                       initOp.getResult());
+        changed = true;
+      }
+    }
+  }
+
+  return changed ? success() : failure();
+}
+
 template <typename Op>
 LogicalResult UseInitOpForPlaintextOperand<Op>::matchAndRewrite(
     Op op, PatternRewriter& rewriter) const {

lib/Transforms/SecretInsertMgmt/SecretInsertMgmtPatterns.h

@@ -5,7 +5,8 @@
 #include "mlir/include/mlir/IR/MLIRContext.h"              // from @llvm-project
 #include "mlir/include/mlir/IR/Operation.h"                // from @llvm-project
 #include "mlir/include/mlir/IR/PatternMatch.h"             // from @llvm-project
-#include "mlir/include/mlir/Support/LLVM.h"                // from @llvm-project
+#include "mlir/include/mlir/Interfaces/ControlFlowInterfaces.h"  // from @llvm-project
+#include "mlir/include/mlir/Support/LLVM.h"  // from @llvm-project
 
 // This file include patterns that are combined and orchestrated in the
 // functions defined in `Pipeline.h`. Warning: use these patterns with care,
@@ -179,6 +180,23 @@ struct UseInitOpForPlaintextOperand : public OpRewritePattern<Op> {
   DataFlowSolver* solver;
 };
 
+/// Insert mgmt.init op for region branch terminators where the other
+/// branch is secret.
+struct UseInitForPlaintextBranchTerminators
+    : public OpInterfaceRewritePattern<RegionBranchTerminatorOpInterface> {
+  UseInitForPlaintextBranchTerminators(MLIRContext* context,
+                                       DataFlowSolver* solver)
+      : OpInterfaceRewritePattern<RegionBranchTerminatorOpInterface>(
+            context, /*benefit=*/1),
+        solver(solver) {}
+
+  LogicalResult matchAndRewrite(RegionBranchTerminatorOpInterface op,
+                                PatternRewriter& rewriter) const override;
+
+ private:
+  DataFlowSolver* solver;
+};
+
 /// when reached a certain depth (water line), bootstrap
 ///
 /// TODO(#1642): make it work with cross-level operation

tests/Transforms/secret_insert_mgmt/ckks/loop_with_trivial_branch.mlir

@@ -0,0 +1,54 @@
+// RUN: heir-opt "--secret-insert-mgmt-ckks=after-mul=true before-mul-include-first-mul=false bootstrap-waterline=10 level-budget=2 slot-number=1024" %s | FileCheck %s
+
+// This test was extracted from a larger matvec example. The main issue is
+// that, when the first loop iteration is peeled, the initializer is replaced
+// by the constant initializer in the else branch of the scf.yield (which just
+// naively returns the iter_arg). In some cases, sccp can resolve this issue
+// and detect that the if statement will not execute its else statement. But in
+// a doubly nested loop (approximated here when the loop condition is computed
+// using a func arg), the else cannot be removed. So we require an additional
+// mgmt.init to be inserted to ensure that the secretness aligns on both
+// branches.
+
+module attributes {backend.lattigo, scheme.ckks} {
+  func.func @loop_with_trivial_if_branch(
+      %arg0: !secret.secret<tensor<1x1024xf32>>,
+      %arg1: tensor<512x1024xf32>,
+      %outer_loop_iv: index
+  ) -> !secret.secret<tensor<1x1024xf32>> {
+    %c23 = arith.constant 23 : index
+    // CHECK: [[CST:%[^ ]*]] = arith.constant dense<0.0{{.*}}> : tensor<1x1024xf32>
+    // CHECK: else
+    // CHECK-NEXT: [[INIT:%[^ ]*]] = mgmt.init [[CST]]
+    // CHECK-NEXT: scf.yield [[INIT]]
+    %cst = arith.constant dense<0.000000e+00> : tensor<1x1024xf32>
+    %c-23 = arith.constant -23 : index
+    %c512 = arith.constant 512 : index
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %0 = secret.generic(%arg0: !secret.secret<tensor<1x1024xf32>>) {
+    ^body(%input0: tensor<1x1024xf32>):
+      %5 = scf.for %arg4 = %c0 to %c23 step %c1 iter_args(%arg5 = %cst) -> (tensor<1x1024xf32>) {
+        %9 = arith.muli %outer_loop_iv, %c23 : index
+        %10 = arith.addi %arg4, %9 : index
+        %11 = arith.cmpi slt, %10, %c512 : index
+        %12 = scf.if %11 -> (tensor<1x1024xf32>) {
+          %extracted_slice = tensor.extract_slice %arg1[%10, 0] [1, 1024] [1, 1] : tensor<512x1024xf32> to tensor<1x1024xf32>
+          %13 = arith.muli %outer_loop_iv, %c-23 : index
+          %14 = tensor_ext.rotate %extracted_slice, %13 : tensor<1x1024xf32>, index
+          %15 = tensor_ext.rotate %input0, %arg4 : tensor<1x1024xf32>, index
+          %16 = arith.mulf %14, %15 : tensor<1x1024xf32>
+          %17 = arith.addf %arg5, %16 : tensor<1x1024xf32>
+          scf.yield %17 : tensor<1x1024xf32>
+        } else {
+          scf.yield %arg5 : tensor<1x1024xf32>
+        }
+        scf.yield %12 : tensor<1x1024xf32>
+      }
+      %6 = arith.muli %outer_loop_iv, %c23 : index
+      %7 = tensor_ext.rotate %5, %6 : tensor<1x1024xf32>, index
+      secret.yield %7 : tensor<1x1024xf32>
+    } -> !secret.secret<tensor<1x1024xf32>>
+    return %0 : !secret.secret<tensor<1x1024xf32>>
+  }
+}