Pipeline Calling Convention 2 (#5291)

* Guarantees that output vectors have length == 0
* Does not guarantee capacity.
* Inputs are either N or n elements, and defines that in-place but
non-selected elements _must_ be valid.

Signed-off-by: Nicholas Gates <[email protected]>

Author: gatesn

vortex-array/src/pipeline/driver/allocation.rs

@@ -8,14 +8,14 @@ use vortex_vector::VectorMut;
 
 use crate::Array;
 use crate::pipeline::driver::{Node, NodeId};
-use crate::pipeline::{N, PipelineVector, VectorId};
+use crate::pipeline::{N, VectorId};
 
 #[derive(Debug)]
 pub struct VectorAllocation {
     /// Where each node writes its output
     pub(crate) output_targets: Vec<OutputTarget>,
     /// The actual allocated vectors
-    pub(crate) vectors: Vec<PipelineVector>,
+    pub(crate) vectors: Vec<VectorMut>,
 }
 
 // TODO(joe): support in-place view operations
@@ -86,7 +86,7 @@ pub(super) fn allocate_vectors(
             .collect(),
         vectors: allocation_types
             .into_iter()
-            .map(|dtype| PipelineVector::Compact(VectorMut::with_capacity(dtype, N)))
+            .map(|dtype| VectorMut::with_capacity(dtype, N))
             .collect(),
     })
 }

vortex-array/src/pipeline/driver/mod.rs

@@ -19,7 +19,7 @@ use crate::pipeline::bit_view::{BitView, BitViewExt};
 use crate::pipeline::driver::allocation::{OutputTarget, allocate_vectors};
 use crate::pipeline::driver::bind::bind_kernels;
 use crate::pipeline::driver::toposort::topological_sort;
-use crate::pipeline::{Kernel, KernelCtx, N, PipelineInputs, PipelineVector};
+use crate::pipeline::{Kernel, KernelCtx, N, PipelineInputs};
 use crate::{Array, ArrayEq, ArrayHash, ArrayOperator, ArrayRef, ArrayVisitor, Precision};
 
 /// A pipeline driver takes a Vortex array and executes it into a canonical vector.
