co-locate optimization with implementation

9 files changed, 113 insertions, 136 deletions

diff --git a/eval/src/tests/tensor/dense_tensor_function_optimizer/dense_tensor_function_optimizer_test.cpp b/eval/src/tests/tensor/dense_tensor_function_optimizer/dense_tensor_function_optimizer_test.cpp
index c09b218508e..269a1e265c5 100644
--- a/eval/src/tests/tensor/dense_tensor_function_optimizer/dense_tensor_function_optimizer_test.cpp
+++ b/eval/src/tests/tensor/dense_tensor_function_optimizer/dense_tensor_function_optimizer_test.cpp
@@ -3,7 +3,6 @@
 #include <vespa/vespalib/testkit/test_kit.h>
 #include <vespa/eval/tensor/dense/dense_dot_product_function.h>
 #include <vespa/eval/tensor/dense/dense_xw_product_function.h>
-#include <vespa/eval/tensor/dense/dense_tensor_function_optimizer.h>
 #include <vespa/eval/eval/operation.h>
 
 using namespace vespalib::eval;
@@ -23,7 +22,7 @@ optimizeDotProduct(const vespalib::string &lhsType,
                                          inject(ValueType::from_spec(rhsType), 3, stash),
                                          Mul::f, stash),
                                     Aggr::SUM, {}, stash);
-    return DenseTensorFunctionOptimizer::optimize(reduceNode, stash);
+    return DenseDotProductFunction::optimize(reduceNode, stash);
 }
 
 void assertParam(const TensorFunction &node, size_t expect_idx) {
@@ -66,7 +65,7 @@ optimizeXWProduct(const vespalib::string &lhsType,
                                          inject(ValueType::from_spec(rhsType), 3, stash),
                                          Mul::f, stash),
                                     Aggr::SUM, {dim}, stash);
-    return DenseTensorFunctionOptimizer::optimize(reduceNode, stash);
+    return DenseXWProductFunction::optimize(reduceNode, stash);
 }
 
 void
diff --git a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
index 9477b36463a..0873d0341fa 100644
--- a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
+++ b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
@@ -7,7 +7,8 @@
 #include "serialization/typed_binary_format.h"
 #include "dense/dense_tensor.h"
 #include "dense/dense_tensor_builder.h"
-#include "dense/dense_tensor_function_optimizer.h"
+#include "dense/dense_dot_product_function.h"
+#include "dense/dense_xw_product_function.h"
 #include <vespa/eval/eval/value.h>
 #include <vespa/eval/eval/tensor_spec.h>
 #include <vespa/eval/eval/simple_tensor_engine.h>
@@ -216,7 +217,8 @@ DefaultTensorEngine::optimize(const TensorFunction &expr, Stash &stash) const
     }
     while (!nodes.empty()) {
         const Child &child = nodes.back();
-        child.set(DenseTensorFunctionOptimizer::optimize(child.get(), stash));
+        child.set(DenseDotProductFunction::optimize(child.get(), stash));
+        child.set(DenseXWProductFunction::optimize(child.get(), stash));
         nodes.pop_back();
     }
     return root.get();
