[flang][openacc] Convert rhs expr to the lhs type on atomic read/write

[clementval]

In some cases the rhs expression scalar type is not the same as the lhs type. A convert op is needed before the acc.atomic.read or acc.atomic.write operation to fit with the requirements of the operations.

[llvmbot]

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

@llvm/pr-subscribers-openacc

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

Changes

In some cases the rhs expression scalar type is not the same as the lhs type. A convert op is needed before the acc.atomic.read or acc.atomic.write operation to fit with the requirements of the operations.

---

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

3 Files Affected:

• (modified) flang/lib/Lower/DirectivesCommon.h (+40)
• (modified) flang/test/Lower/OpenACC/acc-atomic-capture.f90 (+25)
• (modified) flang/test/Lower/OpenACC/acc-atomic-read.f90 (+14)

diff --git a/flang/lib/Lower/DirectivesCommon.h b/flang/lib/Lower/DirectivesCommon.h
index 656fd36099b51ff..496f695e7e01f29 100644
--- a/flang/lib/Lower/DirectivesCommon.h
+++ b/flang/lib/Lower/DirectivesCommon.h
@@ -327,6 +327,36 @@ void genOmpAccAtomicWrite(Fortran::lower::AbstractConverter &converter,
                                 rightHandClauseList);
 }
 
