152 lines
4.8 KiB
Python
152 lines
4.8 KiB
Python
import json
|
|
import os
|
|
import platform
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
@pytest.mark.skipif(platform.system() != "Darwin", reason="macOS backend default")
|
|
def test_dense_attention_defaults_to_sdpa_on_macos():
|
|
env = os.environ.copy()
|
|
env.pop("ATTN_BACKEND", None)
|
|
completed = subprocess.run(
|
|
[
|
|
sys.executable,
|
|
"-c",
|
|
"from trellis2.modules.attention import config; print('BACKEND=' + config.BACKEND)",
|
|
],
|
|
cwd=ROOT,
|
|
env=env,
|
|
text=True,
|
|
capture_output=True,
|
|
check=True,
|
|
)
|
|
assert "BACKEND=sdpa" in completed.stdout
|
|
|
|
|
|
def test_disable_metal_selects_controlled_pytorch_fallback():
|
|
env = os.environ.copy()
|
|
env.update(
|
|
{
|
|
"TRELLIS_DISABLE_METAL": "1",
|
|
"PYTHONPATH": str(ROOT),
|
|
"SPARSE_CONV_BACKEND": "pytorch",
|
|
"SPARSE_ATTN_BACKEND": "sdpa",
|
|
}
|
|
)
|
|
code = """
|
|
import json
|
|
from trellis2.modules.sparse import config
|
|
from trellis2.backends import backend_report
|
|
print('REPORT=' + json.dumps({
|
|
'conv': config.CONV,
|
|
'attention': config.ATTN,
|
|
'backend': backend_report(),
|
|
}, sort_keys=True))
|
|
"""
|
|
completed = subprocess.run(
|
|
[sys.executable, "-c", code],
|
|
cwd=ROOT,
|
|
env=env,
|
|
text=True,
|
|
capture_output=True,
|
|
check=True,
|
|
)
|
|
line = next(line for line in completed.stdout.splitlines() if line.startswith("REPORT="))
|
|
report = json.loads(line.removeprefix("REPORT="))
|
|
assert report["conv"] == "pytorch"
|
|
assert report["attention"] == "sdpa"
|
|
assert report["backend"]["metal_disabled"] is True
|
|
assert report["backend"]["rasterizer"] is None
|
|
assert report["backend"]["mesh"] is None
|
|
|
|
|
|
@pytest.mark.skipif(platform.system() != "Darwin", reason="Metal probe requires macOS")
|
|
def test_metal_sparse_attention_has_its_own_numerical_probe():
|
|
# Pipeline imports performed by other test modules may already have made a
|
|
# controlled fallback choice. Exercise auto-detection in a fresh process,
|
|
# which also mirrors a real CLI invocation.
|
|
code = r"""
|
|
import json
|
|
import torch
|
|
if not torch.backends.mps.is_available():
|
|
raise SystemExit(77)
|
|
from trellis2.modules.sparse import config
|
|
from trellis2.backends import probe_metal_backends
|
|
print('REPORT=' + json.dumps({
|
|
'probe': config.probe_flex_gemm_sparse_attention_on_mps(),
|
|
'conv': config.CONV,
|
|
'attention': config.ATTN,
|
|
'metal': probe_metal_backends(),
|
|
}, sort_keys=True))
|
|
"""
|
|
completed = subprocess.run(
|
|
[sys.executable, "-c", code],
|
|
cwd=ROOT,
|
|
text=True,
|
|
capture_output=True,
|
|
check=False,
|
|
)
|
|
if completed.returncode == 77:
|
|
pytest.skip("MPS is unavailable")
|
|
assert completed.returncode == 0, completed.stderr
|
|
line = next(line for line in completed.stdout.splitlines() if line.startswith("REPORT="))
|
|
report = json.loads(line.removeprefix("REPORT="))
|
|
assert report["probe"]
|
|
assert report["conv"] == "flex_gemm"
|
|
assert report["attention"] == "flex_gemm_sparse_attn"
|
|
assert report["metal"]["rasterizer"]["ok"], report["metal"]["rasterizer"]
|
|
assert report["metal"]["mesh_bvh"]["ok"], report["metal"]["mesh_bvh"]
|
|
|
|
|
|
@pytest.mark.skipif(platform.system() != "Darwin", reason="Metal parity requires macOS")
|
|
def test_pure_pytorch_sparse_conv_matches_flex_gemm_on_mps():
|
|
code = r"""
|
|
import torch
|
|
from trellis2.modules.sparse import SparseTensor, config
|
|
from trellis2.modules.sparse.conv.conv import SparseConv3d
|
|
|
|
if not torch.backends.mps.is_available():
|
|
raise SystemExit(77)
|
|
coords = torch.tensor(
|
|
[[0, x, y, z] for x in range(3) for y in range(3) for z in range(3)],
|
|
dtype=torch.int32,
|
|
device='mps',
|
|
)
|
|
features = torch.randn(27, 4, dtype=torch.float16, device='mps')
|
|
sparse = SparseTensor(
|
|
feats=features,
|
|
coords=coords,
|
|
shape=torch.Size([1, 4]),
|
|
spatial_shape=[3, 3, 3],
|
|
)
|
|
config.set_conv_backend('flex_gemm')
|
|
metal = SparseConv3d(4, 5, 3, bias=True).half().to('mps')
|
|
metal_out = metal(sparse).feats
|
|
|
|
config.set_conv_backend('pytorch')
|
|
fallback = SparseConv3d(4, 5, 3, bias=True).half().to('mps')
|
|
fallback.weight.data.copy_(metal.weight.data)
|
|
fallback.bias.data.copy_(metal.bias.data)
|
|
fallback_out = fallback(sparse).feats
|
|
torch.mps.synchronize()
|
|
print('PARITY=' + str(float((metal_out - fallback_out).abs().max().item())))
|
|
"""
|
|
completed = subprocess.run(
|
|
[sys.executable, "-c", code],
|
|
cwd=ROOT,
|
|
text=True,
|
|
capture_output=True,
|
|
check=False,
|
|
)
|
|
if completed.returncode == 77:
|
|
pytest.skip("MPS is unavailable")
|
|
assert completed.returncode == 0, completed.stderr
|
|
parity_line = next(line for line in completed.stdout.splitlines() if line.startswith("PARITY="))
|
|
assert float(parity_line.removeprefix("PARITY=")) <= 2e-2
|