Skip to content

Commit

Permalink
[NPU] npu msda supports aclnn (#3149)
Browse files Browse the repository at this point in the history
  • Loading branch information
Pr0Wh1teGivee authored Jul 17, 2024
1 parent 8f23a0b commit 44eab26
Showing 1 changed file with 3 additions and 9 deletions.
12 changes: 3 additions & 9 deletions mmcv/ops/csrc/pytorch/npu/ms_deform_attn_npu.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -57,15 +57,9 @@ Tensor ms_deform_attn_forward_npu(const Tensor &value,
value.size(0), sampling_locations.size(1), value.size(2) * value.size(3)};
at::Tensor output = at::zeros(output_size, value_fp32.options());

OpCommand cmd;
cmd.Name("MultiScaleDeformableAttnFunction")
.Input(value_fp32)
.Input(value_spatial_shapes_int32)
.Input(value_level_start_index_int32)
.Input(sampling_locations_fp32)
.Input(attention_weights_fp32)
.Output(output)
.Run();
EXEC_NPU_CMD(aclnnMultiScaleDeformableAttnFunction, value_fp32,
value_spatial_shapes_int32, value_level_start_index_int32,
sampling_locations_fp32, attention_weights_fp32, output)

at::Tensor real_output = output;
if (value.scalar_type() != at::kFloat) {
Expand Down

0 comments on commit 44eab26

Please sign in to comment.