Skip to content

Commit 8e0f4b9

Browse files
akuegelakuegel
authored
[NVPTX] Add support for atomic add for f16 type (#84295)
atom.add.noftz.f16 is supported since SM 7.0
1 parent 85f6669 commit 8e0f4b9

File tree

4 files changed

+146
-0
lines changed

4 files changed

+146
-0
lines changed

‎llvm/lib/Target/NVPTX/NVPTXISelLowering.cpp

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6100,6 +6100,9 @@ NVPTXTargetLowering::shouldExpandAtomicRMWInIR(AtomicRMWInst *AI) const {
61006100

61016101
if (AI->isFloatingPointOperation()) {
61026102
if (AI->getOperation() == AtomicRMWInst::BinOp::FAdd) {
6103+
if (Ty->isHalfTy() && STI.getSmVersion() >= 70 &&
6104+
STI.getPTXVersion() >= 63)
6105+
return AtomicExpansionKind::None;
61036106
if (Ty->isFloatTy())
61046107
return AtomicExpansionKind::None;
61056108
if (Ty->isDoubleTy() && STI.hasAtomAddF64())

‎llvm/lib/Target/NVPTX/NVPTXIntrinsics.td

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1630,6 +1630,13 @@ defm INT_PTX_ATOM_ADD_GEN_64 : F_ATOMIC_2<i64, Int64Regs, "", ".u64", ".add",
16301630
defm INT_PTX_ATOM_ADD_GEN_64_USE_G : F_ATOMIC_2<i64, Int64Regs, ".global", ".u64",
16311631
".add", atomic_load_add_64_gen, i64imm, imm>;
16321632

1633+
defm INT_PTX_ATOM_ADD_G_F16 : F_ATOMIC_2<f16, Int16Regs, ".global", ".f16", ".add.noftz",
1634+
atomic_load_add_g, f16imm, fpimm, [hasSM<70>, hasPTX<63>]>;
1635+
defm INT_PTX_ATOM_ADD_S_F16 : F_ATOMIC_2<f16, Int16Regs, ".shared", ".f16", ".add.noftz",
1636+
atomic_load_add_s, f16imm, fpimm, [hasSM<70>, hasPTX<63>]>;
1637+
defm INT_PTX_ATOM_ADD_GEN_F16 : F_ATOMIC_2<f16, Int16Regs, "", ".f16", ".add.noftz",
1638+
atomic_load_add_gen, f16imm, fpimm, [hasSM<70>, hasPTX<63>]>;
1639+
16331640
defm INT_PTX_ATOM_ADD_G_F32 : F_ATOMIC_2<f32, Float32Regs, ".global", ".f32", ".add",
16341641
atomic_load_add_g, f32imm, fpimm>;
16351642
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,
20072014
SDNode Imm, ValueType ImmTy,
20082015
list<Predicate> Preds> {
20092016
let AddedComplexity = 1 in {
2017+
def : ATOM23_impl<AsmStr, regT, regclass, Preds,
2018+
(ins Int16Regs:$src, regclass:$b),
2019+
(Intr (i16 Int16Regs:$src), (regT regclass:$b))>;
20102020
def : ATOM23_impl<AsmStr, regT, regclass, Preds,
20112021
(ins Int32Regs:$src, regclass:$b),
20122022
(Intr (i32 Int32Regs:$src), (regT regclass:$b))>;
@@ -2017,6 +2027,9 @@ multiclass ATOM2P_impl<string AsmStr, Intrinsic Intr,
20172027
// tablegen can't infer argument types from Intrinsic (though it can
20182028
// from Instruction) so we have to enforce specific type on
20192029
// immediates via explicit cast to ImmTy.
2030+
def : ATOM23_impl<AsmStr, regT, regclass, Preds,
2031+
(ins Int16Regs:$src, ImmType:$b),
2032+
(Intr (i16 Int16Regs:$src), (ImmTy Imm:$b))>;
20202033
def : ATOM23_impl<AsmStr, regT, regclass, Preds,
20212034
(ins Int32Regs:$src, ImmType:$b),
20222035
(Intr (i32 Int32Regs:$src), (ImmTy Imm:$b))>;
@@ -2136,6 +2149,8 @@ multiclass ATOM2_add_impl<string OpStr> {
21362149
defm _s32 : ATOM2S_impl<OpStr, "i", "s32", i32, Int32Regs, i32imm, imm, i32, []>;
21372150
defm _u32 : ATOM2S_impl<OpStr, "i", "u32", i32, Int32Regs, i32imm, imm, i32, []>;
21382151
defm _u64 : ATOM2S_impl<OpStr, "i", "u64", i64, Int64Regs, i64imm, imm, i64, []>;
2152+
defm _f16 : ATOM2S_impl<OpStr, "f", "f16", f16, Int16Regs, f16imm, fpimm, f16,
2153+
[hasSM<70>, hasPTX<63>]>;
21392154
defm _f32 : ATOM2S_impl<OpStr, "f", "f32", f32, Float32Regs, f32imm, fpimm, f32,
21402155
[]>;
21412156
defm _f64 : ATOM2S_impl<OpStr, "f", "f64", f64, Float64Regs, f64imm, fpimm, f64,

‎llvm/test/CodeGen/NVPTX/atomics-sm70.ll

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
2+
; RUN: llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx63 | FileCheck %s --check-prefixes=CHECK
3+
; RUN: llc < %s -march=nvptx64 -mcpu=sm_70 -mattr=+ptx63 | FileCheck %s --check-prefixes=CHECK64
4+
; RUN: llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx62 | FileCheck %s --check-prefixes=CHECKPTX62
5+
; RUN: %if ptxas && !ptxas-12.0 %{ llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx63 | %ptxas-verify -arch=sm_70 %}
6+
; RUN: %if ptxas %{ llc < %s -march=nvptx64 -mcpu=sm_70 -mattr=+ptx63 | %ptxas-verify -arch=sm_70 %}
7+
; RUN: %if ptxas && !ptxas-12.0 %{ llc < %s -march=nvptx -mcpu=sm_70 -mattr=+ptx62 | %ptxas-verify -arch=sm_70 %}
8+
9+
target triple = "nvptx64-nvidia-cuda"
10+
11+
define void @test(ptr %dp0, ptr addrspace(1) %dp1, ptr addrspace(3) %dp3, half %val) {
12+
; CHECK-LABEL: test(
13+
; CHECK: {
14+
; CHECK-NEXT: .reg .b16 %rs<5>;
15+
; CHECK-NEXT: .reg .b32 %r<4>;
16+
; CHECK-EMPTY:
17+
; CHECK-NEXT: // %bb.0:
18+
; CHECK-NEXT: ld.param.u32 %r1, [test_param_0];
19+
; CHECK-NEXT: ld.param.b16 %rs1, [test_param_3];
20+
; CHECK-NEXT: atom.add.noftz.f16 %rs2, [%r1], %rs1;
21+
; CHECK-NEXT: ld.param.u32 %r2, [test_param_1];
22+
; CHECK-NEXT: atom.global.add.noftz.f16 %rs3, [%r2], %rs1;
23+
; CHECK-NEXT: ld.param.u32 %r3, [test_param_2];
24+
; CHECK-NEXT: atom.shared.add.noftz.f16 %rs4, [%r3], %rs1;
25+
; CHECK-NEXT: ret;
26+
;
27+
; CHECK64-LABEL: test(
28+
; CHECK64: {
29+
; CHECK64-NEXT: .reg .b16 %rs<5>;
30+
; CHECK64-NEXT: .reg .b64 %rd<4>;
31+
; CHECK64-EMPTY:
32+
; CHECK64-NEXT: // %bb.0:
33+
; CHECK64-NEXT: ld.param.u64 %rd1, [test_param_0];
34+
; CHECK64-NEXT: ld.param.b16 %rs1, [test_param_3];
35+
; CHECK64-NEXT: atom.add.noftz.f16 %rs2, [%rd1], %rs1;
36+
; CHECK64-NEXT: ld.param.u64 %rd2, [test_param_1];
37+
; CHECK64-NEXT: atom.global.add.noftz.f16 %rs3, [%rd2], %rs1;
38+
; CHECK64-NEXT: ld.param.u64 %rd3, [test_param_2];
39+
; CHECK64-NEXT: atom.shared.add.noftz.f16 %rs4, [%rd3], %rs1;
40+
; CHECK64-NEXT: ret;
41+
;
42+
; CHECKPTX62-LABEL: test(
43+
; CHECKPTX62: {
44+
; CHECKPTX62-NEXT: .reg .pred %p<4>;
45+
; CHECKPTX62-NEXT: .reg .b16 %rs<14>;
46+
; CHECKPTX62-NEXT: .reg .b32 %r<49>;
47+
; CHECKPTX62-EMPTY:
48+
; CHECKPTX62-NEXT: // %bb.0:
49+
; CHECKPTX62-NEXT: ld.param.b16 %rs1, [test_param_3];
50+
; CHECKPTX62-NEXT: ld.param.u32 %r20, [test_param_2];
51+
; CHECKPTX62-NEXT: ld.param.u32 %r19, [test_param_1];
52+
; CHECKPTX62-NEXT: ld.param.u32 %r21, [test_param_0];
53+
; CHECKPTX62-NEXT: and.b32 %r1, %r21, -4;
54+
; CHECKPTX62-NEXT: and.b32 %r22, %r21, 3;
55+
; CHECKPTX62-NEXT: shl.b32 %r2, %r22, 3;
56+
; CHECKPTX62-NEXT: mov.b32 %r23, 65535;
57+
; CHECKPTX62-NEXT: shl.b32 %r24, %r23, %r2;
58+
; CHECKPTX62-NEXT: not.b32 %r3, %r24;
59+
; CHECKPTX62-NEXT: ld.u32 %r46, [%r1];
60+
; CHECKPTX62-NEXT: $L__BB0_1: // %atomicrmw.start
61+
; CHECKPTX62-NEXT: // =>This Inner Loop Header: Depth=1
62+
; CHECKPTX62-NEXT: shr.u32 %r25, %r46, %r2;
63+
; CHECKPTX62-NEXT: cvt.u16.u32 %rs2, %r25;
64+
; CHECKPTX62-NEXT: add.rn.f16 %rs4, %rs2, %rs1;
65+
; CHECKPTX62-NEXT: cvt.u32.u16 %r26, %rs4;
66+
; CHECKPTX62-NEXT: shl.b32 %r27, %r26, %r2;
67+
; CHECKPTX62-NEXT: and.b32 %r28, %r46, %r3;
68+
; CHECKPTX62-NEXT: or.b32 %r29, %r28, %r27;
69+
; CHECKPTX62-NEXT: atom.cas.b32 %r6, [%r1], %r46, %r29;
70+
; CHECKPTX62-NEXT: setp.ne.s32 %p1, %r6, %r46;
71+
; CHECKPTX62-NEXT: mov.u32 %r46, %r6;
72+
; CHECKPTX62-NEXT: @%p1 bra $L__BB0_1;
73+
; CHECKPTX62-NEXT: // %bb.2: // %atomicrmw.end
74+
; CHECKPTX62-NEXT: and.b32 %r7, %r19, -4;
75+
; CHECKPTX62-NEXT: shl.b32 %r30, %r19, 3;
76+
; CHECKPTX62-NEXT: and.b32 %r8, %r30, 24;
77+
; CHECKPTX62-NEXT: shl.b32 %r32, %r23, %r8;
78+
; CHECKPTX62-NEXT: not.b32 %r9, %r32;
79+
; CHECKPTX62-NEXT: ld.global.u32 %r47, [%r7];
80+
; CHECKPTX62-NEXT: $L__BB0_3: // %atomicrmw.start9
81+
; CHECKPTX62-NEXT: // =>This Inner Loop Header: Depth=1
82+
; CHECKPTX62-NEXT: shr.u32 %r33, %r47, %r8;
83+
; CHECKPTX62-NEXT: cvt.u16.u32 %rs6, %r33;
84+
; CHECKPTX62-NEXT: add.rn.f16 %rs8, %rs6, %rs1;
85+
; CHECKPTX62-NEXT: cvt.u32.u16 %r34, %rs8;
86+
; CHECKPTX62-NEXT: shl.b32 %r35, %r34, %r8;
87+
; CHECKPTX62-NEXT: and.b32 %r36, %r47, %r9;
88+
; CHECKPTX62-NEXT: or.b32 %r37, %r36, %r35;
89+
; CHECKPTX62-NEXT: atom.global.cas.b32 %r12, [%r7], %r47, %r37;
90+
; CHECKPTX62-NEXT: setp.ne.s32 %p2, %r12, %r47;
91+
; CHECKPTX62-NEXT: mov.u32 %r47, %r12;
92+
; CHECKPTX62-NEXT: @%p2 bra $L__BB0_3;
93+
; CHECKPTX62-NEXT: // %bb.4: // %atomicrmw.end8
94+
; CHECKPTX62-NEXT: and.b32 %r13, %r20, -4;
95+
; CHECKPTX62-NEXT: shl.b32 %r38, %r20, 3;
96+
; CHECKPTX62-NEXT: and.b32 %r14, %r38, 24;
97+
; CHECKPTX62-NEXT: shl.b32 %r40, %r23, %r14;
98+
; CHECKPTX62-NEXT: not.b32 %r15, %r40;
99+
; CHECKPTX62-NEXT: ld.shared.u32 %r48, [%r13];
100+
; CHECKPTX62-NEXT: $L__BB0_5: // %atomicrmw.start27
101+
; CHECKPTX62-NEXT: // =>This Inner Loop Header: Depth=1
102+
; CHECKPTX62-NEXT: shr.u32 %r41, %r48, %r14;
103+
; CHECKPTX62-NEXT: cvt.u16.u32 %rs10, %r41;
104+
; CHECKPTX62-NEXT: add.rn.f16 %rs12, %rs10, %rs1;
105+
; CHECKPTX62-NEXT: cvt.u32.u16 %r42, %rs12;
106+
; CHECKPTX62-NEXT: shl.b32 %r43, %r42, %r14;
107+
; CHECKPTX62-NEXT: and.b32 %r44, %r48, %r15;
108+
; CHECKPTX62-NEXT: or.b32 %r45, %r44, %r43;
109+
; CHECKPTX62-NEXT: atom.shared.cas.b32 %r18, [%r13], %r48, %r45;
110+
; CHECKPTX62-NEXT: setp.ne.s32 %p3, %r18, %r48;
111+
; CHECKPTX62-NEXT: mov.u32 %r48, %r18;
112+
; CHECKPTX62-NEXT: @%p3 bra $L__BB0_5;
113+
; CHECKPTX62-NEXT: // %bb.6: // %atomicrmw.end26
114+
; CHECKPTX62-NEXT: ret;
115+
%r1 = atomicrmw fadd ptr %dp0, half %val seq_cst
116+
%r2 = atomicrmw fadd ptr addrspace(1) %dp1, half %val seq_cst
117+
%ret = atomicrmw fadd ptr addrspace(3) %dp3, half %val seq_cst
118+
ret void
119+
}
120+
121+
attributes #1 = { argmemonly nounwind }

‎llvm/test/CodeGen/NVPTX/atomics.ll

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -175,6 +175,13 @@ define float @atomicrmw_add_f32_generic(ptr %addr, float %val) {
175175
ret float %ret
176176
}
177177

178+
; CHECK-LABEL: atomicrmw_add_f16_generic
179+
define half @atomicrmw_add_f16_generic(ptr %addr, half %val) {
180+
; CHECK: atom.cas
181+
%ret = atomicrmw fadd ptr %addr, half %val seq_cst
182+
ret half %ret
183+
}
184+
178185
; CHECK-LABEL: atomicrmw_add_f32_addrspace1
179186
define float @atomicrmw_add_f32_addrspace1(ptr addrspace(1) %addr, float %val) {
180187
; CHECK: atom.global.add.f32

0 commit comments

Comments
 (0)