refactor(frontend-python): re-write bit width assignment

zama-ai · Aug 3, 2023 · 994676d · 994676d
1 parent cf4c0f6
commit 994676d
Show file tree

Hide file tree

Showing 7 changed files with 566 additions and 241 deletions.
diff --git a/frontends/concrete-python/.pylintrc b/frontends/concrete-python/.pylintrc
@@ -441,6 +441,7 @@ disable=raw-checker-failed,
         too-many-instance-attributes,
         too-many-lines,
         too-many-locals,
+        too-many-nested-blocks,
         too-many-public-methods,
         too-many-statements,
         unnecessary-lambda-assignment,

diff --git a/frontends/concrete-python/.ruff.toml b/frontends/concrete-python/.ruff.toml
@@ -16,6 +16,7 @@ ignore = [
 "**/__init__.py" = ["F401"]
 "concrete/fhe/compilation/configuration.py" = ["ARG002"]
 "concrete/fhe/mlir/processors/all.py" = ["F401"]
+"concrete/fhe/mlir/processors/assign_bit_widths.py" = ["ARG002"]
 "concrete/fhe/mlir/converter.py" = ["ARG002", "B011", "F403", "F405"]
 "examples/**" = ["PLR2004"]
 "tests/**" = ["PLR2004", "PLW0603", "SIM300", "S311"]
diff --git a/frontends/concrete-python/concrete/fhe/mlir/context.py b/frontends/concrete-python/concrete/fhe/mlir/context.py
@@ -802,7 +802,10 @@ def dot(self, resulting_type: ConversionType, x: Conversion, y: Conversion) -> C
 
             self.error(highlights)
 
-        assert self.is_bit_width_compatible(resulting_type, x, y)
+        if x.is_encrypted and y.is_encrypted:
+            assert self.is_bit_width_compatible(x, y)
+        else:
+            assert self.is_bit_width_compatible(resulting_type, x, y)
 
         if x.is_scalar or y.is_scalar:
             return self.mul(resulting_type, x, y)
@@ -1369,7 +1372,10 @@ def matmul(self, resulting_type: ConversionType, x: Conversion, y: Conversion) -
 
             self.error(highlights)
 
-        assert self.is_bit_width_compatible(resulting_type, x, y)
+        if x.is_encrypted and y.is_encrypted:
+            assert self.is_bit_width_compatible(x, y)
+        else:
+            assert self.is_bit_width_compatible(resulting_type, x, y)
 
         if resulting_type.shape == ():
             if x.is_clear:
@@ -1489,7 +1495,10 @@ def mul(self, resulting_type: ConversionType, x: Conversion, y: Conversion) -> C
 
             self.error(highlights)
 
-        assert self.is_bit_width_compatible(resulting_type, x, y)
+        if x.is_encrypted and y.is_encrypted:
+            assert self.is_bit_width_compatible(x, y)
+        else:
+            assert self.is_bit_width_compatible(resulting_type, x, y)
 
         use_linalg = x.is_tensor or y.is_tensor