Star181
852
代码
代码
Issues161
Pull Requests323
流水线
项目讨论
Wiki
分析
Star181
852
已合并
add logit and logit_backward #2576
AtomGit-Bot创建于 2025年6月4日
add logit and logit_backward #2576
已合并
AtomGit-Bot创建于 2025年6月4日
从refs/pull/2576/head合入到master
共 4 个文件变更+206-0
Mop_plugin/config/op_plugin_functions.yaml+15-0
@@ -6354,6 +6354,21 @@ official:
6354 - func: isin.Tensor_Tensor_out(Tensor elements, Tensor test_elements, *, bool assume_unique=False, bool invert=False, Tensor(a!) out) -> Tensor(a!)6354 - func: isin.Tensor_Tensor_out(Tensor elements, Tensor test_elements, *, bool assume_unique=False, bool invert=False, Tensor(a!) out) -> Tensor(a!)
6355 op_api: [v2.1, newest]6355 op_api: [v2.1, newest]
6356 6356 
6357+ - func: logit(Tensor self, float? eps=None) -> Tensor
6358+ op_api: [v2.1, newest]
6359+ 
6360+ - func: logit_(Tensor(a!) self, float? eps=None) -> Tensor(a!)
6361+ op_api: [v2.1, newest]
6362+ 
6363+ - func: logit.out(Tensor self, float? eps=None, *, Tensor(a!) out) -> Tensor(a!)
6364+ op_api: [v2.1, newest]
6365+ 
6366+ - func: logit_backward.grad_input(Tensor grad_output, Tensor self, float? eps=None, *, Tensor(a!) grad_input) -> Tensor(a!)
6367+ op_api: [v2.1, newest]
6368+ 
6369+ - func: logit_backward(Tensor grad_output, Tensor self, float? eps=None) -> Tensor
6370+ op_api: [v2.1, newest]
6371+ 
6357custom:6372custom:
6358 - func: npu_gather_backward(Tensor grad, SymInt[] self_size, int dim, Tensor index, bool sparse_grad) -> Tensor6373 - func: npu_gather_backward(Tensor grad, SymInt[] self_size, int dim, Tensor index, bool sparse_grad) -> Tensor
6359 op_api: all_version6374 op_api: all_version
Aop_plugin/ops/opapi/LogitBackwardKernelOpApi.cpp+61-0
@@ -0,0 +1,61 @@
1+// Copyright (c) 2025 Huawei Technologies Co., Ltd
2+// All rights reserved.
3+ 
4+#include "op_plugin/OpApiInterface.h"
5+#include "op_plugin/utils/op_api_common.h"
6+#include "torch_npu/csrc/framework/utils/UtilForOpAdapter.h"
7+ 
8+namespace op_api {
9+using npu_preparation = at_npu::native::OpPreparation;
10+ 
11+inline bool _logit_backward_fallback_condition()
12+{
13+ static const bool is_aclnn_kernel_available = check_aclnn_kernel_available("aclnnLogitGrad");
14+ static const bool is_support_soc = (c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1 &&
15+ c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend310B1) ||
16+ (c10_npu::GetSocVersion() > c10_npu::SocVersion::Ascend310B4);
17+ if (!is_aclnn_kernel_available || !is_support_soc) {
18+ TORCH_NPU_WARN_ONCE("CAUTION: The operator aten::logit_backward and aten::logit_backward.out is currently "
19+ "not supported on the NPU backend. Now this operator will fallback to run on the CPU "
20+ "and may have performance implications.");
21+ return true;
22+ }
23+ return false;
24+}
25+ 
26+at::Tensor &logit_backward_out(const at::Tensor &grad_output, const at::Tensor &self, c10::optional<double> eps, at::Tensor &grad_input)
27+{
28+ if (_logit_backward_fallback_condition()) {
29+ at::Tensor grad_output_cpu = grad_output.cpu();
30+ at::Tensor self_cpu = self.cpu();
31+ at::Tensor grad_input_cpu = grad_input.cpu();
32+ grad_input_cpu = at::logit_backward_outf(grad_output_cpu, self_cpu, eps, grad_input_cpu);
33+ grad_input.copy_(grad_input_cpu);
34+ return grad_input;
35+ }
36+ auto eps_value = eps.value_or(-1);
37+ auto output_size_0 = self.sizes();
38+ auto output_dtype_0 = self.scalar_type();
39+ npu_preparation::check_tensor({grad_output, self}, grad_input, output_dtype_0, output_size_0);
40+ EXEC_NPU_CMD(aclnnLogitGrad, grad_output, self, eps_value, grad_input);
41+ return grad_input;
42+}
43+ 
44+at::Tensor logit_backward(const at::Tensor &grad_output, const at::Tensor &self, c10::optional<double> eps)
45+{
46+ if (_logit_backward_fallback_condition()) {
47+ at::Tensor grad_output_cpu = grad_output.cpu();
48+ at::Tensor self_cpu = self.cpu();
49+ at::Tensor out_cpu = at::logit_backward(grad_output_cpu, self_cpu, eps);
50+ return out_cpu.to(grad_output.device());
51+ }
52+ auto eps_value = eps.value_or(-1);
53+ auto output_size_0 = self.sizes();
54+ auto output_dtype_0 = self.scalar_type();
55+ at::Tensor grad_input = npu_preparation::apply_tensor_without_format(output_size_0,
56+ grad_output.options().dtype(output_dtype_0));
57+ EXEC_NPU_CMD(aclnnLogitGrad, grad_output, self, eps_value, grad_input);
58+ return grad_input;
59+}
60+ 
61+}
Aop_plugin/ops/opapi/LogitKernelNpuOpApi.cpp+63-0
@@ -0,0 +1,63 @@
1+// Copyright (c) 2025 Huawei Technologies Co., Ltd
2+// All rights reserved.
3+ 
4+#include "op_plugin/OpApiInterface.h"
5+#include "op_plugin/utils/op_api_common.h"
6+#include "torch_npu/csrc/framework/utils/UtilForOpAdapter.h"
7+ 
8+namespace op_api {
9+using npu_preparation = at_npu::native::OpPreparation;
10+ 
11+inline bool _logit_fallback_condition()
12+{
13+ static const bool is_aclnn_kernel_available = check_aclnn_kernel_available("aclnnLogit");
14+ static const bool is_support_soc = (c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend910B1 &&
15+ c10_npu::GetSocVersion() < c10_npu::SocVersion::Ascend310B1) ||
16+ (c10_npu::GetSocVersion() > c10_npu::SocVersion::Ascend310B4);
17+ if (!is_aclnn_kernel_available || !is_support_soc) {
18+ TORCH_NPU_WARN_ONCE("CAUTION: The operator aten::logit, aten::logit_ and aten::logit.out is currently "
19+ "not supported on the NPU backend. Now this operator will fallback to run on the CPU "
20+ "and may have performance implications.");
21+ return true;
22+ }
23+ return false;
24+}
25+ 
26+at::Tensor logit(const at::Tensor &self, c10::optional<double> eps)
27+{
28+ if (_logit_fallback_condition()) {
29+ at::Tensor self_cpu = self.cpu();
30+ at::Tensor out_cpu = at::native::logit(self_cpu, eps);
31+ return out_cpu.to(self.device());
32+ }
33+ auto eps_value = eps.value_or(-1);
34+ auto output_size_0 = self.sizes();
35+ auto output_dtype_0 = self.scalar_type();
36+ at::Tensor out = npu_preparation::apply_tensor_without_format(output_size_0, self.options().dtype(output_dtype_0));
37+ EXEC_NPU_CMD(aclnnLogit, self, eps_value, out);
38+ return out;
39+}
40+ 
41+at::Tensor &logit_(at::Tensor &self, c10::optional<double> eps)
42+{
43+ return at::native::logit_(self, eps);
44+}
45+ 
46+at::Tensor &logit_out(const at::Tensor &self, c10::optional<double> eps, at::Tensor &out)
47+{
48+ if (_logit_fallback_condition()) {
49+ at::Tensor self_cpu = self.cpu();
50+ at::Tensor out_cpu = out.cpu();
51+ out_cpu = at::native::logit_out(self_cpu, eps, out_cpu);
52+ out.copy_(out_cpu);
53+ return out;
54+ }
55+ auto eps_value = eps.value_or(-1);
56+ auto output_size_0 = self.sizes();
57+ auto output_dtype_0 = self.scalar_type();
58+ npu_preparation::check_tensor({self}, out, output_dtype_0, output_size_0);
59+ EXEC_NPU_CMD(aclnnLogit, self, eps_value, out);
60+ return out;
61+}
62+ 
63+}
Atest/test_base_ops/test_logit.py+67-0
@@ -0,0 +1,67 @@
1+import unittest
2+import numpy as np
3+import torch
4+ 
5+import torch_npu
6+from torch.testing._internal.common_utils import TestCase, run_tests
7+from torch_npu.testing.common_utils import create_common_tensor
8+ 
9+ 
10+class TestLogit(TestCase):
11+ def cpu_op_exec(self, input1, eps=None):
12+ output = torch.logit(input1, eps)
13+ output = output.numpy()
14+ return output
15+
16+ def cpu_backward_op_exec(self, input1, eps=None):
17+ input1.requires_grad_(True)
18+ output = torch.logit(input1, eps)
19+ output.backward(torch.ones_like(input1))
20+ return input1.grad.numpy()
21+ 
22+ def npu_op_exec(self, input1, eps=None):
23+ output = torch.logit(input1, eps)
24+ output = output.to("cpu")
25+ output = output.numpy()
26+ return output
27+ 
28+ def npu_backward_op_exec(self, input1, eps=None):
29+ input1.requires_grad_(True)
30+ output = torch.logit(input1, eps)
31+ output.backward(torch.ones_like(input1))
32+ return input1.grad.cpu().numpy()
33+ 
34+ def npu_op_exec_out(self, input1, out, eps=None):
35+ torch.logit(input1, eps, out=out)
36+ output = out.to("cpu")
37+ output = output.numpy()
38+ return output
39+ 
40+ def test_logit_common_shape_format(self):
41+ shape_format = [
42+ [[np.float32, 0, [3, 4]], 1e-5],
43+ [[np.float32, 0, [3, 128, 256]], 3e-5],
44+ [[np.float32, 0, [3, 256, 128, 8]], None],
45+ ]
46+ for item in shape_format:
47+ cpu_input1, npu_input1 = create_common_tensor(item[0], 0.1, 0.9)
48+ cpu_out, npu_out = create_common_tensor(item[0], 0.1, 0.9)
49+ eps = item[1]
50+ if eps is None:
51+ cpu_output = self.cpu_op_exec(cpu_input1)
52+ npu_output = self.npu_op_exec(npu_input1)
53+ npu_output_out = self.npu_op_exec_out(npu_input1, npu_out)
54+ npu_output_inplace = npu_input1.logit_().cpu().numpy()
55+ else:
56+ cpu_output = self.cpu_op_exec(cpu_input1, eps)
57+ npu_output = self.npu_op_exec(npu_input1, eps)
58+ npu_output_out = self.npu_op_exec_out(npu_input1, npu_out, eps)
59+ npu_output_inplace = npu_input1.logit_(eps).cpu().numpy()
60+ 
61+ self.assertEqual(cpu_output, npu_output)
62+ self.assertEqual(cpu_output, npu_output_out)
63+ self.assertEqual(cpu_output, npu_output_inplace)
64+ 
65+ 
66+if __name__ == "__main__":
67+ run_tests()