Amdsupport

.gitmodules

@@ -0,0 +1,4 @@
+[submodule "crates/rustc_offload_frontend"]
+	path = crates/rustc_offload_frontend
+	url = https://github.com/sa4dus/rustc_offload_frontend
+	branch = main

Cargo.toml

@@ -1,13 +1,28 @@
+cargo-features = ["panic-immediate-abort"]
+
 [package]
 name = "rust_perf"
 version = "0.1.0"
 edition = "2024"
 
 [lib]
 crate-type = ["cdylib"]
+#crate-type = ["rlib", "cdylib"]
+path = "src/lib.rs"
 
 [dependencies]
 libc = { version = "0.2.175", default-features = false }
+rustc_offload_frontend = { path = "crates/rustc_offload_frontend" }
+
+[profile.release]
+lto = "fat"
+panic = "immediate-abort"
+codegen-units = 1
+
+[profile.dev]
+lto = "fat"
+panic = "immediate-abort"
+
 
 [features]
 default = ["all", "f64"]
@@ -40,10 +55,3 @@ vol3d = []
 f64 = []
 a64 = []
 
-[profile.release]
-lto = "fat"
-panic = "abort"
-
-[profile.dev]
-lto = "fat"
-panic = "abort"

crates/rustc_offload_frontend/.gitignore

@@ -0,0 +1 @@
+/target

crates/rustc_offload_frontend/Cargo.toml

@@ -0,0 +1,30 @@
+[package]
+name = "rustc_offload_frontend"
+version = "0.1.0"
+edition = "2024"
+license = "MIT"
+description = "rustc offload feature frontend draft"
+repository = "https://github.com/sa4dus/rustc_offload_frontend"
+readme = "README.md"
+
+[lib]
+crate-type = ["rlib"]
+#crate-type = ["rlib", "cdylib"]
+path = "src/lib.rs"
+
+[[bin]]
+name = "main"
+path = "src/main.rs"
+
+[dependencies]
+libc = { version = "0.2.175", default-features = false }
+
+[features]
+
+[profile.release]
+lto = "fat"
+panic = "abort"
+
+[profile.dev]
+lto = "fat"
+panic = "abort"

crates/rustc_offload_frontend/README.md

@@ -0,0 +1 @@
+# rustc_offload_frontend

crates/rustc_offload_frontend/src/gpu.rs

