From 81ae3b7b7e971d1b82560834a581ff9fb67a6a2e Mon Sep 17 00:00:00 2001 From: Alexander Kalistratov Date: Fri, 24 Nov 2023 01:41:51 +0100 Subject: [PATCH] Re-enable kmeans and add dbscan and gpairs for numba_mlir --- .github/workflows/build_and_run.yml | 1 - .../benchmarks/dbscan/dbscan_numba_mlir_k.py | 154 ++++++++++++++++ .../benchmarks/dbscan/dbscan_numba_mlir_p.py | 151 ++++++++++++++++ .../benchmarks/gpairs/gpairs_numba_mlir_k.py | 165 ++++++++++++++++++ .../benchmarks/gpairs/gpairs_numba_mlir_p.py | 68 ++++++++ .../benchmarks/kmeans/kmeans_numba_mlir_k.py | 8 +- dpbench/configs/bench_info/kmeans.toml | 2 - .../frameworks/numba_mlir_framework.py | 8 +- 8 files changed, 549 insertions(+), 8 deletions(-) create mode 100644 dpbench/benchmarks/dbscan/dbscan_numba_mlir_k.py create mode 100644 dpbench/benchmarks/dbscan/dbscan_numba_mlir_p.py create mode 100644 dpbench/benchmarks/gpairs/gpairs_numba_mlir_k.py create mode 100644 dpbench/benchmarks/gpairs/gpairs_numba_mlir_p.py diff --git a/.github/workflows/build_and_run.yml b/.github/workflows/build_and_run.yml index 4e492634..c60995de 100644 --- a/.github/workflows/build_and_run.yml +++ b/.github/workflows/build_and_run.yml @@ -68,7 +68,6 @@ jobs: shell: bash -l {0} run: | find ./environments -type f | xargs sed -i 's/intel::numpy/numpy/' - find ./environments -type f | xargs sed -i '/numba-mlir/d' find ./environments -type f | xargs sed -i 's/setuptools>=42,<64/setuptools/' - name: Setup miniconda diff --git a/dpbench/benchmarks/dbscan/dbscan_numba_mlir_k.py b/dpbench/benchmarks/dbscan/dbscan_numba_mlir_k.py new file mode 100644 index 00000000..f417d626 --- /dev/null +++ b/dpbench/benchmarks/dbscan/dbscan_numba_mlir_k.py @@ -0,0 +1,154 @@ +# SPDX-FileCopyrightText: 2022 - 2023 Intel Corporation +# +# SPDX-License-Identifier: Apache-2.0 + +import dpctl.tensor as np +import numba_mlir.kernel as nb +import numpy +from numba_mlir import njit + +NOISE = -1 +UNDEFINED = -2 +DEFAULT_QUEUE_CAPACITY = 10 + + +@njit +def _queue_create(capacity): + return (numpy.empty(capacity, dtype=numpy.int64), 0, 0) + + +@njit +def _queue_resize(qu, tail, new_capacity): + tail = min(tail, new_capacity) + + new_qu = numpy.empty(new_capacity, dtype=numpy.int64) + new_qu[:tail] = qu[:tail] + return new_qu, tail + + +@njit +def _queue_push(qu, value, tail, capacity): + if tail == capacity: + capacity = 2 * capacity + qu, tail = _queue_resize(qu, tail, capacity) + + qu[tail] = value + tail += 1 + + return qu, tail, capacity + + +@njit +def _queue_pop(qu, head, tail): + head += 1 + return qu[head - 1], head + + +@njit +def _queue_empty(head, tail): + return head == tail + + +@nb.kernel +def get_neighborhood(n, dim, data, eps, ind_lst, sz_lst, block_size, nblocks): + i = nb.get_global_id(0) + + start = i * block_size + stop = n if i + 1 == nblocks else start + block_size + eps2 = eps * eps + block_size1 = 256 + nblocks1 = n // block_size1 + int(n % block_size1 > 0) + for ii in range(nblocks1): + i1 = ii * block_size1 + i2 = n if ii + 1 == nblocks1 else i1 + block_size1 + for j in range(start, stop): + for k in range(i1, i2): + dist = data.dtype.type(0.0) + for m in range(dim): + diff = data[k * dim + m] - data[j * dim + m] + dist += diff * diff + if dist <= eps2: + size = sz_lst[j] + ind_lst[j * n + size] = k + sz_lst[j] = size + 1 + + +@njit(parallel=False, fastmath=True) +def compute_clusters( # noqa: C901: TODO: can we simplify logic? + n, min_pts, assignments, sizes, indices_list +): + nclusters = 0 + nnoise = 0 + for i in range(n): + if assignments[i] != UNDEFINED: + continue + size = sizes[i] + if size < min_pts: + nnoise += 1 + assignments[i] = NOISE + continue + nclusters += 1 + assignments[i] = nclusters - 1 + + qu_capacity = DEFAULT_QUEUE_CAPACITY + qu, head, tail = _queue_create(qu_capacity) + for j in range(size): + next_point = indices_list[i * n + j] + if assignments[next_point] == NOISE: + nnoise -= 1 + assignments[next_point] = nclusters - 1 + elif assignments[next_point] == UNDEFINED: + assignments[next_point] = nclusters - 1 + qu, tail, qu_capacity = _queue_push( + qu, next_point, tail, qu_capacity + ) + + while not _queue_empty(head, tail): + cur_point, head = _queue_pop(qu, head, tail) + size = sizes[cur_point] + assignments[cur_point] = nclusters - 1 + if size < min_pts: + continue + + for j in range(size): + next_point = indices_list[cur_point * n + j] + if assignments[next_point] == NOISE: + nnoise -= 1 + assignments[next_point] = nclusters - 1 + elif assignments[next_point] == UNDEFINED: + assignments[next_point] = nclusters - 1 + qu, tail, qu_capacity = _queue_push( + qu, next_point, tail, qu_capacity + ) + + return nclusters + + +def dbscan(n_samples, n_features, data, eps, min_pts): + indices_list = np.empty( + shape=n_samples * n_samples, dtype=np.int64, device=data.device + ) + sizes = np.zeros(shape=n_samples, dtype=np.int64, device=data.device) + + get_neighborhood[n_samples,]( + n_samples, + n_features, + data, + eps, + indices_list, + sizes, + 1, + n_samples, + ) + + assignments = numpy.empty(n_samples, dtype=numpy.int64) + for i in range(n_samples): + assignments[i] = UNDEFINED + + return compute_clusters( + n_samples, + min_pts, + assignments, + np.asnumpy(sizes), + np.asnumpy(indices_list), + ) diff --git a/dpbench/benchmarks/dbscan/dbscan_numba_mlir_p.py b/dpbench/benchmarks/dbscan/dbscan_numba_mlir_p.py new file mode 100644 index 00000000..b66741ba --- /dev/null +++ b/dpbench/benchmarks/dbscan/dbscan_numba_mlir_p.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: 2022 - 2023 Intel Corporation +# +# SPDX-License-Identifier: Apache-2.0 + +import dpctl.tensor as np +import numba as nb +import numpy +from numba_mlir import njit + +NOISE = -1 +UNDEFINED = -2 +DEFAULT_QUEUE_CAPACITY = 10 + + +@njit() +def _queue_create(capacity): + return (numpy.empty(capacity, dtype=numpy.int64), 0, 0) + + +@njit() +def _queue_resize(qu, tail, new_capacity): + tail = min(tail, new_capacity) + + new_qu = numpy.empty(new_capacity, dtype=numpy.int64) + new_qu[:tail] = qu[:tail] + return new_qu, tail + + +@njit() +def _queue_push(qu, value, tail, capacity): + if tail == capacity: + capacity = 2 * capacity + qu, tail = _queue_resize(qu, tail, capacity) + + qu[tail] = value + tail += 1 + + return qu, tail, capacity + + +@njit() +def _queue_pop(qu, head, tail): + head += 1 + return qu[head - 1], head + + +@njit() +def _queue_empty(head, tail): + return head == tail + + +@njit +def get_neighborhood(n, dim, data, eps, ind_lst, sz_lst): + block_size = 1 + nblocks = n // block_size + int(n % block_size > 0) + + # TODO: get rid of it once prange supports dtype + # https://github.com/IntelPython/numba-dpex/issues/1063 + float0 = data.dtype.type(0.0) + + for i in nb.prange(nblocks): + start = i * block_size + stop = n if i + 1 == nblocks else start + block_size + eps2 = eps * eps + block_size1 = 256 + nblocks1 = n // block_size1 + int(n % block_size1 > 0) + for ii in range(nblocks1): + i1 = ii * block_size1 + i2 = n if ii + 1 == nblocks1 else i1 + block_size1 + for j in range(start, stop): + for k in range(i1, i2): + dist = float0 + for m in range(dim): + diff = data[k * dim + m] - data[j * dim + m] + dist += diff * diff + if dist <= eps2: + size = sz_lst[j] + ind_lst[j * n + size] = k + sz_lst[j] = size + 1 + + +@njit(parallel=False, fastmath=True) +def compute_clusters( # noqa: C901: TODO: can we simplify logic? + n, min_pts, assignments, sizes, indices_list +): + nclusters = 0 + nnoise = 0 + for i in range(n): + if assignments[i] != UNDEFINED: + continue + size = sizes[i] + if size < min_pts: + nnoise += 1 + assignments[i] = NOISE + continue + nclusters += 1 + assignments[i] = nclusters - 1 + + qu_capacity = DEFAULT_QUEUE_CAPACITY + qu, head, tail = _queue_create(qu_capacity) + for j in range(size): + next_point = indices_list[i * n + j] + if assignments[next_point] == NOISE: + nnoise -= 1 + assignments[next_point] = nclusters - 1 + elif assignments[next_point] == UNDEFINED: + assignments[next_point] = nclusters - 1 + qu, tail, qu_capacity = _queue_push( + qu, next_point, tail, qu_capacity + ) + + while not _queue_empty(head, tail): + cur_point, head = _queue_pop(qu, head, tail) + size = sizes[cur_point] + assignments[cur_point] = nclusters - 1 + if size < min_pts: + continue + + for j in range(size): + next_point = indices_list[cur_point * n + j] + if assignments[next_point] == NOISE: + nnoise -= 1 + assignments[next_point] = nclusters - 1 + elif assignments[next_point] == UNDEFINED: + assignments[next_point] = nclusters - 1 + qu, tail, qu_capacity = _queue_push( + qu, next_point, tail, qu_capacity + ) + + return nclusters + + +def dbscan(n_samples, n_features, data, eps, min_pts): + indices_list = np.empty( + shape=n_samples * n_samples, dtype=np.int64, device=data.device + ) + sizes = np.zeros(shape=n_samples, dtype=np.int64, device=data.device) + + get_neighborhood(n_samples, n_features, data, eps, indices_list, sizes) + + assignments = numpy.empty(n_samples, dtype=np.int64) + for i in range(n_samples): + assignments[i] = UNDEFINED + + return compute_clusters( + n_samples, + min_pts, + assignments, + np.asnumpy(sizes), + np.asnumpy(indices_list), + ) diff --git a/dpbench/benchmarks/gpairs/gpairs_numba_mlir_k.py b/dpbench/benchmarks/gpairs/gpairs_numba_mlir_k.py new file mode 100644 index 00000000..508bbbd5 --- /dev/null +++ b/dpbench/benchmarks/gpairs/gpairs_numba_mlir_k.py @@ -0,0 +1,165 @@ +# SPDX-FileCopyrightText: 2022 - 2023 Intel Corporation +# +# SPDX-License-Identifier: Apache-2.0 + +import numba_dpex as dpex + +# This implementation is numba dpex kernel version with atomics. + + +@dpex.kernel +def count_weighted_pairs_3d_intel_no_slm_ker( + n, + nbins, + slm_hist_size, + private_hist_size, + x0, + y0, + z0, + w0, + x1, + y1, + z1, + w1, + rbins_squared, + result, +): + dtype = x0.dtype + lid0 = dpex.get_local_id(0) + gr0 = dpex.get_group_id(0) + + lid1 = dpex.get_local_id(1) + gr1 = dpex.get_group_id(1) + + lws0 = dpex.get_local_size(0) + lws1 = dpex.get_local_size(1) + + n_wi = 20 + + dsq_mat = dpex.private.array(shape=(20 * 20), dtype=dtype) + w0_vec = dpex.private.array(shape=(20), dtype=dtype) + w1_vec = dpex.private.array(shape=(20), dtype=dtype) + + offset0 = gr0 * n_wi * lws0 + lid0 + offset1 = gr1 * n_wi * lws1 + lid1 + + # work item works on pointer + # j0 = gr0 * n_wi * lws0 + i0 * lws0 + lid0, and + # j1 = gr1 * n_wi * lws1 + i1 * lws1 + lid1 + + j1 = offset1 + i1 = 0 + while (i1 < n_wi) and (j1 < n): + w1_vec[i1] = w1[j1] + i1 += 1 + j1 += lws1 + + # compute (n_wi, n_wi) matrix of squared distances in work-item + j0 = offset0 + i0 = 0 + while (i0 < n_wi) and (j0 < n): + x0v = x0[j0] + y0v = y0[j0] + z0v = z0[j0] + w0_vec[i0] = w0[j0] + + j1 = offset1 + i1 = 0 + while (i1 < n_wi) and (j1 < n): + dx = x0v - x1[j1] + dy = y0v - y1[j1] + dz = z0v - z1[j1] + dsq_mat[i0 * n_wi + i1] = dx * dx + dy * dy + dz * dz + i1 += 1 + j1 += lws1 + + i0 += 1 + j0 += lws0 + + # update slm_hist. Use work-item private buffer of 16 tfloat elements + for k in range(0, slm_hist_size, private_hist_size): + private_hist = dpex.private.array(shape=(16), dtype=dtype) + for p in range(private_hist_size): + private_hist[p] = 0.0 + + j0 = offset0 + i0 = 0 + while (i0 < n_wi) and (j0 < n): + j1 = offset1 + i1 = 0 + while (i1 < n_wi) and (j1 < n): + dsq = dsq_mat[i0 * n_wi + i1] + pw = w0_vec[i0] * w1_vec[i1] + pk = k + for p in range(private_hist_size): + private_hist[p] += ( + pw + if (pk < nbins and dsq <= rbins_squared[pk]) + else dtype.type(0.0) + ) + pk += 1 + + i1 += 1 + j1 += lws1 + + i0 += 1 + j0 += lws0 + + pk = k + for p in range(private_hist_size): + dpex.atomic.add(result, pk, private_hist[p]) + pk += 1 + + +def ceiling_quotient(n, m): + return int((n + m - 1) / m) + + +def gpairs( + nopt, + nbins, + x1, + y1, + z1, + w1, + x2, + y2, + z2, + w2, + rbins, + results, +): + n_wi = 20 + private_hist_size = 16 + lws0 = 16 + lws1 = 16 + + m0 = n_wi * lws0 + m1 = n_wi * lws1 + + n_groups0 = ceiling_quotient(nopt, m0) + n_groups1 = ceiling_quotient(nopt, m1) + + gwsRange = n_groups0 * lws0, n_groups1 * lws1 + lwsRange = lws0, lws1 + + slm_hist_size = ( + ceiling_quotient(nbins, private_hist_size) * private_hist_size + ) + + count_weighted_pairs_3d_intel_no_slm_ker[gwsRange, lwsRange]( + nopt, + nbins, + slm_hist_size, + private_hist_size, + x1, + y1, + z1, + w1, + x2, + y2, + z2, + w2, + rbins, + results, + ) diff --git a/dpbench/benchmarks/gpairs/gpairs_numba_mlir_p.py b/dpbench/benchmarks/gpairs/gpairs_numba_mlir_p.py new file mode 100644 index 00000000..4889903b --- /dev/null +++ b/dpbench/benchmarks/gpairs/gpairs_numba_mlir_p.py @@ -0,0 +1,68 @@ +# SPDX-FileCopyrightText: 2022 - 2023 Intel Corporation +# +# SPDX-License-Identifier: Apache-2.0 + +import dpctl.tensor as np +import numba as nb +from numba_mlir import njit + +# This implementation is numba dpex prange version without atomics. + + +@njit +def count_weighted_pairs_3d_diff_ker( + n, nbins, x1, y1, z1, w1, x2, y2, z2, w2, rbins_squared, result +): + for i in nb.prange(n): + px = x1[i] + py = y1[i] + pz = z1[i] + pw = w1[i] + for j in range(n): + qx = x2[j] + qy = y2[j] + qz = z2[j] + qw = w2[j] + dx = px - qx + dy = py - qy + dz = pz - qz + wprod = pw * qw + dsq = dx * dx + dy * dy + dz * dz + + if dsq <= rbins_squared[nbins - 1]: + for k in range(nbins - 1, -1, -1): + if dsq > rbins_squared[k]: + result[i, k + 1] += wprod + break + if k == 0: + result[i, k] += wprod + break + + for j in range(nbins - 2, -1, -1): + for k in range(j + 1, nbins, 1): + result[i, k] += result[i, j] + + +@njit +def count_weighted_pairs_3d_diff_agg_ker(nbins, result, n): + for col_id in nb.prange(nbins): + for i in range(1, n): + result[0, col_id] += result[i, col_id] + + +def gpairs(nopt, nbins, x1, y1, z1, w1, x2, y2, z2, w2, rbins, results): + # allocate per-work item private result vector in device global memory + results_disjoint = np.zeros( + shape=(nopt, rbins.shape[0]), device=results.device + ) + + # call gpairs compute kernel + count_weighted_pairs_3d_diff_ker( + nopt, nbins, x1, y1, z1, w1, x2, y2, z2, w2, rbins, results_disjoint + ) + + # aggregate the results from the compute kernel + count_weighted_pairs_3d_diff_agg_ker(nbins, results_disjoint, nopt) + + # copy to results vector + results[:] = results_disjoint[0] diff --git a/dpbench/benchmarks/kmeans/kmeans_numba_mlir_k.py b/dpbench/benchmarks/kmeans/kmeans_numba_mlir_k.py index de82d801..a36f0471 100644 --- a/dpbench/benchmarks/kmeans/kmeans_numba_mlir_k.py +++ b/dpbench/benchmarks/kmeans/kmeans_numba_mlir_k.py @@ -56,7 +56,7 @@ def groupByCluster( for lc in range(local_copies): localNewCount[lc, c] = 0 - nb.barrier(nb.CLK_LOCAL_MEM_FENCE) + nb.barrier(nb.LOCAL_MEM_FENCE) for i in range(WorkPI): point_id = grid * WorkPI * local_size + i * local_size + lid @@ -87,7 +87,7 @@ def groupByCluster( if last: arrayPcluster[point_id] = nearest_centroid - nb.barrier(nb.CLK_LOCAL_MEM_FENCE) + nb.barrier(nb.LOCAL_MEM_FENCE) for i in range(lid, num_centroids * dims, local_size): local_centroid_d = dtyp.type(0) @@ -139,7 +139,7 @@ def updateCentroids(diff, arrayC, arrayCnumpoint, NewCentroids, NewCount): max_distance = max(max_distance, distance) local_distance[c] = max_distance - nb.barrier(nb.CLK_LOCAL_MEM_FENCE) + nb.barrier(nb.LOCAL_MEM_FENCE) if lid == 0: max_distance = group_reduce_max(local_distance[c]) @@ -163,7 +163,7 @@ def updateLabels(arrayP, arrayPcluster, arrayC): for i in range(lid, num_centroids * dims, local_size): localCentroids[i % dims, i // dims] = arrayC[i // dims, i % dims] - nb.barrier(nb.CLK_LOCAL_MEM_FENCE) + nb.barrier(nb.LOCAL_MEM_FENCE) for i in range(WorkPI): point_id = grid * WorkPI * local_size + i * local_size + lid diff --git a/dpbench/configs/bench_info/kmeans.toml b/dpbench/configs/bench_info/kmeans.toml index 1fea6d20..9790e04f 100644 --- a/dpbench/configs/bench_info/kmeans.toml +++ b/dpbench/configs/bench_info/kmeans.toml @@ -28,8 +28,6 @@ output_args = [ "arrayC", "arrayCnumpoint", ] -# TODO: remove once fixed. Fails randomly -expected_failure_implementations = ["numba_mlir_k"] [benchmark.parameters.S] npoints = 4096 diff --git a/dpbench/infrastructure/frameworks/numba_mlir_framework.py b/dpbench/infrastructure/frameworks/numba_mlir_framework.py index c2a87170..6fff19b4 100644 --- a/dpbench/infrastructure/frameworks/numba_mlir_framework.py +++ b/dpbench/infrastructure/frameworks/numba_mlir_framework.py @@ -65,6 +65,12 @@ def copy_from_func(self) -> Callable: if self.sycl_device: import dpctl.tensor as dpt - return dpt.asnumpy + def cpy(val): + if isinstance(val, dpt.usm_ndarray): + return dpt.asnumpy(val) + else: + return np.asarray(val) + + return cpy else: return np.copy