add optimization for join between bfloat16 tensor and number

2 files changed, 50 insertions, 30 deletions

diff --git a/eval/src/tests/instruction/join_with_number/join_with_number_function_test.cpp b/eval/src/tests/instruction/join_with_number/join_with_number_function_test.cpp
index a6486de6858..4bdc342b629 100644
--- a/eval/src/tests/instruction/join_with_number/join_with_number_function_test.cpp
+++ b/eval/src/tests/instruction/join_with_number/join_with_number_function_test.cpp
@@ -5,8 +5,8 @@
 #include <vespa/eval/eval/test/eval_fixture.h>
 #include <vespa/eval/eval/test/gen_spec.h>
 #include <vespa/eval/instruction/join_with_number_function.h>
-
 #include <vespa/vespalib/util/stringfmt.h>
+#include <vespa/vespalib/util/unwind_message.h>
 
 using namespace vespalib;
 using namespace vespalib::eval;
@@ -47,16 +47,17 @@ struct FunInfo {
 };
 
 void verify_optimized(const vespalib::string &expr, Primary primary, bool inplace) {
-    // fprintf(stderr, "%s\n", expr.c_str());
+    UNWIND_MSG("optimize %s", expr.c_str());
     const CellTypeSpace stable_types(CellTypeUtils::list_stable_types(), 2);
     FunInfo stable_details{primary, inplace};
     TEST_DO(EvalFixture::verify<FunInfo>(expr, {stable_details}, stable_types));
     const CellTypeSpace unstable_types(CellTypeUtils::list_unstable_types(), 2);
-    TEST_DO(EvalFixture::verify<FunInfo>(expr, {}, unstable_types));
+    FunInfo unstable_details{primary, false};
+    TEST_DO(EvalFixture::verify<FunInfo>(expr, {unstable_details}, unstable_types));
 }
 
 void verify_not_optimized(const vespalib::string &expr) {
-    // fprintf(stderr, "%s\n", expr.c_str());
+    UNWIND_MSG("not: %s", expr.c_str());
     CellTypeSpace all_types(CellTypeUtils::list_types(), 2);
     TEST_DO(EvalFixture::verify<FunInfo>(expr, {}, all_types));
 }
