已合并
[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
已合并
thickhair创建于 7月29日
60 个文件变更+4376-3813
@@ -35,7 +35,7 @@ BreakBeforeTernaryOperators: true
35BreakConstructorInitializersBeforeComma: false35BreakConstructorInitializersBeforeComma: false
36BreakAfterJavaFieldAnnotations: false36BreakAfterJavaFieldAnnotations: false
37BreakStringLiterals: false37BreakStringLiterals: false
38-ColumnLimit: 8038+ColumnLimit: 120
39CommentPragmas: '^ IWYU pragma:'39CommentPragmas: '^ IWYU pragma:'
40CompactNamespaces: false40CompactNamespaces: false
41ConstructorInitializerAllOnOneLineOrOnePerLine: true41ConstructorInitializerAllOnOneLineOrOnePerLine: 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 
59std::tuple<at::Tensor&, at::Tensor&, at::Tensor&> unique_consecutive_out_npu(58std::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 
104std::tuple<at::Tensor, at::Tensor, at::Tensor> unique_consecutive(101std::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_op120} // 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} // namespace78} // 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 
114at::Tensor upsample_bicubic2d_backward(111at::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_op130} // 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} // namespace70} // 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 
109at::Tensor upsample_bicubic2d(106at::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_op129} // 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} // namespace45} // 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 
72at::Tensor upsample_bilinear2d_backward(68at::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_op82} // 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} // namespace48} // 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 
83at::Tensor upsample_bilinear2d(74at::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_op93} // namespace acl_op
@@ -21,82 +21,101 @@ namespace acl_op {
21using npu_preparation = at_npu::native::OpPreparation;21using npu_preparation = at_npu::native::OpPreparation;
22 22 
23namespace {23namespace {
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} // namespace90} // 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 the108+ // Since only NCHW format input is currently supported, first convert the
90- // input grad_output (3 dimensions) to 4 dimensions as the input of npu109+ // 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_op121} // namespace acl_op
@@ -22,30 +22,29 @@ using npu_preparation = at_npu::native::OpPreparation;
22using npu_utils = at_npu::native::NpuUtils;22using npu_utils = at_npu::native::NpuUtils;
23 23 
24namespace {24namespace {
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 
51at::Tensor& upsample_linear1d_out_nocheck(50at::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 the58+ // input self (3 dimensions) to 4 dimensions as the input of npu
60- // input self (3 dimensions) to 4 dimensions as the input of npu59+ 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 scale62+ 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} // namespace84} // 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 
115at::Tensor upsample_linear1d(107at::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_op119} // namespace acl_op
@@ -22,19 +22,19 @@ using npu_preparation = at_npu::native::OpPreparation;
22using npu_utils = at_npu::native::NpuUtils;22using npu_utils = at_npu::native::NpuUtils;
23 23 
24namespace {24namespace {
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 
40at::Tensor& upsample_nearest1d_backward_out_nocheck(40at::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 Resize69+ .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} // namespace80} // 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 
108at::Tensor upsample_nearest1d_backward(102at::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_op113} // namespace acl_op
@@ -22,77 +22,84 @@ using npu_preparation = at_npu::native::OpPreparation;
22using npu_utils = at_npu::native::NpuUtils;22using npu_utils = at_npu::native::NpuUtils;
23 23 
24namespace {24namespace {
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} // namespace76} // 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_op105} // namespace acl_op
@@ -22,54 +22,65 @@ using npu_preparation = at_npu::native::OpPreparation;
22using npu_utils = at_npu::native::NpuUtils;22using npu_utils = at_npu::native::NpuUtils;
23 23 
24namespace {24namespace {
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} // namespace50} // 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_op86} // namespace acl_op
@@ -22,64 +22,74 @@ using npu_preparation = at_npu::native::OpPreparation;
22using npu_utils = at_npu::native::NpuUtils;22using npu_utils = at_npu::native::NpuUtils;
23 23 
24namespace {24namespace {
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} // namespace61} // 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_op95} // 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 
53at::Tensor& upsample_trilinear3d_backward_out_nocheck(53at::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} // namespace73} // 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 
106at::Tensor upsample_trilinear3d_backward(100at::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_op115} // namespace acl_op
@@ -22,66 +22,80 @@ using npu_preparation = at_npu::native::OpPreparation;
22using npu_utils = at_npu::native::NpuUtils;22using npu_utils = at_npu::native::NpuUtils;
23 23 
24namespace {24namespace {
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} // namespace66} // 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_op101} // namespace acl_op
@@ -21,44 +21,33 @@ using npu_preparation = at_npu::native::OpPreparation;
21using npu_utils = at_npu::native::NpuUtils;21using npu_utils = at_npu::native::NpuUtils;
22 22 
23namespace {23namespace {
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} // namespace30} // 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_plugin53+} // 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 
40at::Tensor var(39at::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 
54std::tuple<at::Tensor, at::Tensor> var_mean(52std::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_op65} // 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 same25// 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} // namespace77} // 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_op121} // namespace acl_op
@@ -36,11 +36,11 @@ at::SmallVector<int64_t, SIZE> where_npu_output_size(const at::Tensor& condition
36std::vector<at::Tensor> where(const at::Tensor& condition) {36std::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_op88} // namespace acl_op
@@ -21,118 +21,107 @@ using npu_preparation = at_npu::native::OpPreparation;
21using npu_utils = at_npu::native::NpuUtils;21using npu_utils = at_npu::native::NpuUtils;
22 22 
23namespace {23namespace {
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} // namespace41} // 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_op127} // namespace acl_op
@@ -25,34 +25,39 @@ namespace {
25inline void yolo_boxes_encode_check(25inline 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} // namespace62} // 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_op83} // namespace acl_op
@@ -18,10 +18,9 @@
18#include "op_plugin/utils/custom_functions/aclops/inner_compute.h"18#include "op_plugin/utils/custom_functions/aclops/inner_compute.h"
19 19 
20namespace acl_op {20namespace 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 
27at::Tensor zeros_symint(26at::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- 
38at::Tensor zeros(35at::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_op44} // namespace acl_op
@@ -21,28 +21,23 @@ using npu_preparation = at_npu::native::OpPreparation;
21using npu_utils = at_npu::native::NpuUtils;21using npu_utils = at_npu::native::NpuUtils;
22 22 
23namespace {23namespace {
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} // namespace30} // 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_op43} // namespace acl_op
@@ -17,56 +17,103 @@
17 17 
18namespace atb {18namespace atb {
19namespace {19namespace {
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 
57namespace {104namespace {
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 
65namespace {113namespace {
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;
20namespace atb {20namespace atb {
21 21 
22using GroupTopkParam = atb::infer::GroupTopkParam;22using 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 
48namespace {45namespace {
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 
55namespace {51namespace {
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 atb57} // namespace atb
@@ -20,39 +20,33 @@ using namespace std;
20namespace atb {20namespace atb {
21 21 
22using LinearParam = atb::infer::LinearParam;22using 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 
44namespace {40namespace {
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 
51namespace {46namespace {
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 atb52} // namespace atb
@@ -1,108 +1,216 @@
1#include "op_plugin/OpApiInterface.h"1#include "op_plugin/OpApiInterface.h"
2#include "op_plugin/utils/custom_functions/atb/AtbCommon.h"2#include "op_plugin/utils/custom_functions/atb/AtbCommon.h"
3 3 
4- 
5namespace atb {4namespace atb {
6namespace {5namespace {
7std::unordered_map<c10::string_view, uint16_t> cache_mode_map = {6std::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 
13std::unordered_map<c10::string_view, uint16_t> quant_mode_map = {11std::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 
101namespace {201namespace {
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