Training support

fusor-ml/core/src/compute_graph/resolve/execution.rs

@@ -636,6 +636,10 @@ impl Resolver {
         if !self.execution_graph.contains_node(node_idx) {
             return;
         }
+        let inner_idx = self.execution_graph[node_idx].inner_idx;
+        if self.targets.contains(&inner_idx) {
+            return;
+        }
         if self
             .execution_graph
             .neighbors_directed(node_idx, petgraph::Direction::Outgoing)

fusor-ml/core/src/compute_graph/resolve/run.rs

@@ -534,7 +534,24 @@ impl Resolver {
             let mut dispatch_index = 0usize;
             let mut command_index = 0usize;
             let dispatches_per_pass = dispatches_per_pass(total_kernels);
+            let dispatches_per_submit = dispatches_per_submit(total_kernels, device.backend());
+            let wait_after_chunk_submit = device.backend() == wgpu::Backend::Metal;
+            let mut dispatches_in_submit = 0usize;
+            let mut encoder_has_commands = false;
             while command_index < commands.len() {
+                if encoder_has_commands && dispatches_in_submit >= dispatches_per_submit {
+                    let next_encoder = resolver_command_encoder(&device);
+                    let ready_encoder = std::mem::replace(&mut command_encoder, next_encoder);
+                    submit_resolver_encoder(
+                        &device,
+                        ready_encoder,
+                        wait_after_chunk_submit,
+                        host_trace,
+                        &mut host_profile,
+                    );
+                    encoder_has_commands = false;
+                    dispatches_in_submit = 0;
+                }
                 match &commands[command_index] {
                     CommandRecord::CopyBuffer(copy) => {
                         command_encoder.copy_buffer_to_buffer(
@@ -544,6 +561,7 @@ impl Resolver {
                             copy.destination_offset,
                             copy.size,
                         );
+                        encoder_has_commands = true;
                         command_index += 1;
                     }
                     CommandRecord::Dispatch(_) => {
@@ -568,6 +586,8 @@ impl Resolver {
                                 record.dispatch.run(&mut pass);
                             }
                             dispatch_index += 1;
+                            dispatches_in_submit += 1;
+                            encoder_has_commands = true;
                             command_index += 1;
                             continue;
                         }
@@ -582,6 +602,9 @@ impl Resolver {
                             if pass_dispatches >= dispatches_per_pass {
                                 break;
                             }
+                            if dispatches_in_submit >= dispatches_per_submit {
+                                break;
+                            }
                             let CommandRecord::Dispatch(record) = &commands[command_index] else {
                                 break;
                             };
@@ -595,8 +618,10 @@ impl Resolver {
                                 pass.write_timestamp(query_set, (dispatch_index * 2 + 1) as u32);
                             }
                             dispatch_index += 1;
+                            dispatches_in_submit += 1;
                             command_index += 1;
                             pass_dispatches += 1;
+                            encoder_has_commands = true;
                         }
                     }
                 }
@@ -625,11 +650,13 @@ impl Resolver {
         let tail_result = tail(&data, &mut command_encoder);
 
         // Submit any remaining commands.
-        let submit_start = host_trace.then(Instant::now);
-        device.wgpu_queue().submit(Some(command_encoder.finish()));
-        if let Some(start) = submit_start {
-            host_profile.submit += start.elapsed();
-        }
+        submit_resolver_encoder(
+            &device,
+            command_encoder,
+            false,
+            host_trace,
+            &mut host_profile,
+        );
         #[cfg(not(target_arch = "wasm32"))]
         {
             if let Some((_, _, readback_buffer, raw_query_size)) = &query_resources {
@@ -726,3 +753,45 @@ fn dispatches_per_pass(total_kernels: usize) -> usize {
 
     if total_kernels >= 1024 { 1 } else { usize::MAX }
 }
+
+fn dispatches_per_submit(total_kernels: usize, backend: wgpu::Backend) -> usize {
+    if let Ok(value) = std::env::var("FUSOR_RESOLVE_DISPATCHES_PER_SUBMIT")
+        && let Ok(parsed) = value.parse::<usize>()
+        && parsed > 0
+    {
+        return parsed;
+    }
+
+    // Chunked submits exist to bound in-flight memory on giant training
+    // graphs; small inference graphs must stay a single submit.
+    if backend == wgpu::Backend::Metal && total_kernels >= 1024 {
+        256
+    } else {
+        usize::MAX
+    }
+}
+
+fn resolver_command_encoder(device: &crate::Device) -> wgpu::CommandEncoder {
+    device
+        .wgpu_device()
+        .create_command_encoder(&wgpu::CommandEncoderDescriptor {
+            label: Some("Resolver Encoder"),
+        })
+}
+
+fn submit_resolver_encoder(
+    device: &crate::Device,
+    command_encoder: wgpu::CommandEncoder,
+    wait: bool,
+    host_trace: bool,
+    host_profile: &mut ResolveHostProfile,
+) {
+    let submit_start = host_trace.then(Instant::now);
+    device.wgpu_queue().submit(Some(command_encoder.finish()));
+    if wait {
+        device.poll_wait();
+    }
+    if let Some(start) = submit_start {
+        host_profile.submit += start.elapsed();
+    }
+}

fusor-ml/core/src/nary_direct.rs

@@ -849,6 +849,31 @@ fn emit_function(function: &NaryFunction, values: &mut [(ValueTile, DataTypeEnum
             values[1].0.clone(),
             function.output_type,
         ),
+        NaryOp::Less => values[0].0.clone().compare(
+            tile_ir::TileCompareOp::Lt,
+            values[1].0.clone(),
+            function.output_type,
+        ),
+        NaryOp::Equal => values[0].0.clone().compare(
+            tile_ir::TileCompareOp::Eq,
+            values[1].0.clone(),
+            function.output_type,
+        ),
+        NaryOp::NotEqual => values[0].0.clone().compare(
+            tile_ir::TileCompareOp::Ne,
+            values[1].0.clone(),
+            function.output_type,
+        ),
+        NaryOp::Greater => values[0].0.clone().compare(
+            tile_ir::TileCompareOp::Gt,
+            values[1].0.clone(),
+            function.output_type,
+        ),
+        NaryOp::GreaterEqual => values[0].0.clone().compare(
+            tile_ir::TileCompareOp::Ge,
+            values[1].0.clone(),
+            function.output_type,
+        ),
         NaryOp::AddConst(scalar) => values[0].0.clone().binary(
             tile_ir::TileBinaryOp::Add,
             tile_literal(scalar).cast_to(values[0].1),

fusor-ml/core/src/nary_wise.rs

@@ -86,6 +86,16 @@ pub(crate) enum NaryOp {
     /// index-dependent masking (e.g. composed causal attention compares the
     /// kv position against the query position).
     LessEqual,
+    /// Binary `a < b` comparison producing 1/0 in the output type.
+    Less,
+    /// Binary `a == b` comparison producing 1/0 in the output type.
+    Equal,
+    /// Binary `a != b` comparison producing 1/0 in the output type.
+    NotEqual,
+    /// Binary `a > b` comparison producing 1/0 in the output type.
+    Greater,
+    /// Binary `a >= b` comparison producing 1/0 in the output type.
+    GreaterEqual,
     AddConst(NaryScalar),
     SubConst(NaryScalar),
     RSubConst(NaryScalar),

fusor-ml/core/src/pair_wise.rs

@@ -1,7 +1,7 @@
 use std::ops::{Add, Div, Mul, Sub};
 
 use crate::{
-    Tensor,
+    DataType, Tensor,
     nary_wise::{NaryFunction, NaryOp},
 };
 
@@ -81,3 +81,51 @@ macro_rules! impl_pairwise_method {
 }
 
 impl_pairwise_method!(pow, NaryOp::Pow, "pow", pow_, |a, b| a.pow(&b));
+
+/// Emit a tensor-tensor comparison method producing 1/0 in the output type
+/// `D`, mirroring the scalar comparisons in `element_wise`.
+macro_rules! impl_pairwise_cmp {
+    ($(#[$meta:meta])* $method:ident, $nary_op:expr, $op_name:literal) => {
+        impl Tensor {
+            $(#[$meta])*
+            pub fn $method<D: DataType>(&self, other: &Self) -> Tensor {
+                assert_eq!(self.datatype(), other.datatype());
+                self.binary_nary(
+                    other,
+                    NaryFunction::binary(
+                        Some($op_name.to_string()),
+                        $nary_op,
+                        self.datatype(),
+                        other.datatype(),
+                        D::DATA_TYPE,
+                    ),
+                )
+            }
+        }
+    };
+}
+
+impl_pairwise_cmp!(
+    /// Element-wise `self == other` returning 1 for true and 0 for false.
+    eq_tensor, NaryOp::Equal, "eq"
+);
+impl_pairwise_cmp!(
+    /// Element-wise `self != other` returning 1 for true and 0 for false.
+    ne_tensor, NaryOp::NotEqual, "ne"
+);
+impl_pairwise_cmp!(
+    /// Element-wise `self < other` returning 1 for true and 0 for false.
+    lt_tensor, NaryOp::Less, "lt"
+);
+impl_pairwise_cmp!(
+    /// Element-wise `self <= other` returning 1 for true and 0 for false.
+    lte_tensor, NaryOp::LessEqual, "lte"
+);
+impl_pairwise_cmp!(
+    /// Element-wise `self > other` returning 1 for true and 0 for false.
+    gt_tensor, NaryOp::Greater, "gt"
+);
+impl_pairwise_cmp!(
+    /// Element-wise `self >= other` returning 1 for true and 0 for false.
+    gte_tensor, NaryOp::GreaterEqual, "gte"
+);