diff options
author | Arne Juul <arnej@yahooinc.com> | 2023-06-28 09:40:57 +0000 |
---|---|---|
committer | Arne Juul <arnej@yahooinc.com> | 2023-06-28 13:45:06 +0000 |
commit | 116403b79d7e356d78eb0a1d2237611c46961e6d (patch) | |
tree | 0473d1cac6bc694dd64170df23d2376606d1449f /eval | |
parent | 3a066100ef96ffd3ac73d1d06096c98da077f296 (diff) |
add MixedL2Distance optimizer
Diffstat (limited to 'eval')
-rw-r--r-- | eval/CMakeLists.txt | 1 | ||||
-rw-r--r-- | eval/src/tests/instruction/mixed_l2_distance/CMakeLists.txt | 10 | ||||
-rw-r--r-- | eval/src/tests/instruction/mixed_l2_distance/mixed_l2_distance_test.cpp | 78 | ||||
-rw-r--r-- | eval/src/vespa/eval/eval/optimize_tensor_function.cpp | 4 | ||||
-rw-r--r-- | eval/src/vespa/eval/instruction/CMakeLists.txt | 1 | ||||
-rw-r--r-- | eval/src/vespa/eval/instruction/mixed_l2_distance.cpp | 131 | ||||
-rw-r--r-- | eval/src/vespa/eval/instruction/mixed_l2_distance.h | 21 |
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 ¶m = 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 ¶m = 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 |