[flang][cuda] Fix type mismatch in atomiccas

Author: clementval

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);
 }
 

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