Sync third-party tests for integer tensordot and int8 matmul

dpnp/tests/third_party/cupy/linalg_tests/test_product.py

@@ -548,6 +548,77 @@ def test_kron_empty_via_slice(self, xp, dtype):
         return xp.kron(a, b)
 
 
+@testing.slow
+# @pytest.mark.thread_unsafe(reason="Allocation too large.")
+def test_integer_tensordot_large_indexing():
+    # Dimensions fit int32, but the integer kernel's M * K does not.
+    # This needs just over 4 GiB; uint8 also avoids the int8 cuBLAS path.
+    # At this size the old bounds wrap to a small positive value, keeping
+    # its reads inside the allocation instead of causing invalid accesses.
+    k, m = 2**16 + 1, 2**16
+    a = b = out = None
+    try:
+        try:
+            a = cupy.zeros((2, k), dtype=cupy.uint8)
+            b = cupy.zeros((k, m), dtype=cupy.uint8)
+        except MemoryError:
+            pytest.skip("out of memory in test.")
+
+        # Select the second row, which is zero. The overflowing bound clamps
+        # its reads to the last element of the first row, which is one.
+        # A single contribution avoids masking the error via uint8 wraparound.
+        a[:, 1] = 1
+        b[0, :] = 1
+        out = cupy.tensordot(a, b, axes=1)
+        expected = numpy.zeros((2, m), dtype=numpy.uint8)
+        testing.assert_array_equal(out, expected)
+    finally:
+        del out, b, a
+        # cupy.get_default_memory_pool().free_all_blocks()
+
+
+class TestInt8Tensordot:
+    """Smoke test for cupy.tensordot with int8 dtype via tensordot_core."""
+
+    # def setup_method(self):
+    #     if cupy.cuda.runtime.is_hip:
+    #         pytest.skip("int8 cublasGemmEx path is NVIDIA-only")
+    #     if int(cupy.cuda.Device().compute_capability) < 61:
+    #         pytest.skip(
+    #             "CUBLAS_COMPUTE_32I requires compute capability >= 6.1"
+    #         )
+
+    @testing.numpy_cupy_array_equal()
+    def test_int8_tensordot_aligned(self, xp):
+        """k=16 is IMMA-aligned; exercises the cuBLAS Tensor Core path."""
+        rng = numpy.random.default_rng(seed=7)
+        a = xp.asarray(rng.integers(-5, 5, (8, 16), dtype=numpy.int8))
+        b = xp.asarray(rng.integers(-5, 5, (16, 8), dtype=numpy.int8))
+        return xp.tensordot(a, b, axes=1)
+
+    @testing.numpy_cupy_array_equal()
+    def test_int8_tensordot_unaligned(self, xp):
+        """k=7 is not IMMA-aligned; exercises the alignment-guard fallback."""
+        rng = numpy.random.default_rng(seed=13)
+        a = xp.asarray(rng.integers(-5, 5, (8, 7), dtype=numpy.int8))
+        b = xp.asarray(rng.integers(-5, 5, (7, 8), dtype=numpy.int8))
+        return xp.tensordot(a, b, axes=1)
+
+    @testing.numpy_cupy_array_equal()
+    def test_int8_inner_multidim(self, xp):
+        """ret_shape (2, 3, 3, 2) differs from the 2-D (6, 6) GEMM output."""
+        a = testing.shaped_arange((2, 3, 4), xp, numpy.int8)
+        b = testing.shaped_arange((3, 2, 4), xp, numpy.int8)
+        return xp.inner(a, b)
+
+    @testing.numpy_cupy_array_equal()
+    def test_int8_tensordot_multidim(self, xp):
+        rng = numpy.random.default_rng(seed=21)
+        a = xp.asarray(rng.integers(-5, 5, (2, 3, 8), dtype=numpy.int8))
+        b = xp.asarray(rng.integers(-5, 5, (8, 4, 5), dtype=numpy.int8))
+        return xp.tensordot(a, b, axes=1)
+
+
 @testing.parameterize(
     *testing.product(
         {

dpnp/tests/third_party/cupy/math_tests/test_matmul.py

@@ -1,10 +1,13 @@
 import operator
 import unittest
+import warnings
 
 import numpy
 import pytest
 
 import dpnp as cupy
+
+# from cupy._core import _routines_linalg as _linalg
 from dpnp.tests.helper import has_support_aspect64
 from dpnp.tests.third_party.cupy import testing
 
@@ -376,6 +379,116 @@ def test_cupy_matmul(self, xp):
         return xp.matmul(x1, x2)
 
 
+@testing.parameterize(
+    *testing.product(
+        {
+            "shape_pair": [
+                # k is a multiple of 4 (cuBLAS IMMA path)
+                ((32, 64), (64, 96)),
+                ((64, 32), (32, 96)),
+                # k is NOT a multiple of 4 (alignment guard -> CUDA kernel path)
+                ((32, 7), (7, 16)),
+                ((16, 9), (9, 32)),
+            ],
+        }
+    )
+)
+class TestMatmulInt8(unittest.TestCase):
+    """int8 matmul correctness against NumPy across aligned and unaligned k."""
+
+    # def setUp(self):
+    #     if cupy.cuda.runtime.is_hip:
+    #         pytest.skip("int8 cublasGemmEx path is NVIDIA-only")
+    #     if int(cupy.cuda.Device().compute_capability) < 61:
+    #         pytest.skip(
+    #             "CUBLAS_COMPUTE_32I requires compute capability >= 6.1"
+    #         )
+
+    @testing.numpy_cupy_array_equal()
+    def test_operator_matmul(self, xp):
+        rng = numpy.random.default_rng(seed=42)
+        m, k = self.shape_pair[0]
+        _, n = self.shape_pair[1]
+        x1 = xp.asarray(rng.integers(-10, 10, (m, k), dtype=numpy.int8))
+        x2 = xp.asarray(rng.integers(-10, 10, (k, n), dtype=numpy.int8))
+        return x1 @ x2
+
+    @testing.numpy_cupy_array_equal()
+    def test_transposed_inputs(self, xp):
+        """Exercises _mat_to_cublas_contiguous with non-standard strides."""
+        rng = numpy.random.default_rng(seed=0)
+        m, k = self.shape_pair[0]
+        _, n = self.shape_pair[1]
+        x1 = xp.asarray(
+            numpy.asfortranarray(
+                rng.integers(-10, 10, (m, k), dtype=numpy.int8)
+            )
+        )
+        x2 = xp.asarray(
+            numpy.asfortranarray(
+                rng.integers(-10, 10, (k, n), dtype=numpy.int8)
+            )
+        )
+        return x1 @ x2
+
+    def test_overflow_wraps_like_numpy(self):
+        """int32->int8 cast wraps mod 256, matching NumPy matmul semantics."""
+        a = cupy.full((4, 4), 100, dtype=numpy.int8)
+        b = cupy.full((4, 4), 100, dtype=numpy.int8)
+        np_a = cupy.asnumpy(a)
+        np_b = cupy.asnumpy(b)
+        testing.assert_array_equal(a @ b, np_a @ np_b)
+
+    @pytest.mark.skip("cuBLAS is not supported")
+    # @pytest.mark.thread_unsafe(
+    #     reason="patches cublas.gemmEx via mock.patch.object"
+    # )
+    def test_fallback_emits_performance_warning(self):
+        """When cuBLAS fails, PerformanceWarning is emitted before fallback."""
+        import unittest.mock as mock
+
+        from cupy import _util
+        from cupy_backends.cuda.libs import cublas
+
+        m, k = self.shape_pair[0]
+        _, n = self.shape_pair[1]
+        k_aligned = ((k + 3) // 4) * 4
+        a = cupy.zeros((m, k_aligned), dtype=numpy.int8)
+        b = cupy.zeros((k_aligned, n), dtype=numpy.int8)
+
+        def _failing_gemmEx(*args, **kwargs):
+            raise cublas.CUBLASError(13)  # CUBLAS_STATUS_NOT_SUPPORTED
+
+        with mock.patch.object(cublas, "gemmEx", side_effect=_failing_gemmEx):
+            with warnings.catch_warnings(record=True) as caught:
+                warnings.simplefilter("always")
+                _ = a @ b
+        perf_warnings = [
+            w
+            for w in caught
+            if issubclass(w.category, _util.PerformanceWarning)
+        ]
+        assert len(perf_warnings) >= 1
+
+
+@pytest.mark.skip("set_compute_type() is not supported")
+class TestMatmulInt8SetComputeTypeError(unittest.TestCase):
+    """set_compute_type rejects non-DEFAULT compute types for int dtypes."""
+
+    def test_int8_rejects_fp16_compute_type(self):
+        with pytest.raises(ValueError, match="integer dtypes"):
+            cupy._core.set_compute_type(numpy.int8, _linalg.COMPUTE_TYPE_FP16)
+
+    def test_int8_rejects_fp32_compute_type(self):
+        with pytest.raises(ValueError, match="integer dtypes"):
+            cupy._core.set_compute_type(numpy.int8, _linalg.COMPUTE_TYPE_FP32)
+
+    def test_int8_accepts_default_compute_type(self):
+        old = cupy._core.get_compute_type(numpy.int8)
+        cupy._core.set_compute_type(numpy.int8, _linalg.COMPUTE_TYPE_DEFAULT)
+        cupy._core.set_compute_type(numpy.int8, old)
+
+
 @testing.parameterize(
     *testing.product(
         {