[flang-commits] [flang] [flang][NFC] speedup BoxedProcedure for derived types with many components (PR #86144)

via flang-commits flang-commits at lists.llvm.org
Fri Mar 22 06:23:24 PDT 2024


https://github.com/jeanPerier updated https://github.com/llvm/llvm-project/pull/86144

>From acd6be06216914ae90d1ff819ebea9d75198cf13 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Thu, 21 Mar 2024 08:59:47 -0700
Subject: [PATCH 1/2] [flang][NFC] speedup BoxedProcedure for derived types
 with many components

This patch speeds up the compilation time of the example in
https://github.com/llvm/llvm-project/issues/76478#issuecomment-2011023289
from 2mins with my builds to about 2 seconds.

MLIR timers showed more than 98% of the time was spend in BoxedProcedure
trying to figure out if a type needs to be converted.

This is because walking the fir.type members is very expansive for types
containing many components and/or components with many sub-components.

Increase the caching time of visited types from "the type being visited"
to "the whole pass". Use DenseMap since it is not ok anymore to assume
this container will only have a few elements.
---
 .../lib/Optimizer/CodeGen/BoxedProcedure.cpp  | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
index 746c275f37eaca..7a1cc3536b263b 100644
--- a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
+++ b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
@@ -69,17 +69,20 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
       return false;
     }
     if (auto recTy = ty.dyn_cast<RecordType>()) {
-      if (llvm::is_contained(visitedTypes, recTy))
-        return false;
+      auto visited = visitedTypes.find(ty);
+      if (visited != visitedTypes.end())
+        return visited->second;
+      [[maybe_unused]] auto newIt = visitedTypes.try_emplace(ty, false);
+      assert(newIt.second && "expected ty to not be in the map");
       bool result = false;
-      visitedTypes.push_back(recTy);
       for (auto t : recTy.getTypeList()) {
         if (needsConversion(t.second)) {
           result = true;
           break;
         }
       }
-      visitedTypes.pop_back();
+      // newIt may have been invalidated.
+      visitedTypes.find(ty)->second = result;
       return result;
     }
     if (auto boxTy = ty.dyn_cast<BaseBoxType>())
@@ -140,9 +143,7 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
       if (rec.isFinalized())
         return rec;
       auto it = convertedTypes.try_emplace(ty, rec);
-      if (!it.second) {
-        llvm::errs() << "failed\n" << ty << "\n";
-      }
+      assert(it.second && "expected ty to not be in the map");
       std::vector<RecordType::TypePair> ps = ty.getLenParamList();
       std::vector<RecordType::TypePair> cs;
       for (auto t : ty.getTypeList()) {
@@ -171,11 +172,11 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
   void setLocation(mlir::Location location) { loc = location; }
 
 private:
-  llvm::SmallVector<mlir::Type> visitedTypes;
-  // Map to deal with recursive derived types (avoid infinite loops).
+  // Maps to deal with recursive derived types (avoid infinite loops).
   // Caching is also beneficial for apps with big types (dozens of
   // components and or parent types), so the lifetime of the cache
   // is the whole pass.
+  llvm::DenseMap<mlir::Type, bool> visitedTypes;
   llvm::DenseMap<mlir::Type, mlir::Type> convertedTypes;
   mlir::Location loc;
 };

>From 42b08e942702711cb78eb76b888967604b866201 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Fri, 22 Mar 2024 06:15:15 -0700
Subject: [PATCH 2/2] fix regression with indirect type recursion

The previous update caused t002 type in the added test to not
be rewritten because it was first analyzed when analyzing t001.
Since the reason it needs to be rewritten is because it contains
a t001, but t001 is not fully analyzed when first analyzing t002,
the first analyzed yieled false. This "wrong" result is not an
issue when analyzing t001 (since t002 does not determine t001
result). But caching this invalid result for later usage is invalid.

Add a boolean to only keep analysis for top-level types where the result
is known to be correct.
This restriction has no measurable impact on the speed-up: I still measure
2 seconds.
---
 flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp | 17 +++++++++++++++--
 flang/test/Fir/boxproc-2.fir                   |  7 +++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
