From ef778ac75b74dc656ea16aff8282a9e7bb0c63da Mon Sep 17 00:00:00 2001 From: youben11 Date: Tue, 11 Oct 2022 14:31:03 +0100 Subject: [PATCH] refactor: replace some operands by attrs in bs/ks --- .../Conversion/ConcreteToBConcrete/Pass.h | 2 +- .../Dialect/BConcrete/IR/BConcreteOps.td | 59 +++----- .../Dialect/Concrete/IR/ConcreteOps.td | 10 +- .../concretelang/Dialect/TFHE/IR/TFHEOps.td | 3 +- .../include/concretelang/Runtime/wrappers.h | 2 + .../concretelang/Support/CompilerEngine.h | 2 + .../include/concretelang/Support/Pipeline.h | 2 +- .../ConcreteToBConcrete.cpp | 98 +++++++------- .../lib/Conversion/FHEToTFHE/FHEToTFHE.cpp | 4 +- .../TFHEGlobalParametrization.cpp | 3 +- .../TFHEToConcrete/TFHEToConcrete.cpp | 20 +-- .../BufferizableOpInterfaceImpl.cpp | 126 +++++++++++++----- compiler/lib/Runtime/wrappers.cpp | 2 + compiler/lib/Support/CompilerEngine.cpp | 14 +- compiler/lib/Support/Pipeline.cpp | 6 +- .../apply_lookup_table.mlir | 20 +-- .../apply_lookup_table_cst.mlir | 19 +-- .../ConcreteToBConcrete/gpu_ops.mlir | 31 ----- .../Conversion/ConcreteToLLVM/gpu_ops.mlir | 10 ++ .../FHEToTFHE/apply_univariate.mlir | 2 +- .../FHEToTFHE/apply_univariate_cst.mlir | 2 +- .../TFHEGlobalParametrization/pbs_ks_bs.mlir | 4 +- .../Conversion/TFHEToConcrete/bootstrap.mlir | 10 +- .../check_tests/Dialect/BConcrete/ops.mlir | 12 +- .../check_tests/Dialect/Concrete/ops.mlir | 8 +- 25 files changed, 224 insertions(+), 247 deletions(-) delete mode 100644 compiler/tests/check_tests/Conversion/ConcreteToBConcrete/gpu_ops.mlir create mode 100644 compiler/tests/check_tests/Conversion/ConcreteToLLVM/gpu_ops.mlir diff --git a/compiler/include/concretelang/Conversion/ConcreteToBConcrete/Pass.h b/compiler/include/concretelang/Conversion/ConcreteToBConcrete/Pass.h index e39beaaff..acd3b0a91 100644 --- a/compiler/include/concretelang/Conversion/ConcreteToBConcrete/Pass.h +++ b/compiler/include/concretelang/Conversion/ConcreteToBConcrete/Pass.h @@ -12,7 +12,7 @@ namespace mlir { namespace concretelang { /// Create a pass to convert `Concrete` dialect to `BConcrete` dialect. std::unique_ptr> -createConvertConcreteToBConcretePass(bool loopParallelize, bool emitGPUOps); +createConvertConcreteToBConcretePass(bool loopParallelize); } // namespace concretelang } // namespace mlir diff --git a/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td index 0af28d5d2..3763ed5f4 100644 --- a/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td +++ b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td @@ -78,7 +78,9 @@ def BConcrete_KeySwitchLweBufferOp : BConcrete_Op<"keyswitch_lwe_buffer"> { // LweKeySwitchKeyType:$keyswitch_key, 1DTensorOf<[I64]>:$ciphertext, I32Attr:$level, - I32Attr:$baseLog + I32Attr:$baseLog, + I32Attr:$lwe_dim_in, + I32Attr:$lwe_dim_out ); let results = (outs 1DTensorOf<[I64]>:$result); } @@ -87,12 +89,12 @@ def BConcrete_BootstrapLweBufferOp : BConcrete_Op<"bootstrap_lwe_buffer"> { let arguments = (ins 1DTensorOf<[I64]>:$input_ciphertext, 1DTensorOf<[I64]>:$lookup_table, - I32:$inputLweDim, - I32:$polySize, - I32:$level, - I32:$baseLog, - I32:$glweDimension, - I32:$outPrecision + I32Attr:$inputLweDim, + I32Attr:$polySize, + I32Attr:$level, + I32Attr:$baseLog, + I32Attr:$glweDimension, + I32Attr:$outPrecision ); let results = (outs 1DTensorOf<[I64]>:$result); } @@ -137,12 +139,12 @@ def BConcrete_BootstrapLweBufferAsyncOffloadOp : let arguments = (ins 1DTensorOf<[I64]>:$input_ciphertext, 1DTensorOf<[I64]>:$lookup_table, - I32:$inputLweDim, - I32:$polySize, - I32:$level, - I32:$baseLog, - I32:$glweDimension, - I32:$outPrecision + I32Attr:$inputLweDim, + I32Attr:$polySize, + I32Attr:$level, + I32Attr:$baseLog, + I32Attr:$glweDimension, + I32Attr:$outPrecision ); let results = (outs RT_Future : $result); } @@ -153,35 +155,4 @@ def BConcrete_AwaitFutureOp : let results = (outs 1DTensorOf<[I64]>:$result); } -// This is a different op in BConcrete just because of the way we are lowering to CAPI -// When the CAPI lowering is detached from bufferization, we can remove this op, and lower -// to the appropriate CAPI (gpu or cpu) depending on the emitGPUOps compilation option -def BConcrete_BootstrapLweGPUBufferOp : BConcrete_Op<"bootstrap_lwe_gpu_buffer"> { - let arguments = (ins - 1DTensorOf<[I64]>:$input_ciphertext, - 1DTensorOf<[I64]>:$table, - I32:$inputLweDim, - I32:$polySize, - I32:$level, - I32:$baseLog, - I32:$glweDimension, - I32:$outPrecision - ); - let results = (outs 1DTensorOf<[I64]>:$result); -} - -// This is a different op in BConcrete just because of the way we are lowering to CAPI -// When the CAPI lowering is detached from bufferization, we can remove this op, and lower -// to the appropriate CAPI (gpu or cpu) depending on the emitGPUOps compilation option -def BConcrete_KeySwitchLweGPUBufferOp : BConcrete_Op<"keyswitch_lwe_gpu_buffer"> { - let arguments = (ins - 1DTensorOf<[I64]>:$ciphertext, - I32:$level, - I32:$baseLog, - I32:$lwe_dim_in, - I32:$lwe_dim_out - ); - let results = (outs 1DTensorOf<[I64]>:$result); -} - #endif diff --git a/compiler/include/concretelang/Dialect/Concrete/IR/ConcreteOps.td b/compiler/include/concretelang/Dialect/Concrete/IR/ConcreteOps.td index 06e99fcf2..08ff2a34a 100644 --- a/compiler/include/concretelang/Dialect/Concrete/IR/ConcreteOps.td +++ b/compiler/include/concretelang/Dialect/Concrete/IR/ConcreteOps.td @@ -58,12 +58,10 @@ def Concrete_BootstrapLweOp : Concrete_Op<"bootstrap_lwe"> { let arguments = (ins Concrete_LweCiphertextType:$input_ciphertext, 1DTensorOf<[I64]>:$lookup_table, - I32:$inputLweDim, - I32:$polySize, - I32:$level, - I32:$baseLog, - I32:$glweDimension, - I32:$outPrecision + I32Attr:$level, + I32Attr:$baseLog, + I32Attr:$polySize, + I32Attr:$glweDimension ); let results = (outs Concrete_LweCiphertextType:$result); } diff --git a/compiler/include/concretelang/Dialect/TFHE/IR/TFHEOps.td b/compiler/include/concretelang/Dialect/TFHE/IR/TFHEOps.td index 69f87730b..fec935355 100644 --- a/compiler/include/concretelang/Dialect/TFHE/IR/TFHEOps.td +++ b/compiler/include/concretelang/Dialect/TFHE/IR/TFHEOps.td @@ -99,10 +99,9 @@ def TFHE_BootstrapGLWEOp : TFHE_Op<"bootstrap_glwe"> { let arguments = (ins TFHE_GLWECipherTextType : $ciphertext, 1DTensorOf<[I64]> : $lookup_table, - I32Attr : $inputLweDim, - I32Attr : $polySize, I32Attr : $level, I32Attr : $baseLog, + I32Attr : $polySize, I32Attr : $glweDimension ); diff --git a/compiler/include/concretelang/Runtime/wrappers.h b/compiler/include/concretelang/Runtime/wrappers.h index 720640015..8d8011403 100644 --- a/compiler/include/concretelang/Runtime/wrappers.h +++ b/compiler/include/concretelang/Runtime/wrappers.h @@ -62,6 +62,8 @@ void memref_keyswitch_lwe_u64(uint64_t *out_allocated, uint64_t *out_aligned, uint64_t out_stride, uint64_t *ct0_allocated, uint64_t *ct0_aligned, uint64_t ct0_offset, uint64_t ct0_size, uint64_t ct0_stride, + uint32_t level, uint32_t base_log, + uint32_t input_lwe_dim, uint32_t output_lwe_dim, mlir::concretelang::RuntimeContext *context); void *memref_keyswitch_async_lwe_u64( uint64_t *out_allocated, uint64_t *out_aligned, uint64_t out_offset, diff --git a/compiler/include/concretelang/Support/CompilerEngine.h b/compiler/include/concretelang/Support/CompilerEngine.h index 6ef2f289b..b88687168 100644 --- a/compiler/include/concretelang/Support/CompilerEngine.h +++ b/compiler/include/concretelang/Support/CompilerEngine.h @@ -18,6 +18,8 @@ namespace mlir { namespace concretelang { +bool getEmitGPUOption(); + /// Compilation context that acts as the root owner of LLVM and MLIR /// data structures directly and indirectly referenced by artefacts /// produced by the `CompilerEngine`. diff --git a/compiler/include/concretelang/Support/Pipeline.h b/compiler/include/concretelang/Support/Pipeline.h index ed85f5b76..a5ef12ffb 100644 --- a/compiler/include/concretelang/Support/Pipeline.h +++ b/compiler/include/concretelang/Support/Pipeline.h @@ -47,7 +47,7 @@ lowerTFHEToConcrete(mlir::MLIRContext &context, mlir::ModuleOp &module, mlir::LogicalResult lowerConcreteToBConcrete(mlir::MLIRContext &context, mlir::ModuleOp &module, std::function enablePass, - bool parallelizeLoops, bool emitGPUOps); + bool parallelizeLoops); mlir::LogicalResult optimizeConcrete(mlir::MLIRContext &context, mlir::ModuleOp &module, diff --git a/compiler/lib/Conversion/ConcreteToBConcrete/ConcreteToBConcrete.cpp b/compiler/lib/Conversion/ConcreteToBConcrete/ConcreteToBConcrete.cpp index dc138b2ab..3fe631d6e 100644 --- a/compiler/lib/Conversion/ConcreteToBConcrete/ConcreteToBConcrete.cpp +++ b/compiler/lib/Conversion/ConcreteToBConcrete/ConcreteToBConcrete.cpp @@ -48,12 +48,11 @@ struct ConcreteToBConcretePass : public ConcreteToBConcreteBase { void runOnOperation() final; ConcreteToBConcretePass() = delete; - ConcreteToBConcretePass(bool loopParallelize, bool emitGPUOps) - : loopParallelize(loopParallelize), emitGPUOps(emitGPUOps){}; + ConcreteToBConcretePass(bool loopParallelize) + : loopParallelize(loopParallelize){}; private: bool loopParallelize; - bool emitGPUOps; }; } // namespace @@ -201,43 +200,67 @@ struct LowToBConcrete : public mlir::OpRewritePattern { }; }; -struct KeySwitchToGPU : public mlir::OpRewritePattern< +struct LowerKeySwitch : public mlir::OpRewritePattern< mlir::concretelang::Concrete::KeySwitchLweOp> { - KeySwitchToGPU(::mlir::MLIRContext *context, mlir::PatternBenefit benefit = 1) + LowerKeySwitch(::mlir::MLIRContext *context, mlir::PatternBenefit benefit = 1) : ::mlir::OpRewritePattern( context, benefit) {} ::mlir::LogicalResult - matchAndRewrite(mlir::concretelang::Concrete::KeySwitchLweOp keySwitchOp, + matchAndRewrite(mlir::concretelang::Concrete::KeySwitchLweOp ksOp, ::mlir::PatternRewriter &rewriter) const override { ConcreteToBConcreteTypeConverter converter; - mlir::Value levelCst = rewriter.create( - keySwitchOp.getLoc(), keySwitchOp.level(), 32); - mlir::Value baseLogCst = rewriter.create( - keySwitchOp.getLoc(), keySwitchOp.baseLog(), 32); - - // construct operands for in/out dimensions - mlir::concretelang::Concrete::LweCiphertextType outType = - keySwitchOp.getType(); - auto outDim = outType.getDimension(); - mlir::Value outDimCst = rewriter.create( - keySwitchOp.getLoc(), outDim, 32); + // construct attributes for in/out dimensions + mlir::concretelang::Concrete::LweCiphertextType outType = ksOp.getType(); + auto outDimAttr = rewriter.getI32IntegerAttr(outType.getDimension()); auto inputType = - keySwitchOp.ciphertext() + ksOp.ciphertext() .getType() .cast(); - auto inputDim = inputType.getDimension(); - mlir::Value inputDimCst = rewriter.create( - keySwitchOp.getLoc(), inputDim, 32); + mlir::IntegerAttr inputDimAttr = + rewriter.getI32IntegerAttr(inputType.getDimension()); - mlir::Operation *bKeySwitchGPUOp = rewriter.replaceOpWithNewOp< - mlir::concretelang::BConcrete::KeySwitchLweGPUBufferOp>( - keySwitchOp, outType, keySwitchOp.ciphertext(), levelCst, baseLogCst, - inputDimCst, outDimCst); + mlir::Operation *bKeySwitchOp = rewriter.replaceOpWithNewOp< + mlir::concretelang::BConcrete::KeySwitchLweBufferOp>( + ksOp, outType, ksOp.ciphertext(), ksOp.levelAttr(), ksOp.baseLogAttr(), + inputDimAttr, outDimAttr); mlir::concretelang::convertOperandAndResultTypes( - rewriter, bKeySwitchGPUOp, [&](mlir::MLIRContext *, mlir::Type t) { + rewriter, bKeySwitchOp, [&](mlir::MLIRContext *, mlir::Type t) { + return converter.convertType(t); + }); + + return ::mlir::success(); + }; +}; + +struct LowerBootstrap : public mlir::OpRewritePattern< + mlir::concretelang::Concrete::BootstrapLweOp> { + LowerBootstrap(::mlir::MLIRContext *context, mlir::PatternBenefit benefit = 1) + : ::mlir::OpRewritePattern( + context, benefit) {} + + ::mlir::LogicalResult + matchAndRewrite(mlir::concretelang::Concrete::BootstrapLweOp bsOp, + ::mlir::PatternRewriter &rewriter) const override { + ConcreteToBConcreteTypeConverter converter; + + mlir::concretelang::Concrete::LweCiphertextType outType = bsOp.getType(); + auto inputType = + bsOp.input_ciphertext() + .getType() + .cast(); + auto inputDimAttr = rewriter.getI32IntegerAttr(inputType.getDimension()); + auto outputPrecisionAttr = rewriter.getI32IntegerAttr(outType.getP()); + mlir::Operation *bBootstrapOp = rewriter.replaceOpWithNewOp< + mlir::concretelang::BConcrete::BootstrapLweBufferOp>( + bsOp, outType, bsOp.input_ciphertext(), bsOp.lookup_table(), + inputDimAttr, bsOp.polySizeAttr(), bsOp.levelAttr(), bsOp.baseLogAttr(), + bsOp.glweDimensionAttr(), outputPrecisionAttr); + + mlir::concretelang::convertOperandAndResultTypes( + rewriter, bBootstrapOp, [&](mlir::MLIRContext *, mlir::Type t) { return converter.convertType(t); }); @@ -909,6 +932,7 @@ void ConcreteToBConcretePass::runOnOperation() { // Add patterns to trivialy convert Concrete op to the equivalent // BConcrete op patterns.insert< + LowerBootstrap, LowerKeySwitch, LowToBConcrete, @@ -919,24 +943,6 @@ void ConcreteToBConcretePass::runOnOperation() { LowToBConcrete>(&getContext()); - if (this->emitGPUOps) { - patterns - .insert, - KeySwitchToGPU>(&getContext()); - } else { - patterns.insert< - LowToBConcrete, - LowToBConcrete>( - &getContext()); - } - // Add patterns to rewrite tensor operators that works on encrypted // tensors patterns @@ -1063,8 +1069,8 @@ void ConcreteToBConcretePass::runOnOperation() { namespace mlir { namespace concretelang { std::unique_ptr> -createConvertConcreteToBConcretePass(bool loopParallelize, bool emitGPUOps) { - return std::make_unique(loopParallelize, emitGPUOps); +createConvertConcreteToBConcretePass(bool loopParallelize) { + return std::make_unique(loopParallelize); } } // namespace concretelang } // namespace mlir diff --git a/compiler/lib/Conversion/FHEToTFHE/FHEToTFHE.cpp b/compiler/lib/Conversion/FHEToTFHE/FHEToTFHE.cpp index a28934863..c6d2c473c 100644 --- a/compiler/lib/Conversion/FHEToTFHE/FHEToTFHE.cpp +++ b/compiler/lib/Conversion/FHEToTFHE/FHEToTFHE.cpp @@ -113,7 +113,7 @@ struct ApplyLookupTableEintOpToKeyswitchBootstrapPattern }); // %0 = "TFHE.bootstrap_glwe"(%glwe_ks, %glwe_lut) rewriter.replaceOpWithNewOp( - lutOp, resultTy, glweKs, lutOp.lut(), -1, -1, -1, -1, -1); + lutOp, resultTy, glweKs, lutOp.lut(), -1, -1, -1, -1); return ::mlir::success(); }; }; @@ -146,8 +146,6 @@ struct ApplyLookupTableEintOpToWopPBSPattern matchAndRewrite(FHE::ApplyLookupTableEintOp lutOp, mlir::PatternRewriter &rewriter) const override { FHEToTFHETypeConverter converter; - auto inputTy = converter.convertType(lutOp.a().getType()) - .cast(); auto resultTy = converter.convertType(lutOp.getType()); // %0 = "TFHE.wop_pbs_glwe"(%ct, %lut) // : (!TFHE.glwe<{_,_,_}{2}>, tensor<4xi64>) -> diff --git a/compiler/lib/Conversion/TFHEGlobalParametrization/TFHEGlobalParametrization.cpp b/compiler/lib/Conversion/TFHEGlobalParametrization/TFHEGlobalParametrization.cpp index ce1bcacc2..bac2dac45 100644 --- a/compiler/lib/Conversion/TFHEGlobalParametrization/TFHEGlobalParametrization.cpp +++ b/compiler/lib/Conversion/TFHEGlobalParametrization/TFHEGlobalParametrization.cpp @@ -154,9 +154,8 @@ struct BootstrapGLWEOpPattern auto newOutputTy = converter.convertType(outputTy); auto newOp = rewriter.replaceOpWithNewOp( bsOp, newOutputTy, bsOp.ciphertext(), bsOp.lookup_table(), - cryptoParameters.nSmall, cryptoParameters.getPolynomialSize(), cryptoParameters.brLevel, cryptoParameters.brLogBase, - cryptoParameters.glweDimension); + cryptoParameters.getPolynomialSize(), cryptoParameters.glweDimension); rewriter.startRootUpdate(newOp); newOp.ciphertext().setType(newInputTy); rewriter.finalizeRootUpdate(newOp); diff --git a/compiler/lib/Conversion/TFHEToConcrete/TFHEToConcrete.cpp b/compiler/lib/Conversion/TFHEToConcrete/TFHEToConcrete.cpp index 7796fc4b4..749c09fe3 100644 --- a/compiler/lib/Conversion/TFHEToConcrete/TFHEToConcrete.cpp +++ b/compiler/lib/Conversion/TFHEToConcrete/TFHEToConcrete.cpp @@ -79,25 +79,9 @@ struct BootstrapGLWEOpPattern mlir::PatternRewriter &rewriter) const override { mlir::Type resultType = converter.convertType(bsOp.getType()); - auto precision = bsOp.getType().cast().getP(); - - mlir::Value inputLweDimCst = rewriter.create( - bsOp.getLoc(), bsOp.inputLweDim(), 32); - mlir::Value polySizeCst = rewriter.create( - bsOp.getLoc(), bsOp.polySize(), 32); - mlir::Value levelCst = rewriter.create( - bsOp.getLoc(), bsOp.level(), 32); - mlir::Value baseLogCst = rewriter.create( - bsOp.getLoc(), bsOp.baseLog(), 32); - mlir::Value glweDimCst = rewriter.create( - bsOp.getLoc(), bsOp.glweDimension(), 32); - mlir::Value precisionCst = rewriter.create( - bsOp.getLoc(), precision, 32); - auto newOp = rewriter.replaceOpWithNewOp( - bsOp, resultType, bsOp.ciphertext(), bsOp.lookup_table(), - inputLweDimCst, polySizeCst, levelCst, baseLogCst, glweDimCst, - precisionCst); + bsOp, resultType, bsOp.ciphertext(), bsOp.lookup_table(), bsOp.level(), + bsOp.baseLog(), bsOp.polySize(), bsOp.glweDimension()); rewriter.startRootUpdate(newOp); newOp.input_ciphertext().setType( diff --git a/compiler/lib/Dialect/BConcrete/Transforms/BufferizableOpInterfaceImpl.cpp b/compiler/lib/Dialect/BConcrete/Transforms/BufferizableOpInterfaceImpl.cpp index fcdd84b57..276f4c579 100644 --- a/compiler/lib/Dialect/BConcrete/Transforms/BufferizableOpInterfaceImpl.cpp +++ b/compiler/lib/Dialect/BConcrete/Transforms/BufferizableOpInterfaceImpl.cpp @@ -18,6 +18,7 @@ #include "concretelang/Dialect/BConcrete/IR/BConcreteDialect.h" #include "concretelang/Dialect/BConcrete/IR/BConcreteOps.h" #include "concretelang/Dialect/BConcrete/Transforms/BufferizableOpInterfaceImpl.h" +#include "concretelang/Support/CompilerEngine.h" #include #include #include @@ -110,15 +111,14 @@ mlir::LogicalResult insertForwardDeclarationOfTheCAPI( } else if (funcName == memref_negate_lwe_ciphertext_u64) { funcType = mlir::FunctionType::get(rewriter.getContext(), {memref1DType, memref1DType}, {}); - } else if (funcName == memref_keyswitch_lwe_u64) { - funcType = mlir::FunctionType::get( - rewriter.getContext(), {memref1DType, memref1DType, contextType}, {}); - } else if (funcName == memref_keyswitch_lwe_cuda_u64) { + } else if (funcName == memref_keyswitch_lwe_u64 || + funcName == memref_keyswitch_lwe_cuda_u64) { funcType = mlir::FunctionType::get(rewriter.getContext(), {memref1DType, memref1DType, i32Type, i32Type, i32Type, i32Type, contextType}, {}); - } else if (funcName == memref_bootstrap_lwe_u64) { + } else if (funcName == memref_bootstrap_lwe_u64 || + funcName == memref_bootstrap_lwe_cuda_u64) { funcType = mlir::FunctionType::get(rewriter.getContext(), {memref1DType, memref1DType, memref1DType, i32Type, i32Type, i32Type, @@ -138,12 +138,6 @@ mlir::LogicalResult insertForwardDeclarationOfTheCAPI( funcType = mlir::FunctionType::get( rewriter.getContext(), {memref1DType, futureType, memref1DType, memref1DType}, {}); - } else if (funcName == memref_bootstrap_lwe_cuda_u64) { - funcType = mlir::FunctionType::get(rewriter.getContext(), - {memref1DType, memref1DType, - memref1DType, i32Type, i32Type, i32Type, - i32Type, i32Type, i32Type, contextType}, - {}); } else if (funcName == memref_expand_lut_in_trivial_glwe_ct_u64) { funcType = mlir::FunctionType::get(rewriter.getContext(), { @@ -253,10 +247,10 @@ struct BufferizableWithCallOpInterface } }; -template +template struct BufferizableWithAsyncCallOpInterface : public BufferizableOpInterface::ExternalModel< - BufferizableWithAsyncCallOpInterface, Op> { + BufferizableWithAsyncCallOpInterface, Op> { bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand, const AnalysisState &state) const { return true; @@ -300,7 +294,7 @@ struct BufferizableWithAsyncCallOpInterface } // The first operand is the result - mlir::SmallVector operands{ + mlir::SmallVector operands{ getCastedMemRef(rewriter, loc, *outMemref), }; // For all tensor operand get the corresponding casted buffer @@ -313,10 +307,9 @@ struct BufferizableWithAsyncCallOpInterface operands.push_back(getCastedMemRef(rewriter, loc, memrefOperand)); } } - // Append the context argument - if (withContext) { - operands.push_back(getContextArgument(op)); - } + + // Append additional arguments + pushAdditionalArgs(castOp, operands, rewriter); // Insert forward declaration of the function if (insertForwardDeclarationOfTheCAPI(op, rewriter, funcName).failed()) { @@ -422,6 +415,18 @@ template <> void pushAdditionalArgs(BConcrete::KeySwitchLweBufferOp op, mlir::SmallVector &operands, RewriterBase &rewriter) { + // level + operands.push_back( + rewriter.create(op.getLoc(), op.levelAttr())); + // base_log + operands.push_back( + rewriter.create(op.getLoc(), op.baseLogAttr())); + // lwe_dim_in + operands.push_back(rewriter.create( + op.getLoc(), op.lwe_dim_inAttr())); + // lwe_dim_out + operands.push_back(rewriter.create( + op.getLoc(), op.lwe_dim_outAttr())); // context operands.push_back(getContextArgument(op)); }; @@ -430,6 +435,58 @@ template <> void pushAdditionalArgs(BConcrete::BootstrapLweBufferOp op, mlir::SmallVector &operands, RewriterBase &rewriter) { + // input_lwe_dim + operands.push_back(rewriter.create( + op.getLoc(), op.inputLweDimAttr())); + // poly_size + operands.push_back( + rewriter.create(op.getLoc(), op.polySizeAttr())); + // level + operands.push_back( + rewriter.create(op.getLoc(), op.levelAttr())); + // base_log + operands.push_back( + rewriter.create(op.getLoc(), op.baseLogAttr())); + // glwe_dim + operands.push_back(rewriter.create( + op.getLoc(), op.glweDimensionAttr())); + // out_precision + operands.push_back(rewriter.create( + op.getLoc(), op.outPrecisionAttr())); + // context + operands.push_back(getContextArgument(op)); +}; + +template <> +void pushAdditionalArgs(BConcrete::KeySwitchLweBufferAsyncOffloadOp op, + mlir::SmallVector &operands, + RewriterBase &rewriter) { + // context + operands.push_back(getContextArgument(op)); +}; + +template <> +void pushAdditionalArgs(BConcrete::BootstrapLweBufferAsyncOffloadOp op, + mlir::SmallVector &operands, + RewriterBase &rewriter) { + // input_lwe_dim + operands.push_back(rewriter.create( + op.getLoc(), op.inputLweDimAttr())); + // poly_size + operands.push_back( + rewriter.create(op.getLoc(), op.polySizeAttr())); + // level + operands.push_back( + rewriter.create(op.getLoc(), op.levelAttr())); + // base_log + operands.push_back( + rewriter.create(op.getLoc(), op.baseLogAttr())); + // glwe_dim + operands.push_back(rewriter.create( + op.getLoc(), op.glweDimensionAttr())); + // out_precision + operands.push_back(rewriter.create( + op.getLoc(), op.outPrecisionAttr())); // context operands.push_back(getContextArgument(op)); }; @@ -488,31 +545,32 @@ void mlir::concretelang::BConcrete:: BufferizableWithCallOpInterface>( *ctx); - BConcrete::KeySwitchLweGPUBufferOp::attachInterface< - BufferizableWithCallOpInterface>( - *ctx); - BConcrete::BootstrapLweGPUBufferOp::attachInterface< - BufferizableWithCallOpInterface>( - *ctx); - BConcrete::KeySwitchLweBufferOp::attachInterface< - BufferizableWithCallOpInterface>(*ctx); - BConcrete::BootstrapLweBufferOp::attachInterface< - BufferizableWithCallOpInterface>(*ctx); + if (mlir::concretelang::getEmitGPUOption()) { + BConcrete::KeySwitchLweBufferOp::attachInterface< + BufferizableWithCallOpInterface>(*ctx); + BConcrete::BootstrapLweBufferOp::attachInterface< + BufferizableWithCallOpInterface>(*ctx); + } else { + BConcrete::KeySwitchLweBufferOp::attachInterface< + BufferizableWithCallOpInterface>(*ctx); + BConcrete::BootstrapLweBufferOp::attachInterface< + BufferizableWithCallOpInterface>(*ctx); + } BConcrete::WopPBSCRTLweBufferOp::attachInterface< BufferizableWithCallOpInterface>(*ctx); BConcrete::KeySwitchLweBufferAsyncOffloadOp::attachInterface< BufferizableWithAsyncCallOpInterface< BConcrete::KeySwitchLweBufferAsyncOffloadOp, - memref_keyswitch_async_lwe_u64, true>>(*ctx); + memref_keyswitch_async_lwe_u64>>(*ctx); BConcrete::BootstrapLweBufferAsyncOffloadOp::attachInterface< BufferizableWithAsyncCallOpInterface< BConcrete::BootstrapLweBufferAsyncOffloadOp, - memref_bootstrap_async_lwe_u64, true>>(*ctx); + memref_bootstrap_async_lwe_u64>>(*ctx); BConcrete::AwaitFutureOp::attachInterface< BufferizableWithSyncCallOpInterface>(*ctx); diff --git a/compiler/lib/Runtime/wrappers.cpp b/compiler/lib/Runtime/wrappers.cpp index 98f05c3ed..b9eeb05ee 100644 --- a/compiler/lib/Runtime/wrappers.cpp +++ b/compiler/lib/Runtime/wrappers.cpp @@ -327,6 +327,8 @@ void memref_keyswitch_lwe_u64(uint64_t *out_allocated, uint64_t *out_aligned, uint64_t out_stride, uint64_t *ct0_allocated, uint64_t *ct0_aligned, uint64_t ct0_offset, uint64_t ct0_size, uint64_t ct0_stride, + uint32_t level, uint32_t base_log, + uint32_t input_lwe_dim, uint32_t output_lwe_dim, mlir::concretelang::RuntimeContext *context) { CAPI_ASSERT_ERROR( default_engine_discard_keyswitch_lwe_ciphertext_u64_raw_ptr_buffers( diff --git a/compiler/lib/Support/CompilerEngine.cpp b/compiler/lib/Support/CompilerEngine.cpp index 244f4e331..85ac53261 100644 --- a/compiler/lib/Support/CompilerEngine.cpp +++ b/compiler/lib/Support/CompilerEngine.cpp @@ -44,6 +44,11 @@ namespace mlir { namespace concretelang { +// TODO: should be removed when bufferization is not related to CAPI lowering +// Control whether we should call a cpu of gpu function when lowering +// to CAPI +static bool EMIT_GPU_OPS; +bool getEmitGPUOption() { return EMIT_GPU_OPS; } /// Creates a new compilation context that can be shared across /// compilation engines and results @@ -223,9 +228,13 @@ CompilerEngine::compile(llvm::SourceMgr &sm, Target target, OptionalLib lib) { CompilationResult res(this->compilationContext); - mlir::MLIRContext &mlirContext = *this->compilationContext->getMLIRContext(); CompilationOptions &options = this->compilerOptions; + // enable/disable usage of gpu functions during bufferization + EMIT_GPU_OPS = options.emitGPUOps; + + mlir::MLIRContext &mlirContext = *this->compilationContext->getMLIRContext(); + if (options.verifyDiagnostics) { // Only build diagnostics verifier handler if diagnostics should // be verified in order to avoid diagnostic messages to be @@ -355,8 +364,7 @@ CompilerEngine::compile(llvm::SourceMgr &sm, Target target, OptionalLib lib) { // Concrete -> BConcrete if (mlir::concretelang::pipeline::lowerConcreteToBConcrete( - mlirContext, module, this->enablePass, loopParallelize, - options.emitGPUOps) + mlirContext, module, this->enablePass, loopParallelize) .failed()) { return StreamStringError( "Lowering from Concrete to Bufferized Concrete failed"); diff --git a/compiler/lib/Support/Pipeline.cpp b/compiler/lib/Support/Pipeline.cpp index a58c1a447..72238ec18 100644 --- a/compiler/lib/Support/Pipeline.cpp +++ b/compiler/lib/Support/Pipeline.cpp @@ -242,13 +242,13 @@ optimizeConcrete(mlir::MLIRContext &context, mlir::ModuleOp &module, mlir::LogicalResult lowerConcreteToBConcrete(mlir::MLIRContext &context, mlir::ModuleOp &module, std::function enablePass, - bool parallelizeLoops, bool emitGPUOps) { + bool parallelizeLoops) { mlir::PassManager pm(&context); pipelinePrinting("ConcreteToBConcrete", pm, context); std::unique_ptr conversionPass = - mlir::concretelang::createConvertConcreteToBConcretePass(parallelizeLoops, - emitGPUOps); + mlir::concretelang::createConvertConcreteToBConcretePass( + parallelizeLoops); bool passEnabled = enablePass(conversionPass.get()); diff --git a/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/apply_lookup_table.mlir b/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/apply_lookup_table.mlir index 687fee07b..1569b3418 100644 --- a/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/apply_lookup_table.mlir +++ b/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/apply_lookup_table.mlir @@ -1,24 +1,12 @@ // RUN: concretecompiler --passes concrete-to-bconcrete --action=dump-bconcrete %s 2>&1| FileCheck %s //CHECK: func.func @apply_lookup_table(%[[A0:.*]]: tensor<1025xi64>, %[[A1:.*]]: tensor<16xi64>) -> tensor<1025xi64> { -//CHECK: %[[C1:.*]] = arith.constant 600 : i32 -//CHECK: %[[C2:.*]] = arith.constant 1024 : i32 -//CHECK: %[[C3:.*]] = arith.constant 2 : i32 -//CHECK: %[[C4:.*]] = arith.constant 3 : i32 -//CHECK: %[[C5:.*]] = arith.constant 1 : i32 -//CHECK: %[[C6:.*]] = arith.constant 4 : i32 -//CHECK: %[[V1:.*]] = "BConcrete.keyswitch_lwe_buffer"(%[[A0]]) {baseLog = 2 : i32, inputLweDimension = 1 : i32, level = 3 : i32, outputLweDimension = 600 : i32} : (tensor<1025xi64>) -> tensor<601xi64> -//CHECK: %[[V2:.*]] = "BConcrete.bootstrap_lwe_buffer"(%[[V1]], %arg1, %[[C1]], %[[C2]], %[[C3]], %[[C4]], %[[C5]], %[[C6]]) : (tensor<601xi64>, tensor<16xi64>, i32, i32, i32, i32, i32, i32) -> tensor<1025xi64> +//CHECK: %[[V1:.*]] = "BConcrete.keyswitch_lwe_buffer"(%[[A0]]) {baseLog = 2 : i32, level = 3 : i32, lwe_dim_in = 1024 : i32, lwe_dim_out = 600 : i32} : (tensor<1025xi64>) -> tensor<601xi64> +//CHECK: %[[V2:.*]] = "BConcrete.bootstrap_lwe_buffer"(%[[V1]], %arg1) {baseLog = 2 : i32, glweDimension = 4 : i32, inputLweDim = 600 : i32, level = 3 : i32, outPrecision = 4 : i32, polySize = 1024 : i32} : (tensor<601xi64>, tensor<16xi64>) -> tensor<1025xi64> //CHECK: return %[[V2]] : tensor<1025xi64> //CHECK: } func.func @apply_lookup_table(%arg0: !Concrete.lwe_ciphertext<1024,4>, %arg1: tensor<16xi64>) -> !Concrete.lwe_ciphertext<1024,4> { - %c1 = arith.constant 600 : i32 - %c2 = arith.constant 1024 : i32 - %c3 = arith.constant 2 : i32 - %c4 = arith.constant 3 : i32 - %c5 = arith.constant 1 : i32 - %c6 = arith.constant 4 : i32 - %1 = "Concrete.keyswitch_lwe"(%arg0) {baseLog = 2 : i32, inputLweDimension = 1 : i32, level = 3 : i32, outputLweDimension = 600 : i32} : (!Concrete.lwe_ciphertext<1024,4>) -> !Concrete.lwe_ciphertext<600,4> - %2 = "Concrete.bootstrap_lwe"(%1, %arg1, %c1, %c2, %c3, %c4, %c5, %c6) : (!Concrete.lwe_ciphertext<600,4>, tensor<16xi64>, i32, i32, i32, i32, i32, i32) -> !Concrete.lwe_ciphertext<1024,4> + %1 = "Concrete.keyswitch_lwe"(%arg0) {baseLog = 2 : i32, level = 3 : i32} : (!Concrete.lwe_ciphertext<1024,4>) -> !Concrete.lwe_ciphertext<600,4> + %2 = "Concrete.bootstrap_lwe"(%1, %arg1) {baseLog = 2 : i32, polySize = 1024 : i32, level = 3 : i32, glweDimension = 4 : i32} : (!Concrete.lwe_ciphertext<600,4>, tensor<16xi64> ) -> !Concrete.lwe_ciphertext<1024,4> return %2 : !Concrete.lwe_ciphertext<1024,4> } diff --git a/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/apply_lookup_table_cst.mlir b/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/apply_lookup_table_cst.mlir index 8e221d0b0..93ff17d5e 100644 --- a/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/apply_lookup_table_cst.mlir +++ b/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/apply_lookup_table_cst.mlir @@ -1,25 +1,14 @@ // RUN: concretecompiler --passes concrete-to-bconcrete --action=dump-bconcrete %s 2>&1| FileCheck %s //CHECK: func.func @apply_lookup_table_cst(%[[A0:.*]]: tensor<2049xi64>) -> tensor<2049xi64> { -//CHECK: %[[C1:.*]] = arith.constant 600 : i32 -//CHECK: %[[C2:.*]] = arith.constant 2048 : i32 -//CHECK: %[[C3:.*]] = arith.constant 4 : i32 -//CHECK: %[[C4:.*]] = arith.constant 5 : i32 -//CHECK: %[[C5:.*]] = arith.constant 1 : i32 //CHECK: %cst = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi64> -//CHECK: %[[V1:.*]] = "BConcrete.keyswitch_lwe_buffer"(%[[A0]]) {baseLog = 2 : i32, inputLweDimension = 1 : i32, level = 3 : i32, outputLweDimension = 600 : i32} : (tensor<2049xi64>) -> tensor<601xi64> -//CHECK: %[[V2:.*]] = "BConcrete.bootstrap_lwe_buffer"(%[[V1]], %cst, %[[C1]], %[[C2]], %[[C3]], %[[C4]], %[[C5]], %[[C3]]) : (tensor<601xi64>, tensor<16xi64>, i32, i32, i32, i32, i32, i32) -> tensor<2049xi64> +//CHECK: %[[V1:.*]] = "BConcrete.keyswitch_lwe_buffer"(%[[A0]]) {baseLog = 2 : i32, level = 3 : i32, lwe_dim_in = 2048 : i32, lwe_dim_out = 600 : i32} : (tensor<2049xi64>) -> tensor<601xi64> +//CHECK: %[[V2:.*]] = "BConcrete.bootstrap_lwe_buffer"(%[[V1]], %cst) {baseLog = 2 : i32, glweDimension = 4 : i32, inputLweDim = 600 : i32, level = 3 : i32, outPrecision = 4 : i32, polySize = 2048 : i32} : (tensor<601xi64>, tensor<16xi64>) -> tensor<2049xi64> //CHECK: return %[[V2]] : tensor<2049xi64> //CHECK: } func.func @apply_lookup_table_cst(%arg0: !Concrete.lwe_ciphertext<2048,4>) -> !Concrete.lwe_ciphertext<2048,4> { - %c1 = arith.constant 600 : i32 - %c2 = arith.constant 2048 : i32 - %c3 = arith.constant 4 : i32 - %c4 = arith.constant 5 : i32 - %c5 = arith.constant 1 : i32 - %c6 = arith.constant 4 : i32 %tlu = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi64> - %1 = "Concrete.keyswitch_lwe"(%arg0) {baseLog = 2 : i32, inputLweDimension = 1 : i32, level = 3 : i32, outputLweDimension = 600 : i32} : (!Concrete.lwe_ciphertext<2048,4>) -> !Concrete.lwe_ciphertext<600,4> - %2 = "Concrete.bootstrap_lwe"(%1, %tlu, %c1, %c2, %c3, %c4, %c5, %c6) : (!Concrete.lwe_ciphertext<600,4>, tensor<16xi64>, i32, i32, i32, i32, i32, i32) -> !Concrete.lwe_ciphertext<2048,4> + %1 = "Concrete.keyswitch_lwe"(%arg0) {baseLog = 2 : i32, level = 3 : i32} : (!Concrete.lwe_ciphertext<2048,4>) -> !Concrete.lwe_ciphertext<600,4> + %2 = "Concrete.bootstrap_lwe"(%1, %tlu) {baseLog = 2 : i32, polySize = 2048 : i32, level = 3 : i32, glweDimension = 4 : i32} : (!Concrete.lwe_ciphertext<600,4>, tensor<16xi64>) -> !Concrete.lwe_ciphertext<2048,4> return %2 : !Concrete.lwe_ciphertext<2048,4> } diff --git a/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/gpu_ops.mlir b/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/gpu_ops.mlir deleted file mode 100644 index 435f2ffb7..000000000 --- a/compiler/tests/check_tests/Conversion/ConcreteToBConcrete/gpu_ops.mlir +++ /dev/null @@ -1,31 +0,0 @@ -// RUN: concretecompiler --passes concrete-to-bconcrete --action=dump-bconcrete --emit-gpu-ops %s 2>&1| FileCheck %s - - -//CHECK: func.func @main(%arg0: tensor<1025xi64>) -> tensor<1025xi64> { -//CHECK: %c1_i32 = arith.constant 1 : i32 -//CHECK: %c8_i32 = arith.constant 8 : i32 -//CHECK: %c2_i32 = arith.constant 2 : i32 -//CHECK: %c1024_i32 = arith.constant 1024 : i32 -//CHECK: %c575_i32 = arith.constant 575 : i32 -//CHECK: %cst = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi64> -//CHECK: %c5_i32 = arith.constant 5 : i32 -//CHECK: %c2_i32_0 = arith.constant 2 : i32 -//CHECK: %c575_i32_1 = arith.constant 575 : i32 -//CHECK: %c1024_i32_2 = arith.constant 1024 : i32 -//CHECK: %0 = "BConcrete.keyswitch_lwe_gpu_buffer"(%arg0, %c5_i32, %c2_i32_0, %c1024_i32_2, %c575_i32_1) : (tensor<1025xi64>, i32, i32, i32, i32) -> tensor<576xi64> -//CHECK: %1 = "BConcrete.bootstrap_lwe_gpu_buffer"(%0, %cst, %c575_i32, %c1024_i32, %c2_i32, %c8_i32, %c1_i32, %c2_i32) : (tensor<576xi64>, tensor<4xi64>, i32, i32, i32, i32, i32, i32) -> tensor<1025xi64> -//CHECK: return %1 : tensor<1025xi64> -//CHECK: } -func.func @main(%arg0: !Concrete.lwe_ciphertext<1024,2>) -> !Concrete.lwe_ciphertext<1024,2> { - %c1_i32 = arith.constant 1 : i32 - %c8_i32 = arith.constant 8 : i32 - %c2_i32 = arith.constant 2 : i32 - %c1024_i32 = arith.constant 1024 : i32 - %c575_i32 = arith.constant 575 : i32 - %cst = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi64> - %0 = "Concrete.keyswitch_lwe"(%arg0) {baseLog = 2 : i32, level = 5 : i32} : (!Concrete.lwe_ciphertext<1024,2>) -> !Concrete.lwe_ciphertext<575,2> - %1 = "Concrete.bootstrap_lwe"(%0, %cst, %c575_i32, %c1024_i32, %c2_i32, %c8_i32, %c1_i32, %c2_i32) : (!Concrete.lwe_ciphertext<575,2>, tensor<4xi64>, i32, i32, i32, i32, i32, i32) -> !Concrete.lwe_ciphertext<1024,2> - return %1 : !Concrete.lwe_ciphertext<1024,2> -} - - diff --git a/compiler/tests/check_tests/Conversion/ConcreteToLLVM/gpu_ops.mlir b/compiler/tests/check_tests/Conversion/ConcreteToLLVM/gpu_ops.mlir new file mode 100644 index 000000000..f474a8195 --- /dev/null +++ b/compiler/tests/check_tests/Conversion/ConcreteToLLVM/gpu_ops.mlir @@ -0,0 +1,10 @@ +// RUN: concretecompiler --action=dump-llvm-dialect --emit-gpu-ops %s 2>&1| FileCheck %s + +//CHECK: llvm.call @memref_keyswitch_lwe_cuda_u64 +//CHECK: llvm.call @memref_bootstrap_lwe_cuda_u64 +func.func @main(%arg0: !Concrete.lwe_ciphertext<1024,2>) -> !Concrete.lwe_ciphertext<1024,2> { + %cst = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi64> + %0 = "Concrete.keyswitch_lwe"(%arg0) {baseLog = 2 : i32, level = 5 : i32} : (!Concrete.lwe_ciphertext<1024,2>) -> !Concrete.lwe_ciphertext<575,2> + %1 = "Concrete.bootstrap_lwe"(%0, %cst) {baseLog = 2 : i32, level = 5 : i32, polySize = 1024: i32, glweDimension = 1 : i32} : (!Concrete.lwe_ciphertext<575,2>, tensor<4xi64>) -> !Concrete.lwe_ciphertext<1024,2> + return %1 : !Concrete.lwe_ciphertext<1024,2> +} diff --git a/compiler/tests/check_tests/Conversion/FHEToTFHE/apply_univariate.mlir b/compiler/tests/check_tests/Conversion/FHEToTFHE/apply_univariate.mlir index 32b5b6c06..074d74929 100644 --- a/compiler/tests/check_tests/Conversion/FHEToTFHE/apply_univariate.mlir +++ b/compiler/tests/check_tests/Conversion/FHEToTFHE/apply_univariate.mlir @@ -2,7 +2,7 @@ // CHECK: func.func @apply_lookup_table(%[[A0:.*]]: !TFHE.glwe<{_,_,_}{2}>, %[[LUT:.*]]: tensor<4xi64>) -> !TFHE.glwe<{_,_,_}{3}> { // CHECK-NEXT: %[[V1:.*]] = "TFHE.keyswitch_glwe"(%[[A0]]) {baseLog = -1 : i32, level = -1 : i32} : (!TFHE.glwe<{_,_,_}{2}>) -> !TFHE.glwe<{_,_,_}{2}> -// CHECK-NEXT: %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %[[LUT]]) {baseLog = -1 : i32, glweDimension = -1 : i32, inputLweDim = -1 : i32, level = -1 : i32, polySize = -1 : i32} : (!TFHE.glwe<{_,_,_}{2}>, tensor<4xi64>) -> !TFHE.glwe<{_,_,_}{3}> +// CHECK-NEXT: %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %[[LUT]]) {baseLog = -1 : i32, glweDimension = -1 : i32, level = -1 : i32, polySize = -1 : i32} : (!TFHE.glwe<{_,_,_}{2}>, tensor<4xi64>) -> !TFHE.glwe<{_,_,_}{3}> // CHECK-NEXT: return %[[V2]] : !TFHE.glwe<{_,_,_}{3}> func.func @apply_lookup_table(%arg0: !FHE.eint<2>, %arg1: tensor<4xi64>) -> !FHE.eint<3> { %1 = "FHE.apply_lookup_table"(%arg0, %arg1): (!FHE.eint<2>, tensor<4xi64>) -> (!FHE.eint<3>) diff --git a/compiler/tests/check_tests/Conversion/FHEToTFHE/apply_univariate_cst.mlir b/compiler/tests/check_tests/Conversion/FHEToTFHE/apply_univariate_cst.mlir index 911d118f9..d58363b5d 100644 --- a/compiler/tests/check_tests/Conversion/FHEToTFHE/apply_univariate_cst.mlir +++ b/compiler/tests/check_tests/Conversion/FHEToTFHE/apply_univariate_cst.mlir @@ -3,7 +3,7 @@ //CHECK: func.func @apply_lookup_table_cst(%[[A0:.*]]: !TFHE.glwe<{_,_,_}{7}>) -> !TFHE.glwe<{_,_,_}{7}> { //CHECK-NEXT: %cst = arith.constant dense<"0x00000000000000000100000000000000020000000000000003000000000000000400000000000000050000000000000006000000000000000700000000000000080000000000000009000000000000000A000000000000000B000000000000000C000000000000000D000000000000000E000000000000000F0000000000000010000000000000001100000000000000120000000000000013000000000000001400000000000000150000000000000016000000000000001700000000000000180000000000000019000000000000001A000000000000001B000000000000001C000000000000001D000000000000001E000000000000001F0000000000000020000000000000002100000000000000220000000000000023000000000000002400000000000000250000000000000026000000000000002700000000000000280000000000000029000000000000002A000000000000002B000000000000002C000000000000002D000000000000002E000000000000002F0000000000000030000000000000003100000000000000320000000000000033000000000000003400000000000000350000000000000036000000000000003700000000000000380000000000000039000000000000003A000000000000003B000000000000003C000000000000003D000000000000003E000000000000003F0000000000000040000000000000004100000000000000420000000000000043000000000000004400000000000000450000000000000046000000000000004700000000000000480000000000000049000000000000004A000000000000004B000000000000004C000000000000004D000000000000004E000000000000004F0000000000000050000000000000005100000000000000520000000000000053000000000000005400000000000000550000000000000056000000000000005700000000000000580000000000000059000000000000005A000000000000005B000000000000005C000000000000005D000000000000005E000000000000005F0000000000000060000000000000006100000000000000620000000000000063000000000000006400000000000000650000000000000066000000000000006700000000000000680000000000000069000000000000006A000000000000006B000000000000006C000000000000006D000000000000006E000000000000006F0000000000000070000000000000007100000000000000720000000000000073000000000000007400000000000000750000000000000076000000000000007700000000000000780000000000000079000000000000007A000000000000007B000000000000007C000000000000007D000000000000007E000000000000007F00000000000000"> : tensor<128xi64> //CHECK-NEXT: %[[V1:.*]] = "TFHE.keyswitch_glwe"(%[[A0]]) {baseLog = -1 : i32, level = -1 : i32} : (!TFHE.glwe<{_,_,_}{7}>) -> !TFHE.glwe<{_,_,_}{7}> -//CHECK-NEXT: %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %cst) {baseLog = -1 : i32, glweDimension = -1 : i32, inputLweDim = -1 : i32, level = -1 : i32, polySize = -1 : i32} : (!TFHE.glwe<{_,_,_}{7}>, tensor<128xi64>) -> !TFHE.glwe<{_,_,_}{7}> +//CHECK-NEXT: %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %cst) {baseLog = -1 : i32, glweDimension = -1 : i32, level = -1 : i32, polySize = -1 : i32} : (!TFHE.glwe<{_,_,_}{7}>, tensor<128xi64>) -> !TFHE.glwe<{_,_,_}{7}> //CHECK-NEXT: return %[[V2]] : !TFHE.glwe<{_,_,_}{7}> //CHECK-NEXT: } func.func @apply_lookup_table_cst(%arg0: !FHE.eint<7>) -> !FHE.eint<7> { diff --git a/compiler/tests/check_tests/Conversion/TFHEGlobalParametrization/pbs_ks_bs.mlir b/compiler/tests/check_tests/Conversion/TFHEGlobalParametrization/pbs_ks_bs.mlir index 014ca3d40..6de2ab148 100644 --- a/compiler/tests/check_tests/Conversion/TFHEGlobalParametrization/pbs_ks_bs.mlir +++ b/compiler/tests/check_tests/Conversion/TFHEGlobalParametrization/pbs_ks_bs.mlir @@ -3,12 +3,12 @@ //CHECK: func.func @main(%[[A0:.*]]: !TFHE.glwe<{2048,1,64}{4}>) -> !TFHE.glwe<{2048,1,64}{4}> { //CHECK: %cst = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi64> //CHECK: %[[V1:.*]] = "TFHE.keyswitch_glwe"(%[[A0]]) {baseLog = 4 : i32, level = 3 : i32} : (!TFHE.glwe<{2048,1,64}{4}>) -> !TFHE.glwe<{750,1,64}{4}> -//CHECK: %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %cst) {baseLog = 23 : i32, glweDimension = 2 : i32, inputLweDim = 750 : i32, level = 1 : i32, polySize = 1024 : i32} : (!TFHE.glwe<{750,1,64}{4}>, tensor<16xi64>) -> !TFHE.glwe<{2048,1,64}{4}> +//CHECK: %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %cst) {baseLog = 23 : i32, glweDimension = 2 : i32, level = 1 : i32, polySize = 1024 : i32} : (!TFHE.glwe<{750,1,64}{4}>, tensor<16xi64>) -> !TFHE.glwe<{2048,1,64}{4}> //CHECK: return %[[V2]] : !TFHE.glwe<{2048,1,64}{4}> //CHECK: } func.func @main(%arg0: !TFHE.glwe<{_,_,_}{4}>) -> !TFHE.glwe<{_,_,_}{4}> { %cst = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi64> %1 = "TFHE.keyswitch_glwe"(%arg0) {baseLog = -1 : i32, level = -1 : i32} : (!TFHE.glwe<{_,_,_}{4}>) -> !TFHE.glwe<{_,_,_}{4}> - %2 = "TFHE.bootstrap_glwe"(%1, %cst) {baseLog = -1 : i32, glweDimension = -1 : i32, inputLweDim = -1 : i32, level = -1 : i32, polySize = -1 : i32} : (!TFHE.glwe<{_,_,_}{4}>, tensor<16xi64>) -> !TFHE.glwe<{_,_,_}{4}> + %2 = "TFHE.bootstrap_glwe"(%1, %cst) {baseLog = -1 : i32, glweDimension = -1 : i32, level = -1 : i32, polySize = -1 : i32} : (!TFHE.glwe<{_,_,_}{4}>, tensor<16xi64>) -> !TFHE.glwe<{_,_,_}{4}> return %2 : !TFHE.glwe<{_,_,_}{4}> } diff --git a/compiler/tests/check_tests/Conversion/TFHEToConcrete/bootstrap.mlir b/compiler/tests/check_tests/Conversion/TFHEToConcrete/bootstrap.mlir index d24101b1d..fadf286a6 100644 --- a/compiler/tests/check_tests/Conversion/TFHEToConcrete/bootstrap.mlir +++ b/compiler/tests/check_tests/Conversion/TFHEToConcrete/bootstrap.mlir @@ -2,17 +2,11 @@ //CHECK: func.func @bootstrap_lwe(%[[A0:.*]]: !Concrete.lwe_ciphertext<600,7>) -> !Concrete.lwe_ciphertext<1024,4> { //CHECK: %cst = arith.constant dense<"0x00000000000000000100000000000000020000000000000003000000000000000400000000000000050000000000000006000000000000000700000000000000080000000000000009000000000000000A000000000000000B000000000000000C000000000000000D000000000000000E000000000000000F0000000000000010000000000000001100000000000000120000000000000013000000000000001400000000000000150000000000000016000000000000001700000000000000180000000000000019000000000000001A000000000000001B000000000000001C000000000000001D000000000000001E000000000000001F0000000000000020000000000000002100000000000000220000000000000023000000000000002400000000000000250000000000000026000000000000002700000000000000280000000000000029000000000000002A000000000000002B000000000000002C000000000000002D000000000000002E000000000000002F0000000000000030000000000000003100000000000000320000000000000033000000000000003400000000000000350000000000000036000000000000003700000000000000380000000000000039000000000000003A000000000000003B000000000000003C000000000000003D000000000000003E000000000000003F0000000000000040000000000000004100000000000000420000000000000043000000000000004400000000000000450000000000000046000000000000004700000000000000480000000000000049000000000000004A000000000000004B000000000000004C000000000000004D000000000000004E000000000000004F0000000000000050000000000000005100000000000000520000000000000053000000000000005400000000000000550000000000000056000000000000005700000000000000580000000000000059000000000000005A000000000000005B000000000000005C000000000000005D000000000000005E000000000000005F0000000000000060000000000000006100000000000000620000000000000063000000000000006400000000000000650000000000000066000000000000006700000000000000680000000000000069000000000000006A000000000000006B000000000000006C000000000000006D000000000000006E000000000000006F0000000000000070000000000000007100000000000000720000000000000073000000000000007400000000000000750000000000000076000000000000007700000000000000780000000000000079000000000000007A000000000000007B000000000000007C000000000000007D000000000000007E000000000000007F00000000000000"> : tensor<128xi64> -//CHECK: %[[C1:.*]] = arith.constant 600 : i32 -//CHECK: %[[C2:.*]] = arith.constant 1024 : i32 -//CHECK: %[[C3:.*]] = arith.constant 3 : i32 -//CHECK: %[[C4:.*]] = arith.constant 1 : i32 -//CHECK: %[[C5:.*]] = arith.constant 1 : i32 -//CHECK: %[[C6:.*]] = arith.constant 4 : i32 -//CHECK: %[[V1:.*]] = "Concrete.bootstrap_lwe"(%arg0, %cst, %[[C1]], %[[C2]], %[[C3]], %[[C4]], %[[C5]], %[[C6]]) : (!Concrete.lwe_ciphertext<600,7>, tensor<128xi64>, i32, i32, i32, i32, i32, i32) -> !Concrete.lwe_ciphertext<1024,4> +//CHECK: %[[V1:.*]] = "Concrete.bootstrap_lwe"(%arg0, %cst) {baseLog = 1 : i32, glweDimension = 1 : i32, level = 3 : i32, polySize = 1024 : i32} : (!Concrete.lwe_ciphertext<600,7>, tensor<128xi64>) -> !Concrete.lwe_ciphertext<1024,4> //CHECK: return %[[V1]] : !Concrete.lwe_ciphertext<1024,4> //CHECK: } func.func @bootstrap_lwe(%ciphertext: !TFHE.glwe<{600,1,64}{7}>) -> !TFHE.glwe<{1024,1,64}{4}> { %cst = arith.constant dense<"0x00000000000000000100000000000000020000000000000003000000000000000400000000000000050000000000000006000000000000000700000000000000080000000000000009000000000000000A000000000000000B000000000000000C000000000000000D000000000000000E000000000000000F0000000000000010000000000000001100000000000000120000000000000013000000000000001400000000000000150000000000000016000000000000001700000000000000180000000000000019000000000000001A000000000000001B000000000000001C000000000000001D000000000000001E000000000000001F0000000000000020000000000000002100000000000000220000000000000023000000000000002400000000000000250000000000000026000000000000002700000000000000280000000000000029000000000000002A000000000000002B000000000000002C000000000000002D000000000000002E000000000000002F0000000000000030000000000000003100000000000000320000000000000033000000000000003400000000000000350000000000000036000000000000003700000000000000380000000000000039000000000000003A000000000000003B000000000000003C000000000000003D000000000000003E000000000000003F0000000000000040000000000000004100000000000000420000000000000043000000000000004400000000000000450000000000000046000000000000004700000000000000480000000000000049000000000000004A000000000000004B000000000000004C000000000000004D000000000000004E000000000000004F0000000000000050000000000000005100000000000000520000000000000053000000000000005400000000000000550000000000000056000000000000005700000000000000580000000000000059000000000000005A000000000000005B000000000000005C000000000000005D000000000000005E000000000000005F0000000000000060000000000000006100000000000000620000000000000063000000000000006400000000000000650000000000000066000000000000006700000000000000680000000000000069000000000000006A000000000000006B000000000000006C000000000000006D000000000000006E000000000000006F0000000000000070000000000000007100000000000000720000000000000073000000000000007400000000000000750000000000000076000000000000007700000000000000780000000000000079000000000000007A000000000000007B000000000000007C000000000000007D000000000000007E000000000000007F00000000000000"> : tensor<128xi64> - %bootstraped = "TFHE.bootstrap_glwe"(%ciphertext, %cst) {baseLog = 1 : i32, glweDimension = 1 : i32, inputLweDim = 600 : i32, level = 3 : i32, polySize = 1024 : i32} : (!TFHE.glwe<{600,1,64}{7}>, tensor<128xi64>) -> !TFHE.glwe<{1024,1,64}{4}> + %bootstraped = "TFHE.bootstrap_glwe"(%ciphertext, %cst) {baseLog = 1 : i32, glweDimension = 1 : i32, level = 3 : i32, polySize = 1024 : i32} : (!TFHE.glwe<{600,1,64}{7}>, tensor<128xi64>) -> !TFHE.glwe<{1024,1,64}{4}> return %bootstraped : !TFHE.glwe<{1024,1,64}{4}> } diff --git a/compiler/tests/check_tests/Dialect/BConcrete/ops.mlir b/compiler/tests/check_tests/Dialect/BConcrete/ops.mlir index 84fb12927..5dc8fa7b5 100644 --- a/compiler/tests/check_tests/Dialect/BConcrete/ops.mlir +++ b/compiler/tests/check_tests/Dialect/BConcrete/ops.mlir @@ -72,20 +72,20 @@ func.func @negate_crt_lwe_ciphertext(%arg0: tensor<5x2049xi64>) -> tensor<5x2049 return %0 : tensor<5x2049xi64> } -//CHECK: func.func @bootstrap_lwe(%arg0: tensor<2049xi64>, %arg1: tensor<16xi64>, %arg2: i32, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32) -> tensor<2049xi64> { -//CHECK: %[[V0:.*]] = "BConcrete.bootstrap_lwe_buffer"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7) : (tensor<2049xi64>, tensor<16xi64>, i32, i32, i32, i32, i32, i32) -> tensor<2049xi64> +//CHECK: func.func @bootstrap_lwe(%arg0: tensor<2049xi64>, %arg1: tensor<16xi64>) -> tensor<2049xi64> { +//CHECK: %[[V0:.*]] = "BConcrete.bootstrap_lwe_buffer"(%arg0, %arg1) {baseLog = 2 : i32, glweDimension = 4 : i32, inputLweDim = 600 : i32, level = 3 : i32, outPrecision = 4 : i32, polySize = 2048 : i32} : (tensor<2049xi64>, tensor<16xi64>) -> tensor<2049xi64> //CHECK: return %[[V0]] : tensor<2049xi64> //CHECK: } -func.func @bootstrap_lwe(%arg0: tensor<2049xi64>, %arg1: tensor<16xi64>, %arg2: i32, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32) -> tensor<2049xi64> { - %0 = "BConcrete.bootstrap_lwe_buffer"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7) : (tensor<2049xi64>, tensor<16xi64>, i32, i32, i32, i32, i32, i32) -> (tensor<2049xi64>) +func.func @bootstrap_lwe(%arg0: tensor<2049xi64>, %arg1: tensor<16xi64>) -> tensor<2049xi64> { + %0 = "BConcrete.bootstrap_lwe_buffer"(%arg0, %arg1) {baseLog = 2 : i32, glweDimension = 4 : i32, inputLweDim = 600 : i32, level = 3 : i32, outPrecision = 4 : i32, polySize = 2048 : i32} : (tensor<2049xi64>, tensor<16xi64>) -> (tensor<2049xi64>) return %0 : tensor<2049xi64> } //CHECK: func.func @keyswitch_lwe(%[[A0:.*]]: tensor<2049xi64>) -> tensor<2049xi64> { -//CHECK: %[[V0:.*]] = "BConcrete.keyswitch_lwe_buffer"(%[[A0]]) {baseLog = 2 : i32, inputLweDimension = 1 : i32, level = 3 : i32, outputLweDimension = 1 : i32} : (tensor<2049xi64>) -> tensor<2049xi64> +//CHECK: %[[V0:.*]] = "BConcrete.keyswitch_lwe_buffer"(%[[A0]]) {baseLog = 2 : i32, level = 3 : i32, lwe_dim_in = 2048 : i32, lwe_dim_out = 2048 : i32} : (tensor<2049xi64>) -> tensor<2049xi64> //CHECK: return %[[V0]] : tensor<2049xi64> //CHECK: } func.func @keyswitch_lwe(%arg0: tensor<2049xi64>) -> tensor<2049xi64> { - %0 = "BConcrete.keyswitch_lwe_buffer"(%arg0) {baseLog = 2 : i32, inputLweDimension = 1 : i32, level = 3 : i32, outputLweDimension = 1 : i32} : (tensor<2049xi64>) -> (tensor<2049xi64>) + %0 = "BConcrete.keyswitch_lwe_buffer"(%arg0) {baseLog = 2 : i32, level = 3 : i32, lwe_dim_in = 2048 : i32, lwe_dim_out = 2048 : i32} : (tensor<2049xi64>) -> (tensor<2049xi64>) return %0 : tensor<2049xi64> } diff --git a/compiler/tests/check_tests/Dialect/Concrete/ops.mlir b/compiler/tests/check_tests/Dialect/Concrete/ops.mlir index 6f10384c4..970bb73f0 100644 --- a/compiler/tests/check_tests/Dialect/Concrete/ops.mlir +++ b/compiler/tests/check_tests/Dialect/Concrete/ops.mlir @@ -36,11 +36,11 @@ func.func @negate_lwe_ciphertext(%arg0: !Concrete.lwe_ciphertext<2048,7>) -> !Co return %1: !Concrete.lwe_ciphertext<2048,7> } -// CHECK-LABEL: func.func @bootstrap_lwe(%arg0: !Concrete.lwe_ciphertext<2048,7>, %arg1: tensor<128xi64>, %arg2: i32, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32) -> !Concrete.lwe_ciphertext<2048,7> -func.func @bootstrap_lwe(%arg0: !Concrete.lwe_ciphertext<2048,7>, %arg1: tensor<128xi64>, %arg2: i32, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32) -> !Concrete.lwe_ciphertext<2048,7> { - // CHECK-NEXT: %[[V1:.*]] = "Concrete.bootstrap_lwe"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7) : (!Concrete.lwe_ciphertext<2048,7>, tensor<128xi64>, i32, i32, i32, i32, i32, i32) -> !Concrete.lwe_ciphertext<2048,7> +// CHECK-LABEL: func.func @bootstrap_lwe(%arg0: !Concrete.lwe_ciphertext<2048,7>, %arg1: tensor<128xi64>) -> !Concrete.lwe_ciphertext<2048,7> +func.func @bootstrap_lwe(%arg0: !Concrete.lwe_ciphertext<2048,7>, %arg1: tensor<128xi64>) -> !Concrete.lwe_ciphertext<2048,7> { + // CHECK-NEXT: %[[V1:.*]] = "Concrete.bootstrap_lwe"(%arg0, %arg1) {baseLog = 2 : i32, glweDimension = 4 : i32, level = 3 : i32, polySize = 2048 : i32} : (!Concrete.lwe_ciphertext<2048,7>, tensor<128xi64>) -> !Concrete.lwe_ciphertext<2048,7> // CHECK-NEXT: return %[[V1]] : !Concrete.lwe_ciphertext<2048,7> - %1 = "Concrete.bootstrap_lwe"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7) : (!Concrete.lwe_ciphertext<2048,7>, tensor<128xi64>, i32, i32, i32, i32, i32, i32) -> !Concrete.lwe_ciphertext<2048,7> + %1 = "Concrete.bootstrap_lwe"(%arg0, %arg1) {baseLog = 2 : i32, polySize = 2048 : i32, level = 3 : i32, glweDimension = 4 : i32} : (!Concrete.lwe_ciphertext<2048,7>, tensor<128xi64>) -> !Concrete.lwe_ciphertext<2048,7> return %1: !Concrete.lwe_ciphertext<2048,7> }