mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-21 06:40:23 +08:00
134 lines
3.9 KiB
Python
134 lines
3.9 KiB
Python
import pytest
|
|
import tilelang
|
|
import tilelang.language as T
|
|
import torch
|
|
import tilelang.testing
|
|
|
|
|
|
@tilelang.jit(
|
|
out_idx=[-1],
|
|
pass_configs={
|
|
tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: False,
|
|
},
|
|
)
|
|
def bitwise_reduce(
|
|
M,
|
|
N,
|
|
block_M,
|
|
block_N,
|
|
name,
|
|
func,
|
|
clear=True,
|
|
):
|
|
@T.prim_func
|
|
def reduce_func(
|
|
A: T.Tensor((M, N), T.int32),
|
|
B: T.Tensor((M), T.int32),
|
|
Output: T.Tensor((M), T.int32),
|
|
):
|
|
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as (bx, by):
|
|
A_shared = T.alloc_shared((block_M, block_N), T.int32)
|
|
A_fragment = T.alloc_fragment((block_M, block_N), T.int32)
|
|
B_shared = T.alloc_shared((block_M,), T.int32)
|
|
B_fragment = T.alloc_fragment((block_M), T.int32)
|
|
T.copy(A[by * block_M, bx * block_N], A_shared)
|
|
T.copy(A_shared, A_fragment)
|
|
T.copy(B[by * block_M], B_shared)
|
|
T.copy(B_shared, B_fragment)
|
|
func(A_fragment, B_fragment, clear=clear)
|
|
T.copy(B_fragment, Output[by * block_M])
|
|
|
|
return reduce_func
|
|
|
|
|
|
def run_single_bitwise_reduce(
|
|
name,
|
|
func,
|
|
a,
|
|
clear=True,
|
|
):
|
|
M, N = 32, 32
|
|
block_M, block_N = 32, 32
|
|
kernel = bitwise_reduce(M, N, block_M, block_N, name, func, clear)
|
|
|
|
if name == "reduce_bitand":
|
|
expected = torch.full((M,), -1, device="cuda", dtype=torch.int32)
|
|
elif name == "reduce_bitor" or name == "reduce_bitxor":
|
|
expected = torch.full((M,), 0, device="cuda", dtype=torch.int32)
|
|
else:
|
|
raise ValueError("Invalid name: {}".format(name))
|
|
|
|
output = kernel(a, expected)
|
|
|
|
expected = reference_bitwise_reduce(name, a)
|
|
assert torch.all(output == expected)
|
|
print("✓ {} with clear={} test passed".format(name, clear))
|
|
|
|
|
|
def reference_bitwise_reduce(name, a):
|
|
if name == "reduce_bitand":
|
|
op = torch.bitwise_and
|
|
identity = -1
|
|
elif name == "reduce_bitor":
|
|
op = torch.bitwise_or
|
|
identity = 0
|
|
elif name == "reduce_bitxor":
|
|
op = torch.bitwise_xor
|
|
identity = 0
|
|
else:
|
|
raise ValueError("Invalid name: {}".format(name))
|
|
|
|
reduced = a
|
|
while reduced.shape[1] > 1:
|
|
if reduced.shape[1] % 2:
|
|
padding = torch.full(
|
|
(reduced.shape[0], 1),
|
|
identity,
|
|
device=reduced.device,
|
|
dtype=reduced.dtype,
|
|
)
|
|
reduced = torch.cat([reduced, padding], dim=1)
|
|
reduced = op(reduced[:, 0::2], reduced[:, 1::2])
|
|
return reduced[:, 0]
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def bitwise_reduce_input():
|
|
M, N = 32, 32
|
|
rows = torch.arange(M, dtype=torch.int32)[:, None]
|
|
cols = torch.arange(N, dtype=torch.int32)[None, :]
|
|
|
|
row_pattern = (rows & 0xF) << (rows % 4)
|
|
col_pattern = torch.bitwise_left_shift(torch.ones_like(cols), cols % 31)
|
|
position_factor = (rows * N + cols) % 256
|
|
|
|
a = row_pattern ^ col_pattern ^ position_factor
|
|
|
|
clear_rows = (rows % 4) == 0
|
|
clear_bits = torch.bitwise_left_shift(torch.ones_like(rows), rows // 4)
|
|
a = torch.where(clear_rows, a & torch.bitwise_not(clear_bits), a)
|
|
|
|
set_rows = ((rows % 4) != 0) & ((rows % 2) == 0)
|
|
set_bits = torch.bitwise_left_shift(torch.ones_like(rows), rows // 2)
|
|
a = torch.where(set_rows, a | set_bits, a)
|
|
|
|
return a.to(device="cuda")
|
|
|
|
|
|
BITWISE_REDUCE_OPS = [
|
|
("reduce_bitand", T.reduce_bitand),
|
|
("reduce_bitor", T.reduce_bitor),
|
|
("reduce_bitxor", T.reduce_bitxor),
|
|
]
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(("name", "func"), BITWISE_REDUCE_OPS, ids=[name for name, _ in BITWISE_REDUCE_OPS])
|
|
@pytest.mark.parametrize("clear", [True, False], ids=["clear", "no-clear"])
|
|
def test_bitwise_reduce_ops(bitwise_reduce_input, name, func, clear):
|
|
run_single_bitwise_reduce(name, func, bitwise_reduce_input, clear=clear)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|