+template <typename PtrLikeTy>
+mlir::Value convertRhs(fir::FirOpBuilder &builder, mlir::Location loc,
+                       PtrLikeTy ptrLikeType, mlir::Value rhs,
+                       mlir::Operation *captureOp) {
+  if (ptrLikeType.getElementType() != rhs.getType()) {
+    auto crtPos = builder.saveInsertionPoint();
+    builder.setInsertionPoint(captureOp);
+    mlir::Value convertedRhs =
+        builder.create<fir::ConvertOp>(loc, ptrLikeType.getElementType(), rhs);
+    builder.restoreInsertionPoint(crtPos);
+    return convertedRhs;
+  }
+  return rhs;
+}
+
+static mlir::Value addConversionIfNeeded(fir::FirOpBuilder &builder,
+                                         mlir::Location loc, mlir::Value lhs,
+                                         mlir::Value rhs,
+                                         mlir::Operation *captureOp = nullptr) {
+  if (auto ptrLikeType =
+          mlir::dyn_cast_or_null<mlir::acc::PointerLikeType>(lhs.getType()))
+    return convertRhs<mlir::acc::PointerLikeType>(builder, loc, ptrLikeType,
+                                                  rhs, captureOp);
+  if (auto ptrLikeType =
+          mlir::dyn_cast_or_null<mlir::omp::PointerLikeType>(lhs.getType()))
+    return convertRhs<mlir::omp::PointerLikeType>(builder, loc, ptrLikeType,
+                                                  rhs, captureOp);
+  return rhs;
+}
+
 /// Processes an atomic construct with read clause.
 template <typename AtomicT, typename AtomicListT>
 void genOmpAccAtomicRead(Fortran::lower::AbstractConverter &converter,
@@ -357,6 +387,13 @@ void genOmpAccAtomicRead(Fortran::lower::AbstractConverter &converter,
       fir::getBase(converter.genExprAddr(fromExpr, stmtCtx));
   mlir::Value toAddress = fir::getBase(converter.genExprAddr(
       *Fortran::semantics::GetExpr(assignmentStmtVariable), stmtCtx));
+
+  mlir::Location loc = converter.getCurrentLocation();
+  fir::FirOpBuilder &builder = converter.getFirOpBuilder();
+
+  if (fromAddress.getType() != toAddress.getType())
+    fromAddress =
+        builder.create<fir::ConvertOp>(loc, toAddress.getType(), fromAddress);
   genOmpAccAtomicCaptureStatement(converter, fromAddress, toAddress,
                                   leftHandClauseList, rightHandClauseList,
                                   elementType);
@@ -517,6 +554,9 @@ void genOmpAccAtomicCapture(Fortran::lower::AbstractConverter &converter,
           converter, stmt1RHSArg, stmt1LHSArg,
           /*leftHandClauseList=*/nullptr,
           /*rightHandClauseList=*/nullptr, elementType);
+      stmt2RHSArg =
+          addConversionIfNeeded(firOpBuilder, currentLocation, stmt1RHSArg,
+                                stmt2RHSArg, atomicCaptureOp);
       genOmpAccAtomicWriteStatement<AtomicListT>(
           converter, stmt1RHSArg, stmt2RHSArg,
           /*leftHandClauseList=*/nullptr,
diff --git a/flang/test/Lower/OpenACC/acc-atomic-capture.f90 b/flang/test/Lower/OpenACC/acc-atomic-capture.f90
index 382991cf7221ba7..59c16cf61acf718 100644
--- a/flang/test/Lower/OpenACC/acc-atomic-capture.f90
+++ b/flang/test/Lower/OpenACC/acc-atomic-capture.f90
@@ -97,3 +97,28 @@ subroutine pointers_in_atomic_capture()
         b = a
     !$acc end atomic
 end subroutine
+
+
+subroutine capture_with_convert_f32_to_i32()
+  implicit none
+  integer :: k, v, i
+
+  k = 1
+  v = 0
+
+  !$acc atomic capture
+  v = k
+  k = (i + 1) * 3.14
+  !$acc end atomic
+end subroutine
+
+! CHECK-LABEL: func.func @_QPcapture_with_convert_f32_to_i32()
+! CHECK: %[[K:.*]] = fir.alloca i32 {bindc_name = "k", uniq_name = "_QFcapture_with_convert_f32_to_i32Ek"}
+! CHECK: %[[V:.*]] = fir.alloca i32 {bindc_name = "v", uniq_name = "_QFcapture_with_convert_f32_to_i32Ev"}
+! CHECK: %[[CST:.*]] = arith.constant 3.140000e+00 : f32
+! CHECK: %[[MUL:.*]] = arith.mulf %{{.*}}, %[[CST]] fastmath<contract> : f32
+! CHECK: %[[CONV:.*]] = fir.convert %[[MUL]] : (f32) -> i32
+! CHECK: acc.atomic.capture {
+! CHECK:   acc.atomic.read %[[V]] = %[[K]] : !fir.ref<i32>, i32
+! CHECK:   acc.atomic.write %[[K]] = %[[CONV]] : !fir.ref<i32>, i32
+! CHECK: }
diff --git a/flang/test/Lower/OpenACC/acc-atomic-read.f90 b/flang/test/Lower/OpenACC/acc-atomic-read.f90
index 28f0ce44e6f413d..3a718576124c3bb 100644
--- a/flang/test/Lower/OpenACC/acc-atomic-read.f90
+++ b/flang/test/Lower/OpenACC/acc-atomic-read.f90
@@ -46,3 +46,17 @@ subroutine atomic_read_pointer()
 
   x = y
 end
+
+subroutine atomic_read_with_convert()
+  integer(4) :: x
+  integer(8) :: y
+
+  !$acc atomic read
+  y = x
+end
+
+! CHECK-LABEL: func.func @_QPatomic_read_with_convert() {
+! CHECK: %[[X:.*]] = fir.alloca i32 {bindc_name = "x", uniq_name = "_QFatomic_read_with_convertEx"}
+! CHECK: %[[Y:.*]] = fir.alloca i64 {bindc_name = "y", uniq_name = "_QFatomic_read_with_convertEy"}
+! CHECK: %[[CONV:.*]] = fir.convert %[[X]] : (!fir.ref<i32>) -> !fir.ref<i64>
+! CHECK: acc.atomic.read %[[Y]] = %[[CONV]] : !fir.ref<i64>, i32

[kiranchandramohan]

Haven't looked into this in detail but using the typed assignment rather than the parser::Expr had previously fixed a similar issue for Atomic Write. https://reviews.llvm.org/D154163

[clementval]

Thanks @kiranchandramohan I'll check if we can apply the same fix!

[kiranchandramohan]

If that does not work feel free to move ahead with this patch.

[clementval]

If that does not work feel free to move ahead with this patch.

It might be possible for the capture case but I see extra store for the read one.

[clementval]

@kiranchandramohan I updated the capture code to use the typed assignment. I left the read code as is to avoid extra store operation.

[kiranchandramohan]

Thanks @clementval. LGTM.

[razvanlupusoru]

Thank you!