feat(compiler/python): expose memory usage in bindings

zama-ai · Aug 8, 2023 · 29016bc · 29016bc
1 parent 5c085b2
commit 29016bc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp b/compilers/concrete-compiler/compiler/lib/Bindings/Python/CompilerAPIModule.cpp
@@ -162,7 +162,10 @@ void mlir::concretelang::python::populateCompilerAPISubmodule(
           "crt_decompositions_of_outputs",
           &mlir::concretelang::CompilationFeedback::crtDecompositionsOfOutputs)
       .def_readonly("statistics",
-                    &mlir::concretelang::CompilationFeedback::statistics);
+                    &mlir::concretelang::CompilationFeedback::statistics)
+      .def_readonly(
+          "memory_usage_per_location",
+          &mlir::concretelang::CompilationFeedback::memoryUsagePerLoc);
 
   pybind11::class_<mlir::concretelang::JitCompilationResult>(
       m, "JITCompilationResult");

diff --git a/.../concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/compilation_feedback.py b/.../concrete-compiler/compiler/lib/Bindings/Python/concrete/compiler/compilation_feedback.py
@@ -48,6 +48,7 @@ def __init__(self, compilation_feedback: _CompilationFeedback):
             compilation_feedback.crt_decompositions_of_outputs
         )
         self.statistics = compilation_feedback.statistics
+        self.memory_usage_per_location = compilation_feedback.memory_usage_per_location
 
         super().__init__(compilation_feedback)