Restructure GPU execution to model GPU pipelines in the node graph (#1088)

* Start implementing GpuExecutor for wgpu

* Implement read_output_buffer function

* Implement extraction node in the compiler

* Generate type annotations during shader compilation

* Start adding node wrapprs for graph execution api

* Wrap more of the api in nodes

* Restructure Pipeline to accept arbitrary shader inputs

* Adapt nodes to new trait definitions

* Start implementing gpu-compiler trait

* Adapt shader generation

* Hardstuck on pointer casts

* Pass nodes as references in gpu code to avoid zsts

* Update gcore to compile on the gpu

* Fix color doc tests

* Impl Node for node refs

node-graph/compilation-client/Cargo.toml

@@ -9,9 +9,18 @@ license = "MIT OR Apache-2.0"
 [dependencies]
 #tokio = { version = "1.0", features = ["full"] }
 serde_json = "1.0"
-graph-craft = { version = "0.1.0", path = "../graph-craft", features = ["serde"] }
+graph-craft = { version = "0.1.0", path = "../graph-craft", features = [
+	"serde",
+] }
+gpu-executor = { version = "0.1.0", path = "../gpu-executor" }
 gpu-compiler-bin-wrapper = { version = "0.1.0", path = "../gpu-compiler/gpu-compiler-bin-wrapper" }
 tempfile = "3.3.0"
 anyhow = "1.0.68"
-reqwest = { version = "0.11", features = ["blocking", "serde_json", "json", "rustls", "rustls-tls"] }
-future-executor = {path = "../future-executor"}
+reqwest = { version = "0.11", features = [
+	"blocking",
+	"serde_json",
+	"json",
+	"rustls",
+	"rustls-tls",
+] }
+future-executor = { path = "../future-executor" }

node-graph/compilation-client/src/lib.rs

@@ -1,15 +1,30 @@
 use gpu_compiler_bin_wrapper::CompileRequest;
-use graph_craft::document::*;
+use gpu_executor::ShaderIO;
+use graph_craft::{proto::ProtoNetwork, Type};
 
-pub async fn compile<I, O>(network: NodeNetwork) -> Result<Vec<u8>, reqwest::Error> {
+pub async fn compile(network: ProtoNetwork, inputs: Vec<Type>, output: Type, io: ShaderIO) -> Result<Shader, reqwest::Error> {
 	let client = reqwest::Client::new();
 
-	let compile_request = CompileRequest::new(network, std::any::type_name::<I>().to_owned(), std::any::type_name::<O>().to_owned());
+	let compile_request = CompileRequest::new(network, inputs.clone(), output.clone(), io.clone());
 	let response = client.post("http://localhost:3000/compile/spirv").json(&compile_request).send();
 	let response = response.await?;
-	response.bytes().await.map(|b| b.to_vec())
+	response.bytes().await.map(|b| Shader {
+		spirv_binary: b.windows(4).map(|x| u32::from_le_bytes(x.try_into().unwrap())).collect(),
+		input_types: inputs,
+		output_type: output,
+		io,
+	})
 }
 
-pub fn compile_sync<I: 'static, O: 'static>(network: NodeNetwork) -> Result<Vec<u8>, reqwest::Error> {
-	future_executor::block_on(compile::<I, O>(network))
+pub fn compile_sync(network: ProtoNetwork, inputs: Vec<Type>, output: Type, io: ShaderIO) -> Result<Shader, reqwest::Error> {
+	future_executor::block_on(compile(network, inputs, output, io))
+}
+
+// TODO: should we add the entry point as a field?
+/// A compiled shader with type annotations.
+pub struct Shader {
+	pub spirv_binary: Vec<u32>,
+	pub input_types: Vec<Type>,
+	pub output_type: Type,
+	pub io: ShaderIO,
 }

node-graph/compilation-client/src/main.rs

@@ -1,40 +1,55 @@
 use gpu_compiler_bin_wrapper::CompileRequest;
+use gpu_executor::{ShaderIO, ShaderInput};
 use graph_craft::concrete;
 use graph_craft::document::*;
-
 use graph_craft::*;
+
 use std::borrow::Cow;
+use std::time::Duration;
 
 fn main() {
 	let client = reqwest::blocking::Client::new();
 
-	let network = NodeNetwork {
-		inputs: vec![0],
-		outputs: vec![NodeOutput::new(0, 0)],
-		disabled: vec![],
-		previous_outputs: None,
-		nodes: [(
-			0,
-			DocumentNode {
-				name: "Inc".into(),
-				inputs: vec![NodeInput::Network(concrete!(u32))],
-				implementation: DocumentNodeImplementation::Network(add_network()),
-				metadata: DocumentNodeMetadata::default(),
-			},
-		)]
-		.into_iter()
-		.collect(),
+	// let network = NodeNetwork {
+	// 	inputs: vec![0],
+	// 	outputs: vec![NodeOutput::new(0, 0)],
+	// 	disabled: vec![],
+	// 	previous_outputs: None,
+	// 	nodes: [(
+	// 		0,
+	// 		DocumentNode {
+	// 			name: "Inc".into(),
+	// 			inputs: vec![NodeInput::Network(concrete!(u32))],
+	// 			implementation: DocumentNodeImplementation::Network(add_network()),
+	// 			metadata: DocumentNodeMetadata::default(),
+	// 		},
+	// 	)]
+	// 	.into_iter()
+	// 	.collect(),
+	// };
+	let network = add_network();
+	let compiler = graph_craft::executor::Compiler {};
+	let proto_network = compiler.compile_single(network, true).unwrap();
+
+	let io = ShaderIO {
+		inputs: vec![ShaderInput::StorageBuffer((), concrete!(u32))],
+		output: ShaderInput::OutputBuffer((), concrete!(&mut [u32])),
 	};
 
-	let compile_request = CompileRequest::new(network, "u32".to_owned(), "u32".to_owned());
-	let response = client.post("http://localhost:3000/compile/spirv").json(&compile_request).send().unwrap();
+	let compile_request = CompileRequest::new(proto_network, vec![concrete!(u32)], concrete!(u32), io);
+	let response = client
+		.post("http://localhost:3000/compile/spirv")
+		.timeout(Duration::from_secs(30))
+		.json(&compile_request)
+		.send()
+		.unwrap();
 	println!("response: {:?}", response);
 }
 
 fn add_network() -> NodeNetwork {
 	NodeNetwork {
-		inputs: vec![0],
-		outputs: vec![NodeOutput::new(1, 0)],
+		inputs: vec![],
+		outputs: vec![NodeOutput::new(0, 0)],
 		disabled: vec![],
 		previous_outputs: None,
 		nodes: [
@@ -42,20 +57,20 @@ fn add_network() -> NodeNetwork {
 				0,
 				DocumentNode {
 					name: "Dup".into(),
-					inputs: vec![NodeInput::Network(concrete!(u32))],
+					inputs: vec![NodeInput::value(value::TaggedValue::U32(5u32), false)],
 					metadata: DocumentNodeMetadata::default(),
-					implementation: DocumentNodeImplementation::Unresolved(NodeIdentifier::new("graphene_core::ops::DupNode")),
-				},
-			),
-			(
-				1,
-				DocumentNode {
-					name: "Add".into(),
-					inputs: vec![NodeInput::node(0, 0)],
-					metadata: DocumentNodeMetadata::default(),
-					implementation: DocumentNodeImplementation::Unresolved(NodeIdentifier::new("graphene_core::ops::AddNode")),
+					implementation: DocumentNodeImplementation::Unresolved(NodeIdentifier::new("graphene_core::ops::IdNode")),
 				},
 			),
+			// (
+			// 	1,
+			// 	DocumentNode {
+			// 		name: "Add".into(),
+			// 		inputs: vec![NodeInput::node(0, 0)],
+			// 		metadata: DocumentNodeMetadata::default(),
+			// 		implementation: DocumentNodeImplementation::Unresolved(NodeIdentifier::new("graphene_core::ops::AddNode")),
+			// 	},
+			// ),
 		]
 		.into_iter()
 		.collect(),