Improve pairwise_distance workloads

* Use 2D dispatch in kernel impls instead of huge sequential inner loop.
* Use nested prange in `numba_mlir_p` impl, `numba` and `numba-dpex` doesn't support nested pranges, but `numba-mlir` does.

Author: Hardcode84

dpbench/benchmarks/pairwise_distance/pairwise_distance_numba_dpex_k.py

@@ -9,16 +9,16 @@
 @dpex.kernel
 def _pairwise_distance_kernel(X1, X2, D):
     i = dpex.get_global_id(0)
+    j = dpex.get_global_id(1)
 
-    X2_rows = X2.shape[0]
     X1_cols = X1.shape[1]
-    for j in range(X2_rows):
-        d = X1.dtype.type(0.0)
-        for k in range(X1_cols):
-            tmp = X1[i, k] - X2[j, k]
-            d += tmp * tmp
-        D[i, j] = np.sqrt(d)
+
+    d = X1.dtype.type(0.0)
+    for k in range(X1_cols):
+        tmp = X1[i, k] - X2[j, k]
+        d += tmp * tmp
+    D[i, j] = np.sqrt(d)
 
 
 def pairwise_distance(X1, X2, D):
-    _pairwise_distance_kernel[X1.shape[0],](X1, X2, D)
+    _pairwise_distance_kernel[(X1.shape[0], X2.shape[0]),](X1, X2, D)

dpbench/benchmarks/pairwise_distance/pairwise_distance_numba_mlir_k.py

@@ -9,16 +9,18 @@
 @nb.kernel(gpu_fp64_truncate="auto")
 def _pairwise_distance_kernel(X1, X2, D):
     i = nb.get_global_id(0)
+    j = nb.get_global_id(1)
 
-    X2_rows = X2.shape[0]
     X1_cols = X1.shape[1]
-    for j in range(X2_rows):
-        d = 0.0
-        for k in range(X1_cols):
-            tmp = X1[i, k] - X2[j, k]
-            d += tmp * tmp
-        D[i, j] = np.sqrt(d)
+
+    d = 0.0
+    for k in range(X1_cols):
+        tmp = X1[i, k] - X2[j, k]
+        d += tmp * tmp
+    D[i, j] = np.sqrt(d)
 
 
 def pairwise_distance(X1, X2, D):
-    _pairwise_distance_kernel[X1.shape[0], nb.DEFAULT_LOCAL_SIZE](X1, X2, D)
+    _pairwise_distance_kernel[
+        (X1.shape[0], X2.shape[0]), nb.DEFAULT_LOCAL_SIZE
+    ](X1, X2, D)

dpbench/benchmarks/pairwise_distance/pairwise_distance_numba_mlir_p.py

@@ -25,7 +25,7 @@ def _pairwise_distance(X1, X2, D):
     # Outermost parallel loop over the matrix X1
     for i in numba.prange(X1_rows):
         # Loop over the matrix X2
-        for j in range(X2_rows):
+        for j in numba.prange(X2_rows):
             d = 0.0
             # Compute exclidean distance
             for k in range(X1_cols):