Merge pull request #57 from InfiniTensor/standardize-test-parametrization-and-configuration

Standardize test parametrization and configuration

Author: voltjia

src/ntops/torch/utils.py

@@ -6,10 +6,55 @@
 import ntops
 
 
+class _CachedMakeDefaultConfig:
+    def __init__(self, num_warps=None, num_stages=None, max_num_configs=None):
+        self.num_warps = num_warps
+
+        self.num_stages = num_stages
+
+        self.max_num_configs = max_num_configs
+
+
+_cached_make_default_config = _CachedMakeDefaultConfig()
+
+
+def get_default_num_warps():
+    return _cached_make_default_config.num_warps
+
+
+def set_default_num_warps(num_warps):
+    _cached_make_default_config.num_warps = num_warps
+
+
+def get_default_num_stages():
+    return _cached_make_default_config.num_stages
+
+
+def set_default_num_stages(num_stages):
+    _cached_make_default_config.num_stages = num_stages
+
+
+def get_default_max_num_configs():
+    return _cached_make_default_config.max_num_configs
+
+
+def set_default_max_num_configs(max_num_configs):
+    _cached_make_default_config.max_num_configs = max_num_configs
+
+
 @functools.cache
 def _cached_make(
     premake, *args, num_warps=None, num_stages=None, max_num_configs=None, **keywords
 ):
+    if num_warps is None:
+        num_warps = _cached_make_default_config.num_warps
+
+    if num_stages is None:
+        num_stages = _cached_make_default_config.num_stages
+
+    if max_num_configs is None:
+        max_num_configs = _cached_make_default_config.max_num_configs
+
     return ninetoothed.make(
         *premake(*args, **keywords),
         num_warps=num_warps,

tests/conftest.py

@@ -4,11 +4,15 @@
 import pytest
 import torch
 
+import ntops.torch.utils
+
 
 def pytest_configure():
     torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
     torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
 
+    ntops.torch.utils.set_default_max_num_configs(_DEFAULT_MAX_NUM_CONFIGS)
+
 
 def pytest_collectstart(collector):
     if isinstance(collector, pytest.Module):
@@ -25,6 +29,9 @@ def set_seed_per_test(request):
     _set_random_seed(_hash(_test_case_path_from_request(request)))
 
 
+_DEFAULT_MAX_NUM_CONFIGS = 3
+
+
 def _set_random_seed(seed):
     random.seed(seed)
     torch.manual_seed(seed)

tests/test_abs.py

@@ -8,12 +8,10 @@
 
 @skip_if_cuda_not_available
 @pytest.mark.parametrize(*generate_arguments())
-def test_cuda(shape, dtype, atol, rtol):
-    device = "cuda"
-
+def test_abs(shape, dtype, device, rtol, atol):
     input = torch.randn(shape, dtype=dtype, device=device)
 
     ninetoothed_output = ntops.torch.abs(input)
     reference_output = torch.abs(input)
 
-    assert torch.allclose(ninetoothed_output, reference_output, atol=atol, rtol=rtol)
+    assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol)

tests/test_add.py

@@ -8,14 +8,12 @@
 
 @skip_if_cuda_not_available
 @pytest.mark.parametrize(*generate_arguments())
-def test_cuda(shape, dtype, atol, rtol):
-    device = "cuda"
-
+def test_add(shape, dtype, device, rtol, atol):
     input = torch.randn(shape, dtype=dtype, device=device)
     other = torch.randn(shape, dtype=dtype, device=device)
     alpha = gauss()
 
     ninetoothed_output = ntops.torch.add(input, other, alpha=alpha)
     reference_output = torch.add(input, other, alpha=alpha)
 
-    assert torch.allclose(ninetoothed_output, reference_output, atol=atol, rtol=rtol)
+    assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol)

tests/test_addmm.py

@@ -9,9 +9,7 @@
 
 @skip_if_cuda_not_available
 @pytest.mark.parametrize(*generate_arguments())
-def test_cuda(m, n, k, dtype, atol, rtol):
-    device = "cuda"
-
+def test_addmm(m, n, k, dtype, device, rtol, atol):
     input = torch.randn((m, n), dtype=dtype, device=device)
     x = torch.randn((m, k), dtype=dtype, device=device)
     y = torch.randn((k, n), dtype=dtype, device=device)
@@ -21,4 +19,4 @@ def test_cuda(m, n, k, dtype, atol, rtol):
     ninetoothed_output = ntops.torch.addmm(input, x, y, beta=beta, alpha=alpha)
     reference_output = torch.addmm(input, x, y, beta=beta, alpha=alpha)
 
-    assert torch.allclose(ninetoothed_output, reference_output, atol=atol, rtol=rtol)
+    assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol)

tests/test_bitwise_and.py

@@ -8,9 +8,7 @@
 
 @skip_if_cuda_not_available
 @pytest.mark.parametrize(*generate_arguments(False))
-def test_cuda(shape, dtype, atol, rtol):
-    device = "cuda"
-
+def test_bitwise_and(shape, dtype, device, rtol, atol):
     if dtype == torch.bool:
         prob = 0.5
         input = torch.rand(shape, dtype=torch.float32, device=device) > prob

tests/test_bitwise_not.py

@@ -8,9 +8,7 @@
 
 @skip_if_cuda_not_available
 @pytest.mark.parametrize(*generate_arguments(False))
-def test_cuda(shape, dtype, atol, rtol):
-    device = "cuda"
-
+def test_bitwise_not(shape, dtype, device, rtol, atol):
     if dtype == torch.bool:
         prob = 0.5
         input = torch.rand(shape, dtype=torch.float32, device=device) > prob

tests/test_bitwise_or.py

@@ -8,9 +8,7 @@
 
 @skip_if_cuda_not_available
 @pytest.mark.parametrize(*generate_arguments(False))
-def test_cuda(shape, dtype, atol, rtol):
-    device = "cuda"
-
+def test_bitwise_or(shape, dtype, device, rtol, atol):
     if dtype == torch.bool:
         prob = 0.5
         input = torch.rand(shape, dtype=torch.float32, device=device) > prob

tests/test_bmm.py

@@ -10,14 +10,12 @@
 
 @skip_if_cuda_not_available
 @pytest.mark.parametrize(*generate_arguments())
-def test_cuda(m, n, k, dtype, atol, rtol):
-    device = "cuda"
-
+def test_bmm(m, n, k, dtype, device, rtol, atol):
     b = random.randint(4, 16)
     input = torch.randn((b, m, k), dtype=dtype, device=device)
     other = torch.randn((b, k, n), dtype=dtype, device=device)
 
     ninetoothed_output = ntops.torch.bmm(input, other)
     reference_output = torch.bmm(input, other)
 
-    assert torch.allclose(ninetoothed_output, reference_output, atol=atol, rtol=rtol)
+    assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol)

tests/test_clamp.py

@@ -8,14 +8,12 @@
 
 @skip_if_cuda_not_available
 @pytest.mark.parametrize(*generate_arguments())
-def test_cuda(shape, dtype, atol, rtol):
-    device = "cuda"
-
+def test_clamp(shape, dtype, device, rtol, atol):
     input = torch.randn(shape, dtype=dtype, device=device)
     min = torch.randn(shape, dtype=dtype, device=device)
     max = torch.randn(shape, dtype=dtype, device=device)
 
     ninetoothed_output = ntops.torch.clamp(input, min, max)
     reference_output = torch.clamp(input, min, max)
 
-    assert torch.allclose(ninetoothed_output, reference_output, atol=atol, rtol=rtol)
+    assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol)