add MixedL2Distance optimizer

7 files changed, 245 insertions, 1 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index 822a48bffc3..4e9fd1b27be 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -75,6 +75,7 @@ vespa_define_module(
     src/tests/instruction/mapped_lookup
     src/tests/instruction/mixed_112_dot_product
     src/tests/instruction/mixed_inner_product_function
+    src/tests/instruction/mixed_l2_distance
     src/tests/instruction/mixed_simple_join_function
     src/tests/instruction/pow_as_map_optimizer
     src/tests/instruction/remove_trivial_dimension_optimizer
diff --git a/eval/src/tests/instruction/mixed_l2_distance/CMakeLists.txt b/eval/src/tests/instruction/mixed_l2_distance/CMakeLists.txt
new file mode 100644
index 00000000000..ecbe69a69f4
--- /dev/null
+++ b/eval/src/tests/instruction/mixed_l2_distance/CMakeLists.txt
@@ -0,0 +1,10 @@
+# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+vespa_add_executable(eval_mixed_l2_distance_test_app TEST
+    SOURCES
+    mixed_l2_distance_test.cpp
+    DEPENDS
+    vespaeval
+    GTest::GTest
+)
+vespa_add_test(NAME eval_mixed_l2_distance_test_app COMMAND eval_mixed_l2_distance_test_app)
diff --git a/eval/src/tests/instruction/mixed_l2_distance/mixed_l2_distance_test.cpp b/eval/src/tests/instruction/mixed_l2_distance/mixed_l2_distance_test.cpp
new file mode 100644
index 00000000000..8f651f7a891
--- /dev/null
+++ b/eval/src/tests/instruction/mixed_l2_distance/mixed_l2_distance_test.cpp
@@ -0,0 +1,78 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/eval/eval/fast_value.h>
+#include <vespa/eval/eval/tensor_function.h>
+#include <vespa/eval/eval/test/eval_fixture.h>
+#include <vespa/eval/eval/test/gen_spec.h>
+#include <vespa/eval/instruction/mixed_l2_distance.h>
+#include <vespa/vespalib/util/stash.h>
+#include <vespa/vespalib/util/stringfmt.h>
+
+#include <vespa/vespalib/util/require.h>
+#include <vespa/vespalib/gtest/gtest.h>
+
+using namespace vespalib;
+using namespace vespalib::eval;
+using namespace vespalib::eval::test;
+using namespace vespalib::eval::tensor_function;
+
+struct FunInfo {
+    using LookFor = MixedL2Distance;
+    bool debug_dump;
+    void verify(const LookFor &fun) const {
+        EXPECT_TRUE(fun.result_is_mutable());
+        if (debug_dump) {
+            fprintf(stderr, "%s", fun.as_string().c_str());
+        }
+    }
+};
+
+void verify_optimized(const vespalib::string &expr) {
+    SCOPED_TRACE(expr.c_str());
+    auto diff_types = CellTypeSpace(CellTypeUtils::list_types(), 2).different();
+    EvalFixture::verify<FunInfo>(expr, {}, diff_types);
+    auto same_types = CellTypeSpace(CellTypeUtils::list_types(), 2).same();
+    EvalFixture::verify<FunInfo>(expr, {FunInfo{false}}, same_types);
+}
+
+void verify_not_optimized(const vespalib::string &expr) {
+    SCOPED_TRACE(expr.c_str());
+    CellTypeSpace just_double({CellType::DOUBLE}, 2);
+    EvalFixture::verify<FunInfo>(expr, {}, just_double);
+}
+
+//-----------------------------------------------------------------------------
+
+TEST(MixedL2DistanceTest, squared_l2_distance_can_be_optimized) {
+    verify_optimized("reduce(map(x5-x5y7_2, f(a)(a * a)), sum, x)");
+    verify_optimized("reduce((x5-x5y7_2)^2,sum,x)");
+    verify_optimized("reduce((x5y7_2-x5)^2,sum,x)");    
+    verify_optimized("sqrt(reduce(map(x5-x5y7_2, f(a)(a * a)), sum, x))");
+}
+
+TEST(MixedL2DistanceTest, trivial_dimensions_are_ignored) {
+    verify_optimized("reduce((x5z1-x5y7_2)^2,sum,x)");
+    verify_optimized("reduce((x5-x5y7_2z1)^2,sum,x)");
+    verify_optimized("reduce((x5z1-x5y7_2z1)^2,sum,x)");
+}
+
+TEST(MixedL2DistanceTest, multiple_dimensions_can_be_used) {
+    verify_optimized("reduce((x5z3-x5y7_2z3)^2,sum,x,z)");
+    verify_optimized("reduce((x5-x5y7_2z3_1)^2,sum,x)");
+}
+
+TEST(MixedL2DistanceTest, not_optimizing_close_match) {
+    verify_not_optimized("reduce(map(x5-x5y7_2, f(a)(a * a)), avg, x)");
+    verify_not_optimized("reduce(map(x5-x5y7_2, f(a)(a + a)), sum, x)");    
+}
+
+TEST(MixedL2DistanceTest, result_must_be_sparse) {
+    verify_not_optimized("reduce((x5-x5y7_2)^2,sum,x,y)");    
+    verify_not_optimized("reduce((x5z1-x5y7_2)^2,sum,x,y)");
+    verify_not_optimized("reduce((x5z3-x5y7_2z3)^2,sum,x)");
+    verify_not_optimized("reduce((x5z3-x5y7_2z3)^2,sum,z)");
+}
+
+//-----------------------------------------------------------------------------
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/eval/src/vespa/eval/eval/optimize_tensor_function.cpp b/eval/src/vespa/eval/eval/optimize_tensor_function.cpp
index 41fa550ce07..1d0be47a309 100644
--- a/eval/src/vespa/eval/eval/optimize_tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/optimize_tensor_function.cpp
@@ -34,6 +34,7 @@
 #include <vespa/eval/instruction/dense_tensor_peek_function.h>
 #include <vespa/eval/instruction/dense_hamming_distance.h>
 #include <vespa/eval/instruction/l2_distance.h>
