[AMDGPU] Add sanity check that fixes bad shift operation in AMD backend

[konradkusiak97]

Summary

There is a problem with the SILoadStoreOptimizer::dmasksCanBeCombined() function that can lead to UB.

This boolean function decides if two masks can be combined into 1. The idea here is that the bits which are "on" in one mask, don't overlap with the "on" bits of the other. Consider an example (10 bits for simplicity):

Mask 1: 0101101000
Mask 2: 0000000110

Those can be combined into a single mask: 0101101110.

To check if such an operation is possible, the code takes the mask which is greater and counts how many 0s there are, starting from the LSB and stopping at the first 1. Then, it shifts 1u by this number and compares it with the smaller mask. The problem is that when both masks are 0, the counter will find 32 zeroes in the first mask and will try to do a shift by 32 positions which leads to UB.

The fix is a simple sanity check, if the bigger mask is 0 or not.

[foad]

Seems reasonable.

Stray ] in first line of commit message, and the target is called "AMDGPU" not "AMD".

An alternative fix would be to avoid the shift altogether, by checking something like countl_zero(MinMask) + countr_zero(MaxMask) >= 32.

[JonChesterfield]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

877

Should this be if masks are equal?

Not sure sanity check describes the condition, maybe drop the comment

[konradkusiak97]

Dropped the comment

[konradkusiak97]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

877

Hm, If the masks are equal, the result will be returning false. I think that's the correct behaviour that the author of this function had in mind. That's based on the fact that there is <= and not < sign in if ((1u << AllowedBitsForMin) <= MinMask), so it really checks if the masks overlap - and two equal masks overlap fully.

[JonChesterfield]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

877

If equal masks imply they can't be combined, that should probably be true for the special case of equal masks that are zero.

I haven't looked at the call tree to determine what combining masks means in this context.

[foad]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

877

Firstly, there is no real use case for dmask=0. The resulting instructions would be invalid. All we have to do is handle it gracefully and not crash.

If equal masks imply they can't be combined, that should probably be true for the special case of equal masks that are zero.

I don't agree. This code is trying to check whether the range of bits set in one mask (from the lowest set bit to the highest including any gaps) overlaps with the range of set bits in the other. By that definition, two equal non-zero masks do overlap but two equal zero masks do not.

[JonChesterfield]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

877

It's called dmasksCanBeCombined though. If equal masks return false, and the zero case doesn't matter other than some arithmetic in the compiler, then masks equal zero returning false seems reasonable. Drive by review, not a blocking comment.

[foad]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

877

I still disagree. Returning false for all equal masks is no simpler to implement, and less easy to justify.

[konradkusiak97]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

877

I'm also more directed towards returning true in the case of both masks being 0. Because you can combine two 0 masks. Whereas two equal, non-zero masks can't be trivially combined because they overlap.

[JonChesterfield]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

877

We already return false for all other equal masks. I'm proposing if (!MaxMask) return false.

I'm running on the heuristic that special cases show up in bug reports and there's no test case in this diff. I don't see zero as special other than we're tripping over it in the compiler.

Usually (x64) shift 32 behaves as if masking the value, so the other UB fix is probably &31 in the right place, which would also return false (I think, haven't checked carefully)

[arsenm]

Testcase?

[konradkusiak97]

clang-format

[konradkusiak97]

included all local changes

[arsenm]

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

877

First I'd say countr_zero is just broken if this is UB.

Also needs testcase

[konradkusiak97]

Updated revision

[konradkusiak97]

I updated the fix to return false for this case. I've been trying to create a testcase for this but I'm still trying to figure out exactly how to write it so it triggers this specific case with 0 mask.

[arsenm]

In D155051#4569610, @konradkusiak97 wrote:

I updated the fix to return false for this case. I've been trying to create a testcase for this but I'm still trying to figure out exactly how to write it so it triggers this specific case with 0 mask.

You probably just have to write a MIR test

[konradkusiak97]

Created a testcase

[konradkusiak97]

Fixed the UB behaviour and included a testcase

[konradkusiak97]

Thanks @arsenm. As I don't have the commit access, could you land this patch for me? Please use "Konrad Kusiak konrad.kusiak@codeplay" to commit the change.

[arsenm]

4fa8a5487e3b1a4b2ce743b0008a912026aa3524

llvm/lib/Target/AMDGPU/SILoadStoreOptimizer.cpp

if (Idx != -1 &&
         CI.I->getOperand(Idx).getImm() != Paired.I->getOperand(Idx).getImm())
       return false;
   }
 
   // Check DMask for overlaps.
   unsigned MaxMask = std::max(CI.DMask, Paired.DMask);
   unsigned MinMask = std::min(CI.DMask, Paired.DMask);
 
+  if (!MaxMask)
+    return false;
+
   unsigned AllowedBitsForMin = llvm::countr_zero(MaxMask);
   if ((1u << AllowedBitsForMin) <= MinMask)
     return false;
 
   return true;
 }
 
 static unsigned getBufferFormatWithCompCount(unsigned OldFormat,
                                        unsigned ComponentCount,

llvm/test/CodeGen/AMDGPU/merge-image-load.mir

bb.0.entry:
     %3:sgpr_256 = S_LOAD_DWORDX8_IMM %1, 208, 0
     %4:vgpr_32 = COPY %2.sub3
     %5:vreg_128 = BUFFER_LOAD_DWORDX4_OFFSET %2:sgpr_128, 0, 0, 0, 0, implicit $exec :: (dereferenceable invariant load (s128))
     %6:vgpr_32 = IMAGE_LOAD_V1_V4 %5:vreg_128, %3:sgpr_256, 4, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s32), addrspace 4)
     %7:vreg_96 = IMAGE_LOAD_V3_V4 %5:vreg_128, %3:sgpr_256, 7, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s96), align 16, addrspace 4)
 ...
 ---
 
