[flang-commits] [flang] [flang] Fix invalid iterator erasure in boxed-procedure pass (PR #79830)

[via flang-commits]

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

The code in BoxedProcedureRewrite was erasing the mappings <old type, new type> once "new type" was fully translated. This was done in an invalid way because the map was an llvm::SmallMapVector that do not have stable iterators, and new items were added to the map between the creation of the iterator and its erasure.

It is anyway not needed and expensive to erase the types (see llvm::MapVector note), the cache can be used for the whole pass, which is very beneficial in the context of an apps using "big derived types" (dozens of components/parents).

>From 054ca85fe481f52316a623856971d24df5e32185 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Mon, 29 Jan 2024 06:02:02 -0800
Subject: [PATCH] [flang] Fix invalid iterator erasure in boxed-procedure pass

The code in BoxedProcedureRewrite was erasing the mappings
<old type, new type> once "new type" was fully translated.
This was done in an invalid way because the map was an llvm::SmallMapVector
that do not have stable iterators, and new items were added to the map
between the creation of the iterator and its erasure.

It is anyway not needed and expensive to erase the types (see
llvm::MapVector note), the cache can be used for the whole pass, which
is very beneficial in the context of an apps using "big derived types"
(dozens of components/parents).
---
 flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp | 16 +++++++++++-----
 flang/test/Fir/boxproc-2.fir                   |  6 ++++++
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
index 83333e2bae5b3db..846a78931acba73 100644
--- a/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
+++ b/flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp
@@ -19,7 +19,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
-#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/DenseMap.h"
 
 namespace fir {
 #define GEN_PASS_DEF_BOXEDPROCEDUREPASS
@@ -133,13 +133,16 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
     addConversion([&](RecordType ty) -> mlir::Type {
       if (!needsConversion(ty))
         return ty;
-      if (auto converted = typeInConversion.lookup(ty))
+      if (auto converted = convertedTypes.lookup(ty))
         return converted;
       auto rec = RecordType::get(ty.getContext(),
                                  ty.getName().str() + boxprocSuffix.str());
       if (rec.isFinalized())
         return rec;
-      auto it = typeInConversion.try_emplace(ty, rec);
+      auto it = convertedTypes.try_emplace(ty, rec);
+      if (!it.second) {
+        llvm::errs() << "failed\n" << ty << "\n";
+      }
       std::vector<RecordType::TypePair> ps = ty.getLenParamList();
       std::vector<RecordType::TypePair> cs;
       for (auto t : ty.getTypeList()) {
@@ -149,7 +152,6 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
           cs.emplace_back(t.first, t.second);
       }
       rec.finalize(ps, cs);
-      typeInConversion.erase(it.first);
       return rec;
     });
     addArgumentMaterialization(materializeProcedure);
@@ -170,7 +172,11 @@ class BoxprocTypeRewriter : public mlir::TypeConverter {
 
 private:
   llvm::SmallVector<mlir::Type> visitedTypes;
-  llvm::SmallMapVector<mlir::Type, mlir::Type, 8> typeInConversion;
+  // Map 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, mlir::Type> convertedTypes;
   mlir::Location loc;
 };
 
diff --git a/flang/test/Fir/boxproc-2.fir b/flang/test/Fir/boxproc-2.fir
index b7a5bc69e687e5a..3279bb9bed8a4c4 100644
--- a/flang/test/Fir/boxproc-2.fir
+++ b/flang/test/Fir/boxproc-2.fir
@@ -102,3 +102,9 @@ func.func @box_addr_test(%arg0: !fir.box<!t2>) -> !fir.ref<!t2> {
 }
 // CHECK: func.func @box_addr_test(
 // CHECK:   fir.box_addr %{{.*}} : (!fir.box<!fir.type<t2UnboxProc{i:() -> ()}>>) -> !fir.ref<!fir.type<t2UnboxProc{i:() -> ()}>>
+
+func.func @very_nested_type(%arg0: !fir.type<t0{t01:!fir.type<t01{t02:!fir.type<t02{t3:!fir.type<t3{t4:!fir.type<t4{t5:!fir.type<t5{t6:!fir.type<t6{t7:!fir.type<t7{t8:!fir.type<t8{t9:!fir.type<t9{t10:!fir.type<t10{t11:!fir.type<t11{t12:!fir.type<t12{t13:!fir.type<t13{t14:!fir.type<t14{t15:!fir.type<t15{t16:!fir.type<t16{t17:!fir.type<t17{t18:!fir.type<t18{t19:!fir.type<t19{b:!fir.boxproc<()->()>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>) {
+  return
+}
+// 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:() -> ()}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>)