search

StijnWoestenborghs / gradi-mojo

35 stars 3 forks source link

Vectorizing squared distance calculation #4

Open shashankprasanna opened 9 months ago

shashankprasanna commented 9 months ago

Hi👋 Thank you for this awesome project! I can't wait to try it out!

Just wanted to share that you can further speed up the squared distance calculation by vectorizing your code by operating on simd_width vectors and using the stdlib math functions. For example in the squared_distance calculation you could vectorize 

for d in range(X.cols):
    squared_distance += (X[i, d] - X[j, d])**2

With something like this:

@parameter
fn dist(idx: Int)->None:
    let diff = X.load[simd_width](i*X.cols)-X.load[simd_width](j*cols)
    squared_distance += (diff*diff).reduce_add()
vectorize[simd_width, diff](X.cols)

Other places to vectorize: transpose, gradient_descent

Thanks for sharing this with the Mojo community!

CAClaveau commented 9 months ago

(Hi, sorry, I didn't want to open a separate thread just for the following suggestion: I think it would be super interesting to also include an implementation in Julia and one based on DACE for comparison. In any case, thank you very much for sharing this project and the results you got in the article posted on your website!)