fix bugs on illegal memory access on Release

Signed-off-by: Jinjie Liu <jjliu@baai.ac.cn>

python/triton_tvm_ffi/_ffi_api.py

@@ -23,6 +23,7 @@ if TYPE_CHECKING:
 # fmt: on
 # tvm-ffi-stubgen(end)
 
+
 # tvm-ffi-stubgen(import-object): tvm_ffi.register_object;False;_FFI_REG_OBJ
 # tvm-ffi-stubgen(import-object): ffi.Object;False;_ffi_Object
 @_FFI_REG_OBJ("triton_tvm_ffi.TypedValue")
@@ -35,6 +36,7 @@ class TypedValue(_ffi_Object):
     # fmt: on
     # tvm-ffi-stubgen(end)
 
+
 __all__ = [
     # tvm-ffi-stubgen(begin): __all__
     "LIB",

python/triton_tvm_ffi/driver.py

@@ -2,11 +2,12 @@ from __future__ import annotations
 
 from typing import Any, List, Optional, Type
 from triton.backends.nvidia.driver import CudaDriver
+from triton.runtime import _allocation
 from . import TypedValue, utils, string_to_type
 
 
 class TVMLauncher(object):
-    def __init__(self, src: List[bool], metadata, *args, **kwargs) -> TVMLauncher:
+    def __init__(self, src, metadata, *args, **kwargs) -> TVMLauncher:
         super().__init__(*args, **kwargs)
 
         self.signature: List[str] = src.signature.values()
@@ -32,8 +33,6 @@ class TVMLauncher(object):
         launch_exit_hook,
         *args,
     ):
-        from triton.runtime import _allocation
-
         def allocate_scratch(size, align, allocator):
             if size > 0:
                 grid_size = gridX * gridY * gridZ