add openfhe ckks e2e convolution test

PiperOrigin-RevId: 827608103

Author: asraa

lib/Pipelines/ArithmeticPipelineRegistration.cpp

@@ -35,6 +35,8 @@
 #include "lib/Transforms/ElementwiseToAffine/ElementwiseToAffine.h"
 #include "lib/Transforms/FoldConstantTensors/FoldConstantTensors.h"
 #include "lib/Transforms/FoldPlaintextMasks/FoldPlaintextMasks.h"
+#include "lib/Transforms/ForwardInsertSliceToExtractSlice/ForwardInsertSliceToExtractSlice.h"
+#include "lib/Transforms/ForwardInsertToExtract/ForwardInsertToExtract.h"
 #include "lib/Transforms/FullLoopUnroll/FullLoopUnroll.h"
 #include "lib/Transforms/GenerateParam/GenerateParam.h"
 #include "lib/Transforms/InlineActivations/InlineActivations.h"
@@ -113,6 +115,7 @@ void lowerAssignLayout(OpPassManager& pm, bool unroll = false) {
   pm.addNestedPass<func::FuncOp>(affine::createAffineExpandIndexOpsPass());
   pm.addNestedPass<func::FuncOp>(affine::createSimplifyAffineStructuresPass());
   pm.addNestedPass<func::FuncOp>(affine::createAffineLoopNormalizePass(true));
+  pm.addNestedPass<func::FuncOp>(createForwardInsertSliceToExtractSlice());
 
   // The lowered assign_layout ops involve plaintext operations that are still
   // inside secret.generic, and are not handled well by downstream noise models
@@ -169,6 +172,7 @@ void mlirToSecretArithmeticPipelineBuilder(
   pm.addPass(
       createConvertToCiphertextSemantics(convertToCiphertextSemanticsOptions));
 
+  pm.addPass(createApplyFolders());
   pm.addPass(createCanonicalizerPass());
   pm.addPass(tensor_ext::createImplementRotateAndReduce());
 
@@ -188,6 +192,7 @@ void mlirToSecretArithmeticPipelineBuilder(
   pm.addPass(createAddClientInterface(addClientInterfaceOptions));
 
   // Clean up after lowering assign_layout and various related packing code
+  pm.addPass(createApplyFolders());
   pm.addPass(createFoldConstantTensors());
   pm.addPass(createCanonicalizerPass());
   pm.addPass(createCSEPass());
@@ -385,6 +390,10 @@ void mlirToRLWEPipeline(OpPassManager& pm,
       exit(EXIT_FAILURE);
   }
 
+  pm.addPass(createForwardInsertToExtract());
+  pm.addPass(createCanonicalizerPass());
+  pm.addPass(createCSEPass());
+
   ElementwiseToAffineOptions elementwiseOptions;
   elementwiseOptions.convertDialects = {"ckks", "bgv", "lwe"};
   pm.addPass(createElementwiseToAffine(elementwiseOptions));

lib/Pipelines/BUILD

@@ -134,6 +134,8 @@ cc_library(
         "@heir//lib/Transforms/ElementwiseToAffine",
         "@heir//lib/Transforms/FoldConstantTensors",
         "@heir//lib/Transforms/FoldPlaintextMasks",
+        "@heir//lib/Transforms/ForwardInsertSliceToExtractSlice",
+        "@heir//lib/Transforms/ForwardInsertToExtract",
         "@heir//lib/Transforms/FullLoopUnroll",
         "@heir//lib/Transforms/GenerateParam",
         "@heir//lib/Transforms/InlineActivations",

lib/Target/OpenFhePke/OpenFhePkeEmitter.cpp

@@ -237,8 +237,8 @@ LogicalResult OpenFhePkeEmitter::translate(Operation& op) {
           // Arith ops
           .Case<arith::ConstantOp, arith::ExtSIOp, arith::ExtUIOp,
                 arith::IndexCastOp, arith::ExtFOp, arith::RemSIOp,
-                arith::AddIOp, arith::AndIOp, arith::SubIOp, arith::MulIOp,
-                arith::DivSIOp, arith::CmpIOp, arith::SelectOp>(
+                arith::AddIOp, arith::AddFOp, arith::AndIOp, arith::SubIOp,
+                arith::MulIOp, arith::DivSIOp, arith::CmpIOp, arith::SelectOp>(
               [&](auto op) { return printOperation(op); })
           // SCF ops
           .Case<scf::IfOp, scf::ForOp, scf::YieldOp>(
@@ -248,7 +248,7 @@ LogicalResult OpenFhePkeEmitter::translate(Operation& op) {
                 tensor::InsertSliceOp, tensor::ExtractOp,
                 tensor::ExtractSliceOp, tensor::SplatOp,
                 tensor::CollapseShapeOp, tensor::ExpandShapeOp,
-                tensor::FromElementsOp>(
+                tensor::FromElementsOp, tensor::ConcatOp>(
               [&](auto op) { return printOperation(op); })
           // LWE ops
           .Case<lwe::RLWEDecodeOp, lwe::ReinterpretApplicationDataOp>(
@@ -940,6 +940,10 @@ LogicalResult OpenFhePkeEmitter::printOperation(arith::AddIOp op) {
   return printBinaryOp(op, op.getLhs(), op.getRhs(), "+");
 }
 
+LogicalResult OpenFhePkeEmitter::printOperation(arith::AddFOp op) {
+  return printBinaryOp(op, op.getLhs(), op.getRhs(), "+");
+}
+
 LogicalResult OpenFhePkeEmitter::printOperation(arith::AndIOp op) {
   return printBinaryOp(op, op.getLhs(), op.getRhs(), "&&");
 }
@@ -993,21 +997,20 @@ LogicalResult OpenFhePkeEmitter::printOperation(arith::CmpIOp op) {
 }
 
 LogicalResult OpenFhePkeEmitter::printOperation(tensor::ConcatOp op) {
-  // concat dim(0) %foo, %foo, ...
-  // lower to a loop
-  auto operandType = cast<RankedTensorType>(op.getOperands()[0].getType());
+  // concat dim(0) %value1, %value2, ...
   auto resultType = op.getResult().getType();
   std::string varName = variableNames->getNameForValue(op.getResult());
-  if (resultType.getRank() != 1 || operandType.getRank() != 1) {
-    return failure();
-  }
+
   // std::vector<8192> result;
   if (failed(emitType(resultType, op->getLoc()))) {
     return failure();
   }
   os << " " << varName << ";\n";
 
+  // If all the operands are the same, we can just repeat the operand
+  // insertion in a loop to minimize code size.
   if (llvm::all_equal(op.getOperands())) {
+    auto operandType = cast<RankedTensorType>(op.getOperands()[0].getType());
     std::string operandName =
         variableNames->getNameForValue(op.getOperands()[0]);
     int64_t numRepeats =
@@ -1022,12 +1025,23 @@ LogicalResult OpenFhePkeEmitter::printOperation(tensor::ConcatOp op) {
 
     os.unindent();
     os << "}\n";
+  }
+
+  // If we are concatenating on dimension 0, insert the operands
+  // one by one into the result vector.
+  if (op.getDim() == 0) {
+    for (auto operand : op.getOperands()) {
+      // result.insert(result.end(), foo.begin(), foo.end());
+      std::string operandName = variableNames->getNameForValue(operand);
+      os << varName << ".insert(" << varName << ".end(), " << operandName
+         << ".begin(), " << operandName << ".end());\n";
+    }
     return success();
   }
 
   // More complicated concat ops are not supported yet. The earlier lowerings
-  // should just produce concat for lack of a "repeat" op. Maybe we should make
-  // a tensor_ext.repeat op?
+  // should just produce concat for lack of a "repeat" op. Maybe we should
+  // make a tensor_ext.repeat op?
   return failure();
 }
 
@@ -1065,8 +1079,8 @@ LogicalResult OpenFhePkeEmitter::printOperation(tensor::ExtractOp op) {
 
 LogicalResult OpenFhePkeEmitter::printOperation(
     ::mlir::tensor::CollapseShapeOp op) {
-  // A rank-reduced type will have the same number of elements so collapsing is
-  // a no-op on a flattened tensor.
+  // A rank-reduced type will have the same number of elements so collapsing
+  // is a no-op on a flattened tensor.
   SliceVerificationResult res =
       isRankReducedType(op.getSrcType(), op.getResultType());
   if (res != SliceVerificationResult::Success) {
@@ -1329,8 +1343,9 @@ LogicalResult OpenFhePkeEmitter::printOperation(
   if (failed(resultCC)) return resultCC;
   std::string cc = variableNames->getNameForValue(resultCC.value());
 
-  // In certain conditions, we might end up with the input being tensor<..xi64>
-  // which isn't a valid input type for MakeCKKSPackedPlaintext, so we convert
+  // In certain conditions, we might end up with the input being
+  // tensor<..xi64> which isn't a valid input type for
+  // MakeCKKSPackedPlaintext, so we convert
   if (getElementTypeOrSelf(op.getValue().getType()).isInteger()) {
     // This means we will have created a std::vector<int64_t>
     // but we need a std::vector<double>
@@ -1393,10 +1408,10 @@ FailureOr<std::pair<unsigned, int64_t>> getNonUnitDimension(
 }
 
 LogicalResult OpenFhePkeEmitter::printOperation(lwe::RLWEDecodeOp op) {
-  // In OpenFHE a plaintext is already decoded by decrypt. The internal OpenFHE
-  // implementation is simple enough (and dependent on currently-hard-coded
-  // encoding choices) that we will eventually need to work at a lower level of
-  // the API to support this operation properly.
+  // In OpenFHE a plaintext is already decoded by decrypt. The internal
+  // OpenFHE implementation is simple enough (and dependent on
+  // currently-hard-coded encoding choices) that we will eventually need to
+  // work at a lower level of the API to support this operation properly.
   bool isCKKS = llvm::isa<lwe::InverseCanonicalEncodingAttr>(op.getEncoding());
   auto tensorTy = dyn_cast<RankedTensorType>(op.getResult().getType());
   if (tensorTy) {

lib/Target/OpenFhePke/OpenFhePkeEmitter.h

@@ -96,6 +96,7 @@ class OpenFhePkeEmitter {
   LogicalResult printOperation(::mlir::affine::AffineForOp op);
   LogicalResult printOperation(::mlir::affine::AffineYieldOp op);
   LogicalResult printOperation(::mlir::arith::AddIOp op);
+  LogicalResult printOperation(::mlir::arith::AddFOp op);
   LogicalResult printOperation(::mlir::arith::AndIOp op);
   LogicalResult printOperation(::mlir::arith::CmpIOp op);
   LogicalResult printOperation(::mlir::arith::ConstantOp op);

lib/Transforms/ApplyFolders/ApplyFolders.cpp

@@ -3,16 +3,27 @@
 #include <utility>
 
 #include "mlir/include/mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
-#include "mlir/include/mlir/IR/MLIRContext.h"            // from @llvm-project
-#include "mlir/include/mlir/IR/PatternMatch.h"           // from @llvm-project
+#include "mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h"  // from @llvm-project
+#include "mlir/include/mlir/IR/MLIRContext.h"   // from @llvm-project
+#include "mlir/include/mlir/IR/PatternMatch.h"  // from @llvm-project
 #include "mlir/include/mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
 
+// required for generated patterns
+#include "mlir/include/mlir/IR/Matchers.h"  // from @llvm-project
+
 namespace mlir {
 namespace heir {
 
 #define GEN_PASS_DEF_APPLYFOLDERS
 #include "lib/Transforms/ApplyFolders/ApplyFolders.h.inc"
 
+namespace {
+
+// keep in anonymous namespace
+#include "lib/Transforms/ApplyFolders/Patterns.cpp.inc"
+
+}  // namespace
+
 struct ApplyFolders : impl::ApplyFoldersBase<ApplyFolders> {
   using ApplyFoldersBase::ApplyFoldersBase;
 
@@ -22,6 +33,12 @@ struct ApplyFolders : impl::ApplyFoldersBase<ApplyFolders> {
     tensor::ControlConstantExtractSliceFusionFn controlFn =
         [](tensor::ExtractSliceOp op) { return true; };
     tensor::populateFoldConstantExtractSlicePatterns(patterns, controlFn);
+    tensor::populateFoldTensorSubsetOpPatterns(patterns);
+    tensor::populateDecomposeTensorConcatPatterns(patterns);
+    tensor::populateFoldTensorEmptyPatterns(patterns);
+    tensor::populateDropRedundantInsertSliceRankExpansionPatterns(patterns);
+    tensor::populateMergeConsecutiveInsertExtractSlicePatterns(patterns);
+    populateWithGenerated(patterns);
     // Use the greedy pattern driver to apply folders.
     // TODO (#1221): Investigate whether folding (default: on) can be skipped
     // here.

lib/Transforms/ApplyFolders/BUILD

@@ -1,4 +1,5 @@
 load("@heir//lib/Transforms:transforms.bzl", "add_heir_transforms")
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
 
 package(
@@ -14,15 +15,29 @@ cc_library(
     ],
     deps = [
         ":pass_inc_gen",
+        ":patterns_inc_gen",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:TensorTransforms",
         "@llvm-project//mlir:TransformUtils",
-        "@llvm-project//mlir:Transforms",
     ],
 )
 
 add_heir_transforms(
     generated_target_name = "pass_inc_gen",
     pass_name = "ApplyFolders",
 )
+
+gentbl_cc_library(
+    name = "patterns_inc_gen",
+    tbl_outs = {"Patterns.cpp.inc": ["-gen-rewriters"]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "Patterns.td",
+    deps = [
+        "@heir//lib/Utils/DRR",
+        "@llvm-project//mlir:ArithOpsTdFiles",
+        "@llvm-project//mlir:TensorOpsTdFiles",
+    ],
+)

lib/Transforms/ApplyFolders/Patterns.td

@@ -0,0 +1,25 @@
+#ifndef LIB_TRANSFORMS_APPLYFOLDERS_PATTERNS_TD_
+#define LIB_TRANSFORMS_APPLYFOLDERS_PATTERNS_TD_
+
+include "mlir/Dialect/Arith/IR/ArithOps.td"
+include "mlir/Dialect/Tensor/IR/TensorOps.td"
+include "mlir/IR/PatternBase.td"
+include "lib/Utils/DRR/Utils.td"
+
+def AnyZero : AttrConstraint<
+    CPred<"::mlir::matchPattern($_self, m_AnyZeroFloat())">,
+    "is int or float zero">;
+
+// a + 0.0 = a
+def AddFloatingPointZero : Pat<
+  (Arith_AddFOp $a,
+    (ConstantLikeMatcher AnyZero:$value), $anyAttr),
+  (replaceWithValue $a)>;
+
+// add(empty, a) = a
+def AddEmptyTensor : Pat<
+  (Arith_AddFOp $a,
+    (Tensor_EmptyOp $b), $anyAttr),
+  (replaceWithValue $a)>;
+
+#endif // LIB_TRANSFORMS_APPLYFOLDERS_PATTERNS_TD_

lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp

@@ -540,7 +540,7 @@ struct ConvertLinalgMatvecLayout
     ImplicitLocOpBuilder b(op.getLoc(), rewriter);
     IRMaterializingVisitor visitor(
         b, input.getType(),
-        [&](Operation* createdOp) { setMaterializedAttr(op); });
+        [&](Operation* createdOp) { setMaterializedAttr(createdOp); });
     Value finalOutput = implementedKernel->visit(visitor);
 
     auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));
@@ -656,7 +656,7 @@ struct ConvertLinalgConv2D
     ImplicitLocOpBuilder b(op.getLoc(), rewriter);
     IRMaterializingVisitor visitor(
         b, data.getType(),
-        [&](Operation* createdOp) { setMaterializedAttr(op); });
+        [&](Operation* createdOp) { setMaterializedAttr(createdOp); });
     Value finalOutput = implementedKernel->visit(visitor);
 
     auto layoutAttr = cast<LayoutAttr>(op->getAttr(kLayoutAttrName));
@@ -1834,6 +1834,13 @@ struct ConvertToCiphertextSemantics
       return signalPassFailure();
     }
 
+    // Walk the IR to validate that there are no remaining unrealized conversion
+    // cast ops.
+    module->walk([&](UnrealizedConversionCastOp op) {
+      op->emitError() << "unexpected unrealized conversion cast op found";
+      signalPassFailure();
+    });
+
     clearAttrs(module, kLayoutAttrName);
     clearAttrs(module, kMaterializedAttrName);
   }

tests/Emitter/Openfhe/emit_openfhe_pke.mlir

@@ -293,8 +293,8 @@ module attributes {scheme.bgv} {
 // -----
 
 module attributes {scheme.bgv} {
-  // CHECK: test_concat
-  func.func @test_concat() -> tensor<64xi16> {
+  // CHECK: test_concat_same
+  func.func @test_concat_same() -> tensor<64xi16> {
     // CHECK: std::vector<int16_t> [[v0:.*]] =
     %cst = arith.constant dense<[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]> : tensor<32xi16>
     // CHECK: std::vector<int16_t> [[v1:.*]];
@@ -309,6 +309,41 @@ module attributes {scheme.bgv} {
 
 // -----
 
+module attributes {scheme.bgv} {
+  // CHECK: test_concat
+  func.func @test_concat() -> tensor<64xi16> {
+    // CHECK: std::vector<int16_t> [[c0:.*]] =
+    // CHECK: std::vector<int16_t> [[c1:.*]] =
+    %cst = arith.constant dense<[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]> : tensor<32xi16>
+    %cst0 = arith.constant dense<[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]> : tensor<32xi16>
+    // CHECK: std::vector<int16_t> [[v1:.*]];
+    // CHECK: [[v1]].insert([[v1]].end(), [[c0]].begin(), [[c0]].end());
+    // CHECK: [[v1]].insert([[v1]].end(), [[c1]].begin(), [[c1]].end());
+    %v = tensor.concat dim(0) %cst, %cst0 : (tensor<32xi16>, tensor<32xi16>) -> tensor<64xi16>
+
+    return %v : tensor<64xi16>
+  }
+}
+
+// -----
+
+module attributes {scheme.bgv} {
+  // CHECK: test_concat_multidim
+  func.func @test_concat_multidim() -> tensor<4x16xi16> {
+    // CHECK: std::vector<int16_t> [[c0:.*]] =
+    // CHECK: std::vector<int16_t> [[c1:.*]] =
+    %cst = arith.constant dense<[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]]> : tensor<2x16xi16>
+    %cst0 = arith.constant dense<[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]]> : tensor<2x16xi16>
+    // CHECK: std::vector<int16_t> [[v1:.*]];
+    // CHECK: [[v1]].insert([[v1]].end(), [[c0]].begin(), [[c0]].end());
+    // CHECK: [[v1]].insert([[v1]].end(), [[c1]].begin(), [[c1]].end());
+    %v = tensor.concat dim(0) %cst, %cst0 : (tensor<2x16xi16>, tensor<2x16xi16>) -> tensor<4x16xi16>
+    return %v : tensor<4x16xi16>
+  }
+}
+
+// -----
+
 module attributes {scheme.ckks} {
   // CHECK: test_insert_slice_1d
   // CHECK: std::vector<float> [[v4:[^(]*]](8, 0.100000001);

tests/Examples/openfhe/ckks/convolution/BUILD

@@ -0,0 +1,19 @@
+# See README.md for setup required to run these tests
+
+load("@heir//tests/Examples/openfhe:test.bzl", "openfhe_end_to_end_test")
+
+package(default_applicable_licenses = ["@heir//:license"])
+
+openfhe_end_to_end_test(
+    name = "convolution_test",
+    generated_lib_header = "convolution_testlib.h",
+    heir_opt_flags = [
+        "--annotate-module=backend=openfhe scheme=ckks",
+        "--torch-linalg-to-ckks=ciphertext-degree=1024",
+        "--scheme-to-openfhe",
+    ],
+    heir_translate_flags = [],
+    mlir_src = "convolution.mlir",
+    tags = ["notap"],
+    test_src = "convolution_test.cpp",
+)