feat(compiler): Add support for tiling of optimizer.partition_frontier

zama-ai · Apr 12, 2024 · 2048105 · 2048105
1 parent fcfaaee
commit 2048105
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/...lers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp b/...lers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp
@@ -2021,6 +2021,9 @@ struct TensorPartitionFrontierOpToLinalgGeneric
                                                  outs, maps, iteratorTypes, doc,
                                                  call, bodyBuilder);
 
+    if (pfOp->hasAttr("tile-sizes"))
+      genericOp->setAttr("tile-sizes", pfOp->getAttr("tile-sizes"));
+
     rewriter.replaceOp(pfOp, {genericOp.getResult(0)});
 
     return ::mlir::success();