Offload expose device selection

[Sa4dUs]

get amount of available devices with omp_get_num_devices

r? @ZuseZ4

[rust-log-analyzer]

The job tidy failed! Check out the build log: (web) (plain enhanced) (plain)

Click to see the possible cause of the failure (guessed by this bot)

[TIMING:end] tool::Tidy { compiler: Compiler { stage: 0, host: x86_64-unknown-linux-gnu, forced_compiler: false }, target: x86_64-unknown-linux-gnu } -- 0.000
fmt check
Diff in /checkout/compiler/rustc_codegen_llvm/src/intrinsic.rs:35:
 use crate::builder::Builder;
 use crate::builder::autodiff::{adjust_activity_to_abi, generate_enzyme_call};
 use crate::builder::gpu_offload::{
-    OffloadKernelDims, gen_call_handling, gen_define_handling, register_offload, generate_decl,
+    OffloadKernelDims, gen_call_handling, gen_define_handling, generate_decl, register_offload,
 };
 use crate::context::CodegenCx;
 use crate::declare::declare_raw_fn;
Diff in /checkout/compiler/rustc_hir_analysis/src/check/intrinsic.rs:362:
             ],
---
-          0,
-          vec![],
-          tcx.types.i32,
-        ),
+        sym::offload_get_num_devices => (0, 0, vec![], tcx.types.i32),
         sym::offset => (2, 0, vec![param(0), param(1)], param(0)),
         sym::arith_offset => (
             1,
fmt: checked 6924 files
Bootstrap failed while executing `test src/tools/tidy tidyselftest --extra-checks=py,cpp,js,spellcheck`
Build completed unsuccessfully in 0:00:51

[bjorn3]

Does it actually need to be a rust intrinsic? Does LLVM give omp_get_num_devices any special behavior or is it just a regular runtime call that we could directly call from rust code too?

[ZuseZ4]

@bjorn3 Yes, you could call all of our openmp/offload libraries directly, especially here it's easy enough. Some of the other ones we use for data transfer and kernel launches are just a bit more involved to set up, so we do it in the compiler.

The two problems of exposing it are that you'd need to link in the GCC or LLVM libraries yourself, if you build and use the Offload setup, we do it for users. Also cranelift doesn't provide it, with an intrinsic we could have better fallback behaviour.