From 671679e2d5633d2223a699d865b803a441ea5f00 Mon Sep 17 00:00:00 2001
From: Thomas Alexander <thomasalexander2718@gmail.com>
Date: Mon, 29 Jul 2024 22:40:06 -0300
Subject: [PATCH] Reduce memory usage in pulse sequencing  (#336)

Builds on #335.
---
 include/Utils/SymbolCacheAnalysis.h                       | 8 +++++---
 lib/Dialect/QUIR/Utils/Utils.cpp                          | 3 ++-
 .../notes/symbol-cache-use-densemap-c178e05e59c41bbe.yaml | 4 ++++
 3 files changed, 11 insertions(+), 4 deletions(-)
 create mode 100644 releasenotes/notes/symbol-cache-use-densemap-c178e05e59c41bbe.yaml

diff --git a/include/Utils/SymbolCacheAnalysis.h b/include/Utils/SymbolCacheAnalysis.h
index aa56b19a3..c913508e1 100644
--- a/include/Utils/SymbolCacheAnalysis.h
+++ b/include/Utils/SymbolCacheAnalysis.h
@@ -49,11 +49,11 @@ namespace qssc::utils {
 //                .addToCache<SequenceOp>();
 //
 // This analysis is intended to be used with MLIR's getAnalysis
-// framework. It has been designed to reused the chached value
+// framework. It has been designed to reuse the cached value
 // and will not be invalidated automatically with each pass.
 // If a pass manipulates the symbols that are cached with this
 // analysis then it should use the addCallee method to update the
-// map or call invalidate after appying updates.
+// map or call invalidate after applying updates.
 // Note this analysis should always be used by reference or
 // via a pointer to ensure that updates are applied to the maps
 // stored by the MLIR analysis framework.
@@ -95,6 +95,8 @@ class SymbolCacheAnalysis {
 
     op->walk([&](CalleeOp op) {
       symbolOpsMap[op.getSymName()] = op.getOperation();
+      // Don't recurse symbols
+      return mlir::WalkResult::skip();
     });
     cachedTypes.insert(typeName);
     invalid = false;
@@ -193,7 +195,7 @@ class SymbolCacheAnalysis {
 
 private:
   llvm::StringMap<mlir::Operation *> symbolOpsMap;
-  std::unordered_map<mlir::Operation *, mlir::Operation *> callMap;
+  llvm::DenseMap<mlir::Operation *, mlir::Operation *> callMap;
   std::unordered_set<std::string> cachedTypes;
   mlir::Operation *topOp{nullptr};
   bool invalid{true};
diff --git a/lib/Dialect/QUIR/Utils/Utils.cpp b/lib/Dialect/QUIR/Utils/Utils.cpp
index eb6068929..b958389db 100644
--- a/lib/Dialect/QUIR/Utils/Utils.cpp
+++ b/lib/Dialect/QUIR/Utils/Utils.cpp
@@ -97,7 +97,8 @@ auto getMainFunction(Operation *moduleOperation) -> Operation * {
       mainFunc = funcOp.getOperation();
       return WalkResult::interrupt();
     }
-    return WalkResult::advance();
+    // Don't process nested values
+    return WalkResult::skip();
   });
   if (!mainFunc)
     llvm::errs() << "Error: Main function not found!\n";
diff --git a/releasenotes/notes/symbol-cache-use-densemap-c178e05e59c41bbe.yaml b/releasenotes/notes/symbol-cache-use-densemap-c178e05e59c41bbe.yaml
new file mode 100644
index 000000000..5848491bf
--- /dev/null
+++ b/releasenotes/notes/symbol-cache-use-densemap-c178e05e59c41bbe.yaml
@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    Symbol cache now uses llvm::DenseMap for performance reasons.