[llvm] [NVPTX] Add support for atomic add for f16 type (PR #84295)
Adrian Kuegel via llvm-commits
llvm-commits at lists.llvm.org
Mon Mar 11 00:12:05 PDT 2024
https://github.com/akuegel updated https://github.com/llvm/llvm-project/pull/84295
>From 7bac5a88e37b016d4797072aa40cb791792c8749 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel at google.com>
Date: Thu, 7 Mar 2024 09:56:57 +0000
Subject: [PATCH 1/5] [NVPTX] Add support for atomic add for f16 type.
atom.add.noftz.f16 is supported since SM 7.0
---
llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp | 2 ++
llvm/lib/Target/NVPTX/NVPTXIntrinsics.td | 15 +++++++++++++++
2 files changed, 17 insertions(+)
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index c979c03dc1b835..7e39fd0c81d86c 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -6100,6 +6100,8 @@ NVPTXTargetLowering::shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const {
if (AI->isFloatingPointOperation()) {
if (AI->getOperation() == AtomicRMWInst::BinOp::FAdd) {
+ if (Ty->isHalfTy() && STI.getSmVersion() >= 70)
+ return AtomicExpansionKind::None;
if (Ty->isFloatTy())
return AtomicExpansionKind::None;
if (Ty->isDoubleTy() && STI.hasAtomAddF64())
diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
index 477789a164ead2..31b5338b0d0f5d 100644
--- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
+++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
@@ -1630,6 +1630,13 @@ defm INT_PTX_ATOM_ADD_GEN_64 : F_ATOMIC_2<i64, Int64Regs, "", ".u64", ".add",
defm INT_PTX_ATOM_ADD_GEN_64_USE_G : F_ATOMIC_2<i64, Int64Regs, ".global", ".u64",
".add", atomic_load_add_64_gen, i64imm, imm>;
+defm INT_PTX_ATOM_ADD_G_F16 : F_ATOMIC_2<f16, Int16Regs, ".global", ".f16", ".add.noftz",
+ atomic_load_add_g, f16imm, fpimm, [hasSM<70>]>;
+defm INT_PTX_ATOM_ADD_S_F16 : F_ATOMIC_2<f16, Int16Regs, ".shared", ".f16", ".add.noftz",
+ atomic_load_add_s, f16imm, fpimm, [hasSM<70>]>;
+defm INT_PTX_ATOM_ADD_GEN_F16 : F_ATOMIC_2<f16, Int16Regs, "", ".f16", ".add.noftz",
+ atomic_load_add_gen, f16imm, fpimm, [hasSM<70>]>;
+
defm INT_PTX_ATOM_ADD_G_F32 : F_ATOMIC_2<f32, Float32Regs, ".global", ".f32", ".add",
atomic_load_add_g, f32imm, fpimm>;
defm INT_PTX_ATOM_ADD_S_F32 : F_ATOMIC_2<f32, Float32Regs, ".shared", ".f32", ".add",
@@ -2007,6 +2014,9 @@ multiclass ATOM2P_impl<string AsmStr, Intrinsic Intr,
SDNode Imm, ValueType ImmTy,
list<Predicate> Preds> {
let AddedComplexity = 1 in {
+ def : ATOM23_impl<AsmStr, regT, regclass, Preds,
+ (ins Int16Regs:$src, regclass:$b),
+ (Intr (i16 Int16Regs:$src), (regT regclass:$b))>;
def : ATOM23_impl<AsmStr, regT, regclass, Preds,
(ins Int32Regs:$src, regclass:$b),
(Intr (i32 Int32Regs:$src), (regT regclass:$b))>;
@@ -2017,6 +2027,9 @@ multiclass ATOM2P_impl<string AsmStr, Intrinsic Intr,
// tablegen can't infer argument types from Intrinsic (though it can
// from Instruction) so we have to enforce specific type on
// immediates via explicit cast to ImmTy.
+ def : ATOM23_impl<AsmStr, regT, regclass, Preds,
+ (ins Int16Regs:$src, ImmType:$b),
+ (Intr (i16 Int16Regs:$src), (ImmTy Imm:$b))>;
def : ATOM23_impl<AsmStr, regT, regclass, Preds,
(ins Int32Regs:$src, ImmType:$b),
(Intr (i32 Int32Regs:$src), (ImmTy Imm:$b))>;
@@ -2136,6 +2149,8 @@ multiclass ATOM2_add_impl<string OpStr> {
defm _s32 : ATOM2S_impl<OpStr, "i", "s32", i32, Int32Regs, i32imm, imm, i32, []>;
defm _u32 : ATOM2S_impl<OpStr, "i", "u32", i32, Int32Regs, i32imm, imm, i32, []>;
defm _u64 : ATOM2S_impl<OpStr, "i", "u64", i64, Int64Regs, i64imm, imm, i64, []>;
+ defm _f16 : ATOM2S_impl<OpStr, "f", "f16", f16, Int16Regs, f16imm, fpimm, f16,
+ [hasSM<70>]>;
defm _f32 : ATOM2S_impl<OpStr, "f", "f32", f32, Float32Regs, f32imm, fpimm, f32,
[]>;
defm _f64 : ATOM2S_impl<OpStr, "f", "f64", f64, Float64Regs, f64imm, fpimm, f64,
>From 9aefbbee78df14eedac9ffce4238ed520c7a2fe2 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel at google.com>
Date: Thu, 7 Mar 2024 10:00:19 +0000
Subject: [PATCH 2/5] [NVPTX] Add atomic add test.
---
llvm/test/CodeGen/NVPTX/atomics-sm70.ll | 17 +++++++++++++++++
1 file changed, 17 insertions(+)
create mode 100644 llvm/test/CodeGen/NVPTX/atomics-sm70.ll
diff --git a/llvm/test/CodeGen/NVPTX/atomics-sm70.ll b/llvm/test/CodeGen/NVPTX/atomics-sm70.ll
new file mode 100644
index 00000000000000..71d9d7b62f9107
--- /dev/null
+++ b/llvm/test/CodeGen/NVPTX/atomics-sm70.ll
@@ -0,0 +1,17 @@
+; RUN: llc < %s -march=nvptx -mcpu=sm_70 | FileCheck %s
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_70 | FileCheck %s
+; RUN: %if ptxas && !ptxas-12.0 %{ llc < %s -march=nvptx -mcpu=sm_70 | %ptxas-verify -arch=sm_70 %}
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_70 | %ptxas-verify -arch=sm_70 %}
+
+; CHECK-LABEL: .func test(
+define void @test(ptr %dp0, ptr addrspace(1) %dp1, ptr addrspace(3) %dp3, half %d) {
+; CHECK: atom.add.noftz.f16
+ %r1 = atomicrmw fadd ptr %dp0, half %d seq_cst
+; CHECK: atom.global.add.noftz.f16
+ %r2 = atomicrmw fadd ptr addrspace(1) %dp1, half %d seq_cst
+; CHECK: atom.shared.add.noftz.f16
+ %ret = atomicrmw fadd ptr addrspace(3) %dp3, half %d seq_cst
+ ret void
+}
+
+attributes #1 = { argmemonly nounwind }
>From dbd3f9a2b5707246932a17c3271298ba2e53fa1a Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel at google.com>
Date: Fri, 8 Mar 2024 07:37:35 +0000
Subject: [PATCH 3/5] Also check PTX version.
---
llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp | 3 ++-
llvm/lib/Target/NVPTX/NVPTXIntrinsics.td | 8 ++++----
2 files changed, 6 insertions(+), 5 deletions(-)
diff --git a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
index 7e39fd0c81d86c..c411c8ef9528d7 100644
--- a/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
+++ b/llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp
@@ -6100,7 +6100,8 @@ NVPTXTargetLowering::shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const {
if (AI->isFloatingPointOperation()) {
if (AI->getOperation() == AtomicRMWInst::BinOp::FAdd) {
- if (Ty->isHalfTy() && STI.getSmVersion() >= 70)
+ if (Ty->isHalfTy() && STI.getSmVersion() >= 70 &&
+ STI.getPTXVersion() >= 63)
return AtomicExpansionKind::None;
if (Ty->isFloatTy())
return AtomicExpansionKind::None;
diff --git a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
index 31b5338b0d0f5d..869b13369e87e1 100644
--- a/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
+++ b/llvm/lib/Target/NVPTX/NVPTXIntrinsics.td
@@ -1631,11 +1631,11 @@ defm INT_PTX_ATOM_ADD_GEN_64_USE_G : F_ATOMIC_2<i64, Int64Regs, ".global", ".u64
".add", atomic_load_add_64_gen, i64imm, imm>;
defm INT_PTX_ATOM_ADD_G_F16 : F_ATOMIC_2<f16, Int16Regs, ".global", ".f16", ".add.noftz",
- atomic_load_add_g, f16imm, fpimm, [hasSM<70>]>;
+ atomic_load_add_g, f16imm, fpimm, [hasSM<70>, hasPTX<63>]>;
defm INT_PTX_ATOM_ADD_S_F16 : F_ATOMIC_2<f16, Int16Regs, ".shared", ".f16", ".add.noftz",
- atomic_load_add_s, f16imm, fpimm, [hasSM<70>]>;
+ atomic_load_add_s, f16imm, fpimm, [hasSM<70>, hasPTX<63>]>;
defm INT_PTX_ATOM_ADD_GEN_F16 : F_ATOMIC_2<f16, Int16Regs, "", ".f16", ".add.noftz",
- atomic_load_add_gen, f16imm, fpimm, [hasSM<70>]>;
+ atomic_load_add_gen, f16imm, fpimm, [hasSM<70>, hasPTX<63>]>;
defm INT_PTX_ATOM_ADD_G_F32 : F_ATOMIC_2<f32, Float32Regs, ".global", ".f32", ".add",
atomic_load_add_g, f32imm, fpimm>;
@@ -2150,7 +2150,7 @@ multiclass ATOM2_add_impl<string OpStr> {
defm _u32 : ATOM2S_impl<OpStr, "i", "u32", i32, Int32Regs, i32imm, imm, i32, []>;
defm _u64 : ATOM2S_impl<OpStr, "i", "u64", i64, Int64Regs, i64imm, imm, i64, []>;
defm _f16 : ATOM2S_impl<OpStr, "f", "f16", f16, Int16Regs, f16imm, fpimm, f16,
- [hasSM<70>]>;
+ [hasSM<70>, hasPTX<63>]>;
defm _f32 : ATOM2S_impl<OpStr, "f", "f32", f32, Float32Regs, f32imm, fpimm, f32,
[]>;
defm _f64 : ATOM2S_impl<OpStr, "f", "f64", f64, Float64Regs, f64imm, fpimm, f64,
>From 9765a9a1362abe3536c6888b7d44ab1afbaecfb8 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel at google.com>
Date: Mon, 11 Mar 2024 07:05:11 +0000
Subject: [PATCH 4/5] Update tests and add more tests
---
llvm/test/CodeGen/NVPTX/atomics-sm70.ll | 27 +++++++++++++++++--------
llvm/test/CodeGen/NVPTX/atomics.ll | 7 +++++++
2 files changed, 26 insertions(+), 8 deletions(-)
diff --git a/llvm/test/CodeGen/NVPTX/atomics-sm70.ll b/llvm/test/CodeGen/NVPTX/atomics-sm70.ll
index 71d9d7b62f9107..e43cb6f9dec524 100644
--- a/llvm/test/CodeGen/NVPTX/atomics-sm70.ll
+++ b/llvm/test/CodeGen/NVPTX/atomics-sm70.ll
@@ -1,16 +1,27 @@
-; RUN: llc < %s -march=nvptx -mcpu=sm_70 | FileCheck %s
-; RUN: llc < %s -march=nvptx64 -mcpu=sm_70 | FileCheck %s
-; RUN: %if ptxas && !ptxas-12.0 %{ llc < %s -march=nvptx -mcpu=sm_70 | %ptxas-verify -arch=sm_70 %}
-; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_70 | %ptxas-verify -arch=sm_70 %}
+; RUN: llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx63 | FileCheck %s --dump-input=always
+; RUN: llc < %s -march=nvptx64 -mcpu=sm_70 -mattr=+ptx63 | FileCheck %s
+; RUN: llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx62 | FileCheck %s --check-prefixes=CHECKPTX62
+; RUN: %if ptxas && !ptxas-12.0 %{ llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx63 | %ptxas-verify -arch=sm_70 %}
+; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_70 -mattr=+ptx63 | %ptxas-verify -arch=sm_70 %}
+; RUN: %if ptxas && !ptxas-12.0 %{ llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx62 | %ptxas-verify -arch=sm_70 %}
; CHECK-LABEL: .func test(
-define void @test(ptr %dp0, ptr addrspace(1) %dp1, ptr addrspace(3) %dp3, half %d) {
+define void @test(ptr %dp0, ptr addrspace(1) %dp1, ptr addrspace(3) %dp3, half %val) {
; CHECK: atom.add.noftz.f16
- %r1 = atomicrmw fadd ptr %dp0, half %d seq_cst
+; CHECK-NOT: atom.cas
+; CHECKPTX62: atom.cas
+; CHECKPTX62-NOT: atom.add.noftz.f16
+ %r1 = atomicrmw fadd ptr %dp0, half %val seq_cst
; CHECK: atom.global.add.noftz.f16
- %r2 = atomicrmw fadd ptr addrspace(1) %dp1, half %d seq_cst
+; CHECK-NOT: atom.global.cas
+; CHECKPTX62: atom.global.cas
+; CHECKPTX62-NOT: atom.global.add.noftz.f16
+ %r2 = atomicrmw fadd ptr addrspace(1) %dp1, half %val seq_cst
; CHECK: atom.shared.add.noftz.f16
- %ret = atomicrmw fadd ptr addrspace(3) %dp3, half %d seq_cst
+; CHECK-NOT: atom.shared.cas
+; CHECKPTX62: atom.shared.cas
+; CHECKPTX62-NOT: atom.shared.add.noftz.f16
+ %ret = atomicrmw fadd ptr addrspace(3) %dp3, half %val seq_cst
ret void
}
diff --git a/llvm/test/CodeGen/NVPTX/atomics.ll b/llvm/test/CodeGen/NVPTX/atomics.ll
index e99d0fd05e346b..6f2b5dcf47f13b 100644
--- a/llvm/test/CodeGen/NVPTX/atomics.ll
+++ b/llvm/test/CodeGen/NVPTX/atomics.ll
@@ -175,6 +175,13 @@ define float @atomicrmw_add_f32_generic(ptr %addr, float %val) {
ret float %ret
}
+; CHECK-LABEL: atomicrmw_add_f16_generic
+define half @atomicrmw_add_f16_generic(ptr %addr, half %val) {
+; CHECK: atom.cas
+ %ret = atomicrmw fadd ptr %addr, half %val seq_cst
+ ret half %ret
+}
+
; CHECK-LABEL: atomicrmw_add_f32_addrspace1
define float @atomicrmw_add_f32_addrspace1(ptr addrspace(1) %addr, float %val) {
; CHECK: atom.global.add.f32
>From 3bd7abaa7f1582f6f2706aac14bd8026925d7730 Mon Sep 17 00:00:00 2001
From: Adrian Kuegel <akuegel at google.com>
Date: Mon, 11 Mar 2024 07:11:18 +0000
Subject: [PATCH 5/5] Remove dump-input=always left over from debugging.
---
llvm/test/CodeGen/NVPTX/atomics-sm70.ll | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/llvm/test/CodeGen/NVPTX/atomics-sm70.ll b/llvm/test/CodeGen/NVPTX/atomics-sm70.ll
index e43cb6f9dec524..e1ae9661cb1617 100644
--- a/llvm/test/CodeGen/NVPTX/atomics-sm70.ll
+++ b/llvm/test/CodeGen/NVPTX/atomics-sm70.ll
@@ -1,4 +1,4 @@
-; RUN: llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx63 | FileCheck %s --dump-input=always
+; RUN: llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx63 | FileCheck %s
; RUN: llc < %s -march=nvptx64 -mcpu=sm_70 -mattr=+ptx63 | FileCheck %s
; RUN: llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx62 | FileCheck %s --check-prefixes=CHECKPTX62
; RUN: %if ptxas && !ptxas-12.0 %{ llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx63 | %ptxas-verify -arch=sm_70 %}
More information about the llvm-commits
mailing list