Improve SubroutineCloningPass and MergeCircuitsPass performance (#192)

bcdonovan · web-flow · commit 82883b1fefeb · 2023-11-21T10:58:18.000-05:00
Improves the performance of the SubroutineCloningPassm and
MergeCircuitsPass by replacing the usage of 
SymbolTable::lookupSymbolIn with a llvm::StringMap storing 
the symbols of interest.
diff --git a/cmake/apple-clang.cmake b/cmake/apple-clang.cmake
@@ -39,6 +39,6 @@ if(ENABLE_ADDRESS_SANITIZER OR ENABLE_UNDEFINED_SANITIZER OR ENABLE_THREAD_SANIT
 endif()
 
 set (CMAKE_CXX_FLAGS_DEBUG "-g3 -O0")
-set (CMAKE_CXX_FLAGS_RELEASE "-g -O2 -DNDEBUG")
+set (CMAKE_CXX_FLAGS_RELEASE "-g -O2 -DNOVERIFY")
 
 set (CMAKE_INSTALL_RPATH_USE_LINK_PATH ON)
diff --git a/cmake/linux-gcc.cmake b/cmake/linux-gcc.cmake
@@ -40,7 +40,7 @@ if(ENABLE_THREAD_SANITIZER)
 endif()
 
 set (CMAKE_CXX_FLAGS_DEBUG "-g3 -O0")
-set (CMAKE_CXX_FLAGS_RELEASE "-g -O2 -DNDEBUG")
+set (CMAKE_CXX_FLAGS_RELEASE "-g -O2 -DNOVERIFY")
 
 # set (CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -fuse-ld=gold")
 set (CMAKE_CXX_STANDARD_LIBRARIES "-lstdc++fs -lpthread")
diff --git a/cmake/llvm-clang.cmake b/cmake/llvm-clang.cmake
@@ -42,7 +42,7 @@ if(ENABLE_ADDRESS_SANITIZER OR ENABLE_UNDEFINED_SANITIZER OR ENABLE_THREAD_SANIT
 endif()
 
 set (CMAKE_CXX_FLAGS_DEBUG "-g3 -O0")
-set (CMAKE_CXX_FLAGS_RELEASE "-g -O2 -DNDEBUG")
+set (CMAKE_CXX_FLAGS_RELEASE "-g -O2 -DNOVERIFY")
 
 set (CMAKE_EXE_LINKER_FLAGS "${CMAKE_EXE_LINKER_FLAGS} -fuse-ld=gold")
 set (CMAKE_CXX_STANDARD_LIBRARIES "-lstdc++fs -lpthread")
diff --git a/include/Dialect/QUIR/Transforms/MergeCircuits.h b/include/Dialect/QUIR/Transforms/MergeCircuits.h
@@ -34,10 +34,12 @@ struct MergeCircuitsPass
     : public PassWrapper<MergeCircuitsPass, OperationPass<>> {
   void runOnOperation() override;
 
-  static CircuitOp getCircuitOp(CallCircuitOp callCircuitOp);
-  static LogicalResult mergeCallCircuits(PatternRewriter &rewriter,
-                                         CallCircuitOp callCircuitOp,
-                                         CallCircuitOp nextCallCircuitOp);
+  static CircuitOp getCircuitOp(CallCircuitOp callCircuitOp,
+                                llvm::StringMap<Operation *> *symbolMap);
+  static LogicalResult
+  mergeCallCircuits(PatternRewriter &rewriter, CallCircuitOp callCircuitOp,
+                    CallCircuitOp nextCallCircuitOp,
+                    llvm::StringMap<Operation *> *symbolMap);
 
   llvm::StringRef getArgument() const override;
   llvm::StringRef getDescription() const override;
diff --git a/include/Dialect/QUIR/Transforms/SubroutineCloning.h b/include/Dialect/QUIR/Transforms/SubroutineCloning.h
@@ -35,14 +35,17 @@ class Operation;
 } // namespace mlir
 
 namespace mlir::quir {
+
+using SymbolOpMap = llvm::StringMap<Operation *>;
+
 struct SubroutineCloningPass
     : public PassWrapper<SubroutineCloningPass, OperationPass<>> {
   auto lookupQubitId(const Value val) -> int;
 
   template <class CallLikeOp>
   auto getMangledName(Operation *op) -> std::string;
   template <class CallLikeOp, class FuncLikeOp>
-  void processCallOp(Operation *op);
+  void processCallOp(Operation *op, SymbolOpMap &symbolOpMap);
 
   void runOnOperation() override;
 
diff --git a/lib/API/api.cpp b/lib/API/api.cpp
@@ -85,7 +85,7 @@ static llvm::cl::opt<bool> verifyDiagnostics(
                    "expected-* lines on the corresponding line"),
     llvm::cl::init(false), llvm::cl::cat(qssc::config::getQSSCCategory()));
 
-#ifndef NDEBUG
+#ifndef NOVERIFY
 #define VERIFY_PASSES_DEFAULT true
 #else
 #define VERIFY_PASSES_DEFAULT false
diff --git a/lib/Dialect/QUIR/Transforms/MergeCircuits.cpp b/lib/Dialect/QUIR/Transforms/MergeCircuits.cpp
@@ -52,8 +52,13 @@ bool moveUsers(Operation *curOp, MoveListVec &moveList) {
 
 // This pattern matches on two CallCircuitOps separated by non-quantum ops
 struct CircuitAndCircuitPattern : public OpRewritePattern<CallCircuitOp> {
-  explicit CircuitAndCircuitPattern(MLIRContext *ctx)
-      : OpRewritePattern<CallCircuitOp>(ctx) {}
+  explicit CircuitAndCircuitPattern(MLIRContext *ctx,
+                                    llvm::StringMap<Operation *> &symbolMap)
+      : OpRewritePattern<CallCircuitOp>(ctx) {
+    _symbolMap = &symbolMap;
+  }
+
+  llvm::StringMap<Operation *> *_symbolMap;
 
   LogicalResult matchAndRewrite(CallCircuitOp callCircuitOp,
                                 PatternRewriter &rewriter) const override {
@@ -141,7 +146,7 @@ struct CircuitAndCircuitPattern : public OpRewritePattern<CallCircuitOp> {
       return failure();
 
     return MergeCircuitsPass::mergeCallCircuits(rewriter, callCircuitOp,
-                                                nextCallCircuitOp);
+                                                nextCallCircuitOp, _symbolMap);
   } // matchAndRewrite
 };  // struct CircuitAndCircuitPattern
 
@@ -220,22 +225,25 @@ struct CircuitAndBarrierPattern : public OpRewritePattern<CallCircuitOp> {
 
 } // end anonymous namespace
 
-CircuitOp MergeCircuitsPass::getCircuitOp(CallCircuitOp callCircuitOp) {
-  auto circuitAttr = callCircuitOp->getAttrOfType<FlatSymbolRefAttr>("callee");
-  assert(circuitAttr && "Requires a 'callee' symbol reference attribute");
+CircuitOp
+MergeCircuitsPass::getCircuitOp(CallCircuitOp callCircuitOp,
+                                llvm::StringMap<Operation *> *symbolMap) {
+  // look for func def match
+  assert(symbolMap && "a valid symbolMap pointer must be provided");
+  auto search = symbolMap->find(callCircuitOp.callee());
+
+  assert(search != symbolMap->end() && "matching circuit not found");
 
-  auto circuitOp = SymbolTable::lookupNearestSymbolFrom<CircuitOp>(
-      callCircuitOp, circuitAttr);
+  auto circuitOp = dyn_cast<CircuitOp>(search->second);
   assert(circuitOp && "matching circuit not found");
   return circuitOp;
 }
 
-LogicalResult
-MergeCircuitsPass::mergeCallCircuits(PatternRewriter &rewriter,
-                                     CallCircuitOp callCircuitOp,
-                                     CallCircuitOp nextCallCircuitOp) {
-  auto circuitOp = getCircuitOp(callCircuitOp);
-  auto nextCircuitOp = getCircuitOp(nextCallCircuitOp);
+LogicalResult MergeCircuitsPass::mergeCallCircuits(
+    PatternRewriter &rewriter, CallCircuitOp callCircuitOp,
+    CallCircuitOp nextCallCircuitOp, llvm::StringMap<Operation *> *symbolMap) {
+  auto circuitOp = getCircuitOp(callCircuitOp, symbolMap);
+  auto nextCircuitOp = getCircuitOp(nextCallCircuitOp, symbolMap);
 
   rewriter.setInsertionPointAfter(nextCircuitOp);
 
@@ -365,14 +373,24 @@ MergeCircuitsPass::mergeCallCircuits(PatternRewriter &rewriter,
   rewriter.replaceOp(nextCallCircuitOp,
                      ResultRange(iterSep, newCallOp.result_end()));
 
+  // add new name to symbolMap
+  // do not remove old in case the are multiple calls
+  (*symbolMap)[newName] = newCircuitOp.getOperation();
+
   return success();
 }
 
 void MergeCircuitsPass::runOnOperation() {
   Operation *moduleOperation = getOperation();
 
+  llvm::StringMap<Operation *> circuitOpsMap;
+
+  moduleOperation->walk([&](CircuitOp circuitOp) {
+    circuitOpsMap[circuitOp.sym_name()] = circuitOp.getOperation();
+  });
+
   RewritePatternSet patterns(&getContext());
-  patterns.add<CircuitAndCircuitPattern>(&getContext());
+  patterns.add<CircuitAndCircuitPattern>(&getContext(), circuitOpsMap);
   patterns.add<BarrierAndCircuitPattern>(&getContext());
   patterns.add<CircuitAndBarrierPattern>(&getContext());
 
diff --git a/lib/Dialect/QUIR/Transforms/SubroutineCloning.cpp b/lib/Dialect/QUIR/Transforms/SubroutineCloning.cpp
@@ -82,13 +82,21 @@ auto SubroutineCloningPass::getMangledName(Operation *op) -> std::string {
 } // getMangledName
 
 template <class CallLikeOp, class FuncLikeOp>
-void SubroutineCloningPass::processCallOp(Operation *op) {
+void SubroutineCloningPass::processCallOp(Operation *op,
+                                          SymbolOpMap &symbolOps) {
   auto callOp = dyn_cast<CallLikeOp>(op);
   OpBuilder build(moduleOperation->getRegion(0));
 
   // look for func def match
-  Operation *findOp =
-      SymbolTable::lookupSymbolIn(moduleOperation, callOp.callee());
+  auto search = symbolOps.find(callOp.callee());
+
+  if (search == symbolOps.end()) {
+    callOp->emitOpError() << "No matching function def found for "
+                          << callOp.callee() << "\n";
+    return signalPassFailure();
+  }
+
+  Operation *findOp = search->second;
   if (findOp) {
     std::vector<Value> qOperands;
     qubitCallOperands(callOp, qOperands);
@@ -99,9 +107,7 @@ void SubroutineCloningPass::processCallOp(Operation *op) {
                     FlatSymbolRefAttr::get(&getContext(), mangledName));
 
     // does the mangled function already exist?
-    Operation *mangledOp =
-        SymbolTable::lookupSymbolIn(moduleOperation, mangledName);
-    if (mangledOp) // nothing to do
+    if (symbolOps.find(mangledName) != symbolOps.end())
       return;
 
     // clone the func def with the new name
@@ -127,6 +133,8 @@ void SubroutineCloningPass::processCallOp(Operation *op) {
     // add calls within the new func def to the callWorkList
     newFunc->walk([&](CallLikeOp op) { callWorkList.push_back(op); });
 
+    symbolOps[mangledName] = newFunc.getOperation();
+
   } else { // matching function not found
     callOp->emitOpError() << "No matching function def found for "
                           << callOp.callee() << "\n";
@@ -148,18 +156,33 @@ void SubroutineCloningPass::runOnOperation() {
 
   mainFunc->walk([&](CallSubroutineOp op) { callWorkList.push_back(op); });
 
+  SymbolOpMap symbolOps;
+
+  if (!callWorkList.empty()) {
+    moduleOperation->walk([&](FuncOp functionOp) {
+      symbolOps[functionOp.sym_name()] = functionOp.getOperation();
+    });
+  }
+
   while (!callWorkList.empty()) {
     Operation *op = callWorkList.front();
     callWorkList.pop_front();
-    processCallOp<CallSubroutineOp, FuncOp>(op);
+    processCallOp<CallSubroutineOp, FuncOp>(op, symbolOps);
   }
 
   mainFunc->walk([&](CallCircuitOp op) { callWorkList.push_back(op); });
 
+  if (!callWorkList.empty()) {
+    symbolOps.clear();
+    moduleOperation->walk([&](CircuitOp circuitOp) {
+      symbolOps[circuitOp.sym_name()] = circuitOp.getOperation();
+    });
+  }
+
   while (!callWorkList.empty()) {
     Operation *op = callWorkList.front();
     callWorkList.pop_front();
-    processCallOp<CallCircuitOp, CircuitOp>(op);
+    processCallOp<CallCircuitOp, CircuitOp>(op, symbolOps);
   }
 
   // All subroutine defs that have been cloned are no longer needed
