1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
|
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
#include "dense_lambda_peek_function.h"
#include "index_lookup_table.h"
#include <vespa/eval/eval/value.h>
namespace vespalib::eval {
using namespace tensor_function;
namespace {
struct Self {
const ValueType &result_type;
IndexLookupTable::Token::UP table_token;
Self(const ValueType &result_type_in, const Function &function)
: result_type(result_type_in),
table_token(IndexLookupTable::create(function, result_type_in))
{
assert(table_token->get().size() == result_type.dense_subspace_size());
}
};
template <typename DST_CT, typename SRC_CT>
void my_lambda_peek_op(InterpretedFunction::State &state, uint64_t param) {
const auto &self = unwrap_param<Self>(param);
const std::vector<uint32_t> &lookup_table = self.table_token->get();
auto src_cells = state.peek(0).cells().typify<SRC_CT>();
ArrayRef<DST_CT> dst_cells = state.stash.create_uninitialized_array<DST_CT>(lookup_table.size());
DST_CT *dst = &dst_cells[0];
for (uint32_t idx: lookup_table) {
*dst++ = src_cells[idx];
}
state.pop_push(state.stash.create<DenseValueView>(self.result_type, TypedCells(dst_cells)));
}
struct MyLambdaPeekOp {
template <typename DST_CT, typename SRC_CT>
static auto invoke() { return my_lambda_peek_op<DST_CT, SRC_CT>; }
};
} // namespace <unnamed>
DenseLambdaPeekFunction::DenseLambdaPeekFunction(const ValueType &result_type,
const TensorFunction &child,
std::shared_ptr<Function const> idx_fun)
: Op1(result_type, child),
_idx_fun(std::move(idx_fun))
{
}
DenseLambdaPeekFunction::~DenseLambdaPeekFunction() = default;
InterpretedFunction::Instruction
DenseLambdaPeekFunction::compile_self(const ValueBuilderFactory &, Stash &stash) const
{
const Self &self = stash.create<Self>(result_type(), *_idx_fun);
using MyTypify = TypifyCellType;
auto op = typify_invoke<2,MyTypify,MyLambdaPeekOp>(result_type().cell_type(), child().result_type().cell_type());
assert(child().result_type().is_dense());
return InterpretedFunction::Instruction(op, wrap_param<Self>(self));
}
vespalib::string
DenseLambdaPeekFunction::idx_fun_dump() const {
return _idx_fun->dump_as_lambda();
}
} // namespace
|