forked from ccf-ai-infra/TileOPs-Metax
650 lines
24 KiB
Python
650 lines
24 KiB
Python
"""Tests for special predicate elementwise ops (isnan, isinf, isfinite).
|
|
|
|
Covers L1 smoke correctness (fp16, 1M) and L4 edge cases (fp32, 4K).
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from tests.test_base import FixtureBase, TestBase, exact_compare
|
|
from tileops.ops.elementwise import IsfiniteOp, IsinfOp, IsnanOp
|
|
|
|
|
|
class SpecialFixture(FixtureBase):
|
|
"""Parametrize over shapes / dtypes for special predicate ops."""
|
|
|
|
PARAMS = [
|
|
("n_total, dtype", [
|
|
pytest.param(1_048_576, torch.float16, marks=pytest.mark.smoke),
|
|
pytest.param(1_048_576, torch.bfloat16, marks=pytest.mark.full),
|
|
pytest.param(1_048_576, torch.float32, marks=pytest.mark.full),
|
|
]),
|
|
]
|
|
|
|
|
|
class SpecialEdgeFixture(FixtureBase):
|
|
"""L4 edge-case fixture: fp32, 4K elements."""
|
|
|
|
PARAMS = [
|
|
("n_total, dtype", [
|
|
pytest.param(4096, torch.float32, marks=pytest.mark.smoke),
|
|
]),
|
|
]
|
|
|
|
|
|
class SpecialTest(TestBase):
|
|
"""Generic test harness for special predicate ops."""
|
|
|
|
def __init__(self, n_total: int, dtype: torch.dtype, ref_fn, gen_fn=None):
|
|
self.n_total = n_total
|
|
self.dtype = dtype
|
|
self._ref_fn = ref_fn
|
|
self._gen_fn = gen_fn
|
|
|
|
def gen_inputs(self) -> tuple[torch.Tensor]:
|
|
if self._gen_fn is not None:
|
|
return (self._gen_fn(self.n_total, self.dtype),)
|
|
x = torch.randn(self.n_total, device="cuda", dtype=self.dtype)
|
|
quarter = self.n_total // 4
|
|
x[:quarter] = float("nan")
|
|
x[quarter:2 * quarter] = float("inf")
|
|
x[2 * quarter:3 * quarter] = float("-inf")
|
|
return (x,)
|
|
|
|
def ref_program(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self._ref_fn(x)
|
|
|
|
|
|
def _make_special_test(n_total, dtype, op_cls, ref_fn, gen_fn=None) -> None:
|
|
test = SpecialTest(n_total, dtype, ref_fn=ref_fn, gen_fn=gen_fn)
|
|
op = op_cls(N_total=n_total, dtype=dtype)
|
|
test.check(op, *test.gen_inputs(), compare=exact_compare)
|
|
|
|
|
|
@SpecialFixture
|
|
def test_isnan(n_total: int, dtype: torch.dtype) -> None:
|
|
_make_special_test(n_total, dtype, IsnanOp, torch.isnan)
|
|
|
|
|
|
@SpecialFixture
|
|
def test_isinf(n_total: int, dtype: torch.dtype) -> None:
|
|
_make_special_test(n_total, dtype, IsinfOp, torch.isinf)
|
|
|
|
|
|
@SpecialFixture
|
|
def test_isfinite(n_total: int, dtype: torch.dtype) -> None:
|
|
_make_special_test(n_total, dtype, IsfiniteOp, torch.isfinite)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# L4 edge-case tests (fp32, 4K)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@SpecialEdgeFixture
|
|
def test_isnan_edge(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: all NaN input."""
|
|
def _all_nan(n, dtype):
|
|
return torch.full((n,), float("nan"), device="cuda", dtype=dtype)
|
|
|
|
_make_special_test(n_total, dtype, IsnanOp, torch.isnan, gen_fn=_all_nan)
|
|
|
|
|
|
@SpecialEdgeFixture
|
|
def test_isinf_edge(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: mix of +inf and -inf."""
|
|
def _all_inf(n, dtype):
|
|
x = torch.full((n,), float("inf"), device="cuda", dtype=dtype)
|
|
x[:n // 2] = float("-inf")
|
|
return x
|
|
|
|
_make_special_test(n_total, dtype, IsinfOp, torch.isinf, gen_fn=_all_inf)
|
|
|
|
|
|
@SpecialEdgeFixture
|
|
def test_isfinite_edge(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: all finite input."""
|
|
def _all_finite(n, dtype):
|
|
return torch.randn(n, device="cuda", dtype=dtype)
|
|
|
|
_make_special_test(n_total, dtype, IsfiniteOp, torch.isfinite, gen_fn=_all_finite)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_special_predicates_reject_non_float_dtype() -> None:
|
|
from tileops.kernels.elementwise import IsnanKernel
|
|
|
|
with pytest.raises(ValueError, match="only supports dtypes"):
|
|
IsnanKernel(N_total=16, dtype=torch.int32)
|
|
|
|
|
|
# ===========================================================================
|
|
# Independent special ops: where, clamp, masked_fill, nan_to_num,
|
|
# alibi, sinusoidal
|
|
# ===========================================================================
|
|
|
|
|
|
class IndependentFixture(FixtureBase):
|
|
"""Parametrize over shapes / dtypes for independent custom-signature ops."""
|
|
|
|
PARAMS = [
|
|
("n_total, dtype", [
|
|
pytest.param(1_048_576, torch.float16, marks=pytest.mark.smoke),
|
|
pytest.param(1_048_576, torch.bfloat16, marks=pytest.mark.full),
|
|
pytest.param(1_048_576, torch.float32, marks=pytest.mark.full),
|
|
]),
|
|
]
|
|
|
|
|
|
class IndependentEdgeFixture(FixtureBase):
|
|
"""L4 edge-case fixture: fp32, 4K elements."""
|
|
|
|
PARAMS = [
|
|
("n_total, dtype", [
|
|
pytest.param(4096, torch.float32, marks=pytest.mark.smoke),
|
|
]),
|
|
]
|
|
|
|
|
|
# --- L1: where ---
|
|
|
|
@IndependentFixture
|
|
def test_where(n_total: int, dtype: torch.dtype) -> None:
|
|
from tileops.ops.elementwise import WhereOp
|
|
|
|
cond = torch.randint(0, 2, (n_total,), device="cuda").bool()
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
y = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.where(cond, x, y)
|
|
op = WhereOp(N_total=n_total, dtype=dtype)
|
|
out = op(cond, x, y)
|
|
torch.testing.assert_close(out, ref, atol=0, rtol=0)
|
|
print("All checks passed for WhereOp.")
|
|
|
|
|
|
# --- L1: clamp ---
|
|
|
|
@IndependentFixture
|
|
def test_clamp(n_total: int, dtype: torch.dtype) -> None:
|
|
from tileops.ops.elementwise import ClampOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.clamp(x, -0.5, 0.5)
|
|
op = ClampOp(N_total=n_total, dtype=dtype, min_val=-0.5, max_val=0.5)
|
|
out = op(x)
|
|
if dtype == torch.float16:
|
|
tol = {"atol": 1e-3, "rtol": 1e-3}
|
|
elif dtype == torch.bfloat16:
|
|
tol = {"atol": 1.6e-2, "rtol": 1.6e-2}
|
|
else:
|
|
tol = {"atol": 1e-5, "rtol": 1e-5}
|
|
torch.testing.assert_close(out, ref, **tol)
|
|
print("All checks passed for ClampOp.")
|
|
|
|
|
|
# --- L1: masked_fill ---
|
|
|
|
@IndependentFixture
|
|
def test_masked_fill(n_total: int, dtype: torch.dtype) -> None:
|
|
from tileops.ops.elementwise import MaskedFillOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
mask = torch.randint(0, 2, (n_total,), device="cuda").bool()
|
|
# Use -100.0 to avoid fp16 overflow (fp16 max ~65504)
|
|
fill_value = -100.0
|
|
ref = x.masked_fill(mask, fill_value)
|
|
op = MaskedFillOp(N_total=n_total, dtype=dtype, fill_value=fill_value)
|
|
out = op(x, mask)
|
|
if dtype == torch.float16:
|
|
tol = {"atol": 1e-3, "rtol": 1e-3}
|
|
elif dtype == torch.bfloat16:
|
|
tol = {"atol": 1.6e-2, "rtol": 1.6e-2}
|
|
else:
|
|
tol = {"atol": 1e-5, "rtol": 1e-5}
|
|
torch.testing.assert_close(out, ref, **tol)
|
|
print("All checks passed for MaskedFillOp.")
|
|
|
|
|
|
# --- L1: nan_to_num ---
|
|
|
|
@IndependentFixture
|
|
def test_nan_to_num(n_total: int, dtype: torch.dtype) -> None:
|
|
from tileops.ops.elementwise import NanToNumOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
quarter = n_total // 4
|
|
x[:quarter] = float("nan")
|
|
x[quarter:2 * quarter] = float("inf")
|
|
x[2 * quarter:3 * quarter] = float("-inf")
|
|
ref = torch.nan_to_num(x, nan=0.0, posinf=1e4, neginf=-1e4)
|
|
op = NanToNumOp(N_total=n_total, dtype=dtype, nan_val=0.0, posinf_val=1e4, neginf_val=-1e4)
|
|
out = op(x)
|
|
if dtype == torch.float16:
|
|
tol = {"atol": 1e-3, "rtol": 1e-3}
|
|
elif dtype == torch.bfloat16:
|
|
tol = {"atol": 1.6e-2, "rtol": 1.6e-2}
|
|
else:
|
|
tol = {"atol": 1e-5, "rtol": 1e-5}
|
|
torch.testing.assert_close(out, ref, **tol, equal_nan=True)
|
|
print("All checks passed for NanToNumOp.")
|
|
|
|
|
|
# --- L1: alibi ---
|
|
|
|
class AlibiFixture(FixtureBase):
|
|
PARAMS = [
|
|
("seq_len, num_heads, dtype", [
|
|
pytest.param(128, 8, torch.float16, marks=pytest.mark.smoke),
|
|
pytest.param(128, 8, torch.float32, marks=pytest.mark.full),
|
|
]),
|
|
]
|
|
|
|
|
|
@AlibiFixture
|
|
def test_alibi(seq_len: int, num_heads: int, dtype: torch.dtype) -> None:
|
|
from tileops.ops.elementwise import AlibiOp
|
|
|
|
op = AlibiOp(seq_len=seq_len, num_heads=num_heads, dtype=dtype)
|
|
out = op()
|
|
|
|
# Reference: slope_h = 2^(-8*(h+1)/H), bias = -slope * |i - j|
|
|
positions = torch.arange(seq_len, device="cuda", dtype=torch.float32)
|
|
dist = (positions.unsqueeze(1) - positions.unsqueeze(0)).abs()
|
|
slopes = torch.pow(
|
|
2.0,
|
|
-8.0 * torch.arange(1, num_heads + 1, device="cuda", dtype=torch.float32) / num_heads,
|
|
)
|
|
ref = (-slopes[:, None, None] * dist[None, :, :]).to(dtype)
|
|
|
|
tol = {"atol": 1e-2, "rtol": 1e-2} if dtype == torch.float16 else {"atol": 1e-5, "rtol": 1e-5}
|
|
torch.testing.assert_close(out, ref, **tol)
|
|
print("All checks passed for AlibiOp.")
|
|
|
|
|
|
# --- L1: sinusoidal ---
|
|
|
|
class SinusoidalFixture(FixtureBase):
|
|
PARAMS = [
|
|
("seq_len, d_model, dtype", [
|
|
pytest.param(512, 256, torch.float16, marks=pytest.mark.smoke),
|
|
pytest.param(512, 256, torch.float32, marks=pytest.mark.full),
|
|
]),
|
|
]
|
|
|
|
|
|
@SinusoidalFixture
|
|
def test_sinusoidal(seq_len: int, d_model: int, dtype: torch.dtype) -> None:
|
|
|
|
from tileops.ops.elementwise import SinusoidalOp
|
|
|
|
op = SinusoidalOp(seq_len=seq_len, d_model=d_model, dtype=dtype)
|
|
out = op()
|
|
|
|
# Reference
|
|
pos = torch.arange(seq_len, device="cuda", dtype=torch.float32).unsqueeze(1)
|
|
dim_pairs = torch.arange(0, d_model, 2, device="cuda", dtype=torch.float32)
|
|
angles = pos / torch.pow(10000.0, dim_pairs / d_model)
|
|
ref = torch.zeros(seq_len, d_model, device="cuda", dtype=torch.float32)
|
|
ref[:, 0::2] = torch.sin(angles)
|
|
ref[:, 1::2] = torch.cos(angles)
|
|
ref = ref.to(dtype)
|
|
|
|
if dtype == torch.float16:
|
|
tol = {"atol": 1e-3, "rtol": 1e-3}
|
|
else:
|
|
tol = {"atol": 1e-5, "rtol": 1e-5}
|
|
torch.testing.assert_close(out, ref, **tol)
|
|
print("All checks passed for SinusoidalOp.")
|
|
|
|
|
|
# ===========================================================================
|
|
# L2 — Dtype x Size (4 cases for clamp)
|
|
# ===========================================================================
|
|
|
|
|
|
class ClampDtypeSizeFixture(FixtureBase):
|
|
PARAMS = [
|
|
("n_total, dtype", [
|
|
pytest.param(1_048_576, torch.float32, marks=pytest.mark.smoke),
|
|
pytest.param(1_048_576, torch.bfloat16, marks=pytest.mark.full),
|
|
pytest.param(4096, torch.float16, marks=pytest.mark.full),
|
|
pytest.param(16_777_216, torch.float16, marks=pytest.mark.full),
|
|
]),
|
|
]
|
|
|
|
|
|
@ClampDtypeSizeFixture
|
|
def test_clamp_dtype_size(n_total: int, dtype: torch.dtype) -> None:
|
|
from tileops.ops.elementwise import ClampOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.clamp(x, -0.5, 0.5)
|
|
op = ClampOp(N_total=n_total, dtype=dtype, min_val=-0.5, max_val=0.5)
|
|
out = op(x)
|
|
if dtype == torch.float16:
|
|
tol = {"atol": 1e-3, "rtol": 1e-3}
|
|
elif dtype == torch.bfloat16:
|
|
tol = {"atol": 1.6e-2, "rtol": 1.6e-2}
|
|
else:
|
|
tol = {"atol": 1e-5, "rtol": 1e-5}
|
|
torch.testing.assert_close(out, ref, **tol)
|
|
print("All checks passed for ClampOp dtype/size variant.")
|
|
|
|
|
|
# ===========================================================================
|
|
# L4 — Edge Cases (8 cases, fp32, 4K)
|
|
# ===========================================================================
|
|
|
|
|
|
@IndependentEdgeFixture
|
|
def test_clamp_min_gt_max(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: min > max -- PyTorch clamp semantics: min wins (output = min_val)."""
|
|
from tileops.ops.elementwise import ClampOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
# When min > max, PyTorch clamp returns min_val for all elements
|
|
ref = torch.clamp(x, min=0.5, max=-0.5)
|
|
op = ClampOp(N_total=n_total, dtype=dtype, min_val=0.5, max_val=-0.5)
|
|
out = op(x)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
|
|
print("All checks passed for ClampOp min>max edge case.")
|
|
|
|
|
|
@IndependentEdgeFixture
|
|
def test_clamp_upper_only(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: min=None, max=0.5 (upper bound only)."""
|
|
from tileops.ops.elementwise import ClampOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.clamp(x, min=None, max=0.5)
|
|
op = ClampOp(N_total=n_total, dtype=dtype, min_val=None, max_val=0.5)
|
|
out = op(x)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
|
|
print("All checks passed for ClampOp upper-only edge case.")
|
|
|
|
|
|
@IndependentEdgeFixture
|
|
def test_clamp_lower_only(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: min=-0.5, max=None (lower bound only)."""
|
|
from tileops.ops.elementwise import ClampOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.clamp(x, min=-0.5, max=None)
|
|
op = ClampOp(N_total=n_total, dtype=dtype, min_val=-0.5, max_val=None)
|
|
out = op(x)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
|
|
print("All checks passed for ClampOp lower-only edge case.")
|
|
|
|
|
|
@IndependentEdgeFixture
|
|
def test_masked_fill_all_true(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: all True mask -> all values replaced."""
|
|
from tileops.ops.elementwise import MaskedFillOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
mask = torch.ones(n_total, device="cuda", dtype=torch.bool)
|
|
fill_value = -1e9
|
|
ref = x.masked_fill(mask, fill_value)
|
|
op = MaskedFillOp(N_total=n_total, dtype=dtype, fill_value=fill_value)
|
|
out = op(x, mask)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
|
|
print("All checks passed for MaskedFillOp all-true edge case.")
|
|
|
|
|
|
@IndependentEdgeFixture
|
|
def test_masked_fill_all_false(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: all False mask -> input unchanged."""
|
|
from tileops.ops.elementwise import MaskedFillOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
mask = torch.zeros(n_total, device="cuda", dtype=torch.bool)
|
|
fill_value = -1e9
|
|
ref = x.masked_fill(mask, fill_value)
|
|
op = MaskedFillOp(N_total=n_total, dtype=dtype, fill_value=fill_value)
|
|
out = op(x, mask)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
|
|
print("All checks passed for MaskedFillOp all-false edge case.")
|
|
|
|
|
|
@IndependentEdgeFixture
|
|
def test_where_all_true(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: all True cond -> output = x."""
|
|
from tileops.ops.elementwise import WhereOp
|
|
|
|
cond = torch.ones(n_total, device="cuda", dtype=torch.bool)
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
y = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.where(cond, x, y)
|
|
op = WhereOp(N_total=n_total, dtype=dtype)
|
|
out = op(cond, x, y)
|
|
torch.testing.assert_close(out, ref, atol=0, rtol=0)
|
|
print("All checks passed for WhereOp all-true edge case.")
|
|
|
|
|
|
@IndependentEdgeFixture
|
|
def test_where_all_false(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: all False cond -> output = y."""
|
|
from tileops.ops.elementwise import WhereOp
|
|
|
|
cond = torch.zeros(n_total, device="cuda", dtype=torch.bool)
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
y = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.where(cond, x, y)
|
|
op = WhereOp(N_total=n_total, dtype=dtype)
|
|
out = op(cond, x, y)
|
|
torch.testing.assert_close(out, ref, atol=0, rtol=0)
|
|
print("All checks passed for WhereOp all-false edge case.")
|
|
|
|
|
|
@IndependentEdgeFixture
|
|
def test_nan_to_num_edge(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: explicit [NaN, Inf, -Inf, 1.0] pattern."""
|
|
from tileops.ops.elementwise import NanToNumOp
|
|
|
|
x = torch.zeros(n_total, device="cuda", dtype=dtype)
|
|
# Fill pattern: NaN, Inf, -Inf, 1.0, repeating
|
|
for k in range(0, n_total, 4):
|
|
x[k] = float("nan")
|
|
if k + 1 < n_total:
|
|
x[k + 1] = float("inf")
|
|
if k + 2 < n_total:
|
|
x[k + 2] = float("-inf")
|
|
if k + 3 < n_total:
|
|
x[k + 3] = 1.0
|
|
|
|
ref = torch.nan_to_num(x, nan=0.0, posinf=1e4, neginf=-1e4)
|
|
op = NanToNumOp(N_total=n_total, dtype=dtype, nan_val=0.0, posinf_val=1e4, neginf_val=-1e4)
|
|
out = op(x)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5, equal_nan=True)
|
|
print("All checks passed for NanToNumOp edge case.")
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_independent_special_rejects_non_float_dtype() -> None:
|
|
from tileops.kernels.elementwise import ClampKernel
|
|
with pytest.raises(ValueError, match="only supports dtypes"):
|
|
ClampKernel(N_total=16, dtype=torch.int32)
|
|
|
|
|
|
# ===========================================================================
|
|
# Negative tests: forward() dtype / numel validation
|
|
# ===========================================================================
|
|
|
|
|
|
@pytest.mark.smoke
|
|
@pytest.mark.parametrize("op_cls, kwargs", [
|
|
pytest.param("EluOp", {"alpha": 1.0}, id="elu"),
|
|
pytest.param("HardtanhOp", {"min_val": -1.0, "max_val": 1.0}, id="hardtanh"),
|
|
pytest.param("SoftplusOp", {"beta": 1.0, "threshold": 20.0}, id="softplus"),
|
|
pytest.param("ClampOp", {"min_val": -0.5, "max_val": 0.5}, id="clamp"),
|
|
])
|
|
def test_forward_rejects_wrong_dtype(op_cls: str, kwargs: dict) -> None:
|
|
"""forward() must raise ValueError when input dtype mismatches."""
|
|
import tileops.ops.elementwise as mod
|
|
cls = getattr(mod, op_cls)
|
|
op = cls(N_total=1024, dtype=torch.float16, **kwargs)
|
|
x = torch.randn(1024, device="cuda", dtype=torch.float32)
|
|
with pytest.raises(ValueError, match="dtype"):
|
|
op(x)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
@pytest.mark.parametrize("op_cls, kwargs", [
|
|
pytest.param("EluOp", {"alpha": 1.0}, id="elu"),
|
|
pytest.param("HardtanhOp", {"min_val": -1.0, "max_val": 1.0}, id="hardtanh"),
|
|
pytest.param("SoftplusOp", {"beta": 1.0, "threshold": 20.0}, id="softplus"),
|
|
pytest.param("ClampOp", {"min_val": -0.5, "max_val": 0.5}, id="clamp"),
|
|
])
|
|
def test_forward_rejects_wrong_numel(op_cls: str, kwargs: dict) -> None:
|
|
"""forward() must raise ValueError when input numel mismatches."""
|
|
import tileops.ops.elementwise as mod
|
|
cls = getattr(mod, op_cls)
|
|
op = cls(N_total=1024, dtype=torch.float16, **kwargs)
|
|
x = torch.randn(512, device="cuda", dtype=torch.float16)
|
|
with pytest.raises(ValueError, match="elements"):
|
|
op(x)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
@pytest.mark.parametrize("op_cls, kwargs", [
|
|
pytest.param("MaskedFillOp", {"fill_value": -100.0}, id="masked_fill"),
|
|
])
|
|
def test_masked_fill_forward_rejects_wrong_dtype(op_cls: str, kwargs: dict) -> None:
|
|
"""MaskedFillOp forward() must raise ValueError when input dtype mismatches."""
|
|
import tileops.ops.elementwise as mod
|
|
cls = getattr(mod, op_cls)
|
|
op = cls(N_total=1024, dtype=torch.float16, **kwargs)
|
|
x = torch.randn(1024, device="cuda", dtype=torch.float32)
|
|
mask = torch.ones(1024, device="cuda", dtype=torch.bool)
|
|
with pytest.raises(ValueError, match="dtype"):
|
|
op(x, mask)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
@pytest.mark.parametrize("op_cls, kwargs", [
|
|
pytest.param("MaskedFillOp", {"fill_value": -100.0}, id="masked_fill"),
|
|
])
|
|
def test_masked_fill_forward_rejects_wrong_numel(op_cls: str, kwargs: dict) -> None:
|
|
"""MaskedFillOp forward() must raise ValueError when input numel mismatches."""
|
|
import tileops.ops.elementwise as mod
|
|
cls = getattr(mod, op_cls)
|
|
op = cls(N_total=1024, dtype=torch.float16, **kwargs)
|
|
x = torch.randn(512, device="cuda", dtype=torch.float16)
|
|
mask = torch.ones(512, device="cuda", dtype=torch.bool)
|
|
with pytest.raises(ValueError, match="elements"):
|
|
op(x, mask)
|
|
|
|
|
|
# ===========================================================================
|
|
# Negative tests: __init__() scalar parameter validation
|
|
# ===========================================================================
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_elu_rejects_unrepresentable_alpha() -> None:
|
|
"""EluOp must reject alpha that overflows the kernel dtype."""
|
|
from tileops.ops.elementwise import EluOp
|
|
with pytest.raises((ValueError, TypeError)):
|
|
EluOp(N_total=1024, dtype=torch.float16, alpha=1e6)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_hardtanh_rejects_unrepresentable_min_val() -> None:
|
|
"""HardtanhOp must reject min_val that overflows the kernel dtype."""
|
|
from tileops.ops.elementwise import HardtanhOp
|
|
with pytest.raises((ValueError, TypeError)):
|
|
HardtanhOp(N_total=1024, dtype=torch.float16, min_val=1e6)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_hardtanh_rejects_unrepresentable_max_val() -> None:
|
|
"""HardtanhOp must reject max_val that overflows the kernel dtype."""
|
|
from tileops.ops.elementwise import HardtanhOp
|
|
with pytest.raises((ValueError, TypeError)):
|
|
HardtanhOp(N_total=1024, dtype=torch.float16, max_val=1e6)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_softplus_rejects_unrepresentable_beta() -> None:
|
|
"""SoftplusOp must reject beta that overflows the kernel dtype."""
|
|
from tileops.ops.elementwise import SoftplusOp
|
|
with pytest.raises((ValueError, TypeError)):
|
|
SoftplusOp(N_total=1024, dtype=torch.float16, beta=1e6)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_softplus_rejects_unrepresentable_threshold() -> None:
|
|
"""SoftplusOp must reject threshold that overflows the kernel dtype."""
|
|
from tileops.ops.elementwise import SoftplusOp
|
|
with pytest.raises((ValueError, TypeError)):
|
|
SoftplusOp(N_total=1024, dtype=torch.float16, threshold=1e6)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_clamp_rejects_unrepresentable_min_val() -> None:
|
|
"""ClampOp must reject min_val that overflows the kernel dtype."""
|
|
from tileops.ops.elementwise import ClampOp
|
|
with pytest.raises((ValueError, TypeError)):
|
|
ClampOp(N_total=1024, dtype=torch.float16, min_val=1e6)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_clamp_rejects_unrepresentable_max_val() -> None:
|
|
"""ClampOp must reject max_val that overflows the kernel dtype."""
|
|
from tileops.ops.elementwise import ClampOp
|
|
with pytest.raises((ValueError, TypeError)):
|
|
ClampOp(N_total=1024, dtype=torch.float16, max_val=1e6)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_masked_fill_forward_rejects_cpu_mask() -> None:
|
|
"""MaskedFillOp forward() must raise ValueError when mask is not on CUDA."""
|
|
from tileops.ops.elementwise import MaskedFillOp
|
|
op = MaskedFillOp(N_total=1024, dtype=torch.float16, fill_value=-100.0)
|
|
x = torch.randn(1024, device="cuda", dtype=torch.float16)
|
|
mask = torch.ones(1024, dtype=torch.bool) # CPU mask
|
|
with pytest.raises(ValueError, match="Mask must be a CUDA tensor"):
|
|
op(x, mask)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_masked_fill_forward_rejects_non_bool_mask() -> None:
|
|
"""MaskedFillOp forward() must raise ValueError when mask dtype is not bool."""
|
|
from tileops.ops.elementwise import MaskedFillOp
|
|
op = MaskedFillOp(N_total=1024, dtype=torch.float16, fill_value=-100.0)
|
|
x = torch.randn(1024, device="cuda", dtype=torch.float16)
|
|
mask = torch.ones(1024, device="cuda", dtype=torch.float32) # wrong dtype
|
|
with pytest.raises(ValueError, match="mask.dtype"):
|
|
op(x, mask)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_masked_fill_forward_rejects_wrong_mask_numel() -> None:
|
|
"""MaskedFillOp forward() must raise ValueError when mask numel mismatches."""
|
|
from tileops.ops.elementwise import MaskedFillOp
|
|
op = MaskedFillOp(N_total=1024, dtype=torch.float16, fill_value=-100.0)
|
|
x = torch.randn(1024, device="cuda", dtype=torch.float16)
|
|
mask = torch.ones(512, device="cuda", dtype=torch.bool) # wrong numel
|
|
with pytest.raises(ValueError, match="elements"):
|
|
op(x, mask)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_elu_rejects_infinite_alpha() -> None:
|
|
"""EluOp must reject infinite alpha."""
|
|
from tileops.ops.elementwise import EluOp
|
|
with pytest.raises(ValueError, match="finite"):
|
|
EluOp(N_total=1024, dtype=torch.float32, alpha=float("inf"))
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_softplus_rejects_non_numeric_beta() -> None:
|
|
"""SoftplusOp must reject non-numeric beta."""
|
|
from tileops.ops.elementwise import SoftplusOp
|
|
with pytest.raises(TypeError, match="int/float"):
|
|
SoftplusOp(N_total=1024, dtype=torch.float32, beta="bad")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-vvs"])
|