[flang-commits] [flang] [flang][NFC] Cache derived type translation in lowering (PR #80179)

[via flang-commits]

https://github.com/jeanPerier created https://github.com/llvm/llvm-project/pull/80179

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.

>From 0b397900347c4bb6c9c5f095888bb2da40759a83 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Wed, 31 Jan 2024 08:04:36 -0800
Subject: [PATCH] [flang][NFC] Cache derived type translation in lowering

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 showed that the current caching was incorrect, the type
symbol is the same for kind parametrized derived types regardless of the
kind parameters while instance with different kinds should lower to
different MLIR types. See added test.
Using the type scope fixes the problem.
---
 flang/include/flang/Lower/AbstractConverter.h |  3 ++-
 flang/lib/Lower/ConvertType.cpp               | 26 +++++++------------
 .../Lower/derived-types-kind-params-2.f90     | 14 ++++++++++
 3 files changed, 25 insertions(+), 18 deletions(-)
 create mode 100644 flang/test/Lower/derived-types-kind-params-2.f90

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>>>}>>>}>>