Skip to content

[flang][cuda] Adding atomicadd as a cudadevice intrinsic and converting it LLVM dialect #123840

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jan 23, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions flang/include/flang/Optimizer/Builder/IntrinsicCall.h
Original file line number Diff line number Diff line change
Expand Up @@ -185,6 +185,7 @@ struct IntrinsicLibrary {
mlir::Value genAnint(mlir::Type, llvm::ArrayRef<mlir::Value>);
fir::ExtendedValue genAny(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
mlir::Value genAtanpi(mlir::Type, llvm::ArrayRef<mlir::Value>);
mlir::Value genAtomicAdd(mlir::Type, llvm::ArrayRef<mlir::Value>);
fir::ExtendedValue
genCommandArgumentCount(mlir::Type, llvm::ArrayRef<fir::ExtendedValue>);
mlir::Value genAsind(mlir::Type, llvm::ArrayRef<mlir::Value>);
Expand Down
26 changes: 25 additions & 1 deletion flang/lib/Optimizer/Builder/IntrinsicCall.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -44,14 +44,14 @@
#include "flang/Runtime/iostat-consts.h"
#include "mlir/Dialect/Complex/IR/Complex.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/MathExtras.h"
#include "llvm/Support/raw_ostream.h"
#include <cfenv> // temporary -- only used in genIeeeGetOrSetModesOrStatus
#include <mlir/IR/Value.h>
#include <optional>

#define DEBUG_TYPE "flang-lower-intrinsic"
Expand Down Expand Up @@ -147,6 +147,10 @@ static constexpr IntrinsicHandler handlers[]{
{"atan2pi", &I::genAtanpi},
{"atand", &I::genAtand},
{"atanpi", &I::genAtanpi},
{"atomicaddd", &I::genAtomicAdd, {{{"a", asAddr}, {"v", asValue}}}, false},
{"atomicaddf", &I::genAtomicAdd, {{{"a", asAddr}, {"v", asValue}}}, false},
{"atomicaddi", &I::genAtomicAdd, {{{"a", asAddr}, {"v", asValue}}}, false},
{"atomicaddl", &I::genAtomicAdd, {{{"a", asAddr}, {"v", asValue}}}, false},
{"bessel_jn",
&I::genBesselJn,
{{{"n1", asValue}, {"n2", asValue}, {"x", asValue}}},
Expand Down Expand Up @@ -2574,6 +2578,26 @@ mlir::Value IntrinsicLibrary::genAtanpi(mlir::Type resultType,
return builder.create<mlir::arith::MulFOp>(loc, atan, factor);
}

static mlir::Value genAtomBinOp(fir::FirOpBuilder &builder, mlir::Location &loc,
mlir::LLVM::AtomicBinOp binOp, mlir::Value arg0,
mlir::Value arg1) {
auto llvmPointerType = mlir::LLVM::LLVMPointerType::get(builder.getContext());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Out of curiosity, do you now if FIR alias analysis is clever enough to go through convert from this LLVM type and deduce what the llvm.atomicrmw is reading/writing not (or mainly, to deduce what it is not reading/writing to)?

It makes sense to me that we interoperate with LLVM pointer type to avoid redefining all the intrinsics at the FIR level, I just wonder if we need to improve AA here.

Copy link
Contributor Author

@Renaud-K Renaud-K Jan 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We always skip through fir.convert

        .Case<fir::ConvertOp>([&](auto op) {
          // Skip ConvertOp's and track further through the operand.
          v = op->getOperand(0);
          defOp = v.getDefiningOp();
        })
...

I was trying to make sense of the MemoryEffects of LLVM_AtomicRMWOp but couldn't quite make sense of it:

def LLVM_AtomicRMWOp : LLVM_MemAccessOpBase<"atomicrmw", [
      TypesMatchWith<"result #0 and operand #1 have the same type",
                     "val", "res", "$_self">]> {

We talked about these ops in the office hours today and while we need the hooks in place to create the appropriate instructions we may eventually replace them with higher level atomic operations. TDB.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, the effect description is not straightforward. If found this patch that explains some of it: feb7bea

From what I understand, it is giving atomic operation a global write effects because the synchronization aspect (e.g., could write some global lock).

arg0 = builder.createConvert(loc, llvmPointerType, arg0);
return builder.create<mlir::LLVM::AtomicRMWOp>(
loc, binOp, arg0, arg1, mlir::LLVM::AtomicOrdering::seq_cst);
}

mlir::Value IntrinsicLibrary::genAtomicAdd(mlir::Type resultType,
llvm::ArrayRef<mlir::Value> args) {
assert(args.size() == 2);

mlir::LLVM::AtomicBinOp binOp =
mlir::isa<mlir::IntegerType>(args[1].getType())
? mlir::LLVM::AtomicBinOp::add
: mlir::LLVM::AtomicBinOp::fadd;
return genAtomBinOp(builder, loc, binOp, args[0], args[1]);
}

// ASSOCIATED
fir::ExtendedValue
IntrinsicLibrary::genAssociated(mlir::Type resultType,
Expand Down
28 changes: 27 additions & 1 deletion flang/module/cudadevice.f90
Original file line number Diff line number Diff line change
Expand Up @@ -92,5 +92,31 @@ attributes(device) subroutine threadfence_system()
end function
end interface
public :: __fadd_ru


! Atomic Operations

interface atomicadd
attributes(device) pure integer function atomicaddi(address, val)
!dir$ ignore_tkr (d) address, (d) val
integer, intent(inout) :: address
integer, value :: val
end function
attributes(device) pure real function atomicaddf(address, val)
!dir$ ignore_tkr (d) address, (d) val
real, intent(inout) :: address
real, value :: val
end function
attributes(device) pure real*8 function atomicaddd(address, val)
!dir$ ignore_tkr (d) address, (d) val
real*8, intent(inout) :: address
real*8, value :: val
end function
attributes(device) pure integer(8) function atomicaddl(address, val)
!dir$ ignore_tkr (d) address, (d) val
integer(8), intent(inout) :: address
integer(8), value :: val
end function
end interface
public :: atomicadd

end module
13 changes: 13 additions & 0 deletions flang/test/Lower/CUDA/cuda-device-proc.cuf
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,10 @@
attributes(global) subroutine devsub()
implicit none
integer :: ret
real(4) :: af
real(8) :: ad
integer(4) :: ai
integer(8) :: al

call syncthreads()
call syncwarp(1)
Expand All @@ -14,6 +18,11 @@ attributes(global) subroutine devsub()
ret = syncthreads_and(1)
ret = syncthreads_count(1)
ret = syncthreads_or(1)

ai = atomicadd(ai, 1_4)
al = atomicadd(al, 1_8)
af = atomicadd(af, 1.0_4)
ad = atomicadd(ad, 1.0_8)
end

! CHECK-LABEL: func.func @_QPdevsub() attributes {cuf.proc_attr = #cuf.cuda_proc<global>}
Expand All @@ -25,6 +34,10 @@ end
! CHECK: %{{.*}} = fir.call @llvm.nvvm.barrier0.and(%c1_i32_0) fastmath<contract> : (i32) -> i32
! CHECK: %{{.*}} = fir.call @llvm.nvvm.barrier0.popc(%c1_i32_1) fastmath<contract> : (i32) -> i32
! CHECK: %{{.*}} = fir.call @llvm.nvvm.barrier0.or(%c1_i32_2) fastmath<contract> : (i32) -> i32
! CHECK: %{{.*}} = llvm.atomicrmw add %{{.*}}, %{{.*}} seq_cst : !llvm.ptr, i32
! CHECK: %{{.*}} = llvm.atomicrmw add %{{.*}}, %{{.*}} seq_cst : !llvm.ptr, i64
! CHECK: %{{.*}} = llvm.atomicrmw fadd %{{.*}}, %{{.*}} seq_cst : !llvm.ptr, f32
! CHECK: %{{.*}} = llvm.atomicrmw fadd %{{.*}}, %{{.*}} seq_cst : !llvm.ptr, f64

! CHECK: func.func private @llvm.nvvm.barrier0()
! CHECK: func.func private @__syncwarp(!fir.ref<i32> {cuf.data_attr = #cuf.cuda<device>}) attributes {cuf.proc_attr = #cuf.cuda_proc<device>, fir.bindc_name = "__syncwarp", fir.proc_attrs = #fir.proc_attrs<bind_c>}
Expand Down
9 changes: 9 additions & 0 deletions flang/test/Semantics/cuf-device-procedures01.cuf
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,17 @@ end
! CHECK: threadfence_system (Subroutine): Use from threadfence_system in cudadevice

subroutine host()
real(4) :: af
real(8) :: ad
integer(4) :: ai
integer(8) :: al
call syncthreads()
ai = atomicadd(ai, 1_4)
al = atomicadd(al, 1_8)
af = atomicadd(af, 1.0_4)
ad = atomicadd(ad, 1.0_8)
end subroutine

! CHECK-LABEL: Subprogram scope: host
! CHECK: atomicadd, EXTERNAL: HostAssoc{{$}}
! CHECK: syncthreads, EXTERNAL: HostAssoc{{$}}
Loading