universal tensor dot product

[jobergum]

As we get more useful deep-learned representations, for example ColPali, the universal dot product is in more use, and we should see if there is anything we can do to accelerate it further.

For ColPali with multi-page representations, where a PDF document is represented by a screenshot of all the pages, we have two tensors:

doc tensor<float>(page{}, patch{}, v[128])
query tensor<float>(token{}, v[128])

Typically, pages 5-20, patch 1030, and token 20:

And where we want to find the score per page by

reduce(reduce(reduce(query * doc, sum, v), max, patch), sum, token)

This is obviously compute intensive and currently triggers the universal dot product optimization:

Playground link

[
  {
    "class": "vespalib::eval::tensor_function::Inject",
    "symbol": "<inject_param>"
  },
  {
    "class": "vespalib::eval::tensor_function::Inject",
    "symbol": "<inject_param>"
  },
  {
    "class": "vespalib::eval::UniversalDotProduct",
    "symbol": "void vespalib::eval::(anonymous namespace)::my_universal_dot_product_op<float, float, float, false, true, true>(vespalib::eval::InterpretedFunction::State&, unsigned long)"
  },
  {
    "class": "vespalib::eval::tensor_function::Reduce",
    "symbol": "void vespalib::eval::instruction::(anonymous namespace)::my_generic_reduce_op<float, float, vespalib::eval::aggr::Max<float> >(vespalib::eval::InterpretedFunction::State&, unsigned long)"
  },
  {
    "class": "vespalib::eval::tensor_function::Reduce",
    "symbol": "void vespalib::eval::instruction::(anonymous namespace)::my_generic_reduce_op<float, float, vespalib::eval::aggr::Sum<float> >(vespalib::eval::InterpretedFunction::State&, unsigned long)"
  }
]