yyda
/
Intro-ops
forked from ccf-ai-infra/Intro-ops
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add more shared test cases 

Co-authored-by: wawahejun <hejunlbbc@gmail.com>
This commit is contained in:
yutianyu 2026-05-04 21:15:31 +08:00
parent e9a4c87564
commit a2cda3fda7
9 changed files with 43 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tests/bench/copy.py
View File

@ -27,7 +27,9 @@ def bench_copy(backend: str) -> list[PerformanceResult]:
for case in copy_cases.benchmark_cases():
src = torch.randn(case["shape"], dtype=case["dtype"], device="cuda")
out = torch.empty_like(src)
runtime = cuda_time_ms(lambda: copy(src, backend=backend))
from operator_runtime import prepare_copy
with prepare_copy(out, src, backend=backend) as prepared:
runtime = cuda_time_ms(prepared.run)
torch_ms = cuda_time_ms(lambda: out.copy_(src))
bytes_, flops = _estimate_copy(src)
rows.append(

5
tests/bench/reduce_sum.py
View File

@ -27,7 +27,10 @@ def bench_reduce_sum(backend: str) -> list[PerformanceResult]:
rows: list[PerformanceResult] = []
for case in reduce_sum_cases.benchmark_cases():
src = torch.randn(case["shape"], dtype=case["dtype"], device="cuda")
runtime = cuda_time_ms(lambda: reduce_sum(src, dim=1, backend=backend))
from operator_runtime import prepare_reduce_sum
out = torch.empty(src.shape[0], dtype=src.dtype, device="cuda")
with prepare_reduce_sum(out, src, dim=1, backend=backend) as prepared:
runtime = cuda_time_ms(prepared.run)
torch_ms = cuda_time_ms(lambda: torch.sum(src, dim=1))
bytes_, flops = _estimate_reduce_sum(src)
rows.append(

7
tests/bench/softmax.py
View File

@ -26,8 +26,11 @@ def bench_softmax(backend: str) -> list[PerformanceResult]:
rows: list[PerformanceResult] = []
for case in softmax_cases.benchmark_cases():
src = torch.randn(case["shape"], dtype=case["dtype"], device="cuda")
runtime = cuda_time_ms(lambda: softmax(src, dim=1, backend=backend))
torch_ms = cuda_time_ms(lambda: torch.softmax(src, dim=1))
out = torch.empty_like(src)
from operator_runtime import prepare_softmax
with prepare_softmax(out, src, dim=1, backend=backend) as prepared:
runtime = cuda_time_ms(prepared.run)
torch_ms = cuda_time_ms(lambda: torch.softmax(src, dim=1, out=out))
bytes_, flops = _estimate_softmax(src)
rows.append(
PerformanceResult(

7
tests/bench/vector_add.py
View File

@ -27,8 +27,11 @@ def bench_vector_add(backend: str) -> list[PerformanceResult]:
for case in vector_add_cases.benchmark_cases():
a = torch.randn(case["shape"], dtype=case["dtype"], device="cuda")
b = torch.randn_like(a)
runtime = cuda_time_ms(lambda: vector_add(a, b, backend=backend))
torch_ms = cuda_time_ms(lambda: torch.add(a, b))
out = torch.empty_like(a)
from operator_runtime import prepare_vector_add
with prepare_vector_add(out, a, b, backend=backend) as prepared:
runtime = cuda_time_ms(prepared.run)
torch_ms = cuda_time_ms(lambda: torch.add(a, b, out=out))
bytes_, flops = _estimate_vector_add(a)
rows.append(
PerformanceResult(

9
tests/cases/copy.py
View File

@ -22,5 +22,12 @@ def api_error_cases():
def benchmark_cases():
return [
{"name": "contiguous_1m", "shape": (1 << 20,), "dtype": torch.float16},
{"name": "contiguous_1m_fp16", "shape": (1 << 20,), "dtype": torch.float16},
{"name": "contiguous_4m_fp16", "shape": (4 << 20,), "dtype": torch.float16},
{"name": "contiguous_16m_fp16", "shape": (16 << 20,), "dtype": torch.float16},
{"name": "contiguous_64m_fp16", "shape": (64 << 20,), "dtype": torch.float16},
{"name": "contiguous_1m_fp32", "shape": (1 << 20,), "dtype": torch.float32},
{"name": "contiguous_4m_fp32", "shape": (4 << 20,), "dtype": torch.float32},
{"name": "contiguous_16m_fp32", "shape": (16 << 20,), "dtype": torch.float32},
{"name": "contiguous_64m_fp32", "shape": (64 << 20,), "dtype": torch.float32},
]

3
tests/cases/reduce_sum.py
View File

@ -24,5 +24,8 @@ def api_error_cases():
def benchmark_cases():
return [
{"name": "rowwise_128x1024", "shape": (128, 1024), "dtype": torch.float32},
{"name": "rowwise_1024x1024", "shape": (1024, 1024), "dtype": torch.float32},
{"name": "rowwise_1024x4096", "shape": (1024, 4096), "dtype": torch.float32},
{"name": "rowwise_4096x1024", "shape": (4096, 1024), "dtype": torch.float32},
]

3
tests/cases/softmax.py
View File

@ -24,5 +24,8 @@ def api_error_cases():
def benchmark_cases():
return [
{"name": "rowwise_128x1024", "shape": (128, 1024), "dtype": torch.float32},
{"name": "rowwise_1024x1024", "shape": (1024, 1024), "dtype": torch.float32},
{"name": "rowwise_1024x4096", "shape": (1024, 4096), "dtype": torch.float32},
{"name": "rowwise_4096x1024", "shape": (4096, 1024), "dtype": torch.float32},
]

9
tests/cases/vector_add.py
View File

@ -22,5 +22,12 @@ def api_error_cases():
def benchmark_cases():
return [
{"name": "contiguous_1m", "shape": (1 << 20,), "dtype": torch.float16},
{"name": "contiguous_1m_fp16", "shape": (1 << 20,), "dtype": torch.float16},
{"name": "contiguous_4m_fp16", "shape": (4 << 20,), "dtype": torch.float16},
{"name": "contiguous_16m_fp16", "shape": (16 << 20,), "dtype": torch.float16},
{"name": "contiguous_64m_fp16", "shape": (64 << 20,), "dtype": torch.float16},
{"name": "contiguous_1m_fp32", "shape": (1 << 20,), "dtype": torch.float32},
{"name": "contiguous_4m_fp32", "shape": (4 << 20,), "dtype": torch.float32},
{"name": "contiguous_16m_fp32", "shape": (16 << 20,), "dtype": torch.float32},
{"name": "contiguous_64m_fp32", "shape": (64 << 20,), "dtype": torch.float32},
]

8
tests/test_tilelang_copy_templates.py
View File

@ -16,8 +16,8 @@ pytestmark = pytest.mark.skipif(
@pytest.mark.parametrize(
("copy_fn_name", "copy_out_fn_name", "module_name"),
[
("copy_eager", "copy_eager_", "ops.common.tilelang.eager_copy"),
("copy_lazy_out_idx", "copy_lazy_out_idx_", "ops.common.tilelang.lazy_out_idx_copy"),
("copy_eager", "copy_eager_", "operator_runtime_backends.tilelang.copy_templates"),
("copy_lazy_out_idx", "copy_lazy_out_idx_", "operator_runtime_backends.tilelang.copy_templates"),
],
)
def test_tilelang_copy_templates_match_torch(copy_fn_name, copy_out_fn_name, module_name) -> None:
@ -39,8 +39,8 @@ def test_tilelang_copy_templates_match_torch(copy_fn_name, copy_out_fn_name, mod
@pytest.mark.parametrize(
("copy_out_fn_name", "module_name", "message"),
[
("copy_eager_", "ops.common.tilelang.eager_copy", "matching shapes"),
("copy_lazy_out_idx_", "ops.common.tilelang.lazy_out_idx_copy", "matching shapes"),
("copy_eager_", "operator_runtime_backends.tilelang.copy_templates", "matching shapes"),
("copy_lazy_out_idx_", "operator_runtime_backends.tilelang.copy_templates", "matching shapes"),
],
)
def test_tilelang_copy_templates_reject_mismatched_out(