[flang] Fixed TBAA tags for derived types with descriptors. (#148093)

We cannot attach any "data" or "descriptor" tag to accesses
of derived types that contain descriptors, because this
will make them non-aliasing with any generic "data" or "descriptor"
accesses, which is not correct. We have to skip TBAA tags attachment
for such accesses same way we do it for boxes.

Author: vzakhari

flang/lib/Optimizer/Transforms/AddAliasTags.cpp

@@ -79,6 +79,10 @@ class PassState {
   // For the given fir.declare returns the outermost fir.dummy_scope
   // in the current function.
   fir::DummyScopeOp getOutermostScope(fir::DeclareOp declareOp) const;
+  // Returns true, if the given type of a memref of a FirAliasTagOpInterface
+  // operation is a descriptor or contains a descriptor
+  // (e.g. !fir.ref<!fir.type<Derived{f:!fir.box<!fir.heap<f32>>}>>).
+  bool typeReferencesDescriptor(mlir::Type type);
 
 private:
   mlir::DominanceInfo &domInfo;
@@ -95,6 +99,10 @@ class PassState {
   // to the dominance information.
   llvm::DenseMap<mlir::func::FuncOp, llvm::SmallVector<fir::DummyScopeOp, 16>>
       sortedScopeOperations;
+
+  // Local pass cache for derived types that contain descriptor
+  // member(s), to avoid the cost of isRecordWithDescriptorMember().
+  llvm::DenseSet<mlir::Type> typesContainingDescriptors;
 };
 
 // Process fir.dummy_scope operations in the given func:
@@ -145,6 +153,22 @@ fir::DummyScopeOp PassState::getOutermostScope(fir::DeclareOp declareOp) const {
   return nullptr;
 }
 
+bool PassState::typeReferencesDescriptor(mlir::Type type) {
+  type = fir::unwrapAllRefAndSeqType(type);
+  if (mlir::isa<fir::BaseBoxType>(type))
+    return true;
+
+  if (mlir::isa<fir::RecordType>(type)) {
+    if (typesContainingDescriptors.contains(type))
+      return true;
+    if (fir::isRecordWithDescriptorMember(type)) {
+      typesContainingDescriptors.insert(type);
+      return true;
+    }
+  }
+  return false;
+}
+
 class AddAliasTagsPass : public fir::impl::AddAliasTagsBase<AddAliasTagsPass> {
 public:
   void runOnOperation() override;
@@ -200,11 +224,17 @@ void AddAliasTagsPass::runOnAliasInterface(fir::FirAliasTagOpInterface op,
          "load and store only access one address");
   mlir::Value memref = accessedOperands.front();
 
-  // skip boxes. These get an "any descriptor access" tag in TBAABuilder
-  // (CodeGen). I didn't see any speedup from giving each box a separate TBAA
-  // type.
-  if (mlir::isa<fir::BaseBoxType>(fir::unwrapRefType(memref.getType())))
+  // Skip boxes and derived types that contain descriptors.
+  // The box accesses get an "any descriptor access" tag in TBAABuilder
+  // (CodeGen). The derived types accesses get "any access" tag
+  // (because they access both the data and the descriptor(s)).
+  // Note that it would be incorrect to attach any "data" access
+  // tag to the derived type accesses here, because the tags
+  // attached to the descriptor accesses in CodeGen will make
+  // them non-conflicting with any descriptor accesses.
+  if (state.typeReferencesDescriptor(memref.getType()))
     return;
+
   LLVM_DEBUG(llvm::dbgs() << "Analysing " << op << "\n");
 
   const fir::AliasAnalysis::Source &source = state.getSource(memref);

flang/test/Transforms/tbaa-derived-with-descriptor.fir

@@ -0,0 +1,64 @@
+// Test TBAA tags attachment for accesses of derived types
+// that contain descriptors. They cannot be attached
+// with any reasonable "data" tag by fir-add-alias-tags pass,
+// because this will conflict with "any descriptor"
+// tags attached to the descriptor accesses that may be part
+// of these derived type values.
+// Representative Fortran example:
+// subroutine test
+//   type t
+//      real, allocatable :: f
+//   end type t
+//   associate(a => callee([1.0]))
+//     call ext(a(1)%f)
+//   end associate
+// contains
+//   elemental type(t) function callee(x)
+//     real, intent(in) :: x
+//     callee = t(x)
+//   end function callee
+// end subroutine test
+// RUN: fir-opt --fir-add-alias-tags %s | FileCheck %s
+
+func.func @_QPtest() {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %0 = fir.alloca !fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>> {bindc_name = ".tmp.array"}
+  %1 = fir.convert %0 : (!fir.ref<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>) -> !fir.heap<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>
+  %2 = fir.alloca f32 {bindc_name = ".tmp"}
+  %3 = fir.alloca !fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}> {bindc_name = ".result"}
+  %4 = fir.dummy_scope : !fir.dscope
+  %14 = fir.address_of(@_QQro.1xr4.0) : !fir.ref<!fir.array<1xf32>>
+  %15 = fir.shape %c1 : (index) -> !fir.shape<1>
+  %16 = fir.declare %14(%15) {fortran_attrs = #fir.var_attrs<parameter>, uniq_name = "_QQro.1xr4.0"} : (!fir.ref<!fir.array<1xf32>>, !fir.shape<1>) -> !fir.ref<!fir.array<1xf32>>
+  %17 = fir.declare %1(%15) {uniq_name = ".tmp.array"} : (!fir.heap<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>, !fir.shape<1>) -> !fir.heap<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>
+  fir.do_loop %arg0 = %c1 to %c1 step %c1 unordered {
+    %27 = fir.array_coor %16(%15) %arg0 : (!fir.ref<!fir.array<1xf32>>, !fir.shape<1>, index) -> !fir.ref<f32>
+    %28 = fir.declare %2 {uniq_name = ".tmp"} : (!fir.ref<f32>) -> !fir.ref<f32>
+// CHECK: fir.load %{{.*}} {tbaa = {{.*}}} : !fir.ref<f32>
+    %29 = fir.load %27 : !fir.ref<f32>
+// CHECK: fir.store %{{.*}} {tbaa = {{.*}}} : !fir.ref<f32>
+    fir.store %29 to %28 : !fir.ref<f32>
+    %30 = fir.call @_QFtestPcallee(%28) proc_attrs<elemental, pure> fastmath<fast> : (!fir.ref<f32>) -> !fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>
+    fir.save_result %30 to %3 : !fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>, !fir.ref<!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>
+    %31 = fir.declare %3 {uniq_name = ".tmp.func_result"} : (!fir.ref<!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>) -> !fir.ref<!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>
+    %32 = fir.array_coor %17(%15) %arg0 : (!fir.heap<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>, !fir.shape<1>, index) -> !fir.ref<!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>
+// CHECK: fir.load %{{[0-9]+}} : !fir.ref<!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>
+    %33 = fir.load %31 : !fir.ref<!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>
+// CHECK: fir.store %{{.*}} to %{{[0-9]+}} : !fir.ref<!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>
+    fir.store %33 to %32 : !fir.ref<!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>
+  }
+  %18 = fir.convert %17 : (!fir.heap<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>) -> !fir.ref<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>
+  %19 = fir.declare %18(%15) {uniq_name = "_QFtestEa"} : (!fir.ref<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>, !fir.shape<1>) -> !fir.ref<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>
+  %20 = fir.array_coor %19(%15) %c1 : (!fir.ref<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>, !fir.shape<1>, index) -> !fir.ref<!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>
+  %21 = fir.coordinate_of %20, f : (!fir.ref<!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>) -> !fir.ref<!fir.box<!fir.heap<f32>>>
+// CHECK: fir.load %{{[0-9]+}} : !fir.ref<!fir.box<!fir.heap<f32>>>
+  %22 = fir.load %21 : !fir.ref<!fir.box<!fir.heap<f32>>>
+  %23 = fir.box_addr %22 : (!fir.box<!fir.heap<f32>>) -> !fir.heap<f32>
+  %24 = fir.convert %23 : (!fir.heap<f32>) -> !fir.ref<f32>
+  fir.call @_QPext(%24) fastmath<fast> : (!fir.ref<f32>) -> ()
+  %25 = fir.embox %17(%15) : (!fir.heap<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>, !fir.shape<1>) -> !fir.box<!fir.heap<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>>
+  %26 = fir.convert %25 : (!fir.box<!fir.heap<!fir.array<1x!fir.type<_QFtestTt{f:!fir.box<!fir.heap<f32>>}>>>>) -> !fir.box<none>
+  fir.call @_FortranADestroyWithoutFinalization(%26) : (!fir.box<none>) -> ()
+  return
+}