[Draft/RFC] cuda-oxide (native CUDA) device build for vortx-shaders

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" }

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)]
 

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]

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);

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<T>(thread_id: usize, stride: usize, workspace: &mut [T; WORKGROUP_SIZE])
+#[inline(always)]
+fn reduce_workspace_sum<T>(thread_id: usize, stride: usize, workspace: &mut impl MaybeIndexUnchecked<T>)
 where
     T: Copy + Add<T, Output = T>,
 {
@@ -380,8 +379,8 @@ where
     khal_std::sync::workgroup_memory_barrier_with_group_sync();
 }
 
-#[inline]
-fn reduce_workspace_mul<T>(thread_id: usize, stride: usize, workspace: &mut [T; WORKGROUP_SIZE])
+#[inline(always)]
+fn reduce_workspace_mul<T>(thread_id: usize, stride: usize, workspace: &mut impl MaybeIndexUnchecked<T>)
 where
     T: Copy + Mul<T, Output = T>,
 {

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,