Remove usage of torch.autograd.profiler_legacy for benchmarks

benchmarks/CMakeLists.txt

@@ -4,14 +4,20 @@ set(CMAKE_CXX_STANDARD 20)
 
 project(TritonBenchmark)
 
+option(USE_IPEX "Use IPEX" ON)
+
 if(NOT WIN32)
     list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
 endif()
 
 find_package(Python3 COMPONENTS Interpreter)
 find_package(Torch REQUIRED)
 find_library(TORCH_PYTHON_LIBRARY torch_python PATH "${TORCH_INSTALL_PREFIX}/lib")
-find_package(IPEX REQUIRED)
+
+if(USE_IPEX)
+  string(APPEND CMAKE_CXX_FLAGS " -DUSE_IPEX")
+  find_package(IPEX REQUIRED)
+endif()
 
 
 # add the XeTLA kernel.

benchmarks/setup.py

@@ -8,7 +8,11 @@
 from setuptools import setup
 
 import torch
-import intel_extension_for_pytorch
+
+ipex_cmake_prefix_path = ""
+if os.getenv("USE_IPEX", "1") == "1":
+    import intel_extension_for_pytorch
+    ipex_cmake_prefix_path = f";{intel_extension_for_pytorch.cmake_prefix_path}"
 
 
 class CMakeBuild():
@@ -43,7 +47,7 @@ def build_extension(self):
             "Ninja",  # Ninja is much faster than make
             "-DCMAKE_MAKE_PROGRAM=" +
             ninja_dir,  # Pass explicit path to ninja otherwise cmake may cache a temporary path
-            f"-DCMAKE_PREFIX_PATH={torch.utils.cmake_prefix_path};{intel_extension_for_pytorch.cmake_prefix_path}",
+            f"-DCMAKE_PREFIX_PATH={torch.utils.cmake_prefix_path}{ipex_cmake_prefix_path}",
             "-DCMAKE_EXPORT_COMPILE_COMMANDS=ON",
             "-DCMAKE_ARCHIVE_OUTPUT_DIRECTORY=" + self.extdir,
             "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + self.extdir,

benchmarks/triton_kernels_benchmark/__init__.py

@@ -1,6 +1 @@
-from triton.runtime import driver
-from . import benchmark_driver
 from .benchmark_testing import do_bench, assert_close, perf_report, Benchmark  # type: ignore # noqa: F401
-
-# replace the launcher with the profilier hook.
-driver.active.launcher_cls = benchmark_driver.XPULauncher