已合并
feat: Ascend950比较类算子device校验及平台隔离 #5286
wangqi_ai创建于 6月27日
feat: Ascend950比较类算子device校验及平台隔离 #5286
已合并
共 9 个文件变更+251-0
| @@ -22,6 +22,15 @@ using npu_preparation = at_npu::native::OpPreparation; | |||
| 22 | 22 | ||
| 23 | at::Tensor& eq_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) { | 23 | at::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 | ||
| 41 | at::Tensor eq(const at::Tensor& self, const at::Tensor& other) { | 50 | at::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 | ||
| 95 | at::Tensor& eq_(at::Tensor &self, const at::Tensor &other) { | 114 | at::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 | ||
| 46 | at::Tensor& ge_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) { | 46 | at::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 | ||
| 54 | at::Tensor ge(const at::Tensor& self, const at::Tensor& other) { | 63 | at::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) { | |||
| 76 | at::Tensor& ge_(at::Tensor &self, const at::Tensor &other) | 95 | at::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) | |||
| 53 | at::Tensor& gt_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) | 53 | at::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& | |||
| 72 | at::Tensor gt(const at::Tensor& self, const at::Tensor& other) | 81 | at::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 size | 94 | // 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) | |||
| 94 | at::Tensor& gt_(at::Tensor &self, const at::Tensor &other) | 113 | at::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 | 17 | ||
| 18 | 18 | ||
| 19 | 19 | ||
| 20 | +using npu_preparation = at_npu::native::OpPreparation; | ||
| 21 | + | ||
| 20 | namespace op_api { | 22 | namespace op_api { |
| 21 | 23 | ||
| 22 | at::Tensor &le_out(const at::Tensor &self, const at::Tensor &other, at::Tensor &result) | 24 | at::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 & | |||
| 49 | at::Tensor le(const at::Tensor &self, const at::Tensor &other) | 60 | at::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) | |||
| 84 | at::Tensor &le_(at::Tensor &self, const at::Tensor &other) | 105 | at::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 { | |||
| 21 | using npu_preparation = at_npu::native::OpPreparation; | 21 | using npu_preparation = at_npu::native::OpPreparation; |
| 22 | at::Tensor& lt_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) { | 22 | at::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 | ||
| 40 | at::Tensor lt(const at::Tensor& self, const at::Tensor& other) { | 49 | at::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 size | 61 | // 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 | ||
| 86 | at::Tensor& lt_(at::Tensor& self, const at::Tensor& other) { | 105 | at::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 | 17 | ||
| 18 | 18 | ||
| 19 | 19 | ||
| 20 | +using npu_preparation = at_npu::native::OpPreparation; | ||
| 21 | + | ||
| 20 | namespace op_api { | 22 | namespace op_api { |
| 21 | 23 | ||
| 22 | at::Tensor &maximum_out(const at::Tensor &self, const at::Tensor &other, at::Tensor &result) | 24 | at::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 | |||
| 41 | at::Tensor maximum(const at::Tensor &self, const at::Tensor &other) | 50 | at::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 | 18 | ||
| 19 | 19 | ||
| 20 | 20 | ||
| 21 | +using npu_preparation = at_npu::native::OpPreparation; | ||
| 22 | + | ||
| 21 | namespace op_api { | 23 | namespace op_api { |
| 22 | 24 | ||
| 23 | at::Tensor minimum(const at::Tensor& self, const at::Tensor& other) | 25 | at::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) | |||
| 33 | at::Tensor& minimum_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) | 44 | at::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; | |||
| 24 | at::Tensor& ne_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& result) | 24 | at::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& | |||
| 49 | at::Tensor ne(const at::Tensor& self, const at::Tensor& other) | 58 | at::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) | |||
| 78 | at::Tensor& ne_(at::Tensor& self, const at::Tensor& other) | 97 | at::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 | 852 | ||