pytorch
diff --git a/‎test/profiler/test_cpp_thread.py
Copy file name to clipboard
+157-11Lines changed: 157 additions & 11 deletions b/‎test/profiler/test_cpp_thread.py
Copy file name to clipboard
+157-11Lines changed: 157 additions & 11 deletions
diff --git a/‎test/profiler/test_execution_trace.py
Copy file name to clipboard
+16-7Lines changed: 16 additions & 7 deletions b/‎test/profiler/test_execution_trace.py
Copy file name to clipboard
+16-7Lines changed: 16 additions & 7 deletions
@@ -1,6 +1,7 @@
 # Owner(s): ["oncall: profiler"]
 
 import os
+import unittest
 from unittest import skipIf
 
 import torch
@@ -32,6 +33,7 @@
 KinetoProfiler = None
 IterationCount = 5
 ActivateIteration = 2
+device = None
 
 
 def blueprint(text):
@@ -56,26 +58,29 @@ def onIterationStart(self, iteration: int) -> None:
             KinetoProfiler.step()
 
     def emulateTraining(self, iteration: int, thread_id: int) -> None:
+        global device
         # blueprint(f"training iteration {iteration} in thread {thread_id}")
-        device = torch.device("cuda")
-        # device = torch.device("cpu")
+        torch_device = getattr(torch, device)
+        assert hasattr(torch_device, "synchronize")
+        sync_func = torch_device.synchronize
+
         with torch.autograd.profiler.record_function("user_function"):
             a = torch.ones(1, device=device)
             b = torch.ones(1, device=device)
             torch.add(a, b).cpu()
-            torch.cuda.synchronize()
+            sync_func()
 
 
-class CppThreadTest(TestCase):
+class CppThreadTestCUDA(TestCase):
     ThreadCount = 20  # set to 2 for debugging
     EventHandler = None
     TraceObject = None
 
     @classmethod
     def setUpClass(cls) -> None:
         super(TestCase, cls).setUpClass()
-        CppThreadTest.EventHandler = PythonProfilerEventHandler()
-        cpp.ProfilerEventHandler.Register(CppThreadTest.EventHandler)
+        CppThreadTestCUDA.EventHandler = PythonProfilerEventHandler()
+        cpp.ProfilerEventHandler.Register(CppThreadTestCUDA.EventHandler)
 
     @classmethod
     def tearDownClass(cls):
@@ -85,6 +90,8 @@ def tearDownClass(cls):
     def setUp(self) -> None:
         if not torch.cuda.is_available():
             self.skipTest("Test machine does not have cuda")
+        global device
+        device = "cuda"
 
         # this clears off events from initialization
         self.start_profiler(False)
@@ -103,7 +110,7 @@ def start_profiler(self, profile_memory):
         )
 
     def set_trace(self, trace_obj) -> None:
-        CppThreadTest.TraceObject = trace_obj
+        CppThreadTestCUDA.TraceObject = trace_obj
 
     def assert_text(self, condition, text, msg):
         if condition:
@@ -114,7 +121,7 @@ def assert_text(self, condition, text, msg):
 
     def check_trace(self, expected, mem=False) -> None:
         blueprint("verifying trace")
-        event_list = CppThreadTest.TraceObject.events()
+        event_list = CppThreadTestCUDA.TraceObject.events()
         for key, values in expected.items():
             count = values[0]
             min_count = count * (ActivateIteration - 1)
@@ -160,7 +167,7 @@ def check_trace(self, expected, mem=False) -> None:
         IS_WINDOWS,
         "Failing on windows cuda, see https://github.com/pytorch/pytorch/pull/130037 for slightly more context",
     )
-    def test_with_enable_profiler_in_child_thread(self) -> None:
+    def test_with_enable_profiler_in_child_thread_cuda(self) -> None:
         self.start_profiler(False)
         cpp.start_threads(self.ThreadCount, IterationCount, True)
         self.check_trace(
@@ -174,7 +181,7 @@ def test_with_enable_profiler_in_child_thread(self) -> None:
         IS_WINDOWS,
         "Failing on windows cuda, see https://github.com/pytorch/pytorch/pull/130037 for slightly more context",
     )
-    def test_without_enable_profiler_in_child_thread(self) -> None:
+    def test_without_enable_profiler_in_child_thread_cuda(self) -> None:
         self.start_profiler(False)
         cpp.start_threads(self.ThreadCount, IterationCount, False)
         self.check_trace(
@@ -188,7 +195,146 @@ def test_without_enable_profiler_in_child_thread(self) -> None:
         IS_WINDOWS,
         "Failing on windows cuda, see https://github.com/pytorch/pytorch/pull/130037 for slightly more context",
     )
-    def test_profile_memory(self) -> None:
+    def test_profile_memory_cuda(self) -> None:
+        self.start_profiler(True)
+        cpp.start_threads(self.ThreadCount, IterationCount, True)
+        self.check_trace(
+            {
+                "aten::add": [self.ThreadCount, "CPU"],
+            },
+            mem=True,
+        )
+
+
+# Here duplicate the CppThreadTest to enable the xpu cases because the
+# instantiate_device_type_tests will call class method setUpClass.
+# In function setUpClass, the instantiated class(e.g CppThreadTestCPU, CppThreadTestXPU)
+# needs to be called to get it member EventHandler, while in this period,
+# the input class in argument cls is CppThreadTest, which is not defined any more.
+# We cannot detect which instantiated class is being created in setUpClass, so duplicate here
+# for enabling xpu test cases
+class CppThreadTestXPU(TestCase):
+    ThreadCount = 20  # set to 2 for debugging
+    EventHandler = None
+    TraceObject = None
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        super(TestCase, cls).setUpClass()
+        CppThreadTestXPU.EventHandler = PythonProfilerEventHandler()
+        cpp.ProfilerEventHandler.Register(CppThreadTestXPU.EventHandler)
+
+    @classmethod
+    def tearDownClass(cls):
+        if not is_fbcode():
+            torch.testing._internal.common_utils.remove_cpp_extensions_build_root()
+
+    def setUp(self) -> None:
+        if not torch.xpu.is_available():
+            self.skipTest("Test machine does not have xpu")
+        global device
+        device = "xpu"
+
+        # this clears off events from initialization
+        self.start_profiler(False)
+        cpp.start_threads(1, IterationCount, False)
+
+    def start_profiler(self, profile_memory):
+        global KinetoProfiler
+        KinetoProfiler = torch.profiler.profile(
+            schedule=torch.profiler.schedule(
+                wait=1, warmup=1, active=ActivateIteration, repeat=1
+            ),
+            on_trace_ready=self.set_trace,
+            with_stack=True,
+            profile_memory=profile_memory,
+            record_shapes=True,
+        )
+
+    def set_trace(self, trace_obj) -> None:
+        CppThreadTestXPU.TraceObject = trace_obj
+
+    def assert_text(self, condition, text, msg):
+        if condition:
+            print(f"\33[32m{text}\33[0m")
+        else:
+            print(f"\33[31m{text}\33[0m")
+        self.assertTrue(condition, msg)
+
+    def check_trace(self, expected, mem=False) -> None:
+        blueprint("verifying trace")
+        event_list = CppThreadTestXPU.TraceObject.events()
+        for key, values in expected.items():
+            count = values[0]
+            min_count = count * (ActivateIteration - 1)
+            device = values[1]
+            filtered = filter(
+                lambda ev: ev.name == key
+                and str(ev.device_type) == f"DeviceType.{device}",
+                event_list,
+            )
+
+            if mem:
+                actual = 0
+                for ev in filtered:
+                    sev = str(ev)
+                    has_cuda_memory_usage = (
+                        sev.find("xpu_memory_usage=0 ") < 0
+                        and sev.find("xpu_memory_usage=") > 0
+                    )
+                    if has_cuda_memory_usage:
+                        actual += 1
+                self.assert_text(
+                    actual >= min_count,
+                    f"{key}: {actual} >= {min_count}",
+                    "not enough event with xpu_memory_usage set",
+                )
+            else:
+                actual = len(list(filtered))
+                if count == 1:  # test_without
+                    count *= ActivateIteration
+                    self.assert_text(
+                        actual == count,
+                        f"{key}: {actual} == {count}",
+                        "baseline event count incorrect",
+                    )
+                else:
+                    self.assert_text(
+                        actual >= min_count,
+                        f"{key}: {actual} >= {min_count}",
+                        "not enough event recorded",
+                    )
+
+    @unittest.skip(
+        reason="The XPU Profiler will not cover this case for now. Will support it in next period."
+    )
+    def test_with_enable_profiler_in_child_thread_xpu(self) -> None:
+        self.start_profiler(False)
+        cpp.start_threads(self.ThreadCount, IterationCount, True)
+        self.check_trace(
+            {
+                "aten::add": [self.ThreadCount, "CPU"],
+                "user_function": [self.ThreadCount, "XPU"],
+            }
+        )
+
+    @unittest.skip(
+        reason="The XPU Profiler will not cover this case for now. Will support it in next period."
+    )
+    def test_without_enable_profiler_in_child_thread_xpu(self) -> None:
+        self.start_profiler(False)
+        cpp.start_threads(self.ThreadCount, IterationCount, False)
+        self.check_trace(
+            {
+                "aten::add": [1, "CPU"],
+                "user_function": [1, "XPU"],
+            }
+        )
+
+    @unittest.skip(
+        reason="The XPU Profiler will not cover this case for now. Will support it in next period."
+    )
+    def test_profile_memory_xpu(self) -> None:
         self.start_profiler(True)
         cpp.start_threads(self.ThreadCount, IterationCount, True)
         self.check_trace(
 
@@ -41,6 +41,7 @@
     skipIfHpu,
     skipIfTorchDynamo,
     TEST_HPU,
+    TEST_XPU,
     TestCase,
 )
 from torch.utils._triton import has_triton
