diff --git a/.gitmodules b/.gitmodules index c8a359670..d5b545545 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,7 +1,7 @@ [submodule "3rdparty/tvm"] path = 3rdparty/tvm url = https://github.com/TileLang/tvm.git - branch = tilelang + branch = upstream [submodule "3rdparty/cutlass"] path = 3rdparty/cutlass url = https://github.com/TileLang/cutlass diff --git a/3rdparty/tvm b/3rdparty/tvm index 27078affb..e1c5b0897 160000 --- a/3rdparty/tvm +++ b/3rdparty/tvm @@ -1 +1 @@ -Subproject commit 27078affbe26b65d690d505f67178734d5c52629 +Subproject commit e1c5b089737e47a3849afa87df2432c13b633594 diff --git a/bitblas/ops/base_scheduler.py b/bitblas/ops/base_scheduler.py index acdc057d6..f18c98026 100644 --- a/bitblas/ops/base_scheduler.py +++ b/bitblas/ops/base_scheduler.py @@ -70,3 +70,13 @@ def common_header(self): # TODO(lei): For HIP Backend it should be different common_header = "#include \n" return common_header + + +# Decorator to simplify the output of a function +def simplify_prim_func(func: Callable): + + def wrapper(*args, **kwargs): + stmt: Union[PrimFunc, IRModule] = (func)(*args, **kwargs) + return BaseScheduler.Simplify(stmt) + + return wrapper diff --git a/bitblas/ops/general_matmul/tilelang/dense/matmul_tensorcore.py b/bitblas/ops/general_matmul/tilelang/dense/matmul_tensorcore.py index eea256fd9..13658aab4 100644 --- a/bitblas/ops/general_matmul/tilelang/dense/matmul_tensorcore.py +++ b/bitblas/ops/general_matmul/tilelang/dense/matmul_tensorcore.py @@ -424,7 +424,9 @@ def apply_config( threads = warp_size * (block_row_warps * block_col_warps) # Calculate local fragment sizes for tensor core - local_size = (micro_size_x * micro_size_y) // warp_size + local_size_a = (micro_size_x * micro_size_k) // warp_size + local_size_b = (micro_size_y * micro_size_k) // warp_size + local_size_c = (micro_size_x * micro_size_y) // warp_size warp_rows = warp_row_tiles // micro_size_x warp_cols = warp_col_tiles // micro_size_y @@ -459,9 +461,9 @@ def main( A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope=shared_scope) B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope=shared_scope) C_shared = T.alloc_shared(C_shared_shape, out_dtype, scope=shared_scope) - A_local = T.alloc_local((warp_rows * local_size), in_dtype) - B_local = T.alloc_local((warp_cols * local_size), in_dtype) - C_local = T.alloc_local((warp_rows * warp_cols * local_size), accum_dtype) + A_local = T.alloc_local((warp_rows * local_size_a), in_dtype) + B_local = T.alloc_local((warp_cols * local_size_b), in_dtype) + C_local = T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype) # Thread-level parallelism for Tensor Cores thread_bindings = T.thread_binding(0, threads, "threadIdx.x") diff --git a/bitblas/ops/general_matmul/tilelang/dequantize/finegrained_primitive_tensorcore.py b/bitblas/ops/general_matmul/tilelang/dequantize/finegrained_primitive_tensorcore.py index d755ba2f8..d57951455 100644 --- a/bitblas/ops/general_matmul/tilelang/dequantize/finegrained_primitive_tensorcore.py +++ b/bitblas/ops/general_matmul/tilelang/dequantize/finegrained_primitive_tensorcore.py @@ -231,7 +231,9 @@ def apply_config( block_K = chunk threads = warp_size * (block_row_warps * block_col_warps) - fragement_size = (micro_size_x * micro_size_y) // warp_size + fragement_size_a = (micro_size_x * micro_size_k) // warp_size + fragement_size_b = (micro_size_y * micro_size_k) // warp_size + fragement_size_c = (micro_size_x * micro_size_y) // warp_size warp_rows = warp_row_tiles // micro_size_x warp_cols = warp_col_tiles // micro_size_y @@ -318,9 +320,9 @@ def general_dequant_matmul( B_dequantize_shared = T.alloc_shared(B_dequantize_shared_shape, in_dtype) C_shared = T.alloc_shared(C_shared_shape, out_dtype) - A_frag = T.alloc_local((warp_rows * fragement_size), in_dtype) - B_frag = T.alloc_local((warp_cols * fragement_size), in_dtype) - C_frag = T.alloc_local((warp_rows * warp_cols * fragement_size), accum_dtype) + A_frag = T.alloc_local((warp_rows * fragement_size_a), in_dtype) + B_frag = T.alloc_local((warp_cols * fragement_size_b), in_dtype) + C_frag = T.alloc_local((warp_rows * warp_cols * fragement_size_c), accum_dtype) B_local = T.alloc_local([local_size_compressed], storage_dtype) B_dequantize_local = T.alloc_local([local_size], in_dtype) diff --git a/bitblas/ops/general_matmul/tilelang/dequantize/ladder_weight_transform_tensorcore.py b/bitblas/ops/general_matmul/tilelang/dequantize/ladder_weight_transform_tensorcore.py index bb463e59a..7f8920575 100644 --- a/bitblas/ops/general_matmul/tilelang/dequantize/ladder_weight_transform_tensorcore.py +++ b/bitblas/ops/general_matmul/tilelang/dequantize/ladder_weight_transform_tensorcore.py @@ -71,7 +71,9 @@ def apply_config( block_K = chunk threads = warp_size * (block_row_warps * block_col_warps) - fragement_size = (micro_size_x * micro_size_y) // warp_size + fragement_size_a = (micro_size_x * micro_size_k) // warp_size + fragement_size_b = (micro_size_y * micro_size_k) // warp_size + fragement_size_c = (micro_size_x * micro_size_y) // warp_size warp_rows = warp_row_tiles // micro_size_x warp_cols = warp_col_tiles // micro_size_y @@ -173,11 +175,11 @@ def general_dequant_matmul( B_shared = T.alloc_shared(B_shared_shape, storage_dtype) C_shared = T.alloc_shared(C_shared_shape, out_dtype) - A_frag = T.alloc_local((warp_rows * fragement_size), in_dtype) - B_frag = T.alloc_local((warp_cols * fragement_size // num_elems_per_byte), + A_frag = T.alloc_local((warp_rows * fragement_size_a), in_dtype) + B_frag = T.alloc_local((warp_cols * fragement_size_b // num_elems_per_byte), storage_dtype) - B_dequantize_frag = T.alloc_local((warp_cols * fragement_size), in_dtype) - C_frag = T.alloc_local((warp_rows * warp_cols * fragement_size), accum_dtype) + B_dequantize_frag = T.alloc_local((warp_cols * fragement_size_b), in_dtype) + C_frag = T.alloc_local((warp_rows * warp_cols * fragement_size_c), accum_dtype) tx = T.thread_binding(0, threads, thread="threadIdx.x") diff --git a/bitblas/tl/macro_generator.py b/bitblas/tl/macro_generator.py index 0f7adb791..63433a52a 100644 --- a/bitblas/tl/macro_generator.py +++ b/bitblas/tl/macro_generator.py @@ -33,19 +33,21 @@ class TensorCoreIntrinEmitter(object): "e5m2_float8": "e5m2", } - def __init__(self, - a_dtype="float16", - b_dtype="float16", - accum_dtype="float16", - a_transposed=False, - b_transposed=False, - block_row_warps=2, - block_col_warps=2, - warp_row_tiles=8, - warp_col_tiles=8, - chunk=16, - reduce_k=1, - num_elems_per_byte=1): + def __init__( + self, + a_dtype="float16", + b_dtype="float16", + accum_dtype="float16", + a_transposed=False, + b_transposed=False, + block_row_warps=2, + block_col_warps=2, + warp_row_tiles=8, + warp_col_tiles=8, + chunk=16, + reduce_k=1, + num_elems_per_byte=1, + ): self.a_dtype = a_dtype self.b_dtype = b_dtype self.accum_dtype = accum_dtype @@ -65,7 +67,7 @@ def __init__(self, self.warp_rows = warp_row_tiles // self.micro_size_x self.warp_cols = warp_col_tiles // self.micro_size_y self.reduce_k = reduce_k - self.threads = self.WARP_SIZE * (block_row_warps * block_col_warps) * reduce_k + self.threads = (self.WARP_SIZE * (block_row_warps * block_col_warps) * reduce_k) self.num_elems_per_byte = num_elems_per_byte def _initialize_k_dim(self, a_dtype="float16"): @@ -96,130 +98,172 @@ def _initialize_micro_size(self, m_dim=16, n_dim=16, k_dim=16): self.micro_size_y = n_dim self.micro_size_k = k_dim - @T.macro - def _warp_ldmatrix_a( - inst, - A_local_buf, - A_shared_buf, - ki, - thread_bindings, - rk=0, - ): - stride = A_shared_buf.shape[-1] - tx = thread_bindings % inst.WARP_SIZE - ty = (thread_bindings // inst.WARP_SIZE) % inst.block_row_warps - - for i in T.serial(inst.warp_rows): - T.ptx_ldmatrix( - inst.a_dtype, - T.bool(False), - 4, - ".b16", - A_local_buf.data, - i * inst.local_size_a, - T.address_of(A_shared_buf[ - ty * inst.warp_row_tiles + i * inst.micro_size_x, - rk * inst.chunk + ki * inst.micro_size_k, - ]), - get_ldmatrix_offset("A", tx, 0, stride, inst.a_dtype, inst.a_transposed), - ) - - @T.macro - def _warp_ldmatrix_b( - inst, - B_local_buf, - B_shared_buf, - ki, - thread_bindings, - rk=0, - ): - stride = B_shared_buf.shape[-1] - tx = thread_bindings % inst.WARP_SIZE - tz = (thread_bindings // (inst.WARP_SIZE * inst.block_row_warps)) % inst.block_col_warps - - for j in T.serial(inst.warp_cols): - # Assign B_shared_elem - ri, rj = tz * inst.warp_col_tiles + j * inst.micro_size_y, rk * inst.chunk + ki * inst.micro_size_k - B_shared_elem = B_shared_buf[ri, rj] - - T.ptx_ldmatrix( - inst.b_dtype, - T.bool(False), # TODO(lei): should be optimized - 4, - ".b16", - B_local_buf.data, - j * inst.local_size_b, - T.address_of(B_shared_elem), - get_ldmatrix_offset("B", tx, 0, stride, inst.b_dtype, inst.b_transposed), - ) - - @T.macro - def _warp_mma(inst, A_local_buf, B_local_buf, C_local_buf): - for i, j in T.grid(inst.warp_rows, inst.warp_cols): - T.ptx_mma( - inst.accum_dtype, - inst.mma_prefix, - "row", - "col", - inst.a_dtype_abbrv, - inst.b_dtype_abbrv, - inst.accum_dtype_abbrv, - A_local_buf.data, - i * inst.local_size_a, - B_local_buf.data, - j * inst.local_size_b, - C_local_buf.data, - i * inst.warp_cols * inst.local_size_out + j * inst.local_size_out, - T.bool(False), - ) - - T.ptx_mma( - inst.accum_dtype, - inst.mma_prefix, - "row", - "col", - inst.a_dtype_abbrv, - inst.b_dtype_abbrv, - inst.accum_dtype_abbrv, - A_local_buf.data, - i * inst.local_size_a, - B_local_buf.data, - j * inst.local_size_b + lift(inst.local_size_b) // 2, - C_local_buf.data, - i * inst.warp_cols * inst.local_size_out + j * inst.local_size_out + - lift(inst.local_size_out) // 2, - T.bool(False), - ) - - # STS - # MMA Store must be in simulated instead of TVM Intrins - # As TVM Intrins is like a hack that the threadIdx.x should be always - # equal to the warp_size - @T.macro - def _warp_stmatrix(inst, C_local_buf, C_shared_buf, thread_bindings): - tx = thread_bindings % inst.WARP_SIZE - ty = (thread_bindings // inst.WARP_SIZE) % inst.block_row_warps - tz = (thread_bindings // (inst.WARP_SIZE * inst.block_row_warps)) % inst.block_col_warps - for i, j in T.grid(inst.warp_rows, inst.warp_cols): - for local_id_o in T.serial(inst.local_size_out // 2): - for local_id_i in T.vectorized(2): - local_id = local_id_o * 2 + local_id_i - row, col = T.meta_var(mma_store_index_map(tx, local_id)) - C_shared_buf[ty * inst.warp_rows + i, tz * inst.warp_cols + j, row, - col] = C_local_buf[i * (inst.warp_cols * inst.local_size_out) + - j * inst.local_size_out + local_id] - def ldmatrix_a(self, A_local_buf, A_shared_buf, ki, thread_bindings, rk=0): - return self._warp_ldmatrix_a(self, A_local_buf, A_shared_buf, ki, thread_bindings, rk) + WARP_SIZE = self.WARP_SIZE + block_row_warps = self.block_row_warps + warp_row_tiles = self.warp_row_tiles + warp_rows = self.warp_rows + chunk = self.chunk + micro_size_x = self.micro_size_x + micro_size_k = self.micro_size_k + a_dtype = self.a_dtype + a_transposed = self.a_transposed + local_size_a = self.local_size_a + + @T.macro + def _warp_ldmatrix_a( + A_local_buf, + A_shared_buf, + ki, + thread_bindings, + rk=0, + ): + stride = A_shared_buf.shape[-1] + tx = thread_bindings % WARP_SIZE + ty = (thread_bindings // WARP_SIZE) % block_row_warps + + for i in T.serial(warp_rows): + T.ptx_ldmatrix( + a_dtype, + T.bool(False), + 4, + ".b16", + A_local_buf.data, + i * local_size_a, + T.address_of(A_shared_buf[ + ty * warp_row_tiles + i * micro_size_x, + rk * chunk + ki * micro_size_k, + ]), + get_ldmatrix_offset("A", tx, 0, stride, a_dtype, a_transposed), + ) + + return _warp_ldmatrix_a(A_local_buf, A_shared_buf, ki, thread_bindings, rk) def ldmatrix_b(self, B_local_buf, B_shared_buf, ki, thread_bindings, rk=0): - return self._warp_ldmatrix_b(self, B_local_buf, B_shared_buf, ki, thread_bindings, rk) + + WARP_SIZE = self.WARP_SIZE + block_row_warps = self.block_row_warps + block_col_warps = self.block_col_warps + warp_col_tiles = self.warp_col_tiles + warp_cols = self.warp_cols + chunk = self.chunk + micro_size_y = self.micro_size_y + micro_size_k = self.micro_size_k + local_size_b = self.local_size_b + b_dtype = self.b_dtype + b_transposed = self.b_transposed + + @T.macro + def _warp_ldmatrix_b( + B_local_buf, + B_shared_buf, + ki, + thread_bindings, + rk=0, + ): + stride = B_shared_buf.shape[-1] + tx = thread_bindings % WARP_SIZE + tz = (thread_bindings // (WARP_SIZE * block_row_warps)) % block_col_warps + + for j in T.serial(warp_cols): + # Assign B_shared_elem + ri, rj = ( + tz * warp_col_tiles + j * micro_size_y, + rk * chunk + ki * micro_size_k, + ) + B_shared_elem = B_shared_buf[ri, rj] + + T.ptx_ldmatrix( + b_dtype, + T.bool(False), # TODO(lei): should be optimized + 4, + ".b16", + B_local_buf.data, + j * local_size_b, + T.address_of(B_shared_elem), + get_ldmatrix_offset("B", tx, 0, stride, b_dtype, b_transposed), + ) + + return _warp_ldmatrix_b(B_local_buf, B_shared_buf, ki, thread_bindings, rk) def mma(self, A_local_buf, B_local_buf, C_local_buf): - return self._warp_mma(self, A_local_buf, B_local_buf, C_local_buf) + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + accum_dtype = self.accum_dtype + accum_dtype_abbrv = self.accum_dtype_abbrv + mma_prefix = self.mma_prefix + + @T.macro + def _warp_mma(A_local_buf, B_local_buf, C_local_buf): + for i, j in T.grid(warp_rows, warp_cols): + T.ptx_mma( + accum_dtype, + mma_prefix, + "row", + "col", + a_dtype_abbrv, + b_dtype_abbrv, + accum_dtype_abbrv, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out, + T.bool(False), + ) + + T.ptx_mma( + accum_dtype, + mma_prefix, + "row", + "col", + a_dtype_abbrv, + b_dtype_abbrv, + accum_dtype_abbrv, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b + lift(local_size_b) // 2, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out + lift(local_size_out) // 2, + T.bool(False), + ) + + return _warp_mma(A_local_buf, B_local_buf, C_local_buf) def stmatrix(self, C_local_buf, C_shared_buf, thread_bindings): - return self._warp_stmatrix(self, C_local_buf, C_shared_buf, thread_bindings) + WARP_SIZE = self.WARP_SIZE + block_row_warps = self.block_row_warps + block_col_warps = self.block_col_warps + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_out = self.local_size_out + + # STS + # MMA Store must be in simulated instead of TVM Intrins + # As TVM Intrins is like a hack that the threadIdx.x should be always + # equal to the warp_size + @T.macro + def _warp_stmatrix(C_local_buf, C_shared_buf, thread_bindings): + tx = thread_bindings % WARP_SIZE + ty = (thread_bindings // WARP_SIZE) % block_row_warps + tz = (thread_bindings // (WARP_SIZE * block_row_warps)) % block_col_warps + for i, j in T.grid(warp_rows, warp_cols): + for local_id_o in T.serial(local_size_out // 2): + for local_id_i in T.vectorized(2): + local_id = local_id_o * 2 + local_id_i + row, col = T.meta_var(mma_store_index_map(tx, local_id)) + C_shared_buf[ty * warp_rows + i, tz * warp_cols + j, row, + col] = C_local_buf[i * (warp_cols * local_size_out) + + j * local_size_out + local_id] + + return _warp_stmatrix(C_local_buf, C_shared_buf, thread_bindings) class TensorCoreIntrinEmitterWithLadderTransform(TensorCoreIntrinEmitter): @@ -307,91 +351,124 @@ def _initialize_transform_kind(self, transform_kind_a, transform_kind_b): assert transform_kind_b in [0, 3], "Currently only support 0 and 3" - @T.macro - def _warp_ldmatrix_b( - inst, - B_local_buf, - B_shared_buf, - ki, - thread_bindings, - rk=0, - ): - stride = B_shared_buf.shape[-1] - tx = thread_bindings % inst.WARP_SIZE - tz = (thread_bindings // (inst.WARP_SIZE * inst.block_row_warps)) % inst.block_col_warps - - if inst.transform_kind_b < TransformKind.LDMatrixTransform: - for j in T.serial(inst.warp_cols): - # Assign B_shared_elem - ri, rj = tz * inst.warp_col_tiles + j * inst.micro_size_y, rk * inst.chunk + ki * inst.micro_size_k - ni, nj, nii, njj = (ri) // inst.micro_size_y, (rj) // inst.micro_size_k, ( - ri) % inst.micro_size_y, (rj) % inst.micro_size_k - args = (ni, nj, nii, njj) if inst.transform_kind_b > 0 else (ri, rj) - B_shared_elem = B_shared_buf[args] + def ldmatrix_b(self, B_local_buf, B_shared_buf, ki, thread_bindings, rk=0): + WARP_SIZE = self.WARP_SIZE + block_row_warps = self.block_row_warps + block_col_warps = self.block_col_warps + warp_col_tiles = self.warp_col_tiles + warp_cols = self.warp_cols + chunk = self.chunk + micro_size_y = self.micro_size_y + micro_size_k = self.micro_size_k + local_size_b = self.local_size_b + b_dtype = self.b_dtype + transform_kind_b = self.transform_kind_b + b_transposed = self.b_transposed + num_elems_per_byte = self.num_elems_per_byte + + @T.macro + def _warp_ldmatrix_b( + B_local_buf, + B_shared_buf, + ki, + thread_bindings, + rk=0, + ): + stride = B_shared_buf.shape[-1] + tx = thread_bindings % WARP_SIZE + tz = (thread_bindings // (WARP_SIZE * block_row_warps)) % block_col_warps + + if transform_kind_b < TransformKind.LDMatrixTransform: + for j in T.serial(warp_cols): + # Assign B_shared_elem + ri, rj = ( + tz * warp_col_tiles + j * micro_size_y, + rk * chunk + ki * micro_size_k, + ) + ni, nj, nii, njj = ( + (ri) // micro_size_y, + (rj) // micro_size_k, + (ri) % micro_size_y, + (rj) % micro_size_k, + ) + args = ((ni, nj, nii, njj) if transform_kind_b > 0 else (ri, rj)) + B_shared_elem = B_shared_buf[args] + + T.ptx_ldmatrix( + b_dtype, + T.bool(False), # TODO(lei): should be optimized + 4, + ".b16", + B_local_buf.data, + j * local_size_b, + T.address_of(B_shared_elem), + get_ldmatrix_offset("B", tx, 0, stride, b_dtype, b_transposed), + ) + else: + local_size_dequantize = local_size_b // num_elems_per_byte + for j in T.serial(warp_cols): + for local_id in T.vectorized(local_size_dequantize): + # Assign B_shared_elem + ri, rj = ( + tz * warp_cols + j, + rk * (chunk // micro_size_k) + ki, + ) + rii, rjj = (tx * local_size_dequantize + + local_id) // (micro_size_k // num_elems_per_byte), ( + tx * local_size_dequantize + local_id) % ( + micro_size_k // num_elems_per_byte) + B_local_buf[j * local_size_dequantize + local_id] = ( + B_shared_buf[ri, rj, rii, rjj]) + + return _warp_ldmatrix_b(B_local_buf, B_shared_buf, ki, thread_bindings, rk) - T.ptx_ldmatrix( - inst.b_dtype, - T.bool(False), # TODO(lei): should be optimized - 4, - ".b16", + def mma(self, A_local_buf, B_local_buf, C_local_buf): + warp_rows = self.warp_rows + warp_cols = self.warp_cols + local_size_a = self.local_size_a + local_size_b = self.local_size_b + local_size_out = self.local_size_out + a_dtype_abbrv = self.a_dtype_abbrv + b_dtype_abbrv = self.b_dtype_abbrv + accum_dtype = self.accum_dtype + accum_dtype_abbrv = self.accum_dtype_abbrv + mma_prefix = self.mma_prefix + + @T.macro + def _warp_mma(A_local_buf, B_local_buf, C_local_buf): + for i, j in T.grid(warp_rows, warp_cols): + T.ptx_mma( + accum_dtype, + mma_prefix, + "row", + "col", + a_dtype_abbrv, + b_dtype_abbrv, + accum_dtype_abbrv, + A_local_buf.data, + i * local_size_a, B_local_buf.data, - j * inst.local_size_b, - T.address_of(B_shared_elem), - get_ldmatrix_offset("B", tx, 0, stride, inst.b_dtype, inst.b_transposed), + j * local_size_b, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out, + T.bool(False), ) - else: - local_size_dequantize = inst.local_size_b // inst.num_elems_per_byte - for j in T.serial(inst.warp_cols): - for local_id in T.vectorized(local_size_dequantize): - # Assign B_shared_elem - ri, rj = tz * inst.warp_cols + j, rk * (inst.chunk // inst.micro_size_k) + ki - rii, rjj = (tx * local_size_dequantize + - local_id) // (inst.micro_size_k // inst.num_elems_per_byte), ( - tx * local_size_dequantize + local_id) % ( - inst.micro_size_k // inst.num_elems_per_byte) - B_local_buf[j * local_size_dequantize + local_id] = B_shared_buf[ri, rj, rii, - rjj] - - @T.macro - def _warp_mma(inst, A_local_buf, B_local_buf, C_local_buf): - for i, j in T.grid(inst.warp_rows, inst.warp_cols): - T.ptx_mma( - inst.accum_dtype, - inst.mma_prefix, - "row", - "col", - inst.a_dtype_abbrv, - inst.b_dtype_abbrv, - inst.accum_dtype_abbrv, - A_local_buf.data, - i * inst.local_size_a, - B_local_buf.data, - j * inst.local_size_b, - C_local_buf.data, - i * inst.warp_cols * inst.local_size_out + j * inst.local_size_out, - T.bool(False), - ) - - T.ptx_mma( - inst.accum_dtype, - inst.mma_prefix, - "row", - "col", - inst.a_dtype_abbrv, - inst.b_dtype_abbrv, - inst.accum_dtype_abbrv, - A_local_buf.data, - i * inst.local_size_a, - B_local_buf.data, - j * inst.local_size_b + lift(inst.local_size_b) // 2, - C_local_buf.data, - i * inst.warp_cols * inst.local_size_out + j * inst.local_size_out + - lift(inst.local_size_out) // 2, - T.bool(False), - ) - def ldmatrix_b(self, B_local_buf, B_shared_buf, ki, thread_bindings, rk=0): - return self._warp_ldmatrix_b(self, B_local_buf, B_shared_buf, ki, thread_bindings, rk) + T.ptx_mma( + accum_dtype, + mma_prefix, + "row", + "col", + a_dtype_abbrv, + b_dtype_abbrv, + accum_dtype_abbrv, + A_local_buf.data, + i * local_size_a, + B_local_buf.data, + j * local_size_b + lift(local_size_b) // 2, + C_local_buf.data, + i * warp_cols * local_size_out + j * local_size_out + lift(local_size_out) // 2, + T.bool(False), + ) - def mma(self, A_local_buf, B_local_buf, C_local_buf): - return self._warp_mma(self, A_local_buf, B_local_buf, C_local_buf) + return _warp_mma(A_local_buf, B_local_buf, C_local_buf) diff --git a/bitblas/tl/mma_layout.py b/bitblas/tl/mma_layout.py index 8be21a1d1..719885be5 100644 --- a/bitblas/tl/mma_layout.py +++ b/bitblas/tl/mma_layout.py @@ -14,15 +14,15 @@ def ldmatrix_trans_32x8_to_shared_16x16_layout(thread_id, local_id): return row, col -def ldmatrix_32x16_to_shared_16x32_layout_a(thread_id, local_id): +def ldmatrix_16x32_to_shared_16x32_layout_a(thread_id, local_id): row = thread_id % 16 - col = local_id + (thread_id // 16) * 16 + col = 16 * (thread_id // 16) + local_id % 16 return row, col -def ldmatrix_32x16_to_shared_16x32_layout_b(thread_id, local_id): - row = (thread_id // 16) * 8 + (thread_id % 8) - col = local_id + 16 * ((thread_id % 16) // 8) +def ldmatrix_16x32_to_shared_16x32_layout_b(thread_id, local_id): + row = 8 * (thread_id // 16) + (thread_id % 8) + col = 16 * ((thread_id % 16) // 8) + local_id % 16 return row, col diff --git a/bitblas/tl/utils.py b/bitblas/tl/utils.py index 4b8b4cf6e..053dbe4d5 100644 --- a/bitblas/tl/utils.py +++ b/bitblas/tl/utils.py @@ -8,8 +8,8 @@ from .mma_layout import ( ldmatrix_32x8_to_shared_16x16_layout, ldmatrix_trans_32x8_to_shared_16x16_layout, - ldmatrix_32x16_to_shared_16x32_layout_a, - ldmatrix_32x16_to_shared_16x32_layout_b, + ldmatrix_16x32_to_shared_16x32_layout_a, + ldmatrix_16x32_to_shared_16x32_layout_b, mma_store_32x8_to_shared_16x16_layout, ) @@ -70,28 +70,40 @@ def get_swizzle_layout(row_idx, col_idx, row_size, dtype: Union[DataType, str]): return row_idx, ana.simplify(new_col_idx_outer * bank_elems + col_idx_inner) +# the original implementation and insight is from the following code snippet +# 3rdparty/tvm/python/tvm/tir/tensor_intrin/cuda.py#get_ldmatrix_intrin def get_ldmatrix_offset( matrix: Literal["A", "B"], row_idx, col_idx, stride, dtype: Literal["float16", "int8"] = "float16", - transpose: bool = False, + transposed: bool = False, ): assert matrix in ["A", "B"], "matrix should be either A or B" - transform_func = ( - ldmatrix_32x8_to_shared_16x16_layout - if dtype in ["float16", "bfloat16"] else ldmatrix_32x16_to_shared_16x32_layout_b) - transform_func_trans = ( - ldmatrix_trans_32x8_to_shared_16x16_layout - if dtype in ["float16", "bfloat16"] else ldmatrix_32x16_to_shared_16x32_layout_a) - if matrix == "A": - assert not transpose, "A matrix should not be transposed" - new_row_idx, new_col_idx = transform_func(row_idx, col_idx) - return new_row_idx * stride + new_col_idx + dtype_bits = DataType(dtype).bits + if dtype_bits == 16: + transform_func = ldmatrix_32x8_to_shared_16x16_layout + transform_func_trans = ldmatrix_trans_32x8_to_shared_16x16_layout + if transposed: + new_row_idx, new_col_idx = transform_func_trans(row_idx, col_idx) + return new_row_idx * stride + new_col_idx + else: + new_row_idx, new_col_idx = transform_func(row_idx, col_idx) + return new_row_idx * stride + new_col_idx + elif dtype_bits == 8: + if matrix == "B" and transposed: + transform_func = ldmatrix_16x32_to_shared_16x32_layout_b + new_row_idx, new_col_idx = transform_func(row_idx, col_idx) + return new_row_idx * stride + new_col_idx + elif matrix == "A" and not transposed: + transform_func = ldmatrix_16x32_to_shared_16x32_layout_a + new_row_idx, new_col_idx = transform_func(row_idx, col_idx) + return new_row_idx * stride + new_col_idx + else: + raise ValueError("ldmatrix only supports B transposed and A non-transposed for int8") else: - new_row_idx, new_col_idx = transform_func_trans(row_idx, col_idx) - return new_row_idx * stride + new_col_idx + raise ValueError(f"Unsupported dtype {dtype}") def mma_store_index_map(*args, **kwargs): diff --git a/format.sh b/format.sh index c5e81a1ef..5d3056123 100755 --- a/format.sh +++ b/format.sh @@ -148,7 +148,7 @@ echo 'bitblas codespell: Done' echo 'bitblas ruff: Check Start' # Lint specified files lint() { - ruff "$@" + ruff check "$@" } # Lint files that differ from main branch. Ignores dirs that are not slated @@ -170,7 +170,7 @@ lint_changed() { if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \ - ruff + ruff check fi } diff --git a/requirements-dev.txt b/requirements-dev.txt index 0b09c0856..de7f9d340 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -2,7 +2,7 @@ yapf==0.40.2 toml==0.10.2 tomli==2.0.1 -ruff==0.1.5 +ruff==0.6.5 codespell==2.3.0 cffi diff --git a/requirements-test.txt b/requirements-test.txt index 13fd3d1af..a06a6dd87 100644 --- a/requirements-test.txt +++ b/requirements-test.txt @@ -2,7 +2,7 @@ yapf==0.40.2 toml==0.10.2 tomli==2.0.1 -ruff==0.1.5 +ruff==0.6.5 codespell==2.3.0 cffi diff --git a/testing/python/tilelang/test_tilelang_macro_gemm.py b/testing/python/tilelang/test_tilelang_macro_gemm.py index 4d1318960..cc4839568 100644 --- a/testing/python/tilelang/test_tilelang_macro_gemm.py +++ b/testing/python/tilelang/test_tilelang_macro_gemm.py @@ -14,6 +14,7 @@ TensorCoreIntrinEmitterWithLadderTransform, ) from bitblas.gpu.intrin.lop3 import decode_i4_to_f16 +from bitblas.ops.base_scheduler import simplify_prim_func torch.manual_seed(0) @@ -33,6 +34,7 @@ def transform_func(i, j): return T.Layout(shape, transform_func) +@simplify_prim_func def tl_matmul( M, N, @@ -61,7 +63,8 @@ def tl_matmul( block_col_warps = 1 warp_row_tiles = 16 warp_col_tiles = 16 - chunk = 32 if in_dtype == "float16" else 64 + # chunk = 32 if in_dtype == "float16" else 64 + chunk = 32 shared_scope = "shared.dyn" # Pipeline Stage @@ -84,7 +87,9 @@ def tl_matmul( warp_size = 32 threads = warp_size * (block_row_warps * block_col_warps) - local_size = (micro_size_x * micro_size_y) // warp_size + local_size_a = (micro_size_x * micro_size_k) // warp_size + local_size_b = (micro_size_y * micro_size_k) // warp_size + local_size_c = (micro_size_x * micro_size_y) // warp_size warp_rows = warp_row_tiles // micro_size_x warp_cols = warp_col_tiles // micro_size_y @@ -113,9 +118,9 @@ def main( A_shared = T.alloc_shared(A_shared_shape, in_dtype, scope=shared_scope) B_shared = T.alloc_shared(B_shared_shape, in_dtype, scope=shared_scope) C_shared = T.alloc_shared(C_shared_shape, out_dtype, scope=shared_scope) - A_local = T.alloc_local((warp_rows * local_size), in_dtype) - B_local = T.alloc_local((warp_cols * local_size), in_dtype) - C_local = T.alloc_local((warp_rows * warp_cols * local_size), accum_dtype) + A_local = T.alloc_local((warp_rows * local_size_a), in_dtype) + B_local = T.alloc_local((warp_cols * local_size_b), in_dtype) + C_local = T.alloc_local((warp_rows * warp_cols * local_size_c), accum_dtype) thread_bindings = T.thread_binding(0, threads, "threadIdx.x") @@ -181,15 +186,18 @@ def main( def assert_tl_matmul_correctness(M, N, K, in_dtype, out_dtype, accum_dtype): matmul = tl_matmul(M, N, K, in_dtype, out_dtype, accum_dtype) - mod, params = TL.lower(matmul) src_code = mod.imported_modules[0].get_source() - # src_code is the generated cuda source assert src_code is not None - A = torch.rand(M, K, device="cuda", dtype=getattr(torch, in_dtype)) - B = torch.rand(N, K, device="cuda", dtype=getattr(torch, in_dtype)) + if in_dtype == "int8": + A = torch.randint(-128, 127, (M, K), device="cuda", dtype=torch.int8) + B = torch.randint(-128, 127, (N, K), device="cuda", dtype=torch.int8) + else: + A = torch.rand(M, K, device="cuda", dtype=getattr(torch, in_dtype)) + B = torch.rand(N, K, device="cuda", dtype=getattr(torch, in_dtype)) + C = torch.zeros(M, N, device="cuda", dtype=getattr(torch, accum_dtype)) mod = TL.Profiler(mod, params, [], TL.TensorSupplyType.Integer) @@ -202,7 +210,9 @@ def assert_tl_matmul_correctness(M, N, K, in_dtype, out_dtype, accum_dtype): assert latency is not None # Get Reference Result - ref_c = torch.matmul(A, B.T).to(getattr(torch, accum_dtype)) + ref_c = torch.matmul(A.to(torch.float32), B.T.to(torch.float32)).to(getattr(torch, accum_dtype)) + print(C) + print(ref_c) torch.testing.assert_close(C, ref_c, rtol=1e-2, atol=1e-2) @@ -873,6 +883,7 @@ def assert_tl_matmul_with_ladder_weight_only_transform_block_reduce_int4_correct def test_assert_tl_matmul(): assert_tl_matmul_correctness(128, 128, 128, "float16", "float16", "float16") assert_tl_matmul_correctness(128, 256, 256, "float16", "float32", "float32") + assert_tl_matmul_correctness(128, 256, 256, "int8", "int32", "int32") def test_assert_tl_matmul_with_block_reduce():