Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion mmcv/ops/csrc/common/pytorch_npu_helper.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
#ifndef PYTORCH_NPU_HELPER_HPP_
#define PYTORCH_NPU_HELPER_HPP_

#include <torch_npu/csrc/aten/CustomFunctions.h>
#include <torch_npu/csrc/framework/utils/CalcuOpUtil.h>
#include <torch_npu/csrc/framework/utils/OpAdapter.h>

Expand Down
16 changes: 14 additions & 2 deletions mmcv/ops/csrc/pytorch/npu/focal_loss_npu.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -99,8 +99,20 @@ void softmax_focal_loss_forward_npu(Tensor input, Tensor target, Tensor weight,
c10::SmallVector<int64_t, 2> offsets = {0, 0};
c10::SmallVector<int64_t, 2> sizes = {n_batch, 1};
at::IntArrayRef offset = at::IntArrayRef(offsets);
at::IntArrayRef size = at::IntArrayRef(sizes);
at_npu::native::custom_ops::npu_slice_out(op_output, offset, size, output);
at::IntArrayRef size_array = at::IntArrayRef(sizes);
c10::SmallVector<int64_t, N> output_size;
for (uint64_t i = 0; i < size_array.size(); i++) {
output_size.emplace_back(size_array[i]);
}
at::Tensor result = at::empty(output_size, op_output.options());
c10::SmallVector<int64_t, N> offsetVec = array_to_small_vector(offset);
c10::SmallVector<int64_t, N> sizeVec = array_to_small_vector(size_array);
cmd.Name("Slice")
.Input(op_output)
.Input(offsetVec)
.Input(sizeVec)
.Output(output)
.Run();
}

void softmax_focal_loss_forward_impl(Tensor input, Tensor target, Tensor weight,
Expand Down