+# GFX9-LABEL: name: image_load_dmask_zero_not_merged
+# GFX9: %{{[0-9]+}}:vgpr_32 = IMAGE_LOAD_V1_V4 %5, %3, 0, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s32), addrspace 4)
+# GFX9: %{{[0-9]+}}:vreg_96 = IMAGE_LOAD_V3_V4 %5, %3, 0, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s96), align 16, addrspace 4)
+
+name: image_load_dmask_zero_not_merged
+body:             |
+  bb.0.entry:
+    %0:sgpr_64 = COPY $sgpr0_sgpr1
+    %1:sreg_64_xexec = S_LOAD_DWORDX2_IMM %0, 36, 0
+    %2:sgpr_128 = COPY $sgpr96_sgpr97_sgpr98_sgpr99
+    %3:sgpr_256 = S_LOAD_DWORDX8_IMM %1, 208, 0
+    %4:vgpr_32 = COPY %2.sub3
+    %5:vreg_128 = BUFFER_LOAD_DWORDX4_OFFSET %2:sgpr_128, 0, 0, 0, 0, implicit $exec :: (dereferenceable invariant load (s128))
+    %6:vgpr_32 = IMAGE_LOAD_V1_V4 %5:vreg_128, %3:sgpr_256, 0, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s32), addrspace 4)
+    %7:vreg_96 = IMAGE_LOAD_V3_V4 %5:vreg_128, %3:sgpr_256, 0, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s96), align 16, addrspace 4)
+...
+---
+
 # GFX9-LABEL: name: image_load_dmask_not_disjoint_not_merged
 # GFX9: %{{[0-9]+}}:vgpr_32 = IMAGE_LOAD_V1_V4 %5, %3, 4, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s32), addrspace 4)
 # GFX9: %{{[0-9]+}}:vreg_96 = IMAGE_LOAD_V3_V4 %5, %3, 11, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s96), align 16, addrspace 4)
 
 name: image_load_dmask_not_disjoint_not_merged
 body:             |
   bb.0.entry:
     %0:sgpr_64 = COPY $sgpr0_sgpr1

llvm/test/CodeGen/AMDGPU/merge-image-sample.mir

bb.0.entry:
     %3:sgpr_256 = S_LOAD_DWORDX8_IMM %1, 208, 0
     %4:vgpr_32 = COPY %2.sub3
     %5:vreg_128 = BUFFER_LOAD_DWORDX4_OFFSET %2:sgpr_128, 0, 0, 0, 0, implicit $exec :: (dereferenceable invariant load (s128))
     %6:vgpr_32 = IMAGE_SAMPLE_L_V1_V4 %5:vreg_128, %3:sgpr_256, %2:sgpr_128, 4, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s32), addrspace 4)
     %7:vreg_96 = IMAGE_SAMPLE_L_V3_V4 %5:vreg_128, %3:sgpr_256, %2:sgpr_128, 7, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s96), align 16, addrspace 4)
 ...
 ---
 
+# GFX9-LABEL: name: image_sample_l_dmask_zero_not_merged
+# GFX9: %{{[0-9]+}}:vgpr_32 = IMAGE_SAMPLE_L_V1_V4 %5, %3, %2, 0, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s32), addrspace 4)
+# GFX9: %{{[0-9]+}}:vreg_96 = IMAGE_SAMPLE_L_V3_V4 %5, %3, %2, 0, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s96), align 16, addrspace 4)
+
+name: image_sample_l_dmask_zero_not_merged
+body:             |
+  bb.0.entry:
+    %0:sgpr_64 = COPY $sgpr0_sgpr1
+    %1:sreg_64_xexec = S_LOAD_DWORDX2_IMM %0, 36, 0
+    %2:sgpr_128 = COPY $sgpr96_sgpr97_sgpr98_sgpr99
+    %3:sgpr_256 = S_LOAD_DWORDX8_IMM %1, 208, 0
+    %4:vgpr_32 = COPY %2.sub3
+    %5:vreg_128 = BUFFER_LOAD_DWORDX4_OFFSET %2:sgpr_128, 0, 0, 0, 0, implicit $exec :: (dereferenceable invariant load (s128))
+    %6:vgpr_32 = IMAGE_SAMPLE_L_V1_V4 %5:vreg_128, %3:sgpr_256, %2:sgpr_128, 0, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s32), addrspace 4)
+    %7:vreg_96 = IMAGE_SAMPLE_L_V3_V4 %5:vreg_128, %3:sgpr_256, %2:sgpr_128, 0, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s96), align 16, addrspace 4)
+...
+---
+
 # GFX9-LABEL: name: image_sample_l_dmask_not_disjoint_not_merged
 # GFX9: %{{[0-9]+}}:vgpr_32 = IMAGE_SAMPLE_L_V1_V4 %5, %3, %2, 4, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s32), addrspace 4)
 # GFX9: %{{[0-9]+}}:vreg_96 = IMAGE_SAMPLE_L_V3_V4 %5, %3, %2, 11, 0, 0, 0, 0, 0, -1, 0, implicit $exec :: (dereferenceable load (s96), align 16, addrspace 4)
 
 name: image_sample_l_dmask_not_disjoint_not_merged
 body:             |
   bb.0.entry:
     %0:sgpr_64 = COPY $sgpr0_sgpr1