tilelang-rocm 0.1.4.post11__cp310-cp310-manylinux1_x86_64.whl → 0.1.4.post13__cp310-cp310-manylinux1_x86_64.whl

tilelang/VERSION

@@ -1 +1 @@
-0.1.4.post11
+0.1.4.post13

tilelang/autotuner/__init__.py

@@ -100,6 +100,7 @@ class AutoTuner:
     compile_args = CompileArgs()
     profile_args = ProfileArgs()
 
+    _kernel_parameters: Optional[Tuple[str, ...]] = None
     _lock = threading.Lock()  # For thread safety
     _memory_cache = {}  # In-memory cache dictionary
     cache_dir: Path = Path(TILELANG_CACHE_DIR)
@@ -167,7 +168,7 @@ class AutoTuner:
                          max_mismatched_ratio: float = 0.01,
                          skip_check: bool = False,
                          manual_check_prog: Callable = None,
-                         cache_input_tensors: bool = True):
+                         cache_input_tensors: bool = False):
         """Set profiling arguments for the auto-tuner.
 
         Args:
@@ -209,12 +210,26 @@ class AutoTuner:
 
         return self
 
-    def generate_cache_key(self) -> Optional[AutotuneResult]:
+    def set_kernel_parameters(self, parameters: Tuple[str, ...]):
+        # for cache key generation
+        self._kernel_parameters = parameters
+
+    def generate_cache_key(self, parameters: Dict[str, Any]) -> Optional[AutotuneResult]:
         """Generate a cache key for the auto-tuning process.
         """
+        # extract parameters from the function signature
+        op_parameters = []
+        for _, default_value in parameters.items():
+            if default_value.default is not inspect.Parameter.empty:
+                op_parameters.append(default_value.default)
+        
+        if self._kernel_parameters is not None:
+            op_parameters += self._kernel_parameters
+
         func_source = inspect.getsource(self.fn)
         key_data = {
             "version": __version__,
+            "op_parameters": tuple(op_parameters),
             "func_source": func_source,
             "configs": self.configs,
             "compile_args": hash(self.compile_args),
@@ -244,12 +259,17 @@ class AutoTuner:
         """
         _init_logger_handlers()
 
-        key = self.generate_cache_key()
+
+        sig = inspect.signature(self.fn)
+        parameters = sig.parameters
+
+        key = self.generate_cache_key(parameters)
+
         with self._lock:
             if is_cache_enabled():
                 # First check in-memory cache
                 if key in self._memory_cache:
