Merge pull request #16894 from vespa-engine/arnej/map-prep-celltypes

Arnej/map prep celltypes

8 files changed, 80 insertions, 80 deletions

diff --git a/eval/src/tests/instruction/generic_map/generic_map_test.cpp b/eval/src/tests/instruction/generic_map/generic_map_test.cpp
index bfa7154968d..54ba0ea2a31 100644
--- a/eval/src/tests/instruction/generic_map/generic_map_test.cpp
+++ b/eval/src/tests/instruction/generic_map/generic_map_test.cpp
@@ -35,11 +35,13 @@ const std::vector<GenSpec> map_layouts = {
 
 TensorSpec perform_generic_map(const TensorSpec &a, map_fun_t func, const ValueBuilderFactory &factory)
 {
+    Stash stash;
     auto lhs = value_from_spec(a, factory);
     auto res_type = lhs->type().map();
-    auto my_op = GenericMap::make_instruction(res_type, lhs->type(), func);
+    auto my_op = GenericMap::make_instruction(res_type, lhs->type(), func, stash);
     InterpretedFunction::EvalSingle single(factory, my_op);
-    return spec_from_value(single.eval(std::vector<Value::CREF>({*lhs})));
+    const auto & v = single.eval(std::vector<Value::CREF>({*lhs}));
+    return spec_from_value(v);
 }
 
 void test_generic_map_with(const ValueBuilderFactory &factory) {
diff --git a/eval/src/tests/instruction/mixed_map_function/mixed_map_function_test.cpp b/eval/src/tests/instruction/mixed_map_function/mixed_map_function_test.cpp
index dae147da216..6a6d569e8a4 100644
--- a/eval/src/tests/instruction/mixed_map_function/mixed_map_function_test.cpp
+++ b/eval/src/tests/instruction/mixed_map_function/mixed_map_function_test.cpp
@@ -16,6 +16,7 @@ const ValueBuilderFactory &prod_factory = FastValueBuilderFactory::get();
 EvalFixture::ParamRepo make_params() {
     return EvalFixture::ParamRepo()
         .add("a", GenSpec(1.5))
+        .add_mutable("@a", GenSpec(1.5))
         .add("b", GenSpec(2.5))
         .add_variants("sparse", GenSpec().map("x", {"a"}))
         .add_variants("mixed", GenSpec().map("x", {"a"}).idx("y", 5))
@@ -23,7 +24,7 @@ EvalFixture::ParamRepo make_params() {
 }
 EvalFixture::ParamRepo param_repo = make_params();
 
-void verify_optimized(const vespalib::string &expr, bool inplace) {
+void verify_optimized(const vespalib::string &expr) {
     EvalFixture slow_fixture(prod_factory, expr, param_repo, false);
     EvalFixture fixture(prod_factory, expr, param_repo, true, true);
     EXPECT_EQ(fixture.result(), EvalFixture::ref(expr, param_repo));
@@ -31,13 +32,10 @@ void verify_optimized(const vespalib::string &expr, bool inplace) {
     auto info = fixture.find_all<MixedMapFunction>();
     ASSERT_EQ(info.size(), 1u);
     EXPECT_TRUE(info[0]->result_is_mutable());
-    EXPECT_EQ(info[0]->inplace(), inplace);
+    EXPECT_EQ(info[0]->inplace(), true);
     ASSERT_EQ(fixture.num_params(), 1);
-    if (inplace) {
-        EXPECT_EQ(fixture.get_param(0), fixture.result());
-    } else {
-        EXPECT_TRUE(!(fixture.get_param(0) == fixture.result()));
-    }
+    // inplace check:
+    EXPECT_EQ(fixture.get_param(0), fixture.result());
 }
 
 void verify_not_optimized(const vespalib::string &expr) {
@@ -47,36 +45,29 @@ void verify_not_optimized(const vespalib::string &expr) {
     EXPECT_EQ(fixture.result(), slow_fixture.result());
     auto info = fixture.find_all<MixedMapFunction>();
     EXPECT_TRUE(info.empty());
+    EXPECT_TRUE(!(fixture.get_param(0) == fixture.result()));
 }
 
-TEST(MapTest, dense_map_is_optimized) {
-    verify_optimized("map(x5y3,f(x)(x+10))", false);
-    verify_optimized("map(x5y3_f,f(x)(x+10))", false);
-}
-
-TEST(MapTest, simple_dense_map_can_be_inplace) {
-    verify_optimized("map(@x5y3,f(x)(x+10))", true);
-    verify_optimized("map(@x5y3_f,f(x)(x+10))", true);
+TEST(MapTest, dense_map_can_be_optimized) {
+    verify_not_optimized("map(x5y3,f(x)(x+10))");
+    verify_not_optimized("map(x5y3_f,f(x)(x+10))");
+    verify_optimized("map(@x5y3,f(x)(x+10))");
+    verify_optimized("map(@x5y3_f,f(x)(x+10))");
 }
 
 TEST(MapTest, scalar_map_is_not_optimized) {
     verify_not_optimized("map(a,f(x)(x+10))");
+    verify_not_optimized("map(@a,f(x)(x+10))");
 }
 
-TEST(MapTest, sparse_map_is_optimized) {
-    verify_optimized("map(sparse,f(x)(x+10))", false);
-}
-
-TEST(MapTest, sparse_map_can_be_inplace) {
-    verify_optimized("map(@sparse,f(x)(x+10))", true);
-}
-
-TEST(MapTest, mixed_map_is_optimized) {
-    verify_optimized("map(mixed,f(x)(x+10))", false);
+TEST(MapTest, sparse_map_can_be_optimized) {
+    verify_not_optimized("map(sparse,f(x)(x+10))");
+    verify_optimized("map(@sparse,f(x)(x+10))");
 }
 
-TEST(MapTest, mixed_map_can_be_inplace) {
-    verify_optimized("map(@mixed,f(x)(x+10))", true);
+TEST(MapTest, mixed_map_can_be_optimized) {
+    verify_not_optimized("map(mixed,f(x)(x+10))");
+    verify_optimized("map(@mixed,f(x)(x+10))");
 }
 
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/eval/src/vespa/eval/eval/tensor_function.cpp b/eval/src/vespa/eval/eval/tensor_function.cpp
index fd8b20bfed0..c0cd7280212 100644
--- a/eval/src/vespa/eval/eval/tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/tensor_function.cpp
@@ -158,9 +158,9 @@ Reduce::visit_self(vespalib::ObjectVisitor &visitor) const
 //-----------------------------------------------------------------------------
 
 Instruction
-Map::compile_self(const ValueBuilderFactory &, Stash &) const
+Map::compile_self(const ValueBuilderFactory &, Stash &stash) const
 {
-    return instruction::GenericMap::make_instruction(result_type(), child().result_type(), _function);
+    return instruction::GenericMap::make_instruction(result_type(), child().result_type(), _function, stash);
 }
 
 void
diff --git a/eval/src/vespa/eval/eval/test/reference_operations.cpp b/eval/src/vespa/eval/eval/test/reference_operations.cpp
index 58c90de65a2..34452c1b2ae 100644
--- a/eval/src/vespa/eval/eval/test/reference_operations.cpp
+++ b/eval/src/vespa/eval/eval/test/reference_operations.cpp
@@ -164,7 +164,7 @@ TensorSpec ReferenceOperations::join(const TensorSpec &in_a, const TensorSpec &i
 
 TensorSpec ReferenceOperations::map(const TensorSpec &in_a, map_fun_t func) {
     auto a = in_a.normalize();
-    ValueType res_type = ValueType::from_spec(a.type());
+    ValueType res_type = ValueType::from_spec(a.type()).map();
     TensorSpec result(res_type.to_spec());
     if (res_type.is_error()) {
         return result;
diff --git a/eval/src/vespa/eval/instruction/generic_map.cpp b/eval/src/vespa/eval/instruction/generic_map.cpp
index 4f6780c2276..33f7d565157 100644
--- a/eval/src/vespa/eval/instruction/generic_map.cpp
+++ b/eval/src/vespa/eval/instruction/generic_map.cpp
@@ -1,8 +1,9 @@
 // Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "generic_map.h"
-#include <vespa/eval/eval/value.h>
 #include <vespa/eval/eval/inline_operation.h>
+#include <vespa/eval/eval/value.h>
+#include <vespa/eval/eval/wrap_param.h>
 #include <vespa/vespalib/util/stash.h>
 #include <vespa/vespalib/util/typify.h>
 #include <cassert>
@@ -11,54 +12,70 @@ namespace vespalib::eval::instruction {
 
 using State = InterpretedFunction::State;
 using Instruction = InterpretedFunction::Instruction;
+using vespalib::eval::tensor_function::unwrap_param;
+using vespalib::eval::tensor_function::wrap_param;
 
-namespace {
+struct MapParam {
+    const ValueType res_type;
+    const map_fun_t function;
+    MapParam(const ValueType &r, map_fun_t f) : res_type(r), function(f) {}
+};
 
-uint64_t to_param(map_fun_t value) { return (uint64_t)value; }
-map_fun_t to_map_fun(uint64_t param) { return (map_fun_t)param; }
+namespace {
 
-template <typename CT, typename Func>
+template <typename ICT, typename OCT, typename Func>
 void my_generic_map_op(State &state, uint64_t param_in) {
-    Func function(to_map_fun(param_in));
+    const auto &param = unwrap_param<MapParam>(param_in);
+    Func function(param.function);
     const Value &a = state.peek(0);
-    auto input_cells = a.cells().typify<CT>();
-    auto output_cells = state.stash.create_uninitialized_array<CT>(input_cells.size());
+    auto input_cells = a.cells().typify<ICT>();
+    auto output_cells = state.stash.create_uninitialized_array<OCT>(input_cells.size());
     auto pos = output_cells.begin();
-    for (CT value : input_cells) {
-        *pos++ = (CT) function(value);
+    if constexpr (std::is_same<ICT,OCT>::value) {
+        apply_op1_vec(pos, input_cells.begin(), output_cells.size(), function);
+    } else {
+        for (ICT value : input_cells) {
+            *pos++ = (OCT) function(value);
+        }
+        assert(pos == output_cells.end());
     }
-    assert(pos == output_cells.end());
-    Value &result_ref = state.stash.create<ValueView>(a.type(), a.index(), TypedCells(output_cells));
+    Value &result_ref = state.stash.create<ValueView>(param.res_type, a.index(), TypedCells(output_cells));
     state.pop_push(result_ref);
 }
 
 template <typename Func>
 void my_double_map_op(State &state, uint64_t param_in) {
-    Func fun(to_map_fun(param_in));
+    const auto &param = unwrap_param<MapParam>(param_in);
+    Func fun(param.function);
     state.pop_push(state.stash.create<DoubleValue>(fun(state.peek(0).as_double())));
 }
 
 struct SelectGenericMapOp {
-    template <typename CT, typename Func> static auto invoke(const ValueType &type) {
-        if (type.is_double()) {
-            assert((std::is_same_v<CT,double>));
+    template <typename ICM, typename Func> static auto invoke() {
+        if constexpr (ICM::value.is_scalar) {
             return my_double_map_op<Func>;
+        } else {
+            using ICT = CellValueType<ICM::value.cell_type>;
+            using OCT = CellValueType<ICM::value.map().cell_type>;
+            return my_generic_map_op<ICT, OCT, Func>;
         }
-        return my_generic_map_op<CT, Func>;
     }
 };
 
 } // namespace <unnamed>
 
-using MapTypify = TypifyValue<TypifyCellType,operation::TypifyOp1>;
+using MapTypify = TypifyValue<TypifyCellMeta,operation::TypifyOp1>;
  
 InterpretedFunction::Instruction
 GenericMap::make_instruction(const ValueType &result_type,
-                             const ValueType &input_type, map_fun_t function)
+                             const ValueType &input_type,
+                             map_fun_t function,
+                             Stash &stash)
 {
+    const auto &param = stash.create<MapParam>(result_type, function);
     assert(result_type == input_type.map());
-    auto op = typify_invoke<2,MapTypify,SelectGenericMapOp>(input_type.cell_type(), function, input_type);
-    return Instruction(op, to_param(function));
+    auto op = typify_invoke<2,MapTypify,SelectGenericMapOp>(input_type.cell_meta(), function);
+    return Instruction(op, wrap_param<MapParam>(param));
 }
 
 } // namespace
diff --git a/eval/src/vespa/eval/instruction/generic_map.h b/eval/src/vespa/eval/instruction/generic_map.h
index f6c01760898..6878546018d 100644
--- a/eval/src/vespa/eval/instruction/generic_map.h
+++ b/eval/src/vespa/eval/instruction/generic_map.h
@@ -15,7 +15,9 @@ using map_fun_t = operation::op1_t;
 struct GenericMap {
     static InterpretedFunction::Instruction
     make_instruction(const ValueType &result_type,
-                     const ValueType &input_type, map_fun_t function);
+                     const ValueType &input_type,
+                     map_fun_t function,
+                     Stash &stash);
 };
 
 } // namespace
diff --git a/eval/src/vespa/eval/instruction/mixed_map_function.cpp b/eval/src/vespa/eval/instruction/mixed_map_function.cpp
index 69917ae94e0..23d5809fbca 100644
--- a/eval/src/vespa/eval/instruction/mixed_map_function.cpp
+++ b/eval/src/vespa/eval/instruction/mixed_map_function.cpp
@@ -8,47 +8,32 @@
 
 namespace vespalib::eval {
 
-using vespalib::ArrayRef;
+using operation::TypifyOp1;
+using tensor_function::map_fun_t;
 
-using namespace operation;
-using namespace tensor_function;
-
-using op_function = InterpretedFunction::op_function;
 using Instruction = InterpretedFunction::Instruction;
 using State = InterpretedFunction::State;
 
 namespace {
 
-template <typename CT, bool inplace>
-ArrayRef<CT> make_dst_cells(ConstArrayRef<CT> src_cells, Stash &stash) {
-    if (inplace) {
-        return unconstify(src_cells);
-    } else {
-        return stash.create_uninitialized_array<CT>(src_cells.size());
-    }
-}
-
-template <typename CT, typename Fun, bool inplace>
-void my_simple_map_op(State &state, uint64_t param) {
+template <typename CT, typename Fun>
+void my_inplace_map_op(State &state, uint64_t param) {
     Fun my_fun((map_fun_t)param);
     auto const &child = state.peek(0);
     auto src_cells = child.cells().typify<CT>();
-    auto dst_cells = make_dst_cells<CT, inplace>(src_cells, state.stash);
+    auto dst_cells = unconstify(src_cells);
     apply_op1_vec(dst_cells.begin(), src_cells.begin(), dst_cells.size(), my_fun);
-    if (!inplace) {
-        state.pop_push(state.stash.create<ValueView>(child.type(), child.index(), TypedCells(dst_cells)));
-    }
 }
 
 //-----------------------------------------------------------------------------
 
 struct MyGetFun {
-    template <typename R1, typename R2, typename R3> static auto invoke() {
-        return my_simple_map_op<R1, R2, R3::value>;
+    template <typename CT, typename Fun> static auto invoke() {
+        return my_inplace_map_op<CT, Fun>;
     }
 };
 
-using MyTypify = TypifyValue<TypifyCellType,TypifyOp1,TypifyBool>;
+using MyTypify = TypifyValue<TypifyCellType,TypifyOp1>;
 
 } // namespace vespalib::eval::<unnamed>
 
@@ -66,7 +51,7 @@ MixedMapFunction::~MixedMapFunction() = default;
 Instruction
 MixedMapFunction::compile_self(const ValueBuilderFactory &, Stash &) const
 {
-    auto op = typify_invoke<3,MyTypify,MyGetFun>(result_type().cell_type(), function(), inplace());
+    auto op = typify_invoke<2,MyTypify,MyGetFun>(result_type().cell_type(), function());
     static_assert(sizeof(uint64_t) == sizeof(function()));
     return Instruction(op, (uint64_t)(function()));
 }
@@ -75,7 +60,10 @@ const TensorFunction &
 MixedMapFunction::optimize(const TensorFunction &expr, Stash &stash)
 {
     if (auto map = as<Map>(expr)) {
-        if (! map->child().result_type().is_double()) {
+        if ((map->result_type() == map->child().result_type())
+            && (! map->child().result_type().is_double())
+            && map->child().result_is_mutable())
+        {
             return stash.create<MixedMapFunction>(map->result_type(), map->child(), map->function());
         }
     }
diff --git a/eval/src/vespa/eval/instruction/mixed_map_function.h b/eval/src/vespa/eval/instruction/mixed_map_function.h
index 3f3b821d15e..2e7ca7c9422 100644
--- a/eval/src/vespa/eval/instruction/mixed_map_function.h
+++ b/eval/src/vespa/eval/instruction/mixed_map_function.h
@@ -7,7 +7,7 @@
 namespace vespalib::eval {
 
 /**
- * Tensor function optimizing map operations on tensors.
+ * Tensor function optimizing in-place map operations on tensors.
  **/
 class MixedMapFunction : public tensor_function::Map
 {