diff --git a/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp b/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp index 1fd921355d..9e826db230 100644 --- a/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp +++ b/compilers/concrete-compiler/compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp @@ -305,6 +305,9 @@ struct FHELinalgOpToLinalgGeneric : public mlir::OpRewritePattern { 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 index f2a327dc2b..9865e06bf5 100644 --- 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( op) || - llvm::isa(op)) { + llvm::isa(op) || + llvm::isa(op)) { op->setAttr("tile-sizes", tileAttr); } });