diff --git a/compiler/include/concretelang/Dialect/BConcrete/CMakeLists.txt b/compiler/include/concretelang/Dialect/BConcrete/CMakeLists.txt new file mode 100644 index 000000000..f33061b2d --- /dev/null +++ b/compiler/include/concretelang/Dialect/BConcrete/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteDialect.h b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteDialect.h new file mode 100644 index 000000000..51c86df38 --- /dev/null +++ b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteDialect.h @@ -0,0 +1,18 @@ +// Part of the Concrete Compiler Project, under the BSD3 License with Zama +// Exceptions. See +// https://github.com/zama-ai/concrete-compiler-internal/blob/master/LICENSE.txt +// for license information. + +#ifndef ZAMALANG_DIALECT_BConcrete_IR_BConcreteDIALECT_H +#define ZAMALANG_DIALECT_BConcrete_IR_BConcreteDIALECT_H + +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Dialect.h" + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Dialect.h" + +#include "concretelang/Dialect/BConcrete/IR/BConcreteOpsDialect.h.inc" + +#endif diff --git a/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteDialect.td b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteDialect.td new file mode 100644 index 000000000..df4524c16 --- /dev/null +++ b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteDialect.td @@ -0,0 +1,15 @@ +#ifndef ZAMALANG_DIALECT_BConcrete_IR_BConcrete_DIALECT +#define ZAMALANG_DIALECT_BConcrete_IR_BConcrete_DIALECT + +include "mlir/IR/OpBase.td" + +def BConcrete_Dialect : Dialect { + let name = "BConcrete"; + let summary = "Bufferized concrete dialect"; + let description = [{ + A dialect for representation of bufferized concrete operations on fully homomorphic ciphertext. + }]; + let cppNamespace = "::mlir::concretelang::BConcrete"; +} + +#endif diff --git a/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.h b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.h new file mode 100644 index 000000000..f3ed6efa4 --- /dev/null +++ b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.h @@ -0,0 +1,20 @@ +// Part of the Concrete Compiler Project, under the BSD3 License with Zama +// Exceptions. See +// https://github.com/zama-ai/concrete-compiler-internal/blob/master/LICENSE.txt +// for license information. + +#ifndef ZAMALANG_DIALECT_BConcrete_BConcrete_OPS_H +#define ZAMALANG_DIALECT_BConcrete_BConcrete_OPS_H + +#include +#include +#include +#include +#include + +#include "concretelang/Dialect/Concrete/IR/ConcreteTypes.h" + +#define GET_OP_CLASSES +#include "concretelang/Dialect/BConcrete/IR/BConcreteOps.h.inc" + +#endif diff --git a/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td new file mode 100644 index 000000000..44ff5ca92 --- /dev/null +++ b/compiler/include/concretelang/Dialect/BConcrete/IR/BConcreteOps.td @@ -0,0 +1,65 @@ +#ifndef ZAMALANG_DIALECT_BConcrete_IR_BConcrete_OPS +#define ZAMALANG_DIALECT_BConcrete_IR_BConcrete_OPS + +include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/Interfaces/ControlFlowInterfaces.td" +include "mlir/IR/BuiltinTypes.td" +include "mlir/Dialect/MemRef/IR/MemRefBase.td" + +include "concretelang/Dialect/BConcrete/IR/BConcreteDialect.td" +include "concretelang/Dialect/Concrete/IR/ConcreteTypes.td" + +class BConcrete_Op traits = []> : + Op; + +def AddLweBuffersOp : BConcrete_Op<"add_lwe_buffer"> { + let arguments = (ins + 1DTensorOf<[I64]>:$result, + 1DTensorOf<[I64]>:$lhs, + 1DTensorOf<[I64]>:$rhs + ); + let results = (outs); +} + +def AddPlaintextLweBufferOp : BConcrete_Op<"add_plaintext_lwe_buffer"> { + let arguments = (ins 1DTensorOf<[I64]>:$result, 1DTensorOf<[I64]>:$lhs, PlaintextType:$rhs); + let results = (outs); +} + +def MulCleartextLweBufferOp : BConcrete_Op<"mul_cleartext_lwe_buffer"> { + let arguments = (ins 1DTensorOf<[I64]>:$result, 1DTensorOf<[I64]>:$lhs, CleartextType:$rhs); + let results = (outs); +} + +def NegateLweBufferOp : BConcrete_Op<"negate_lwe_buffer"> { + let arguments = (ins 1DTensorOf<[I64]>:$result, 1DTensorOf<[I64]>:$ciphertext); + let results = (outs); +} + +def KeySwitchLweBufferOp : BConcrete_Op<"keyswitch_lwe_buffer"> { + let arguments = (ins + 1DTensorOf<[I64]>:$result, + // LweKeySwitchKeyType:$keyswitch_key, + 1DTensorOf<[I64]>:$ciphertext, + I32Attr:$level, + I32Attr:$baseLog + ); + let results = (outs); +} + +def BootstrapLweBufferOp : BConcrete_Op<"bootstrap_lwe_buffer"> { + let arguments = (ins + 1DTensorOf<[I64]>:$result, + // LweBootstrapKeyType:$bootstrap_key, + 1DTensorOf<[I64]>:$input_ciphertext, + GlweCiphertextType:$accumulator, + I32Attr:$glweDimension, + I32Attr:$polynomialSize, + I32Attr:$level, + I32Attr:$baseLog + ); + let results = (outs); +} + + +#endif diff --git a/compiler/include/concretelang/Dialect/BConcrete/IR/CMakeLists.txt b/compiler/include/concretelang/Dialect/BConcrete/IR/CMakeLists.txt new file mode 100644 index 000000000..b96f45f0e --- /dev/null +++ b/compiler/include/concretelang/Dialect/BConcrete/IR/CMakeLists.txt @@ -0,0 +1,9 @@ +set(LLVM_TARGET_DEFINITIONS BConcreteOps.td) +mlir_tablegen(BConcreteOps.h.inc -gen-op-decls) +mlir_tablegen(BConcreteOps.cpp.inc -gen-op-defs) +mlir_tablegen(BConcreteOpsTypes.h.inc -gen-typedef-decls -typedefs-dialect=BConcrete) +mlir_tablegen(BConcreteOpsTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=BConcrete) +mlir_tablegen(BConcreteOpsDialect.h.inc -gen-dialect-decls -dialect=BConcrete) +mlir_tablegen(BConcreteOpsDialect.cpp.inc -gen-dialect-defs -dialect=BConcrete) +add_public_tablegen_target(MLIRBConcreteOpsIncGen) +add_dependencies(mlir-headers MLIRBConcreteOpsIncGen) \ No newline at end of file diff --git a/compiler/include/concretelang/Dialect/CMakeLists.txt b/compiler/include/concretelang/Dialect/CMakeLists.txt index 7075a2e7f..ed435f204 100644 --- a/compiler/include/concretelang/Dialect/CMakeLists.txt +++ b/compiler/include/concretelang/Dialect/CMakeLists.txt @@ -2,4 +2,5 @@ add_subdirectory(FHE) add_subdirectory(FHELinalg) add_subdirectory(TFHE) add_subdirectory(Concrete) +add_subdirectory(BConcrete) add_subdirectory(RT) diff --git a/compiler/lib/Dialect/BConcrete/CMakeLists.txt b/compiler/lib/Dialect/BConcrete/CMakeLists.txt new file mode 100644 index 000000000..f33061b2d --- /dev/null +++ b/compiler/lib/Dialect/BConcrete/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/compiler/lib/Dialect/BConcrete/IR/BConcreteDialect.cpp b/compiler/lib/Dialect/BConcrete/IR/BConcreteDialect.cpp new file mode 100644 index 000000000..d9bd81c0f --- /dev/null +++ b/compiler/lib/Dialect/BConcrete/IR/BConcreteDialect.cpp @@ -0,0 +1,26 @@ +// Part of the Concrete Compiler Project, under the BSD3 License with Zama +// Exceptions. See +// https://github.com/zama-ai/concrete-compiler-internal/blob/master/LICENSE.txt +// for license information. + +#include "concretelang/Dialect/BConcrete/IR/BConcreteDialect.h" +#include "concretelang/Dialect/BConcrete/IR/BConcreteOps.h" + +#define GET_TYPEDEF_CLASSES +#include "concretelang/Dialect/BConcrete/IR/BConcreteOpsTypes.cpp.inc" + +#include "concretelang/Dialect/BConcrete/IR/BConcreteOpsDialect.cpp.inc" + +using namespace mlir::concretelang::BConcrete; + +void BConcreteDialect::initialize() { + addOperations< +#define GET_OP_LIST +#include "concretelang/Dialect/BConcrete/IR/BConcreteOps.cpp.inc" + >(); + + addTypes< +#define GET_TYPEDEF_LIST +#include "concretelang/Dialect/BConcrete/IR/BConcreteOpsTypes.cpp.inc" + >(); +} diff --git a/compiler/lib/Dialect/BConcrete/IR/BConcreteOps.cpp b/compiler/lib/Dialect/BConcrete/IR/BConcreteOps.cpp new file mode 100644 index 000000000..edea56d9e --- /dev/null +++ b/compiler/lib/Dialect/BConcrete/IR/BConcreteOps.cpp @@ -0,0 +1,9 @@ +// Part of the Concrete Compiler Project, under the BSD3 License with Zama +// Exceptions. See +// https://github.com/zama-ai/concrete-compiler-internal/blob/master/LICENSE.txt +// for license information. + +#include "concretelang/Dialect/BConcrete/IR/BConcreteOps.h" + +#define GET_OP_CLASSES +#include "concretelang/Dialect/BConcrete/IR/BConcreteOps.cpp.inc" diff --git a/compiler/lib/Dialect/BConcrete/IR/CMakeLists.txt b/compiler/lib/Dialect/BConcrete/IR/CMakeLists.txt new file mode 100644 index 000000000..c233fc380 --- /dev/null +++ b/compiler/lib/Dialect/BConcrete/IR/CMakeLists.txt @@ -0,0 +1,14 @@ +add_mlir_dialect_library(BConcreteDialect + BConcreteDialect.cpp + BConcreteOps.cpp + + ADDITIONAL_HEADER_DIRS + ${PROJECT_SOURCE_DIR}/include/concretelang/Dialect/BConcrete + + DEPENDS + MLIRBConcreteOpsIncGen + + LINK_LIBS PUBLIC + MLIRIR) + +target_link_libraries(BConcreteDialect PUBLIC MLIRIR) diff --git a/compiler/lib/Dialect/CMakeLists.txt b/compiler/lib/Dialect/CMakeLists.txt index 22156f5be..2a809ceab 100644 --- a/compiler/lib/Dialect/CMakeLists.txt +++ b/compiler/lib/Dialect/CMakeLists.txt @@ -2,4 +2,5 @@ add_subdirectory(FHELinalg) add_subdirectory(FHE) add_subdirectory(TFHE) add_subdirectory(Concrete) +add_subdirectory(BConcrete) add_subdirectory(RT) diff --git a/compiler/lib/Support/CompilerEngine.cpp b/compiler/lib/Support/CompilerEngine.cpp index 42b8eea55..274ef0cd3 100644 --- a/compiler/lib/Support/CompilerEngine.cpp +++ b/compiler/lib/Support/CompilerEngine.cpp @@ -19,6 +19,7 @@ #include #include +#include #include #include #include @@ -61,6 +62,8 @@ mlir::MLIRContext *CompilationContext::getMLIRContext() { ->getOrLoadDialect(); this->mlirContext ->getOrLoadDialect(); + this->mlirContext + ->getOrLoadDialect(); this->mlirContext->getOrLoadDialect(); this->mlirContext->getOrLoadDialect(); this->mlirContext->getOrLoadDialect(); diff --git a/compiler/src/CMakeLists.txt b/compiler/src/CMakeLists.txt index 168fdab9f..1c8ae4938 100644 --- a/compiler/src/CMakeLists.txt +++ b/compiler/src/CMakeLists.txt @@ -12,6 +12,7 @@ if(CONCRETELANG_PARALLEL_EXECUTION_ENABLED) ${conversion_libs} MLIRTransforms + BConcreteDialect ConcreteDialect TFHEDialect FHEDialect @@ -37,6 +38,7 @@ else() ${conversion_libs} MLIRTransforms + BConcreteDialect ConcreteDialect TFHEDialect FHEDialect diff --git a/compiler/tests/Dialect/BConcrete/ops.mlir b/compiler/tests/Dialect/BConcrete/ops.mlir new file mode 100644 index 000000000..e522007d1 --- /dev/null +++ b/compiler/tests/Dialect/BConcrete/ops.mlir @@ -0,0 +1,61 @@ +// RUN: concretecompiler --action=roundtrip %s 2>&1| FileCheck %s + +// CHECK-LABEL: func @add_lwe_ciphertexts(%arg0: tensor<2049xi64>, %arg1: tensor<2049xi64>) -> tensor<2049xi64> +func @add_lwe_ciphertexts(%arg0: tensor<2049xi64>, %arg1: tensor<2049xi64>) -> tensor<2049xi64> { + // CHECK-NEXT: %[[V0:.*]] = linalg.init_tensor [2049] : tensor<2049xi64> + // CHECK-NEXT: "BConcrete.add_lwe_buffer"(%[[V0]], %arg0, %arg1) : (tensor<2049xi64>, tensor<2049xi64>, tensor<2049xi64>) -> () + // CHECK-NEXT: return %[[V0]] : tensor<2049xi64> + %0 = linalg.init_tensor [2049] : tensor<2049xi64> + "BConcrete.add_lwe_buffer"(%0, %arg0, %arg1) : (tensor<2049xi64>, tensor<2049xi64>, tensor<2049xi64>) -> () + return %0 : tensor<2049xi64> +} + +// CHECK-LABEL: func @add_plaintext_lwe_ciphertext(%arg0: tensor<2049xi64>, %arg1: !Concrete.plaintext<5>) -> tensor<2049xi64> +func @add_plaintext_lwe_ciphertext(%arg0: tensor<2049xi64>, %arg1: !Concrete.plaintext<5>) -> tensor<2049xi64> { + // CHECK-NEXT: %[[V0:.*]] = linalg.init_tensor [2049] : tensor<2049xi64> + // CHECK-NEXT: "BConcrete.add_plaintext_lwe_buffer"(%[[V0]], %arg0, %arg1) : (tensor<2049xi64>, tensor<2049xi64>, !Concrete.plaintext<5>) -> () + // CHECK-NEXT: return %[[V0]] : tensor<2049xi64> + %0 = linalg.init_tensor [2049] : tensor<2049xi64> + "BConcrete.add_plaintext_lwe_buffer"(%0, %arg0, %arg1) : (tensor<2049xi64>, tensor<2049xi64>, !Concrete.plaintext<5>) -> () + return %0 : tensor<2049xi64> +} + +// CHECK-LABEL: func @mul_cleartext_lwe_ciphertext(%arg0: tensor<2049xi64>, %arg1: !Concrete.cleartext<7>) -> tensor<2049xi64> +func @mul_cleartext_lwe_ciphertext(%arg0: tensor<2049xi64>, %arg1: !Concrete.cleartext<7>) -> tensor<2049xi64> { + // CHECK-NEXT: %[[V0:.*]] = linalg.init_tensor [2049] : tensor<2049xi64> + // CHECK-NEXT: "BConcrete.mul_cleartext_lwe_buffer"(%[[V0]], %arg0, %arg1) : (tensor<2049xi64>, tensor<2049xi64>, !Concrete.cleartext<7>) -> () + // CHECK-NEXT: return %[[V0]] : tensor<2049xi64> + %0 = linalg.init_tensor [2049] : tensor<2049xi64> + "BConcrete.mul_cleartext_lwe_buffer"(%0, %arg0, %arg1) : (tensor<2049xi64>, tensor<2049xi64>, !Concrete.cleartext<7>) -> () + return %0 : tensor<2049xi64> +} + +// CHECK-LABEL: func @negate_lwe_ciphertext(%arg0: tensor<2049xi64>) -> tensor<2049xi64> +func @negate_lwe_ciphertext(%arg0: tensor<2049xi64>) -> tensor<2049xi64> { + // CHECK-NEXT: %[[V0:.*]] = linalg.init_tensor [2049] : tensor<2049xi64> + // CHECK-NEXT: "BConcrete.negate_lwe_buffer"(%[[V0]], %arg0) : (tensor<2049xi64>, tensor<2049xi64>) -> () + // CHECK-NEXT: return %[[V0]] : tensor<2049xi64> + %0 = linalg.init_tensor [2049] : tensor<2049xi64> + "BConcrete.negate_lwe_buffer"(%0, %arg0) : (tensor<2049xi64>, tensor<2049xi64>) -> () + return %0 : tensor<2049xi64> +} + +// CHECK-LABEL: func @bootstrap_lwe(%arg0: tensor<2049xi64>, %arg1: !Concrete.glwe_ciphertext) -> tensor<2049xi64> +func @bootstrap_lwe(%arg0: tensor<2049xi64>, %arg1: !Concrete.glwe_ciphertext) -> tensor<2049xi64> { + // CHECK-NEXT: %[[V0:.*]] = linalg.init_tensor [2049] : tensor<2049xi64> + // CHECK-NEXT: "BConcrete.bootstrap_lwe_buffer"(%[[V0]], %arg0, %arg1) {baseLog = -1 : i32, glweDimension = 1 : i32, level = -1 : i32, polynomialSize = 1024 : i32} : (tensor<2049xi64>, tensor<2049xi64>, !Concrete.glwe_ciphertext) -> () + // CHECK-NEXT: return %[[V0]] : tensor<2049xi64> + %0 = linalg.init_tensor [2049] : tensor<2049xi64> + "BConcrete.bootstrap_lwe_buffer"(%0, %arg0, %arg1) {baseLog = -1 : i32, glweDimension = 1 : i32, level = -1 : i32, polynomialSize = 1024 : i32} : (tensor<2049xi64>, tensor<2049xi64>, !Concrete.glwe_ciphertext) -> () + return %0 : tensor<2049xi64> +} + +// CHECK-LABEL: func @keyswitch_lwe(%arg0: tensor<2049xi64>) -> tensor<2049xi64> +func @keyswitch_lwe(%arg0: tensor<2049xi64>) -> tensor<2049xi64> { + // CHECK-NEXT: %[[V0:.*]] = linalg.init_tensor [2049] : tensor<2049xi64> + // CHECK-NEXT: "BConcrete.keyswitch_lwe_buffer"(%[[V0]], %arg0) {baseLog = 2 : i32, inputLweDimension = 1 : i32, level = 3 : i32, outputLweDimension = 1 : i32} : (tensor<2049xi64>, tensor<2049xi64>) -> () + // CHECK-NEXT: return %[[V0]] : tensor<2049xi64> + %0 = linalg.init_tensor [2049] : tensor<2049xi64> + "BConcrete.keyswitch_lwe_buffer"(%0, %arg0) {baseLog = 2 : i32, inputLweDimension = 1 : i32, level = 3 : i32, outputLweDimension = 1 : i32} : (tensor<2049xi64>, tensor<2049xi64>) -> () + return %0 : tensor<2049xi64> +}