MergeModules utility function (#3449)

* * Adding a mergeModule utility function to merge 2 modules
* Avoid using walk as it is wrong. Filtering artifacts by their IR node
  type is incorrect.
* Adding unit tests

co-authored by: Eric Schweitz <[email protected]>
Signed-off-by: Sachin Pisal <[email protected]>

* formatting

Signed-off-by: Sachin Pisal <[email protected]>

* * Moving mergeModules definition to Factory.cpp
* Updating the unittests name

Signed-off-by: Sachin Pisal <[email protected]>

---------

Signed-off-by: Sachin Pisal <[email protected]>

Author: sacpis

include/cudaq/Optimizer/Builder/Factory.h

@@ -294,6 +294,7 @@ std::pair<mlir::func::FuncOp, /*alreadyDefined=*/bool>
 getOrAddFunc(mlir::Location loc, mlir::StringRef funcName,
              mlir::FunctionType funcTy, mlir::ModuleOp module);
 
+void mergeModules(mlir::ModuleOp into, mlir::ModuleOp from);
 } // namespace factory
 
 std::size_t getDataSize(llvm::DataLayout &dataLayout, mlir::Type ty);

lib/Optimizer/Builder/Factory.cpp

@@ -774,4 +774,17 @@ factory::getOrAddFunc(mlir::Location loc, mlir::StringRef funcName,
   return {func, /*defined=*/false};
 }
 
+void factory::mergeModules(ModuleOp into, ModuleOp from) {
+  for (Operation &op : *from.getBody()) {
+    auto sym = dyn_cast<SymbolOpInterface>(op);
+    if (!sym)
+      continue; // Only merge named symbols, avoids duplicating anonymous ops.
+
+    // If `into` already has a symbol with this name, skip it.
+    if (SymbolTable::lookupSymbolIn(into, sym.getName()))
+      continue;
+
+    into.push_back(op.clone());
+  }
+}
 } // namespace cudaq::opt

python/extension/CUDAQuantumExtension.cpp

@@ -311,11 +311,9 @@ PYBIND11_MODULE(_quakeDialects, m) {
         auto ctx = unwrap(mod).getContext();
         auto moduleB = mlir::parseSourceString<mlir::ModuleOp>(code, ctx);
         auto moduleA = unwrap(mod);
-        moduleB->walk([&moduleA](mlir::func::FuncOp op) {
-          if (!moduleA.lookupSymbol<mlir::func::FuncOp>(op.getName()))
-            moduleA.push_back(op.clone());
-          return mlir::WalkResult::advance();
-        });
+
+        // Merge symbols from moduleB into moduleA.
+        cudaq::opt::factory::mergeModules(moduleA, *moduleB);
         return kName;
       },
       "Given a python module name like `mod1.mod2.func`, see if there is a "

unittests/Optimizer/CMakeLists.txt

@@ -8,12 +8,14 @@
 
 include(HandleLLVMOptions)
 
-add_executable(OptimizerUnitTests HermitianTrait.cpp)
+add_executable(OptimizerUnitTests HermitianTrait.cpp FactoryMergeModuleTest.cpp)
 
 target_link_libraries(OptimizerUnitTests
   PRIVATE
+    MLIRParser
     QuakeDialect
     gtest_main
+    OptimBuilder
 )
 
 gtest_discover_tests(OptimizerUnitTests)

unittests/Optimizer/FactoryMergeModuleTest.cpp

@@ -0,0 +1,99 @@
+/*******************************************************************************
+ * Copyright (c) 2025 NVIDIA Corporation & Affiliates.                         *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+#include "cudaq/Optimizer/Builder/Factory.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/DialectRegistry.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Parser/Parser.h"
+#include <gtest/gtest.h>
+
+using namespace mlir;
+using cudaq::opt::factory::mergeModules;
+
+static mlir::OwningOpRef<mlir::ModuleOp> parse(MLIRContext &ctx,
+                                               llvm::StringRef ir) {
+  return mlir::parseSourceString<mlir::ModuleOp>(ir, &ctx);
+}
+
+TEST(FactoryMergeModuleTest, CopiesMissingFunction) {
+  DialectRegistry registry;
+  registry.insert<mlir::func::FuncDialect, mlir::LLVM::LLVMDialect>();
+  MLIRContext ctx(registry);
+
+  auto dst = parse(ctx, R"mlir(
+        module {
+            func.func @alreadyThere() { return }
+        }
+    )mlir");
+  ASSERT_TRUE(dst);
+
+  auto src = parse(ctx, R"mlir(
+        module {
+            func.func @newFunc() { return }
+            func.func @alreadyThere() { return }
+        }
+    )mlir");
+  ASSERT_TRUE(src);
+
+  mergeModules(*dst, *src);
+
+  auto newFunc = dst->lookupSymbol<mlir::func::FuncOp>("newFunc");
+  EXPECT_TRUE(newFunc);
+
+  int countAlreadyThere = 0;
+  dst->walk([&](mlir::func::FuncOp f) {
+    if (f.getSymName() == "alreadyThere")
+      countAlreadyThere++;
+  });
+  EXPECT_EQ(countAlreadyThere, 1);
+}
+
+TEST(FactoryMergeModuleTest, RetainOriginalModuleSymbols) {
+  DialectRegistry registry;
+  registry.insert<mlir::func::FuncDialect, mlir::LLVM::LLVMDialect>();
+  MLIRContext ctx(registry);
+
+  auto dst = parse(ctx, R"mlir(
+        module attributes { test.attr = "keepme" } {
+            func.func @a() { return }
+        }
+    )mlir");
+  ASSERT_TRUE(dst);
+
+  auto src = parse(ctx, R"mlir(
+        module {
+            func.func @b() { return }
+        }
+    )mlir");
+  ASSERT_TRUE(src);
+
+  mergeModules(*dst, *src);
+
+  // Verify the destination attribute remains
+  auto sattr =
+      dst->getOperation()->getAttrOfType<mlir::StringAttr>("test.attr");
+  ASSERT_TRUE(sattr);
+  ASSERT_EQ(sattr.getValue(), "keepme");
+
+  // Both symbols exists exactly once
+  EXPECT_TRUE(dst->lookupSymbol<mlir::func::FuncOp>("a"));
+  EXPECT_TRUE(dst->lookupSymbol<mlir::func::FuncOp>("b"));
+
+  int countA = 0, countB = 0;
+  dst->walk([&](mlir::func::FuncOp f) {
+    if (f.getSymName() == "a")
+      countA++;
+    if (f.getSymName() == "b")
+      countB++;
+  });
+  EXPECT_EQ(countA, 1);
+  EXPECT_EQ(countB, 1);
+}