added dense remove dimension optimizer

11 files changed, 200 insertions, 5 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index faa94fd6719..f1f559a129d 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -29,6 +29,7 @@ vespa_define_module(
     src/tests/tensor/dense_fast_rename_optimizer
     src/tests/tensor/dense_inplace_join_function
     src/tests/tensor/dense_inplace_map_function
+    src/tests/tensor/dense_remove_dimension_optimizer
     src/tests/tensor/dense_replace_type_function
     src/tests/tensor/dense_tensor_address_combiner
     src/tests/tensor/dense_tensor_builder
diff --git a/eval/src/tests/tensor/dense_remove_dimension_optimizer/CMakeLists.txt b/eval/src/tests/tensor/dense_remove_dimension_optimizer/CMakeLists.txt
new file mode 100644
index 00000000000..c945bd31609
--- /dev/null
+++ b/eval/src/tests/tensor/dense_remove_dimension_optimizer/CMakeLists.txt
@@ -0,0 +1,8 @@
+# Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(eval_dense_remove_dimension_optimizer_test_app TEST
+    SOURCES
+    dense_remove_dimension_optimizer_test.cpp
+    DEPENDS
+    vespaeval
+)
+vespa_add_test(NAME eval_dense_remove_dimension_optimizer_test_app COMMAND eval_dense_remove_dimension_optimizer_test_app)
diff --git a/eval/src/tests/tensor/dense_remove_dimension_optimizer/dense_remove_dimension_optimizer_test.cpp b/eval/src/tests/tensor/dense_remove_dimension_optimizer/dense_remove_dimension_optimizer_test.cpp
new file mode 100644
index 00000000000..b87393c2980
--- /dev/null
+++ b/eval/src/tests/tensor/dense_remove_dimension_optimizer/dense_remove_dimension_optimizer_test.cpp
@@ -0,0 +1,83 @@
+// Copyright 2018 Yahoo Holdings. 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/eval/eval/tensor_function.h>
+#include <vespa/eval/eval/simple_tensor.h>
+#include <vespa/eval/eval/simple_tensor_engine.h>
+#include <vespa/eval/tensor/default_tensor_engine.h>
+#include <vespa/eval/tensor/dense/dense_replace_type_function.h>
+#include <vespa/eval/tensor/dense/dense_fast_rename_optimizer.h>
+#include <vespa/eval/tensor/dense/dense_tensor.h>
+#include <vespa/eval/eval/test/tensor_model.hpp>
+#include <vespa/eval/eval/test/eval_fixture.h>
+
+#include <vespa/vespalib/util/stringfmt.h>
+#include <vespa/vespalib/util/stash.h>
+
+using namespace vespalib;
+using namespace vespalib::eval;
+using namespace vespalib::eval::test;
+using namespace vespalib::tensor;
+using namespace vespalib::eval::tensor_function;
+
+const TensorEngine &prod_engine = DefaultTensorEngine::ref();
+
+EvalFixture::ParamRepo make_params() {
+    return EvalFixture::ParamRepo()
+        .add("x1y5z1", spec({x(1),y(5),z(1)}, N()))
+        .add("x1y1z1", spec({x(1),y(1),z(1)}, N()))
+        .add("x1y5z1_u", spec({x(1),y(5),z(1)}, N()), "tensor(x[1],y[5],z[])")
+        .add("x1y5z_m", spec({x(1),y(5),z({"a"})}, N()));
+}
+EvalFixture::ParamRepo param_repo = make_params();
+
+void verify_optimized(const vespalib::string &expr) {
+    EvalFixture fixture(prod_engine, expr, param_repo, true);
+    EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
+    auto info = fixture.find_all<DenseReplaceTypeFunction>();
+    EXPECT_EQUAL(info.size(), 1u);
+}
+
+void verify_not_optimized(const vespalib::string &expr) {
+    EvalFixture fixture(prod_engine, expr, param_repo, true);
+    EXPECT_EQUAL(fixture.result(), EvalFixture::ref(expr, param_repo));
+    auto info = fixture.find_all<DenseReplaceTypeFunction>();
+    EXPECT_TRUE(info.empty());
+}
+
+TEST("require that dimension removal can be optimized for appropriate aggregators") {
+    TEST_DO(verify_optimized("reduce(x1y5z1,avg,x)"));
+    TEST_DO(verify_not_optimized("reduce(x1y5z1,count,x)")); // NB
+    TEST_DO(verify_optimized("reduce(x1y5z1,prod,x)"));
+    TEST_DO(verify_optimized("reduce(x1y5z1,sum,x)"));
+    TEST_DO(verify_optimized("reduce(x1y5z1,max,x)"));
+    TEST_DO(verify_optimized("reduce(x1y5z1,min,x)"));    
+}
+
+TEST("require that multi-dimension removal can be optimized") {
+    TEST_DO(verify_optimized("reduce(x1y5z1,sum,x,z)"));
+}
+
+TEST("require that chained dimension removal can be optimized (and compacted)") {
+    TEST_DO(verify_optimized("reduce(reduce(x1y5z1,sum,x),sum,z)"));
+}
+
+TEST("require that reducing non-trivial dimension is not optimized") {
+    TEST_DO(verify_not_optimized("reduce(x1y5z1,sum,y)"));
+    TEST_DO(verify_not_optimized("reduce(x1y5z1,sum,x,y)"));
+    TEST_DO(verify_not_optimized("reduce(x1y5z1,sum,y,z)"));
+}
+
+TEST("require that full reduce is not optimized") {
+    TEST_DO(verify_not_optimized("reduce(x1y1z1,sum)"));
+    TEST_DO(verify_not_optimized("reduce(x1y1z1,sum,x,y,z)"));
+}
+
+TEST("require that inappropriate tensor types cannot be optimized") {
+    TEST_DO(verify_not_optimized("reduce(x1y5z1_u,sum,x)"));
+    TEST_DO(verify_not_optimized("reduce(x1y5z1_u,sum,z)"));
+    TEST_DO(verify_not_optimized("reduce(x1y5z_m,sum,x)"));
+    TEST_DO(verify_not_optimized("reduce(x1y5z_m,sum,z)"));
+}
+
+TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/eval/src/tests/tensor/dense_replace_type_function/dense_replace_type_function_test.cpp b/eval/src/tests/tensor/dense_replace_type_function/dense_replace_type_function_test.cpp
index 0e774be3dc3..7e6b933d7c6 100644
--- a/eval/src/tests/tensor/dense_replace_type_function/dense_replace_type_function_test.cpp
+++ b/eval/src/tests/tensor/dense_replace_type_function/dense_replace_type_function_test.cpp
@@ -64,4 +64,16 @@ TEST_F("require that DenseReplaceTypeFunction works as expected", Fixture()) {
     fprintf(stderr, "%s\n", f1.my_fun.as_string().c_str());
 }
 