-                    self.logger.warning("Found kernel in memory cache. For better performance," \
+                    logger.warning("Found kernel in memory cache. For better performance," \
                                         " consider using `@tilelang.autotune` instead of direct AutoTuner.from_kernel.")
                     return self._memory_cache[key]
 
@@ -260,8 +280,6 @@ class AutoTuner:
                     self._memory_cache[key] = result
                     return result
 
-        sig = inspect.signature(self.fn)
-        parameters = sig.parameters
         best_latency: float = 1e8
         best_config: Optional[Dict[str, Any]] = None
         best_kernel: Optional[tilelang.JITKernel] = None
@@ -345,9 +363,13 @@ class AutoTuner:
         config_args = []
         for config in self.configs:
             new_kwargs = {}
+            keys = config.keys()
             for name, _ in parameters.items():
                 if name in config:
                     new_kwargs[name] = config[name]
+            unused_keys = set(keys) - set(new_kwargs.keys())
+            if len(unused_keys) > 0:
+                raise ValueError(f"Unused keys in config: {unused_keys}")
             config_args.append(new_kwargs)
 
         num_workers = max(1, int(get_available_cpu_count() * 0.9))
@@ -463,8 +485,26 @@ class _AutoTunerImplementation:
     rep: int = 100
     timeout: int = 100
     configs: Any = None
-
-    def __init__(self, configs: Any, warmup: int = 25, rep: int = 100, timeout: int = 100) -> None:
+    supply_type: tilelang.TensorSupplyType = tilelang.TensorSupplyType.Auto
+    ref_prog: Callable = None
+    supply_prog: Callable = None
+    rtol: float = 1e-2
+    atol: float = 1e-2
+    max_mismatched_ratio: float = 0.01
+    skip_check: bool = False
+    manual_check_prog: Callable = None
+    cache_input_tensors: bool = False
+
+    def __init__(self, configs: Any, warmup: int = 25, rep: int = 100, timeout: int = 100,
+                 supply_type: tilelang.TensorSupplyType = tilelang.TensorSupplyType.Auto,
+                 ref_prog: Callable = None,
+                 supply_prog: Callable = None,
+                 rtol: float = 1e-2,
+                 atol: float = 1e-2,
+                 max_mismatched_ratio: float = 0.01,
+                 skip_check: bool = False,
+                 manual_check_prog: Callable = None,
+                 cache_input_tensors: bool = False) -> None:
         """Initialize the AutoTunerImplementation.
 
         Args:
@@ -509,8 +549,20 @@ class _AutoTunerImplementation:
                 def jit_compile(**config_arg):
                     return fn(*args, **kwargs, __tune_params=config_arg)
 
-                autotuner = AutoTuner(fn, configs=configs)
+                autotuner = AutoTuner(fn, configs=configs).set_profile_args(
+                    supply_type=self.supply_type,
+                    ref_prog=self.ref_prog,
+                    supply_prog=self.supply_prog,
+                    rtol=self.rtol,
+                    atol=self.atol,
+                    max_mismatched_ratio=self.max_mismatched_ratio,
+                    skip_check=self.skip_check,
+                    manual_check_prog=self.manual_check_prog,
+                    cache_input_tensors=self.cache_input_tensors,
+                )
                 autotuner.jit_compile = jit_compile
+                autotuner.set_kernel_parameters(key)
+
                 autotuner.run = partial(autotuner.run, warmup, rep, timeout)
 
                 artifact = autotuner.run()
@@ -525,9 +577,21 @@ def autotune(  # This is the new public interface
         func: Union[Callable[_P, _RProg], PrimFunc, None] = None,
         *,  # Indicates subsequent arguments are keyword-only
         configs: Any,
+        # profile arguments
         warmup: int = 25,
         rep: int = 100,
-        timeout: int = 100):
+        timeout: int = 100,
+        # compile arguments
+        supply_type: tilelang.TensorSupplyType = tilelang.TensorSupplyType.Auto,
+        ref_prog: Callable = None,
+        supply_prog: Callable = None,
+        rtol: float = 1e-2,
+        atol: float = 1e-2,
+        max_mismatched_ratio: float = 0.01,
+        skip_check: bool = False,
+        manual_check_prog: Callable = None,
+        cache_input_tensors: bool = False,
+    ):
     """
     Just-In-Time (JIT) compiler decorator for TileLang functions.
 
@@ -571,5 +635,15 @@ def autotune(  # This is the new public interface
         # Create a _AutoTunerImplementation instance with the provided/defaulted arguments.
         # This instance is a decorator that will be applied to the function later.
         configured_decorator = _AutoTunerImplementation(
-            configs=configs, warmup=warmup, rep=rep, timeout=timeout)
+            configs=configs, warmup=warmup, rep=rep, timeout=timeout,
+            supply_type=supply_type,
+            ref_prog=ref_prog,
+            supply_prog=supply_prog,
+            rtol=rtol,
+            atol=atol,
+            max_mismatched_ratio=max_mismatched_ratio,
+            skip_check=skip_check,
+            manual_check_prog=manual_check_prog,
+            cache_input_tensors=cache_input_tensors,
+        )
         return configured_decorator

tilelang/engine/phase.py

@@ -142,6 +142,7 @@ def OptimizeForTarget(mod: IRModule, target: Target) -> IRModule:
     mod = tilelang.transform.AnnotateDeviceRegions()(mod)
     mod = tir.transform.SplitHostDevice()(mod)
 
+    # mod = tir.transform.MergeSharedMemoryAllocations()(mod)
     if allow_warp_specialized(pass_ctx=pass_ctx, target=target):
         # This is a workaround to avoid the bug in the MergeSharedMemoryAllocations pass
         # when warp specialization is enabled, as different warp threads may access different

tilelang/intrinsics/mfma_macro_generator.py

@@ -358,6 +358,8 @@ class MatrixCoreIntrinEmitter(object):
         BLOCK_M = block_row_warps * warp_rows
         BLOCK_N = block_col_warps * warp_cols
         M_DIM, N_DIM = self.M_DIM, self.N_DIM
+        C_buf_dims = len(C_buf.shape)
+        assert C_buf_dims in {2, 4}, "C_buf should be 2D or 4D"
 
         # STS
         # MMA Store must be in simulated instead of TVM Intrins
@@ -369,9 +371,15 @@ class MatrixCoreIntrinEmitter(object):
             for i, j in T.grid(warp_rows, warp_cols):
                 for local_id in T.vectorized(local_size_out):
                     row, col = T.meta_var(mfma_store_index_map(tx, local_id))
-                    C_buf[warp_m * warp_rows + i, warp_n * warp_cols + j, row,
-                          col] = C_local_buf[i * warp_cols * local_size_out + j * local_size_out +
-                                             local_id]
+                    if C_buf_dims == 2:
+                        C_buf[(warp_m * warp_rows + i) * M_DIM + row,
+                                  (warp_n * warp_cols + j) * N_DIM +
+                                  col] = C_local_buf[i * (warp_cols * local_size_out) +
+                                                     j * local_size_out + local_id]
+                    else:
+                        C_buf[warp_m * warp_rows + i, warp_n * warp_cols + j, row,
+                            col] = C_local_buf[i * warp_cols * local_size_out + j * local_size_out +
+                                                local_id]
 
         @T.macro
         def _warp_stmatrix_global(C_local_buf, C_buf, thread_binding):

tilelang/jit/adapter/cython/adapter.py

@@ -432,8 +432,8 @@ class CythonKernelAdapter(BaseKernelAdapter):
     def _convert_torch_func(self) -> Callable:
         """Returns a PyTorch-compatible function wrapper for the kernel."""
 
-        def lambda_forward(*args, stream: int = -1):
-            return self.cython_wrapper.forward([*args], stream=stream)
+        def lambda_forward(*args, stream: int = -1, skip_check: bool = False):
+            return self.cython_wrapper.forward([*args], stream=stream, skip_check=skip_check)
 
         return lambda_forward
 

tilelang/jit/adapter/cython/cython_wrapper.pyx

@@ -66,7 +66,7 @@ cdef class CythonKernelWrapper:
         self.buffer_device_map = buffer_device_map
         return self
 
-    cpdef forward(self, list inputs, int64_t stream = -1):
+    cpdef forward(self, list inputs, int64_t stream = -1, bint skip_check = False):
         # Validate input dimensions and prepare for kernel execution
         cdef int total_params = len(self.params)
         cdef int total_inputs = len(inputs)
@@ -135,29 +135,30 @@ cdef class CythonKernelWrapper:
                 raise ValueError(f"Unsupported tensor type: {type(tensor)}")
 
         # Check buffer device
-        # cdef str tensor_list_device_type = tensor_list[0].device.type
-        if isinstance(tensor_list[0], torch.Tensor):
-            tensor_list_device_type = tensor_list[0].device.type
-        for param, (buffer_idx, device) in self.buffer_device_map.items():
-            if isinstance(tensor_list[buffer_idx], torch.Tensor):
-                tensor_device = tensor_list[buffer_idx].device
-                # Compare device types and indices separately to handle both string and torch.device objects            
-                if (tensor_list_device_type != device.type or 
-                    (tensor_device.index is not None and device.index is not None and tensor_device.index != device.index)):
-                    raise ValueError(f"Buffer device mismatch for parameter {param}: expected {device}, got {tensor_device}")
-
-        # Check buffer dtype map
-        for param, (buffer_idx, torch_dtype) in self.buffer_dtype_map.items():
-            if isinstance(tensor_list[buffer_idx], torch.Tensor):
-                if tensor_list[buffer_idx].dtype != torch_dtype:
-                    raise ValueError(f"Buffer dtype mismatch for parameter {param}: expected {torch_dtype}, got {tensor_list[buffer_idx].dtype}")
-        
-        # Check static shape map
-        for param, (buffer_idx, shape_list) in self.static_shape_map.items():
-            if isinstance(tensor_list[buffer_idx], torch.Tensor):
-                for shape_idx, shape in shape_list:
-                    if tensor_list[buffer_idx].shape[shape_idx] != shape:
-                        raise ValueError(f"Static shape mismatch for parameter {param}: expected {shape} at index {shape_idx}, got {tensor_list[buffer_idx].shape}")
+        if not skip_check:
+            # cdef str tensor_list_device_type = tensor_list[0].device.type
+            if isinstance(tensor_list[0], torch.Tensor):
+                tensor_list_device_type = tensor_list[0].device.type
+            for param, (buffer_idx, device) in self.buffer_device_map.items():
+                if isinstance(tensor_list[buffer_idx], torch.Tensor):
+                    tensor_device = tensor_list[buffer_idx].device
+                    # Compare device types and indices separately to handle both string and torch.device objects            
+                    if (tensor_list_device_type != device.type or 
+                        (tensor_device.index is not None and device.index is not None and tensor_device.index != device.index)):
+                        raise ValueError(f"Buffer device mismatch for parameter {param}: expected {device}, got {tensor_device}")
+
+            # Check buffer dtype map
+            for param, (buffer_idx, torch_dtype) in self.buffer_dtype_map.items():
+                if isinstance(tensor_list[buffer_idx], torch.Tensor):
+                    if tensor_list[buffer_idx].dtype != torch_dtype:
+                        raise ValueError(f"Buffer dtype mismatch for parameter {param}: expected {torch_dtype}, got {tensor_list[buffer_idx].dtype}")
+            
+            # Check static shape map
+            for param, (buffer_idx, shape_list) in self.static_shape_map.items():
+                if isinstance(tensor_list[buffer_idx], torch.Tensor):
+                    for shape_idx, shape in shape_list:
+                        if tensor_list[buffer_idx].shape[shape_idx] != shape:
+                            raise ValueError(f"Static shape mismatch for parameter {param}: expected {shape} at index {shape_idx}, got {tensor_list[buffer_idx].shape}")
 
         # Add dynamic dimension values to kernel arguments
         for _, (buffer_idx, shape_idx) in self.dynamic_symbolic_map.items():

tilelang/jit/adapter/libgen.py

@@ -31,6 +31,8 @@ class LibraryGenerator(object):
     def load_lib(self, lib_path: Optional[str] = None):
         if lib_path is None:
             lib_path = self.libpath
+        else:
+            self.libpath = lib_path
         return ctypes.CDLL(lib_path)
 
     def compile_lib(self, timeout: float = None):

tilelang/src/tl_templates/hip/common.h

@@ -99,3 +99,16 @@ TL_DEVICE unsigned __pack_half2(const half_t x, const half_t y) {
   unsigned v1 = *((unsigned short *)&y);
   return (v1 << 16) | v0;
 }
+
+// Pack two bfloat16_t values.
+TL_DEVICE unsigned __pack_bfloat162(const bfloat16_t x, const bfloat16_t y) {
+  unsigned v0 = *((unsigned short *)&x);
+  unsigned v1 = *((unsigned short *)&y);
+  return (v1 << 16) | v0;
+}
+
+
+template <typename T1, typename T2>
+TL_DEVICE void AtomicAdd(T1 *address, T2 val) {
+  atomicAdd(reinterpret_cast<T1 *>(address), static_cast<T1>(val));
+}

{tilelang_rocm-0.1.4.post11.dist-info → tilelang_rocm-0.1.4.post13.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tilelang-rocm
-Version: 0.1.4.post11
+Version: 0.1.4.post13
 Summary: A tile level programming language to generate high performance code.
 Home-page: https://github.com/tile-ai/tilelang
 Author: Microsoft Research