[flang-commits] [flang] 34ed7db - [Flang] Fix ALLOCATE with MOLD where MOLD is a scalar

[Peter Steinfeld via flang-commits]

Author: Peter Steinfeld
Date: 2023-03-09T06:07:50-08:00
New Revision: 34ed7db9e1f65ef641776b3c5dff594d5183cab3

URL: https://github.com/llvm/llvm-project/commit/34ed7db9e1f65ef641776b3c5dff594d5183cab3
DIFF: https://github.com/llvm/llvm-project/commit/34ed7db9e1f65ef641776b3c5dff594d5183cab3.diff

LOG: [Flang] Fix ALLOCATE with MOLD where MOLD is a scalar

We were failing tests where an ALLOCATE statement that allocated an
array had a non-character scalar MOLD argument.

I fixed this by merging the code for ALLOCATE statements with MOLD and
SOURCE arguments.

Differential Revision: https://reviews.llvm.org/D145418

Added: 
    

Modified: 
    flang/lib/Lower/Allocatable.cpp
    flang/test/Lower/allocate-mold.f90

Removed: 
    


################################################################################
diff  --git a/flang/lib/Lower/Allocatable.cpp b/flang/lib/Lower/Allocatable.cpp
index 44956d65fc65..090d9f25f087 100644
--- a/flang/lib/Lower/Allocatable.cpp
+++ b/flang/lib/Lower/Allocatable.cpp
@@ -371,13 +371,12 @@ class AllocateStmtHelper {
     fir::MutableBoxValue boxAddr =
         genMutableBoxValue(converter, loc, alloc.getAllocObj());
 
-    if (sourceExpr) {
-      genSourceAllocation(alloc, boxAddr);
-    } else if (moldExpr) {
-      genMoldAllocation(alloc, boxAddr);
-    } else {
+    if (sourceExpr)
+      genSourceMoldAllocation(alloc, boxAddr, /*isSource=*/true);
+    else if (moldExpr)
+      genSourceMoldAllocation(alloc, boxAddr, /*isSource=*/false);
+    else
       genSimpleAllocation(alloc, boxAddr);
-    }
   }
 
   static bool lowerBoundsAreOnes(const Allocation &alloc) {
@@ -557,8 +556,10 @@ class AllocateStmtHelper {
     }
   }
 
