Use entry points to help determine the env.

Author: maarquitos14

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

@@ -272,14 +272,37 @@ getExecutionModel(const SPIRVSubtarget &STI, const Function &F) {
   // precise enough. For now, we'll rely instead on `isLogicalSPIRV()`, but this
   // should be changed when `isOpenCLEnv()` and `isVulkanEnv()` cannot be true
   // at the same time.
-  if (!STI.isLogicalSPIRV())
+  if (STI.isOpenCLEnv())
     return SPIRV::ExecutionModel::Kernel;
 
+  if (STI.isVulkanEnv()) {
+    auto attribute = F.getFnAttribute("hlsl.shader");
+    if (!attribute.isValid()) {
+      report_fatal_error(
+          "This entry point lacks mandatory hlsl.shader attribute.");
+    }
+
+    const auto value = attribute.getValueAsString();
+    if (value == "compute")
+      return SPIRV::ExecutionModel::GLCompute;
+
+    report_fatal_error(
+        "This HLSL entry point is not supported by this backend.");
+  }
+
+  assert(STI.getEnv() == Unknown);
+  // Can we rely on "hlsl.shader" attribute? Is it mandatory for Vulkan env? If
+  // so, we can set Env to Vulkan whenever we find it, and to OpenCL otherwise.
+
+  // We will now change the Env based on the attribute, so we need to strip
+  // `const` out of the ref to STI.
+  SPIRVSubtarget *NonConstSTI = const_cast<SPIRVSubtarget *>(&STI);
   auto attribute = F.getFnAttribute("hlsl.shader");
   if (!attribute.isValid()) {
-    report_fatal_error(
-        "This entry point lacks mandatory hlsl.shader attribute.");
+    NonConstSTI->setEnv(SPIRVSubtarget::OpenCL);
+    return SPIRV::ExecutionModel::Kernel;
   }
+  NonConstSTI->setEnv(SPIRVSubtarget::Vulkan);
 
   const auto value = attribute.getValueAsString();
   if (value == "compute")
@@ -444,6 +467,11 @@ bool SPIRVCallLowering::lowerFormalArguments(MachineIRBuilder &MIRBuilder,
 
   // Handle entry points and function linkage.
   if (isEntryPoint(F)) {
+    // EntryPoints can help us to determine the environment we're working on.
+    // Therefore, we need a non-const pointer to SPIRVSubtarget to update the
+    // environment if we need to.
+    const SPIRVSubtarget *ST =
+        static_cast<const SPIRVSubtarget *>(&MIRBuilder.getMF().getSubtarget());
     auto MIB = MIRBuilder.buildInstr(SPIRV::OpEntryPoint)
                    .addImm(static_cast<uint32_t>(getExecutionModel(*ST, F)))
                    .addUse(FuncVReg);

llvm/lib/Target/SPIRV/SPIRVModuleAnalysis.cpp

@@ -151,11 +151,11 @@ void SPIRVModuleAnalysis::setBaseInfo(const Module &M) {
     // TODO: Add support for VulkanMemoryModel.
     // FIXME: At the moment, there's a possibility that both `isOpenCLEnv()` and
     // `isVulkanEnv()` return true. This is because the Triple is not always
-    // precise enough. For now, we'll rely instead on `isLogicalSPIRV()`, but this
-    // should be changed when `isOpenCLEnv()` and `isVulkanEnv()` cannot be true
-    // at the same time.
+    // precise enough. For now, we'll rely instead on `isLogicalSPIRV()`, but
+    // this should be changed when `isOpenCLEnv()` and `isVulkanEnv()` cannot be
+    // true at the same time.
     MAI.Mem = !ST->isLogicalSPIRV() ? SPIRV::MemoryModel::OpenCL
-                                : SPIRV::MemoryModel::GLSL450;
+                                    : SPIRV::MemoryModel::GLSL450;
     if (MAI.Mem == SPIRV::MemoryModel::OpenCL) {
       unsigned PtrSize = ST->getPointerSize();
       MAI.Addr = PtrSize == 32   ? SPIRV::AddressingModel::Physical32
@@ -1808,7 +1808,11 @@ void addInstrRequirements(const MachineInstr &MI,
     // not allowed to produce
     // StorageImageReadWithoutFormat/StorageImageWriteWithoutFormat, see
     // https://github.com/KhronosGroup/SPIRV-Headers/issues/487
-    if (isImageTypeWithUnknownFormat(TypeDef) && !ST.isOpenCLEnv())
+
+    // FIXME: For now, `isOpenCLEnv()` is not precise enough. Instead, we're
+    // using `isLogicalSPIRV()`, but we should change this when `isOpenCLEnv()`
+    // is precise enough.
+    if (isImageTypeWithUnknownFormat(TypeDef) && ST.isLogicalSPIRV())
       Reqs.addCapability(SPIRV::Capability::StorageImageReadWithoutFormat);
     break;
   }
@@ -1821,7 +1825,11 @@ void addInstrRequirements(const MachineInstr &MI,
     // not allowed to produce
     // StorageImageReadWithoutFormat/StorageImageWriteWithoutFormat, see
     // https://github.com/KhronosGroup/SPIRV-Headers/issues/487
-    if (isImageTypeWithUnknownFormat(TypeDef) && !ST.isOpenCLEnv())
+
+    // FIXME: For now, `isOpenCLEnv()` is not precise enough. Instead, we're
+    // using `isLogicalSPIRV()`, but we should change this when `isOpenCLEnv()`
+    // is precise enough.
+    if (isImageTypeWithUnknownFormat(TypeDef) && ST.isLogicalSPIRV())
       Reqs.addCapability(SPIRV::Capability::StorageImageWriteWithoutFormat);
     break;
   }

llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp

@@ -83,6 +83,14 @@ SPIRVSubtarget::SPIRVSubtarget(const Triple &TT, const std::string &CPU,
   }
   OpenCLVersion = VersionTuple(2, 2);
 
+  // Set the environment based on the target triple.
+  if (TargetTriple.getOS() == Triple::Vulkan)
+    Env = Vulkan;
+  else if (TargetTriple.getEnvironment() == Triple::OpenCL)
+    Env = OpenCL;
+  else
+    Env = Unknown;
+
   // The order of initialization is important.
   initAvailableExtensions(Extensions);
   initAvailableExtInstSets();

llvm/lib/Target/SPIRV/SPIRVSubtarget.h

@@ -36,8 +36,10 @@ class StringRef;
 class SPIRVTargetMachine;
 
 class SPIRVSubtarget : public SPIRVGenSubtargetInfo {
+public:
   // Enum for the SPIR-V environment: OpenCL, Vulkan or Unkwnown.
   enum SPIRVEnvType { OpenCL, Vulkan, Unknown };
+
 private:
   const unsigned PointerSize;
   VersionTuple SPIRVVersion;
@@ -81,28 +83,21 @@ class SPIRVSubtarget : public SPIRVGenSubtargetInfo {
   unsigned getPointerSize() const { return PointerSize; }
   unsigned getBound() const { return GR->getBound(); }
   bool canDirectlyComparePointers() const;
-  SPIRVEnvType getEnv() const {
-    if (TargetTriple.getOS() == Triple::Vulkan)
-      return Vulkan;
-    if (TargetTriple.getEnvironment() == Triple::OpenCL)
-      return OpenCL;
-    return Unknown;
-  }
-  // FIXME: For now, we rely only on the triple to determine the environment.
-  // However, a lot of frontends emit unknown OS or environment, which makes it
-  // difficult to determine the environment. We should consider adding other
-  // methods. For now, we will return true for both OpenCL and Unknown env.
-  bool isOpenCLEnv() const {
-    return getEnv() == OpenCL || getEnv() == Unknown;
-  }
-  // FIXME: For now, we rely only on the triple to determine the environment.
-  // However, a lot of frontends emit unknown OS or environment, which makes it
-  // difficult to determine the environment. We should consider adding other
-  // methods. For now, we will return true for both Vulkan and Unknown env.
-  bool isVulkanEnv() const {
-    return getEnv() == Vulkan || getEnv() == Unknown;
+  void setEnv(SPIRVEnvType E) {
+    assert(E != Unknown && "Unknown environment is not allowed");
+    assert(Env == Unknown && "Environment is already set");
+
+    Env = E;
   }
-  bool isLogicalSPIRV() const { return TargetTriple.getArch() == Triple::spirv; }
+  SPIRVEnvType getEnv() const { return Env; }
+  bool isOpenCLEnv() const { return getEnv() == OpenCL; }
+  bool isVulkanEnv() const { return getEnv() == Vulkan; }
+  // FIXME: This should check the triple arch instead, but a lot of places use
+  // this method now instead of `is[OpenCL/Vulkan]Env()`, and this is a
+  // shortcut to make sure `is[OpenCL/Vulkan]Env()` works as expected. When we
+  // change back all uses of `isLogicalSPIRV()` to `is[OpenCL/Vulkan]Env()`, we
+  // can implement this correctly again.
+  bool isLogicalSPIRV() const { return isVulkanEnv(); }
   const std::string &getTargetTripleAsStr() const { return TargetTriple.str(); }
   VersionTuple getSPIRVVersion() const { return SPIRVVersion; };
   bool isAtLeastSPIRVVer(VersionTuple VerToCompareTo) const;