From c099ca4e45dbccd58dfa1964fa1f304c6055958d Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier@nvidia.com>
Date: Wed, 2 Feb 2022 09:21:44 +0100
Subject: [PATCH] [flang][optimizer] support aggregate types inside tuple and
 record type

This patch allows:
 - fir.box type to be a member of tuple<> or fir.type<> types,
 - tuple<> type to be a member of tuple<> type.

When a fir.box types are nested in tuple<> or fir.type<>, it is translated
to the struct type of a Fortran runtime descriptor, and not a
pointer to a descriptor. This is because the fir.box is owned by the tuple
or fir.type.

FIR type translation was also flattening nested tuple while lowering to LLVM
dialect types. There does not seem to be a deep reason for doing that
and doing it causes issues in fir.coordinate_of generated on such tuple
(a fir.coordinate_of getting tuple<B, C> in tuple<A, tuple<B, C>>
ended-up lowered to an LLVM GEP getting B).

Differential Revision: https://reviews.llvm.org/D118701
---
 flang/lib/Optimizer/CodeGen/TypeConverter.h | 27 +++++++++++++++++----
 flang/lib/Optimizer/Dialect/FIRType.cpp     | 10 ++++----
 flang/test/Fir/types-to-llvm.fir            | 18 ++++++++++++++
 3 files changed, 45 insertions(+), 10 deletions(-)

diff --git a/flang/lib/Optimizer/CodeGen/TypeConverter.h b/flang/lib/Optimizer/CodeGen/TypeConverter.h
index 478c2592b10b..38105e18e62a 100644
--- a/flang/lib/Optimizer/CodeGen/TypeConverter.h
+++ b/flang/lib/Optimizer/CodeGen/TypeConverter.h
@@ -111,11 +111,15 @@ public:
     });
     addConversion([&](mlir::TupleType tuple) {
       LLVM_DEBUG(llvm::dbgs() << "type convert: " << tuple << '\n');
-      llvm::SmallVector<mlir::Type> inMembers;
-      tuple.getFlattenedTypes(inMembers);
       llvm::SmallVector<mlir::Type> members;
-      for (auto mem : inMembers)
-        members.push_back(convertType(mem).cast<mlir::Type>());
+      for (auto mem : tuple.getTypes()) {
+        // Prevent fir.box from degenerating to a pointer to a descriptor in the
+        // context of a tuple type.
+        if (auto box = mem.dyn_cast<fir::BoxType>())
+          members.push_back(convertBoxTypeAsStruct(box));
+        else
+          members.push_back(convertType(mem).cast<mlir::Type>());
+      }
       return mlir::LLVM::LLVMStructType::getLiteral(&getContext(), members,
                                                     /*isPacked=*/false);
     });
@@ -140,7 +144,12 @@ public:
     }
     llvm::SmallVector<mlir::Type> members;
     for (auto mem : derived.getTypeList()) {
-      members.push_back(convertType(mem.second).cast<mlir::Type>());
+      // Prevent fir.box from degenerating to a pointer to a descriptor in the
+      // context of a record type.
+      if (auto box = mem.second.dyn_cast<fir::BoxType>())
+        members.push_back(convertBoxTypeAsStruct(box));
+      else
+        members.push_back(convertType(mem.second).cast<mlir::Type>());
     }
     if (mlir::failed(st.setBody(members, /*isPacked=*/false)))
       return failure();
@@ -227,6 +236,14 @@ public:
                                                /*isPacked=*/false));
   }
 
+  /// Convert fir.box type to the corresponding llvm struct type instead of a
+  /// pointer to this struct type.
+  mlir::Type convertBoxTypeAsStruct(BoxType box) {
+    return convertBoxType(box)
+        .cast<mlir::LLVM::LLVMPointerType>()
+        .getElementType();
+  }
+
   unsigned characterBitsize(fir::CharacterType charTy) {
     return kindMapping.getCharacterBitsize(charTy.getFKind());
   }
diff --git a/flang/lib/Optimizer/Dialect/FIRType.cpp b/flang/lib/Optimizer/Dialect/FIRType.cpp
index 7095af8a7170..4126c7f3c525 100644
--- a/flang/lib/Optimizer/Dialect/FIRType.cpp
+++ b/flang/lib/Optimizer/Dialect/FIRType.cpp
@@ -61,11 +61,11 @@ static bool isaIntegerType(mlir::Type ty) {
 }
 
 bool verifyRecordMemberType(mlir::Type ty) {
-  return !(ty.isa<BoxType>() || ty.isa<BoxCharType>() ||
-           ty.isa<BoxProcType>() || ty.isa<ShapeType>() ||
-           ty.isa<ShapeShiftType>() || ty.isa<ShiftType>() ||
-           ty.isa<SliceType>() || ty.isa<FieldType>() || ty.isa<LenType>() ||
-           ty.isa<ReferenceType>() || ty.isa<TypeDescType>());
+  return !(ty.isa<BoxCharType>() || ty.isa<BoxProcType>() ||
+           ty.isa<ShapeType>() || ty.isa<ShapeShiftType>() ||
+           ty.isa<ShiftType>() || ty.isa<SliceType>() || ty.isa<FieldType>() ||
+           ty.isa<LenType>() || ty.isa<ReferenceType>() ||
+           ty.isa<TypeDescType>());
 }
 
 bool verifySameLists(llvm::ArrayRef<RecordType::TypePair> a1,
diff --git a/flang/test/Fir/types-to-llvm.fir b/flang/test/Fir/types-to-llvm.fir
index 599c7924cc73..f47410dd4cc5 100644
--- a/flang/test/Fir/types-to-llvm.fir
+++ b/flang/test/Fir/types-to-llvm.fir
@@ -400,3 +400,21 @@ func private @foo2(%arg : !fir.tdesc<f32>)
 func private @foo3(%arg : !fir.tdesc<!fir.type<derived7{f1:f32,f2:f32}>>)
 // CHECK-LABEL: foo3
 // CHECK-SAME:  !llvm.ptr<i8>
+
+// -----
+
+// Test nested tuple types
+func private @foo0(%arg0: tuple<i64, tuple<f32, i64>>)
+// CHECK-LABEL: foo0
+// CHECK-SAME: !llvm.struct<(i64, struct<(f32, i64)>)>
+
+// -----
+
+// Test that fir.box inside tuple and derived type are lowered to struct type.
+func private @foo0(%arg0: tuple<i64, !fir.box<i32>>)
+// CHECK-LABEL: foo0
+// CHECK-SAME: !llvm.struct<(i64, struct<(ptr<i32>, i{{.*}})>)>
+
+func private @foo1(%arg0: !fir.type<derived8{a:i64,b:!fir.box<i32>}>)
+// CHECK-LABEL: foo1
+// CHECK-SAME: !llvm.struct<"derived8", (i64, struct<(ptr<i32>, i{{.*}})>)>