From 3b00274a0219d42eaeb46e77485e6788138fb850 Mon Sep 17 00:00:00 2001
From: Andi Drebes <andi.drebes@zama.ai>
Date: Tue, 7 Mar 2023 16:03:00 +0100
Subject: [PATCH] refactor(compiler): Use FoldAdaptor for FHE and FHELinalg
 folders

As per
https://discourse.llvm.org/t/psa-new-improved-fold-method-signature-has-landed-please-update-your-downstream-projects/67618,
attribute-based folders are now deprecated and their use generates a
warning during compilation.

This patch replaces the raw attribute-based folders with folders using
`FoldAdaptor`.
---
 .../concretelang/Dialect/FHE/IR/FHEDialect.td     |  1 -
 .../Dialect/FHELinalg/IR/FHELinalgDialect.td      |  1 -
 .../compiler/lib/Dialect/FHE/IR/FHEOps.cpp        | 15 ++++++---------
 .../lib/Dialect/FHELinalg/IR/FHELinalgOps.cpp     | 15 ++++++---------
 4 files changed, 12 insertions(+), 20 deletions(-)

diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Dialect/FHE/IR/FHEDialect.td b/compilers/concrete-compiler/compiler/include/concretelang/Dialect/FHE/IR/FHEDialect.td
index 66249cc49..eabcdd301 100644
--- a/compilers/concrete-compiler/compiler/include/concretelang/Dialect/FHE/IR/FHEDialect.td
+++ b/compilers/concrete-compiler/compiler/include/concretelang/Dialect/FHE/IR/FHEDialect.td
@@ -19,7 +19,6 @@ def FHE_Dialect : Dialect {
   }];
   let cppNamespace = "::mlir::concretelang::FHE";
   let useDefaultTypePrinterParser = 1;
-  let useFoldAPI = kEmitRawAttributesFolder;
 }
 
 #endif
diff --git a/compilers/concrete-compiler/compiler/include/concretelang/Dialect/FHELinalg/IR/FHELinalgDialect.td b/compilers/concrete-compiler/compiler/include/concretelang/Dialect/FHELinalg/IR/FHELinalgDialect.td
index 32fd003b3..0bcf37c52 100644
--- a/compilers/concrete-compiler/compiler/include/concretelang/Dialect/FHELinalg/IR/FHELinalgDialect.td
+++ b/compilers/concrete-compiler/compiler/include/concretelang/Dialect/FHELinalg/IR/FHELinalgDialect.td
@@ -10,7 +10,6 @@ def FHELinalg_Dialect : Dialect {
     A dialect for representation of high level linalg operations on fully homomorphic ciphertexts.
   }];
   let cppNamespace = "::mlir::concretelang::FHELinalg";
-  let useFoldAPI = kEmitRawAttributesFolder;
 }
 
 #endif
diff --git a/compilers/concrete-compiler/compiler/lib/Dialect/FHE/IR/FHEOps.cpp b/compilers/concrete-compiler/compiler/lib/Dialect/FHE/IR/FHEOps.cpp
index 82355f185..468f35b6a 100644
--- a/compilers/concrete-compiler/compiler/lib/Dialect/FHE/IR/FHEOps.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Dialect/FHE/IR/FHEOps.cpp
@@ -300,9 +300,8 @@ mlir::LogicalResult RoundEintOp::verify() {
 }
 
 /// Avoid addition with constant 0
-OpFoldResult AddEintIntOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2);
-  auto toAdd = operands[1].dyn_cast_or_null<mlir::IntegerAttr>();
+OpFoldResult AddEintIntOp::fold(FoldAdaptor operands) {
+  auto toAdd = operands.getB().dyn_cast_or_null<mlir::IntegerAttr>();
   if (toAdd != nullptr) {
     auto intToAdd = toAdd.getInt();
     if (intToAdd == 0) {
@@ -313,9 +312,8 @@ OpFoldResult AddEintIntOp::fold(ArrayRef<Attribute> operands) {
 }
 
 /// Avoid subtraction with constant 0
-OpFoldResult SubEintIntOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2);
-  auto toSub = operands[1].dyn_cast_or_null<mlir::IntegerAttr>();
+OpFoldResult SubEintIntOp::fold(FoldAdaptor operands) {
+  auto toSub = operands.getB().dyn_cast_or_null<mlir::IntegerAttr>();
   if (toSub != nullptr) {
     auto intToSub = toSub.getInt();
     if (intToSub == 0) {
@@ -326,9 +324,8 @@ OpFoldResult SubEintIntOp::fold(ArrayRef<Attribute> operands) {
 }
 
 /// Avoid multiplication with constant 1
-OpFoldResult MulEintIntOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2);
-  auto toMul = operands[1].dyn_cast_or_null<mlir::IntegerAttr>();
+OpFoldResult MulEintIntOp::fold(FoldAdaptor operands) {
+  auto toMul = operands.getB().dyn_cast_or_null<mlir::IntegerAttr>();
   if (toMul != nullptr) {
     auto intToMul = toMul.getInt();
     if (intToMul == 1) {
diff --git a/compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/IR/FHELinalgOps.cpp b/compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/IR/FHELinalgOps.cpp
index e48db6ebc..aa8c703f1 100644
--- a/compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/IR/FHELinalgOps.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/IR/FHELinalgOps.cpp
@@ -1356,9 +1356,8 @@ mlir::LogicalResult ToUnsignedOp::verify() {
 }
 
 /// Avoid addition with constant tensor of 0s
-OpFoldResult AddEintIntOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2);
-  auto toAdd = operands[1].dyn_cast_or_null<mlir::DenseIntElementsAttr>();
+OpFoldResult AddEintIntOp::fold(FoldAdaptor operands) {
+  auto toAdd = operands.getRhs().dyn_cast_or_null<mlir::DenseIntElementsAttr>();
   if (toAdd == nullptr)
     return nullptr;
   for (auto it = toAdd.begin(); it != toAdd.end(); it++) {
@@ -1370,9 +1369,8 @@ OpFoldResult AddEintIntOp::fold(ArrayRef<Attribute> operands) {
 }
 
 /// Avoid subtraction with constant tensor of 0s
-OpFoldResult SubEintIntOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2);
-  auto toSub = operands[1].dyn_cast_or_null<mlir::DenseIntElementsAttr>();
+OpFoldResult SubEintIntOp::fold(FoldAdaptor operands) {
+  auto toSub = operands.getRhs().dyn_cast_or_null<mlir::DenseIntElementsAttr>();
   if (toSub == nullptr)
     return nullptr;
   for (auto it = toSub.begin(); it != toSub.end(); it++) {
@@ -1384,9 +1382,8 @@ OpFoldResult SubEintIntOp::fold(ArrayRef<Attribute> operands) {
 }
 
 /// Avoid multiplication with constant tensor of 1s
-OpFoldResult MulEintIntOp::fold(ArrayRef<Attribute> operands) {
-  assert(operands.size() == 2);
-  auto toMul = operands[1].dyn_cast_or_null<mlir::DenseIntElementsAttr>();
+OpFoldResult MulEintIntOp::fold(FoldAdaptor operands) {
+  auto toMul = operands.getRhs().dyn_cast_or_null<mlir::DenseIntElementsAttr>();
   if (toMul == nullptr)
     return nullptr;
   for (auto it = toMul.begin(); it != toMul.end(); it++) {