Skip to content

Commit

Permalink
[Feature] Support MsDeformAttnForward with fast kernel
Browse files Browse the repository at this point in the history
  • Loading branch information
DanieeelLiu committed Aug 2, 2024
1 parent bfdd1f9 commit f186407
Show file tree
Hide file tree
Showing 3 changed files with 873 additions and 1 deletion.
23 changes: 23 additions & 0 deletions mmcv/ops/csrc/common/mlu/ms_deform_attn_fast_mlu_kernel.hpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
/*************************************************************************
* Copyright (C) [2024] by Cambricon, Inc.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
* OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
* MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
* IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
* CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
* TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*************************************************************************/
#ifndef MS_DEFORM_ATTN_FORWARD_FAST_MLU_KERNEL_HPP_
#define MS_DEFORM_ATTN_FORWARD_FAST_MLU_KERNEL_HPP_
void KernelMsDeformAttnForwardFast(
cnrtDim3_t k_dim, cnrtFunctionType_t k_type, cnrtQueue_t queue,
const cnrtDataType_t d_type, const char *data_value_gdram,
const char *data_spatial_shapes_gdram,
const char *data_level_start_index_gdram,
const char *data_sampling_loc_gdram, const char *data_attn_weight_gdram,
const int32_t batch_size, const int32_t num_keys, const int32_t num_heads,
const int32_t channels, const int32_t num_levels, const int32_t num_queries,
const int32_t num_points, char *data_col_gdram);
#endif // MS_DEFORM_ATTN_FORWARD_FAST_MLU_KERNEL_HPP_
Loading

0 comments on commit f186407

Please sign in to comment.