+#include <vespa/eval/instruction/mixed_l2_distance.h>
 #include <vespa/eval/instruction/simple_join_count.h>
 #include <vespa/eval/instruction/mapped_lookup.h>
 
@@ -73,7 +74,8 @@ const TensorFunction &optimize_for_factory(const ValueBuilderFactory &, const Te
                           child.set(Sparse112DotProduct::optimize(child.get(), stash));
                           child.set(Mixed112DotProduct::optimize(child.get(), stash));
                           child.set(BestSimilarityFunction::optimize(child.get(), stash));
-                          child.set(L2Distance::optimize(child.get(), stash));                          
+                          child.set(L2Distance::optimize(child.get(), stash));
+                          child.set(MixedL2Distance::optimize(child.get(), stash));                                                    
                       });
     run_optimize_pass(root, [&stash](const Child &child)
                       {
diff --git a/eval/src/vespa/eval/instruction/CMakeLists.txt b/eval/src/vespa/eval/instruction/CMakeLists.txt
index 7df3f745e79..02bbfec5dd3 100644
--- a/eval/src/vespa/eval/instruction/CMakeLists.txt
+++ b/eval/src/vespa/eval/instruction/CMakeLists.txt
@@ -34,6 +34,7 @@ vespa_add_library(eval_instruction OBJECT
     mapped_lookup.cpp
     mixed_112_dot_product.cpp
     mixed_inner_product_function.cpp
+    mixed_l2_distance.cpp
     mixed_simple_join_function.cpp
     pow_as_map_optimizer.cpp
     remove_trivial_dimension_optimizer.cpp
diff --git a/eval/src/vespa/eval/instruction/mixed_l2_distance.cpp b/eval/src/vespa/eval/instruction/mixed_l2_distance.cpp
new file mode 100644
index 00000000000..d11bf704df5
--- /dev/null
+++ b/eval/src/vespa/eval/instruction/mixed_l2_distance.cpp
@@ -0,0 +1,131 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "mixed_l2_distance.h"
+#include <vespa/eval/eval/operation.h>
+#include <vespa/eval/eval/value.h>
+#include <vespa/vespalib/hwaccelrated/iaccelrated.h>
+#include <vespa/vespalib/util/require.h>
+
+#include <vespa/log/log.h>
+LOG_SETUP(".eval.instruction.mixed_l2_distance");
+
+namespace vespalib::eval {
+
+using namespace tensor_function;
+
+namespace {
+
+static const auto &hw = hwaccelrated::IAccelrated::getAccelerator();
+
+template <typename T>
+double h_sq_l2(const T *a, const T *b, size_t len) {
+    return hw.squaredEuclideanDistance(a, b, len);
+}
+
+template <>
+double h_sq_l2<Int8Float>(const Int8Float *a, const Int8Float *b, size_t len) {
+    return hw.squaredEuclideanDistance((const int8_t *)a, (const int8_t *)b, len);
+}
+
+template <>
+double h_sq_l2<BFloat16>(const BFloat16 *a, const BFloat16 *b, size_t len) {
+    float sum = 0.0;
+    for (size_t i = 0; i < len; ++i) {
+        float x = a[i];
+        float y = b[i];
+        float d = (x - y);
+        sum += d * d;
+    }
+    return sum;
+}
+
+struct MixedSqL2Param {
+    const ValueType res_type;
+    const size_t vec_len;
+    MixedSqL2Param(const ValueType &r, size_t vl) : res_type(r), vec_len(vl) {}
+};
+
+template <typename ICT, typename OCT>
+void mixed_squared_l2_distance_op(InterpretedFunction::State &state, uint64_t param_in) {
+    const auto &param = unwrap_param<MixedSqL2Param>(param_in);
+    const Value &vec = state.peek(0);
+    const Value &mix = state.peek(1);
+    size_t output_size = mix.index().size();
+    auto output_cells = state.stash.create_uninitialized_array<OCT>(output_size);
+    auto vec_cells = (const ICT *) vec.cells().data;
+    auto mix_cells = (const ICT *) mix.cells().data;
+    for (size_t i = 0; i < output_size; ++i) {
+        output_cells[i] = h_sq_l2<ICT>(vec_cells, mix_cells, param.vec_len);
+        mix_cells += param.vec_len;
+    }
+    Value &result_ref = state.stash.create<ValueView>(param.res_type, mix.index(), TypedCells(output_cells));
+    state.pop_pop_push(result_ref);
+}
+
+struct MultiSelectOp {
+    template <typename ICM>
+    static InterpretedFunction::op_function invoke() {
+        using ICT = CellValueType<ICM::value.cell_type>;
+        constexpr CellMeta ocm = ICM::value.decay();
+        using OCT = CellValueType<ocm.cell_type>;
+        return mixed_squared_l2_distance_op<ICT, OCT>;
+    }
+};
+
+bool mixed_compatible_types(const ValueType &mix, const ValueType &vec, const ValueType &res) {
+    return ((mix.cell_type() == vec.cell_type()) &&
+	   vec.is_dense() &&
+	   res.nontrivial_indexed_dimensions().empty() &&
+	   (res.mapped_dimensions().size() > 0) &&
+            (mix.nontrivial_indexed_dimensions() == vec.nontrivial_indexed_dimensions()) &&
+	   (mix.mapped_dimensions() == res.mapped_dimensions()));
+}
+
+
+} // namespace <unnamed>
+
+MixedL2Distance::MixedL2Distance(const ValueType &result_type, const TensorFunction &mix_in, const TensorFunction &vec_in)
+    : tensor_function::Op2(result_type, mix_in, vec_in)
+{
+}
+
+InterpretedFunction::Instruction
+MixedL2Distance::compile_self(const ValueBuilderFactory &, Stash &stash) const
+{
+    auto mix_t = lhs().result_type();
+    auto vec_t = rhs().result_type();
+    REQUIRE_EQ(mix_t.cell_type(), vec_t.cell_type());
+    REQUIRE_EQ(mix_t.dense_subspace_size(), vec_t.dense_subspace_size());
+    const auto &param = stash.create<MixedSqL2Param>(result_type(), mix_t.dense_subspace_size());
+    auto mix_cm = mix_t.cell_meta().not_scalar();
+    auto res_cm = mix_t.cell_meta().decay();
+    REQUIRE_EQ(res_cm.cell_type, result_type().cell_type());
+    auto op = typify_invoke<1, TypifyCellMeta, MultiSelectOp>(mix_cm);
+    return InterpretedFunction::Instruction(op, wrap_param<MixedSqL2Param>(param));
+}
+
+const TensorFunction &
+MixedL2Distance::optimize(const TensorFunction &expr, Stash &stash)
+{
+    auto reduce = as<Reduce>(expr);
+    if (reduce && (reduce->aggr() == Aggr::SUM)) {
+        auto map = as<Map>(reduce->child());
+        if (map && (map->function() == operation::Square::f)) {
+            auto join = as<Join>(map->child());
+            if (join && (join->function() == operation::Sub::f)) {
+	       const auto & res_type = expr.result_type();
+	       const auto & left_type = join->lhs().result_type();
+	       const auto & right_type = join->rhs().result_type();
+	       if (mixed_compatible_types(left_type, right_type, res_type)) {
+                    return stash.create<MixedL2Distance>(res_type, join->lhs(), join->rhs());
+	       }
+	       if (mixed_compatible_types(right_type, left_type, res_type)) {
+                    return stash.create<MixedL2Distance>(res_type, join->rhs(), join->lhs());
+	       }
+            }
+        }
+    }
+    return expr;
+}
+
+} // namespace
diff --git a/eval/src/vespa/eval/instruction/mixed_l2_distance.h b/eval/src/vespa/eval/instruction/mixed_l2_distance.h
new file mode 100644
index 00000000000..84f6f14d3c9
--- /dev/null
+++ b/eval/src/vespa/eval/instruction/mixed_l2_distance.h
@@ -0,0 +1,21 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/eval/eval/tensor_function.h>
+
+namespace vespalib::eval {
+
+/**
+ * Tensor function for a squared euclidean distance producing a sparse result.
+ **/
+class MixedL2Distance : public tensor_function::Op2
+{
+public:
+    MixedL2Distance(const ValueType &result_type, const TensorFunction &mix_in, const TensorFunction &vec_in);
+    InterpretedFunction::Instruction compile_self(const ValueBuilderFactory &factory, Stash &stash) const override;
+    bool result_is_mutable() const override { return true; }
+    static const TensorFunction &optimize(const TensorFunction &expr, Stash &stash);
+};
+
+} // namespace