aboutsummaryrefslogtreecommitdiffstats
path: root/eval/src/tests/instruction/generic_peek/generic_peek_test.cpp
blob: 7f25edd50c8f7172ea8bfd5041c47aacf8f0f1f7 (plain) (blame)
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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

#include <vespa/eval/eval/simple_value.h>
#include <vespa/eval/eval/fast_value.h>
#include <vespa/eval/eval/tensor_function.h>
#include <vespa/eval/eval/value_codec.h>
#include <vespa/eval/instruction/generic_peek.h>
#include <vespa/eval/eval/interpreted_function.h>
#include <vespa/eval/eval/test/reference_operations.h>
#include <vespa/eval/eval/test/gen_spec.h>
#include <vespa/vespalib/util/stringfmt.h>
#include <vespa/vespalib/util/overload.h>
#include <vespa/vespalib/gtest/gtest.h>
#include <vespa/vespalib/stllike/asciistream.h>
#include <stdlib.h>
#include <variant>

using namespace vespalib;
using namespace vespalib::eval;
using namespace vespalib::eval::instruction;
using namespace vespalib::eval::test;

using vespalib::make_string_short::fmt;

GenSpec G() { return GenSpec(); }

const std::vector<GenSpec> peek_layouts = {
    G().idx("x", 4),
    G().idx("x", 4).idx("y", 5),
    G().idx("x", 4).idx("y", 5).idx("z", 3),
    G().map("x", {"-1","0","2"}),
    G().map("x", {"-1","0","2"}).map("y", {"-2","0","1"}).map("z", {"-2","-1","0","1","2"}),
    G().idx("x", 4).map("y", {"-2","0","1"}).idx("z", 3),
    G().map("x", {"-1","0","2"}).idx("y", 5).map("z", {"-2","-1","0","1","2"})
};

using PeekSpec = GenericPeek::SpecMap;

TensorSpec reference_peek(const TensorSpec &param, const PeekSpec &spec) {
    std::vector<TensorSpec> children;
    children.push_back(param);
    PeekSpec with_indexes;
    for (const auto & [dim_name, label_or_child] : spec) {
        const vespalib::string &dim = dim_name;
        std::visit(vespalib::overload
                   {
                       [&](const TensorSpec::Label &label) {
                           with_indexes.emplace(dim, label);
                       },
                       [&](const size_t &child_value) {
                           // here, label_or_child is a size_t specifying the value
                           // we pretend a child produced
                           size_t child_idx = children.size();
                           TensorSpec child("double");
                           // (but cast to signed first, to allow labels like the string "-2")
                           child.add({}, ssize_t(child_value));
                           children.push_back(child);
                           with_indexes.emplace(dim, child_idx);
                       }
                   }, label_or_child);
    }
    return ReferenceOperations::peek(with_indexes, children);
}

TensorSpec perform_generic_peek(const TensorSpec &a, const ValueType &result_type,
                                PeekSpec spec, const ValueBuilderFactory &factory)
{
    auto param = value_from_spec(a, factory);
    EXPECT_FALSE(param->type().is_error());
    EXPECT_FALSE(result_type.is_error());
    Stash stash;
    std::vector<Value::CREF> my_stack;
    my_stack.push_back(*param);
    size_t child_idx = 1;
    for (auto & [dim_name, label_or_child] : spec) {
        if (std::holds_alternative<size_t>(label_or_child)) {
            // here, label_or_child is a size_t specifying the value
            // this child should produce (but cast to signed first,
            // to allow negative values)
            ssize_t child_value = std::get<size_t>(label_or_child);
            my_stack.push_back(stash.create<DoubleValue>(child_value));
            // overwrite label_or_child, now it should be the index of
            // the child for make_instruction
            label_or_child = child_idx++;
        }
    }
    auto my_op = GenericPeek::make_instruction(result_type, param->type(), spec, factory, stash);
    InterpretedFunction::EvalSingle single(factory, my_op);
    return spec_from_value(single.eval(my_stack));
}

