[flang][NFC] Cache derived type translation in lowering

[jeanPerier]

Derived type translation is proving expensive in modern fortran apps with many big derived types with dozens of components and parents.

Extending the cache that prevent recursion is proving to have little cost on apps with small derived types and significant gain (can divide compile time by 2) on modern fortran apps.

It is legal since the cache lifetime is less than the MLIRContext lifetime that owns the cached mlir::Type.

Doing so also exposed that the current caching was incorrect, the type symbol is the same for kind parametrized derived types regardless of the kind parameters. Instances with different kinds should lower to different MLIR types. See added test.
Using the type scopes fixes the problem.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: None (jeanPerier)

Changes

Derived type translation is proving expensive in modern fortran apps with many big derived types with dozens of components and parents.

Extending the cache that prevent recursion is proving to have little cost on apps with small derived types and significant gain (can divide compile time by 2) on modern fortran apps.

It is legal since the cache lifetime is less than the MLIRContext lifetime that owns the cached mlir::Type.

Doing so also exposed that the current caching was incorrect, the type symbol is the same for kind parametrized derived types regardless of the kind parameters. Instances with different kinds should lower to different MLIR types. See added test.
Using the type scopes fixes the problem.

---

Full diff: https://github.com/llvm/llvm-project/pull/80179.diff

3 Files Affected:

• (modified) flang/include/flang/Lower/AbstractConverter.h (+2-1)
• (modified) flang/lib/Lower/ConvertType.cpp (+9-17)
• (added) flang/test/Lower/derived-types-kind-params-2.f90 (+14)

diff --git a/flang/include/flang/Lower/AbstractConverter.h b/flang/include/flang/Lower/AbstractConverter.h
index c19dcbdcdb390..796933a4eb5f6 100644
--- a/flang/include/flang/Lower/AbstractConverter.h
+++ b/flang/include/flang/Lower/AbstractConverter.h
@@ -47,6 +47,7 @@ class CharBlock;
 }
 namespace semantics {
 class Symbol;
+class Scope;
 class DerivedTypeSpec;
 } // namespace semantics
 
@@ -59,7 +60,7 @@ struct Variable;
 using SomeExpr = Fortran::evaluate::Expr<Fortran::evaluate::SomeType>;
 using SymbolRef = Fortran::common::Reference<const Fortran::semantics::Symbol>;
 using TypeConstructionStack =
-    llvm::SmallVector<std::pair<const Fortran::lower::SymbolRef, mlir::Type>>;
+    llvm::DenseMap<const Fortran::semantics::Scope *, mlir::Type>;
 class StatementContext;
 
 using ExprToValueMap = llvm::DenseMap<const SomeExpr *, mlir::Value>;
diff --git a/flang/lib/Lower/ConvertType.cpp b/flang/lib/Lower/ConvertType.cpp
index 8caafb72e472a..21564e8b81d70 100644
--- a/flang/lib/Lower/ConvertType.cpp
+++ b/flang/lib/Lower/ConvertType.cpp
@@ -374,19 +374,20 @@ struct TypeBuilderImpl {
   mlir::Type genDerivedType(const Fortran::semantics::DerivedTypeSpec &tySpec) {
     std::vector<std::pair<std::string, mlir::Type>> ps;
     std::vector<std::pair<std::string, mlir::Type>> cs;
-    const Fortran::semantics::Symbol &typeSymbol = tySpec.typeSymbol();
-    if (mlir::Type ty = getTypeIfDerivedAlreadyInConstruction(typeSymbol))
-      return ty;
-
     if (tySpec.IsVectorType()) {
       return genVectorType(tySpec);
     }
 
+    const Fortran::semantics::Symbol &typeSymbol = tySpec.typeSymbol();
     const Fortran::semantics::Scope &derivedScope = DEREF(tySpec.GetScope());
+    if (mlir::Type ty = getTypeIfDerivedAlreadyInConstruction(derivedScope))
+      return ty;
 
     auto rec = fir::RecordType::get(context, converter.mangleName(tySpec));
-    // Maintain the stack of types for recursive references.
-    derivedTypeInConstruction.emplace_back(typeSymbol, rec);
+    // Maintain the stack of types for recursive references and to speed-up
+    // the derived type constructions that can be expensive for derived type
+    // with dozens of components/parents (modern Fortran).
+    derivedTypeInConstruction.try_emplace(&derivedScope, rec);
 
     // Gather the record type fields.
     // (1) The data components.
@@ -446,7 +447,6 @@ struct TypeBuilderImpl {
       }
 
     rec.finalize(ps, cs);
-    popDerivedTypeInConstruction();
 
     if (!ps.empty()) {
       // TODO: this type is a PDT (parametric derived type) with length
@@ -552,16 +552,8 @@ struct TypeBuilderImpl {
   /// type `t` have type `t`. This helper returns `t` if it is already being
   /// lowered to avoid infinite loops.
   mlir::Type getTypeIfDerivedAlreadyInConstruction(
-      const Fortran::lower::SymbolRef derivedSym) const {
-    for (const auto &[sym, type] : derivedTypeInConstruction)
-      if (sym == derivedSym)
-        return type;
-    return {};
-  }
-
-  void popDerivedTypeInConstruction() {
-    assert(!derivedTypeInConstruction.empty());
-    derivedTypeInConstruction.pop_back();
+      const Fortran::semantics::Scope &derivedScope) const {
+    return derivedTypeInConstruction.lookup(&derivedScope);
   }
 
   /// Stack derived type being processed to avoid infinite loops in case of
diff --git a/flang/test/Lower/derived-types-kind-params-2.f90 b/flang/test/Lower/derived-types-kind-params-2.f90
new file mode 100644
index 0000000000000..5833079901b9c
--- /dev/null
+++ b/flang/test/Lower/derived-types-kind-params-2.f90
@@ -0,0 +1,14 @@
+! This is a crazy program, recursive derived types with recursive kind
+! parameters are a terrible idea if they do not converge quickly.
+
+! RUN: bbc -emit-hlfir -o - -I nw %s | FileCheck %s
+
+subroutine foo(x)
+  type t(k)
+    integer, kind :: k
+    type(t(modulo(k+1,2))), pointer :: p
+  end type
+  type(t(1)) :: x
+end subroutine
+! CHECK-LABEL: func.func @_QPfoo(
+! CHECK-SAME: !fir.ref<!fir.type<_QFfooTtK1{p:!fir.box<!fir.ptr<!fir.type<_QFfooTtK0{p:!fir.box<!fir.ptr<!fir.type<_QFfooTtK1>>>}>>>}>>

[clementval]

LGTM!

[clementval]

Do we have an open question on the MLIR discourse about why the MLIR infra is not doing it for us?

[jeanPerier]

Do we have an open question on the MLIR discourse about why the MLIR infra is not doing it for us?

Not yet, I want to gather a few facts here first about what is happening when mlir::ParametricStorageUniquer::getOrCreate is used deep in fir::RecordType::get. It is possible that the issue may lie in how the RecordTypeStorage is defined.

[clementval]

Do we have an open question on the MLIR discourse about why the MLIR infra is not doing it for us?

Not yet, I want to gather a few facts here first about what is happening when mlir::ParametricStorageUniquer::getOrCreate is used deep in fir::RecordType::get. It is possible that the issue may lie in how the RecordTypeStorage is defined.

That makes sense! Thanks for investigating!