[flang] TBAA for memory accesses of derived type values.

[vzakhari]

Since HLFIR bufferization can introduce shallow copies of derived
type values we have to be careful not to treat these load/store
operations as data-only-accesses. If a derived type has descriptor
members, we attach any-access tag now.

[vzakhari]

This is needed after #68040

[llvmbot]

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

@llvm/pr-subscribers-flang-codegen

Changes

Since HLFIR bufferization can introduce shallow copies of derived
type values we have to be careful not to treat these load/store
operations as data-only-accesses. If a derived type has descriptor
members, we attach any-access tag now.

---

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

5 Files Affected:

• (modified) flang/include/flang/Optimizer/CodeGen/TBAABuilder.h (+3)
• (modified) flang/include/flang/Optimizer/Dialect/FIRType.h (+3)
• (modified) flang/lib/Optimizer/CodeGen/TBAABuilder.cpp (+16-2)
• (modified) flang/lib/Optimizer/Dialect/FIRType.cpp (+11)
• (modified) flang/test/Fir/tbaa.fir (+17)

diff --git a/flang/include/flang/Optimizer/CodeGen/TBAABuilder.h b/flang/include/flang/Optimizer/CodeGen/TBAABuilder.h
index b94782bb9ea9845..3aeb124f911867e 100644
--- a/flang/include/flang/Optimizer/CodeGen/TBAABuilder.h
+++ b/flang/include/flang/Optimizer/CodeGen/TBAABuilder.h
@@ -188,6 +188,9 @@ class TBAABuilder {
   // Returns TBAATagAttr representing access tag:
   //   < <any data access>, <any data access>, 0 >
   mlir::LLVM::TBAATagAttr getAnyDataAccessTag();
+  // Returns TBAATagAttr representing access tag:
+  //   < <any access>, <any access>, 0 >
+  mlir::LLVM::TBAATagAttr getAnyAccessTag();
 
   // Returns TBAATagAttr representing access tag described by the base and
   // access FIR types and the LLVM::GepOp representing the access in terms of
diff --git a/flang/include/flang/Optimizer/Dialect/FIRType.h b/flang/include/flang/Optimizer/Dialect/FIRType.h
index 77807ea2a308c67..df078fc0f8598f2 100644
--- a/flang/include/flang/Optimizer/Dialect/FIRType.h
+++ b/flang/include/flang/Optimizer/Dialect/FIRType.h
@@ -349,6 +349,9 @@ mlir::Type unwrapInnerType(mlir::Type ty);
 /// Return true iff `ty` is a RecordType with members that are allocatable.
 bool isRecordWithAllocatableMember(mlir::Type ty);
 
+/// Return true iff `ty` is a RecordType with members that are descriptors.
+bool isRecordWithDescriptorMember(mlir::Type ty);
+
 /// Return true iff `ty` is a RecordType with type parameters.
 inline bool isRecordWithTypeParameters(mlir::Type ty) {
   if (auto recTy = ty.dyn_cast_or_null<fir::RecordType>())
diff --git a/flang/lib/Optimizer/CodeGen/TBAABuilder.cpp b/flang/lib/Optimizer/CodeGen/TBAABuilder.cpp
index 04259c0e4c72133..ae5167fc1941320 100644
--- a/flang/lib/Optimizer/CodeGen/TBAABuilder.cpp
+++ b/flang/lib/Optimizer/CodeGen/TBAABuilder.cpp
@@ -92,6 +92,10 @@ TBAATagAttr TBAABuilder::getDataAccessTag(Type baseFIRType, Type accessFIRType,
   return getAnyDataAccessTag();
 }
 
+TBAATagAttr TBAABuilder::getAnyAccessTag() {
+  return getAccessTag(anyAccessTypeDesc, anyAccessTypeDesc, /*offset=*/0);
+}
+
 void TBAABuilder::attachTBAATag(AliasAnalysisOpInterface op, Type baseFIRType,
                                 Type accessFIRType, GEPOp gep) {
   if (!enableTBAA)
@@ -106,10 +110,20 @@ void TBAABuilder::attachTBAATag(AliasAnalysisOpInterface op, Type baseFIRType,
                           << "\n");
 
   TBAATagAttr tbaaTagSym;
-  if (baseFIRType.isa<fir::BaseBoxType>())
+  if (auto derivedTy = mlir::dyn_cast<fir::RecordType>(baseFIRType)) {
+    // A memory access that addresses an aggregate that contains
+    // a mix of data members and descriptor members may alias
+    // with both data and descriptor accesses.
+    // Conservatively set any-access tag if there is any descriptor member.
+    if (fir::isRecordWithDescriptorMember(derivedTy))
+      tbaaTagSym = getAnyAccessTag();
+    else
+      tbaaTagSym = getDataAccessTag(baseFIRType, accessFIRType, gep);
+  } else if (baseFIRType.isa<fir::BaseBoxType>()) {
     tbaaTagSym = getBoxAccessTag(baseFIRType, accessFIRType, gep);
-  else
+  } else {
     tbaaTagSym = getDataAccessTag(baseFIRType, accessFIRType, gep);
+  }
 
   if (!tbaaTagSym)
     return;
diff --git a/flang/lib/Optimizer/Dialect/FIRType.cpp b/flang/lib/Optimizer/Dialect/FIRType.cpp
index dd79aa27645452f..48aec4444f6952a 100644
--- a/flang/lib/Optimizer/Dialect/FIRType.cpp
+++ b/flang/lib/Optimizer/Dialect/FIRType.cpp
@@ -360,6 +360,17 @@ bool isRecordWithAllocatableMember(mlir::Type ty) {
   return false;
 }
 
+bool isRecordWithDescriptorMember(mlir::Type ty) {
+  if (auto recTy = ty.dyn_cast<fir::RecordType>())
+    for (auto [field, memTy] : recTy.getTypeList()) {
+      if (mlir::isa<fir::BaseBoxType>(memTy))
+        return true;
+      if (memTy.isa<fir::RecordType>() && isRecordWithDescriptorMember(memTy))
+        return true;
+    }
+  return false;
+}
+
 mlir::Type unwrapAllRefAndSeqType(mlir::Type ty) {
   while (true) {
     mlir::Type nt = unwrapSequenceType(unwrapRefType(ty));
diff --git a/flang/test/Fir/tbaa.fir b/flang/test/Fir/tbaa.fir
index f8a52b2e98db0cc..0be25a248481441 100644
--- a/flang/test/Fir/tbaa.fir
+++ b/flang/test/Fir/tbaa.fir
@@ -351,3 +351,20 @@ func.func @tbaa(%arg0: !fir.box<!fir.array<?xi32>>) {
 // CHECK:           %[[VAL_14:.*]] = llvm.bitcast %[[VAL_13]] : !llvm.ptr<i8> to !llvm.ptr<i32>
 // CHECK:           llvm.return
 // CHECK:         }
+
+// -----
+
+// Check that the aggregate load/store with a descriptor member
+// is mapped to any-access.
+// CHECK-DAG:     #[[ROOT:.*]] = #llvm.tbaa_root<id = "Flang Type TBAA Root">
+// CHECK-DAG:     #[[ANYACC:.*]] = #llvm.tbaa_type_desc<id = "any access", members = {<#[[ROOT]], 0>}>
+// CHECK-DAG:     #[[$ANYT:.*]] = #llvm.tbaa_tag<base_type = #[[ANYACC]], access_type = #[[ANYACC]], offset = 0>
+
+func.func @tbaa(%arg0: !fir.ref<!fir.type<_QMtypesTt{x:!fir.box<!fir.heap<f32>>}>>, %arg1: !fir.ref<!fir.type<_QMtypesTt{x:!fir.box<!fir.heap<f32>>}>>) {
+  %0 = fir.load %arg0 : !fir.ref<!fir.type<_QMtypesTt{x:!fir.box<!fir.heap<f32>>}>>
+  fir.store %0 to %arg1 : !fir.ref<!fir.type<_QMtypesTt{x:!fir.box<!fir.heap<f32>>}>>
+  return
+}
+// CHECK-LABEL:   llvm.func @tbaa(
+// CHECK: llvm.load{{.*}}{tbaa = [#[[$ANYT]]]}
+// CHECK: llvm.store{{.*}}{tbaa = [#[[$ANYT]]]}

[jeanPerier]

Nit, LGTM otherwise but please wait for @tblah feedback since he is working on TBAA.

[jeanPerier]

Could also be a fir.array<fir.type<... {fir.box}>>, so you should likely unwrap any SequenceType here.

[vzakhari]

Good catch! Thank you, Jean!

[tblah]

Looks good to me, thanks!