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/gpu-compiler/src/lib.rs

@@ -1,6 +1,8 @@
-use std::path::Path;
-
+use gpu_executor::{GPUConstant, ShaderIO, ShaderInput, SpirVCompiler};
 use graph_craft::proto::*;
+use graphene_core::Cow;
+
+use std::path::{Path, PathBuf};
 use tera::Context;
 
 fn create_cargo_toml(metadata: &Metadata) -> Result<String, tera::Error> {
@@ -24,10 +26,10 @@ impl Metadata {
 	}
 }
 
-pub fn create_files(matadata: &Metadata, network: &ProtoNetwork, compile_dir: &Path, input_type: &str, output_type: &str) -> anyhow::Result<()> {
+pub fn create_files(metadata: &Metadata, network: &ProtoNetwork, compile_dir: &Path, io: &ShaderIO) -> anyhow::Result<()> {
 	let src = compile_dir.join("src");
 	let cargo_file = compile_dir.join("Cargo.toml");
-	let cargo_toml = create_cargo_toml(matadata)?;
+	let cargo_toml = create_cargo_toml(metadata)?;
 	std::fs::write(cargo_file, cargo_toml)?;
 
 	let toolchain_file = compile_dir.join("rust-toolchain.toml");
@@ -44,26 +46,100 @@ pub fn create_files(matadata: &Metadata, network: &ProtoNetwork, compile_dir: &P
 		}
 	}
 	let lib = src.join("lib.rs");
-	let shader = serialize_gpu(network, input_type, output_type)?;
-	println!("{}", shader);
+	let shader = serialize_gpu(network, io)?;
+	eprintln!("{}", shader);
 	std::fs::write(lib, shader)?;
 	Ok(())
 }
 
