<commit_before>"""
Benchmarks for the CUDA backend.
"""

from __future__ import division

import math
import sys

import numpy as np

from numba import cuda, float32, float64


def addmul(x, y, out):
    i = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    if i >= x.shape[0]:
        return
    out[i] = x[i] + y[i] * math.fabs(x[i])

addmul_f32 = cuda.jit(argtypes=(float32[:], float32[:], float32[:]))(addmul)
addmul_f64 = cuda.jit(argtypes=(float64[:], float64[:], float64[:]))(addmul)


@cuda.jit(argtypes=())
def no_op():
    pass


# N-body simulation.  We actually only run the step which computes the
# accelerations from the positions and weights of the bodies (updating
# speeds and positions is relatively uninteresting).

# CUDA version adapted from http://http.developer.nvidia.com/GPUGems3/gpugems3_ch31.html

eps_2 = np.float32(1e-6)
zero = np.float32(0.0)
one = np.float32(1.0)

@cuda.jit(argtypes=(float32, float32,
                    float32, float32, float32,
                    float32, float32),
          device=True, inline=True)
def body_body_interaction(xi, yi, xj, yj, wj, axi, ayi):
    """
    Compute the influence of body j on the acceleration of body i.
    """
    rx = xj - xi
    ry = yj - yi
    sqr_dist = rx * rx + ry * ry + eps_2
    sixth_dist = sqr_dist * sqr_dist * sqr_dist
    inv_dist_cube = one / math.sqrt(sixth_dist)
    s = wj * inv_dist_cube
    axi += rx * s
    ayi += ry * s
    return axi, ayi

@cuda.jit(argtypes=(float32, float32, float32, float32,
                    float32[:,:], float32[:]), device=True, inline=True)
def tile_calculation(xi, yi, axi, ayi, positions, weights):
    """
    Compute the contribution of this block's tile to the acceleration
    of body i.
    """
    for j in range(cuda.blockDim.x):
        xj = positions[j,0]
        yj = positions[j,1]
        wj = weights[j]
        axi, ayi = body_body_interaction(xi, yi, xj, yj, wj, axi, ayi)
    return axi, ayi


tile_size = 128

@cuda.jit(argtypes=(float32[:,:], float32[:], float32[:,:]))
def calculate_forces(positions, weights, accelerations):
    """
    Calculate accelerations produced on all bodies by mutual gravitational
    forces.
    """
    sh_positions = cuda.shared.array((tile_size, 2), float32)
    sh_weights = cuda.shared.array(tile_size, float32)
    i = cuda.grid(1)
    axi = 0.0
    ayi = 0.0
    xi = positions[i,0]
    yi = positions[i,1]
    for j in range(0, len(weights), tile_size):
        index = (j // tile_size) * cuda.blockDim.x + cuda.threadIdx.x
        sh_index = cuda.threadIdx.x
        sh_positions[sh_index,0] = positions[index,0]
        sh_positions[sh_index,1] = positions[index,1]
        sh_weights[sh_index] = weights[index]
        cuda.syncthreads()
        axi, ayi = tile_calculation(xi, yi, axi, ayi,
                                    sh_positions, sh_weights)
        cuda.syncthreads()
    accelerations[i,0] = axi
    accelerations[i,1] = ayi


class NBodyCUDARunner:

    def __init__(self, positions, weights):
        self.accelerations = np.zeros_like(positions)
        self.n_bodies = len(weights)
        self.stream = cuda.stream()
        self.d_pos = cuda.to_device(positions, self.stream)
        self.d_wei = cuda.to_device(weights, self.stream)
        self.d_acc = cuda.to_device(self.accelerations, self.stream)
        self.stream.synchronize()

    def run(self):
        blockdim = tile_size
        griddim = int(math.ceil(self.n_bodies / blockdim))
        calculate_forces[griddim, blockdim, self.stream](
            self.d_pos, self.d_wei, self.d_acc)
        self.stream.synchronize()

    def results(self):
        self.d_acc.copy_to_host(self.accelerations, self.stream)
        self.stream.synchronize()
        return self.accelerations


def run_cpu_nbody(positions, weights):
    accelerations = np.zeros_like(positions)
    n = weights.size
    for j in range(n):
        # Compute influence of j'th body on all bodies
        r = positions[j] - positions
        rx = r[:,0]
        ry = r[:,1]
        sqr_dist = rx * rx + ry * ry + eps_2
        sixth_dist = sqr_dist * sqr_dist * sqr_dist
        inv_dist_cube = one / np.sqrt(sixth_dist)
        s = weights[j] * inv_dist_cube
        accelerations += (r.transpose() * s).transpose()
    return accelerations

def make_nbody_samples(n_bodies):
    positions = np.random.RandomState(0).uniform(-1.0, 1.0, (n_bodies, 2))
    weights = np.random.RandomState(0).uniform(1.0, 2.0, n_bodies)
    return positions.astype(np.float32), weights.astype(np.float32)

<commit_msg>Fix bench_cuda to run partially despite the macro expansion bug on recent Numbas<commit_after>"""
Benchmarks for the CUDA backend.
"""

from __future__ import division

import math
import sys

import numpy as np

from numba import cuda, float32, float64


def addmul(x, y, out):
    i = cuda.threadIdx.x + cuda.blockIdx.x * cuda.blockDim.x
    if i >= x.shape[0]:
        return
    out[i] = x[i] + y[i] * math.fabs(x[i])

addmul_f32 = cuda.jit(argtypes=(float32[:], float32[:], float32[:]))(addmul)
addmul_f64 = cuda.jit(argtypes=(float64[:], float64[:], float64[:]))(addmul)


@cuda.jit(argtypes=())
def no_op():
    pass


# N-body simulation.  We actually only run the step which computes the
# accelerations from the positions and weights of the bodies (updating
# speeds and positions is relatively uninteresting).

# CUDA version adapted from http://http.developer.nvidia.com/GPUGems3/gpugems3_ch31.html

eps_2 = np.float32(1e-6)
zero = np.float32(0.0)
one = np.float32(1.0)

@cuda.jit(argtypes=(float32, float32,
                    float32, float32, float32,
                    float32, float32),
          device=True, inline=True)
def body_body_interaction(xi, yi, xj, yj, wj, axi, ayi):
    """
    Compute the influence of body j on the acceleration of body i.
    """
    rx = xj - xi
    ry = yj - yi
    sqr_dist = rx * rx + ry * ry + eps_2
    sixth_dist = sqr_dist * sqr_dist * sqr_dist
    inv_dist_cube = one / math.sqrt(sixth_dist)
    s = wj * inv_dist_cube
    axi += rx * s
    ayi += ry * s
    return axi, ayi

@cuda.jit(argtypes=(float32, float32, float32, float32,
                    float32[:,:], float32[:]), device=True, inline=True)
def tile_calculation(xi, yi, axi, ayi, positions, weights):
    """
    Compute the contribution of this block's tile to the acceleration
    of body i.
    """
    for j in range(cuda.blockDim.x):
        xj = positions[j,0]
        yj = positions[j,1]
        wj = weights[j]
        axi, ayi = body_body_interaction(xi, yi, xj, yj, wj, axi, ayi)
    return axi, ayi


tile_size = 128

# Don't JIT this function at the top-level as it breaks until Numba 0.16.
def calculate_forces(positions, weights, accelerations):
    """
    Calculate accelerations produced on all bodies by mutual gravitational
    forces.
    """
    sh_positions = cuda.shared.array((tile_size, 2), float32)
    sh_weights = cuda.shared.array(tile_size, float32)
    i = cuda.grid(1)
    axi = 0.0
    ayi = 0.0
    xi = positions[i,0]
    yi = positions[i,1]
    for j in range(0, len(weights), tile_size):
        index = (j // tile_size) * cuda.blockDim.x + cuda.threadIdx.x
        sh_index = cuda.threadIdx.x
        sh_positions[sh_index,0] = positions[index,0]
        sh_positions[sh_index,1] = positions[index,1]
        sh_weights[sh_index] = weights[index]
        cuda.syncthreads()
        axi, ayi = tile_calculation(xi, yi, axi, ayi,
                                    sh_positions, sh_weights)
        cuda.syncthreads()
    accelerations[i,0] = axi
    accelerations[i,1] = ayi


class NBodyCUDARunner:

    def __init__(self, positions, weights):
        self.calculate_forces = cuda.jit(
            argtypes=(float32[:,:], float32[:], float32[:,:])
            )(calculate_forces)
        self.accelerations = np.zeros_like(positions)
        self.n_bodies = len(weights)
        self.stream = cuda.stream()
        self.d_pos = cuda.to_device(positions, self.stream)
        self.d_wei = cuda.to_device(weights, self.stream)
        self.d_acc = cuda.to_device(self.accelerations, self.stream)
        self.stream.synchronize()

    def run(self):
        blockdim = tile_size
        griddim = int(math.ceil(self.n_bodies / blockdim))
        self.calculate_forces[griddim, blockdim, self.stream](
            self.d_pos, self.d_wei, self.d_acc)
        self.stream.synchronize()

    def results(self):
        self.d_acc.copy_to_host(self.accelerations, self.stream)
        self.stream.synchronize()
        return self.accelerations


def run_cpu_nbody(positions, weights):
    accelerations = np.zeros_like(positions)
    n = weights.size
    for j in range(n):
        # Compute influence of j'th body on all bodies
        r = positions[j] - positions
        rx = r[:,0]
        ry = r[:,1]
        sqr_dist = rx * rx + ry * ry + eps_2
        sixth_dist = sqr_dist * sqr_dist * sqr_dist
        inv_dist_cube = one / np.sqrt(sixth_dist)
        s = weights[j] * inv_dist_cube
        accelerations += (r.transpose() * s).transpose()
    return accelerations

def make_nbody_samples(n_bodies):
    positions = np.random.RandomState(0).uniform(-1.0, 1.0, (n_bodies, 2))
    weights = np.random.RandomState(0).uniform(1.0, 2.0, n_bodies)
    return positions.astype(np.float32), weights.astype(np.float32)


{'fraction_non_alphanumeric': 0.12305879832922781, 'fraction_numerical': 0.024526079040376996, 'mean_word_length': 2.777508090614887, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6477110', 'n_tokens_mistral': 3512, 'n_tokens_neox': 3325, 'n_words': 774}