Refactor hipcc.py to support target specification in tvm_callback_hip…

…_compile
TileLang · Oct 29, 2024 · 27078af · 27078af
1 parent 072a5a1
commit 27078af
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/tvm/contrib/hipcc.py b/python/tvm/contrib/hipcc.py
@@ -106,8 +106,8 @@ def compile_hip(code, target_format="hsaco", arch=None, options=None, path_targe
         return data
 
 
-@tvm._ffi.register_func("tvm_callback_hip_compile")
-def tvm_callback_hip_compile(code):
+@tvm._ffi.register_func
+def tvm_callback_hip_compile(code, target):
     """use hipcc to generate fatbin code for better optimization"""
     hsaco = compile_hip(code, target_format="hsaco")
     return hsaco