tilelang-rocm 0.1.4.post5__cp310-cp310-manylinux1_x86_64.whl → 0.1.4.post9__cp310-cp310-manylinux1_x86_64.whl

tilelang/3rdparty/tvm/python/tvm/_ffi/runtime_ctypes.py

@@ -67,7 +67,7 @@ class DataTypeCode(object):
     BFLOAT = 4
     E4M3Float = 6
     E5M2Float = 7
-
+    E4M3FnuzFloat = 23
 
 class DataType(ctypes.Structure):
     """TVM datatype structure"""
@@ -81,6 +81,7 @@ class DataType(ctypes.Structure):
         DataTypeCode.BFLOAT: "bfloat",
         DataTypeCode.E4M3Float: "e4m3_float",
         DataTypeCode.E5M2Float: "e5m2_float",
+        DataTypeCode.E4M3FnuzFloat: "e4m3fnuz_float",
     }
     NUMPY2STR = {
         np.dtype(np.bool_): "bool",
@@ -110,6 +111,8 @@ class DataType(ctypes.Structure):
         "uint32": {"type_code": DataTypeCode.UINT, "bits": 32, "lanes": 1},
         "uint64": {"type_code": DataTypeCode.UINT, "bits": 64, "lanes": 1},
         "e4m3_float8": {"type_code": DataTypeCode.E4M3Float, "bits": 8, "lanes": 1},
+        "e4m3fnuz_float8": {"type_code": DataTypeCode.E4M3FnuzFloat, "bits": 8, "lanes": 1},
+        "float8_e4m3fnuz": {"type_code": DataTypeCode.E4M3FnuzFloat, "bits": 8, "lanes": 1},
         "e5m2_float8": {"type_code": DataTypeCode.E5M2Float, "bits": 8, "lanes": 1},
         "float16": {"type_code": DataTypeCode.FLOAT, "bits": 16, "lanes": 1},
         "float32": {"type_code": DataTypeCode.FLOAT, "bits": 32, "lanes": 1},
@@ -167,6 +170,9 @@ class DataType(ctypes.Structure):
         elif head.startswith("e5m2_float"):
             self.type_code = DataTypeCode.E5M2Float
             head = head[10:]
+        elif head.startswith("e4m3fnuz_float"):
+            self.type_code = DataTypeCode.E4M3FnuzFloat
+            head = head[14:]
         elif head.startswith("custom"):
             # pylint: disable=import-outside-toplevel
             import tvm.runtime._ffi_api

tilelang/3rdparty/tvm/python/tvm/contrib/tvmjs.py

@@ -307,6 +307,13 @@ def load_ndarray_cache(cachepath: str, device: tvm.runtime.Device):
                     raise RuntimeError(
                         "ml_dtypes is not installed, cannot convert e4m3_float8 array to numpy."
                     )
+            if dtype == "e4m3fnuz_float8":
+                if ml_dtypes is not None:
+                    dtype = ml_dtypes.float8_e4m3fnuz
+                else:
+                    raise RuntimeError(
+                        "ml_dtypes is not installed, cannot convert e4m3fnuz_float8 array to numpy."
+                    )
             if dtype == "e5m2_float8":
                 if ml_dtypes is not None:
                     dtype = ml_dtypes.float8_e5m2

tilelang/3rdparty/tvm/python/tvm/runtime/ndarray.py

@@ -247,6 +247,14 @@ class NDArray(NDArrayBase):
                 raise RuntimeError(
                     "ml_dtypes is not installed, cannot convert e4m3_float8 array to numpy."
                 )
+        if dtype == "e4m3fnuz_float8":
+            if ml_dtypes is not None:
+                dtype = ml_dtypes.float8_e4m3fnuz
+            else:
+                raise RuntimeError(
+                    "ml_dtypes is not installed, cannot convert e4m3fnuz_float8 array to numpy."
+                )
+            
         if dtype == "e5m2_float8":
             if ml_dtypes is not None:
                 dtype = ml_dtypes.float8_e5m2

tilelang/3rdparty/tvm/python/tvm/script/ir_builder/tir/ir.py

@@ -1455,6 +1455,13 @@ e5m2_float8x16 = func_gen(("E5M2Float8x16"))
 e5m2_float8x32 = func_gen(("E5M2Float8x32"))
 e5m2_float8x64 = func_gen(("E5M2Float8x64"))
 
+e4m3fnuz_float8 = func_gen(("E4M3FNUZFloat8"))
+e4m3fnuz_float8x4 = func_gen(("E4M3FNUZFloat8x4"))
+e4m3fnuz_float8x8 = func_gen(("E4M3FNUZFloat8x8"))
+e4m3fnuz_float8x16 = func_gen(("E4M3FNUZFloat8x16"))
+e4m3fnuz_float8x32 = func_gen(("E4M3FNUZFloat8x32"))
+e4m3fnuz_float8x64 = func_gen(("E4M3FNUZFloat8x64"))
+
 # pylint: enable=invalid-name
 
 

tilelang/3rdparty/tvm/src/target/llvm/codegen_llvm.cc

@@ -586,7 +586,8 @@ llvm::Type* CodeGenLLVM::DTypeToLLVMType(const DataType& dtype) const {
       default:
         LOG(FATAL) << "do not support " << dtype;
     }
-  } else if (dtype.code() == DataType::kFloat8_e4m3fn || dtype.code() == DataType::kFloat8_e5m2) {
+  } else if (dtype.code() == DataType::kFloat8_e4m3fn || dtype.code() == DataType::kFloat8_e5m2
+             || dtype.code() == DataType::kFloat8_e4m3fnuz) {
     etype = llvm::Type::getInt8Ty(*ctx);
   }
   if (!dtype.is_scalar()) {

tilelang/3rdparty/tvm/src/target/source/codegen_cuda.cc

@@ -60,6 +60,8 @@ std::string GetFP8Type(DataType type) {
   }
   if (type.code() == DataType::kFloat8_e4m3fn) {
     stream << "fp8_e4" << vec << "_t";
+  } else if (type.code() == DataType::kFloat8_e4m3fnuz) {
+    stream << "fp8_e4" << vec << "_t";
   } else if (type.code() == DataType::kFloat8_e5m2) {
     stream << "fp8_e5" << vec << "_t";
   } else {
@@ -691,7 +693,8 @@ void CodeGenCUDA::VisitExpr_(const CastNode* op, std::ostream& os) {
   if (from_ty.is_scalar()) return CodeGenC::VisitExpr_(op, os);
 
   if (target_ty.code() == DataType::kFloat8_e4m3fn || target_ty.code() == DataType::kFloat8_e5m2 ||
-      from_ty.code() == DataType::kFloat8_e4m3fn || from_ty.code() == DataType::kFloat8_e5m2) {
+      from_ty.code() == DataType::kFloat8_e4m3fn || from_ty.code() == DataType::kFloat8_e5m2
+      || target_ty.code() == DataType::kFloat8_e4m3fnuz) {
     std::ostringstream val;
     val << "(";
     PrintType(target_ty, val);

tilelang/3rdparty/tvm/src/target/source/ptx.cc

@@ -55,6 +55,7 @@ enum class DataType : int {
   kInt64 = 8,
   kUInt64 = 9,
   kFloat8_e4m3 = 10,
+  kFloat8_e4m3fnuz = 23,
   kFloat8_e5m2 = 11,
   kFloat16 = 12,
   kBFloat16 = 13,
@@ -102,6 +103,8 @@ inline DataType DTypeFromString(const std::string str) {
     return DataType::kUInt64;
   } else if (str == "e4m3" || str == ".e4m3") {
     return DataType::kFloat8_e4m3;
+  } else if (str == "e4m3fnuz" || str == ".e4m3fnu") {
+    return DataType::kFloat8_e4m3fnuz;
   } else if (str == "e5m2" || str == ".e5m2") {
     return DataType::kFloat8_e5m2;
   } else if (str == "float16" || str == "fp16" || str == ".f16") {

tilelang/3rdparty/tvm/src/tir/ir/index_map.cc

@@ -93,7 +93,8 @@ std::pair<IndexMap, PrimExpr> IndexMapInverseImpl(const IndexMap& self,
                                        /*check_level=*/check_level, analyzer,
                                        /*simplify_trivial_iterators=*/false);
   CHECK(padded_iter_map->errors.empty()) << "Could not parse mapping as sum of iterators.  "
-                                         << "Error: " << padded_iter_map->errors[0];
+                                         << "\nIndex map: " << self->initial_indices << " -> " << self->final_indices
+                                         << "\nError: " << padded_iter_map->errors[0];
 
   // Determine expressions for the input variables, in terms of the
   // output variables.

tilelang/3rdparty/tvm/src/tir/ir/stmt.cc

@@ -219,12 +219,6 @@ TVM_REGISTER_NODE_TYPE(ProducerStoreNode);
 // Allocate
 Allocate::Allocate(Var buffer_var, DataType dtype, Array<PrimExpr> extents, PrimExpr condition,
                    Stmt body, Map<String, ObjectRef> annotations, Span span) {
-  CHECK(IsPointerType(buffer_var->type_annotation, dtype) ||
-        (dtype.is_bool() && IsPointerType(buffer_var->type_annotation, DataType::Int(8))))
-      << "The allocated data type (" << dtype
-      << ") does not match the type annotation of the buffer " << buffer_var << " ("
-      << buffer_var->type_annotation
-      << "). The data type should be an element of the pointer type.";
 
   for (size_t i = 0; i < extents.size(); ++i) {
     ICHECK(extents[i].defined());

tilelang/3rdparty/tvm/src/tir/op/op.cc

@@ -201,6 +201,12 @@ void BinaryOpMatchTypes(PrimExpr& lhs, PrimExpr& rhs, Span span) {  // NOLINT(*)
   } else if (ltype.is_float8() && !rtype.is_float8()) {
     // Cast int->float8 for rhs when lhs is a float8
     rhs = cast(ltype, rhs);
+  } else if (!ltype.is_float8_e4m3fnuz() && rtype.is_float8_e4m3fnuz()) {
+    // Cast int->float8_e4m3fnuz for rhs when lhs is a float8_e4m3fnuz
+    rhs = cast(ltype, rhs);
+  } else if (ltype.is_float8_e4m3fnuz() && !rtype.is_float8_e4m3fnuz()) {
+    // Cast int->float8_e4m3fnuz for lhs when rhs is a float8_e4m3fnuz
+    lhs = cast(rtype, lhs);
   } else if ((ltype.is_int() && rtype.is_int()) || (ltype.is_uint() && rtype.is_uint())) {
     // Promote int to higher bits e.g. int8 + int16 --> int16 + int16
     if (ltype.bits() < rtype.bits()) {

tilelang/3rdparty/tvm/src/tir/transforms/dtype_conversion.h

@@ -124,6 +124,9 @@ class FloatConfig {
         case DataType::kFloat8_e4m3fn:
           // E4M3 format, not consistent with IEEE-754
           return FloatConfig(4, 3, 7, InftyStyle::kNone, NaNStyle::kAllOnes);
+        case DataType::kFloat8_e4m3fnuz:
+          // E4M3 format, not consistent with IEEE-754
+          return FloatConfig(4, 3, 7, InftyStyle::kNone, NaNStyle::kNone);
         default:
           // E5M2 format, consistent with IEEE-754
           return FloatConfig(5, 2, 15, InftyStyle::kIEEE, NaNStyle::kIEEE);

tilelang/README.md

@@ -126,8 +126,10 @@ import tilelang.language as T
 from tilelang.intrinsics import (
     make_mma_swizzle_layout as make_swizzle_layout,)
 
+# add decorator @tilelang.jit if you want to return a torch function
+# @tilelang.jit
 def matmul(M, N, K, block_M, block_N, block_K, dtype="float16", accum_dtype="float"):
-    # add decorator @tilelang.jit if you want to return a torch function
+
     @T.prim_func
     def main(
         A: T.Tensor((M, K), dtype),

tilelang/VERSION

@@ -1 +1 @@
-0.1.4.post5
+0.1.4.post9

tilelang/__init__.py

@@ -4,7 +4,6 @@
 import sys
 import os
 import ctypes
-import subprocess
 
 import logging
 from tqdm import tqdm
@@ -62,6 +61,7 @@ from .env import enable_cache, disable_cache, is_cache_enabled  # noqa: F401
 
 import tvm
 import tvm._ffi.base
+from tvm import DataType  # noqa: F401
 
 from . import libinfo
 
@@ -77,67 +77,6 @@ def _load_tile_lang_lib():
     lib_path = libinfo.find_lib_path(lib_name, optional=False)
     return ctypes.CDLL(lib_path[0]), lib_path[0]
 
-def check_and_install_rocm_torch():
-    """检查并自动安装ROCm PyTorch"""
-    try:
-        # 检查torch是否已安装且为ROCm版本
-        import torch
-        if hasattr(torch.version, 'hip') and torch.version.hip:
-            print(f"ROCm PyTorch {torch.version.__version__} is already installed")
-            return True
-    except ImportError:
-        pass
-    
-    print("ROCm PyTorch not found. Installing...")
-    install_rocm_torch()
-    
-    # 强制重新加载torch模块
-    if 'torch' in sys.modules:
-        del sys.modules['torch']
-    
-    try:
-        import torch
-        print(f"Successfully installed ROCm PyTorch {torch.version.__version__}")
-        return True
-    except ImportError:
-        print("Failed to import torch after installation")
-        return False
-
-def install_rocm_torch():
-    """安装ROCm PyTorch"""
-    print("Installing ROCm PyTorch components...")
-    
-    wheels = [
-        "https://repo.radeon.com/rocm/manylinux/rocm-rel-6.3.4/torch-2.4.0%2Brocm6.3.4.git7cecbf6d-cp310-cp310-linux_x86_64.whl",
-        "https://repo.radeon.com/rocm/manylinux/rocm-rel-6.3.4/torchvision-0.19.0%2Brocm6.3.4.gitfab84886-cp310-cp310-linux_x86_64.whl",
-        "https://repo.radeon.com/rocm/manylinux/rocm-rel-6.3.4/pytorch_triton_rocm-3.0.0%2Brocm6.3.4.git75cc27c2-cp310-cp310-linux_x86_64.whl",
-        "https://repo.radeon.com/rocm/manylinux/rocm-rel-6.3.4/torchaudio-2.4.0%2Brocm6.3.4.git69d40773-cp310-cp310-linux_x86_64.whl"
-    ]
-    
-    # 卸载现有的torch相关包
-    print("Uninstalling existing PyTorch...")
-    try:
-        subprocess.run([sys.executable, "-m", "pip", "uninstall", 
-                       "torch", "torchvision", "pytorch-triton-rocm", "torchaudio", "-y"],
-                      capture_output=True)
-    except subprocess.CalledProcessError:
-        print("No existing PyTorch installation found")
-    
-    # 安装ROCm版本
-    print("Downloading and installing ROCm PyTorch...")
-    subprocess.run([sys.executable, "-m", "pip", "install"] + wheels, check=True)
-
-# 在模块导入时自动检查并安装
-_rocm_torch_checked = False
-
-def ensure_rocm_torch():
-    global _rocm_torch_checked
-    if not _rocm_torch_checked:
-        check_and_install_rocm_torch()
-        _rocm_torch_checked = True
-
-# 在模块导入时立即检查
-ensure_rocm_torch()
 
 # only load once here
 if SKIP_LOADING_TILELANG_SO == "0":
@@ -157,10 +96,10 @@ from .layout import (
 )
 from . import (
     transform,  # noqa: F401
-    autotuner,  # noqa: F401
     language,  # noqa: F401
     engine,  # noqa: F401
 )
+from .autotuner import autotune  # noqa: F401
 from .transform import PassConfigKey  # noqa: F401
 
 from .engine import lower, register_cuda_postproc, register_hip_postproc  # noqa: F401