已合并
[feat]Bucketize新增Aicore通路支持 #4711
jinpenghe创建于 4月13日
[feat]Bucketize新增Aicore通路支持 #4711
已合并
共 2 个文件变更+39-5
| @@ -23,7 +23,19 @@ at::Tensor bucketize(const at::Tensor& self, const at::Tensor& boundaries, bool | |||
| 23 | { | 23 | { |
| 24 | TORCH_CHECK(boundaries.dim() == 1, "boundaries tensor must be 1 dimension, but got dim(", | 24 | TORCH_CHECK(boundaries.dim() == 1, "boundaries tensor must be 1 dimension, but got dim(", |
| 25 | boundaries.dim(), ")" + OPS_ERROR(ErrCode::PARAM)); | 25 | boundaries.dim(), ")" + OPS_ERROR(ErrCode::PARAM)); |
| 26 | - return op_api::searchsorted(boundaries, self, out_int32, right, c10::nullopt, c10::nullopt); | 26 | + |
| 27 | + static bool isRegBaseSoc = c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950; | ||
| 28 | + bool typeSupport = self.dtype() != at::kDouble && boundaries.dtype() != at::kDouble; | ||
| 29 | + static const bool is_bucketize_available = check_aclnn_kernel_available("aclnnBucketize"); | ||
| 30 | + if (isRegBaseSoc && typeSupport && is_bucketize_available) { | ||
| 31 | + at::ScalarType expectedType = out_int32 ? at::kInt : at::kLong; | ||
| 32 | + at::TensorOptions opts = self.options().dtype(expectedType); | ||
| 33 | + at::Tensor result = at_npu::native::OpPreparation::apply_tensor_without_format(self.sizes(), opts); | ||
| 34 | + EXEC_NPU_CMD(aclnnBucketize, self, boundaries, out_int32, right, result); | ||
| 35 | + return result; | ||
| 36 | + } else { | ||
| 37 | + return op_api::searchsorted(boundaries, self, out_int32, right, c10::nullopt, c10::nullopt); | ||
| 38 | + } | ||
| 27 | } | 39 | } |
| 28 | 40 | ||
| 29 | at::Tensor bucketize(const at::Scalar& self, const at::Tensor& boundaries, bool out_int32, bool right) | 41 | at::Tensor bucketize(const at::Scalar& self, const at::Tensor& boundaries, bool out_int32, bool right) |
| @@ -42,6 +54,19 @@ at::Tensor &bucketize_out( | |||
| 42 | { | 54 | { |
| 43 | TORCH_CHECK(boundaries.dim() == 1, "boundaries tensor must be 1 dimension, but got dim(", boundaries.dim(), | 55 | TORCH_CHECK(boundaries.dim() == 1, "boundaries tensor must be 1 dimension, but got dim(", boundaries.dim(), |
| 44 | ")" + OPS_ERROR(ErrCode::PARAM)); | 56 | ")" + OPS_ERROR(ErrCode::PARAM)); |
| 45 | - return op_api::searchsorted_out(boundaries, self, out_int32, right, c10::nullopt, c10::nullopt, out); | 57 | + |
 | |||
| 58 | + static bool isRegBaseSoc = c10_npu::GetSocVersion() >= c10_npu::SocVersion::Ascend950; | ||
| 59 | + bool typeSupport = self.dtype() != at::kDouble && boundaries.dtype() != at::kDouble; | ||
| 60 | + static const bool is_bucketize_available = check_aclnn_kernel_available("aclnnBucketize"); | ||
| 61 | + if (isRegBaseSoc && typeSupport && is_bucketize_available) { | ||
| 62 | + at_npu::native::OpPreparation::check_tensor({self, boundaries}, | ||
| 63 | + out, | ||
| 64 | + out.scalar_type(), | ||
| 65 | + self.sizes()); | ||
| 66 | + EXEC_NPU_CMD(aclnnBucketize, self, boundaries, out_int32, right, out); | ||
| 67 | + return out; | ||
| 68 | + } else { | ||
| 69 | + return op_api::searchsorted_out(boundaries, self, out_int32, right, c10::nullopt, c10::nullopt, out); | ||
| 70 | + } | ||
| 46 | } | 71 | } |
| 47 | } | 72 | } |
| @@ -8,9 +8,18 @@ from torch_npu.testing.testcase import TestCase, run_tests | |||
| 8 | 8 | ||
| 9 | class TestBucketize(TestCase): | 9 | class TestBucketize(TestCase): |
| 10 | 10 | ||
| 11 | - def test_bucketize(self): | 11 | + def test_bucketize_aclnn_search(self): |
| 12 | - v = torch.tensor([[3, 6, 9], [3, 6, 9]]) | 12 | + v = torch.tensor([[3, 6, 9], [3, 6, 9]], dtype=torch.float64) |
| 13 | - boundaries = torch.tensor([1, 3, 5, 7, 9]) | 13 | + boundaries = torch.tensor([1, 3, 5, 7, 9], dtype=torch.float64) |
| 14 | + | ||
| 15 | + cpu_output = torch.bucketize(v, boundaries) | ||
| 16 | + npu_output = torch.bucketize(v.npu(), boundaries.npu()) | ||
| 17 | + | ||
| 18 | + self.assertRtolEqual(cpu_output, npu_output) | ||
| 19 | + | ||
| 20 | + def test_bucketize_aclnn_bucketize(self): | ||
| 21 | + v = torch.tensor([[3, 6, 9], [3, 6, 9]], dtype=torch.int32) | ||
| 22 | + boundaries = torch.tensor([1, 3, 5, 7, 9], dtype=torch.int32) | ||
| 14 | 23 | ||
| 15 | cpu_output = torch.bucketize(v, boundaries) | 24 | cpu_output = torch.bucketize(v, boundaries) |
| 16 | npu_output = torch.bucketize(v.npu(), boundaries.npu()) | 25 | npu_output = torch.bucketize(v.npu(), boundaries.npu()) |
新增业务逻辑，补充ut