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

tilelang/VERSION

@@ -1 +1 @@
-0.1.4.post9
+0.1.4.post10

tilelang/intrinsics/mfma_layout.py

@@ -64,6 +64,12 @@ def shared_16x16_to_local_64x4_layout_B(i, j):
     return thread_id, local
 
 
+shared_16x16_to_local_64x4_layout_m_n = shared_16x16_to_local_64x4_layout_A
+shared_16x16_to_local_64x4_layout_n_k = shared_16x16_to_local_64x4_layout_A
+shared_16x16_to_local_64x4_layout_n_m = shared_16x16_to_local_64x4_layout_B
+shared_16x16_to_local_64x4_layout_k_n = shared_16x16_to_local_64x4_layout_B
+
+
 def thread_id_shared_access_64x4_to_16x16_layout_C_m_n(thread_id, local_id):
     i = local_id + (thread_id // 16) * 4
     j = thread_id % 16
@@ -100,6 +106,30 @@ def shared_16x32_to_local_64x8_layout_B(i, j):
     return thread_id, local
 
 
+def thread_id_shared_access_64x16_to_16x64_layout_A(thread_id, local_id):
+    i = thread_id % 16
+    j = local_id + (thread_id // 16) * 16
+    return i, j
+
+
+def shared_16x64_to_local_64x16_layout_A(i, j):
+    thread_id = i + 16 * (j // 16)
+    local = (j % 16)
+    return thread_id, local
+
+
+def thread_id_shared_access_64x16_to_16x64_layout_B(thread_id, local_id):
+    i = local_id + (thread_id // 16) * 16
+    j = thread_id % 16
+    return i, j
+
+
+def shared_16x64_to_local_64x16_layout_B(i, j):
+    thread_id = i + 16 * (j // 16)
+    local = (j % 16)
+    return thread_id, local
+
+
 def make_mfma_swizzle_layout(shared_buf, vecSize=8):
     dtype = shared_buf.dtype
     shape = shared_buf.shape

tilelang/intrinsics/mfma_macro_generator.py

@@ -150,12 +150,16 @@ class MatrixCoreIntrinEmitter(object):
             shared_16x16_to_local_64x4_layout_B,
             shared_16x32_to_local_64x8_layout_A,
             shared_16x32_to_local_64x8_layout_B,
+            shared_16x64_to_local_64x16_layout_A,
+            shared_16x64_to_local_64x16_layout_B,
             thread_id_shared_access_64x1_to_16x4_layout_A,
             thread_id_shared_access_64x1_to_4x16_layout_B,
             thread_id_shared_access_64x4_to_16x16_layout_A,
             thread_id_shared_access_64x4_to_16x16_layout_B,
             thread_id_shared_access_64x8_to_16x32_layout_A,
             thread_id_shared_access_64x8_to_16x32_layout_B,
+            thread_id_shared_access_64x16_to_16x64_layout_A,
+            thread_id_shared_access_64x16_to_16x64_layout_B,
         )
 
         k_dim = self.k_dim * self.k_pack
@@ -180,8 +184,15 @@ class MatrixCoreIntrinEmitter(object):
             if is_b:
                 index_map = shared_16x32_to_local_64x8_layout_A if transposed else shared_16x32_to_local_64x8_layout_B
                 reverse_index_map = thread_id_shared_access_64x8_to_16x32_layout_A if transposed else thread_id_shared_access_64x8_to_16x32_layout_B
+        elif k_dim == 64:
+            index_map = shared_16x64_to_local_64x16_layout_B if transposed else shared_16x64_to_local_64x16_layout_A
+            reverse_index_map = thread_id_shared_access_64x16_to_16x64_layout_B if transposed else thread_id_shared_access_64x16_to_16x64_layout_A
+
+            if is_b:
+                index_map = shared_16x64_to_local_64x16_layout_A if transposed else shared_16x64_to_local_64x16_layout_B
+                reverse_index_map = thread_id_shared_access_64x16_to_16x64_layout_A if transposed else thread_id_shared_access_64x16_to_16x64_layout_B
         else:
-            raise ValueError("k_dim must be 4 or 16 currently")
+            raise ValueError("k_dim must be 4 or 16 or 32 or 64 currently")
 
         return index_map, reverse_index_map
 
@@ -240,7 +251,7 @@ class MatrixCoreIntrinEmitter(object):
                 for i in T.serial(warp_rows):
                     for local_id in T.vectorized(k_pack * local_size_a):
                         row, col = T.meta_var(reverse_index_map(tx, local_id))
-                        l, r = (rk * chunk + ki * micro_size_k,
+                        l, r = (rk * chunk + ki * (k_pack * micro_size_k),
                                 warp_m * warp_row_tiles + i * micro_size_x)
                         A_local_buf[i * k_pack * local_size_a + local_id] = A_shared_buf[l + row,
                                                                                          r + col]
@@ -249,7 +260,7 @@ class MatrixCoreIntrinEmitter(object):
                     for local_id in T.vectorized(k_pack * local_size_a):
                         row, col = T.meta_var(reverse_index_map(tx, local_id))
                         l, r = (warp_m * warp_row_tiles + i * micro_size_x,
-                                rk * chunk + ki * micro_size_k)
+                                rk * chunk + ki * (k_pack * micro_size_k))
                         A_local_buf[i * k_pack * local_size_a + local_id] = A_shared_buf[l + row,
                                                                                          r + col]
 
@@ -284,7 +295,7 @@ class MatrixCoreIntrinEmitter(object):
                         row, col = T.meta_var(reverse_index_map(tx, local_id))
                         l, r = (
                             warp_n * warp_col_tiles + j * micro_size_y,
-                            rk * chunk + ki * micro_size_k,
+                            rk * chunk + ki * (k_pack * micro_size_k),
                         )
                         B_local_buf[j * k_pack * local_size_b + local_id] = B_shared_buf[l + row,
                                                                                          r + col]
@@ -293,7 +304,7 @@ class MatrixCoreIntrinEmitter(object):
                     for local_id in T.vectorized(k_pack * local_size_b):
                         row, col = T.meta_var(reverse_index_map(tx, local_id))
                         l, r = (
-                            rk * chunk + ki * micro_size_k,
+                            rk * chunk + ki * (k_pack * micro_size_k),
                             warp_n * warp_col_tiles + j * micro_size_y,
                         )
                         B_local_buf[j * k_pack * local_size_b + local_id] = B_shared_buf[l + row,
@@ -356,7 +367,7 @@ class MatrixCoreIntrinEmitter(object):
         def _warp_stmatrix_shared(C_local_buf, C_buf, thread_binding):
             tx, warp_n, warp_m = self.extract_thread_binding(thread_binding)
             for i, j in T.grid(warp_rows, warp_cols):
-                for local_id in T.serial(local_size_out):
+                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 +
@@ -366,7 +377,7 @@ class MatrixCoreIntrinEmitter(object):
         def _warp_stmatrix_global(C_local_buf, C_buf, thread_binding):
             tx, warp_n, warp_m = self.extract_thread_binding(thread_binding)
             for i, j in T.grid(warp_rows, warp_cols):
-                for local_id in T.serial(local_size_out):
+                for local_id in T.vectorized(local_size_out):
                     row, col = T.meta_var(mfma_store_index_map(tx, local_id))
                     C_buf[(pid_m * BLOCK_M + warp_m * warp_rows + i) * M_DIM + row,
                           (pid_n * BLOCK_N + warp_n * warp_cols + j) * N_DIM +

tilelang/jit/__init__.py

@@ -181,6 +181,9 @@ class _JitImplementation:
                 else:
                     raise ValueError(f"Invalid function type: {type(program_result_source)}")
 
+                if self.verbose:
+                    logger.info(f"Verbose: Compiling for program \n {program_result.script()}")
+
                 kernel_result = compile(
                     program_result,
                     out_idx=self.out_idx,

tilelang/src/tl_templates/hip/common.h

@@ -66,7 +66,7 @@ using float16x16 =
 
 using half_t = float16_t;
 
-using bfloat16_t = __hip_bfloat16;
+using bfloat16_t = hip_bfloat16;
 
 struct bfloat16x2 {
   bfloat16_t data[2];

tilelang/src/tl_templates/hip/gemm.h

@@ -19,16 +19,16 @@ template <> struct MfmaTraits<half> {
   }
 };
 
-// Specialization for __hip_bfloat16
-template <> struct MfmaTraits<__hip_bfloat16> {
+// Specialization for bfloat16_t
+template <> struct MfmaTraits<bfloat16_t> {
   template <typename AccType>
-  static TL_DEVICE void mfma_op(const __hip_bfloat16 *b,
-                                const __hip_bfloat16 *a, AccType *c) {
+  static TL_DEVICE void mfma_op(const bfloat16_t *b, const bfloat16_t *a,
+                                AccType *c) {
     bfloat16x4_vec b_vec, a_vec;
 
     // Reinterpret the pointers
-    short *b_short = reinterpret_cast<short *>(const_cast<__hip_bfloat16 *>(b));
-    short *a_short = reinterpret_cast<short *>(const_cast<__hip_bfloat16 *>(a));
+    short *b_short = reinterpret_cast<short *>(const_cast<bfloat16_t *>(b));
+    short *a_short = reinterpret_cast<short *>(const_cast<bfloat16_t *>(a));
 
     // Copy the data
     for (int i = 0; i < 4; ++i) {

{tilelang_rocm-0.1.4.post9.dist-info → tilelang_rocm-0.1.4.post10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tilelang-rocm
-Version: 0.1.4.post9
+Version: 0.1.4.post10
 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.post9.dist-info → tilelang_rocm-0.1.4.post10.dist-info}/RECORD

@@ -1,7 +1,7 @@
 tilelang/CMakeLists.txt,sha256=xJhnusYZI4UhD_fzseGH3Tn2BeovUzz3aWUwPq-WU0Y,7010
 tilelang/LICENSE,sha256=v9fVeAgRKQXc5ySwTns767gj0-dHN9XYPpGURkAVAXs,1127
 tilelang/README.md,sha256=1RC_2IUBY-p0BR-d2xkNXC8zrva8-U3AVkmCozkssbY,11924
-tilelang/VERSION,sha256=jHyORphLiyy1Y7K0gWps7qDNJc-NPepcM6_C-_a_2bI,12
+tilelang/VERSION,sha256=Il_uPi-DubG_kcKlTV5qP75Dq4zOcrugFwCAygQ4SQU,13
 tilelang/__init__.py,sha256=yH0BknCRnFQN-E7d6p1HPNbeY4o3COqG7XzR_EJpbTo,3215
 tilelang/_ffi_api.py,sha256=D-HfDxx8EZq6qItftg-ejOhpC_smIZLN-pWPVCNX_UM,243
 tilelang/config.cmake,sha256=370i6N3wwi7-LPGZDBtiiu54UWp39ndD-9lCurLhHwI,14330
@@ -6331,12 +6331,12 @@ tilelang/engine/lower.py,sha256=OJX7d_qk6WXA9VDF-fTTKw1q-ZnBkn4tuD9lA4TPdxk,8923
 tilelang/engine/param.py,sha256=5eWc48aao84WIrbtaLuAYGrb3RE1SyiG6hIy8DwlryI,3943
 tilelang/engine/phase.py,sha256=SPruc1tsC_yFD9q93pilkFmzY3UZ2YJTCRH06YPDI1Q,7154
 tilelang/intrinsics/__init__.py,sha256=ymvtsKjVY0f_9k-QIMtO4CEh6hEnG7H4NiW3buNLVQg,501
-tilelang/intrinsics/mfma_layout.py,sha256=KSA2TFy9XB9lC-mxFgqJX7ZuHNrsJLLJriVYbf8-wSc,3168
-tilelang/intrinsics/mfma_macro_generator.py,sha256=ws3Tyd0Q7laSbIcv-0Ff1YF80P_G4mooM4ir_aqiPTs,16783
+tilelang/intrinsics/mfma_layout.py,sha256=O5jh8gfILH-ASDss0pvTZhJq2jgmcyHwRPzigxGoGW8,4056
+tilelang/intrinsics/mfma_macro_generator.py,sha256=axkPrTN-Lplt0xFMk-ckQu8Qynve-lNl6thQCIwY1Q0,17653
 tilelang/intrinsics/mma_layout.py,sha256=eHFiNKd3zKzNFuRrpZdEQx0apbHWj8Ak6Q3e9_CeDiM,5090
 tilelang/intrinsics/mma_macro_generator.py,sha256=BVx3Bt3K67XSTp-Op425OrPxmtD43jyLre0wY8AeW-w,44478
 tilelang/intrinsics/utils.py,sha256=dbQpWOy0F4rg3WotzHQToPtJgY2BLtRy1CKkSnrs--k,4243
-tilelang/jit/__init__.py,sha256=KYrkfRFdQOyyF55iJKSo_IGOwXZNegLMV3Bi7WXOhLY,11901
+tilelang/jit/__init__.py,sha256=kUDaEwa4Xmr2lvaEI6SvFH4R8wWZWkTLbfgb0iPaEoE,12031
 tilelang/jit/env.py,sha256=Ih2qp1fTLcA0zuAJfkq4-yAQSXu3VCil_XjL-ndKYtc,2604
 tilelang/jit/kernel.py,sha256=B-j3H8kFYalZ3SENKQ3SKbN16Grlr0CKFB2RBZlawA4,15284
 tilelang/jit/param.py,sha256=GIIvPUyXSwi8mY_TMrZGOB9pcg7-hycz8PQNjNQ-JpU,1218
@@ -6383,10 +6383,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=nz9QoBmoqWOZAn68IDvMhMBk_c_etYJf9IathdtijJE,5040880
-tilelang/lib/libtilelang_module.so,sha256=FzKrdpHtsQGwLGD4pQKQp9-jhROAkrNh3JtslXs1EnY,5040880
-tilelang/lib/libtvm.so,sha256=OVNSaJ79WYV-AyY4NNqW4PAnQjS_4uCh6ZibrlzTN_c,83980800
-tilelang/lib/libtvm_runtime.so,sha256=77cdY-xJBotDchdDrqJEIT0Bhs5qw4XkxqGhORhpa3A,4794616
+tilelang/lib/libtilelang.so,sha256=gWKpADoERYHrDu-eOXjYtG8ak-dvv_2cs8WDiFYnc8o,5040880
+tilelang/lib/libtilelang_module.so,sha256=eStBjzb0e9y3lZVc21EefE1_nryt0Cs0wlYyz32QifI,5040880
+tilelang/lib/libtvm.so,sha256=7qy1XeUSGm4YXJC22LUmLZMfaGEPdTCXVxq3pkqS_ZY,83982112
+tilelang/lib/libtvm_runtime.so,sha256=BgUg1x-f2q3urmCbmBrbUtVFmS3SUyI6dBSAVP_yi8k,4794680
 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
@@ -6416,10 +6416,10 @@ tilelang/src/tl_templates/cuda/gemm_sm90.h,sha256=S3v63snxR_3TEg8LLCoR6cqGVgii8k
 tilelang/src/tl_templates/cuda/ldsm.h,sha256=TxCxYVzUK4tvUNVqULCL5HEaAuW9vOv0_-QYmoRFUkM,5053
 tilelang/src/tl_templates/cuda/reduce.h,sha256=U9mKcHSttin1FQ0BohpaP0bHvgPvb3t-czwTuDeK5-8,4394
 tilelang/src/tl_templates/cuda/threadblock_swizzle.h,sha256=GIXQwC1gzwUhnq4CzORHh5hA_QHVfMrOcUeGTy1Fon8,1945
-tilelang/src/tl_templates/hip/common.h,sha256=_JyiapMcqlUuwYaqNbgO52_mZgQxWz6cqW9au5aKzN0,3593
+tilelang/src/tl_templates/hip/common.h,sha256=KzMntW2OlABMwYjlxNqKiBDZQXQNrYApeVf8F9548s4,3591
 tilelang/src/tl_templates/hip/copy.h,sha256=fGHkbe4ReXoEtIWrgQ-mlCycaIL65SvNGWK1OJZdUQo,3324
 tilelang/src/tl_templates/hip/debug.h,sha256=9xGr4ka5x_nvY55XwbgTJFFwEnd09ta9jAZwjHyQau0,8231
-tilelang/src/tl_templates/hip/gemm.h,sha256=opIRbjhIXRYmzdi4vw8vqoVCZYEQXPBM0ckakgiezNg,11634
+tilelang/src/tl_templates/hip/gemm.h,sha256=lYeOjV8OG2oZbcS7ByzOudE7i0FQJ71mrUcImkfhTrg,11610
 tilelang/src/tl_templates/hip/hip_fp8.h,sha256=JYGiuuroLQH7CXT7IdKcpNUECmGOTe8DjIjcS9eLc0U,377
 tilelang/src/tl_templates/hip/ldsm.h,sha256=gRx_bSdsCsgcVumwUJwOnv4HuHruU2kC9TE9x_jo8k0,106
 tilelang/src/tl_templates/hip/reduce.h,sha256=-VKpG-TNbzPHIqsSReYpqZoM-oXFzIx6fMeBieV26Kc,1372
@@ -6437,8 +6437,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=SZ4ewoJ-Mq3zg8zIHS7-XLUmYDdlNwh841yUkjnQtNU,12573
-tilelang_rocm-0.1.4.post9.dist-info/licenses/LICENSE,sha256=v9fVeAgRKQXc5ySwTns767gj0-dHN9XYPpGURkAVAXs,1127
-tilelang_rocm-0.1.4.post9.dist-info/METADATA,sha256=0JNlrV6aGKHvKKiSYB1qgu8R4cnfceXp3lxj70wX5Ng,13075
-tilelang_rocm-0.1.4.post9.dist-info/WHEEL,sha256=0-G7woG4LgutcYzUGJCOYFgoh749-FtfhSMeIPLVGS0,104
-tilelang_rocm-0.1.4.post9.dist-info/top_level.txt,sha256=qvMq-AYkDVggI-9VIAzCe5CXHl66IEWj7J29-JbuFsI,21
-tilelang_rocm-0.1.4.post9.dist-info/RECORD,,
+tilelang_rocm-0.1.4.post10.dist-info/licenses/LICENSE,sha256=v9fVeAgRKQXc5ySwTns767gj0-dHN9XYPpGURkAVAXs,1127
+tilelang_rocm-0.1.4.post10.dist-info/METADATA,sha256=aBSwulXbZ0_bU2cIYKNCuRsUM2qlAsJ6O14B7fPJUoQ,13076
+tilelang_rocm-0.1.4.post10.dist-info/WHEEL,sha256=0-G7woG4LgutcYzUGJCOYFgoh749-FtfhSMeIPLVGS0,104
+tilelang_rocm-0.1.4.post10.dist-info/top_level.txt,sha256=qvMq-AYkDVggI-9VIAzCe5CXHl66IEWj7J29-JbuFsI,21
+tilelang_rocm-0.1.4.post10.dist-info/RECORD,,