Lay groundwork for directly rendering to the canvas without a cpu roundrip (#1291)

* Add Texture handle type

* Add Texture View to shader inputs

* Implement basic rendering pipeline

* Render first texture using render pipeline

* Fix output color space

* Precompute the rendering pipeline

* Move gpu context creation to editor api

* Port gpu-executor nodes to node registry

* Fix canvas nodes and make code compile for non wasm targets

* Pin wasm-bindgen version

* Disable miri temoporarily for better ci times

* Fix formatting

* Remove unsafe block

* Bump wasm-pack version

* Bump wasm-bindgen version

* Add gpu feature guard for push node

* Make Into node async

node-graph/gstd/Cargo.toml

@@ -9,16 +9,12 @@ license = "MIT OR Apache-2.0"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [features]
-default = ["wgpu"]
-gpu = [
-	"graphene-core/gpu",
-	"gpu-compiler-bin-wrapper",
-	"compilation-client",
-	"gpu-executor",
-]
+default = ["wasm"]
+gpu = ["graphene-core/gpu", "gpu-compiler-bin-wrapper", "compilation-client", "gpu-executor"]
 vulkan = ["gpu", "vulkan-executor"]
 wgpu = ["gpu", "wgpu-executor"]
 quantization = ["autoquant"]
+wasm = ["wasm-bindgen", "web-sys", "js-sys"]
 
 
 [dependencies]
@@ -29,6 +25,7 @@ graphene-core = { path = "../gcore", features = [
 	"async",
 	"std",
 	"serde",
+	"alloc",
 ], default-features = false }
 dyn-any = { path = "../../libraries/dyn-any", features = ["derive"] }
 graph-craft = { path = "../graph-craft" }
@@ -53,8 +50,25 @@ xxhash-rust = { workspace = true }
 serde_json = "1.0.96"
 reqwest = { version = "0.11.17", features = ["rustls", "rustls-tls"] }
 futures = "0.3.28"
+wasm-bindgen = { version = "0.2.84", optional = true }
+js-sys = { version = "0.3.55", optional = true }
+wgpu-types = "0.16.0"
+wgpu = "0.16.1"
 
 [dependencies.serde]
 version = "1.0"
 optional = true
 features = ["derive"]
+
+
+[dependencies.web-sys]
+version = "0.3.4"
+optional = true
+features = [
+	"Window",
+	"CanvasRenderingContext2d",
+	"ImageData",
+	"Document",
+	"HtmlCanvasElement",
+	"ImageBitmapRenderingContext",
+]

node-graph/gstd/src/gpu_nodes.rs

@@ -6,10 +6,12 @@ use graph_craft::document::*;
 use graph_craft::proto::*;
 use graphene_core::raster::*;
 use graphene_core::*;
-use wgpu_executor::NewExecutor;
+use wgpu_executor::WgpuExecutor;
 
 use std::sync::Arc;
 
+use crate::wasm_application_io::WasmApplicationIo;
+
 pub struct GpuCompiler<TypingContext, ShaderIO> {
 	typing_context: TypingContext,
 	io: ShaderIO,
@@ -37,12 +39,13 @@ async fn compile_gpu(node: &'input DocumentNode, mut typing_context: TypingConte
 	compilation_client::compile(proto_networks, input_types, output_types, io).await.unwrap()
 }
 
-pub struct MapGpuNode<Node> {
+pub struct MapGpuNode<Node, EditorApi> {
 	node: Node,
+	editor_api: EditorApi,
 }
 
 #[node_macro::node_fn(MapGpuNode)]
-async fn map_gpu(image: ImageFrame<Color>, node: DocumentNode) -> ImageFrame<Color> {
+async fn map_gpu<'a: 'input>(image: ImageFrame<Color>, node: DocumentNode, editor_api: graphene_core::application_io::EditorApi<'a, WasmApplicationIo>) -> ImageFrame<Color> {
 	log::debug!("Executing gpu node");
 	let compiler = graph_craft::graphene_compiler::Compiler {};
 	let inner_network = NodeNetwork::value_network(node);
@@ -121,7 +124,24 @@ async fn map_gpu(image: ImageFrame<Color>, node: DocumentNode) -> ImageFrame<Col
 	//return ImageFrame::empty();
 	let len: usize = image.image.data.len();
 
-	let executor = NewExecutor::new().await.unwrap();
+	let executor = &editor_api.application_io.gpu_executor.as_ref().unwrap();
+
+	/*
+	let canvas = editor_api.application_io.create_surface();
+
+	let surface = unsafe { executor.create_surface(canvas) }.unwrap();
+	//log::debug!("id: {:?}", surface);
+	let surface_id = surface.surface_id;
+
+	let texture = executor.create_texture_buffer(image.image.clone(), TextureBufferOptions::Texture).unwrap();
+
+	//executor.create_render_pass(texture, surface).unwrap();
+
+	let frame = SurfaceFrame {
+		surface_id,
+		transform: image.transform,
+	};
+	return frame;*/
 	log::debug!("creating buffer");
 	let width_uniform = executor.create_uniform_buffer(image.image.width).unwrap();
 	let storage_buffer = executor
@@ -152,7 +172,6 @@ async fn map_gpu(image: ImageFrame<Color>, node: DocumentNode) -> ImageFrame<Col
 		io: shader.io,
 	};
 	log::debug!("loading shader");
-	log::debug!("shader: {:?}", shader.source);
 	let shader = executor.load_shader(shader).unwrap();
 	log::debug!("loaded shader");
 	let pipeline = PipelineLayout {
@@ -333,7 +352,7 @@ async fn blend_gpu_image(foreground: ImageFrame<Color>, background: ImageFrame<C
 	.unwrap();
 	let len = background.image.data.len();
 
-	let executor = NewExecutor::new()
+	let executor = WgpuExecutor::new()
 		.await
 		.expect("Failed to create wgpu executor. Please make sure that webgpu is enabled for your browser.");
 	log::debug!("creating buffer");

node-graph/gstd/src/lib.rs

@@ -12,7 +12,7 @@ pub mod http;
 pub mod any;
 
 #[cfg(feature = "gpu")]
-pub mod executor;
+pub mod gpu_nodes;
 
 #[cfg(feature = "quantization")]
 pub mod quantization;
@@ -22,3 +22,6 @@ pub use graphene_core::*;
 pub mod image_segmentation;
 
 pub mod brush;
+
+#[cfg(feature = "wasm")]
+pub mod wasm_application_io;

node-graph/gstd/src/wasm_application_io.rs

@@ -0,0 +1,151 @@
+use std::cell::RefCell;
+
+use dyn_any::StaticType;
+use graphene_core::application_io::{ApplicationIo, SurfaceHandle, SurfaceHandleFrame, SurfaceId};
+use graphene_core::{
+	raster::{color::SRGBA8, ImageFrame},
+	Node,
+};
+use js_sys::{Object, Reflect};
+use std::sync::Arc;
+use wasm_bindgen::{Clamped, JsCast, JsValue};
+use web_sys::{window, CanvasRenderingContext2d, HtmlCanvasElement};
+#[cfg(feature = "wgpu")]
+use wgpu_executor::WgpuExecutor;
+
+pub struct Canvas(CanvasRenderingContext2d);
+
+#[derive(Debug, Default)]
+pub struct WasmApplicationIo {
+	ids: RefCell<u64>,
+	#[cfg(feature = "wgpu")]
+	pub(crate) gpu_executor: Option<WgpuExecutor>,
+}
+
+impl WasmApplicationIo {
+	pub async fn new() -> Self {
+		Self {
+			ids: RefCell::new(0),
+			#[cfg(feature = "wgpu")]
+			gpu_executor: WgpuExecutor::new().await,
+		}
+	}
+}
+
+unsafe impl StaticType for WasmApplicationIo {
+	type Static = WasmApplicationIo;
+}
+
+impl<'a> From<WasmEditorApi<'a>> for &'a WasmApplicationIo {
+	fn from(editor_api: WasmEditorApi<'a>) -> Self {
+		editor_api.application_io
+	}
+}
+#[cfg(feature = "wgpu")]
+impl<'a> From<&'a WasmApplicationIo> for &'a WgpuExecutor {
+	fn from(app_io: &'a WasmApplicationIo) -> Self {
+		app_io.gpu_executor.as_ref().unwrap()
+	}
+}
+
+pub type WasmEditorApi<'a> = graphene_core::application_io::EditorApi<'a, WasmApplicationIo>;
+
+impl ApplicationIo for WasmApplicationIo {
+	type Surface = HtmlCanvasElement;
+	#[cfg(feature = "wgpu")]
+	type Executor = WgpuExecutor;
+	#[cfg(not(feature = "wgpu"))]
+	type Executor = ();
+
+	fn create_surface(&self) -> SurfaceHandle<Self::Surface> {
+		let wrapper = || {
+			let document = window().expect("should have a window in this context").document().expect("window should have a document");
+
+			let canvas: HtmlCanvasElement = document.create_element("canvas")?.dyn_into::<HtmlCanvasElement>()?;
+			let mut guard = self.ids.borrow_mut();
+			let id = SurfaceId(*guard);
+			*guard += 1;
+			// store the canvas in the global scope so it doesn't get garbage collected
+			let window = window().expect("should have a window in this context");
+			let window = Object::from(window);
+
+			let image_canvases_key = JsValue::from_str("imageCanvases");
+
+			let mut canvases = Reflect::get(&window, &image_canvases_key);
+			if canvases.is_err() {
+				Reflect::set(&JsValue::from(web_sys::window().unwrap()), &image_canvases_key, &Object::new()).unwrap();
+				canvases = Reflect::get(&window, &image_canvases_key);
+			}
+
+			// Convert key and value to JsValue
+			let js_key = JsValue::from_str(format!("canvas{}", id.0).as_str());
+			let js_value = JsValue::from(canvas.clone());
+
+			let canvases = Object::from(canvases.unwrap());
+
+			// Use Reflect API to set property
+			Reflect::set(&canvases, &js_key, &js_value)?;
+			Ok::<_, JsValue>(SurfaceHandle { surface_id: id, surface: canvas })
+		};
+
+		wrapper().expect("should be able to set canvas in global scope")
+	}
+
+	fn destroy_surface(&self, surface_id: SurfaceId) {
+		let window = window().expect("should have a window in this context");
+		let window = Object::from(window);
+
+		let image_canvases_key = JsValue::from_str("imageCanvases");
+
+		let wrapper = || {
+			if let Ok(canvases) = Reflect::get(&window, &image_canvases_key) {
+				// Convert key and value to JsValue
+				let js_key = JsValue::from_str(format!("canvas{}", surface_id.0).as_str());
+
+				// Use Reflect API to set property
+				Reflect::delete_property(&canvases.into(), &js_key)?;
+			}
+			Ok::<_, JsValue>(())
+		};
+
+		wrapper().expect("should be able to set canvas in global scope")
+	}
+
+	#[cfg(feature = "wgpu")]
+	fn gpu_executor(&self) -> Option<&Self::Executor> {
+		self.gpu_executor.as_ref()
+	}
+}
+
+pub type WasmSurfaceHandle = SurfaceHandle<HtmlCanvasElement>;
+pub type WasmSurfaceHandleFrame = SurfaceHandleFrame<HtmlCanvasElement>;
+
+pub struct CreateSurfaceNode {}
+
+#[node_macro::node_fn(CreateSurfaceNode)]
+async fn create_surface_node<'a: 'input>(editor: WasmEditorApi<'a>) -> Arc<SurfaceHandle<HtmlCanvasElement>> {
+	editor.application_io.create_surface().into()
+}
+
+pub struct DrawImageFrameNode<Surface> {
+	surface_handle: Surface,
+}
+
+#[node_macro::node_fn(DrawImageFrameNode)]
+async fn draw_image_frame_node<'a: 'input>(image: ImageFrame<SRGBA8>, surface_handle: Arc<SurfaceHandle<HtmlCanvasElement>>) -> SurfaceHandleFrame<HtmlCanvasElement> {
+	let image_data = image.image.data;
+	let array: Clamped<&[u8]> = Clamped(bytemuck::cast_slice(image_data.as_slice()));
+	if image.image.width > 0 && image.image.height > 0 {
+		let canvas = &surface_handle.surface;
+		canvas.set_width(image.image.width);
+		canvas.set_height(image.image.height);
+		// TODO: replace "2d" with "bitmaprenderer" once we switch to ImageBitmap (lives on gpu) from ImageData (lives on cpu)
+		let context = canvas.get_context("2d").unwrap().unwrap().dyn_into::<CanvasRenderingContext2d>().unwrap();
+		let image_data = web_sys::ImageData::new_with_u8_clamped_array_and_sh(array, image.image.width as u32, image.image.height as u32).expect("Failed to construct ImageData");
+		context.put_image_data(&image_data, 0.0, 0.0).unwrap();
+	}
+	SurfaceHandleFrame {
+		surface_handle,
+		transform: image.transform,
+	}
+}