Merge pull request #16101 from vespa-engine/arnej/general-pow-as-map

extend PowAsMap to handle mixed tensors

8 files changed, 25 insertions, 28 deletions

diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt
index 00618f2c838..3c94639d0aa 100644
--- a/eval/CMakeLists.txt
+++ b/eval/CMakeLists.txt
@@ -56,7 +56,7 @@ vespa_define_module(
     src/tests/instruction/add_trivial_dimension_optimizer
     src/tests/instruction/fast_rename_optimizer
     src/tests/instruction/dense_inplace_join_function
-    src/tests/instruction/dense_pow_as_map_optimizer
+    src/tests/instruction/pow_as_map_optimizer
     src/tests/instruction/remove_trivial_dimension_optimizer
     src/tests/instruction/dense_replace_type_function
     src/tests/instruction/dense_simple_join_function
diff --git a/eval/src/tests/instruction/dense_pow_as_map_optimizer/CMakeLists.txt b/eval/src/tests/instruction/dense_pow_as_map_optimizer/CMakeLists.txt
deleted file mode 100644
index d6ce9f1924c..00000000000
--- a/eval/src/tests/instruction/dense_pow_as_map_optimizer/CMakeLists.txt
+++ /dev/null
@@ -1,9 +0,0 @@
-# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-vespa_add_executable(eval_dense_pow_as_map_optimizer_test_app TEST
-    SOURCES
-    dense_pow_as_map_optimizer_test.cpp
-    DEPENDS
-    vespaeval
-    GTest::GTest
-)
-vespa_add_test(NAME eval_dense_pow_as_map_optimizer_test_app COMMAND eval_dense_pow_as_map_optimizer_test_app)
diff --git a/eval/src/tests/instruction/pow_as_map_optimizer/CMakeLists.txt b/eval/src/tests/instruction/pow_as_map_optimizer/CMakeLists.txt
new file mode 100644
index 00000000000..25eeb73df37
--- /dev/null
+++ b/eval/src/tests/instruction/pow_as_map_optimizer/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_pow_as_map_optimizer_test_app TEST
+    SOURCES
+    pow_as_map_optimizer_test.cpp
+    DEPENDS
+    vespaeval
+    GTest::GTest
+)
+vespa_add_test(NAME eval_pow_as_map_optimizer_test_app COMMAND eval_pow_as_map_optimizer_test_app)
diff --git a/eval/src/tests/instruction/dense_pow_as_map_optimizer/dense_pow_as_map_optimizer_test.cpp b/eval/src/tests/instruction/pow_as_map_optimizer/pow_as_map_optimizer_test.cpp
index 67567b4e289..920198f67d7 100644
--- a/eval/src/tests/instruction/dense_pow_as_map_optimizer/dense_pow_as_map_optimizer_test.cpp
+++ b/eval/src/tests/instruction/pow_as_map_optimizer/pow_as_map_optimizer_test.cpp
@@ -19,7 +19,7 @@ EvalFixture::ParamRepo make_params() {
     return EvalFixture::ParamRepo()
         .add("a", spec(1.5))
         .add("b", spec(2.5))
-        .add("sparse", spec({x({"a"})}, N()))
+        .add("sparse", spec({x({"a","b"})}, N()))
         .add("mixed", spec({x({"a"}),y(5)}, N()))
         .add_matrix("x", 5, "y", 3);
 }
