Implement experimental WebGPU support (#1238)

* Web gpu execution MVP

Ready infrastructure for wgpu experimentation

Start implementing simple gpu test case

Fix Extract Node not working with nested networks

Convert inputs for extracted node to network inputs

Fix missing cors headers

Feature gate gcore to make it once again no-std compatible

Add skeleton structure gpu shader

Work on gpu node graph output saving

Fix Get and Set nodes

Fix storage nodes

Fix shader construction errors -> spirv errors

Add unsafe version

Add once cell node

Web gpu execution MVP
This commit is contained in:
Dennis Kobert 2023-05-27 19:27:46 +02:00 committed by Keavon Chambers
parent 4bd9fbd073
commit 0586d52f3a
33 changed files with 1080 additions and 239 deletions

View file

@ -603,6 +603,7 @@ dependencies = [
"num-traits",
"once_cell",
"rand_chacha",
"rustybuzz",
"serde",
"specta",
"spin",
@ -1153,6 +1154,22 @@ dependencies = [
"serde",
]
[[package]]
name = "rustybuzz"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ab9e34ecf6900625412355a61bda0bd68099fe674de707c67e5e4aed2c05e489"
dependencies = [
"bitflags",
"bytemuck",
"smallvec",
"ttf-parser",
"unicode-bidi-mirroring",
"unicode-ccc",
"unicode-general-category",
"unicode-script",
]
[[package]]
name = "ryu"
version = "1.0.12"
@ -1486,6 +1503,12 @@ dependencies = [
"once_cell",
]
[[package]]
name = "ttf-parser"
version = "0.17.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "375812fa44dab6df41c195cd2f7fecb488f6c09fbaafb62807488cefab642bff"
[[package]]
name = "typenum"
version = "1.16.0"
@ -1557,12 +1580,36 @@ dependencies = [
"unic-common",
]
[[package]]
name = "unicode-bidi-mirroring"
version = "0.1.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "56d12260fb92d52f9008be7e4bca09f584780eb2266dc8fecc6a192bec561694"
[[package]]
name = "unicode-ccc"
version = "0.1.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "cc2520efa644f8268dce4dcd3050eaa7fc044fca03961e9998ac7e2e92b77cf1"
[[package]]
name = "unicode-general-category"
version = "0.6.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "2281c8c1d221438e373249e065ca4989c4c36952c211ff21a0ee91c44a3869e7"
[[package]]
name = "unicode-ident"
version = "1.0.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "84a22b9f218b40614adcb3f4ff08b703773ad44fa9423e4e0d346d5db86e4ebc"
[[package]]
name = "unicode-script"
version = "0.5.5"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "7d817255e1bed6dfd4ca47258685d14d2bdcfbc64fdc9e3819bd5848057b8ecc"
[[package]]
name = "unicode-width"
version = "0.1.10"

View file

@ -6,7 +6,7 @@ use std::io::Write;
pub fn compile_spirv(request: &CompileRequest, compile_dir: Option<&str>, manifest_path: &str) -> anyhow::Result<Vec<u8>> {
let serialized_graph = serde_json::to_string(&gpu_executor::CompileRequest {
network: request.network.clone(),
networks: request.networks.clone(),
io: request.shader_io.clone(),
})?;
@ -43,23 +43,23 @@ pub fn compile_spirv(request: &CompileRequest, compile_dir: Option<&str>, manife
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)]
pub struct CompileRequest {
network: graph_craft::proto::ProtoNetwork,
networks: Vec<graph_craft::proto::ProtoNetwork>,
input_types: Vec<Type>,
output_type: Type,
output_types: Vec<Type>,
shader_io: ShaderIO,
}
impl CompileRequest {
pub fn new(network: ProtoNetwork, input_types: Vec<Type>, output_type: Type, io: ShaderIO) -> Self {
pub fn new(networks: Vec<ProtoNetwork>, input_types: Vec<Type>, output_types: Vec<Type>, io: ShaderIO) -> Self {
// TODO: add type checking
// for (input, buffer) in input_types.iter().zip(io.inputs.iter()) {
// assert_eq!(input, &buffer.ty());
// }
// assert_eq!(output_type, io.output.ty());
Self {
network,
networks,
input_types,
output_type,
output_types,
shader_io: io,
}
}

View file

@ -26,7 +26,7 @@ impl Metadata {
}
}
pub fn create_files(metadata: &Metadata, network: &ProtoNetwork, compile_dir: &Path, io: &ShaderIO) -> anyhow::Result<()> {
pub fn create_files(metadata: &Metadata, networks: &[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(metadata)?;
@ -46,7 +46,7 @@ pub fn create_files(metadata: &Metadata, network: &ProtoNetwork, compile_dir: &P
}
}
let lib = src.join("lib.rs");
let shader = serialize_gpu(network, io)?;
let shader = serialize_gpu(networks, io)?;
eprintln!("{}", shader);
std::fs::write(lib, shader)?;
Ok(())
@ -67,20 +67,21 @@ fn constant_attribute(constant: &GPUConstant) -> &'static str {
}
}
pub fn construct_argument(input: &ShaderInput<()>, position: u32) -> String {
match input {
ShaderInput::Constant(constant) => format!("#[spirv({})] i{}: {},", constant_attribute(constant), position, constant.ty()),
pub fn construct_argument(input: &ShaderInput<()>, position: u32, binding_offset: u32) -> String {
let line = 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,)
format!("#[spirv(uniform, descriptor_set = 0, binding = {})] i{}: &[{}]", position + binding_offset, position, ty,)
}
ShaderInput::StorageBuffer(_, ty) | ShaderInput::ReadBackBuffer(_, ty) => {
format!("#[spirv(storage_buffer, descriptor_set = 0, binding = {})] i{}: &[{}]", position, position, ty,)
format!("#[spirv(storage_buffer, descriptor_set = 0, binding = {})] i{}: &[{}]", position + binding_offset, position, ty,)
}
ShaderInput::OutputBuffer(_, ty) => {
format!("#[spirv(storage_buffer, descriptor_set = 0, binding = {})] i{}: &mut[{}]", position, position, ty,)
format!("#[spirv(storage_buffer, descriptor_set = 0, binding = {})] o{}: &mut[{}]", position + binding_offset, position, ty,)
}
ShaderInput::WorkGroupMemory(_, ty) => format!("#[spirv(workgroup_memory] i{}: {}", position, ty,),
}
};
line.replace("glam::u32::uvec3::UVec3", "spirv_std::glam::UVec3")
}
struct GpuCompiler {
@ -88,10 +89,10 @@ struct GpuCompiler {
}
impl SpirVCompiler for GpuCompiler {
fn compile(&self, network: ProtoNetwork, io: &ShaderIO) -> anyhow::Result<gpu_executor::Shader> {
fn compile(&self, networks: &[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)?;
create_files(&metadata, networks, &self.compile_dir, io)?;
let result = compile(&self.compile_dir)?;
let bytes = std::fs::read(result.module.unwrap_single())?;
@ -105,50 +106,80 @@ impl SpirVCompiler for GpuCompiler {
}
}
pub fn serialize_gpu(network: &ProtoNetwork, io: &ShaderIO) -> anyhow::Result<String> {
pub fn serialize_gpu(networks: &[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 inputs = io
.inputs
.iter()
.filter(|x| !x.is_output())
.enumerate()
.map(|(i, input)| construct_argument(input, i as u32, 0))
.collect::<Vec<_>>();
let offset = inputs.len() as u32;
inputs.extend(io.inputs.iter().filter(|x| x.is_output()).enumerate().map(|(i, input)| construct_argument(input, i as u32, offset)));
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 {
let mut output_nodes = Vec::new();
for network in networks {
dbg!(&network);
//assert_eq!(network.inputs.len(), io.inputs.iter().filter(|x| !x.is_output()).count());
#[derive(serde::Serialize, Debug)]
struct Node {
id: String,
index: usize,
fqn: String,
args: Vec<String>,
}
for (i, id) in network.inputs.iter().enumerate() {
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);
let node = Node {
id: id.clone(),
index: i,
fqn: fqn.to_string().split('<').next().unwrap().to_owned(),
args: node.construction_args.new_function_args(),
};
dbg!(&node);
if !io.inputs[i].is_output() {
if input_nodes.iter().any(|x: &Node| x.id == id) {
continue;
}
input_nodes.push(node);
}
}
let fqn = &node.identifier.name;
let id = nid(id);
for (ref id, node) in network.nodes.iter() {
if network.inputs.contains(id) {
continue;
}
nodes.push(Node {
id,
fqn: fqn.to_string().split("<").next().unwrap().to_owned(),
args: node.construction_args.new_function_args(),
});
let fqn = &node.identifier.name;
let id = nid(id);
if nodes.iter().any(|x: &Node| x.id == id) {
continue;
}
nodes.push(Node {
id,
index: 0,
fqn: fqn.to_string().split("<").next().unwrap().to_owned(),
args: node.construction_args.new_function_args(),
});
}
let output = nid(&network.output);
output_nodes.push(output);
}
dbg!(&input_nodes);
let template = include_str!("templates/spirv-template.rs");
let mut tera = tera::Tera::default();
@ -156,8 +187,8 @@ pub fn serialize_gpu(network: &ProtoNetwork, io: &ShaderIO) -> anyhow::Result<St
let mut context = Context::new();
context.insert("inputs", &inputs);
context.insert("input_nodes", &input_nodes);
context.insert("output_nodes", &output_nodes);
context.insert("nodes", &nodes);
context.insert("last_node", &nid(&network.output));
context.insert("compute_threads", &64);
Ok(tera.render("spirv", &context)?)
}
@ -171,9 +202,9 @@ pub fn compile(dir: &Path) -> Result<spirv_builder::CompileResult, spirv_builder
.preserve_bindings(true)
.release(true)
.spirv_metadata(SpirvMetadata::Full)
.extra_arg("no-early-report-zombies")
.extra_arg("no-infer-storage-classes")
.extra_arg("spirt-passes=qptr")
//.extra_arg("no-early-report-zombies")
//.extra_arg("no-infer-storage-classes")
//.extra_arg("spirt-passes=qptr")
.build()?;
Ok(result)

View file

@ -13,7 +13,7 @@ fn main() -> anyhow::Result<()> {
let metadata = compiler::Metadata::new("project".to_owned(), vec!["test@example.com".to_owned()]);
compiler::create_files(&metadata, &request.network, &compile_dir, &request.io)?;
compiler::create_files(&metadata, &request.networks, &compile_dir, &request.io)?;
let result = compiler::compile(&compile_dir)?;
let bytes = std::fs::read(result.module.unwrap_single())?;

View file

@ -4,32 +4,38 @@
#[cfg(target_arch = "spirv")]
extern crate spirv_std;
#[cfg(target_arch = "spirv")]
pub mod gpu {
use super::*;
//#[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_index: UVec3,
{% for input in inputs %}
{{input}}
{{input}},
{% endfor %}
) {
use graphene_core::Node;
/*
{% for input in input_nodes %}
let i{{loop.index0}} = graphene_core::value::CopiedNode::new(i{{loop.index0}});
let i{{input.index}} = graphene_core::value::CopiedNode::new(i{{input.index}});
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}});
let {{input.id}} = graphene_core::structural::ComposeNode::new(i{{input.index}}, _{{input.id}});
{% endfor %}
*/
{% for node in nodes %}
let {{node.id}} = {{node.fqn}}::new({% for arg in node.args %}{{arg}}, {% endfor %});
{% endfor %}
let output = {{last_node}}.eval(());
// TODO: Write output to buffer
{% for output in output_nodes %}
let v = {{output}}.eval(());
o{{loop.index0}}[_global_index.x as usize] = v;
{% endfor %}
// TODO: Write output to buffer
}
}
//}