@@ -278,56 +278,38 @@ impl Pipeline {
             // take the intermediate vector and write into that.
             match &self.output_targets[node_idx] {
                 OutputTarget::ExternalOutput => {
-                    assert!(
-                        output.capacity() >= N,
-                        "Insufficient capacity in external output vector"
-                    );
-
-                    let prev_output_len = output.len();
-                    kernel.step(&self.ctx, selection, output)?;
-
-                    let added_len = output.len() - prev_output_len;
-                    match added_len {
-                        N => {
-                            // If the kernel added N elements, the output is in-place.
-                            // TODO(ngates): we need to filter if the true count is not N.
-                        }
-                        _ if added_len == selection.true_count() => {
-                            // If the kernel added exactly the number of selected elements,
-                            // the output is already compacted into the start of the vector.
-                        }
-                        _ => vortex_bail!(
-                            "Kernel produced incorrect number of output elements, expected to append either {} or {}, got {}",
+                    // We split off the next N elements of capacity from the external output vector.
+                    let mut tail = output.split_off(output.len());
+                    assert!(tail.is_empty());
+
+                    kernel.step(&self.ctx, selection, &mut tail)?;
+                    if tail.len() != N && tail.len() != selection.true_count() {
+                        vortex_bail!(
+                            "Kernel produced incorrect number of output elements, expected either {} or {}, got {}",
                             N,
                             selection.true_count(),
-                            added_len
-                        ),
+                            tail.len()
+                        );
                     }
+
+                    // Now we append the produced output back to the main output vector.
+                    output.unsplit(tail);
                 }
                 OutputTarget::IntermediateVector(vector_id) => {
-                    let mut out_vector = VectorMut::from(self.ctx.take_output(vector_id));
+                    let mut out_vector = self.ctx.take_output(vector_id);
                     out_vector.clear();
-                    assert!(
-                        out_vector.capacity() >= N,
-                        "Insufficient capacity in intermediate vector"
-                    );
 
+                    assert!(out_vector.is_empty());
                     kernel.step(&self.ctx, selection, &mut out_vector)?;
 
                     match out_vector.len() {
-                        N => {
+                        // Valid cases are all N elements, or only the selected elements.
+                        n if n == N || n == selection.true_count() => {
                             // If the kernel added N elements, the output is in-place.
-                            self.ctx
-                                .replace_output(vector_id, PipelineVector::InPlace(out_vector));
-                        }
-                        _ if out_vector.len() == selection.true_count() => {
-                            // If the kernel added exactly the number of selected elements,
-                            // the output is already compacted into the start of the vector.
-                            self.ctx
-                                .replace_output(vector_id, PipelineVector::Compact(out_vector));
+                            self.ctx.replace_output(vector_id, out_vector);
                         }
                         _ => vortex_bail!(
-                            "Kernel produced incorrect number of output elements, expected to append either {} or {}, got {}",
+                            "Kernel produced incorrect number of output elements, expected either {} or {}, got {}",
                             N,
                             selection.true_count(),
                             out_vector.len()

vortex-array/src/pipeline/mod.rs

@@ -67,19 +67,19 @@ pub trait BindContext {
 /// Each step of the kernel processes zero or more input vectors, and writes output to a
 /// pre-allocated mutable output vector.
 ///
-/// Input vectors are provided via the [`KernelCtx`] and indicate the position of their elements
-/// as either [`PipelineVector::InPlace`] or [`PipelineVector::Compact`] based on whether
-/// the selected elements are in their original positions or compacted at the start of the vector
-/// respectively.
+/// Input vectors will either have length [`N`], indicating that all elements from the step are
+/// present. Or they will have length equal to the [`BitView::true_count`] of the selection mask,
+/// in which case only the selected elements are present.
 ///
-/// The provided mutable output vector is guaranteed to have at least `N` elements of capacity.
-/// The kernel **must** append either [`BitView::true_count`] elements to the output vector (in
-/// which case the output elements are considered to be in the "Compact" position), or it must
-/// append `N` elements (in which case the output elements are considered to be in their "InPlace"
-/// positions). The pipeline driver will assert these conditions after each step.
+/// Output vectors will always be passed with length zero.
 ///
-/// Note that the output vector may not be empty at the start of the step. The kernel must append
-/// its output to the existing contents of the output vector, rather than replacing it.
+/// Kernels may choose to output either all `N` elements in their original positions, or output
+/// only the selected elements to the first `true_count` positions of the output vector. When
+/// emitting `N` elements in-place, the kernel may omit expensive computations over the unselected
+/// elements, provided that the output elements in those positions are still valid (i.e. typically
+/// zeroed, rather than undefined).
+///
+/// The pipeline driver will verify these conditions before and after each step.
 pub trait Kernel: Send {
     /// Perform a single step of the kernel.
     fn step(
@@ -92,11 +92,11 @@ pub trait Kernel: Send {
 
 /// The context provided to kernels during execution to access input vectors.
 pub struct KernelCtx {
-    vectors: Vec<Option<PipelineVector>>,
+    vectors: Vec<Option<VectorMut>>,
 }
 
 impl KernelCtx {
-    fn new(vectors: Vec<PipelineVector>) -> Self {
+    fn new(vectors: Vec<VectorMut>) -> Self {
         Self {
             vectors: vectors.into_iter().map(Some).collect(),
         }
@@ -112,21 +112,21 @@ impl KernelCtx {
     ///
     /// If the input vector at the given index is not available (typically because the vector
     /// happens to be currently borrowed as an output vector!).
-    pub fn input(&mut self, id: VectorId) -> &PipelineVector {
+    pub fn input(&mut self, id: VectorId) -> &VectorMut {
         self.vectors[id.0]
             .as_ref()
             .vortex_expect("Input vector at index is not available")
     }
 
     #[inline]
-    fn take_output(&mut self, id: &VectorId) -> PipelineVector {
+    fn take_output(&mut self, id: &VectorId) -> VectorMut {
         self.vectors[id.0]
             .take()
             .vortex_expect("Output vector at index is not available")
     }
 
     #[inline]
-    fn replace_output(&mut self, id: &VectorId, vec: PipelineVector) {
+    fn replace_output(&mut self, id: &VectorId, vec: VectorMut) {
         self.vectors[id.0] = Some(vec);
     }
 }
@@ -140,23 +140,3 @@ impl VectorId {
         VectorId(idx)
     }
 }
-
-/// A pipeline vector passed into and out of pipeline kernels.
-#[derive(Debug)]
-pub enum PipelineVector {
-    /// `InPlace` indicates that elements are in their original positions, where the selected
-    /// elements are identified by true values in the selection mask.
-    InPlace(VectorMut),
-    /// Compact indicates that the selected elements are compacted at the start of the vector in
-    /// positions `0..true_count`.
-    Compact(VectorMut),
-}
-
-impl From<PipelineVector> for VectorMut {
-    fn from(value: PipelineVector) -> Self {
-        match value {
-            PipelineVector::InPlace(vec) => vec,
-            PipelineVector::Compact(vec) => vec,
-        }
-    }
-}