opensearch-project / k-NN

🆕 Find the k-nearest neighbors (k-NN) for your vector data
https://opensearch.org/docs/latest/search-plugins/knn/index/
Apache License 2.0
152 stars 113 forks source link

Add AVX512 support to k-nn FAISS #2069

Closed akashsha1 closed 1 week ago

akashsha1 commented 2 weeks ago

Description

This change adds support to speed up vector search and indexing in faiss using AVX512 hardware accelerator.

Related Issues

Resolves #2056

Check List

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license. For more information on following Developer Certificate of Origin and signing off your commits, please check here.

naveentatikonda commented 2 weeks ago

@akashsha1 Can you also update .github/workflows/CI.yml and .github/workflows/test_security.yml to detect if avx512 is supported or not and pass the flag to gradle.

Also, pls add an entry in CHANGELOG for this PR and sign off the PR to fix DCO check. Thanks!

naveentatikonda commented 2 weeks ago

Can you also please update build script to build all 3 versions of faiss lib for the distribution. Pls change this block of code like shown below

if [ "$PLATFORM" != "windows" ] && [ "$ARCHITECTURE" = "x64" ]; then
  echo "Building k-NN library after enabling AVX2"
  # Skip applying patches as patches were applied already from previous :buildJniLib task
  # If we apply patches again, it fails with conflict
  ./gradlew :buildJniLib -Davx2.enabled=true -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false

  echo "Building k-NN library after enabling AVX512"
  ./gradlew :buildJniLib -Davx512.enabled=true -Dbuild.lib.commit_patches=false -Dbuild.lib.apply_patches=false
fi
naveentatikonda commented 2 weeks ago

@akashsha1 You also need to add the avx512 faiss library in plugin-security.policy to provide runtime permissions for loading the library. Pls test end to end after making all these changes. Thanks!

naveentatikonda commented 2 weeks ago

@akashsha1 Looks like some of the commits were not properly signed causing the DCO to fail. Do you want to rebase and squash everything into a single commit and sign that commit ?

image

naveentatikonda commented 2 weeks ago

@akashsha1 We need to add the CPU flag detection logic in these 2 workflows and pass the gradle flags to fix the failing checks https://github.com/opensearch-project/k-NN/blob/main/.github/workflows/backwards_compatibility_tests_workflow.yml https://github.com/opensearch-project/k-NN/blob/main/.github/workflows/test_security.yml

akashsha1 commented 1 week ago

Hi @naveentatikonda - please consider this PR: https://github.com/opensearch-project/k-NN/pull/2110 in favor of the current one. The rebasing due to DCO sign-off got messy, so I created a new PR instead. It has all feedback addressed to this point, except the workflow failures which needs to be resolved.

akashsha1 commented 1 week ago

@akashsha1 We need to add the CPU flag detection logic in these 2 workflows and pass the gradle flags to fix the failing checks https://github.com/opensearch-project/k-NN/blob/main/.github/workflows/backwards_compatibility_tests_workflow.yml https://github.com/opensearch-project/k-NN/blob/main/.github/workflows/test_security.yml

This needs to be addressed in https://github.com/opensearch-project/k-NN/pull/2110 ..