Offload expose device selection

compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs

@@ -197,6 +197,16 @@ fn generate_launcher<'ll>(cx: &CodegenCx<'ll, '_>) -> (&'ll llvm::Value, &'ll ll
     (tgt_decl, tgt_fn_ty)
 }
 
+pub(crate) fn generate_decl<'ll>(cx: &CodegenCx<'ll, '_>) -> (&'ll llvm::Value, &'ll llvm::Type) {
+    let ti32 = cx.type_i32();
+    let tgt_fn_ty = cx.type_func(&[], ti32);
+    let name = "omp_get_num_devices";
+    let tgt_decl = declare_offload_fn(&cx, name, tgt_fn_ty);
+    let nounwind = llvm::AttributeKind::NoUnwind.create_attr(cx.llcx);
+    attributes::apply_to_llfn(tgt_decl, Function, &[nounwind]);
+    (tgt_decl, tgt_fn_ty)
+}
+
 // What is our @1 here? A magic global, used in our data_{begin/update/end}_mapper:
 // @0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
 // @1 = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 22, ptr @0 }, align 8

compiler/rustc_codegen_llvm/src/intrinsic.rs

@@ -35,7 +35,7 @@ use crate::abi::FnAbiLlvmExt;
 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,
+    OffloadKernelDims, gen_call_handling, gen_define_handling, register_offload, generate_decl,
 };
 use crate::context::CodegenCx;
 use crate::declare::declare_raw_fn;
@@ -243,6 +243,13 @@ impl<'ll, 'tcx> IntrinsicCallBuilderMethods<'tcx> for Builder<'_, 'll, 'tcx> {
                 // offload *has* a return type, but somehow works without mentioning the place
                 return IntrinsicResult::WroteIntoPlace;
             }
+            sym::offload_get_num_devices => {
+                let (fn_decl, fn_ty) = generate_decl(self.cx);
+
+                let llval = self.call(fn_ty, None, None, fn_decl, &[], None, None);
+
+                return IntrinsicResult::Operand(OperandValue::Immediate(llval));
+            },
             sym::is_val_statically_known => {
                 if let OperandValue::Immediate(imm) = args[0].val {
                     self.call_intrinsic(

compiler/rustc_codegen_ssa/src/mir/intrinsic.rs

@@ -124,6 +124,7 @@ impl<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> FunctionCx<'a, 'tcx, Bx> {
                 | sym::atomic_fence
                 | sym::atomic_singlethreadfence
                 | sym::caller_location
+                | sym::offload_get_num_devices
                 | sym::return_address => {}
                 _ => {
                     span_bug!(

compiler/rustc_hir_analysis/src/check/intrinsic.rs

@@ -164,6 +164,7 @@ fn intrinsic_operation_unsafety(tcx: TyCtxt<'_>, intrinsic_id: LocalDefId) -> hi
         | sym::mul_with_overflow
         | sym::needs_drop
         | sym::offload
+        | sym::offload_get_num_devices
         | sym::offset_of
         | sym::overflow_checks
         | sym::powf16
@@ -361,6 +362,12 @@ pub(crate) fn check_intrinsic_type(
             ],
             param(2),
         ),
+        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,

compiler/rustc_span/src/symbol.rs

@@ -1425,6 +1425,7 @@ symbols! {
         of,
         off,
         offload,
+        offload_get_num_devices,
         offload_kernel,
         offset,
         offset_of,

library/core/src/intrinsics/mod.rs

@@ -3573,6 +3573,11 @@ pub const fn offload<F, T: crate::marker::Tuple, R>(
     args: T,
 ) -> R;
 
+// TODO(Sa4dUs): add docs
+#[rustc_nounwind]
+#[rustc_intrinsic]
+pub const fn offload_get_num_devices() -> i32;
+
 /// Inform Miri that a given pointer definitely has a certain alignment.
 #[cfg(miri)]
 #[rustc_allow_const_fn_unstable(const_eval_select)]