updates for onemkl interfaces

Author: Vahid Tavanashad

dpnp/backend/extensions/blas/blas_py.cpp

@@ -146,14 +146,14 @@ PYBIND11_MODULE(_blas_impl, m)
 
     {
         m.def(
-            "_row_major_is_available",
-            [](void) {
-#if defined(USE_ONEMKL_CUBLAS)
-                return false;
-#else
+            "_using_onemkl_interfaces",
+            []() {
+#ifdef USE_ONEMKL_INTERFACES
                 return true;
-#endif // USE_ONEMKL_CUBLAS
+#else
+                return false;
+#endif
             },
-            "Check if the onemkl::blas::row_major can be used.");
+            "Check if the OneMKL interfaces are being used.");
     }
 }

dpnp/dpnp_utils/dpnp_utils_linearalgebra.py

@@ -503,6 +503,28 @@ def _gemm_matmul(exec_q, x1, x2, res):
     return res
 
 
+def _gemm_special_case(x1, x2, res_dtype, call_flag):
+    """
+    `gemm` and `gemm_batch` support these special cases of data types
+    while `gemv` does not.
+
+    """
+
+    # TODO: replace with dpnp.int8 when it is added
+    x1_is_int8 = dpnp.issubdtype(x1.dtype, numpy.int8)
+    x2_is_int8 = dpnp.issubdtype(x2.dtype, numpy.int8)
+    res_is_int32 = dpnp.issubdtype(res_dtype, dpnp.int32)
+    res_is_float32 = dpnp.issubdtype(res_dtype, dpnp.float32)
+
+    flag = x1_is_int8 and x2_is_int8 and (res_is_int32 or res_is_float32)
+    flag = flag and call_flag in ["gemm", "gemm_batch"]
+
+    # onemkl_interfaces does not support these data types
+    onemkl_interfaces = bi._using_onemkl_interfaces()
+
+    return flag and not onemkl_interfaces
+
+
 def _shape_error(shape1, shape2, func, err_msg):
     """Validate the shapes of input and output arrays."""
 
@@ -1008,21 +1030,18 @@ def dpnp_multiplication(
         elif x1.size == 0 or x2.size == 0:
             result.fill(0)
         else:
-            # TODO: replace with dpnp.int8 when it is added
-            x1_is_int8 = dpnp.issubdtype(x1.dtype, numpy.int8)
-            x2_is_int8 = dpnp.issubdtype(x2.dtype, numpy.int8)
-            res_is_int32 = dpnp.issubdtype(res_dtype, dpnp.int32)
-            special_case = x1_is_int8 and x2_is_int8 and res_is_int32
-            special_case = special_case and call_flag == "gemm"
-            if special_case:
-                # OneMath supports this special case
+            if _gemm_special_case(x1, x2, res_dtype, call_flag):
                 x1 = _copy_array(
                     x1, copy_flag=not x1_contig_flag, order=res_order
                 )
                 x2 = _copy_array(
                     x2, copy_flag=not x2_contig_flag, order=res_order
                 )
-                result = _gemm_matmul(exec_q, x1, x2, result)
+                if call_flag == "gemm":
+                    result = _gemm_matmul(exec_q, x1, x2, result)
+                else:
+                    assert call_flag == "gemm_batch"
+                    result = _gemm_batch_matmul(exec_q, x1, x2, result)
             elif dpnp.issubdtype(res_dtype, dpnp.inexact):
                 # copying is needed if dtypes of input arrays are different or
                 # their base (last 2-dimensions) is not c-contiguous or f-contiguous

dpnp/tests/test_product.py

@@ -834,7 +834,7 @@ def test_dtype_matrix(self, dt_in1, dt_in2, dt_out, shape1, shape2):
             # while NumPy give slightly different results. NumPy result obtained
             # with `dtype` is much closer to dpnp (a smaller `tol`) while the
             # result obtained with `out` needs a larger `tol` to match dpnp
-            assert_allclose(result, expected, rtol=1e-6, atol=1e-6)
+            assert_allclose(result, expected, rtol=1e-5, atol=1e-5)
         else:
             assert_raises(TypeError, dpnp.matmul, ia, ib, dtype=dt_out)
             assert_raises(TypeError, numpy.matmul, a, b, dtype=dt_out)
@@ -1151,23 +1151,29 @@ def test_large_values(self, dtype):
         expected = numpy.matmul(a, b)
         assert_dtype_allclose(result, expected)
 
-    def test_special_case(self):
+    @pytest.mark.parametrize("dt_out", [numpy.int32, numpy.float32])
+    @pytest.mark.parametrize(
+        "shape1, shape2",
+        [((2, 4), (4, 3)), ((4, 2, 3), (4, 3, 5))],
+        ids=["gemm", "gemm_batch"],
+    )
+    def test_special_case(self, dt_out, shape1, shape2):
         # Although inputs are int, gemm will be used for calculation
-        a = numpy.ones((3, 4), dtype=numpy.int8)
-        b = numpy.ones((4, 5), dtype=numpy.int8)
+        a = numpy.ones(shape1, dtype=numpy.int8)
+        b = numpy.ones(shape2, dtype=numpy.int8)
         ia, ib = dpnp.array(a), dpnp.array(b)
 
-        result = dpnp.matmul(ia, ib, dtype=numpy.int32)
-        expected = numpy.matmul(a, b, dtype=numpy.int32)
+        result = dpnp.matmul(ia, ib, dtype=dt_out)
+        expected = numpy.matmul(a, b, dtype=dt_out)
         assert_dtype_allclose(result, expected)
 
-        iout = dpnp.empty((3, 5), dtype=numpy.int32)
+        iout = dpnp.empty(result.shape, dtype=dt_out)
         result = dpnp.matmul(ia, ib, out=iout)
         assert_dtype_allclose(result, expected)
 
     def test_bool(self):
-        a = numpy.ones((3, 4), dtype=numpy.bool)
-        b = numpy.ones((4, 5), dtype=numpy.bool)
+        a = numpy.ones((3, 4), dtype=dpnp.bool)
+        b = numpy.ones((4, 5), dtype=dpnp.bool)
         ia, ib = dpnp.array(a), dpnp.array(b)
 
         # the output is (3, 4) array filled with 4