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

tilelang/VERSION

@@ -1 +1 @@
-0.1.4.post4
+0.1.4.post5

tilelang/__init__.py

@@ -4,6 +4,7 @@
 import sys
 import os
 import ctypes
+import subprocess
 
 import logging
 from tqdm import tqdm
@@ -76,6 +77,67 @@ 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":

{tilelang_rocm-0.1.4.post4.dist-info → tilelang_rocm-0.1.4.post5.dist-info}/METADATA

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

{tilelang_rocm-0.1.4.post4.dist-info → tilelang_rocm-0.1.4.post5.dist-info}/RECORD

@@ -1,8 +1,8 @@
 tilelang/CMakeLists.txt,sha256=xJhnusYZI4UhD_fzseGH3Tn2BeovUzz3aWUwPq-WU0Y,7010
 tilelang/LICENSE,sha256=v9fVeAgRKQXc5ySwTns767gj0-dHN9XYPpGURkAVAXs,1127
 tilelang/README.md,sha256=bk1No4TTq4QGnHFvCryFAV2ZkqizZZy-4qSMwfAgM-4,11911
-tilelang/VERSION,sha256=WVgIvJCERXBRt77au2fBkozvMuBc6pwDsPZyewlakhw,12
-tilelang/__init__.py,sha256=IS1RDxLiuQkqvtMzTm73d2se5Fa_ryYeMr9YxtrWhCU,3159
+tilelang/VERSION,sha256=Z_-tN0NhQNsklCicBtnNgmKd2vtDbgzBQbNqXG__Rn0,12
+tilelang/__init__.py,sha256=uCqLGeGMgs0ik-NxKq5S-xC__caiMXzKTVFbsvQLKcM,5503
 tilelang/_ffi_api.py,sha256=D-HfDxx8EZq6qItftg-ejOhpC_smIZLN-pWPVCNX_UM,243
 tilelang/config.cmake,sha256=370i6N3wwi7-LPGZDBtiiu54UWp39ndD-9lCurLhHwI,14330
 tilelang/env.py,sha256=naNDQ5Icc7wx09-TtK4bC416IZd4NK4QEFvE1JdfPSA,8391
@@ -7326,10 +7326,10 @@ tilelang/layout/__init__.py,sha256=F1wr9yBG9GW84h8KWXz-hRJFfqyZuY0EKSrG08KyrWQ,2
 tilelang/layout/fragment.py,sha256=zTv9P96lsYi9BWc5pxR4PA2Z5RSDGP7D5uJCiNw7_oc,8445
 tilelang/layout/layout.py,sha256=20CWxz_S8k_WNvWiR4gdIrEsQ36e5bsnOEqmu4zGk_c,4311
 tilelang/layout/swizzle.py,sha256=PMqu_s1sNCh9uo8eDs5qmLKXnDqZwv34GT3H9D4YDO0,438
-tilelang/lib/libtilelang.so,sha256=bkRIVkuti7c91AYraa4qCDQSP2cPzb2oc3TNOZ2fvww,4760416
-tilelang/lib/libtilelang_module.so,sha256=tFH0b1ay8k5BxLvBlPNpJbyNXIGnbaRVOHRGIq7s0es,4760416
-tilelang/lib/libtvm.so,sha256=DqL8HQezwF2SLmeQvoQuQ51o7E7_0EBg2P-lbmTf6DU,80820960
-tilelang/lib/libtvm_runtime.so,sha256=U7r_vMlPJIYT120_w2KtJg5-X4xYXdcaK3Ck1-WV1kk,4759920
+tilelang/lib/libtilelang.so,sha256=Pv-058HPWNR37AMY1uyTk7oKItNvmYY3Lm4cWE84a6s,4760416
+tilelang/lib/libtilelang_module.so,sha256=_S-AaWcyKqQJYYhLhKO5riz83PsDIGXHYG4mAKrO2XA,4760416
+tilelang/lib/libtvm.so,sha256=Pfs0fV-CUGps28TWQXQaqwUKTmB8Zk10ix1BcG_i7js,80820960
+tilelang/lib/libtvm_runtime.so,sha256=QOK7ropRHXIxXxwoHzJT85Qw76c28F8yJpizx1XVK8I,4759920
 tilelang/math/__init__.py,sha256=JC4fqrU_LV_wDErti-wHNr4j6_mqP1PsK0qqkhaSzRU,209
 tilelang/primitives/__init__.py,sha256=10gQN3QWUFM1nkGXY46QFcWUXxwsKMsVn23JdyFHil4,167
 tilelang/primitives/gemm/__init__.py,sha256=j62ObmbL5Q6m3lSouNBQDk1hZZRnSp4UNNCCaSlKYXU,1658
@@ -7375,8 +7375,8 @@ tilelang/utils/deprecated.py,sha256=CiZ9y_76_dZ24SFDdasDiLmibwi6xO2Gdj6WzTWU0Qg,
 tilelang/utils/language.py,sha256=KUzUZ8Z2x1np0Hu_MrjWOIcRrVAZHX90li1Xw9fYZXY,3291
 tilelang/utils/target.py,sha256=P-74pdCLWcp2MZMQUoPIFwKF1NZ1QT-L0VroIL8m2to,2486
 tilelang/utils/tensor.py,sha256=8ZnSw8E8cBesbVI71iEQWmVHIFUwcf6GyjblluX2C4c,12523
-tilelang_rocm-0.1.4.post4.dist-info/licenses/LICENSE,sha256=v9fVeAgRKQXc5ySwTns767gj0-dHN9XYPpGURkAVAXs,1127
-tilelang_rocm-0.1.4.post4.dist-info/METADATA,sha256=AQV52vbAcmP0exrGveVIFcPJ6-sHomU1H_9yPI0bGjI,13062
-tilelang_rocm-0.1.4.post4.dist-info/WHEEL,sha256=tjHrT7-JLZ0qMFFxr_WeEipEQugV2Q-OnKugRM4coRw,104
-tilelang_rocm-0.1.4.post4.dist-info/top_level.txt,sha256=qvMq-AYkDVggI-9VIAzCe5CXHl66IEWj7J29-JbuFsI,21
-tilelang_rocm-0.1.4.post4.dist-info/RECORD,,
+tilelang_rocm-0.1.4.post5.dist-info/licenses/LICENSE,sha256=v9fVeAgRKQXc5ySwTns767gj0-dHN9XYPpGURkAVAXs,1127
+tilelang_rocm-0.1.4.post5.dist-info/METADATA,sha256=uwPUMUQofy-72iLum1mjIBSaDJhEfZkUmiKFTcdM5yA,13062
+tilelang_rocm-0.1.4.post5.dist-info/WHEEL,sha256=tjHrT7-JLZ0qMFFxr_WeEipEQugV2Q-OnKugRM4coRw,104
+tilelang_rocm-0.1.4.post5.dist-info/top_level.txt,sha256=qvMq-AYkDVggI-9VIAzCe5CXHl66IEWj7J29-JbuFsI,21
+tilelang_rocm-0.1.4.post5.dist-info/RECORD,,