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

The change improves the code in general and, as a side effect, avoids crashing
on an impossible address space casts guarded by `__isGlobal/__isShared`, which
partially fixes llvm#112760
It's still possible to trigger the issue by using explicit AS casts w/o
AS checks, but LLVM should no longer crash on valid code.

Author: Artem-B

llvm/include/llvm/Support/NVPTXAddrSpace.h

@@ -0,0 +1,33 @@
+//===---------------- NVPTXAddrSpace.h -------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// NVPTX address space definition
+///
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SUPPORT_NVPTXADDRSPACE_H
+#define LLVM_SUPPORT_NVPTXADDRSPACE_H
+
+namespace llvm {
+namespace NVPTXAS {
+enum AddressSpace : unsigned {
+  ADDRESS_SPACE_GENERIC = 0,
+  ADDRESS_SPACE_GLOBAL = 1,
+  ADDRESS_SPACE_SHARED = 3,
+  ADDRESS_SPACE_CONST = 4,
+  ADDRESS_SPACE_LOCAL = 5,
+
+  ADDRESS_SPACE_PARAM = 101,
+};
+} // end namespace NVPTXAS
+
+} // end namespace llvm
+
+#endif // LLVM_SUPPORT_NVPTXADDRSPACE_H

llvm/lib/Target/NVPTX/MCTargetDesc/NVPTXBaseInfo.h

@@ -16,18 +16,10 @@
 #ifndef LLVM_LIB_TARGET_NVPTX_MCTARGETDESC_NVPTXBASEINFO_H
 #define LLVM_LIB_TARGET_NVPTX_MCTARGETDESC_NVPTXBASEINFO_H
 
+#include "llvm/Support/NVPTXAddrSpace.h"
 namespace llvm {
 
-enum AddressSpace {
-  ADDRESS_SPACE_GENERIC = 0,
-  ADDRESS_SPACE_GLOBAL = 1,
-  ADDRESS_SPACE_SHARED = 3,
-  ADDRESS_SPACE_CONST = 4,
-  ADDRESS_SPACE_LOCAL = 5,
-
-  // NVVM Internal
-  ADDRESS_SPACE_PARAM = 101
-};
+using namespace NVPTXAS;
 
 namespace NVPTXII {
 enum {

llvm/lib/Target/NVPTX/NVPTXTargetTransformInfo.cpp

@@ -15,10 +15,12 @@
 #include "llvm/CodeGen/CostTable.h"
 #include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/IR/Constants.h"
+#include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/IntrinsicsNVPTX.h"
 #include "llvm/IR/Value.h"
 #include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Transforms/InstCombine/InstCombiner.h"
 #include <optional>
 using namespace llvm;
@@ -117,7 +119,8 @@ bool NVPTXTTIImpl::isSourceOfDivergence(const Value *V) {
 }
 
 // Convert NVVM intrinsics to target-generic LLVM code where possible.
-static Instruction *simplifyNvvmIntrinsic(IntrinsicInst *II, InstCombiner &IC) {
+static Instruction *convertNvvmIntrinsicToLlvm(InstCombiner &IC,
+                                               IntrinsicInst *II) {
   // Each NVVM intrinsic we can simplify can be replaced with one of:
   //
   //  * an LLVM intrinsic,
@@ -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.
+    if (AS == NVPTXAS::ADDRESS_SPACE_GENERIC)
+      if (auto *ASCO = dyn_cast<AddrSpaceCastOperator>(Op0))
+        AS = ASCO->getOperand(0)->getType()->getPointerAddressSpace();
+
+    if (std::optional<bool> Answer = CheckASMatch(IID, AS))
+      return IC.replaceInstUsesWith(II, ConstantInt::get(Ty, *Answer));
+    return nullptr; // Don't know the answer, got to check at run time.
+  }
+  default:
+    return std::nullopt;
+  }
+}
+
 std::optional<Instruction *>
 NVPTXTTIImpl::instCombineIntrinsic(InstCombiner &IC, IntrinsicInst &II) const {
-  if (Instruction *I = simplifyNvvmIntrinsic(&II, IC)) {
+  if (std::optional<Instruction *> I = handleSpaceCheckIntrinsics(IC, II))
+    return *I;
+  if (Instruction *I = convertNvvmIntrinsicToLlvm(IC, &II))
     return I;
-  }
+
   return std::nullopt;
 }