From f3bf844d2ff0a2984ca9bf976014decc0241d2b4 Mon Sep 17 00:00:00 2001
From: Asher Mancinelli <ashermancinelli@gmail.com>
Date: Wed, 16 Apr 2025 17:15:15 -0700
Subject: [PATCH] [flang] Unwrap sequence types when checking for descriptor
 members (#136039)

The TBAA generation gives conservative TBAA metadata when handling an
access of a record type with a descriptor member, since the access may
be a regular data access OR another descriptor. Array members were being
incorrectly identified as non-descriptor-members, and were giving
incorrect TBAA metadata which led to bugs showing up in the optimizer
when LLVM encountered mismatching TBAA.

`fir::isRecordWithDescriptorMember` now unwraps sequence types before
checking for descriptor members.
---
 flang/lib/Optimizer/Dialect/FIRType.cpp |  1 +
 flang/test/Fir/tbaa-codegen-records.fir | 30 +++++++++++++++++++++++++
 2 files changed, 31 insertions(+)
 create mode 100644 flang/test/Fir/tbaa-codegen-records.fir

diff --git a/flang/lib/Optimizer/Dialect/FIRType.cpp b/flang/lib/Optimizer/Dialect/FIRType.cpp
index 4d40c2618a0d..b76856d72a01 100644
--- a/flang/lib/Optimizer/Dialect/FIRType.cpp
+++ b/flang/lib/Optimizer/Dialect/FIRType.cpp
@@ -435,6 +435,7 @@ bool isRecordWithDescriptorMember(mlir::Type ty) {
   ty = unwrapSequenceType(ty);
   if (auto recTy = mlir::dyn_cast<fir::RecordType>(ty))
     for (auto [field, memTy] : recTy.getTypeList()) {
+      memTy = unwrapSequenceType(memTy);
       if (mlir::isa<fir::BaseBoxType>(memTy))
         return true;
       if (mlir::isa<fir::RecordType>(memTy) &&
diff --git a/flang/test/Fir/tbaa-codegen-records.fir b/flang/test/Fir/tbaa-codegen-records.fir
new file mode 100644
index 000000000000..336354098f0f
--- /dev/null
+++ b/flang/test/Fir/tbaa-codegen-records.fir
@@ -0,0 +1,30 @@
+// RUN: fir-opt --split-input-file --pass-pipeline="builtin.module(fir-to-llvm-ir{apply-tbaa=true})" %s | FileCheck %s
+
+// Ensure that records with array members are identified as having descriptor members,
+// as reflected by the TBAA metadata.
+
+func.func @record_array_member(%arg0 : !fir.ref<!fir.type<_QFTt2{y:!fir.array<1x!fir.type<_QFTt{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>}>>) {
+  %c0_i64 = arith.constant 0 : i64
+  %c1 = arith.constant 1 : index
+  %1 = fir.alloca !fir.type<_QFTt2{y:!fir.array<1x!fir.type<_QFTt{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>}> {bindc_name = "z", uniq_name = "_QFB1Ez"}
+  fir.copy %arg0 to %1 no_overlap : !fir.ref<!fir.type<_QFTt2{y:!fir.array<1x!fir.type<_QFTt{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>}>>, !fir.ref<!fir.type<_QFTt2{y:!fir.array<1x!fir.type<_QFTt{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>}>>
+  %3 = fir.coordinate_of %1, y : (!fir.ref<!fir.type<_QFTt2{y:!fir.array<1x!fir.type<_QFTt{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>}>>) -> !fir.ref<!fir.array<1x!fir.type<_QFTt{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>>
+  %4 = fircg.ext_array_coor %3(%c1)<%c1> : (!fir.ref<!fir.array<1x!fir.type<_QFTt{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>>, index, index) -> !fir.ref<!fir.type<_QFTt{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>
+  %5 = fir.coordinate_of %4, x : (!fir.ref<!fir.type<_QFTt{x:!fir.box<!fir.heap<!fir.array<?xi32>>>}>>) -> !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  %6 = fir.load %5 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xi32>>>>
+  return
+}
+
+// CHECK: #[[$ATTR_0:.+]] = #llvm.tbaa_root<id = "Flang function root record_array_member">
+// CHECK: #[[$ATTR_1:.+]] = #llvm.tbaa_type_desc<id = "any access", members = {<#tbaa_root, 0>}>
+// CHECK: #[[$ATTR_2:.+]] = #llvm.tbaa_tag<base_type = #tbaa_type_desc, access_type = #tbaa_type_desc, offset = 0>
+// CHECK: #[[$ATTR_3:.+]] = #llvm.tbaa_type_desc<id = "descriptor member", members = {<#tbaa_type_desc, 0>}>
+// CHECK: #[[$ATTR_4:.+]] = #llvm.tbaa_tag<base_type = #tbaa_type_desc1, access_type = #tbaa_type_desc1, offset = 0>
+
+// CHECK-LABEL:   llvm.func @record_array_member(
+// CHECK-SAME:                                   %[[ARG0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: !llvm.ptr) {
+// CHECK:           %[[X_VAL:.*]] = llvm.alloca %{{.+}} x !llvm.struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)> {alignment = 8 : i64} : (i32) -> !llvm.ptr
+// CHECK:           %[[Z_VAL:.*]] = llvm.alloca %{{.+}} x !llvm.struct<"_QFTt2", (array<1 x struct<"_QFTt", (struct<(ptr, i64, i32, i8, i8, i8, i8, array<1 x array<3 x i64>>)>)>>)> {bindc_name = "z"} : (i64) -> !llvm.ptr
+// CHECK:           "llvm.intr.memcpy"(%[[Z_VAL]], %[[ARG0]], %{{.+}}) <{isVolatile = false, tbaa = [#[[$ATTR_2]]]}> : (!llvm.ptr, !llvm.ptr, i64) -> ()
+
+// CHECK:           "llvm.intr.memcpy"(%[[X_VAL]], %{{.+}}, %{{.+}}) <{isVolatile = false, tbaa = [#[[$ATTR_4]]]}> : (!llvm.ptr, !llvm.ptr, i32) -> ()