From 6e1e2e50fbadd9dc303e0b4b844b6ef1908d254c Mon Sep 17 00:00:00 2001 From: Andi Drebes Date: Mon, 26 Feb 2024 15:29:27 +0100 Subject: [PATCH] feat(compiler): Add support for tiling of FHELinalg.apply_lookup_table --- .../Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp | 3 +++ .../compiler/lib/Dialect/FHELinalg/Transforms/Tiling.cpp | 8 ++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp b/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp index 3dc83a8ca..1fd921355 100644 --- a/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp +++ b/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp @@ -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(); 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 67cf34138..f2a327dc2 100644 --- a/compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/Transforms/Tiling.cpp +++ b/compilers/concrete-compiler/compiler/lib/Dialect/FHELinalg/Transforms/Tiling.cpp @@ -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( + op) || + llvm::isa(op)) { + op->setAttr("tile-sizes", tileAttr); + } }); }