properly integrate float16

sivukhin · sivukhin · commit 69bce662bf96 · 2024-08-22T00:48:30.000+04:00
diff --git a/libsql-sqlite3/src/vector.c b/libsql-sqlite3/src/vector.c
@@ -122,6 +122,8 @@ float vectorDistanceCos(const Vector *pVector1, const Vector *pVector2){
       return vector1BitDistanceHamming(pVector1, pVector2);
     case VECTOR_TYPE_FLOAT8:
       return vectorF8DistanceCos(pVector1, pVector2);
+    case VECTOR_TYPE_FLOAT16:
+      return vectorF16DistanceCos(pVector1, pVector2);
     default:
       assert(0);
   }
@@ -137,6 +139,8 @@ float vectorDistanceL2(const Vector *pVector1, const Vector *pVector2){
       return vectorF64DistanceL2(pVector1, pVector2);
     case VECTOR_TYPE_FLOAT8:
       return vectorF8DistanceL2(pVector1, pVector2);
+    case VECTOR_TYPE_FLOAT16:
+      return vectorF16DistanceL2(pVector1, pVector2);
     default:
       assert(0);
   }
@@ -303,6 +307,13 @@ static int vectorParseMeta(const unsigned char *pBlob, size_t nBlobSize, int *pT
     nTrailingBytes = pBlob[nBlobSize - 1];
     *pDims = (nBlobSize - 2) - sizeof(float) - sizeof(float) - nTrailingBytes;
     *pDataSize = nBlobSize - 2;
+  }else if( *pType == VECTOR_TYPE_FLOAT16 ){
+    if( nBlobSize % 2 != 0 ){
+      *pzErrMsg = sqlite3_mprintf("vector: float16 vector blob length must be divisible by 2 (excluding 'type'-byte): length=%d", nBlobSize);
+      return SQLITE_ERROR;
+    }
+    *pDims = nBlobSize / sizeof(u16);
+    *pDataSize = nBlobSize;
   }else{
     *pzErrMsg = sqlite3_mprintf("vector: unexpected binary type: %d", *pType);
     return SQLITE_ERROR;
@@ -351,6 +362,9 @@ int vectorParseSqliteBlobWithType(
     case VECTOR_TYPE_FLOAT8:
       vectorF8DeserializeFromBlob(pVector, pBlob, nDataSize);
       return 0;
+    case VECTOR_TYPE_FLOAT16:
+      vectorF16DeserializeFromBlob(pVector, pBlob, nDataSize);
+      return 0;
     default: 
       assert(0);
   }
@@ -452,6 +466,9 @@ void vectorDump(const Vector *pVector){
     case VECTOR_TYPE_FLOAT8:
       vectorF8Dump(pVector);
       break;
+    case VECTOR_TYPE_FLOAT16:
+      vectorF16Dump(pVector);
+      break;
     default:
       assert(0);
   }
@@ -477,7 +494,7 @@ static int vectorMetaSize(VectorType type, VectorDims dims){
   int nDataSize;
   if( type == VECTOR_TYPE_FLOAT32 ){
     return 0;
-  }else if( type == VECTOR_TYPE_FLOAT64 ){
+  }else if( type == VECTOR_TYPE_FLOAT64 || type == VECTOR_TYPE_FLOAT16 ){
     return 1;
   }else if( type == VECTOR_TYPE_FLOAT1BIT ){
     nDataSize = vectorDataSize(type, dims);
@@ -496,14 +513,14 @@ static int vectorMetaSize(VectorType type, VectorDims dims){
 static void vectorSerializeMeta(const Vector *pVector, size_t nDataSize, unsigned char *pBlob, size_t nBlobSize){
   if( pVector->type == VECTOR_TYPE_FLOAT32 ){
     // no meta for f32 type as this is "default" vector type
-  }else if( pVector->type == VECTOR_TYPE_FLOAT64 ){
+  }else if( pVector->type == VECTOR_TYPE_FLOAT64 || pVector->type == VECTOR_TYPE_FLOAT16 ){
     assert( nDataSize % 2 == 0 );
     assert( nBlobSize == nDataSize + 1 );
-    pBlob[nBlobSize - 1] = VECTOR_TYPE_FLOAT64;
+    pBlob[nBlobSize - 1] = pVector->type;
   }else if( pVector->type == VECTOR_TYPE_FLOAT1BIT ){
     assert( nBlobSize % 2 == 1 );
     assert( nBlobSize >= 3 );
-    pBlob[nBlobSize - 1] = VECTOR_TYPE_FLOAT1BIT;
+    pBlob[nBlobSize - 1] = pVector->type;
     pBlob[nBlobSize - 2] = 8 * (nBlobSize - 1) - pVector->dims;
     if( vectorMetaSize(pVector->type, pVector->dims) == 3 ){
       pBlob[nBlobSize - 3] = 0;
@@ -512,8 +529,9 @@ static void vectorSerializeMeta(const Vector *pVector, size_t nDataSize, unsigne
     assert( nBlobSize % 2 == 1 );
     assert( nDataSize % 2 == 0 );
     assert( nBlobSize == nDataSize + 3 );
-    pBlob[nBlobSize - 1] = VECTOR_TYPE_FLOAT8;
+    pBlob[nBlobSize - 1] = pVector->type;
     pBlob[nBlobSize - 2] = ALIGN(pVector->dims, sizeof(float)) - pVector->dims;
+    pBlob[nBlobSize - 3] = 0;
   }else{
     assert( 0 );
   }
@@ -561,6 +579,9 @@ void vectorSerializeToBlob(const Vector *pVector, unsigned char *pBlob, size_t n
     case VECTOR_TYPE_FLOAT8:
       vectorF8SerializeToBlob(pVector, pBlob, nBlobSize);
       break;
+    case VECTOR_TYPE_FLOAT16:
+      vectorF16SerializeToBlob(pVector, pBlob, nBlobSize);
+      break;
     default:
       assert(0);
   }
@@ -576,6 +597,7 @@ static void vectorConvertFromF32(const Vector *pFrom, Vector *pTo){
 
   u8 *dst1Bit;
   double *dstF64;
+  u16 *dstF16;
 
   assert( pFrom->dims == pTo->dims );
   assert( pFrom->type != pTo->type );
@@ -597,6 +619,11 @@ static void vectorConvertFromF32(const Vector *pFrom, Vector *pTo){
         dst1Bit[i / 8] |= (1 << (i & 7));
       }
     }
+  }else if( pTo->type == VECTOR_TYPE_FLOAT16 ){
+    dstF16 = pTo->data;
+    for(i = 0; i < pFrom->dims; i++){
+      dstF16[i] = vectorF16FromFloat(src[i]);
+    }
   }else{
     assert( 0 );
   }
@@ -608,6 +635,7 @@ static void vectorConvertFromF64(const Vector *pFrom, Vector *pTo){
 
   u8 *dst1Bit;
   float *dstF32;
+  u16 *dstF16;
 
   assert( pFrom->dims == pTo->dims );
   assert( pFrom->type != pTo->type );
@@ -629,6 +657,11 @@ static void vectorConvertFromF64(const Vector *pFrom, Vector *pTo){
         dst1Bit[i / 8] |= (1 << (i & 7));
       }
     }
+  }else if( pTo->type == VECTOR_TYPE_FLOAT16 ){
+    dstF16 = pTo->data;
+    for(i = 0; i < pFrom->dims; i++){
+      dstF16[i] = vectorF16FromFloat(src[i]);
+    }
   }else{
     assert( 0 );
   }
@@ -640,6 +673,7 @@ static void vectorConvertFrom1Bit(const Vector *pFrom, Vector *pTo){
 
   float *dstF32;
   double *dstF64;
+  u16 *dstF16;
 
   assert( pFrom->dims == pTo->dims );
   assert( pFrom->type != pTo->type );
@@ -664,6 +698,17 @@ static void vectorConvertFrom1Bit(const Vector *pFrom, Vector *pTo){
         dstF64[i] = -1;
       }
     }
+  }else if( pTo->type == VECTOR_TYPE_FLOAT16 ){
+    u16 positive = vectorF16FromFloat(+1);
+    u16 negative = vectorF16FromFloat(-1);
+    dstF16 = pTo->data;
+    for(i = 0; i < pFrom->dims; i++){
+      if( ((src[i / 8] >> (i & 7)) & 1) == 1 ){
+        dstF16[i] = positive;
+      }else{
+        dstF16[i] = negative;
+      }
+    }
   }else{
     assert( 0 );
   }
@@ -677,6 +722,7 @@ static void vectorConvertFromF8(const Vector *pFrom, Vector *pTo){
   float *dstF32;
   double *dstF64;
   u8 *dst1Bit;
+  u16 *dstF16;
 
   assert( pFrom->dims == pTo->dims );
   assert( pFrom->type != pTo->type );
@@ -705,6 +751,49 @@ static void vectorConvertFromF8(const Vector *pFrom, Vector *pTo){
         dst1Bit[i / 8] |= (1 << (i & 7));
       }
     }
+  }else if( pTo->type == VECTOR_TYPE_FLOAT16 ){
+    dstF16 = pTo->data;
+    for(i = 0; i < pFrom->dims; i++){
+      dstF16[i] = vectorF16FromFloat(alpha * src[i] + shift);
+    }
+  }else{
+    assert( 0 );
+  }
+}
+
+static void vectorConvertFromF16(const Vector *pFrom, Vector *pTo){
+  int i;
+  u16 *src;
+
+  float *dstF32;
+  double *dstF64;
+  u8 *dst1Bit;
+
+  assert( pFrom->dims == pTo->dims );
+  assert( pFrom->type != pTo->type );
+  assert( pFrom->type == VECTOR_TYPE_FLOAT16 );
+
+  src = pFrom->data;
+  if( pTo->type == VECTOR_TYPE_FLOAT32 ){
+    dstF32 = pTo->data;
+    for(i = 0; i < pFrom->dims; i++){
+      dstF32[i] = vectorF16ToFloat(src[i]);
+    }
+  }else if( pTo->type == VECTOR_TYPE_FLOAT64 ){
+    dstF64 = pTo->data;
+    for(i = 0; i < pFrom->dims; i++){
+      dstF64[i] = vectorF16ToFloat(src[i]);
+    }
+  }else if( pTo->type == VECTOR_TYPE_FLOAT1BIT ){
+    dst1Bit = pTo->data;
+    for(i = 0; i < pFrom->dims; i += 8){
+      dst1Bit[i / 8] = 0;
+    }
+    for(i = 0; i < pFrom->dims; i++){
+      if( vectorF16ToFloat(src[i]) > 0 ){
+        dst1Bit[i / 8] |= (1 << (i & 7));
+      }
+    }
   }else{
     assert( 0 );
   }
@@ -730,6 +819,7 @@ static void vectorConvertToF8(const Vector *pFrom, Vector *pTo){
   float *srcF32;
   double *srcF64;
   u8 *src1Bit;
+  u16 *srcF16;
 
   assert( pFrom->dims == pTo->dims );
   assert( pFrom->type != pTo->type );
@@ -766,6 +856,16 @@ static void vectorConvertToF8(const Vector *pFrom, Vector *pTo){
     for(i = 0; i < pFrom->dims; i++){
       dst[i] = clip(((((src1Bit[i / 8] >> (i & 7)) & 1) ? +1 : -1) - shift) / alpha, 0, 255);
     }
+  }else if( pFrom->type == VECTOR_TYPE_FLOAT16 ){
+    srcF16 = pFrom->data;
+    for(i = 0; i < pFrom->dims; i++){
+      MINMAX(i, vectorF16ToFloat(srcF16[i]), minF, maxF);
+    }
+    shift = minF;
+    alpha = (maxF - minF) / 255;
+    for(i = 0; i < pFrom->dims; i++){
+      dst[i] = clip((vectorF16ToFloat(srcF16[i]) - shift) / alpha, 0, 255);
+    }
   }else{
     assert( 0 );
   }
@@ -791,6 +891,8 @@ void vectorConvert(const Vector *pFrom, Vector *pTo){
     vectorConvertFrom1Bit(pFrom, pTo);
   }else if( pFrom->type == VECTOR_TYPE_FLOAT8 ){
     vectorConvertFromF8(pFrom, pTo);
+  }else if( pFrom->type == VECTOR_TYPE_FLOAT16 ){
+    vectorConvertFromF16(pFrom, pTo);
   }else{
     assert( 0 );
   }
@@ -875,6 +977,14 @@ static void vector8Func(
   vectorFuncHintedType(context, argc, argv, VECTOR_TYPE_FLOAT8);
 }
 
+static void vector16Func(
+  sqlite3_context *context,
+  int argc,
+  sqlite3_value **argv
+){
+  vectorFuncHintedType(context, argc, argv, VECTOR_TYPE_FLOAT16);
+}
+
 static void vector1BitFunc(
   sqlite3_context *context,
   int argc,
@@ -1033,6 +1143,7 @@ void sqlite3RegisterVectorFunctions(void){
     FUNCTION(vector64,            1, 0, 0, vector64Func),
     FUNCTION(vector1bit,          1, 0, 0, vector1BitFunc),
     FUNCTION(vector8,             1, 0, 0, vector8Func),
+    FUNCTION(vector16,            1, 0, 0, vector16Func),
     FUNCTION(vector_extract,      1, 0, 0, vectorExtractFunc),
     FUNCTION(vector_distance_cos, 2, 0, 0, vectorDistanceCosFunc),
     FUNCTION(vector_distance_l2,  2, 0, 0, vectorDistanceL2Func),
diff --git a/libsql-sqlite3/src/vectorIndex.c b/libsql-sqlite3/src/vectorIndex.c
@@ -385,6 +385,8 @@ static struct VectorColumnType VECTOR_COLUMN_TYPES[] = {
   { "F1BIT_BLOB", VECTOR_TYPE_FLOAT1BIT },
   { "FLOAT8",     VECTOR_TYPE_FLOAT8 },
   { "F8_BLOB",    VECTOR_TYPE_FLOAT8 },
+  { "FLOAT16",    VECTOR_TYPE_FLOAT16 },
+  { "F16_BLOB",   VECTOR_TYPE_FLOAT16 },
 };
 
 /*
@@ -405,6 +407,7 @@ static struct VectorParamName VECTOR_PARAM_NAMES[] = {
   { "metric",             VECTOR_METRIC_TYPE_PARAM_ID,        0, "l2",        VECTOR_METRIC_TYPE_L2 },
   { "compress_neighbors", VECTOR_COMPRESS_NEIGHBORS_PARAM_ID, 0, "float1bit", VECTOR_TYPE_FLOAT1BIT },
   { "compress_neighbors", VECTOR_COMPRESS_NEIGHBORS_PARAM_ID, 0, "float8",    VECTOR_TYPE_FLOAT8 },
+  { "compress_neighbors", VECTOR_COMPRESS_NEIGHBORS_PARAM_ID, 0, "float16",   VECTOR_TYPE_FLOAT16 },
   { "compress_neighbors", VECTOR_COMPRESS_NEIGHBORS_PARAM_ID, 0, "float32",   VECTOR_TYPE_FLOAT32 },
   { "alpha",              VECTOR_PRUNING_ALPHA_PARAM_ID, 2, 0, 0 },
   { "search_l",           VECTOR_SEARCH_L_PARAM_ID,      1, 0, 0 },
diff --git a/libsql-sqlite3/src/vectorInt.h b/libsql-sqlite3/src/vectorInt.h
@@ -150,6 +150,9 @@ void vector1BitDeserializeFromBlob(Vector *, const unsigned char *, size_t);
 void vectorInitStatic(Vector *, VectorType, VectorDims, void *);
 void vectorInitFromBlob(Vector *, const unsigned char *, size_t);
 
+u16 vectorF16FromFloat(float);
+float vectorF16ToFloat(u16);
+
 void vectorConvert(const Vector *, Vector *);
 
 /* Detect type and dimension of vector provided with first parameter of sqlite3_value * type */
diff --git a/libsql-sqlite3/src/vectorfloat16.c b/libsql-sqlite3/src/vectorfloat16.c
@@ -23,6 +23,8 @@
 ******************************************************************************
 **
 ** 16-bit (FLOAT16) floating point vector format utilities.
+**
+** See https://en.wikipedia.org/wiki/Half-precision_floating-point_format
 */
 #ifndef SQLITE_OMIT_VECTOR
 #include "sqliteInt.h"
@@ -40,7 +42,7 @@
 // f16: [ffffffffffeeeees]
 //       0123456789012345
 
-static float vectorF16ToFloat(u16 f16){
+float vectorF16ToFloat(u16 f16){
   u32 f32;
   // sng: [0000000000000000000000000000000s]
   u32 sgn = ((u32)f16 & 0x8000) << 16;   
@@ -72,7 +74,7 @@ static float vectorF16ToFloat(u16 f16){
   return *((float*)&f32);
 }
 
-static u16 vectorF16FromFloat(float f){
+u16 vectorF16FromFloat(float f){
   u32 i = *((u32*)&f);
 
   // sng: [000000000000000s]
@@ -160,7 +162,7 @@ float vectorF16DistanceL2(const Vector *v1, const Vector *v2){
   int i;
   float sum = 0;
   float value1, value2;
-  u8 *data1 = v1->data, *data2 = v2->data;
+  u16 *data1 = v1->data, *data2 = v2->data;
 
   assert( v1->dims == v2->dims );
   assert( v1->type == VECTOR_TYPE_FLOAT16 );
diff --git a/libsql-sqlite3/test/libsql_vector.test b/libsql-sqlite3/test/libsql_vector.test
diff --git a/libsql-sqlite3/test/libsql_vector_index.test b/libsql-sqlite3/test/libsql_vector_index.test
