feat(compiler): Add support for tiling of element-wise FHELinalg oper…

…ations
zama-ai · Apr 9, 2024 · e6ae713 · e6ae713
1 parent f506f5f
commit e6ae713
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/...lers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp b/...lers/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
@@ -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);
       }
     });