From f6e4dbf3e3b6655a46658fe43fa2c1491fb2611a Mon Sep 17 00:00:00 2001 From: Dennis Date: Sun, 24 Apr 2022 11:58:31 +0200 Subject: [PATCH] Add GPU support for node graph --- node-graph/Cargo.lock | 54 ++++++++++++++++++++ node-graph/gcore/Cargo.toml | 10 ++-- node-graph/gcore/src/generic.rs | 22 ++++----- node-graph/gcore/src/lib.rs | 8 ++- node-graph/gcore/src/ops.rs | 79 ++++++++++++++++++++++++++---- node-graph/gcore/src/structural.rs | 70 ++++++++++++++++++++++++-- node-graph/gcore/src/value.rs | 15 ++---- node-graph/gstd/src/lib.rs | 4 +- node-graph/gstd/src/main.rs | 14 +++--- node-graph/gstd/src/memo.rs | 2 +- node-graph/gstd/src/value.rs | 6 +-- 11 files changed, 232 insertions(+), 52 deletions(-) diff --git a/node-graph/Cargo.lock b/node-graph/Cargo.lock index b9886480..fbc223f0 100644 --- a/node-graph/Cargo.lock +++ b/node-graph/Cargo.lock @@ -243,6 +243,15 @@ version = "0.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7ab85b9b05e3978cc9a9cf8fea7f01b494e1a09ed3037e16ba39edc7a29eb61a" +[[package]] +name = "glam" +version = "0.20.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f43e957e744be03f5801a55472f593d43fabdebf25a4585db250f04d86b1675f" +dependencies = [ + "num-traits", +] + [[package]] name = "graph-proc-macros" version = "0.1.0" @@ -259,6 +268,7 @@ name = "graphene-core" version = "0.1.0" dependencies = [ "dyn-any", + "spirv-std", ] [[package]] @@ -358,6 +368,12 @@ version = "0.2.121" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "efaa7b300f3b5fe8eb6bf21ce3895e1751d9665086af2d64b42f19701015ff4f" +[[package]] +name = "libm" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33a33a362ce288760ec6a508b94caaec573ae7d3bbbd91b87aa0bad4456839db" + [[package]] name = "lock_api" version = "0.4.7" @@ -407,6 +423,16 @@ dependencies = [ "windows-sys 0.28.0", ] +[[package]] +name = "num-traits" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a64b1ec5cda2586e284722486d802acf1f7dbdc623e2bfc57e65ca1cd099290" +dependencies = [ + "autocfg", + "libm", +] + [[package]] name = "num_cpus" version = "1.13.1" @@ -1092,6 +1118,34 @@ dependencies = [ "serde", ] +[[package]] +name = "spirv-std" +version = "0.4.0-alpha.12" +source = "git+https://github.com/EmbarkStudios/rust-gpu#685c79a9725f9fb68ce17dd00e742b98a03f3870" +dependencies = [ + "bitflags", + "glam", + "num-traits", + "spirv-std-macros", + "spirv-types", +] + +[[package]] +name = "spirv-std-macros" +version = "0.4.0-alpha.12" +source = "git+https://github.com/EmbarkStudios/rust-gpu#685c79a9725f9fb68ce17dd00e742b98a03f3870" +dependencies = [ + "proc-macro2", + "quote", + "spirv-types", + "syn", +] + +[[package]] +name = "spirv-types" +version = "0.4.0-alpha.12" +source = "git+https://github.com/EmbarkStudios/rust-gpu#685c79a9725f9fb68ce17dd00e742b98a03f3870" + [[package]] name = "storage-map" version = "0.3.0" diff --git a/node-graph/gcore/Cargo.toml b/node-graph/gcore/Cargo.toml index b88d474a..1b635347 100644 --- a/node-graph/gcore/Cargo.toml +++ b/node-graph/gcore/Cargo.toml @@ -8,8 +8,12 @@ authors = ["Dennis Kobert "] # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [features] -#std = [] -#default = ["std"] +std = ["dyn-any"] +default = ["gpu"] +gpu = ["spirv-std"] +nightly = [] [dependencies] -dyn-any = {path = "../dyn-any", features = ["derive"]} +dyn-any = {path = "../dyn-any", features = ["derive"], optional = true} + +spirv-std = { git = "https://github.com/EmbarkStudios/rust-gpu", features = ["glam"] , optional = true} diff --git a/node-graph/gcore/src/generic.rs b/node-graph/gcore/src/generic.rs index 7d4e9c73..39822d53 100644 --- a/node-graph/gcore/src/generic.rs +++ b/node-graph/gcore/src/generic.rs @@ -1,38 +1,38 @@ -use core::{borrow::Borrow, marker::PhantomData}; +use core::marker::PhantomData; use crate::Node; -pub struct FnNode O, In, O>(T, PhantomData, PhantomData); -impl<'n, T: Fn(&In) -> O, In, O: 'n> Node<'n, In> for FnNode { +pub struct FnNode O, In, O>(T, PhantomData, PhantomData); +impl<'n, T: Fn(In) -> O, In, O: 'n> Node<'n, In> for FnNode { type Output = O; - fn eval(&'n self, input: &'n In) -> Self::Output { - self.0(input.borrow()) + fn eval(&'n self, input: In) -> Self::Output { + self.0(input) } } -impl O, In, O> FnNode { +impl O, In, O> FnNode { pub fn new(f: T) -> Self { FnNode(f, PhantomData::default(), PhantomData::default()) } } -pub struct FnNodeWithState O, In, O, State>( +pub struct FnNodeWithState O, In, O, State>( T, State, PhantomData, PhantomData, ); -impl<'n, T: Fn(&In, &State) -> O, In, O: 'n, State> Node<'n, In> +impl<'n, T: Fn(In, &State) -> O, In, O: 'n, State> Node<'n, In> for FnNodeWithState { type Output = O; - fn eval(&'n self, input: &'n In) -> Self::Output { - self.0(input.borrow(), &self.1) + fn eval(&'n self, input: In) -> Self::Output { + self.0(input, &self.1) } } -impl O, In, O, State> FnNodeWithState { +impl O, In, O, State> FnNodeWithState { pub fn new(f: T, state: State) -> Self { FnNodeWithState(f, state, PhantomData::default(), PhantomData::default()) } diff --git a/node-graph/gcore/src/lib.rs b/node-graph/gcore/src/lib.rs index 74ec48ff..3e43e6eb 100644 --- a/node-graph/gcore/src/lib.rs +++ b/node-graph/gcore/src/lib.rs @@ -1,4 +1,5 @@ #![no_std] +#![cfg_attr(target_arch = "spirv", feature(register_attr), register_attr(spirv))] pub mod generic; pub mod ops; @@ -9,12 +10,13 @@ pub mod value; pub trait Node< 'n, Input> { type Output : 'n; - fn eval(&'n self, input: &'n Input) -> Self::Output; + fn eval(&'n self, input: Input) -> Self::Output; } +// TODO: Fix exec trait pub trait Exec<'n>: Node<'n, ()> { fn exec(&'n self) -> Self::Output { - self.eval(&()) + self.eval(()) } } impl<'n, T: Node<'n, ()>> Exec<'n> for T {} @@ -23,7 +25,9 @@ pub trait Cache { fn clear(&mut self); } +#[cfg(not(feature = "gpu"))] extern crate alloc; +#[cfg(not(feature = "gpu"))] impl<'n, I, O: 'n> Node<'n, I> for alloc::boxed::Box> { type Output = O; diff --git a/node-graph/gcore/src/ops.rs b/node-graph/gcore/src/ops.rs index 75145e73..c6444ff6 100644 --- a/node-graph/gcore/src/ops.rs +++ b/node-graph/gcore/src/ops.rs @@ -1,35 +1,94 @@ -use core::{borrow::Borrow, marker::PhantomData, ops::Add}; +use core::{marker::PhantomData, ops::Add}; use crate::Node; +#[repr(C)] #[derive(Default)] pub struct AddNode(PhantomData); impl<'n, T: Add + Copy + 'n> Node<'n, (T, T)> for AddNode { type Output = ::Output; - fn eval(&'n self, input: &'n (T, T)) -> T::Output { - let (ref a, ref b) = input.borrow(); - *a + *b + fn eval(&'n self, input: (T, T)) -> T::Output { + let (a, b) = input; + a + b } } +#[repr(C)] #[derive(Default)] /// Destructures a Tuple of two values and returns the first one pub struct FstNode(PhantomData, PhantomData); impl<'n, T: Copy + 'n, U> Node<'n, (T, U)> for FstNode { - type Output = &'n T; - fn eval(&'n self, input: &'n (T, U)) -> Self::Output { - let &(ref a, _) = input.borrow(); + type Output = T; + fn eval(&'n self, input: (T, U)) -> Self::Output { + let (a, _) = input; a } } +#[repr(C)] #[derive(Default)] /// Destructures a Tuple of two values and returns the first one pub struct SndNode(PhantomData, PhantomData); impl<'n, T, U: Copy + 'n> Node<'n, (T, U)> for SndNode { - type Output = &'n U; - fn eval(&'n self, input: &'n (T, U)) -> Self::Output { - let &(_, ref b) = input.borrow(); + type Output = U; + fn eval(&'n self, input: (T, U)) -> Self::Output { + let (_, b) = input; b } } +#[repr(C)] +#[derive(Default)] +/// Destructures a Tuple of two values and returns the first one +pub struct DupNode(PhantomData); +impl<'n, T: Copy + 'n> Node<'n, T> for DupNode { + type Output = (T, T); + fn eval(&'n self, input: T) -> Self::Output { + (input, input) + } +} + +#[cfg(target_arch = "spirv")] +pub mod gpu { + //#![deny(warnings)] + #[repr(C)] + pub struct PushConsts { + n: u32, + node: u32, + } + use super::*; + use crate::{structural::ComposeNodeOwned, Node}; + //use crate::Node; + use spirv_std::glam::UVec3; + const ADD: AddNode = AddNode(PhantomData); + const OPERATION: ComposeNodeOwned<'_, (u32, u32), u32, FstNode, DupNode> = + ComposeNodeOwned::new(FstNode(PhantomData, PhantomData), DupNode(PhantomData)); + + #[allow(unused)] + #[spirv(compute(threads(64)))] + pub fn spread( + #[spirv(global_invocation_id)] global_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] a: &[(u32, u32)], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] y: &mut [(u32, u32)], + #[spirv(push_constant)] push_consts: &PushConsts, + ) { + let gid = global_id.x as usize; + // Only process up to n, which is the length of the buffers. + if global_id.x < push_consts.n { + y[gid] = OPERATION.eval(a[gid]); + } + } + #[allow(unused)] + #[spirv(compute(threads(64)))] + pub fn add( + #[spirv(global_invocation_id)] global_id: UVec3, + #[spirv(storage_buffer, descriptor_set = 0, binding = 0)] a: &[(u32, u32)], + #[spirv(storage_buffer, descriptor_set = 0, binding = 1)] y: &mut [u32], + #[spirv(push_constant)] push_consts: &PushConsts, + ) { + let gid = global_id.x as usize; + // Only process up to n, which is the length of the buffers. + if global_id.x < push_consts.n { + y[gid] = ADD.eval(a[gid]); + } + } +} diff --git a/node-graph/gcore/src/structural.rs b/node-graph/gcore/src/structural.rs index d1acf51b..f960f817 100644 --- a/node-graph/gcore/src/structural.rs +++ b/node-graph/gcore/src/structural.rs @@ -12,20 +12,36 @@ pub struct ComposeNode<'n, Input, Inter, FIRST, SECOND> { impl<'n, Input: 'n, Inter: 'n, First, Second> Node<'n, Input> for ComposeNode<'n, Input, Inter, First, Second> where - First: Node<'n, Input, Output = &'n Inter>, + First: Node<'n, Input, Output = Inter>, Second: Node<'n, Inter>, /*+ Node<>::Output<'n>>*/ { type Output = >::Output; - fn eval(&'n self, input: &'n Input) -> Self::Output { + fn eval(&'n self, input: Input) -> Self::Output { // evaluate the first node with the given input // and then pipe the result from the first computation // into the second node - let arg: &Inter = self.first.eval(input); + let arg: Inter = self.first.eval(input); self.second.eval(arg) } } +#[cfg(feature = "nightly")] +impl<'n, Input, Inter, FIRST, SECOND> ComposeNode<'n, Input, Inter, FIRST, SECOND> +where + FIRST: Node<'n, Input>, +{ + pub const fn new(first: &'n FIRST, second: &'n SECOND) -> Self { + ComposeNode::<'n, Input, Inter, FIRST, SECOND> { + first, + second, + _phantom: PhantomData, + _phantom2: PhantomData, + } + } +} + +#[cfg(not(feature = "nightly"))] impl<'n, Input, Inter, FIRST, SECOND> ComposeNode<'n, Input, Inter, FIRST, SECOND> where FIRST: Node<'n, Input>, @@ -39,6 +55,54 @@ where } } } +#[repr(C)] +pub struct ComposeNodeOwned<'n, Input, Inter, FIRST, SECOND> { + first: FIRST, + second: SECOND, + _phantom: PhantomData<&'n Input>, + _phantom2: PhantomData, +} + +impl<'n, Input: 'n, Inter: 'n, First, Second> Node<'n, Input> + for ComposeNodeOwned<'n, Input, Inter, First, Second> +where + First: Node<'n, Input, Output = Inter>, + Second: Node<'n, Inter>, +{ + type Output = >::Output; + + fn eval(&'n self, input: Input) -> Self::Output { + // evaluate the first node with the given input + // and then pipe the result from the first computation + // into the second node + let arg: Inter = self.first.eval(input); + self.second.eval(arg) + } +} + +impl<'n, Input, Inter, First: 'n, Second> ComposeNodeOwned<'n, Input, Inter, First, Second> +where + First: Node<'n, Input, Output = Inter>, +{ + #[cfg(feature = "nightly")] + pub const fn new(first: First, second: Second) -> Self { + ComposeNodeOwned::<'n, Input, Inter, First, Second> { + first, + second, + _phantom: PhantomData, + _phantom2: PhantomData, + } + } + #[cfg(not(feature = "nightly"))] + pub fn new(first: First, second: Second) -> Self { + ComposeNodeOwned::<'n, Input, Inter, First, Second> { + first, + second, + _phantom: PhantomData, + _phantom2: PhantomData, + } + } +} pub trait After: Sized { fn after<'n, First: Node<'n, I>>( diff --git a/node-graph/gcore/src/value.rs b/node-graph/gcore/src/value.rs index 696291df..11bb761d 100644 --- a/node-graph/gcore/src/value.rs +++ b/node-graph/gcore/src/value.rs @@ -5,25 +5,20 @@ use crate::Node; pub struct IntNode; impl<'n, const N: u32> Node<'n, ()> for IntNode { type Output = u32; - fn eval(&self, _input: &()) -> u32 { + fn eval(&self, _input: ()) -> u32 { N } } -use dyn_any::{DynAny, StaticType, StaticTypeSized}; #[derive(Default)] pub struct ValueNode<'n, T>(T, PhantomData<&'n ()>); impl<'n, T: 'n> Node<'n, ()> for ValueNode<'n, T> { type Output = &'n T; - fn eval(&self, _input: &()) -> &T { + fn eval(&self, _input: ()) -> &T { &self.0 } } -impl<'n, T: StaticTypeSized> StaticType for ValueNode<'n, T> { - type Static = ValueNode<'static, ::Static>; -} - impl<'n, T> ValueNode<'n, T> { pub const fn new(value: T) -> ValueNode<'n, T> { ValueNode(value, PhantomData) @@ -34,7 +29,7 @@ impl<'n, T> ValueNode<'n, T> { pub struct DefaultNode(PhantomData); impl<'n, T: Default + 'n> Node<'n, ()> for DefaultNode { type Output = T; - fn eval(&self, _input: &()) -> T { + fn eval(&self, _input: ()) -> T { T::default() } } @@ -47,8 +42,8 @@ impl DefaultNode { pub struct DefaultRefNode<'n, T>(ValueNode<'n, T>); impl<'n, T: 'n> Node<'n, ()> for DefaultRefNode<'n, T> { type Output = &'n T; - fn eval(&'n self, _input: &'n ()) -> &'n T { - self.0.eval(&()) + fn eval(&'n self, _input: ()) -> &'n T { + self.0.eval(()) } } impl<'n, T: Default> Default for DefaultRefNode<'n, T> { diff --git a/node-graph/gstd/src/lib.rs b/node-graph/gstd/src/lib.rs index f436c9fb..22ba75eb 100644 --- a/node-graph/gstd/src/lib.rs +++ b/node-graph/gstd/src/lib.rs @@ -17,11 +17,11 @@ pub trait DynamicInput<'n> { fn set_arg_by_index(&mut self, index: usize, value: DynAnyNode<'n>); } -pub trait AnyRef<'n, I: StaticType>: Node<'n, I> { +pub trait AnyRef<'n, I: 'n + StaticType>: Node<'n, &'n I> { fn any(&'n self, input: &'n dyn DynAny<'n>) -> Self::Output; } -impl<'n, T: Node<'n, I>, I: StaticType + 'n> AnyRef<'n, I> for T { +impl<'n, N: Node<'n, &'n I>, I: StaticType + 'n> AnyRef<'n, I> for N { fn any(&'n self, input: &'n dyn DynAny<'n>) -> Self::Output { self.eval(downcast_ref::(input).unwrap_or_else(|| { panic!( diff --git a/node-graph/gstd/src/main.rs b/node-graph/gstd/src/main.rs index b5e9a6b3..ea7ef50b 100644 --- a/node-graph/gstd/src/main.rs +++ b/node-graph/gstd/src/main.rs @@ -27,11 +27,11 @@ mod mul { } impl<'n> Node<'n, ()> for MulNodeAnyProxy<'n> { type Output = MulNodeInput<'n>; - fn eval(&'n self, _input: &'n ()) -> >::Output { - let a = self.a.unwrap().eval(&()); + fn eval(&'n self, _input: ()) -> >::Output { + let a = self.a.unwrap().eval(()); let a: &f32 = self .a - .map(|v| downcast_ref(v.eval(&())).unwrap()) + .map(|v| downcast_ref(v.eval(())).unwrap()) .unwrap_or(&1.); /*let b: &f32 = self .b @@ -43,9 +43,9 @@ mod mul { } impl<'n> Node<'n, ()> for MulNodeTypedProxy<'n> { type Output = MulNodeInput<'n>; - fn eval(&'n self, _input: &'n ()) -> >::Output { - let a = self.a.unwrap().eval(&()); - let b = self.b.unwrap().eval(&()); + fn eval(&'n self, _input: ()) -> >::Output { + let a = self.a.unwrap().eval(()); + let b = self.b.unwrap().eval(()); MulNodeInput { a, b } } } @@ -98,7 +98,7 @@ fn main() { //let node = unsafe { stack.get(0) }; //let boxed = Box::new(StorageNode::new(node)); //unsafe { stack.push(boxed) }; - let result = unsafe { &stack.get()[0] }.eval(&()); + let result = unsafe { &stack.get()[0] }.eval(()); /*unsafe { stack .push(Box::new(AnyRefNode::new(stack.get(0).as_ref())) diff --git a/node-graph/gstd/src/memo.rs b/node-graph/gstd/src/memo.rs index b5e382fd..1982765a 100644 --- a/node-graph/gstd/src/memo.rs +++ b/node-graph/gstd/src/memo.rs @@ -8,7 +8,7 @@ pub struct CacheNode<'n, CachedNode: Node<'n, Input>, Input> { } impl<'n, CashedNode: Node<'n, Input>, Input> Node<'n, Input> for CacheNode<'n, CashedNode, Input> { type Output = &'n CashedNode::Output; - fn eval(&'n self, input: &'n Input) -> Self::Output { + fn eval(&'n self, input: Input) -> Self::Output { self.cache.get_or_init(|| self.node.eval(input)) } } diff --git a/node-graph/gstd/src/value.rs b/node-graph/gstd/src/value.rs index 8ba1d446..3330aac6 100644 --- a/node-graph/gstd/src/value.rs +++ b/node-graph/gstd/src/value.rs @@ -14,7 +14,7 @@ impl<'n, N: Node<'n, I, Output = &'n O>, I, O: DynAny<'n>> Node<'n, I> for AnyRefNode<'n, N, I, &'n O> { type Output = &'n (dyn DynAny<'n>); - fn eval(&'n self, input: &'n I) -> Self::Output { + fn eval(&'n self, input: I) -> Self::Output { let value: &O = self.0.eval(input); value } @@ -29,7 +29,7 @@ pub struct StorageNode<'n>(&'n dyn Node<'n, (), Output = &'n dyn DynAny<'n>>); impl<'n> Node<'n, ()> for StorageNode<'n> { type Output = &'n (dyn DynAny<'n>); - fn eval(&'n self, input: &'n ()) -> Self::Output { + fn eval(&'n self, input: ()) -> Self::Output { let value = self.0.eval(input); value } @@ -44,7 +44,7 @@ impl<'n> StorageNode<'n> { pub struct AnyValueNode<'n, T>(T, PhantomData<&'n ()>); impl<'n, T: 'n + DynAny<'n>> Node<'n, ()> for AnyValueNode<'n, T> { type Output = &'n dyn DynAny<'n>; - fn eval(&'n self, _input: &()) -> &'n dyn DynAny<'n> { + fn eval(&'n self, _input: ()) -> &'n dyn DynAny<'n> { &self.0 } }