feat(compiler): Add support for tiling of fhelinalg.apply_multi_looku…

…p_table
zama-ai · Apr 9, 2024 · 866752a · 866752a
1 parent 2db84ae
commit 866752a
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 deletions.
diff --git a/...lers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp b/...lers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp
@@ -598,6 +598,9 @@ struct FHELinalgApplyMultiLookupTableToLinalgGeneric
             fheLinalgLutOp.getLoc(), resTypes, ins, outs, maps, iteratorTypes,
             doc, call, bodyBuilder);
 
+    if (fheLinalgLutOp->hasAttr("tile-sizes"))
+      genericOp->setAttr("tile-sizes", fheLinalgLutOp->getAttr("tile-sizes"));
+
     rewriter.replaceOp(fheLinalgLutOp, {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
@@ -186,9 +186,9 @@ class FHELinalgTilingMarkerPass
         mlir::Builder(&this->getContext()).getI64ArrayAttr(tileSizes);
 
     op->walk([&](mlir::Operation *op) {
-      if (llvm::isa<
-              mlir::concretelang::FHELinalg::ApplyLookupTableEintOp,
-              mlir::concretelang::FHELinalg::ApplyMappedLookupTableEintOp>(
+      if (llvm::isa<mlir::concretelang::FHELinalg::ApplyLookupTableEintOp,
+                    mlir::concretelang::FHELinalg::ApplyMappedLookupTableEintOp,
+                    mlir::concretelang::FHELinalg::ApplyMultiLookupTableEintOp>(
               op) ||
           llvm::isa<mlir::concretelang::FHELinalg::MatMulEintIntOp,
                     mlir::concretelang::FHELinalg::MatMulIntEintOp,

diff --git a/compilers/concrete-compiler/compiler/tests/check_tests/Dialect/FHELinalg/tiling.mlir b/compilers/concrete-compiler/compiler/tests/check_tests/Dialect/FHELinalg/tiling.mlir
@@ -85,3 +85,11 @@ func.func @apply_mapped_lookup_table(
   %0 = "FHELinalg.apply_mapped_lookup_table"(%input, %luts, %map) { "tile-sizes" = [2,3,2] } : (tensor<2x3x4x!FHE.eint<7>>, tensor<10x128xi64>, tensor<2x3x4xindex>) -> (tensor<2x3x4x!FHE.eint<7>>)
   return %0: tensor<2x3x4x!FHE.eint<7>>
 }
+
+// -----
+
+// CHECK: %[[res:.*]] = linalg.generic {indexing_maps = [#[[map:.*]], #[[map2:.*]]], iterator_types = ["parallel", "parallel"]} ins(%[[extracted_slice:.*]] : tensor<3x3x!FHE.eint<2>>) outs(%[[extracted_slice_0:.*]] : tensor<3x3x!FHE.eint<2>>) attrs =  {"tile-sizes" = [3, 3]} {
+func.func @main(%arg0: tensor<3x3x!FHE.eint<2>>, %arg1: tensor<3x3x4xi8>) -> tensor<3x3x!FHE.eint<2>> {
+  %1 = "FHELinalg.apply_multi_lookup_table"(%arg0, %arg1) { "tile-sizes" = [3, 3] }: (tensor<3x3x!FHE.eint<2>>, tensor<3x3x4xi8>) -> tensor<3x3x!FHE.eint<2>>
+  return %1: tensor<3x3x!FHE.eint<2>>
+}