TensorSpec tensor_function_peek(const TensorSpec &a, const ValueType &result_type,
                                PeekSpec spec, const ValueBuilderFactory &factory)
{
    Stash stash;
    auto param = value_from_spec(a, factory);
    EXPECT_FALSE(param->type().is_error());
    EXPECT_FALSE(result_type.is_error());
    std::vector<Value::CREF> my_stack;
    my_stack.push_back(*param);
    const auto &func_double = tensor_function::inject(ValueType::double_type(), 1, stash);
    std::map<vespalib::string, std::variant<TensorSpec::Label, TensorFunction::CREF>> func_spec;
    for (auto & [dim_name, label_or_child] : spec) {
        if (std::holds_alternative<size_t>(label_or_child)) {
            // here, label_or_child is a size_t specifying the value
            // this child should produce (but cast to signed first,
            // to allow negative values)
            ssize_t child_value = std::get<size_t>(label_or_child);
            my_stack.push_back(stash.create<DoubleValue>(double(child_value)));
            func_spec.emplace(dim_name, func_double);
        } else {
            auto label = std::get<TensorSpec::Label>(label_or_child);
            func_spec.emplace(dim_name, label);
        }
    }
    const auto &func_param = tensor_function::inject(param->type(), 0, stash);
    const auto &peek_node = tensor_function::peek(func_param, func_spec, stash);
    auto my_op = peek_node.compile_self(factory, stash);
    InterpretedFunction::EvalSingle single(factory, my_op);
    return spec_from_value(single.eval(my_stack));
}

vespalib::string to_str(const PeekSpec &spec) {
    vespalib::asciistream os;
    os << "{ ";
    for (const auto & [dim, label_or_index] : spec) {
        os << dim << " : ";
        if (std::holds_alternative<size_t>(label_or_index)) {
            os << "[" << ssize_t(std::get<size_t>(label_or_index)) << "] ";
        } else {
            auto label = std::get<TensorSpec::Label>(label_or_index);
            if (label.is_mapped()) {
                os << "'" << label.name << "' ";
            } else {
                os << "(" << ssize_t(label.index) << ") ";
            }
        }
    }
    os << "}";
    return os.str();
}

void verify_peek_equal(const TensorSpec &input,
                       const PeekSpec &spec,
                       const ValueBuilderFactory &factory)
{
    ValueType param_type = ValueType::from_spec(input.type());
    std::vector<vespalib::string> reduce_dims;
    for (const auto & [dim_name, ignored] : spec) {
        reduce_dims.push_back(dim_name);
    }
    if (reduce_dims.empty()) return;
    ValueType result_type = param_type.peek(reduce_dims);
    auto expect = reference_peek(input, spec);
    SCOPED_TRACE(fmt("peek input: %s\n  peek spec: %s\n  peek result %s\n",
                     input.to_string().c_str(),
                     to_str(spec).c_str(),
                     expect.to_string().c_str()));
    auto actual = perform_generic_peek(input, result_type, spec, factory);
    EXPECT_EQ(actual, expect);
    auto from_func = tensor_function_peek(input, result_type, spec, factory);
    EXPECT_EQ(from_func, expect);
}

void fill_dims_and_check(const TensorSpec &input,
                         PeekSpec spec,
                         std::vector<ValueType::Dimension> dimensions,
                         const ValueBuilderFactory &factory)
{
    if (dimensions.empty()) {
        verify_peek_equal(input, spec, factory);
        return;
    }
    auto dim = dimensions.back();
    dimensions.pop_back();
    fill_dims_and_check(input, spec, dimensions, factory);
    for (int64_t label_value : {-2, -1, 0, 1, 3}) {
        if (dim.is_indexed()) {
            size_t index = label_value;
            if (index >= dim.size) continue;
            TensorSpec::Label label(index);
            spec.insert_or_assign(dim.name, label);
        } else {
            TensorSpec::Label label(make_string("%" PRId64, label_value));
            spec.insert_or_assign(dim.name, label);
        }
        fill_dims_and_check(input, spec, dimensions, factory);
    }
    for (int64_t child_value : {-2, -1, 0, 1, 3}) {
        spec.insert_or_assign(dim.name, size_t(child_value));
        fill_dims_and_check(input, spec, dimensions, factory);
    }
}

void test_generic_peek_with(const ValueBuilderFactory &factory) {
    for (const auto &layout : peek_layouts) {
        for (CellType ct : CellTypeUtils::list_types()) {
            auto input = layout.cpy().cells(ct);
            ValueType input_type = input.type();
            const auto &dims = input_type.dimensions();
            PeekSpec spec;
            fill_dims_and_check(input, spec, dims, factory);
        }
    }
}

TEST(GenericPeekTest, generic_peek_works_for_simple_values) {
    test_generic_peek_with(SimpleValueBuilderFactory::get());
}

TEST(GenericPeekTest, generic_peek_works_for_fast_values) {
    test_generic_peek_with(FastValueBuilderFactory::get());
}

GTEST_MAIN_RUN_ALL_TESTS()