已合并
TORCH MAIN SYNC : unified kernel macro (AMP) && checkpoint error message improvement #11520
dilililiwhy创建于 2024年5月6日
TORCH MAIN SYNC : unified kernel macro (AMP) && checkpoint error message improvement #11520
已合并
dilililiwhy创建于 2024年5月6日
refs/pull/11520/head合入到master
4 个文件变更+187-175
Mtest/npu/test_public_bindings.py+4-0
@@ -219,6 +219,7 @@ class TestPublicBindings(TestCase):
219 "Generator",219 "Generator",
220 "GeneratorType",220 "GeneratorType",
221 "get_autocast_cpu_dtype",221 "get_autocast_cpu_dtype",
222+ "get_autocast_dtype",
DY
dilililiwhydilililiwhy2024年5月6日

get_autocast_dtype、set_autocast_dtype对外公开吗;test_public_bindings.py中的temp_filter是临时过滤,后续还要整改的

这里同步的社区改动,这两个是社区新增接口不是PTA的

likedislike
Yyuhaiyan82024年5月6日

get_autocast_dtype、set_autocast_dtype对外公开吗;test_public_bindings.py中的temp_filter是临时过滤,后续还要整改的

likedislike
222 "get_autocast_ipu_dtype",223 "get_autocast_ipu_dtype",
223 "get_default_dtype",224 "get_default_dtype",
224 "get_num_interop_threads",225 "get_num_interop_threads",
@@ -268,6 +269,7 @@ class TestPublicBindings(TestCase):
268 "NumberType",269 "NumberType",
269 "OperatorInfo",270 "OperatorInfo",
270 "OptionalType",271 "OptionalType",
272+ "OutOfMemoryError",
271 "ParameterDict",273 "ParameterDict",
272 "parse_ir",274 "parse_ir",
273 "parse_schema",275 "parse_schema",
@@ -295,6 +297,7 @@ class TestPublicBindings(TestCase):
295 "set_anomaly_enabled",297 "set_anomaly_enabled",
296 "set_autocast_cache_enabled",298 "set_autocast_cache_enabled",
297 "set_autocast_cpu_dtype",299 "set_autocast_cpu_dtype",
300+ "set_autocast_dtype",
298 "set_autocast_ipu_dtype",301 "set_autocast_ipu_dtype",
299 "set_autocast_cpu_enabled",302 "set_autocast_cpu_enabled",
300 "set_autocast_ipu_enabled",303 "set_autocast_ipu_enabled",
@@ -307,6 +310,7 @@ class TestPublicBindings(TestCase):
307 "StaticModule",310 "StaticModule",
308 "Stream",311 "Stream",
309 "StreamObjType",312 "StreamObjType",
313+ "Event",
310 "StringType",314 "StringType",
311 "SUM",315 "SUM",
312 "SymFloat",316 "SymFloat",
Mtest/test_utils.py+1-1
@@ -148,7 +148,7 @@ class TestCheckpoint(TestCase):
148 chunks = 2148 chunks = 2
149 modules = list(model.children())149 modules = list(model.children())
150 out = checkpoint_sequential(modules, chunks, input_var, use_reentrant=True)150 out = checkpoint_sequential(modules, chunks, input_var, use_reentrant=True)
151- with self.assertRaisesRegex(RuntimeError, "Checkpointing is not compatible"):151+ with self.assertRaisesRegex(RuntimeError, "torch.utils.checkpoint is incompatible"):
152 torch.autograd.grad(152 torch.autograd.grad(
153 outputs=[out], grad_outputs=[torch.ones(1, 5)], inputs=[input_var], create_graph=True153 outputs=[out], grad_outputs=[torch.ones(1, 5)], inputs=[input_var], create_graph=True
154 )154 )
Mtorch_npu/csrc/aten/AutoCastOps.cpp+142-139
@@ -9,150 +9,153 @@ using namespace at::autocast;
9Banned functions9Banned functions
10*******************************/10*******************************/
11 11 
12-at::Tensor binary_cross_entropy_banned(const at::Tensor &, const at::Tensor &, const c10::optional<at::Tensor>&, int64_t) {12+at::Tensor binary_cross_entropy_banned(const at::Tensor &, const at::Tensor &, const c10::optional<at::Tensor>&, int64_t)
13- AT_ERROR("torch.nn.functional.binary_cross_entropy and torch.nn.BCELoss are unsafe to autocast.\n"13+{
14- "Many models use a sigmoid layer right before the binary cross entropy layer.\n"14+ AT_ERROR("torch.nn.functional.binary_cross_entropy and torch.nn.BCELoss are unsafe to autocast.\n"
15- "In this case, combine the two layers using torch.nn.functional.binary_cross_entropy_with_logits\n"15+ "Many models use a sigmoid layer right before the binary cross entropy layer.\n"
16- "or torch.nn.BCEWithLogitsLoss. binary_cross_entropy_with_logits and BCEWithLogits are\n"16+ "In this case, combine the two layers using torch.nn.functional.binary_cross_entropy_with_logits\n"
17- "safe to autocast.");17+ "or torch.nn.BCEWithLogitsLoss. binary_cross_entropy_with_logits and BCEWithLogits are\n"
18+ "safe to autocast.");
18}19}
19-TORCH_LIBRARY_IMPL(_, AutocastPrivateUse1, m) {20+TORCH_LIBRARY_IMPL(_, AutocastPrivateUse1, m)
20- m.fallback(torch::CppFunction::makeFallthrough());21+{
22+ m.fallback(torch::CppFunction::makeFallthrough());
21}23}
22 24 
23-TORCH_LIBRARY_IMPL(aten, AutocastPrivateUse1, m) {25+TORCH_LIBRARY_IMPL(aten, AutocastPrivateUse1, m)
24- // lower_precision_fp26+{
25- KERNEL_PRIVATEUSEONE2(_convolution, deprecated, lower_precision_fp)27+ // lower_precision_fp
26- KERNEL_PRIVATEUSEONE(_convolution, lower_precision_fp)28+ KERNEL_PRIVATEUSEONE(_convolution, deprecated, lower_precision_fp)
27- KERNEL_PRIVATEUSEONE(conv1d, lower_precision_fp)29+ KERNEL_PRIVATEUSEONE(_convolution, lower_precision_fp)
28- KERNEL_PRIVATEUSEONE(conv2d, lower_precision_fp)30+ KERNEL_PRIVATEUSEONE(conv1d, lower_precision_fp)
29- KERNEL_PRIVATEUSEONE(conv3d, lower_precision_fp)31+ KERNEL_PRIVATEUSEONE(conv2d, lower_precision_fp)
30- KERNEL_PRIVATEUSEONE(conv_tbc, lower_precision_fp)32+ KERNEL_PRIVATEUSEONE(conv3d, lower_precision_fp)
31- KERNEL_PRIVATEUSEONE(conv_transpose1d, lower_precision_fp)33+ KERNEL_PRIVATEUSEONE(conv_tbc, lower_precision_fp)
32- KERNEL_PRIVATEUSEONE2(conv_transpose2d, input, lower_precision_fp)34+ KERNEL_PRIVATEUSEONE(conv_transpose1d, lower_precision_fp)
33- KERNEL_PRIVATEUSEONE2(conv_transpose3d, input, lower_precision_fp)35+ KERNEL_PRIVATEUSEONE(conv_transpose2d, input, lower_precision_fp)
34- KERNEL_PRIVATEUSEONE(convolution, lower_precision_fp)36+ KERNEL_PRIVATEUSEONE(conv_transpose3d, input, lower_precision_fp)
35- KERNEL_PRIVATEUSEONE(cudnn_convolution, lower_precision_fp)37+ KERNEL_PRIVATEUSEONE(convolution, lower_precision_fp)
36- KERNEL_PRIVATEUSEONE(cudnn_convolution_transpose, lower_precision_fp)38+ KERNEL_PRIVATEUSEONE(cudnn_convolution, lower_precision_fp)
37- KERNEL_PRIVATEUSEONE(prelu, lower_precision_fp)39+ KERNEL_PRIVATEUSEONE(cudnn_convolution_transpose, lower_precision_fp)
38- KERNEL_PRIVATEUSEONE(addmm, lower_precision_fp)40+ KERNEL_PRIVATEUSEONE(prelu, lower_precision_fp)
39- KERNEL_PRIVATEUSEONE(addmv, lower_precision_fp)41+ KERNEL_PRIVATEUSEONE(addmm, lower_precision_fp)
40- KERNEL_PRIVATEUSEONE(addr, lower_precision_fp)42+ KERNEL_PRIVATEUSEONE(addmv, lower_precision_fp)
41- KERNEL_PRIVATEUSEONE(matmul, lower_precision_fp)43+ KERNEL_PRIVATEUSEONE(addr, lower_precision_fp)
42- KERNEL_PRIVATEUSEONE(einsum, lower_precision_fp)44+ KERNEL_PRIVATEUSEONE(matmul, lower_precision_fp)
43- KERNEL_PRIVATEUSEONE(mm, lower_precision_fp)45+ KERNEL_PRIVATEUSEONE(einsum, lower_precision_fp)
44- KERNEL_PRIVATEUSEONE(mv, lower_precision_fp)46+ KERNEL_PRIVATEUSEONE(mm, lower_precision_fp)
45- KERNEL_PRIVATEUSEONE(linear, lower_precision_fp)47+ KERNEL_PRIVATEUSEONE(mv, lower_precision_fp)
46- KERNEL_PRIVATEUSEONE(addbmm, lower_precision_fp)48+ KERNEL_PRIVATEUSEONE(linear, lower_precision_fp)
47- KERNEL_PRIVATEUSEONE(baddbmm, lower_precision_fp)49+ KERNEL_PRIVATEUSEONE(addbmm, lower_precision_fp)
48- KERNEL_PRIVATEUSEONE(bmm, lower_precision_fp)50+ KERNEL_PRIVATEUSEONE(baddbmm, lower_precision_fp)
49- KERNEL_PRIVATEUSEONE(chain_matmul, lower_precision_fp)51+ KERNEL_PRIVATEUSEONE(bmm, lower_precision_fp)
50- KERNEL_PRIVATEUSEONE(linalg_multi_dot, lower_precision_fp)52+ KERNEL_PRIVATEUSEONE(chain_matmul, lower_precision_fp)
51- KERNEL_PRIVATEUSEONE(_thnn_fused_lstm_cell, lower_precision_fp)53+ KERNEL_PRIVATEUSEONE(linalg_multi_dot, lower_precision_fp)
52- KERNEL_PRIVATEUSEONE(_thnn_fused_gru_cell, lower_precision_fp)54+ KERNEL_PRIVATEUSEONE(_thnn_fused_lstm_cell, lower_precision_fp)
53- KERNEL_PRIVATEUSEONE(lstm_cell, lower_precision_fp)55+ KERNEL_PRIVATEUSEONE(_thnn_fused_gru_cell, lower_precision_fp)
54- KERNEL_PRIVATEUSEONE(gru_cell, lower_precision_fp)56+ KERNEL_PRIVATEUSEONE(lstm_cell, lower_precision_fp)
55- KERNEL_PRIVATEUSEONE(rnn_tanh_cell, lower_precision_fp)57+ KERNEL_PRIVATEUSEONE(gru_cell, lower_precision_fp)
56- KERNEL_PRIVATEUSEONE(rnn_relu_cell, lower_precision_fp)58+ KERNEL_PRIVATEUSEONE(rnn_tanh_cell, lower_precision_fp)
57- KERNEL_PRIVATEUSEONE(_scaled_dot_product_flash_attention, lower_precision_fp)59+ KERNEL_PRIVATEUSEONE(rnn_relu_cell, lower_precision_fp)
58- KERNEL_PRIVATEUSEONE(scaled_dot_product_attention, lower_precision_fp)60+ KERNEL_PRIVATEUSEONE(_scaled_dot_product_flash_attention, lower_precision_fp)
61+ KERNEL_PRIVATEUSEONE(scaled_dot_product_attention, lower_precision_fp)
59 62 
60- // fp3263+ // fp32
61- KERNEL_PRIVATEUSEONE(acos, fp32)64+ KERNEL_PRIVATEUSEONE(acos, fp32)
62- KERNEL_PRIVATEUSEONE(asin, fp32)65+ KERNEL_PRIVATEUSEONE(asin, fp32)
63- KERNEL_PRIVATEUSEONE(cosh, fp32)66+ KERNEL_PRIVATEUSEONE(cosh, fp32)
64- KERNEL_PRIVATEUSEONE(erfinv, fp32)67+ KERNEL_PRIVATEUSEONE(erfinv, fp32)
65- KERNEL_PRIVATEUSEONE(exp, fp32)68+ KERNEL_PRIVATEUSEONE(exp, fp32)
66- KERNEL_PRIVATEUSEONE(expm1, fp32)69+ KERNEL_PRIVATEUSEONE(expm1, fp32)
67- KERNEL_PRIVATEUSEONE(log, fp32)70+ KERNEL_PRIVATEUSEONE(log, fp32)
68- KERNEL_PRIVATEUSEONE(log10, fp32)71+ KERNEL_PRIVATEUSEONE(log10, fp32)
69- KERNEL_PRIVATEUSEONE(log2, fp32)72+ KERNEL_PRIVATEUSEONE(log2, fp32)
70- KERNEL_PRIVATEUSEONE(log1p, fp32)73+ KERNEL_PRIVATEUSEONE(log1p, fp32)
71- KERNEL_PRIVATEUSEONE(reciprocal, fp32)74+ KERNEL_PRIVATEUSEONE(reciprocal, fp32)
72- KERNEL_PRIVATEUSEONE(rsqrt, fp32)75+ KERNEL_PRIVATEUSEONE(rsqrt, fp32)
73- KERNEL_PRIVATEUSEONE(sinh, fp32)76+ KERNEL_PRIVATEUSEONE(sinh, fp32)
74- KERNEL_PRIVATEUSEONE(tan, fp32)77+ KERNEL_PRIVATEUSEONE(tan, fp32)
75- KERNEL_PRIVATEUSEONE2(pow, Tensor_Scalar, fp32)78+ KERNEL_PRIVATEUSEONE(pow, Tensor_Scalar, fp32)
76- KERNEL_PRIVATEUSEONE2(pow, Tensor_Tensor, fp32)79+ KERNEL_PRIVATEUSEONE(pow, Tensor_Tensor, fp32)
77- KERNEL_PRIVATEUSEONE2(pow, Scalar, fp32)80+ KERNEL_PRIVATEUSEONE(pow, Scalar, fp32)
78- KERNEL_PRIVATEUSEONE(softplus, fp32)81+ KERNEL_PRIVATEUSEONE(softplus, fp32)
79- KERNEL_PRIVATEUSEONE(layer_norm, fp32)82+ KERNEL_PRIVATEUSEONE(layer_norm, fp32)
80- KERNEL_PRIVATEUSEONE(native_layer_norm, fp32)83+ KERNEL_PRIVATEUSEONE(native_layer_norm, fp32)
81- KERNEL_PRIVATEUSEONE(group_norm, fp32)84+ KERNEL_PRIVATEUSEONE(group_norm, fp32)
82- KERNEL_PRIVATEUSEONE2(frobenius_norm, dim, fp32)85+ KERNEL_PRIVATEUSEONE(frobenius_norm, dim, fp32)
83- KERNEL_PRIVATEUSEONE(nuclear_norm, fp32)86+ KERNEL_PRIVATEUSEONE(nuclear_norm, fp32)
84- KERNEL_PRIVATEUSEONE2(nuclear_norm, dim, fp32)87+ KERNEL_PRIVATEUSEONE(nuclear_norm, dim, fp32)
85- KERNEL_PRIVATEUSEONE(cosine_similarity, fp32)88+ KERNEL_PRIVATEUSEONE(cosine_similarity, fp32)
86- KERNEL_PRIVATEUSEONE(poisson_nll_loss, fp32)89+ KERNEL_PRIVATEUSEONE(poisson_nll_loss, fp32)
87- KERNEL_PRIVATEUSEONE(cosine_embedding_loss, fp32)90+ KERNEL_PRIVATEUSEONE(cosine_embedding_loss, fp32)
88- KERNEL_PRIVATEUSEONE(nll_loss, fp32)91+ KERNEL_PRIVATEUSEONE(nll_loss, fp32)
89- KERNEL_PRIVATEUSEONE(nll_loss2d, fp32)92+ KERNEL_PRIVATEUSEONE(nll_loss2d, fp32)
90- KERNEL_PRIVATEUSEONE(hinge_embedding_loss, fp32)93+ KERNEL_PRIVATEUSEONE(hinge_embedding_loss, fp32)
91- KERNEL_PRIVATEUSEONE(kl_div, fp32)94+ KERNEL_PRIVATEUSEONE(kl_div, fp32)
92- KERNEL_PRIVATEUSEONE(l1_loss, fp32)95+ KERNEL_PRIVATEUSEONE(l1_loss, fp32)
93- KERNEL_PRIVATEUSEONE(smooth_l1_loss, fp32)96+ KERNEL_PRIVATEUSEONE(smooth_l1_loss, fp32)
94- KERNEL_PRIVATEUSEONE(huber_loss, fp32)97+ KERNEL_PRIVATEUSEONE(huber_loss, fp32)
95- KERNEL_PRIVATEUSEONE(mse_loss, fp32)98+ KERNEL_PRIVATEUSEONE(mse_loss, fp32)
96- KERNEL_PRIVATEUSEONE(margin_ranking_loss, fp32)99+ KERNEL_PRIVATEUSEONE(margin_ranking_loss, fp32)
97- KERNEL_PRIVATEUSEONE(multilabel_margin_loss, fp32)100+ KERNEL_PRIVATEUSEONE(multilabel_margin_loss, fp32)
98- KERNEL_PRIVATEUSEONE(soft_margin_loss, fp32)101+ KERNEL_PRIVATEUSEONE(soft_margin_loss, fp32)
99- KERNEL_PRIVATEUSEONE(triplet_margin_loss, fp32)102+ KERNEL_PRIVATEUSEONE(triplet_margin_loss, fp32)
100- KERNEL_PRIVATEUSEONE(multi_margin_loss, fp32)103+ KERNEL_PRIVATEUSEONE(multi_margin_loss, fp32)
101- KERNEL_PRIVATEUSEONE(binary_cross_entropy_with_logits, fp32)104+ KERNEL_PRIVATEUSEONE(binary_cross_entropy_with_logits, fp32)
102- KERNEL_PRIVATEUSEONE(dist, fp32)105+ KERNEL_PRIVATEUSEONE(dist, fp32)
103- KERNEL_PRIVATEUSEONE(pdist, fp32)106+ KERNEL_PRIVATEUSEONE(pdist, fp32)
104- KERNEL_PRIVATEUSEONE(cdist, fp32)107+ KERNEL_PRIVATEUSEONE(cdist, fp32)
105- KERNEL_PRIVATEUSEONE(renorm, fp32)108+ KERNEL_PRIVATEUSEONE(renorm, fp32)
106- KERNEL_PRIVATEUSEONE(logsumexp, fp32)109+ KERNEL_PRIVATEUSEONE(logsumexp, fp32)
107- // fp32_set_opt_dtype110+ // fp32_set_opt_dtype
108- KERNEL_PRIVATEUSEONE(prod, fp32_set_opt_dtype)111+ KERNEL_PRIVATEUSEONE(prod, fp32_set_opt_dtype)
109- KERNEL_PRIVATEUSEONE2(prod, dim_int, fp32_set_opt_dtype)112+ KERNEL_PRIVATEUSEONE(prod, dim_int, fp32_set_opt_dtype)
110- KERNEL_PRIVATEUSEONE2(prod, dim_Dimname, fp32_set_opt_dtype)113+ KERNEL_PRIVATEUSEONE(prod, dim_Dimname, fp32_set_opt_dtype)
111- KERNEL_PRIVATEUSEONE2(softmax, int, fp32_set_opt_dtype)114+ KERNEL_PRIVATEUSEONE(softmax, int, fp32_set_opt_dtype)
112- KERNEL_PRIVATEUSEONE2(softmax, Dimname, fp32_set_opt_dtype)115+ KERNEL_PRIVATEUSEONE(softmax, Dimname, fp32_set_opt_dtype)
113- KERNEL_PRIVATEUSEONE2(log_softmax, int, fp32_set_opt_dtype)116+ KERNEL_PRIVATEUSEONE(log_softmax, int, fp32_set_opt_dtype)
114- KERNEL_PRIVATEUSEONE2(log_softmax, Dimname, fp32_set_opt_dtype)117+ KERNEL_PRIVATEUSEONE(log_softmax, Dimname, fp32_set_opt_dtype)
115- KERNEL_PRIVATEUSEONE(cumprod, fp32_set_opt_dtype)118+ KERNEL_PRIVATEUSEONE(cumprod, fp32_set_opt_dtype)
116- KERNEL_PRIVATEUSEONE2(cumprod, dimname, fp32_set_opt_dtype)119+ KERNEL_PRIVATEUSEONE(cumprod, dimname, fp32_set_opt_dtype)
117- KERNEL_PRIVATEUSEONE(cumsum, fp32_set_opt_dtype)120+ KERNEL_PRIVATEUSEONE(cumsum, fp32_set_opt_dtype)
118- KERNEL_PRIVATEUSEONE2(cumsum, dimname, fp32_set_opt_dtype)121+ KERNEL_PRIVATEUSEONE(cumsum, dimname, fp32_set_opt_dtype)
119- KERNEL_PRIVATEUSEONE(linalg_vector_norm, fp32_set_opt_dtype)122+ KERNEL_PRIVATEUSEONE(linalg_vector_norm, fp32_set_opt_dtype)
120- KERNEL_PRIVATEUSEONE(linalg_matrix_norm, fp32_set_opt_dtype)123+ KERNEL_PRIVATEUSEONE(linalg_matrix_norm, fp32_set_opt_dtype)
121- KERNEL_PRIVATEUSEONE2(linalg_matrix_norm, str_ord, fp32_set_opt_dtype)124+ KERNEL_PRIVATEUSEONE(linalg_matrix_norm, str_ord, fp32_set_opt_dtype)
122- // commenting these out because they accept an explicit (not-optional) dtype, and we shouldn't try to flip that even125+ // commenting these out because they accept an explicit (not-optional) dtype, and we shouldn't try to flip that even
123- // when autocasting.126+ // when autocasting.
124- KERNEL_PRIVATEUSEONE(sum, fp32_set_opt_dtype)127+ KERNEL_PRIVATEUSEONE(sum, fp32_set_opt_dtype)
125- KERNEL_PRIVATEUSEONE2(sum, dim_IntList, fp32_set_opt_dtype)128+ KERNEL_PRIVATEUSEONE(sum, dim_IntList, fp32_set_opt_dtype)
126- KERNEL_PRIVATEUSEONE2(sum, dim_DimnameList, fp32_set_opt_dtype)129+ KERNEL_PRIVATEUSEONE(sum, dim_DimnameList, fp32_set_opt_dtype)
127- // fp32_append_dtype130+ // fp32_append_dtype
128- // The fp32_append_dtype wrapper overrides implicit promotion behavior.131+ // The fp32_append_dtype wrapper overrides implicit promotion behavior.
129- // norm does not implicitly promote, but be aware when adding new ops to this policy.132+ // norm does not implicitly promote, but be aware when adding new ops to this policy.
130- KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_PRIVATEUSEONE(ADD_NS(norm),133+ KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_PRIVATEUSEONE(ADD_NS(norm),
131- "norm.Scalar", at::Tensor (const at::Tensor &, const c10::Scalar&),134+ "norm.Scalar", at::Tensor (const at::Tensor &, const c10::Scalar&),
132- at::Tensor (const at::Tensor &, const c10::optional<c10::Scalar>&, at::ScalarType),135+ at::Tensor (const at::Tensor &, const c10::optional<c10::Scalar>&, at::ScalarType),
133- fp32_append_dtype)136+ fp32_append_dtype)
134- KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_PRIVATEUSEONE(ADD_NS(norm), "norm.ScalarOpt_dim",137+ KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_PRIVATEUSEONE(ADD_NS(norm), "norm.ScalarOpt_dim",
135- at::Tensor (const at::Tensor &, const c10::optional<c10::Scalar>&, at::IntArrayRef, bool),138+ at::Tensor (const at::Tensor &, const c10::optional<c10::Scalar>&, at::IntArrayRef, bool),
136- at::Tensor (const at::Tensor &, const c10::optional<c10::Scalar>&, at::IntArrayRef, bool, at::ScalarType),139+ at::Tensor (const at::Tensor &, const c10::optional<c10::Scalar>&, at::IntArrayRef, bool, at::ScalarType),
137- fp32_append_dtype)140+ fp32_append_dtype)
138- KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_PRIVATEUSEONE(ADD_NS(norm), "norm.names_ScalarOpt_dim",141+ KERNEL_DIFFERENT_REDISPATCH_SIGNATURE_PRIVATEUSEONE(ADD_NS(norm), "norm.names_ScalarOpt_dim",
139- at::Tensor (const at::Tensor &, const c10::optional<c10::Scalar>&, at::DimnameList, bool),142+ at::Tensor (const at::Tensor &, const c10::optional<c10::Scalar>&, at::DimnameList, bool),
140- at::Tensor (const at::Tensor &, const c10::optional<c10::Scalar>&, at::DimnameList, bool, at::ScalarType),143+ at::Tensor (const at::Tensor &, const c10::optional<c10::Scalar>&, at::DimnameList, bool, at::ScalarType),
141- fp32_append_dtype)144+ fp32_append_dtype)
142- // promote145+ // promote
143- KERNEL_PRIVATEUSEONE(addcdiv, promote)146+ KERNEL_PRIVATEUSEONE(addcdiv, promote)
144- KERNEL_PRIVATEUSEONE(addcmul, promote)147+ KERNEL_PRIVATEUSEONE(addcmul, promote)
145- KERNEL_PRIVATEUSEONE(atan2, promote)148+ KERNEL_PRIVATEUSEONE(atan2, promote)
146- KERNEL_PRIVATEUSEONE(bilinear, promote)149+ KERNEL_PRIVATEUSEONE(bilinear, promote)
147- KERNEL_PRIVATEUSEONE(cross, promote)150+ KERNEL_PRIVATEUSEONE(cross, promote)
148- KERNEL_PRIVATEUSEONE(dot, promote)151+ KERNEL_PRIVATEUSEONE(dot, promote)
149- KERNEL_PRIVATEUSEONE(grid_sampler, promote)152+ KERNEL_PRIVATEUSEONE(grid_sampler, promote)
150- KERNEL_PRIVATEUSEONE(index_put, promote)153+ KERNEL_PRIVATEUSEONE(index_put, promote)
151- KERNEL_PRIVATEUSEONE(tensordot, promote)154+ KERNEL_PRIVATEUSEONE(tensordot, promote)
152- KERNEL_PRIVATEUSEONE(scatter_add, promote)155+ KERNEL_PRIVATEUSEONE(scatter_add, promote)
153 156 
154- m.impl(TORCH_SELECTIVE_NAME("aten::binary_cross_entropy"),157+ m.impl(TORCH_SELECTIVE_NAME("aten::binary_cross_entropy"),
155- TORCH_FN((&binary_cross_entropy_banned)));158+ TORCH_FN((&binary_cross_entropy_banned)));
156}159}
157 160 
158}161}
Mtorch_npu/csrc/utils/AutocastMode.cpp+40-35
@@ -8,45 +8,49 @@
8namespace torch_npu {8namespace torch_npu {
9namespace autocast {9namespace autocast {
10 10 
11-static PyObject* set_autocast_enabled(PyObject* _unused, PyObject* arg) {11+static PyObject* set_autocast_enabled(PyObject* _unused, PyObject* arg)
12- HANDLE_TH_ERRORS12+{
13- if (!PyBool_Check(arg)) {13+ HANDLE_TH_ERRORS
14- throw torch::TypeError("enabled must be a bool (got %s)", Py_TYPE(arg)->tp_name);14+ if (!PyBool_Check(arg)) {
15- }15+ throw torch::TypeError("enabled must be a bool (got %s)", Py_TYPE(arg)->tp_name);
16- at::autocast::set_privateuseone_enabled(arg == Py_True);16+ }
17- Py_RETURN_NONE;17+ at::autocast::set_autocast_enabled(at::kPrivateUse1, arg == Py_True);
18- END_HANDLE_TH_ERRORS18+ Py_RETURN_NONE;
19+ END_HANDLE_TH_ERRORS
19}20}
20 21 
21-static PyObject* is_autocast_enabled(PyObject* _unused, PyObject* arg) {22+static PyObject* is_autocast_enabled(PyObject* _unused, PyObject* arg)
22- HANDLE_TH_ERRORS23+{
23- if (at::autocast::is_privateuseone_enabled()) {24+ HANDLE_TH_ERRORS
24- Py_RETURN_TRUE;25+ if (at::autocast::is_autocast_enabled(at::kPrivateUse1)) {
25- } else {26+ Py_RETURN_TRUE;
26- Py_RETURN_FALSE;27+ } else {
27- }28+ Py_RETURN_FALSE;
28- END_HANDLE_TH_ERRORS29+ }
30+ END_HANDLE_TH_ERRORS
29}31}
30 32 
31-static PyObject* set_autocast_dtype(PyObject* _unused, PyObject* arg) {33+static PyObject* set_autocast_dtype(PyObject* _unused, PyObject* arg)
32- HANDLE_TH_ERRORS34+{
33- if (!THPDtype_Check(arg)) {35+ HANDLE_TH_ERRORS
34- throw torch::TypeError(36+ if (!THPDtype_Check(arg)) {
35- "dtype must be a torch.dtype (got %s)", Py_TYPE(arg)->tp_name);37+ throw torch::TypeError(
36- }38+ "dtype must be a torch.dtype (got %s)", Py_TYPE(arg)->tp_name);
37- at::ScalarType targetType = reinterpret_cast<THPDtype*>(arg)->scalar_type;39+ }
38- at::autocast::set_autocast_privateuseone_dtype(targetType);40+ at::ScalarType targetType = reinterpret_cast<THPDtype*>(arg)->scalar_type;
39- Py_RETURN_NONE;41+ at::autocast::set_autocast_dtype(at::kPrivateUse1, targetType);
40- END_HANDLE_TH_ERRORS42+ Py_RETURN_NONE;
43+ END_HANDLE_TH_ERRORS
41}44}
42 45 
43-static PyObject* get_autocast_dtype(PyObject* _unused, PyObject* arg) {46+static PyObject* get_autocast_dtype(PyObject* _unused, PyObject* arg)
44- HANDLE_TH_ERRORS47+{
45- at::ScalarType current_dtype = at::autocast::get_autocast_privateuseone_dtype();48+ HANDLE_TH_ERRORS
46- auto dtype = (PyObject*)torch::getTHPDtype(current_dtype);49+ at::ScalarType current_dtype = at::autocast::get_autocast_dtype(at::kPrivateUse1);
47- Py_INCREF(dtype);50+ auto dtype = (PyObject*)torch::getTHPDtype(current_dtype);
48- return dtype;51+ Py_INCREF(dtype);
49- END_HANDLE_TH_ERRORS52+ return dtype;
53+ END_HANDLE_TH_ERRORS
50}54}
51 55 
52// autocast methods on torch._C56// autocast methods on torch._C
@@ -57,8 +61,9 @@ static PyMethodDef methods[] = { // NOLINT
57 {"get_autocast_dtype", get_autocast_dtype, METH_NOARGS, nullptr},61 {"get_autocast_dtype", get_autocast_dtype, METH_NOARGS, nullptr},
58 {nullptr, nullptr, 0, nullptr}};62 {nullptr, nullptr, 0, nullptr}};
59 63 
60-PyMethodDef* autocast_mode_functions() {64+PyMethodDef* autocast_mode_functions()
61- return methods;65+{
66+ return methods;
62}67}
63 68 
64}69}