[flang-commits] [flang] [flang][OpenMP] Fix fir.convert in omp.atomic.update region (PR #138397)

[via flang-commits]

https://github.com/NimishMishra created https://github.com/llvm/llvm-project/pull/138397

Region generation in omp.atomic.update currently emits a direct `fir.convert`. This crashes when the RHS expression involves complex type but the LHS variable is primitive type (say `f32`), since a `fir.convert` from `complex<f32>` to `f32` is emitted, which is illegal. This PR adds a conditional check to emit an additional `ExtractValueOp` in case RHS expression has a complex type.

Fixes https://github.com/llvm/llvm-project/issues/138396

>From e667941c15865bba517fde0e0d0882e6b92c78cb Mon Sep 17 00:00:00 2001
From: Nimish Mishra <neelam.nimish at gmail.com>
Date: Sat, 3 May 2025 18:54:15 +0530
Subject: [PATCH] [flang][OpenMP] Fix fir.convert in omp.atomic.update region

---
 flang/lib/Lower/OpenMP/OpenMP.cpp         | 20 +++++++++++++++++---
 flang/test/Lower/OpenMP/atomic-update.f90 | 18 ++++++++++++++++++
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 47e7c266ff7d3..0910d0003356c 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -2833,9 +2833,23 @@ static void genAtomicUpdateStatement(
     lower::StatementContext atomicStmtCtx;
     mlir::Value rhsExpr = fir::getBase(converter.genExprValue(
         *semantics::GetExpr(assignmentStmtExpr), atomicStmtCtx));
-    mlir::Value convertResult =
-        firOpBuilder.createConvert(currentLocation, varType, rhsExpr);
-    firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, convertResult);
+    mlir::Type exprType = fir::unwrapRefType(rhsExpr.getType());
+    if (fir::isa_complex(exprType) && !fir::isa_complex(varType)) {
+      // Emit an additional `ExtractValueOp` if the expression is of complex
+      // type
+      auto extract = firOpBuilder.create<fir::ExtractValueOp>(
+          currentLocation,
+          mlir::cast<mlir::ComplexType>(exprType).getElementType(), rhsExpr,
+          firOpBuilder.getArrayAttr(
+              firOpBuilder.getIntegerAttr(firOpBuilder.getIndexType(), 0)));
+      mlir::Value convertResult = firOpBuilder.create<fir::ConvertOp>(
+          currentLocation, varType, extract);
+      firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, convertResult);
+    } else {
+      mlir::Value convertResult =
+          firOpBuilder.createConvert(currentLocation, varType, rhsExpr);
+      firOpBuilder.create<mlir::omp::YieldOp>(currentLocation, convertResult);
+    }
     converter.resetExprOverrides();
   }
   firOpBuilder.setInsertionPointAfter(atomicUpdateOp);
diff --git a/flang/test/Lower/OpenMP/atomic-update.f90 b/flang/test/Lower/OpenMP/atomic-update.f90
index 31bf447006930..257ae8fb497ff 100644
--- a/flang/test/Lower/OpenMP/atomic-update.f90
+++ b/flang/test/Lower/OpenMP/atomic-update.f90
@@ -20,6 +20,8 @@ program OmpAtomicUpdate
 !CHECK: %[[VAL_C_DECLARE:.*]]:2 = hlfir.declare %[[VAL_C_ADDRESS]] {{.*}}
 !CHECK: %[[VAL_D_ADDRESS:.*]] = fir.address_of(@_QFEd) : !fir.ref<i32>
 !CHECK: %[[VAL_D_DECLARE:.*]]:2 = hlfir.declare %[[VAL_D_ADDRESS]] {{.}}
+!CHECK: %[[VAL_G_ADDRESS:.*]] = fir.alloca complex<f32> {bindc_name = "g", uniq_name = "_QFEg"}
+!CHECK: %[[VAL_G_DECLARE:.*]]:2 = hlfir.declare %[[VAL_G_ADDRESS]] {uniq_name = "_QFEg"} : (!fir.ref<complex<f32>>) -> (!fir.ref<complex<f32>>, !fir.ref<complex<f32>>)
 !CHECK: %[[VAL_i1_ALLOCA:.*]] = fir.alloca i8 {bindc_name = "i1", uniq_name = "_QFEi1"}
 !CHECK: %[[VAL_i1_DECLARE:.*]]:2 = hlfir.declare %[[VAL_i1_ALLOCA]] {{.*}}
 !CHECK: %[[VAL_c5:.*]] = arith.constant 5 : index
@@ -40,6 +42,7 @@ program OmpAtomicUpdate
     integer, target :: c, d
     integer(1) :: i1
     integer, dimension(5) :: k
+    complex :: g
 
 !CHECK: %[[EMBOX:.*]] = fir.embox %[[VAL_C_DECLARE]]#0 : (!fir.ref<i32>) -> !fir.box<!fir.ptr<i32>>
 !CHECK: fir.store %[[EMBOX]] to %[[VAL_A_DECLARE]]#0 : !fir.ref<!fir.box<!fir.ptr<i32>>>
@@ -200,4 +203,19 @@ program OmpAtomicUpdate
 !CHECK:  }
   !$omp atomic update
     x = x + sum([ (y+2, y=1, z) ])
+
+!CHECK: %[[LOAD:.*]] = fir.load %[[VAL_G_DECLARE]]#0 : !fir.ref<complex<f32>>
+!CHECK: omp.atomic.update %[[VAL_W_DECLARE]]#0 : !fir.ref<i32> {
+!CHECK: ^bb0(%[[ARG:.*]]: i32):
+!CHECK: %[[CVT:.*]] = fir.convert %[[ARG]] : (i32) -> f32
+!CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f32
+!CHECK: %[[UNDEF:.*]] = fir.undefined complex<f32>
+!CHECK: %[[IDX1:.*]] = fir.insert_value %[[UNDEF]], %[[CVT]], [0 : index] : (complex<f32>, f32) -> complex<f32>
+!CHECK: %[[IDX2:.*]] = fir.insert_value %[[IDX1]], %[[CST]], [1 : index] : (complex<f32>, f32) -> complex<f32>
+!CHECK: %[[ADD:.*]] = fir.addc %[[IDX2]], %[[LOAD]] {fastmath = #arith.fastmath<contract>} : complex<f32>
+!CHECK: %[[EXT:.*]] = fir.extract_value %[[ADD]], [0 : index] : (complex<f32>) -> f32
+!CHECK: %[[RESULT:.*]] = fir.convert %[[EXT]] : (f32) -> i32
+!CHECK: omp.yield(%[[RESULT]] : i32)
+  !$omp atomic update
+    w = w + g  
 end program OmpAtomicUpdate