diff --git a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
index 1fa839ca4b2..3bd81ff8df3 100644
--- a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
+++ b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
@@ -8,7 +8,6 @@ vespa_add_library(eval_tensor_dense OBJECT
     dense_tensor_address_combiner.cpp
     dense_tensor_builder.cpp
     dense_tensor_cells_iterator.cpp
-    dense_tensor_function_optimizer.cpp
     dense_tensor_view.cpp
     dense_tensor_reduce.cpp
     mutable_dense_tensor_view.cpp
diff --git a/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.cpp
index 58c2db85627..7985b59da56 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.cpp
@@ -3,25 +3,23 @@
 #include "dense_dot_product_function.h"
 #include "dense_tensor.h"
 #include "dense_tensor_view.h"
+#include <vespa/eval/eval/operation.h>
 #include <vespa/eval/eval/value.h>
 #include <vespa/eval/tensor/tensor.h>
 
 namespace vespalib::tensor {
 
 using CellsRef = DenseTensorView::CellsRef;
-
-DenseDotProductFunction::DenseDotProductFunction(const eval::TensorFunction &lhs_in,
-                                                 const eval::TensorFunction &rhs_in)
-    : eval::tensor_function::Op2(eval::ValueType::double_type(), lhs_in, rhs_in),
-      _hwAccelerator(hwaccelrated::IAccelrated::getAccelrator())
-{
-}
+using eval::ValueType;
+using eval::TensorFunction;
+using eval::as;
+using eval::Aggr;
+using namespace eval::tensor_function;
+using namespace eval::operation;
 
 namespace {
 
-CellsRef
-getCellsRef(const eval::Value &value)
-{
+CellsRef getCellsRef(const eval::Value &value) {
     const DenseTensorView &denseTensor = static_cast<const DenseTensorView &>(value);
     return denseTensor.cellsRef();
 }
@@ -35,6 +33,24 @@ void my_op(eval::InterpretedFunction::State &state, uint64_t param) {
     state.pop_pop_push(state.stash.create<eval::DoubleValue>(result));
 }
 
+bool is1dDenseTensor(const ValueType &type) {
+    return (type.is_dense() && (type.dimensions().size() == 1));
+}
+
+bool isDenseDotProduct(const ValueType &res, const ValueType &lhsType, const ValueType &rhsType) {
+    return (res.is_double() &&
+            is1dDenseTensor(lhsType) &&
+            is1dDenseTensor(rhsType) &&
+            (lhsType.dimensions()[0].name == rhsType.dimensions()[0].name));
+}
+
+} // namespace vespalib::tensor::<unnamed>
+
+DenseDotProductFunction::DenseDotProductFunction(const eval::TensorFunction &lhs_in,
+                                                 const eval::TensorFunction &rhs_in)
+    : eval::tensor_function::Op2(eval::ValueType::double_type(), lhs_in, rhs_in),
+      _hwAccelerator(hwaccelrated::IAccelrated::getAccelrator())
+{
 }
 
 eval::InterpretedFunction::Instruction
@@ -43,5 +59,21 @@ DenseDotProductFunction::compile_self(Stash &) const
     return eval::InterpretedFunction::Instruction(my_op, (uint64_t)(_hwAccelerator.get()));
 }
 
+const TensorFunction &
+DenseDotProductFunction::optimize(const eval::TensorFunction &expr, Stash &stash)
+{
+    const Reduce *reduce = as<Reduce>(expr);
+    if (reduce && (reduce->aggr() == Aggr::SUM)) {
+        const Join *join = as<Join>(reduce->child());
+        if (join && (join->function() == Mul::f)) {
+            const TensorFunction &lhs = join->lhs();
+            const TensorFunction &rhs = join->rhs();
+            if (isDenseDotProduct(expr.result_type(), lhs.result_type(), rhs.result_type())) {
+                return stash.create<DenseDotProductFunction>(lhs, rhs);
+            }
+        }
+    }
+    return expr;
 }
 
+} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.h b/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.h
index 1c38654c014..eec7448f041 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_dot_product_function.h
@@ -19,7 +19,7 @@ public:
     DenseDotProductFunction(const eval::TensorFunction &lhs_in,
                             const eval::TensorFunction &rhs_in);
     eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
 };
 
