Graphene: Fine-grained context caching (#2500)

* RFC: Fine Grained Context Caching

* Fix typos

* Fix label

* Add description of inject traits

* Explicitly support context modification

* Start implementation of context invalidation

* Add inject trait variants
* Route Extract / Inject traits to the proto nodes

* Implement context dependency analysis

* Implement context modification node insertion

* Fix erronous force graph run message

* Fix Extract* Inject* annotations in the nodes

* Require Hash implementation for VarArgs

* Fix nullification node insertion

* Cross of done items unresolved questions section

* Update Cargo.lock

* Fix context features propagation

* Update demo artwork

* Remove BondlessFootprint and FreezeRealTime nodes

* Fix migration

* Add migrations for adding context features to old networks

* Always update real time regardless of animation state

* Cargo fmt

* Fix tests

* Readd sed command to hopefully fix profile result parsing

* Add debug output to profiling pr

* Use new totals instead of summaries for for iai results

* Even more debugging

* Use correct debug metrics (hopefully)

* Add more MemoNode implementations

* Add context features annotation to shader node macro

* Cleanup

* Time -> RealTime

* Code review

---------

Co-authored-by: Keavon Chambers <keavon@keavon.com>

node-graph/graph-craft/src/document.rs

@@ -7,7 +7,7 @@ use glam::IVec2;
 use graphene_core::memo::MemoHashGuard;
 pub use graphene_core::uuid::NodeId;
 pub use graphene_core::uuid::generate_uuid;
-use graphene_core::{Context, Cow, MemoHash, ProtoNodeIdentifier, Type};
+use graphene_core::{Context, ContextDependencies, Cow, MemoHash, ProtoNodeIdentifier, Type};
 use log::Metadata;
 use rustc_hash::FxHashMap;
 use std::collections::HashMap;
@@ -60,6 +60,9 @@ pub struct DocumentNode {
 	/// The path to this node and its inputs and outputs as of when [`NodeNetwork::generate_node_paths`] was called.
 	#[serde(skip)]
 	pub original_location: OriginalLocation,
+	/// List of Extract and Inject annotations for the Context.
+	#[serde(default)]
+	pub context_features: ContextDependencies,
 }
 
 /// Represents the original location of a node input/output when [`NodeNetwork::generate_node_paths`] was called, allowing the types and errors to be derived.
@@ -92,6 +95,7 @@ impl Default for DocumentNode {
 			visible: true,
 			skip_deduplication: Default::default(),
 			original_location: OriginalLocation::default(),
+			context_features: Default::default(),
 		}
 	}
 }
@@ -159,6 +163,7 @@ impl DocumentNode {
 			construction_args: args,
 			original_location: self.original_location,
 			skip_deduplication: self.skip_deduplication,
+			context_features: self.context_features,
 		}
 	}
 }

node-graph/graph-craft/src/document/value.rs

@@ -16,7 +16,7 @@ use graphene_core::uuid::NodeId;
 use graphene_core::vector::Vector;
 use graphene_core::vector::style::Fill;
 use graphene_core::vector::style::GradientStops;
-use graphene_core::{Artboard, Color, Graphic, MemoHash, Node, Type};
+use graphene_core::{Artboard, Color, ContextFeatures, Graphic, MemoHash, Node, Type};
 use graphene_svg_renderer::RenderMetadata;
 use std::fmt::Display;
 use std::hash::Hash;