@@ -0,0 +1,191 @@
+#[derive(Clone, Copy)]
+pub struct Dim3 {
+    pub x: usize,
+    pub y: usize,
+    pub z: usize,
+}
+
+#[cfg(target_arch = "amdgpu")]
+use core::intrinsics::gpu::amdgpu_dispatch_ptr;
+
+/// Get the packet for this dispatch.
+///
+/// Get a reference to the packet that was used to dispatch this kernel.
+/// The dispatch packet contains information like the workgroup size and dispatch size.
+///
+/// # Example
+///
+/// ```rust
+/// # #![no_std]
+/// # extern crate alloc;
+/// # fn main() {
+/// use amdgpu_device_libs::prelude::*;
+///
+/// let dispatch = dispatch_ptr();
+/// println!("Workgroup size {}x{}x{}", dispatch.workgroup_size_x, dispatch.workgroup_size_y, dispatch.workgroup_size_z);
+/// # }
+/// ```
+#[cfg(target_arch = "amdgpu")]
+#[inline]
+pub fn dispatch_ptr() -> &'static HsaKernelDispatchPacket {
+    unsafe {
+        &*core::mem::transmute::<*const (), *const HsaKernelDispatchPacket>(
+            core::intrinsics::gpu::amdgpu_dispatch_ptr(),
+        )
+    }
+}
+/// HSA packet to dispatch a kernel.
+///
+/// A pointer to the packet that was used to dispatch the currently running kernel can be obtained with [`dispatch_ptr`].
+#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
+#[repr(C)]
+pub struct HsaKernelDispatchPacket {
+    /// Packet header. Used to configure multiple packet parameters such as the
+    /// packet type. The parameters are described by hsa_packet_header_t.
+    pub header: u16,
+    /// Dispatch setup parameters. Used to configure kernel dispatch parameters
+    /// such as the number of dimensions in the grid. The parameters are described
+    /// by hsa_kernel_dispatch_packet_setup_t.
+    pub setup: u16,
+    /// X dimension of work-group, in work-items. Must be greater than 0.
+    pub workgroup_size_x: u16,
+    /// Y dimension of work-group, in work-items. Must be greater than
+    /// 0. If the grid has 1 dimension, the only valid value is 1.
+    pub workgroup_size_y: u16,
+    /// Z dimension of work-group, in work-items. Must be greater than
+    /// 0. If the grid has 1 or 2 dimensions, the only valid value is 1.
+    pub workgroup_size_z: u16,
+    /// Reserved. Must be 0.
+    pub reserved0: u16,
+    /// X dimension of grid, in work-items. Must be greater than 0. Must
+    /// not be smaller than @a workgroup_size_x.
+    pub grid_size_x: u32,
+    /// Y dimension of grid, in work-items. Must be greater than 0. If the grid has
+    /// 1 dimension, the only valid value is 1. Must not be smaller than @a
+    /// workgroup_size_y.
+    pub grid_size_y: u32,
+    /// Z dimension of grid, in work-items. Must be greater than 0. If the grid has
+    /// 1 or 2 dimensions, the only valid value is 1. Must not be smaller than @a
+    /// workgroup_size_z.
+    pub grid_size_z: u32,
+    /// Size in bytes of private memory allocation request (per work-item).
+    pub private_segment_size: u32,
+    /// Size in bytes of group memory allocation request (per work-group). Must not
+    /// be less than the sum of the group memory used by the kernel (and the
+    /// functions it calls directly or indirectly) and the dynamically allocated
+    /// group segment variables.
+    pub group_segment_size: u32,
+    /// Opaque handle to a code object that includes an implementation-defined
+    /// executable code for the kernel.
+    pub kernel_object: u64,
+    /// Pointer to the kernel arguments.
+    pub kernarg_address: *mut core::ffi::c_void,
+    /// Reserved. Must be 0.
+    pub reserved2: u64,
+    /// Signal used to indicate completion of the job. The application can use the
+    /// special signal handle 0 to indicate that no signal is used.
+    pub completion_signal: u64,
+}
+// Handle to an HSA signal.
+//#[cfg(feature = "device_libs")]
+//#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
+//#[repr(C)]
+//pub struct HsaSignal {
+//    /// The internal representation of an HSA signal.
+//    pub handle: u64,
+//}
+
+pub(crate) fn global_thread_dim() -> Dim3 {
+    #[cfg(target_arch = "nvptx64")]
+    unsafe {
+        use core::arch::nvptx::*;
+        Dim3 {
+            x: (_block_idx_x() * _block_dim_x() + _thread_idx_x()) as usize,
+            y: (_block_idx_y() * _block_dim_y() + _thread_idx_y()) as usize,
+            z: (_block_idx_z() * _block_dim_z() + _thread_idx_z()) as usize,
+        }
+    }
+    #[cfg(target_arch = "amdgpu")]
+    unsafe {
+        use core::arch::amdgpu::*;
+        let dispatch = dispatch_ptr();
+
+        let x = (workgroup_id_x() * (*dispatch).workgroup_size_x as u32 + workitem_id_x()) as usize;
+        let y = (workgroup_id_y() * (*dispatch).workgroup_size_y as u32 + workitem_id_y()) as usize;
+        let z = (workgroup_id_z() * (*dispatch).workgroup_size_z as u32 + workitem_id_z()) as usize;
+        Dim3 { x, y, z }
+    }
+    #[cfg(target_os = "linux")]
+    Dim3 { x: 0, y: 0, z: 0 }
+}
+
+pub(crate) fn block_idx() -> Dim3 {
+    #[cfg(target_arch = "nvptx64")]
+    unsafe {
+        use core::arch::nvptx::*;
+        Dim3 {
+            x: _block_idx_x() as usize,
+            y: _block_idx_y() as usize,
+            z: _block_idx_z() as usize,
+        }
+    }
+    #[cfg(target_arch = "amdgpu")]
+    unsafe {
+        use core::arch::amdgpu::*;
+        let dispatch = dispatch_ptr();
+
+        let x = (workgroup_id_x()) as usize;
+        let y = (workgroup_id_y()) as usize;
+        let z = (workgroup_id_z()) as usize;
+        Dim3 { x, y, z }
+    }
+    #[cfg(target_os = "linux")]
+    Dim3 { x: 0, y: 0, z: 0 }
+}
+
+pub(crate) fn block_dim() -> Dim3 {
+    #[cfg(target_arch = "nvptx64")]
+    unsafe {
+        use core::arch::nvptx::*;
+        Dim3 {
+            x: _block_dim_x() as usize,
+            y: _block_dim_y() as usize,
+            z: _block_dim_z() as usize,
+        }
+    }
+    #[cfg(target_arch = "amdgpu")]
+    unsafe {
+        use core::arch::amdgpu::*;
+        let dispatch = dispatch_ptr();
+
+        let x = dispatch.workgroup_size_x as usize;
+        let y = dispatch.workgroup_size_y as usize;
+        let z = dispatch.workgroup_size_z as usize;
+        Dim3 { x, y, z }
+    }
+    #[cfg(target_os = "linux")]
+    Dim3 { x: 0, y: 0, z: 0 }
+}
+
+pub(crate) fn thread_idx() -> Dim3 {
+    #[cfg(target_arch = "nvptx64")]
+    unsafe {
+        use core::arch::nvptx::*;
+        Dim3 {
+            x: _thread_idx_x() as usize,
+            y: _thread_idx_y() as usize,
+            z: _thread_idx_z() as usize,
+        }
+    }
+    #[cfg(target_arch = "amdgpu")]
+    unsafe {
+        use core::arch::amdgpu::*;
+
+        let x = (workitem_id_x()) as usize;
+        let y = (workitem_id_y()) as usize;
+        let z = (workitem_id_z()) as usize;
+        Dim3 { x, y, z }
+    }
+    #[cfg(target_os = "linux")]
+    Dim3 { x: 0, y: 0, z: 0 }
+}

