feat(frontend): compile using in-memory module

zama-ai · Jul 20, 2023 · 636c338 · 636c338
1 parent 6fcbb6d
commit 636c338
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 14 deletions.
diff --git a/frontends/concrete-python/concrete/fhe/compilation/circuit.py b/frontends/concrete-python/concrete/fhe/compilation/circuit.py
@@ -8,6 +8,7 @@
 
 import numpy as np
 from concrete.compiler import SimulatedValueDecrypter, SimulatedValueExporter
+from mlir.ir import Module as MlirModule
 
 from ..internal.utils import assert_that
 from ..representation import Graph
@@ -29,13 +30,13 @@ class Circuit:
     configuration: Configuration
 
     graph: Graph
-    mlir: str
+    mlir: MlirModule
 
     client: Client
     server: Server
     simulator: Server
 
-    def __init__(self, graph: Graph, mlir: str, configuration: Optional[Configuration] = None):
+    def __init__(self, graph: Graph, mlir: MlirModule, configuration: Optional[Configuration] = None):
         self.configuration = configuration if configuration is not None else Configuration()
 
         self.graph = graph

diff --git a/frontends/concrete-python/concrete/fhe/compilation/compiler.py b/frontends/concrete-python/concrete/fhe/compilation/compiler.py
@@ -446,9 +446,12 @@ def compile(
                 message = "Function you are trying to compile cannot be compiled\n\n" + fmtd_graph
                 raise RuntimeError(message)
 
-            mlir = GraphConverter().convert(self.graph, self.configuration)
+            # in-memory MLIR module
+            mlir_module = GraphConverter().convert(self.graph, self.configuration)
+            # textual representation of the MLIR module
+            mlir_str = str(mlir_module).strip()
             if self.artifacts is not None:
-                self.artifacts.add_mlir_to_compile(mlir)
+                self.artifacts.add_mlir_to_compile(mlir_str)
 
             show_graph = (
                 self.configuration.show_graph
@@ -475,7 +478,7 @@ def compile(
                 )
 
                 longest_graph_line = max(len(line) for line in graph.split("\n"))
-                longest_mlir_line = max(len(line) for line in mlir.split("\n"))
+                longest_mlir_line = max(len(line) for line in mlir_str.split("\n"))
                 longest_line = max(longest_graph_line, longest_mlir_line)
 
                 try:  # pragma: no cover
@@ -506,7 +509,7 @@ def compile(
 
                     print("MLIR")
                     print("-" * columns)
-                    print(mlir)
+                    print(mlir_str)
                     print("-" * columns)
 
                     print()
@@ -517,7 +520,7 @@ def compile(
                     print("Optimizer")
                     print("-" * columns)
 
-            circuit = Circuit(self.graph, mlir, self.configuration)
+            circuit = Circuit(self.graph, mlir_module, self.configuration)
 
             if hasattr(circuit, "client"):
                 client_parameters = circuit.client.specs.client_parameters

diff --git a/frontends/concrete-python/concrete/fhe/compilation/server.py b/frontends/concrete-python/concrete/fhe/compilation/server.py
@@ -25,6 +25,7 @@
     PublicArguments,
 )
 from mlir._mlir_libs._concretelang._compiler import OptimizerStrategy
+from mlir.ir import Module as MlirModule
 
 from ..internal.utils import assert_that
 from .configuration import (
@@ -81,12 +82,12 @@ def __init__(
         )
 
     @staticmethod
-    def create(mlir: str, configuration: Configuration, is_simulated: bool = False) -> "Server":
+    def create(mlir: MlirModule, configuration: Configuration, is_simulated: bool = False) -> "Server":
         """
         Create a server using MLIR and output sign information.
 
         Args:
-            mlir (str):
+            mlir (MlirModule):
                 mlir to compile
 
             configuration (Configuration):
@@ -182,7 +183,7 @@ def create(mlir: str, configuration: Configuration, is_simulated: bool = False)
         )
 
         # pylint: disable=protected-access
-        result._mlir = mlir
+        result._mlir = str(mlir).strip()
         result._configuration = configuration
         # pylint: enable=protected-access
 

diff --git a/frontends/concrete-python/concrete/fhe/mlir/converter.py b/frontends/concrete-python/concrete/fhe/mlir/converter.py
@@ -34,7 +34,7 @@ class Converter:
     Converter class, to convert a computation graph to MLIR.
     """
 
-    def convert(self, graph: Graph, configuration: Configuration) -> str:
+    def convert(self, graph: Graph, configuration: Configuration) -> MlirModule:
         """
         Convert a computation graph to MLIR.
 
@@ -46,8 +46,8 @@ def convert(self, graph: Graph, configuration: Configuration) -> str:
                 configuration to use
 
         Return:
-            str:
-                MLIR corresponding to graph
+            MlirModule:
+                In-memory MLIR module corresponding to the graph
         """
 
         graph = self.process(graph, configuration)
@@ -88,7 +88,7 @@ def main(*args):
 
                     return tuple(outputs)
 
-        return str(module).strip()
+        return module
 
     @staticmethod
     def stdout_with_ansi_support() -> bool: