From 98ef51004fc7063be44a3121f7854653b3618299 Mon Sep 17 00:00:00 2001
From: Leo Guo <58431564+ZibinGuo@users.noreply.github.com>
Date: Thu, 25 Apr 2024 18:49:06 +0800
Subject: [PATCH] [XPU] Added a Python interface to Context::set_debug_level in
 XDNN for setting the XPUAPI_DEBUG level. (#62210)

---
 paddle/fluid/platform/device/xpu/xpu_info.cc |  4 ++++
 paddle/fluid/platform/device/xpu/xpu_info.h  |  2 ++
 paddle/fluid/pybind/place.cc                 |  2 ++
 paddle/phi/backends/xpu/xpu_info.cc          |  9 ++++++++
 paddle/phi/backends/xpu/xpu_info.h           |  1 +
 python/paddle/device/xpu/__init__.py         | 23 ++++++++++++++++++++
 6 files changed, 41 insertions(+)

diff --git a/paddle/fluid/platform/device/xpu/xpu_info.cc b/paddle/fluid/platform/device/xpu/xpu_info.cc
index cc7388df4c22f4..9919243ea664fd 100644
--- a/paddle/fluid/platform/device/xpu/xpu_info.cc
+++ b/paddle/fluid/platform/device/xpu/xpu_info.cc
@@ -94,6 +94,10 @@ phi::backends::xpu::XPUVersion get_xpu_version(int dev_id) {
   return phi::backends::xpu::get_xpu_version(dev_id);
 }
 
+void set_xpu_debug_level(int level) {
+  phi::backends::xpu::set_xpu_debug_level(level);
+}
+
 /**************************** XPU Allocator **************************/
 size_t XPUMinChunkSize() { return 1 << 6; }
 
diff --git a/paddle/fluid/platform/device/xpu/xpu_info.h b/paddle/fluid/platform/device/xpu/xpu_info.h
index 163c3a3dcc8800..14ed9d7caf208c 100644
--- a/paddle/fluid/platform/device/xpu/xpu_info.h
+++ b/paddle/fluid/platform/device/xpu/xpu_info.h
@@ -70,6 +70,8 @@ using XPUDeviceGuard = phi::backends::xpu::XPUDeviceGuard;
 
 phi::backends::xpu::XPUVersion get_xpu_version(int dev_id);
 
+void set_xpu_debug_level(int level);
+
 //! Get the minimum chunk size for XPU allocator.
 size_t XPUMinChunkSize();
 
diff --git a/paddle/fluid/pybind/place.cc b/paddle/fluid/pybind/place.cc
index e6c25413988b8a..4b22798735e2d0 100644
--- a/paddle/fluid/pybind/place.cc
+++ b/paddle/fluid/pybind/place.cc
@@ -462,6 +462,8 @@ void BindPlace(pybind11::module &m) {  // NOLINT
       .value("XPU3", phi::backends::xpu::XPUVersion::XPU3)
       .export_values();
   m.def("get_xpu_device_count", platform::GetXPUDeviceCount);
+  m.def("set_xpu_debug_level",
+        [](int level) { platform::set_xpu_debug_level(level); });
   m.def("get_xpu_device_version",
         [](int device_id) { return platform::get_xpu_version(device_id); });
 #ifdef PADDLE_WITH_XPU_KP
diff --git a/paddle/phi/backends/xpu/xpu_info.cc b/paddle/phi/backends/xpu/xpu_info.cc
index 6916c10d9f1e7a..bab545020e9c05 100644
--- a/paddle/phi/backends/xpu/xpu_info.cc
+++ b/paddle/phi/backends/xpu/xpu_info.cc
@@ -21,6 +21,8 @@ limitations under the License. */
 #include "paddle/phi/backends/xpu/xpu_header.h"
 #include "paddle/phi/common/place.h"
 
+#include "paddle/phi/api/lib/kernel_dispatch.h"
+
 // TODO(wilber): The phi computing library requires a component to manage
 // flags.
 #include "paddle/common/flags.h"
@@ -204,6 +206,13 @@ XPUVersion get_xpu_version(int dev_id) {
   }
 }
 
+void set_xpu_debug_level(int level) {
+  auto* dev_ctx =
+      paddle::experimental::GetDeviceContextByBackend(phi::Backend::XPU);
+  auto* xpu_ctx = static_cast<const phi::XPUContext*>(dev_ctx);
+  PADDLE_ENFORCE_XPU_SUCCESS(xpu_ctx->x_context()->set_debug_level(level));
+}
+
 int get_xpu_max_ptr_size(int dev_id) {
   auto xpu_version = get_xpu_version(dev_id);
   int max_ptr_size = 0;
diff --git a/paddle/phi/backends/xpu/xpu_info.h b/paddle/phi/backends/xpu/xpu_info.h
index e47a45e9b24514..a151dec26ef4e9 100644
--- a/paddle/phi/backends/xpu/xpu_info.h
+++ b/paddle/phi/backends/xpu/xpu_info.h
@@ -99,6 +99,7 @@ class XPUDeviceGuard {
 
 enum XPUVersion { XPU1, XPU2, XPU3 };
 XPUVersion get_xpu_version(int dev_id);
+void set_xpu_debug_level(int level);
 
 int get_xpu_max_ptr_size(int dev_id);
 
diff --git a/python/paddle/device/xpu/__init__.py b/python/paddle/device/xpu/__init__.py
index 706213becd78f5..de756d07a942ab 100644
--- a/python/paddle/device/xpu/__init__.py
+++ b/python/paddle/device/xpu/__init__.py
@@ -82,3 +82,26 @@ def device_count():
     )
 
     return num_xpus
+
+
+def set_debug_level(level=1):
+    '''
+    Set the debug level of XPUs' api.
+
+    Parameters:
+        int: debug level of XPUs available.
+        |level        |name       |usage
+        |0            |stop       |stop the debug mode
+        |0x1          |trace      |Print the invocation of the interface
+        |0x10         |checksum   |Print the checksum of the tensor
+        |0x100        |dump       |Save the tensor as a file in npy format
+        |0x1000       |profiling  |Record the execution time of each operator
+
+    Examples:
+        .. code-block:: python
+
+            >>> import paddle
+
+            >>> paddle.device.xpu.set_debug_level(0x1)
+    '''
+    core.set_xpu_debug_level(level)