-}
-
+} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_function_optimizer.cpp b/eval/src/vespa/eval/tensor/dense/dense_tensor_function_optimizer.cpp
deleted file mode 100644
index 48eefa3eed4..00000000000
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_function_optimizer.cpp
+++ /dev/null
@@ -1,91 +0,0 @@
-// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-
-#include "dense_dot_product_function.h"
-#include "dense_xw_product_function.h"
-#include "dense_tensor_function_optimizer.h"
-#include <vespa/eval/eval/operation.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::tensor {
-
-namespace {
-
-bool is1dDenseTensor(const ValueType &type) {
-    return (type.is_dense() && (type.dimensions().size() == 1));
-}
-
-bool isConcreteDenseTensor(const ValueType &type, size_t d) {
-    return (type.is_dense() && (type.dimensions().size() == d) && !type.is_abstract());
-}
-
-bool isDenseDotProduct(const ValueType &res, const ValueType &lhsType, const ValueType &rhsType) {
-    return (res.is_double() &&
-            is1dDenseTensor(lhsType) &&
-            is1dDenseTensor(rhsType) &&
-            (lhsType.dimensions()[0].name == rhsType.dimensions()[0].name));
-}
-
-bool isDenseXWProduct(const ValueType &res, const ValueType &vec, const ValueType &mat) {
-    if (isConcreteDenseTensor(res, 1) &&
-        isConcreteDenseTensor(vec, 1) &&
-        isConcreteDenseTensor(mat, 2))
-    {
-        size_t res_idx = mat.dimension_index(res.dimensions()[0].name);
-        size_t vec_idx = mat.dimension_index(vec.dimensions()[0].name);
-        size_t npos = ValueType::Dimension::npos;
-        if ((res_idx != npos) && (vec_idx != npos) && (res_idx != vec_idx)) {
-            return ((mat.dimensions()[res_idx].size == res.dimensions()[0].size) &&
-                    (mat.dimensions()[vec_idx].size == vec.dimensions()[0].size));
-        }
-    }
-    return false;
-}
-
-const TensorFunction &createDenseXWProduct(const ValueType &res, const TensorFunction &vec, const TensorFunction &mat, Stash &stash) {
-    bool common_is_inner = (mat.result_type().dimension_index(vec.result_type().dimensions()[0].name) == 1);
-    return stash.create<DenseXWProductFunction>(res, vec, mat,
-                                                vec.result_type().dimensions()[0].size,
-                                                res.dimensions()[0].size,
-                                                common_is_inner);
-}
-
-struct InnerProductFunctionOptimizer
-{
-    static const TensorFunction &optimize(const TensorFunction &expr, Stash &stash) {
-        const Reduce *reduce = as<Reduce>(expr);
-        if (reduce && (reduce->aggr() == Aggr::SUM)) {
-            const ValueType &result_type = reduce->result_type();
-            const Join *join = as<Join>(reduce->child());
-            if (join && (join->function() == Mul::f)) {
-                const TensorFunction &lhs = join->lhs();
-                const TensorFunction &rhs = join->rhs();
-                if (isDenseDotProduct(result_type, lhs.result_type(), rhs.result_type())) {
-                    return stash.create<DenseDotProductFunction>(lhs, rhs);
-                }
-                if (isDenseXWProduct(result_type, lhs.result_type(), rhs.result_type())) {
-                    return createDenseXWProduct(result_type, lhs, rhs, stash);
-                }
-                if (isDenseXWProduct(result_type, rhs.result_type(), lhs.result_type())) {
-                    return createDenseXWProduct(result_type, rhs, lhs, stash);
-                }
-            }
-        }
-        return expr;
-    }
-};
-
-}
-
-const TensorFunction &
-DenseTensorFunctionOptimizer::optimize(const eval::TensorFunction &expr, Stash &stash)
-{
-    return InnerProductFunctionOptimizer::optimize(expr, stash);
-}
-
-}
-
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_function_optimizer.h b/eval/src/vespa/eval/tensor/dense/dense_tensor_function_optimizer.h
deleted file mode 100644
index 2478447ca48..00000000000
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_function_optimizer.h
+++ /dev/null
@@ -1,21 +0,0 @@
-// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-
-#pragma once
-
-#include <vespa/eval/eval/tensor_function.h>
-
-namespace vespalib { class Stash; }
-
-namespace vespalib::tensor {
-
-/**
- * Class that recognizes calculations over dense tensors (in tensor function intermediate representation)
- * and optimizes this into an explicit tensor function.
- */
-struct DenseTensorFunctionOptimizer
-{
-    static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
-};
-
-}
-
diff --git a/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.cpp
index 91b16b0fd7d..fa07028ef27 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.cpp
@@ -4,17 +4,24 @@
 #include "dense_tensor.h"
 #include "dense_tensor_view.h"
 #include <vespa/eval/eval/value.h>
