Merge pull request #16927 from vespa-engine/havardpe/single-dense-reduce-revisit

revisit dense_single_reduce_function with test

5 files changed, 94 insertions, 107 deletions

diff --git a/eval/src/tests/eval/gen_spec/gen_spec_test.cpp b/eval/src/tests/eval/gen_spec/gen_spec_test.cpp
index 30ef42d7ebb..304eac1eb09 100644
--- a/eval/src/tests/eval/gen_spec/gen_spec_test.cpp
+++ b/eval/src/tests/eval/gen_spec/gen_spec_test.cpp
@@ -237,10 +237,10 @@ TEST(GenSpecTest, gen_spec_can_be_implicitly_converted_to_tensor_spec) {
 //-----------------------------------------------------------------------------
 
 TEST(GenSpecFromDescTest, dim_spec_and_gen_spec_can_be_created_from_desc) {
-    // 'a2b3#7c5#'
-    auto expect = GenSpec().idx("a", 2).map("b", 3, 7).map("c", 5).gen();
-    auto dim_desc = GenSpec().desc("a2").desc("b3#7").desc("c5#").gen();
-    auto gen_desc = GenSpec::from_desc("a2b3#7c5#").gen();
+    // 'a2b3_7'
+    auto expect = GenSpec().idx("a", 2).map("b", 3, 7).gen();
+    auto dim_desc = GenSpec().desc("a2").desc("b3_7").gen();
+    auto gen_desc = GenSpec::from_desc("a2b3_7").gen();
     EXPECT_EQ(dim_desc, expect);
     EXPECT_EQ(gen_desc, expect);
 }
diff --git a/eval/src/tests/instruction/dense_single_reduce_function/dense_single_reduce_function_test.cpp b/eval/src/tests/instruction/dense_single_reduce_function/dense_single_reduce_function_test.cpp
index a2733b1f7cb..c26f46f5e11 100644
--- a/eval/src/tests/instruction/dense_single_reduce_function/dense_single_reduce_function_test.cpp
+++ b/eval/src/tests/instruction/dense_single_reduce_function/dense_single_reduce_function_test.cpp
@@ -14,23 +14,10 @@ using namespace vespalib;
 using namespace vespalib::eval;
 using namespace vespalib::eval::test;
 using namespace vespalib::eval::tensor_function;
+using vespalib::make_string_short::fmt;
 
 const ValueBuilderFactory &prod_factory = FastValueBuilderFactory::get();
 
-EvalFixture::ParamRepo make_params() {
-    return EvalFixture::ParamRepo()
-        .add_variants("a2b3c4d5", GenSpec().idx("a", 2).idx("b", 3).idx("c", 4).idx("d", 5))
-        .add_variants("a9b9c9d9", GenSpec().idx("a", 9).idx("b", 9).idx("c", 9).idx("d", 9))
-        .add_variants("a2b1c1", GenSpec().idx("a", 2).idx("b", 1).idx("c", 1))
-        .add_variants("a1b2c1", GenSpec().idx("a", 1).idx("b", 2).idx("c", 1))
-        .add_variants("a1b1c2", GenSpec().idx("a", 1).idx("b", 1).idx("c", 2))
-        .add_variants("a1b1c1", GenSpec().idx("a", 1).idx("b", 1).idx("c", 1))
-        .add_variants("a10", GenSpec().idx("a", 10))
-        .add("xy_mapped", GenSpec().map("x", {"a", "b"}).map("y", {"x", "y"}))
-        .add("xyz_mixed", GenSpec().map("x", {"a", "b"}).map("y", {"x", "y"}).idx("z", 3));
-}
-EvalFixture::ParamRepo param_repo = make_params();
-
 struct ReduceSpec {
     size_t outer_size;
     size_t reduce_size;
@@ -38,37 +25,49 @@ struct ReduceSpec {
     Aggr aggr;
 };
 
-void verify_optimized_impl(const vespalib::string &expr, const std::vector<ReduceSpec> &spec_list) {
-    EvalFixture slow_fixture(prod_factory, expr, param_repo, false);
-    EvalFixture fixture(prod_factory, expr, param_repo, true);
-    EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
-    EXPECT_EQUAL(fixture.result(), slow_fixture.result());
-    auto info = fixture.find_all<DenseSingleReduceFunction>();
-    ASSERT_EQUAL(info.size(), spec_list.size());
-    for (size_t i = 0; i < spec_list.size(); ++i) {
-        EXPECT_TRUE(info[i]->result_is_mutable());
-        EXPECT_EQUAL(info[i]->outer_size(), spec_list[i].outer_size);
-        EXPECT_EQUAL(info[i]->reduce_size(), spec_list[i].reduce_size);
-        EXPECT_EQUAL(info[i]->inner_size(), spec_list[i].inner_size);
-        EXPECT_EQUAL(int(info[i]->aggr()), int(spec_list[i].aggr));
+void verify_impl(const vespalib::string &expr,
+                 const std::vector<ReduceSpec> &spec_list,
+                 const std::vector<CellType> &with_cell_types)
+{
+    auto fun = Function::parse(expr);
+    ASSERT_EQUAL(fun->num_params(), 1u);
+    vespalib::string param_name = fun->param_name(0);
+    const auto param_spec = GenSpec::from_desc(param_name);
+    for (CellType ct: with_cell_types) {
+        EvalFixture::ParamRepo param_repo;
+        param_repo.add(param_name, param_spec.cpy().cells(ct));
+        EvalFixture slow_fixture(prod_factory, expr, param_repo, false);
+        EvalFixture fixture(prod_factory, expr, param_repo, true);
+        EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
+        EXPECT_EQUAL(fixture.result(), slow_fixture.result());
+        auto info = fixture.find_all<DenseSingleReduceFunction>();
+        ASSERT_EQUAL(info.size(), spec_list.size());
+        for (size_t i = 0; i < spec_list.size(); ++i) {
+            EXPECT_TRUE(info[i]->result_is_mutable());
+            EXPECT_EQUAL(info[i]->outer_size(), spec_list[i].outer_size);
+            EXPECT_EQUAL(info[i]->reduce_size(), spec_list[i].reduce_size);
+            EXPECT_EQUAL(info[i]->inner_size(), spec_list[i].inner_size);
+            EXPECT_EQUAL(int(info[i]->aggr()), int(spec_list[i].aggr));
+        }
     }
 }
 
-void verify_optimized(const vespalib::string &expr, const ReduceSpec &spec) {
-    verify_optimized_impl(expr, {spec});
+void verify_not_optimized(const vespalib::string &expr,
+                          std::vector<CellType> with_cell_types = {CellType::DOUBLE})
+{
+    verify_impl(expr, {}, with_cell_types);
 }
 
-void verify_optimized(const vespalib::string &expr, const ReduceSpec &spec1, const ReduceSpec &spec2) {
-    verify_optimized_impl(expr, {spec1, spec2});
+void verify_optimized(const vespalib::string &expr, const ReduceSpec &spec,
+                      std::vector<CellType> with_cell_types = CellTypeUtils::list_types())
+{
+    verify_impl(expr, {spec}, with_cell_types);
 }
 
-void verify_not_optimized(const vespalib::string &expr) {
-    EvalFixture slow_fixture(prod_factory, expr, param_repo, false);
-    EvalFixture fixture(prod_factory, expr, param_repo, true);
-    EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
-    EXPECT_EQUAL(fixture.result(), slow_fixture.result());
-    auto info = fixture.find_all<DenseSingleReduceFunction>();
-    EXPECT_TRUE(info.empty());
+void verify_optimized(const vespalib::string &expr, const ReduceSpec &spec1, const ReduceSpec &spec2,
+                      std::vector<CellType> with_cell_types = CellTypeUtils::list_types())
+{
+    verify_impl(expr, {spec1, spec2}, with_cell_types);
 }
 
 TEST("require that reduce to scalar is not optimized") {
@@ -77,14 +76,14 @@ TEST("require that reduce to scalar is not optimized") {
 }
 
 TEST("require that sparse reduce is not optimized") {
-    TEST_DO(verify_not_optimized("reduce(xy_mapped,sum,x)"));
-    TEST_DO(verify_not_optimized("reduce(xy_mapped,sum,y)"));
+    TEST_DO(verify_not_optimized("reduce(x2_1y2_1,sum,x)"));
+    TEST_DO(verify_not_optimized("reduce(x2_1y2_1,sum,y)"));
 }
 
 TEST("require that mixed reduce is not optimized") {
-    TEST_DO(verify_not_optimized("reduce(xyz_mixed,sum,x)"));
-    TEST_DO(verify_not_optimized("reduce(xyz_mixed,sum,y)"));
-    TEST_DO(verify_not_optimized("reduce(xyz_mixed,sum,z)"));
+    TEST_DO(verify_not_optimized("reduce(x2_1y2_1z3,sum,x)"));
+    TEST_DO(verify_not_optimized("reduce(x2_1y2_1z3,sum,y)"));
+    TEST_DO(verify_not_optimized("reduce(x2_1y2_1z3,sum,z)"));
 }
 
 TEST("require that reducing trivial dimensions is not optimized") {
@@ -98,11 +97,11 @@ TEST("require that reducing trivial dimensions is not optimized") {
 }
 
 TEST("require that atleast_8 dense single reduce works") {
-    TEST_DO(verify_optimized("reduce(a9b9c9d9,avg,a)", {1, 9, 729, Aggr::AVG}));
-    TEST_DO(verify_optimized("reduce(a9b9c9d9,avg,b)", {9, 9, 81, Aggr::AVG}));
-    TEST_DO(verify_optimized("reduce(a9b9c9d9,avg,c)", {81, 9, 9, Aggr::AVG}));
-    TEST_DO(verify_optimized("reduce(a9b9c9d9,avg,d)", {729, 9, 1, Aggr::AVG}));
-    TEST_DO(verify_optimized("reduce(a9b9c9d9,sum,c,d)", {81, 81, 1, Aggr::SUM}));
+    TEST_DO(verify_optimized("reduce(a9b9c9d9,avg,a)", {1, 9, 729, Aggr::AVG}, {CellType::FLOAT}));
+    TEST_DO(verify_optimized("reduce(a9b9c9d9,avg,b)", {9, 9, 81, Aggr::AVG}, {CellType::FLOAT}));
+    TEST_DO(verify_optimized("reduce(a9b9c9d9,avg,c)", {81, 9, 9, Aggr::AVG}, {CellType::FLOAT}));
+    TEST_DO(verify_optimized("reduce(a9b9c9d9,avg,d)", {729, 9, 1, Aggr::AVG}, {CellType::FLOAT}));
+    TEST_DO(verify_optimized("reduce(a9b9c9d9,sum,c,d)", {81, 81, 1, Aggr::SUM}, {CellType::FLOAT}));
 }
 
 TEST("require that simple aggregators can be decomposed into multiple reduce operations") {
@@ -123,17 +122,11 @@ TEST("require that non-simple aggregators cannot be decomposed into multiple red
     TEST_DO(verify_not_optimized("reduce(a2b3c4d5,median,a,c)"));
 }
 
-vespalib::string make_expr(const vespalib::string &arg, const vespalib::string &dim, bool float_cells, Aggr aggr) {
-    return make_string("reduce(%s%s,%s,%s)", arg.c_str(), float_cells ? "_f" : "", AggrNames::name_of(aggr)->c_str(), dim.c_str());
-}
-
 void verify_optimized_multi(const vespalib::string &arg, const vespalib::string &dim, size_t outer_size, size_t reduce_size, size_t inner_size) {
-    for (bool float_cells: {false, true}) {
-        for (Aggr aggr: Aggregator::list()) {
-            if (aggr != Aggr::PROD) {
-                auto expr = make_expr(arg, dim, float_cells, aggr);
-                TEST_DO(verify_optimized(expr, {outer_size, reduce_size, inner_size, aggr}));
-            }
+    for (Aggr aggr: Aggregator::list()) {
+        if (aggr != Aggr::PROD) {
+            auto expr = fmt("reduce(%s,%s,%s)", arg.c_str(), AggrNames::name_of(aggr)->c_str(), dim.c_str());
+            TEST_DO(verify_optimized(expr, {outer_size, reduce_size, inner_size, aggr}));
         }
     }
 }
diff --git a/eval/src/vespa/eval/eval/test/gen_spec.cpp b/eval/src/vespa/eval/eval/test/gen_spec.cpp
index 87e49c376d1..1f7c658ff79 100644
--- a/eval/src/vespa/eval/eval/test/gen_spec.cpp
+++ b/eval/src/vespa/eval/eval/test/gen_spec.cpp
@@ -54,8 +54,7 @@ DimSpec::make_dict(size_t size, size_t stride, const vespalib::string &prefix)
 }
 
 // 'a2' -> DimSpec("a", 2);
-// 'b2#' -> DimSpec("b", make_dict(2, 1, ""));
-// 'c2#3' -> DimSpec("c", make_dict(2, 3, ""));
+// 'b2_3' -> DimSpec("b", make_dict(2, 3, ""));
 DimSpec
 DimSpec::from_desc(const vespalib::string &desc)
 {
@@ -63,7 +62,7 @@ DimSpec::from_desc(const vespalib::string &desc)
     vespalib::string name;
     auto is_num = [](char c) { return ((c >= '0') && (c <= '9')); };
     auto as_num = [](char c) { return size_t(c - '0'); };
-    auto is_map_tag = [](char c) { return (c == '#'); };
+    auto is_map_tag = [](char c) { return (c == '_'); };
     auto is_dim_name = [](char c) { return ((c >= 'a') && (c <= 'z')); };
     auto extract_number = [&]() {
         assert(idx < desc.size());
@@ -82,11 +81,8 @@ DimSpec::from_desc(const vespalib::string &desc)
     if (idx < desc.size()) {
         // mapped
         assert(is_map_tag(desc[idx++]));
-        size_t stride = 1;
-        if (idx < desc.size()) {
-            stride = extract_number();
-            assert(idx == desc.size());
-        }
+        size_t stride = extract_number();
+        assert(idx == desc.size());
         return {name, make_dict(size, stride, "")};
     } else {
         // indexed
@@ -95,7 +91,7 @@ DimSpec::from_desc(const vespalib::string &desc)
 }
 
 // 'a2b12c5' -> GenSpec().idx("a", 2).idx("b", 12).idx("c", 5);
-// 'a2#b3#2c5#' -> GenSpec().map("a", 2).map("b", 3, 2).map("c", 5);
+// 'a2_1b3_2c5_1' -> GenSpec().map("a", 2).map("b", 3, 2).map("c", 5);
 GenSpec
 GenSpec::from_desc(const vespalib::string &desc)
 {
diff --git a/eval/src/vespa/eval/eval/test/gen_spec.h b/eval/src/vespa/eval/eval/test/gen_spec.h
index 6943ba6c98a..def24e6711f 100644
--- a/eval/src/vespa/eval/eval/test/gen_spec.h
+++ b/eval/src/vespa/eval/eval/test/gen_spec.h
@@ -81,8 +81,7 @@ public:
     // (first character is used as dimension name)
     //
     // 'a2' -> DimSpec("a", 2);
-    // 'b2#' -> DimSpec("b", make_dict(2, 1, ""));
-    // 'c2#3' -> DimSpec("c", make_dict(2, 3, ""));
+    // 'b2_3' -> DimSpec("b", make_dict(2, 3, ""));
     static DimSpec from_desc(const vespalib::string &desc);
 };
 
@@ -109,7 +108,7 @@ public:
     // (dimension names must be single characters a-z)
     //
     // 'a2b12c5' -> GenSpec().idx("a", 2).idx("b", 12).idx("c", 5);
-    // 'a2#b3#2c5#' -> GenSpec().map("a", 2).map("b", 3, 2).map("c", 5);
+    // 'a2_1b3_2c5_1' -> GenSpec().map("a", 2).map("b", 3, 2).map("c", 5);
     static GenSpec from_desc(const vespalib::string &desc);
 
     GenSpec(GenSpec &&other);
diff --git a/eval/src/vespa/eval/instruction/dense_single_reduce_function.cpp b/eval/src/vespa/eval/instruction/dense_single_reduce_function.cpp
index 492b6380ad2..014f36eaa2c 100644
--- a/eval/src/vespa/eval/instruction/dense_single_reduce_function.cpp
+++ b/eval/src/vespa/eval/instruction/dense_single_reduce_function.cpp
@@ -21,8 +21,8 @@ struct Params {
         : result_type(result_type_in), outer_size(outer_size_in), reduce_size(reduce_size_in), inner_size(inner_size_in) {}
 };
 
-template <typename CT, typename AGGR>
-CT reduce_cells(const CT *src, size_t reduce_size, size_t stride) {
+template <typename ICT, typename AGGR>
+auto reduce_cells(const ICT *src, size_t reduce_size, size_t stride) {
     AGGR aggr(*src);
     for (size_t i = 1; i < reduce_size; ++i) {
         src += stride;
@@ -54,38 +54,38 @@ auto reduce_cells_atleast_8(size_t n, GET &&get) {
     return aggrs[0].result();
 }
 
-template <typename CT, typename AGGR>
-auto reduce_cells_atleast_8(const CT *src, size_t n) {
+template <typename ICT, typename AGGR>
+auto reduce_cells_atleast_8(const ICT *src, size_t n) {
     return reduce_cells_atleast_8<AGGR>(n, [&](size_t idx) { return src[idx]; });
 }
 
-template <typename CT, typename AGGR>
-auto reduce_cells_atleast_8(const CT *src, size_t n, size_t stride) {
+template <typename ICT, typename AGGR>
+auto reduce_cells_atleast_8(const ICT *src, size_t n, size_t stride) {
     return reduce_cells_atleast_8<AGGR>(n, [&](size_t idx) { return src[idx * stride]; });
 }
 
-template <typename CT, typename AGGR, bool atleast_8, bool is_inner>
-void trace_reduce_impl(const Params &params, const CT *src, CT *dst) {
+template <typename ICT, typename OCT, typename AGGR, bool atleast_8, bool is_inner>
+void trace_reduce_impl(const Params &params, const ICT *src, OCT *dst) {
     constexpr bool aggr_is_complex = is_complex(AGGR::enum_value());
     const size_t block_size = (params.reduce_size * params.inner_size);
     for (size_t outer = 0; outer < params.outer_size; ++outer) {
         for (size_t inner = 0; inner < params.inner_size; ++inner) {
             if (atleast_8 && !aggr_is_complex) {
                 if (is_inner) {
-                    *dst++ = reduce_cells_atleast_8<CT, AGGR>(src + inner, params.reduce_size);
+                    *dst++ = reduce_cells_atleast_8<ICT, AGGR>(src + inner, params.reduce_size);
                 } else {
-                    *dst++ = reduce_cells_atleast_8<CT, AGGR>(src + inner, params.reduce_size, params.inner_size);
+                    *dst++ = reduce_cells_atleast_8<ICT, AGGR>(src + inner, params.reduce_size, params.inner_size);
                 }
             } else {
-                *dst++ = reduce_cells<CT, AGGR>(src + inner, params.reduce_size, params.inner_size);
+                *dst++ = reduce_cells<ICT, AGGR>(src + inner, params.reduce_size, params.inner_size);
             }
         }
         src += block_size;
     }
 }
 
-template <typename CT, typename AGGR>
-void fold_reduce_impl(const Params &params, const CT *src, CT *dst) {
+template <typename ICT, typename OCT, typename AGGR>
+void fold_reduce_impl(const Params &params, const ICT *src, OCT *dst) {
     for (size_t outer = 0; outer < params.outer_size; ++outer) {
         auto saved_dst = dst;
         for (size_t inner = 0; inner < params.inner_size; ++inner) {
@@ -101,27 +101,28 @@ void fold_reduce_impl(const Params &params, const CT *src, CT *dst) {
     }
 }
 
-template <typename CT, typename AGGR, bool atleast_8, bool is_inner>
+template <typename ICT, typename OCT, typename AGGR, bool atleast_8, bool is_inner>
 void my_single_reduce_op(InterpretedFunction::State &state, uint64_t param) {
-    static_assert(std::is_same_v<CT,typename AGGR::value_type>);
+    static_assert(std::is_same_v<OCT,typename AGGR::value_type>);
     constexpr bool aggr_is_simple = is_simple(AGGR::enum_value());
     const auto &params = unwrap_param<Params>(param);
-    const CT *src = state.peek(0).cells().typify<CT>().cbegin();
-    auto dst_cells = state.stash.create_uninitialized_array<CT>(params.outer_size * params.inner_size);
-    CT *dst = dst_cells.begin();
+    const ICT *src = state.peek(0).cells().typify<ICT>().cbegin();
+    auto dst_cells = state.stash.create_uninitialized_array<OCT>(params.outer_size * params.inner_size);
+    OCT *dst = dst_cells.begin();
     if constexpr (aggr_is_simple && !is_inner) {
-        fold_reduce_impl<CT, AGGR>(params, src, dst);
+        fold_reduce_impl<ICT, OCT, AGGR>(params, src, dst);
     } else {
-        trace_reduce_impl<CT,AGGR,atleast_8,is_inner>(params, src, dst);
+        trace_reduce_impl<ICT, OCT, AGGR, atleast_8, is_inner>(params, src, dst);
     }
     state.pop_push(state.stash.create<DenseValueView>(params.result_type, TypedCells(dst_cells)));
 }
 
 struct MyGetFun {
-    template <typename R1, typename R2, typename R3, typename R4> static auto invoke() {
-        using OCT = CellValueType<R1::value.decay().cell_type>;
-        using AggrType = typename R2::template templ<OCT>;
-        return my_single_reduce_op<OCT, AggrType, R3::value, R4::value>;
+    template <typename ICM, typename AGGR, typename GE8, typename I1> static auto invoke() {
+        using ICT = CellValueType<ICM::value.cell_type>;
+        using OCT = CellValueType<CellMeta::reduce(ICM::value.cell_type, false).cell_type>;
+        using AggrType = typename AGGR::template templ<OCT>;
+        return my_single_reduce_op<ICT, OCT, AggrType, GE8::value, I1::value>;
     }
 };
 
@@ -229,7 +230,7 @@ DenseSingleReduceFunction::~DenseSingleReduceFunction() = default;
 InterpretedFunction::Instruction
 DenseSingleReduceFunction::compile_self(const ValueBuilderFactory &, Stash &stash) const
 {
-    auto op = typify_invoke<4,MyTypify,MyGetFun>(child().result_type().cell_meta().limit().not_scalar(),
+    auto op = typify_invoke<4,MyTypify,MyGetFun>(child().result_type().cell_meta().not_scalar(),
                                                  _aggr,
                                                  (_reduce_size >= 8), (_inner_size == 1));
     auto &params = stash.create<Params>(result_type(), _outer_size, _reduce_size, _inner_size);
@@ -241,15 +242,13 @@ DenseSingleReduceFunction::optimize(const TensorFunction &expr, Stash &stash)
 {
     if (auto reduce = as<Reduce>(expr)) {
         const auto &child = reduce->child();
-        if (reduce->result_type().cell_meta().eq(child.result_type().cell_meta())) {
-            auto spec_list = make_dense_single_reduce_list(child.result_type(), reduce->aggr(), reduce->dimensions()); 
-            if (!spec_list.empty()) {
-                const auto *prev = &child;
-                for (const auto &spec: spec_list) {
-                    prev = &stash.create<DenseSingleReduceFunction>(spec, *prev);
-                }
-                return *prev;
+        auto spec_list = make_dense_single_reduce_list(child.result_type(), reduce->aggr(), reduce->dimensions()); 
+        if (!spec_list.empty()) {
+            const auto *prev = &child;
+            for (const auto &spec: spec_list) {
+                prev = &stash.create<DenseSingleReduceFunction>(spec, *prev);
             }
+            return *prev;
         }
     }
     return expr;