Add remaining pipeline plumbing for loop support

PiperOrigin-RevId: 878645040

Author: j2kun

lib/Dialect/Secret/Conversions/BUILD

@@ -24,6 +24,7 @@ cc_library(
         "@llvm-project//mlir:DialectUtils",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
         "@llvm-project//mlir:TransformUtils",

lib/Dialect/Secret/Conversions/Patterns.cpp

@@ -22,6 +22,7 @@
 #include "mlir/include/mlir/Dialect/Affine/IR/AffineOps.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/SCF/IR/SCF.h"        // from @llvm-project
 #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
 #include "mlir/include/mlir/Dialect/Utils/StaticValueUtils.h"  // from @llvm-project
 #include "mlir/include/mlir/IR/Attributes.h"          // from @llvm-project
@@ -516,15 +517,12 @@ void addSecretToSchemeDefaultConversionTargetsAndPatterns(
     ContextAwareTypeConverter& typeConverter) {
   target.addLegalDialect<lwe::LWEDialect, arith::ArithDialect,
                          tensor::TensorDialect>();
-  target.addLegalOp<ModuleOp>();
-
   target.addIllegalDialect<secret::SecretDialect>();
-  target.addIllegalOp<mgmt::ModReduceOp, mgmt::RelinearizeOp,
-                      secret::GenericOp>();
+  target.addIllegalOp<mgmt::ModReduceOp, mgmt::RelinearizeOp>();
 
-  target.addDynamicallyLegalOp<affine::AffineForOp, affine::AffineYieldOp>(
-      [&](Operation* op) { return typeConverter.isLegal(op); });
-  target.addDynamicallyLegalOp<func::CallOp>(
+  target.addDynamicallyLegalOp<affine::AffineForOp, affine::AffineYieldOp,
+                               affine::AffineIfOp, scf::ForOp, scf::IfOp,
+                               scf::YieldOp, func::CallOp>(
       [&](Operation* op) { return typeConverter.isLegal(op); });
   target.markUnknownOpDynamicallyLegal(
       [&](Operation* op) { return !hasSecretOperandsOrResults(op); });
@@ -538,7 +536,11 @@ void addSecretToSchemeDefaultConversionTargetsAndPatterns(
                SecretGenericOpConversion<tensor::EmptyOp, tensor::EmptyOp>,
                SecretGenericFuncCallConversion, ConvertExtractSlice,
                ConvertInsertSlice, ConvertAnyContextAware<affine::AffineForOp>,
+               ConvertAnyContextAware<affine::AffineIfOp>,
                ConvertAnyContextAware<affine::AffineYieldOp>,
+               ConvertAnyContextAware<scf::ForOp>,
+               ConvertAnyContextAware<scf::IfOp>,
+               ConvertAnyContextAware<scf::YieldOp>,
                ConvertAnyContextAware<tensor::ExtractOp>,
                ConvertAnyContextAware<tensor::InsertOp>,
                ConvertAnyContextAware<func::CallOp>>(typeConverter,

lib/Pipelines/ArithmeticPipelineRegistration.cpp

@@ -10,6 +10,7 @@
 #include "lib/Dialect/LWE/Conversions/LWEToLattigo/LWEToLattigo.h"
 #include "lib/Dialect/LWE/Conversions/LWEToOpenfhe/LWEToOpenfhe.h"
 #include "lib/Dialect/LWE/Transforms/AddDebugPort.h"
+#include "lib/Dialect/LWE/Transforms/ImplementTrivialEncryptionAsAddition.h"
 #include "lib/Dialect/Lattigo/Transforms/AllocToInPlace.h"
 #include "lib/Dialect/Lattigo/Transforms/ConfigureCryptoContext.h"
 #include "lib/Dialect/Openfhe/Transforms/AllocToInPlace.h"
@@ -412,6 +413,9 @@ void mlirToRLWEPipeline(OpPassManager& pm,
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createCSEPass());
 
+  // TODO(#2554): skip this pass if the backend supports trivial encryption
+  pm.addPass(lwe::createImplementTrivialEncryptionAsAddition());
+
   // Add a __preprocessed helper for offline pre-packing of plaintexts
   auto splitPreprocessingOptions = SplitPreprocessingOptions{};
   splitPreprocessingOptions.maxReturnValues = options.splitPreprocessing;

lib/Pipelines/BUILD

@@ -108,6 +108,7 @@ cc_library(
         "@heir//lib/Dialect/LWE/Conversions/LWEToOpenfhe",
         "@heir//lib/Dialect/LWE/Conversions/LWEToPolynomial",
         "@heir//lib/Dialect/LWE/Transforms:AddDebugPort",
+        "@heir//lib/Dialect/LWE/Transforms:ImplementTrivialEncryptionAsAddition",
         "@heir//lib/Dialect/Lattigo/Transforms:AllocToInPlace",
         "@heir//lib/Dialect/Lattigo/Transforms:ConfigureCryptoContext",
         "@heir//lib/Dialect/Openfhe/Transforms:AllocToInPlace",

lib/Utils/ContextAwareConversionUtils.h

@@ -181,8 +181,15 @@ class SecretGenericOpConversion
     // OperandAndResultAttrInterface, and this gives special names to the
     // attributes, so it doesn't make sense to copy those names to the
     // converted op.
-    convertArrayOfDicts(op.getAllOperandAttrsAttr(), attrsToPreserve);
+    //
+    // Note that these attributes are deduped, so a for example, if there is a
+    // mgmt.mgmt attribute on the operands and results of a generic, only the
+    // result mgmt.mgmt attribute will be copied to be the mgmt.mgmt attribute
+    // associated with the new op result. This is what we want, but if the order
+    // of these calls were swapped, the _operand_ mgmt attrs would become
+    // associated with the new _result_ SSA value, which is incorrect.
     convertArrayOfDicts(op.getAllResultAttrsAttr(), attrsToPreserve);
+    convertArrayOfDicts(op.getAllOperandAttrsAttr(), attrsToPreserve);
     DenseSet<StringRef> seenNames;
     SmallVector<NamedAttribute> dedupedAttrsToPreserve;
     for (NamedAttribute preservedAttr : attrsToPreserve) {