-pub fn serialize_gpu(network: &ProtoNetwork, input_type: &str, output_type: &str) -> anyhow::Result<String> {
-	assert_eq!(network.inputs.len(), 1);
+fn constant_attribute(constant: &GPUConstant) -> &'static str {
+	match constant {
+		GPUConstant::SubGroupId => "subgroup_id",
+		GPUConstant::SubGroupInvocationId => "subgroup_local_invocation_id",
+		GPUConstant::SubGroupSize => todo!(),
+		GPUConstant::NumSubGroups => "num_subgroups",
+		GPUConstant::WorkGroupId => "workgroup_id",
+		GPUConstant::WorkGroupInvocationId => "local_invocation_id",
+		GPUConstant::WorkGroupSize => todo!(),
+		GPUConstant::NumWorkGroups => "num_workgroups",
+		GPUConstant::GlobalInvocationId => "global_invocation_id",
+		GPUConstant::GlobalSize => todo!(),
+	}
+}
+
+pub fn construct_argument(input: &ShaderInput<()>, position: u32) -> String {
+	match input {
+		ShaderInput::Constant(constant) => format!("#[spirv({})] i{}: {},", constant_attribute(constant), position, constant.ty()),
+		ShaderInput::UniformBuffer(_, ty) => {
+			format!("#[spirv(uniform, descriptor_set = 0, binding = {})] i{}: &[{}]", position, position, ty,)
+		}
+		ShaderInput::StorageBuffer(_, ty) | ShaderInput::ReadBackBuffer(_, ty) => {
+			format!("#[spirv(storage_buffer, descriptor_set = 0, binding = {})] i{}: &[{}]", position, position, ty,)
+		}
+		ShaderInput::OutputBuffer(_, ty) => {
+			format!("#[spirv(storage_buffer, descriptor_set = 0, binding = {})] i{}: &mut[{}]", position, position, ty,)
+		}
+		ShaderInput::WorkGroupMemory(_, ty) => format!("#[spirv(workgroup_memory] i{}: {}", position, ty,),
+	}
+}
+
+struct GpuCompiler {
+	compile_dir: PathBuf,
+}
+
+impl SpirVCompiler for GpuCompiler {
+	fn compile(&self, network: ProtoNetwork, io: &ShaderIO) -> anyhow::Result<gpu_executor::Shader> {
+		let metadata = Metadata::new("project".to_owned(), vec!["test@example.com".to_owned()]);
+
+		create_files(&metadata, &network, &self.compile_dir, io)?;
+		let result = compile(&self.compile_dir)?;
+
+		let bytes = std::fs::read(result.module.unwrap_single())?;
+		let words = bytes.chunks(4).map(|chunk| u32::from_ne_bytes(chunk.try_into().unwrap())).collect::<Vec<_>>();
+
+		Ok(gpu_executor::Shader {
+			source: Cow::Owned(words),
+			name: "",
+			io: io.clone(),
+		})
+	}
+}
+
+pub fn serialize_gpu(network: &ProtoNetwork, io: &ShaderIO) -> anyhow::Result<String> {
 	fn nid(id: &u64) -> String {
 		format!("n{id}")
 	}
 
+	dbg!(&network);
+	dbg!(&io);
+	let inputs = io.inputs.iter().enumerate().map(|(i, input)| construct_argument(input, i as u32)).collect::<Vec<_>>();
+
 	let mut nodes = Vec::new();
+	let mut input_nodes = Vec::new();
 	#[derive(serde::Serialize)]
 	struct Node {
 		id: String,
 		fqn: String,
 		args: Vec<String>,
 	}
+	for id in network.inputs.iter() {
+		let Some((_, node)) = network.nodes.iter().find(|(i, _)| i == id) else {
+			anyhow::bail!("Input node not found");
+		};
+		let fqn = &node.identifier.name;
+		let id = nid(id);
+		input_nodes.push(Node {
+			id,
+			fqn: fqn.to_string().split("<").next().unwrap().to_owned(),
+			args: node.construction_args.new_function_args(),
+		});
+	}
+
 	for (ref id, node) in network.nodes.iter() {
+		if network.inputs.contains(id) {
+			continue;
+		}
+
 		let fqn = &node.identifier.name;
 		let id = nid(id);
 
@@ -78,8 +154,8 @@ pub fn serialize_gpu(network: &ProtoNetwork, input_type: &str, output_type: &str
 	let mut tera = tera::Tera::default();
 	tera.add_raw_template("spirv", template)?;
 	let mut context = Context::new();
-	context.insert("input_type", &input_type);
-	context.insert("output_type", &output_type);
+	context.insert("inputs", &inputs);
+	context.insert("input_nodes", &input_nodes);
 	context.insert("nodes", &nodes);
 	context.insert("last_node", &nid(&network.output));
 	context.insert("compute_threads", &64);
@@ -89,14 +165,15 @@ pub fn serialize_gpu(network: &ProtoNetwork, input_type: &str, output_type: &str
 use spirv_builder::{MetadataPrintout, SpirvBuilder, SpirvMetadata};
 pub fn compile(dir: &Path) -> Result<spirv_builder::CompileResult, spirv_builder::SpirvBuilderError> {
 	dbg!(&dir);
-	let result = SpirvBuilder::new(dir, "spirv-unknown-spv1.5")
+	let result = SpirvBuilder::new(dir, "spirv-unknown-vulkan1.2")
 		.print_metadata(MetadataPrintout::DependencyOnly)
 		.multimodule(false)
 		.preserve_bindings(true)
 		.release(true)
-		//.relax_struct_store(true)
-		//.relax_block_layout(true)
 		.spirv_metadata(SpirvMetadata::Full)
+		.extra_arg("no-early-report-zombies")
+		.extra_arg("no-infer-storage-classes")
+		.extra_arg("spirt-passes=qptr")
 		.build()?;
 
 	Ok(result)
@@ -104,7 +181,6 @@ pub fn compile(dir: &Path) -> Result<spirv_builder::CompileResult, spirv_builder
 
 #[cfg(test)]
 mod test {
-
 	#[test]
 	fn test_create_cargo_toml() {
 		let cargo_toml = super::create_cargo_toml(&super::Metadata {

node-graph/gpu-compiler/src/main.rs

@@ -1,4 +1,5 @@
 use gpu_compiler as compiler;
+use gpu_executor::CompileRequest;
 use graph_craft::document::NodeNetwork;
 use std::io::Write;
 
@@ -6,17 +7,13 @@ fn main() -> anyhow::Result<()> {
 	println!("Starting GPU Compiler!");
 	let mut stdin = std::io::stdin();
 	let mut stdout = std::io::stdout();
-	let input_type = std::env::args().nth(1).expect("input type arg missing");
-	let output_type = std::env::args().nth(2).expect("output type arg missing");
-	let compile_dir = std::env::args().nth(3).map(|x| std::path::PathBuf::from(&x)).unwrap_or(tempfile::tempdir()?.into_path());
-	let network: NodeNetwork = serde_json::from_reader(&mut stdin)?;
-	let compiler = graph_craft::executor::Compiler {};
-	let proto_network = compiler.compile_single(network, true).unwrap();
+	let compile_dir = std::env::args().nth(1).map(|x| std::path::PathBuf::from(&x)).unwrap_or(tempfile::tempdir()?.into_path());
+	let request: CompileRequest = serde_json::from_reader(&mut stdin)?;
 	dbg!(&compile_dir);
 
 	let metadata = compiler::Metadata::new("project".to_owned(), vec!["test@example.com".to_owned()]);
 
-	compiler::create_files(&metadata, &proto_network, &compile_dir, &input_type, &output_type)?;
+	compiler::create_files(&metadata, &request.network, &compile_dir, &request.io)?;
 	let result = compiler::compile(&compile_dir)?;
 
 	let bytes = std::fs::read(result.module.unwrap_single())?;

node-graph/gpu-compiler/src/templates/Cargo-template.toml

@@ -1,7 +1,7 @@
 [package]
+authors = [{% for author in authors %}"{{author}}", {% endfor %}]
 name = "{{name}}-node"
 version = "0.1.0"
-authors = [{%for author in authors%}"{{author}}", {%endfor%}]
 edition = "2021"
 license = "MIT OR Apache-2.0"
 publish = false
@@ -13,5 +13,7 @@ crate-type = ["dylib", "lib"]
 libm = { git = "https://github.com/rust-lang/libm", tag = "0.2.5" }
 
 [dependencies]
-spirv-std = { version = "0.5" , features= ["glam"]}
-graphene-core = {path = "{{gcore_path}}", default-features = false, features = ["gpu"]}
+spirv-std = { version = "0.7" }
+graphene-core = { path = "{{gcore_path}}", default-features = false, features = [
+	"gpu",
+] }

node-graph/gpu-compiler/src/templates/rust-toolchain.toml

@@ -1,3 +1,10 @@
 [toolchain]
-channel = "nightly-2022-12-18"
-components = ["rust-src", "rustc-dev", "llvm-tools-preview", "clippy", "cargofmt", "rustc"]
+channel = "nightly-2023-03-04"
+components = [
+	"rust-src",
+	"rustc-dev",
+	"llvm-tools-preview",
+	"clippy",
+	"rustfmt",
+	"rustc",
+]

node-graph/gpu-compiler/src/templates/spirv-template.rs

@@ -1,6 +1,5 @@
 #![no_std]
 #![feature(unchecked_math)]
-#![deny(warnings)]
 
 #[cfg(target_arch = "spirv")]
 extern crate spirv_std;
@@ -14,25 +13,23 @@ pub mod gpu {
 	#[allow(unused)]
 	#[spirv(compute(threads({{compute_threads}})))]
 	pub fn eval (
-		#[spirv(global_invocation_id)] global_id: UVec3,
-		#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] a: &[{{input_type}}],
-		#[spirv(storage_buffer, descriptor_set = 0, binding = 1)] y: &mut [{{output_type}}],
-		//#[spirv(push_constant)] push_consts: &graphene_core::gpu::PushConstants,
+		{% for input in inputs %}
+		{{input}}
+		{% endfor %}
 	) {
-		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] = node_graph(a[gid]);
-		//}
-	}
-
-	fn node_graph(input: {{input_type}}) -> {{output_type}} {
 		use graphene_core::Node;
 
+		{% for input in input_nodes %}
+		let i{{loop.index0}} = graphene_core::value::CopiedNode::new(i{{loop.index0}});
+		let _{{input.id}} = {{input.fqn}}::new({% for arg in input.args %}{{arg}}, {% endfor %});
+		let {{input.id}} = graphene_core::structural::ComposeNode::new(i{{loop.index0}}, _{{input.id}});
+		{% endfor %}
+
 		{% for node in nodes %}
 		let {{node.id}} = {{node.fqn}}::new({% for arg in node.args %}{{arg}}, {% endfor %});
 		{% endfor %}
-		{{last_node}}.eval(input)
-	}
+		let output = {{last_node}}.eval(());
+		// TODO: Write output to buffer
 
+	}
 }