fix data_op device for gpu pinned tensor (#60357)

PaddlePaddle · Dec 27, 2023 · 430894e · 430894e
1 parent 9952846
commit 430894e
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 18 deletions.
diff --git a/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc b/paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
@@ -689,7 +689,12 @@ phi::KernelKey GetKernelKey(
     auto data_place =
         op->attributes().at("place").dyn_cast<PlaceAttribute>().data();
 
-    auto backend = paddle::experimental::ParseBackend(data_place);
+    phi::Backend backend;
+    if (data_place.GetType() == AllocationType::GPUPINNED) {
+      backend = phi::Backend::CPU;
+    } else {
+      backend = paddle::experimental::ParseBackend(data_place);
+    }
 
     return {backend,
             phi::DataLayout::ANY,

diff --git a/python/paddle/jit/sot/opcode_translator/executor/function_graph.py b/python/paddle/jit/sot/opcode_translator/executor/function_graph.py
@@ -242,9 +242,7 @@ def collect(inp):
     def guard_fn(self) -> Guard:
         with tmp_name_guard():
             guards = []
-            with EventGuard(
-                "guard_fn: find vars and make stringify guard", event_level=1
-            ):
+            with EventGuard("guard_fn: find vars and make stringify guard"):
                 for variable in find_traceable_vars(
                     self.input_variables + list(self._global_guarded_variables)
                 ):

diff --git a/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py b/python/paddle/jit/sot/opcode_translator/executor/opcode_executor.py
@@ -562,7 +562,7 @@ def step(self, instr: Instruction):
 
         opname = instr.opname if instr.opname != "PRECALL" else "PRECALL__CALL"
         assert opname != "CALL", "CALL should fused with PRECALL"
-        with EventGuard(f"{opname}", event_level=1):
+        with EventGuard(f"{opname}", event_level=2):
             return getattr(self, opname)(instr)  # run single step.
 
     def indexof(self, instr: Instruction):

diff --git a/python/paddle/jit/sot/profiler.py b/python/paddle/jit/sot/profiler.py
@@ -18,7 +18,7 @@
 
 from paddle.framework import core
 
-_event_level = int(os.environ.get("EVENT_LEVEL", "-1"))
+_event_level = int(os.environ.get("EVENT_LEVEL", "0"))
 
 
 class SotProfiler:
@@ -37,7 +37,7 @@ def disable(self):
 
 
 @contextmanager
-def EventGuard(event_name, event_level=0):
+def EventGuard(event_name, event_level=1):
     try:
         global _event_level
         need_pop = False
@@ -50,20 +50,11 @@ def EventGuard(event_name, event_level=0):
             core.nvprof_nvtx_pop()
 
 
-if _event_level == -1:
-
-    @contextmanager
-    def _EmptyEventGuard(event_name, event_level=0):
-        yield
-
-    EventGuard = _EmptyEventGuard  # noqa: F811
-
-
-def event_register(event_name, event_level=0):
+def event_register(event_name, event_level=1):
     def event_wrapper(func):
         @wraps(func)
         def call_with_event(*args, **kwargs):
-            with EventGuard(event_name, event_level=0):
+            with EventGuard(event_name, event_level=event_level):
                 return func(*args, **kwargs)
 
         return call_with_event