[llvm] [InferAddressSpaces] Handle unconverted ptrmask (PR #140802)

[Robert Imschweiler via llvm-commits]

https://github.com/ro-i created https://github.com/llvm/llvm-project/pull/140802

In case a ptrmask cannot be converted to the new address space due to an unknown mask value, this needs to be detcted and an addrspacecast is needed to not hinder a future use of the unconverted return value of ptrmask. Otherwise, users of this value will become invalid by receiving a nullptr as an operand.

This LLVM defect was identified via the AMD Fuzzing project.

(See https://reviews.llvm.org/D80129 for an explanation of why some ptrmasks are impossible to convert to other addrspaces.)

>From 7e2f064c3323933a18f1f71e3bc998fcce90cd4a Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Tue, 20 May 2025 15:55:36 -0500
Subject: [PATCH] [InferAddressSpaces] Handle unconverted ptrmask

In case a ptrmask cannot be converted to the new address space due to an
unknown mask value, this needs to be detcted and an addrspacecast is
needed to not hinder a future use of the unconverted return value of
ptrmask. Otherwise, users of this value will become invalid by receiving
a nullptr as an operand.

This LLVM defect was identified via the AMD Fuzzing project.
---
 .../Transforms/Scalar/InferAddressSpaces.cpp   | 15 ++++++++++++++-
 .../InferAddressSpaces/AMDGPU/ptrmask.ll       | 18 ++++++++++++++++++
 2 files changed, 32 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
index d3771c0903456..4f2e8bbd1102a 100644
--- a/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
+++ b/llvm/lib/Transforms/Scalar/InferAddressSpaces.cpp
@@ -1338,7 +1338,20 @@ bool InferAddressSpacesImpl::rewriteWithNewAddressSpaces(
 
     unsigned OperandNo = PoisonUse->getOperandNo();
     assert(isa<PoisonValue>(NewV->getOperand(OperandNo)));
-    NewV->setOperand(OperandNo, ValueWithNewAddrSpace.lookup(PoisonUse->get()));
+    WeakTrackingVH NewOp = ValueWithNewAddrSpace.lookup(PoisonUse->get());
+    if (NewOp) {
+      NewV->setOperand(OperandNo, NewOp);
+    } else {
+      // Something went wrong while converting the instruction defining the new
+      // operand value.  -> Replace the poison value with the previous operand
+      // value combined with an addrspace case.
+      Value *PoisonOp = NewV->getOperand(OperandNo);
+      Value *OldOp = V->getOperand(OperandNo);
+      Value *AddrSpaceCast =
+          new AddrSpaceCastInst(OldOp, PoisonOp->getType(), "",
+                                cast<Instruction>(NewV)->getIterator());
+      NewV->setOperand(OperandNo, AddrSpaceCast);
+    }
   }
 
   SmallVector<Instruction *, 16> DeadInstructions;
diff --git a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/ptrmask.ll b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/ptrmask.ll
index 6ef926f935830..1c1d1df79520d 100644
--- a/llvm/test/Transforms/InferAddressSpaces/AMDGPU/ptrmask.ll
+++ b/llvm/test/Transforms/InferAddressSpaces/AMDGPU/ptrmask.ll
@@ -343,6 +343,24 @@ define i8 @ptrmask_cast_local_to_flat_load_range_mask(ptr addrspace(3) %src.ptr,
   ret i8 %load
 }
 
+; Non-const masks with no known range should not prevent other ptr-manipulating
+; instructions (such as gep) from being converted.
+define i8 @ptrmask_cast_local_to_flat_unknown_mask(ptr addrspace(3) %src.ptr, i64 %mask, i64 %idx) {
+; CHECK-LABEL: @ptrmask_cast_local_to_flat_unknown_mask(
+; CHECK-NEXT:    [[CAST:%.*]] = addrspacecast ptr addrspace(3) [[SRC_PTR:%.*]] to ptr
+; CHECK-NEXT:    [[MASKED:%.*]] = call ptr @llvm.ptrmask.p0.i64(ptr [[CAST]], i64 [[MASK:%.*]])
+; CHECK-NEXT:    [[TMP1:%.*]] = addrspacecast ptr [[MASKED]] to ptr addrspace(3)
+; CHECK-NEXT:    [[GEP:%.*]] = getelementptr i8, ptr addrspace(3) [[TMP1]], i64 [[IDX:%.*]]
+; CHECK-NEXT:    [[LOAD:%.*]] = load i8, ptr addrspace(3) [[GEP]], align 1
+; CHECK-NEXT:    ret i8 [[LOAD]]
+;
+  %cast = addrspacecast ptr addrspace(3) %src.ptr to ptr
+  %masked = call ptr @llvm.ptrmask.p0.i64(ptr %cast, i64 %mask)
+  %gep = getelementptr i8, ptr %masked, i64 %idx
+  %load = load i8, ptr %gep
+  ret i8 %load
+}
+
 declare ptr @llvm.ptrmask.p0.i64(ptr, i64) #0
 declare ptr addrspace(5) @llvm.ptrmask.p5.i32(ptr addrspace(5), i32) #0
 declare ptr addrspace(3) @llvm.ptrmask.p3.i32(ptr addrspace(3), i32) #0