github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/vector/hnsw/distancer/l2_amd64_test.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package distancer
    13  
    14  import (
    15  	"fmt"
    16  	"math/rand"
    17  	"testing"
    18  
    19  	"github.com/stretchr/testify/assert"
    20  	"github.com/weaviate/weaviate/adapters/repos/db/vector/hnsw/distancer/asm"
    21  	"golang.org/x/sys/cpu"
    22  )
    23  
    24  func L2PureGo(a, b []float32) float32 {
    25  	var sum float32
    26  
    27  	for i := range a {
    28  		diff := a[i] - b[i]
    29  		sum += diff * diff
    30  	}
    31  
    32  	return sum
    33  }
    34  
    35  func Test_L2_DistanceImplementation(t *testing.T) {
    36  	lengths := []int{1, 4, 16, 31, 32, 35, 64, 67, 128, 130, 256, 260, 384, 390, 768, 777}
    37  
    38  	for _, length := range lengths {
    39  		t.Run(fmt.Sprintf("with vector l=%d", length), func(t *testing.T) {
    40  			x := make([]float32, length)
    41  			y := make([]float32, length)
    42  			for i := range x {
    43  				x[i] = rand.Float32()
    44  				y[i] = rand.Float32()
    45  			}
    46  
    47  			control := L2PureGo(x, y)
    48  
    49  			asmResult := asm.L2AVX256(x, y)
    50  			assert.InEpsilon(t, control, asmResult, 0.01)
    51  
    52  			if cpu.X86.HasAVX512 {
    53  				asmResult = asm.L2AVX512(x, y)
    54  				assert.InEpsilon(t, control, asmResult, 0.01)
    55  			}
    56  		})
    57  	}
    58  }
    59  
    60  func Test_L2_DistanceImplementation_OneNegativeValue(t *testing.T) {
    61  	lengths := []int{1, 4, 16, 31, 32, 35, 64, 67, 128, 130, 256, 260, 384, 390, 768, 777}
    62  
    63  	for _, length := range lengths {
    64  		t.Run(fmt.Sprintf("with vector l=%d", length), func(t *testing.T) {
    65  			x := make([]float32, length)
    66  			y := make([]float32, length)
    67  			for i := range x {
    68  				x[i] = -rand.Float32()
    69  				y[i] = rand.Float32()
    70  			}
    71  
    72  			control := L2PureGo(x, y)
    73  
    74  			asmResult := asm.L2AVX256(x, y)
    75  			assert.InEpsilon(t, control, asmResult, 0.01)
    76  
    77  			if cpu.X86.HasAVX512 {
    78  				asmResult = asm.L2AVX512(x, y)
    79  				assert.InEpsilon(t, control, asmResult, 0.01)
    80  			}
    81  		})
    82  	}
    83  }
    84  
    85  func Benchmark_L2(b *testing.B) {
    86  	r := getRandomSeed()
    87  	lengths := []int{
    88  		1,
    89  		2,
    90  		3,
    91  		4,
    92  		5,
    93  		6,
    94  		8,
    95  		10,
    96  		12,
    97  		16,
    98  		24,
    99  		30,
   100  		31,
   101  		32,
   102  		64,
   103  		67,
   104  		128,
   105  		256,
   106  		260,
   107  		299,
   108  		300,
   109  		384,
   110  		390,
   111  		600,
   112  		768,
   113  		777,
   114  		784,
   115  		1024,
   116  		1536,
   117  	}
   118  	for _, length := range lengths {
   119  		b.Run(fmt.Sprintf("vector dim=%d", length), func(b *testing.B) {
   120  			x := make([]float32, length)
   121  			y := make([]float32, length)
   122  			for i := range x {
   123  				x[i] = -r.Float32()
   124  				y[i] = r.Float32()
   125  			}
   126  
   127  			b.Run("pure go", func(b *testing.B) {
   128  				for i := 0; i < b.N; i++ {
   129  					L2PureGo(x, y)
   130  				}
   131  			})
   132  
   133  			b.Run("asm AVX", func(b *testing.B) {
   134  				for i := 0; i < b.N; i++ {
   135  					asm.L2(x, y)
   136  				}
   137  			})
   138  
   139  			b.Run("asm AVX512", func(b *testing.B) {
   140  				for i := 0; i < b.N; i++ {
   141  					asm.L2AVX512(x, y)
   142  				}
   143  			})
   144  		})
   145  	}
   146  }