@@ -131,6 +132,7 @@ def trace_handler(p):
 
         use_device = (
             torch.profiler.ProfilerActivity.CUDA
+            or torch.profiler.ProfilerActivity.XPU in supported_activities()
             or torch.profiler.ProfilerActivity.HPU in supported_activities()
         )
         # Create a temp file to save execution trace and kineto data.
@@ -202,6 +204,7 @@ def test_execution_trace_alone(self, device):
         use_device = (
             torch.profiler.ProfilerActivity.CUDA
             or torch.profiler.ProfilerActivity.HPU in supported_activities()
+            or torch.profiler.ProfilerActivity.XPU in supported_activities()
         )
         # Create a temp file to save execution trace data.
         fp = tempfile.NamedTemporaryFile("w+t", suffix=".et.json", delete=False)
@@ -241,8 +244,10 @@ def test_execution_trace_alone(self, device):
     @unittest.skipIf(
         sys.version_info >= (3, 12), "torch.compile is not supported on python 3.12+"
     )
-    @unittest.skipIf(not TEST_CUDA or not has_triton(), "need CUDA and triton to run")
-    @skipIfHpu
+    @unittest.skipIf(
+        (not has_triton()) or (not TEST_CUDA and not TEST_XPU),
+        "need triton and device(CUDA or XPU) availability to run",
+    )
     def test_execution_trace_with_pt2(self, device):
         @torchdynamo.optimize("inductor")
         def fn(a, b, c):
@@ -290,6 +295,7 @@ def fn(a, b, c):
     def test_execution_trace_start_stop(self, device):
         use_device = (
             torch.profiler.ProfilerActivity.CUDA
+            or torch.profiler.ProfilerActivity.XPU in supported_activities()
             or torch.profiler.ProfilerActivity.HPU in supported_activities()
         )
         # Create a temp file to save execution trace data.
@@ -328,6 +334,7 @@ def test_execution_trace_start_stop(self, device):
     def test_execution_trace_repeat_in_loop(self, device):
         use_device = (
             torch.profiler.ProfilerActivity.CUDA
+            or torch.profiler.ProfilerActivity.XPU in supported_activities()
             or torch.profiler.ProfilerActivity.HPU in supported_activities()
         )
         iter_list = {3, 4, 6, 8}
@@ -361,8 +368,7 @@ def test_execution_trace_repeat_in_loop(self, device):
             assert found_root_node
         assert event_count == expected_loop_events
 
-    @skipIfHpu
-    def test_execution_trace_no_capture(self, device):
+    def test_execution_trace_no_capture(self):
         fp = tempfile.NamedTemporaryFile("w+t", suffix=".et.json", delete=False)
         fp.close()
         et = ExecutionTraceObserver().register_callback(fp.name)
@@ -377,8 +383,7 @@ def test_execution_trace_no_capture(self, device):
         assert found_root_node
 
     @skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/124500")
-    @skipIfHpu
-    def test_execution_trace_nested_tensor(self, device):
+    def test_execution_trace_nested_tensor(self):
         fp = tempfile.NamedTemporaryFile("w+t", suffix=".et.json", delete=False)
         fp.close()
 
@@ -404,9 +409,13 @@ def fn(nt):
 
 
 devices = ["cpu", "cuda"]
+if TEST_XPU:
+    devices.append("xpu")
 if TEST_HPU:
     devices.append("hpu")
-instantiate_device_type_tests(TestExecutionTrace, globals(), only_for=devices)
+instantiate_device_type_tests(
+    TestExecutionTrace, globals(), allow_xpu="xpu" in devices, only_for=devices
+)
 
 if __name__ == "__main__":
     run_tests()