crates/rustc_offload_frontend/src/lib.rs

@@ -0,0 +1,70 @@
+#![allow(internal_features)]
+#![allow(linker_messages)]
+#![allow(improper_ctypes)]
+#![allow(improper_gpu_kernel_arg)]
+#![allow(improper_ctypes_definitions)]
+#![feature(gpu_offload, core_intrinsics, gpu_intrinsics, offload)]
+#![cfg_attr(target_arch = "nvptx64", feature(stdarch_nvptx))]
+#![cfg_attr(target_arch = "amdgpu", feature(stdarch_amdgpu))]
+#![cfg_attr(any(target_arch = "nvptx64", target_arch = "amdgpu"), no_std)]
+
+pub use core::offload::offload_kernel;
+
+pub mod gpu;
+pub mod partition;
+
+#[macro_export]
+macro_rules! offload {
+    ( $($field:ident = $val:expr),* $(,)? ) => {
+        $crate::offload!(@munch
+            [ $($field = $val),* ];
+            kernel = NONE;
+            grid_dim = ([1, 1, 1]);
+            block_dim = ([1, 1, 1]);
+            dyn_cache = (0);
+            args = NONE
+        );
+    };
+
+    (@munch [kernel = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; grid_dim = $g:tt; block_dim = $b:tt; dyn_cache = $d:tt; args = $a:tt) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = (SOME $val); grid_dim = $g; block_dim = $b; dyn_cache = $d; args = $a);
+    };
+    (@munch [grid_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; grid_dim = $g:tt; block_dim = $b:tt; dyn_cache = $d:tt; args = $a:tt) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; grid_dim = ($val); block_dim = $b; dyn_cache = $d; args = $a);
+    };
+    (@munch [block_dim = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; grid_dim = $g:tt; block_dim = $b:tt; dyn_cache = $d:tt; args = $a:tt) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; grid_dim = $g; block_dim = ($val); dyn_cache = $d; args = $a);
+    };
+    (@munch [dyn_cache = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; grid_dim = $g:tt; block_dim = $b:tt; dyn_cache = $d:tt; args = $a:tt) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; grid_dim = $g; block_dim = $b; dyn_cache = ($val); args = $a);
+    };
+    (@munch [args = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; grid_dim = $g:tt; block_dim = $b:tt; dyn_cache = $d:tt; args = $a:tt) => {
+        $crate::offload!(@munch [$($rest_f = $rest_v),*]; kernel = $k; grid_dim = $g; block_dim = $b; dyn_cache = $d; args = (SOME $val));
+    };
+
+    (@munch [$invalid:ident = $val:expr $(, $rest_f:ident = $rest_v:expr)*]; kernel = $k:tt; grid_dim = $g:tt; block_dim = $b:tt; dyn_cache = $d:tt; args = $a:tt) => {
+        compile_error!(concat!("unknown field ", stringify!($invalid)));
+    };
+
+    (@munch []; kernel = NONE; grid_dim = $g:tt; block_dim = $b:tt; dyn_cache = $d:tt; args = $a:tt) => {
+        compile_error!("missing `kernel`");
+    };
+    (@munch []; kernel = $k:tt; grid_dim = $g:tt; block_dim = $b:tt; dyn_cache = $d:tt; args = NONE) => {
+        compile_error!("missing `args`");
+    };
+    (@munch []; kernel = (SOME $kernel:expr); grid_dim = ($grid_dim:expr); block_dim = ($block_dim:expr); dyn_cache = ($dyn_cache:expr); args = (SOME $args:expr)) => {
+        core::intrinsics::offload::<_, _, ()>(
+            $kernel,
+            $grid_dim,
+            $block_dim,
+            $dyn_cache,
+            $args,
+        )
+    };
+}
+
+#[cfg(target_arch = "nvptx64")]
+#[panic_handler]
+fn panic(_: &core::panic::PanicInfo) -> ! {
+    loop {}
+}