mirror of
https://github.com/tile-ai/tilelang.git
synced 2026-10-26 07:22:54 +08:00
Sam Sui
e8b9da92d2
* fix(cuda): support int4x2 and uint4x2 codegen (#2994, #2995) * fix(cuda): handle packed int4 lane access * fix(cuda): handle packed 4-bit x2 construction * fix(cuda): load packed x2 gathers by logical index * fix(cuda): validate packed int4 vector alignment * fix(cuda): preserve packed vector load fallback
657 lines
19 KiB
Python
657 lines
19 KiB
Python
import pytest
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.testing
|
|
from tilelang import language as T, tvm
|
|
from tvm import tirx
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _disable_tilelang_cache():
|
|
tilelang.disable_cache()
|
|
try:
|
|
yield
|
|
finally:
|
|
tilelang.enable_cache()
|
|
|
|
|
|
def _generated_call_lines(source, function):
|
|
return [line for line in source.splitlines() if f"{function}(" in line and not line.lstrip().startswith("TL_DEVICE")]
|
|
|
|
|
|
def _generated_call_count(source, function):
|
|
return sum(line.count(f"{function}(") for line in _generated_call_lines(source, function))
|
|
|
|
|
|
def _packed_input(dtype, n):
|
|
values = [(((byte // 8) % 8) << 4) | (byte % 8) for byte in range(n // 2)]
|
|
packed = torch.tensor(values, dtype=torch.uint8, device="cuda")
|
|
return packed if dtype == "uint4" else packed.view(torch.int8)
|
|
|
|
|
|
def _copy_kernel(dtype, n, threads):
|
|
@T.prim_func
|
|
def kernel(A: T.Tensor((n,), dtype), B: T.Tensor((n,), dtype)):
|
|
with T.Kernel(1, threads=threads):
|
|
T.copy(A, B)
|
|
|
|
return kernel
|
|
|
|
|
|
def _packed_x2_select_source(dtype):
|
|
on_true = tirx.Var("on_true", f"{dtype}x2")
|
|
on_false = tirx.Var("on_false", f"{dtype}x2")
|
|
ramp = tirx.Ramp(0, 1, 2)
|
|
limit = tirx.Broadcast(tirx.const(1, "int32"), 2)
|
|
selected = tirx.Select(ramp < limit, on_true, on_false)
|
|
func = tirx.PrimFunc([on_true, on_false], tirx.Evaluate(selected))
|
|
func = func.with_attr("global_symbol", "packed_x2_select")
|
|
func = func.with_attr("calling_conv", tvm.ir.CallingConv.DEVICE_KERNEL_LAUNCH)
|
|
mod = tvm.IRModule({"packed_x2_select": func})
|
|
build = tvm.get_global_func("target.build.tilelang_cuda")
|
|
return build(mod, tvm.target.Target("cuda")).inspect_source()
|
|
|
|
|
|
def _packed_x2_constructor_source(dtype, constructor):
|
|
name = f"packed_{dtype}_x2_{constructor}"
|
|
if constructor == "broadcast":
|
|
value = -3 if dtype == "int4" else 13
|
|
func = tirx.PrimFunc([], tirx.Evaluate(tirx.Broadcast(tirx.const(value, dtype), 2)))
|
|
elif constructor == "scalar_load":
|
|
buffer = tirx.decl_buffer((4,), dtype=dtype, name="input")
|
|
index = tirx.Ramp(tirx.const(0, "int32"), tirx.const(2, "int32"), 2)
|
|
func = tirx.PrimFunc(
|
|
[buffer.data],
|
|
tirx.Evaluate(tirx.BufferLoad(buffer, [index])),
|
|
buffer_map={buffer.data: buffer},
|
|
)
|
|
elif constructor == "shuffle":
|
|
low = tirx.Var("low", dtype)
|
|
high = tirx.Var("high", dtype)
|
|
func = tirx.PrimFunc([low, high], tirx.Evaluate(tirx.Shuffle([low, high], [0, 1])))
|
|
elif constructor == "ramp":
|
|
ramp = tirx.Ramp(tirx.const(0, dtype), tirx.const(1, dtype), 2)
|
|
func = tirx.PrimFunc([], tirx.Evaluate(ramp))
|
|
else:
|
|
raise ValueError(f"Unknown constructor: {constructor}")
|
|
|
|
func = func.with_attr("global_symbol", name)
|
|
func = func.with_attr("calling_conv", tvm.ir.CallingConv.DEVICE_KERNEL_LAUNCH)
|
|
mod = tvm.IRModule({name: func})
|
|
build = tvm.get_global_func("target.build.tilelang_cuda")
|
|
return build(mod, tvm.target.Target("cuda")).inspect_source()
|
|
|
|
|
|
def _packed_vector_load_source(
|
|
dtype,
|
|
base,
|
|
stride,
|
|
*,
|
|
dynamic_base=False,
|
|
dynamic_scale=1,
|
|
lanes=2,
|
|
):
|
|
name = f"packed_{dtype}_x{lanes}_load"
|
|
buffer = tirx.decl_buffer((64,), dtype=dtype, name="input")
|
|
params = [buffer.data]
|
|
|
|
if dynamic_base:
|
|
dynamic = tirx.Var("base", "int32")
|
|
params.append(dynamic)
|
|
base_expr = dynamic * tirx.const(dynamic_scale, "int32") + tirx.const(
|
|
base,
|
|
"int32",
|
|
)
|
|
else:
|
|
base_expr = tirx.const(base, "int32")
|
|
|
|
index = tirx.Ramp(
|
|
base_expr,
|
|
tirx.const(stride, "int32"),
|
|
lanes,
|
|
)
|
|
func = tirx.PrimFunc(
|
|
params,
|
|
tirx.Evaluate(tirx.BufferLoad(buffer, [index])),
|
|
buffer_map={buffer.data: buffer},
|
|
)
|
|
func = func.with_attr("global_symbol", name)
|
|
func = func.with_attr(
|
|
"calling_conv",
|
|
tvm.ir.CallingConv.DEVICE_KERNEL_LAUNCH,
|
|
)
|
|
mod = tvm.IRModule({name: func})
|
|
build = tvm.get_global_func("target.build.tilelang_cuda")
|
|
return build(mod, tvm.target.Target("cuda")).inspect_source()
|
|
|
|
|
|
def _packed_vector_store_source(
|
|
dtype,
|
|
base,
|
|
stride,
|
|
*,
|
|
dynamic_base=False,
|
|
dynamic_scale=1,
|
|
lanes=2,
|
|
):
|
|
name = f"packed_{dtype}_x{lanes}_store"
|
|
buffer = tirx.decl_buffer((64,), dtype=dtype, name="output")
|
|
value = tirx.Var("value", f"{dtype}x{lanes}")
|
|
params = [buffer.data, value]
|
|
|
|
if dynamic_base:
|
|
dynamic = tirx.Var("base", "int32")
|
|
params.append(dynamic)
|
|
base_expr = dynamic * tirx.const(dynamic_scale, "int32") + tirx.const(
|
|
base,
|
|
"int32",
|
|
)
|
|
else:
|
|
base_expr = tirx.const(base, "int32")
|
|
|
|
index = tirx.Ramp(
|
|
base_expr,
|
|
tirx.const(stride, "int32"),
|
|
lanes,
|
|
)
|
|
func = tirx.PrimFunc(
|
|
params,
|
|
tirx.BufferStore(buffer, value, [index]),
|
|
buffer_map={buffer.data: buffer},
|
|
)
|
|
func = func.with_attr("global_symbol", name)
|
|
func = func.with_attr(
|
|
"calling_conv",
|
|
tvm.ir.CallingConv.DEVICE_KERNEL_LAUNCH,
|
|
)
|
|
mod = tvm.IRModule({name: func})
|
|
build = tvm.get_global_func("target.build.tilelang_cuda")
|
|
return build(mod, tvm.target.Target("cuda")).inspect_source()
|
|
|
|
|
|
def _packed_x2_broadcast_kernel(dtype, value):
|
|
@T.prim_func
|
|
def kernel(output: T.Tensor((2,), dtype)):
|
|
with T.Kernel(1, threads=1):
|
|
for i in T.vectorized(2):
|
|
output[i] = T.cast(value, dtype)
|
|
|
|
return kernel
|
|
|
|
|
|
def _packed_x2_ramp_kernel(dtype, offset):
|
|
@T.prim_func
|
|
def kernel(output: T.Tensor((2,), dtype)):
|
|
with T.Kernel(1, threads=1):
|
|
for i in T.vectorized(2):
|
|
output[i] = T.cast(i + offset, dtype)
|
|
|
|
return kernel
|
|
|
|
|
|
def _packed_x2_gather_kernel(dtype):
|
|
@T.prim_func
|
|
def kernel(
|
|
source: T.Tensor((8,), dtype),
|
|
output: T.Tensor((2,), dtype),
|
|
):
|
|
with T.Kernel(1, threads=1):
|
|
output[T.Ramp(0, 1, 2)] = source[T.Ramp(0, 2, 2)]
|
|
|
|
return kernel
|
|
|
|
|
|
def _packed_vector_gather_kernel(dtype, lanes, base, stride):
|
|
@T.prim_func
|
|
def kernel(
|
|
source: T.Tensor((64,), dtype),
|
|
output: T.Tensor((lanes,), dtype),
|
|
):
|
|
with T.Kernel(1, threads=1):
|
|
output[T.Ramp(0, 1, lanes)] = source[T.Ramp(base, stride, lanes)]
|
|
|
|
return kernel
|
|
|
|
|
|
def _packed_x2_store_kernel(dtype):
|
|
@T.prim_func
|
|
def kernel(
|
|
source: T.Tensor((2,), dtype),
|
|
output: T.Tensor((8,), dtype),
|
|
):
|
|
with T.Kernel(1, threads=1):
|
|
output[T.Ramp(2, 1, 2)] = source[T.Ramp(0, 1, 2)]
|
|
|
|
return kernel
|
|
|
|
|
|
def _packed_int4_scalar_load_kernel():
|
|
@T.prim_func
|
|
def kernel(
|
|
source: T.Tensor((4,), "int4"),
|
|
output: T.Tensor((4,), "int32"),
|
|
):
|
|
with T.Kernel(1, threads=1):
|
|
for i in T.serial(4):
|
|
output[i] = T.cast(source[i], "int32")
|
|
|
|
return kernel
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(("dtype", "storage_type"), [("int4", "int8_t"), ("uint4", "uint8_t")])
|
|
def test_packed_int4x2_copy_round_trip(dtype, storage_type):
|
|
n = 256
|
|
kernel = _copy_kernel(dtype, n, threads=128)
|
|
compiled = tilelang.compile(kernel, out_idx=[1])
|
|
|
|
assert storage_type in compiled.get_kernel_source()
|
|
source = _packed_input(dtype, n)
|
|
result = compiled(source)
|
|
assert torch.equal(result.view(torch.uint8), source.view(torch.uint8))
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("dtype", ["int4", "uint4"])
|
|
def test_packed_int4x2_lane_select_codegen(dtype):
|
|
source = _packed_x2_select_source(dtype)
|
|
assignments = [line for line in source.splitlines() if "on_true" in line and "on_false" in line and "=" in line]
|
|
|
|
assert len(assignments) == 2
|
|
assert all("((const unsigned char*)" in line for line in assignments)
|
|
assert all(line.startswith(" ((unsigned char*)") for line in assignments)
|
|
assert all(f"on_true.{lane}" not in source for lane in "xy")
|
|
assert all(f"on_false.{lane}" not in source for lane in "xy")
|
|
assert ">> 0" in assignments[0] and "<< 0" in assignments[0]
|
|
assert ">> 4" in assignments[1] and "<< 4" in assignments[1]
|
|
assert all(("^ 8) - 8" in line) == (dtype == "int4") for line in assignments)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("constructor", ["broadcast", "scalar_load", "shuffle", "ramp"])
|
|
@pytest.mark.parametrize("dtype", ["int4", "uint4"])
|
|
def test_packed_int4x2_constructor_codegen(dtype, constructor):
|
|
source = _packed_x2_constructor_source(dtype, constructor)
|
|
helper = "tl_pack_uint4x2" if dtype == "uint4" else "tl_pack_int4x2"
|
|
calls = [line for line in source.splitlines() if f"{helper}(" in line and not line.lstrip().startswith("TL_DEVICE")]
|
|
|
|
assert len(calls) == 1
|
|
assert "make_int8_t(" not in source
|
|
assert "make_uint8_t(" not in source
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("dtype", ["int4", "uint4"])
|
|
@pytest.mark.parametrize(
|
|
("base", "stride", "dynamic_base"),
|
|
[
|
|
(0, 2, False),
|
|
(1, 1, False),
|
|
(0, 1, True),
|
|
],
|
|
)
|
|
def test_packed_int4x2_gather_uses_logical_indices(
|
|
dtype,
|
|
base,
|
|
stride,
|
|
dynamic_base,
|
|
):
|
|
source = _packed_vector_load_source(
|
|
dtype,
|
|
base,
|
|
stride,
|
|
dynamic_base=dynamic_base,
|
|
)
|
|
|
|
pack_lines = _generated_call_lines(source, f"tl_pack_{dtype}x2")
|
|
assert _generated_call_count(source, f"tl_{dtype}_packed_load") == 2
|
|
assert len(pack_lines) == 1
|
|
pack_line = pack_lines[0]
|
|
assert pack_line.count("v_.x") == 1
|
|
assert pack_line.count("v_.y") == 1
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(
|
|
("dtype", "storage_type"),
|
|
[
|
|
("int4", "int8_t"),
|
|
("uint4", "uint8_t"),
|
|
],
|
|
)
|
|
def test_packed_int4x2_aligned_load_keeps_direct_path(dtype, storage_type):
|
|
source = _packed_vector_load_source(dtype, base=2, stride=1)
|
|
|
|
assert f"*((({storage_type}*)input) + 1);" in source
|
|
assert f"tl_{dtype}_packed_load(" not in source
|
|
assert f"tl_pack_{dtype}x2(" not in source
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(
|
|
("dtype", "storage_type"),
|
|
[
|
|
("int4", "int8_t"),
|
|
("uint4", "uint8_t"),
|
|
],
|
|
)
|
|
def test_packed_int4x2_symbolic_aligned_load_keeps_direct_path(dtype, storage_type):
|
|
source = _packed_vector_load_source(
|
|
dtype,
|
|
base=0,
|
|
stride=1,
|
|
dynamic_base=True,
|
|
dynamic_scale=2,
|
|
)
|
|
|
|
assert f"(({storage_type}*)input)" in source
|
|
assert not _generated_call_lines(source, f"tl_{dtype}_packed_load")
|
|
assert not _generated_call_lines(source, f"tl_pack_{dtype}x2")
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(
|
|
("dtype", "storage_type"),
|
|
[
|
|
("int4", "int8_t"),
|
|
("uint4", "uint8_t"),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
("base", "dynamic_base", "dynamic_scale"),
|
|
[
|
|
(2, False, 1),
|
|
(0, True, 2),
|
|
(2, True, 4),
|
|
],
|
|
)
|
|
def test_packed_int4x2_aligned_store_keeps_direct_path(
|
|
dtype,
|
|
storage_type,
|
|
base,
|
|
dynamic_base,
|
|
dynamic_scale,
|
|
):
|
|
source = _packed_vector_store_source(
|
|
dtype,
|
|
base,
|
|
stride=1,
|
|
dynamic_base=dynamic_base,
|
|
dynamic_scale=dynamic_scale,
|
|
)
|
|
|
|
assert f"(({storage_type}*)output)" in source
|
|
assert f"tl_{dtype}_packed_store(" not in source
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("dtype", ["int4", "uint4"])
|
|
@pytest.mark.parametrize(
|
|
("base", "stride", "dynamic_base", "dynamic_scale"),
|
|
[
|
|
(1, 1, False, 1),
|
|
(0, 2, False, 1),
|
|
(0, 1, True, 1),
|
|
(1, 1, True, 2),
|
|
],
|
|
)
|
|
def test_packed_int4x2_unsafe_store_rejected(
|
|
dtype,
|
|
base,
|
|
stride,
|
|
dynamic_base,
|
|
dynamic_scale,
|
|
):
|
|
with pytest.raises(tvm.TVMError, match="provably divisible by the vector lane count"):
|
|
_packed_vector_store_source(
|
|
dtype,
|
|
base,
|
|
stride,
|
|
dynamic_base=dynamic_base,
|
|
dynamic_scale=dynamic_scale,
|
|
)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("dtype", ["int4", "uint4"])
|
|
@pytest.mark.parametrize("lanes", [4, 8, 16, 32])
|
|
@pytest.mark.parametrize("access", ["load", "store"])
|
|
def test_packed_int4_wide_aligned_access_keeps_direct_path(
|
|
dtype,
|
|
lanes,
|
|
access,
|
|
):
|
|
if access == "load":
|
|
constant_source = _packed_vector_load_source(
|
|
dtype,
|
|
base=lanes,
|
|
stride=1,
|
|
lanes=lanes,
|
|
)
|
|
symbolic_source = _packed_vector_load_source(
|
|
dtype,
|
|
base=0,
|
|
stride=1,
|
|
dynamic_base=True,
|
|
dynamic_scale=lanes,
|
|
lanes=lanes,
|
|
)
|
|
else:
|
|
constant_source = _packed_vector_store_source(
|
|
dtype,
|
|
base=lanes,
|
|
stride=1,
|
|
lanes=lanes,
|
|
)
|
|
symbolic_source = _packed_vector_store_source(
|
|
dtype,
|
|
base=0,
|
|
stride=1,
|
|
dynamic_base=True,
|
|
dynamic_scale=lanes,
|
|
lanes=lanes,
|
|
)
|
|
|
|
carrier_type = {
|
|
("int4", 4): "int16_t",
|
|
("uint4", 4): "uint16_t",
|
|
("int4", 8): "int",
|
|
("uint4", 8): "uint",
|
|
("int4", 16): "int2",
|
|
("uint4", 16): "uint2",
|
|
("int4", 32): "int4",
|
|
("uint4", 32): "uint4",
|
|
}[(dtype, lanes)]
|
|
buffer_name = "input" if access == "load" else "output"
|
|
expected_access = f"*((({carrier_type}*){buffer_name}) + 1)"
|
|
|
|
if access == "load":
|
|
assert f"{expected_access};" in constant_source
|
|
else:
|
|
assert f"{expected_access} = value;" in constant_source
|
|
assert f"tl_{dtype}_packed_load(" not in constant_source
|
|
assert f"tl_{dtype}_packed_store(" not in constant_source
|
|
|
|
helper = f"tl_{dtype}_packed_load" if access == "load" else f"tl_{dtype}_packed_store"
|
|
assert f"(({carrier_type}*){buffer_name})" in symbolic_source
|
|
assert not _generated_call_lines(symbolic_source, helper)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("dtype", ["int4", "uint4"])
|
|
@pytest.mark.parametrize("lanes", [4, 8, 16, 32])
|
|
def test_packed_int4_wide_unsafe_store_rejected(dtype, lanes):
|
|
unsafe_cases = [
|
|
(1, 1, False, 1),
|
|
(lanes // 2, 1, False, 1),
|
|
(0, 2, False, 1),
|
|
(0, 1, True, 1),
|
|
]
|
|
for base, stride, dynamic_base, dynamic_scale in unsafe_cases:
|
|
with pytest.raises(
|
|
tvm.TVMError,
|
|
match="provably divisible by the vector lane count",
|
|
):
|
|
_packed_vector_store_source(
|
|
dtype,
|
|
base,
|
|
stride,
|
|
dynamic_base=dynamic_base,
|
|
dynamic_scale=dynamic_scale,
|
|
lanes=lanes,
|
|
)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("dtype", ["int4", "uint4"])
|
|
@pytest.mark.parametrize("lanes", [4, 8, 16, 32])
|
|
@pytest.mark.parametrize(
|
|
("base", "stride", "dynamic_base"),
|
|
[
|
|
(1, 1, False),
|
|
(0, 2, False),
|
|
(0, 1, True),
|
|
],
|
|
)
|
|
def test_packed_int4_wide_unaligned_load_uses_logical_indices(
|
|
dtype,
|
|
lanes,
|
|
base,
|
|
stride,
|
|
dynamic_base,
|
|
):
|
|
source = _packed_vector_load_source(
|
|
dtype,
|
|
base,
|
|
stride,
|
|
dynamic_base=dynamic_base,
|
|
lanes=lanes,
|
|
)
|
|
|
|
assert _generated_call_count(source, f"tl_{dtype}_packed_load") == lanes
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("dtype", ["int4", "uint4"])
|
|
@pytest.mark.parametrize("lanes", [4, 8])
|
|
def test_packed_int4_wide_strided_gather_runtime(dtype, lanes):
|
|
base = 1
|
|
stride = 2
|
|
compiled = tilelang.compile(
|
|
_packed_vector_gather_kernel(dtype, lanes, base, stride),
|
|
out_idx=[1],
|
|
)
|
|
logical = [i % 16 for i in range(64)]
|
|
packed = [logical[i] | (logical[i + 1] << 4) for i in range(0, 64, 2)]
|
|
source = torch.tensor(packed, dtype=torch.uint8, device="cuda")
|
|
if dtype == "int4":
|
|
source = source.view(torch.int8)
|
|
|
|
result = compiled(source)
|
|
gathered = [logical[base + i * stride] for i in range(lanes)]
|
|
expected = torch.tensor(
|
|
[gathered[i] | (gathered[i + 1] << 4) for i in range(0, lanes, 2)],
|
|
dtype=torch.uint8,
|
|
device="cuda",
|
|
)
|
|
|
|
assert torch.equal(result.view(torch.uint8), expected)
|
|
assert _generated_call_count(compiled.get_kernel_source(), f"tl_{dtype}_packed_load") == lanes
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(
|
|
("dtype", "values", "expected"),
|
|
[
|
|
("int4", [0xF8, 0x72, 0x05, 0x00], 0x28),
|
|
("uint4", [0x21, 0x43, 0x65, 0x07], 0x31),
|
|
],
|
|
)
|
|
def test_packed_int4x2_strided_gather_runtime(dtype, values, expected):
|
|
compiled = tilelang.compile(_packed_x2_gather_kernel(dtype), out_idx=[1])
|
|
source = torch.tensor(values, dtype=torch.uint8, device="cuda")
|
|
if dtype == "int4":
|
|
source = source.view(torch.int8)
|
|
|
|
result = compiled(source)
|
|
kernel_source = compiled.get_kernel_source()
|
|
|
|
assert _generated_call_count(kernel_source, f"tl_{dtype}_packed_load") == 2
|
|
assert len(_generated_call_lines(kernel_source, f"tl_pack_{dtype}x2")) == 1
|
|
assert result.view(torch.uint8).item() == expected
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize("dtype", ["int4", "uint4"])
|
|
def test_packed_int4x2_aligned_store_preserves_neighboring_bytes(dtype):
|
|
compiled = tilelang.compile(_packed_x2_store_kernel(dtype))
|
|
source = torch.tensor([0x21], dtype=torch.uint8, device="cuda")
|
|
output = torch.tensor(
|
|
[0xBA, 0xDC, 0xFE, 0x98],
|
|
dtype=torch.uint8,
|
|
device="cuda",
|
|
)
|
|
if dtype == "int4":
|
|
source = source.view(torch.int8)
|
|
output = output.view(torch.int8)
|
|
|
|
compiled(source, output)
|
|
|
|
expected = torch.tensor(
|
|
[0xBA, 0x21, 0xFE, 0x98],
|
|
dtype=torch.uint8,
|
|
device="cuda",
|
|
)
|
|
assert torch.equal(output.view(torch.uint8), expected)
|
|
assert f"tl_{dtype}_packed_store(" not in compiled.get_kernel_source()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
def test_packed_int4_scalar_load_sign_extension():
|
|
compiled = tilelang.compile(_packed_int4_scalar_load_kernel(), out_idx=[1])
|
|
source = torch.tensor(
|
|
[0xF8, 0x70],
|
|
dtype=torch.uint8,
|
|
device="cuda",
|
|
).view(torch.int8)
|
|
|
|
result = compiled(source)
|
|
expected = torch.tensor(
|
|
[-8, -1, 0, 7],
|
|
dtype=torch.int32,
|
|
device="cuda",
|
|
)
|
|
|
|
assert torch.equal(result, expected)
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(("dtype", "value"), [("int4", -3), ("uint4", 13)])
|
|
def test_packed_int4x2_broadcast_runtime(dtype, value):
|
|
compiled = tilelang.compile(_packed_x2_broadcast_kernel(dtype, value))
|
|
storage_dtype = torch.uint8 if dtype == "uint4" else torch.int8
|
|
output = torch.empty((1,), dtype=storage_dtype, device="cuda")
|
|
|
|
compiled(output)
|
|
|
|
assert output.view(torch.uint8).item() == 0xDD
|
|
assert f"tl_pack_{dtype}x2(" in compiled.get_kernel_source()
|
|
|
|
|
|
@tilelang.testing.requires_cuda
|
|
@pytest.mark.parametrize(("dtype", "offset", "expected"), [("int4", -3, 0xED), ("uint4", 2, 0x32)])
|
|
def test_packed_int4x2_distinct_lane_runtime(dtype, offset, expected):
|
|
compiled = tilelang.compile(_packed_x2_ramp_kernel(dtype, offset))
|
|
storage_dtype = torch.uint8 if dtype == "uint4" else torch.int8
|
|
output = torch.empty((1,), dtype=storage_dtype, device="cuda")
|
|
|
|
compiled(output)
|
|
|
|
assert output.view(torch.uint8).item() == expected
|
|
assert f"tl_pack_{dtype}x2(" in compiled.get_kernel_source()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tilelang.testing.main()
|