Restructure GPU compilation execution pipeline (#903)

* Restructure gpu compilation execution pipeline

* Add compilation server/client infrastructure

* Add wgpu executor

node-graph/gpu-compiler/Cargo.toml

@@ -0,0 +1,31 @@
+[package]
+name = "gpu-compiler"
+version = "0.1.0"
+edition = "2021"
+license = "MIT OR Apache-2.0"
+
+[features]
+default = []
+profiling = ["nvtx"]
+serde = ["graphene-core/serde", "glam/serde"]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+graphene-core = { path = "../gcore", features = ["async", "std", "alloc"] }
+graph-craft = {path = "../graph-craft", features = ["serde"] }
+dyn-any = { path = "../../libraries/dyn-any", features = ["log-bad-types", "rc", "glam"] }
+num-traits = "0.2"
+rand_chacha = "0.3.1"
+log = "0.4"
+serde = { version = "1", features = ["derive", "rc"]}
+glam = { version = "0.22" }
+base64 = "0.13"
+
+bytemuck = { version = "1.8" }
+nvtx = { version = "1.1.1", optional = true }
+tempfile = "3"
+spirv-builder = { version = "0.4",  default-features = false, features=["use-installed-tools"] }
+tera = { version = "1.17.1" }
+anyhow = "1.0.66"
+serde_json = "1.0.91"

node-graph/gpu-compiler/gpu-compiler-bin-wrapper/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "gpu-compiler-bin-wrapper"
+version = "0.1.0"
+edition = "2021"
+license = "MIT OR Apache-2.0"
+
+[features]
+default = []
+profiling = []
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+graph-craft = {path = "../../graph-craft" , features = ["serde"]}
+log = "0.4"
+anyhow = "1.0.66"
+serde_json = "1.0.91"
+serde = { version = "1.0", features = ["derive"] }

node-graph/gpu-compiler/gpu-compiler-bin-wrapper/src/lib.rs

@@ -0,0 +1,53 @@
+use serde::{Deserialize, Serialize};
+use std::io::Write;
+
+pub fn compile_spirv(network: &graph_craft::document::NodeNetwork, input_type: &str, output_type: &str, compile_dir: Option<&str>, manifest_path: &str) -> anyhow::Result<Vec<u8>> {
+	let serialized_graph = serde_json::to_string(&network)?;
+	let features = "";
+	#[cfg(feature = "profiling")]
+	let features = "profiling";
+
+	println!("calling cargo run!");
+	let non_cargo_env_vars = std::env::vars().filter(|(k, _)| k.starts_with("PATH")).collect::<Vec<_>>();
+	let mut cargo_command = std::process::Command::new("/usr/bin/cargo")
+		.arg("run")
+		.arg("--release")
+		.arg("--manifest-path")
+		.arg(manifest_path)
+        .current_dir(manifest_path.replace("Cargo.toml", ""))
+        .env_clear()
+        .envs(non_cargo_env_vars)
+		.arg("--features")
+		.arg(features)
+        .arg("--")
+        .arg(input_type)
+        .arg(output_type)
+        // TODO: handle None case properly
+        .arg(compile_dir.unwrap())
+		.stdin(std::process::Stdio::piped())
+		.stdout(std::process::Stdio::piped())
+        .spawn()?;
+
+	cargo_command.stdin.as_mut().unwrap().write_all(serialized_graph.as_bytes())?;
+	let output = cargo_command.wait_with_output()?;
+	if !output.status.success() {
+		return Err(anyhow::anyhow!("cargo failed: {}", String::from_utf8_lossy(&output.stderr)));
+	}
+	Ok(output.stdout)
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
+pub struct CompileRequest {
+	network: graph_craft::document::NodeNetwork,
+	input_type: String,
+	output_type: String,
+}
+
+impl CompileRequest {
+	pub fn new(network: graph_craft::document::NodeNetwork, input_type: String, output_type: String) -> Self {
+		Self { network, input_type, output_type }
+	}
+	pub fn compile(&self, compile_dir: &str, manifest_path: &str) -> anyhow::Result<Vec<u8>> {
+		compile_spirv(&self.network, &self.input_type, &self.output_type, Some(compile_dir), manifest_path)
+	}
+}

node-graph/gpu-compiler/rust-toolchain

@@ -0,0 +1,3 @@
+[toolchain]
+channel = "nightly-2022-10-29"
+components = ["rust-src", "rustc-dev", "llvm-tools-preview", "clippy", "cargofmt", "rustc"]

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

@@ -0,0 +1,136 @@
+use std::path::Path;
+
+use graph_craft::proto::*;
+use tera::Context;
+
+fn create_cargo_toml(metadata: &Metadata) -> Result<String, tera::Error> {
+	let mut tera = tera::Tera::default();
+	tera.add_raw_template("cargo_toml", include_str!("templates/Cargo-template.toml"))?;
+	let mut context = Context::new();
+	context.insert("name", &metadata.name);
+	context.insert("authors", &metadata.authors);
+	context.insert("gcore_path", &format!("{}{}", env!("CARGO_MANIFEST_DIR"), "/../gcore"));
+	tera.render("cargo_toml", &context)
+}
+
+pub struct Metadata {
+	name: String,
+	authors: Vec<String>,
+}
+
+impl Metadata {
+	pub fn new(name: String, authors: Vec<String>) -> Self {
+		Self { name, authors }
+	}
+}
+
+pub fn create_files(matadata: &Metadata, network: &ProtoNetwork, compile_dir: &Path, input_type: &str, output_type: &str) -> anyhow::Result<()> {
+	let src = compile_dir.join("src");
+	let cargo_file = compile_dir.join("Cargo.toml");
+	let cargo_toml = create_cargo_toml(matadata)?;
+	std::fs::write(cargo_file, cargo_toml)?;
+
+	let toolchain_file = compile_dir.join("rust-toolchain");
+	let toolchain = include_str!("templates/rust-toolchain");
+	std::fs::write(toolchain_file, toolchain)?;
+
+	// create src dir
+	match std::fs::create_dir(&src) {
+		Ok(_) => {}
+		Err(e) => {
+			if e.kind() != std::io::ErrorKind::AlreadyExists {
+				return Err(e.into());
+			}
+		}
+	}
+	let lib = src.join("lib.rs");
+	let shader = serialize_gpu(network, input_type, output_type)?;
+	println!("{}", 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 nid(id: &u64) -> String {
+		format!("n{id}")
+	}
+
+	let mut nodes = Vec::new();
+	#[derive(serde::Serialize)]
+	struct Node {
+		id: String,
+		fqn: String,
+		args: Vec<String>,
+	}
+	for (ref id, node) in network.nodes.iter() {
+		let fqn = &node.identifier.name;
+		let id = nid(id);
+
+		nodes.push(Node {
+			id,
+			fqn: fqn.to_string(),
+			args: node.construction_args.new_function_args(),
+		});
+	}
+
+	let template = include_str!("templates/spirv-template.rs");
+	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("nodes", &nodes);
+	context.insert("last_node", &nid(&network.output));
+	context.insert("compute_threads", &64);
+	Ok(tera.render("spirv", &context)?)
+}
+
+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")
+        .print_metadata(MetadataPrintout::DependencyOnly)
+        .multimodule(false)
+        .preserve_bindings(true)
+        .release(true)
+        //.relax_struct_store(true)
+        //.relax_block_layout(true)
+        .spirv_metadata(SpirvMetadata::Full)
+        .build()?;
+
+	Ok(result)
+}
+
+#[cfg(test)]
+mod test {
+
+	#[test]
+	fn test_create_cargo_toml() {
+		let cargo_toml = super::create_cargo_toml(&super::Metadata {
+			name: "project".to_owned(),
+			authors: vec!["Example <john.smith@example.com>".to_owned(), "smith.john@example.com".to_owned()],
+		});
+		let cargo_toml = cargo_toml.expect("failed to build carog toml template");
+		let lines = cargo_toml.split('\n').collect::<Vec<_>>();
+		let cargo_toml = lines[..lines.len() - 2].join("\n");
+		let reference = r#"[package]
+name = "project-node"
+version = "0.1.0"
+authors = ["Example <john.smith@example.com>", "smith.john@example.com", ]
+edition = "2021"
+license = "MIT OR Apache-2.0"
+publish = false
+
+[lib]
+crate-type = ["dylib", "lib"]
+
+[patch.crates-io]
+libm = { git = "https://github.com/rust-lang/libm", tag = "0.2.5" }
+
+[dependencies]
+spirv-std = { git = "https://github.com/EmbarkStudios/rust-gpu" , features= ["glam"]}"#;
+
+		assert_eq!(cargo_toml, reference);
+	}
+}

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

@@ -0,0 +1,26 @@
+use graph_craft::document::NodeNetwork;
+use gpu_compiler as compiler;
+use std::io::Write;
+
+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(network, true);
+    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)?;
+	let result = compiler::compile(&compile_dir)?;
+
+	let bytes = std::fs::read(result.module.unwrap_single())?;
+	stdout.write_all(&bytes)?;
+
+	Ok(())
+}

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

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

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

@@ -0,0 +1,3 @@
+[toolchain]
+channel = "nightly-2022-10-29"
+components = ["rust-src", "rustc-dev", "llvm-tools-preview", "clippy", "cargofmt", "rustc"]

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

@@ -0,0 +1,38 @@
+#![no_std]
+#![feature(unchecked_math)]
+#![deny(warnings)]
+
+#[cfg(target_arch = "spirv")]
+extern crate spirv_std;
+
+#[cfg(target_arch = "spirv")]
+pub mod gpu {
+    use super::*;
+    use spirv_std::spirv;
+    use spirv_std::glam::UVec3;
+
+    #[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,
+    ) {
+        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 node in nodes %}
+        let {{node.id}} = {{node.fqn}}::new({% for arg in node.args %}{{arg}}, {% endfor %});
+        {% endfor %}
+		{{last_node}}.eval(input)
+    }
+
+}