[AMDGPU][LowerBufferFatPointers] Fix lack of rewrite when loading/storing null (#154128)

Fixes #154056.

The fat buffer lowering pass was erroniously detecting that it did not
need to run on functions that only load/store to the null constant (or
other such constants). We thought this would be covered by specializing
constants out to instructions, but that doesn't account foc trivial
constants like null. Therefore, we check the operands of instructions
for buffer fat pointers in order to find such constants and ensure the
pass runs.

---------

Co-authored-by: Nikita Popov <[email protected]>

Author: krzysz00

llvm/lib/Target/AMDGPU/AMDGPULowerBufferFatPointers.cpp

@@ -2366,8 +2366,12 @@ static bool containsBufferFatPointers(const Function &F,
                                       BufferFatPtrToStructTypeMap *TypeMap) {
   bool HasFatPointers = false;
   for (const BasicBlock &BB : F)
-    for (const Instruction &I : BB)
+    for (const Instruction &I : BB) {
       HasFatPointers |= (I.getType() != TypeMap->remapType(I.getType()));
+      // Catch null pointer constants in loads, stores, etc.
+      for (const Value *V : I.operand_values())
+        HasFatPointers |= (V->getType() != TypeMap->remapType(V->getType()));
+    }
   return HasFatPointers;
 }
 

llvm/test/CodeGen/AMDGPU/lower-buffer-fat-pointers-constants.ll

@@ -223,3 +223,43 @@ define i32 @fancy_zero() {
   ptr addrspace(7) addrspacecast (ptr addrspace(8) @buf to ptr addrspace(7))
   to i32)
 }
+
+define i32 @load_null() {
+; CHECK-LABEL: define i32 @load_null
+; CHECK-SAME: () #[[ATTR0]] {
+; CHECK-NEXT:    [[X:%.*]] = call i32 @llvm.amdgcn.raw.ptr.buffer.load.i32(ptr addrspace(8) align 4 null, i32 0, i32 0, i32 0)
+; CHECK-NEXT:    ret i32 [[X]]
+;
+  %x = load i32, ptr addrspace(7) null, align 4
+  ret i32 %x
+}
+
+define void @store_null() {
+; CHECK-LABEL: define void @store_null
+; CHECK-SAME: () #[[ATTR0]] {
+; CHECK-NEXT:    call void @llvm.amdgcn.raw.ptr.buffer.store.i32(i32 0, ptr addrspace(8) align 4 null, i32 0, i32 0, i32 0)
+; CHECK-NEXT:    ret void
+;
+  store i32 0, ptr addrspace(7) null, align 4
+  ret void
+}
+
+define i32 @load_poison() {
+; CHECK-LABEL: define i32 @load_poison
+; CHECK-SAME: () #[[ATTR0]] {
+; CHECK-NEXT:    [[X:%.*]] = call i32 @llvm.amdgcn.raw.ptr.buffer.load.i32(ptr addrspace(8) align 4 poison, i32 poison, i32 0, i32 0)
+; CHECK-NEXT:    ret i32 [[X]]
+;
+  %x = load i32, ptr addrspace(7) poison, align 4
+  ret i32 %x
+}
+
+define void @store_poison() {
+; CHECK-LABEL: define void @store_poison
+; CHECK-SAME: () #[[ATTR0]] {
+; CHECK-NEXT:    call void @llvm.amdgcn.raw.ptr.buffer.store.i32(i32 0, ptr addrspace(8) align 4 poison, i32 poison, i32 0, i32 0)
+; CHECK-NEXT:    ret void
+;
+  store i32 0, ptr addrspace(7) poison, align 4
+  ret void
+}