[Mlir-commits] [mlir] [mlir] Speed up FuncToLLVM: CallOpLowering using a hashtable (PR #68082)

[llvmlistbot at llvm.org]

https://github.com/tdanyluk updated https://github.com/llvm/llvm-project/pull/68082

>From 078df36699ef3db1131bc9266b25fc43dbd49551 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tam=C3=A1s=20Danyluk?= <tdanyluk at google.com>
Date: Tue, 3 Oct 2023 11:07:20 +0200
Subject: [PATCH] [mlir] Speed up FuncToLLVM by using a hashmap

We have a project where this saves 23% of the compilation time.
---
 mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp | 71 +++++++++++++++----
 1 file changed, 58 insertions(+), 13 deletions(-)

diff --git a/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp b/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
index 7de7f3cb9e36b06..d37f8c6991ae380 100644
--- a/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
+++ b/mlir/lib/Conversion/FuncToLLVM/FuncToLLVM.cpp
@@ -39,6 +39,7 @@
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/IRBuilder.h"
@@ -601,19 +602,37 @@ struct CallOpInterfaceLowering : public ConvertOpToLLVMPattern<CallOpType> {
   }
 };
 
-struct CallOpLowering : public CallOpInterfaceLowering<func::CallOp> {
-  using Super::Super;
+class CallOpLowering : public CallOpInterfaceLowering<func::CallOp> {
+public:
+  CallOpLowering(const LLVMTypeConverter &typeConverter,
+                 // Can be nullptr.
+                 const StringSet<> *functionsHavingBarePtrAttribute,
+                 PatternBenefit benefit = 1)
+      : CallOpInterfaceLowering<func::CallOp>(typeConverter, benefit),
+        barePtrCallConvForcedByOptions(
+            typeConverter.getOptions().useBarePtrCallConv),
+        functionsHavingBarePtrAttribute(functionsHavingBarePtrAttribute) {}
 
   LogicalResult
   matchAndRewrite(func::CallOp callOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     bool useBarePtrCallConv = false;
-    if (Operation *callee = SymbolTable::lookupNearestSymbolFrom(
+    if (barePtrCallConvForcedByOptions) {
+      useBarePtrCallConv = true;
+    } else if (functionsHavingBarePtrAttribute != nullptr) {
+      useBarePtrCallConv = functionsHavingBarePtrAttribute->contains(
+          callOp.getCalleeAttr().getValue());
+    } else if ( // Warning: This is a linear lookup.
+        Operation *callee = SymbolTable::lookupNearestSymbolFrom(
             callOp, callOp.getCalleeAttr())) {
-      useBarePtrCallConv = shouldUseBarePtrCallConv(callee, getTypeConverter());
+      useBarePtrCallConv = callee->hasAttr(barePtrAttrName);
     }
     return matchAndRewriteImpl(callOp, adaptor, rewriter, useBarePtrCallConv);
   }
+
+private:
+  bool barePtrCallConvForcedByOptions = false;
+  const StringSet<> *functionsHavingBarePtrAttribute = nullptr;
 };
 
 struct CallIndirectOpLowering
@@ -728,16 +747,24 @@ void mlir::populateFuncToLLVMFuncOpConversionPattern(
   patterns.add<FuncOpConversion>(converter);
 }
 
+namespace {
+void populateFuncToLLVMConversionPatternsInternal(
+    LLVMTypeConverter &converter, RewritePatternSet &patterns,
+    // Can be nullptr.
+    const StringSet<> *functionsHavingBarePtrAttribute) {
+  populateFuncToLLVMFuncOpConversionPattern(converter, patterns);
+  patterns.add<CallIndirectOpLowering>(converter);
+  patterns.add<CallOpLowering>(converter, functionsHavingBarePtrAttribute);
+  patterns.add<ConstantOpLowering>(converter);
+  patterns.add<ReturnOpLowering>(converter);
+}
+} // namespace
+
 void mlir::populateFuncToLLVMConversionPatterns(LLVMTypeConverter &converter,
                                                 RewritePatternSet &patterns) {
-  populateFuncToLLVMFuncOpConversionPattern(converter, patterns);
-  // clang-format off
-  patterns.add<
-      CallIndirectOpLowering,
-      CallOpLowering,
-      ConstantOpLowering,
-      ReturnOpLowering>(converter);
-  // clang-format on
+  populateFuncToLLVMConversionPatternsInternal(
+      converter, patterns,
+      /*functionsHavingBarePtrAttribute=*/nullptr);
 }
 
 namespace {
@@ -765,6 +792,23 @@ struct ConvertFuncToLLVMPass
 
     const auto &dataLayoutAnalysis = getAnalysis<DataLayoutAnalysis>();
 
+    StringSet<> functionsHavingBarePtrAttribute;
+    for (Region &region : m->getRegions()) {
+      for (Block &block : region.getBlocks()) {
+        for (Operation &op : block.getOperations()) {
+          if (auto funcOp = dyn_cast<func::FuncOp>(op);
+              funcOp != nullptr && funcOp->hasAttr(barePtrAttrName)) {
+            bool inserted =
+                functionsHavingBarePtrAttribute.insert(funcOp.getSymName())
+                    .second;
+            (void)inserted;
+            assert(inserted &&
+                   "expected module to contain uniquely named funcOps");
+          }
+        }
+      }
+    }
+
     LowerToLLVMOptions options(&getContext(),
                                dataLayoutAnalysis.getAtOrAbove(m));
     options.useBarePtrCallConv = useBarePtrCallConv;
@@ -777,7 +821,8 @@ struct ConvertFuncToLLVMPass
                                     &dataLayoutAnalysis);
 
     RewritePatternSet patterns(&getContext());
-    populateFuncToLLVMConversionPatterns(typeConverter, patterns);
+    populateFuncToLLVMConversionPatternsInternal(
+        typeConverter, patterns, &functionsHavingBarePtrAttribute);
 
     // TODO: Remove these in favor of their dedicated conversion passes.
     arith::populateArithToLLVMConversionPatterns(typeConverter, patterns);