enable local hook just when cuda profiling is enabled

wffpy · Nov 24, 2024 · 1288637 · 1288637
1 parent a36b60e
commit 1288637
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 2 deletions.
diff --git a/lib/dump/DumpHook.cpp b/lib/dump/DumpHook.cpp
@@ -204,9 +204,9 @@ std::string LaunchInfo::to_string(const TensorMap& tensor_map, const int64_t cap
         DLOG() << "  prointer_level: " << parsed_type.prointer_level;
         // memory ptr
         if (parsed_type.prointer_level > 0) {
+            char* cpu_ptr = ((char*)params_.data()) + pos;
             // input memory
             if (parsed_type.is_const) {
-                char* cpu_ptr = ((char*)params_.data()) + pos;
                 LOG() << "input memory " << index << ": " << std::hex << *((int64_t*)cpu_ptr);
                 uint64_t dev_ptr = *((uint64_t*)cpu_ptr);
                 int64_t bytes = 0;
@@ -242,6 +242,8 @@ std::string LaunchInfo::to_string(const TensorMap& tensor_map, const int64_t cap
 
                     free(host_data);
                 }
+            } else {
+                LOG() << "output memory " << index << ": " << std::hex << *((int64_t*)cpu_ptr);
             }
             pos = pos + 8;
         } else if (parsed_type.base_type == "int64_t" || parsed_type.base_type == "long") {

diff --git a/lib/hook/LocalHook/LocalHook.cpp b/lib/hook/LocalHook/LocalHook.cpp
@@ -12,8 +12,10 @@
 #include <sys/mman.h>
 #include <unistd.h>
 
+#ifdef USE_CAPSTONE
 #include "capstone/capstone.h"
 #include "capstone/x86.h"
+#endif
 
 #include "hook/CFuncHook.h"
 #include "utils/Log/Log.h"
@@ -73,6 +75,7 @@ uintptr_t find_free_address(uintptr_t aligned_addr, size_t size) {
 //     return mmap_addr;
 // }
 
+#ifdef USE_CAPSTONE
 struct X64Instructions {
     cs_insn *instructions;
     uint32_t numInstructions;
@@ -821,7 +824,12 @@ void install_local_hook(void *hooked_func, void *payload_func,
         }
     }
 }
+#else
+
+void install_local_hook(void *hooked_func, void *payload_func,
+                        void **trampoline_ptr) {}
 
+#endif
 void LocalHookRegistrar::add(LocalHookInfo info) { hooks.emplace_back(info); }
 
 void LocalHookRegistrar::install() {

diff --git a/python/module_logging/CMakeLists.txt b/python/module_logging/CMakeLists.txt
@@ -11,7 +11,9 @@ target_link_libraries(Hook PRIVATE GpuHook)
 target_link_libraries(Hook PRIVATE LocalHook)
 target_link_libraries(Hook PRIVATE Dump)
 
-target_link_libraries(Hook PRIVATE capstone)
+if (USE_CAPSTONE STREQUAL TRUE)
+    target_link_libraries(Hook PRIVATE capstone)
+endif()
 
 target_include_directories(Hook PUBLIC ${Python3_INCLUDE_DIRS})
 

diff --git a/setup.py b/setup.py
@@ -54,8 +54,10 @@ def build_extension(self, ext):
             # '-DPYTHON_LIBRARY=' + os.environ.get('PYTHON_LIBRARY'),
             # f'-B {os.path.join(script_dir, "build")}',
         ]
+        # up to now, local hook is just used for cuda Profiling 
         if enable_cuda == "true":
             cmake_args.append("-DCUDA_DEV=TRUE")
+            cmake_args.append("-DUSE_CAPSTONE=TRUE")
         if enable_xpu == "true":
             cmake_args.append("-DXPU_DEV=TRUE")