@@ -217,6 +217,7 @@ tagged_value! {
 	BrushStrokes(Vec<BrushStroke>),
 	BrushCache(BrushCache),
 	DocumentNode(DocumentNode),
+	ContextFeatures(ContextFeatures),
 	Curve(graphene_raster_nodes::curve::Curve),
 	Footprint(graphene_core::transform::Footprint),
 	VectorModification(Box<graphene_core::vector::VectorModification>),

node-graph/graph-craft/src/graphene_compiler.rs

@@ -17,7 +17,7 @@ impl Compiler {
 		let proto_networks = network.into_proto_networks();
 
 		proto_networks.map(move |mut proto_network| {
-			proto_network.resolve_inputs()?;
+			proto_network.insert_context_nullification_nodes()?;
 			proto_network.generate_stable_node_ids();
 			Ok(proto_network)
 		})

node-graph/graph-craft/src/proto.rs

@@ -1,3 +1,4 @@
+use crate::document::value::TaggedValue;
 use crate::document::{InlineRust, value};
 use crate::document::{NodeId, OriginalLocation};
 pub use graphene_core::registry::*;
@@ -132,6 +133,7 @@ pub struct ProtoNode {
 	pub identifier: ProtoNodeIdentifier,
 	pub original_location: OriginalLocation,
 	pub skip_deduplication: bool,
+	pub(crate) context_features: ContextDependencies,
 }
 
 impl Default for ProtoNode {
@@ -142,6 +144,7 @@ impl Default for ProtoNode {
 			call_argument: concrete!(()),
 			original_location: OriginalLocation::default(),
 			skip_deduplication: false,
+			context_features: Default::default(),
 		}
 	}
 }
@@ -181,6 +184,7 @@ impl ProtoNode {
 				..Default::default()
 			},
 			skip_deduplication: false,
+			context_features: Default::default(),
 		}
 	}
 
@@ -290,15 +294,137 @@ impl ProtoNetwork {
 		(inwards_edges, id_map)
 	}
 
