support typedvalue

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

python/triton_tvm_ffi/driver.py

@@ -1,9 +1,9 @@
 from __future__ import annotations
 
 from ctypes import c_void_p
-from typing import List, Mapping, Tuple, Type
+from typing import Any, List, Mapping, Optional, Tuple, Type
 from triton.backends.nvidia.driver import CudaDriver
-from .utils import get_device_properties, launch, load_binary
+from . import TypedValue, get_device_properties, launch, load_binary, string_to_type
 
 
 class TVMFFIUtils(object):
@@ -53,7 +53,7 @@ class TVMLauncher(object):
     def __init__(self, src: List[bool], metadata, *args, **kwargs) -> TVMLauncher:
         super().__init__(*args, **kwargs)
 
-        self.mask: List[bool] = [annotation != "constexpr" for annotation in src.signature.values()]
+        self.signature = src.signature.values()
         self.num_ctas = getattr(metadata, "num_ctas", 1)
         self.launch = launch
         self.global_scratch_size = metadata.global_scratch_size
@@ -63,12 +63,6 @@ class TVMLauncher(object):
         self.launch_cooperative_grid = metadata.launch_cooperative_grid
         self.launch_pdl = metadata.launch_pdl
 
-        # We assume the global Triton allocator is not enabled: `_allocator` must be a NullAllocator.
-        # This module depends on NullAllocator behavior; ensure no other code replaces the allocator.
-        from triton.runtime._allocation import _allocator, NullAllocator
-
-        assert isinstance(_allocator.get(), NullAllocator)
-
     def __call__(
         self,
         gridX,
@@ -102,8 +96,15 @@ class TVMLauncher(object):
         )
         assert not self.launch_cooperative_grid
         assert not self.launch_pdl
-        assert len(self.mask) == len(args)
-        args = [arg for arg, m in zip(args, self.mask) if m]
+        assert len(self.signature) == len(args)
+
+        def canonicalize(arg: Any, sig: str) -> TypedValue:
+            ty: Optional[int] = string_to_type(sig)
+            assert ty is not None, sig
+            return TypedValue(ty, arg)
+
+        args = [canonicalize(arg, sig) for arg, sig in zip(args, self.signature)]
+
         return launch(
             gridX,
             gridY,