[AMDGPU] Check for NaN when folding output modifiers

[bcahoon]

Summary

Hardware ignores output modifiers if the IEEE mode is set. If the
no NaN flag is set on an instruction, then the output modifier
folding optimization can assume that the IEEE mode is not set.

[arsenm]

Oh whoops, I forgot the hardware totally ignores the modifiers with IEEE mode on (not that this behavior makes any sense). We can't actually do this

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

for (I = MBB->begin(); I != MBB->end(); I = Next) {
       if (MI.mayLoad() && tryFoldLoad(MI))
         continue;
 
       if (!TII->isFoldableCopy(MI)) {
         // Saw an unknown clobber of m0, so we no longer know what it is.
         if (CurrentKnownM0Val && MI.modifiesRegister(AMDGPU::M0, TRI))
           CurrentKnownM0Val = nullptr;
 
+        // The IEEE mode value can be overridden, which is detected by checking
+        // the NaN flag on the instruction.
         // TODO: Omod might be OK if there is NSZ only on the source
         // instruction, and not the omod multiply.
-        if (IsIEEEMode || (!HasNSZ && !MI.getFlag(MachineInstr::FmNsz)) ||
-            !tryFoldOMod(MI))
+        if ((IsIEEEMode && !MI.getFlag(MachineInstr::FmNoNans)) ||
+            (!HasNSZ && !MI.getFlag(MachineInstr::FmNsz)) || !tryFoldOMod(MI))
           tryFoldClamp(MI);
 
         continue;
       }
 
       // Specially track simple redefs of m0 to the same value in a block, so we
       // can erase the later ones.
       if (MI.getOperand(0).getReg() == AMDGPU::M0) {

llvm/test/CodeGen/AMDGPU/omod.ll

 ; GCN: v_add_f64  v{{\[[0-9]+:[0-9]+\]}}, v{{\[[0-9]+:[0-9]+\]}}, 1.0 div:2{{$}}
 define amdgpu_ps void @v_omod_div2_f64(double %a) #5 {
   %add = fadd nsz double %a, 1.0
   %div2 = fmul nsz double %add, 0.5
   store double %div2, double addrspace(1)* undef
   ret void
 }
 
+; GCN-LABEL: {{^}}v_omod_div2_f64_nnan:
+; GCN: v_add_f64  v{{\[[0-9]+:[0-9]+\]}}, s{{\[[0-9]+:[0-9]+\]}}, 1.0 div:2{{$}}
+define amdgpu_kernel void @v_omod_div2_f64_nnan(double %a) #5 {
+  %add = fadd nnan nsz double %a, 1.0
+  %div2 = fmul nnan nsz double %add, 0.5
+  store double %div2, double addrspace(1)* undef
+  ret void
+}
+
 ; GCN-LABEL: {{^}}v_omod_mul2_f32:
 ; GCN: v_add_f32_e64 v{{[0-9]+}}, v0, 1.0 mul:2{{$}}
 define amdgpu_ps void @v_omod_mul2_f32(float %a) #0 {
   %add = fadd float %a, 1.0
   %div2 = fmul float %add, 2.0
   store float %div2, float addrspace(1)* undef
   ret void
 }
 
 ; GCN-LABEL: {{^}}v_omod_mul2_f64:
 ; GCN: v_add_f64  v{{\[[0-9]+:[0-9]+\]}}, v{{\[[0-9]+:[0-9]+\]}}, 1.0 mul:2{{$}}
 define amdgpu_ps void @v_omod_mul2_f64(double %a) #5 {
   %add = fadd nsz double %a, 1.0
   %div2 = fmul nsz double %add, 2.0
   store double %div2, double addrspace(1)* undef
   ret void
 }
 
+; GCN-LABEL: {{^}}v_omod_mul2_f64_nnan:
+; GCN: v_add_f64  v{{\[[0-9]+:[0-9]+\]}}, s{{\[[0-9]+:[0-9]+\]}}, 1.0 mul:2{{$}}
+define amdgpu_kernel void @v_omod_mul2_f64_nnan(double %a) #5 {
+  %add = fadd nnan nsz double %a, 1.0
+  %div2 = fmul nnan nsz double %add, 2.0
+  store double %div2, double addrspace(1)* undef
+  ret void
+}
+
 ; GCN-LABEL: {{^}}v_omod_mul4_f32:
 ; GCN: v_add_f32_e64 v{{[0-9]+}}, v0, 1.0 mul:4{{$}}
 define amdgpu_ps void @v_omod_mul4_f32(float %a) #0 {
   %add = fadd float %a, 1.0
   %div2 = fmul float %add, 4.0
   store float %div2, float addrspace(1)* undef
   ret void
 }
 
 ; GCN-LABEL: {{^}}v_omod_mul4_f64:
 ; GCN: v_add_f64  v{{\[[0-9]+:[0-9]+\]}}, v{{\[[0-9]+:[0-9]+\]}}, 1.0 mul:4{{$}}
 define amdgpu_ps void @v_omod_mul4_f64(double %a) #5 {
   %add = fadd nsz double %a, 1.0
   %div2 = fmul nsz double %add, 4.0
   store double %div2, double addrspace(1)* undef
   ret void
 }
 
+; GCN-LABEL: {{^}}v_omod_mul4_f64_nnan:
+; GCN: v_add_f64  v{{\[[0-9]+:[0-9]+\]}}, s{{\[[0-9]+:[0-9]+\]}}, 1.0 mul:4{{$}}
+define amdgpu_kernel void @v_omod_mul4_f64_nnan(double %a) #5 {
+  %add = fadd nnan nsz double %a, 1.0
+  %div2 = fmul nnan nsz double %add, 4.0
+  store double %div2, double addrspace(1)* undef
+  ret void
+}
+
 ; GCN-LABEL: {{^}}v_omod_mul4_multi_use_f32:
 ; GCN: v_add_f32_e32 [[ADD:v[0-9]+]], 1.0, v0{{$}}
 ; GCN: v_mul_f32_e32 v{{[0-9]+}}, 4.0, [[ADD]]{{$}}
 define amdgpu_ps void @v_omod_mul4_multi_use_f32(float %a) #0 {
   %add = fadd float %a, 1.0
   %div2 = fmul float %add, 4.0
   store float %div2, float addrspace(1)* undef
   store volatile float %add, float addrspace(1)* undef