Add dot product aggregation

SConstruct

@@ -933,6 +933,8 @@ envDict = dict(BUILD_ROOT=buildDir,
 env = Environment(variables=env_vars, **envDict)
 del envDict
 
+env.Append(CXXFLAGS='/arch:AVX')
+
 env.AddMethod(env_os_is_wrapper, 'TargetOSIs')
 env.AddMethod(env_get_os_name_wrapper, 'GetTargetOSName')
 

src/mongo/db/pipeline/SConscript

@@ -170,6 +170,7 @@ env.Library(
     target='expression',
     source=[
         'expression.cpp',
+        'expression_simd.cpp',
         ],
     LIBDEPS=[
         'dependencies',

src/mongo/db/pipeline/expression.h

@@ -1065,6 +1065,24 @@ class ExpressionMultiply final : public ExpressionVariadic<ExpressionMultiply> {
 };
 
 
+class ExpressionVectorDot final : public ExpressionVariadic<ExpressionVectorDot> {
+public:
+    explicit ExpressionVectorDot(const boost::intrusive_ptr<ExpressionContext>& expCtx)
+        : ExpressionVariadic<ExpressionVectorDot>(expCtx) {}
+
+    Value evaluateInternal(Variables* vars) const final;
+    const char* getOpName() const final;
+
+    bool isAssociative() const final {
+        return true;
+    }
+
+    bool isCommutative() const final {
+        return true;
+    }
+};
+
+
 class ExpressionMonth final : public ExpressionFixedArity<ExpressionMonth, 1> {
 public:
     explicit ExpressionMonth(const boost::intrusive_ptr<ExpressionContext>& expCtx)

src/mongo/db/pipeline/expression_simd.cpp

@@ -0,0 +1,189 @@
+/**
+ * Copyright (c) 2011 10gen Inc.
+ *
+ * This program is free software: you can redistribute it and/or  modify
+ * it under the terms of the GNU Affero General Public License, version 3,
+ * as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ * As a special exception, the copyright holders give permission to link the
+ * code of portions of this program with the OpenSSL library under certain
+ * conditions as described in each individual source file and distribute
+ * linked combinations including the program with the OpenSSL library. You
+ * must comply with the GNU Affero General Public License in all respects for
+ * all of the code used other than as permitted herein. If you modify file(s)
+ * with this exception, you may extend this exception to your version of the
+ * file(s), but you are not obligated to do so. If you do not wish to do so,
+ * delete this exception statement from your version. If you delete this
+ * exception statement from all source files in the program, then also delete
+ * it in the license file.
+ */
+
+
+#include "mongo/platform/basic.h"
+
+#include "mongo/db/pipeline/expression.h"
+
+#include <algorithm>
+#include <boost/algorithm/string.hpp>
+#include <cstdio>
+#include <vector>
+
+#include <immintrin.h>
+
+#include "mongo/db/jsobj.h"
+#include "mongo/db/pipeline/document.h"
+#include "mongo/db/pipeline/expression_context.h"
+#include "mongo/db/pipeline/value.h"
+#include "mongo/platform/bits.h"
+#include "mongo/platform/decimal128.h"
+#include "mongo/util/mongoutils/str.h"
+#include "mongo/util/string_map.h"
+#include "mongo/util/summation.h"
+
+namespace mongo {
+using Parser = Expression::Parser;
+
+using namespace mongoutils;
+
+using boost::intrusive_ptr;
+using std::map;
+using std::move;
+using std::pair;
+using std::set;
+using std::string;
+using std::vector;
+
+/// Helper function to easily wrap constants with $const.
+static Value serializeConstant(Value val) {
+    if (val.missing()) {
+        return Value("$$REMOVE"_sd);
+    }
+
+    return Value(DOC("$const" << val));
+}
+
+/* ------------------------- ExpressionVectorDot ----------------------------- */
+
+struct AlignedVector {
+    alignas(32) float v[8];
+};
+
+Value ExpressionVectorDot::evaluateInternal(Variables* vars) const {
+    const size_t n = vpOperand.size();
+    if (n != 2) {
+        uasserted(40418,
+                  str::stream() << "$vector_dot requires two arguments but got "
+                                << n);
+    }
+
+    Value valLeft = vpOperand[0]->evaluateInternal(vars);
+    Value valRight = vpOperand[1]->evaluateInternal(vars);
+
+    if (valLeft.nullish() || valRight.nullish()) {
+        return Value(BSONNULL);
+    }
+
+    const auto type = valLeft.getType();
+    if (type != valRight.getType()) {
+        uasserted(40419,
+                  str::stream() << "$vector_dot requires that both operands are of the same type");
+    }
+
+    // db.local.insert({_id:0, 'a':BinData(128, "AACAPwAAAEAAAIBAAAAAQQAAgEEAAABCAACAQgAAAEM=")})
+    // db.local.aggregate([{$match:{a:{$type:4}}}, {$project:{result: {$vector_dot:['$a', '$a']}}}])
+
+    if (type == BinData) {
+        const auto left = valLeft.coerceToBinData();
+        const auto right = valRight.coerceToBinData();
+
+        const size_t m = left.size();
+        if (m != right.size()) {
+            uasserted(40420,
+                      str::stream() << "$vector_dot requires that both vectors are of the same length");
+        }
+
+        if ((m & 3) != 0) {
+            uasserted(40421,
+                      str::stream() << "$vector_dot requires that the vector length is a multiple of four");
+        }
+
+        const auto* leftData =  reinterpret_cast<const float*>(left.rawData());
+        const auto* rightData = reinterpret_cast<const float*>(right.rawData());
+        const auto left32byteAligned = (reinterpret_cast<uintptr_t>(leftData) & 31) == 0;
+        const auto right32byteAligned = (reinterpret_cast<uintptr_t>(rightData) & 31) == 0;
+        const auto not32byteAligned = !(left32byteAligned && right32byteAligned);
+
+        size_t i = 0;
+        const size_t elems = m / 4;
+        const size_t stride = 8;
+        auto sum = _mm256_setzero_ps();
+
+        // process blockwise
+        if (elems >= stride) {
+            if (not32byteAligned) {
+                for (; i <= elems - stride; i += stride) {
+                    const auto a = _mm256_loadu_ps(leftData);
+                    const auto b = _mm256_loadu_ps(rightData);
+                    leftData += stride;
+                    rightData += stride;
+                    const auto partialSum = _mm256_dp_ps(a, b, 0xff);
+                    sum = _mm256_add_ps(sum, partialSum);
+                }
+            }
+            else {
+                // The internal pointer appears to always be aligned to
+                // 16 byte boundaries, thus this code might never be reached.
+                // Forcing the alignment on Value or ValueStorage using
+                // alignas() triggers size assertions all over the place.
+                // I leave this code here in case of lucky circumstances.
+                for (; i <= elems - stride; i += stride) {
+                    const auto a = _mm256_load_ps(leftData);
+                    const auto b = _mm256_load_ps(rightData);
+                    leftData += stride;
+                    rightData += stride;
+                    const auto partialSum = _mm256_dp_ps(a, b, 0xff);
+                    sum = _mm256_add_ps(sum, partialSum);
+                }
+            }
+        }
+
+        // process remaining
+        if (i < elems) {
+            AlignedVector remainingLeft = {0.0f};
+            AlignedVector remainingRight = {0.0f};
+            for (size_t j = 0; i < elems; ++i, ++j) {
+                remainingLeft.v[j] = *(leftData++);
+                remainingRight.v[j] = *(rightData++);
+            }
+            const auto a = _mm256_load_ps(remainingLeft.v);
+            const auto b = _mm256_load_ps(remainingRight.v);
+            const auto partialSum = _mm256_dp_ps(a, b, 0xff);
+            sum = _mm256_add_ps(sum, partialSum);
+        }
+
+        const auto s = reinterpret_cast<float*>(&sum);
+        const auto dotProduct = static_cast<double>(s[0] + s[4]);
+        return Value(dotProduct);
+    } else {
+        uasserted(40422,
+                  str::stream() << "$vector_dot only supports BinData types, not "
+                                << typeName(valLeft.getType())
+                                << " and "
+                                << typeName(valRight.getType()));
+    }
+}
+
+REGISTER_EXPRESSION(vector_dot, ExpressionVectorDot::parse);
+const char* ExpressionVectorDot::getOpName() const {
+    return "$vector_dot";
+}
+
+}  // namespace mongo

src/mongo/db/pipeline/value.cpp

@@ -623,6 +623,18 @@ static string tmToISODateString(const tm& time) {
     return buf;
 }
 
+StringData Value::coerceToBinData() const {
+    switch (getType()) {
+        case BinData:
+            return getStringData();
+        default:
+            uassert(40423,
+                    str::stream() << "can't convert from BSON type " << typeName(getType())
+                                  << " to BinData",
+                    false);
+    }
+}
+
 string Value::coerceToString() const {
     switch (getType()) {
         case NumberDouble:

src/mongo/db/pipeline/value.h

@@ -223,6 +223,7 @@ class Value {
      *  TODO: decided how to handle unsupported types.
      */
     std::string coerceToString() const;
+    StringData coerceToBinData() const;
     int coerceToInt() const;
     long long coerceToLong() const;
     double coerceToDouble() const;