github/concrete
1
1
Fork 0
mirror of https://github.com/zama-ai/concrete.git synced 2026-04-17 03:00:54 -04:00
Code Issues Packages Projects Releases Wiki Activity

feat(compiler): Add support for tiling of FHELinalg.apply_lookup_table

Browse Source
This commit is contained in:
Andi Drebes
2024-02-26 15:29:27 +01:00
parent 6a2db92f8f
commit 59ded11911
2 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp
View File

@@ -697,6 +697,9 @@ struct FHELinalgApplyLookupTableToLinalgGeneric
outs, maps, iteratorTypes, doc,
call, bodyBuilder);
if (lutOp->hasAttr("tile-sizes"))
genericOp->setAttr("tile-sizes", lutOp->getAttr("tile-sizes"));
rewriter.replaceOp(lutOp, {genericOp.getResult(0)});
return ::mlir::success();

8
compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/Transforms/Tiling.cpp
View File

@@ -185,8 +185,12 @@ public:
mlir::ArrayAttr tileAttr =
mlir::Builder(&this->getContext()).getI64ArrayAttr(tileSizes);
op->walk([&](mlir::concretelang::FHELinalg::MatMulEintIntOp matmulOp) {
matmulOp.getOperation()->setAttr("tile-sizes", tileAttr);
op->walk([&](mlir::Operation *op) {
if (llvm::isa<mlir::concretelang::FHELinalg::ApplyLookupTableEintOp>(
op) ||
llvm::isa<mlir::concretelang::FHELinalg::MatMulEintIntOp>(op)) {
op->setAttr("tile-sizes", tileAttr);
}
});
}