refactor(frontend-python): reduce memory usage for table construction…

… of non-multi table lookups

frontends/concrete-python/concrete/fhe/mlir/converter.py

@@ -541,6 +541,7 @@ def tlu(self, ctx: Context, node: Node, preds: List[Conversion]) -> Conversion:
 
         if len(tables) == 1:
             table = tables[0][0]
+            assert tables[0][1] is None
 
             # The reduction on 63b is to avoid problems like doing a TLU of
             # the form T[j] = 2<<j, for j which is supposed to be 7b as per
@@ -563,6 +564,8 @@ def tlu(self, ctx: Context, node: Node, preds: List[Conversion]) -> Conversion:
             for i, (table, indices) in enumerate(tables):
                 assert len(table) == individual_table_size
                 lut_values[i, :] = table
+
+                assert indices is not None
                 for index in indices:
                     map_values[index] = i
 

frontends/concrete-python/concrete/fhe/mlir/utils.py

@@ -132,11 +132,14 @@ def construct_table(node: Node, preds: List[Node]) -> List[Any]:
     np.seterr(divide="ignore")
 
     inputs: List[Any] = [pred() if pred.operation == Operation.Constant else None for pred in preds]
-    table: List[Optional[Union[np.bool_, np.integer, np.floating, np.ndarray]]] = []
+    table: List[Optional[Union[int, np.bool_, np.integer, np.floating, np.ndarray]]] = []
     for value in values:
         try:
             inputs[variable_input_index] = np.ones(variable_input_shape, dtype=np.int64) * value
-            table.append(node(*inputs))
+            evaluation = node(*inputs)
+            table.append(
+                evaluation if evaluation.min() != evaluation.max() else int(evaluation.min())
+            )
         except Exception:  # pylint: disable=broad-except
             # here we try our best to fill the table
             # if it fails, we append None and let flooding algoritm replace None values below
@@ -152,7 +155,7 @@ def construct_table(node: Node, preds: List[Node]) -> List[Any]:
 def construct_deduplicated_tables(
     node: Node,
     preds: List[Node],
-) -> Tuple[Tuple[np.ndarray, List[Tuple[int, ...]]], ...]:
+) -> Tuple[Tuple[np.ndarray, Optional[List[Tuple[int, ...]]]], ...]:
     """
     Construct lookup tables for each cell of the input for an Operation.Generic node.
 
@@ -187,8 +190,22 @@ def construct_deduplicated_tables(
                 [ [5, 8, 6, 7][input[2, 0]] , [3, 1, 2, 4][input[2, 1]] ]
     """
 
+    raw_table = construct_table(node, preds)
+    if all(isinstance(value, int) for value in raw_table):
+        return ((np.array(raw_table), None),)
+
     node_complete_table = np.concatenate(
-        tuple(np.expand_dims(array, -1) for array in construct_table(node, preds)),
+        tuple(
+            np.expand_dims(
+                (
+                    array
+                    if isinstance(array, np.ndarray)
+                    else np.broadcast_to(array, node.output.shape)
+                ),
+                -1,
+            )
+            for array in raw_table
+        ),
         axis=-1,
     )