Merge branch 'develop' into i41

Author: co63oc

paddle/fluid/distributed/collective/deep_ep/include/event_pool.h

@@ -22,7 +22,7 @@ namespace deep_ep::detail {
 
 class EventPool {
  public:
-  EventPool() = default;
+  EventPool();
   EventPool(const EventPool&) = delete;
   EventPool(EventPool&&) = delete;
   ~EventPool();

paddle/fluid/distributed/collective/deep_ep/src/event_pool.cc

@@ -22,6 +22,16 @@ EventPool &EventPool::Instance() {
   return pool;
 }
 
+EventPool::EventPool() {
+  for (size_t i = 0; i < 1000; ++i) {
+    cudaEvent_t new_event;
+    CUDA_CHECK(cudaEventCreate(&new_event));
+
+    cudaEventRecord(new_event, 0);
+    incomplished_events_.push(new_event);
+  }
+}
+
 EventPool::~EventPool() {
   const auto &DestroyEvent = [](cudaEvent_t event) {
     cudaError_t e = cudaEventDestroy(event);

paddle/fluid/eager/auto_code_generator/generator/python_c_gen.py

@@ -503,7 +503,6 @@ def GeneratePythonCFunction(self, no_input_out_tensor=False):
         get_input_out_str = ""
         if (
             not no_input_out_tensor
-            and not forward_inplace_map
             and len(self.forward_outputs_position_map) == 1
             and next(iter(self.forward_outputs_position_map.values()))[0]
             == "Tensor"
@@ -573,6 +572,7 @@ def GeneratePythonCFunction(self, no_input_out_tensor=False):
                 namespace,
                 GetForwardFunctionName(inplaced_forward_api_name),
             )
+            dygraph_function_call_str = ",".join(dygraph_function_call_list)
 
             inplace_noamp_dygraph_function_str = (
                 NOAMP_DYGRAPH_FUNCTION_TEMPLATE.format(

paddle/fluid/framework/ir/graph_pattern_detector.cc

@@ -2342,7 +2342,9 @@ PDNode *patterns::QuantConv::operator()(const std::string &conv_type) {
   auto conv_op = pattern->NewNode(conv_op_repr())->assert_is_op(conv_type);
   conv_op->assert_more([&](Node *node) {
     return node->Op()->GetAttrIfExists<std::string>("mkldnn_data_type") ==
-           "bfloat16";
+               "bfloat16" ||
+           node->Op()->GetAttrIfExists<std::string>("onednn_data_type") ==
+               "bfloat16";
   });
 
   quant_op->LinksFrom({quant_in}).LinksTo({conv_in});
@@ -3172,7 +3174,8 @@ PDNode *patterns::QuantizePlacement::operator()(
   auto *op =
       pattern->NewNode(op_repr())->assert_is_ops(quantize_enabled_op_types);
   op->assert_more([&](Node *node) {
-    return node->Op()->GetAttrIfExists<bool>("use_mkldnn");
+    return node->Op()->GetAttrIfExists<bool>("use_mkldnn") ||
+           node->Op()->GetAttrIfExists<bool>("use_onednn");
   });
   return op;
 }
@@ -3218,6 +3221,7 @@ PDNode *patterns::Bfloat16Placement::operator()(
   auto *op = pattern->NewNode(op_repr())->assert_is_ops(supported_op_types);
   op->assert_more([&](Node *node) {
     return node->Op()->GetAttrIfExists<bool>("use_mkldnn") ||
+           node->Op()->GetAttrIfExists<bool>("use_onednn") ||
            node->Op()->Type() == "reshape2";
   });
   op->LinksFrom({op_in});
@@ -3227,25 +3231,35 @@ PDNode *patterns::Bfloat16Placement::operator()(
 PDNode *patterns::OrphanedBfloat16::operator()() {
   auto *prev_op = pattern->NewNode(prev_op_repr())->assert_is_op();
   prev_op->assert_more([&](Node *node) {
-    bool data_type_is_missing = !node->Op()->HasAttr("mkldnn_data_type");
-    bool data_type_is_fp32 = node->Op()->GetAttrIfExists<std::string>(
-                                 "mkldnn_data_type") == "float32";
+    bool data_type_is_missing = !node->Op()->HasAttr("mkldnn_data_type") &&
+                                !node->Op()->HasAttr("onednn_data_type");
+    bool data_type_is_fp32 =
+        node->Op()->GetAttrIfExists<std::string>("mkldnn_data_type") ==
+            "float32" ||
+        node->Op()->GetAttrIfExists<std::string>("onednn_data_type") ==
+            "float32";
     return data_type_is_missing || data_type_is_fp32;
   });
   auto *prev_out = pattern->NewNode(prev_out_repr())->AsOutput();
 
   auto *op = pattern->NewNode(op_repr())->assert_is_op();
   op->assert_more([&](Node *node) {
     return node->Op()->GetAttrIfExists<std::string>("mkldnn_data_type") ==
-           "bfloat16";
+               "bfloat16" ||
+           node->Op()->GetAttrIfExists<std::string>("onednn_data_type") ==
+               "bfloat16";
   });
   auto *op_out = pattern->NewNode(op_out_repr())->AsOutput();
 
   auto *next_op = pattern->NewNode(next_op_repr())->assert_is_op();
   next_op->assert_more([&](Node *node) {
-    bool data_type_is_missing = !node->Op()->HasAttr("mkldnn_data_type");
-    bool data_type_is_fp32 = node->Op()->GetAttrIfExists<std::string>(
-                                 "mkldnn_data_type") == "float32";
+    bool data_type_is_missing = !node->Op()->HasAttr("mkldnn_data_type") &&
+                                !node->Op()->HasAttr("onednn_data_type");
+    bool data_type_is_fp32 =
+        node->Op()->GetAttrIfExists<std::string>("mkldnn_data_type") ==
+            "float32" ||
+        node->Op()->GetAttrIfExists<std::string>("onednn_data_type") ==
+            "float32";
     return data_type_is_missing || data_type_is_fp32;
   });
 
@@ -3258,14 +3272,17 @@ PDNode *patterns::OrphanedBfloat16::operator()() {
 PDNode *patterns::UnsupportedBfloat16::operator()() {
   auto *prev_op = pattern->NewNode(prev_op_repr())->assert_is_op();
   prev_op->assert_more([&](Node *node) {
-    return node->Op()->HasAttr("mkldnn_data_type") == false;
+    return node->Op()->HasAttr("mkldnn_data_type") == false &&
+           node->Op()->HasAttr("onednn_data_type") == false;
   });
   auto *prev_out = pattern->NewNode(prev_out_repr())->AsOutput();
 
   auto *op = pattern->NewNode(op_repr())->assert_is_op();
   op->assert_more([&](Node *node) {
     return node->Op()->GetAttrIfExists<std::string>("mkldnn_data_type") ==
-           "bfloat16";
+               "bfloat16" ||
+           node->Op()->GetAttrIfExists<std::string>("onednn_data_type") ==
+               "bfloat16";
   });
   prev_op->LinksTo({prev_out});
   op->LinksFrom({prev_out});
@@ -3276,7 +3293,9 @@ PDNode *patterns::Bloat16Ops::operator()() {
   auto op = pattern->NewNode(op_repr())->assert_is_op();
   op->assert_more([&](Node *node) {
     return node->Op()->GetAttrIfExists<std::string>("mkldnn_data_type") ==
-           "bfloat16";
+               "bfloat16" ||
+           node->Op()->GetAttrIfExists<std::string>("onednn_data_type") ==
+               "bfloat16";
   });
   return op;
 }
@@ -3298,8 +3317,8 @@ PDNode *patterns::ONEDNNInPlace::operator()() {
   auto next_op = pattern->NewNode(next_op_repr())->assert_is_op();
   auto next_output = pattern->NewNode(next_op_out_repr())->AsOutput();
 
-  // Check if op is MKL-DNN enabled
-  possible_inplace_op->assert_op_attr("use_mkldnn", true);
+  // Check if op is ONE-DNN enabled
+  possible_inplace_op->assert_op_attr_or("use_mkldnn", "use_onednn", true);
 
   // linked structure
   possible_inplace_op->LinksTo({output});

paddle/fluid/framework/ir/graph_pattern_detector.h

@@ -168,6 +168,20 @@ struct PDNode {
     return this;
   }
 
+  template <typename T>
+  PDNode* assert_op_attr_or(const std::string& attr_name1,
+                            const std::string& attr_name2,
+                            const T& attr) {
+    asserts_.emplace_back([=](Node* x) {
+      return x && x->IsOp() &&
+             ((x->Op()->HasAttr(attr_name1) &&
+               PADDLE_GET_CONST(T, x->Op()->GetAttr(attr_name1)) == attr) ||
+              (x->Op()->HasAttr(attr_name2) &&
+               PADDLE_GET_CONST(T, x->Op()->GetAttr(attr_name2)) == attr));
+    });
+    return this;
+  }
+
  private:
   PDNode(PDPattern* pattern,
          const std::string& name = "",

paddle/phi/kernels/impl/slogdeterminant_grad_kernel_impl.h

@@ -82,8 +82,54 @@ void SlogDeterminantGradKernel(const Context& dev_ctx,
   inverse_A.Resize(x.dims());
   dev_ctx.template Alloc<T>(&inverse_A);
 
-  phi::funcs::MatrixInverseFunctor<Context, T> mat_inv;
-  mat_inv(dev_ctx, x, &inverse_A);
+  const auto& mat_dims = x.dims();
+  const int rank = mat_dims.size();
+  int n = mat_dims[rank - 1];
+  int64_t total_batch_size = rank > 2 ? x.numel() / (n * n) : 1;
+
+  // Divide the batch into chunks because of cublasMatInv limitation
+  if (total_batch_size <= 65536) {
+    phi::funcs::MatrixInverseFunctor<Context, T> mat_inv;
+    mat_inv(dev_ctx, x, &inverse_A);
+  } else {
+    constexpr int64_t max_batch_size = 65536;
+    int64_t processed = 0;
+
+    VLOG(3) << "Large batch size detected (" << total_batch_size
+            << "), processing in chunks of " << max_batch_size;
+
+    while (processed < total_batch_size) {
+      int64_t current_batch =
+          std::min(max_batch_size, total_batch_size - processed);
+
+      // Extract current batch data
+      DenseTensor x_batch;
+      x_batch.ShareDataWith(x);
+      x_batch.Resize({total_batch_size, n, n});
+      x_batch = x_batch.Slice(processed, processed + current_batch);
+      x_batch.Resize({current_batch, n, n});
+
+      DenseTensor inverse_batch;
+      inverse_batch.Resize({current_batch, n, n});
+      dev_ctx.template Alloc<T>(&inverse_batch);
+
+      // Compute the inverse matrix for the current batch
+      phi::funcs::MatrixInverseFunctor<Context, T> mat_inv;
+      mat_inv(dev_ctx, x_batch, &inverse_batch);
+
+      // Copy the result to the output tensor
+      DenseTensor output_slice;
+      output_slice.ShareDataWith(inverse_A);
+      output_slice.Resize({total_batch_size, n, n});
+      output_slice = output_slice.Slice(processed, processed + current_batch);
+      output_slice.Resize({current_batch, n, n});
+
+      phi::Copy(
+          dev_ctx, inverse_batch, dev_ctx.GetPlace(), false, &output_slice);
+
+      processed += current_batch;
+    }
+  }
 
   VLOG(3) << "inverse(A) dims: " << inverse_A.dims();
 

paddle/phi/ops/yaml/op_compat.yaml

@@ -1765,7 +1765,7 @@
   attrs :
     {scale_data : Scale_data, shift_data : Shift_data, scale_weights : Scale_weights}
   extra :
-    attrs : [bool use_mkldnn = true, bool use_onednn = false, str mkldnn_data_type = "float32"]
+    attrs : [bool use_mkldnn = false, bool use_onednn = false, str mkldnn_data_type = "float32"]
 
 - op : fusion_repeated_fc_relu
   inputs :

python/paddle/__init__.py

@@ -129,6 +129,12 @@
     tensor as tensor,
     utils as utils,
 )
+from .amp import (
+    get_autocast_cpu_dtype,
+    get_autocast_dtype,
+    get_autocast_gpu_dtype,
+    is_autocast_enabled,
+)
 from .autograd import (
     enable_grad,
     grad,
@@ -1233,6 +1239,10 @@
     'nan',
     'pi',
     'e',
+    'is_autocast_enabled',
+    'get_autocast_dtype',
+    'get_autocast_cpu_dtype',
+    'get_autocast_gpu_dtype',
 ]
 
 import os

python/paddle/amp/__init__.py

@@ -33,6 +33,8 @@
     amp_guard,
     auto_cast,
     decorate,
+    get_autocast_dtype,
+    is_autocast_enabled,
 )
 from .grad_scaler import (  # noqa: F401
     AmpScaler,
@@ -46,8 +48,15 @@
     'decorate',
     'is_float16_supported',
     'is_bfloat16_supported',
+    'is_autocast_enabled',
+    'get_autocast_dtype',
+    'get_autocast_cpu_dtype',
+    'get_autocast_gpu_dtype',
 ]
 
+get_autocast_cpu_dtype = get_autocast_dtype
+get_autocast_gpu_dtype = get_autocast_dtype
+
 
 def is_float16_supported(device: str | None = None) -> bool:
     """

python/paddle/amp/auto_cast.py

@@ -48,6 +48,7 @@
     from typing_extensions import TypeAlias, TypeGuard
 
     from paddle import Tensor
+    from paddle._typing import PlaceLike
     from paddle._typing.dtype_like import _DTypeLiteral
     from paddle.nn import Layer
     from paddle.nn.layer.layers import _StateDict
@@ -1322,3 +1323,73 @@ def decorate(
             master_grad,
             excluded_layers,
         )
+
+
+def is_autocast_enabled(device_type: PlaceLike | None = None) -> bool:
+    """
+    Check whether auto-mixed-precision is enabled in the current context.
+
+    Args:
+        device_type (PlaceLike, optional): The device type to check. This argument is ignored for all devices sharing the same AMP state in paddlepaddle.
+
+    Returns:
+        bool: True if auto-mixed-precision is enabled, False otherwise.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> # Demo1: Check if auto-mixed-precision is enabled by default
+            >>> import paddle
+            >>> paddle.device.set_device('gpu')
+            >>> print(paddle.is_autocast_enabled())
+            False
+
+            >>> # Demo2: Enable auto-mixed-precision and check again
+            >>> with paddle.amp.auto_cast():
+            ...     print(paddle.is_autocast_enabled())
+            True
+    """
+    if in_pir_mode():
+        amp_attrs = core._get_amp_attrs()
+        return amp_attrs._amp_level != AMP_LEVEL.O0
+    else:
+        tracer = _dygraph_tracer()
+        if tracer:
+            return tracer._amp_level != core.AmpLevel.O0
+        return False
+
+
+def get_autocast_dtype(device_type: PlaceLike | None = None) -> _DTypeLiteral:
+    """
+    Get the auto-mixed-precision dtype in the current context if autocast is enabled else default AMP dtype(float16).
+
+    Args:
+        device_type (PlaceLike, optional): The device type to check. This argument is ignored for all devices sharing the same AMP state in paddlepaddle.
+
+    Returns:
+        _DTypeLiteral: The current AMP dtype.
+
+    Examples:
+        .. code-block:: python
+
+            >>> # doctest: +REQUIRES(env:GPU)
+            >>> # Demo1: Get default auto-mixed-precision dtype
+            >>> import paddle
+            >>> paddle.device.set_device('gpu')
+            >>> print(paddle.get_autocast_dtype())
+            float16
+
+            >>> # Demo2: Enable auto-mixed-precision and get the dtype
+            >>> with paddle.amp.auto_cast():
+            ...     print(paddle.get_autocast_dtype())
+            float16
+    """
+    if not is_autocast_enabled():
+        return "float16"
+    if in_pir_mode():
+        amp_attrs = core._get_amp_attrs()
+        return amp_attrs._amp_dtype
+    else:
+        tracer = _dygraph_tracer()
+        return tracer._amp_dtype