[flang] [llvm] [OpenMP][Flang] Fix atomic operations on complex types (PR #165366)

[Krish Gupta via llvm-commits]

https://github.com/KrxGu updated https://github.com/llvm/llvm-project/pull/165366

>From b0b04fddd8ece561d41bb930e0fd8c91a2f0b98a Mon Sep 17 00:00:00 2001
From: Krish Gupta <krishgupta at Krishs-MacBook-Air.local>
Date: Tue, 28 Oct 2025 15:33:40 +0530
Subject: [PATCH] [OpenMP][Flang] Fix atomic size for complex types

Fixes #165184

Use element type instead of pointer type when computing size for
atomic operations on struct types.
---
 .../test/Lower/OpenMP/atomic-read-complex.f90 | 37 +++++++++++++++++++
 .../Lower/OpenMP/atomic-write-complex.f90     | 37 +++++++++++++++++++
 llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp     | 10 ++---
 3 files changed, 78 insertions(+), 6 deletions(-)
 create mode 100644 flang/test/Lower/OpenMP/atomic-read-complex.f90
 create mode 100644 flang/test/Lower/OpenMP/atomic-write-complex.f90

diff --git a/flang/test/Lower/OpenMP/atomic-read-complex.f90 b/flang/test/Lower/OpenMP/atomic-read-complex.f90
new file mode 100644
index 0000000000000..3720ddac211a3
--- /dev/null
+++ b/flang/test/Lower/OpenMP/atomic-read-complex.f90
@@ -0,0 +1,37 @@
+! Test lowering of atomic read to LLVM IR for complex types.
+! This is a regression test for issue #165184.
+
+! REQUIRES: x86-registered-target || aarch64-registered-target
+
+! RUN: %flang_fc1 -triple x86_64-unknown-linux-gnu -emit-llvm -fopenmp %s -o - | FileCheck %s
+! RUN: %flang_fc1 -triple aarch64-unknown-linux-gnu -emit-llvm -fopenmp %s -o - | FileCheck %s
+
+! Test that atomic read operations with complex types emit the correct
+! size parameter to __atomic_load:
+! - complex(4) (8 bytes total): should call __atomic_load(i64 8, ...)
+! - complex(8) (16 bytes total): should call __atomic_load(i64 16, ...)
+
+program atomic_read_complex
+  implicit none
+
+  ! Test complex(4) - single precision (8 bytes)
+  complex(4) :: c41, c42
+  ! Test complex(8) - double precision (16 bytes)
+  complex(8) :: c81, c82
+  
+  c42 = (1.0_4, 1.0_4)
+  c82 = (1.0_8, 1.0_8)
+
+  ! CHECK-LABEL: define {{.*}} @_QQmain
+
+  ! Single precision complex: 8 bytes
+  ! CHECK: call void @__atomic_load(i64 8, ptr {{.*}}, ptr {{.*}}, i32 {{.*}})
+!$omp atomic read
+  c41 = c42
+  
+  ! Double precision complex: 16 bytes (this was broken before the fix)
+  ! CHECK: call void @__atomic_load(i64 16, ptr {{.*}}, ptr {{.*}}, i32 {{.*}})
+!$omp atomic read
+  c81 = c82
+
+end program atomic_read_complex
diff --git a/flang/test/Lower/OpenMP/atomic-write-complex.f90 b/flang/test/Lower/OpenMP/atomic-write-complex.f90
new file mode 100644
index 0000000000000..7fb96762e3788
--- /dev/null
+++ b/flang/test/Lower/OpenMP/atomic-write-complex.f90
@@ -0,0 +1,37 @@
+! Test lowering of atomic write to LLVM IR for complex types.
+! This is a regression test for issue #165184.
+
+! REQUIRES: x86-registered-target || aarch64-registered-target
+
+! RUN: %flang_fc1 -triple x86_64-unknown-linux-gnu -emit-llvm -fopenmp %s -o - | FileCheck %s
+! RUN: %flang_fc1 -triple aarch64-unknown-linux-gnu -emit-llvm -fopenmp %s -o - | FileCheck %s
+
+! Test that atomic write operations with complex types emit the correct
+! size parameter to __atomic_store:
+! - complex(4) (8 bytes total): should call __atomic_store(i64 8, ...)
+! - complex(8) (16 bytes total): should call __atomic_store(i64 16, ...)
+
+program atomic_write_complex
+  implicit none
+
+  ! Test complex(4) - single precision (8 bytes)
+  complex(4) :: c41, c42
+  ! Test complex(8) - double precision (16 bytes)  
+  complex(8) :: c81, c82
+  
+  c42 = (1.0_4, 1.0_4)
+  c82 = (1.0_8, 1.0_8)
+
+  ! CHECK-LABEL: define {{.*}} @_QQmain
+  
+  ! Single precision complex: 8 bytes
+  ! CHECK: call void @__atomic_store(i64 8, ptr {{.*}}, ptr {{.*}}, i32 {{.*}})
+!$omp atomic write
+  c41 = c42
+  
+  ! Double precision complex: 16 bytes (this was broken before the fix)
+  ! CHECK: call void @__atomic_store(i64 16, ptr {{.*}}, ptr {{.*}}, i32 {{.*}})
+!$omp atomic write
+  c81 = c82
+
+end program atomic_write_complex
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 286ed039b1214..c71b4a1b6eaa1 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -9338,9 +9338,8 @@ OpenMPIRBuilder::createAtomicRead(const LocationDescription &Loc,
     // target does not support `atomicrmw` of the size of the struct
     LoadInst *OldVal = Builder.CreateLoad(XElemTy, X.Var, "omp.atomic.read");
     OldVal->setAtomic(AO);
-    const DataLayout &LoadDL = OldVal->getModule()->getDataLayout();
-    unsigned LoadSize =
-        LoadDL.getTypeStoreSize(OldVal->getPointerOperand()->getType());
+    const DataLayout &DL = OldVal->getModule()->getDataLayout();
+    unsigned LoadSize = DL.getTypeStoreSize(XElemTy);
     OpenMPIRBuilder::AtomicInfo atomicInfo(
         &Builder, XElemTy, LoadSize * 8, LoadSize * 8, OldVal->getAlign(),
         OldVal->getAlign(), true /* UseLibcall */, AllocaIP, X.Var);
@@ -9384,9 +9383,8 @@ OpenMPIRBuilder::createAtomicWrite(const LocationDescription &Loc,
     XSt->setAtomic(AO);
   } else if (XElemTy->isStructTy()) {
     LoadInst *OldVal = Builder.CreateLoad(XElemTy, X.Var, "omp.atomic.read");
-    const DataLayout &LoadDL = OldVal->getModule()->getDataLayout();
-    unsigned LoadSize =
-        LoadDL.getTypeStoreSize(OldVal->getPointerOperand()->getType());
+    const DataLayout &DL = OldVal->getModule()->getDataLayout();
+    unsigned LoadSize = DL.getTypeStoreSize(XElemTy);
     OpenMPIRBuilder::AtomicInfo atomicInfo(
         &Builder, XElemTy, LoadSize * 8, LoadSize * 8, OldVal->getAlign(),
         OldVal->getAlign(), true /* UseLibcall */, AllocaIP, X.Var);