Add tensor function compiler for dot product between two 1d dense tensors.

Currently, the function is executing the dot product using the tensor::Tensor API.

13 files changed, 264 insertions, 0 deletions

diff --git a/vespalib/CMakeLists.txt b/vespalib/CMakeLists.txt
index 2dfbe88f977..b6d07ca2cf1 100644
--- a/vespalib/CMakeLists.txt
+++ b/vespalib/CMakeLists.txt
@@ -75,6 +75,7 @@ vespa_define_module(
     src/tests/tensor/sparse_tensor_builder
     src/tests/tensor/dense_tensor_address_combiner
     src/tests/tensor/dense_tensor_builder
+    src/tests/tensor/dense_tensor_function_compiler
     src/tests/tensor/tensor_address
     src/tests/tensor/tensor_conformance
     src/tests/tensor/tensor_mapper
diff --git a/vespalib/src/tests/tensor/dense_tensor_function_compiler/CMakeLists.txt b/vespalib/src/tests/tensor/dense_tensor_function_compiler/CMakeLists.txt
new file mode 100644
index 00000000000..a34b39abb70
--- /dev/null
+++ b/vespalib/src/tests/tensor/dense_tensor_function_compiler/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(vespalib_dense_tensor_function_compiler_test_app TEST
+    SOURCES
+    dense_tensor_function_compiler_test.cpp
+    DEPENDS
+    vespalib
+    vespalib_vespalib_tensor
+)
+vespa_add_test(NAME vespalib_dense_tensor_function_compiler_test_app COMMAND vespalib_dense_tensor_function_compiler_test_app)
diff --git a/vespalib/src/tests/tensor/dense_tensor_function_compiler/FILES b/vespalib/src/tests/tensor/dense_tensor_function_compiler/FILES
new file mode 100644
index 00000000000..3c4ec2f1753
--- /dev/null
+++ b/vespalib/src/tests/tensor/dense_tensor_function_compiler/FILES
@@ -0,0 +1 @@
+dense_tensor_function_compiler_test.cpp
diff --git a/vespalib/src/tests/tensor/dense_tensor_function_compiler/dense_tensor_function_compiler_test.cpp b/vespalib/src/tests/tensor/dense_tensor_function_compiler/dense_tensor_function_compiler_test.cpp
new file mode 100644
index 00000000000..a49bea805c9
--- /dev/null
+++ b/vespalib/src/tests/tensor/dense_tensor_function_compiler/dense_tensor_function_compiler_test.cpp
@@ -0,0 +1,66 @@
+// Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/vespalib/testkit/test_kit.h>
+#include <vespa/vespalib/tensor/dense/dense_dot_product_function.h>
+#include <vespa/vespalib/tensor/dense/dense_tensor_function_compiler.h>
+
+using namespace vespalib::eval;
+using namespace vespalib::eval::operation;
+using namespace vespalib::eval::tensor_function;
+using namespace vespalib::tensor;
+
+template <typename T>
+const T *as(const TensorFunction &function) { return dynamic_cast<const T *>(&function); }
+
+TensorFunction::UP
+compileDotProduct(const vespalib::string &lhsType,
+                  const vespalib::string &rhsType)
+{
+    Node_UP reduceNode = reduce(apply(Mul(),
+                                      inject(ValueType::from_spec(lhsType), 1),
+                                      inject(ValueType::from_spec(rhsType), 3)),
+                                Add(), {});
+    return DenseTensorFunctionCompiler::compile(std::move(reduceNode));
+}
+
+void
+assertCompiledDotProduct(const vespalib::string &lhsType,
+                         const vespalib::string &rhsType)
+{
+    TensorFunction::UP func = compileDotProduct(lhsType, rhsType);
+    const DenseDotProductFunction *dotProduct = as<DenseDotProductFunction>(*func);
+    ASSERT_TRUE(dotProduct);
+    EXPECT_EQUAL(ValueType::from_spec(lhsType), dotProduct->lhsTensor().result_type);
+    EXPECT_EQUAL(1u, dotProduct->lhsTensor().tensor_id);
+    EXPECT_EQUAL(ValueType::from_spec(rhsType), dotProduct->rhsTensor().result_type);
+    EXPECT_EQUAL(3u, dotProduct->rhsTensor().tensor_id);
+}
+
+void
+assertNotCompiledDotProduct(const vespalib::string &lhsType,
+                            const vespalib::string &rhsType)
+{
+    TensorFunction::UP func = compileDotProduct(lhsType, rhsType);
+    const Reduce *reduce = as<Reduce>(*func);
+    EXPECT_TRUE(reduce);
+}
+
+TEST("require that dot product with compatible dimensions is compiled")
+{
+    TEST_DO(assertCompiledDotProduct("tensor(x[5])", "tensor(x[5])"));
+    TEST_DO(assertCompiledDotProduct("tensor(x[3])", "tensor(x[5])"));
+    TEST_DO(assertCompiledDotProduct("tensor(x[5])", "tensor(x[3])"));
+}
+
+TEST("require that dot product with incompatible dimensions is NOT compiled")
+{
+    TEST_DO(assertNotCompiledDotProduct("tensor(x[5])",      "tensor(y[5])"));
+    TEST_DO(assertNotCompiledDotProduct("tensor(y[5])",      "tensor(x[5])"));
+    TEST_DO(assertNotCompiledDotProduct("tensor(x[])",       "tensor(x[5])"));
+    TEST_DO(assertNotCompiledDotProduct("tensor(x[5])",      "tensor(x[])"));
+    TEST_DO(assertNotCompiledDotProduct("tensor(x[])",       "tensor(x[])"));
+    TEST_DO(assertNotCompiledDotProduct("tensor(x[5])",      "tensor(x[5],y[7])"));
+    TEST_DO(assertNotCompiledDotProduct("tensor(x[5],y[7])", "tensor(x[5],y[7])"));
+}
+
+TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/vespalib/src/vespa/vespalib/eval/operation.h b/vespalib/src/vespa/vespalib/eval/operation.h
index 242d338a539..bd730b5bce6 100644
--- a/vespalib/src/vespa/vespalib/eval/operation.h
+++ b/vespalib/src/vespa/vespalib/eval/operation.h
@@ -24,6 +24,12 @@ struct Operation {
     virtual ~Operation() {}
 };
 
+/**
+ * Simple typecasting utility.
+ */
+template <typename T>
+const T *as(const Operation &op) { return dynamic_cast<const T *>(&op); }
+
 //-----------------------------------------------------------------------------
 
 /**
diff --git a/vespalib/src/vespa/vespalib/eval/tensor_function.h b/vespalib/src/vespa/vespalib/eval/tensor_function.h
index 253042ebc45..37e17e64d8a 100644
--- a/vespalib/src/vespa/vespalib/eval/tensor_function.h
+++ b/vespalib/src/vespa/vespalib/eval/tensor_function.h
@@ -85,6 +85,12 @@ struct Node : public TensorFunction
 };
 using Node_UP = std::unique_ptr<Node>;
 
+/**
+ * Simple typecasting utility.
+ */
+template <typename T>
+const T *as(const Node &node) { return dynamic_cast<const T *>(&node); }
+
 struct Inject : Node {
     size_t tensor_id;
     Inject(const ValueType &result_type_in,
diff --git a/vespalib/src/vespa/vespalib/tensor/default_tensor_engine.cpp b/vespalib/src/vespa/vespalib/tensor/default_tensor_engine.cpp
index a2bc118c00b..7b218d80a85 100644
--- a/vespalib/src/vespa/vespalib/tensor/default_tensor_engine.cpp
+++ b/vespalib/src/vespa/vespalib/tensor/default_tensor_engine.cpp
@@ -7,6 +7,7 @@
 #include <vespa/vespalib/eval/operation_visitor.h>
 #include "tensor.h"
 #include "dense/dense_tensor_builder.h"
+#include "dense/dense_tensor_function_compiler.h"
 #include "default_tensor.h"
 
 namespace vespalib {
@@ -55,6 +56,12 @@ DefaultTensorEngine::to_spec(const Tensor &tensor) const
     return my_tensor.toSpec();
 }
 
+eval::TensorFunction::UP
+DefaultTensorEngine::compile(eval::tensor_function::Node_UP expr)
+{
+    return DenseTensorFunctionCompiler::compile(std::move(expr));
+}
+
 struct IsAddOperation : public eval::DefaultOperationVisitor {
     bool result = false;
     void visitDefault(const eval::Operation &) override {}
diff --git a/vespalib/src/vespa/vespalib/tensor/default_tensor_engine.h b/vespalib/src/vespa/vespalib/tensor/default_tensor_engine.h
index 7e1bd903626..ea8cf2836d0 100644
--- a/vespalib/src/vespa/vespalib/tensor/default_tensor_engine.h
+++ b/vespalib/src/vespa/vespalib/tensor/default_tensor_engine.h
@@ -24,6 +24,8 @@ public:
     vespalib::string to_string(const Tensor &tensor) const override;
     TensorSpec to_spec(const Tensor &tensor) const override;
 
+    virtual eval::TensorFunction::UP compile(eval::tensor_function::Node_UP expr);
+
     std::unique_ptr<Tensor> create(const TensorSpec &spec) const override;
     const Value &reduce(const Tensor &tensor, const BinaryOperation &op, const std::vector<vespalib::string> &dimensions, Stash &stash) const override;
     const Value &map(const UnaryOperation &op, const Tensor &a, Stash &stash) const override;
diff --git a/vespalib/src/vespa/vespalib/tensor/dense/CMakeLists.txt b/vespalib/src/vespa/vespalib/tensor/dense/CMakeLists.txt
index b7ca224fb90..4941e360e17 100644
--- a/vespalib/src/vespa/vespalib/tensor/dense/CMakeLists.txt
+++ b/vespalib/src/vespa/vespalib/tensor/dense/CMakeLists.txt
@@ -2,10 +2,12 @@
 vespa_add_library(vespalib_vespalib_tensor_dense OBJECT
     SOURCES
     direct_dense_tensor_builder.cpp
+    dense_dot_product_function.cpp
     dense_tensor.cpp
     dense_tensor_address_combiner.cpp
     dense_tensor_builder.cpp
     dense_tensor_cells_iterator.cpp
+    dense_tensor_function_compiler.cpp
     dense_tensor_view.cpp
     DEPENDS
 )
diff --git a/vespalib/src/vespa/vespalib/tensor/dense/dense_dot_product_function.cpp b/vespalib/src/vespa/vespalib/tensor/dense/dense_dot_product_function.cpp
new file mode 100644
index 00000000000..1be9dd64bec
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/tensor/dense/dense_dot_product_function.cpp
@@ -0,0 +1,38 @@
+// Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/fastos/fastos.h>
+#include "dense_dot_product_function.h"
+#include <vespa/vespalib/eval/value.h>
+#include <vespa/vespalib/tensor/tensor.h>
+
+namespace vespalib {
+namespace tensor {
+
+DenseDotProductFunction::DenseDotProductFunction(InjectUP lhsTensor_, InjectUP rhsTensor_)
+    : _lhsTensor(std::move(lhsTensor_)),
+      _rhsTensor(std::move(rhsTensor_))
+{
+}
+
+const eval::Value &
+DenseDotProductFunction::eval(const Input &input, Stash &stash) const
+{
+    const eval::Value &lhsValue = input.get_tensor(_lhsTensor->tensor_id);
+    const eval::Value &rhsValue = input.get_tensor(_rhsTensor->tensor_id);
+    assert(lhsValue.is_tensor());
+    assert(rhsValue.is_tensor());
+    const Tensor *lhsTensor = dynamic_cast<const Tensor *>(lhsValue.as_tensor());
+    const Tensor *rhsTensor = dynamic_cast<const Tensor *>(rhsValue.as_tensor());
+    assert(lhsTensor);
+    assert(rhsTensor);
+    double result;
+    if (_lhsTensor->result_type == _rhsTensor->result_type) {
+        result = lhsTensor->match(*rhsTensor)->sum();
+    } else {
+        result = lhsTensor->multiply(*rhsTensor)->sum();
+    }
+    return stash.create<eval::DoubleValue>(result);
+}
+
+} // namespace tensor
+} // namespace vespalib
diff --git a/vespalib/src/vespa/vespalib/tensor/dense/dense_dot_product_function.h b/vespalib/src/vespa/vespalib/tensor/dense/dense_dot_product_function.h
new file mode 100644
index 00000000000..b2b87c5b1d6
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/tensor/dense/dense_dot_product_function.h
@@ -0,0 +1,29 @@
+// Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/vespalib/eval/tensor_function.h>
+
+namespace vespalib {
+namespace tensor {
+
+/**
+ * Tensor function for a dot product between two 1-dimensional dense tensors.
+ */
+class DenseDotProductFunction : public eval::TensorFunction
+{
+private:
+    using InjectUP = std::unique_ptr<eval::tensor_function::Inject>;
+
+    InjectUP _lhsTensor;
+    InjectUP _rhsTensor;
+
+public:
+    DenseDotProductFunction(InjectUP lhsTensor_, InjectUP rhsTensor_);
+    const eval::tensor_function::Inject &lhsTensor() const { return *_lhsTensor; }
+    const eval::tensor_function::Inject &rhsTensor() const { return *_rhsTensor; }
+    virtual const eval::Value &eval(const Input &input, Stash &stash) const override;
+};
+
+} // namespace tensor
+} // namespace vespalib
diff --git a/vespalib/src/vespa/vespalib/tensor/dense/dense_tensor_function_compiler.cpp b/vespalib/src/vespa/vespalib/tensor/dense/dense_tensor_function_compiler.cpp
new file mode 100644
index 00000000000..ed3f937578b
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/tensor/dense/dense_tensor_function_compiler.cpp
@@ -0,0 +1,77 @@
+// Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/fastos/fastos.h>
+#include "dense_dot_product_function.h"
+#include "dense_tensor_function_compiler.h"
+#include <vespa/vespalib/eval/operation_visitor.h>
+#include <vespa/vespalib/eval/operation_visitor.h>
+#include <vespa/vespalib/test/insertion_operators.h>
+#include <iostream>
+
+using namespace vespalib::eval;
+using namespace vespalib::eval::tensor_function;
+using namespace vespalib::eval::operation;
+
+namespace vespalib {
+namespace tensor {
+
+namespace {
+
+template <typename T>
+bool    
+isType(const BinaryOperation &op)
+{
+    return (as<T>(op) != nullptr);
+}
+
+bool
+willReduceAllDimensions(const std::vector<vespalib::string> &dimensions)
+{
+    return (dimensions.empty() || (dimensions.size() == 1));
+}
+
+bool
+is1dBoundDenseTensor(const ValueType &type)
+{
+    return (type.is_dense() && (type.dimensions().size() == 1) && type.dimensions()[0].is_bound());
+}
+
+bool
+isCompatibleTensorsForDotProduct(const ValueType &lhsType, const ValueType &rhsType)
+{
+    return (is1dBoundDenseTensor(lhsType) &&
+            is1dBoundDenseTensor(rhsType) &&
+            (lhsType.dimensions()[0].name == rhsType.dimensions()[0].name));
+}
+
+struct DotProductFunctionCompiler
+{
+    static TensorFunction::UP compile(Node_UP expr) {
+        const Reduce *reduce = as<Reduce>(*expr);
+        if (reduce && isType<Add>(*reduce->op) && willReduceAllDimensions(reduce->dimensions)) {
+            const Apply *apply = as<Apply>(*reduce->tensor);
+            if (apply && isType<Mul>(*apply->op)) {
+                const Inject *lhsTensor = as<Inject>(*apply->lhs_tensor);
+                const Inject *rhsTensor = as<Inject>(*apply->rhs_tensor);
+                if (lhsTensor && rhsTensor &&
+                    isCompatibleTensorsForDotProduct(lhsTensor->result_type, rhsTensor->result_type))
+                {
+                    return std::make_unique<DenseDotProductFunction>(std::make_unique<Inject>(lhsTensor->result_type, lhsTensor->tensor_id),
+                                                                     std::make_unique<Inject>(rhsTensor->result_type, rhsTensor->tensor_id));
+                }
+            }
+        }
+        return std::move(expr);
+    }
+};
+
+}
+
+TensorFunction::UP
+DenseTensorFunctionCompiler::compile(Node_UP expr)
+{
+    return DotProductFunctionCompiler::compile(std::move(expr));
+}
+
+} // namespace tensor
+} // namespace vespalib
diff --git a/vespalib/src/vespa/vespalib/tensor/dense/dense_tensor_function_compiler.h b/vespalib/src/vespa/vespalib/tensor/dense/dense_tensor_function_compiler.h
new file mode 100644
index 00000000000..9d05d414bf1
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/tensor/dense/dense_tensor_function_compiler.h
@@ -0,0 +1,20 @@
+// Copyright 2016 Yahoo Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/vespalib/eval/tensor_function.h>
+
+namespace vespalib {
+namespace tensor {
+
+/**
+ * Class that recognizes calculations over dense tensors (in tensor function intermediate representation)
+ * and compiles this into an explicit tensor function.
+ */
+struct DenseTensorFunctionCompiler
+{
+    static eval::TensorFunction::UP compile(eval::tensor_function::Node_UP expr);
+};
+
+} // namespace tensor
+} // namespace vespalib