From 77b9f322787b0e698a70aa6dec9682118016b1f5 Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval@gmail.com>
Date: Mon, 24 Feb 2025 11:04:55 -0800
Subject: [PATCH] [flang][cuda] Fix type mismatch in atomiccas

---
 flang/lib/Optimizer/Builder/IntrinsicCall.cpp | 11 +++++++++--
 flang/test/Lower/CUDA/cuda-device-proc.cuf    | 10 ++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
index 28fbe83defb61..406c443e5b52d 100644
--- a/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
+++ b/flang/lib/Optimizer/Builder/IntrinsicCall.cpp
@@ -2728,15 +2728,22 @@ mlir::Value IntrinsicLibrary::genAtomicOr(mlir::Type resultType,
 mlir::Value IntrinsicLibrary::genAtomicCas(mlir::Type resultType,
                                            llvm::ArrayRef<mlir::Value> args) {
   assert(args.size() == 3);
-  assert(args[1].getType() == args[2].getType());
   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];
+  if (arg1.getType() != arg2.getType()) {
+    // arg1 and arg2 need to have the same type in AtomicCmpXchgOp.
+    arg2 = builder.createConvert(loc, arg1.getType(), arg2);
+  }
+
   auto address =
       builder.create<mlir::UnrealizedConversionCastOp>(loc, llvmPtrTy, args[0])
           .getResult(0);
   auto cmpxchg = builder.create<mlir::LLVM::AtomicCmpXchgOp>(
-      loc, address, args[1], args[2], successOrdering, failureOrdering);
+      loc, address, arg1, arg2, successOrdering, failureOrdering);
   return builder.create<mlir::LLVM::ExtractValueOp>(loc, cmpxchg, 1);
 }
 
diff --git a/flang/test/Lower/CUDA/cuda-device-proc.cuf b/flang/test/Lower/CUDA/cuda-device-proc.cuf
index fcfcc2e537039..6487edeb07004 100644
--- a/flang/test/Lower/CUDA/cuda-device-proc.cuf
+++ b/flang/test/Lower/CUDA/cuda-device-proc.cuf
@@ -162,3 +162,13 @@ end subroutine
 ! 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
+
+attributes(device) subroutine testAtomic2()
+  integer(8) :: a, i, istat
+  istat = atomiccas(a, i, 14)
+end subroutine
+
+! CHECK-LABEL: func.func @_QPtestatomic2()
+! CHECK: %[[VAL:.*]] = fir.convert %c14{{.*}} : (i32) -> i64
+! CHECK: %[[ADDR:.*]] = builtin.unrealized_conversion_cast %{{.*}}#1 : !fir.ref<i64> to !llvm.ptr
+! CHECK: llvm.cmpxchg %{{.*}}, %{{.*}}, %[[VAL]] acq_rel monotonic : !llvm.ptr, i64