optimize embedding #1891
optimize embedding #1891
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR optimizes the embedding operation for index select kernels on Intel XPU by introducing specialized embedding kernel functors. The optimization addresses performance issues where the previous implementation used incorrect index kernel configurations leading to low GPU occupancy.
- Adds dedicated
EmbeddingKernelFunctorandEmbeddingKernelSLMFunctorclasses for optimized embedding operations - Implements dynamic kernel selection based on shared local memory availability
- Integrates the new embedding path into the existing index select kernel for specific tensor configurations
Reviewed Changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| src/ATen/native/xpu/sycl/Indexing.h | Adds new embedding kernel functor classes with optimized launch configurations |
| src/ATen/native/xpu/sycl/Indexing.cpp | Implements embedding function and integrates it into index select kernel for 2D contiguous tensors |
| for (auto thread_id = item.get_global_linear_id(); | ||
| thread_id < indices_length_ * embedding_dim_; | ||
| thread_id += item.get_local_range(0) * item.get_group_range(0)) { | ||
| SYCL_KERNEL_ASSERT(index_[thread_id / embedding_dim_] < num_embeddings_); |
There was a problem hiding this comment.
The assertion index_[thread_id / embedding_dim_] < num_embeddings_ is evaluated inside the tight loop for every thread iteration. Consider moving this validation outside the performance-critical path or using a debug-only assertion to avoid overhead in production builds.
Copilot uses AI. Check for mistakes.
| for (auto thread_id = item.get_global_linear_id(); | ||
| thread_id < indices_length_ * embedding_dim_; | ||
| thread_id += item.get_local_range(0) * item.get_group_range(0)) { | ||
| SYCL_KERNEL_ASSERT(index_[thread_id / embedding_dim_] < num_embeddings_); |
There was a problem hiding this comment.
Similar to the regular embedding kernel, this assertion inside the loop may impact performance. The same index bounds check is duplicated between both kernel functors and could benefit from optimization or debug-only execution.
| SYCL_KERNEL_ASSERT(index_[thread_id / embedding_dim_] < num_embeddings_); | |
| #ifdef DEBUG | |
| if (index_[thread_id / embedding_dim_] >= num_embeddings_) { | |
| SYCL_KERNEL_ASSERT(false && "Index out of bounds in EmbeddingKernelSLMFunctor"); | |
| } | |
| #endif |
Copilot uses AI. Check for mistakes.
| src_info, dst_info, index_info, new_indexing_dim); | ||
| else | ||
| if (dst.is_contiguous() && indices.is_contiguous()) { | ||
| if (src.dim() == 2 && indices.dim() == 1 && src.is_contiguous()) { |
There was a problem hiding this comment.
The condition src.is_contiguous() is redundant here since it's already checked in the parent condition on line 244. This duplicate check adds unnecessary complexity to the conditional logic.
| if (src.dim() == 2 && indices.dim() == 1 && src.is_contiguous()) { | |
| if (src.dim() == 2 && indices.dim() == 1) { |
Copilot uses AI. Check for mistakes.
Co-authored-by: Copilot <[email protected]>
Current index select kernel uses index kernel config, which sets wrong batch and problem batch, result to wrong launch configs that may lead to low occupancy. Also, inside kernel, half of thread is skipped in the following case:
index: 56103808
embedding table: [1683, 1]
previous on pvc: 19ms
now on pvc: 6.7ms
We will try to improve index kernel config next step. @yucai-intel