forked from ccf-ai-infra/TileOPs-Metax
667 lines
28 KiB
Python
667 lines
28 KiB
Python
"""Correctness tests for softmax-family ops (softmax, log_softmax, logsumexp).
|
||
|
||
Tests cover fp32/fp16/bf16 dtypes, 1D-4D inputs, non-contiguous tensors,
|
||
power-of-2 and non-power-of-2 hidden dims, tail-M cases, and validate
|
||
against PyTorch reference implementations.
|
||
|
||
Smoke tests (1 per function, first param) use small data for quick CI.
|
||
Full tests use small data for config breadth + large data for stress.
|
||
|
||
All operators use the spec-conformant interface:
|
||
SoftmaxFwdOp(dim=dim)
|
||
LogSoftmaxFwdOp(dim=dim)
|
||
LogSumExpFwdOp(dtype=dtype, dim=dim, keepdim=keepdim)
|
||
"""
|
||
|
||
import pytest
|
||
import torch
|
||
import torch.nn.functional as F
|
||
|
||
from tests.test_base import FixtureBase, TestBase
|
||
from tileops.ops.reduction.softmax import LogSoftmaxFwdOp, LogSumExpFwdOp, SoftmaxFwdOp
|
||
from workloads.reduction import LogSoftmaxTest as _LogSoftmaxTestWorkload
|
||
from workloads.reduction import LogSumExpTest as _LogSumExpTestWorkload
|
||
from workloads.reduction import SoftmaxTest as _SoftmaxTestWorkload
|
||
|
||
# Tolerances (from docs/design/testing.md)
|
||
|
||
|
||
def _get_tolerances(dtype: torch.dtype) -> tuple[float, float]:
|
||
if dtype == torch.float32:
|
||
return 1e-5, 1e-5
|
||
elif dtype == torch.float16:
|
||
return 1e-3, 1e-3
|
||
else: # bfloat16
|
||
return 1.6e-2, 1.6e-2
|
||
|
||
|
||
# Softmax — spec-conformant interface (shape, dim, dtype)
|
||
|
||
|
||
class SoftmaxFixture(FixtureBase):
|
||
PARAMS = [
|
||
(
|
||
"shape, dim, dtype, tune",
|
||
[
|
||
# Smoke: 2D, dim=-1, fp32, pow2
|
||
pytest.param((32, 256), -1, torch.float32, False, marks=[pytest.mark.smoke, pytest.mark.packaging]),
|
||
pytest.param((32, 256), -1, torch.float16, False, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), -1, torch.bfloat16, False, marks=pytest.mark.smoke),
|
||
# tune=True regression: kernel must be built before autotune runs
|
||
pytest.param((32, 256), -1, torch.float16, True, marks=pytest.mark.full),
|
||
# dim=-1 (default path): dtypes x pow2/non-pow2
|
||
pytest.param((32, 300), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((32, 300), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((32, 300), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, tail-M (non-aligned M)
|
||
pytest.param((33, 256), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((33, 256), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((33, 256), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, 3D input
|
||
pytest.param((2, 16, 256), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((2, 16, 256), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((2, 16, 256), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, 4D input
|
||
pytest.param((2, 4, 8, 256), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((2, 4, 8, 256), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((2, 4, 8, 256), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, large-N (triggers N-tiling path)
|
||
pytest.param((4, 32768), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((4, 32768), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, M×N both non-aligned (single-tile path)
|
||
pytest.param((33, 300), -1, torch.float32, False, marks=pytest.mark.full),
|
||
# dim=-1, M×N both non-aligned (multi-tile, masked loads)
|
||
pytest.param((33, 33000), -1, torch.float16, False, marks=pytest.mark.full),
|
||
# dim=-1, non-aligned M + large-N tiled path
|
||
pytest.param((33, 32768), -1, torch.float16, False, marks=pytest.mark.full),
|
||
# dim=0 (reduce along first dim — different M/N split)
|
||
pytest.param((256, 32), 0, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((256, 32), 0, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((256, 32), 0, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=1 (middle dim for 3D)
|
||
pytest.param((2, 256, 16), 1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((2, 256, 16), 1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((2, 256, 16), 1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
class SoftmaxTest(_SoftmaxTestWorkload, TestBase):
|
||
def ref_program(self, x: torch.Tensor) -> torch.Tensor:
|
||
return F.softmax(x.float(), dim=self.dim).to(x.dtype)
|
||
|
||
def __init__(self, shape: tuple, dtype: torch.dtype, dim: int = -1):
|
||
super().__init__(shape, dtype)
|
||
self.dim = dim
|
||
|
||
|
||
@SoftmaxFixture
|
||
def test_softmax_op(shape: tuple, dim: int, dtype: torch.dtype, tune: bool) -> None:
|
||
test = SoftmaxTest(shape, dtype, dim=dim)
|
||
op = SoftmaxFwdOp(dim=dim, tune=tune)
|
||
atol, rtol = _get_tolerances(dtype)
|
||
test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol)
|
||
|
||
|
||
# Softmax — non-contiguous input (spec interface)
|
||
|
||
|
||
class SoftmaxNonContigFixture(FixtureBase):
|
||
PARAMS = [
|
||
(
|
||
"shape, dtype",
|
||
[
|
||
pytest.param((32, 256), torch.float32, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), torch.float16, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), torch.bfloat16, marks=pytest.mark.smoke),
|
||
pytest.param((32, 300), torch.float32, marks=pytest.mark.full),
|
||
pytest.param((32, 300), torch.float16, marks=pytest.mark.full),
|
||
pytest.param((32, 300), torch.bfloat16, marks=pytest.mark.full),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
@SoftmaxNonContigFixture
|
||
def test_softmax_non_contiguous(shape: tuple, dtype: torch.dtype) -> None:
|
||
"""Test softmax with non-contiguous input (sliced tensor)."""
|
||
m, n = shape
|
||
x_full = torch.randn(m, n * 2, dtype=dtype, device="cuda")
|
||
x = x_full[:, :n] # non-contiguous slice
|
||
|
||
op = SoftmaxFwdOp(dim=-1)
|
||
|
||
y_ref = F.softmax(x.float().contiguous(), dim=-1).to(dtype)
|
||
y = op(x)
|
||
atol, rtol = _get_tolerances(dtype)
|
||
assert torch.allclose(y, y_ref, atol=atol, rtol=rtol), (
|
||
f"Non-contiguous softmax failed, max err: {(y - y_ref).abs().max()}"
|
||
)
|
||
|
||
|
||
# Softmax — 1D input (spec interface)
|
||
|
||
|
||
class Softmax1DFixture(FixtureBase):
|
||
PARAMS = [
|
||
(
|
||
"n, dtype",
|
||
[
|
||
pytest.param(256, torch.float32, marks=pytest.mark.smoke),
|
||
pytest.param(256, torch.float16, marks=pytest.mark.smoke),
|
||
pytest.param(256, torch.bfloat16, marks=pytest.mark.smoke),
|
||
pytest.param(300, torch.float32, marks=pytest.mark.full),
|
||
pytest.param(300, torch.float16, marks=pytest.mark.full),
|
||
pytest.param(300, torch.bfloat16, marks=pytest.mark.full),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
@Softmax1DFixture
|
||
def test_softmax_1d(n: int, dtype: torch.dtype) -> None:
|
||
"""Test softmax with 1D input (single row)."""
|
||
x = torch.randn(n, dtype=dtype, device="cuda")
|
||
op = SoftmaxFwdOp(dim=-1)
|
||
|
||
y_ref = F.softmax(x.float(), dim=-1).to(dtype)
|
||
y = op(x)
|
||
atol, rtol = _get_tolerances(dtype)
|
||
assert torch.allclose(y, y_ref, atol=atol, rtol=rtol), (
|
||
f"1D softmax failed, max err: {(y - y_ref).abs().max()}"
|
||
)
|
||
|
||
|
||
# LogSoftmax — spec-conformant interface (shape, dim, dtype)
|
||
|
||
|
||
class LogSoftmaxFixture(FixtureBase):
|
||
PARAMS = [
|
||
(
|
||
"shape, dim, dtype, tune",
|
||
[
|
||
# Smoke: 2D, dim=-1, fp32, pow2
|
||
pytest.param((32, 256), -1, torch.float32, False, marks=[pytest.mark.smoke, pytest.mark.packaging]),
|
||
pytest.param((32, 256), -1, torch.float16, False, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), -1, torch.bfloat16, False, marks=pytest.mark.smoke),
|
||
# tune=True regression: kernel must be built before autotune runs
|
||
pytest.param((32, 256), -1, torch.float16, True, marks=pytest.mark.full),
|
||
# dim=-1 (default path): dtypes x pow2/non-pow2
|
||
pytest.param((32, 300), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((32, 300), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((32, 300), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, tail-M
|
||
pytest.param((33, 256), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((33, 256), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((33, 256), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, 3D input
|
||
pytest.param((2, 16, 256), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((2, 16, 256), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((2, 16, 256), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, 4D input
|
||
pytest.param((2, 4, 8, 256), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((2, 4, 8, 256), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((2, 4, 8, 256), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, large-N (triggers N-tiling path)
|
||
pytest.param((4, 32768), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((4, 32768), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, M×N both non-aligned (single-tile path)
|
||
pytest.param((33, 300), -1, torch.float32, False, marks=pytest.mark.full),
|
||
# dim=-1, M×N both non-aligned (multi-tile, masked loads)
|
||
pytest.param((33, 33000), -1, torch.float16, False, marks=pytest.mark.full),
|
||
# dim=-1, non-aligned M + large-N tiled path
|
||
pytest.param((33, 32768), -1, torch.float16, False, marks=pytest.mark.full),
|
||
# dim=0 (reduce along first dim)
|
||
pytest.param((256, 32), 0, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((256, 32), 0, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((256, 32), 0, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=1 (middle dim for 3D)
|
||
pytest.param((2, 256, 16), 1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((2, 256, 16), 1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((2, 256, 16), 1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
class LogSoftmaxTest(_LogSoftmaxTestWorkload, TestBase):
|
||
def ref_program(self, x: torch.Tensor) -> torch.Tensor:
|
||
return F.log_softmax(x.float(), dim=self.dim).to(x.dtype)
|
||
|
||
def __init__(self, shape: tuple, dtype: torch.dtype, dim: int = -1):
|
||
super().__init__(shape, dtype)
|
||
self.dim = dim
|
||
|
||
|
||
@LogSoftmaxFixture
|
||
def test_log_softmax_op(shape: tuple, dim: int, dtype: torch.dtype, tune: bool) -> None:
|
||
test = LogSoftmaxTest(shape, dtype, dim=dim)
|
||
op = LogSoftmaxFwdOp(dim=dim, tune=tune)
|
||
atol, rtol = _get_tolerances(dtype)
|
||
test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol)
|
||
|
||
|
||
# LogSumExp — spec-conformant interface (shape, dim, keepdim, dtype)
|
||
|
||
|
||
class LogSumExpFixture(FixtureBase):
|
||
PARAMS = [
|
||
(
|
||
"shape, dim, dtype, tune",
|
||
[
|
||
# Smoke: 2D, dim=-1, fp32, pow2
|
||
pytest.param((32, 256), -1, torch.float32, False, marks=[pytest.mark.smoke, pytest.mark.packaging]),
|
||
pytest.param((32, 256), -1, torch.float16, False, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), -1, torch.bfloat16, False, marks=pytest.mark.smoke),
|
||
# tune=True regression: kernel must be built before autotune runs
|
||
pytest.param((32, 256), -1, torch.float16, True, marks=pytest.mark.full),
|
||
# dim=-1: dtypes x pow2/non-pow2
|
||
pytest.param((32, 300), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((32, 300), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((32, 300), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, tail-M
|
||
pytest.param((33, 256), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((33, 256), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((33, 256), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, 3D input
|
||
pytest.param((2, 16, 256), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((2, 16, 256), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((2, 16, 256), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, 4D input
|
||
pytest.param((2, 4, 8, 256), -1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((2, 4, 8, 256), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((2, 4, 8, 256), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, large-N (triggers N-tiling path)
|
||
pytest.param((4, 32768), -1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((4, 32768), -1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=-1, M×N both non-aligned (single-tile path)
|
||
pytest.param((33, 300), -1, torch.float32, False, marks=pytest.mark.full),
|
||
# dim=-1, M×N both non-aligned (multi-tile, masked loads)
|
||
pytest.param((33, 33000), -1, torch.float16, False, marks=pytest.mark.full),
|
||
# dim=-1, non-aligned M + large-N tiled path
|
||
pytest.param((33, 32768), -1, torch.float16, False, marks=pytest.mark.full),
|
||
# dim=0
|
||
pytest.param((256, 32), 0, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((256, 32), 0, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((256, 32), 0, torch.bfloat16, False, marks=pytest.mark.full),
|
||
# dim=1 (middle dim for 3D)
|
||
pytest.param((2, 256, 16), 1, torch.float32, False, marks=pytest.mark.full),
|
||
pytest.param((2, 256, 16), 1, torch.float16, False, marks=pytest.mark.full),
|
||
pytest.param((2, 256, 16), 1, torch.bfloat16, False, marks=pytest.mark.full),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
class LogSumExpTest(_LogSumExpTestWorkload, TestBase):
|
||
def ref_program(self, x: torch.Tensor) -> torch.Tensor:
|
||
return torch.logsumexp(x.float(), dim=self.dim).to(x.dtype)
|
||
|
||
def __init__(self, shape: tuple, dtype: torch.dtype, dim: int = -1):
|
||
super().__init__(shape, dtype)
|
||
self.dim = dim
|
||
|
||
|
||
@LogSumExpFixture
|
||
def test_logsumexp_op(shape: tuple, dim: int, dtype: torch.dtype, tune: bool) -> None:
|
||
test = LogSumExpTest(shape, dtype, dim=dim)
|
||
op = LogSumExpFwdOp(dtype=dtype, dim=dim, tune=tune)
|
||
atol, rtol = _get_tolerances(dtype)
|
||
test.check(op, *test.gen_inputs(), atol=atol, rtol=rtol)
|
||
|
||
|
||
# LogSumExp — keepdim=True (exercises _reshape_output keepdim path)
|
||
|
||
|
||
class LogSumExpKeepdimFixture(FixtureBase):
|
||
PARAMS = [
|
||
(
|
||
"shape, dim, dtype",
|
||
[
|
||
# dim=-1 (last dim, no transpose)
|
||
pytest.param((32, 256), -1, torch.float32, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), -1, torch.float16, marks=pytest.mark.smoke),
|
||
pytest.param((2, 16, 256), -1, torch.float32, marks=pytest.mark.full),
|
||
# dim=0 (non-last dim, exercises transpose + keepdim)
|
||
pytest.param((256, 32), 0, torch.float32, marks=pytest.mark.full),
|
||
pytest.param((256, 32), 0, torch.float16, marks=pytest.mark.full),
|
||
# dim=1 (middle dim, 3D)
|
||
pytest.param((2, 256, 16), 1, torch.float32, marks=pytest.mark.full),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
@LogSumExpKeepdimFixture
|
||
def test_logsumexp_keepdim(shape: tuple, dim: int, dtype: torch.dtype) -> None:
|
||
"""Test logsumexp with keepdim=True — output retains reduced dim as size 1."""
|
||
x = torch.randn(*shape, dtype=dtype, device="cuda")
|
||
op = LogSumExpFwdOp(dtype=dtype, dim=dim, keepdim=True)
|
||
|
||
y_ref = torch.logsumexp(x.float(), dim=dim, keepdim=True).to(dtype)
|
||
y = op(x)
|
||
assert y.shape == y_ref.shape, f"Shape mismatch: {y.shape} vs {y_ref.shape}"
|
||
atol, rtol = _get_tolerances(dtype)
|
||
assert torch.allclose(y, y_ref, atol=atol, rtol=rtol), (
|
||
f"keepdim logsumexp failed, max err: {(y - y_ref).abs().max()}"
|
||
)
|
||
|
||
|
||
# Non-contiguous input tests (spec interface)
|
||
|
||
|
||
class LogSoftmaxNonContigFixture(FixtureBase):
|
||
PARAMS = [
|
||
(
|
||
"shape, dtype",
|
||
[
|
||
pytest.param((32, 256), torch.float32, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), torch.float16, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), torch.bfloat16, marks=pytest.mark.smoke),
|
||
pytest.param((32, 300), torch.float32, marks=pytest.mark.full),
|
||
pytest.param((32, 300), torch.float16, marks=pytest.mark.full),
|
||
pytest.param((32, 300), torch.bfloat16, marks=pytest.mark.full),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
@LogSoftmaxNonContigFixture
|
||
def test_log_softmax_non_contiguous(shape: tuple, dtype: torch.dtype) -> None:
|
||
"""Test log_softmax with non-contiguous input (sliced tensor)."""
|
||
m, n = shape
|
||
x_full = torch.randn(m, n * 2, dtype=dtype, device="cuda")
|
||
x = x_full[:, :n]
|
||
|
||
op = LogSoftmaxFwdOp(dim=-1)
|
||
|
||
y_ref = F.log_softmax(x.float().contiguous(), dim=-1).to(dtype)
|
||
y = op(x)
|
||
atol, rtol = _get_tolerances(dtype)
|
||
assert torch.allclose(y, y_ref, atol=atol, rtol=rtol), (
|
||
f"Non-contiguous log_softmax failed, max err: {(y - y_ref).abs().max()}"
|
||
)
|
||
|
||
|
||
class LogSumExpNonContigFixture(FixtureBase):
|
||
PARAMS = [
|
||
(
|
||
"shape, dtype",
|
||
[
|
||
pytest.param((32, 256), torch.float32, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), torch.float16, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), torch.bfloat16, marks=pytest.mark.smoke),
|
||
pytest.param((32, 300), torch.float32, marks=pytest.mark.full),
|
||
pytest.param((32, 300), torch.float16, marks=pytest.mark.full),
|
||
pytest.param((32, 300), torch.bfloat16, marks=pytest.mark.full),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
@LogSumExpNonContigFixture
|
||
def test_logsumexp_non_contiguous(shape: tuple, dtype: torch.dtype) -> None:
|
||
"""Test logsumexp with non-contiguous input."""
|
||
m, n = shape
|
||
x_full = torch.randn(m, n * 2, dtype=dtype, device="cuda")
|
||
x = x_full[:, :n]
|
||
|
||
op = LogSumExpFwdOp(dtype=dtype, dim=-1)
|
||
|
||
y_ref = torch.logsumexp(x.float().contiguous(), dim=-1).to(dtype)
|
||
y = op(x)
|
||
atol, rtol = _get_tolerances(dtype)
|
||
assert torch.allclose(y, y_ref, atol=atol, rtol=rtol), (
|
||
f"Non-contiguous logsumexp failed, max err: {(y - y_ref).abs().max()}"
|
||
)
|
||
|
||
|
||
# 1D input tests (spec interface)
|
||
|
||
|
||
class LogSoftmax1DFixture(FixtureBase):
|
||
PARAMS = [
|
||
(
|
||
"n, dtype",
|
||
[
|
||
pytest.param(256, torch.float32, marks=pytest.mark.smoke),
|
||
pytest.param(256, torch.float16, marks=pytest.mark.smoke),
|
||
pytest.param(256, torch.bfloat16, marks=pytest.mark.smoke),
|
||
pytest.param(300, torch.float32, marks=pytest.mark.full),
|
||
pytest.param(300, torch.float16, marks=pytest.mark.full),
|
||
pytest.param(300, torch.bfloat16, marks=pytest.mark.full),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
@LogSoftmax1DFixture
|
||
def test_log_softmax_1d(n: int, dtype: torch.dtype) -> None:
|
||
"""Test log_softmax with 1D input."""
|
||
x = torch.randn(n, dtype=dtype, device="cuda")
|
||
op = LogSoftmaxFwdOp(dim=-1)
|
||
|
||
y_ref = F.log_softmax(x.float(), dim=-1).to(dtype)
|
||
y = op(x)
|
||
atol, rtol = _get_tolerances(dtype)
|
||
assert torch.allclose(y, y_ref, atol=atol, rtol=rtol), (
|
||
f"1D log_softmax failed, max err: {(y - y_ref).abs().max()}"
|
||
)
|
||
|
||
|
||
class LogSumExp1DFixture(FixtureBase):
|
||
PARAMS = [
|
||
(
|
||
"n, dtype",
|
||
[
|
||
pytest.param(256, torch.float32, marks=pytest.mark.smoke),
|
||
pytest.param(256, torch.float16, marks=pytest.mark.smoke),
|
||
pytest.param(256, torch.bfloat16, marks=pytest.mark.smoke),
|
||
pytest.param(300, torch.float32, marks=pytest.mark.full),
|
||
pytest.param(300, torch.float16, marks=pytest.mark.full),
|
||
pytest.param(300, torch.bfloat16, marks=pytest.mark.full),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
@LogSumExp1DFixture
|
||
def test_logsumexp_1d(n: int, dtype: torch.dtype) -> None:
|
||
"""Test logsumexp with 1D input -- output should be a scalar."""
|
||
x = torch.randn(n, dtype=dtype, device="cuda")
|
||
op = LogSumExpFwdOp(dtype=dtype, dim=-1)
|
||
|
||
y_ref = torch.logsumexp(x.float(), dim=-1).to(dtype)
|
||
y = op(x)
|
||
atol, rtol = _get_tolerances(dtype)
|
||
assert y.shape == y_ref.shape, f"Shape mismatch: {y.shape} vs {y_ref.shape}"
|
||
assert torch.allclose(y, y_ref, atol=atol, rtol=rtol), (
|
||
f"1D logsumexp failed, max err: {(y - y_ref).abs().max()}"
|
||
)
|
||
|
||
|
||
# Multi-dim guard tests: SoftmaxFwdOp and LogSoftmaxFwdOp must reject
|
||
# list/tuple dims eagerly (before kernel build/execute).
|
||
|
||
|
||
@pytest.mark.smoke
|
||
def test_softmax_rejects_multidim_before_kernel() -> None:
|
||
"""SoftmaxFwdOp must raise ValueError for list dim before touching the kernel."""
|
||
x = torch.randn(4, 8, device="cuda", dtype=torch.float32)
|
||
op = SoftmaxFwdOp(dim=[-1, 0])
|
||
with pytest.raises(ValueError, match="does not support multi-dim"):
|
||
op(x)
|
||
# Verify no kernel was built (cache must remain empty).
|
||
assert len(op._kernel_cache) == 0
|
||
|
||
|
||
@pytest.mark.smoke
|
||
def test_log_softmax_rejects_multidim_before_kernel() -> None:
|
||
"""LogSoftmaxFwdOp must raise ValueError for list dim before touching the kernel."""
|
||
x = torch.randn(4, 8, device="cuda", dtype=torch.float32)
|
||
op = LogSoftmaxFwdOp(dim=[-1, 0])
|
||
with pytest.raises(ValueError, match="does not support multi-dim"):
|
||
op(x)
|
||
assert len(op._kernel_cache) == 0
|
||
|
||
|
||
@pytest.mark.smoke
|
||
def test_logsumexp_accepts_multidim() -> None:
|
||
"""LogSumExpFwdOp must accept list dim without error (multi-dim is supported)."""
|
||
x = torch.randn(4, 8, device="cuda", dtype=torch.float32)
|
||
op = LogSumExpFwdOp(dtype=torch.float32, dim=[0, 1])
|
||
y = op(x)
|
||
y_ref = torch.logsumexp(x.float(), dim=[0, 1])
|
||
assert torch.allclose(y, y_ref, atol=1e-5, rtol=1e-5)
|
||
|
||
|
||
class SoftmaxImplicitDimFixture(FixtureBase):
|
||
# Smoke covers each ndim branch (1D, 2D, 3D) and each dtype at least once.
|
||
PARAMS = [
|
||
(
|
||
"shape, dtype",
|
||
[
|
||
pytest.param((256,), torch.float32, marks=pytest.mark.smoke),
|
||
pytest.param((32, 256), torch.float16, marks=pytest.mark.smoke),
|
||
pytest.param((4, 16, 32), torch.bfloat16, marks=pytest.mark.smoke),
|
||
],
|
||
),
|
||
]
|
||
|
||
|
||
def _expected_implicit_dim(ndim: int) -> int:
|
||
return 0 if ndim in (0, 1, 3) else 1
|
||
|
||
|
||
@SoftmaxImplicitDimFixture
|
||
def test_softmax_dim_none_implicit_axis(shape: tuple, dtype: torch.dtype) -> None:
|
||
"""SoftmaxFwdOp(dim=None) must match F.softmax(x, dim=None) and warn."""
|
||
import warnings as _warnings
|
||
x = torch.randn(*shape, dtype=dtype, device="cuda")
|
||
op = SoftmaxFwdOp(dim=None)
|
||
|
||
with _warnings.catch_warnings(record=True) as caught:
|
||
_warnings.simplefilter("always")
|
||
y = op(x)
|
||
assert any(
|
||
issubclass(w.category, UserWarning) and "Implicit dimension choice" in str(w.message)
|
||
for w in caught
|
||
), f"Expected implicit-dim UserWarning, got {[str(w.message) for w in caught]}"
|
||
|
||
with _warnings.catch_warnings():
|
||
_warnings.simplefilter("ignore", UserWarning)
|
||
y_ref = F.softmax(x.float(), dim=None).to(dtype)
|
||
|
||
atol, rtol = _get_tolerances(dtype)
|
||
assert torch.allclose(y, y_ref, atol=atol, rtol=rtol), (
|
||
f"dim=None softmax (shape={shape}, dtype={dtype}) failed, "
|
||
f"max err: {(y - y_ref).abs().max()}"
|
||
)
|
||
|
||
|
||
@SoftmaxImplicitDimFixture
|
||
def test_log_softmax_dim_none_implicit_axis(shape: tuple, dtype: torch.dtype) -> None:
|
||
"""LogSoftmaxFwdOp(dim=None) must match F.log_softmax(x, dim=None) and warn."""
|
||
import warnings as _warnings
|
||
x = torch.randn(*shape, dtype=dtype, device="cuda")
|
||
op = LogSoftmaxFwdOp(dim=None)
|
||
|
||
with _warnings.catch_warnings(record=True) as caught:
|
||
_warnings.simplefilter("always")
|
||
y = op(x)
|
||
assert any(
|
||
issubclass(w.category, UserWarning) and "Implicit dimension choice" in str(w.message)
|
||
for w in caught
|
||
), f"Expected implicit-dim UserWarning, got {[str(w.message) for w in caught]}"
|
||
|
||
with _warnings.catch_warnings():
|
||
_warnings.simplefilter("ignore", UserWarning)
|
||
y_ref = F.log_softmax(x.float(), dim=None).to(dtype)
|
||
|
||
atol, rtol = _get_tolerances(dtype)
|
||
assert torch.allclose(y, y_ref, atol=atol, rtol=rtol), (
|
||
f"dim=None log_softmax (shape={shape}, dtype={dtype}) failed, "
|
||
f"max err: {(y - y_ref).abs().max()}"
|
||
)
|
||
|
||
|
||
@pytest.mark.smoke
|
||
def test_softmax_dim_none_reused_across_ranks() -> None:
|
||
"""SoftmaxFwdOp(dim=None) must re-resolve per call across input ranks."""
|
||
import warnings as _warnings
|
||
op = SoftmaxFwdOp(dim=None)
|
||
|
||
x1 = torch.randn(4, dtype=torch.float32, device="cuda")
|
||
x2 = torch.randn(2, 4, dtype=torch.float32, device="cuda")
|
||
x3 = torch.randn(4, 3, 5, dtype=torch.float32, device="cuda")
|
||
|
||
with _warnings.catch_warnings():
|
||
_warnings.simplefilter("ignore", UserWarning)
|
||
y1 = op(x1)
|
||
y2 = op(x2)
|
||
y3 = op(x3)
|
||
y1_ref = F.softmax(x1.float(), dim=None)
|
||
y2_ref = F.softmax(x2.float(), dim=None)
|
||
y3_ref = F.softmax(x3.float(), dim=None)
|
||
|
||
assert op.dim is None, f"op.dim was mutated to {op.dim!r}; expected None"
|
||
|
||
atol, rtol = _get_tolerances(torch.float32)
|
||
assert torch.allclose(y1, y1_ref, atol=atol, rtol=rtol)
|
||
assert torch.allclose(y2, y2_ref, atol=atol, rtol=rtol)
|
||
assert torch.allclose(y3, y3_ref, atol=atol, rtol=rtol)
|
||
|
||
|
||
@pytest.mark.smoke
|
||
def test_log_softmax_dim_none_reused_across_ranks() -> None:
|
||
"""LogSoftmaxFwdOp(dim=None) must re-resolve per call (no self.dim mutation)."""
|
||
import warnings as _warnings
|
||
op = LogSoftmaxFwdOp(dim=None)
|
||
|
||
x1 = torch.randn(4, dtype=torch.float32, device="cuda")
|
||
x2 = torch.randn(2, 4, dtype=torch.float32, device="cuda")
|
||
x3 = torch.randn(4, 3, 5, dtype=torch.float32, device="cuda")
|
||
|
||
with _warnings.catch_warnings():
|
||
_warnings.simplefilter("ignore", UserWarning)
|
||
y1 = op(x1)
|
||
y2 = op(x2)
|
||
y3 = op(x3)
|
||
y1_ref = F.log_softmax(x1.float(), dim=None)
|
||
y2_ref = F.log_softmax(x2.float(), dim=None)
|
||
y3_ref = F.log_softmax(x3.float(), dim=None)
|
||
|
||
assert op.dim is None, f"op.dim was mutated to {op.dim!r}; expected None"
|
||
|
||
atol, rtol = _get_tolerances(torch.float32)
|
||
assert torch.allclose(y1, y1_ref, atol=atol, rtol=rtol)
|
||
assert torch.allclose(y2, y2_ref, atol=atol, rtol=rtol)
|
||
assert torch.allclose(y3, y3_ref, atol=atol, rtol=rtol)
|
||
|
||
|
||
# Roofline regression: LogSoftmax FLOPs must equal 5 * M * N (not 6 * M * N).
|
||
# Direct construction — no manifest-string indirection.
|
||
|
||
|
||
@pytest.mark.smoke
|
||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||
def test_log_softmax_eval_roofline_flops_5mn() -> None:
|
||
"""LogSoftmaxFwdOp.eval_roofline() must report flops == 5 * M * N."""
|
||
M, N = 64, 256
|
||
dtype = torch.float16
|
||
op = LogSoftmaxFwdOp(dim=-1)
|
||
x = torch.randn(M, N, dtype=dtype, device="cuda")
|
||
op(x) # bind dynamic shape
|
||
flops, mem_bytes = op.eval_roofline()
|
||
elem_bytes = dtype.itemsize
|
||
assert flops == 5 * M * N, f"LogSoftmax flops {flops} != 5 * M * N = {5 * M * N}"
|
||
assert mem_bytes == 2 * M * N * elem_bytes, (
|
||
f"LogSoftmax bytes {mem_bytes} != 2 * M * N * elem_bytes = "
|
||
f"{2 * M * N * elem_bytes}"
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
pytest.main([__file__, "-vvs"])
|