已合并
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
已合并
从refs/pull/11520/head合入到master
共 4 个文件变更+187-175
| @@ -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", | ||
| 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", |
| @@ -148,7 +148,7 @@ class TestCheckpoint(TestCase): | |||
| 148 | chunks = 2 | 148 | 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=True | 153 | outputs=[out], grad_outputs=[torch.ones(1, 5)], inputs=[input_var], create_graph=True |
| 154 | ) | 154 | ) |
| @@ -9,150 +9,153 @@ using namespace at::autocast; | |||
| 9 | Banned functions | 9 | Banned 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_fp | 26 | +{ |
| 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 | - // fp32 | 63 | + // 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_dtype | 110 | + // 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 even | 125 | + // 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_dtype | 130 | + // 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 | - // promote | 145 | + // 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 | } |
| @@ -8,45 +8,49 @@ | |||
| 8 | namespace torch_npu { | 8 | namespace torch_npu { |
| 9 | namespace autocast { | 9 | namespace 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_ERRORS | 12 | +{ |
| 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_ERRORS | 18 | + 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_ERRORS | 23 | +{ |
| 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_ERRORS | 29 | + } |
| 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_ERRORS | 34 | +{ |
| 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_ERRORS | 42 | + 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_ERRORS | 47 | +{ |
| 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_ERRORS | 52 | + return dtype; |
| 53 | + END_HANDLE_TH_ERRORS | ||
| 50 | } | 54 | } |
| 51 | 55 | ||
| 52 | // autocast methods on torch._C | 56 | // 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 | } |


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