Star181
850
代码
代码
Issues157
Pull Requests329
流水线
项目讨论
Wiki
分析
Star181
850
已合并
[fix]split_with_sizes_copy performance reduction #3466
zhangqiongwen创建于 2025年11月11日
[fix]split_with_sizes_copy performance reduction #3466
已合并
zhangqiongwen创建于 2025年11月11日
从zqwenn/op-pluginmaster_split_fix合入到Ascend/op-pluginmaster
共 5 个文件变更+40-52
Mop_plugin/config/op_plugin_functions.yaml+0-5
@@ -5415,13 +5415,8 @@ official:
5415 - func: unfold_backward(Tensor grad_in, SymInt[] input_sizes, int dim, int size, int step) -> Tensor5415 - func: unfold_backward(Tensor grad_in, SymInt[] input_sizes, int dim, int size, int step) -> Tensor
5416 op_api: [v2.1, newest]5416 op_api: [v2.1, newest]
5417 5417 
5418- - func: split_with_sizes_copy(Tensor self, SymInt[] split_sizes, int dim=0) -> Tensor[]
5419- op_api: [ v2.1, newest ]
5420- 
5421 - func: split_with_sizes_copy.out(Tensor self, SymInt[] split_sizes, int dim=0, *, Tensor(a!)[] out) -> ()5418 - func: split_with_sizes_copy.out(Tensor self, SymInt[] split_sizes, int dim=0, *, Tensor(a!)[] out) -> ()
5422 op_api: [ v2.1, newest ]5419 op_api: [ v2.1, newest ]
5423- gen_opapi:
5424- exec: aclnnSplitWithSize
5425 5420 
5426autograd:5421autograd:
5427 - func: _thnn_fused_gru_cell(Tensor input_gates, Tensor hidden_gates, Tensor hx, Tensor? input_bias=None, Tensor? hidden_bias=None) -> (Tensor, Tensor)5422 - func: _thnn_fused_gru_cell(Tensor input_gates, Tensor hidden_gates, Tensor hx, Tensor? input_bias=None, Tensor? hidden_bias=None) -> (Tensor, Tensor)
Mop_plugin/ops/opapi/SplitWithSizesCopyNpuOpApi.cpp+40-10
@@ -17,20 +17,50 @@
17namespace op_api {17namespace op_api {
18using npu_preparation = at_npu::native::OpPreparation;18using npu_preparation = at_npu::native::OpPreparation;
19 19 
20+const static int64_t LEN_MIN = 32;
21+const static int64_t LEN_MAX = 64;
22+const static int64_t DIM_2 = 2;
23+const static int64_t DIM_MIN_16 = 16;
24+const static int64_t DIM_MIN_8 = 8;
25+const static int64_t DIM_MAX = 65536;
20 26 
21-std::vector<at::Tensor> split_with_sizes_copy(const at::Tensor & self, at::IntArrayRef split_sizes, int64_t dim)27+bool is_fused_op_optim(const at::Tensor& self, at::IntArrayRef split_sizes)
22{28{
23- auto output_shapes = op_infer::split_with_sizes_copy_output_size(op_infer::array_to_small_vector(self.sizes()), split_sizes, dim);29+ if (!op_plugin::utils::is_gte_cann_version_830rc1()) {
24- auto output_dtype = self.scalar_type();30+ return false;
25- 
26- std::vector<at::Tensor> result;
27- for (const auto& shape : output_shapes) {
28- result.push_back(npu_preparation::apply_tensor_without_format(shape, self.options().dtype(output_dtype)));
29 }31 }
30 32 
31- at::TensorList result_ = at::TensorList(result);33+ if (self.dim() != DIM_2) {
32- EXEC_NPU_CMD(aclnnSplitWithSize, self, split_sizes, dim, result_);34+ return false;
33- return result;35+ }
36+ 
37+ int64_t len = split_sizes.size();
38+ if ((len <= LEN_MIN) || (len > LEN_MAX)) {
39+ return false;
40+ }
41+ 
42+ at::ScalarType dtype = self.scalar_type();
43+ 
44+ int64_t dim0 = self.size(0);
45+ int64_t dim1 = self.size(1);
46+ 
47+ if (dtype == at::ScalarType::Half || dtype == at::ScalarType::BFloat16) {
48+ return (dim0 <= DIM_MIN_16) && (dim1 > DIM_MAX);
49+ } else if (dtype == at::ScalarType::Float) {
50+ return (dim0 <= DIM_MIN_8) && (dim1 > DIM_MAX);
51+ } else {
52+ return false;
53+ }
54+}
55+ 
56+ 
57+void split_with_sizes_copy_out(const at::Tensor & self, at::IntArrayRef split_sizes, int64_t dim, at::TensorList out)
58+{
59+ if (is_fused_op_optim(self, split_sizes)) {
60+ EXEC_NPU_CMD(aclnnSplitWithSize, self, split_sizes, dim, out);
61+ } else {
62+ at::native::split_with_sizes_copy_out(self, split_sizes, dim, out);
O
OopenLiBingCI2025年11月11日

代码可维护性: 注释掉的代码应该被清理，这些代码在当前版本中不再使用，但保留在代码库中会造成混淆和维护困难。

问题类型: 代码可维护性 文件路径: op_plugin/ops/opapi/SplitWithSizesCopyNpuOpApi.cpp 行号: 55 问题代码:

// auto array = self.split_with_sizes(split_sizes, dim);
// copy_tensor_array_to_out("split_with_sizes_copy_out", array, out);

修改建议:

删除注释掉的代码，或者如果这些代码有参考价值，应该移动到专门的文档中

此评论由代码审查工具自动生成

likedislike
63+ }
34}64}
35 65 
36} // namespace op_api66} // namespace op_api
Mop_plugin/utils/KernelNpuOutputSize.cpp+0-33
@@ -2152,39 +2152,6 @@ c10::SmallVector<int64_t, SIZE> npu_moe_token_unpermute_grad_probs_out_size(cons
2152 }2152 }
2153}2153}
2154 2154 
2155-c10::SmallVector<c10::SmallVector<int64_t, SIZE>, SIZE> split_with_sizes_copy_output_size(
2156- const c10::SmallVector<int64_t, SIZE>& input_shape,
2157- const c10::IntArrayRef split_sizes,
2158- int64_t dim)
2159-{
2160- const int64_t ndim = static_cast<int64_t>(input_shape.size());
2161- 
2162- // Wrap negative dim
2163- if (dim < 0) {
2164- dim += ndim;
2165- }
2166- 
2167- TORCH_CHECK(dim >= 0 && dim < ndim, "Invalid dimension: ", dim, OPS_ERROR(ErrCode::PARAM));
2168- 
2169- // Check total size
2170- int64_t total = 0;
2171- for (int64_t size : split_sizes) {
2172- TORCH_CHECK(size >= 0, "split_sizes must be non-negative", OPS_ERROR(ErrCode::PARAM));
2173- total += size;
2174- }
2175- TORCH_CHECK(total == input_shape[dim],
2176- "Sum of split_sizes (", total, ") must equal input_shape[dim] (", input_shape[dim], ")", OPS_ERROR(ErrCode::PARAM));
2177- 
2178- // Construct output shapes
2179- c10::SmallVector<c10::SmallVector<int64_t, SIZE>, SIZE> output_shapes;
2180- for (int64_t size : split_sizes) {
2181- c10::SmallVector<int64_t, SIZE> shape = input_shape;
2182- shape[dim] = size;
2183- output_shapes.emplace_back(std::move(shape));
2184- }
2185- 
2186- return output_shapes;
2187-}
2188 2155 
2189std::tuple<c10::SmallVector<int64_t, SIZE>, c10::SmallVector<int64_t, SIZE>> triangular_solve_output_size(const at::Tensor& self, const at::Tensor& A)2156std::tuple<c10::SmallVector<int64_t, SIZE>, c10::SmallVector<int64_t, SIZE>> triangular_solve_output_size(const at::Tensor& self, const at::Tensor& A)
2190{2157{
Mop_plugin/utils/KernelNpuOutputSize.h+0-1
@@ -376,7 +376,6 @@ OP_PLUGIN_HIDDEN c10::SmallVector<int64_t, SIZE> npu_moe_token_unpermute_out_siz
376OP_PLUGIN_HIDDEN c10::SmallVector<int64_t, SIZE> npu_moe_token_permute_grad_out_size(const at::Tensor &tokens, const at::Tensor &permuted_tokens, const at::Tensor &indices, const at::Tensor &sorted_indices);376OP_PLUGIN_HIDDEN c10::SmallVector<int64_t, SIZE> npu_moe_token_permute_grad_out_size(const at::Tensor &tokens, const at::Tensor &permuted_tokens, const at::Tensor &indices, const at::Tensor &sorted_indices);
377OP_PLUGIN_HIDDEN c10::SmallVector<int64_t, SIZE> npu_moe_token_unpermute_grad_permuted_tokens_out_size(const at::Tensor &permuted_tokens, const at::Tensor &grad_unpermuted_tokens, const at::Tensor &sorted_indices, const c10::optional<at::Tensor> &probs);377OP_PLUGIN_HIDDEN c10::SmallVector<int64_t, SIZE> npu_moe_token_unpermute_grad_permuted_tokens_out_size(const at::Tensor &permuted_tokens, const at::Tensor &grad_unpermuted_tokens, const at::Tensor &sorted_indices, const c10::optional<at::Tensor> &probs);
378OP_PLUGIN_HIDDEN c10::SmallVector<int64_t, SIZE> npu_moe_token_unpermute_grad_probs_out_size(const at::Tensor &permuted_tokens, const at::Tensor &grad_unpermuted_tokens, const at::Tensor &sorted_indices, const c10::optional<at::Tensor> &probs);378OP_PLUGIN_HIDDEN c10::SmallVector<int64_t, SIZE> npu_moe_token_unpermute_grad_probs_out_size(const at::Tensor &permuted_tokens, const at::Tensor &grad_unpermuted_tokens, const at::Tensor &sorted_indices, const c10::optional<at::Tensor> &probs);
379-OP_PLUGIN_HIDDEN c10::SmallVector<c10::SmallVector<int64_t, SIZE>, SIZE> split_with_sizes_copy_output_size(const c10::SmallVector<int64_t, SIZE>& input_shape, const c10::IntArrayRef split_sizes, int64_t dim);
380 379 
381OP_PLUGIN_HIDDEN std::tuple<c10::SmallVector<int64_t, SIZE>, c10::SmallVector<int64_t, SIZE>> triangular_solve_output_size(const at::Tensor& self, const at::Tensor& A);380OP_PLUGIN_HIDDEN std::tuple<c10::SmallVector<int64_t, SIZE>, c10::SmallVector<int64_t, SIZE>> triangular_solve_output_size(const at::Tensor& self, const at::Tensor& A);
382 381 
Mtest/core_tests/torch_npu_OpApi_schema_all.json+0-3
@@ -4301,9 +4301,6 @@
4301 "func: unfold_backward(Tensor grad_in, SymInt[] input_sizes, int dim, int size, int step) -> Tensor": {4301 "func: unfold_backward(Tensor grad_in, SymInt[] input_sizes, int dim, int size, int step) -> Tensor": {
4302 "version": ["all_version"]4302 "version": ["all_version"]
4303 },4303 },
4304- "func: split_with_sizes_copy(Tensor self, SymInt[] split_sizes, int dim=0) -> Tensor[]": {
4305- "version": ["all_version"]
4306- },
4307 "func: split_with_sizes_copy.out(Tensor self, SymInt[] split_sizes, int dim=0, *, Tensor(a!)[] out) -> ()": {4304 "func: split_with_sizes_copy.out(Tensor self, SymInt[] split_sizes, int dim=0, *, Tensor(a!)[] out) -> ()": {
4308 "version": ["all_version"]4305 "version": ["all_version"]
4309 },4306 },