Do not rename external func decls during --inline-nested-modules (#2244)

**Context:**
One of the first steps in the compilation pipeline is
`--inline-nested-modules`. This pass lifts the per-qnode module up into
the global module of the qjit.

Because (usually) the symbol table of a module is isolated from above,
multiple modules in parallel may contain the same names inside.
Therefore when inlining child modules into the global module, we need to
rename them to unique names.

However, there's a small problem. Each qnode module has their own
specific transform sequence, so the quantum passes (applied during
`--apply-transform-sequence`) have to be resolved before
`--inline-nested-module`. If a pass generates a function declaration to
an external API (e.g. runtime functions), that name must not be altered.

**Description of the Change:**
Do not perform renaming on external function declarations from within
the qnode modules during `--inline-nested-modules`.
Only inline the first occurrence of such func decls from the qnode
modules into the global qjit module.

**Benefits:**
Quantum passes that generate calls to other APIs can work.  
The flipside is also true! API developers do not have to maintain
multiple aliases for their API functions (e.g. gridsynth pass #2140 )

[sc-105020]

---------

Co-authored-by: Joseph Lee <[email protected]>

Author: paul0403

doc/releases/changelog-dev.md

@@ -263,6 +263,13 @@
 
   This fix enables automatic qubit management to be used with gradients.
 
+* The `--inline-nested-module` pass no longer renames external function declarations.
+  [(#2244)](https://github.com/PennyLaneAI/catalyst/pull/2244)
+
+  This pass inlines the qnode MLIR modules into the global QJIT MLIR module. If a qnode module
+  contains function declarations to external APIs, the names of these declarations must
+  stay unchanged. This change enables quantum passes to generate calls to external APIs.
+
 <h3>Internal changes ⚙️</h3>
 
 * Resource tracking now writes out at device destruction time instead of qubit deallocation

mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp

@@ -65,6 +65,8 @@
 
 #include <deque>
 
+#include "llvm/ADT/SmallSet.h"
+
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
@@ -183,14 +185,17 @@ struct AnnotateWithFullyQualifiedName : public OpInterfaceRewritePattern<SymbolO
 
 struct RenameFunctionsPattern : public RewritePattern {
     /// This overload constructs a pattern that matches any operation type.
-    RenameFunctionsPattern(MLIRContext *context, SmallVector<Operation *> *symbolTables)
-        : RewritePattern(MatchAnyOpTypeTag(), 1, context), _symbolTables(symbolTables)
+    RenameFunctionsPattern(MLIRContext *context, SmallVector<Operation *> *symbolTables,
+                           llvm::SmallSet<StringRef, 8> *externalFuncDeclNames)
+        : RewritePattern(MatchAnyOpTypeTag(), 1, context), _symbolTables(symbolTables),
+          _externalFuncDeclNames(externalFuncDeclNames)
     {
     }
 
     LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override;
 
     SmallVector<Operation *> *_symbolTables;
+    llvm::SmallSet<StringRef, 8> *_externalFuncDeclNames;
 };
 
 static constexpr llvm::StringRef hasBeenRenamedAttrName = "catalyst.unique_names";
@@ -228,8 +233,21 @@ LogicalResult RenameFunctionsPattern::matchAndRewrite(Operation *child,
     for (auto &region : child->getRegions()) {
         for (auto &block : region.getBlocks()) {
             for (auto &op : block) {
-                if (!isa<SymbolOpInterface>(op))
+                if (!isa<SymbolOpInterface>(op)) {
                     continue;
+                }
+
+                // We should not rename external function declarations, as they can be
+                // names required by other APIs.
+                // We record these external func decls during the rename pattern.
+                // Then during the actual inlining stage, only the first occurrence of the
+                // per-module func decls of these external decls should be inlined.
+                if (auto f = dyn_cast<func::FuncOp>(op)) {
+                    if (f.isExternal()) {
+                        _externalFuncDeclNames->insert(f.getName());
+                        continue;
+                    }
+                }
 
                 if (failed(childSymTab.renameToUnique(&op, raw_tables))) {
                     // TODO: Check for error in one of the tests.
@@ -249,7 +267,23 @@ LogicalResult RenameFunctionsPattern::matchAndRewrite(Operation *child,
 
 struct InlineNestedModule : public RewritePattern {
     /// This overload constructs a pattern that matches any operation type.
-    InlineNestedModule(MLIRContext *context) : RewritePattern(MatchAnyOpTypeTag(), 1, context) {}
+    InlineNestedModule(MLIRContext *context,
+                       const llvm::SmallSet<StringRef, 8> &externalFuncDeclNames)
+        : RewritePattern(MatchAnyOpTypeTag(), 1, context),
+          _externalFuncDeclNames(externalFuncDeclNames)
+    {
+    }
+
+    const llvm::SmallSet<StringRef, 8> &_externalFuncDeclNames;
+
+    // Note: mlir expects pattern objects to be const.
+    // In other words, repeated applications of a rewrite pattern should not have dependency on each
+    // other.
+    // This --inline-nested-module pass is breaking this assumption.
+    //
+    // TODO: refactor this pass to not use the pattern rewriter, but just raw logic in a
+    // `runOnOperation()`
+    mutable llvm::SmallSet<StringRef, 8> alreadyInlinedFuncDeclNames;
 
     LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override
     {
@@ -261,6 +295,31 @@ struct InlineNestedModule : public RewritePattern {
         }
 
         auto parent = op->getParentOp();
+        assert(parent->hasTrait<OpTrait::SymbolTable>() &&
+               "the direct parent of a qnode module must be a module op");
+
+        // Look for the func decls in the current qnode module
+        // If it is a recorded external func decl, erase it if it already has been inlined.
+        SmallVector<Operation *> _erasureWorklist;
+        for (auto &region : op->getRegions()) {
+            auto funcOps = region.getOps<func::FuncOp>();
+            for (auto f : funcOps) {
+                StringRef funcName = f.getName();
+                if (f.isExternal() && _externalFuncDeclNames.contains(funcName)) {
+                    if (alreadyInlinedFuncDeclNames.contains(funcName)) {
+                        _erasureWorklist.push_back(f);
+                    }
+                    else {
+                        alreadyInlinedFuncDeclNames.insert(funcName);
+                    }
+                }
+            }
+        }
+
+        for (auto _op : _erasureWorklist) {
+            rewriter.eraseOp(_op);
+        }
+
         // Can't generalize getting a region other than the zero-th one.
         rewriter.inlineRegionBefore(op->getRegion(0), &parent->getRegion(0).back());
         Block *inlinedBlock = &parent->getRegion(0).front();
@@ -360,6 +419,7 @@ struct CleanupPattern : public RewritePattern {
             return failure();
         }
         rewriter.modifyOpInPlace(op, [&] { op->removeAttr(fullyQualifiedNameAttr); });
+
         return success();
     }
 };
@@ -426,15 +486,16 @@ struct InlineNestedSymbolTablePass : PassWrapper<InlineNestedSymbolTablePass, Op
             return WalkResult::skip();
         });
 
-        renameFunctions.add<RenameFunctionsPattern>(context, &symbolTables);
+        llvm::SmallSet<StringRef, 8> externalFuncDeclNames;
+        renameFunctions.add<RenameFunctionsPattern>(context, &symbolTables, &externalFuncDeclNames);
 
         bool run = _stopAfterStep >= 2 || _stopAfterStep == 0;
         if (run && failed(applyPatternsGreedily(symbolTable, std::move(renameFunctions), config))) {
             signalPassFailure();
         }
 
         RewritePatternSet inlineNested(context);
-        inlineNested.add<InlineNestedModule>(context);
+        inlineNested.add<InlineNestedModule>(context, externalFuncDeclNames);
         run = _stopAfterStep >= 3 || _stopAfterStep == 0;
         if (run && failed(applyPatternsGreedily(symbolTable, std::move(inlineNested), config))) {
             signalPassFailure();

mlir/test/Catalyst/NestedModule2.mlir

@@ -61,4 +61,3 @@ module @outer {
     }
   }
 }
-

mlir/test/Catalyst/NestedModule5.mlir

@@ -29,3 +29,35 @@ module @outer {
   catalyst.launch_kernel @inner::@f() : () -> ()
 }
 
+// -----
+
+// Test external API func decl names are preserved
+module @global {
+
+  // CHECK-DAG: func.func private @f()
+  // CHECK-DAG: func.func @main_0() {
+  // CHECK-DAG:   call @f() : () -> ()
+  // CHECK-DAG:   return
+  // CHECK-DAG: }
+  // CHECK-DAG: func.func @main_1() {
+  // CHECK-DAG:   call @f() : () -> ()
+  // CHECK-DAG:   return
+  // CHECK-DAG: }
+
+  module @local0 {
+    func.func private @f()
+    func.func @main() {
+      func.call @f() : () -> ()
+      return
+    }
+  }
+
+  module @local1 {
+    func.func private @f()
+    func.func @main() {
+      func.call @f() : () -> ()
+      return
+    }
+  }
+
+}