[llvm] [SPIR-V] Fix environment resolution causing legalization crash (PR #179052)
Dmitry Sidorov via llvm-commits
llvm-commits at lists.llvm.org
Sat Jan 31 09:01:52 PST 2026
https://github.com/MrSidims created https://github.com/llvm/llvm-project/pull/179052
When the triple is spirv-unknown-unknown, SPIRVSubtarget::Env starts as Unknown and was set via const_cast in SPIRVCallLowering when the first entry point was lowered. This is too late: SPIRVLegalizerInfo has already been constructed with, for example, the wrong vector size limits, causing a crash (at best) or invalid SPIR-V generation.
Resolve the environment early in SPIRVPrepareFunctions::runOnModule() by scanning the module for "hlsl.shader" attributes. Reinitialize the legalizer and extended instruction sets after resolution. Remove the const_cast lazy setting from SPIRVCallLowering.
Fixes: https://github.com/llvm/llvm-project/issues/171898
>From 602a2879e9466c6d7efd966494482e482c2f58f4 Mon Sep 17 00:00:00 2001
From: Dmitry Sidorov <Dmitry.Sidorov at amd.com>
Date: Sat, 31 Jan 2026 17:58:21 +0100
Subject: [PATCH] [SPIR-V] Fix environment resolution causing legalization
crash
When the triple is spirv-unknown-unknown, SPIRVSubtarget::Env starts as
Unknown and was set via const_cast in SPIRVCallLowering when the
first entry point was lowered. This is too late: SPIRVLegalizerInfo has
already been constructed with, for example, the wrong vector size limits,
causing a crash (at best) or invalid SPIR-V generation.
Resolve the environment early in SPIRVPrepareFunctions::runOnModule() by
scanning the module for "hlsl.shader" attributes. Reinitialize the
legalizer and extended instruction sets after resolution. Remove the
const_cast lazy setting from SPIRVCallLowering.
---
llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp | 39 +++----------------
.../Target/SPIRV/SPIRVPrepareFunctions.cpp | 6 +++
llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp | 19 +++++++++
llvm/lib/Target/SPIRV/SPIRVSubtarget.h | 13 ++++---
llvm/lib/Target/SPIRV/SPIRVTargetMachine.h | 2 +
llvm/test/CodeGen/SPIRV/is-shader-env.ll | 25 ++++++++++++
6 files changed, 65 insertions(+), 39 deletions(-)
create mode 100644 llvm/test/CodeGen/SPIRV/is-shader-env.ll
diff --git a/llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp b/llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp
index 23c5798f9d0af..dff9b4a48564b 100644
--- a/llvm/lib/Target/SPIRV/SPIRVCallLowering.cpp
+++ b/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);
diff --git a/llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp b/llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp
index a3425704f050d..1aedab4e94048 100644
--- a/llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp
+++ b/llvm/lib/Target/SPIRV/SPIRVPrepareFunctions.cpp
@@ -656,6 +656,12 @@ 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);
diff --git a/llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp b/llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp
index ad6c9cd421b7c..7e1a41dcb2276 100644
--- a/llvm/lib/Target/SPIRV/SPIRVSubtarget.cpp
+++ b/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) {
diff --git a/llvm/lib/Target/SPIRV/SPIRVSubtarget.h b/llvm/lib/Target/SPIRV/SPIRVSubtarget.h
index ad3e38d296ed7..01888adf88fbb 100644
--- a/llvm/lib/Target/SPIRV/SPIRVSubtarget.h
+++ b/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;
}
diff --git a/llvm/lib/Target/SPIRV/SPIRVTargetMachine.h b/llvm/lib/Target/SPIRV/SPIRVTargetMachine.h
index 9c59d021dfc1b..ea09fe98c55ee 100644
--- a/llvm/lib/Target/SPIRV/SPIRVTargetMachine.h
+++ b/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;
diff --git a/llvm/test/CodeGen/SPIRV/is-shader-env.ll b/llvm/test/CodeGen/SPIRV/is-shader-env.ll
new file mode 100644
index 0000000000000..47a323ded0df5
--- /dev/null
+++ b/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" }
More information about the llvm-commits
mailing list