-
Notifications
You must be signed in to change notification settings - Fork 49
add vectorization path on maxpool forward channel last #1883
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
jianyizh
wants to merge
3
commits into
main
Choose a base branch
from
jianyi/maxpool
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
Show all changes
3 commits
Select commit
Hold shift + click to select a range
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change | ||||
|---|---|---|---|---|---|---|
|
|
@@ -11,7 +11,9 @@ | |||||
|
|
||||||
| #include <ATen/native/xpu/sycl/Atomics.h> | ||||||
| #include <ATen/native/xpu/sycl/BatchKernel.h> | ||||||
| #include <ATen/native/xpu/sycl/MemoryAccess.h> | ||||||
| #include <ATen/native/xpu/sycl/NumericLimits.h> | ||||||
|
|
||||||
| #include <comm/Runtime.h> | ||||||
| #include <comm/SYCLHelpers.h> | ||||||
|
|
||||||
|
|
@@ -151,6 +153,119 @@ struct MaxPool2dKernelFunctor { | |||||
| BatchKernelConfig cfg_; | ||||||
| }; | ||||||
|
|
||||||
| template <typename scalar_t, typename vec_t, int vec_size> | ||||||
| struct MaxPool2dChannelLastVec { | ||||||
| void operator()(sycl::nd_item<1> item) const { | ||||||
| for (auto outputIndex = item.get_global_linear_id(); | ||||||
| outputIndex < numBatch_ * stride_ / vec_size; | ||||||
| outputIndex += item.get_local_range(0) * item.get_group_range(0)) { | ||||||
| int batch = outputIndex / (stride_ / vec_size); | ||||||
| int plane, outputH, outputW; | ||||||
| int64_t load_offset, store_offset; | ||||||
| plane = outputIndex % (numPlane_ / vec_size); | ||||||
| outputH = | ||||||
| outputIndex / (numPlane_ / vec_size) / outputSizeW_ % outputSizeH_; | ||||||
| outputW = outputIndex / (numPlane_ / vec_size) % outputSizeW_; | ||||||
| store_offset = outputIndex; | ||||||
|
|
||||||
| vec_t maxVal_vec; | ||||||
| #pragma unroll | ||||||
| for (int i = 0; i < vec_size; i++) { | ||||||
| maxVal_vec[i] = at::numeric_limits<scalar_t>::lower_bound(); | ||||||
| } | ||||||
| int64_t maxIndex[vec_size]; | ||||||
| for (int i = 0; i < vec_size; i++) { | ||||||
| maxIndex[i] = int64_t(-1); | ||||||
| } | ||||||
| int StartH = outputH * dH_ - padH_; | ||||||
| int StartW = outputW * dW_ - padW_; | ||||||
| int EndH = std::min(StartH + (kH_ - 1) * dilationH_ + 1, inputSizeH_); | ||||||
| int EndW = std::min(StartW + (kW_ - 1) * dilationW_ + 1, inputSizeW_); | ||||||
| while (StartH < 0) | ||||||
| StartH += dilationH_; | ||||||
| while (StartW < 0) | ||||||
| StartW += dilationW_; | ||||||
| for (int h = StartH; h < EndH; h += dilationH_) { | ||||||
| for (int w = StartW; w < EndW; w += dilationW_) { | ||||||
| load_offset = batch * inputSizeH_*inputSizeW_*numPlane_ / vec_size + plane + | ||||||
| h * inputSizeW_ * numPlane_ / vec_size + w * numPlane_ / vec_size; | ||||||
| vec_t val_vec = input_vec_[load_offset]; | ||||||
| #pragma unroll | ||||||
| for (int i = 0; i < vec_size; i++) { | ||||||
| if ((static_cast<scalar_t>(val_vec[i]) > maxVal_vec[i]) || | ||||||
| at::_isnan(val_vec[i])) { | ||||||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Using the private '_isnan' function instead of the standard 'std::isnan'. Consider using 'std::isnan' for better portability and standards compliance. Copilot uses AI. Check for mistakes. Positive FeedbackNegative Feedback |
||||||
| maxIndex[i] = h * inputSizeW_ + w; | ||||||
| maxVal_vec[i] = static_cast<scalar_t>(val_vec[i]); | ||||||
| } | ||||||
| } | ||||||
| } | ||||||
| } | ||||||
| #pragma unroll | ||||||
| for (int i = 0; i < vec_size; i++) { | ||||||
| indices_[store_offset * vec_size + i] = maxIndex[i]; | ||||||
| } | ||||||
| output_vec_[store_offset] = maxVal_vec; | ||||||
| } | ||||||
| } | ||||||
| MaxPool2dChannelLastVec( | ||||||
| vec_t* output_vec, | ||||||
| int64_t* indices, | ||||||
| const vec_t* input_vec, | ||||||
| int numBatch, | ||||||
| int numPlane, | ||||||
| int inputSizeH, | ||||||
| int inputSizeW, | ||||||
| int outputSizeH, | ||||||
| int outputSizeW, | ||||||
| int kH, | ||||||
| int kW, | ||||||
| int dH, | ||||||
| int dW, | ||||||
| int padH, | ||||||
| int padW, | ||||||
| int dilationH, | ||||||
| int dilationW, | ||||||
| int stride) | ||||||
| : output_vec_(output_vec), | ||||||
| indices_(indices), | ||||||
| input_vec_(input_vec), | ||||||
| numBatch_(numBatch), | ||||||
| numPlane_(numPlane), | ||||||
| inputSizeH_(inputSizeH), | ||||||
| inputSizeW_(inputSizeW), | ||||||
| outputSizeH_(outputSizeH), | ||||||
| outputSizeW_(outputSizeW), | ||||||
| kH_(kH), | ||||||
| kW_(kW), | ||||||
| dH_(dH), | ||||||
| dW_(dW), | ||||||
| padH_(padH), | ||||||
| padW_(padW), | ||||||
| dilationH_(dilationH), | ||||||
| dilationW_(dilationW), | ||||||
| stride_(stride) {} | ||||||
|
|
||||||
| private: | ||||||
| vec_t* output_vec_; | ||||||
| int64_t* indices_; | ||||||
| const vec_t* input_vec_; | ||||||
| int numBatch_; | ||||||
| int numPlane_; | ||||||
| int inputSizeH_; | ||||||
| int inputSizeW_; | ||||||
| int outputSizeH_; | ||||||
| int outputSizeW_; | ||||||
| int kH_; | ||||||
| int kW_; | ||||||
| int dH_; | ||||||
| int dW_; | ||||||
| int padH_; | ||||||
| int padW_; | ||||||
| int dilationH_; | ||||||
| int dilationW_; | ||||||
| int stride_; | ||||||
| }; | ||||||
|
|
||||||
| template <typename scalar_t, bool is_channels_last> | ||||||
| struct MaxPool2dBackwardKernelFunctor { | ||||||
| void operator()(sycl::nd_item<2> item) const { | ||||||
|
|
@@ -349,6 +464,56 @@ struct MaxPool2dBackwardDeterministicKernelFunctor { | |||||
| BatchKernelConfig cfg_; | ||||||
| }; | ||||||
|
|
||||||
| #define LAUNCH_MAXPOOL_CHANNEL_LAST_VEC( \ | ||||||
| scalar_t, \ | ||||||
| vec_size, \ | ||||||
| num_wg, \ | ||||||
| wg_size, \ | ||||||
| queue, \ | ||||||
| output, \ | ||||||
| indices, \ | ||||||
| input, \ | ||||||
| numBatch, \ | ||||||
| numPlane, \ | ||||||
| inputSizeH, \ | ||||||
| inputSizeW, \ | ||||||
| outputSizeH, \ | ||||||
| outputSizeW, \ | ||||||
| kH, \ | ||||||
| kW, \ | ||||||
| dH, \ | ||||||
| dW, \ | ||||||
| padH, \ | ||||||
| padW, \ | ||||||
| dilationH, \ | ||||||
| dilationW, \ | ||||||
| stride) \ | ||||||
| { \ | ||||||
| using vec_t = memory::aligned_vector<scalar_t, vec_size>; \ | ||||||
| vec_t* output_vec = reinterpret_cast<vec_t*>(output); \ | ||||||
| const vec_t* input_vec = reinterpret_cast<const vec_t*>(input); \ | ||||||
| auto kfn = MaxPool2dChannelLastVec<scalar_t, vec_t, vec_size>( \ | ||||||
| output_vec, \ | ||||||
| indices, \ | ||||||
| input_vec, \ | ||||||
| numBatch, \ | ||||||
| numPlane, \ | ||||||
| inputSizeH, \ | ||||||
| inputSizeW, \ | ||||||
| outputSizeH, \ | ||||||
| outputSizeW, \ | ||||||
| kH, \ | ||||||
| kW, \ | ||||||
| dH, \ | ||||||
| dW, \ | ||||||
| padH, \ | ||||||
| padW, \ | ||||||
| dilationH, \ | ||||||
| dilationW, \ | ||||||
| stride); \ | ||||||
| sycl_kernel_submit(num_wg * wg_size, wg_size, queue, kfn); \ | ||||||
| } | ||||||
|
|
||||||
| template <typename scalar_t, bool is_channels_last> | ||||||
| void launch_max_pool2d_kernel( | ||||||
| scalar_t* output, | ||||||
|
|
@@ -368,11 +533,114 @@ void launch_max_pool2d_kernel( | |||||
| int padW, | ||||||
| int dilationH, | ||||||
| int dilationW) { | ||||||
| using KernelClass = MaxPool2dKernelFunctor<scalar_t, is_channels_last>; | ||||||
|
|
||||||
| auto& queue = at::xpu::getCurrentSYCLQueue(); | ||||||
| int outputSize = numBatch * numPlane * outputSizeH * outputSizeW; | ||||||
| int stride = numPlane * outputSizeH * outputSizeW; | ||||||
| int vec_size = 1; | ||||||
| int thread_slots = syclGpuEuCount() * syclGpuHWThreadsPerEU(); | ||||||
| int num_sub_wg; | ||||||
| auto wg_size = syclDeviceMaxWorkGroupSize(); | ||||||
| int64_t num_wg; | ||||||
| if constexpr (is_channels_last) { | ||||||
| for (vec_size = | ||||||
| std::min(8, memory::can_vectorize_up_to<scalar_t>((char*)input)); | ||||||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. C-style cast to 'char*' should be replaced with a safer C++ cast like 'reinterpret_cast<char*>(input)' for better type safety and clarity.
Suggested change
Copilot uses AI. Check for mistakes. Positive FeedbackNegative Feedback |
||||||
| vec_size >= 1; | ||||||
| vec_size /= 2) { | ||||||
| if (numPlane % vec_size != 0) { | ||||||
| continue; | ||||||
| } | ||||||
| num_sub_wg = outputSize / vec_size / syclMaxSubGroupSize(); | ||||||
| if (2 * num_sub_wg > thread_slots) { | ||||||
| int total_thread = outputSize / vec_size; | ||||||
| num_wg = (total_thread + wg_size - 1) / wg_size; | ||||||
| break; | ||||||
| } | ||||||
| } | ||||||
| switch (vec_size) { | ||||||
| case 8: | ||||||
| LAUNCH_MAXPOOL_CHANNEL_LAST_VEC( | ||||||
| scalar_t, | ||||||
| 8, | ||||||
| num_wg, | ||||||
| wg_size, | ||||||
| queue, | ||||||
| output, | ||||||
| indices, | ||||||
| input, | ||||||
| numBatch, | ||||||
| numPlane, | ||||||
| inputSizeH, | ||||||
| inputSizeW, | ||||||
| outputSizeH, | ||||||
| outputSizeW, | ||||||
| kH, | ||||||
| kW, | ||||||
| dH, | ||||||
| dW, | ||||||
| padH, | ||||||
| padW, | ||||||
| dilationH, | ||||||
| dilationW, | ||||||
| stride); | ||||||
| return; | ||||||
| case 4: | ||||||
| LAUNCH_MAXPOOL_CHANNEL_LAST_VEC( | ||||||
| scalar_t, | ||||||
| 4, | ||||||
| num_wg, | ||||||
| wg_size, | ||||||
| queue, | ||||||
| output, | ||||||
| indices, | ||||||
| input, | ||||||
| numBatch, | ||||||
| numPlane, | ||||||
| inputSizeH, | ||||||
| inputSizeW, | ||||||
| outputSizeH, | ||||||
| outputSizeW, | ||||||
| kH, | ||||||
| kW, | ||||||
| dH, | ||||||
| dW, | ||||||
| padH, | ||||||
| padW, | ||||||
| dilationH, | ||||||
| dilationW, | ||||||
| stride); | ||||||
| return; | ||||||
| case 2: | ||||||
| LAUNCH_MAXPOOL_CHANNEL_LAST_VEC( | ||||||
| scalar_t, | ||||||
| 2, | ||||||
| num_wg, | ||||||
| wg_size, | ||||||
| queue, | ||||||
| output, | ||||||
| indices, | ||||||
| input, | ||||||
| numBatch, | ||||||
| numPlane, | ||||||
| inputSizeH, | ||||||
| inputSizeW, | ||||||
| outputSizeH, | ||||||
| outputSizeW, | ||||||
| kH, | ||||||
| kW, | ||||||
| dH, | ||||||
| dW, | ||||||
| padH, | ||||||
| padW, | ||||||
| dilationH, | ||||||
| dilationW, | ||||||
| stride); | ||||||
| return; | ||||||
| default: | ||||||
| break; | ||||||
| }; | ||||||
| } | ||||||
| using KernelClass = MaxPool2dKernelFunctor<scalar_t, is_channels_last>; | ||||||
|
|
||||||
| BatchKernelConfig cfg = BatchKernelConfig::make_config<KernelClass>( | ||||||
| 1, outputSize, 1, 1, true, BatchKernelConfig::Policy::pAdaptive); | ||||||
| auto kfn = KernelClass( | ||||||
|
|
@@ -704,6 +972,6 @@ void max_pool2d_with_indices_backward_kernel( | |||||
| } | ||||||
|
|
||||||
| } // namespace at::native::xpu | ||||||
|
|
||||||
| #undef LAUNCH_MAXPOOL_CHANNEL_LAST_VEC | ||||||
| #pragma GCC diagnostic pop | ||||||
| #pragma clang diagnostic pop | ||||||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[nitpick] The complex offset calculation spans multiple lines and lacks spacing around operators. Consider breaking this into intermediate variables or adding consistent spacing (e.g., 'inputSizeH_ * inputSizeW_ * numPlane_').
Copilot uses AI. Check for mistakes.