diff --git a/eval/src/vespa/eval/instruction/join_with_number_function.cpp b/eval/src/vespa/eval/instruction/join_with_number_function.cpp
index c574e3f8ad9..fbaca49997e 100644
--- a/eval/src/vespa/eval/instruction/join_with_number_function.cpp
+++ b/eval/src/vespa/eval/instruction/join_with_number_function.cpp
@@ -14,39 +14,55 @@ namespace vespalib::eval {
 
 using Instruction = InterpretedFunction::Instruction;
 using State = InterpretedFunction::State;
+using vespalib::eval::tensor_function::unwrap_param;
+using vespalib::eval::tensor_function::wrap_param;
 
 namespace {
 
-template <typename CT, bool inplace>
-ArrayRef<CT> make_dst_cells(ConstArrayRef<CT> src_cells, Stash &stash) {
-    if (inplace) {
+struct JoinWithNumberParam {
+    const ValueType res_type;
+    const join_fun_t function;
+    JoinWithNumberParam(const ValueType &r, join_fun_t f) : res_type(r), function(f) {}
+};
+
+template <typename ICT, typename OCT, bool inplace>
+ArrayRef<OCT> make_dst_cells(ConstArrayRef<ICT> src_cells, Stash &stash) {
+    if constexpr (inplace) {
+        static_assert(std::is_same_v<ICT,OCT>);
         return unconstify(src_cells);
     } else {
-        return stash.create_uninitialized_array<CT>(src_cells.size());
+        return stash.create_uninitialized_array<OCT>(src_cells.size());
     }
 }
 
-template <typename CT, typename Fun, bool inplace, bool swap>
-void my_number_join_op(State &state, uint64_t param) {
+template <typename ICT, typename OCT, typename Fun, bool inplace, bool swap>
+void my_number_join_op(State &state, uint64_t param_in) {
+    const auto &param = unwrap_param<JoinWithNumberParam>(param_in);
     using OP = typename std::conditional<swap,SwapArgs2<Fun>,Fun>::type;
-    OP my_op((join_fun_t)param);
+    OP my_op(param.function);
     const Value &tensor = state.peek(swap ? 0 : 1);
-    CT number = state.peek(swap ? 1 : 0).as_double();
-    auto src_cells = tensor.cells().typify<CT>();
-    auto dst_cells = make_dst_cells<CT, inplace>(src_cells, state.stash);
+    OCT number = state.peek(swap ? 1 : 0).as_double();
+    auto src_cells = tensor.cells().typify<ICT>();
+    auto dst_cells = make_dst_cells<ICT, OCT, inplace>(src_cells, state.stash);
     apply_op2_vec_num(dst_cells.begin(), src_cells.begin(), number, dst_cells.size(), my_op);
     if (inplace) {
         state.pop_pop_push(tensor);
     } else {
-        state.pop_pop_push(state.stash.create<ValueView>(tensor.type(), tensor.index(), TypedCells(dst_cells)));
+        state.pop_pop_push(state.stash.create<ValueView>(param.res_type, tensor.index(), TypedCells(dst_cells)));
     }
 }
 
 struct SelectJoinWithNumberOp {
-    template<typename CT, typename Fun,
-             typename InputIsMutable, typename NumberWasLeft>
+    template<typename CM, typename Fun,
+             typename Inplace, typename NumberWasLeft>
     static auto invoke() {
-        return my_number_join_op<CT, Fun, InputIsMutable::value, NumberWasLeft::value>;
+        constexpr CellMeta icm = CM::value;
+        constexpr CellMeta num(CellType::DOUBLE, true);
+        constexpr CellMeta ocm = CellMeta::join(icm, num); 
+        using ICT = CellValueType<icm.cell_type>;
+        using OCT = CellValueType<ocm.cell_type>;
+        constexpr bool inplace = (Inplace::value && std::is_same_v<ICT,OCT>);
+        return my_number_join_op<ICT, OCT, Fun, inplace, NumberWasLeft::value>;
     }
 };
 
@@ -64,22 +80,27 @@ JoinWithNumberFunction::~JoinWithNumberFunction() = default;
 bool
 JoinWithNumberFunction::inplace() const {
     if (_primary == Primary::LHS) {
-        return lhs().result_is_mutable();
+        return lhs().result_type() == result_type()
+            && lhs().result_is_mutable();
     } else {
-        return rhs().result_is_mutable();
+        return rhs().result_type() == result_type()
+            && rhs().result_is_mutable();
     }
 }
 
-using MyTypify = TypifyValue<TypifyCellType,vespalib::TypifyBool,operation::TypifyOp2>;
+using MyTypify = TypifyValue<TypifyCellMeta,vespalib::TypifyBool,operation::TypifyOp2>;
 
 InterpretedFunction::Instruction
-JoinWithNumberFunction::compile_self(const ValueBuilderFactory &, Stash &) const
+JoinWithNumberFunction::compile_self(const ValueBuilderFactory &, Stash &stash) const
 {
-    auto op = typify_invoke<4,MyTypify,SelectJoinWithNumberOp>(result_type().cell_type(),
+    const auto &param = stash.create<JoinWithNumberParam>(result_type(), _function);
+    auto input_type = (_primary == Primary::LHS) ? lhs().result_type() : rhs().result_type();
+    assert(result_type() == input_type.map());
+    auto op = typify_invoke<4,MyTypify,SelectJoinWithNumberOp>(input_type.cell_meta(),
                                                                _function,
                                                                inplace(),
                                                                (_primary == Primary::RHS));
-    return Instruction(op, (uint64_t)(_function));
+    return Instruction(op, wrap_param<JoinWithNumberParam>(param));
 }
 
 void
@@ -98,14 +119,12 @@ JoinWithNumberFunction::optimize(const TensorFunction &expr, Stash &stash)
             const ValueType &result_type = join->result_type();
             const TensorFunction &lhs = join->lhs();
             const TensorFunction &rhs = join->rhs();
-            if (lhs.result_type().is_double() &&
-                (result_type == rhs.result_type()))
-            {
+            if (lhs.result_type().is_double()) {
+                assert(result_type.dimensions() == rhs.result_type().dimensions());
                 return stash.create<JoinWithNumberFunction>(*join, true);
             }
-            if (rhs.result_type().is_double() &&
-                (result_type == lhs.result_type()))
-            {
+            if (rhs.result_type().is_double()) {
+                assert(result_type.dimensions() == lhs.result_type().dimensions());
                 return stash.create<JoinWithNumberFunction>(*join, false);
             }
         }