Merge 10ba3ef into 3b87a8e

Author: azevaykin

ydb/library/yql/udfs/common/knn/knn-distance.h

@@ -0,0 +1,41 @@
+#pragma once
+
+#include <library/cpp/dot_product/dot_product.h>
+
+// UDF strings are 8-byte aligned.
+// SSE prefers 16-byte alignment.
+// So, there are four cases:
+// 1. lhs % 16 == 0, rhs % 16 == 0
+// 2. lhs % 16 == 8, rhs % 16 == 0
+// 3. lhs % 16 == 0, rhs % 16 == 8
+// 4. lhs % 16 == 8, rhs % 16 == 8
+
+// In case 4 we divide aligned and unaligned elements.
+// So, we make two DotProduct calls. It can impove performance.
+
+// Case 1 is perfectly aligned for SSE. Read will be aligned.
+// Cases 2,3 are badly aligned. Read will be unaligned.
+
+inline float KnnDotProduct(const float* lhs, const float* rhs, size_t length) {
+    if ((size_t)lhs % 16 == 8 && (size_t)rhs % 16 == 8) {
+        const size_t numUnaligned = 8 / sizeof(float);
+        const float resUnaligned = DotProduct(lhs, rhs, numUnaligned);
+        const float resAligned = DotProduct(lhs + numUnaligned, rhs + numUnaligned, length - numUnaligned);
+        return resUnaligned + resAligned;
+    } else {
+        return DotProduct(lhs, rhs, length);
+    }
+}
+
+static inline TTriWayDotProduct<float> KnnTriWayDotProduct(const float* lhs, const float* rhs, size_t length) {
+    if ((size_t)lhs % 16 == 8 && (size_t)rhs % 16 == 8) {
+        const size_t numUnaligned = 8 / sizeof(float);
+        const auto resUnaligned = TriWayDotProduct(lhs, rhs, numUnaligned);
+        const auto resAligned = TriWayDotProduct(lhs + numUnaligned, rhs + numUnaligned, length - numUnaligned);
+        return { resUnaligned.LL + resAligned.LL, 
+                 resUnaligned.LR + resAligned.LR, 
+                 resUnaligned.RR + resAligned.RR };
+    } else {
+        return TriWayDotProduct(lhs, rhs, length);
+    }  
+}

ydb/library/yql/udfs/common/knn/knn-serializer.h

@@ -11,22 +11,23 @@
 using namespace NYql;
 using namespace NYql::NUdf;
 
