ICE when using function as an argument

[nazar-pc]

I need to reuse a sorting function to sort by different keys depending on use case, so I made the following changes to my project:

@@ -17,6 +17,7 @@
     bit_position: u32,
     block_size: usize,
     local_bucket: &mut [PositionR; ELEMENTS_PER_THREAD],
+    cmp_key: fn(&PositionR) -> u32,
 ) {
     // For local swaps within a thread's register data, iterate over half the elements to form pairs
     // `(a_offset, b_offset)` where indices differ only at `bit_position` and swaps them in-place
@@ -40,7 +41,7 @@
         let select_smaller =
             ((lane_id as usize * ELEMENTS_PER_THREAD + a_offset) & block_size) == 0;
 
-        let (final_a, final_b) = if (a.position <= b.position) == select_smaller {
+        let (final_a, final_b) = if (cmp_key(&a) <= cmp_key(&b)) == select_smaller {
             (a, b)
         } else {
             (b, a)
@@ -57,6 +58,7 @@
     bit_position: u32,
     block_size: usize,
     local_bucket: &mut [PositionR; ELEMENTS_PER_THREAD],
+    cmp_key: fn(&PositionR) -> u32,
 ) {
     // For cross-subgroup swaps, compute partner lane differing at `lane_bit_position`
     let lane_bit_position = bit_position - ELEMENTS_PER_THREAD.ilog2();
@@ -84,7 +86,7 @@
             PositionR { position, r }
         };
 
-        let selected_value = if (a.position <= b.position) == select_smaller {
+        let selected_value = if (cmp_key(&a) <= cmp_key(&b)) == select_smaller {
             a
         } else {
             b
@@ -121,6 +123,7 @@
 fn sort_local_bucket<const ELEMENTS_PER_THREAD: usize>(
     lane_id: u32,
     local_bucket: &mut [PositionR; ELEMENTS_PER_THREAD],
+    cmp_key: fn(&PositionR) -> u32,
 ) {
     // Iterate over merger stages, doubling block_size each time
     let mut block_size = 2;
@@ -132,10 +135,22 @@
         for bit_position in (0..merger_stage).rev() {
             if bit_position < ELEMENTS_PER_THREAD.ilog2() {
                 // Local swaps within thread's registers, no synchronization needed
-                perform_local_compare_swap(lane_id, bit_position, block_size, local_bucket);
+                perform_local_compare_swap(
+                    lane_id,
+                    bit_position,
+                    block_size,
+                    local_bucket,
+                    cmp_key,
+                );
             } else {
                 // Cross-lane swaps using subgroup shuffles for communication
-                perform_cross_compare_swap(lane_id, bit_position, block_size, local_bucket);
+                perform_cross_compare_swap(
+                    lane_id,
+                    bit_position,
+                    block_size,
+                    local_bucket,
+                    cmp_key,
+                );
             }
         }
         block_size *= 2;
@@ -172,7 +187,12 @@
 
     load_into_local_bucket(lane_id, bucket_size, bucket, &mut local_bucket);
 
-    sort_local_bucket(lane_id, &mut local_bucket);
+    sort_local_bucket(
+        lane_id,
+        &mut local_bucket,
+        #[inline(always)]
+        |position_r| position_r.position,
+    );
 
     store_from_local_bucket(lane_id, bucket, &local_bucket);
 }

Unfortunately, rustc didn't like that at all:

ICE

thread 'rustc' panicked at /home/nazar-pc/.cargo/git/checkouts/rust-gpu-d06d15e2ba0f0ae2/6a9ea3b/crates/rustc_codegen_spirv/src/linker/simple_passes.rs:240:44:
no entry found for key
stack backtrace:
   0:     0x7f4c944a3b85 - std::backtrace::Backtrace::create::h32a51249d78657b9
   1:     0x7f4c944a3ad5 - std::backtrace::Backtrace::force_capture::h59fbd3beb0b7ebd4
   2:     0x7f4c93584f52 - std[47a95048f41649f2]::panicking::update_hook::<alloc[a82d94b8c3b79735]::boxed::Box<rustc_driver_impl[5ab17e6b36e59dd4]::install_ice_hook::{closure#1}>>::{closure#0}
   3:     0x7f4c935847e3 - std[47a95048f41649f2]::panicking::update_hook::<alloc[a82d94b8c3b79735]::boxed::Box<rustc_driver_impl[5ab17e6b36e59dd4]::install_ice_hook::{closure#1}>>::{closure#0}
   4:     0x7f4c944bdbbb - std::panicking::rust_panic_with_hook::h25a3025a927f112d
   5:     0x7f4c944bd8ba - std::panicking::begin_panic_handler::{{closure}}::h21266e60c12bf697
   6:     0x7f4c944b9ee9 - std::sys::backtrace::__rust_end_short_backtrace::hef199fe2823162f9
   7:     0x7f4c944bd59d - __rustc[a3a962cb80ef9725]::rust_begin_unwind
   8:     0x7f4c90ca4570 - core::panicking::panic_fmt::hb5cc6ca3635c7279
   9:     0x7f4c9200c5fb - core::option::expect_failed::h1d79806550377b1e
  10:     0x7f4c82768e60 - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::visit::hf4b4918e200b98b5
  11:     0x7f4c827689d7 - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::visit::hf4b4918e200b98b5
  12:     0x7f4c827689d7 - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::visit::hf4b4918e200b98b5
  13:     0x7f4c827689d7 - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::visit::hf4b4918e200b98b5
  14:     0x7f4c827689d7 - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::visit::hf4b4918e200b98b5
  15:     0x7f4c82767ffe - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::hef947097e42244f9
  16:     0x7f4c82627c5f - rustc_codegen_spirv::linker::link::hf9bff09557bd6aa7
  17:     0x7f4c82752f6c - rustc_codegen_spirv::link::link::h4ba4c007671007c2
  18:     0x7f4c829dd035 - <rustc_codegen_spirv::SpirvCodegenBackend as rustc_codegen_spirv::maybe_pqp_cg_ssa::traits::backend::CodegenBackend>::link::h2e0b043a12ab3da6
  19:     0x7f4c95d09633 - <rustc_interface[9a4f103755b13c1e]::queries::Linker>::link
  20:     0x7f4c95e24219 - rustc_interface[9a4f103755b13c1e]::interface::run_compiler::<(), rustc_driver_impl[5ab17e6b36e59dd4]::run_compiler::{closure#0}>::{closure#1}
  21:     0x7f4c95de0638 - std[47a95048f41649f2]::sys::backtrace::__rust_begin_short_backtrace::<rustc_interface[9a4f103755b13c1e]::util::run_in_thread_with_globals<rustc_interface[9a4f103755b13c1e]::util::run_in_thread_pool_with_globals<rustc_interface[9a4f103755b13c1e]::interface::run_compiler<(), rustc_driver_impl[5ab17e6b36e59dd4]::run_compiler::{closure#0}>::{closure#1}, ()>::{closure#0}, ()>::{closure#0}::{closure#0}, ()>
  22:     0x7f4c95de0316 - <<std[47a95048f41649f2]::thread::Builder>::spawn_unchecked_<rustc_interface[9a4f103755b13c1e]::util::run_in_thread_with_globals<rustc_interface[9a4f103755b13c1e]::util::run_in_thread_pool_with_globals<rustc_interface[9a4f103755b13c1e]::interface::run_compiler<(), rustc_driver_impl[5ab17e6b36e59dd4]::run_compiler::{closure#0}>::{closure#1}, ()>::{closure#0}, ()>::{closure#0}::{closure#0}, ()>::{closure#1} as core[28534e9693d03804]::ops::function::FnOnce<()>>::call_once::{shim:vtable#0}
  23:     0x7f4c95dde93d - std::sys::pal::unix::thread::Thread::new::thread_start::hf0d6fc1bcf013f18
  24:     0x7f4c8f69caa4 - start_thread
                               at ./nptl/pthread_create.c:447:8
  25:     0x7f4c8f729c6c - clone3
                               at ./misc/../sysdeps/unix/sysv/linux/x86_64/clone3.S:78:0
  26:                0x0 - <unknown>


rustc version: 1.90.0-nightly (35f603652 2025-06-29)
platform: x86_64-unknown-linux-gnuthread 'rustc' panicked at /home/nazar-pc/.cargo/git/checkouts/rust-gpu-d06d15e2ba0f0ae2/6a9ea3b/crates/rustc_codegen_spirv/src/linker/simple_passes.rs:240:44:
no entry found for key
stack backtrace:
   0:     0x7f4c944a3b85 - std::backtrace::Backtrace::create::h32a51249d78657b9
   1:     0x7f4c944a3ad5 - std::backtrace::Backtrace::force_capture::h59fbd3beb0b7ebd4
   2:     0x7f4c93584f52 - std[47a95048f41649f2]::panicking::update_hook::<alloc[a82d94b8c3b79735]::boxed::Box<rustc_driver_impl[5ab17e6b36e59dd4]::install_ice_hook::{closure#1}>>::{closure#0}
   3:     0x7f4c944bdbbb - std::panicking::rust_panic_with_hook::h25a3025a927f112d
   4:     0x7f4c944bd8ba - std::panicking::begin_panic_handler::{{closure}}::h21266e60c12bf697
   5:     0x7f4c944b9ee9 - std::sys::backtrace::__rust_end_short_backtrace::hef199fe2823162f9
   6:     0x7f4c944bd59d - __rustc[a3a962cb80ef9725]::rust_begin_unwind
   7:     0x7f4c90ca4570 - core::panicking::panic_fmt::hb5cc6ca3635c7279
   8:     0x7f4c9200c5fb - core::option::expect_failed::h1d79806550377b1e
   9:     0x7f4c82768e60 - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::visit::hf4b4918e200b98b5
  10:     0x7f4c827689d7 - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::visit::hf4b4918e200b98b5
  11:     0x7f4c827689d7 - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::visit::hf4b4918e200b98b5
  12:     0x7f4c827689d7 - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::visit::hf4b4918e200b98b5
  13:     0x7f4c827689d7 - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::visit::hf4b4918e200b98b5
  14:     0x7f4c82767ffe - rustc_codegen_spirv::linker::simple_passes::check_fragment_insts::hef947097e42244f9
  15:     0x7f4c82627c5f - rustc_codegen_spirv::linker::link::hf9bff09557bd6aa7
  16:     0x7f4c82752f6c - rustc_codegen_spirv::link::link::h4ba4c007671007c2
  17:     0x7f4c829dd035 - <rustc_codegen_spirv::SpirvCodegenBackend as rustc_codegen_spirv::maybe_pqp_cg_ssa::traits::backend::CodegenBackend>::link::h2e0b043a12ab3da6
  18:     0x7f4c95d09633 - <rustc_interface[9a4f103755b13c1e]::queries::Linker>::link
  19:     0x7f4c95e24219 - rustc_interface[9a4f103755b13c1e]::interface::run_compiler::<(), rustc_driver_impl[5ab17e6b36e59dd4]::run_compiler::{closure#0}>::{closure#1}
  20:     0x7f4c95de0638 - std[47a95048f41649f2]::sys::backtrace::__rust_begin_short_backtrace::<rustc_interface[9a4f103755b13c1e]::util::run_in_thread_with_globals<rustc_interface[9a4f103755b13c1e]::util::run_in_thread_pool_with_globals<rustc_interface[9a4f103755b13c1e]::interface::run_compiler<(), rustc_driver_impl[5ab17e6b36e59dd4]::run_compiler::{closure#0}>::{closure#1}, ()>::{closure#0}, ()>::{closure#0}::{closure#0}, ()>
  21:     0x7f4c95de0316 - <<std[47a95048f41649f2]::thread::Builder>::spawn_unchecked_<rustc_interface[9a4f103755b13c1e]::util::run_in_thread_with_globals<rustc_interface[9a4f103755b13c1e]::util::run_in_thread_pool_with_globals<rustc_interface[9a4f103755b13c1e]::interface::run_compiler<(), rustc_driver_impl[5ab17e6b36e59dd4]::run_compiler::{closure#0}>::{closure#1}, ()>::{closure#0}, ()>::{closure#0}::{closure#0}, ()>::{closure#1} as core[28534e9693d03804]::ops::function::FnOnce<()>>::call_once::{shim:vtable#0}
  22:     0x7f4c95dde93d - std::sys::pal::unix::thread::Thread::new::thread_start::hf0d6fc1bcf013f18
  23:     0x7f4c8f69caa4 - start_thread
                               at ./nptl/pthread_create.c:447:8
  24:     0x7f4c8f729c6c - clone3
                               at ./misc/../sysdeps/unix/sysv/linux/x86_64/clone3.S:78:0
  25:                0x0 - <unknown>

This is unexpected, especially because I want it to be literally inlined by the compiler so that generated code is no different to before.

Currently using 6a9ea3b.

[nazar-pc]

Replacing fn() with a generic works, but it is unnecessarily verbose:

@@ -12,12 +12,15 @@
 pub const WORKGROUP_SIZE: u32 = 256;
 
 #[inline(always)]
-fn perform_local_compare_swap<const ELEMENTS_PER_THREAD: usize>(
+fn perform_local_compare_swap<const ELEMENTS_PER_THREAD: usize, CmpKey>(
     lane_id: u32,
     bit_position: u32,
     block_size: usize,
     local_bucket: &mut [PositionR; ELEMENTS_PER_THREAD],
-) {
+    cmp_key: CmpKey,
+) where
+    CmpKey: Fn(&PositionR) -> u32,
+{
     // For local swaps within a thread's register data, iterate over half the elements to form pairs
     // `(a_offset, b_offset)` where indices differ only at `bit_position` and swaps them in-place
     for pair_id in 0..ELEMENTS_PER_THREAD / 2 {
@@ -40,7 +43,7 @@
         let select_smaller =
             ((lane_id as usize * ELEMENTS_PER_THREAD + a_offset) & block_size) == 0;
 
-        let (final_a, final_b) = if (a.position <= b.position) == select_smaller {
+        let (final_a, final_b) = if (cmp_key(&a) <= cmp_key(&b)) == select_smaller {
             (a, b)
         } else {
             (b, a)
@@ -52,12 +55,15 @@
 }
 
 #[inline(always)]
-fn perform_cross_compare_swap<const ELEMENTS_PER_THREAD: usize>(
+fn perform_cross_compare_swap<const ELEMENTS_PER_THREAD: usize, CmpKey>(
     lane_id: u32,
     bit_position: u32,
     block_size: usize,
     local_bucket: &mut [PositionR; ELEMENTS_PER_THREAD],
-) {
+    cmp_key: &CmpKey,
+) where
+    CmpKey: Fn(&PositionR) -> u32,
+{
     // For cross-subgroup swaps, compute partner lane differing at `lane_bit_position`
     let lane_bit_position = bit_position - ELEMENTS_PER_THREAD.ilog2();
     let partner_lane_id = lane_id ^ (1 << lane_bit_position);
@@ -84,7 +90,7 @@
             PositionR { position, r }
         };
 
-        let selected_value = if (a.position <= b.position) == select_smaller {
+        let selected_value = if (cmp_key(&a) <= cmp_key(&b)) == select_smaller {
             a
         } else {
             b
@@ -118,10 +124,13 @@
 }
 
 #[inline(always)]
-fn sort_local_bucket<const ELEMENTS_PER_THREAD: usize>(
+fn sort_local_bucket<const ELEMENTS_PER_THREAD: usize, CmpKey>(
     lane_id: u32,
     local_bucket: &mut [PositionR; ELEMENTS_PER_THREAD],
-) {
+    cmp_key: CmpKey,
+) where
+    CmpKey: Fn(&PositionR) -> u32,
+{
     // Iterate over merger stages, doubling block_size each time
     let mut block_size = 2;
     let mut merger_stage = 1;
@@ -132,10 +141,22 @@
         for bit_position in (0..merger_stage).rev() {
             if bit_position < ELEMENTS_PER_THREAD.ilog2() {
                 // Local swaps within thread's registers, no synchronization needed
-                perform_local_compare_swap(lane_id, bit_position, block_size, local_bucket);
+                perform_local_compare_swap(
+                    lane_id,
+                    bit_position,
+                    block_size,
+                    local_bucket,
+                    &cmp_key,
+                );
             } else {
                 // Cross-lane swaps using subgroup shuffles for communication
-                perform_cross_compare_swap(lane_id, bit_position, block_size, local_bucket);
+                perform_cross_compare_swap(
+                    lane_id,
+                    bit_position,
+                    block_size,
+                    local_bucket,
+                    &cmp_key,
+                );
             }
         }
         block_size *= 2;
@@ -172,7 +193,12 @@
 
     load_into_local_bucket(lane_id, bucket_size, bucket, &mut local_bucket);
 
-    sort_local_bucket(lane_id, &mut local_bucket);
+    sort_local_bucket(
+        lane_id,
+        &mut local_bucket,
+        #[inline(always)]
+        |position_r| position_r.position,
+    );
 
     store_from_local_bucket(lane_id, bucket, &local_bucket);
 }

[LegNeato]

This is the failure point:

rust-gpu/crates/rustc_codegen_spirv/src/linker/simple_passes.rs

Line 240 in 808974a

let callee = func_id_to_idx[&inst.operands[0].unwrap_id_ref()];

[LegNeato]

#453 fixes the ICE, but I don't think indirect fn pointers can be supported until @eddyb 's qptr work lands.

[Firestar99]

You could use cmp_key: &impl Fn(&PositionR) -> u32 to shorten it a bit.

But yea, we don't support function pointers at all right now, since they could come from a dynamic source. Like, you could dynamically select from which function to create a pointer, which is not allowed in spirv (and pretty much all GPU ISAs afaik), since each invocation could have a different function pointer the GPU wouldn't know where to jump to. But closures work, since they're generics monomorphization will just "inline" them for free for us.

[eddyb]

The issue # in the PR wasn't preceded by the verb "fixes" so this issue didn't get closed automatically.

[eddyb]

#453 fixes the ICE, but I don't think indirect fn pointers can be supported until @eddyb 's qptr work lands.

Small correction: the recursion/function pointer emulator is, in theory, separate from all the qptr work (although in the most general case, some interactions are needed to fully legalize everything, e.g. if you have pointer values that need to be written on the stack so they can be restored later, etc.).

What's missing for something like this "I wanted it inlined anyway" usecase, is turning an indirect call into a direct one (presumably with some amount of eliminating local variables, maybe constant-folding, etc.).
If the situation is simple enough, this might be doable with the code we have already, but if not, it might need to wait for a SPIR-T inliner (which doesn't exist, not even in the braches holding my ridiculous tower of hacks).

---

Also, for a bit more visibility, I'm duplicating this comment I left on the fix PR, here:

The underlying bug¹ is that OpFunctionCall is still used with function pointers - the correct instruction is OpFunctionPointerCallINTEL (from the same SPIR-V extension that adds function pointers through a new CodeSectionINTEL Storage Class).

_{¹(which one of my branches fixes, been meaning to open a PR for that for a while, already landed most of its prerequisites earlier this year)}