已合并
feat: Ascend950比较类算子device校验及平台隔离 #5286
feat: Ascend950比较类算子device校验及平台隔离 #5286
已合并
wangqi_ai创建于 6月27日
9 个文件变更+251-0
@@ -22,6 +22,15 @@ using npu_preparation = at_npu::native::OpPreparation;
22 22 
23at::Tensor& eq_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) {23at::Tensor& eq_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) {
24 DO_COMPATIBILITY(aclnnEqTensor, acl_op::eq_out(self, other, result));24 DO_COMPATIBILITY(aclnnEqTensor, acl_op::eq_out(self, other, result));
25+ if (is_ascend950_path()) {
26+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
27+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
28+ auto output_size = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
29+ npu_preparation::check_tensor({self_device, other_device}, result, output_size);
30+ EXEC_NPU_CMD(aclnnEqTensor, self_device, other_device, result);
31+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
32+ return result;
33+ }
25 std::vector<at::Tensor> tensor_list = {self, other};34 std::vector<at::Tensor> tensor_list = {self, other};
26 auto maybe_names = op_plugin::utils::compute_names_npu(tensor_list);35 auto maybe_names = op_plugin::utils::compute_names_npu(tensor_list);
27 36 
@@ -40,6 +49,16 @@ at::Tensor& eq_out(const at::Tensor& self, const at::Tensor& other, at::Tensor&
40 49 
41at::Tensor eq(const at::Tensor& self, const at::Tensor& other) {50at::Tensor eq(const at::Tensor& self, const at::Tensor& other) {
42 DO_COMPATIBILITY(aclnnEqTensor, acl_op::eq(self, other));51 DO_COMPATIBILITY(aclnnEqTensor, acl_op::eq(self, other));
52+ if (is_ascend950_path()) {
53+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
54+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
55+ auto output_size = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
56+ at::Tensor result = npu_preparation::apply_tensor_without_format(
57+ output_size, self_device.options().dtype(at::kBool));
58+ EXEC_NPU_CMD(aclnnEqTensor, self_device, other_device, result);
59+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
60+ return result;
61+ }
43 std::vector<at::Tensor> tensor_list = {self, other};62 std::vector<at::Tensor> tensor_list = {self, other};
44 auto maybe_names = op_plugin::utils::compute_names_npu(tensor_list);63 auto maybe_names = op_plugin::utils::compute_names_npu(tensor_list);
45 64 
@@ -94,6 +113,17 @@ at::Tensor& eq_out(const at::Tensor& self, const at::Scalar& other, at::Tensor&
94 113 
95at::Tensor& eq_(at::Tensor &self, const at::Tensor &other) {114at::Tensor& eq_(at::Tensor &self, const at::Tensor &other) {
96 DO_COMPATIBILITY(aclnnInplaceEqTensor, acl_op::eq_(self, other));115 DO_COMPATIBILITY(aclnnInplaceEqTensor, acl_op::eq_(self, other));
116+ if (is_ascend950_path()) {
117+ TORCH_CHECK(torch_npu::utils::is_npu(self),
118+ "inplace eq_ requires self to be NPU tensor", OPS_ERROR(ErrCode::PARAM));
119+ at::Tensor other_device = other;
120+ if (!torch_npu::utils::is_npu(other)) {
121+ other_device = other.to(self.device());
122+ }
123+ npu_preparation::CheckMemory({self, other_device}, {self});
124+ EXEC_NPU_CMD(aclnnInplaceEqTensor, self, other_device);
125+ return self;
126+ }
97 127 
98 const std::initializer_list<at::Tensor> inputs = {self, other};128 const std::initializer_list<at::Tensor> inputs = {self, other};
99 const std::initializer_list<at::Tensor> outputs = {self};129 const std::initializer_list<at::Tensor> outputs = {self};
@@ -45,6 +45,15 @@ at::Tensor& ge_(at::Tensor& self, const at::Scalar& other) {
45 45 
46at::Tensor& ge_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) {46at::Tensor& ge_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) {
47 DO_COMPATIBILITY(aclnnGeTensor, acl_op::ge_out(self, other, result));47 DO_COMPATIBILITY(aclnnGeTensor, acl_op::ge_out(self, other, result));
48+ if (is_ascend950_path()) {
49+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
50+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
51+ auto output_size = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
52+ npu_preparation::check_tensor({self_device, other_device}, result, output_size);
53+ EXEC_NPU_CMD(aclnnGeTensor, self_device, other_device, result);
54+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
55+ return result;
56+ }
48 auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);57 auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);
49 npu_preparation::check_tensor({self, other}, result, output_size);58 npu_preparation::check_tensor({self, other}, result, output_size);
50 EXEC_NPU_CMD(aclnnGeTensor, self, other, result);59 EXEC_NPU_CMD(aclnnGeTensor, self, other, result);
@@ -53,6 +62,16 @@ at::Tensor& ge_out(const at::Tensor& self, const at::Tensor& other, at::Tensor&
53 62 
54at::Tensor ge(const at::Tensor& self, const at::Tensor& other) {63at::Tensor ge(const at::Tensor& self, const at::Tensor& other) {
55 DO_COMPATIBILITY(aclnnGeTensor, acl_op::ge(self, other));64 DO_COMPATIBILITY(aclnnGeTensor, acl_op::ge(self, other));
65+ if (is_ascend950_path()) {
66+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
67+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
68+ auto output_size = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
69+ at::Tensor result = npu_preparation::apply_tensor_without_format(
70+ output_size, self_device.options().dtype(at::kBool));
71+ EXEC_NPU_CMD(aclnnGeTensor, self_device, other_device, result);
72+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
73+ return result;
74+ }
56 if (other.dim() == 0 && !torch_npu::utils::is_npu(other)) {75 if (other.dim() == 0 && !torch_npu::utils::is_npu(other)) {
57 DO_COMPATIBILITY(aclnnGeScalar, acl_op::ge(self, other));76 DO_COMPATIBILITY(aclnnGeScalar, acl_op::ge(self, other));
58 at::Tensor result = npu_preparation::apply_tensor_without_format(self.sizes(), self.options().dtype(at::kBool));77 at::Tensor result = npu_preparation::apply_tensor_without_format(self.sizes(), self.options().dtype(at::kBool));
@@ -76,6 +95,17 @@ at::Tensor ge(const at::Tensor& self, const at::Tensor& other) {
76at::Tensor& ge_(at::Tensor &self, const at::Tensor &other)95at::Tensor& ge_(at::Tensor &self, const at::Tensor &other)
77{96{
78 DO_COMPATIBILITY(aclnnInplaceGeTensor, acl_op::ge_(self, other));97 DO_COMPATIBILITY(aclnnInplaceGeTensor, acl_op::ge_(self, other));
98+ if (is_ascend950_path()) {
99+ TORCH_CHECK(torch_npu::utils::is_npu(self),
100+ "inplace ge_ requires self to be NPU tensor", OPS_ERROR(ErrCode::PARAM));
101+ at::Tensor other_device = other;
102+ if (!torch_npu::utils::is_npu(other)) {
103+ other_device = other.to(self.device());
104+ }
105+ npu_preparation::CheckMemory({self, other_device}, {self});
106+ EXEC_NPU_CMD(aclnnInplaceGeTensor, self, other_device);
107+ return self;
108+ }
79 if (other.dim() == 0 && !torch_npu::utils::is_npu(other)) {109 if (other.dim() == 0 && !torch_npu::utils::is_npu(other)) {
80 return op_api::ge_(self, other.item());110 return op_api::ge_(self, other.item());
81 } else {111 } else {
@@ -53,6 +53,15 @@ at::Tensor& gt_(at::Tensor& self, const at::Scalar& other)
53at::Tensor& gt_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result)53at::Tensor& gt_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result)
54{54{
55 DO_COMPATIBILITY(aclnnGtTensor, acl_op::gt_out(self, other, result));55 DO_COMPATIBILITY(aclnnGtTensor, acl_op::gt_out(self, other, result));
56+ if (is_ascend950_path()) {
57+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
58+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
59+ auto output_size = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
60+ npu_preparation::check_tensor({self_device, other_device}, result, output_size);
61+ EXEC_NPU_CMD(aclnnGtTensor, self_device, other_device, result);
62+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
63+ return result;
64+ }
56 auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);65 auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);
57 66 
58 npu_preparation::check_tensor({self, other}, result, output_size);67 npu_preparation::check_tensor({self, other}, result, output_size);
@@ -72,6 +81,16 @@ at::Tensor& gt_out(const at::Tensor& self, const at::Tensor& other, at::Tensor&
72at::Tensor gt(const at::Tensor& self, const at::Tensor& other)81at::Tensor gt(const at::Tensor& self, const at::Tensor& other)
73{82{
74 DO_COMPATIBILITY(aclnnGtTensor, acl_op::gt(self, other));83 DO_COMPATIBILITY(aclnnGtTensor, acl_op::gt(self, other));
84+ if (is_ascend950_path()) {
85+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
86+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
87+ auto output_size = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
88+ at::Tensor result = npu_preparation::apply_tensor_without_format(
89+ output_size, self_device.options().dtype(at::kBool));
90+ EXEC_NPU_CMD(aclnnGtTensor, self_device, other_device, result);
91+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
92+ return result;
93+ }
75 // calculate the output size94 // calculate the output size
76 auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);95 auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);
77 96 
@@ -94,6 +113,17 @@ at::Tensor gt(const at::Tensor& self, const at::Tensor& other)
94at::Tensor& gt_(at::Tensor &self, const at::Tensor &other)113at::Tensor& gt_(at::Tensor &self, const at::Tensor &other)
95{114{
96 DO_COMPATIBILITY(aclnnInplaceGtTensor, acl_op::gt_(self, other));115 DO_COMPATIBILITY(aclnnInplaceGtTensor, acl_op::gt_(self, other));
116+ if (is_ascend950_path()) {
117+ TORCH_CHECK(torch_npu::utils::is_npu(self),
118+ "inplace gt_ requires self to be NPU tensor", OPS_ERROR(ErrCode::PARAM));
119+ at::Tensor other_device = other;
120+ if (!torch_npu::utils::is_npu(other)) {
121+ other_device = other.to(self.device());
122+ }
123+ npu_preparation::CheckMemory({self, other_device}, {self});
124+ EXEC_NPU_CMD(aclnnInplaceGtTensor, self, other_device);
125+ return self;
126+ }
97 if (other.dim() == 0 && !torch_npu::utils::is_npu(other)) {127 if (other.dim() == 0 && !torch_npu::utils::is_npu(other)) {
98 return op_api::gt_(self, other.item());128 return op_api::gt_(self, other.item());
99 } else {129 } else {
@@ -17,11 +17,22 @@
17#include "op_plugin/OpApiInterface.h"17#include "op_plugin/OpApiInterface.h"
18#include "op_plugin/utils/op_api_common.h"18#include "op_plugin/utils/op_api_common.h"
19 19 
20+using npu_preparation = at_npu::native::OpPreparation;
21+ 
20namespace op_api {22namespace op_api {
21 23 
22at::Tensor &le_out(const at::Tensor &self, const at::Tensor &other, at::Tensor &result)24at::Tensor &le_out(const at::Tensor &self, const at::Tensor &other, at::Tensor &result)
23{25{
24 DO_COMPATIBILITY(aclnnLeTensor, acl_op::le_out(self, other, result));26 DO_COMPATIBILITY(aclnnLeTensor, acl_op::le_out(self, other, result));
27+ if (is_ascend950_path()) {
28+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
29+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
30+ auto outputSize = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
31+ npu_preparation::check_tensor({self_device, other_device}, result, outputSize);
32+ EXEC_NPU_CMD(aclnnLeTensor, self_device, other_device, result);
33+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
34+ return result;
35+ }
25 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);36 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);
26 at_npu::native::OpPreparation::check_tensor({self}, result, result.scalar_type(), outputSize);37 at_npu::native::OpPreparation::check_tensor({self}, result, result.scalar_type(), outputSize);
27 if (at_npu::native::OpPreparation::IsCPUScalar(self)) {38 if (at_npu::native::OpPreparation::IsCPUScalar(self)) {
@@ -49,6 +60,16 @@ at::Tensor &le_out(const at::Tensor &self, const at::Scalar &other, at::Tensor &
49at::Tensor le(const at::Tensor &self, const at::Tensor &other)60at::Tensor le(const at::Tensor &self, const at::Tensor &other)
50{61{
51 DO_COMPATIBILITY(aclnnLeTensor, acl_op::le(self, other));62 DO_COMPATIBILITY(aclnnLeTensor, acl_op::le(self, other));
63+ if (is_ascend950_path()) {
64+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
65+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
66+ auto outputSize = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
67+ at::Tensor result = npu_preparation::apply_tensor_without_format(
68+ outputSize, self_device.options().dtype(at::kBool));
69+ EXEC_NPU_CMD(aclnnLeTensor, self_device, other_device, result);
70+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
71+ return result;
72+ }
52 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);73 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);
53 at::Tensor result =74 at::Tensor result =
54 at_npu::native::OpPreparation::apply_tensor_without_format(outputSize, self.options().dtype(at::kBool));75 at_npu::native::OpPreparation::apply_tensor_without_format(outputSize, self.options().dtype(at::kBool));
@@ -84,6 +105,17 @@ at::Tensor &le_(at::Tensor &self, const at::Scalar &other)
84at::Tensor &le_(at::Tensor &self, const at::Tensor &other)105at::Tensor &le_(at::Tensor &self, const at::Tensor &other)
85{106{
86 DO_COMPATIBILITY(aclnnInplaceLeTensor, acl_op::le_(self, other));107 DO_COMPATIBILITY(aclnnInplaceLeTensor, acl_op::le_(self, other));
108+ if (is_ascend950_path()) {
109+ TORCH_CHECK(torch_npu::utils::is_npu(self),
110+ "inplace le_ requires self to be NPU tensor", OPS_ERROR(ErrCode::PARAM));
111+ at::Tensor other_device = other;
112+ if (!torch_npu::utils::is_npu(other)) {
113+ other_device = other.to(self.device());
114+ }
115+ npu_preparation::CheckMemory({self, other_device}, {self});
116+ EXEC_NPU_CMD(aclnnInplaceLeTensor, self, other_device);
117+ return self;
118+ }
87 if (other.dim() == 0 && !torch_npu::utils::is_npu(other)) {119 if (other.dim() == 0 && !torch_npu::utils::is_npu(other)) {
88 return op_api::le_(self, other.item());120 return op_api::le_(self, other.item());
89 } else {121 } else {
@@ -21,6 +21,15 @@ namespace op_api {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22at::Tensor& lt_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) {22at::Tensor& lt_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) {
23 DO_COMPATIBILITY(aclnnLtTensor, acl_op::lt_out(self, other, result));23 DO_COMPATIBILITY(aclnnLtTensor, acl_op::lt_out(self, other, result));
24+ if (is_ascend950_path()) {
25+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
26+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
27+ auto outputSize = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
28+ npu_preparation::check_tensor({self_device, other_device}, result, outputSize);
29+ EXEC_NPU_CMD(aclnnLtTensor, self_device, other_device, result);
30+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
31+ return result;
32+ }
24 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);33 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);
25 34 
26 at_npu::native::OpPreparation::check_tensor({self}, result, result.scalar_type(), outputSize);35 at_npu::native::OpPreparation::check_tensor({self}, result, result.scalar_type(), outputSize);
@@ -39,6 +48,16 @@ at::Tensor& lt_out(const at::Tensor& self, const at::Tensor& other, at::Tensor&
39 48 
40at::Tensor lt(const at::Tensor& self, const at::Tensor& other) {49at::Tensor lt(const at::Tensor& self, const at::Tensor& other) {
41 DO_COMPATIBILITY(aclnnLtTensor, acl_op::lt(self, other));50 DO_COMPATIBILITY(aclnnLtTensor, acl_op::lt(self, other));
51+ if (is_ascend950_path()) {
52+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
53+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
54+ auto outputSize = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
55+ at::Tensor result = npu_preparation::apply_tensor_without_format(
56+ outputSize, self_device.options().dtype(at::kBool));
57+ EXEC_NPU_CMD(aclnnLtTensor, self_device, other_device, result);
58+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
59+ return result;
60+ }
42 // calculate the output size61 // calculate the output size
43 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);62 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);
44 63 
@@ -85,6 +104,17 @@ at::Tensor lt(const at::Tensor &self, const at::Scalar& other)
85 104 
86at::Tensor& lt_(at::Tensor& self, const at::Tensor& other) {105at::Tensor& lt_(at::Tensor& self, const at::Tensor& other) {
87 DO_COMPATIBILITY(aclnnInplaceLtTensor, acl_op::lt_(self, other));106 DO_COMPATIBILITY(aclnnInplaceLtTensor, acl_op::lt_(self, other));
107+ if (is_ascend950_path()) {
108+ TORCH_CHECK(torch_npu::utils::is_npu(self),
109+ "inplace lt_ requires self to be NPU tensor", OPS_ERROR(ErrCode::PARAM));
110+ at::Tensor other_device = other;
111+ if (!torch_npu::utils::is_npu(other)) {
112+ other_device = other.to(self.device());
113+ }
114+ npu_preparation::CheckMemory({self, other_device}, {self});
115+ EXEC_NPU_CMD(aclnnInplaceLtTensor, self, other_device);
116+ return self;
117+ }
88 if (npu_preparation::IsCPUScalar(other)) {118 if (npu_preparation::IsCPUScalar(other)) {
89 return op_api::lt_(self, other.item());119 return op_api::lt_(self, other.item());
90 } else {120 } else {
@@ -17,11 +17,20 @@
17#include "op_plugin/OpApiInterface.h"17#include "op_plugin/OpApiInterface.h"
18#include "op_plugin/utils/op_api_common.h"18#include "op_plugin/utils/op_api_common.h"
19 19 
20+using npu_preparation = at_npu::native::OpPreparation;
21+ 
20namespace op_api {22namespace op_api {
21 23 
22at::Tensor &maximum_out(const at::Tensor &self, const at::Tensor &other, at::Tensor &result)24at::Tensor &maximum_out(const at::Tensor &self, const at::Tensor &other, at::Tensor &result)
23{25{
24 DO_COMPATIBILITY(aclnnMaximum, acl_op::maximum_out(self, other, result));26 DO_COMPATIBILITY(aclnnMaximum, acl_op::maximum_out(self, other, result));
27+ if (is_ascend950_path()) {
28+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
29+ auto output_size = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
30+ npu_preparation::check_tensor({self_device, other_device}, result, output_size);
31+ EXEC_NPU_CMD(aclnnMaximum, self_device, other_device, result);
32+ return result;
33+ }
25 at::Tensor cp_self = self;34 at::Tensor cp_self = self;
26 if (at_npu::native::OpPreparation::IsCPUScalar(self)) {35 if (at_npu::native::OpPreparation::IsCPUScalar(self)) {
27 at::Scalar scalar = self.item();36 at::Scalar scalar = self.item();
@@ -41,6 +50,15 @@ at::Tensor &maximum_out(const at::Tensor &self, const at::Tensor &other, at::Ten
41at::Tensor maximum(const at::Tensor &self, const at::Tensor &other)50at::Tensor maximum(const at::Tensor &self, const at::Tensor &other)
42{51{
43 DO_COMPATIBILITY(aclnnMaximum, acl_op::maximum(self, other));52 DO_COMPATIBILITY(aclnnMaximum, acl_op::maximum(self, other));
53+ if (is_ascend950_path()) {
54+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
55+ at::ScalarType result_dtype = at::native::result_type(self_device, other_device);
56+ auto output_size = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
57+ at::Tensor result = npu_preparation::apply_tensor_without_format(
58+ output_size, self_device.options().dtype(result_dtype));
59+ EXEC_NPU_CMD(aclnnMaximum, self_device, other_device, result);
60+ return result;
61+ }
44 at::Tensor cp_self = self;62 at::Tensor cp_self = self;
45 if (at_npu::native::OpPreparation::IsCPUScalar(self)) {63 if (at_npu::native::OpPreparation::IsCPUScalar(self)) {
46 at::Scalar scalar = self.item();64 at::Scalar scalar = self.item();
@@ -18,11 +18,22 @@
18#include "op_plugin/OpApiInterface.h"18#include "op_plugin/OpApiInterface.h"
19#include "op_plugin/utils/op_api_common.h"19#include "op_plugin/utils/op_api_common.h"
20 20 
21+using npu_preparation = at_npu::native::OpPreparation;
22+ 
21namespace op_api {23namespace op_api {
22 24 
23at::Tensor minimum(const at::Tensor& self, const at::Tensor& other)25at::Tensor minimum(const at::Tensor& self, const at::Tensor& other)
24{26{
25 DO_COMPATIBILITY(aclnnMinimum, acl_op::minimum(self, other));27 DO_COMPATIBILITY(aclnnMinimum, acl_op::minimum(self, other));
28+ if (is_ascend950_path()) {
29+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
30+ at::ScalarType result_dtype = at::native::result_type(self_device, other_device);
31+ auto output_size = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
32+ at::Tensor result = npu_preparation::apply_tensor_without_format(
33+ output_size, self_device.options().dtype(result_dtype));
34+ EXEC_NPU_CMD(aclnnMinimum, self_device, other_device, result);
35+ return result;
36+ }
26 auto result_type = at::result_type(self, other);37 auto result_type = at::result_type(self, other);
27 auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);38 auto output_size = op_infer::broadcast_ops_npu_output_size(self, other);
28 at::Tensor result =39 at::Tensor result =
@@ -33,6 +44,13 @@ at::Tensor minimum(const at::Tensor& self, const at::Tensor& other)
33at::Tensor& minimum_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result)44at::Tensor& minimum_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result)
34{45{
35 DO_COMPATIBILITY(aclnnMinimum, acl_op::minimum_out(self, other, result));46 DO_COMPATIBILITY(aclnnMinimum, acl_op::minimum_out(self, other, result));
47+ if (is_ascend950_path()) {
48+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
49+ auto output_size = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
50+ npu_preparation::check_tensor({self_device, other_device}, result, output_size);
51+ EXEC_NPU_CMD(aclnnMinimum, self_device, other_device, result);
52+ return result;
53+ }
36 at::Tensor cp_other = other;54 at::Tensor cp_other = other;
37 at::Tensor cp_self = self;55 at::Tensor cp_self = self;
38 if (at_npu::native::OpPreparation::IsCPUScalar(other)) {56 if (at_npu::native::OpPreparation::IsCPUScalar(other)) {
@@ -24,6 +24,15 @@ using npu_preparation = at_npu::native::OpPreparation;
24at::Tensor& ne_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result)24at::Tensor& ne_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result)
25{25{
26 DO_COMPATIBILITY(aclnnNeTensor, acl_op::ne_out(self, other, result));26 DO_COMPATIBILITY(aclnnNeTensor, acl_op::ne_out(self, other, result));
27+ if (is_ascend950_path()) {
28+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
29+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
30+ auto outputSize = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
31+ npu_preparation::check_tensor({self_device, other_device}, result, result.scalar_type(), at::IntArrayRef(outputSize));
32+ EXEC_NPU_CMD(aclnnNeTensor, self_device, other_device, result);
33+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
34+ return result;
35+ }
27 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);36 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);
28 npu_preparation::check_tensor({self, other}, result, result.scalar_type(), at::IntArrayRef(outputSize));37 npu_preparation::check_tensor({self, other}, result, result.scalar_type(), at::IntArrayRef(outputSize));
29 if (npu_preparation::IsCPUScalar(self)) {38 if (npu_preparation::IsCPUScalar(self)) {
@@ -49,6 +58,16 @@ at::Tensor& ne_out(const at::Tensor& self, const at::Scalar& other, at::Tensor&
49at::Tensor ne(const at::Tensor& self, const at::Tensor& other)58at::Tensor ne(const at::Tensor& self, const at::Tensor& other)
50{59{
51 DO_COMPATIBILITY(aclnnNeTensor, acl_op::ne(self, other));60 DO_COMPATIBILITY(aclnnNeTensor, acl_op::ne(self, other));
61+ if (is_ascend950_path()) {
62+ auto [self_device, other_device] = prepare_binary_tensors(self, other);
63+ auto maybe_names = op_plugin::utils::compute_names_npu({self, other});
64+ auto outputSize = op_infer::broadcast_ops_npu_output_size(self_device, other_device);
65+ at::Tensor result = npu_preparation::apply_tensor_without_format(
66+ outputSize, self_device.options().dtype(at::kBool));
67+ EXEC_NPU_CMD(aclnnNeTensor, self_device, other_device, result);
68+ at::namedinference::propagate_names_if_nonempty(result, maybe_names);
69+ return result;
70+ }
52 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);71 auto outputSize = op_infer::broadcast_ops_npu_output_size(self, other);
53 at::Tensor result =72 at::Tensor result =
54 npu_preparation::apply_tensor_without_format(outputSize, self.options().dtype(at::kBool));73 npu_preparation::apply_tensor_without_format(outputSize, self.options().dtype(at::kBool));
@@ -78,6 +97,17 @@ at::Tensor ne(const at::Tensor& self, const at::Scalar& other)
78at::Tensor& ne_(at::Tensor& self, const at::Tensor& other)97at::Tensor& ne_(at::Tensor& self, const at::Tensor& other)
79{98{
80 DO_COMPATIBILITY(aclnnInplaceNeTensor, acl_op::ne_(self, other));99 DO_COMPATIBILITY(aclnnInplaceNeTensor, acl_op::ne_(self, other));
100+ if (is_ascend950_path()) {
101+ TORCH_CHECK(torch_npu::utils::is_npu(self),
102+ "inplace ne_ requires self to be NPU tensor", OPS_ERROR(ErrCode::PARAM));
103+ at::Tensor other_device = other;
104+ if (!torch_npu::utils::is_npu(other)) {
105+ other_device = other.to(self.device());
106+ }
107+ npu_preparation::CheckMemory({self, other_device}, {self});
108+ EXEC_NPU_CMD(aclnnInplaceNeTensor, self, other_device);
109+ return self;
110+ }
81 npu_preparation::check_memory({self, other}, {self});111 npu_preparation::check_memory({self, other}, {self});
82 if (npu_preparation::IsCPUScalar(other)) {112 if (npu_preparation::IsCPUScalar(other)) {
83 return op_api::ne_(self, other.item());113 return op_api::ne_(self, other.item());
@@ -816,4 +816,37 @@ std::vector<EnumType> convert_int_to_enum(ArrayType& v) {
816 }816 }
817 return converted;817 return converted;
818}818}
819+ 
820+/**
821+ * @brief Prepare binary op tensors by ensuring both are on NPU device.
822+ * @param self First input tensor (NPU or CPU).
823+ * @param other Second input tensor (NPU or CPU).
824+ * @return Tuple of tensors both on NPU device.
825+ * @note If self is NPU: other can be CPU or NPU (CPU tensors are converted to NPU).
826+ * If self is CPU: other must be NPU, and self is converted to NPU.
827+ */
828+inline std::tuple<at::Tensor, at::Tensor> prepare_binary_tensors(
829+ const at::Tensor &self, const at::Tensor &other)
830+{
831+ at::Tensor self_device = self;
832+ at::Tensor other_device = other;
833+ if (torch_npu::utils::is_npu(self)) {
834+ if (!torch_npu::utils::is_npu(other)) {
835+ other_device = other.to(self.device());
836+ }
837+ } else {
838+ TORCH_CHECK(torch_npu::utils::is_npu(other),
839+ "when self is a CPU tensor, other must be an NPU tensor", OPS_ERROR(ErrCode::PARAM));
840+ self_device = self.to(other.device());
841+ }
842+ return std::make_tuple(self_device, other_device);
843+}
844+ 
845+// Check if current platform is Ascend950 with CANN >= 9.1.0.
846+inline bool is_ascend950_path()
847+{
848+ return c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950 &&
849+ op_plugin::utils::is_gte_cann_version_910();
850+}
851+ 
819#endif // TORCHNPU_TORCH_NPU_CSRC_ATEN_OPS_OP_API_PTA_COMMON_H_852#endif // TORCHNPU_TORCH_NPU_CSRC_ATEN_OPS_OP_API_PTA_COMMON_H_