Skip to content

Commit 54fb06e

Browse files
committed
torch_npu support aclnn
1 parent 3718a0d commit 54fb06e

File tree

2 files changed

+10
-2
lines changed

2 files changed

+10
-2
lines changed

mmcv/ops/csrc/common/pytorch_npu_helper.hpp

-1
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
#ifndef PYTORCH_NPU_HELPER_HPP_
1919
#define PYTORCH_NPU_HELPER_HPP_
2020

21-
#include <torch_npu/csrc/aten/CustomFunctions.h>
2221
#include <torch_npu/csrc/framework/utils/CalcuOpUtil.h>
2322
#include <torch_npu/csrc/framework/utils/OpAdapter.h>
2423

mmcv/ops/csrc/pytorch/npu/focal_loss_npu.cpp

+10-1
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,16 @@ void softmax_focal_loss_forward_npu(Tensor input, Tensor target, Tensor weight,
100100
c10::SmallVector<int64_t, 2> sizes = {n_batch, 1};
101101
at::IntArrayRef offset = at::IntArrayRef(offsets);
102102
at::IntArrayRef size = at::IntArrayRef(sizes);
103-
at_npu::native::custom_ops::npu_slice_out(op_output, offset, size, output);
103+
at::IntArrayRef size_array = at::IntArrayRef(sizes);
104+
c10::SmallVector<int64_t, N> offsetVec = array_to_small_vector(offset);
105+
c10::SmallVector<int64_t, N> sizeVec = array_to_small_vector(size_array);
106+
OpCommand cmd2;
107+
cmd2.Name("Slice")
108+
.Input(op_output)
109+
.Input(offsetVec)
110+
.Input(sizeVec)
111+
.Output(output)
112+
.Run();
104113
}
105114

106115
void softmax_focal_loss_forward_impl(Tensor input, Tensor target, Tensor weight,

0 commit comments

Comments
 (0)