onnx wrapper

7 files changed, 448 insertions, 0 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index 3e81521550a..fe9d9985c6a 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -52,6 +52,7 @@ vespa_define_module(
     src/tests/tensor/direct_dense_tensor_builder
     src/tests/tensor/direct_sparse_tensor_builder
     src/tests/tensor/index_lookup_table
+    src/tests/tensor/onnx_wrapper
     src/tests/tensor/tensor_add_operation
     src/tests/tensor/tensor_address
     src/tests/tensor/tensor_conformance
diff --git a/eval/src/tests/tensor/onnx_wrapper/CMakeLists.txt b/eval/src/tests/tensor/onnx_wrapper/CMakeLists.txt
new file mode 100644
index 00000000000..9c92f75476c
--- /dev/null
+++ b/eval/src/tests/tensor/onnx_wrapper/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(eval_onnx_wrapper_test_app TEST
+    SOURCES
+    onnx_wrapper_test.cpp
+    DEPENDS
+    vespaeval
+    GTest::GTest
+)
+vespa_add_test(NAME eval_onnx_wrapper_test_app COMMAND eval_onnx_wrapper_test_app)
diff --git a/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp b/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp
new file mode 100644
index 00000000000..07e065f9e39
--- /dev/null
+++ b/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp
@@ -0,0 +1,119 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/eval/tensor/dense/onnx_wrapper.h>
+#include <vespa/eval/tensor/dense/dense_tensor_view.h>
+#include <vespa/eval/tensor/dense/mutable_dense_tensor_view.h>
+#include <vespa/vespalib/util/stringfmt.h>
+#include <vespa/vespalib/gtest/gtest.h>
+
+using namespace vespalib::eval;
+using namespace vespalib::tensor;
+
+using vespalib::make_string_short::fmt;
+
+std::string get_source_dir() {
+    const char *dir = getenv("SOURCE_DIRECTORY");
+    return (dir ? dir : ".");
+}
+std::string source_dir = get_source_dir();
+std::string vespa_dir = source_dir + "/" + "../../../../..";
+std::string simple_model = vespa_dir + "/" + "model-integration/src/test/models/onnx/simple/simple.onnx";
+
+vespalib::string to_str(const std::vector<size_t> &dim_sizes) {
+    vespalib::string res;
+    for (size_t dim_size: dim_sizes) {
+        if (dim_size == 0) {
+            res += "[]";
+        } else {
+            res += fmt("[%zu]", dim_size);
+        }
+    }
+    return res;
+}
+
+vespalib::string to_str(OnnxWrapper::TensorInfo::ElementType element_type) {
+    if (element_type == OnnxWrapper::TensorInfo::ElementType::FLOAT) {
+        return "float";
+    }
+    if (element_type == OnnxWrapper::TensorInfo::ElementType::DOUBLE) {
+        return "double";
+    }
+    return "???";
+}
+
+void dump_info(const char *ctx, const std::vector<OnnxWrapper::TensorInfo> &info) {
+    fprintf(stderr, "%s:\n", ctx);
+    for (size_t i = 0; i < info.size(); ++i) {
+        fprintf(stderr, "  %s[%zu]: '%s' %s%s\n", ctx, i, info[i].name.c_str(),
+                to_str(info[i].elements).c_str(),to_str(info[i].dimensions).c_str());
+    }
+}
+
+TEST(OnnxWrapperTest, onnx_model_can_be_inspected)
+{
+    OnnxWrapper wrapper(simple_model, OnnxWrapper::Optimize::DISABLE);
+    dump_info("inputs", wrapper.inputs());
+    dump_info("outputs", wrapper.outputs());
+    ASSERT_EQ(wrapper.inputs().size(), 3);
+    ASSERT_EQ(wrapper.outputs().size(), 1);
+    //-------------------------------------------------------------------------
+    EXPECT_EQ(       wrapper.inputs()[0].name,        "query_tensor");
+    EXPECT_EQ(to_str(wrapper.inputs()[0].dimensions), "[1][4]");
+    EXPECT_EQ(to_str(wrapper.inputs()[0].elements),   "float");
+    //-------------------------------------------------------------------------
+    EXPECT_EQ(       wrapper.inputs()[1].name,        "attribute_tensor");
+    EXPECT_EQ(to_str(wrapper.inputs()[1].dimensions), "[4][1]");
+    EXPECT_EQ(to_str(wrapper.inputs()[1].elements),   "float");
+    //-------------------------------------------------------------------------
+    EXPECT_EQ(       wrapper.inputs()[2].name,        "bias_tensor");
+    EXPECT_EQ(to_str(wrapper.inputs()[2].dimensions), "[1][1]");
+    EXPECT_EQ(to_str(wrapper.inputs()[2].elements),   "float");
+    //-------------------------------------------------------------------------
+    EXPECT_EQ(       wrapper.outputs()[0].name,        "output");
+    EXPECT_EQ(to_str(wrapper.outputs()[0].dimensions), "[1][1]");
+    EXPECT_EQ(to_str(wrapper.outputs()[0].elements),   "float");
+}
+
+TEST(OnnxWrapperTest, onnx_model_can_be_evaluated)
+{
+    OnnxWrapper wrapper(simple_model, OnnxWrapper::Optimize::ENABLE);
+
+    ValueType query_type = ValueType::from_spec("tensor<float>(a[1],b[4])");
+    std::vector<float> query_values({1.0, 2.0, 3.0, 4.0});
+    DenseTensorView query(query_type, TypedCells(query_values));
+    EXPECT_TRUE(wrapper.inputs()[0].is_compatible(query_type));
+    EXPECT_FALSE(wrapper.inputs()[1].is_compatible(query_type));
+    EXPECT_FALSE(wrapper.inputs()[2].is_compatible(query_type));
+
+    ValueType attribute_type = ValueType::from_spec("tensor<float>(a[4],b[1])");
+    std::vector<float> attribute_values({5.0, 6.0, 7.0, 8.0});
+    DenseTensorView attribute(attribute_type, TypedCells(attribute_values));
+    EXPECT_FALSE(wrapper.inputs()[0].is_compatible(attribute_type));
+    EXPECT_TRUE(wrapper.inputs()[1].is_compatible(attribute_type));
+    EXPECT_FALSE(wrapper.inputs()[2].is_compatible(attribute_type));
+
+    ValueType bias_type = ValueType::from_spec("tensor<float>(a[1],b[1])");
+    std::vector<float> bias_values({9.0});
+    DenseTensorView bias(bias_type, TypedCells(bias_values));
+    EXPECT_FALSE(wrapper.inputs()[0].is_compatible(bias_type));
+    EXPECT_FALSE(wrapper.inputs()[1].is_compatible(bias_type));
+    EXPECT_TRUE(wrapper.inputs()[2].is_compatible(bias_type));
+
+    MutableDenseTensorView output(wrapper.outputs()[0].make_compatible_type());
+    EXPECT_EQ(output.fast_type().to_spec(), "tensor<float>(d0[1],d1[1])");
+
+    OnnxWrapper::Params params;
+    params.bind(0, query);
+    params.bind(1, attribute);
+    params.bind(2, bias);
+    auto result = wrapper.eval(params);
+
+    EXPECT_EQ(result.num_values(), 1);
+    result.get(0, output);
+    auto cells = output.cellsRef();
+    EXPECT_EQ(cells.type, ValueType::CellType::FLOAT);
+    EXPECT_EQ(cells.size, 1);
+    EXPECT_EQ(cells.get(0), 79.0);
+}
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/eval/src/vespa/eval/CMakeLists.txt b/eval/src/vespa/eval/CMakeLists.txt
index c28643e605e..04f151f7ced 100644
--- a/eval/src/vespa/eval/CMakeLists.txt
+++ b/eval/src/vespa/eval/CMakeLists.txt
@@ -12,6 +12,7 @@ vespa_add_library(vespaeval
     $<TARGET_OBJECTS:eval_tensor_sparse>
     INSTALL lib64
     DEPENDS
+    onnxruntime
     ${VESPA_LLVM_LIB}
 )
 
diff --git a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
index c4b8138148c..b4e849a1dde 100644
--- a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
+++ b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
@@ -30,6 +30,7 @@ vespa_add_library(eval_tensor_dense OBJECT
     dense_xw_product_function.cpp
     index_lookup_table.cpp
     mutable_dense_tensor_view.cpp
+    onnx_wrapper.cpp
     typed_cells.cpp
     typed_dense_tensor_builder.cpp
     vector_from_doubles_function.cpp
diff --git a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp
new file mode 100644
index 00000000000..fa0379473c9
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp
@@ -0,0 +1,233 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "onnx_wrapper.h"
+#include <vespa/eval/eval/value_type.h>
+#include "dense_tensor_view.h"
+#include "mutable_dense_tensor_view.h"
+#include <vespa/vespalib/util/arrayref.h>
+#include <vespa/vespalib/util/stringfmt.h>
+#include <assert.h>
+#include <cmath>
+#include <stdlib.h>
+#include <stdio.h>
+
+using vespalib::eval::ValueType;
+using vespalib::make_string_short::fmt;
+
+namespace vespalib::tensor {
+
+namespace {
+
+ValueType::CellType as_cell_type(OnnxWrapper::TensorInfo::ElementType type) {
+    if (type == OnnxWrapper::TensorInfo::ElementType::FLOAT) {
+        return ValueType::CellType::FLOAT;
+    }
+    if (type == OnnxWrapper::TensorInfo::ElementType::DOUBLE) {
+        return ValueType::CellType::DOUBLE;
+    }
+    abort();
+}
+
+auto convert_optimize(OnnxWrapper::Optimize optimize) {
+    if (optimize == OnnxWrapper::Optimize::ENABLE) {
+        return ORT_ENABLE_ALL;
+    } else {
+        assert(optimize == OnnxWrapper::Optimize::DISABLE);
+        return ORT_DISABLE_ALL;
+    }
+}
+
+class OnnxString {
+private:
+    static Ort::AllocatorWithDefaultOptions _alloc;
+    char *_str;
+    void cleanup() {
+        if (_str != nullptr) {
+            _alloc.Free(_str);
+            _str = nullptr;
+        }
+    }
+    OnnxString(char *str) : _str(str) {}
+public:
+    OnnxString(const OnnxString &rhs) = delete;
+    OnnxString &operator=(const OnnxString &rhs) = delete;
+    OnnxString(OnnxString &&rhs) : _str(rhs._str) {
+        rhs._str = nullptr;
+    }
+    OnnxString &operator=(OnnxString &&rhs) {
+        cleanup();
+        _str = rhs._str;
+        rhs._str = nullptr;
+        return *this;
+    }
+    const char *get() const { return _str; }
+    ~OnnxString() { cleanup(); }
+    static OnnxString get_input_name(const Ort::Session &session, size_t idx) {
+        return OnnxString(session.GetInputName(idx, _alloc));
+    }
+    static OnnxString get_output_name(const Ort::Session &session, size_t idx) {
+        return OnnxString(session.GetOutputName(idx, _alloc));
+    }
+};
+Ort::AllocatorWithDefaultOptions OnnxString::_alloc;
+
+std::vector<size_t> make_dimensions(const std::vector<int64_t> &shape) {
+    std::vector<size_t> result;
+    for (int64_t size: shape) {
+        result.push_back(std::max(size, 0L));
+    } 
+    return result;
+}
+
+OnnxWrapper::TensorInfo::ElementType make_element_type(ONNXTensorElementDataType element_type) {
+    if (element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) {
+        return OnnxWrapper::TensorInfo::ElementType::FLOAT;
+    } else if (element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE) {
+        return OnnxWrapper::TensorInfo::ElementType::DOUBLE;
+    } else {
+        return OnnxWrapper::TensorInfo::ElementType::UNKNOWN;
+    }
+}
+
+OnnxWrapper::TensorInfo make_tensor_info(const OnnxString &name, const Ort::TypeInfo &type_info) {
+    auto tensor_info = type_info.GetTensorTypeAndShapeInfo();
+    auto shape = tensor_info.GetShape();
+    auto element_type = tensor_info.GetElementType();
+    return OnnxWrapper::TensorInfo{vespalib::string(name.get()), make_dimensions(shape), make_element_type(element_type)};
+}
+
+}
+
+bool
+OnnxWrapper::TensorInfo::is_compatible(const eval::ValueType &type) const
+{
+    if ((elements == ElementType::UNKNOWN) || dimensions.empty()) {
+        return false;
+    }
+    if (type.cell_type() != as_cell_type(elements)) {
+        return false;
+    }
+    if (type.dimensions().size() != dimensions.size()) {
+        return false;
+    }
+    for (size_t i = 0; i < dimensions.size(); ++i) {
+        if (type.dimensions()[i].size != dimensions[i]) {
+            return false;
+        }
+    }
+    return true;
+}
+
+eval::ValueType
+OnnxWrapper::TensorInfo::make_compatible_type() const
+{
+    if ((elements == ElementType::UNKNOWN) || dimensions.empty()) {
+        return ValueType::error_type();
+    }
+    std::vector<ValueType::Dimension> dim_list;
+    for (size_t dim_size: dimensions) {
+        if ((dim_size == 0) || (dim_list.size() > 9)) {
+            return ValueType::error_type();
+        }
+        dim_list.emplace_back(fmt("d%zu", dim_list.size()), dim_size);
+    }
+    return ValueType::tensor_type(std::move(dim_list), as_cell_type(elements));
+}
+
+OnnxWrapper::TensorInfo::~TensorInfo() = default;
+
+OnnxWrapper::Shared::Shared()
+    : _env(ORT_LOGGING_LEVEL_WARNING, "vespa-onnx-wrapper")
+{
+}
+
+void
+OnnxWrapper::Params::bind(size_t idx, const DenseTensorView &src)
+{
+    assert(idx == values.size());
+    std::vector<int64_t> dim_sizes;
+    for (const auto &dim: src.fast_type().dimensions()) {
+        dim_sizes.push_back(dim.size);
+    }
+    auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
+    if (src.fast_type().cell_type() == ValueType::CellType::FLOAT) {
+        // NB: create requires non-const input
+        auto cells = unconstify(src.cellsRef().typify<float>());
+        values.push_back(Ort::Value::CreateTensor<float>(memory_info, cells.begin(), cells.size(), dim_sizes.data(), dim_sizes.size()));
+    } else if (src.fast_type().cell_type() == ValueType::CellType::DOUBLE) {
+        // NB: create requires non-const input
+        auto cells = unconstify(src.cellsRef().typify<double>());
+        values.push_back(Ort::Value::CreateTensor<double>(memory_info, cells.begin(), cells.size(), dim_sizes.data(), dim_sizes.size()));
+    }
+}
+
+void
+OnnxWrapper::Result::get(size_t idx, MutableDenseTensorView &dst)
+{
+    assert(values[idx].IsTensor());
+    auto meta = values[idx].GetTensorTypeAndShapeInfo();
+    if (dst.fast_type().cell_type() == ValueType::CellType::FLOAT) {
+        assert(meta.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT);
+        ConstArrayRef<float> cells(values[idx].GetTensorMutableData<float>(), meta.GetElementCount());
+        dst.setCells(TypedCells(cells));
+    } else if (dst.fast_type().cell_type() == ValueType::CellType::DOUBLE) {
+        assert(meta.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE);
+        ConstArrayRef<double> cells(values[idx].GetTensorMutableData<double>(), meta.GetElementCount());
+        dst.setCells(TypedCells(cells));
+    }
+}
+
+OnnxWrapper::Shared &
+OnnxWrapper::Shared::get() {
+    static Shared shared;
+    return shared;
+}
+
+void
+OnnxWrapper::extract_meta_data()
+{
+    Ort::AllocatorWithDefaultOptions allocator;
+    size_t num_inputs = _session.GetInputCount();
+    for (size_t i = 0; i < num_inputs; ++i) {
+        _inputs.push_back(make_tensor_info(OnnxString::get_input_name(_session, i), _session.GetInputTypeInfo(i)));
+    }
+    size_t num_outputs = _session.GetOutputCount();
+    for (size_t i = 0; i < num_outputs; ++i) {
+        _outputs.push_back(make_tensor_info(OnnxString::get_output_name(_session, i), _session.GetOutputTypeInfo(i)));
+    }
+    for (const auto &input: _inputs) {
+        _input_name_refs.push_back(input.name.c_str());
+    }
+    for (const auto &output: _outputs) {
+        _output_name_refs.push_back(output.name.c_str());
+    }
+}
+
+OnnxWrapper::OnnxWrapper(const vespalib::string &model_file, Optimize optimize)
+    : _shared(Shared::get()),
+      _options(),
+      _session(nullptr),
+      _inputs(),
+      _outputs(),
+      _input_name_refs(),
+      _output_name_refs()
+{
+    _options.SetIntraOpNumThreads(1);
+    _options.SetInterOpNumThreads(1);
+    _options.SetGraphOptimizationLevel(convert_optimize(optimize));
+    _session = Ort::Session(_shared.env(), model_file.c_str(), _options);
+    extract_meta_data();
+}
+
+OnnxWrapper::~OnnxWrapper() = default;
+
+OnnxWrapper::Result
+OnnxWrapper::eval(const Params &params)
+{
+    assert(params.values.size() == _inputs.size());
+    Ort::RunOptions run_opts(nullptr);
+    return Result(_session.Run(run_opts, _input_name_refs.data(), params.values.data(), _inputs.size(),
+                               _output_name_refs.data(), _outputs.size()));
+}
+
+}
diff --git a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h
new file mode 100644
index 00000000000..67a64f2d318
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h
@@ -0,0 +1,84 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <onnxruntime/onnxruntime_cxx_api.h>
+#include <vespa/vespalib/stllike/string.h>
+#include <vespa/eval/eval/value_type.h>
+#include <vector>
+
+namespace vespalib::tensor {
+
+class DenseTensorView;
+class MutableDenseTensorView;
+
+/**
+ * Wrapper around an ONNX model handeled by onnxruntime.
+ **/
+class OnnxWrapper {
+public:
+    // model optimization
+    enum class Optimize { ENABLE, DISABLE };
+
+    // information about a single input or output tensor
+    struct TensorInfo {
+        enum class ElementType { FLOAT, DOUBLE, UNKNOWN };
+        vespalib::string name;
+        std::vector<size_t> dimensions;
+        ElementType elements;
+        bool is_compatible(const eval::ValueType &type) const;
+        eval::ValueType make_compatible_type() const;
+        ~TensorInfo();
+    };
+
+    // used to build model parameters
+    class Params {
+        friend class OnnxWrapper;
+    private:
+        std::vector<Ort::Value> values;
+    public:
+        Params() : values() {}
+        void bind(size_t idx, const DenseTensorView &src);
+    };
+
+    // used to inspect model results
+    class Result {
+        friend class OnnxWrapper;
+    private:
+        std::vector<Ort::Value> values;
+        Result(std::vector<Ort::Value> values_in) : values(std::move(values_in)) {}
+    public:
+        size_t num_values() const { return values.size(); }
+        void get(size_t idx, MutableDenseTensorView &dst);
+    };
+
+private:
+    // common stuff shared between model sessions
+    class Shared {
+    private:
+        Ort::Env _env;
+        Shared();
+    public:
+        static Shared &get();
+        Ort::Env &env() { return _env; }
+    };
+
+    Shared                   &_shared;
+    Ort::SessionOptions       _options;
+    Ort::Session              _session;
+    std::vector<TensorInfo>   _inputs;
+    std::vector<TensorInfo>   _outputs;
+    std::vector<const char *> _input_name_refs;
+    std::vector<const char *> _output_name_refs;
+
+    void extract_meta_data();
+
+public:
+    OnnxWrapper(const vespalib::string &model_file, Optimize optimize);
+    ~OnnxWrapper();
+    const std::vector<TensorInfo> &inputs() const { return _inputs; }
+    const std::vector<TensorInfo> &outputs() const { return _outputs; }
+    Result eval(const Params &params); // NB: Run requires non-const session
+};
+
+}