[SPIR-V] Fix environment resolution causing legalization crash

llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp

@@ -221,41 +221,17 @@ static SPIRVType *getArgSPIRVType(const Function &F, unsigned ArgIdx,
 
 static SPIRV::ExecutionModel::ExecutionModel
 getExecutionModel(const SPIRVSubtarget &STI, const Function &F) {
+  assert(STI.getEnv() != SPIRVSubtarget::Unknown &&
+         "Environment must be resolved before lowering entry points. ");
+
   if (STI.isKernel())
     return SPIRV::ExecutionModel::Kernel;
 
-  if (STI.isShader()) {
-    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;
-    if (value == "vertex")
-      return SPIRV::ExecutionModel::Vertex;
-    if (value == "pixel")
-      return SPIRV::ExecutionModel::Fragment;
-
-    report_fatal_error(
-        "This HLSL entry point is not supported by this backend.");
-  }
-
-  assert(STI.getEnv() == SPIRVSubtarget::Unknown);
-  // "hlsl.shader" attribute is mandatory for Vulkan, so we can set Env to
-  // Shader whenever we find it, and to Kernel 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()) {
-    NonConstSTI->setEnv(SPIRVSubtarget::Kernel);
-    return SPIRV::ExecutionModel::Kernel;
+    report_fatal_error(
+        "This entry point lacks mandatory hlsl.shader attribute.");
   }
-  NonConstSTI->setEnv(SPIRVSubtarget::Shader);
 
   const auto value = attribute.getValueAsString();
   if (value == "compute")
@@ -432,11 +408,6 @@ 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/SPIRVPrepareFunctions.cpp

@@ -656,6 +656,13 @@ bool SPIRVPrepareFunctions::removeAggregateTypesFromCalls(Function *F) {
 }
 
 bool SPIRVPrepareFunctions::runOnModule(Module &M) {
+  // Resolve the SPIR-V environment from module content before any
+  // function-level processing. This must happen before legalization so that
+  // isShader()/isKernel() return correct values.
+  const_cast<SPIRVTargetMachine &>(TM)
+      .getMutableSubtargetImpl()
+      ->resolveEnvFromModule(M);
+
   bool Changed = false;
   for (Function &F : M) {
     Changed |= substituteIntrinsicCalls(&F);

llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp

@@ -173,6 +173,25 @@ void SPIRVSubtarget::initAvailableExtInstSets() {
   accountForAMDShaderTrinaryMinmax();
 }
 
+void SPIRVSubtarget::resolveEnvFromModule(const Module &M) {
+  if (Env != Unknown)
+    return;
+
+  bool HasShaderAttr = false;
+  for (const Function &F : M) {
+    if (F.getFnAttribute("hlsl.shader").isValid()) {
+      HasShaderAttr = true;
+      break;
+    }
+  }
+
+  Env = HasShaderAttr ? Shader : Kernel;
+
+  // Reinitialize Env-dependent state: ext inst sets and legalizer info.
+  initAvailableExtInstSets();
+  Legalizer = std::make_unique<SPIRVLegalizerInfo>(*this);
+}
+
 // Set available extensions after SPIRVSubtarget is created.
 void SPIRVSubtarget::initAvailableExtensions(
     const std::set<SPIRV::Extension::Extension> &AllowedExtIds) {

llvm/lib/Target/SPIRV/SPIRVSubtarget.h

@@ -25,6 +25,7 @@
 #include "llvm/CodeGen/SelectionDAGTargetInfo.h"
 #include "llvm/CodeGen/TargetSubtargetInfo.h"
 #include "llvm/IR/DataLayout.h"
+#include "llvm/IR/Module.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/TargetParser/Triple.h"
 
@@ -62,9 +63,6 @@ class SPIRVSubtarget : public SPIRVGenSubtargetInfo {
   std::unique_ptr<InstructionSelector> InstSelector;
   std::unique_ptr<InlineAsmLowering> InlineAsmInfo;
 
-  // TODO: Initialise the available extensions, extended instruction sets
-  // based on the environment settings.
-  void initAvailableExtInstSets();
   void accountForAMDShaderTrinaryMinmax();
 
 public:
@@ -76,6 +74,11 @@ class SPIRVSubtarget : public SPIRVGenSubtargetInfo {
 
   void initAvailableExtensions(
       const std::set<SPIRV::Extension::Extension> &AllowedExtIds);
+  void initAvailableExtInstSets();
+
+  // If Env is Unknown, scan module for "hlsl.shader" attributes to resolve it.
+  // Must be called before any pass that depends on isShader()/isKernel().
+  void resolveEnvFromModule(const Module &M);
 
   // Parses features string setting specified subtarget options.
   // The definition of this function is auto generated by tblgen.
@@ -86,8 +89,8 @@ class SPIRVSubtarget : public SPIRVGenSubtargetInfo {
   void setEnv(SPIRVEnvType E) {
     if (E == Unknown)
       report_fatal_error("Unknown environment is not allowed.");
-    if (Env != Unknown)
-      report_fatal_error("Environment is already set.");
+    if (Env != Unknown && Env != E)
+      report_fatal_error("Environment is already set to a different value.");
 
     Env = E;
   }

llvm/lib/Target/SPIRV/SPIRVTargetMachine.h

@@ -35,6 +35,8 @@ class SPIRVTargetMachine : public CodeGenTargetMachineImpl {
     return &Subtarget;
   }
 
+  SPIRVSubtarget *getMutableSubtargetImpl() { return &Subtarget; }
+
   TargetTransformInfo getTargetTransformInfo(const Function &F) const override;
 
   TargetPassConfig *createPassConfig(PassManagerBase &PM) override;

llvm/test/CodeGen/SPIRV/is-shader-env.ll

@@ -0,0 +1,25 @@
+; RUN: llc -O0 -mtriple=spirv-unknown-unknown %s -o - | FileCheck %s
+; RUN: %if spirv-tools %{ llc -O0 -mtriple=spirv-unknown-unknown %s -o - -filetype=obj | spirv-val %}
+
+; Regression test for https://github.com/llvm/llvm-project/issues/171898
+; When triple is spirv-unknown-unknown and a non-entry-point function using
+; wide vectors (e.g. <8 x i32>) appears before the entry point with
+; hlsl.shader attribute, the environment must be resolved early enough that
+; legalization uses the correct vector size limits.
+
+; CHECK-DAG: OpCapability Shader
+; CHECK-DAG: OpEntryPoint GLCompute %[[#entry:]] "main"
+
+define <4 x i32> @helper(<4 x i32> %a, <4 x i32> %b) {
+entry:
+  %result = add <4 x i32> %a, %b
+  ret <4 x i32> %result
+}
+
+define void @main() #0 {
+entry:
+  %a = call <4 x i32> @helper(<4 x i32> <i32 1, i32 2, i32 3, i32 4>, <4 x i32> <i32 5, i32 6, i32 7, i32 8>)
+  ret void
+}
+
+attributes #0 = { "hlsl.numthreads"="1,1,1" "hlsl.shader"="compute" }