Execute tensor product functions with RowFn

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

vortex-tensor/benches/cosine_similarity.rs

@@ -21,6 +21,7 @@ use vortex_array::VortexSessionExecute;
 use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::ExtensionArray;
 use vortex_array::arrays::FixedSizeListArray;
+use vortex_array::arrays::MaskedArray;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::dtype::DType;
 use vortex_array::dtype::Nullability;
@@ -41,15 +42,11 @@ fn main() {
 }
 
 /// Total `f64` elements per operand, held constant across widths: the row count is
-/// `ELEMENTS / width`. This budget is well below the one the other tensor benches use, because
-/// the constant arms recompute the broadcast vector's norm per row and cost roughly ten times the
-/// column arms per element. It is what keeps every arm inside the 1 ms per-iteration limit from
-/// `docs/developer-guide/benchmarking.md`, measured against CodSpeed's CPU simulation.
+/// `ELEMENTS / width`. The smaller budget keeps the wider cosine kernels inside the 1 ms
+/// per-iteration limit from `docs/developer-guide/benchmarking.md` under CodSpeed simulation.
 const ELEMENTS: usize = 1_024;
 
-/// Widths chosen to separate the two costs, as in `l2_norm.rs`: the redundant norm pass is
-/// `O(rows * width)`, one third of the closure's arithmetic, so wide tensors show the hoist
-/// while a narrow one is dominated by per-row framework costs.
+/// Widths that expose both fixed row-framework costs and the `O(width)` kernel work.
 const WIDTHS: &[usize] = &[2, 32, 256];
 
 /// `ELEMENTS / width` vectors of `width` `f64` elements, non-nullable. `seed` offsets the values so
@@ -106,6 +103,22 @@ fn column_x_constant(bencher: Bencher, width: usize) {
     bench_cosine(bencher, vectors(width, 0), constant_vector(width));
 }
 
+/// The lhs is a broadcast query vector, whose norm is the same in every row.
+#[divan::bench(args = WIDTHS)]
+fn constant_x_column(bencher: Bencher, width: usize) {
+    bench_cosine(bencher, constant_vector(width), vectors(width, 31));
+}
+
+/// A nullable broadcast rhs exercises constant preparation and output validity together.
+#[divan::bench(args = WIDTHS)]
+fn column_x_nullable_constant(bencher: Bencher, width: usize) {
+    let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0));
+    let rhs = MaskedArray::try_new(constant_vector(width), validity)
+        .unwrap()
+        .into_array();
+    bench_cosine(bencher, vectors(width, 0), rhs);
+}
+
 /// One query vector represented as an extension array over constant storage.
 fn extension_constant_vector(width: usize) -> ArrayRef {
     let ext_dtype = vectors(width, 0).dtype().as_extension().clone();

vortex-tensor/benches/inner_product.rs

@@ -15,11 +15,17 @@ use divan::Bencher;
 use divan::counter::ItemsCount;
 use mimalloc::MiMalloc;
 use vortex_array::ArrayRef;
+use vortex_array::EmptyMetadata;
 use vortex_array::IntoArray;
 use vortex_array::VortexSessionExecute;
+use vortex_array::arrays::ConstantArray;
 use vortex_array::arrays::FixedSizeListArray;
 use vortex_array::arrays::MaskedArray;
 use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::Nullability;
+use vortex_array::dtype::PType;
+use vortex_array::scalar::Scalar;
 use vortex_array::validity::Validity;
 use vortex_buffer::Buffer;
 use vortex_tensor::scalar_fns::inner_product::InnerProduct;
@@ -56,6 +62,16 @@ fn vectors(width: usize, seed: usize) -> ArrayRef {
     Vector::try_new_vector_array(storage).unwrap()
 }
 
+fn constant_vector(width: usize) -> ArrayRef {
+    let element_dtype = DType::Primitive(PType::F64, Nullability::NonNullable);
+    let children = (0..width)
+        .map(|i| Scalar::primitive(((i % 97) as f64) - 48.0, Nullability::NonNullable))
+        .collect();
+    let storage = Scalar::fixed_size_list(element_dtype, children, Nullability::NonNullable);
+    let vector = Scalar::extension::<Vector>(EmptyMetadata, storage);
+    ConstantArray::new(vector, ELEMENTS / width).into_array()
+}
+
 fn bench_inner_product(bencher: Bencher, lhs: ArrayRef, rhs: ArrayRef) {
     let session = vortex_array::array_session();
     bencher
@@ -84,3 +100,22 @@ fn nullable(bencher: Bencher, width: usize) {
         .into_array();
     bench_inner_product(bencher, lhs, vectors(width, 31));
 }
+
+#[divan::bench(args = WIDTHS)]
+fn column_x_constant(bencher: Bencher, width: usize) {
+    bench_inner_product(bencher, vectors(width, 0), constant_vector(width));
+}
+
+#[divan::bench(args = WIDTHS)]
+fn constant_x_column(bencher: Bencher, width: usize) {
+    bench_inner_product(bencher, constant_vector(width), vectors(width, 31));
+}
+
+#[divan::bench(args = WIDTHS)]
+fn column_x_nullable_constant(bencher: Bencher, width: usize) {
+    let validity = Validity::from_iter((0..ELEMENTS / width).map(|i| i % 8 != 0));
+    let rhs = MaskedArray::try_new(constant_vector(width), validity)
+        .unwrap()
+        .into_array();
+    bench_inner_product(bencher, vectors(width, 0), rhs);
+}

vortex-tensor/src/encodings/normalized/mod.rs

@@ -21,7 +21,6 @@ pub use array::NormalizedSlots;
 mod compress;
 pub use compress::NormalizedScheme;
 pub use compress::normalize;
-pub(crate) use compress::try_build_constant_normalized;
 
 mod execute;