已合并
fix linalg_cross bug #4338
zzhongmin创建于 2月28日
fix linalg_cross bug #4338
已合并
共 2 个文件变更+65-7
| @@ -2924,17 +2924,10 @@ official: | |||
| 2924 | - func: linalg_cross(Tensor self, Tensor other, *, int dim=-1) -> Tensor | 2924 | - func: linalg_cross(Tensor self, Tensor other, *, int dim=-1) -> Tensor |
| 2925 | acl_op: all_version | 2925 | acl_op: all_version |
| 2926 | op_api: all_version | 2926 | op_api: all_version |
| 2927 | - gen_opapi: | ||
| 2928 | - structured_inherit: linalg_cross.out | ||
| 2929 | 2927 | ||
| 2930 | - func: linalg_cross.out(Tensor self, Tensor other, *, int dim=-1, Tensor(a!) out) -> Tensor(a!) | 2928 | - func: linalg_cross.out(Tensor self, Tensor other, *, int dim=-1, Tensor(a!) out) -> Tensor(a!) |
| 2931 | acl_op: all_version | 2929 | acl_op: all_version |
| 2932 | op_api: all_version | 2930 | op_api: all_version |
| 2933 | - gen_opapi: | ||
| 2934 | - out: | ||
| 2935 | - size: broadcast_ops_npu_output_size(self, other) | ||
| 2936 | - dtype: self | ||
| 2937 | - exec: aclnnLinalgCross | ||
| 2938 | 2931 | ||
| 2939 | - func: linalg_qr(Tensor self, str mode='reduced') -> (Tensor Q, Tensor R) | 2932 | - func: linalg_qr(Tensor self, str mode='reduced') -> (Tensor Q, Tensor R) |
| 2940 | acl_op: all_version | 2933 | acl_op: all_version |
| @@ -0,0 +1,65 @@ | |||
| 1 | +// Copyright (c) 2026 Huawei Technologies Co., Ltd | ||
| 2 | +// All rights reserved. | ||
| 3 | +// | ||
| 4 | +// Licensed under the BSD 3-Clause License (the "License"); | ||
| 5 | +// you may not use this file except in compliance with the License. | ||
| 6 | +// You may obtain a copy of the License at | ||
| 7 | +// | ||
| 8 | +// https://opensource.org/licenses/BSD-3-Clause | ||
| 9 | +// | ||
| 10 | +// Unless required by applicable law or agreed to in writing, software | ||
| 11 | +// distributed under the License is distributed on an "AS IS" BASIS, | ||
| 12 | +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| 13 | +// See the License for the specific language governing permissions and | ||
| 14 | +// limitations under the License. | ||
| 15 | + | ||
| 16 | + | ||
| 17 | + | ||
| 18 | + | ||
| 19 | + | ||
| 20 | + | ||
| 21 | +namespace op_api { | ||
| 22 | +using npu_preparation = at_npu::native::OpPreparation; | ||
| 23 | +using namespace op_infer; | ||
| 24 | + | ||
| 25 | +namespace { | ||
| 26 | +constexpr int64_t kCrossVectorSize = 3; | ||
| 27 | + | ||
| 28 | +inline void linalg_cross_check(const at::Tensor& self, const at::Tensor& other, int64_t dim) | ||
| 29 | +{ | ||
| 30 | + auto x_d = self.dim(); | ||
| 31 | + auto y_d = other.dim(); | ||
| 32 | + TORCH_CHECK(x_d == y_d, "linalg.cross: inputs must have the same number of dimensions.", | ||
| 33 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 34 | + auto wrap_dim = at::maybe_wrap_dim(dim, x_d); | ||
| 35 | + TORCH_CHECK(self.size(wrap_dim) == kCrossVectorSize && other.size(wrap_dim) == kCrossVectorSize, | ||
C | |||
| 36 | + "linalg.cross: inputs dimension ", wrap_dim, " must have length ", kCrossVectorSize, ". Got ", self.size(wrap_dim), | ||
| 37 | + " and ", other.size(wrap_dim), OPS_ERROR(ErrCode::PARAM)); | ||
| 38 | +} | ||
| 39 | +} // namespace | ||
| 40 | + | ||
| 41 | +at::Tensor linalg_cross(const at::Tensor& self, const at::Tensor& other, int64_t dim) | ||
| 42 | +{ | ||
| 43 | + DO_COMPATIBILITY(aclnnLinalgCross, acl_op::linalg_cross(self, other, dim)); | ||
| 44 | + linalg_cross_check(self, other, dim); | ||
| 45 | + auto wrap_dim = at::maybe_wrap_dim(dim, self.dim()); | ||
| 46 | + auto output_size_0 = broadcast_ops_npu_output_size(self, other); | ||
| 47 | + auto output_dtype_0 = self.scalar_type(); | ||
| 48 | + at::Tensor out = npu_preparation::apply_tensor_without_format(output_size_0, | ||
| 49 | + self.options().dtype(output_dtype_0)); | ||
| 50 | + EXEC_NPU_CMD(aclnnLinalgCross, self, other, wrap_dim, out); | ||
| 51 | + return out; | ||
| 52 | +} | ||
| 53 | + | ||
| 54 | +at::Tensor& linalg_cross_out(const at::Tensor& self, const at::Tensor& other, int64_t dim, at::Tensor& out) | ||
| 55 | +{ | ||
| 56 | + DO_COMPATIBILITY(aclnnLinalgCross, acl_op::linalg_cross_out(self, other, dim, out)); | ||
| 57 | + linalg_cross_check(self, other, dim); | ||
| 58 | + auto wrap_dim = at::maybe_wrap_dim(dim, self.dim()); | ||
| 59 | + auto output_size_0 = broadcast_ops_npu_output_size(self, other); | ||
| 60 | + auto output_dtype_0 = self.scalar_type(); | ||
| 61 | + npu_preparation::check_tensor({self, other}, out, output_dtype_0, output_size_0); | ||
| 62 | + EXEC_NPU_CMD(aclnnLinalgCross, self, other, wrap_dim, out); | ||
| 63 | + return out; | ||
| 64 | +} | ||
| 65 | +} // namespace op_api | ||
是因为需要加这个check，才需要手动对接？