[llvm] [NVPTX] Respect FTZ flag when lowering atomicrmw fadd. (PR #200732)

[via llvm-commits]

================
@@ -1737,152 +1739,290 @@ define i64 @usub_sat_acq_rel_i64_global_cta(ptr addrspace(1) %addr, i64 %val) {
 }
 
 define float @fadd_acq_rel_float_global_cta(ptr addrspace(1) %addr, float %val) {
-; SM90-LABEL: fadd_acq_rel_float_global_cta(
-; SM90:       {
-; SM90-NEXT:    .reg .b32 %r<3>;
-; SM90-NEXT:    .reg .b64 %rd<2>;
-; SM90-EMPTY:
-; SM90-NEXT:  // %bb.0:
-; SM90-NEXT:    ld.param.b64 %rd1, [fadd_acq_rel_float_global_cta_param_0];
-; SM90-NEXT:    ld.param.b32 %r1, [fadd_acq_rel_float_global_cta_param_1];
-; SM90-NEXT:    atom.acq_rel.cta.global.add.f32 %r2, [%rd1], %r1;
-; SM90-NEXT:    st.param.b32 [func_retval0], %r2;
-; SM90-NEXT:    ret;
+; SM90-NOFTZ-LABEL: fadd_acq_rel_float_global_cta(
+; SM90-NOFTZ:       {
+; SM90-NOFTZ-NEXT:    .reg .pred %p<2>;
+; SM90-NOFTZ-NEXT:    .reg .b32 %r<5>;
+; SM90-NOFTZ-NEXT:    .reg .b64 %rd<2>;
+; SM90-NOFTZ-EMPTY:
+; SM90-NOFTZ-NEXT:  // %bb.0:
+; SM90-NOFTZ-NEXT:    ld.param.b32 %r2, [fadd_acq_rel_float_global_cta_param_1];
+; SM90-NOFTZ-NEXT:    ld.param.b64 %rd1, [fadd_acq_rel_float_global_cta_param_0];
+; SM90-NOFTZ-NEXT:    fence.release.cta;
+; SM90-NOFTZ-NEXT:    ld.relaxed.cta.global.b32 %r4, [%rd1];
+; SM90-NOFTZ-NEXT:  $L__BB60_1: // %atomicrmw.start
+; SM90-NOFTZ-NEXT:    // =>This Inner Loop Header: Depth=1
+; SM90-NOFTZ-NEXT:    add.rn.f32 %r3, %r4, %r2;
+; SM90-NOFTZ-NEXT:    atom.relaxed.cta.global.cas.b32 %r1, [%rd1], %r4, %r3;
+; SM90-NOFTZ-NEXT:    setp.ne.b32 %p1, %r1, %r4;
+; SM90-NOFTZ-NEXT:    mov.b32 %r4, %r1;
+; SM90-NOFTZ-NEXT:    @%p1 bra $L__BB60_1;
+; SM90-NOFTZ-NEXT:  // %bb.2: // %atomicrmw.end
+; SM90-NOFTZ-NEXT:    fence.acquire.cta;
+; SM90-NOFTZ-NEXT:    st.param.b32 [func_retval0], %r1;
+; SM90-NOFTZ-NEXT:    ret;
+;
+; SM90-FTZ-LABEL: fadd_acq_rel_float_global_cta(
+; SM90-FTZ:       {
+; SM90-FTZ-NEXT:    .reg .b32 %r<3>;
+; SM90-FTZ-NEXT:    .reg .b64 %rd<2>;
+; SM90-FTZ-EMPTY:
+; SM90-FTZ-NEXT:  // %bb.0:
+; SM90-FTZ-NEXT:    ld.param.b64 %rd1, [fadd_acq_rel_float_global_cta_param_0];
+; SM90-FTZ-NEXT:    ld.param.b32 %r1, [fadd_acq_rel_float_global_cta_param_1];
+; SM90-FTZ-NEXT:    atom.acq_rel.cta.global.add.f32 %r2, [%rd1], %r1;
+; SM90-FTZ-NEXT:    st.param.b32 [func_retval0], %r2;
+; SM90-FTZ-NEXT:    ret;
         %retval = atomicrmw fadd ptr  addrspace(1) %addr, float %val syncscope("block") acq_rel
         ret float %retval
 }
 
 define float @fsub_acq_rel_float_global_cta(ptr addrspace(1) %addr, float %val) {
-; SM90-LABEL: fsub_acq_rel_float_global_cta(
-; SM90:       {
-; SM90-NEXT:    .reg .pred %p<2>;
-; SM90-NEXT:    .reg .b32 %r<5>;
-; SM90-NEXT:    .reg .b64 %rd<2>;
-; SM90-EMPTY:
-; SM90-NEXT:  // %bb.0:
-; SM90-NEXT:    ld.param.b32 %r2, [fsub_acq_rel_float_global_cta_param_1];
-; SM90-NEXT:    ld.param.b64 %rd1, [fsub_acq_rel_float_global_cta_param_0];
-; SM90-NEXT:    fence.release.cta;
-; SM90-NEXT:    ld.relaxed.cta.global.b32 %r4, [%rd1];
-; SM90-NEXT:  $L__BB61_1: // %atomicrmw.start
-; SM90-NEXT:    // =>This Inner Loop Header: Depth=1
-; SM90-NEXT:    sub.rn.f32 %r3, %r4, %r2;
-; SM90-NEXT:    atom.relaxed.cta.global.cas.b32 %r1, [%rd1], %r4, %r3;
-; SM90-NEXT:    setp.ne.b32 %p1, %r1, %r4;
-; SM90-NEXT:    mov.b32 %r4, %r1;
-; SM90-NEXT:    @%p1 bra $L__BB61_1;
-; SM90-NEXT:  // %bb.2: // %atomicrmw.end
-; SM90-NEXT:    fence.acquire.cta;
-; SM90-NEXT:    st.param.b32 [func_retval0], %r1;
-; SM90-NEXT:    ret;
+; SM90-NOFTZ-LABEL: fsub_acq_rel_float_global_cta(
+; SM90-NOFTZ:       {
+; SM90-NOFTZ-NEXT:    .reg .pred %p<2>;
+; SM90-NOFTZ-NEXT:    .reg .b32 %r<5>;
+; SM90-NOFTZ-NEXT:    .reg .b64 %rd<2>;
+; SM90-NOFTZ-EMPTY:
+; SM90-NOFTZ-NEXT:  // %bb.0:
+; SM90-NOFTZ-NEXT:    ld.param.b32 %r2, [fsub_acq_rel_float_global_cta_param_1];
+; SM90-NOFTZ-NEXT:    ld.param.b64 %rd1, [fsub_acq_rel_float_global_cta_param_0];
+; SM90-NOFTZ-NEXT:    fence.release.cta;
+; SM90-NOFTZ-NEXT:    ld.relaxed.cta.global.b32 %r4, [%rd1];
+; SM90-NOFTZ-NEXT:  $L__BB61_1: // %atomicrmw.start
+; SM90-NOFTZ-NEXT:    // =>This Inner Loop Header: Depth=1
+; SM90-NOFTZ-NEXT:    sub.rn.f32 %r3, %r4, %r2;
+; SM90-NOFTZ-NEXT:    atom.relaxed.cta.global.cas.b32 %r1, [%rd1], %r4, %r3;
+; SM90-NOFTZ-NEXT:    setp.ne.b32 %p1, %r1, %r4;
+; SM90-NOFTZ-NEXT:    mov.b32 %r4, %r1;
+; SM90-NOFTZ-NEXT:    @%p1 bra $L__BB61_1;
+; SM90-NOFTZ-NEXT:  // %bb.2: // %atomicrmw.end
+; SM90-NOFTZ-NEXT:    fence.acquire.cta;
+; SM90-NOFTZ-NEXT:    st.param.b32 [func_retval0], %r1;
+; SM90-NOFTZ-NEXT:    ret;
+;
+; SM90-FTZ-LABEL: fsub_acq_rel_float_global_cta(
+; SM90-FTZ:       {
+; SM90-FTZ-NEXT:    .reg .pred %p<2>;
+; SM90-FTZ-NEXT:    .reg .b32 %r<5>;
+; SM90-FTZ-NEXT:    .reg .b64 %rd<2>;
+; SM90-FTZ-EMPTY:
+; SM90-FTZ-NEXT:  // %bb.0:
+; SM90-FTZ-NEXT:    ld.param.b32 %r2, [fsub_acq_rel_float_global_cta_param_1];
+; SM90-FTZ-NEXT:    ld.param.b64 %rd1, [fsub_acq_rel_float_global_cta_param_0];
+; SM90-FTZ-NEXT:    fence.release.cta;
+; SM90-FTZ-NEXT:    ld.relaxed.cta.global.b32 %r4, [%rd1];
+; SM90-FTZ-NEXT:  $L__BB61_1: // %atomicrmw.start
+; SM90-FTZ-NEXT:    // =>This Inner Loop Header: Depth=1
+; SM90-FTZ-NEXT:    sub.rn.ftz.f32 %r3, %r4, %r2;
----------------
gonzalobg wrote:

@akshayrdeodhar this is correct, but it probably would be better if this were to lower to `atom.add.global.f32` instead which performs ftz. Agree? (Not a request to do this in this PR, that should happen in a separate PR).

https://github.com/llvm/llvm-project/pull/200732