-  void genSourceAllocation(const Allocation &alloc,
-                           const fir::MutableBoxValue &box) {
+  void genSourceMoldAllocation(const Allocation &alloc,
+                               const fir::MutableBoxValue &box, bool isSource) {
+    fir::ExtendedValue exv = isSource ? sourceExv : moldExv;
+    ;
     // Generate a sequence of runtime calls.
     errorManager.genStatCheck(builder, loc);
     genAllocateObjectInit(box);
@@ -568,24 +569,17 @@ class AllocateStmtHelper {
     // from source for the deferred length parameter.
     if (lenParams.empty() && box.isCharacter() &&
         !box.hasNonDeferredLenParams())
-      lenParams.push_back(fir::factory::readCharLen(builder, loc, sourceExv));
-    if (alloc.type.IsPolymorphic())
-      genRuntimeAllocateApplyMold(builder, loc, box, sourceExv,
+      lenParams.push_back(fir::factory::readCharLen(builder, loc, exv));
+    if (!isSource || alloc.type.IsPolymorphic())
+      genRuntimeAllocateApplyMold(builder, loc, box, exv,
                                   alloc.getSymbol().Rank());
     genSetDeferredLengthParameters(alloc, box);
     genAllocateObjectBounds(alloc, box);
-    mlir::Value stat =
-        genRuntimeAllocateSource(builder, loc, box, sourceExv, errorManager);
-    fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
-    errorManager.assignStat(builder, loc, stat);
-  }
-
-  void genMoldAllocation(const Allocation &alloc,
-                         const fir::MutableBoxValue &box) {
-    genRuntimeAllocateApplyMold(builder, loc, box, moldExv,
-                                alloc.getSymbol().Rank());
-    errorManager.genStatCheck(builder, loc);
-    mlir::Value stat = genRuntimeAllocate(builder, loc, box, errorManager);
+    mlir::Value stat;
+    if (isSource)
+      stat = genRuntimeAllocateSource(builder, loc, box, exv, errorManager);
+    else
+      stat = genRuntimeAllocate(builder, loc, box, errorManager);
     fir::factory::syncMutableBoxFromIRBox(builder, loc, box);
     errorManager.assignStat(builder, loc, stat);
   }

diff  --git a/flang/test/Lower/allocate-mold.f90 b/flang/test/Lower/allocate-mold.f90
index 6a1b0aaf7381..245141d94ca4 100644
--- a/flang/test/Lower/allocate-mold.f90
+++ b/flang/test/Lower/allocate-mold.f90
@@ -17,3 +17,27 @@ subroutine scalar_mold_allocation()
 ! CHECK: %{{.*}} = fir.call @_FortranAAllocatableApplyMold(%[[A_REF_BOX_NONE1]], %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32) -> none
 ! CHECK: %[[A_REF_BOX_NONE2:.*]] = fir.convert %[[A]] : (!fir.ref<!fir.box<!fir.heap<i32>>>) -> !fir.ref<!fir.box<none>>
 ! CHECK: %{{.*}} = fir.call @_FortranAAllocatableAllocate(%[[A_REF_BOX_NONE2]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32
+
+subroutine array_scalar_mold_allocation()
+  real, allocatable :: a(:)
+
+  allocate (a(10), mold=3.0)
+end subroutine array_scalar_mold_allocation
+
+! CHECK-LABEL: func.func @_QParray_scalar_mold_allocation() {
+! CHECK: %[[A:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", uniq_name = "_QFarray_scalar_mold_allocationEa"}
+! CHECK: %[[HEAP_A:.*]] = fir.alloca !fir.heap<!fir.array<?xf32>> {uniq_name = "_QFarray_scalar_mold_allocationEa.addr"}
+! CHECK: %[[EXT0:.*]] = fir.alloca index {uniq_name = "_QFarray_scalar_mold_allocationEa.ext0"}
+! CHECK: %[[ZERO:.*]] = fir.zero_bits !fir.heap<!fir.array<?xf32>>
+! CHECK: fir.store %[[ZERO]] to %[[HEAP_A]] : !fir.ref<!fir.heap<!fir.array<?xf32>>>
+! CHECK: %[[LOADED_A:.*]] = fir.load %[[HEAP_A]] : !fir.ref<!fir.heap<!fir.array<?xf32>>>
+! CHECK: %[[SHAPESHIFT:.*]] = fir.shape_shift {{.*}}, {{.*}} : (index, index) -> !fir.shapeshift<1>
+! CHECK: %[[BOX_SHAPESHIFT:.*]] = fir.embox %[[LOADED_A]](%[[SHAPESHIFT]]) : (!fir.heap<!fir.array<?xf32>>, !fir.shapeshift<1>) -> !fir.box<!fir.heap<!fir.array<?xf32>>>
+! CHECK: fir.store %[[BOX_SHAPESHIFT]] to %[[A]] : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>
+! CHECK: %[[REF_BOX_A0:.*]] = fir.convert %1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %{{.*}} = fir.call @_FortranAAllocatableApplyMold(%[[REF_BOX_A0]], {{.*}}, {{.*}}) fastmath<contract> : (!fir.ref<!fir.box<none>>, !fir.box<none>, i32) -> none
+! CHECK: %[[C10:.*]] = arith.constant 10 : i32
+! CHECK: %[[REF_BOX_A1:.*]] = fir.convert %1 : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %{{.*}} = fir.call @_FortranAAllocatableSetBounds(%[[REF_BOX_A1]], {{.*}},{{.*}}, {{.*}}) fastmath<contract> : (!fir.ref<!fir.box<none>>, i32, i64, i64) -> none
+! CHECK: %[[REF_BOX_A2:.*]] = fir.convert %[[A]] : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> !fir.ref<!fir.box<none>>
+! CHECK: %{{.*}} = fir.call @_FortranAAllocatableAllocate(%[[REF_BOX_A2]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}) {{.*}} : (!fir.ref<!fir.box<none>>, i1, !fir.box<none>, !fir.ref<i8>, i32) -> i32