use syrk for int dtypes when possible

Author: Vahid Tavanashad

dpnp/dpnp_utils/dpnp_utils_linearalgebra.py

@@ -50,29 +50,6 @@
 ]
 
 
-def _call_syrk(x1, x2):
-    """
-    Check to see if `syrk` can be called instead of `gemm`.
-
-    It is assumed that x1 and x2 are usm_ndarray objects. These arrays have
-    already been validated to be 2-dimensional and contiguous. Therefore, this
-    function only verifies the following: Both arrays reference the same
-    memory. The number of rows in x1 equals the number of columns in x2. If one
-    array is C-contiguous, the other must be F-contiguous.
-
-    """
-    call_syrk = False
-    if (
-        x1._pointer == x2._pointer
-        and x1.shape[0] == x2.shape[1]
-        and x1.flags.c_contiguous != x2.flags.c_contiguous
-        and x1.flags.f_contiguous != x2.flags.f_contiguous
-    ):
-        call_syrk = True
-
-    return call_syrk
-
-
 def _compute_res_dtype(*arrays, dtype=None, out=None, casting="no"):
     """
     Determines the output array data type.
@@ -541,6 +518,29 @@ def _get_signature(func):
     return signature, distinct_core
 
 
+def _is_syrk_compatible(x1, x2):
+    """
+    Check to see if `syrk` can be called instead of `gemm`.
+    Input arrays have already been validated to be 2-dimensional.
+
+    """
+    # Must share data (same base buffer)
+    if dpnp.get_usm_ndarray(x1)._pointer != dpnp.get_usm_ndarray(x2)._pointer:
+        return False
+
+    # Result must be square
+    if x1.shape[0] != x2.shape[1]:
+        return False
+
+    # Strides must match transpose pattern
+    x1_strides = x1.strides
+    x2_strides = x2.strides
+    if x1_strides[0] != x2_strides[1] or x1_strides[1] != x2_strides[0]:
+        return False
+
+    return True
+
+
 def _shape_error(shape1, shape2, func, err_msg):
     """Validate the shapes of input and output arrays."""
 
@@ -983,6 +983,11 @@ def dpnp_multiplication(
         x1 = dpnp.reshape(x1, x1_shape[-2:])
         x2 = dpnp.reshape(x2, x2_shape[-2:])
         res_shape = (x1_shape[-2], x2_shape[-1])
+        if _is_syrk_compatible(x1, x2):
+            call_flag = "syrk"
+            res_dtype_orig = res_dtype
+            if dpnp.issubdtype(res_dtype, dpnp.integer):
+                res_dtype = dpnp.default_float_type(x1.device)
     elif x1_base_is_1D:
         # TODO: implement gemv_batch to use it here with transpose
         call_flag = "gemm_batch"
@@ -1088,21 +1093,17 @@ def dpnp_multiplication(
                         depends=_manager.submitted_events,
                     )
                     _manager.add_event_pair(ht_ev, gemv_ev)
+                elif call_flag == "syrk":
+                    _manager = dpu.SequentialOrderManager[exec_q]
+                    ht_ev, gemv_ev = bi._syrk(
+                        exec_q,
+                        dpnp.get_usm_ndarray(x1),
+                        dpnp.get_usm_ndarray(result),
+                        depends=_manager.submitted_events,
+                    )
+                    _manager.add_event_pair(ht_ev, gemv_ev)
                 elif call_flag == "gemm":
-                    x1_usm = dpnp.get_usm_ndarray(x1)
-                    x2_usm = dpnp.get_usm_ndarray(x2)
-                    call_syrk = _call_syrk(x1_usm, x2_usm)
-                    if call_syrk:
-                        _manager = dpu.SequentialOrderManager[exec_q]
-                        ht_ev, gemv_ev = bi._syrk(
-                            exec_q,
-                            x1_usm,
-                            dpnp.get_usm_ndarray(result),
-                            depends=_manager.submitted_events,
-                        )
-                        _manager.add_event_pair(ht_ev, gemv_ev)
-                    else:
-                        result = _gemm_matmul(exec_q, x1_usm, x2_usm, result)
+                    result = _gemm_matmul(exec_q, x1, x2, result)
                 else:
                     assert call_flag == "gemm_batch"
                     result = _gemm_batch_matmul(exec_q, x1, x2, result)
@@ -1130,6 +1131,9 @@ def dpnp_multiplication(
     elif res_shape != result_shape:
         result = dpnp.reshape(result, result_shape)
 
+    if call_flag == "syrk" and res_dtype_orig != res_dtype:
+        result = result.astype(res_dtype_orig)
+
     if out is None:
         if axes is not None:
             # Move the data back to the appropriate axes of the result array

dpnp/tests/test_product.py

@@ -12,7 +12,6 @@
     assert_dtype_allclose,
     generate_random_numpy_array,
     get_all_dtypes,
-    get_float_complex_dtypes,
     numpy_version,
 )
 from .third_party.cupy import testing
@@ -1184,7 +1183,7 @@ def test_special_case(self, dt_out, shape1, shape2):
         result = dpnp.matmul(ia, ib, out=iout)
         assert_dtype_allclose(result, expected)
 
-    @pytest.mark.parametrize("dt", get_float_complex_dtypes())
+    @pytest.mark.parametrize("dt", get_all_dtypes())
     def test_syrk(self, dt):
         a = generate_random_numpy_array((6, 9), dtype=dt)
         ia = dpnp.array(a)
@@ -1202,6 +1201,21 @@ def test_syrk(self, dt):
         expected = a.T @ a
         assert_dtype_allclose(result, expected)
 
+    @pytest.mark.parametrize("dt", [dpnp.int32, dpnp.float32])
+    def test_syrk_strided(self, dt):
+        a = generate_random_numpy_array((20, 30), dtype=dt)
+        ia = dpnp.array(a)
+        a = a[::2, ::2]
+        ia = ia[::2, ::2]
+
+        result = dpnp.matmul(ia, ia.mT)
+        expected = numpy.matmul(a, a.T)
+        assert_dtype_allclose(result, expected)
+
+        result = ia.mT @ ia
+        expected = a.T @ a
+        assert_dtype_allclose(result, expected)
+
     @pytest.mark.parametrize(
         "order, out_order",
         [("C", "C"), ("C", "F"), ("F", "C"), ("F", "F")],
@@ -1226,6 +1240,18 @@ def test_syrk_order(self, order):
         result = dpnp.matmul(ia, ia.mT)
         assert_dtype_allclose(result, expected)
 
+    # added for coverage
+    def test_not_syrk(self):
+        a = generate_random_numpy_array((20, 20), low=-5, high=5)
+        ia = dpnp.array(a)
+
+        # Result must be square
+        b = a.mT[:, ::2]
+        ib = ia.mT[:, ::2]
+        expected = numpy.matmul(a, b)
+        result = dpnp.matmul(ia, ib)
+        assert_dtype_allclose(result, expected)
+
     def test_bool(self):
         a = generate_random_numpy_array((3, 4), dtype=dpnp.bool)
         b = generate_random_numpy_array((4, 5), dtype=dpnp.bool)