Merge pull request #3327 from google:alex/affine-fix

PiperOrigin-RevId: 962877358

Author: copybara-github

lib/Transforms/SplitPreprocessing/BUILD

@@ -18,6 +18,8 @@ cc_library(
         "@heir//lib/Dialect/Preprocessing/IR:Dialect",
         "@heir//lib/Utils:AttributeUtils",
         "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:AffineTransforms",
         "@llvm-project//mlir:ArithDialect",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",

lib/Transforms/SplitPreprocessing/SplitPreprocessing.cpp

@@ -10,8 +10,10 @@
 #include "lib/Dialect/Preprocessing/IR/PreprocessingOps.h"
 #include "lib/Dialect/Preprocessing/IR/PreprocessingTypes.h"
 #include "lib/Utils/AttributeUtils.h"
-#include "llvm/include/llvm/ADT/STLExtras.h"             // from @llvm-project
-#include "llvm/include/llvm/ADT/SmallVector.h"           // from @llvm-project
+#include "llvm/include/llvm/ADT/STLExtras.h"    // from @llvm-project
+#include "llvm/include/llvm/ADT/SmallVector.h"  // from @llvm-project
+#include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.h"  // from @llvm-project
+#include "mlir/include/mlir/Dialect/Affine/Transforms/Passes.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Arith/IR/Arith.h"    // from @llvm-project
 #include "mlir/include/mlir/Dialect/Func/IR/FuncOps.h"   // from @llvm-project
 #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
@@ -27,6 +29,7 @@
 #include "mlir/include/mlir/IR/MLIRContext.h"            // from @llvm-project
 #include "mlir/include/mlir/IR/OpDefinition.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/IR/Region.h"                 // from @llvm-project
 #include "mlir/include/mlir/IR/TypeUtilities.h"          // from @llvm-project
 #include "mlir/include/mlir/IR/Types.h"                  // from @llvm-project
@@ -119,13 +122,106 @@ static bool isAllowedPlaintextType(Type type) {
   return false;
 }
 
+// Rebuild each affine.for in `funcOp` without the iter_args whose region
+// argument and loop result are both unused.
+//
+// When cloning a loop into the preprocessing function we leave behind a dummy
+// initializer for any loop-carried value that isn't part of the plaintext
+// slice (e.g. a ciphertext accumulator), relying on remove-dead-values to drop
+// the now-unused iter_arg. That works for scf.for, but remove-dead-values only
+// poisons (it does not structurally remove) dead affine.for iter_args, so a
+// ub.poison-typed loop-carried value would otherwise survive into backend
+// lowering, where it cannot be converted or emitted. This performs the removal
+// remove-dead-values cannot, after which the dummy initializers become dead and
+// are cleaned up normally.
+static void removeDeadAffineForIterArgs(func::FuncOp funcOp) {
+  IRRewriter rewriter(funcOp.getContext());
+
+  SmallVector<affine::AffineForOp> loops;
+  funcOp.walk([&](affine::AffineForOp forOp) { loops.push_back(forOp); });
+
+  for (affine::AffineForOp forOp : loops) {
+    unsigned numIterArgs = forOp.getNumIterOperands();
+    if (numIterArgs == 0) continue;
+
+    SmallVector<unsigned> keptIndices;
+    for (unsigned i = 0; i < numIterArgs; ++i) {
+      if (!forOp.getRegionIterArgs()[i].use_empty() ||
+          !forOp.getResult(i).use_empty()) {
+        keptIndices.push_back(i);
+      }
+    }
+    if (keptIndices.size() == numIterArgs) continue;  // nothing dead
+
+    rewriter.setInsertionPoint(forOp);
+    SmallVector<Value> keptInits;
+    for (unsigned i : keptIndices) keptInits.push_back(forOp.getInits()[i]);
+
+    auto newLoop = affine::AffineForOp::create(
+        rewriter, forOp.getLoc(), forOp.getLowerBoundOperands(),
+        forOp.getLowerBoundMap(), forOp.getUpperBoundOperands(),
+        forOp.getUpperBoundMap(), forOp.getStepAsInt(), keptInits);
+
+    // Trim the existing terminator down to the kept loop-carried values.
+    auto yieldOp =
+        cast<affine::AffineYieldOp>(forOp.getBody()->getTerminator());
+    SmallVector<Value> keptYields;
+    for (unsigned i : keptIndices) keptYields.push_back(yieldOp.getOperand(i));
+    rewriter.modifyOpInPlace(
+        yieldOp, [&]() { yieldOp.getOperandsMutable().assign(keptYields); });
+
+    // With no kept iter_args the builder added a default terminator; drop it so
+    // the merged (trimmed) affine.yield is the loop's only terminator.
+    if (keptInits.empty()) {
+      rewriter.eraseOp(newLoop.getBody()->getTerminator());
+    }
+
+    // Map the old block arguments onto the new loop: induction var, then each
+    // iter_arg. Kept ones map to the new region args; dead ones are unused, so
+    // their (type-matched, dominating) original initializer is a safe
+    // placeholder that is never actually referenced.
+    SmallVector<Value> blockArgReplacements;
+    blockArgReplacements.push_back(newLoop.getInductionVar());
+    unsigned keptCursor = 0;
+    for (unsigned i = 0; i < numIterArgs; ++i) {
+      if (keptCursor < keptIndices.size() && keptIndices[keptCursor] == i) {
+        blockArgReplacements.push_back(
+            newLoop.getRegionIterArgs()[keptCursor++]);
+      } else {
+        blockArgReplacements.push_back(forOp.getInits()[i]);
+      }
+    }
+    rewriter.mergeBlocks(forOp.getBody(), newLoop.getBody(),
+                         blockArgReplacements);
+
+    for (auto [newIdx, oldIdx] : llvm::enumerate(keptIndices)) {
+      rewriter.replaceAllUsesWith(forOp.getResult(oldIdx),
+                                  newLoop.getResult(newIdx));
+    }
+    rewriter.eraseOp(forOp);
+  }
+}
+
 struct SplitPreprocessingPass
     : impl::SplitPreprocessingBase<SplitPreprocessingPass> {
   using SplitPreprocessingBase::SplitPreprocessingBase;
 
   void runOnOperation() override {
     Operation* root = getOperation();
 
+    // The storage layout sizes each site by the enclosing loops' trip counts,
+    // while the store/load indices are those loops' induction variables.
+    // Normalize affine loops before capturing those indices so a later loop
+    // normalization cannot rewrite a captured index to an affine.apply of the
+    // original non-zero-based, non-unit-step induction variable.
+    OpPassManager normalizeLoops("builtin.module");
+    normalizeLoops.addNestedPass<func::FuncOp>(
+        affine::createAffineLoopNormalizePass(true));
+    if (failed(runPipeline(normalizeLoops, root))) {
+      signalPassFailure();
+      return;
+    }
+
     // Annotate each encode op with a stable site id
     int32_t encodeId = 0;
     root->walk([&](PlaintextEncodeOpInterface op) {
@@ -180,6 +276,15 @@ struct SplitPreprocessingPass
     (void)runPipeline(pipeline, preprocessingFuncOp);
     (void)runPipeline(pipeline, preprocessedFuncOp);
     (void)runPipeline(pipeline, funcOp);
+
+    // remove-dead-values poisons but cannot structurally strip a dead
+    // affine.for iter_arg (it only does so for scf.for), so the dummy
+    // ciphertext iter_arg left when cloning a loop survives as a ub.poison
+    // loop-carried value that later backend lowering can neither convert nor
+    // emit. Strip those dead iter_args now, then re-run the cleanup so the
+    // orphaned ub.poison initializers are removed too.
+    removeDeadAffineForIterArgs(preprocessingFuncOp);
+    (void)runPipeline(pipeline, preprocessingFuncOp);
   }
 
   void updateOriginalFunc(FuncOp funcOp, FuncOp preprocessingFuncOp,

tests/Transforms/split_preprocessing/dead_ciphertext_iter_arg.mlir

@@ -0,0 +1,43 @@
+// RUN: heir-opt %s --split-preprocessing | FileCheck %s
+
+// A loop carries a ciphertext iter_arg (the running accumulator) and also
+// encodes a plaintext per iteration. The ciphertext value is not part of the
+// plaintext slice, so when the loop is cloned into the preprocessing function
+// its iter_arg becomes dead. remove-dead-values only poisons (it cannot
+// structurally strip) a dead affine.for iter_arg, so split-preprocessing must
+// remove it itself -- otherwise a ub.poison-typed loop-carried value survives
+// and later backend lowering can neither convert nor emit it.
+
+!Z36028797017456641_i64 = !mod_arith.int<36028797017456641 : i64>
+!Z35184371138561_i64 = !mod_arith.int<35184371138561 : i64>
+!Z35184372121601_i64 = !mod_arith.int<35184372121601 : i64>
+#inverse_canonical_encoding = #lwe.inverse_canonical_encoding<scaling_factor = 0>
+#key = #lwe.key<>
+#ring_f64_1_x1024 = #polynomial.ring<coefficientType = f64, polynomialModulus = <1 + x**1024>>
+!rns_L2 = !rns.rns<!Z36028797017456641_i64, !Z35184371138561_i64, !Z35184372121601_i64>
+!pt = !lwe.lwe_plaintext<plaintext_space = <ring = #ring_f64_1_x1024, encoding = #inverse_canonical_encoding>>
+#ring_rns_L2_1_x1024 = #polynomial.ring<coefficientType = !rns_L2, polynomialModulus = <1 + x**1024>>
+#ciphertext_space_L2 = #lwe.ciphertext_space<ring = #ring_rns_L2_1_x1024, encryption_type = mix>
+!ct_L2 = !lwe.lwe_ciphertext<plaintext_space = <ring = #ring_f64_1_x1024, encoding = #inverse_canonical_encoding>, ciphertext_space = #ciphertext_space_L2, key = #key, modulus_chain = #lwe.modulus_chain<elements = <36028797017456641 : i64, 35184371138561 : i64, 35184372121601 : i64>, current = 2>>
+
+// The preprocessing loop must be a pure store loop: no iter_args, no ub.poison.
+// CHECK:       func.func @f__preprocessing() -> !preprocessing.storage<!pt>
+// CHECK-NOT:     ub.poison
+// CHECK:         affine.for %[[I:.*]] = 0 to 4 {
+// CHECK-NOT:     iter_args
+// CHECK:           %[[PT:.*]] = lwe.rlwe_encode
+// CHECK:           preprocessing.store %[[PT]], %{{.*}}[%[[I]]] site 0<!pt> : !pt, <!pt>
+// CHECK:         return
+
+module attributes {backend.openfhe, ckks.schemeParam = #ckks.scheme_param<logN = 14, Q = [36028797017456641, 35184371138561, 35184372121601], P = [1152921504607338497, 1152921504608747521], logDefaultScale = 45>, scheme.ckks} {
+  func.func @f(%arg0: tensor<1x!ct_L2>) -> tensor<1x!ct_L2> {
+    %cst = arith.constant dense<1.0> : tensor<1024xf32>
+    %0 = affine.for %i = 0 to 4 iter_args(%sum = %arg0) -> (tensor<1x!ct_L2>) {
+      %pt = lwe.rlwe_encode %cst {encoding = #inverse_canonical_encoding, ring = #ring_f64_1_x1024} : tensor<1024xf32> -> !pt
+      %from = tensor.from_elements %pt : tensor<1x!pt>
+      %1 = ckks.add_plain %sum, %from : (tensor<1x!ct_L2>, tensor<1x!pt>) -> tensor<1x!ct_L2>
+      affine.yield %1 : tensor<1x!ct_L2>
+    }
+    return %0 : tensor<1x!ct_L2>
+  }
+}

tests/Transforms/split_preprocessing/noncanonical_affine_loop.mlir

@@ -0,0 +1,30 @@
+// RUN: heir-opt %s --split-preprocessing --affine-loop-normalize | FileCheck %s
+
+// split-preprocessing uses enclosing loop induction variables as storage
+// indices. A non-zero lower bound or non-unit step must therefore be normalized
+// before splitting, so the storage is indexed by iteration number rather than
+// by the original induction variable values.
+
+!Z36028797017456641_i64 = !mod_arith.int<36028797017456641 : i64>
+!Z35184371138561_i64 = !mod_arith.int<35184371138561 : i64>
+!Z35184372121601_i64 = !mod_arith.int<35184372121601 : i64>
+#inverse_canonical_encoding = #lwe.inverse_canonical_encoding<scaling_factor = 0>
+#ring_f64_1_x1024 = #polynomial.ring<coefficientType = f64, polynomialModulus = <1 + x**1024>>
+!rns_L2 = !rns.rns<!Z36028797017456641_i64, !Z35184371138561_i64, !Z35184372121601_i64>
+!pt = !lwe.lwe_plaintext<plaintext_space = <ring = #ring_f64_1_x1024, encoding = #inverse_canonical_encoding>>
+
+// CHECK:       func.func @noncanonical_affine_loop__preprocessing
+// CHECK:         affine.for %[[I:.*]] = 0 to 2 {
+// CHECK-NOT:       affine.apply
+// CHECK:           preprocessing.store %{{.*}}, %{{.*}}[%[[I]]] site 0
+// CHECK:         return
+
+module attributes {backend.openfhe, ckks.schemeParam = #ckks.scheme_param<logN = 14, Q = [36028797017456641, 35184371138561, 35184372121601], P = [1152921504607338497, 1152921504608747521], logDefaultScale = 45>, scheme.ckks} {
+  func.func @noncanonical_affine_loop() {
+    %cst = arith.constant dense<1.0> : tensor<1024xf32>
+    affine.for %i = 1 to 7 step 3 {
+      %pt = lwe.rlwe_encode %cst {encoding = #inverse_canonical_encoding, ring = #ring_f64_1_x1024} : tensor<1024xf32> -> !pt
+    }
+    return
+  }
+}