f 1d44d89

zama-ai · Aug 28, 2023 · 86233ae · 86233ae
1 parent e168369
commit 86233ae
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 18 deletions.
diff --git a/compilers/concrete-compiler/compiler/lib/Dialect/Concrete/Analysis/MemoryUsage.cpp b/compilers/concrete-compiler/compiler/lib/Dialect/Concrete/Analysis/MemoryUsage.cpp
@@ -199,7 +199,8 @@ static std::optional<StringError> on_enter(mlir::Operation *op,
 
     // the search would be faster if we use an unsorted_set, but we need a hash
     // function for mlir::Value
-    if (std::count(visited.begin(), visited.end(), lastVisitedBuffer) == 0) {
+    if (std::find(visited.begin(), visited.end(), lastVisitedBuffer) !=
+        visited.end()) {
       visited.push_back(lastVisitedBuffer);
 
       auto maybeBufferSize =

diff --git a/compilers/concrete-compiler/compiler/lib/Support/CompilationFeedback.cpp b/compilers/concrete-compiler/compiler/lib/Support/CompilationFeedback.cpp
@@ -93,14 +93,11 @@ llvm::json::Value toJSON(const mlir::concretelang::CompilationFeedback &v) {
       {"crtDecompositionsOfOutputs", v.crtDecompositionsOfOutputs},
   };
 
-  auto memoryUsageJsonArray = llvm::json::Array();
+  auto memoryUsageObject = llvm::json::Object();
   for (auto key : v.memoryUsagePerLoc) {
-    auto memoryUsageJson = llvm::json::Object();
-    memoryUsageJson.insert({"location", key.first});
-    memoryUsageJson.insert({"usage", key.second});
-    memoryUsageJsonArray.push_back(std::move(memoryUsageJson));
+    memoryUsageObject.insert({key.first, key.second});
   }
-  object.insert({"memoryUsagePerLoc", std::move(memoryUsageJsonArray)});
+  object.insert({"memoryUsagePerLoc", std::move(memoryUsageObject)});
 
   auto statisticsJson = llvm::json::Array();
   for (auto statistic : v.statistics) {
@@ -186,21 +183,17 @@ bool fromJSON(const llvm::json::Value j,
     return false;
   }
 
-  auto memoryUsageArray = object->getArray("memoryUsagePerLoc");
-  if (!memoryUsageArray) {
+  auto memoryUsageObject = object->getObject("memoryUsagePerLoc");
+  if (!memoryUsageObject) {
     return false;
   }
-  for (auto memoryUsageValue : *memoryUsageArray) {
-    auto memoryUsage = memoryUsageValue.getAsObject();
-    if (!memoryUsage) {
+  for (auto entry : *memoryUsageObject) {
+    auto loc = entry.getFirst().str();
+    auto maybeUsage = entry.getSecond().getAsInteger();
+    if (!maybeUsage.has_value()) {
       return false;
     }
-    auto loc = memoryUsage->getString("location");
-    auto usage = memoryUsage->getInteger("usage");
-    if (!loc || !usage) {
-      return false;
-    }
-    v.memoryUsagePerLoc[loc->str()] = *usage;
+    v.memoryUsagePerLoc[loc] = *maybeUsage;
   }
 
   auto statistics = object->getArray("statistics");