[flang][cuda] Fix atmoicxor lowering to accept arrays

[clementval]

The first agrument can be an address of a scalare, an array element or even just the address of the first element of an array. Update lowering to not trigger elemental lowering.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

The first agrument can be an address of a scalare, an array element or even just the address of the first element of an array. Update lowering to not trigger elemental lowering.

---

Full diff: https://github.com/llvm/llvm-project/pull/130331.diff

3 Files Affected:

• (modified) flang/include/flang/Optimizer/Builder/IntrinsicCall.h (+2-1)
• (modified) flang/lib/Optimizer/Builder/IntrinsicCall.cpp (+6-6)
• (modified) flang/test/Lower/CUDA/cuda-device-proc.cuf (+10-3)

diff --git a/flang/include/flang/Optimizer/Builder/IntrinsicCall.h b/flang/include/flang/Optimizer/Builder/IntrinsicCall.h
index c82e5265970c5..778fcf1d93d8b 100644
--- a/flang/include/flang/Optimizer/Builder/IntrinsicCall.h
+++ b/flang/include/flang/Optimizer/Builder/IntrinsicCall.h
@@ -195,7 +195,8 @@ struct IntrinsicLibrary {
   mlir::Value genAtomicMin(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genAtomicOr(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genAtomicSub(mlir::Type, llvm::ArrayRef<mlir::Value>);
-  mlir::Value genAtomicXor(mlir::Type, llvm::ArrayRef<mlir::Value>);
+  fir::ExtendedValue genAtomicXor(mlir::Type,
+                                  llvm::ArrayRef<fir::ExtendedValue>);
   fir::ExtendedValue
       genCommandArgumentCount(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genAsind(mlir::Type, llvm::ArrayRef<mlir::Value>);
diff --git a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
index ede3be074a820..0a70ed3c6bdf1 100644
--- a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
+++ b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
@@ -2827,13 +2827,13 @@ mlir::Value IntrinsicLibrary::genAtomicMin(mlir::Type resultType,
 }
 
 // ATOMICXOR
-mlir::Value IntrinsicLibrary::genAtomicXor(mlir::Type resultType,
-                                           llvm::ArrayRef<mlir::Value> args) {
+fir::ExtendedValue
+IntrinsicLibrary::genAtomicXor(mlir::Type resultType,
+                               llvm::ArrayRef<fir::ExtendedValue> args) {
   assert(args.size() == 2);
-  assert(mlir::isa<mlir::IntegerType>(args[1].getType()));
-
-  mlir::LLVM::AtomicBinOp binOp = mlir::LLVM::AtomicBinOp::_xor;
-  return genAtomBinOp(builder, loc, binOp, args[0], args[1]);
+  mlir::Value arg0 = fir::getBase(args[0]);
+  mlir::Value arg1 = fir::getBase(args[1]);
+  return genAtomBinOp(builder, loc, mlir::LLVM::AtomicBinOp::_xor, arg0, arg1);
 }
 
 // ASSOCIATED
diff --git a/flang/test/Lower/CUDA/cuda-device-proc.cuf b/flang/test/Lower/CUDA/cuda-device-proc.cuf
index 5f39f78f8ecae..5c7f334bae8ba 100644
--- a/flang/test/Lower/CUDA/cuda-device-proc.cuf
+++ b/flang/test/Lower/CUDA/cuda-device-proc.cuf
@@ -150,21 +150,28 @@ end subroutine
 ! CHECK: fir.convert %{{.*}} : (f64) -> i64
 ! CHECK: fir.call @llvm.nvvm.match.any.sync.i64p
 
-attributes(device) subroutine testAtomic()
-  integer :: a, istat, j
+attributes(device) subroutine testAtomic(aa, n)
+  integer :: aa(*)
+  integer, intent(in) :: n
+  integer :: a, istat, j, i
   real :: r
   istat = atomicexch(a,0)
   istat = atomicexch(r, 0.0)
   istat = atomicxor(a, j)
   istat = atomiccas(a, i, 14)
+  do i = 1, n
+    istat = atomicxor(aa, i)
+  end do
 end subroutine
 
-! CHECK-LABEL: func.func @_QPtestatomic()
+! CHECK-LABEL: func.func @_QPtestatomic
 ! CHECK: llvm.atomicrmw xchg %{{.*}}, %c0{{.*}} seq_cst : !llvm.ptr, i32
 ! CHECK: llvm.atomicrmw xchg %{{.*}}, %{{.*}} seq_cst : !llvm.ptr, f32
 ! CHECK: llvm.atomicrmw _xor %{{.*}}, %{{.*}} seq_cst : !llvm.ptr, i32
 ! CHECK: %[[ADDR:.*]] = builtin.unrealized_conversion_cast %{{.*}}#1 : !fir.ref<i32> to !llvm.ptr
 ! CHECK: llvm.cmpxchg %[[ADDR]], %{{.*}}, %c14{{.*}} acq_rel monotonic : !llvm.ptr, i32
+! CHECK: fir.do_loop
+! CHECK: llvm.atomicrmw _xor %{{.*}}, %{{.*}} seq_cst : !llvm.ptr, i32
 
 attributes(device) subroutine testAtomic2()
   integer(8) :: a, i, istat