+#include <vespa/eval/eval/operation.h>
 #include <vespa/eval/tensor/tensor.h>
 #include <vespa/vespalib/util/exceptions.h>
 #include <assert.h>
 
 namespace vespalib::tensor {
 
+using CellsRef = DenseTensorView::CellsRef;
+using eval::ValueType;
+using eval::TensorFunction;
+using eval::as;
+using eval::Aggr;
+using namespace eval::tensor_function;
+using namespace eval::operation;
+
 namespace {
 
-DenseTensorView::CellsRef
-getCellsRef(const eval::Value &value)
-{
+CellsRef getCellsRef(const eval::Value &value) {
     const DenseTensorView &denseTensor = static_cast<const DenseTensorView &>(value);
     return denseTensor.cellsRef();
 }
@@ -54,8 +61,8 @@ template <bool commonDimensionInnermost>
 void my_op(eval::InterpretedFunction::State &state, uint64_t param) {
     DenseXWProductFunction::Self *self = (DenseXWProductFunction::Self *)(param);
 
-    DenseTensorView::CellsRef vectorCells = getCellsRef(state.peek(1));
-    DenseTensorView::CellsRef matrixCells = getCellsRef(state.peek(0));
+    CellsRef vectorCells = getCellsRef(state.peek(1));
+    CellsRef matrixCells = getCellsRef(state.peek(0));
 
     ArrayRef<double> outputCells = state.stash.create_array<double>(self->_resultSize);
 
@@ -67,6 +74,34 @@ void my_op(eval::InterpretedFunction::State &state, uint64_t param) {
     state.pop_pop_push(state.stash.create<DenseTensorView>(self->_resultType, outputCells));
 }
 
+bool isConcreteDenseTensor(const ValueType &type, size_t d) {
+    return (type.is_dense() && (type.dimensions().size() == d) && !type.is_abstract());
+}
+
+bool isDenseXWProduct(const ValueType &res, const ValueType &vec, const ValueType &mat) {
+    if (isConcreteDenseTensor(res, 1) &&
+        isConcreteDenseTensor(vec, 1) &&
+        isConcreteDenseTensor(mat, 2))
+    {
+        size_t res_idx = mat.dimension_index(res.dimensions()[0].name);
+        size_t vec_idx = mat.dimension_index(vec.dimensions()[0].name);
+        size_t npos = ValueType::Dimension::npos;
+        if ((res_idx != npos) && (vec_idx != npos) && (res_idx != vec_idx)) {
+            return ((mat.dimensions()[res_idx].size == res.dimensions()[0].size) &&
+                    (mat.dimensions()[vec_idx].size == vec.dimensions()[0].size));
+        }
+    }
+    return false;
+}
+
+const TensorFunction &createDenseXWProduct(const ValueType &res, const TensorFunction &vec, const TensorFunction &mat, Stash &stash) {
+    bool common_is_inner = (mat.result_type().dimension_index(vec.result_type().dimensions()[0].name) == 1);
+    return stash.create<DenseXWProductFunction>(res, vec, mat,
+                                                vec.result_type().dimensions()[0].size,
+                                                res.dimensions()[0].size,
+                                                common_is_inner);
+}
+
 } // namespace vespalib::tensor::<unnamed>
 
 DenseXWProductFunction::Self::Self(const eval::ValueType &resultType,
@@ -98,4 +133,25 @@ DenseXWProductFunction::compile_self(Stash &stash) const
     return eval::InterpretedFunction::Instruction(op, (uint64_t)(&self));
 }
 
+const TensorFunction &
+DenseXWProductFunction::optimize(const eval::TensorFunction &expr, Stash &stash)
+{
+        const Reduce *reduce = as<Reduce>(expr);
+        if (reduce && (reduce->aggr() == Aggr::SUM)) {
+            const ValueType &result_type = reduce->result_type();
+            const Join *join = as<Join>(reduce->child());
+            if (join && (join->function() == Mul::f)) {
+                const TensorFunction &lhs = join->lhs();
+                const TensorFunction &rhs = join->rhs();
+                if (isDenseXWProduct(result_type, lhs.result_type(), rhs.result_type())) {
+                    return createDenseXWProduct(result_type, lhs, rhs, stash);
+                }
+                if (isDenseXWProduct(result_type, rhs.result_type(), lhs.result_type())) {
+                    return createDenseXWProduct(result_type, rhs, lhs, stash);
+                }
+            }
+        }
+        return expr;
+}
+
 } // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.h b/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.h
index 4fc5e209948..221c3891775 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_xw_product_function.h
@@ -49,7 +49,8 @@ public:
     bool matrixHasCommonDimensionInnermost() const { return _commonDimensionInnermost; }
 
     eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
+    static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
 };
 
-}
+} // namespace vespalib::tensor