diff options
author | Arne Juul <arnej@verizonmedia.com> | 2021-03-26 13:40:54 +0000 |
---|---|---|
committer | Arne Juul <arnej@verizonmedia.com> | 2021-04-06 12:41:23 +0000 |
commit | de080e77a2fd2730e37fcd7c7ae02f9bf36c43c6 (patch) | |
tree | 6e1a26138762d0256d8ff28edbee138fe22939f4 /eval/src | |
parent | aaba78b6f46590c4326fefe205e51c9a026644f1 (diff) |
add optimization for join between bfloat16 tensor and number
Diffstat (limited to 'eval/src')
-rw-r--r-- | eval/src/tests/instruction/join_with_number/join_with_number_function_test.cpp | 9 | ||||
-rw-r--r-- | eval/src/vespa/eval/instruction/join_with_number_function.cpp | 71 |
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 ¶m = 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 ¶m = 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); } } |