index 7a1cc3536b263b..dccb642d5f79f0 100644
--- a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
+++ b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
@@ -74,6 +74,8 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
         return visited->second;
       [[maybe_unused]] auto newIt = visitedTypes.try_emplace(ty, false);
       assert(newIt.second && "expected ty to not be in the map");
+      bool wasAlreadyVisitingRecordType = needConversionIsVisitingRecordType;
+      needConversionIsVisitingRecordType = true;
       bool result = false;
       for (auto t : recTy.getTypeList()) {
         if (needsConversion(t.second)) {
@@ -81,8 +83,18 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
           break;
         }
       }
-      // newIt may have been invalidated.
-      visitedTypes.find(ty)->second = result;
+      // Only keep the result cached if the fir.type visited was a "top-level
+      // type". Nested types with a recursive reference to the "top-level type"
+      // may incorrectly have been resolved as not needed conversions because it
+      // had not been determined yet if the "top-level type" needed conversion.
+      // This is not an issue to determine the "top-level type" need of
+      // conversion, but the result should not be kept and later used in other
+      // contexts.
+      needConversionIsVisitingRecordType = wasAlreadyVisitingRecordType;
+      if (needConversionIsVisitingRecordType)
+        visitedTypes.erase(ty);
+      else
+        visitedTypes.find(ty)->second = result;
       return result;
     }
     if (auto boxTy = ty.dyn_cast<BaseBoxType>())
@@ -177,6 +189,7 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
   // components and or parent types), so the lifetime of the cache
   // is the whole pass.
   llvm::DenseMap<mlir::Type, bool> visitedTypes;
+  bool needConversionIsVisitingRecordType = false;
   llvm::DenseMap<mlir::Type, mlir::Type> convertedTypes;
   mlir::Location loc;
 };
diff --git a/flang/test/Fir/boxproc-2.fir b/flang/test/Fir/boxproc-2.fir
index 3279bb9bed8a4c..84132f89afebb2 100644
--- a/flang/test/Fir/boxproc-2.fir
+++ b/flang/test/Fir/boxproc-2.fir
@@ -108,3 +108,10 @@ func.func @very_nested_type(%arg0: !fir.type<t0{t01:!fir.type<t01{t02:!fir.type<
 }
 // CHECK-LABEL: func.func @very_nested_type(
 // CHECK-SAME: !fir.type<t0UnboxProc{t01:!fir.type<t01UnboxProc{t02:!fir.type<t02UnboxProc{t3:!fir.type<t3UnboxProc{t4:!fir.type<t4UnboxProc{t5:!fir.type<t5UnboxProc{t6:!fir.type<t6UnboxProc{t7:!fir.type<t7UnboxProc{t8:!fir.type<t8UnboxProc{t9:!fir.type<t9UnboxProc{t10:!fir.type<t10UnboxProc{t11:!fir.type<t11UnboxProc{t12:!fir.type<t12UnboxProc{t13:!fir.type<t13UnboxProc{t14:!fir.type<t14UnboxProc{t15:!fir.type<t15UnboxProc{t16:!fir.type<t16UnboxProc{t17:!fir.type<t17UnboxProc{t18:!fir.type<t18UnboxProc{t19:!fir.type<t19UnboxProc{b:() -> ()}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>)
+
+func.func @test_indirect_type_recursion() {
+    %0 = fir.zero_bits !fir.ptr<!fir.type<t001{x:!fir.ptr<!fir.type<t002{x:!fir.ptr<!fir.type<t001>>}>>,p:!fir.boxproc<() -> ()>}>>
+  return
+}
+// CHECK-LABEL: func.func @test_indirect_type_recursion(
+// CHECK: fir.zero_bits !fir.ptr<!fir.type<t001UnboxProc{x:!fir.ptr<!fir.type<t002UnboxProc{x:!fir.ptr<!fir.type<t001UnboxProc>>}>>,p:() -> ()}>>



More information about the flang-commits mailing list