Spaces:
Running
Running
File size: 3,337 Bytes
e313085 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """Exercise the same dispatch guard on CPU tensors without compiling candidate kernels."""
import pytest
torch = pytest.importorskip("torch")
from kernel_env.server.worker import HostOperationError, _evaluate_candidate, kernel_only_mode
@pytest.mark.parametrize("operation", [
lambda x: x + 1,
lambda x: x - 1,
lambda x: x * 2,
lambda x: x / 2,
lambda x: 1.0 / x,
lambda x: x @ x.t(),
lambda x: x.sum(),
lambda x: torch.relu(x),
lambda x: x.add_(1),
lambda x: torch.add(x, 1, out=torch.empty_like(x)),
lambda x: x[0, 0].item(),
])
def test_host_tensor_compute_is_rejected(operation):
x = torch.ones(2, 2)
with pytest.raises(HostOperationError, match="use a custom kernel"):
with kernel_only_mode("cpu"):
operation(x)
def test_allocation_layout_copies_and_scalar_grid_math_are_allowed():
x = torch.ones(2, 3)
with kernel_only_mode("cpu"):
y = torch.empty_like(x)
y.copy_(x)
view = y.t().contiguous().view(6)
converted = view.to(torch.float64)
zeros = torch.zeros_like(x)
grid = (x.numel() + 1023) // 1024
assert grid == 1 and torch.equal(converted, torch.ones(6, dtype=torch.float64))
assert torch.equal(zeros, torch.zeros(2, 3))
def test_cpu_autotuner_statistics_are_not_blocked_by_cuda_guard():
with kernel_only_mode("cuda"):
assert torch.quantile(torch.tensor([1.0, 2.0, 3.0]), 0.5) == 2
@pytest.mark.parametrize("copy", [lambda x: x.to("cpu"), lambda x: torch.empty(2).copy_(x)])
def test_off_device_copy_cannot_bypass_the_guard(copy):
# Meta exercises device routing without allocating a CUDA tensor or running any kernel.
x = torch.empty(2, device="meta")
with pytest.raises(HostOperationError):
with kernel_only_mode("meta"):
copy(x)
def test_audited_host_arithmetic_fails_even_though_ast_accepts_it():
from kernel_env.server.static_checks import check_source
from kernels import with_forward
assert check_source(with_forward("return x * (1.0 / x)")) == []
x = torch.ones(2, 2)
with pytest.raises(HostOperationError):
with kernel_only_mode("cpu"):
x * (1.0 / x)
@pytest.mark.parametrize("compute_on_call", [1, 3])
def test_worker_rejects_host_compute_in_correctness_or_benchmark(tmp_path, monkeypatch, compute_on_call):
# This is a trusted, hand-written test module, not model-generated source.
path = tmp_path / "candidate.py"
path.write_text(f"""import torch
class ModelNew(torch.nn.Module):
def __init__(self):
super().__init__()
self.calls = 0
def forward(self, x):
self.calls += 1
if self.calls >= {compute_on_call}:
return x * 0
return torch.zeros_like(x)
""")
monkeypatch.setattr(torch.cuda, "synchronize", lambda: None)
inputs = lambda seed: [torch.ones(2, 2)]
def bench(model, args):
model(*args)
return 1.0
result = _evaluate_candidate(
{"code_path": str(path), "seeds": [1, 2], "atol": 0.01, "rtol": 0.01},
torch, torch.device("cpu"), lambda x: torch.zeros_like(x), [], inputs, bench,
{"t_ref_ms": 1.0, "gpu": "cpu-test"},
)
assert result["stage"] == "format" and result["correct"] is False
assert "use a custom kernel" in result["feedback"]
|