forked from ccf-ai-infra/TileOPs-Metax
258 lines
7.9 KiB
Python
258 lines
7.9 KiB
Python
"""Tests for bitwise elementwise ops (bitwise_and, bitwise_or, bitwise_xor, bitwise_not).
|
|
|
|
Bitwise ops operate on integer inputs. We use int32 tensors for testing
|
|
binary bitwise ops, and all bool/integer dtypes for bitwise_not.
|
|
Covers L1 smoke correctness.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tests.test_base import FixtureBase, TestBase, exact_compare
|
|
from tileops.ops.elementwise import (
|
|
BitwiseAndFwdOp,
|
|
BitwiseNotFwdOp,
|
|
BitwiseOrFwdOp,
|
|
BitwiseXorFwdOp,
|
|
)
|
|
|
|
# Shared helpers
|
|
|
|
|
|
def _exact_compare(output: torch.Tensor, output_ref: torch.Tensor) -> None:
|
|
"""Exact comparison for integer outputs."""
|
|
assert torch.equal(output, output_ref), (
|
|
f"Mismatch: {(output != output_ref).sum().item()} elements differ"
|
|
)
|
|
|
|
|
|
class BitwiseTest(TestBase):
|
|
"""Reusable test body for bitwise ops."""
|
|
|
|
def __init__(self, n_total: int, ref_fn):
|
|
self.n_total = n_total
|
|
self.dtype = torch.int32
|
|
self.ref_fn = ref_fn
|
|
|
|
def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor]:
|
|
a = torch.randint(-1000, 1000, (self.n_total,), dtype=torch.int32, device="cuda")
|
|
b = torch.randint(-1000, 1000, (self.n_total,), dtype=torch.int32, device="cuda")
|
|
return a, b
|
|
|
|
def ref_program(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
|
|
return self.ref_fn(a, b)
|
|
|
|
|
|
# BitwiseAnd op
|
|
|
|
|
|
class BitwiseAndFixture(FixtureBase):
|
|
PARAMS = [
|
|
("n_total", [
|
|
pytest.param(4_096, marks=pytest.mark.smoke),
|
|
pytest.param(16_384, marks=pytest.mark.full),
|
|
]),
|
|
]
|
|
|
|
|
|
@BitwiseAndFixture
|
|
def test_bitwise_and_op(n_total: int) -> None:
|
|
test = BitwiseTest(n_total, torch.bitwise_and)
|
|
shape = (n_total,)
|
|
op = BitwiseAndFwdOp(a_shape=shape, b_shape=shape, dtype=torch.int32)
|
|
test.check(op, *test.gen_inputs(), compare=_exact_compare)
|
|
|
|
|
|
# BitwiseOr op
|
|
|
|
|
|
class BitwiseOrFixture(FixtureBase):
|
|
PARAMS = [
|
|
("n_total", [
|
|
pytest.param(4_096, marks=pytest.mark.smoke),
|
|
pytest.param(16_384, marks=pytest.mark.full),
|
|
]),
|
|
]
|
|
|
|
|
|
@BitwiseOrFixture
|
|
def test_bitwise_or_op(n_total: int) -> None:
|
|
test = BitwiseTest(n_total, torch.bitwise_or)
|
|
shape = (n_total,)
|
|
op = BitwiseOrFwdOp(a_shape=shape, b_shape=shape, dtype=torch.int32)
|
|
test.check(op, *test.gen_inputs(), compare=_exact_compare)
|
|
|
|
|
|
# BitwiseXor op
|
|
|
|
|
|
class BitwiseXorFixture(FixtureBase):
|
|
PARAMS = [
|
|
("n_total", [
|
|
pytest.param(4_096, marks=pytest.mark.smoke),
|
|
pytest.param(16_384, marks=pytest.mark.full),
|
|
]),
|
|
]
|
|
|
|
|
|
@BitwiseXorFixture
|
|
def test_bitwise_xor_op(n_total: int) -> None:
|
|
test = BitwiseTest(n_total, torch.bitwise_xor)
|
|
shape = (n_total,)
|
|
op = BitwiseXorFwdOp(a_shape=shape, b_shape=shape, dtype=torch.int32)
|
|
test.check(op, *test.gen_inputs(), compare=_exact_compare)
|
|
|
|
|
|
# Broadcast pattern tests for binary bitwise ops (L3)
|
|
|
|
_BROADCAST_PATTERNS = [
|
|
((2, 64, 128), (1, 1, 128)), # bias-add
|
|
((2, 64, 128), (2, 64, 1)), # row broadcast
|
|
((64, 128), (1, 1)), # scalar broadcast
|
|
]
|
|
|
|
_BITWISE_OPS = [
|
|
("bitwise_and", BitwiseAndFwdOp, torch.bitwise_and),
|
|
("bitwise_or", BitwiseOrFwdOp, torch.bitwise_or),
|
|
("bitwise_xor", BitwiseXorFwdOp, torch.bitwise_xor),
|
|
]
|
|
|
|
|
|
class BitwiseBroadcastFixture(FixtureBase):
|
|
PARAMS = [
|
|
("op_name, op_cls, ref_fn, a_shape, b_shape", [
|
|
pytest.param(name, cls, ref, a_s, b_s,
|
|
marks=pytest.mark.smoke if i == 0 and j == 0
|
|
else pytest.mark.full)
|
|
for j, (name, cls, ref) in enumerate(_BITWISE_OPS)
|
|
for i, (a_s, b_s) in enumerate(_BROADCAST_PATTERNS)
|
|
]),
|
|
]
|
|
|
|
|
|
@BitwiseBroadcastFixture
|
|
def test_bitwise_broadcast(
|
|
op_name, op_cls, ref_fn, a_shape, b_shape,
|
|
) -> None:
|
|
a = torch.randint(-1000, 1000, a_shape, dtype=torch.int32, device="cuda")
|
|
b = torch.randint(-1000, 1000, b_shape, dtype=torch.int32, device="cuda")
|
|
op = op_cls(a_shape=a_shape, b_shape=b_shape, dtype=torch.int32)
|
|
ref = ref_fn(a, b)
|
|
with torch.no_grad():
|
|
out = op(a, b)
|
|
_exact_compare(out, ref)
|
|
|
|
|
|
class BoolBitwiseFixture(FixtureBase):
|
|
PARAMS = [
|
|
("op_name, op_cls, ref_fn, a_shape, b_shape", [
|
|
pytest.param(
|
|
name, cls, ref, a_s, b_s,
|
|
marks=pytest.mark.smoke if a_s == b_s else pytest.mark.full,
|
|
)
|
|
for a_s, b_s in [
|
|
((2048, 4096), (2048, 4096)),
|
|
((2, 512, 768), (1, 1, 768)),
|
|
]
|
|
for name, cls, ref in _BITWISE_OPS
|
|
]),
|
|
]
|
|
|
|
|
|
@BoolBitwiseFixture
|
|
def test_bool_bitwise_fast_path(
|
|
op_name, op_cls, ref_fn, a_shape, b_shape,
|
|
) -> None:
|
|
a = torch.randint(0, 2, a_shape, device="cuda").bool()
|
|
b = torch.randint(0, 2, b_shape, device="cuda").bool()
|
|
op = op_cls(a_shape=a_shape, b_shape=b_shape, dtype=torch.bool)
|
|
ref = ref_fn(a, b)
|
|
with torch.no_grad():
|
|
out = op(a, b)
|
|
assert out.dtype == torch.bool
|
|
_exact_compare(out, ref)
|
|
|
|
|
|
# BitwiseNot op
|
|
|
|
|
|
class BitwiseFixture(FixtureBase):
|
|
"""Parametrize over torch-supported bitwise_not dtypes."""
|
|
|
|
PARAMS = [
|
|
("n_total, dtype", [
|
|
pytest.param(1_048_576, torch.bool, marks=pytest.mark.smoke),
|
|
pytest.param(1_048_576, torch.uint8, marks=pytest.mark.smoke),
|
|
pytest.param(1_048_576, torch.int8, marks=pytest.mark.smoke),
|
|
pytest.param(1_048_576, torch.int16, marks=pytest.mark.smoke),
|
|
pytest.param(1_048_576, torch.int32, marks=pytest.mark.smoke),
|
|
pytest.param(1_048_576, torch.int64, marks=pytest.mark.smoke),
|
|
]),
|
|
]
|
|
|
|
|
|
class BitwiseNotTest(TestBase):
|
|
"""Test fixture for bitwise_not."""
|
|
|
|
def __init__(self, n_total: int, dtype: torch.dtype):
|
|
self.n_total = n_total
|
|
self.dtype = dtype
|
|
|
|
def gen_inputs(self) -> tuple[torch.Tensor]:
|
|
if self.dtype == torch.bool:
|
|
x = torch.rand(self.n_total, device="cuda") > 0.5
|
|
elif self.dtype == torch.uint8:
|
|
x = torch.randint(0, 256, (self.n_total,), device="cuda", dtype=self.dtype)
|
|
else:
|
|
x = torch.randint(-128, 128, (self.n_total,), device="cuda", dtype=self.dtype)
|
|
return (x,)
|
|
|
|
def ref_program(self, x: torch.Tensor) -> torch.Tensor:
|
|
return torch.bitwise_not(x)
|
|
|
|
|
|
@BitwiseFixture
|
|
def test_bitwise_not(n_total: int, dtype: torch.dtype) -> None:
|
|
test = BitwiseNotTest(n_total, dtype)
|
|
op = BitwiseNotFwdOp(N_total=n_total, dtype=dtype)
|
|
test.check(op, *test.gen_inputs(), compare=exact_compare)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [
|
|
pytest.param(torch.float16, marks=pytest.mark.smoke),
|
|
pytest.param(torch.bfloat16, marks=pytest.mark.smoke),
|
|
pytest.param(torch.float32, marks=pytest.mark.smoke),
|
|
])
|
|
def test_bitwise_not_rejects_float_dtype(dtype: torch.dtype) -> None:
|
|
from tileops.kernels.elementwise import BitwiseNotFwdKernel
|
|
|
|
with pytest.raises(ValueError, match="only supports dtypes"):
|
|
BitwiseNotFwdKernel(N_total=16, dtype=dtype)
|
|
|
|
|
|
# Dtype rejection tests for binary bitwise ops
|
|
|
|
|
|
class BitwiseBinaryRejectFixture(FixtureBase):
|
|
PARAMS = [
|
|
("op_cls, dtype", [
|
|
pytest.param(BitwiseAndFwdOp, torch.float16, marks=pytest.mark.smoke),
|
|
pytest.param(BitwiseAndFwdOp, torch.bfloat16, marks=pytest.mark.smoke),
|
|
pytest.param(BitwiseAndFwdOp, torch.float32, marks=pytest.mark.smoke),
|
|
pytest.param(BitwiseOrFwdOp, torch.float16, marks=pytest.mark.full),
|
|
pytest.param(BitwiseXorFwdOp, torch.float16, marks=pytest.mark.full),
|
|
]),
|
|
]
|
|
|
|
|
|
@BitwiseBinaryRejectFixture
|
|
def test_bitwise_binary_rejects_float_dtype(op_cls, dtype: torch.dtype) -> None:
|
|
"""Binary bitwise ops only support integer dtypes; floats must be rejected."""
|
|
shape = (16,)
|
|
with pytest.raises(ValueError, match="does not support dtype"):
|
|
op_cls(a_shape=shape, b_shape=shape, dtype=dtype)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-vvs"])
|