diff --git a/batch_invariant_ops/batch_invariant_ops.py b/batch_invariant_ops/batch_invariant_ops.py index b9021bb..2741ab6 100644 --- a/batch_invariant_ops/batch_invariant_ops.py +++ b/batch_invariant_ops/batch_invariant_ops.py @@ -75,7 +75,7 @@ def matmul_kernel_persistent( offs_k_for_mask = tl.arange(0, BLOCK_SIZE_K) num_pid_in_group = GROUP_SIZE_M * num_pid_n - for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True): + for tile_id in range(start_pid, num_tiles, NUM_SMS): pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M, NUM_SMS) start_m = pid_m * BLOCK_SIZE_M start_n = pid_n * BLOCK_SIZE_N @@ -340,6 +340,18 @@ def mean_kernel( Kernel for computing mean along a single dimension. Input is viewed as (M, N, K) where N is the dimension being reduced. """ + # Handle edge case: empty reduction dimension returns NaN + if N == 0: + pid = tl.program_id(0) + m_idx = pid // K + k_idx = pid % K + if m_idx >= M or k_idx >= K: + return + mean_val = float("nan") + output_idx = m_idx * output_stride0 + k_idx * output_stride1 + tl.store(output_ptr + output_idx, mean_val) + return + # Program ID gives us which output element we're computing pid = tl.program_id(0) @@ -390,6 +402,16 @@ def mean_dim( assert -input.ndim <= dim < input.ndim, ( f"Invalid dimension {dim} for tensor with {input.ndim} dimensions" ) + + # Handle empty dimension case explicitly to avoid NaN + if input.shape[dim] == 0: + # Return a tensor of NaN values with the appropriate shape + shape = list(input.shape) + if keepdim: + shape[dim] = 1 + else: + shape.pop(dim) + return torch.full(shape, float("nan"), dtype=dtype or input.dtype, device=input.device) # Handle negative dim if dim < 0: diff --git a/tests/__pycache__/conftest.cpython-312-pytest-7.4.4.pyc b/tests/__pycache__/conftest.cpython-312-pytest-7.4.4.pyc new file mode 100644 index 0000000..e695439 Binary files /dev/null and b/tests/__pycache__/conftest.cpython-312-pytest-7.4.4.pyc differ diff --git a/tests/__pycache__/test_batch_invariant_ops.cpython-312-pytest-7.4.4.pyc b/tests/__pycache__/test_batch_invariant_ops.cpython-312-pytest-7.4.4.pyc new file mode 100644 index 0000000..fce3f22 Binary files /dev/null and b/tests/__pycache__/test_batch_invariant_ops.cpython-312-pytest-7.4.4.pyc differ diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..73bfe13 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,9 @@ +""" +Pytest configuration for batch-invariant ops tests. +""" +import sys +from pathlib import Path + +# Add the project root to the path +project_root = Path(__file__).parent.parent +sys.path.insert(0, str(project_root)) diff --git a/tests/test_batch_invariant_ops.py b/tests/test_batch_invariant_ops.py new file mode 100644 index 0000000..917033f --- /dev/null +++ b/tests/test_batch_invariant_ops.py @@ -0,0 +1,255 @@ +""" +Test suite for batch-invariant operations. + +This test suite verifies that the batch-invariant property holds for the following operations: +- mm (matrix multiplication) +- addmm (matrix multiplication with bias) +- log_softmax +- mean + +The batch-invariant property states that: + op(x[:1], y) == op(x, y)[:1] + +This means computing an operation on a single batch element should give the same result +as computing it on the full batch and then taking the first element. +""" + +import pytest +import torch +from batch_invariant_ops import ( + set_batch_invariant_mode, + matmul_persistent, + addmm_batch_invariant, + log_softmax, + mean_dim, +) + + +def get_device(): + """Get the current accelerator device.""" + device_type = getattr(torch.accelerator.current_accelerator(), "type", "cpu") + if device_type == "cpu": + return "cpu" + elif device_type in ("cuda", "xpu"): + return device_type + return "cpu" + + +DEVICE = get_device() +DTYPES = [torch.float32, torch.float16, torch.bfloat16] if DEVICE != "cpu" else [torch.float32] + + +class TestBatchInvariantMM: + """Tests for mm (matrix multiplication) batch invariance.""" + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mm_batch_invariant_basic(self, dtype): + """Test basic batch invariance for mm operation.""" + B, M, K, N = 4, 16, 32, 16 + + x = torch.randn(B, M, K, dtype=dtype, device=DEVICE) + y = torch.randn(K, N, dtype=dtype, device=DEVICE) + + # Single batch element + out_single = torch.mm(x[:1], y) + # Full batch, take first + out_full = torch.mm(x, y)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mm_batch_invariant_large(self, dtype): + """Test batch invariance with larger matrices.""" + B, M, K, N = 8, 128, 256, 128 + + x = torch.randn(B, M, K, dtype=dtype, device=DEVICE) + y = torch.randn(K, N, dtype=dtype, device=DEVICE) + + out_single = torch.mm(x[:1], y) + out_full = torch.mm(x, y)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mm_batch_invariant_single_row(self, dtype): + """Test with single row in batch dimension.""" + B, M, K, N = 1, 8, 16, 8 + + x = torch.randn(B, M, K, dtype=dtype, device=DEVICE) + y = torch.randn(K, N, dtype=dtype, device=DEVICE) + + out_single = torch.mm(x[:1], y) + out_full = torch.mm(x, y)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + +class TestBatchInvariantAddMM: + """Tests for addmm (matrix multiplication with bias) batch invariance.""" + + @pytest.mark.parametrize("dtype", DTYPES) + def test_addmm_batch_invariant_basic(self, dtype): + """Test basic batch invariance for addmm operation.""" + B, M, K, N = 4, 16, 32, 16 + + x = torch.randn(B, M, K, dtype=dtype, device=DEVICE) + y = torch.randn(K, N, dtype=dtype, device=DEVICE) + bias = torch.randn(N, dtype=dtype, device=DEVICE) + + # Single batch element + out_single = torch.addmm(bias, x[:1], y) + # Full batch, take first + out_full = torch.addmm(bias, x, y)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_addmm_batch_invariant_large(self, dtype): + """Test batch invariance with larger matrices.""" + B, M, K, N = 8, 128, 256, 128 + + x = torch.randn(B, M, K, dtype=dtype, device=DEVICE) + y = torch.randn(K, N, dtype=dtype, device=DEVICE) + bias = torch.randn(N, dtype=dtype, device=DEVICE) + + out_single = torch.addmm(bias, x[:1], y) + out_full = torch.addmm(bias, x, y)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + +class TestBatchInvariantLogSoftmax: + """Tests for log_softmax batch invariance.""" + + @pytest.mark.parametrize("dtype", DTYPES) + def test_log_softmax_batch_invariant_basic(self, dtype): + """Test basic batch invariance for log_softmax operation.""" + B, S, V = 4, 16, 32 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + # Single batch element + out_single = torch.log_softmax(x[:1], dim=-1) + # Full batch, take first + out_full = torch.log_softmax(x, dim=-1)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_log_softmax_batch_invariant_2d(self, dtype): + """Test batch invariance for 2D tensors.""" + B, V = 4, 64 + + x = torch.randn(B, V, dtype=dtype, device=DEVICE) + + out_single = torch.log_softmax(x[:1], dim=-1) + out_full = torch.log_softmax(x, dim=-1)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_log_softmax_batch_invariant_large(self, dtype): + """Test batch invariance with larger sequences.""" + B, S, V = 8, 128, 512 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + out_single = torch.log_softmax(x[:1], dim=-1) + out_full = torch.log_softmax(x, dim=-1)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + +class TestBatchInvariantMean: + """Tests for mean operation batch invariance.""" + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mean_batch_invariant_basic(self, dtype): + """Test basic batch invariance for mean operation.""" + B, S, V = 4, 16, 32 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + # Single batch element + out_single = torch.mean(x[:1], dim=1) + # Full batch, take first + out_full = torch.mean(x, dim=1)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mean_batch_invariant_dim0(self, dtype): + """Test batch invariance for mean along dim 0.""" + B, S, V = 4, 16, 32 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + out_single = torch.mean(x[:1], dim=0) + out_full = torch.mean(x, dim=0)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mean_batch_invariant_keepdim(self, dtype): + """Test batch invariance for mean with keepdim=True.""" + B, S, V = 4, 16, 32 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + out_single = torch.mean(x[:1], dim=1, keepdim=True) + out_full = torch.mean(x, dim=1, keepdim=True)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + @pytest.mark.parametrize("dtype", DTYPES) + def test_mean_batch_invariant_large(self, dtype): + """Test batch invariance with larger tensors.""" + B, S, V = 8, 64, 128 + + x = torch.randn(B, S, V, dtype=dtype, device=DEVICE) + + out_single = torch.mean(x[:1], dim=1) + out_full = torch.mean(x, dim=1)[:1] + + torch.testing.assert_close(out_single, out_full, atol=1e-3, rtol=1e-3) + + +class TestBatchInvariantMode: + """Tests for batch invariant mode context manager.""" + + def test_batch_invariant_mode_enable_disable(self): + """Test enabling and disabling batch invariant mode.""" + from batch_invariant_ops import ( + is_batch_invariant_mode_enabled, + enable_batch_invariant_mode, + disable_batch_invariant_mode, + ) + + # Initially disabled + assert not is_batch_invariant_mode_enabled() + + # Enable + enable_batch_invariant_mode() + assert is_batch_invariant_mode_enabled() + + # Disable + disable_batch_invariant_mode() + assert not is_batch_invariant_mode_enabled() + + def test_batch_invariant_mode_context_manager(self): + """Test context manager for batch invariant mode.""" + from batch_invariant_ops import is_batch_invariant_mode_enabled + + # Initially disabled + assert not is_batch_invariant_mode_enabled() + + # Enable via context manager + with set_batch_invariant_mode(True): + assert is_batch_invariant_mode_enabled() + + # Should be disabled after context + assert not is_batch_invariant_mode_enabled() + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])