[flang-commits] [flang] [flang] Retrieve shape from selector when generating assoc sym type (PR #137117)

[Kareem Ergawy via flang-commits]

https://github.com/ergawy updated https://github.com/llvm/llvm-project/pull/137117

>From c559af9e86c7ecc832a3b6bab7ef24c0e636d2d7 Mon Sep 17 00:00:00 2001
From: ergawy <kareem.ergawy at amd.com>
Date: Thu, 24 Apr 2025 00:34:44 -0500
Subject: [PATCH] [flang] Retrieve shape from selector when generating assoc
 sym type

This PR extends `genSymbolType` so that the type of an associating symbol
carries the shape of the selector expression, if any. This is a fix for
a bug that triggered when an associating symbol is used in a locality
specifier. For example, given the following input:

```fortran
  associate(a => aa(4:))
    do concurrent (i = 4:11) local(a)
      a(i) = 0
    end do
  end associate
```
before the changes in the PR, flang would assert that we are casting
between incompatible types. The issue happened since for the associating
symbol (`a`), flang generated its type as `f32` rather than
`!fir.array<8xf32>` as it should be in this case.
---
 flang/lib/Lower/ConvertType.cpp               | 17 ++++++++++++++
 .../do_concurrent_local_assoc_entity.f90      | 22 +++++++++++++++++++
 2 files changed, 39 insertions(+)
 create mode 100644 flang/test/Lower/do_concurrent_local_assoc_entity.f90

diff --git a/flang/lib/Lower/ConvertType.cpp b/flang/lib/Lower/ConvertType.cpp
index d45f9e7c0bf1b..875bdba6cc6ba 100644
--- a/flang/lib/Lower/ConvertType.cpp
+++ b/flang/lib/Lower/ConvertType.cpp
@@ -279,6 +279,23 @@ struct TypeBuilderImpl {
     bool isPolymorphic = (Fortran::semantics::IsPolymorphic(symbol) ||
                           Fortran::semantics::IsUnlimitedPolymorphic(symbol)) &&
                          !Fortran::semantics::IsAssumedType(symbol);
+    if (const auto *assocDetails =
+            ultimate.detailsIf<Fortran::semantics::AssocEntityDetails>()) {
+      const auto &selector = assocDetails->expr();
+
+      if (selector && selector->Rank() > 0) {
+        auto shapeExpr = Fortran::evaluate::GetShape(
+            converter.getFoldingContext(), selector);
+
+        fir::SequenceType::Shape shape;
+        // If there is no shapExpr, this is an assumed-rank, and the empty shape
+        // will build the desired fir.array<*:T> type.
+        if (shapeExpr)
+          translateShape(shape, std::move(*shapeExpr));
+        ty = fir::SequenceType::get(shape, ty);
+      }
+    }
+
     if (ultimate.IsObjectArray()) {
       auto shapeExpr =
           Fortran::evaluate::GetShape(converter.getFoldingContext(), ultimate);
diff --git a/flang/test/Lower/do_concurrent_local_assoc_entity.f90 b/flang/test/Lower/do_concurrent_local_assoc_entity.f90
new file mode 100644
index 0000000000000..ca16ecaa5c137
--- /dev/null
+++ b/flang/test/Lower/do_concurrent_local_assoc_entity.f90
@@ -0,0 +1,22 @@
+! RUN: %flang_fc1 -emit-hlfir -o - %s | FileCheck %s
+
+subroutine local_assoc
+  implicit none
+  integer  i
+  real, dimension(2:11) :: aa
+
+  associate(a => aa(4:))
+    do concurrent (i = 4:11) local(a)
+      a(i) = 0
+    end do
+  end associate
+end subroutine local_assoc
+
+! CHECK: %[[C8:.*]] = arith.constant 8 : index
+
+! CHECK: fir.do_loop {{.*}} unordered {
+! CHECK:   %[[LOCAL_ALLOC:.*]] = fir.alloca !fir.array<8xf32> {bindc_name = "a", pinned, uniq_name = "{{.*}}local_assocEa"}
+! CHECK:   %[[LOCAL_SHAPE:.*]] = fir.shape %[[C8]] :
+! CHECK:   %[[LOCAL_DECL:.*]]:2 = hlfir.declare %[[LOCAL_ALLOC]](%[[LOCAL_SHAPE]])
+! CHECK:   hlfir.designate %[[LOCAL_DECL]]#0 (%{{.*}})
+! CHECK: }