@@ -30,11 +30,10 @@ void verify_optimized(const vespalib::string &expr, op1_t op1, bool inplace = fa
     EvalFixture fixture(prod_factory, expr, param_repo, true, true);
     EXPECT_EQ(fixture.result(), EvalFixture::ref(expr, param_repo));
     EXPECT_EQ(fixture.result(), slow_fixture.result());
-    auto info = fixture.find_all<DenseSimpleMapFunction>();
+    auto info = fixture.find_all<tensor_function::Map>();
     ASSERT_EQ(info.size(), 1u);
     EXPECT_TRUE(info[0]->result_is_mutable());
     EXPECT_EQ(info[0]->function(), op1);
-    EXPECT_EQ(info[0]->inplace(), inplace);
     ASSERT_EQ(fixture.num_params(), 1);
     if (inplace) {
         EXPECT_EQ(fixture.get_param(0), fixture.result());
@@ -76,16 +75,16 @@ TEST(PowAsMapTest, hypercubed_dense_tensor_is_not_optimized) {
     verify_not_optimized("join(x5y3,4.0,f(x,y)(pow(x,y)))");
 }
 
-TEST(PowAsMapTest, scalar_join_is_not_optimized) {
-    verify_not_optimized("join(a,2.0,f(x,y)(pow(x,y)))");
+TEST(PowAsMapTest, scalar_join_is_optimized) {
+    verify_optimized("join(a,2.0,f(x,y)(pow(x,y)))", Square::f);
 }
 
-TEST(PowAsMapTest, sparse_join_is_not_optimized) {
-    verify_not_optimized("join(sparse,2.0,f(x,y)(pow(x,y)))");
+TEST(PowAsMapTest, sparse_join_is_optimized) {
+    verify_optimized("join(sparse,2.0,f(x,y)(pow(x,y)))", Square::f);
 }
 
-TEST(PowAsMapTest, mixed_join_is_not_optimized) {
-    verify_not_optimized("join(mixed,2.0,f(x,y)(pow(x,y)))");
+TEST(PowAsMapTest, mixed_join_is_optimized) {
+    verify_optimized("join(mixed,2.0,f(x,y)(pow(x,y)))", Square::f);
 }
 
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/eval/src/vespa/eval/eval/optimize_tensor_function.cpp b/eval/src/vespa/eval/eval/optimize_tensor_function.cpp
index 515b48b3693..02d6bdbf0f3 100644
--- a/eval/src/vespa/eval/eval/optimize_tensor_function.cpp
+++ b/eval/src/vespa/eval/eval/optimize_tensor_function.cpp
@@ -16,7 +16,7 @@
 #include <vespa/eval/instruction/dense_simple_expand_function.h>
 #include <vespa/eval/instruction/dense_simple_join_function.h>
 #include <vespa/eval/instruction/join_with_number_function.h>
-#include <vespa/eval/instruction/dense_pow_as_map_optimizer.h>
+#include <vespa/eval/instruction/pow_as_map_optimizer.h>
 #include <vespa/eval/instruction/dense_simple_map_function.h>
 #include <vespa/eval/instruction/vector_from_doubles_function.h>
 #include <vespa/eval/instruction/dense_tensor_create_function.h>
@@ -65,7 +65,7 @@ const TensorFunction &optimize_for_factory(const ValueBuilderFactory &factory, c
             child.set(DenseTensorPeekFunction::optimize(child.get(), stash));
             child.set(DenseLambdaPeekOptimizer::optimize(child.get(), stash));
             child.set(FastRenameOptimizer::optimize(child.get(), stash));
-            child.set(DensePowAsMapOptimizer::optimize(child.get(), stash));
+            child.set(PowAsMapOptimizer::optimize(child.get(), stash));
             child.set(DenseSimpleMapFunction::optimize(child.get(), stash));
             child.set(DenseSimpleJoinFunction::optimize(child.get(), stash));
             child.set(JoinWithNumberFunction::optimize(child.get(), stash));
diff --git a/eval/src/vespa/eval/instruction/CMakeLists.txt b/eval/src/vespa/eval/instruction/CMakeLists.txt
index 1317a2c2cf6..452ed51610a 100644
--- a/eval/src/vespa/eval/instruction/CMakeLists.txt
+++ b/eval/src/vespa/eval/instruction/CMakeLists.txt
@@ -9,7 +9,7 @@ vespa_add_library(eval_instruction OBJECT
     dense_lambda_peek_optimizer.cpp
     dense_matmul_function.cpp
     dense_multi_matmul_function.cpp
-    dense_pow_as_map_optimizer.cpp
+    pow_as_map_optimizer.cpp
     remove_trivial_dimension_optimizer.cpp
     dense_simple_expand_function.cpp
     dense_simple_join_function.cpp
diff --git a/eval/src/vespa/eval/instruction/dense_pow_as_map_optimizer.cpp b/eval/src/vespa/eval/instruction/pow_as_map_optimizer.cpp
index 61ef2243480..00fbc1c642b 100644
--- a/eval/src/vespa/eval/instruction/dense_pow_as_map_optimizer.cpp
+++ b/eval/src/vespa/eval/instruction/pow_as_map_optimizer.cpp
@@ -1,6 +1,6 @@
 // Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
-#include "dense_pow_as_map_optimizer.h"
+#include "pow_as_map_optimizer.h"
 #include "dense_simple_map_function.h"
 #include <vespa/eval/eval/operation.h>
 
@@ -10,14 +10,13 @@ using namespace tensor_function;
 using namespace operation;
 
 const TensorFunction &
-DensePowAsMapOptimizer::optimize(const TensorFunction &expr, Stash &stash)
+PowAsMapOptimizer::optimize(const TensorFunction &expr, Stash &stash)
 {
     if (auto join = as<Join>(expr)) {
         const TensorFunction &lhs = join->lhs();
         const TensorFunction &rhs = join->rhs();
         if ((join->function() == Pow::f) &&
-            lhs.result_type().is_dense() &&
-            rhs.result_type().is_double())
+            rhs.result_type().is_scalar())
         {
             if (auto const_value = as<ConstValue>(rhs)) {
                 if (const_value->value().as_double() == 2.0) {
diff --git a/eval/src/vespa/eval/instruction/dense_pow_as_map_optimizer.h b/eval/src/vespa/eval/instruction/pow_as_map_optimizer.h
index e61069b87b0..df9735f76ca 100644
--- a/eval/src/vespa/eval/instruction/dense_pow_as_map_optimizer.h
+++ b/eval/src/vespa/eval/instruction/pow_as_map_optimizer.h
@@ -10,9 +10,8 @@ namespace vespalib::eval {
  * Tensor function optimizer for converting join expressions on the
  * form 'join(tensor,<small integer constant>,f(x,y)(pow(x,y))' to
  * expressions on the form 'map(tensor,f(x)(x*x...))'.
- * TODO: extend to mixed tensors.
  **/
-struct DensePowAsMapOptimizer {
+struct PowAsMapOptimizer {
     static const TensorFunction &optimize(const TensorFunction &expr, Stash &stash);
 };