github/concrete
1
1
Fork 0
mirror of https://github.com/zama-ai/concrete.git synced 2026-02-09 12:15:09 -05:00
Code Issues Packages Projects Releases Wiki Activity

refactor: lower our conv2d to custom linalg named op

Browse Source 
we can now generate linalg named op with custom operation for add/mul to
handle our types
This commit is contained in:
youben11
2022-06-23 10:06:48 +01:00
committed by Ayoub Benaissa
parent 83f2095af5
commit f1f1db923d
1 changed files with 9 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
compiler/lib/Conversion/FHETensorOpsToLinalg/TensorOpsToLinalg.cpp
View File

@@ -1602,11 +1602,17 @@ struct FHELinalgConv2dToLinalgConv2d
auto stridesAttr = rewriter.getI64VectorAttr(stridesInts);
auto dilationsAttr = rewriter.getI64VectorAttr(dilationsInts);
rewriter.replaceOpWithNewOp<
mlir::concretelang::FHELinalg::FhelinalgConv2DNchwFchwOp>(
auto addOpAttr = rewriter.getNamedAttr(
"add", rewriter.getStringAttr(
mlir::concretelang::FHE::AddEintOp::getOperationName()));
auto mulOpAttr = rewriter.getNamedAttr(
"mul", rewriter.getStringAttr(
mlir::concretelang::FHE::MulEintIntOp::getOperationName()));
rewriter.replaceOpWithNewOp<mlir::linalg::Conv2DNchwFchwOp>(
conv2dOp, biasInitTensor.getType(),
mlir::ValueRange{paddedInput, weight}, biasInitTensor, stridesAttr,
dilationsAttr);
dilationsAttr,
llvm::ArrayRef<mlir::NamedAttribute>({addOpAttr, mulOpAttr}));
return mlir::success();
};
};
@@ -1630,9 +1636,6 @@ void FHETensorOpsToLinalg::runOnOperation() {
target.addLegalDialect<mlir::arith::ArithmeticDialect>();
target.addIllegalOp<mlir::concretelang::FHELinalg::Dot>();
target.addIllegalDialect<mlir::concretelang::FHELinalg::FHELinalgDialect>();
// TODO: this should be removed when we no longer need a custom generated op
// for conv that works on tensors of custom types
target.addLegalOp<mlir::concretelang::FHELinalg::FhelinalgConv2DNchwFchwOp>();
target.addLegalOp<bufferization::AllocTensorOp>();