[flang] Fix invalid iterator erasure in boxed-procedure pass

[jeanPerier]

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).

[llvmbot]

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

@llvm/pr-subscribers-flang-codegen

Author: None (jeanPerier)

Changes

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).

---

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

2 Files Affected:

• (modified) flang/lib/Optimizer/CodeGen/BoxedProcedure.cpp (+11-5)
• (modified) flang/test/Fir/boxproc-2.fir (+6)

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:() -> ()}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>}>)