From e6ae7134e791983caf4691c57333b4ff97d59a5e Mon Sep 17 00:00:00 2001
From: Andi Drebes <andi.drebes@zama.ai>
Date: Tue, 9 Apr 2024 06:47:27 +0200
Subject: [PATCH] feat(compiler): Add support for tiling of element-wise
 FHELinalg operations

---
 .../FHETensorOpsToLinalg/TensorOpsToLinalg.cpp     |  6 ++++++
 .../lib/Dialect/FHELinalg/Transforms/Tiling.cpp    | 14 +++++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp b/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp
index 1fd921355d..9e826db230 100644
--- a/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp
@@ -305,6 +305,9 @@ struct FHELinalgOpToLinalgGeneric : public mlir::OpRewritePattern<FHELinalgOp> {
                                                  ins, outs, maps, iteratorTypes,
                                                  doc, call, bodyBuilder);
 
+    if (linalgOp->hasAttr("tile-sizes"))
+      genericOp->setAttr("tile-sizes", linalgOp->getAttr("tile-sizes"));
+
     rewriter.replaceOp(linalgOp, {genericOp.getResult(0)});
 
     return ::mlir::success();
@@ -1934,6 +1937,9 @@ struct FHELinalgUnaryOpToLinalgGeneric
                                                  ins, outs, maps, iteratorTypes,
                                                  doc, call, bodyBuilder);
 
+    if (linalgOp->hasAttr("tile-sizes"))
+      genericOp->setAttr("tile-sizes", linalgOp->getAttr("tile-sizes"));
+
     rewriter.replaceOp(linalgOp, {genericOp.getResult(0)});
 
     return ::mlir::success();
diff --git a/compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/Transforms/Tiling.cpp b/compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/Transforms/Tiling.cpp
index f2a327dc2b..9865e06bf5 100644
--- a/compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/Transforms/Tiling.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/Transforms/Tiling.cpp
@@ -188,7 +188,19 @@ class FHELinalgTilingMarkerPass
     op->walk([&](mlir::Operation *op) {
       if (llvm::isa<mlir::concretelang::FHELinalg::ApplyLookupTableEintOp>(
               op) ||
-          llvm::isa<mlir::concretelang::FHELinalg::MatMulEintIntOp>(op)) {
+          llvm::isa<mlir::concretelang::FHELinalg::MatMulEintIntOp>(op) ||
+          llvm::isa<mlir::concretelang::FHELinalg::AddEintOp,
+                    mlir::concretelang::FHELinalg::AddEintIntOp,
+                    mlir::concretelang::FHELinalg::SubIntEintOp,
+                    mlir::concretelang::FHELinalg::SubEintIntOp,
+                    mlir::concretelang::FHELinalg::SubEintOp,
+                    mlir::concretelang::FHELinalg::MulEintIntOp,
+                    mlir::concretelang::FHELinalg::MulEintOp,
+                    mlir::concretelang::FHELinalg::LsbEintOp,
+                    mlir::concretelang::FHELinalg::NegEintOp,
+                    mlir::concretelang::FHELinalg::ToSignedOp,
+                    mlir::concretelang::FHELinalg::ToUnsignedOp,
+                    mlir::concretelang::FHELinalg::RoundOp>(op)) {
         op->setAttr("tile-sizes", tileAttr);
       }
     });