[llvm] [NVPTX] instcombine known pointer AS checks. (PR #112964)

[Alex MacLean via llvm-commits]

================
@@ -413,11 +416,65 @@ static Instruction *simplifyNvvmIntrinsic(IntrinsicInst *II, InstCombiner &IC) {
   llvm_unreachable("All SpecialCase enumerators should be handled in switch.");
 }
 
+// Returns an instruction pointer (may be nullptr if we do not know the answer).
+// Returns nullopt if `II` is not one of the `isspacep` intrinsics.
+static std::optional<Instruction *>
+handleSpaceCheckIntrinsics(InstCombiner &IC, IntrinsicInst &II) {
+  Value *Op0 = II.getArgOperand(0);
+  // Returns true/false when we know the answer, nullopt otherwise.
+  auto CheckASMatch = [](unsigned IID, unsigned AS) -> std::optional<bool> {
+    if (AS == NVPTXAS::ADDRESS_SPACE_GENERIC ||
+        AS == NVPTXAS::ADDRESS_SPACE_PARAM)
+      return std::nullopt; // Got to check at run-time.
+    switch (IID) {
+    case Intrinsic::nvvm_isspacep_global:
+      return AS == NVPTXAS::ADDRESS_SPACE_GLOBAL;
+    case Intrinsic::nvvm_isspacep_local:
+      return AS == NVPTXAS::ADDRESS_SPACE_LOCAL;
+    case Intrinsic::nvvm_isspacep_shared:
+      return AS == NVPTXAS::ADDRESS_SPACE_SHARED;
+    case Intrinsic::nvvm_isspacep_shared_cluster:
+      // We can't tell shared from shared_cluster at compile time from AS alone,
+      // but it can't be either is AS is not shared.
+      return AS == NVPTXAS::ADDRESS_SPACE_SHARED ? std::nullopt
+                                                 : std::optional{false};
+    case Intrinsic::nvvm_isspacep_const:
+      return AS == NVPTXAS::ADDRESS_SPACE_CONST;
+    default:
+      llvm_unreachable("Unexpected intrinsic");
+    }
+  };
+
+  switch (auto IID = II.getIntrinsicID()) {
+  case Intrinsic::nvvm_isspacep_global:
+  case Intrinsic::nvvm_isspacep_local:
+  case Intrinsic::nvvm_isspacep_shared:
+  case Intrinsic::nvvm_isspacep_shared_cluster:
+  case Intrinsic::nvvm_isspacep_const: {
+    auto *Ty = II.getType();
+    unsigned AS = Op0->getType()->getPointerAddressSpace();
+    // Peek through ASC to generic AS.
+    // TODO: we could dig deeper through both ASCs and GEPs.
----------------
AlexMaclean wrote:

I specifically mean that this same logic should also be happening in `rewriteIntrinsicWithAddressSpace` when we propagate a new AS into one of these isspacep intrinsics. I agree it can be in a separate patch though. 

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