forked from ccf-ai-infra/TileOPs-Metax
739 lines
26 KiB
Python
739 lines
26 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 (
|
|
ClampScalarFwdOp,
|
|
EluFwdOp,
|
|
HardtanhFwdOp,
|
|
IsfiniteFwdOp,
|
|
IsinfFwdOp,
|
|
IsnanFwdOp,
|
|
SoftplusFwdOp,
|
|
)
|
|
|
|
|
|
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.smoke),
|
|
pytest.param(1_048_576, torch.float32, marks=pytest.mark.smoke),
|
|
]),
|
|
]
|
|
|
|
|
|
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 fixture 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, IsnanFwdOp, torch.isnan)
|
|
|
|
|
|
@SpecialFixture
|
|
def test_isinf(n_total: int, dtype: torch.dtype) -> None:
|
|
_make_special_test(n_total, dtype, IsinfFwdOp, torch.isinf)
|
|
|
|
|
|
@SpecialFixture
|
|
def test_isfinite(n_total: int, dtype: torch.dtype) -> None:
|
|
_make_special_test(n_total, dtype, IsfiniteFwdOp, 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, IsnanFwdOp, 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, IsinfFwdOp, 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, IsfiniteFwdOp, torch.isfinite, gen_fn=_all_finite)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_special_predicates_reject_non_float_dtype() -> None:
|
|
from tileops.kernels.elementwise import IsnanFwdKernel
|
|
|
|
with pytest.raises(ValueError, match="only supports dtypes"):
|
|
IsnanFwdKernel(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.smoke),
|
|
pytest.param(1_048_576, torch.float32, marks=pytest.mark.smoke),
|
|
]),
|
|
]
|
|
|
|
|
|
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 WhereFwdOp
|
|
|
|
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 = WhereFwdOp(condition=(n_total,), input=(n_total,), other=(n_total,), dtype=dtype)
|
|
out = op(cond, x, y)
|
|
torch.testing.assert_close(out, ref, atol=0, rtol=0)
|
|
|
|
|
|
# --- L1: clamp ---
|
|
|
|
@IndependentFixture
|
|
def test_clamp(n_total: int, dtype: torch.dtype) -> None:
|
|
from tileops.ops.elementwise import ClampScalarFwdOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.clamp(x, -0.5, 0.5)
|
|
op = ClampScalarFwdOp(input=(n_total,), min=-0.5, max=0.5, dtype=dtype)
|
|
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)
|
|
|
|
|
|
# --- L1: masked_fill ---
|
|
|
|
@IndependentFixture
|
|
def test_masked_fill(n_total: int, dtype: torch.dtype) -> None:
|
|
from tileops.ops.elementwise import MaskedFillScalarFwdOp
|
|
|
|
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 = MaskedFillScalarFwdOp(input=(n_total,), mask=(n_total,), value=fill_value, dtype=dtype)
|
|
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)
|
|
|
|
|
|
# --- L1: nan_to_num ---
|
|
|
|
@IndependentFixture
|
|
def test_nan_to_num(n_total: int, dtype: torch.dtype) -> None:
|
|
from tileops.ops.elementwise import NanToNumFwdOp
|
|
|
|
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 = NanToNumFwdOp(N_total=n_total, dtype=dtype, nan=0.0, posinf=1e4, neginf=-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)
|
|
|
|
|
|
# --- 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.smoke),
|
|
]),
|
|
]
|
|
|
|
|
|
@AlibiFixture
|
|
def test_alibi(seq_len: int, num_heads: int, dtype: torch.dtype) -> None:
|
|
from tileops.ops.elementwise import AlibiFwdOp
|
|
|
|
op = AlibiFwdOp(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)
|
|
|
|
|
|
# --- 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.smoke),
|
|
]),
|
|
]
|
|
|
|
|
|
@SinusoidalFixture
|
|
def test_sinusoidal(seq_len: int, d_model: int, dtype: torch.dtype) -> None:
|
|
|
|
from tileops.ops.elementwise import SinusoidalFwdOp
|
|
|
|
op = SinusoidalFwdOp(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)
|
|
|
|
|
|
# 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.smoke),
|
|
pytest.param(4096, torch.float16, marks=pytest.mark.smoke),
|
|
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 ClampScalarFwdOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.clamp(x, -0.5, 0.5)
|
|
op = ClampScalarFwdOp(input=(n_total,), min=-0.5, max=0.5, dtype=dtype)
|
|
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)
|
|
|
|
|
|
# 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 ClampScalarFwdOp
|
|
|
|
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 = ClampScalarFwdOp(input=(n_total,), min=0.5, max=-0.5, dtype=dtype)
|
|
out = op(x)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@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 ClampScalarFwdOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.clamp(x, min=None, max=0.5)
|
|
op = ClampScalarFwdOp(input=(n_total,), min=None, max=0.5, dtype=dtype)
|
|
out = op(x)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@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 ClampScalarFwdOp
|
|
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
ref = torch.clamp(x, min=-0.5, max=None)
|
|
op = ClampScalarFwdOp(input=(n_total,), min=-0.5, max=None, dtype=dtype)
|
|
out = op(x)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@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 MaskedFillScalarFwdOp
|
|
|
|
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 = MaskedFillScalarFwdOp(input=(n_total,), mask=(n_total,), value=fill_value, dtype=dtype)
|
|
out = op(x, mask)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@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 MaskedFillScalarFwdOp
|
|
|
|
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 = MaskedFillScalarFwdOp(input=(n_total,), mask=(n_total,), value=fill_value, dtype=dtype)
|
|
out = op(x, mask)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5)
|
|
|
|
|
|
@IndependentEdgeFixture
|
|
def test_where_all_true(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: all True cond -> output = x."""
|
|
from tileops.ops.elementwise import WhereFwdOp
|
|
|
|
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 = WhereFwdOp(condition=(n_total,), input=(n_total,), other=(n_total,), dtype=dtype)
|
|
out = op(cond, x, y)
|
|
torch.testing.assert_close(out, ref, atol=0, rtol=0)
|
|
|
|
|
|
@IndependentEdgeFixture
|
|
def test_where_all_false(n_total: int, dtype: torch.dtype) -> None:
|
|
"""Edge: all False cond -> output = y."""
|
|
from tileops.ops.elementwise import WhereFwdOp
|
|
|
|
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 = WhereFwdOp(condition=(n_total,), input=(n_total,), other=(n_total,), dtype=dtype)
|
|
out = op(cond, x, y)
|
|
torch.testing.assert_close(out, ref, atol=0, rtol=0)
|
|
|
|
|
|
@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 NanToNumFwdOp
|
|
|
|
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 = NanToNumFwdOp(N_total=n_total, dtype=dtype, nan=0.0, posinf=1e4, neginf=-1e4)
|
|
out = op(x)
|
|
torch.testing.assert_close(out, ref, atol=1e-5, rtol=1e-5, equal_nan=True)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_independent_special_rejects_non_float_dtype() -> None:
|
|
from tileops.kernels.elementwise import ClampFwdKernel
|
|
with pytest.raises(ValueError, match="only supports dtypes"):
|
|
ClampFwdKernel(N_total=16, dtype=torch.int32)
|
|
|
|
|
|
# Negative tests: forward() dtype / numel validation
|
|
|
|
|
|
@pytest.mark.smoke
|
|
@pytest.mark.parametrize("op_cls, kwargs", [
|
|
pytest.param("EluFwdOp", {"alpha": 1.0}, id="elu"),
|
|
pytest.param("HardtanhFwdOp", {"min_val": -1.0, "max_val": 1.0}, id="hardtanh"),
|
|
pytest.param("SoftplusFwdOp", {"beta": 1.0, "threshold": 20.0}, id="softplus"),
|
|
pytest.param("ClampScalarFwdOp", {"min": -0.5, "max": 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)
|
|
if cls.__name__ == "ClampScalarFwdOp":
|
|
op = cls(input=(1024,), dtype=torch.float16, **kwargs)
|
|
else:
|
|
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("EluFwdOp", {"alpha": 1.0}, id="elu"),
|
|
pytest.param("HardtanhFwdOp", {"min_val": -1.0, "max_val": 1.0}, id="hardtanh"),
|
|
pytest.param("SoftplusFwdOp", {"beta": 1.0, "threshold": 20.0}, id="softplus"),
|
|
])
|
|
def test_forward_rejects_wrong_numel(op_cls: str, kwargs: dict) -> None:
|
|
"""forward() must raise ValueError when input numel mismatches.
|
|
|
|
ClampScalarFwdOp validates the full input.shape (not just numel), so
|
|
its mismatch case is covered by
|
|
test_clamp_scalar_rejects_same_numel_wrong_shape in
|
|
tests/ops/test_special_elementwise_conformance.py.
|
|
"""
|
|
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("MaskedFillScalarFwdOp", {"value": -100.0}, id="masked_fill"),
|
|
])
|
|
def test_masked_fill_forward_rejects_wrong_dtype(op_cls: str, kwargs: dict) -> None:
|
|
"""MaskedFillFwdOp forward() must raise ValueError when input dtype mismatches."""
|
|
import tileops.ops.elementwise as mod
|
|
cls = getattr(mod, op_cls)
|
|
op = cls(input=(1024,), mask=(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("MaskedFillScalarFwdOp", {"value": -100.0}, id="masked_fill"),
|
|
])
|
|
def test_masked_fill_forward_rejects_wrong_numel(op_cls: str, kwargs: dict) -> None:
|
|
"""MaskedFillFwdOp forward() must raise ValueError when input shape mismatches."""
|
|
import tileops.ops.elementwise as mod
|
|
cls = getattr(mod, op_cls)
|
|
op = cls(input=(1024,), mask=(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="input.shape"):
|
|
op(x, mask)
|
|
|
|
|
|
# Negative tests: __init__() scalar parameter validation
|
|
|
|
|
|
@pytest.mark.smoke
|
|
@pytest.mark.parametrize("make_op", [
|
|
pytest.param(lambda: EluFwdOp(N_total=1024, dtype=torch.float16, alpha=1e6),
|
|
id="elu-alpha"),
|
|
pytest.param(lambda: HardtanhFwdOp(N_total=1024, dtype=torch.float16, min_val=1e6),
|
|
id="hardtanh-min_val"),
|
|
pytest.param(lambda: HardtanhFwdOp(N_total=1024, dtype=torch.float16, max_val=1e6),
|
|
id="hardtanh-max_val"),
|
|
pytest.param(lambda: SoftplusFwdOp(N_total=1024, dtype=torch.float16, beta=1e6),
|
|
id="softplus-beta"),
|
|
pytest.param(lambda: SoftplusFwdOp(N_total=1024, dtype=torch.float16, threshold=1e6),
|
|
id="softplus-threshold"),
|
|
pytest.param(lambda: ClampScalarFwdOp(input=(1024,), min=1e6, dtype=torch.float16),
|
|
id="clamp-min"),
|
|
pytest.param(lambda: ClampScalarFwdOp(input=(1024,), max=1e6, dtype=torch.float16),
|
|
id="clamp-max"),
|
|
])
|
|
def test_scalar_param_rejects_unrepresentable(make_op) -> None:
|
|
"""__init__ must reject scalar params that overflow the kernel dtype."""
|
|
with pytest.raises((ValueError, TypeError)):
|
|
make_op()
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_masked_fill_forward_rejects_cpu_mask() -> None:
|
|
"""MaskedFillFwdOp forward() must raise ValueError when mask is not on CUDA."""
|
|
from tileops.ops.elementwise import MaskedFillScalarFwdOp
|
|
op = MaskedFillScalarFwdOp(input=(1024,), mask=(1024,), value=-100.0, dtype=torch.float16)
|
|
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:
|
|
"""MaskedFillFwdOp forward() must raise ValueError when mask dtype is not bool."""
|
|
from tileops.ops.elementwise import MaskedFillScalarFwdOp
|
|
op = MaskedFillScalarFwdOp(input=(1024,), mask=(1024,), value=-100.0, dtype=torch.float16)
|
|
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:
|
|
"""MaskedFillFwdOp forward() must raise ValueError when mask numel mismatches."""
|
|
from tileops.ops.elementwise import MaskedFillScalarFwdOp
|
|
op = MaskedFillScalarFwdOp(input=(1024,), mask=(1024,), value=-100.0, dtype=torch.float16)
|
|
x = torch.randn(1024, device="cuda", dtype=torch.float16)
|
|
mask = torch.ones(512, device="cuda", dtype=torch.bool) # wrong shape
|
|
with pytest.raises(ValueError, match="mask.shape"):
|
|
op(x, mask)
|
|
|
|
|
|
# MaskedFillScalar: int / uint / bool dtype coverage
|
|
|
|
|
|
_MASKED_FILL_INT_DTYPES = [
|
|
torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64,
|
|
]
|
|
|
|
|
|
def _masked_fill_int_inputs(n_total: int, dtype: torch.dtype):
|
|
iinfo = torch.iinfo(dtype)
|
|
lo = max(iinfo.min, -1000)
|
|
hi = min(iinfo.max, 1000) + 1
|
|
x = torch.randint(lo, hi, (n_total,), device="cuda", dtype=dtype)
|
|
mask = torch.randint(0, 2, (n_total,), device="cuda").bool()
|
|
return x, mask
|
|
|
|
|
|
@pytest.mark.smoke
|
|
@pytest.mark.parametrize("dtype", _MASKED_FILL_INT_DTYPES)
|
|
def test_masked_fill_int_dtypes(dtype: torch.dtype) -> None:
|
|
"""L1: each manifest int dtype matches PyTorch on a representative fill."""
|
|
from tileops.ops.elementwise import MaskedFillScalarFwdOp
|
|
|
|
n_total = 4096
|
|
fill_value = 7 # arbitrary in-range value; the contract is parity with PyTorch.
|
|
x, mask = _masked_fill_int_inputs(n_total, dtype)
|
|
ref = x.masked_fill(mask, fill_value)
|
|
op = MaskedFillScalarFwdOp(
|
|
input=(n_total,), mask=(n_total,), value=fill_value, dtype=dtype,
|
|
)
|
|
out = op(x, mask)
|
|
torch.testing.assert_close(out, ref, atol=0, rtol=0)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_masked_fill_uint8_wraps_negative_int() -> None:
|
|
"""uint8 wraps a negative Python int via two's complement (PyTorch: -1 -> 255)."""
|
|
from tileops.ops.elementwise import MaskedFillScalarFwdOp
|
|
|
|
n_total = 4096
|
|
x, mask = _masked_fill_int_inputs(n_total, torch.uint8)
|
|
ref = x.masked_fill(mask, -1)
|
|
op = MaskedFillScalarFwdOp(
|
|
input=(n_total,), mask=(n_total,), value=-1, dtype=torch.uint8,
|
|
)
|
|
torch.testing.assert_close(op(x, mask), ref, atol=0, rtol=0)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_masked_fill_int_truncates_fractional_float() -> None:
|
|
"""Integer dtypes truncate a float fill toward zero (PyTorch: 1.5 -> 1)."""
|
|
from tileops.ops.elementwise import MaskedFillScalarFwdOp
|
|
|
|
n_total = 4096
|
|
x, mask = _masked_fill_int_inputs(n_total, torch.int32)
|
|
ref = x.masked_fill(mask, 1.5)
|
|
op = MaskedFillScalarFwdOp(
|
|
input=(n_total,), mask=(n_total,), value=1.5, dtype=torch.int32,
|
|
)
|
|
torch.testing.assert_close(op(x, mask), ref, atol=0, rtol=0)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
@pytest.mark.parametrize("fill_value", [True, False])
|
|
def test_masked_fill_bool(fill_value) -> None:
|
|
"""L1: bool masked_fill coerces non-zero -> True via uint8 storage view."""
|
|
from tileops.ops.elementwise import MaskedFillScalarFwdOp
|
|
|
|
n_total = 4096
|
|
x = torch.randint(0, 2, (n_total,), device="cuda").bool()
|
|
mask = torch.randint(0, 2, (n_total,), device="cuda").bool()
|
|
ref = x.masked_fill(mask, fill_value)
|
|
op = MaskedFillScalarFwdOp(
|
|
input=(n_total,), mask=(n_total,), value=fill_value, dtype=torch.bool,
|
|
)
|
|
out = op(x, mask)
|
|
torch.testing.assert_close(out, ref, atol=0, rtol=0)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
@pytest.mark.parametrize("dtype, fill_value", [
|
|
pytest.param(torch.float16, float("inf"), id="fp16-inf"),
|
|
pytest.param(torch.bfloat16, float("-inf"), id="bf16-neg-inf"),
|
|
pytest.param(torch.float32, float("nan"), id="fp32-nan"),
|
|
])
|
|
def test_masked_fill_float_nonfinite(dtype: torch.dtype, fill_value: float) -> None:
|
|
"""L4: +/-Inf and NaN fill values pass through unchanged (no clamp)."""
|
|
from tileops.ops.elementwise import MaskedFillScalarFwdOp
|
|
|
|
n_total = 4096
|
|
x = torch.randn(n_total, device="cuda", dtype=dtype)
|
|
mask = torch.randint(0, 2, (n_total,), device="cuda").bool()
|
|
ref = x.masked_fill(mask, fill_value)
|
|
op = MaskedFillScalarFwdOp(
|
|
input=(n_total,), mask=(n_total,), value=fill_value, dtype=dtype,
|
|
)
|
|
out = op(x, mask)
|
|
torch.testing.assert_close(out, ref, atol=0, rtol=0, equal_nan=True)
|
|
|
|
|
|
_MASKED_FILL_REJECT_CASES = [
|
|
pytest.param(torch.int8, 200, id="signed-int-overflow"),
|
|
pytest.param(torch.int8, 127.5, id="signed-int-float-just-over"),
|
|
pytest.param(torch.uint8, -256, id="uint8-int-wrap-too-low"),
|
|
pytest.param(torch.uint8, -1.0, id="uint8-float-negative"),
|
|
pytest.param(torch.int32, float("inf"), id="int-inf"),
|
|
pytest.param(torch.int32, float("nan"), id="int-nan"),
|
|
]
|
|
|
|
|
|
@pytest.mark.smoke
|
|
@pytest.mark.parametrize("dtype, fill_value", _MASKED_FILL_REJECT_CASES)
|
|
def test_masked_fill_rejects_when_pytorch_rejects(
|
|
dtype: torch.dtype, fill_value,
|
|
) -> None:
|
|
"""Op must reject every scalar that PyTorch's own masked_fill rejects.
|
|
|
|
The contract is parity, not the error message; assert both call sites
|
|
raise, leaving wording to the implementation.
|
|
"""
|
|
from tileops.ops.elementwise import MaskedFillScalarFwdOp
|
|
|
|
pytorch_mask = torch.tensor([True], device="cuda")
|
|
pytorch_tensor = torch.zeros(1, device="cuda", dtype=dtype)
|
|
with pytest.raises(Exception): # noqa: B017
|
|
pytorch_tensor.masked_fill(pytorch_mask, fill_value)
|
|
|
|
with pytest.raises(Exception): # noqa: B017
|
|
MaskedFillScalarFwdOp(
|
|
input=(1024,), mask=(1024,), value=fill_value, dtype=dtype,
|
|
)
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_elu_rejects_infinite_alpha() -> None:
|
|
"""EluFwdOp must reject infinite alpha."""
|
|
from tileops.ops.elementwise import EluFwdOp
|
|
with pytest.raises(ValueError, match="finite"):
|
|
EluFwdOp(N_total=1024, dtype=torch.float32, alpha=float("inf"))
|
|
|
|
|
|
@pytest.mark.smoke
|
|
def test_softplus_rejects_non_numeric_beta() -> None:
|
|
"""SoftplusFwdOp must reject non-numeric beta."""
|
|
from tileops.ops.elementwise import SoftplusFwdOp
|
|
with pytest.raises(TypeError, match="int/float"):
|
|
SoftplusFwdOp(N_total=1024, dtype=torch.float32, beta="bad")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-vvs"])
|