From 8198a79deba3b072ca8b289e8d8fb02da10d560f Mon Sep 17 00:00:00 2001
From: Umut <umutsahin@protonmail.com>
Date: Mon, 24 Jan 2022 12:45:59 +0300
Subject: [PATCH] feat: introduce maximum bit width constant, use it where
 appropriate, export it in python bindings

---
 compiler/include/concretelang/Support/Constants.h   | 1 +
 compiler/lib/Bindings/Python/ConcretelangModule.cpp | 6 +++++-
 compiler/lib/Dialect/FHE/IR/FHEDialect.cpp          | 7 +++++--
 compiler/lib/Dialect/TFHE/IR/TFHEDialect.cpp        | 9 ++++++---
 compiler/tests/python/test_fhe_dialect.py           | 6 +++++-
 5 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/compiler/include/concretelang/Support/Constants.h b/compiler/include/concretelang/Support/Constants.h
index 24eb4d624..e1f24c2f9 100644
--- a/compiler/include/concretelang/Support/Constants.h
+++ b/compiler/include/concretelang/Support/Constants.h
@@ -9,6 +9,7 @@
 namespace mlir {
 namespace concretelang {
 constexpr unsigned DEFAULT_PATTERN_BENEFIT = 1;
+constexpr unsigned MAXIMUM_BIT_WIDTH = 7;
 } // namespace concretelang
 } // namespace mlir
 
diff --git a/compiler/lib/Bindings/Python/ConcretelangModule.cpp b/compiler/lib/Bindings/Python/ConcretelangModule.cpp
index aea806056..3a194d1f2 100644
--- a/compiler/lib/Bindings/Python/ConcretelangModule.cpp
+++ b/compiler/lib/Bindings/Python/ConcretelangModule.cpp
@@ -8,6 +8,7 @@
 
 #include "concretelang-c/Dialect/FHE.h"
 #include "concretelang-c/Dialect/FHELinalg.h"
+#include "concretelang/Support/Constants.h"
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/Registration.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
@@ -23,6 +24,9 @@ PYBIND11_MODULE(_concretelang, m) {
   llvm::sys::PrintStackTraceOnErrorSignal(/*argv=*/"");
   LLVMEnablePrettyStackTrace();
 
+  m.attr("MAXIMUM_BIT_WIDTH") =
+      pybind11::int_(mlir::concretelang::MAXIMUM_BIT_WIDTH);
+
   m.def(
       "register_dialects",
       [](py::object capsule) {
@@ -45,4 +49,4 @@ PYBIND11_MODULE(_concretelang, m) {
 
   py::module api = m.def_submodule("_compiler", "Compiler API");
   mlir::concretelang::python::populateCompilerAPISubmodule(api);
-}
\ No newline at end of file
+}
diff --git a/compiler/lib/Dialect/FHE/IR/FHEDialect.cpp b/compiler/lib/Dialect/FHE/IR/FHEDialect.cpp
index dc5681952..ca8efa73f 100644
--- a/compiler/lib/Dialect/FHE/IR/FHEDialect.cpp
+++ b/compiler/lib/Dialect/FHE/IR/FHEDialect.cpp
@@ -12,6 +12,8 @@
 
 #include "concretelang/Dialect/FHE/IR/FHEOpsDialect.cpp.inc"
 
+#include "concretelang/Support/Constants.h"
+
 using namespace mlir::concretelang::FHE;
 
 void FHEDialect::initialize() {
@@ -51,8 +53,9 @@ void FHEDialect::printType(::mlir::Type type,
 
 mlir::LogicalResult EncryptedIntegerType::verify(
     llvm::function_ref<::mlir::InFlightDiagnostic()> emitError, unsigned p) {
-  if (p == 0 || p > 7) {
-    emitError() << "FHE.eint support only precision in ]0;7]";
+  if (p == 0 || p > mlir::concretelang::MAXIMUM_BIT_WIDTH) {
+    emitError() << "FHE.eint support only precision in ]0;"
+                << mlir::concretelang::MAXIMUM_BIT_WIDTH << "]";
     return mlir::failure();
   }
   return mlir::success();
diff --git a/compiler/lib/Dialect/TFHE/IR/TFHEDialect.cpp b/compiler/lib/Dialect/TFHE/IR/TFHEDialect.cpp
index 875baa909..53a7ad74e 100644
--- a/compiler/lib/Dialect/TFHE/IR/TFHEDialect.cpp
+++ b/compiler/lib/Dialect/TFHE/IR/TFHEDialect.cpp
@@ -12,6 +12,8 @@
 
 #include "concretelang/Dialect/TFHE/IR/TFHEOpsDialect.cpp.inc"
 
+#include "concretelang/Support/Constants.h"
+
 using namespace mlir::concretelang::TFHE;
 
 void TFHEDialect::initialize() {
@@ -47,7 +49,7 @@ void TFHEDialect::printType(::mlir::Type type,
 
 /// Verify that GLWE parameter are consistant
 /// - The bits parameter is 64 (we support only this for v0)
-/// - The p parameter is ]0;7]
+/// - The p parameter is ]0;MAXIMUM_BIT_WIDTH]
 ::mlir::LogicalResult GLWECipherTextType::verify(
     ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError,
     signed dimension, signed polynomialSize, signed bits, signed p) {
@@ -55,8 +57,9 @@ void TFHEDialect::printType(::mlir::Type type,
     emitError() << "GLWE bits parameter can only be 64";
     return ::mlir::failure();
   }
-  if (p == 0 || p > 7) {
-    emitError() << "GLWE p parameter can only be in ]0;7]";
+  if (p == 0 || p > (signed)mlir::concretelang::MAXIMUM_BIT_WIDTH) {
+    emitError() << "GLWE p parameter can only be in ]0;"
+                << mlir::concretelang::MAXIMUM_BIT_WIDTH << "]";
     return mlir::failure();
   }
   return ::mlir::success();
diff --git a/compiler/tests/python/test_fhe_dialect.py b/compiler/tests/python/test_fhe_dialect.py
index ab7c21d48..2f2c6e829 100644
--- a/compiler/tests/python/test_fhe_dialect.py
+++ b/compiler/tests/python/test_fhe_dialect.py
@@ -1,9 +1,13 @@
 import pytest
 from mlir.ir import Context, RankedTensorType, Location
-from concrete.lang import register_dialects
+from concrete.lang import register_dialects, MAXIMUM_BIT_WIDTH
 from concrete.lang.dialects import fhe
 
 
+def test_constants():
+    assert MAXIMUM_BIT_WIDTH == 7
+
+
 @pytest.mark.parametrize("width", list(range(1, 8)))
 def test_eint(width):
     ctx = Context()