No need to rename and undo rename! Just inline the first occurance.

paul0403 · paul0403 · commit d5b7c32aa10e · 2025-12-01T14:53:45.000-05:00
diff --git a/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp b/mlir/lib/Catalyst/Transforms/InlineNestedModules.cpp
@@ -185,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";
@@ -236,14 +239,14 @@ LogicalResult RenameFunctionsPattern::matchAndRewrite(Operation *child,
 
                 // We should not rename external function declarations, as they can be
                 // names required by other APIs.
-                // During inlining, on-the-fly they still need to be renamed, otherwise module
-                // verifier complains. So we save the original external API name and rename
-                // them back after inlining is finished.
+                // We record these external func decls during the rename pattern.
+                // Then during the actual inlining stage, only the first occurance of the per-module
+                // func decls of these external decls should be inlined.
                 if (isa<func::FuncOp>(op)) {
                     auto f = cast<func::FuncOp>(op);
                     if (f.isExternal()) {
-                        op.setAttr("original_external_API_name",
-                                   rewriter.getStringAttr(f.getName()));
+                        _externalFuncDeclNames->insert(f.getName());
+                        continue;
                     }
                 }
 
@@ -265,7 +268,14 @@ 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)
+    {
+    }
+
+    mutable llvm::SmallSet<StringRef, 8> alreadyInlinedFuncDeclNames;
 
     LogicalResult matchAndRewrite(Operation *op, PatternRewriter &rewriter) const override
     {
@@ -277,6 +287,28 @@ 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;
+        op->walk([&](func::FuncOp f) {
+            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();
@@ -288,6 +320,8 @@ struct InlineNestedModule : public RewritePattern {
 
         return success();
     }
+
+    llvm::SmallSet<StringRef, 8> _externalFuncDeclNames;
 };
 
 struct SymbolReplacerPattern
@@ -365,60 +399,6 @@ struct NestedToFlatCallPattern : public OpRewritePattern<catalyst::LaunchKernelO
     const DenseMap<SymbolRefAttr, SymbolRefAttr> *_map;
 };
 
-struct RestoreExternalFuncDeclNamePattern : public OpRewritePattern<func::FuncOp> {
-    using OpRewritePattern<func::FuncOp>::OpRewritePattern;
-
-    RestoreExternalFuncDeclNamePattern(MLIRContext *context)
-        : OpRewritePattern<func::FuncOp>::OpRewritePattern(context)
-    {
-    }
-
-    mutable llvm::SmallSet<StringRef, 8> CreatedAPIFuncNames;
-
-    LogicalResult matchAndRewrite(func::FuncOp op, PatternRewriter &rewriter) const override
-    {
-        if (!op->hasAttr("original_external_API_name")) {
-            return failure();
-        }
-
-        // Create func decl that matches the original name for the external API.
-        // This must happen only once.
-        StringRef APIFuncName =
-            cast<StringAttr>(op->getAttr("original_external_API_name")).getValue();
-        if (!CreatedAPIFuncNames.contains(APIFuncName)) {
-            auto APIFuncDecl = rewriter.create<func::FuncOp>(
-                op->getLoc(), rewriter.getStringAttr(APIFuncName), op.getFunctionType());
-            APIFuncDecl.setPrivate();
-            CreatedAPIFuncNames.insert(APIFuncName);
-            return success();
-        }
-        return failure();
-    }
-};
-
-struct UpdateCalleeToExternalAPINamesPattern : public OpRewritePattern<func::CallOp> {
-    using OpRewritePattern<func::CallOp>::OpRewritePattern;
-
-    UpdateCalleeToExternalAPINamesPattern(
-        MLIRContext *context, mlir::DenseMap<StringRef, StringAttr> uniqued_names_to_APIName)
-        : OpRewritePattern<func::CallOp>::OpRewritePattern(context),
-          _uniqued_names_to_APIName(uniqued_names_to_APIName)
-    {
-    }
-
-    mlir::DenseMap<StringRef, StringAttr> _uniqued_names_to_APIName;
-
-    LogicalResult matchAndRewrite(func::CallOp op, PatternRewriter &rewriter) const override
-    {
-        if (_uniqued_names_to_APIName.contains(op.getCallee())) {
-            op.setCallee(_uniqued_names_to_APIName.at(op.getCallee()));
-            return success();
-        }
-
-        return failure();
-    }
-};
-
 struct CleanupPattern : public RewritePattern {
     /// This overload constructs a pattern that matches any operation type.
     CleanupPattern(MLIRContext *context) : RewritePattern(MatchAnyOpTypeTag(), 1, context) {}
@@ -431,9 +411,6 @@ struct CleanupPattern : public RewritePattern {
         }
         rewriter.modifyOpInPlace(op, [&] { op->removeAttr(fullyQualifiedNameAttr); });
 
-        if (op->hasAttr("original_external_API_name")) {
-            op->erase();
-        }
         return success();
     }
 };
@@ -472,34 +449,6 @@ struct InlineNestedSymbolTablePass : PassWrapper<InlineNestedSymbolTablePass, Op
     int _stopAfterStep;
     InlineNestedSymbolTablePass(int stopAfter) : _stopAfterStep(stopAfter) {}
 
-    void restoreExternalFuncDeclName(Operation *symbolTable, MLIRContext *context,
-                                     GreedyRewriteConfig config)
-    {
-        mlir::DenseMap<StringRef, StringAttr> uniqued_names_to_APIName;
-        symbolTable->walk([&](func::FuncOp funcOp) {
-            if (funcOp->hasAttr("original_external_API_name")) {
-                uniqued_names_to_APIName[funcOp.getSymName()] =
-                    cast<StringAttr>(funcOp->getAttr("original_external_API_name"));
-            }
-        });
-        RewritePatternSet restoreExternalFuncDeclName(context);
-        restoreExternalFuncDeclName.add<RestoreExternalFuncDeclNamePattern>(context);
-        bool run = _stopAfterStep >= 5 || _stopAfterStep == 0;
-        if (run && failed(applyPatternsGreedily(symbolTable, std::move(restoreExternalFuncDeclName),
-                                                config))) {
-            signalPassFailure();
-        }
-
-        RewritePatternSet updateCalleeToExternalAPINames(context);
-        updateCalleeToExternalAPINames.add<UpdateCalleeToExternalAPINamesPattern>(
-            context, uniqued_names_to_APIName);
-        run = _stopAfterStep >= 5 || _stopAfterStep == 0;
-        if (run && failed(applyPatternsGreedily(
-                       symbolTable, std::move(updateCalleeToExternalAPINames), config))) {
-            signalPassFailure();
-        }
-    }
-
     void runOnOperation() override
     {
         // Here we are in a root module/symbol table
@@ -528,15 +477,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();
@@ -568,9 +518,6 @@ struct InlineNestedSymbolTablePass : PassWrapper<InlineNestedSymbolTablePass, Op
             signalPassFailure();
         }
 
-        // Restore external API func decl names and update calls.
-        restoreExternalFuncDeclName(symbolTable, context, config);
-
         RewritePatternSet cleanup(context);
         cleanup.add<CleanupPattern>(context);
         run = _stopAfterStep >= 5 || _stopAfterStep == 0;
