def count_bins(i, x, y, z, h, eps, qmax, rmax, keys, bin_counts, sort_offsets): c = declare('matrix(3, "int")') find_cell_id(x[i], y[i], z[i], h, eps, c) key = flatten(c[0], c[1], c[2], qmax, rmax) keys[i] = key idx = atomic_inc(bin_counts[key]) sort_offsets[i] = idx
def map_dst_to_src(i, dst_to_src, cids_dst, cid_to_idx_dst, keys_dst, keys_src, cids_src, num_particles_src, max_cid_src): idx = cid_to_idx_dst[27 * i] key = keys_dst[idx] idx_src = find_idx(keys_src, num_particles_src, key) dst_to_src[i] = atomic_inc( max_cid_src[0]) if (idx_src == -1) else cids_src[idx_src]
def count_cells(i, keys, cell_size, cell_lengths, sort_offsets): key = keys[i] idx = atomic_inc(cell_lengths[key]) sort_offsets[i] = idx