Restructure GPU compilation execution pipeline (#903)

* Restructure gpu compilation execution pipeline

* Add compilation server/client infrastructure

* Add wgpu executor
This commit is contained in:
Dennis Kobert 2022-12-31 02:52:04 +01:00 committed by Keavon Chambers
parent be32f7949f
commit 79ad3e7908
43 changed files with 2744 additions and 482 deletions

View file

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

View file

@ -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(())
}

View file

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

View file

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

View file

@ -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)
}
}