From b115e95f043367365f9528d78d2897e99416ef5e Mon Sep 17 00:00:00 2001
From: Christian Legnitto <christian@legnitto.com>
Date: Sun, 3 May 2026 18:28:54 -0700
Subject: [PATCH] Keep `#[spirv(...)]` entry-points from being DCE'd by rustc.

When building a shader as a `dylib`, rustc's `cross_crate_inlinable`
heuristic can mark a `pub fn` entry-point as inlinable, which
internalizes it and drops its `OpEntryPoint` before codegen ever runs.

Override the query to return `false` for entry-points.

Fixes #590.
---
 crates/rustc_codegen_spirv/src/attr.rs        | 16 ++++++++++++
 .../ui/dis/issue-590-multiple-entry-points.rs | 24 ++++++++++++++++++
 .../issue-590-multiple-entry-points.stderr    | 25 +++++++++++++++++++
 3 files changed, 65 insertions(+)
 create mode 100644 tests/compiletests/ui/dis/issue-590-multiple-entry-points.rs
 create mode 100644 tests/compiletests/ui/dis/issue-590-multiple-entry-points.stderr

diff --git a/crates/rustc_codegen_spirv/src/attr.rs b/crates/rustc_codegen_spirv/src/attr.rs
index d1a8e5718a9..a8ff7df0a1e 100644
--- a/crates/rustc_codegen_spirv/src/attr.rs
+++ b/crates/rustc_codegen_spirv/src/attr.rs
@@ -522,6 +522,22 @@ pub(crate) fn provide(providers: &mut Providers) {
                 .check_mod_attrs)(tcx, module_def_id);
             check_mod_attrs(tcx, module_def_id);
         },
+        // HACK(LegNeato) keep `#[spirv(...)]` entry-points from getting
+        // internalized and DCE'd before codegen. See issue #590.
+        cross_crate_inlinable: |tcx, def_id| {
+            let sym = Symbols::get();
+            let path = [sym.rust_gpu, sym.spirv_attr_with_version];
+            let attrs: Vec<_> = tcx.get_attrs_by_path(def_id.to_def_id(), &path).collect();
+            let is_entry_point = parse_attrs_for_checking(&sym, attrs)
+                .any(|result| matches!(result, Ok((_, SpirvAttribute::Entry(_)))));
+            if is_entry_point {
+                false
+            } else {
+                (rustc_interface::DEFAULT_QUERY_PROVIDERS
+                    .queries
+                    .cross_crate_inlinable)(tcx, def_id)
+            }
+        },
         ..*providers
     };
 }
diff --git a/tests/compiletests/ui/dis/issue-590-multiple-entry-points.rs b/tests/compiletests/ui/dis/issue-590-multiple-entry-points.rs
new file mode 100644
index 00000000000..0ca76bbc637
--- /dev/null
+++ b/tests/compiletests/ui/dis/issue-590-multiple-entry-points.rs
@@ -0,0 +1,24 @@
+// Tests that small entry-points aren't internalized by rustc and dropped
+// before reaching codegen. See https://github.com/Rust-GPU/rust-gpu/issues/590.
+
+// build-pass
+// compile-flags: -C debuginfo=0 -C llvm-args=--disassemble-globals
+// normalize-stderr-test "OpCapability VulkanMemoryModel\n" -> ""
+// normalize-stderr-test "OpSource .*\n" -> ""
+// normalize-stderr-test "OpExtension .SPV_KHR_vulkan_memory_model.\n" -> ""
+// normalize-stderr-test "OpMemoryModel Logical Vulkan" -> "OpMemoryModel Logical Simple"
+
+// HACK(eddyb) `compiletest` handles `ui\dis\`, but not `ui\\dis\\`, on Windows.
+// normalize-stderr-test "ui/dis/" -> "$$DIR/"
+
+use spirv_std::glam::{UVec3, Vec4};
+use spirv_std::spirv;
+
+#[spirv(vertex)]
+pub fn main_vs(#[spirv(vertex_index)] _v: i32, #[spirv(position)] _pos: &mut Vec4) {}
+
+#[spirv(fragment)]
+pub fn main_fs(_out: &mut Vec4) {}
+
+#[spirv(compute(threads(64)))]
+pub fn main_cs(#[spirv(global_invocation_id)] _id: UVec3) {}
diff --git a/tests/compiletests/ui/dis/issue-590-multiple-entry-points.stderr b/tests/compiletests/ui/dis/issue-590-multiple-entry-points.stderr
new file mode 100644
index 00000000000..38159def44d
--- /dev/null
+++ b/tests/compiletests/ui/dis/issue-590-multiple-entry-points.stderr
@@ -0,0 +1,25 @@
+OpCapability Shader
+OpMemoryModel Logical Simple
+OpEntryPoint GLCompute %1 "main_cs" %2
+OpEntryPoint Fragment %3 "main_fs"
+OpEntryPoint Vertex %4 "main_vs" %5
+OpExecutionMode %1 LocalSize 64 1 1
+OpExecutionMode %3 OriginUpperLeft
+%6 = OpString "$DIR/issue-590-multiple-entry-points.rs"
+OpName %2 "_id"
+OpName %5 "_v"
+OpName %7 "issue_590_multiple_entry_points::main_cs"
+OpName %8 "issue_590_multiple_entry_points::main_fs"
+OpName %9 "issue_590_multiple_entry_points::main_vs"
+OpDecorate %2 BuiltIn GlobalInvocationId
+OpDecorate %5 BuiltIn VertexIndex
+%10 = OpTypeInt 32 0
+%11 = OpTypeVector %10 3
+%12 = OpTypePointer Input %11
+%13 = OpTypeVoid
+%14 = OpTypeFunction %13
+%2 = OpVariable  %12  Input
+%15 = OpTypeFunction %13 %11
+%16 = OpTypeInt 32 1
+%17 = OpTypePointer Input %16
+%5 = OpVariable  %17  Input