-enum EFormat : ui32 {
+enum EFormat : ui8 {
     FloatVector = 1
 };
 
+static constexpr size_t HeaderLen = sizeof(ui8);
 
 class TFloatVectorSerializer {
 public:
     static TUnboxedValue Serialize(const IValueBuilder* valueBuilder, const TUnboxedValue x) {
         auto serialize = [&x] (IOutputStream& outStream) {
-            const EFormat format = EFormat::FloatVector;
-            outStream.Write(&format, sizeof(ui32));
             EnumerateVector(x,  [&outStream] (float element) { outStream.Write(&element, sizeof(float)); });
+            const EFormat format = EFormat::FloatVector;
+            outStream.Write(&format, HeaderLen);
         };
 
         if (x.HasFastListLength()) {
-            auto str = valueBuilder->NewStringNotFilled(sizeof(ui32) + x.GetListLength() * sizeof(float));
+            auto str = valueBuilder->NewStringNotFilled(HeaderLen + x.GetListLength() * sizeof(float));
             auto strRef = str.AsStringRef();
             TMemoryOutput memoryOutput(strRef.Data(), strRef.Size());
 
@@ -42,9 +43,8 @@ class TFloatVectorSerializer {
     }
 
     static TUnboxedValue Deserialize(const IValueBuilder *valueBuilder, const TStringRef& str) {
-        //skip format header, it was already read
-        const char* buf = str.Data() + sizeof(ui32);
-        const size_t len = str.Size() - sizeof(ui32);
+        const char* buf = str.Data();
+        const size_t len = str.Size() - HeaderLen;
 
         if (len % sizeof(float) != 0)    
             return {};
@@ -66,9 +66,8 @@ class TFloatVectorSerializer {
     }
 
     static const TArrayRef<const float> GetArray(const TStringRef& str) {
-        //skip format header, it was already read
-        const char* buf = str.Data() + sizeof(ui32);
-        const size_t len = str.Size() - sizeof(ui32);
+        const char* buf = str.Data();
+        const size_t len = str.Size() - HeaderLen;
 
         if (len % sizeof(float) != 0)    
             return {};
@@ -95,8 +94,8 @@ class TSerializerFacade {
         if (str.Size() == 0)
             return {};
 
-        const ui32* format = reinterpret_cast<const ui32*>(str.Data());
-        switch (*format) {
+        const ui8 format = str.Data()[str.Size() - HeaderLen];
+        switch (format) {
             case EFormat::FloatVector:
                 return TFloatVectorSerializer::Deserialize(valueBuilder, str);
             default:
@@ -108,8 +107,8 @@ class TSerializerFacade {
         if (str.Size() == 0)
             return {};
 
-        const ui32* format = reinterpret_cast<const ui32*>(str.Data());
-        switch (*format) {
+        const ui8 format = str.Data()[str.Size() - HeaderLen];
+        switch (format) {
             case EFormat::FloatVector:
                 return TFloatVectorSerializer::GetArray(str);
             default:

ydb/library/yql/udfs/common/knn/knn.cpp

@@ -1,9 +1,9 @@
+#include "knn-distance.h"
 #include "knn-enumerator.h"
 #include "knn-serializer.h"
 
 #include <ydb/library/yql/public/udf/udf_helpers.h>
 
-#include <library/cpp/dot_product/dot_product.h>
 #include <util/generic/buffer.h>
 #include <util/stream/format.h>
 
@@ -33,7 +33,7 @@ SIMPLE_STRICT_UDF(TInnerProductSimilarity, TOptional<float>(TAutoMap<const char*
     if (vector1.size() != vector2.size() || vector1.empty() || vector2.empty())
         return {};
 
-    const float dotProduct = DotProduct(vector1.data(), vector2.data(), vector1.size());
+    const float dotProduct = KnnDotProduct(vector1.data(), vector2.data(), vector1.size());
     return TUnboxedValuePod{dotProduct};
 }
 
@@ -46,7 +46,7 @@ SIMPLE_STRICT_UDF(TCosineSimilarity, TOptional<float>(TAutoMap<const char*>, TAu
     if (vector1.size() != vector2.size() || vector1.empty() || vector2.empty())
         return {};    
 
-    const auto [ll, lr, rr] = TriWayDotProduct(vector1.data(), vector2.data(), vector1.size());
+    const auto [ll, lr, rr] = KnnTriWayDotProduct(vector1.data(), vector2.data(), vector1.size());
     const float cosine = lr / std::sqrt(ll * rr);
     return TUnboxedValuePod{cosine};
 }
@@ -60,7 +60,7 @@ SIMPLE_STRICT_UDF(TCosineDistance, TOptional<float>(TAutoMap<const char*>, TAuto
     if (vector1.size() != vector2.size() || vector1.empty() || vector2.empty())
         return {};    
 
-    const auto [ll, lr, rr] = TriWayDotProduct(vector1.data(), vector2.data(), vector1.size());
+    const auto [ll, lr, rr] = KnnTriWayDotProduct(vector1.data(), vector2.data(), vector1.size());
     const float cosine = lr / std::sqrt(ll * rr);
     return TUnboxedValuePod{1 - cosine};
 }

ydb/library/yql/udfs/common/knn/test/canondata/test.test_LazyListSerialization_/results.txt

@@ -20,7 +20,7 @@
                 "Data" = [
                     [
                         [
-                            "AQAAAAAAgD8AAABAAABAQAAAgEAAAKBA"
+                            "AACAPwAAAEAAAEBAAACAQAAAoEAB"
                         ]
                     ]
                 ]

ydb/library/yql/udfs/common/knn/test/canondata/test.test_ListSerialization_/results.txt

@@ -20,7 +20,7 @@
                 "Data" = [
                     [
                         [
-                            "AQAAAJqZmT8zMxNAmplZQAAAkEAzM7NA"
+                            "mpmZPzMzE0CamVlAAACQQDMzs0AB"
                         ]
                     ]
                 ]