-	/// Performs topological sort and reorders ids.
-	pub fn resolve_inputs(&mut self) -> Result<(), String> {
+	/// Inserts context nullification nodes to optimize caching.
+	/// This analysis is performed after topological sorting to ensure proper dependency tracking.
+	pub fn insert_context_nullification_nodes(&mut self) -> Result<(), String> {
 		// Perform topological sort once
 		self.reorder_ids()?;
 
+		self.find_context_dependencies(self.output);
+
+		// Perform topological sort a second time to integrate the new nodes
+		self.reorder_ids()?;
+
 		Ok(())
 	}
 
-	/// Update all of the references to a node ID in the graph with a new ID named `replacement_node_id`.
+	fn insert_context_nullification_node(&mut self, node_id: NodeId, context_deps: ContextFeatures) -> NodeId {
+		let (_, node) = &self.nodes[node_id.0 as usize];
+		let mut path = node.original_location.path.clone();
+
+		// Add a path extension with a placeholder value which should not conflict with existing paths
+		if let Some(p) = path.as_mut() {
+			p.push(NodeId(10))
+		}
+
+		let memo_node_id = NodeId(self.nodes.len() as u64);
+
+		self.nodes.push((
+			memo_node_id,
+			ProtoNode {
+				construction_args: ConstructionArgs::Nodes(vec![node_id]),
+				call_argument: concrete!(Context),
+				identifier: graphene_core::memo::memo::IDENTIFIER,
+				original_location: OriginalLocation {
+					path: path.clone(),
+					..Default::default()
+				},
+				..Default::default()
+			},
+		));
+
+		let nullification_value_node_id = NodeId(self.nodes.len() as u64);
+
+		self.nodes.push((
+			nullification_value_node_id,
+			ProtoNode {
+				construction_args: ConstructionArgs::Value(MemoHash::new(TaggedValue::ContextFeatures(context_deps))),
+				call_argument: concrete!(Context),
+				identifier: ProtoNodeIdentifier::new("graphene_core::value::ClonedNode"),
+				original_location: OriginalLocation {
+					path: path.clone(),
+					..Default::default()
+				},
+				..Default::default()
+			},
+		));
+		let nullification_node_id = NodeId(self.nodes.len() as u64);
+		self.nodes.push((
+			nullification_node_id,
+			ProtoNode {
+				construction_args: ConstructionArgs::Nodes(vec![memo_node_id, nullification_value_node_id]),
+				call_argument: concrete!(Context),
+				identifier: graphene_core::context_modification::context_modification::IDENTIFIER,
+				original_location: OriginalLocation {
+					path: path.clone(),
+					..Default::default()
+				},
+				..Default::default()
+			},
+		));
+		nullification_node_id
+	}
+
+	fn find_context_dependencies(&mut self, id: NodeId) -> (ContextFeatures, Option<NodeId>) {
+		let mut branch_dependencies = Vec::new();
+		let mut combined_deps = ContextFeatures::default();
+		let node_index = id.0 as usize;
+
+		let context_features = self.nodes[node_index].1.context_features;
+
+		let mut inputs = match &self.nodes[node_index].1.construction_args {
+			// We pretend like we have already placed context modification nodes after ourselves because value nodes don't need to be cached
+			ConstructionArgs::Value(_) => return (context_features.extract, Some(id)),
+			ConstructionArgs::Nodes(items) => items.clone(),
+			ConstructionArgs::Inline(_) => return (context_features.extract, Some(id)),
+		};
+
+		// Compute the dependencies for each branch and combine all of them
+		for &node in &inputs {
+			let branch = self.find_context_dependencies(node);
+
+			branch_dependencies.push(branch);
+			combined_deps |= branch.0;
+		}
+		let mut new_deps = combined_deps;
+
+		// Remove requirements which this node provides
+		new_deps &= !context_features.inject;
+		// Add requirements we have
+		new_deps |= context_features.extract;
+
+		// If we either introduce new dependencies, we can cache all children which don't yet need that dependency
+		let we_introduce_new_deps = !combined_deps.contains(new_deps);
+
+		// For diverging branches, we can add a cache node for all branches which don't reqire all dependencies
+		for (child_node, (deps, new_id)) in inputs.iter_mut().zip(branch_dependencies.into_iter()) {
+			if let Some(new_id) = new_id {
+				*child_node = new_id;
+			} else if we_introduce_new_deps || deps != combined_deps {
+				*child_node = self.insert_context_nullification_node(*child_node, deps);
+			}
+		}
+		self.nodes[node_index].1.construction_args = ConstructionArgs::Nodes(inputs);
+
+		// Which dependencies do we supply (and don't need ourselves)?
+		let net_injections = context_features.inject.difference(context_features.extract);
+
+		// Which dependencies still need to be met after this node?
+		let remaining_deps_from_children = combined_deps.difference(net_injections);
+
+		// Do we satisfy any existing dependencies?
+		let we_supply_existing_deps = !combined_deps.difference(remaining_deps_from_children).is_empty();
+
+		let mut new_id = None;
+		if we_supply_existing_deps {
+			// Our set of context dependencies has shrunk so we can add a cache node after the current node
+			new_id = Some(self.insert_context_nullification_node(id, new_deps));
+		}
+
+		(new_deps, new_id)
+	}
+
+	/// Update all of the references to a node ID in the graph with a new ID named `compose_node_id`.
 	fn replace_node_id(&mut self, outwards_edges: &HashMap<NodeId, Vec<NodeId>>, node_id: NodeId, replacement_node_id: NodeId) {
 		// Update references in other nodes to use the new node
 		if let Some(referring_nodes) = outwards_edges.get(&node_id) {
@@ -801,7 +927,9 @@ mod test {
 	#[test]
 	fn stable_node_id_generation() {
 		let mut construction_network = test_network();
-		construction_network.resolve_inputs().expect("Error when calling 'resolve_inputs' on 'construction_network.");
+		construction_network
+			.insert_context_nullification_nodes()
+			.expect("Error when calling 'insert_context_nullification_nodes' on 'construction_network.");
 		construction_network.generate_stable_node_ids();
 		assert_eq!(construction_network.nodes[0].1.identifier.name.as_ref(), "value");
 		let ids: Vec<_> = construction_network.nodes.iter().map(|(id, _)| *id).collect();