From 8f8a57d220a1978f8b6fa62de292d58fe8793613 Mon Sep 17 00:00:00 2001
From: Umut <umutsahin@protonmail.com>
Date: Tue, 21 Jun 2022 09:40:51 +0200
Subject: [PATCH] feat: implement MANP calculation for tensor.insert op

---
 compiler/lib/Dialect/FHE/Analysis/MANP.cpp | 26 ++++++++++++++++++++++
 1 file changed, 26 insertions(+)

diff --git a/compiler/lib/Dialect/FHE/Analysis/MANP.cpp b/compiler/lib/Dialect/FHE/Analysis/MANP.cpp
index c7708833d..2c78a03cb 100644
--- a/compiler/lib/Dialect/FHE/Analysis/MANP.cpp
+++ b/compiler/lib/Dialect/FHE/Analysis/MANP.cpp
@@ -966,6 +966,20 @@ static llvm::APInt getSqMANP(
                    operandMANPs[1]->getValue().getMANP().getValue());
 }
 
+static llvm::APInt getSqMANP(
+    mlir::tensor::InsertOp op,
+    llvm::ArrayRef<mlir::LatticeElement<MANPLatticeValue> *> operandMANPs) {
+
+  assert(
+      operandMANPs.size() >= 2 &&
+      operandMANPs[0]->getValue().getMANP().hasValue() &&
+      operandMANPs[1]->getValue().getMANP().hasValue() &&
+      "Missing squared Minimal Arithmetic Noise Padding for encrypted operand");
+
+  return APIntUMax(operandMANPs[0]->getValue().getMANP().getValue(),
+                   operandMANPs[1]->getValue().getMANP().getValue());
+}
+
 static llvm::APInt getSqMANP(
     mlir::tensor::CollapseShapeOp op,
     llvm::ArrayRef<mlir::LatticeElement<MANPLatticeValue> *> operandMANPs) {
@@ -1274,6 +1288,18 @@ struct MANPAnalysis : public mlir::ForwardDataFlowAnalysis<MANPLatticeValue> {
         isDummy = true;
       }
     }
+    // InsertOp
+    else if (auto insertOp = llvm::dyn_cast<mlir::tensor::InsertOp>(op)) {
+      if (insertOp.result()
+              .getType()
+              .cast<mlir::TensorType>()
+              .getElementType()
+              .isa<mlir::concretelang::FHE::EncryptedIntegerType>()) {
+        norm2SqEquiv = getSqMANP(insertOp, operands);
+      } else {
+        isDummy = true;
+      }
+    }
     // InsertSliceOp
     else if (auto insertSliceOp =
                  llvm::dyn_cast<mlir::tensor::InsertSliceOp>(op)) {