[flang][cuda] Allow array pointer for atomicexch and atomiccas (#130363)

Author: clementval

flang/include/flang/Optimizer/Builder/IntrinsicCall.h

@@ -187,9 +187,11 @@ struct IntrinsicLibrary {
   mlir::Value genAtanpi(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genAtomicAdd(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genAtomicAnd(mlir::Type, llvm::ArrayRef<mlir::Value>);
-  mlir::Value genAtomicCas(mlir::Type, llvm::ArrayRef<mlir::Value>);
+  fir::ExtendedValue genAtomicCas(mlir::Type,
+                                  llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genAtomicDec(mlir::Type, llvm::ArrayRef<mlir::Value>);
-  mlir::Value genAtomicExch(mlir::Type, llvm::ArrayRef<mlir::Value>);
+  fir::ExtendedValue genAtomicExch(mlir::Type,
+                                   llvm::ArrayRef<fir::ExtendedValue>);
   mlir::Value genAtomicInc(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genAtomicMax(mlir::Type, llvm::ArrayRef<mlir::Value>);
   mlir::Value genAtomicMin(mlir::Type, llvm::ArrayRef<mlir::Value>);

flang/lib/Optimizer/Builder/IntrinsicCall.cpp

@@ -2868,15 +2868,17 @@ mlir::Value IntrinsicLibrary::genAtomicOr(mlir::Type resultType,
 }
 
 // ATOMICCAS
-mlir::Value IntrinsicLibrary::genAtomicCas(mlir::Type resultType,
-                                           llvm::ArrayRef<mlir::Value> args) {
+fir::ExtendedValue
+IntrinsicLibrary::genAtomicCas(mlir::Type resultType,
+                               llvm::ArrayRef<fir::ExtendedValue> args) {
   assert(args.size() == 3);
   auto successOrdering = mlir::LLVM::AtomicOrdering::acq_rel;
   auto failureOrdering = mlir::LLVM::AtomicOrdering::monotonic;
   auto llvmPtrTy = mlir::LLVM::LLVMPointerType::get(resultType.getContext());
 
-  mlir::Value arg1 = args[1];
-  mlir::Value arg2 = args[2];
+  mlir::Value arg0 = fir::getBase(args[0]);
+  mlir::Value arg1 = fir::getBase(args[1]);
+  mlir::Value arg2 = fir::getBase(args[2]);
 
   auto bitCastFloat = [&](mlir::Value arg) -> mlir::Value {
     if (mlir::isa<mlir::Float32Type>(arg.getType()))
@@ -2897,7 +2899,7 @@ mlir::Value IntrinsicLibrary::genAtomicCas(mlir::Type resultType,
   }
 
   auto address =
-      builder.create<mlir::UnrealizedConversionCastOp>(loc, llvmPtrTy, args[0])
+      builder.create<mlir::UnrealizedConversionCastOp>(loc, llvmPtrTy, arg0)
           .getResult(0);
   auto cmpxchg = builder.create<mlir::LLVM::AtomicCmpXchgOp>(
       loc, address, arg1, arg2, successOrdering, failureOrdering);
@@ -2914,13 +2916,16 @@ mlir::Value IntrinsicLibrary::genAtomicDec(mlir::Type resultType,
 }
 
 // ATOMICEXCH
-mlir::Value IntrinsicLibrary::genAtomicExch(mlir::Type resultType,
-                                            llvm::ArrayRef<mlir::Value> args) {
+fir::ExtendedValue
+IntrinsicLibrary::genAtomicExch(mlir::Type resultType,
+                                llvm::ArrayRef<fir::ExtendedValue> args) {
   assert(args.size() == 2);
-  assert(args[1].getType().isIntOrFloat());
+  mlir::Value arg0 = fir::getBase(args[0]);
+  mlir::Value arg1 = fir::getBase(args[1]);
+  assert(arg1.getType().isIntOrFloat());
 
   mlir::LLVM::AtomicBinOp binOp = mlir::LLVM::AtomicBinOp::xchg;
-  return genAtomBinOp(builder, loc, binOp, args[0], args[1]);
+  return genAtomBinOp(builder, loc, binOp, arg0, arg1);
 }
 
 mlir::Value IntrinsicLibrary::genAtomicInc(mlir::Type resultType,

flang/test/Lower/CUDA/cuda-device-proc.cuf

@@ -161,6 +161,8 @@ attributes(device) subroutine testAtomic(aa, n)
   istat = atomiccas(a, i, 14)
   do i = 1, n
     istat = atomicxor(aa, i)
+    istat = atomiccas(aa, i, 14)
+    istat = atomicexch(aa, 0)
   end do
 end subroutine
 
@@ -172,6 +174,9 @@ end subroutine
 ! CHECK: llvm.cmpxchg %[[ADDR]], %{{.*}}, %c14{{.*}} acq_rel monotonic : !llvm.ptr, i32
 ! CHECK: fir.do_loop
 ! CHECK: llvm.atomicrmw _xor %{{.*}}, %{{.*}} seq_cst : !llvm.ptr, i32
+! CHECK: %[[ADDR:.*]] = builtin.unrealized_conversion_cast %{{.*}}#1 : !fir.ref<!fir.array<?xi32>> to !llvm.ptr 
+! CHECK: llvm.cmpxchg %[[ADDR]], %{{.*}}, %c14{{.*}} acq_rel monotonic : !llvm.ptr, i32
+! CHECK: llvm.atomicrmw xchg %{{.*}}, %c0{{.*}} seq_cst : !llvm.ptr, i32
 
 attributes(device) subroutine testAtomic2()
   integer(8) :: a, i, istat