已合并
[master][Fix] Fix static check errors detected by clang-format #5595
thickhair创建于 7月29日
[master][Fix] Fix static check errors detected by clang-format #5595
已合并
共 60 个文件变更+4376-3813
| @@ -35,7 +35,7 @@ BreakBeforeTernaryOperators: true | |||
| 35 | BreakConstructorInitializersBeforeComma: false | 35 | BreakConstructorInitializersBeforeComma: false |
| 36 | BreakAfterJavaFieldAnnotations: false | 36 | BreakAfterJavaFieldAnnotations: false |
| 37 | BreakStringLiterals: false | 37 | BreakStringLiterals: false |
| 38 | -ColumnLimit: 80 | 38 | +ColumnLimit: 120 |
| 39 | CommentPragmas: '^ IWYU pragma:' | 39 | CommentPragmas: '^ IWYU pragma:' |
| 40 | CompactNamespaces: false | 40 | CompactNamespaces: false |
| 41 | ConstructorInitializerAllOnOneLineOrOnePerLine: true | 41 | ConstructorInitializerAllOnOneLineOrOnePerLine: true |
| @@ -28,33 +28,32 @@ std::tuple<at::Tensor&, at::Tensor&, at::Tensor&> unique_consecutive_out_nocheck | |||
| 28 | const at::Tensor& self, | 28 | const at::Tensor& self, |
| 29 | const bool return_inverse, | 29 | const bool return_inverse, |
| 30 | const bool return_counts, | 30 | const bool return_counts, |
| 31 | - c10::optional<int64_t> dim) | 31 | + c10::optional<int64_t> dim) { |
| 32 | -{ | 32 | + at::Tensor self_copy = self; |
| 33 | - at::Tensor self_copy = self; | 33 | + if (self.scalar_type() == at::ScalarType::Half) { |
| 34 | - if (self.scalar_type() == at::ScalarType::Half) { | 34 | + self_copy = at_npu::native::custom_ops::_npu_dtype_cast(self, at::ScalarType::Float); |
| 35 | - self_copy = at_npu::native::custom_ops::_npu_dtype_cast(self, at::ScalarType::Float); | 35 | + output = at_npu::native::custom_ops::_npu_dtype_cast(output, at::ScalarType::Float); |
| 36 | - output = at_npu::native::custom_ops::_npu_dtype_cast(output, at::ScalarType::Float); | 36 | + } |
| 37 | - } | 37 | + c10::SmallVector<int64_t, N> output_sync_idx = {0, 2}; |
| 38 | - c10::SmallVector<int64_t, N> output_sync_idx = {0, 2}; | 38 | + at_npu::native::OpCommand cmd; |
| 39 | - at_npu::native::OpCommand cmd; | 39 | + cmd.Sync(output_sync_idx) |
| 40 | - cmd.Sync(output_sync_idx) | 40 | + .Name("UniqueConsecutive") |
| 41 | - .Name("UniqueConsecutive") | 41 | + .Input(self_copy) |
| 42 | - .Input(self_copy) | 42 | + .Output(output) |
| 43 | - .Output(output) | 43 | + .Output(inverse_indices) |
| 44 | - .Output(inverse_indices) | 44 | + .Output(counts) |
| 45 | - .Output(counts) | 45 | + .Attr("return_idx", return_inverse) |
| 46 | - .Attr("return_idx", return_inverse) | 46 | + .Attr("return_counts", return_counts); |
| 47 | - .Attr("return_counts", return_counts); | 47 | + if (dim.has_value()) { |
| 48 | - if (dim.has_value()) { | 48 | + cmd.Attr("axis", dim.value()); |
| 49 | - cmd.Attr("axis", dim.value()); | 49 | + } |
| 50 | - } | 50 | + cmd.Run(); |
| 51 | - cmd.Run(); | 51 | + if (self.scalar_type() == at::ScalarType::Half) { |
| 52 | - if (self.scalar_type() == at::ScalarType::Half) { | 52 | + output = at_npu::native::custom_ops::_npu_dtype_cast(output, at::ScalarType::Half); |
| 53 | - output = at_npu::native::custom_ops::_npu_dtype_cast(output, at::ScalarType::Half); | 53 | + } |
| 54 | - } | 54 | + return std::tie(output, inverse_indices, counts); |
| 55 | - return std::tie(output, inverse_indices, counts); | ||
| 56 | -} | ||
| 57 | } | 55 | } |
| 56 | +} // namespace | ||
| 58 | 57 | ||
| 59 | std::tuple<at::Tensor&, at::Tensor&, at::Tensor&> unique_consecutive_out_npu( | 58 | std::tuple<at::Tensor&, at::Tensor&, at::Tensor&> unique_consecutive_out_npu( |
| 60 | at::Tensor& output, | 59 | at::Tensor& output, |
| @@ -63,60 +62,59 @@ std::tuple<at::Tensor&, at::Tensor&, at::Tensor&> unique_consecutive_out_npu( | |||
| 63 | const at::Tensor& self, | 62 | const at::Tensor& self, |
| 64 | const bool return_inverse, | 63 | const bool return_inverse, |
| 65 | const bool return_counts, | 64 | const bool return_counts, |
| 66 | - c10::optional<int64_t> dim) | 65 | + c10::optional<int64_t> dim) { |
| 67 | -{ | 66 | + if (dim.has_value()) { |
| 68 | - if (dim.has_value()) { | 67 | + npu_preparation::CheckOut({self}, output, self); |
| 69 | - npu_preparation::CheckOut({self}, output, self); | 68 | + npu_preparation::CheckOut({self}, inverse_indices, ACL_FORMAT_ND, self.scalar_type(), self.size(dim.value())); |
| 70 | - npu_preparation::CheckOut({self}, inverse_indices, ACL_FORMAT_ND, self.scalar_type(), self.size(dim.value())); | 69 | + npu_preparation::CheckOut({self}, counts, ACL_FORMAT_ND, self.scalar_type(), self.size(dim.value())); |
| 71 | - npu_preparation::CheckOut({self}, counts, ACL_FORMAT_ND, self.scalar_type(), self.size(dim.value())); | 70 | + } else { |
| 72 | - } else { | 71 | + npu_preparation::CheckOut({self}, output, self, {self.numel()}); |
| 73 | - npu_preparation::CheckOut({self}, output, self, {self.numel()}); | 72 | + npu_preparation::CheckOut({self}, inverse_indices, ACL_FORMAT_ND, self.scalar_type(), self.sizes()); |
| 74 | - npu_preparation::CheckOut({self}, inverse_indices, ACL_FORMAT_ND, self.scalar_type(), self.sizes()); | 73 | + npu_preparation::CheckOut({self}, counts, ACL_FORMAT_ND, self.scalar_type(), self.numel()); |
| 75 | - npu_preparation::CheckOut({self}, counts, ACL_FORMAT_ND, self.scalar_type(), self.numel()); | 74 | + } |
| 76 | - } | ||
| 77 | 75 | ||
| 78 | - bool output_match = npu_utils::check_match(&output); | 76 | + bool output_match = npu_utils::check_match(&output); |
| 79 | - bool indices_match = npu_utils::check_match(&inverse_indices); | 77 | + bool indices_match = npu_utils::check_match(&inverse_indices); |
| 80 | - bool counts_match = npu_utils::check_match(&counts); | 78 | + bool counts_match = npu_utils::check_match(&counts); |
| 81 | - if (output_match && indices_match && counts_match) { | 79 | + if (output_match && indices_match && counts_match) { |
| 82 | - unique_consecutive_out_nocheck( | 80 | + unique_consecutive_out_nocheck(output, inverse_indices, counts, self, return_inverse, return_counts, dim); |
| 83 | - output, inverse_indices, counts, self, return_inverse, return_counts, dim); | 81 | + } else { |
| 84 | - } else { | 82 | + at::Tensor contig_output = output_match ? output : npu_utils::format_contiguous(output); |
| 85 | - at::Tensor contig_output = output_match ? output : npu_utils::format_contiguous(output); | 83 | + at::Tensor contig_indices = indices_match ? inverse_indices : npu_utils::format_contiguous(inverse_indices); |
| 86 | - at::Tensor contig_indices = indices_match ? inverse_indices : npu_utils::format_contiguous(inverse_indices); | 84 | + at::Tensor contig_counts = counts_match ? counts : npu_utils::format_contiguous(counts); |
| 87 | - at::Tensor contig_counts = counts_match ? counts : npu_utils::format_contiguous(counts); | 85 | + unique_consecutive_out_nocheck( |
| 88 | - unique_consecutive_out_nocheck( | 86 | + contig_output, contig_indices, contig_counts, self, return_inverse, return_counts, dim); |
| 89 | - contig_output, contig_indices, contig_counts, self, return_inverse, return_counts, dim); | 87 | + if (!output_match) { |
| 90 | - if (!output_match) { | 88 | + npu_utils::format_fresh_view(output, contig_output); |
| 91 | - npu_utils::format_fresh_view(output, contig_output); | ||
| 92 | - } | ||
| 93 | - if (!indices_match) { | ||
| 94 | - npu_utils::format_fresh_view(inverse_indices, contig_indices); | ||
| 95 | - } | ||
| 96 | - if (!counts_match) { | ||
| 97 | - npu_utils::format_fresh_view(counts, contig_counts); | ||
| 98 | - } | ||
| 99 | } | 89 | } |
| 90 | + if (!indices_match) { | ||
| 91 | + npu_utils::format_fresh_view(inverse_indices, contig_indices); | ||
| 92 | + } | ||
| 93 | + if (!counts_match) { | ||
| 94 | + npu_utils::format_fresh_view(counts, contig_counts); | ||
| 95 | + } | ||
| 96 | + } | ||
| 100 | 97 | ||
| 101 | - return std::tie(output, inverse_indices, counts); | 98 | + return std::tie(output, inverse_indices, counts); |
| 102 | } | 99 | } |
| 103 | 100 | ||
| 104 | std::tuple<at::Tensor, at::Tensor, at::Tensor> unique_consecutive( | 101 | std::tuple<at::Tensor, at::Tensor, at::Tensor> unique_consecutive( |
| 105 | const at::Tensor& self, | 102 | const at::Tensor& self, |
| 106 | bool return_inverse, | 103 | bool return_inverse, |
| 107 | bool return_counts, | 104 | bool return_counts, |
| 108 | - c10::optional<int64_t> dim) | 105 | + c10::optional<int64_t> dim) { |
| 109 | -{ | 106 | + at::Tensor output = |
| 110 | - at::Tensor output = (dim.has_value()) ? | 107 | + (dim.has_value()) ? npu_preparation::apply_tensor(self) : npu_preparation::apply_tensor(self, {self.numel()}); |
| 111 | - npu_preparation::apply_tensor(self) : npu_preparation::apply_tensor(self, {self.numel()}); | 108 | + at::Tensor inverse_indices = (dim.has_value()) |
| 112 | - at::Tensor inverse_indices = (dim.has_value()) ? | 109 | + ? npu_preparation::apply_tensor_with_format( |
| 113 | - npu_preparation::apply_tensor_with_format(self.size(dim.value()), self.options().dtype(at::kLong), ACL_FORMAT_ND) : | 110 | + self.size(dim.value()), self.options().dtype(at::kLong), ACL_FORMAT_ND) |
| 114 | - npu_preparation::apply_tensor_with_format(self.sizes(), self.options().dtype(at::kLong), ACL_FORMAT_ND); | 111 | + : npu_preparation::apply_tensor_with_format(self.sizes(), self.options().dtype(at::kLong), ACL_FORMAT_ND); |
| 115 | - at::Tensor counts = (dim.has_value()) ? | 112 | + at::Tensor counts = (dim.has_value()) |
| 116 | - npu_preparation::apply_tensor_with_format(self.size(dim.value()), self.options().dtype(at::kLong), ACL_FORMAT_ND) : | 113 | + ? npu_preparation::apply_tensor_with_format( |
| 117 | - npu_preparation::apply_tensor_with_format({self.numel()}, self.options().dtype(at::kLong), ACL_FORMAT_ND); | 114 | + self.size(dim.value()), self.options().dtype(at::kLong), ACL_FORMAT_ND) |
| 118 | - unique_consecutive_out_nocheck(output, inverse_indices, counts, self, return_inverse, return_counts, dim); | 115 | + : npu_preparation::apply_tensor_with_format({self.numel()}, self.options().dtype(at::kLong), ACL_FORMAT_ND); |
| 119 | - return std::tie(output, inverse_indices, counts); | 116 | + unique_consecutive_out_nocheck(output, inverse_indices, counts, self, return_inverse, return_counts, dim); |
| 117 | + return std::tie(output, inverse_indices, counts); | ||
| 120 | } | 118 | } |
| 121 | 119 | ||
| 122 | } // namespace acl_op | 120 | } // namespace acl_op |
| @@ -29,35 +29,37 @@ at::Tensor& upsample_bicubic2d_backward_out_nocheck( | |||
| 29 | at::IntArrayRef input_size, | 29 | at::IntArrayRef input_size, |
| 30 | bool align_corners, | 30 | bool align_corners, |
| 31 | c10::optional<double> scales_h, | 31 | c10::optional<double> scales_h, |
| 32 | - c10::optional<double> scales_w) | 32 | + c10::optional<double> scales_w) { |
| 33 | -{ | 33 | + TORCH_CHECK( |
| 34 | - TORCH_CHECK(output_size.size() == 2, | 34 | + output_size.size() == 2, |
| 35 | - "It is expected output_size equals to 2, but got size ", | 35 | + "It is expected output_size equals to 2, but got size ", |
| 36 | - output_size.size(), OPS_ERROR(ErrCode::PARAM)); | 36 | + output_size.size(), |
| 37 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 37 | 38 | ||
| 38 | - TORCH_CHECK(input_size.size() == 4, | 39 | + TORCH_CHECK( |
| 39 | - "It is expected input_size equals to 4, but got size ", | 40 | + input_size.size() == 4, |
| 40 | - input_size.size(), OPS_ERROR(ErrCode::PARAM)); | 41 | + "It is expected input_size equals to 4, but got size ", |
| 42 | + input_size.size(), | ||
| 43 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 41 | 44 | ||
| 42 | - float temp_h = 0.0; | 45 | + float temp_h = 0.0; |
| 43 | - float temp_w = 0.0; | 46 | + float temp_w = 0.0; |
| 44 | - temp_h = scales_h.has_value() ? static_cast<float>(scales_h.value()) : temp_h; | 47 | + temp_h = scales_h.has_value() ? static_cast<float>(scales_h.value()) : temp_h; |
| 45 | - temp_w = scales_w.has_value() ? static_cast<float>(scales_w.value()) : temp_w; | 48 | + temp_w = scales_w.has_value() ? static_cast<float>(scales_w.value()) : temp_w; |
| 46 | - c10::SmallVector<float, N> scales = {temp_h, temp_w}; | 49 | + c10::SmallVector<float, N> scales = {temp_h, temp_w}; |
| 47 | - c10::SmallVector<float, N> roi = {}; | 50 | + c10::SmallVector<float, N> roi = {}; |
| 48 | - string coordinate_transformation_mode = | 51 | + string coordinate_transformation_mode = align_corners ? "align_corners" : "half_pixel"; |
| 49 | - align_corners ? "align_corners" : "half_pixel"; | ||
| 50 | 52 | ||
| 51 | - float cu = -0.75; | 53 | + float cu = -0.75; |
| 52 | - int64_t ex = 0; | 54 | + int64_t ex = 0; |
| 53 | - float ext = 0.0; | 55 | + float ext = 0.0; |
| 54 | - string mode = "cubic"; | 56 | + string mode = "cubic"; |
| 55 | - string ne = "round_prefer_floor"; | 57 | + string ne = "round_prefer_floor"; |
| 56 | - string data_format = "HWNC"; | 58 | + string data_format = "HWNC"; |
| 57 | - at_npu::native::OpCommand cmd; | 59 | + at_npu::native::OpCommand cmd; |
| 58 | - at::Tensor grad_output_transpose = grad_output.permute({2, 3, 0, 1}); | 60 | + at::Tensor grad_output_transpose = grad_output.permute({2, 3, 0, 1}); |
| 59 | 61 | ||
| 60 | - cmd.Name("ResizeGradD") | 62 | + cmd.Name("ResizeGradD") |
| 61 | .Input(grad_output_transpose, "grads") | 63 | .Input(grad_output_transpose, "grads") |
| 62 | .Output(grad_input, "y") | 64 | .Output(grad_input, "y") |
| 63 | .Attr("scales", scales) | 65 | .Attr("scales", scales) |
| @@ -71,7 +73,7 @@ at::Tensor& upsample_bicubic2d_backward_out_nocheck( | |||
| 71 | .Attr("nearest_mode", ne) | 73 | .Attr("nearest_mode", ne) |
| 72 | .Attr("data_format", data_format) | 74 | .Attr("data_format", data_format) |
| 73 | .Run(); | 75 | .Run(); |
| 74 | - return grad_input; | 76 | + return grad_input; |
| 75 | } | 77 | } |
| 76 | } // namespace | 78 | } // namespace |
| 77 | 79 | ||
| @@ -82,33 +84,28 @@ at::Tensor& upsample_bicubic2d_backward_out( | |||
| 82 | bool align_corners, | 84 | bool align_corners, |
| 83 | c10::optional<double> scales_h, | 85 | c10::optional<double> scales_h, |
| 84 | c10::optional<double> scales_w, | 86 | c10::optional<double> scales_w, |
| 85 | - at::Tensor& grad_input) | 87 | + at::Tensor& grad_input) { |
| 86 | -{ | 88 | + auto op_infer_output_size = op_infer::upsample_bicubic2d_backward_npu_output_size(input_size); |
| 87 | - auto op_infer_output_size = op_infer::upsample_bicubic2d_backward_npu_output_size(input_size); | ||
| 88 | 89 | ||
| 89 | - npu_preparation::CheckOut( | 90 | + npu_preparation::CheckOut({grad_output}, grad_input, grad_output, op_infer_output_size); |
| 90 | - {grad_output}, | ||
| 91 | - grad_input, | ||
| 92 | - grad_output, | ||
| 93 | - op_infer_output_size); | ||
| 94 | 91 | ||
| 95 | - if (!npu_utils::check_match(&grad_input)) { | 92 | + if (!npu_utils::check_match(&grad_input)) { |
| 96 | - at::Tensor contiguous_result = npu_utils::format_contiguous(grad_input); | 93 | + at::Tensor contiguous_result = npu_utils::format_contiguous(grad_input); |
| 97 | - upsample_bicubic2d_backward_out_nocheck( | 94 | + upsample_bicubic2d_backward_out_nocheck( |
| 98 | - contiguous_result, grad_output, output_size, input_size, align_corners, scales_h, scales_w); | 95 | + contiguous_result, grad_output, output_size, input_size, align_corners, scales_h, scales_w); |
| 99 | - npu_utils::format_fresh_view(grad_input, contiguous_result); | 96 | + npu_utils::format_fresh_view(grad_input, contiguous_result); |
| 100 | - } else { | 97 | + } else { |
| 101 | - upsample_bicubic2d_backward_out_nocheck( | 98 | + upsample_bicubic2d_backward_out_nocheck( |
| 102 | - grad_input, grad_output, output_size, input_size, align_corners, scales_h, scales_w); | 99 | + grad_input, grad_output, output_size, input_size, align_corners, scales_h, scales_w); |
| 103 | - } | 100 | + } |
| 104 | 101 | ||
| 105 | - int64_t N = grad_output.size(0); | 102 | + int64_t N = grad_output.size(0); |
| 106 | - int64_t C = grad_output.size(1); | 103 | + int64_t C = grad_output.size(1); |
| 107 | - int64_t H = input_size[2]; | 104 | + int64_t H = input_size[2]; |
| 108 | - int64_t W = input_size[3]; | 105 | + int64_t W = input_size[3]; |
| 109 | 106 | ||
| 110 | - grad_input = grad_input.reshape({H, W, N, C}).permute({2, 3, 0, 1}); | 107 | + grad_input = grad_input.reshape({H, W, N, C}).permute({2, 3, 0, 1}); |
| 111 | - return grad_input; | 108 | + return grad_input; |
| 112 | } | 109 | } |
| 113 | 110 | ||
| 114 | at::Tensor upsample_bicubic2d_backward( | 111 | at::Tensor upsample_bicubic2d_backward( |
| @@ -117,18 +114,17 @@ at::Tensor upsample_bicubic2d_backward( | |||
| 117 | at::IntArrayRef input_size, | 114 | at::IntArrayRef input_size, |
| 118 | bool align_corners, | 115 | bool align_corners, |
| 119 | c10::optional<double> scales_h, | 116 | c10::optional<double> scales_h, |
| 120 | - c10::optional<double> scales_w) | 117 | + c10::optional<double> scales_w) { |
| 121 | -{ | 118 | + auto op_infer_output_size = op_infer::upsample_bicubic2d_backward_npu_output_size(input_size); |
| 122 | - auto op_infer_output_size = op_infer::upsample_bicubic2d_backward_npu_output_size(input_size); | 119 | + at::Tensor result = npu_preparation::apply_tensor(grad_output, op_infer_output_size); |
| 123 | - at::Tensor result = npu_preparation::apply_tensor(grad_output, op_infer_output_size); | 120 | + upsample_bicubic2d_backward_out_nocheck( |
| 124 | - upsample_bicubic2d_backward_out_nocheck(result, grad_output, output_size, input_size, | 121 | + result, grad_output, output_size, input_size, align_corners, scales_h, scales_w); |
| 125 | - align_corners, scales_h, scales_w); | ||
| 126 | 122 | ||
| 127 | - int64_t N = grad_output.size(0); | 123 | + int64_t N = grad_output.size(0); |
| 128 | - int64_t C = grad_output.size(1); | 124 | + int64_t C = grad_output.size(1); |
| 129 | - int64_t H = input_size[2]; | 125 | + int64_t H = input_size[2]; |
| 130 | - int64_t W = input_size[3]; | 126 | + int64_t W = input_size[3]; |
| 131 | - result = result.reshape({H, W, N, C}).permute({2, 3, 0, 1}); | 127 | + result = result.reshape({H, W, N, C}).permute({2, 3, 0, 1}); |
| 132 | - return result; | 128 | + return result; |
| 133 | } | 129 | } |
| 134 | } // namespace acl_op | 130 | } // namespace acl_op |
| @@ -28,44 +28,44 @@ at::Tensor& upsample_bicubic2d_out_nocheck( | |||
| 28 | at::IntArrayRef output_size, | 28 | at::IntArrayRef output_size, |
| 29 | bool align_corners, | 29 | bool align_corners, |
| 30 | c10::optional<double> scales_h, | 30 | c10::optional<double> scales_h, |
| 31 | - c10::optional<double> scales_w) | 31 | + c10::optional<double> scales_w) { |
| 32 | -{ | 32 | + TORCH_CHECK( |
| 33 | - TORCH_CHECK( | 33 | + output_size.size() == 2, |
| 34 | - output_size.size() == 2, | 34 | + "It is expected output_size equals to 2, but got size ", |
| 35 | - "It is expected output_size equals to 2, but got size ", | 35 | + output_size.size(), |
| 36 | - output_size.size(), OPS_ERROR(ErrCode::PARAM)); | 36 | + OPS_ERROR(ErrCode::PARAM)); |
| 37 | 37 | ||
| 38 | - float temp_h = 0.0; | 38 | + float temp_h = 0.0; |
| 39 | - float temp_w = 0.0; | 39 | + float temp_w = 0.0; |
| 40 | - if (scales_h.has_value()) { | 40 | + if (scales_h.has_value()) { |
| 41 | - temp_h = (float)scales_h.value(); | 41 | + temp_h = (float)scales_h.value(); |
| 42 | - } | 42 | + } |
| 43 | - if (scales_w.has_value()) { | 43 | + if (scales_w.has_value()) { |
| 44 | - temp_w = (float)scales_w.value(); | 44 | + temp_w = (float)scales_w.value(); |
| 45 | - } | 45 | + } |
| 46 | - c10::SmallVector<float, SIZE> scales = {temp_h, temp_w}; | 46 | + c10::SmallVector<float, SIZE> scales = {temp_h, temp_w}; |
| 47 | - c10::SmallVector<float, SIZE> roi = {}; | 47 | + c10::SmallVector<float, SIZE> roi = {}; |
| 48 | - string coordinate_transformation_mode = "half_pixel"; | 48 | + string coordinate_transformation_mode = "half_pixel"; |
| 49 | - if (align_corners == true) { | 49 | + if (align_corners == true) { |
| 50 | - coordinate_transformation_mode = "align_corners"; | 50 | + coordinate_transformation_mode = "align_corners"; |
| 51 | - } | 51 | + } |
| 52 | 52 | ||
| 53 | - at_npu::native::OpCommand cmd; | 53 | + at_npu::native::OpCommand cmd; |
| 54 | - cmd.Name("ResizeD") | 54 | + cmd.Name("ResizeD") |
| 55 | - .Input(self, "X") | 55 | + .Input(self, "X") |
| 56 | - .Output(result, "y") | 56 | + .Output(result, "y") |
| 57 | - .Attr("sizes", output_size) | 57 | + .Attr("sizes", output_size) |
| 58 | - .Attr("scales", scales) | 58 | + .Attr("scales", scales) |
| 59 | - .Attr("roi", roi) | 59 | + .Attr("roi", roi) |
| 60 | - .Attr("coordinate_transformation_mode", coordinate_transformation_mode) | 60 | + .Attr("coordinate_transformation_mode", coordinate_transformation_mode) |
| 61 | - .Attr("cubic_coeff_a", (float)-0.75) | 61 | + .Attr("cubic_coeff_a", (float)-0.75) |
| 62 | - .Attr("exclude_outside", (int64_t)0) | 62 | + .Attr("exclude_outside", (int64_t)0) |
| 63 | - .Attr("extrapolation_value", (float)0.0) | 63 | + .Attr("extrapolation_value", (float)0.0) |
| 64 | - .Attr("mode", (string)"cubic") | 64 | + .Attr("mode", (string) "cubic") |
| 65 | - .Attr("nearest_mode", (string)"round_prefer_floor") | 65 | + .Attr("nearest_mode", (string) "round_prefer_floor") |
| 66 | - .Run(); | 66 | + .Run(); |
| 67 | 67 | ||
| 68 | - return result; | 68 | + return result; |
| 69 | } | 69 | } |
| 70 | } // namespace | 70 | } // namespace |
| 71 | 71 | ||
| @@ -75,35 +75,32 @@ at::Tensor& upsample_bicubic2d_out( | |||
| 75 | bool align_corners, | 75 | bool align_corners, |
| 76 | c10::optional<double> scales_h, | 76 | c10::optional<double> scales_h, |
| 77 | c10::optional<double> scales_w, | 77 | c10::optional<double> scales_w, |
| 78 | - at::Tensor& result) | 78 | + at::Tensor& result) { |
| 79 | -{ | 79 | + TORCH_CHECK( |
| 80 | - TORCH_CHECK(self.dim() >= 2, "The self shoud be at least 2D, but self got", self.dim(), | 80 | + self.dim() >= 2, "The self shoud be at least 2D, but self got", self.dim(), "D" + OPS_ERROR(ErrCode::PARAM)); |
| 81 | - "D" + OPS_ERROR(ErrCode::PARAM)); | 81 | + TORCH_CHECK( |
| 82 | - TORCH_CHECK(output_size.size() == 2, | 82 | + output_size.size() == 2, |
| 83 | - "It is expected output_size equals to 2, but got size ", | 83 | + "It is expected output_size equals to 2, but got size ", |
| 84 | - output_size.size(), OPS_ERROR(ErrCode::PARAM)); | 84 | + output_size.size(), |
| 85 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 85 | 86 | ||
| 86 | - int64_t N = self.size(0); | 87 | + int64_t N = self.size(0); |
| 87 | - int64_t C = self.size(1); | 88 | + int64_t C = self.size(1); |
| 88 | - int64_t H = output_size[0]; | 89 | + int64_t H = output_size[0]; |
| 89 | - int64_t W = output_size[1]; | 90 | + int64_t W = output_size[1]; |
| 90 | 91 | ||
| 91 | - c10::SmallVector<int64_t, SIZE> op_infer_output_size = {N, C, H, W}; | 92 | + c10::SmallVector<int64_t, SIZE> op_infer_output_size = {N, C, H, W}; |
| 92 | - npu_preparation::CheckOut( | 93 | + npu_preparation::CheckOut({self}, result, self, op_infer_output_size); |
| 93 | - {self}, | ||
| 94 | - result, | ||
| 95 | - self, | ||
| 96 | - op_infer_output_size); | ||
| 97 | 94 | ||
| 98 | - if (!npu_utils::check_match(&result)) { | 95 | + if (!npu_utils::check_match(&result)) { |
| 99 | - at::Tensor contiguous_result = npu_utils::format_contiguous(result); | 96 | + at::Tensor contiguous_result = npu_utils::format_contiguous(result); |
| 100 | - upsample_bicubic2d_out_nocheck(contiguous_result, self, output_size, align_corners, scales_h, scales_w); | 97 | + upsample_bicubic2d_out_nocheck(contiguous_result, self, output_size, align_corners, scales_h, scales_w); |
| 101 | - npu_utils::format_fresh_view(result, contiguous_result); | 98 | + npu_utils::format_fresh_view(result, contiguous_result); |
| 102 | - } else { | 99 | + } else { |
| 103 | - upsample_bicubic2d_out_nocheck(result, self, output_size, align_corners, scales_h, scales_w); | 100 | + upsample_bicubic2d_out_nocheck(result, self, output_size, align_corners, scales_h, scales_w); |
| 104 | - } | 101 | + } |
| 105 | 102 | ||
| 106 | - return result; | 103 | + return result; |
| 107 | } | 104 | } |
| 108 | 105 | ||
| 109 | at::Tensor upsample_bicubic2d( | 106 | at::Tensor upsample_bicubic2d( |
| @@ -111,22 +108,22 @@ at::Tensor upsample_bicubic2d( | |||
| 111 | at::IntArrayRef output_size, | 108 | at::IntArrayRef output_size, |
| 112 | bool align_corners, | 109 | bool align_corners, |
| 113 | c10::optional<double> scales_h, | 110 | c10::optional<double> scales_h, |
| 114 | - c10::optional<double> scales_w) | 111 | + c10::optional<double> scales_w) { |
| 115 | -{ | 112 | + TORCH_CHECK(self.dim() >= 2, "The self shoud be at least 2D, but self got", self.dim(), OPS_ERROR(ErrCode::PARAM)); |
| 116 | - TORCH_CHECK(self.dim() >= 2, "The self shoud be at least 2D, but self got", self.dim(), | 113 | + TORCH_CHECK( |
| 117 | - OPS_ERROR(ErrCode::PARAM)); | 114 | + output_size.size() == 2, |
| 118 | - TORCH_CHECK(output_size.size() == 2, | 115 | + "It is expected output_size equals to 2, but got size ", |
| 119 | - "It is expected output_size equals to 2, but got size ", | 116 | + output_size.size(), |
| 120 | - output_size.size(), OPS_ERROR(ErrCode::PARAM)); | 117 | + OPS_ERROR(ErrCode::PARAM)); |
| 121 | 118 | ||
| 122 | - int64_t N = self.size(0); | 119 | + int64_t N = self.size(0); |
| 123 | - int64_t C = self.size(1); | 120 | + int64_t C = self.size(1); |
| 124 | - int64_t H = output_size[0]; | 121 | + int64_t H = output_size[0]; |
| 125 | - int64_t W = output_size[1]; | 122 | + int64_t W = output_size[1]; |
| 126 | - c10::SmallVector<int64_t, SIZE> op_infer_output_size = {N, C, H, W}; | 123 | + c10::SmallVector<int64_t, SIZE> op_infer_output_size = {N, C, H, W}; |
| 127 | - at::Tensor result = npu_preparation::apply_tensor(self, op_infer_output_size); | 124 | + at::Tensor result = npu_preparation::apply_tensor(self, op_infer_output_size); |
| 128 | - upsample_bicubic2d_out_nocheck(result, self, output_size, align_corners, scales_h, scales_w); | 125 | + upsample_bicubic2d_out_nocheck(result, self, output_size, align_corners, scales_h, scales_w); |
| 129 | 126 | ||
| 130 | - return result; | 127 | + return result; |
| 131 | } | 128 | } |
| 132 | } // namespace acl_op | 129 | } // namespace acl_op |
| @@ -30,17 +30,17 @@ at::Tensor& upsample_bilinear2d_backward_out_nocheck( | |||
| 30 | bool align_corners, | 30 | bool align_corners, |
| 31 | c10::optional<double> scales_h, | 31 | c10::optional<double> scales_h, |
| 32 | c10::optional<double> scales_w) { | 32 | c10::optional<double> scales_w) { |
| 33 | - at_npu::native::OpCommand cmd; | 33 | + at_npu::native::OpCommand cmd; |
| 34 | - at::Tensor original_image = npu_preparation::apply_tensor(grad_output, input_size); | 34 | + at::Tensor original_image = npu_preparation::apply_tensor(grad_output, input_size); |
| 35 | - bool half_pixel_centers = !align_corners; | 35 | + bool half_pixel_centers = !align_corners; |
| 36 | - cmd.Name("ResizeBilinearV2Grad") | 36 | + cmd.Name("ResizeBilinearV2Grad") |
| 37 | - .Input(grad_output, "grads") | 37 | + .Input(grad_output, "grads") |
| 38 | - .Input(original_image, "original_image") | 38 | + .Input(original_image, "original_image") |
| 39 | - .Output(grad_input, "y") | 39 | + .Output(grad_input, "y") |
| 40 | - .Attr("align_corners", align_corners) | 40 | + .Attr("align_corners", align_corners) |
| 41 | - .Attr("half_pixel_centers", half_pixel_centers) | 41 | + .Attr("half_pixel_centers", half_pixel_centers) |
| 42 | - .Run(); | 42 | + .Run(); |
| 43 | - return grad_input; | 43 | + return grad_input; |
| 44 | } | 44 | } |
| 45 | } // namespace | 45 | } // namespace |
| 46 | 46 | ||
| @@ -52,21 +52,17 @@ at::Tensor& upsample_bilinear2d_backward_out( | |||
| 52 | c10::optional<double> scales_h, | 52 | c10::optional<double> scales_h, |
| 53 | c10::optional<double> scales_w, | 53 | c10::optional<double> scales_w, |
| 54 | at::Tensor& grad_input) { | 54 | at::Tensor& grad_input) { |
| 55 | - npu_preparation::CheckOut( | 55 | + npu_preparation::CheckOut({grad_output}, grad_input, grad_output, input_size); |
| 56 | - {grad_output}, | 56 | + if (!npu_utils::check_match(&grad_input)) { |
| 57 | - grad_input, | 57 | + at::Tensor contiguous_result = npu_utils::format_contiguous(grad_input); |
| 58 | - grad_output, | 58 | + upsample_bilinear2d_backward_out_nocheck( |
| 59 | - input_size); | 59 | + contiguous_result, grad_output, output_size, input_size, align_corners, scales_h, scales_w); |
| 60 | - if (!npu_utils::check_match(&grad_input)) { | 60 | + npu_utils::format_fresh_view(grad_input, contiguous_result); |
| 61 | - at::Tensor contiguous_result = npu_utils::format_contiguous(grad_input); | 61 | + } else { |
| 62 | - upsample_bilinear2d_backward_out_nocheck( | 62 | + upsample_bilinear2d_backward_out_nocheck( |
| 63 | - contiguous_result, grad_output, output_size, input_size, align_corners, scales_h, scales_w); | 63 | + grad_input, grad_output, output_size, input_size, align_corners, scales_h, scales_w); |
| 64 | - npu_utils::format_fresh_view(grad_input, contiguous_result); | 64 | + } |
| 65 | - } else { | 65 | + return grad_input; |
| 66 | - upsample_bilinear2d_backward_out_nocheck( | ||
| 67 | - grad_input, grad_output, output_size, input_size, align_corners, scales_h, scales_w); | ||
| 68 | - } | ||
| 69 | - return grad_input; | ||
| 70 | } | 66 | } |
| 71 | 67 | ||
| 72 | at::Tensor upsample_bilinear2d_backward( | 68 | at::Tensor upsample_bilinear2d_backward( |
| @@ -76,11 +72,11 @@ at::Tensor upsample_bilinear2d_backward( | |||
| 76 | bool align_corners, | 72 | bool align_corners, |
| 77 | c10::optional<double> scales_h, | 73 | c10::optional<double> scales_h, |
| 78 | c10::optional<double> scales_w) { | 74 | c10::optional<double> scales_w) { |
| 79 | - auto op_infer_output_size = input_size; | 75 | + auto op_infer_output_size = input_size; |
| 80 | - at::Tensor grad_input = npu_preparation::apply_tensor(grad_output, op_infer_output_size); | 76 | + at::Tensor grad_input = npu_preparation::apply_tensor(grad_output, op_infer_output_size); |
| 81 | 77 | ||
| 82 | - upsample_bilinear2d_backward_out_nocheck( | 78 | + upsample_bilinear2d_backward_out_nocheck( |
| 83 | - grad_input, grad_output, output_size, input_size, align_corners, scales_h, scales_w); | 79 | + grad_input, grad_output, output_size, input_size, align_corners, scales_h, scales_w); |
| 84 | - return grad_input; | 80 | + return grad_input; |
| 85 | } | 81 | } |
| 86 | } // namespace acl_op | 82 | } // namespace acl_op |
| @@ -28,23 +28,22 @@ at::Tensor& upsample_bilinear2d_out_nocheck( | |||
| 28 | at::IntArrayRef output_size, | 28 | at::IntArrayRef output_size, |
| 29 | bool align_corners, | 29 | bool align_corners, |
| 30 | c10::optional<double> scales_h, | 30 | c10::optional<double> scales_h, |
| 31 | - c10::optional<double> scales_w) | 31 | + c10::optional<double> scales_w) { |
| 32 | -{ | 32 | + at_npu::native::OpCommand cmd; |
| 33 | - at_npu::native::OpCommand cmd; | 33 | + bool half_pixel_centers = !align_corners; |
| 34 | - bool half_pixel_centers = !align_corners; | 34 | + TORCH_CHECK( |
| 35 | - TORCH_CHECK(output_size.size() >= 2, "The dim input tensor [output_size] must be at least 2." | 35 | + output_size.size() >= 2, "The dim input tensor [output_size] must be at least 2." + OPS_ERROR(ErrCode::PARAM)); |
| 36 | - + OPS_ERROR(ErrCode::PARAM)); | 36 | + int64_t H = output_size[0]; |
| 37 | - int64_t H = output_size[0]; | 37 | + int64_t W = output_size[1]; |
| 38 | - int64_t W = output_size[1]; | 38 | + at::SmallVector<int64_t, N> attr_size = {H, W}; |
| 39 | - at::SmallVector<int64_t, N> attr_size = {H, W}; | 39 | + cmd.Name("ResizeBilinearV2") |
| 40 | - cmd.Name("ResizeBilinearV2") | 40 | + .Input(self, "x") |
| 41 | - .Input(self, "x") | 41 | + .Input(attr_size, at::kInt) |
| 42 | - .Input(attr_size, at::kInt) | 42 | + .Output(result, "y") |
| 43 | - .Output(result, "y") | 43 | + .Attr("align_corners", align_corners) |
| 44 | - .Attr("align_corners", align_corners) | 44 | + .Attr("half_pixel_centers", half_pixel_centers) |
| 45 | - .Attr("half_pixel_centers", half_pixel_centers) | 45 | + .Run(); |
| 46 | - .Run(); | 46 | + return result; |
| 47 | - return result; | ||
| 48 | } | 47 | } |
| 49 | } // namespace | 48 | } // namespace |
| 50 | 49 | ||
| @@ -54,30 +53,22 @@ at::Tensor& upsample_bilinear2d_out( | |||
| 54 | bool align_corners, | 53 | bool align_corners, |
| 55 | c10::optional<double> scales_h, | 54 | c10::optional<double> scales_h, |
| 56 | c10::optional<double> scales_w, | 55 | c10::optional<double> scales_w, |
| 57 | - at::Tensor& result) | 56 | + at::Tensor& result) { |
| 58 | -{ | 57 | + at::Tensor self_apply = self; |
| 59 | - at::Tensor self_apply = self; | 58 | + if (self_apply.scalar_type() != at::ScalarType::Float) { |
| 60 | - if (self_apply.scalar_type() != at::ScalarType::Float) { | 59 | + self_apply = at_npu::native::custom_ops::_npu_dtype_cast(self_apply, at::ScalarType::Float); |
| 61 | - self_apply = at_npu::native::custom_ops::_npu_dtype_cast(self_apply, at::ScalarType::Float); | 60 | + } |
| 62 | - } | 61 | + auto op_infer_output_size = op_infer::upsample_bilinear2d_npu_output_size(self_apply, output_size); |
| 63 | - auto op_infer_output_size = op_infer::upsample_bilinear2d_npu_output_size( | ||
| 64 | - self_apply, output_size); | ||
| 65 | 62 | ||
| 66 | - npu_preparation::CheckOut( | 63 | + npu_preparation::CheckOut({self_apply}, result, self_apply, op_infer_output_size); |
| 67 | - {self_apply}, | 64 | + if (!npu_utils::check_match(&result)) { |
| 68 | - result, | 65 | + at::Tensor contiguous_result = npu_utils::format_contiguous(result); |
| 69 | - self_apply, | 66 | + upsample_bilinear2d_out_nocheck(contiguous_result, self_apply, output_size, align_corners, scales_h, scales_w); |
| 70 | - op_infer_output_size); | 67 | + npu_utils::format_fresh_view(result, contiguous_result); |
| 71 | - if (!npu_utils::check_match(&result)) { | 68 | + } else { |
| 72 | - at::Tensor contiguous_result = npu_utils::format_contiguous(result); | 69 | + upsample_bilinear2d_out_nocheck(result, self_apply, output_size, align_corners, scales_h, scales_w); |
| 73 | - upsample_bilinear2d_out_nocheck( | 70 | + } |
| 74 | - contiguous_result, self_apply, output_size, align_corners, scales_h, scales_w); | 71 | + return result; |
| 75 | - npu_utils::format_fresh_view(result, contiguous_result); | ||
| 76 | - } else { | ||
| 77 | - upsample_bilinear2d_out_nocheck( | ||
| 78 | - result, self_apply, output_size, align_corners, scales_h, scales_w); | ||
| 79 | - } | ||
| 80 | - return result; | ||
| 81 | } | 72 | } |
| 82 | 73 | ||
| 83 | at::Tensor upsample_bilinear2d( | 74 | at::Tensor upsample_bilinear2d( |
| @@ -85,21 +76,18 @@ at::Tensor upsample_bilinear2d( | |||
| 85 | at::IntArrayRef output_size, | 76 | at::IntArrayRef output_size, |
| 86 | bool align_corners, | 77 | bool align_corners, |
| 87 | c10::optional<double> scales_h, | 78 | c10::optional<double> scales_h, |
| 88 | - c10::optional<double> scales_w) | 79 | + c10::optional<double> scales_w) { |
| 89 | -{ | 80 | + at::Tensor self_apply = self; |
| 90 | - at::Tensor self_apply = self; | 81 | + if (self_apply.scalar_type() != at::ScalarType::Float) { |
| 91 | - if (self_apply.scalar_type() != at::ScalarType::Float) { | 82 | + self_apply = at_npu::native::custom_ops::_npu_dtype_cast(self_apply, at::ScalarType::Float); |
| 92 | - self_apply = at_npu::native::custom_ops::_npu_dtype_cast(self_apply, at::ScalarType::Float); | 83 | + } |
| 93 | - } | 84 | + auto op_infer_output_size = op_infer::upsample_bilinear2d_npu_output_size(self_apply, output_size); |
| 94 | - auto op_infer_output_size = op_infer::upsample_bilinear2d_npu_output_size( | 85 | + at::Tensor result = npu_preparation::apply_tensor(self_apply, op_infer_output_size); |
| 95 | - self_apply, output_size); | ||
| 96 | - at::Tensor result = npu_preparation::apply_tensor(self_apply, op_infer_output_size); | ||
| 97 | 86 | ||
| 98 | - upsample_bilinear2d_out_nocheck( | 87 | + upsample_bilinear2d_out_nocheck(result, self_apply, output_size, align_corners, scales_h, scales_w); |
| 99 | - result, self_apply, output_size, align_corners, scales_h, scales_w); | 88 | + if (result.dtype() != self.dtype()) { |
| 100 | - if (result.dtype() != self.dtype()) { | 89 | + result = at_npu::native::custom_ops::_npu_dtype_cast(result, self.scalar_type()); |
| 101 | - result = at_npu::native::custom_ops::_npu_dtype_cast(result, self.scalar_type()); | 90 | + } |
| 102 | - } | 91 | + return result; |
| 103 | - return result; | ||
| 104 | } | 92 | } |
| 105 | } // namespace acl_op | 93 | } // namespace acl_op |
| @@ -21,82 +21,101 @@ namespace acl_op { | |||
| 21 | using npu_preparation = at_npu::native::OpPreparation; | 21 | using npu_preparation = at_npu::native::OpPreparation; |
| 22 | 22 | ||
| 23 | namespace { | 23 | namespace { |
| 24 | -inline void upsample_linear1d_backward_check(const at::Tensor &grad_output, at::IntArrayRef output_size, | 24 | +inline void upsample_linear1d_backward_check( |
| 25 | - at::IntArrayRef input_size) | 25 | + const at::Tensor& grad_output, |
| 26 | -{ | 26 | + at::IntArrayRef output_size, |
| 27 | - TORCH_CHECK(output_size.size() == 1, "It is expected output_size equals to 1, but got size ", output_size.size(), | 27 | + at::IntArrayRef input_size) { |
| 28 | - OPS_ERROR(ErrCode::PARAM)); | 28 | + TORCH_CHECK( |
| 29 | + output_size.size() == 1, | ||
| 30 | + "It is expected output_size equals to 1, but got size ", | ||
| 31 | + output_size.size(), | ||
| 32 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 29 | 33 | ||
| 30 | - TORCH_CHECK(input_size.size() == 3, "It is expected input_size equals to 3, but got size ", input_size.size(), | 34 | + TORCH_CHECK( |
| 31 | - OPS_ERROR(ErrCode::PARAM)); | 35 | + input_size.size() == 3, |
| 36 | + "It is expected input_size equals to 3, but got size ", | ||
| 37 | + input_size.size(), | ||
| 38 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 32 | 39 | ||
| 33 | - TORCH_CHECK(grad_output.dim() >= 3, "grad_output dim must larger than 3 ", grad_output.sizes(), | 40 | + TORCH_CHECK( |
| 34 | - OPS_ERROR(ErrCode::PARAM)); | 41 | + grad_output.dim() >= 3, "grad_output dim must larger than 3 ", grad_output.sizes(), OPS_ERROR(ErrCode::PARAM)); |
| 35 | 42 | ||
| 36 | - int64_t output_width = grad_output.size(2); | 43 | + int64_t output_width = grad_output.size(2); |
| 37 | - int64_t input_width = input_size[2]; | 44 | + int64_t input_width = input_size[2]; |
| 38 | 45 | ||
| 39 | - TORCH_CHECK(output_width > 0 && input_width > 0, | 46 | + TORCH_CHECK( |
| 40 | - "Input and output sizes should be greater than 0, but got input (W: ", input_width, | 47 | + output_width > 0 && input_width > 0, |
| 41 | - ") and output (W: ", output_width, ")" + OPS_ERROR(ErrCode::VALUE)); | 48 | + "Input and output sizes should be greater than 0, but got input (W: ", |
| 49 | + input_width, | ||
| 50 | + ") and output (W: ", | ||
| 51 | + output_width, | ||
| 52 | + ")" + OPS_ERROR(ErrCode::VALUE)); | ||
| 42 | } | 53 | } |
| 43 | 54 | ||
| 44 | -at::Tensor &upsample_linear1d_backward_out_nocheck(at::Tensor &result, const at::Tensor &grad_output, | 55 | +at::Tensor& upsample_linear1d_backward_out_nocheck( |
| 45 | - at::IntArrayRef input_size, bool align_corners, | 56 | + at::Tensor& result, |
| 46 | - c10::optional<double> scales) | 57 | + const at::Tensor& grad_output, |
| 47 | -{ | 58 | + at::IntArrayRef input_size, |
| 48 | - c10::SmallVector<float, N> sc = {}; | 59 | + bool align_corners, |
| 49 | - TORCH_CHECK(input_size.size() == 3 && input_size[2] != 0, "It is expected input_size equals to 3, but got size ", | 60 | + c10::optional<double> scales) { |
| 50 | - input_size.size(), OPS_ERROR(ErrCode::PARAM)); | 61 | + c10::SmallVector<float, N> sc = {}; |
| 62 | + TORCH_CHECK( | ||
| 63 | + input_size.size() == 3 && input_size[2] != 0, | ||
| 64 | + "It is expected input_size equals to 3, but got size ", | ||
| 65 | + input_size.size(), | ||
| 66 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 51 | 67 | ||
| 52 | - TORCH_CHECK(grad_output.dim() >= 3, "grad_output dim must larger than 3 ", grad_output.sizes(), | 68 | + TORCH_CHECK( |
| 53 | - OPS_ERROR(ErrCode::PARAM)); | 69 | + grad_output.dim() >= 3, "grad_output dim must larger than 3 ", grad_output.sizes(), OPS_ERROR(ErrCode::PARAM)); |
| 54 | 70 | ||
| 55 | - if (scales.has_value()) { | 71 | + if (scales.has_value()) { |
| 56 | - sc.push_back(scales.value()); | 72 | + sc.push_back(scales.value()); |
| 57 | - } else { | 73 | + } else { |
| 58 | - float temp = float(grad_output.size(3)) / float(input_size[2]); | 74 | + float temp = float(grad_output.size(3)) / float(input_size[2]); |
| 59 | - sc.push_back(temp); | 75 | + sc.push_back(temp); |
| 60 | - } | 76 | + } |
| 61 | - string coordinate_transformation_mode = align_corners ? "align_corners" : "half_pixel"; | 77 | + string coordinate_transformation_mode = align_corners ? "align_corners" : "half_pixel"; |
| 62 | 78 | ||
| 63 | - at_npu::native::OpCommand cmd; | 79 | + at_npu::native::OpCommand cmd; |
| 64 | - cmd.Name("ResizeGradD") | 80 | + cmd.Name("ResizeGradD") |
| 65 | - .Input(grad_output, "grads") | 81 | + .Input(grad_output, "grads") |
| 66 | - .Output(result, "y") | 82 | + .Output(result, "y") |
| 67 | - .Attr("original_size", input_size) | 83 | + .Attr("original_size", input_size) |
| 68 | - .Attr("scales", sc) | 84 | + .Attr("scales", sc) |
| 69 | - .Attr("coordinate_transformation_mode", coordinate_transformation_mode) | 85 | + .Attr("coordinate_transformation_mode", coordinate_transformation_mode) |
| 70 | - .Attr("mode", static_cast<string>("linear")) | 86 | + .Attr("mode", static_cast<string>("linear")) |
| 71 | - .Run(); | 87 | + .Run(); |
| 72 | - return result; | 88 | + return result; |
| 73 | } | 89 | } |
| 74 | } // namespace | 90 | } // namespace |
| 75 | 91 | ||
| 76 | -at::Tensor upsample_linear1d_backward(const at::Tensor &grad_output, at::IntArrayRef output_size, | 92 | +at::Tensor upsample_linear1d_backward( |
| 77 | - at::IntArrayRef input_size, bool align_corners, c10::optional<double> scales) | 93 | + const at::Tensor& grad_output, |
| 78 | -{ | 94 | + at::IntArrayRef output_size, |
| 79 | - upsample_linear1d_backward_check(grad_output, output_size, input_size); | 95 | + at::IntArrayRef input_size, |
| 80 | - at::Tensor grad_output_cp = grad_output; | 96 | + bool align_corners, |
| 81 | - if (grad_output.scalar_type() != at::ScalarType::Float) { | 97 | + c10::optional<double> scales) { |
| 82 | - grad_output_cp = at_npu::native::custom_ops::_npu_dtype_cast(grad_output_cp, at::ScalarType::Float); | 98 | + upsample_linear1d_backward_check(grad_output, output_size, input_size); |
| 83 | - } | 99 | + at::Tensor grad_output_cp = grad_output; |
| 84 | - int64_t N = grad_output_cp.size(0); | 100 | + if (grad_output.scalar_type() != at::ScalarType::Float) { |
| 85 | - int64_t C = grad_output_cp.size(1); | 101 | + grad_output_cp = at_npu::native::custom_ops::_npu_dtype_cast(grad_output_cp, at::ScalarType::Float); |
| 86 | - int64_t W = input_size[2]; | 102 | + } |
| 87 | - c10::SmallVector<int64_t, SIZE> output_sizes = {N, C, W}; | 103 | + int64_t N = grad_output_cp.size(0); |
| 104 | + int64_t C = grad_output_cp.size(1); | ||
| 105 | + int64_t W = input_size[2]; | ||
| 106 | + c10::SmallVector<int64_t, SIZE> output_sizes = {N, C, W}; | ||
| 88 | 107 | ||
| 89 | - // Since only NCHW format input is currently supported, first convert the | 108 | + // Since only NCHW format input is currently supported, first convert the |
| 90 | - // input grad_output (3 dimensions) to 4 dimensions as the input of npu | 109 | + // input grad_output (3 dimensions) to 4 dimensions as the input of npu |
| 91 | - auto grad_output_4dim = grad_output_cp.unsqueeze(2); | 110 | + auto grad_output_4dim = grad_output_cp.unsqueeze(2); |
| 92 | 111 | ||
| 93 | - at::Tensor result = npu_preparation::apply_tensor(grad_output_cp, output_sizes); | 112 | + at::Tensor result = npu_preparation::apply_tensor(grad_output_cp, output_sizes); |
| 94 | - upsample_linear1d_backward_out_nocheck(result, grad_output_4dim, input_size, align_corners, scales); | 113 | + upsample_linear1d_backward_out_nocheck(result, grad_output_4dim, input_size, align_corners, scales); |
| 95 | 114 | ||
| 96 | - if (result.dtype() != grad_output.dtype()) { | 115 | + if (result.dtype() != grad_output.dtype()) { |
| 97 | - result = result.to(grad_output.dtype()); | 116 | + result = result.to(grad_output.dtype()); |
| 98 | - } | 117 | + } |
| 99 | 118 | ||
| 100 | - return result; | 119 | + return result; |
| 101 | } | 120 | } |
| 102 | } // namespace acl_op | 121 | } // namespace acl_op |
| @@ -22,30 +22,29 @@ using npu_preparation = at_npu::native::OpPreparation; | |||
| 22 | using npu_utils = at_npu::native::NpuUtils; | 22 | using npu_utils = at_npu::native::NpuUtils; |
| 23 | 23 | ||
| 24 | namespace { | 24 | namespace { |
| 25 | -inline void upsample_linear1d_check( | 25 | +inline void upsample_linear1d_check(const at::Tensor& self, at::IntArrayRef output_size) { |
| 26 | - const at::Tensor& self, | 26 | + TORCH_CHECK( |
| 27 | - at::IntArrayRef output_size) | 27 | + output_size.size() == 1, |
| 28 | -{ | 28 | + "It is expected output_size equals to 1, but got size ", |
| 29 | - TORCH_CHECK( | 29 | + output_size.size(), |
| 30 | - output_size.size() == 1, | 30 | + OPS_ERROR(ErrCode::PARAM)); |
| 31 | - "It is expected output_size equals to 1, but got size ", | ||
| 32 | - output_size.size(), OPS_ERROR(ErrCode::PARAM)); | ||
| 33 | 31 | ||
| 34 | - TORCH_CHECK( | 32 | + TORCH_CHECK( |
| 35 | - (self.size(1) != 0 && self.size(2) != 0) && self.dim() == 3, | 33 | + (self.size(1) != 0 && self.size(2) != 0) && self.dim() == 3, |
| 36 | - "Non-empty 3D data tensor expected but got a tensor with sizes ", | 34 | + "Non-empty 3D data tensor expected but got a tensor with sizes ", |
| 37 | - self.sizes(), OPS_ERROR(ErrCode::PARAM)); | 35 | + self.sizes(), |
| 36 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 38 | 37 | ||
| 39 | - int64_t input_width = self.size(2); | 38 | + int64_t input_width = self.size(2); |
| 40 | - int64_t output_width = output_size[0]; | 39 | + int64_t output_width = output_size[0]; |
| 41 | 40 | ||
| 42 | - TORCH_CHECK( | 41 | + TORCH_CHECK( |
| 43 | - input_width > 0 && output_width > 0, | 42 | + input_width > 0 && output_width > 0, |
| 44 | - "Input and output sizes should be greater than 0, but got input (W: ", | 43 | + "Input and output sizes should be greater than 0, but got input (W: ", |
| 45 | - input_width, | 44 | + input_width, |
| 46 | - ") and output (W: ", | 45 | + ") and output (W: ", |
| 47 | - output_width, | 46 | + output_width, |
| 48 | - ")" + OPS_ERROR(ErrCode::VALUE)); | 47 | + ")" + OPS_ERROR(ErrCode::VALUE)); |
| 49 | } | 48 | } |
| 50 | 49 | ||
| 51 | at::Tensor& upsample_linear1d_out_nocheck( | 50 | at::Tensor& upsample_linear1d_out_nocheck( |
| @@ -53,35 +52,34 @@ at::Tensor& upsample_linear1d_out_nocheck( | |||
| 53 | const at::Tensor& self, | 52 | const at::Tensor& self, |
| 54 | at::IntArrayRef output_size, | 53 | at::IntArrayRef output_size, |
| 55 | bool align_corners, | 54 | bool align_corners, |
| 56 | - c10::optional<double> scales) | 55 | + c10::optional<double> scales) { |
| 57 | -{ | 56 | + upsample_linear1d_check(self, output_size); |
| 58 | - upsample_linear1d_check(self, output_size); | 57 | + // Since only NCHW format input is currently supported, first convert the |
| 59 | - // Since only NCHW format input is currently supported, first convert the | 58 | + // input self (3 dimensions) to 4 dimensions as the input of npu |
| 60 | - // input self (3 dimensions) to 4 dimensions as the input of npu | 59 | + at::Tensor selfcp = self.unsqueeze(2); |
| 61 | - at::Tensor selfcp = self.unsqueeze(2); | 60 | + TORCH_CHECK(selfcp.size(3) != 0, "selfcp.size(3) == 0." + OPS_ERROR(ErrCode::PARAM)); |
| 62 | - TORCH_CHECK(selfcp.size(3) != 0, "selfcp.size(3) == 0." + OPS_ERROR(ErrCode::PARAM)); | 61 | + // to calculate the value of scale |
| 63 | - // to calculate the value of scale | 62 | + c10::SmallVector<float, N> sc = {}; |
| 64 | - c10::SmallVector<float, N> sc = {}; | 63 | + if (scales.has_value()) { |
| 65 | - if (scales.has_value()) { | 64 | + sc.push_back(scales.value()); |
| 66 | - sc.push_back(scales.value()); | 65 | + } else { |
| 67 | - } else { | 66 | + float temp = float(output_size[0]) / float(selfcp.size(3)); |
| 68 | - float temp = float(output_size[0]) / float(selfcp.size(3)); | 67 | + sc.push_back(temp); |
| 69 | - sc.push_back(temp); | 68 | + } |
| 70 | - } | 69 | + string coordinate_transformation_mode = align_corners ? "align_corners" : "half_pixel"; |
| 71 | - string coordinate_transformation_mode = align_corners ? "align_corners" : "half_pixel"; | 70 | + string mode = "linear"; |
| 72 | - string mode = "linear"; | ||
| 73 | 71 | ||
| 74 | - at_npu::native::OpCommand cmd; | 72 | + at_npu::native::OpCommand cmd; |
| 75 | - cmd.Name("ResizeD") | 73 | + cmd.Name("ResizeD") |
| 76 | - .Input(selfcp, "X") | 74 | + .Input(selfcp, "X") |
| 77 | - .Output(result, "y") | 75 | + .Output(result, "y") |
| 78 | - .Attr("sizes", output_size) | 76 | + .Attr("sizes", output_size) |
| 79 | - .Attr("coordinate_transformation_mode", coordinate_transformation_mode) | 77 | + .Attr("coordinate_transformation_mode", coordinate_transformation_mode) |
| 80 | - .Attr("mode", mode) | 78 | + .Attr("mode", mode) |
| 81 | - .Attr("scales", sc) | 79 | + .Attr("scales", sc) |
| 82 | - .Run(); | 80 | + .Run(); |
| 83 | 81 | ||
| 84 | - return result; | 82 | + return result; |
| 85 | } | 83 | } |
| 86 | } // namespace | 84 | } // namespace |
| 87 | 85 | ||
| @@ -90,40 +88,32 @@ at::Tensor& upsample_linear1d_out( | |||
| 90 | at::IntArrayRef output_size, | 88 | at::IntArrayRef output_size, |
| 91 | bool align_corners, | 89 | bool align_corners, |
| 92 | c10::optional<double> scales, | 90 | c10::optional<double> scales, |
| 93 | - at::Tensor& out) | 91 | + at::Tensor& out) { |
| 94 | -{ | 92 | + auto output_sizes = op_infer::upsample_linear1d_npu_output_size(self, output_size); |
| 95 | - auto output_sizes = op_infer::upsample_linear1d_npu_output_size( | ||
| 96 | - self, output_size); | ||
| 97 | 93 | ||
| 98 | - npu_preparation::CheckOut( | 94 | + npu_preparation::CheckOut({self}, out, self, output_sizes); |
| 99 | - {self}, | ||
| 100 | - out, | ||
| 101 | - self, | ||
| 102 | - output_sizes); | ||
| 103 | 95 | ||
| 104 | - if (!npu_utils::check_match(&out)) { | 96 | + if (!npu_utils::check_match(&out)) { |
| 105 | - at::Tensor contiguous_result = npu_utils::format_contiguous(out); | 97 | + at::Tensor contiguous_result = npu_utils::format_contiguous(out); |
| 106 | - upsample_linear1d_out_nocheck(contiguous_result, self, output_size, align_corners, scales); | 98 | + upsample_linear1d_out_nocheck(contiguous_result, self, output_size, align_corners, scales); |
| 107 | - npu_utils::format_fresh_view(out, contiguous_result); | 99 | + npu_utils::format_fresh_view(out, contiguous_result); |
| 108 | - } else { | 100 | + } else { |
| 109 | - upsample_linear1d_out_nocheck(out, self, output_size, align_corners, scales); | 101 | + upsample_linear1d_out_nocheck(out, self, output_size, align_corners, scales); |
| 110 | - } | 102 | + } |
| 111 | 103 | ||
| 112 | - return out; | 104 | + return out; |
| 113 | } | 105 | } |
| 114 | 106 | ||
| 115 | at::Tensor upsample_linear1d( | 107 | at::Tensor upsample_linear1d( |
| 116 | const at::Tensor& self, | 108 | const at::Tensor& self, |
| 117 | at::IntArrayRef output_size, | 109 | at::IntArrayRef output_size, |
| 118 | bool align_corners, | 110 | bool align_corners, |
| 119 | - c10::optional<double> scales) | 111 | + c10::optional<double> scales) { |
| 120 | -{ | 112 | + auto output_sizes = op_infer::upsample_linear1d_npu_output_size(self, output_size); |
| 121 | - auto output_sizes = op_infer::upsample_linear1d_npu_output_size( | 113 | + at::Tensor result = npu_preparation::apply_tensor(self, output_sizes); |
| 122 | - self, output_size); | ||
| 123 | - at::Tensor result = npu_preparation::apply_tensor(self, output_sizes); | ||
| 124 | 114 | ||
| 125 | - upsample_linear1d_out_nocheck(result, self, output_size, align_corners, scales); | 115 | + upsample_linear1d_out_nocheck(result, self, output_size, align_corners, scales); |
| 126 | 116 | ||
| 127 | - return result; | 117 | + return result; |
| 128 | } | 118 | } |
| 129 | } // namespace acl_op | 119 | } // namespace acl_op |
| @@ -22,19 +22,19 @@ using npu_preparation = at_npu::native::OpPreparation; | |||
| 22 | using npu_utils = at_npu::native::NpuUtils; | 22 | using npu_utils = at_npu::native::NpuUtils; |
| 23 | 23 | ||
| 24 | namespace { | 24 | namespace { |
| 25 | -c10::SmallVector<int64_t, SIZE> upsample_nearest1d_backward_infer_size(at::IntArrayRef input_size) | 25 | +c10::SmallVector<int64_t, SIZE> upsample_nearest1d_backward_infer_size(at::IntArrayRef input_size) { |
| 26 | -{ | 26 | + TORCH_CHECK( |
| 27 | - TORCH_CHECK( | 27 | + input_size.size() == 3, |
| 28 | - input_size.size() == 3, | 28 | + "It is expected input_size equals to 3, but got size ", |
| 29 | - "It is expected input_size equals to 3, but got size ", | 29 | + input_size.size(), |
| 30 | - input_size.size(), OPS_ERROR(ErrCode::PARAM)); | 30 | + OPS_ERROR(ErrCode::PARAM)); |
| 31 | 31 | ||
| 32 | - c10::SmallVector<int64_t, SIZE> output_size; | 32 | + c10::SmallVector<int64_t, SIZE> output_size; |
| 33 | - int64_t N = input_size[0]; | 33 | + int64_t N = input_size[0]; |
| 34 | - int64_t C = input_size[1]; | 34 | + int64_t C = input_size[1]; |
| 35 | - int64_t W = input_size[2]; | 35 | + int64_t W = input_size[2]; |
| 36 | - output_size = {N, C, 1, W}; | 36 | + output_size = {N, C, 1, W}; |
| 37 | - return output_size; | 37 | + return output_size; |
| 38 | } | 38 | } |
| 39 | 39 | ||
| 40 | at::Tensor& upsample_nearest1d_backward_out_nocheck( | 40 | at::Tensor& upsample_nearest1d_backward_out_nocheck( |
| @@ -42,41 +42,40 @@ at::Tensor& upsample_nearest1d_backward_out_nocheck( | |||
| 42 | const at::Tensor& grad_output, | 42 | const at::Tensor& grad_output, |
| 43 | at::IntArrayRef output_size, | 43 | at::IntArrayRef output_size, |
| 44 | at::IntArrayRef input_size, | 44 | at::IntArrayRef input_size, |
| 45 | - c10::optional<double> scales) | 45 | + c10::optional<double> scales) { |
| 46 | -{ | 46 | + at::Tensor grad_cp = grad_output.unsqueeze(2); |
| 47 | - at::Tensor grad_cp = grad_output.unsqueeze(2); | 47 | + at_npu::native::OpCommand cmd; |
| 48 | - at_npu::native::OpCommand cmd; | 48 | + if (grad_output.scalar_type() == at::kFloat || grad_output.scalar_type() == at::kHalf) { |
| 49 | - if (grad_output.scalar_type() == at::kFloat || grad_output.scalar_type() == at::kHalf) { | 49 | + c10::SmallVector<int64_t, SIZE> result_size = {1, input_size[2]}; |
| 50 | - c10::SmallVector<int64_t, SIZE> result_size = {1, input_size[2]}; | 50 | + cmd.Name("ResizeNearestNeighborV2Grad") |
| 51 | - cmd.Name("ResizeNearestNeighborV2Grad") | 51 | + .Input(grad_cp) |
| 52 | - .Input(grad_cp) | 52 | + .Input(result_size, at::kInt) |
| 53 | - .Input(result_size, at::kInt) | 53 | + .Output(grad_input) |
| 54 | - .Output(grad_input) | 54 | + .Attr("align_corners", false) |
| 55 | - .Attr("align_corners", false) | 55 | + .Attr("half_pixel_centers", false) |
| 56 | - .Attr("half_pixel_centers", false) | 56 | + .Run(); |
| 57 | - .Run(); | 57 | + } else { |
| 58 | - } else { | 58 | + TORCH_CHECK( |
| 59 | - TORCH_CHECK(output_size[0] != 0, "output_size must not equals to 0, but got ", output_size[0], | 59 | + output_size[0] != 0, "output_size must not equals to 0, but got ", output_size[0], OPS_ERROR(ErrCode::PARAM)); |
| 60 | - OPS_ERROR(ErrCode::PARAM)); | 60 | + c10::SmallVector<int64_t, SIZE> origin_size = upsample_nearest1d_backward_infer_size(input_size); |
| 61 | - c10::SmallVector<int64_t, SIZE> origin_size = upsample_nearest1d_backward_infer_size(input_size); | 61 | + at::Scalar scales_cp = scales.has_value() ? scales.value() : -1; |
| 62 | - at::Scalar scales_cp = scales.has_value() ? scales.value() : -1; | 62 | + cmd.Name("ResizeGrad") |
| 63 | - cmd.Name("ResizeGrad") | 63 | + .Input(grad_cp) |
| 64 | - .Input(grad_cp) | 64 | + .Input(scales_cp, at::kFloat) |
| 65 | - .Input(scales_cp, at::kFloat) | 65 | + .Input(scales_cp, at::kFloat) |
| 66 | - .Input(scales_cp, at::kFloat) | 66 | + .Input(origin_size, at::kLong) |
| 67 | - .Input(origin_size, at::kLong) | 67 | + .Output(grad_input) |
| 68 | - .Output(grad_input) | 68 | + // Default value of Resize |
| 69 | - // Default value of Resize | 69 | + .Attr("coordinate_transformation_mode", (string) "pytorch_half_pixel") |
| 70 | - .Attr("coordinate_transformation_mode", (string)"pytorch_half_pixel") | 70 | + .Attr("cubic_coeff_a", (float)-0.75) |
| 71 | - .Attr("cubic_coeff_a", (float)-0.75) | 71 | + .Attr("exclude_outside", (int64_t)0) |
| 72 | - .Attr("exclude_outside", (int64_t)0) | 72 | + .Attr("extrapolation_value", (float)0.0) |
| 73 | - .Attr("extrapolation_value", (float)0.0) | 73 | + .Attr("mode", (string) "nearest") |
| 74 | - .Attr("mode", (string)"nearest") | 74 | + .Attr("nearest_mode", (string) "floor") |
| 75 | - .Attr("nearest_mode", (string)"floor") | 75 | + .Run(); |
| 76 | - .Run(); | 76 | + } |
| 77 | - } | 77 | + grad_input = grad_input.squeeze(2); |
| 78 | - grad_input = grad_input.squeeze(2); | 78 | + return grad_input; |
| 79 | - return grad_input; | ||
| 80 | } | 79 | } |
| 81 | } // namespace | 80 | } // namespace |
| 82 | 81 | ||
| @@ -85,37 +84,30 @@ at::Tensor& upsample_nearest1d_backward_out( | |||
| 85 | at::IntArrayRef output_size, | 84 | at::IntArrayRef output_size, |
| 86 | at::IntArrayRef input_size, | 85 | at::IntArrayRef input_size, |
| 87 | c10::optional<double> scales, | 86 | c10::optional<double> scales, |
| 88 | - at::Tensor& grad_input) | 87 | + at::Tensor& grad_input) { |
| 89 | -{ | 88 | + c10::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest1d_backward_infer_size(input_size); |
| 90 | - c10::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest1d_backward_infer_size(input_size); | 89 | + npu_preparation::CheckOut({grad_output}, grad_input, grad_output, op_infer_output_size); |
| 91 | - npu_preparation::CheckOut( | ||
| 92 | - {grad_output}, | ||
| 93 | - grad_input, | ||
| 94 | - grad_output, | ||
| 95 | - op_infer_output_size); | ||
| 96 | 90 | ||
| 97 | - if (!npu_utils::check_match(&grad_input)) { | 91 | + if (!npu_utils::check_match(&grad_input)) { |
| 98 | - at::Tensor contiguous_result = npu_utils::format_contiguous(grad_input); | 92 | + at::Tensor contiguous_result = npu_utils::format_contiguous(grad_input); |
| 99 | - upsample_nearest1d_backward_out_nocheck(contiguous_result, grad_output, output_size, input_size, scales); | 93 | + upsample_nearest1d_backward_out_nocheck(contiguous_result, grad_output, output_size, input_size, scales); |
| 100 | - npu_utils::format_fresh_view(grad_input, contiguous_result); | 94 | + npu_utils::format_fresh_view(grad_input, contiguous_result); |
| 101 | - } else { | 95 | + } else { |
| 102 | - upsample_nearest1d_backward_out_nocheck(grad_input, grad_output, output_size, input_size, scales); | 96 | + upsample_nearest1d_backward_out_nocheck(grad_input, grad_output, output_size, input_size, scales); |
| 103 | - } | 97 | + } |
| 104 | 98 | ||
| 105 | - return grad_input; | 99 | + return grad_input; |
| 106 | } | 100 | } |
| 107 | 101 | ||
| 108 | at::Tensor upsample_nearest1d_backward( | 102 | at::Tensor upsample_nearest1d_backward( |
| 109 | const at::Tensor& grad_output, | 103 | const at::Tensor& grad_output, |
| 110 | at::IntArrayRef output_size, | 104 | at::IntArrayRef output_size, |
| 111 | at::IntArrayRef input_size, | 105 | at::IntArrayRef input_size, |
| 112 | - c10::optional<double> scales) | 106 | + c10::optional<double> scales) { |
| 113 | -{ | 107 | + c10::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest1d_backward_infer_size(input_size); |
| 114 | - c10::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest1d_backward_infer_size(input_size); | 108 | + at::Tensor grad_input = npu_preparation::apply_tensor(grad_output, op_infer_output_size); |
| 115 | - at::Tensor grad_input = npu_preparation::apply_tensor(grad_output, op_infer_output_size); | ||
| 116 | 109 | ||
| 117 | - upsample_nearest1d_backward_out_nocheck( | 110 | + upsample_nearest1d_backward_out_nocheck(grad_input, grad_output, output_size, input_size, scales); |
| 118 | - grad_input, grad_output, output_size, input_size, scales); | 111 | + return grad_input; |
| 119 | - return grad_input; | ||
| 120 | } | 112 | } |
| 121 | } // namespace acl_op | 113 | } // namespace acl_op |
| @@ -22,77 +22,84 @@ using npu_preparation = at_npu::native::OpPreparation; | |||
| 22 | using npu_utils = at_npu::native::NpuUtils; | 22 | using npu_utils = at_npu::native::NpuUtils; |
| 23 | 23 | ||
| 24 | namespace { | 24 | namespace { |
| 25 | -c10::SmallVector<int64_t, SIZE> upsample_nearest1d_infer_size(const at::Tensor &input, at::IntArrayRef output_size) | 25 | +c10::SmallVector<int64_t, SIZE> upsample_nearest1d_infer_size(const at::Tensor& input, at::IntArrayRef output_size) { |
| 26 | -{ | 26 | + TORCH_CHECK( |
| 27 | - TORCH_CHECK(input.dim() == 3 && (input.size(1) != 0 && input.size(2) != 0), | 27 | + input.dim() == 3 && (input.size(1) != 0 && input.size(2) != 0), |
| 28 | - "Non-empty 3D data tensor expected but got a tensor with sizes ", input.sizes(), | 28 | + "Non-empty 3D data tensor expected but got a tensor with sizes ", |
| 29 | - OPS_ERROR(ErrCode::PARAM)); | 29 | + input.sizes(), |
| 30 | - TORCH_CHECK(output_size.size() == 1, "The length of output_size should be equal to 1, but got ", | 30 | + OPS_ERROR(ErrCode::PARAM)); |
| 31 | - output_size.size(), OPS_ERROR(ErrCode::PARAM)); | 31 | + TORCH_CHECK( |
| 32 | + output_size.size() == 1, | ||
| 33 | + "The length of output_size should be equal to 1, but got ", | ||
| 34 | + output_size.size(), | ||
| 35 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 32 | 36 | ||
| 33 | - c10::SmallVector<int64_t, SIZE> output_sizes; | 37 | + c10::SmallVector<int64_t, SIZE> output_sizes; |
| 34 | - int64_t N = input.size(0); | 38 | + int64_t N = input.size(0); |
| 35 | - int64_t C = input.size(1); | 39 | + int64_t C = input.size(1); |
| 36 | - int64_t W = output_size[0]; | 40 | + int64_t W = output_size[0]; |
| 37 | - output_sizes = {N, C, 1, W}; | 41 | + output_sizes = {N, C, 1, W}; |
| 38 | - return output_sizes; | 42 | + return output_sizes; |
| 39 | } | 43 | } |
| 40 | 44 | ||
| 41 | -at::Tensor &upsample_nearest1d_out_nocheck(at::Tensor &out, const at::Tensor &self, at::IntArrayRef output_size, | 45 | +at::Tensor& upsample_nearest1d_out_nocheck( |
| 42 | - c10::optional<double> scales) | 46 | + at::Tensor& out, |
| 43 | -{ | 47 | + const at::Tensor& self, |
| 44 | - at::Tensor self_cp = self.unsqueeze(2); | 48 | + at::IntArrayRef output_size, |
| 45 | - at_npu::native::OpCommand cmd; | 49 | + c10::optional<double> scales) { |
| 46 | - if (self.scalar_type() == at::kFloat || self.scalar_type() == at::kHalf) { | 50 | + at::Tensor self_cp = self.unsqueeze(2); |
| 47 | - c10::SmallVector<int64_t, SIZE> result_size = {1, output_size[0]}; | 51 | + at_npu::native::OpCommand cmd; |
| 48 | - cmd.Name("ResizeNearestNeighborV2") | 52 | + if (self.scalar_type() == at::kFloat || self.scalar_type() == at::kHalf) { |
| 49 | - .Input(self_cp) | 53 | + c10::SmallVector<int64_t, SIZE> result_size = {1, output_size[0]}; |
| 50 | - .Input(result_size, at::kInt) | 54 | + cmd.Name("ResizeNearestNeighborV2") |
| 51 | - .Output(out) | 55 | + .Input(self_cp) |
| 52 | - .Attr("align_corners", false) | 56 | + .Input(result_size, at::kInt) |
| 53 | - .Attr("half_pixel_centers", false) | 57 | + .Output(out) |
| 54 | - .Run(); | 58 | + .Attr("align_corners", false) |
| 55 | - } else { | 59 | + .Attr("half_pixel_centers", false) |
| 56 | - cmd.Name("Resize") | 60 | + .Run(); |
| 57 | - .Input(self_cp) | 61 | + } else { |
| 58 | - .Input(output_size, at::kFloat) | 62 | + cmd.Name("Resize") |
| 59 | - .Input(output_size, at::kFloat) | 63 | + .Input(self_cp) |
| 60 | - .Input(out.sizes(), at::kLong) | 64 | + .Input(output_size, at::kFloat) |
| 61 | - .Output(out) | 65 | + .Input(output_size, at::kFloat) |
| 62 | - .Attr("mode", static_cast<string>("nearest")) | 66 | + .Input(out.sizes(), at::kLong) |
| 63 | - .Attr("nearest_mode", static_cast<string>("floor")) | 67 | + .Output(out) |
| 64 | - .Attr("coordinate_transformation_mode", (string) "pytorch_half_pixel") | 68 | + .Attr("mode", static_cast<string>("nearest")) |
| 65 | - .Run(); | 69 | + .Attr("nearest_mode", static_cast<string>("floor")) |
| 66 | - } | 70 | + .Attr("coordinate_transformation_mode", (string) "pytorch_half_pixel") |
| 67 | - out = out.squeeze(2); | 71 | + .Run(); |
| 68 | - return out; | 72 | + } |
| 73 | + out = out.squeeze(2); | ||
| 74 | + return out; | ||
| 69 | } | 75 | } |
| 70 | } // namespace | 76 | } // namespace |
| 71 | 77 | ||
| 72 | -at::Tensor &upsample_nearest1d_out(const at::Tensor &self, at::IntArrayRef output_size, c10::optional<double> scales, | 78 | +at::Tensor& upsample_nearest1d_out( |
| 73 | - at::Tensor &out) | 79 | + const at::Tensor& self, |
| 74 | -{ | 80 | + at::IntArrayRef output_size, |
| 75 | - c10::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest1d_infer_size(self, output_size); | 81 | + c10::optional<double> scales, |
| 82 | + at::Tensor& out) { | ||
| 83 | + c10::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest1d_infer_size(self, output_size); | ||
| 76 | 84 | ||
| 77 | - npu_preparation::CheckOut({self}, out, self, op_infer_output_size); | 85 | + npu_preparation::CheckOut({self}, out, self, op_infer_output_size); |
| 78 | 86 | ||
| 79 | - if (!npu_utils::check_match(&out)) { | 87 | + if (!npu_utils::check_match(&out)) { |
| 80 | - at::Tensor contiguous_result = npu_utils::format_contiguous(out); | 88 | + at::Tensor contiguous_result = npu_utils::format_contiguous(out); |
| 81 | - upsample_nearest1d_out_nocheck(contiguous_result, self, output_size, scales); | 89 | + upsample_nearest1d_out_nocheck(contiguous_result, self, output_size, scales); |
| 82 | - npu_utils::format_fresh_view(out, contiguous_result); | 90 | + npu_utils::format_fresh_view(out, contiguous_result); |
| 83 | - } else { | 91 | + } else { |
| 84 | - upsample_nearest1d_out_nocheck(out, self, output_size, scales); | 92 | + upsample_nearest1d_out_nocheck(out, self, output_size, scales); |
| 85 | - } | 93 | + } |
| 86 | 94 | ||
| 87 | - return out; | 95 | + return out; |
| 88 | } | 96 | } |
| 89 | 97 | ||
| 90 | -at::Tensor upsample_nearest1d(const at::Tensor &self, at::IntArrayRef output_size, c10::optional<double> scales) | 98 | +at::Tensor upsample_nearest1d(const at::Tensor& self, at::IntArrayRef output_size, c10::optional<double> scales) { |
| 91 | -{ | 99 | + c10::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest1d_infer_size(self, output_size); |
| 92 | - c10::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest1d_infer_size(self, output_size); | 100 | + at::Tensor out = npu_preparation::apply_tensor(self, op_infer_output_size); |
| 93 | - at::Tensor out = npu_preparation::apply_tensor(self, op_infer_output_size); | ||
| 94 | 101 | ||
| 95 | - upsample_nearest1d_out_nocheck(out, self, output_size, scales); | 102 | + upsample_nearest1d_out_nocheck(out, self, output_size, scales); |
| 96 | - return out; | 103 | + return out; |
| 97 | } | 104 | } |
| 98 | } // namespace acl_op | 105 | } // namespace acl_op |
| @@ -22,54 +22,65 @@ using npu_preparation = at_npu::native::OpPreparation; | |||
| 22 | using npu_utils = at_npu::native::NpuUtils; | 22 | using npu_utils = at_npu::native::NpuUtils; |
| 23 | 23 | ||
| 24 | namespace { | 24 | namespace { |
| 25 | -at::Tensor &upsample_nearest2d_backward_out_nocheck(at::Tensor &y, const at::Tensor &grads, at::IntArrayRef output_size, | 25 | +at::Tensor& upsample_nearest2d_backward_out_nocheck( |
| 26 | - at::IntArrayRef input_size, c10::optional<double> scales_h, | 26 | + at::Tensor& y, |
| 27 | - c10::optional<double> scales_w) | 27 | + const at::Tensor& grads, |
| 28 | -{ | 28 | + at::IntArrayRef output_size, |
| 29 | - TORCH_CHECK(input_size.size() == 4, "The length of input_size should be equal to 4, but got ", input_size.size(), | 29 | + at::IntArrayRef input_size, |
| 30 | - OPS_ERROR(ErrCode::PARAM)); | 30 | + c10::optional<double> scales_h, |
| 31 | + c10::optional<double> scales_w) { | ||
| 32 | + TORCH_CHECK( | ||
| 33 | + input_size.size() == 4, | ||
| 34 | + "The length of input_size should be equal to 4, but got ", | ||
| 35 | + input_size.size(), | ||
| 36 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 31 | 37 | ||
| 32 | - at::SmallVector<int64_t, N> output_sizes = {input_size[2], input_size[3]}; | 38 | + at::SmallVector<int64_t, N> output_sizes = {input_size[2], input_size[3]}; |
| 33 | - at_npu::native::OpCommand cmd; | 39 | + at_npu::native::OpCommand cmd; |
| 34 | - cmd.Name("ResizeNearestNeighborV2Grad") | 40 | + cmd.Name("ResizeNearestNeighborV2Grad") |
| 35 | - .Input(grads, "grads") | 41 | + .Input(grads, "grads") |
| 36 | - .Input(output_sizes, at::kInt) | 42 | + .Input(output_sizes, at::kInt) |
| 37 | - .Output(y, "y") | 43 | + .Output(y, "y") |
| 38 | - .Attr("align_corners", false) | 44 | + .Attr("align_corners", false) |
| 39 | - .Attr("half_pixel_centers", false) | 45 | + .Attr("half_pixel_centers", false) |
| 40 | - .Run(); | 46 | + .Run(); |
| 41 | 47 | ||
| 42 | - return y; | 48 | + return y; |
| 43 | } | 49 | } |
| 44 | } // namespace | 50 | } // namespace |
| 45 | 51 | ||
| 46 | -at::Tensor &upsample_nearest2d_backward_out(const at::Tensor &grads, at::IntArrayRef output_size, | 52 | +at::Tensor& upsample_nearest2d_backward_out( |
| 47 | - at::IntArrayRef input_size, c10::optional<double> scales_h, | 53 | + const at::Tensor& grads, |
| 48 | - c10::optional<double> scales_w, at::Tensor &y) | 54 | + at::IntArrayRef output_size, |
| 49 | -{ | 55 | + at::IntArrayRef input_size, |
| 50 | - npu_preparation::CheckOut({grads}, y, npu_preparation::get_tensor_npu_format(y), grads.scalar_type(), input_size); | 56 | + c10::optional<double> scales_h, |
| 57 | + c10::optional<double> scales_w, | ||
| 58 | + at::Tensor& y) { | ||
| 59 | + npu_preparation::CheckOut({grads}, y, npu_preparation::get_tensor_npu_format(y), grads.scalar_type(), input_size); | ||
| 51 | 60 | ||
| 52 | - if (!npu_utils::check_match(&y)) { | 61 | + if (!npu_utils::check_match(&y)) { |
| 53 | - at::Tensor contiguous_y = npu_utils::format_contiguous(y); | 62 | + at::Tensor contiguous_y = npu_utils::format_contiguous(y); |
| 54 | - upsample_nearest2d_backward_out_nocheck(contiguous_y, grads, output_size, input_size, scales_h, scales_w); | 63 | + upsample_nearest2d_backward_out_nocheck(contiguous_y, grads, output_size, input_size, scales_h, scales_w); |
| 55 | - npu_utils::format_fresh_view(y, contiguous_y); | 64 | + npu_utils::format_fresh_view(y, contiguous_y); |
| 56 | - } else { | 65 | + } else { |
| 57 | - upsample_nearest2d_backward_out_nocheck(y, grads, output_size, input_size, scales_h, scales_w); | 66 | + upsample_nearest2d_backward_out_nocheck(y, grads, output_size, input_size, scales_h, scales_w); |
| 58 | - } | 67 | + } |
| 59 | 68 | ||
| 60 | - return y; | 69 | + return y; |
| 61 | } | 70 | } |
| 62 | 71 | ||
| 63 | -at::Tensor upsample_nearest2d_backward(const at::Tensor &grad_output, at::IntArrayRef output_size, | 72 | +at::Tensor upsample_nearest2d_backward( |
| 64 | - at::IntArrayRef input_size, c10::optional<double> scales_h, | 73 | + const at::Tensor& grad_output, |
| 65 | - c10::optional<double> scales_w) | 74 | + at::IntArrayRef output_size, |
| 66 | -{ | 75 | + at::IntArrayRef input_size, |
| 67 | - at::Tensor grads = grad_output; | 76 | + c10::optional<double> scales_h, |
| 68 | - if (grad_output.scalar_type() != at::ScalarType::Float) { | 77 | + c10::optional<double> scales_w) { |
| 69 | - grads = at_npu::native::custom_ops::_npu_dtype_cast(grad_output, at::kFloat); | 78 | + at::Tensor grads = grad_output; |
| 70 | - } | 79 | + if (grad_output.scalar_type() != at::ScalarType::Float) { |
| 71 | - at::Tensor grad_input = npu_preparation::apply_tensor(input_size, grads.options(), grad_output); | 80 | + grads = at_npu::native::custom_ops::_npu_dtype_cast(grad_output, at::kFloat); |
| 72 | - upsample_nearest2d_backward_out_nocheck(grad_input, grads, output_size, input_size, scales_h, scales_w); | 81 | + } |
| 73 | - return grad_input; | 82 | + at::Tensor grad_input = npu_preparation::apply_tensor(input_size, grads.options(), grad_output); |
| 83 | + upsample_nearest2d_backward_out_nocheck(grad_input, grads, output_size, input_size, scales_h, scales_w); | ||
| 84 | + return grad_input; | ||
| 74 | } | 85 | } |
| 75 | } // namespace acl_op | 86 | } // namespace acl_op |
| @@ -22,64 +22,74 @@ using npu_preparation = at_npu::native::OpPreparation; | |||
| 22 | using npu_utils = at_npu::native::NpuUtils; | 22 | using npu_utils = at_npu::native::NpuUtils; |
| 23 | 23 | ||
| 24 | namespace { | 24 | namespace { |
| 25 | -at::SmallVector<int64_t, SIZE> upsample_nearest2d_infer_size(const at::Tensor &input, at::IntArrayRef output_size) | 25 | +at::SmallVector<int64_t, SIZE> upsample_nearest2d_infer_size(const at::Tensor& input, at::IntArrayRef output_size) { |
| 26 | -{ | 26 | + TORCH_CHECK(input.dim() == 4, "The input should be 4D, but got ", input.dim(), "D" + OPS_ERROR(ErrCode::PARAM)); |
| 27 | - TORCH_CHECK(input.dim() == 4, "The input should be 4D, but got ", input.dim(), "D" + OPS_ERROR(ErrCode::PARAM)); | 27 | + TORCH_CHECK( |
| 28 | - TORCH_CHECK(output_size.size() == 2, "The length of output_size should be equal to 2, but got ", | 28 | + output_size.size() == 2, |
| 29 | - output_size.size(), OPS_ERROR(ErrCode::PARAM)); | 29 | + "The length of output_size should be equal to 2, but got ", |
| 30 | + output_size.size(), | ||
| 31 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 30 | 32 | ||
| 31 | - int64_t N = input.size(0); | 33 | + int64_t N = input.size(0); |
| 32 | - int64_t C = input.size(1); | 34 | + int64_t C = input.size(1); |
| 33 | - int64_t H = output_size[0]; | 35 | + int64_t H = output_size[0]; |
| 34 | - int64_t W = output_size[1]; | 36 | + int64_t W = output_size[1]; |
| 35 | - at::SmallVector<int64_t, SIZE> output_sizes = {N, C, H, W}; | 37 | + at::SmallVector<int64_t, SIZE> output_sizes = {N, C, H, W}; |
| 36 | 38 | ||
| 37 | - return output_sizes; | 39 | + return output_sizes; |
| 38 | } | 40 | } |
| 39 | 41 | ||
| 40 | -at::Tensor &upsample_nearest2d_out_nocheck(at::Tensor &result, const at::Tensor &self, at::IntArrayRef output_size, | 42 | +at::Tensor& upsample_nearest2d_out_nocheck( |
| 41 | - c10::optional<double> scales_h, c10::optional<double> scales_w) | 43 | + at::Tensor& result, |
| 42 | -{ | 44 | + const at::Tensor& self, |
| 43 | - at::SmallVector<int64_t, N> output_size_vec = op_infer::array_to_small_vector(output_size); | 45 | + at::IntArrayRef output_size, |
| 46 | + c10::optional<double> scales_h, | ||
| 47 | + c10::optional<double> scales_w) { | ||
| 48 | + at::SmallVector<int64_t, N> output_size_vec = op_infer::array_to_small_vector(output_size); | ||
| 44 | 49 | ||
| 45 | - at_npu::native::OpCommand cmd; | 50 | + at_npu::native::OpCommand cmd; |
| 46 | - cmd.Name("ResizeNearestNeighborV2") | 51 | + cmd.Name("ResizeNearestNeighborV2") |
| 47 | - .Input(self, "x") | 52 | + .Input(self, "x") |
| 48 | - .Input(output_size_vec, at::kInt) | 53 | + .Input(output_size_vec, at::kInt) |
| 49 | - .Output(result, "y") | 54 | + .Output(result, "y") |
| 50 | - .Attr("align_corners", false) | 55 | + .Attr("align_corners", false) |
| 51 | - .Attr("half_pixel_centers", false) | 56 | + .Attr("half_pixel_centers", false) |
| 52 | - .Run(); | 57 | + .Run(); |
| 53 | 58 | ||
| 54 | - return result; | 59 | + return result; |
| 55 | } | 60 | } |
| 56 | } // namespace | 61 | } // namespace |
| 57 | 62 | ||
| 58 | -at::Tensor &upsample_nearest2d_out(const at::Tensor &self, at::IntArrayRef output_size, c10::optional<double> scales_h, | 63 | +at::Tensor& upsample_nearest2d_out( |
| 59 | - c10::optional<double> scales_w, at::Tensor &result) | 64 | + const at::Tensor& self, |
| 60 | -{ | 65 | + at::IntArrayRef output_size, |
| 61 | - at::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest2d_infer_size(self, output_size); | 66 | + c10::optional<double> scales_h, |
| 62 | - npu_preparation::CheckOut({self}, result, npu_preparation::get_tensor_npu_format(result), self.scalar_type(), | 67 | + c10::optional<double> scales_w, |
| 63 | - op_infer_output_size); | 68 | + at::Tensor& result) { |
| 69 | + at::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest2d_infer_size(self, output_size); | ||
| 70 | + npu_preparation::CheckOut( | ||
| 71 | + {self}, result, npu_preparation::get_tensor_npu_format(result), self.scalar_type(), op_infer_output_size); | ||
| 64 | 72 | ||
| 65 | - if (!npu_utils::check_match(&result)) { | 73 | + if (!npu_utils::check_match(&result)) { |
| 66 | - at::Tensor contiguous_result = npu_utils::format_contiguous(result); | 74 | + at::Tensor contiguous_result = npu_utils::format_contiguous(result); |
| 67 | - upsample_nearest2d_out_nocheck(contiguous_result, self, output_size, scales_h, scales_w); | 75 | + upsample_nearest2d_out_nocheck(contiguous_result, self, output_size, scales_h, scales_w); |
| 68 | - npu_utils::format_fresh_view(result, contiguous_result); | 76 | + npu_utils::format_fresh_view(result, contiguous_result); |
| 69 | - } else { | 77 | + } else { |
| 70 | - upsample_nearest2d_out_nocheck(result, self, output_size, scales_h, scales_w); | 78 | + upsample_nearest2d_out_nocheck(result, self, output_size, scales_h, scales_w); |
| 71 | - } | 79 | + } |
| 72 | 80 | ||
| 73 | - return result; | 81 | + return result; |
| 74 | } | 82 | } |
| 75 | 83 | ||
| 76 | -at::Tensor upsample_nearest2d(const at::Tensor &self, at::IntArrayRef output_size, c10::optional<double> scales_h, | 84 | +at::Tensor upsample_nearest2d( |
| 77 | - c10::optional<double> scales_w) | 85 | + const at::Tensor& self, |
| 78 | -{ | 86 | + at::IntArrayRef output_size, |
| 79 | - at::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest2d_infer_size(self, output_size); | 87 | + c10::optional<double> scales_h, |
| 80 | - at::Tensor result = npu_preparation::apply_tensor(self, op_infer_output_size); | 88 | + c10::optional<double> scales_w) { |
| 81 | - upsample_nearest2d_out_nocheck(result, self, output_size, scales_h, scales_w); | 89 | + at::SmallVector<int64_t, SIZE> op_infer_output_size = upsample_nearest2d_infer_size(self, output_size); |
| 90 | + at::Tensor result = npu_preparation::apply_tensor(self, op_infer_output_size); | ||
| 91 | + upsample_nearest2d_out_nocheck(result, self, output_size, scales_h, scales_w); | ||
| 82 | 92 | ||
| 83 | - return result; | 93 | + return result; |
| 84 | } | 94 | } |
| 85 | } // namespace acl_op | 95 | } // namespace acl_op |
| @@ -27,27 +27,27 @@ at::SmallVector<int64_t, SIZE> upsample_trilinear3d_backward_infer_size( | |||
| 27 | at::IntArrayRef input_size, | 27 | at::IntArrayRef input_size, |
| 28 | c10::optional<double> scales_d, | 28 | c10::optional<double> scales_d, |
| 29 | c10::optional<double> scales_h, | 29 | c10::optional<double> scales_h, |
| 30 | - c10::optional<double> scales_w) | 30 | + c10::optional<double> scales_w) { |
| 31 | -{ | 31 | + TORCH_CHECK( |
| 32 | - TORCH_CHECK( | 32 | + output_size.size() == 3, |
| 33 | - output_size.size() == 3, | 33 | + "It is expected output_size equals to 3, but got size ", |
| 34 | - "It is expected output_size equals to 3, but got size ", | 34 | + output_size.size(), |
| 35 | - output_size.size(), OPS_ERROR(ErrCode::PARAM)); | 35 | + OPS_ERROR(ErrCode::PARAM)); |
| 36 | 36 | ||
| 37 | - TORCH_CHECK( | 37 | + TORCH_CHECK( |
| 38 | - input_size.size() == 5, | 38 | + input_size.size() == 5, |
| 39 | - "It is expected input_size equals to 5, but got size ", | 39 | + "It is expected input_size equals to 5, but got size ", |
| 40 | - input_size.size(), OPS_ERROR(ErrCode::PARAM)); | 40 | + input_size.size(), |
| 41 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 41 | 42 | ||
| 42 | - int64_t nbatch = input_size[0]; | 43 | + int64_t nbatch = input_size[0]; |
| 43 | - int64_t channels = input_size[1]; | 44 | + int64_t channels = input_size[1]; |
| 44 | - int64_t input_depth = input_size[2]; | 45 | + int64_t input_depth = input_size[2]; |
| 45 | - int64_t input_height = input_size[3]; | 46 | + int64_t input_height = input_size[3]; |
| 46 | - int64_t input_width = input_size[4]; | 47 | + int64_t input_width = input_size[4]; |
| 47 | 48 | ||
| 48 | - at::SmallVector<int64_t, SIZE> output_sizes = | 49 | + at::SmallVector<int64_t, SIZE> output_sizes = {nbatch, channels, input_depth, input_height, input_width}; |
| 49 | - {nbatch, channels, input_depth, input_height, input_width}; | 50 | + return output_sizes; |
| 50 | - return output_sizes; | ||
| 51 | } | 51 | } |
| 52 | 52 | ||
| 53 | at::Tensor& upsample_trilinear3d_backward_out_nocheck( | 53 | at::Tensor& upsample_trilinear3d_backward_out_nocheck( |
| @@ -58,18 +58,17 @@ at::Tensor& upsample_trilinear3d_backward_out_nocheck( | |||
| 58 | bool align_corners, | 58 | bool align_corners, |
| 59 | c10::optional<double> scales_d, | 59 | c10::optional<double> scales_d, |
| 60 | c10::optional<double> scales_h, | 60 | c10::optional<double> scales_h, |
| 61 | - c10::optional<double> scales_w) | 61 | + c10::optional<double> scales_w) { |
| 62 | -{ | 62 | + at_npu::native::OpCommand cmd; |
| 63 | - at_npu::native::OpCommand cmd; | 63 | + cmd.Name("UpsampleTrilinear3dGrad") |
| 64 | - cmd.Name("UpsampleTrilinear3dGrad") | 64 | + .Input(grad_output) |
| 65 | - .Input(grad_output) | 65 | + .Output(out) |
| 66 | - .Output(out) | 66 | + .Attr("input_size", input_size) |
| 67 | - .Attr("input_size", input_size) | 67 | + .Attr("output_size", output_size) |
| 68 | - .Attr("output_size", output_size) | 68 | + .Attr("align_corners", align_corners) |
| 69 | - .Attr("align_corners", align_corners) | 69 | + .Run(); |
| 70 | - .Run(); | ||
| 71 | 70 | ||
| 72 | - return out; | 71 | + return out; |
| 73 | } | 72 | } |
| 74 | } // namespace | 73 | } // namespace |
| 75 | 74 | ||
| @@ -81,26 +80,21 @@ at::Tensor& upsample_trilinear3d_backward_out( | |||
| 81 | c10::optional<double> scales_d, | 80 | c10::optional<double> scales_d, |
| 82 | c10::optional<double> scales_h, | 81 | c10::optional<double> scales_h, |
| 83 | c10::optional<double> scales_w, | 82 | c10::optional<double> scales_w, |
| 84 | - at::Tensor& grad_input) | 83 | + at::Tensor& grad_input) { |
| 85 | -{ | 84 | + auto op_infer_output_size = |
| 86 | - auto op_infer_output_size = upsample_trilinear3d_backward_infer_size( | 85 | + upsample_trilinear3d_backward_infer_size(output_size, input_size, scales_d, scales_h, scales_w); |
| 87 | - output_size, input_size, scales_d, scales_h, scales_w); | 86 | + npu_preparation::CheckOut({grad_output}, grad_input, grad_output, op_infer_output_size); |
| 88 | - npu_preparation::CheckOut( | ||
| 89 | - {grad_output}, | ||
| 90 | - grad_input, | ||
| 91 | - grad_output, | ||
| 92 | - op_infer_output_size); | ||
| 93 | 87 | ||
| 94 | - if (!npu_utils::check_match(&grad_input)) { | 88 | + if (!npu_utils::check_match(&grad_input)) { |
| 95 | - auto contiguous_out = npu_utils::format_contiguous(grad_input); | 89 | + auto contiguous_out = npu_utils::format_contiguous(grad_input); |
| 96 | - upsample_trilinear3d_backward_out_nocheck( | 90 | + upsample_trilinear3d_backward_out_nocheck( |
| 97 | - grad_input, grad_output, output_size, input_size, align_corners, scales_d, scales_h, scales_w); | 91 | + grad_input, grad_output, output_size, input_size, align_corners, scales_d, scales_h, scales_w); |
| 98 | - npu_utils::format_fresh_view(grad_input, contiguous_out); | 92 | + npu_utils::format_fresh_view(grad_input, contiguous_out); |
| 99 | - } else { | 93 | + } else { |
| 100 | - upsample_trilinear3d_backward_out_nocheck( | 94 | + upsample_trilinear3d_backward_out_nocheck( |
| 101 | - grad_input, grad_output, output_size, input_size, align_corners, scales_d, scales_h, scales_w); | 95 | + grad_input, grad_output, output_size, input_size, align_corners, scales_d, scales_h, scales_w); |
| 102 | - } | 96 | + } |
| 103 | - return grad_input; | 97 | + return grad_input; |
| 104 | } | 98 | } |
| 105 | 99 | ||
| 106 | at::Tensor upsample_trilinear3d_backward( | 100 | at::Tensor upsample_trilinear3d_backward( |
| @@ -110,13 +104,12 @@ at::Tensor upsample_trilinear3d_backward( | |||
| 110 | bool align_corners, | 104 | bool align_corners, |
| 111 | c10::optional<double> scales_d, | 105 | c10::optional<double> scales_d, |
| 112 | c10::optional<double> scales_h, | 106 | c10::optional<double> scales_h, |
| 113 | - c10::optional<double> scales_w) | 107 | + c10::optional<double> scales_w) { |
| 114 | -{ | 108 | + auto op_infer_output_size = |
| 115 | - auto op_infer_output_size = upsample_trilinear3d_backward_infer_size( | 109 | + upsample_trilinear3d_backward_infer_size(output_size, input_size, scales_d, scales_h, scales_w); |
| 116 | - output_size, input_size, scales_d, scales_h, scales_w); | 110 | + at::Tensor result = npu_preparation::apply_tensor(grad_output, op_infer_output_size); |
| 117 | - at::Tensor result = npu_preparation::apply_tensor(grad_output, op_infer_output_size); | 111 | + upsample_trilinear3d_backward_out_nocheck( |
| 118 | - upsample_trilinear3d_backward_out_nocheck( | 112 | + result, grad_output, output_size, input_size, align_corners, scales_d, scales_h, scales_w); |
| 119 | - result, grad_output, output_size, input_size, align_corners, scales_d, scales_h, scales_w); | 113 | + return result; |
| 120 | - return result; | ||
| 121 | } | 114 | } |
| 122 | } // namespace acl_op | 115 | } // namespace acl_op |
| @@ -22,66 +22,80 @@ using npu_preparation = at_npu::native::OpPreparation; | |||
| 22 | using npu_utils = at_npu::native::NpuUtils; | 22 | using npu_utils = at_npu::native::NpuUtils; |
| 23 | 23 | ||
| 24 | namespace { | 24 | namespace { |
| 25 | -at::SmallVector<int64_t, SIZE> upsample_trilinear3d_infer_size(const at::Tensor &input, at::IntArrayRef output_size, | 25 | +at::SmallVector<int64_t, SIZE> upsample_trilinear3d_infer_size( |
| 26 | - c10::optional<double> scales_d, | 26 | + const at::Tensor& input, |
| 27 | - c10::optional<double> scales_h, | 27 | + at::IntArrayRef output_size, |
| 28 | - c10::optional<double> scales_w) | 28 | + c10::optional<double> scales_d, |
| 29 | -{ | 29 | + c10::optional<double> scales_h, |
| 30 | - TORCH_CHECK(input.dim() == 5, "The input should be 5D, but got ", input.dim(), "D" + OPS_ERROR(ErrCode::PARAM)); | 30 | + c10::optional<double> scales_w) { |
| 31 | - TORCH_CHECK(output_size.size() == 3, "The length of output_size should be equal to 3, but got ", | 31 | + TORCH_CHECK(input.dim() == 5, "The input should be 5D, but got ", input.dim(), "D" + OPS_ERROR(ErrCode::PARAM)); |
| 32 | - output_size.size(), OPS_ERROR(ErrCode::PARAM)); | 32 | + TORCH_CHECK( |
| 33 | + output_size.size() == 3, | ||
| 34 | + "The length of output_size should be equal to 3, but got ", | ||
| 35 | + output_size.size(), | ||
| 36 | + OPS_ERROR(ErrCode::PARAM)); | ||
| 33 | 37 | ||
| 34 | - int64_t output_depth = output_size[0]; | 38 | + int64_t output_depth = output_size[0]; |
| 35 | - int64_t output_height = output_size[1]; | 39 | + int64_t output_height = output_size[1]; |
| 36 | - int64_t output_width = output_size[2]; | 40 | + int64_t output_width = output_size[2]; |
| 37 | 41 | ||
| 38 | - int64_t nbatch = input.size(0); | 42 | + int64_t nbatch = input.size(0); |
| 39 | - int64_t channels = input.size(1); | 43 | + int64_t channels = input.size(1); |
| 40 | 44 | ||
| 41 | - at::SmallVector<int64_t, SIZE> output_sizes = {nbatch, channels, output_depth, output_height, output_width}; | 45 | + at::SmallVector<int64_t, SIZE> output_sizes = {nbatch, channels, output_depth, output_height, output_width}; |
| 42 | - return output_sizes; | 46 | + return output_sizes; |
| 43 | } | 47 | } |
| 44 | 48 | ||
| 45 | -at::Tensor &upsample_trilinear3d_out_nocheck(at::Tensor &result, const at::Tensor &input, at::IntArrayRef output_size, | 49 | +at::Tensor& upsample_trilinear3d_out_nocheck( |
| 46 | - bool align_corners, c10::optional<double> scales_d, | 50 | + at::Tensor& result, |
| 47 | - c10::optional<double> scales_h, c10::optional<double> scales_w) | 51 | + const at::Tensor& input, |
| 48 | -{ | 52 | + at::IntArrayRef output_size, |
| 49 | - at_npu::native::OpCommand cmd; | 53 | + bool align_corners, |
| 50 | - cmd.Name("UpsampleTrilinear3d") | 54 | + c10::optional<double> scales_d, |
| 51 | - .Input(input) | 55 | + c10::optional<double> scales_h, |
| 52 | - .Output(result) | 56 | + c10::optional<double> scales_w) { |
| 53 | - .Attr("output_size", output_size) | 57 | + at_npu::native::OpCommand cmd; |
| 54 | - .Attr("align_corners", align_corners) | 58 | + cmd.Name("UpsampleTrilinear3d") |
| 55 | - .Run(); | 59 | + .Input(input) |
| 56 | - return result; | 60 | + .Output(result) |
| 61 | + .Attr("output_size", output_size) | ||
| 62 | + .Attr("align_corners", align_corners) | ||
| 63 | + .Run(); | ||
| 64 | + return result; | ||
| 57 | } | 65 | } |
| 58 | } // namespace | 66 | } // namespace |
| 59 | 67 | ||
| 60 | -at::Tensor &upsample_trilinear3d_out(const at::Tensor &input, at::IntArrayRef output_size, bool align_corners, | 68 | +at::Tensor& upsample_trilinear3d_out( |
| 61 | - c10::optional<double> scales_d, c10::optional<double> scales_h, | 69 | + const at::Tensor& input, |
| 62 | - c10::optional<double> scales_w, at::Tensor &result) | 70 | + at::IntArrayRef output_size, |
| 63 | -{ | 71 | + bool align_corners, |
| 64 | - auto op_infer_output_size = upsample_trilinear3d_infer_size(input, output_size, scales_d, scales_h, scales_w); | 72 | + c10::optional<double> scales_d, |
| 65 | - npu_preparation::CheckOut({input}, result, input, op_infer_output_size); | 73 | + c10::optional<double> scales_h, |
| 74 | + c10::optional<double> scales_w, | ||
| 75 | + at::Tensor& result) { | ||
| 76 | + auto op_infer_output_size = upsample_trilinear3d_infer_size(input, output_size, scales_d, scales_h, scales_w); | ||
| 77 | + npu_preparation::CheckOut({input}, result, input, op_infer_output_size); | ||
| 66 | 78 | ||
| 67 | - if (!npu_utils::check_match(&result)) { | 79 | + if (!npu_utils::check_match(&result)) { |
| 68 | - auto contiguous_out = npu_utils::format_contiguous(result); | 80 | + auto contiguous_out = npu_utils::format_contiguous(result); |
| 69 | - upsample_trilinear3d_out_nocheck(contiguous_out, input, output_size, align_corners, scales_d, scales_h, | 81 | + upsample_trilinear3d_out_nocheck(contiguous_out, input, output_size, align_corners, scales_d, scales_h, scales_w); |
| 70 | - scales_w); | 82 | + npu_utils::format_fresh_view(result, contiguous_out); |
| 71 | - npu_utils::format_fresh_view(result, contiguous_out); | 83 | + } else { |
| 72 | - } else { | 84 | + upsample_trilinear3d_out_nocheck(result, input, output_size, align_corners, scales_d, scales_h, scales_w); |
| 73 | - upsample_trilinear3d_out_nocheck(result, input, output_size, align_corners, scales_d, scales_h, scales_w); | 85 | + } |
| 74 | - } | 86 | + return result; |
| 75 | - return result; | ||
| 76 | } | 87 | } |
| 77 | 88 | ||
| 78 | -at::Tensor upsample_trilinear3d(const at::Tensor &input, at::IntArrayRef output_size, bool align_corners, | 89 | +at::Tensor upsample_trilinear3d( |
| 79 | - c10::optional<double> scales_d, c10::optional<double> scales_h, | 90 | + const at::Tensor& input, |
| 80 | - c10::optional<double> scales_w) | 91 | + at::IntArrayRef output_size, |
| 81 | -{ | 92 | + bool align_corners, |
| 82 | - auto op_infer_output_size = upsample_trilinear3d_infer_size(input, output_size, scales_d, scales_h, scales_w); | 93 | + c10::optional<double> scales_d, |
| 83 | - at::Tensor result = npu_preparation::apply_tensor(input, op_infer_output_size); | 94 | + c10::optional<double> scales_h, |
| 84 | - upsample_trilinear3d_out_nocheck(result, input, output_size, align_corners, scales_d, scales_h, scales_w); | 95 | + c10::optional<double> scales_w) { |
| 85 | - return result; | 96 | + auto op_infer_output_size = upsample_trilinear3d_infer_size(input, output_size, scales_d, scales_h, scales_w); |
| 97 | + at::Tensor result = npu_preparation::apply_tensor(input, op_infer_output_size); | ||
| 98 | + upsample_trilinear3d_out_nocheck(result, input, output_size, align_corners, scales_d, scales_h, scales_w); | ||
| 99 | + return result; | ||
| 86 | } | 100 | } |
| 87 | } // namespace acl_op | 101 | } // namespace acl_op |
| @@ -21,44 +21,33 @@ using npu_preparation = at_npu::native::OpPreparation; | |||
| 21 | using npu_utils = at_npu::native::NpuUtils; | 21 | using npu_utils = at_npu::native::NpuUtils; |
| 22 | 22 | ||
| 23 | namespace { | 23 | namespace { |
| 24 | -at::Tensor& vdot_out_npu_nocheck(at::Tensor& result, const at::Tensor& self, const at::Tensor& other) | 24 | +at::Tensor& vdot_out_npu_nocheck(at::Tensor& result, const at::Tensor& self, const at::Tensor& other) { |
| 25 | -{ | 25 | + at_npu::native::OpCommand cmd; |
| 26 | - at_npu::native::OpCommand cmd; | 26 | + cmd.Name("Dot").Input(self).Input(other).Output(result).Run(); |
| 27 | - cmd.Name("Dot") | ||
| 28 | - .Input(self) | ||
| 29 | - .Input(other) | ||
| 30 | - .Output(result) | ||
| 31 | - .Run(); | ||
| 32 | 27 | ||
| 33 | - return result; | 28 | + return result; |
| 34 | } | 29 | } |
| 35 | } // namespace | 30 | } // namespace |
| 36 | 31 | ||
| 37 | -at::Tensor& vdot_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& out) | 32 | +at::Tensor& vdot_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& out) { |
| 38 | -{ | 33 | + c10::SmallVector<int64_t, N> output_size = {}; |
| 39 | - c10::SmallVector<int64_t, N> output_size = {}; | 34 | + npu_preparation::CheckOut({self, other}, out, self, output_size); |
| 40 | - npu_preparation::CheckOut( | ||
| 41 | - {self, other}, | ||
| 42 | - out, | ||
| 43 | - self, | ||
| 44 | - output_size); | ||
| 45 | 35 | ||
| 46 | - if (!npu_utils::check_match(&out)) { | 36 | + if (!npu_utils::check_match(&out)) { |
| 47 | - at::Tensor contiguous_out = npu_utils::format_contiguous(out); | 37 | + at::Tensor contiguous_out = npu_utils::format_contiguous(out); |
| 48 | - vdot_out_npu_nocheck(contiguous_out, self, other); | 38 | + vdot_out_npu_nocheck(contiguous_out, self, other); |
| 49 | - npu_utils::format_fresh_view(out, contiguous_out); | 39 | + npu_utils::format_fresh_view(out, contiguous_out); |
| 50 | - } else { | 40 | + } else { |
| 51 | - vdot_out_npu_nocheck(out, self, other); | 41 | + vdot_out_npu_nocheck(out, self, other); |
| 52 | - } | 42 | + } |
| 53 | 43 | ||
| 54 | - return out; | 44 | + return out; |
| 55 | } | 45 | } |
| 56 | 46 | ||
| 57 | -at::Tensor vdot(const at::Tensor& self, const at::Tensor& other) | 47 | +at::Tensor vdot(const at::Tensor& self, const at::Tensor& other) { |
| 58 | -{ | 48 | + c10::SmallVector<int64_t, N> output_size = {}; |
| 59 | - c10::SmallVector<int64_t, N> output_size = {}; | 49 | + at::Tensor result = npu_preparation::apply_tensor(self, output_size); |
| 60 | - at::Tensor result = npu_preparation::apply_tensor(self, output_size); | 50 | + vdot_out_npu_nocheck(result, self, other); |
| 61 | - vdot_out_npu_nocheck(result, self, other); | 51 | + return result; |
| 62 | - return result; | ||
| 63 | } | 52 | } |
| 64 | -} // op_plugin | 53 | +} // namespace acl_op |
| @@ -24,47 +24,42 @@ at::Tensor& var_out( | |||
| 24 | at::OptionalIntArrayRef dim, | 24 | at::OptionalIntArrayRef dim, |
| 25 | const c10::optional<c10::Scalar>& correction, | 25 | const c10::optional<c10::Scalar>& correction, |
| 26 | bool keepdim, | 26 | bool keepdim, |
| 27 | - at::Tensor& result) | 27 | + at::Tensor& result) { |
| 28 | -{ | 28 | + if (!correction_fits_aclnn_int64(correction)) { |
| 29 | - if (!correction_fits_aclnn_int64(correction)) { | 29 | + at::Tensor cpu_out = result.cpu(); |
| 30 | - at::Tensor cpu_out = result.cpu(); | 30 | + at::var_out(cpu_out, self.cpu(), dim, correction, keepdim); |
| 31 | - at::var_out(cpu_out, self.cpu(), dim, correction, keepdim); | 31 | + result.copy_(cpu_out); |
| 32 | - result.copy_(cpu_out); | 32 | + return result; |
| 33 | - return result; | 33 | + } |
| 34 | - } | 34 | + bool unbiased = !(correction.has_value() && correction.value().toLong() == 0); |
| 35 | - bool unbiased = !(correction.has_value() && correction.value().toLong() == 0); | 35 | + int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1; |
| 36 | - int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1; | 36 | + return cal_var_out(self, dim.value_or(at::IntArrayRef{}), real_correction, unbiased, keepdim, result); |
| 37 | - return cal_var_out(self, dim.value_or(at::IntArrayRef{}), real_correction, unbiased, keepdim, result); | ||
| 38 | } | 37 | } |
| 39 | 38 | ||
| 40 | at::Tensor var( | 39 | at::Tensor var( |
| 41 | const at::Tensor& self, | 40 | const at::Tensor& self, |
| 42 | at::OptionalIntArrayRef dim, | 41 | at::OptionalIntArrayRef dim, |
| 43 | const c10::optional<c10::Scalar>& correction, | 42 | const c10::optional<c10::Scalar>& correction, |
| 44 | - bool keepdim) | 43 | + bool keepdim) { |
| 45 | -{ | 44 | + if (!correction_fits_aclnn_int64(correction)) { |
| 46 | - if (!correction_fits_aclnn_int64(correction)) { | 45 | + return at::var(self.cpu(), dim, correction, keepdim).to(self.options()); |
| 47 | - return at::var(self.cpu(), dim, correction, keepdim).to(self.options()); | 46 | + } |
| 48 | - } | 47 | + bool unbiased = !(correction.has_value() && correction.value().toLong() == 0); |
| 49 | - bool unbiased = !(correction.has_value() && correction.value().toLong() == 0); | 48 | + int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1; |
| 50 | - int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1; | 49 | + return cal_var(self, dim.value_or(at::IntArrayRef{}), real_correction, unbiased, keepdim); |
| 51 | - return cal_var(self, dim.value_or(at::IntArrayRef{}), real_correction, unbiased, keepdim); | ||
| 52 | } | 50 | } |
| 53 | 51 | ||
| 54 | std::tuple<at::Tensor, at::Tensor> var_mean( | 52 | std::tuple<at::Tensor, at::Tensor> var_mean( |
| 55 | const at::Tensor& self, | 53 | const at::Tensor& self, |
| 56 | at::OptionalIntArrayRef dim, | 54 | at::OptionalIntArrayRef dim, |
| 57 | const c10::optional<c10::Scalar>& correction, | 55 | const c10::optional<c10::Scalar>& correction, |
| 58 | - bool keepdim) | 56 | + bool keepdim) { |
| 59 | -{ | 57 | + if (!correction_fits_aclnn_int64(correction)) { |
| 60 | - if (!correction_fits_aclnn_int64(correction)) { | 58 | + auto cpu_tup = at::var_mean(self.cpu(), dim, correction, keepdim); |
| 61 | - auto cpu_tup = at::var_mean(self.cpu(), dim, correction, keepdim); | 59 | + return std::make_tuple(std::get<0>(cpu_tup).to(self.options()), std::get<1>(cpu_tup).to(self.options())); |
| 62 | - return std::make_tuple( | 60 | + } |
| 63 | - std::get<0>(cpu_tup).to(self.options()), | 61 | + bool unbiased = !(correction.has_value() && correction.value().toLong() == 0); |
| 64 | - std::get<1>(cpu_tup).to(self.options())); | 62 | + int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1; |
| 65 | - } | 63 | + return cal_var_mean(self, dim.value_or(at::IntArrayRef{}), unbiased, real_correction, keepdim); |
| 66 | - bool unbiased = !(correction.has_value() && correction.value().toLong() == 0); | ||
| 67 | - int64_t real_correction = correction.has_value() ? correction.value().toLong() : 1; | ||
| 68 | - return cal_var_mean(self, dim.value_or(at::IntArrayRef{}), unbiased, real_correction, keepdim); | ||
| 69 | } | 64 | } |
| 70 | } // namespace acl_op | 65 | } // namespace acl_op |
| @@ -24,100 +24,98 @@ namespace { | |||
| 24 | // format are base format (the format of src and dst are all nchw now) | 24 | // format are base format (the format of src and dst are all nchw now) |
| 25 | // dtype are same | 25 | // dtype are same |
| 26 | // so the view_value and ReflushDescBySelf are base on the hypothesis above. | 26 | // so the view_value and ReflushDescBySelf are base on the hypothesis above. |
| 27 | -bool AicoreValid(at::Tensor &self, const at::Tensor &src) | 27 | +bool AicoreValid(at::Tensor& self, const at::Tensor& src) { |
| 28 | -{ | 28 | + const auto& dst_storage_sizes = torch_npu::NPUBridge::GetNpuStorageImpl(self)->npu_desc_.storage_sizes_; |
| 29 | - const auto &dst_storage_sizes = torch_npu::NPUBridge::GetNpuStorageImpl(self)->npu_desc_.storage_sizes_; | 29 | + auto self_size = self.sizes(); |
| 30 | - auto self_size = self.sizes(); | 30 | + auto self_stride = self.strides(); |
| 31 | - auto self_stride = self.strides(); | 31 | + auto dst_storage_size_len = dst_storage_sizes.size(); |
| 32 | - auto dst_storage_size_len = dst_storage_sizes.size(); | 32 | + auto self_size_len = self_size.size(); |
| 33 | - auto self_size_len = self_size.size(); | ||
| 34 | 33 | ||
| 35 | - // count the difference between dst_storage and dst_size. | 34 | + // count the difference between dst_storage and dst_size. |
| 36 | - auto diff = dst_storage_size_len - self_size_len; | 35 | + auto diff = dst_storage_size_len - self_size_len; |
| 37 | - if (diff < 0 || diff > 1) { | 36 | + if (diff < 0 || diff > 1) { |
| 37 | + return false; | ||
| 38 | + } | ||
| 39 | + | ||
| 40 | + // record the index of the difference. | ||
| 41 | + auto diff_index = self_size_len; | ||
| 42 | + for (uint64_t i = 0; i < self_size_len; i++) { | ||
| 43 | + if (dst_storage_sizes[i] != self_size[i]) { | ||
| 44 | + ++diff; | ||
| 45 | + if (diff > 1) { | ||
| 38 | return false; | 46 | return false; |
| 47 | + } | ||
| 48 | + // differece should be 1. | ||
| 49 | + diff_index = i; | ||
| 39 | } | 50 | } |
| 51 | + } | ||
| 40 | 52 | ||
| 41 | - // record the index of the difference. | 53 | + // if diff or diff_index equals 0, no need viewcopy. |
| 42 | - auto diff_index = self_size_len; | 54 | + if (diff == 0 || diff_index == 0) { |
| 43 | - for (uint64_t i = 0; i < self_size_len; i++) { | 55 | + return false; |
| 44 | - if (dst_storage_sizes[i] != self_size[i]) { | 56 | + } |
| 45 | - ++diff; | 57 | + |
| 46 | - if (diff > 1) { | 58 | + const auto& dst_base_stride = torch_npu::NPUBridge::GetNpuStorageImpl(self)->npu_desc_.base_strides_; |
| 47 | - return false; | 59 | + // dst_base_stride should be equal to dst_storage_stride except for diff_index |
| 48 | - } | 60 | + if (self_stride.size() > dst_base_stride.size()) { |
| 49 | - // differece should be 1. | 61 | + return false; |
| 50 | - diff_index = i; | 62 | + } |
| 51 | - } | 63 | + |
| 64 | + for (uint64_t i = 0; i < self_stride.size(); i++) { | ||
| 65 | + if (dst_base_stride[i] != self_stride[i] && i != diff_index) { | ||
| 66 | + return false; | ||
| 52 | } | 67 | } |
| 68 | + } | ||
| 53 | 69 | ||
| 54 | - // if diff or diff_index equals 0, no need viewcopy. | 70 | + // dtype cannot be double and dst_size has to be equal with src_size. |
| 55 | - if (diff == 0 || diff_index == 0) { | 71 | + if (self.dtype() == at::ScalarType::Double || self_size != src.sizes()) { |
| 56 | - return false; | 72 | + return false; |
| 57 | - } | 73 | + } |
| 58 | 74 | ||
| 59 | - const auto &dst_base_stride = torch_npu::NPUBridge::GetNpuStorageImpl(self)->npu_desc_.base_strides_; | 75 | + return true; |
| 60 | - // dst_base_stride should be equal to dst_storage_stride except for diff_index | ||
| 61 | - if (self_stride.size() > dst_base_stride.size()) { | ||
| 62 | - return false; | ||
| 63 | - } | ||
| 64 | - | ||
| 65 | - for (uint64_t i = 0; i < self_stride.size(); i++) { | ||
| 66 | - if (dst_base_stride[i] != self_stride[i] && i != diff_index) { | ||
| 67 | - return false; | ||
| 68 | - } | ||
| 69 | - } | ||
| 70 | - | ||
| 71 | - // dtype cannot be double and dst_size has to be equal with src_size. | ||
| 72 | - if (self.dtype() == at::ScalarType::Double || self_size != src.sizes()) { | ||
| 73 | - return false; | ||
| 74 | - } | ||
| 75 | - | ||
| 76 | - return true; | ||
| 77 | } | 76 | } |
| 78 | } // namespace | 77 | } // namespace |
| 79 | 78 | ||
| 80 | -at::Tensor &npu_view_copy(at::Tensor &self, const at::Tensor &other, bool non_blocking) | 79 | +at::Tensor& npu_view_copy(at::Tensor& self, const at::Tensor& other, bool non_blocking) { |
| 81 | -{ | 80 | + auto self_size = self.sizes(); |
| 82 | - auto self_size = self.sizes(); | 81 | + auto self_stride = self.strides(); |
| 83 | - auto self_stride = self.strides(); | 82 | + auto src_size = other.sizes(); |
| 84 | - auto src_size = other.sizes(); | 83 | + auto src_stride = other.strides(); |
| 85 | - auto src_stride = other.strides(); | ||
| 86 | 84 | ||
| 87 | - at_npu::native::OpCommand cmd; | 85 | + at_npu::native::OpCommand cmd; |
| 88 | - if (AicoreValid(self, other)) { | 86 | + if (AicoreValid(self, other)) { |
| 89 | - at::Tensor contiguous_src(other); | 87 | + at::Tensor contiguous_src(other); |
| 90 | - if (!npu_utils::check_match(&contiguous_src)) { | 88 | + if (!npu_utils::check_match(&contiguous_src)) { |
| 91 | - contiguous_src = npu_utils::format_contiguous(contiguous_src); | 89 | + contiguous_src = npu_utils::format_contiguous(contiguous_src); |
| 92 | - } | ||
| 93 | - src_stride = contiguous_src.strides(); | ||
| 94 | - | ||
| 95 | - cmd.Name("ViewCopy") | ||
| 96 | - .InputWithoutContiguous(self) | ||
| 97 | - .Input(self_size, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 98 | - .Input(self_stride, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 99 | - .Input(at::Scalar(0), at::kLong) | ||
| 100 | - .InputWithoutContiguous(contiguous_src) | ||
| 101 | - .Input(src_size, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 102 | - .Input(src_stride, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 103 | - .Input(at::Scalar(0), at::kLong) | ||
| 104 | - .Output(self) | ||
| 105 | - .Run(); | ||
| 106 | - } else { | ||
| 107 | - cmd.Name("ViewCopy") | ||
| 108 | - .InputWithoutContiguous(self) | ||
| 109 | - .Input(self_size, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 110 | - .Input(self_stride, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 111 | - .Input(at::Scalar(0), at::kLong) | ||
| 112 | - .InputWithoutContiguous(other) | ||
| 113 | - .Input(src_size, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 114 | - .Input(src_stride, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 115 | - .Input(at::Scalar(0), at::kLong) | ||
| 116 | - .Output(self) | ||
| 117 | - .Attr("_exclude_engines", static_cast<string>("AiCore")) | ||
| 118 | - .Run(); | ||
| 119 | } | 90 | } |
| 91 | + src_stride = contiguous_src.strides(); | ||
| 120 | 92 | ||
| 121 | - return self; | 93 | + cmd.Name("ViewCopy") |
| 94 | + .InputWithoutContiguous(self) | ||
| 95 | + .Input(self_size, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 96 | + .Input(self_stride, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 97 | + .Input(at::Scalar(0), at::kLong) | ||
| 98 | + .InputWithoutContiguous(contiguous_src) | ||
| 99 | + .Input(src_size, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 100 | + .Input(src_stride, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 101 | + .Input(at::Scalar(0), at::kLong) | ||
| 102 | + .Output(self) | ||
| 103 | + .Run(); | ||
| 104 | + } else { | ||
| 105 | + cmd.Name("ViewCopy") | ||
| 106 | + .InputWithoutContiguous(self) | ||
| 107 | + .Input(self_size, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 108 | + .Input(self_stride, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 109 | + .Input(at::Scalar(0), at::kLong) | ||
| 110 | + .InputWithoutContiguous(other) | ||
| 111 | + .Input(src_size, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 112 | + .Input(src_stride, at::kLong, npu_compile_type::MEMORY_HOST_COMPILE_INDEPENDENT) | ||
| 113 | + .Input(at::Scalar(0), at::kLong) | ||
| 114 | + .Output(self) | ||
| 115 | + .Attr("_exclude_engines", static_cast<string>("AiCore")) | ||
| 116 | + .Run(); | ||
| 117 | + } | ||
| 118 | + | ||
| 119 | + return self; | ||
| 122 | } | 120 | } |
| 123 | } // namespace acl_op | 121 | } // namespace acl_op |
| @@ -36,11 +36,11 @@ at::SmallVector<int64_t, SIZE> where_npu_output_size(const at::Tensor& condition | |||
| 36 | std::vector<at::Tensor> where(const at::Tensor& condition) { | 36 | std::vector<at::Tensor> where(const at::Tensor& condition) { |
| 37 | at::Tensor format_cast_of_condition = condition; | 37 | at::Tensor format_cast_of_condition = condition; |
| 38 | if (npu_preparation::get_tensor_npu_format(condition) != ACL_FORMAT_ND) { | 38 | if (npu_preparation::get_tensor_npu_format(condition) != ACL_FORMAT_ND) { |
| 39 | - format_cast_of_condition = | 39 | + format_cast_of_condition = at_npu::native::custom_ops::npu_format_cast(format_cast_of_condition, ACL_FORMAT_ND); |
| 40 | - at_npu::native::custom_ops::npu_format_cast(format_cast_of_condition, ACL_FORMAT_ND); | ||
| 41 | } | 40 | } |
| 42 | if (condition.scalar_type() == at::ScalarType::Half) { | 41 | if (condition.scalar_type() == at::ScalarType::Half) { |
| 43 | - format_cast_of_condition = at_npu::native::custom_ops::_npu_dtype_cast(format_cast_of_condition, at::ScalarType::Float); | 42 | + format_cast_of_condition = |
| 43 | + at_npu::native::custom_ops::_npu_dtype_cast(format_cast_of_condition, at::ScalarType::Float); | ||
| 44 | } | 44 | } |
| 45 | 45 | ||
| 46 | auto output_size = where_npu_output_size(format_cast_of_condition); | 46 | auto output_size = where_npu_output_size(format_cast_of_condition); |
| @@ -48,10 +48,7 @@ std::vector<at::Tensor> where(const at::Tensor& condition) { | |||
| 48 | output_size, format_cast_of_condition.options().dtype(at::kLong), ACL_FORMAT_ND); | 48 | output_size, format_cast_of_condition.options().dtype(at::kLong), ACL_FORMAT_ND); |
| 49 | 49 | ||
| 50 | at_npu::native::OpCommand cmd; | 50 | at_npu::native::OpCommand cmd; |
| 51 | - cmd.Name("NonZero") | 51 | + cmd.Name("NonZero").Input(format_cast_of_condition).Output(result).Run(); |
| 52 | - .Input(format_cast_of_condition) | ||
| 53 | - .Output(result) | ||
| 54 | - .Run(); | ||
| 55 | result = result.transpose(1, 0); | 52 | result = result.transpose(1, 0); |
| 56 | std::vector<at::Tensor> chunk_result = result.chunk(result.size(0), 0); | 53 | std::vector<at::Tensor> chunk_result = result.chunk(result.size(0), 0); |
| 57 | std::vector<at::Tensor> squeeze_result; | 54 | std::vector<at::Tensor> squeeze_result; |
| @@ -62,40 +59,30 @@ std::vector<at::Tensor> where(const at::Tensor& condition) { | |||
| 62 | return squeeze_result; | 59 | return squeeze_result; |
| 63 | } | 60 | } |
| 64 | 61 | ||
| 65 | -at::Tensor& where_out( | 62 | +at::Tensor& where_out(const at::Tensor& condition, const at::Tensor& self, const at::Tensor& other, at::Tensor& out) { |
| 66 | - const at::Tensor& condition, | 63 | + at::Tensor b_condition; |
| 67 | - const at::Tensor& self, | 64 | + at::Tensor b_self; |
| 68 | - const at::Tensor& other, | 65 | + at::Tensor b_other; |
| 69 | - at::Tensor& out) { | 66 | + std::tie(b_condition, b_self, b_other) = npu_expand_outplace(condition, self, other, "where_npu"); |
| 70 | - at::Tensor b_condition; | 67 | + npu_preparation::CheckOut({condition, self, other}, out, b_self); |
| 71 | - at::Tensor b_self; | 68 | + if (!npu_utils::check_match(&out)) { |
| 72 | - at::Tensor b_other; | 69 | + at::Tensor contiguous_out = npu_utils::format_contiguous(out); |
| 73 | - std::tie(b_condition, b_self, b_other) = npu_expand_outplace(condition, self, other, "where_npu"); | 70 | + where_out_nocheck(contiguous_out, condition, self, other); |
| 74 | - npu_preparation::CheckOut( | 71 | + npu_utils::format_fresh_view(out, contiguous_out); |
| 75 | - {condition, self, other}, | 72 | + } else { |
| 76 | - out, | 73 | + where_out_nocheck(out, condition, self, other); |
| 77 | - b_self); | 74 | + } |
| 78 | - if (!npu_utils::check_match(&out)) { | ||
| 79 | - at::Tensor contiguous_out = npu_utils::format_contiguous(out); | ||
| 80 | - where_out_nocheck(contiguous_out, condition, self, other); | ||
| 81 | - npu_utils::format_fresh_view(out, contiguous_out); | ||
| 82 | - } else { | ||
| 83 | - where_out_nocheck(out, condition, self, other); | ||
| 84 | - } | ||
| 85 | 75 | ||
| 86 | - return out; | 76 | + return out; |
| 87 | } | 77 | } |
| 88 | 78 | ||
| 89 | -at::Tensor where( | 79 | +at::Tensor where(const at::Tensor& condition, const at::Tensor& self, const at::Tensor& other) { |
| 90 | - const at::Tensor& condition, | 80 | + at::Tensor b_condition; |
| 91 | - const at::Tensor& self, | 81 | + at::Tensor b_self; |
| 92 | - const at::Tensor& other) { | 82 | + at::Tensor b_other; |
| 93 | - at::Tensor b_condition; | 83 | + std::tie(b_condition, b_self, b_other) = npu_expand_outplace(condition, self, other, "where_npu"); |
| 94 | - at::Tensor b_self; | 84 | + at::Tensor ret = npu_preparation::apply_tensor(b_self); |
| 95 | - at::Tensor b_other; | 85 | + where_out_nocheck(ret, b_condition, b_self, b_other); |
| 96 | - std::tie(b_condition, b_self, b_other) = npu_expand_outplace(condition, self, other, "where_npu"); | 86 | + return ret; |
| 97 | - at::Tensor ret = npu_preparation::apply_tensor(b_self); | ||
| 98 | - where_out_nocheck(ret, b_condition, b_self, b_other); | ||
| 99 | - return ret; | ||
| 100 | } | 87 | } |
| 101 | } // namespace acl_op | 88 | } // namespace acl_op |
| @@ -21,118 +21,107 @@ using npu_preparation = at_npu::native::OpPreparation; | |||
| 21 | using npu_utils = at_npu::native::NpuUtils; | 21 | using npu_utils = at_npu::native::NpuUtils; |
| 22 | 22 | ||
| 23 | namespace { | 23 | namespace { |
| 24 | -at::Tensor &xlogy_out_npu_nocheck(at::Tensor &result, const at::Tensor &self, const at::Tensor &other) | 24 | +at::Tensor& xlogy_out_npu_nocheck(at::Tensor& result, const at::Tensor& self, const at::Tensor& other) { |
| 25 | -{ | 25 | + at_npu::native::OpCommand cmd; |
| 26 | - at_npu::native::OpCommand cmd; | 26 | + cmd.Name("Xlogy").Input(self).Input(other).Output(result).Run(); |
| 27 | - cmd.Name("Xlogy").Input(self).Input(other).Output(result).Run(); | 27 | + return result; |
| 28 | - return result; | ||
| 29 | } | 28 | } |
| 30 | 29 | ||
| 31 | -at::Tensor &xlogy_out_npu_nocheck(at::Tensor &result, const at::Tensor &self, const at::Scalar &other) | 30 | +at::Tensor& xlogy_out_npu_nocheck(at::Tensor& result, const at::Tensor& self, const at::Scalar& other) { |
| 32 | -{ | 31 | + at_npu::native::OpCommand cmd; |
| 33 | - at_npu::native::OpCommand cmd; | 32 | + cmd.Name("Xlogy").Input(self).Input(other, self.scalar_type()).Output(result).Run(); |
| 34 | - cmd.Name("Xlogy").Input(self).Input(other, self.scalar_type()).Output(result).Run(); | 33 | + return result; |
| 35 | - return result; | ||
| 36 | } | 34 | } |
| 37 | 35 | ||
| 38 | -at::Tensor &xlogy_out_npu_nocheck(at::Tensor &result, const at::Scalar &self, const at::Tensor &other) | 36 | +at::Tensor& xlogy_out_npu_nocheck(at::Tensor& result, const at::Scalar& self, const at::Tensor& other) { |
| 39 | -{ | 37 | + at_npu::native::OpCommand cmd; |
| 40 | - at_npu::native::OpCommand cmd; | 38 | + cmd.Name("Xlogy").Input(self, other.scalar_type()).Input(other).Output(result).Run(); |
| 41 | - cmd.Name("Xlogy").Input(self, other.scalar_type()).Input(other).Output(result).Run(); | 39 | + return result; |
| 42 | - return result; | ||
| 43 | } | 40 | } |
| 44 | } // namespace | 41 | } // namespace |
| 45 | 42 | ||
| 46 | -at::Tensor &xlogy_out(const at::Tensor &self, const at::Tensor &other, at::Tensor &out) | 43 | +at::Tensor& xlogy_out(const at::Tensor& self, const at::Tensor& other, at::Tensor& out) { |
| 47 | -{ | 44 | + at::Tensor format_cast_of_self = npu_preparation::CastBackToOriFormat(self); |
| 48 | - at::Tensor format_cast_of_self = npu_preparation::CastBackToOriFormat(self); | 45 | + at::Tensor format_cast_of_other = npu_preparation::CastBackToOriFormat(other); |
| 49 | - at::Tensor format_cast_of_other = npu_preparation::CastBackToOriFormat(other); | 46 | + auto output_size = op_infer::broadcast_ops_npu_output_size(self, other); |
| 50 | - auto output_size = op_infer::broadcast_ops_npu_output_size(self, other); | 47 | + npu_preparation::CheckOut( |
| 51 | - npu_preparation::CheckOut({self, other}, out, npu_preparation::get_tensor_npu_format(format_cast_of_self), | 48 | + {self, other}, out, npu_preparation::get_tensor_npu_format(format_cast_of_self), out.scalar_type(), output_size); |
| 52 | - out.scalar_type(), output_size); | 49 | + if (!npu_utils::check_match(&out)) { |
| 53 | - if (!npu_utils::check_match(&out)) { | 50 | + at::Tensor contiguous_result = npu_utils::format_contiguous(out); |
| 54 | - at::Tensor contiguous_result = npu_utils::format_contiguous(out); | 51 | + xlogy_out_npu_nocheck(contiguous_result, format_cast_of_self, format_cast_of_other); |
| 55 | - xlogy_out_npu_nocheck(contiguous_result, format_cast_of_self, format_cast_of_other); | 52 | + npu_utils::format_fresh_view(out, contiguous_result); |
| 56 | - npu_utils::format_fresh_view(out, contiguous_result); | 53 | + } else { |
| 57 | - } else { | 54 | + xlogy_out_npu_nocheck(out, format_cast_of_self, format_cast_of_other); |
| 58 | - xlogy_out_npu_nocheck(out, format_cast_of_self, format_cast_of_other); | 55 | + } |
| 59 | - } | 56 | + return out; |
| 60 | - return out; | ||
| 61 | } | 57 | } |
| 62 | 58 | ||
| 63 | -at::Tensor &xlogy_out(const at::Tensor &self, const at::Scalar &other, at::Tensor &out) | 59 | +at::Tensor& xlogy_out(const at::Tensor& self, const at::Scalar& other, at::Tensor& out) { |
| 64 | -{ | 60 | + npu_preparation::CheckOut({self}, out, self); |
| 65 | - npu_preparation::CheckOut({self}, out, self); | 61 | + if (!npu_utils::check_match(&out)) { |
| 66 | - if (!npu_utils::check_match(&out)) { | 62 | + at::Tensor contiguous_result = npu_utils::format_contiguous(out); |
| 67 | - at::Tensor contiguous_result = npu_utils::format_contiguous(out); | 63 | + xlogy_out_npu_nocheck(contiguous_result, self, other); |
| 68 | - xlogy_out_npu_nocheck(contiguous_result, self, other); | 64 | + npu_utils::format_fresh_view(out, contiguous_result); |
| 69 | - npu_utils::format_fresh_view(out, contiguous_result); | 65 | + } else { |
| 70 | - } else { | 66 | + xlogy_out_npu_nocheck(out, self, other); |
| 71 | - xlogy_out_npu_nocheck(out, self, other); | 67 | + } |
| 72 | - } | 68 | + return out; |
| 73 | - return out; | ||
| 74 | } | 69 | } |
| 75 | 70 | ||
| 76 | -at::Tensor &xlogy_out(const at::Scalar &self, const at::Tensor &other, at::Tensor &out) | 71 | +at::Tensor& xlogy_out(const at::Scalar& self, const at::Tensor& other, at::Tensor& out) { |
| 77 | -{ | 72 | + npu_preparation::CheckOut( |
| 78 | - npu_preparation::CheckOut({other}, out, npu_preparation::get_tensor_npu_format(other), other.scalar_type(), | 73 | + {other}, out, npu_preparation::get_tensor_npu_format(other), other.scalar_type(), other.sizes()); |
| 79 | - other.sizes()); | 74 | + if (!npu_utils::check_match(&out)) { |
| 80 | - if (!npu_utils::check_match(&out)) { | 75 | + at::Tensor contiguous_result = npu_utils::format_contiguous(out); |
| 81 | - at::Tensor contiguous_result = npu_utils::format_contiguous(out); | 76 | + xlogy_out_npu_nocheck(contiguous_result, self, other); |
| 82 | - xlogy_out_npu_nocheck(contiguous_result, self, other); | 77 | + npu_utils::format_fresh_view(out, contiguous_result); |
| 83 | - npu_utils::format_fresh_view(out, contiguous_result); | 78 | + } else { |
| 84 | - } else { | 79 | + xlogy_out_npu_nocheck(out, self, other); |
| 85 | - xlogy_out_npu_nocheck(out, self, other); | 80 | + } |
| 86 | - } | 81 | + return out; |
| 87 | - return out; | ||
| 88 | } | 82 | } |
| 89 | 83 | ||
| 90 | -at::Tensor xlogy(const at::Tensor &self, const at::Tensor &other) | 84 | +at::Tensor xlogy(const at::Tensor& self, const at::Tensor& other) { |
| 91 | -{ | 85 | + bool is_self_wrapped = npu_preparation::is_scalar_wrapped_to_tensor(self); |
| 92 | - bool is_self_wrapped = npu_preparation::is_scalar_wrapped_to_tensor(self); | 86 | + at::Tensor output_tensor = is_self_wrapped ? other : self; |
| 93 | - at::Tensor output_tensor = is_self_wrapped ? other : self; | 87 | + auto output_size = op_infer::broadcast_ops_npu_output_size(self, other); |
| 94 | - auto output_size = op_infer::broadcast_ops_npu_output_size(self, other); | 88 | + at::Tensor result = npu_preparation::apply_tensor(output_tensor, output_size); |
| 95 | - at::Tensor result = npu_preparation::apply_tensor(output_tensor, output_size); | 89 | + xlogy_out_npu_nocheck(result, self, other); |
| 96 | - xlogy_out_npu_nocheck(result, self, other); | 90 | + return result; |
| 97 | - return result; | ||
| 98 | } | 91 | } |
| 99 | 92 | ||
| 100 | -at::Tensor xlogy(const at::Tensor &self, const at::Scalar &other) | 93 | +at::Tensor xlogy(const at::Tensor& self, const at::Scalar& other) { |
| 101 | -{ | 94 | + at::Tensor result = npu_preparation::apply_tensor(self); |
| 102 | - at::Tensor result = npu_preparation::apply_tensor(self); | 95 | + xlogy_out_npu_nocheck(result, self, other); |
| 103 | - xlogy_out_npu_nocheck(result, self, other); | 96 | + return result; |
| 104 | - return result; | ||
| 105 | } | 97 | } |
| 106 | 98 | ||
| 107 | -at::Tensor xlogy(const at::Scalar &self, const at::Tensor &other) | 99 | +at::Tensor xlogy(const at::Scalar& self, const at::Tensor& other) { |
| 108 | -{ | 100 | + at::Tensor result = npu_preparation::apply_tensor(other); |
| 109 | - at::Tensor result = npu_preparation::apply_tensor(other); | 101 | + xlogy_out_npu_nocheck(result, self, other); |
| 110 | - xlogy_out_npu_nocheck(result, self, other); | 102 | + return result; |
| 111 | - return result; | ||
| 112 | } | 103 | } |
| 113 | 104 | ||
| 114 | -at::Tensor &xlogy_(at::Tensor &self, const at::Tensor &other) | 105 | +at::Tensor& xlogy_(at::Tensor& self, const at::Tensor& other) { |
| 115 | -{ | 106 | + npu_preparation::CheckMemory({self, other}, {self}); |
| 116 | - npu_preparation::CheckMemory({self, other}, {self}); | 107 | + if (!npu_utils::check_match(&self)) { |
| 117 | - if (!npu_utils::check_match(&self)) { | 108 | + at::Tensor contiguous_self = npu_utils::format_contiguous(self); |
| 118 | - at::Tensor contiguous_self = npu_utils::format_contiguous(self); | 109 | + at::Tensor result = xlogy_out_npu_nocheck(contiguous_self, contiguous_self, other); |
| 119 | - at::Tensor result = xlogy_out_npu_nocheck(contiguous_self, contiguous_self, other); | 110 | + npu_utils::format_fresh_view(self, contiguous_self); |
| 120 | - npu_utils::format_fresh_view(self, contiguous_self); | 111 | + } else { |
| 121 | - } else { | 112 | + xlogy_out_npu_nocheck(self, self, other); |
| 122 | - xlogy_out_npu_nocheck(self, self, other); | 113 | + } |
| 123 | - } | 114 | + return self; |
| 124 | - return self; | ||
| 125 | } | 115 | } |
| 126 | 116 | ||
| 127 | -at::Tensor &xlogy_(at::Tensor &self, const at::Scalar &other) | 117 | +at::Tensor& xlogy_(at::Tensor& self, const at::Scalar& other) { |
| 128 | -{ | 118 | + if (!npu_utils::check_match(&self)) { |
| 129 | - if (!npu_utils::check_match(&self)) { | 119 | + at::Tensor contiguous_self = npu_utils::format_contiguous(self); |
| 130 | - at::Tensor contiguous_self = npu_utils::format_contiguous(self); | 120 | + xlogy_out_npu_nocheck(contiguous_self, contiguous_self, other); |
| 131 | - xlogy_out_npu_nocheck(contiguous_self, contiguous_self, other); | 121 | + npu_utils::format_fresh_view(self, contiguous_self); |
| 132 | - npu_utils::format_fresh_view(self, contiguous_self); | 122 | + } else { |
| 133 | - } else { | 123 | + xlogy_out_npu_nocheck(self, self, other); |
| 134 | - xlogy_out_npu_nocheck(self, self, other); | 124 | + } |
| 135 | - } | 125 | + return self; |
| 136 | - return self; | ||
| 137 | } | 126 | } |
| 138 | } // namespace acl_op | 127 | } // namespace acl_op |
| @@ -25,34 +25,39 @@ namespace { | |||
| 25 | inline void yolo_boxes_encode_check( | 25 | inline void yolo_boxes_encode_check( |
| 26 | const at::Tensor& anchor_boxes, | 26 | const at::Tensor& anchor_boxes, |
| 27 | const at::Tensor& gt_bboxes, | 27 | const at::Tensor& gt_bboxes, |
| 28 | - const at::Tensor& stride) | 28 | + const at::Tensor& stride) { |
| 29 | -{ | 29 | + TORCH_CHECK( |
| 30 | - TORCH_CHECK( | 30 | + anchor_boxes.dim() == 2 && anchor_boxes.size(1) == 4, |
| 31 | - anchor_boxes.dim() == 2 && anchor_boxes.size(1) == 4, | 31 | + "Non-empty 2D anchor_boxes tensor expected but got a tensor with sizes ", |
| 32 | - "Non-empty 2D anchor_boxes tensor expected but got a tensor with sizes ", | 32 | + anchor_boxes.sizes(), |
| 33 | - anchor_boxes.sizes(), OPS_ERROR(ErrCode::PARAM)); | 33 | + OPS_ERROR(ErrCode::PARAM)); |
| 34 | - TORCH_CHECK( | 34 | + TORCH_CHECK( |
| 35 | - anchor_boxes.size(0) <= MAX_ANCHOR_BOX_SIZE, | 35 | + anchor_boxes.size(0) <= MAX_ANCHOR_BOX_SIZE, |
| 36 | - "anchor_boxes only support max [20480] num, but got num ", | 36 | + "anchor_boxes only support max [20480] num, but got num ", |
| 37 | - anchor_boxes.size(0), OPS_ERROR(ErrCode::PARAM)); | 37 | + anchor_boxes.size(0), |
| 38 | - TORCH_CHECK( | 38 | + OPS_ERROR(ErrCode::PARAM)); |
| 39 | - gt_bboxes.dim() == 2 && gt_bboxes.size(1) == 4, | 39 | + TORCH_CHECK( |
| 40 | - "Non-empty 2D gt_bboxes tensor expected but got a tensor with sizes ", | 40 | + gt_bboxes.dim() == 2 && gt_bboxes.size(1) == 4, |
| 41 | - gt_bboxes.sizes(), OPS_ERROR(ErrCode::PARAM)); | 41 | + "Non-empty 2D gt_bboxes tensor expected but got a tensor with sizes ", |
| 42 | - TORCH_CHECK( | 42 | + gt_bboxes.sizes(), |
| 43 | - stride.dim() == 1, | 43 | + OPS_ERROR(ErrCode::PARAM)); |
| 44 | - "Non-empty 1D stride tensor expected but got a tensor with sizes ", | 44 | + TORCH_CHECK( |
| 45 | - stride.sizes(), OPS_ERROR(ErrCode::PARAM)); | 45 | + stride.dim() == 1, |
| 46 | - TORCH_CHECK( | 46 | + "Non-empty 1D stride tensor expected but got a tensor with sizes ", |
| 47 | - stride.size(0) == gt_bboxes.size(0), | 47 | + stride.sizes(), |
| 48 | - "stride's length should be equal gt_bboxes' num, but got stride length ", | 48 | + OPS_ERROR(ErrCode::PARAM)); |
| 49 | - stride.size(0), | 49 | + TORCH_CHECK( |
| 50 | - "gt_bboxes num ", | 50 | + stride.size(0) == gt_bboxes.size(0), |
| 51 | - gt_bboxes.size(0), OPS_ERROR(ErrCode::PARAM)); | 51 | + "stride's length should be equal gt_bboxes' num, but got stride length ", |
| 52 | - TORCH_CHECK( | 52 | + stride.size(0), |
| 53 | - at::isIntegralType(stride.scalar_type(), true) && stride.scalar_type() != at::ScalarType::Long, | 53 | + "gt_bboxes num ", |
| 54 | - "int32 strdie tensor expected but got a tensor with dtype: ", | 54 | + gt_bboxes.size(0), |
| 55 | - stride.scalar_type(), OPS_ERROR(ErrCode::TYPE)); | 55 | + OPS_ERROR(ErrCode::PARAM)); |
| 56 | + TORCH_CHECK( | ||
| 57 | + at::isIntegralType(stride.scalar_type(), true) && stride.scalar_type() != at::ScalarType::Long, | ||
| 58 | + "int32 strdie tensor expected but got a tensor with dtype: ", | ||
| 59 | + stride.scalar_type(), | ||
| 60 | + OPS_ERROR(ErrCode::TYPE)); | ||
| 56 | } | 61 | } |
| 57 | } // namespace | 62 | } // namespace |
| 58 | 63 | ||
| @@ -60,20 +65,19 @@ at::Tensor npu_yolo_boxes_encode( | |||
| 60 | const at::Tensor& self, | 65 | const at::Tensor& self, |
| 61 | const at::Tensor& gt_bboxes, | 66 | const at::Tensor& gt_bboxes, |
| 62 | const at::Tensor& stride, | 67 | const at::Tensor& stride, |
| 63 | - bool performance_mode) | 68 | + bool performance_mode) { |
| 64 | -{ | 69 | + yolo_boxes_encode_check(self, gt_bboxes, stride); |
| 65 | - yolo_boxes_encode_check(self, gt_bboxes, stride); | 70 | + at::Tensor result = npu_preparation::apply_tensor(gt_bboxes); |
| 66 | - at::Tensor result = npu_preparation::apply_tensor(gt_bboxes); | 71 | + string impl_mode_str = performance_mode ? "high_performance" : "high_precision"; |
| 67 | - string impl_mode_str = performance_mode ? "high_performance" : "high_precision"; | 72 | + at::Tensor stride_cp = at_npu::native::custom_ops::_npu_dtype_cast(stride, at::ScalarType::Int); |
| 68 | - at::Tensor stride_cp = at_npu::native::custom_ops::_npu_dtype_cast(stride, at::ScalarType::Int); | 73 | + at_npu::native::OpCommand cmd; |
| 69 | - at_npu::native::OpCommand cmd; | 74 | + cmd.Name("YoloBoxesEncode") |
| 70 | - cmd.Name("YoloBoxesEncode") | 75 | + .Input(self) |
| 71 | - .Input(self) | 76 | + .Input(gt_bboxes) |
| 72 | - .Input(gt_bboxes) | 77 | + .Input(stride_cp) |
| 73 | - .Input(stride_cp) | 78 | + .Output(result) |
| 74 | - .Output(result) | 79 | + .Attr("performance_mode", impl_mode_str) |
| 75 | - .Attr("performance_mode", impl_mode_str) | 80 | + .Run(); |
| 76 | - .Run(); | 81 | + return result; |
| 77 | - return result; | ||
| 78 | } | 82 | } |
| 79 | } // namespace acl_op | 83 | } // namespace acl_op |
| @@ -18,10 +18,9 @@ | |||
| 18 | 18 | ||
| 19 | 19 | ||
| 20 | namespace acl_op { | 20 | namespace acl_op { |
| 21 | -at::Tensor& zeros_out(at::IntArrayRef size, at::Tensor& out) | 21 | +at::Tensor& zeros_out(at::IntArrayRef size, at::Tensor& out) { |
| 22 | -{ | 22 | + out.resize_(size); |
| 23 | - out.resize_(size); | 23 | + return out.zero_(); |
| 24 | - return out.zero_(); | ||
| 25 | } | 24 | } |
| 26 | 25 | ||
| 27 | at::Tensor zeros_symint( | 26 | at::Tensor zeros_symint( |
| @@ -29,20 +28,17 @@ at::Tensor zeros_symint( | |||
| 29 | c10::optional<at::ScalarType> dtype, | 28 | c10::optional<at::ScalarType> dtype, |
| 30 | c10::optional<at::Layout> layout, | 29 | c10::optional<at::Layout> layout, |
| 31 | c10::optional<at::Device> device, | 30 | c10::optional<at::Device> device, |
| 32 | - c10::optional<bool> pin_memory) | 31 | + c10::optional<bool> pin_memory) { |
| 33 | -{ | 32 | + return zeros_common_nocheck(c10::asIntArrayRefUnchecked(size), dtype, layout, device, pin_memory); |
| 34 | - return zeros_common_nocheck(c10::asIntArrayRefUnchecked(size), dtype, layout, device, pin_memory); | ||
| 35 | } | 33 | } |
| 36 | 34 | ||
| 37 | - | ||
| 38 | at::Tensor zeros( | 35 | at::Tensor zeros( |
| 39 | at::IntArrayRef size, | 36 | at::IntArrayRef size, |
| 40 | c10::optional<at::DimnameList> names, | 37 | c10::optional<at::DimnameList> names, |
| 41 | c10::optional<at::ScalarType> dtype, | 38 | c10::optional<at::ScalarType> dtype, |
| 42 | c10::optional<at::Layout> layout, | 39 | c10::optional<at::Layout> layout, |
| 43 | c10::optional<at::Device> device, | 40 | c10::optional<at::Device> device, |
| 44 | - c10::optional<bool> pin_memory) | 41 | + c10::optional<bool> pin_memory) { |
| 45 | -{ | 42 | + return zeros_common_nocheck(size, dtype, layout, device, pin_memory); |
| 46 | - return zeros_common_nocheck(size, dtype, layout, device, pin_memory); | ||
| 47 | } | 43 | } |
| 48 | } // namespace acl_op | 44 | } // namespace acl_op |
| @@ -21,28 +21,23 @@ using npu_preparation = at_npu::native::OpPreparation; | |||
| 21 | using npu_utils = at_npu::native::NpuUtils; | 21 | using npu_utils = at_npu::native::NpuUtils; |
| 22 | 22 | ||
| 23 | namespace { | 23 | namespace { |
| 24 | -at::Tensor& zeros_like_out_npu_nocheck(at::Tensor& result, const at::Tensor& self) | 24 | +at::Tensor& zeros_like_out_npu_nocheck(at::Tensor& result, const at::Tensor& self) { |
| 25 | -{ | 25 | + at_npu::native::OpCommand cmd; |
| 26 | - at_npu::native::OpCommand cmd; | 26 | + cmd.Name("ZerosLike").Input(self).Output(result).Run(); |
| 27 | - cmd.Name("ZerosLike") | ||
| 28 | - .Input(self) | ||
| 29 | - .Output(result) | ||
| 30 | - .Run(); | ||
| 31 | 27 | ||
| 32 | - return result; | 28 | + return result; |
| 33 | } | 29 | } |
| 34 | } // namespace | 30 | } // namespace |
| 35 | 31 | ||
| 36 | -at::Tensor& zero_(at::Tensor& self) | 32 | +at::Tensor& zero_(at::Tensor& self) { |
| 37 | -{ | 33 | + if (!npu_utils::check_match(&self)) { |
| 38 | - if (!npu_utils::check_match(&self)) { | 34 | + at::Tensor contiguous_self = npu_utils::format_contiguous(self); |
| 39 | - at::Tensor contiguous_self = npu_utils::format_contiguous(self); | 35 | + zeros_like_out_npu_nocheck(contiguous_self, contiguous_self); |
| 40 | - zeros_like_out_npu_nocheck(contiguous_self, contiguous_self); | 36 | + npu_utils::format_fresh_view(self, contiguous_self); |
| 41 | - npu_utils::format_fresh_view(self, contiguous_self); | 37 | + } else { |
| 42 | - } else { | 38 | + zeros_like_out_npu_nocheck(self, self); |
| 43 | - zeros_like_out_npu_nocheck(self, self); | 39 | + } |
| 44 | - } | ||
| 45 | 40 | ||
| 46 | - return self; | 41 | + return self; |
| 47 | } | 42 | } |
| 48 | } // namespace acl_op | 43 | } // namespace acl_op |
| @@ -17,56 +17,103 @@ | |||
| 17 | 17 | ||
| 18 | namespace atb { | 18 | namespace atb { |
| 19 | namespace { | 19 | namespace { |
| 20 | -std::unordered_map<c10::string_view, int> activation_type_map = { | 20 | +std::unordered_map<c10::string_view, int> activation_type_map = {{"activation_sigmoid", 8}}; |
| 21 | - {"activation_sigmoid", 8} | ||
| 22 | -}; | ||
| 23 | 21 | ||
| 24 | -int get_fused_add_div_mode(c10::optional<c10::string_view> activation_type_opt) | 22 | +int get_fused_add_div_mode(c10::optional<c10::string_view> activation_type_opt) { |
| 25 | -{ | 23 | + int activation_type = |
| 26 | - int activation_type = atb::utils::get_op_mode( | 24 | + atb::utils::get_op_mode(activation_type_map, activation_type_opt, "activation_sigmoid", "activation_type"); |
| 27 | - activation_type_map, activation_type_opt, "activation_sigmoid", "activation_type"); | 25 | + return activation_type; |
| 28 | - return activation_type; | ||
| 29 | } | 26 | } |
| 27 | +} // namespace | ||
| 28 | + | ||
| 29 | +std::tuple<at::Tensor, at::Tensor> npu_fused_add_topk_div( | ||
| 30 | + const at::Tensor& x, | ||
| 31 | + const at::Tensor& add_num, | ||
| 32 | + const c10::optional<at::Tensor>& mapping_num, | ||
| 33 | + const c10::optional<at::Tensor>& mapping_table, | ||
| 34 | + c10::optional<c10::string_view> activation_type_opt, | ||
| 35 | + int64_t group_num, | ||
| 36 | + int64_t group_topk, | ||
| 37 | + int64_t n, | ||
| 38 | + int64_t k, | ||
| 39 | + bool is_norm, | ||
| 40 | + double scale, | ||
| 41 | + bool enable_expert_mapping) { | ||
| 42 | + const c10::OptionalDeviceGuard device_guard(device_of(x)); | ||
| 43 | + int64_t a = x.size(0); | ||
| 44 | + at::Tensor y = at::empty({a, k}, x.options().dtype(c10::ScalarType::Float)); | ||
| 45 | + at::Tensor indices = at::empty({a, k}, x.options().dtype(c10::ScalarType::Int)); | ||
| 46 | + float scale_float = static_cast<float>(scale); | ||
| 47 | + auto activation_type = get_fused_add_div_mode(activation_type_opt); | ||
| 48 | + EXEC_ATB_CMD( | ||
| 49 | + AtbFusedAddTopkDiv, | ||
| 50 | + x, | ||
| 51 | + add_num, | ||
| 52 | + mapping_num, | ||
| 53 | + mapping_table, | ||
| 54 | + group_num, | ||
| 55 | + group_topk, | ||
| 56 | + n, | ||
| 57 | + k, | ||
| 58 | + activation_type, | ||
| 59 | + is_norm, | ||
| 60 | + scale_float, | ||
| 61 | + enable_expert_mapping, | ||
| 62 | + y, | ||
| 63 | + indices); | ||
| 64 | + return std::make_tuple(y, indices); | ||
| 30 | } | 65 | } |
| 31 | 66 | ||
| 32 | -std::tuple<at::Tensor, at::Tensor> npu_fused_add_topk_div(const at::Tensor &x, const at::Tensor &add_num, const c10::optional<at::Tensor> &mapping_num, const c10::optional<at::Tensor> &mapping_table, | 67 | +std::tuple<at::Tensor&, at::Tensor&> npu_fused_add_topk_div_out( |
| 33 | - c10::optional<c10::string_view> activation_type_opt, int64_t group_num, int64_t group_topk, int64_t n, int64_t k, bool is_norm, double scale, bool enable_expert_mapping) | 68 | + const at::Tensor& x, |
| 34 | -{ | 69 | + const at::Tensor& add_num, |
| 35 | - const c10::OptionalDeviceGuard device_guard(device_of(x)); | 70 | + const c10::optional<at::Tensor>& mapping_num, |
| 36 | - int64_t a = x.size(0); | 71 | + const c10::optional<at::Tensor>& mapping_table, |
| 37 | - at::Tensor y = at::empty({a, k}, x.options().dtype(c10::ScalarType::Float)); | 72 | + c10::optional<c10::string_view> activation_type_opt, |
| 38 | - at::Tensor indices = at::empty({a, k}, x.options().dtype(c10::ScalarType::Int)); | 73 | + int64_t group_num, |
| 39 | - float scale_float = static_cast<float>(scale); | 74 | + int64_t group_topk, |
| 40 | - auto activation_type = get_fused_add_div_mode(activation_type_opt); | 75 | + int64_t n, |
| 41 | - EXEC_ATB_CMD(AtbFusedAddTopkDiv, x, add_num, mapping_num, mapping_table, group_num, group_topk, n, k, activation_type, is_norm, scale_float, enable_expert_mapping, y, indices); | 76 | + int64_t k, |
| 42 | - return std::make_tuple(y, indices); | 77 | + bool is_norm, |
| 43 | -} | 78 | + double scale, |
| 44 | - | 79 | + bool enable_expert_mapping, |
| 45 | -std::tuple<at::Tensor&, at::Tensor&> npu_fused_add_topk_div_out(const at::Tensor &x, const at::Tensor &add_num, const c10::optional<at::Tensor> &mapping_num, const c10::optional<at::Tensor> &mapping_table, | 80 | + at::Tensor& y, |
| 46 | - c10::optional<c10::string_view> activation_type_opt, int64_t group_num, int64_t group_topk, int64_t n, int64_t k, bool is_norm, double scale, bool enable_expert_mapping, | 81 | + at::Tensor& indices) { |
| 47 | - at::Tensor &y, | 82 | + const c10::OptionalDeviceGuard device_guard(device_of(x)); |
| 48 | - at::Tensor &indices) | 83 | + float scale_float = static_cast<float>(scale); |
| 49 | -{ | 84 | + auto activation_type = get_fused_add_div_mode(activation_type_opt); |
| 50 | - const c10::OptionalDeviceGuard device_guard(device_of(x)); | 85 | + EXEC_ATB_CMD( |
| 51 | - float scale_float = static_cast<float>(scale); | 86 | + AtbFusedAddTopkDiv, |
| 52 | - auto activation_type = get_fused_add_div_mode(activation_type_opt); | 87 | + x, |
| 53 | - EXEC_ATB_CMD(AtbFusedAddTopkDiv, x, add_num, mapping_num, mapping_table, group_num, group_topk, n, k, activation_type, is_norm, scale_float, enable_expert_mapping, y, indices); | 88 | + add_num, |
| 54 | - return std::forward_as_tuple(y, indices); | 89 | + mapping_num, |
| 90 | + mapping_table, | ||
| 91 | + group_num, | ||
| 92 | + group_topk, | ||
| 93 | + n, | ||
| 94 | + k, | ||
| 95 | + activation_type, | ||
| 96 | + is_norm, | ||
| 97 | + scale_float, | ||
| 98 | + enable_expert_mapping, | ||
| 99 | + y, | ||
| 100 | + indices); | ||
| 101 | + return std::forward_as_tuple(y, indices); | ||
| 55 | } | 102 | } |
| 56 | 103 | ||
| 57 | namespace { | 104 | namespace { |
| 58 | -TORCH_LIBRARY_FRAGMENT(atb, m) | 105 | +TORCH_LIBRARY_FRAGMENT(atb, m) { |
| 59 | -{ | 106 | + m.def( |
| 60 | - m.def("npu_fused_add_topk_div(Tensor x, Tensor add_num, *, Tensor? mapping_num=None, Tensor? mapping_table=None, str? activation_type=None, int group_num=1, int group_topk=1, int n=1, int k=1, bool is_norm=True, float scale=1, bool enable_expert_mapping=False) -> (Tensor, Tensor)"); | 107 | + "npu_fused_add_topk_div(Tensor x, Tensor add_num, *, Tensor? mapping_num=None, Tensor? mapping_table=None, str? activation_type=None, int group_num=1, int group_topk=1, int n=1, int k=1, bool is_norm=True, float scale=1, bool enable_expert_mapping=False) -> (Tensor, Tensor)"); |
| 61 | - m.def("npu_fused_add_topk_div.out(Tensor x, Tensor add_num, *, Tensor? mapping_num=None, Tensor? mapping_table=None, str? activation_type=None, int group_num=1, int group_topk=1, int n=1, int k=1, bool is_norm=True, float scale=1, bool enable_expert_mapping=False, Tensor(a!) y, Tensor(b!) indices) -> (Tensor(a!), Tensor(b!))"); | 108 | + m.def( |
| 62 | -} | 109 | + "npu_fused_add_topk_div.out(Tensor x, Tensor add_num, *, Tensor? mapping_num=None, Tensor? mapping_table=None, str? activation_type=None, int group_num=1, int group_topk=1, int n=1, int k=1, bool is_norm=True, float scale=1, bool enable_expert_mapping=False, Tensor(a!) y, Tensor(b!) indices) -> (Tensor(a!), Tensor(b!))"); |
| 63 | } | 110 | } |
| 111 | +} // namespace | ||
| 64 | 112 | ||
| 65 | namespace { | 113 | namespace { |
| 66 | -TORCH_LIBRARY_IMPL(atb, PrivateUse1, m) | 114 | +TORCH_LIBRARY_IMPL(atb, PrivateUse1, m) { |
| 67 | -{ | 115 | + m.impl("npu_fused_add_topk_div", TORCH_FN(atb::npu_fused_add_topk_div)); |
| 68 | - m.impl("npu_fused_add_topk_div", TORCH_FN(atb::npu_fused_add_topk_div)); | 116 | + m.impl("npu_fused_add_topk_div.out", TORCH_FN(atb::npu_fused_add_topk_div_out)); |
| 69 | - m.impl("npu_fused_add_topk_div.out", TORCH_FN(atb::npu_fused_add_topk_div_out)); | ||
| 70 | -} | ||
| 71 | -} | ||
| 72 | } | 117 | } |
| 118 | +} // namespace | ||
| 119 | +} // namespace atb | ||
| @@ -20,43 +20,38 @@ using namespace std; | |||
| 20 | namespace atb { | 20 | namespace atb { |
| 21 | 21 | ||
| 22 | using GroupTopkParam = atb::infer::GroupTopkParam; | 22 | using GroupTopkParam = atb::infer::GroupTopkParam; |
| 23 | -void _npu_group_topk(const at::Tensor &self, int64_t k, int64_t group_num, int64_t n) | 23 | +void _npu_group_topk(const at::Tensor& self, int64_t k, int64_t group_num, int64_t n) { |
| 24 | -{ | 24 | + const c10::OptionalDeviceGuard device_guard(device_of(self)); |
| 25 | - const c10::OptionalDeviceGuard device_guard(device_of(self)); | 25 | + OpParamCache<GroupTopkParam>& GroupTopkParamCache = OpParamCache<GroupTopkParam>::getInstance(); |
| 26 | - OpParamCache<GroupTopkParam>& GroupTopkParamCache = OpParamCache<GroupTopkParam>::getInstance(); | 26 | + GroupTopkParam GroupTopkParam; |
| 27 | - GroupTopkParam GroupTopkParam; | 27 | + GroupTopkParam.groupNum = static_cast<int32_t>(group_num); |
| 28 | - GroupTopkParam.groupNum = static_cast<int32_t>(group_num); | 28 | + GroupTopkParam.k = static_cast<int32_t>(k); |
| 29 | - GroupTopkParam.k = static_cast<int32_t>(k); | 29 | + GroupTopkParam.n = static_cast<uint16_t>(n); |
| 30 | - GroupTopkParam.n = static_cast<uint16_t>(n); | 30 | + GroupTopkParam.groupMultiFlag = static_cast<GroupTopkParam::GroupMultiFlag>(0); |
| 31 | - GroupTopkParam.groupMultiFlag = static_cast<GroupTopkParam::GroupMultiFlag>(0); | 31 | + if (n > 1) { |
| 32 | - if (n > 1) { | 32 | + GroupTopkParam.groupMultiFlag = static_cast<GroupTopkParam::GroupMultiFlag>(1); |
| 33 | - GroupTopkParam.groupMultiFlag = static_cast<GroupTopkParam::GroupMultiFlag>(1); | 33 | + } |
| 34 | - } | ||
| 35 | 34 | ||
| 36 | - at::Tensor out = self; | 35 | + at::Tensor out = self; |
| 37 | - auto idx = at::arange(1024, self.options().device(at::Device(at::kPrivateUse1)).dtype(at::kInt)); | 36 | + auto idx = at::arange(1024, self.options().device(at::Device(at::kPrivateUse1)).dtype(at::kInt)); |
| 38 | 37 | ||
| 39 | - ParamSetter paramsetter; | 38 | + ParamSetter paramsetter; |
| 40 | - paramsetter.Input(self, true) | 39 | + paramsetter.Input(self, true).Input(idx, true).Output(out); |
| 41 | - .Input(idx, true) | 40 | + auto opGroupTopk = GroupTopkParamCache.getOperation(GroupTopkParam, "GroupTopkOperation"); |
| 42 | - .Output(out); | 41 | + RunAtbCmd(opGroupTopk, paramsetter, "GroupTopkOperation"); |
| 43 | - auto opGroupTopk = GroupTopkParamCache.getOperation(GroupTopkParam, "GroupTopkOperation"); | 42 | + return; |
| 44 | - RunAtbCmd(opGroupTopk, paramsetter, "GroupTopkOperation"); | ||
| 45 | - return; | ||
| 46 | } | 43 | } |
| 47 | 44 | ||
| 48 | namespace { | 45 | namespace { |
| 49 | -TORCH_LIBRARY_FRAGMENT(atb, m) | 46 | +TORCH_LIBRARY_FRAGMENT(atb, m) { |
| 50 | -{ | 47 | + m.def("_npu_group_topk(Tensor self, int k=0, int group_num=1, int n=1) -> ()"); |
| 51 | - m.def("_npu_group_topk(Tensor self, int k=0, int group_num=1, int n=1) -> ()"); | ||
| 52 | -} | ||
| 53 | } | 48 | } |
| 49 | +} // namespace | ||
| 54 | 50 | ||
| 55 | namespace { | 51 | namespace { |
| 56 | -TORCH_LIBRARY_IMPL(atb, PrivateUse1, m) | 52 | +TORCH_LIBRARY_IMPL(atb, PrivateUse1, m) { |
| 57 | -{ | 53 | + m.impl("_npu_group_topk", TORCH_FN(atb::_npu_group_topk)); |
| 58 | - m.impl("_npu_group_topk", TORCH_FN(atb::_npu_group_topk)); | ||
| 59 | -} | ||
| 60 | } | 54 | } |
| 55 | +} // namespace | ||
| 61 | 56 | ||
| 62 | } // namespace atb | 57 | } // namespace atb |
| @@ -20,39 +20,33 @@ using namespace std; | |||
| 20 | namespace atb { | 20 | namespace atb { |
| 21 | 21 | ||
| 22 | using LinearParam = atb::infer::LinearParam; | 22 | using LinearParam = atb::infer::LinearParam; |
| 23 | -void _npu_matmul_add_fp32(const at::Tensor &x, const at::Tensor &weight, at::Tensor & C) | 23 | +void _npu_matmul_add_fp32(const at::Tensor& x, const at::Tensor& weight, at::Tensor& C) { |
| 24 | -{ | 24 | + const c10::OptionalDeviceGuard device_guard(device_of(x)); |
| 25 | - const c10::OptionalDeviceGuard device_guard(device_of(x)); | 25 | + OpParamCache<LinearParam>& linearParamCache = OpParamCache<LinearParam>::getInstance(); |
| 26 | - OpParamCache<LinearParam>& linearParamCache = OpParamCache<LinearParam>::getInstance(); | 26 | + LinearParam linearParam; |
| 27 | - LinearParam linearParam; | 27 | + linearParam.transposeA = true; // 是否转置A矩阵 |
| 28 | - linearParam.transposeA = true; // 是否转置A矩阵 | 28 | + linearParam.transposeB = false; // 是否转置B矩阵 |
| 29 | - linearParam.transposeB = false; // 是否转置B矩阵 | 29 | + linearParam.hasBias = false; |
| 30 | - linearParam.hasBias = false; | 30 | + linearParam.enAccum = true; |
| 31 | - linearParam.enAccum = true; | ||
| 32 | 31 | ||
| 33 | - auto opLinear = linearParamCache.getOperation(linearParam, "LinearOperation"); | 32 | + auto opLinear = linearParamCache.getOperation(linearParam, "LinearOperation"); |
| 34 | - ParamSetter paramsetter; | 33 | + ParamSetter paramsetter; |
| 35 | - paramsetter.Input(x) | 34 | + paramsetter.Input(x).Input(weight).Input(C).Output(C); |
| 36 | - .Input(weight) | ||
| 37 | - .Input(C) | ||
| 38 | - .Output(C); | ||
| 39 | 35 | ||
| 40 | - RunAtbCmd(opLinear, paramsetter, "LinearOperation"); | 36 | + RunAtbCmd(opLinear, paramsetter, "LinearOperation"); |
| 41 | - return ; | 37 | + return; |
| 42 | } | 38 | } |
| 43 | 39 | ||
| 44 | namespace { | 40 | namespace { |
| 45 | -TORCH_LIBRARY_FRAGMENT(atb, m) | 41 | +TORCH_LIBRARY_FRAGMENT(atb, m) { |
| 46 | -{ | 42 | + m.def("_npu_matmul_add_fp32(Tensor x, Tensor weight, Tensor(a!) C) -> ()"); |
| 47 | - m.def("_npu_matmul_add_fp32(Tensor x, Tensor weight, Tensor(a!) C) -> ()"); | ||
| 48 | -} | ||
| 49 | } | 43 | } |
| 44 | +} // namespace | ||
| 50 | 45 | ||
| 51 | namespace { | 46 | namespace { |
| 52 | -TORCH_LIBRARY_IMPL(atb, PrivateUse1, m) | 47 | +TORCH_LIBRARY_IMPL(atb, PrivateUse1, m) { |
| 53 | -{ | 48 | + m.impl("_npu_matmul_add_fp32", TORCH_FN(atb::_npu_matmul_add_fp32)); |
| 54 | - m.impl("_npu_matmul_add_fp32", TORCH_FN(atb::_npu_matmul_add_fp32)); | ||
| 55 | -} | ||
| 56 | } | 49 | } |
| 50 | +} // namespace | ||
| 57 | 51 | ||
| 58 | } // namespace atb | 52 | } // namespace atb |
| @@ -1,108 +1,216 @@ | |||
| 1 | 1 | ||
| 2 | 2 | ||
| 3 | 3 | ||
| 4 | - | ||
| 5 | namespace atb { | 4 | namespace atb { |
| 6 | namespace { | 5 | namespace { |
| 7 | std::unordered_map<c10::string_view, uint16_t> cache_mode_map = { | 6 | std::unordered_map<c10::string_view, uint16_t> cache_mode_map = { |
| 8 | {"krope_ctkv", 1}, | 7 | {"krope_ctkv", 1}, |
| 9 | {"int8_nzcache", 2}, | 8 | {"int8_nzcache", 2}, |
| 10 | - {"nzcache", 3} | 9 | + {"nzcache", 3}}; |
| 11 | -}; | ||
| 12 | 10 | ||
| 13 | std::unordered_map<c10::string_view, uint16_t> quant_mode_map = { | 11 | std::unordered_map<c10::string_view, uint16_t> quant_mode_map = { |
| 14 | {"per_tensor_quant_asymm", 0}, | 12 | {"per_tensor_quant_asymm", 0}, |
| 15 | {"per_token_quant_symm", 1}, | 13 | {"per_token_quant_symm", 1}, |
| 16 | }; | 14 | }; |
| 15 | +} // namespace | ||
| 16 | + | ||
| 17 | +std::tuple<at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&> npu_mla_preprocess_out( | ||
| 18 | + const at::Tensor& input, | ||
| 19 | + const at::Tensor& gamma0, | ||
| 20 | + const at::Tensor& beta0, | ||
| 21 | + const at::Tensor& wdqkv, | ||
| 22 | + const at::Tensor& descale0, | ||
| 23 | + const at::Tensor& gamma1, | ||
| 24 | + const at::Tensor& beta1, | ||
| 25 | + const at::Tensor& wuq, | ||
| 26 | + const at::Tensor& descale1, | ||
| 27 | + const at::Tensor& gamma2, | ||
| 28 | + const at::Tensor& cos, | ||
| 29 | + const at::Tensor& sin, | ||
| 30 | + const at::Tensor& wuk, | ||
| 31 | + const at::Tensor& kv_cache, | ||
| 32 | + const at::Tensor& kv_cache_rope, | ||
| 33 | + const at::Tensor& slotmapping, | ||
| 34 | + const c10::optional<at::Tensor>& quant_scale0, | ||
| 35 | + const c10::optional<at::Tensor>& quant_offset0, | ||
| 36 | + const c10::optional<at::Tensor>& bias0, | ||
| 37 | + const c10::optional<at::Tensor>& quant_scale1, | ||
| 38 | + const c10::optional<at::Tensor>& quant_offset1, | ||
| 39 | + const c10::optional<at::Tensor>& bias1, | ||
| 40 | + const c10::optional<at::Tensor>& ctkv_scale, | ||
| 41 | + const c10::optional<at::Tensor>& q_nope_scale, | ||
| 42 | + c10::optional<c10::string_view> cache_mode_opt, | ||
| 43 | + c10::optional<c10::string_view> quant_mode_opt, | ||
| 44 | + at::Tensor& q_out0, | ||
| 45 | + at::Tensor& kv_cache_out0, | ||
| 46 | + at::Tensor& q_out1, | ||
| 47 | + at::Tensor& kv_cache_out1) { | ||
| 48 | + const c10::OptionalDeviceGuard device_guard(device_of(input)); | ||
| 49 | + uint32_t wdq_dim = 0; | ||
| 50 | + uint32_t q_rope_dim = 0; | ||
| 51 | + uint32_t k_rope_dim = 0; | ||
| 52 | + float epsilon = 1e-5; | ||
| 53 | + uint32_t q_rotary_coeff = 2; | ||
| 54 | + uint32_t k_rotary_coeff = 2; | ||
| 55 | + bool transpose_wdq = true; | ||
| 56 | + bool transpose_wuq = true; | ||
| 57 | + bool transpose_wuk = true; | ||
| 58 | + auto cache_mode = atb::utils::get_op_mode(cache_mode_map, cache_mode_opt, "krope_ctkv", "cache_mode"); | ||
| 59 | + auto quant_mode = atb::utils::get_op_mode(quant_mode_map, quant_mode_opt, "per_token_quant_symm", "quant_mode"); | ||
| 60 | + EXEC_ATB_CMD( | ||
| 61 | + AtbMLAPreprocess, | ||
| 62 | + input, | ||
| 63 | + gamma0, | ||
| 64 | + beta0, | ||
| 65 | + quant_scale0, | ||
| 66 | + quant_offset0, | ||
| 67 | + wdqkv, | ||
| 68 | + descale0, | ||
| 69 | + bias0, | ||
| 70 | + gamma1, | ||
| 71 | + beta1, | ||
| 72 | + quant_scale1, | ||
| 73 | + quant_offset1, | ||
| 74 | + wuq, | ||
| 75 | + descale1, | ||
| 76 | + bias1, | ||
| 77 | + gamma2, | ||
| 78 | + cos, | ||
| 79 | + sin, | ||
| 80 | + wuk, | ||
| 81 | + kv_cache, | ||
| 82 | + kv_cache_rope, | ||
| 83 | + slotmapping, | ||
| 84 | + ctkv_scale, | ||
| 85 | + q_nope_scale, | ||
| 86 | + wdq_dim, | ||
| 87 | + q_rope_dim, | ||
| 88 | + k_rope_dim, | ||
| 89 | + epsilon, | ||
| 90 | + q_rotary_coeff, | ||
| 91 | + k_rotary_coeff, | ||
| 92 | + transpose_wdq, | ||
| 93 | + transpose_wuq, | ||
| 94 | + transpose_wuk, | ||
| 95 | + cache_mode, | ||
| 96 | + quant_mode, | ||
| 97 | + q_out0, | ||
| 98 | + kv_cache_out0, | ||
| 99 | + q_out1, | ||
| 100 | + kv_cache_out1); | ||
| 101 | + return std::forward_as_tuple(q_out0, kv_cache_out0, q_out1, kv_cache_out1); | ||
| 17 | } | 102 | } |
| 18 | 103 | ||
| 19 | -std::tuple<at::Tensor&, at::Tensor&, at::Tensor&, at::Tensor&> npu_mla_preprocess_out(const at::Tensor &input, const at::Tensor &gamma0, const at::Tensor &beta0, | 104 | +std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_mla_preprocess( |
| 20 | - const at::Tensor &wdqkv, const at::Tensor &descale0, const at::Tensor &gamma1, const at::Tensor &beta1, | 105 | + const at::Tensor& input, |
| 21 | - const at::Tensor &wuq, const at::Tensor &descale1, const at::Tensor &gamma2, const at::Tensor &cos, const at::Tensor &sin, const at::Tensor &wuk, | 106 | + const at::Tensor& gamma0, |
| 22 | - const at::Tensor &kv_cache, const at::Tensor &kv_cache_rope, const at::Tensor &slotmapping, | 107 | + const at::Tensor& beta0, |
| 23 | - const c10::optional<at::Tensor> &quant_scale0, const c10::optional<at::Tensor> &quant_offset0, const c10::optional<at::Tensor> &bias0, | 108 | + const at::Tensor& wdqkv, |
| 24 | - const c10::optional<at::Tensor> &quant_scale1, const c10::optional<at::Tensor> &quant_offset1, const c10::optional<at::Tensor> &bias1, | 109 | + const at::Tensor& descale0, |
| 25 | - const c10::optional<at::Tensor> &ctkv_scale, const c10::optional<at::Tensor> &q_nope_scale, | 110 | + const at::Tensor& gamma1, |
| 26 | - c10::optional<c10::string_view> cache_mode_opt, c10::optional<c10::string_view> quant_mode_opt, | 111 | + const at::Tensor& beta1, |
| 27 | - at::Tensor &q_out0, | 112 | + const at::Tensor& wuq, |
| 28 | - at::Tensor &kv_cache_out0, | 113 | + const at::Tensor& descale1, |
| 29 | - at::Tensor &q_out1, | 114 | + const at::Tensor& gamma2, |
| 30 | - at::Tensor &kv_cache_out1) | 115 | + const at::Tensor& cos, |
| 31 | -{ | 116 | + const at::Tensor& sin, |
| 32 | - const c10::OptionalDeviceGuard device_guard(device_of(input)); | 117 | + const at::Tensor& wuk, |
| 33 | - uint32_t wdq_dim = 0; | 118 | + const at::Tensor& kv_cache, |
| 34 | - uint32_t q_rope_dim = 0; | 119 | + const at::Tensor& kv_cache_rope, |
| 35 | - uint32_t k_rope_dim = 0; | 120 | + const at::Tensor& slotmapping, |
| 36 | - float epsilon = 1e-5; | 121 | + const c10::optional<at::Tensor>& quant_scale0, |
| 37 | - uint32_t q_rotary_coeff = 2; | 122 | + const c10::optional<at::Tensor>& quant_offset0, |
| 38 | - uint32_t k_rotary_coeff = 2; | 123 | + const c10::optional<at::Tensor>& bias0, |
| 39 | - bool transpose_wdq = true; | 124 | + const c10::optional<at::Tensor>& quant_scale1, |
| 40 | - bool transpose_wuq = true; | 125 | + const c10::optional<at::Tensor>& quant_offset1, |
| 41 | - bool transpose_wuk = true; | 126 | + const c10::optional<at::Tensor>& bias1, |
| 42 | - auto cache_mode = atb::utils::get_op_mode(cache_mode_map, cache_mode_opt, "krope_ctkv", "cache_mode"); | 127 | + const c10::optional<at::Tensor>& ctkv_scale, |
| 43 | - auto quant_mode = atb::utils::get_op_mode(quant_mode_map, quant_mode_opt, "per_token_quant_symm", "quant_mode"); | 128 | + const c10::optional<at::Tensor>& q_nope_scale, |
| 44 | - EXEC_ATB_CMD(AtbMLAPreprocess, input, gamma0, beta0, quant_scale0, quant_offset0, wdqkv, descale0, bias0, gamma1, beta1, quant_scale1, quant_offset1, wuq, descale1, bias1, gamma2, cos, sin, wuk, kv_cache, | 129 | + c10::optional<c10::string_view> cache_mode_opt, |
| 45 | - kv_cache_rope, slotmapping, ctkv_scale, q_nope_scale, | 130 | + c10::optional<c10::string_view> quant_mode_opt) { |
| 46 | - wdq_dim, q_rope_dim, k_rope_dim, epsilon, q_rotary_coeff, k_rotary_coeff, transpose_wdq, transpose_wuq, transpose_wuk, cache_mode, quant_mode, | 131 | + const c10::OptionalDeviceGuard device_guard(device_of(input)); |
| 47 | - q_out0, kv_cache_out0, q_out1, kv_cache_out1); | 132 | + uint32_t wdq_dim = 0; |
| 48 | - return std::forward_as_tuple(q_out0, kv_cache_out0, q_out1, kv_cache_out1); | 133 | + uint32_t q_rope_dim = 0; |
| 49 | -} | 134 | + uint32_t k_rope_dim = 0; |
| 50 | - | 135 | + float epsilon = 1e-5; |
| 51 | -std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> npu_mla_preprocess(const at::Tensor &input, const at::Tensor &gamma0, const at::Tensor &beta0, | 136 | + uint32_t q_rotary_coeff = 2; |
| 52 | - const at::Tensor &wdqkv, const at::Tensor &descale0, const at::Tensor &gamma1, const at::Tensor &beta1, | 137 | + uint32_t k_rotary_coeff = 2; |
| 53 | - const at::Tensor &wuq, const at::Tensor &descale1, const at::Tensor &gamma2, const at::Tensor &cos, const at::Tensor &sin, const at::Tensor &wuk, | 138 | + bool transpose_wdq = true; |
| 54 | - const at::Tensor &kv_cache, const at::Tensor &kv_cache_rope, const at::Tensor &slotmapping, | 139 | + bool transpose_wuq = true; |
| 55 | - const c10::optional<at::Tensor> &quant_scale0, const c10::optional<at::Tensor> &quant_offset0, const c10::optional<at::Tensor> &bias0, | 140 | + bool transpose_wuk = true; |
| 56 | - const c10::optional<at::Tensor> &quant_scale1, const c10::optional<at::Tensor> &quant_offset1, const c10::optional<at::Tensor> &bias1, | 141 | + int token_num = input.size(0); |
| 57 | - const c10::optional<at::Tensor> &ctkv_scale, const c10::optional<at::Tensor> &q_nope_scale, | 142 | + int head_num = wuk.size(0); |
| 58 | - c10::optional<c10::string_view> cache_mode_opt, c10::optional<c10::string_view> quant_mode_opt) | 143 | + at::Tensor q_out0 = at::empty({token_num, head_num, 512}, kv_cache.options()); |
| 59 | -{ | 144 | + at::Tensor kv_cache_out0; |
| 60 | - const c10::OptionalDeviceGuard device_guard(device_of(input)); | 145 | + at::Tensor q_out1 = at::empty({token_num, head_num, 64}, input.options()); |
| 61 | - uint32_t wdq_dim = 0; | 146 | + at::Tensor kv_cache_out1; |
| 62 | - uint32_t q_rope_dim = 0; | 147 | + auto cache_mode = atb::utils::get_op_mode(cache_mode_map, cache_mode_opt, "krope_ctkv", "cache_mode"); |
| 63 | - uint32_t k_rope_dim = 0; | 148 | + auto quant_mode = atb::utils::get_op_mode(quant_mode_map, quant_mode_opt, "per_token_quant_symm", "quant_mode"); |
| 64 | - float epsilon = 1e-5; | 149 | + if (cache_mode == 2 || cache_mode == 3) { |
| 65 | - uint32_t q_rotary_coeff = 2; | 150 | + kv_cache_out0 = at_npu::native::empty_with_format(kv_cache.sizes(), kv_cache.options(), ACL_FORMAT_FRACTAL_NZ); |
| 66 | - uint32_t k_rotary_coeff = 2; | 151 | + kv_cache_out1 = |
| 67 | - bool transpose_wdq = true; | 152 | + at_npu::native::empty_with_format(kv_cache_rope.sizes(), kv_cache_rope.options(), ACL_FORMAT_FRACTAL_NZ); |
| 68 | - bool transpose_wuq = true; | 153 | + } else { |
| 69 | - bool transpose_wuk = true; | 154 | + kv_cache_out0 = at::empty(kv_cache.sizes(), kv_cache.options()); |
| 70 | - int token_num = input.size(0); | 155 | + kv_cache_out1 = at::empty(kv_cache_rope.sizes(), kv_cache_rope.options()); |
| 71 | - int head_num = wuk.size(0); | 156 | + } |
| 72 | - at::Tensor q_out0 = at::empty({token_num, head_num, 512}, kv_cache.options()); | 157 | + EXEC_ATB_CMD( |
| 73 | - at::Tensor kv_cache_out0; | 158 | + AtbMLAPreprocess, |
| 74 | - at::Tensor q_out1 = at::empty({token_num, head_num, 64}, input.options()); | 159 | + input, |
| 75 | - at::Tensor kv_cache_out1; | 160 | + gamma0, |
| 76 | - auto cache_mode = atb::utils::get_op_mode(cache_mode_map, cache_mode_opt, "krope_ctkv", "cache_mode"); | 161 | + beta0, |
| 77 | - auto quant_mode = atb::utils::get_op_mode(quant_mode_map, quant_mode_opt, "per_token_quant_symm", "quant_mode"); | 162 | + quant_scale0, |
| 78 | - if (cache_mode == 2 || cache_mode == 3) { | 163 | + quant_offset0, |
| 79 | - kv_cache_out0 = at_npu::native::empty_with_format(kv_cache.sizes(), kv_cache.options(), ACL_FORMAT_FRACTAL_NZ); | 164 | + wdqkv, |
| 80 | - kv_cache_out1 = at_npu::native::empty_with_format(kv_cache_rope.sizes(), kv_cache_rope.options(), ACL_FORMAT_FRACTAL_NZ); | 165 | + descale0, |
| 81 | - } else { | 166 | + bias0, |
| 82 | - kv_cache_out0 = at::empty(kv_cache.sizes(), kv_cache.options()); | 167 | + gamma1, |
| 83 | - kv_cache_out1 = at::empty(kv_cache_rope.sizes(), kv_cache_rope.options()); | 168 | + beta1, |
| 84 | - } | 169 | + quant_scale1, |
| 85 | - EXEC_ATB_CMD(AtbMLAPreprocess, input, gamma0, beta0, quant_scale0, quant_offset0, wdqkv, descale0, bias0, gamma1, beta1, quant_scale1, quant_offset1, wuq, descale1, bias1, gamma2, cos, sin, wuk, kv_cache, | 170 | + quant_offset1, |
| 86 | - kv_cache_rope, slotmapping, ctkv_scale, q_nope_scale, | 171 | + wuq, |
| 87 | - wdq_dim, q_rope_dim, k_rope_dim, epsilon, q_rotary_coeff, k_rotary_coeff, transpose_wdq, transpose_wuq, transpose_wuk, cache_mode, quant_mode, | 172 | + descale1, |
| 88 | - q_out0, kv_cache_out0, q_out1, kv_cache_out1); | 173 | + bias1, |
| 89 | - return std::make_tuple(q_out0, kv_cache_out0, q_out1, kv_cache_out1); | 174 | + gamma2, |
| 90 | -} | 175 | + cos, |
| 91 | - | 176 | + sin, |
| 92 | - | 177 | + wuk, |
| 93 | -namespace { | 178 | + kv_cache, |
| 94 | -TORCH_LIBRARY_FRAGMENT(atb, m) | 179 | + kv_cache_rope, |
| 95 | -{ | 180 | + slotmapping, |
| 96 | - m.def("npu_mla_preprocess(Tensor input, Tensor gamma0, Tensor beta0, Tensor wdqkv, Tensor descale0, Tensor gamma1, Tensor beta1, Tensor wuq, Tensor descale1, Tensor gamma2, Tensor cos, Tensor sin, Tensor wuk, Tensor kv_cache, Tensor kv_cache_rope, Tensor slotmapping, *, Tensor? quant_scale0=None, Tensor? quant_offset0=None, Tensor? bias0=None, Tensor? quant_scale1=None, Tensor? quant_offset1=None, Tensor? bias1=None, Tensor? ctkv_scale=None, Tensor? q_nope_scale=None, str? cache_mode=None, str? quant_mode=None) -> (Tensor, Tensor, Tensor, Tensor)"); | 181 | + ctkv_scale, |
| 97 | - m.def("npu_mla_preprocess.out(Tensor input, Tensor gamma0, Tensor beta0, Tensor wdqkv, Tensor descale0, Tensor gamma1, Tensor beta1, Tensor wuq, Tensor descale1, Tensor gamma2, Tensor cos, Tensor sin, Tensor wuk, Tensor kv_cache, Tensor kv_cache_rope, Tensor slotmapping, *, Tensor? quant_scale0=None, Tensor? quant_offset0=None, Tensor? bias0=None, Tensor? quant_scale1=None, Tensor? quant_offset1=None, Tensor? bias1=None, Tensor? ctkv_scale=None, Tensor? q_nope_scale=None, str? cache_mode=None, str? quant_mode=None, Tensor(a!) q_out0, Tensor(b!) kv_cache_out0, Tensor(c!) q_out1, Tensor(d!) kv_cache_out1) -> (Tensor(a!), Tensor(b!), Tensor(c!), Tensor(d!))"); | 182 | + q_nope_scale, |
| 98 | -} | 183 | + wdq_dim, |
| 184 | + q_rope_dim, | ||
| 185 | + k_rope_dim, | ||
| 186 | + epsilon, | ||
| 187 | + q_rotary_coeff, | ||
| 188 | + k_rotary_coeff, | ||
| 189 | + transpose_wdq, | ||
| 190 | + transpose_wuq, | ||
| 191 | + transpose_wuk, | ||
| 192 | + cache_mode, | ||
| 193 | + quant_mode, | ||
| 194 | + q_out0, | ||
| 195 | + kv_cache_out0, | ||
| 196 | + q_out1, | ||
| 197 | + kv_cache_out1); | ||
| 198 | + return std::make_tuple(q_out0, kv_cache_out0, q_out1, kv_cache_out1); | ||
| 99 | } | 199 | } |
| 100 | 200 | ||
| 101 | namespace { | 201 | namespace { |
| 102 | -TORCH_LIBRARY_IMPL(atb, PrivateUse1, m) | 202 | +TORCH_LIBRARY_FRAGMENT(atb, m) { |
| 103 | -{ | 203 | + m.def( |
| 104 | - m.impl("npu_mla_preprocess", TORCH_FN(atb::npu_mla_preprocess)); | 204 | + "npu_mla_preprocess(Tensor input, Tensor gamma0, Tensor beta0, Tensor wdqkv, Tensor descale0, Tensor gamma1, Tensor beta1, Tensor wuq, Tensor descale1, Tensor gamma2, Tensor cos, Tensor sin, Tensor wuk, Tensor kv_cache, Tensor kv_cache_rope, Tensor slotmapping, *, Tensor? quant_scale0=None, Tensor? quant_offset0=None, Tensor? bias0=None, Tensor? quant_scale1=None, Tensor? quant_offset1=None, Tensor? bias1=None, Tensor? ctkv_scale=None, Tensor? q_nope_scale=None, str? cache_mode=None, str? quant_mode=None) -> (Tensor, Tensor, Tensor, Tensor)"); |
| 105 | - m.impl("npu_mla_preprocess.out", TORCH_FN(atb::npu_mla_preprocess_out)); | 205 | + m.def( |
| 106 | -} | 206 | + "npu_mla_preprocess.out(Tensor input, Tensor gamma0, Tensor beta0, Tensor wdqkv, Tensor descale0, Tensor gamma1, Tensor beta1, Tensor wuq, Tensor descale1, Tensor gamma2, Tensor cos, Tensor sin, Tensor wuk, Tensor kv_cache, Tensor kv_cache_rope, Tensor slotmapping, *, Tensor? quant_scale0=None, Tensor? quant_offset0=None, Tensor? bias0=None, Tensor? quant_scale1=None, Tensor? quant_offset1=None, Tensor? bias1=None, Tensor? ctkv_scale=None, Tensor? q_nope_scale=None, str? cache_mode=None, str? quant_mode=None, Tensor(a!) q_out0, Tensor(b!) kv_cache_out0, Tensor(c!) q_out1, Tensor(d!) kv_cache_out1) -> (Tensor(a!), Tensor(b!), Tensor(c!), Tensor(d!))"); |
| 107 | } | 207 | } |
| 208 | +} // namespace | ||
| 209 | + | ||
| 210 | +namespace { | ||
| 211 | +TORCH_LIBRARY_IMPL(atb, PrivateUse1, m) { | ||
| 212 | + m.impl("npu_mla_preprocess", TORCH_FN(atb::npu_mla_preprocess)); | ||
| 213 | + m.impl("npu_mla_preprocess.out", TORCH_FN(atb::npu_mla_preprocess_out)); | ||
| 108 | } | 214 | } |
| 215 | +} // namespace | ||
| 216 | +} // namespace atb | ||