+TEST("require that create_compact will collapse duplicate replace operations") {
+    Stash stash;
+    ValueType type = ValueType::double_type();
+    ChildMock leaf(type);
+    const DenseReplaceTypeFunction &a = DenseReplaceTypeFunction::create_compact(type, leaf, stash);
+    const DenseReplaceTypeFunction &b = DenseReplaceTypeFunction::create_compact(type, a, stash);
+    EXPECT_EQUAL(a.result_type(), type);
+    EXPECT_EQUAL(&a.child(), &leaf);
+    EXPECT_EQUAL(b.result_type(), type);
+    EXPECT_EQUAL(&b.child(), &leaf);
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
index 768ed4b1907..7f49fbc4acd 100644
--- a/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
+++ b/eval/src/vespa/eval/tensor/default_tensor_engine.cpp
@@ -10,6 +10,7 @@
 #include "dense/dense_dot_product_function.h"
 #include "dense/dense_xw_product_function.h"
 #include "dense/dense_fast_rename_optimizer.h"
+#include "dense/dense_remove_dimension_optimizer.h"
 #include "dense/dense_inplace_join_function.h"
 #include "dense/dense_inplace_map_function.h"
 #include "dense/vector_from_doubles_function.h"
@@ -229,6 +230,7 @@ DefaultTensorEngine::optimize(const TensorFunction &expr, Stash &stash) const
         child.set(DenseDotProductFunction::optimize(child.get(), stash));
         child.set(DenseXWProductFunction::optimize(child.get(), stash));
         child.set(DenseFastRenameOptimizer::optimize(child.get(), stash));
+        child.set(DenseRemoveDimensionOptimizer::optimize(child.get(), stash));
         child.set(DenseInplaceMapFunction::optimize(child.get(), stash));
         child.set(DenseInplaceJoinFunction::optimize(child.get(), stash));
         nodes.pop_back();
diff --git a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
index fc4f7922410..49675158ad7 100644
--- a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
+++ b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt
@@ -5,6 +5,7 @@ vespa_add_library(eval_tensor_dense OBJECT
     dense_fast_rename_optimizer.cpp
     dense_inplace_join_function.cpp
     dense_inplace_map_function.cpp
+    dense_remove_dimension_optimizer.cpp
     dense_replace_type_function.cpp
     dense_tensor.cpp
     dense_tensor_address_combiner.cpp
diff --git a/eval/src/vespa/eval/tensor/dense/dense_fast_rename_optimizer.cpp b/eval/src/vespa/eval/tensor/dense/dense_fast_rename_optimizer.cpp
index dbed0cfa5d0..4bc2458c29d 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_fast_rename_optimizer.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_fast_rename_optimizer.cpp
@@ -48,11 +48,7 @@ DenseFastRenameOptimizer::optimize(const eval::TensorFunction &expr, Stash &stas
         const ValueType &from_type = rename->child().result_type();
         const ValueType &to_type = expr.result_type();
         if (is_concrete_dense_stable_rename(from_type, to_type, rename->from(), rename->to())) {
-            if (auto replace = as<DenseReplaceTypeFunction>(rename->child())) {
-                return stash.create<DenseReplaceTypeFunction>(to_type, replace->child());
-            } else {
-                return stash.create<DenseReplaceTypeFunction>(to_type, rename->child());
-            }
+            return DenseReplaceTypeFunction::create_compact(to_type, rename->child(), stash);
         }
     }
     return expr;
diff --git a/eval/src/vespa/eval/tensor/dense/dense_remove_dimension_optimizer.cpp b/eval/src/vespa/eval/tensor/dense/dense_remove_dimension_optimizer.cpp
new file mode 100644
index 00000000000..302c81080dd
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_remove_dimension_optimizer.cpp
@@ -0,0 +1,58 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "dense_remove_dimension_optimizer.h"
+#include "dense_replace_type_function.h"
+#include <vespa/eval/eval/value_type.h>
+
+namespace vespalib::tensor {
+
+using eval::Aggr;
+using eval::ValueType;
+using eval::TensorFunction;
+using eval::as;
+using namespace eval::tensor_function;
+
+namespace {
+
+bool is_concrete_dense_tensor(const ValueType &type) {
+    return (type.is_dense() && !type.is_abstract());
+}
+
+bool is_ident_aggr(Aggr aggr) {
+    return ((aggr == Aggr::AVG)  ||
+            (aggr == Aggr::PROD) ||
+            (aggr == Aggr::SUM)  ||
+            (aggr == Aggr::MAX)  ||
+            (aggr == Aggr::MIN));
+}
+
+bool is_trivial_dim_list(const ValueType &type, const std::vector<vespalib::string> &dim_list) {
+    size_t npos = ValueType::Dimension::npos;
+    for (const vespalib::string &dim: dim_list) {
+        size_t idx = type.dimension_index(dim);
+        if ((idx == npos) || (type.dimensions()[idx].size != 1)) {
+            return false;
+        }
+    }
+    return true;
+}
+
+} // namespace vespalib::tensor::<unnamed>
+
+const TensorFunction &
+DenseRemoveDimensionOptimizer::optimize(const eval::TensorFunction &expr, Stash &stash)
+{
+    if (auto reduce = as<Reduce>(expr)) {
+        const TensorFunction &child = reduce->child();
+        if (is_concrete_dense_tensor(expr.result_type()) &&
+            is_concrete_dense_tensor(child.result_type()) &&
+            is_ident_aggr(reduce->aggr()) &&
+            is_trivial_dim_list(child.result_type(), reduce->dimensions()))
+        {
+            return DenseReplaceTypeFunction::create_compact(expr.result_type(), child, stash);
+        }
+    }
+    return expr;
+}
+
+} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_remove_dimension_optimizer.h b/eval/src/vespa/eval/tensor/dense/dense_remove_dimension_optimizer.h
new file mode 100644
index 00000000000..64b057a62d8
--- /dev/null
+++ b/eval/src/vespa/eval/tensor/dense/dense_remove_dimension_optimizer.h
@@ -0,0 +1,17 @@
+// Copyright 2018 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::tensor {
+
+/**
+ * Tensor function optimizer for efficient removal of dimensions with
+ * size 1 for dense tensors.
+ **/
+struct DenseRemoveDimensionOptimizer {
+    static const eval::TensorFunction &optimize(const eval::TensorFunction &expr, Stash &stash);
+};
+
+} // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.cpp b/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.cpp
index dd6d5635ef9..9deac5437e0 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.cpp
+++ b/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.cpp
@@ -10,6 +10,8 @@ using CellsRef = DenseTensorView::CellsRef;
 using eval::Value;
 using eval::ValueType;
 using eval::TensorFunction;
+using eval::as;
+using namespace eval::tensor_function;
 
 namespace {
 
@@ -42,4 +44,16 @@ DenseReplaceTypeFunction::compile_self(Stash &) const
     return eval::InterpretedFunction::Instruction(my_replace_type_op, (uint64_t)&(result_type()));
 }
 
+const DenseReplaceTypeFunction &
+DenseReplaceTypeFunction::create_compact(const eval::ValueType &result_type,
+                                         const eval::TensorFunction &child,
+                                         Stash &stash)
+{
+    if (auto replace = as<DenseReplaceTypeFunction>(child)) {
+        return stash.create<DenseReplaceTypeFunction>(result_type, replace->child());
+    } else {
+        return stash.create<DenseReplaceTypeFunction>(result_type, child);
+    }
+}
+
 } // namespace vespalib::tensor
diff --git a/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.h b/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.h
index 67a221b101e..4ad1f4c1cee 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_replace_type_function.h
@@ -18,6 +18,9 @@ public:
     ~DenseReplaceTypeFunction();
     eval::InterpretedFunction::Instruction compile_self(Stash &stash) const override;
     bool result_is_mutable() const override { return child().result_is_mutable(); }
+    static const DenseReplaceTypeFunction &create_compact(const eval::ValueType &result_type,
+            const eval::TensorFunction &child,
+            Stash &stash);
 };
 
 } // namespace vespalib::tensor