From 05772ce02c78417b6c74a079f45000d4c30ec1c1 Mon Sep 17 00:00:00 2001 From: Sa4dUs Date: Wed, 17 Jun 2026 16:19:15 +0200 Subject: [PATCH] get number of devices intrinisc --- compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs | 10 ++++++++++ compiler/rustc_codegen_llvm/src/intrinsic.rs | 9 ++++++++- compiler/rustc_codegen_ssa/src/mir/intrinsic.rs | 1 + compiler/rustc_hir_analysis/src/check/intrinsic.rs | 7 +++++++ compiler/rustc_span/src/symbol.rs | 1 + library/core/src/intrinsics/mod.rs | 5 +++++ 6 files changed, 32 insertions(+), 1 deletion(-) diff --git a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs b/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs index 0b009321802cf..26a5ecc9ae805 100644 --- a/compiler/rustc_codegen_llvm/src/builder/gpu_offload.rs +++ b/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 diff --git a/compiler/rustc_codegen_llvm/src/intrinsic.rs b/compiler/rustc_codegen_llvm/src/intrinsic.rs index 1c7b415fd04c7..caebab6714ce5 100644 --- a/compiler/rustc_codegen_llvm/src/intrinsic.rs +++ b/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( diff --git a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs b/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs index ac6c6a0a52efa..e3cb345f401c9 100644 --- a/compiler/rustc_codegen_ssa/src/mir/intrinsic.rs +++ b/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!( diff --git a/compiler/rustc_hir_analysis/src/check/intrinsic.rs b/compiler/rustc_hir_analysis/src/check/intrinsic.rs index a4d6056e47bdc..b8870400ccbf3 100644 --- a/compiler/rustc_hir_analysis/src/check/intrinsic.rs +++ b/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, diff --git a/compiler/rustc_span/src/symbol.rs b/compiler/rustc_span/src/symbol.rs index 7263680c302f1..3bc79fd138fbe 100644 --- a/compiler/rustc_span/src/symbol.rs +++ b/compiler/rustc_span/src/symbol.rs @@ -1425,6 +1425,7 @@ symbols! { of, off, offload, + offload_get_num_devices, offload_kernel, offset, offset_of, diff --git a/library/core/src/intrinsics/mod.rs b/library/core/src/intrinsics/mod.rs index 78d7314c58110..111d31c072efb 100644 --- a/library/core/src/intrinsics/mod.rs +++ b/library/core/src/intrinsics/mod.rs @@ -3573,6 +3573,11 @@ pub const fn offload( 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)]