Move isspacep intrinsics processing to NVPTX backend

Author: Artem-B

llvm/lib/Analysis/InstructionSimplify.cpp

@@ -6367,32 +6367,6 @@ static Value *simplifyUnaryIntrinsic(Function *F, Value *Op0,
 
     break;
   }
-  case Intrinsic::nvvm_isspacep_global:
-  case Intrinsic::nvvm_isspacep_local:
-  case Intrinsic::nvvm_isspacep_shared:
-  case Intrinsic::nvvm_isspacep_const: {
-    auto *Ty = F->getReturnType();
-    unsigned AS = Op0->getType()->getPointerAddressSpace();
-    if (AS == NVPTXAS::ADDRESS_SPACE_GENERIC) {
-      if (auto *ASC = dyn_cast<AddrSpaceCastInst>(Op0))
-        AS = ASC->getSrcAddressSpace();
-      else if (auto *ASCO = dyn_cast<AddrSpaceCastOperator>(Op0))
-        AS = ASCO->getOperand(0)->getType()->getPointerAddressSpace();
-    }
-    if (AS == NVPTXAS::ADDRESS_SPACE_GENERIC ||
-        AS == NVPTXAS::ADDRESS_SPACE_PARAM)
-      return nullptr; // Got to check at run-time.
-    bool ASMatches = (AS == NVPTXAS::ADDRESS_SPACE_GLOBAL &&
-                      IID == Intrinsic::nvvm_isspacep_global) ||
-                     (AS == NVPTXAS::ADDRESS_SPACE_LOCAL &&
-                      IID == Intrinsic::nvvm_isspacep_local) ||
-                     (AS == NVPTXAS::ADDRESS_SPACE_SHARED &&
-                      IID == Intrinsic::nvvm_isspacep_shared) ||
-                     (AS == NVPTXAS::ADDRESS_SPACE_CONST &&
-                      IID == Intrinsic::nvvm_isspacep_const);
-    return ConstantInt::get(Ty, ASMatches);
-    break;
-  }
   default:
     break;
   }

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;
 }