NVIDIA/cutlass
1
0
Fork 0
mirror of https://github.com/NVIDIA/cutlass.git synced 2026-07-18 01:27:18 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
main
cutlass/operators/test/unit/test_reference_check.py
Junkai-Wu 7cd295cece v4.6 tag release. (#3362)
2026-07-06 22:05:33 -04:00

478 lines
21 KiB
Python
Raw Permalink Blame History

# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""
Unit tests for assert_close_with_reference_conversion function. Does not require GPU.
"""
import jax.numpy as jnp
import numpy as np
import pytest
import torch
import cutlass
from test_utils.reference_check import ClampMode, assert_close_with_reference_conversion
class TestAssertCloseWithReferenceConversionTorch:
def test_basic_float32_match(self):
"""Test basic comparison with matching float32 tensors."""
reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
result = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float32
)
def test_float32_with_tolerance(self):
"""Test comparison with rtols and atols parameters."""
reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
result = torch.tensor([1.001, 2.002, 3.003], dtype=torch.float32)
# Test that values within tolerance pass
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Float32,
rtols=[1e-2],
atols=[1e-2],
)
# Test that values outside tolerance raise AssertionError
# This test passes if AssertionError is raised (expected behavior)
result_outside_tolerance = torch.tensor([1.1, 2.1, 3.1], dtype=torch.float32)
with pytest.raises(AssertionError):
assert_close_with_reference_conversion(
result_outside_tolerance,
reference,
output_dtypes=cutlass.Float32,
rtols=[1e-2],
atols=[1e-2],
)
def test_cutlass_float16_dtype_conversion(self):
"""Test dtype conversion from float32 reference to cutlass.Float16 result."""
reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
result = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float16)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float16
)
def test_torch_float16_dtype_conversion(self):
"""Test dtype conversion from float32 reference to torch.float16 result."""
reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
result = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float16)
assert_close_with_reference_conversion(
result, reference, output_dtypes=torch.float16
)
def test_int_dtype(self):
"""Test comparison with integer dtype."""
reference = torch.tensor([1, 2, 3], dtype=torch.int32)
result = torch.tensor([1, 2, 3], dtype=torch.int32)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Int32
)
def test_multiple_tensors(self):
"""Test comparison with multiple tensors."""
ref1 = torch.tensor([1.0, 2.0], dtype=torch.float32)
ref2 = torch.tensor([3.0, 4.0], dtype=torch.float32)
res1 = torch.tensor([1.0, 2.0], dtype=torch.float32)
res2 = torch.tensor([3.0, 4.0], dtype=torch.float32)
assert_close_with_reference_conversion(
[res1, res2],
[ref1, ref2],
output_dtypes=[cutlass.Float32, cutlass.Float32],
)
def test_multiple_tensors_different_dtypes(self):
"""Test multiple tensors with different dtypes."""
ref1 = torch.tensor([1.0, 2.0], dtype=torch.float32)
ref2 = torch.tensor([3, 4], dtype=torch.int32)
res1 = torch.tensor([1.0, 2.0], dtype=torch.float16)
res2 = torch.tensor([3, 4], dtype=torch.int32)
assert_close_with_reference_conversion(
[res1, res2],
[ref1, ref2],
output_dtypes=[cutlass.Float16, cutlass.Int32],
)
def test_overflow_saturation_float16(self):
"""Test that overflow values are clamped to dtype range."""
# Create reference with values outside float16 range
reference = torch.tensor([1e10, -1e10, 1.0], dtype=torch.float32)
# Result should be saturated to float16 max/min
result = torch.tensor([65504.0, -65504.0, 1.0], dtype=torch.float16)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float16, clamp=True
)
def test_custom_message(self):
"""Test that custom error message is passed through."""
reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
result = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
# Should not raise, but message should be set
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Float32,
msg="Custom test message",
)
def test_rtols_atols_expansion(self):
"""Test that single rtols/atols values are expanded to match tensor count."""
ref1 = torch.tensor([1.0], dtype=torch.float32)
ref2 = torch.tensor([2.0], dtype=torch.float32)
res1 = torch.tensor([1.001], dtype=torch.float32)
res2 = torch.tensor([2.001], dtype=torch.float32)
# Single rtols/atols should be expanded to both tensors
assert_close_with_reference_conversion(
[res1, res2],
[ref1, ref2],
output_dtypes=[cutlass.Float32, cutlass.Float32],
rtols=[1e-2], # Single value
atols=[1e-2], # Single value
)
def test_mismatched_dtype_raises(self):
"""Test that mismatched result dtype raises ValueError."""
reference = torch.tensor([1.0, 2.0], dtype=torch.float32)
result = torch.tensor([1.0, 2.0], dtype=torch.float32) # Wrong dtype in result
with pytest.raises(ValueError, match="does not match target dtype"):
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Float16, # not float32
)
def test_wrong_number_of_tensors_raises(self):
"""Test that mismatched number of tensors raises ValueError."""
reference = torch.tensor([1.0], dtype=torch.float32)
result1 = torch.tensor([1.0], dtype=torch.float32)
result2 = torch.tensor([2.0], dtype=torch.float32)
with pytest.raises(
ValueError,
match="number of result tensors, references, and output dtypes must be the same",
):
assert_close_with_reference_conversion(
[result1, result2], reference, output_dtypes=cutlass.Float32
)
def test_wrong_number_of_rtols_raises(self):
"""Test that wrong number of rtols values raises ValueError."""
ref1 = torch.tensor([1.0], dtype=torch.float32)
ref2 = torch.tensor([2.0], dtype=torch.float32)
res1 = torch.tensor([1.0], dtype=torch.float32)
res2 = torch.tensor([2.0], dtype=torch.float32)
with pytest.raises(ValueError, match="number of rtol values must be the same"):
assert_close_with_reference_conversion(
[res1, res2],
[ref1, ref2],
output_dtypes=[cutlass.Float32, cutlass.Float32],
rtols=[1e-2, 1e-2, 1e-2], # Wrong count
)
def test_wrong_number_of_atols_raises(self):
"""Test that wrong number of atols values raises ValueError."""
ref1 = torch.tensor([1.0], dtype=torch.float32)
ref2 = torch.tensor([2.0], dtype=torch.float32)
res1 = torch.tensor([1.0], dtype=torch.float32)
res2 = torch.tensor([2.0], dtype=torch.float32)
with pytest.raises(ValueError, match="number of atol values must be the same"):
assert_close_with_reference_conversion(
[res1, res2],
[ref1, ref2],
output_dtypes=[cutlass.Float32, cutlass.Float32],
atols=[1e-2, 1e-2, 1e-2], # Wrong count
)
def test_invalid_output_dtype_raises(self):
"""Test that invalid output dtype type raises ValueError."""
reference = torch.tensor([1.0, 2.0], dtype=torch.float32)
result = torch.tensor([1.0, 2.0], dtype=torch.float32)
with pytest.raises(
ValueError, match="Output dtype requires a torch.dtype or cutlass.Numeric"
):
assert_close_with_reference_conversion(
result, reference, output_dtypes=["invalid"]
)
class TestAssertCloseWithReferenceConversionNumpy:
def test_np_array_basic(self):
"""Test basic comparison with numpy arrays."""
reference = np.array([1.0, 2.0, 3.0], dtype=np.float32)
result = np.array([1.0, 2.0, 3.0], dtype=np.float32)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float32
)
def test_cutlass_float16_dtype_conversion(self):
"""Test dtype conversion from float32 reference to float16 result."""
reference = np.array([1.0, 2.0, 3.0], dtype=np.float32)
result = np.array([1.0, 2.0, 3.0], dtype=np.float16)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float16
)
def test_overflow_float16(self):
"""Test that overflow values are not clamped with non-narrow datatype."""
# Create reference with values outside float16 range
reference = np.array([1e10, -1e10, 1.0], dtype=np.float32)
result = np.array([float("inf"), float("-inf"), 1.0], dtype=np.float16)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float16
)
def test_overflow_saturation_float8_e4m3fn(self):
"""Test that overflow values are clamped to dtype range."""
# Create reference with values outside float8_e4m3fn range
reference = np.array([1e10, -1e10, 1.0], dtype=np.float32)
# Result should be saturated to float8_e4m3fn max/min
result = np.array([448.0, -448.0, 1.0], dtype=jnp.float8_e4m3fn)
assert_close_with_reference_conversion(
result, reference, output_dtypes=jnp.float8_e4m3fn
)
def test_numpy_float16_dtype_conversion(self):
"""Test using np.dtype directly instead of cutlass.Numeric."""
reference = np.array([1.0, 2.0], dtype=np.float32)
result = np.array([1.0, 2.0], dtype=np.float16)
assert_close_with_reference_conversion(
result, reference, output_dtypes=np.float16
)
def test_mismatched_tensor_types_raises(self):
"""Test that mixing numpy and torch tensors raises ValueError."""
reference = np.array([1.0, 2.0], dtype=np.float32)
result = torch.tensor([1.0, 2.0], dtype=torch.float32)
with pytest.raises(
ValueError, match="jax array and result tensor is a torch tensor"
):
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float32
)
class TestAssertCloseWithReferenceConversionJax:
def test_jax_array_float8_e4m3fn_with_cutlass_dtype(self):
"""Test jax float8_e4m3fn arrays with cutlass.Float8E4M3FN output dtype."""
reference = jnp.array([1.0, -1.0, 0.5], dtype=jnp.float32)
result = jnp.array([1.0, -1.0, 0.5], dtype=jnp.float8_e4m3fn)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float8E4M3FN
)
def test_jax_array_float8_e5m2_with_cutlass_dtype(self):
"""Test jax float8_e5m2 arrays with cutlass.Float8E5M2 output dtype."""
reference = jnp.array([1.0, -1.0, 0.5], dtype=jnp.float32)
result = jnp.array([1.0, -1.0, 0.5], dtype=jnp.float8_e5m2)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float8E5M2
)
def test_jax_array_overflow_clamp_float8_e4m3fn(self):
"""Test overflow clamping for jax float8_e4m3fn arrays."""
reference = jnp.array([500.0, -500.0, 1.0], dtype=jnp.float32)
result = jnp.array([448.0, -448.0, 1.0], dtype=jnp.float8_e4m3fn)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float8E4M3FN
)
class TestAssertCloseWithReferenceConversionEdgeCases:
def test_none_rtols_atols_defaults(self):
"""Test that None rtols/atols default to 0.0."""
reference = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
result = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)
# Should work with exact match when rtols/atols are None
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Float32,
rtols=None,
atols=None,
)
result_outside_tolerance = torch.tensor([1.1, 2.1, 3.1], dtype=torch.float32)
with pytest.raises(AssertionError):
assert_close_with_reference_conversion(
result_outside_tolerance,
reference,
output_dtypes=cutlass.Float32,
rtols=None,
atols=None,
)
def test_empty_tensor(self):
"""Test comparison with empty tensors."""
reference = torch.tensor([], dtype=torch.float32)
result = torch.tensor([], dtype=torch.float32)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float32
)
def test_torch_equals_nan(self):
"""Test NaN equality does not raise on assertion."""
reference = torch.tensor([1.0, float("nan"), 3.0], dtype=torch.float16)
result = torch.tensor([1.0, float("nan"), 3.0], dtype=torch.float16)
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Float16,
clamp=False,
equal_nan=True,
)
def test_numpy_equals_nan(self):
"""Test NaN equality does not raise on assertion."""
reference = np.array([1.0, float("nan"), 3.0], dtype=np.float32)
result = np.array([1.0, float("nan"), 3.0], dtype=np.float32)
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Float32,
clamp=False,
equal_nan=True,
)
def test_clamp_mode_enum_clamp(self):
"""Test that ClampMode.Clamp enum value works correctly."""
# Create reference with values outside float16 range
reference = torch.tensor([1e10, -1e10, 1.0], dtype=torch.float32)
# Result should be saturated to float16 max/min
result = torch.tensor([65504.0, -65504.0, 1.0], dtype=torch.float16)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float16, clamp=ClampMode.Clamp
)
def test_clamp_mode_enum_no_clamp(self):
"""Test that ClampMode.NoClamp enum value works correctly."""
reference = torch.tensor([1.0, float("inf"), 3.0], dtype=torch.float32)
result = torch.tensor([1.0, float("inf"), 3.0], dtype=torch.float32)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float32, clamp=ClampMode.NoClamp
)
def test_restricted_clamp_mode_clamp_on_float16(self):
"""Test that ClampMode.ClampOnRestrictedDtypes doesn't clamp on float16."""
# float16 is a restricted dtype, so should not clamp
reference = torch.tensor([1.0, float("inf"), 3.0], dtype=torch.float32)
result = torch.tensor([1.0, float("inf"), 3.0], dtype=torch.float16)
# Should not raise ValueError when clamp=ClampMode.ClampOnRestrictedDtypes with float16
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Float16,
clamp=ClampMode.ClampOnRestrictedDtypes,
)
def test_restricted_clamp_mode_clamp_on_int8(self):
"""Test that ClampMode.ClampOnRestrictedDtypes doesn't clamp on int8."""
# int8 is a restricted dtype, so should not clamp
reference = torch.tensor([1000, -1000, 1], dtype=torch.int32)
# When not clamped, overflow values will wrap or become invalid
result = torch.tensor([-24, 24, 1], dtype=torch.int8)
# Should not clamp overflow values for int8
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Int8,
clamp=ClampMode.ClampOnRestrictedDtypes,
)
def test_restricted_clamp_mode_clamp_on_uint8(self):
"""Test that ClampMode.ClampOnRestrictedDtypes doesn't clamp on uint8."""
# uint8 is a restricted dtype, so should not clamp
reference = torch.tensor([1000, 500, 1], dtype=torch.int32)
# When not clamped, overflow values will wrap
result = torch.tensor([232, 244, 1], dtype=torch.uint8)
# Should not clamp overflow values for uint8
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Uint8,
clamp=ClampMode.ClampOnRestrictedDtypes,
)
def test_restricted_clamp_mode_clamp_on_float8_e4m3fn(self):
"""Test that ClampMode.ClampOnRestrictedDtypes DOES clamp on float8_e4m3fn (not in NOCLAMP list)."""
# float8_e4m3fn is NOT in NOCLAMP list, so should clamp
reference = torch.tensor([500.0, -500.0, 1.0], dtype=torch.float32)
# Result should be clamped to float8_e4m3fn range
result = torch.tensor([448.0, -448.0, 1.0], dtype=torch.float8_e4m3fn)
# Should clamp overflow values for float8_e4m3fn
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Float8E4M3FN,
clamp=ClampMode.ClampOnRestrictedDtypes,
)
def test_restricted_clamp_mode_clamp_on_float8_e5m2(self):
"""Test that ClampMode.ClampOnRestrictedDtypes DOES clamp on float8_e5m2 (not in NOCLAMP list)."""
# float8_e5m2 is NOT in NOCLAMP list, so should clamp
reference = torch.tensor([100000.0, -100000.0, 1.0], dtype=torch.float32)
# Result should be clamped to float8_e5m2 range
result = torch.tensor([57344.0, -57344.0, 1.0], dtype=torch.float8_e5m2)
# Should clamp overflow values for float8_e5m2
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Float8E5M2,
clamp=ClampMode.ClampOnRestrictedDtypes,
)
def test_clamp_mode_bool_true(self):
"""Test that ClampMode can be converted from boolean True."""
reference = torch.tensor([100000.0, -100000.0, 1.0], dtype=torch.float32)
result = torch.tensor([57344.0, -57344.0, 1.0], dtype=torch.float8_e5m2)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float8E5M2, clamp=True
)
def test_clamp_mode_bool_false(self):
"""Test that ClampMode can be converted from boolean False."""
reference = torch.tensor([1000.0, -1000.0, 1.0], dtype=torch.float32)
result = torch.tensor(
[float("nan"), -float("nan"), 1.0], dtype=torch.float8_e4m3fn
)
assert_close_with_reference_conversion(
result,
reference,
output_dtypes=cutlass.Float8E4M3FN,
clamp=False,
equal_nan=True,
)
def test_numpy_reference_jax_result_float32(self):
"""Test numpy reference array with jax result array."""
reference = np.array([1.0, 2.0, 3.0], dtype=np.float32)
result = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
assert_close_with_reference_conversion(
result, reference, output_dtypes=cutlass.Float32
)
Reference in New Issue View Git Blame Copy Permalink