diff --git a/Cargo.toml b/Cargo.toml index 1373ea1..d4b9569 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,8 +29,18 @@ unsafe_remove_boundchecks = [] version = "0.3.0" [workspace.dependencies] +<<<<<<< HEAD +<<<<<<< HEAD khal-std = "0.2" khal = { version = "0.2", features = ["derive"]} +======= +khal-std = { version = "0.1", default-features = false } +khal = { version = "0.1", features = ["derive"]} +>>>>>>> 28ae1be (vortx-shaders: build cleanly for the nvptx64 (cuda-oxide) device target) +======= +khal-std = { version = "0.2", default-features = false } +khal = { version = "0.2", features = ["derive"]} +>>>>>>> a7bfba8 (vortx: cuda-oxide native-CUDA feature wiring (sm_120)) [dependencies] bytemuck = "1" @@ -53,13 +63,19 @@ anyhow = "1" wgpu = "29" [build-dependencies] +<<<<<<< HEAD khal-builder = "0.2" +======= +khal-builder = "0.2.0" +>>>>>>> a7bfba8 (vortx: cuda-oxide native-CUDA feature wiring (sm_120)) # To build the shader from the dependency instead of local path. vortx-shaders = { version = "0.3", path = "./vortx-shaders" } [patch.crates-io] -#khal-builder = { path = "../khal/crates/khal-builder" } -#khal = { path = "../khal/crates/khal" } -#khal-std = { path = "../khal/crates/khal-std" } -#khal-derive = { path = "../khal/crates/khal-derive" } +# Local khal forks — required for the standalone cuda-oxide cubin build +# (crates.io khal-std lacks the `cuda-oxide` feature). +khal-builder = { path = "../khal/crates/khal-builder" } +khal = { path = "../khal/crates/khal" } +khal-std = { path = "../khal/crates/khal-std" } +khal-derive = { path = "../khal/crates/khal-derive" } #glamx = { git = "https://github.com/dimforge/glamx", branch = "bytemuck" } diff --git a/vortx-shaders/src/lib.rs b/vortx-shaders/src/lib.rs index 6ab2b13..96a05ca 100644 --- a/vortx-shaders/src/lib.rs +++ b/vortx-shaders/src/lib.rs @@ -3,7 +3,7 @@ //! This crate contains GPU shaders for tensor operations, geometry transformations, //! and linear algebra primitives, written for rust-gpu. -#![cfg_attr(target_arch = "spirv", no_std)] +#![cfg_attr(any(target_arch = "spirv", target_arch = "nvptx64"), no_std)] #![allow(unexpected_cfgs)] #![allow(clippy::too_many_arguments)] diff --git a/vortx-shaders/src/linalg/contiguous.rs b/vortx-shaders/src/linalg/contiguous.rs index cd64375..7516444 100644 --- a/vortx-shaders/src/linalg/contiguous.rs +++ b/vortx-shaders/src/linalg/contiguous.rs @@ -3,7 +3,6 @@ use super::shape::Shape; #[cfg(feature = "push_constants")] use super::shape::Shapes1; -use crate::utils::iterators::StepRng; use crate::utils::limits::MAX_NUM_WORKGROUPS; use glamx::UVec3; use khal_std::{ @@ -63,7 +62,7 @@ fn contiguous_impl( src: &[u32], offset: u32, ) { - for thread_id in StepRng::new(invocation_id.x..shape_src.len(), MAX_NUM_THREADS) { + for thread_id in (invocation_id.x..shape_src.len()).step_by(MAX_NUM_THREADS as usize) { // Compute the corresponding (i, j, k, l) indices for the out tensor. // A contiguous row-major tensor with dimensions [n,c,h,w] has strides // [c * h * w, h * w, w, 1] diff --git a/vortx-shaders/src/linalg/op_assign.rs b/vortx-shaders/src/linalg/op_assign.rs index f589446..0402b9f 100644 --- a/vortx-shaders/src/linalg/op_assign.rs +++ b/vortx-shaders/src/linalg/op_assign.rs @@ -3,7 +3,6 @@ use super::shape::Shape; #[cfg(feature = "push_constants")] use super::shape::Shapes2; -use crate::utils::iterators::StepRng; use crate::utils::limits::MAX_NUM_WORKGROUPS; use glamx::UVec3; use khal_std::{ @@ -17,7 +16,10 @@ const MAX_NUM_THREADS: u32 = MAX_NUM_WORKGROUPS * WORKGROUP_SIZE; /// Binary operation offsets. #[repr(C)] #[derive(Clone, Copy)] -#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] pub struct BinOpOffsets { pub a: u32, pub b: u32, @@ -45,7 +47,7 @@ pub fn gpu_add( #[cfg(feature = "push_constants")] let (shape_a, shape_b) = (&shapes.shape_a, &shapes.shape_b); - for thread_id in StepRng::new(invocation_id.x..shape_a.len(), MAX_NUM_THREADS) { + for thread_id in (invocation_id.x..shape_a.len()).step_by(MAX_NUM_THREADS as usize) { let id = shape_a.decompose(thread_id); let ia = shape_a.it_vec(id) as usize; let ib = shape_b.it_vec(id) as usize; @@ -73,7 +75,7 @@ pub fn gpu_sub( #[cfg(feature = "push_constants")] let (shape_a, shape_b) = (&shapes.shape_a, &shapes.shape_b); - for thread_id in StepRng::new(invocation_id.x..shape_a.len(), MAX_NUM_THREADS) { + for thread_id in (invocation_id.x..shape_a.len()).step_by(MAX_NUM_THREADS as usize) { let id = shape_a.decompose(thread_id); let ia = shape_a.it_vec(id) as usize; let ib = shape_b.it_vec(id) as usize; @@ -101,7 +103,7 @@ pub fn gpu_mul( #[cfg(feature = "push_constants")] let (shape_a, shape_b) = (&shapes.shape_a, &shapes.shape_b); - for thread_id in StepRng::new(invocation_id.x..shape_a.len(), MAX_NUM_THREADS) { + for thread_id in (invocation_id.x..shape_a.len()).step_by(MAX_NUM_THREADS as usize) { let id = shape_a.decompose(thread_id); let ia = shape_a.it_vec(id) as usize; let ib = shape_b.it_vec(id) as usize; @@ -129,7 +131,7 @@ pub fn gpu_div( #[cfg(feature = "push_constants")] let (shape_a, shape_b) = (&shapes.shape_a, &shapes.shape_b); - for thread_id in StepRng::new(invocation_id.x..shape_a.len(), MAX_NUM_THREADS) { + for thread_id in (invocation_id.x..shape_a.len()).step_by(MAX_NUM_THREADS as usize) { let id = shape_a.decompose(thread_id); let ia = shape_a.it_vec(id) as usize; let ib = shape_b.it_vec(id) as usize; @@ -157,7 +159,7 @@ pub fn gpu_copy( #[cfg(feature = "push_constants")] let (shape_a, shape_b) = (&shapes.shape_a, &shapes.shape_b); - for thread_id in StepRng::new(invocation_id.x..shape_a.len(), MAX_NUM_THREADS) { + for thread_id in (invocation_id.x..shape_a.len()).step_by(MAX_NUM_THREADS as usize) { let id = shape_a.decompose(thread_id); let ia = shape_a.it_vec(id) as usize; let ib = shape_b.it_vec(id) as usize; @@ -178,7 +180,7 @@ pub fn gpu_copy_with_offsets( #[spirv(storage_buffer, descriptor_set = 0, binding = 3)] a: &mut [f32], #[spirv(storage_buffer, descriptor_set = 0, binding = 4)] b: &[f32], ) { - for thread_id in StepRng::new(invocation_id.x..shape_a.len(), MAX_NUM_THREADS) { + for thread_id in (invocation_id.x..shape_a.len()).step_by(MAX_NUM_THREADS as usize) { let id = shape_a.decompose(thread_id); let ia = shape_a.it_vec(id); let ib = shape_b.it_vec(id); diff --git a/vortx-shaders/src/linalg/reduce.rs b/vortx-shaders/src/linalg/reduce.rs index 21e007a..e2984d8 100644 --- a/vortx-shaders/src/linalg/reduce.rs +++ b/vortx-shaders/src/linalg/reduce.rs @@ -3,7 +3,6 @@ use super::shape::Shape; #[cfg(feature = "push_constants")] use super::shape::Shapes1; -use crate::utils::iterators::StepRng; use core::ops::{Add, Mul}; use glamx::UVec3; use khal_std::{ @@ -41,7 +40,7 @@ macro_rules! decl_reductions { let thread_id = local_id.x as usize; workspace.write(thread_id, $zero); - for i in StepRng::new(thread_id as u32..shape.w, WORKGROUP_SIZE as u32) { + for i in (thread_id as u32..shape.w).step_by(WORKGROUP_SIZE) { // TODO: support tensors that are not just vectors. // We'd reduce along the last dimension only, and use the // workgroup_id to compute on each axis in parallel. @@ -99,7 +98,7 @@ macro_rules! decl_reductions { let thread_id = local_id.x as usize; workspace.write(thread_id, $one); - for i in StepRng::new(thread_id as u32..shape.w, WORKGROUP_SIZE as u32) { + for i in (thread_id as u32..shape.w).step_by(WORKGROUP_SIZE) { // TODO: support tensors that are not just vectors. // We'd reduce along the last dimension only, and use the // workgroup_id to compute on each axis in parallel. @@ -158,7 +157,7 @@ macro_rules! decl_reductions { let thread_id = local_id.x as usize; workspace.write(thread_id, $max); - for i in StepRng::new(thread_id as u32..shape.w, WORKGROUP_SIZE as u32) { + for i in (thread_id as u32..shape.w).step_by(WORKGROUP_SIZE) { // TODO: support tensors that are not just vectors. // We'd reduce along the last dimension only, and use the // workgroup_id to compute on each axis in parallel. @@ -173,11 +172,11 @@ macro_rules! decl_reductions { #[cfg(not(feature = "subgroup_ops"))] { - #[inline] + #[inline(always)] fn reduce_workspace_min( thread_id: usize, stride: usize, - workspace: &mut [$T; WORKGROUP_SIZE], + workspace: &mut impl MaybeIndexUnchecked<$T>, ) { if thread_id < stride { workspace.write( @@ -233,7 +232,7 @@ macro_rules! decl_reductions { let thread_id = local_id.x as usize; workspace.write(thread_id, $min); - for i in StepRng::new(thread_id as u32..shape.w, WORKGROUP_SIZE as u32) { + for i in (thread_id as u32..shape.w).step_by(WORKGROUP_SIZE) { // TODO: support tensors that are not just vectors. // We'd reduce along the last dimension only, and use the // workgroup_id to compute on each axis in parallel. @@ -249,11 +248,11 @@ macro_rules! decl_reductions { #[cfg(not(feature = "subgroup_ops"))] { - #[inline] + #[inline(always)] fn reduce_workspace_max( thread_id: usize, stride: usize, - workspace: &mut [$T; WORKGROUP_SIZE], + workspace: &mut impl MaybeIndexUnchecked<$T>, ) { if thread_id < stride { workspace.write( @@ -329,7 +328,7 @@ pub fn reduce_sq_norm( let thread_id = local_id.x as usize; workspace.write(thread_id, 0.0); - for i in StepRng::new(thread_id as u32..shape.w, WORKGROUP_SIZE as u32) { + for i in (thread_id as u32..shape.w).step_by(WORKGROUP_SIZE) { // TODO: support tensors that are not just vectors. // We'd reduce along the last dimension only, and use the // workgroup_id to compute on each axis in parallel. @@ -366,8 +365,8 @@ pub fn reduce_sq_norm( } } -#[inline] -fn reduce_workspace_sum(thread_id: usize, stride: usize, workspace: &mut [T; WORKGROUP_SIZE]) +#[inline(always)] +fn reduce_workspace_sum(thread_id: usize, stride: usize, workspace: &mut impl MaybeIndexUnchecked) where T: Copy + Add, { @@ -380,8 +379,8 @@ where khal_std::sync::workgroup_memory_barrier_with_group_sync(); } -#[inline] -fn reduce_workspace_mul(thread_id: usize, stride: usize, workspace: &mut [T; WORKGROUP_SIZE]) +#[inline(always)] +fn reduce_workspace_mul(thread_id: usize, stride: usize, workspace: &mut impl MaybeIndexUnchecked) where T: Copy + Mul, { diff --git a/vortx-shaders/src/linalg/shape.rs b/vortx-shaders/src/linalg/shape.rs index fa9a3aa..369f335 100644 --- a/vortx-shaders/src/linalg/shape.rs +++ b/vortx-shaders/src/linalg/shape.rs @@ -8,7 +8,10 @@ use glamx::UVec4; /// (Samples, Channels, Height, Width), where height is the row count, and width the column count. #[repr(C)] #[derive(Clone, Copy)] -#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] pub struct Shape { /// Number of rows in each matrix of the tensor. pub n: u32, @@ -100,7 +103,10 @@ pub fn div_ceil4(a: u32) -> u32 { #[cfg(feature = "push_constants")] #[repr(C)] #[derive(Clone, Copy)] -#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] pub struct Shapes2 { /// First shape (typically output or left operand). pub shape_a: Shape, @@ -112,7 +118,10 @@ pub struct Shapes2 { #[cfg(feature = "push_constants")] #[repr(C)] #[derive(Clone, Copy)] -#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] pub struct Shapes3 { /// Output shape. pub shape_out: Shape, @@ -126,7 +135,10 @@ pub struct Shapes3 { #[cfg(feature = "push_constants")] #[repr(C)] #[derive(Clone, Copy)] -#[cfg_attr(not(target_arch_is_gpu), derive(bytemuck::Pod, bytemuck::Zeroable))] +#[cfg_attr( + not(any(target_arch = "spirv", target_arch = "nvptx64")), + derive(bytemuck::Pod, bytemuck::Zeroable) +)] pub struct Shapes1 { /// The shape. pub shape: Shape,