import numpy as np

from .common import Benchmark, get_data

UFUNCS = [obj for obj in np._core.umath.__dict__.values() if

isinstance(obj, np.ufunc)]

UFUNCS_UNARY = [uf for uf in UFUNCS if "O->O" in uf.types]

class _AbstractBinary(Benchmark):

params = []

param_names = ['ufunc', 'stride_in0', 'stride_in1', 'stride_out', 'dtype']

timeout = 10

arrlen = 1000000

data_finite = True

data_denormal = False

data_zeros = False

def setup(self, ufunc, stride_in0, stride_in1, stride_out, dtype):

ufunc_insig = f'{dtype}{dtype}->'

if ufunc_insig + dtype not in ufunc.types:

for st_sig in (ufunc_insig, dtype):

test = [sig for sig in ufunc.types if sig.startswith(st_sig)]

if test:

break

if not test:

raise NotImplementedError(

f"Ufunc {ufunc} doesn't support "

f"binary input of dtype {dtype}"

) from None

tin, tout = test[0].split('->')

else:

tin = dtype + dtype

tout = dtype

self.ufunc_args = []

for i, (dt, stride) in enumerate(zip(tin, (stride_in0, stride_in1))):

self.ufunc_args += [get_data(

self.arrlen * stride, dt, i,

zeros=self.data_zeros,

finite=self.data_finite,

denormal=self.data_denormal,

)[::stride]]

for dt in tout:

self.ufunc_args += [

np.empty(stride_out * self.arrlen, dt)[::stride_out]

]

np.seterr(all='ignore')

def time_binary(self, ufunc, stride_in0, stride_in1, stride_out,

dtype):

ufunc(*self.ufunc_args)

def time_binary_scalar_in0(self, ufunc, stride_in0, stride_in1,

stride_out, dtype):

ufunc(self.ufunc_args[0][0], *self.ufunc_args[1:])

def time_binary_scalar_in1(self, ufunc, stride_in0, stride_in1,

stride_out, dtype):

ufunc(self.ufunc_args[0], self.ufunc_args[1][0], *self.ufunc_args[2:])

class _AbstractUnary(Benchmark):

params = []

param_names = ['ufunc', 'stride_in', 'stride_out', 'dtype']

timeout = 10

arrlen = 1000000

data_finite = True

data_denormal = False

data_zeros = False

def setup(self, ufunc, stride_in, stride_out, dtype):

arr_in = get_data(

stride_in * self.arrlen, dtype,

zeros=self.data_zeros,

finite=self.data_finite,

denormal=self.data_denormal,

)

self.ufunc_args = [arr_in[::stride_in]]

ufunc_insig = f'{dtype}->'

if ufunc_insig + dtype not in ufunc.types:

test = [sig for sig in ufunc.types if sig.startswith(ufunc_insig)]

if not test:

raise NotImplementedError(

f"Ufunc {ufunc} doesn't support "

f"unary input of dtype {dtype}"

) from None

tout = test[0].split('->')[1]

else:

tout = dtype

for dt in tout:

self.ufunc_args += [

np.empty(stride_out * self.arrlen, dt)[::stride_out]

]

np.seterr(all='ignore')

def time_unary(self, ufunc, stride_in, stride_out, dtype):

ufunc(*self.ufunc_args)

class UnaryFP(_AbstractUnary):

params = [[uf for uf in UFUNCS_UNARY

if uf not in (np.invert, np.bitwise_count)],

[1, 4],

[1, 2],

['e', 'f', 'd']]

def setup(self, ufunc, stride_in, stride_out, dtype):

_AbstractUnary.setup(self, ufunc, stride_in, stride_out, dtype)

if (ufunc.__name__ == 'arccosh'):

self.ufunc_args[0] += 1.0

class UnaryFPSpecial(UnaryFP):

data_finite = False

data_denormal = True

data_zeros = True

class BinaryFP(_AbstractBinary):

params = [

[np.maximum, np.minimum, np.fmax, np.fmin, np.ldexp],

[1, 2], [1, 4], [1, 2, 4], ['f', 'd']

]

class BinaryFPSpecial(BinaryFP):

data_finite = False

data_denormal = True

data_zeros = True

class BinaryComplex(_AbstractBinary):

params = [

[np.add, np.subtract, np.multiply, np.divide],

[1, 2, 4], [1, 2, 4], [1, 2, 4],

['F', 'D']

]

class UnaryComplex(_AbstractUnary):

params = [

[np.reciprocal, np.absolute, np.square, np.conjugate],

[1, 2, 4], [1, 2, 4], ['F', 'D']

]

class BinaryInt(_AbstractBinary):

arrlen = 100000

params = [

[np.maximum, np.minimum],

[1, 2], [1, 2], [1, 2],

['b', 'B', 'h', 'H', 'i', 'I', 'l', 'L', 'q', 'Q']

]

class BinaryIntContig(_AbstractBinary):

params = [

[getattr(np, uf) for uf in (

'add', 'subtract', 'multiply', 'bitwise_and', 'bitwise_or',

'bitwise_xor', 'logical_and', 'logical_or', 'logical_xor',

'right_shift', 'left_shift',

)],

[1], [1], [1],

['b', 'B', 'h', 'H', 'i', 'I', 'l', 'L', 'q', 'Q']

]

class UnaryIntContig(_AbstractUnary):

arrlen = 100000

params = [

[getattr(np, uf) for uf in (

'positive', 'square', 'reciprocal', 'conjugate', 'logical_not',

'invert', 'isnan', 'isinf', 'isfinite',

'absolute', 'sign', 'bitwise_count'

)],

[1], [1],

['b', 'B', 'h', 'H', 'i', 'I', 'l', 'L', 'q', 'Q']

]

class Mandelbrot(Benchmark):

def f(self, z):

return np.abs(z) < 4.0

def g(self, z, c):

return np.sum(np.multiply(z, z) + c)

def mandelbrot_numpy(self, c, maxiter):

output = np.zeros(c.shape, np.int32)

z = np.empty(c.shape, np.complex64)

for it in range(maxiter):

notdone = self.f(z)

output[notdone] = it

z[notdone] = self.g(z[notdone], c[notdone])

output[output == maxiter - 1] = 0

return output

def mandelbrot_set(self, xmin, xmax, ymin, ymax, width, height, maxiter):

r1 = np.linspace(xmin, xmax, width, dtype=np.float32)

r2 = np.linspace(ymin, ymax, height, dtype=np.float32)

c = r1 + r2[:, None] * 1j

n3 = self.mandelbrot_numpy(c, maxiter)

return (r1, r2, n3.T)

def time_mandel(self):

self.mandelbrot_set(-0.74877, -0.74872, 0.06505, 0.06510, 1000, 1000, 2048)

class LogisticRegression(Benchmark):

param_names = ['dtype']

params = [np.float32, np.float64]

timeout = 1000

def train(self, max_epoch):

for epoch in range(max_epoch):

z = np.matmul(self.X_train, self.W)

A = 1 / (1 + np.exp(-z)) # sigmoid(z)

Y_train = self.Y_train

loss = -np.mean(Y_train * np.log(A) + (1 - Y_train) * np.log(1 - A))

dz = A - Y_train

dw = (1 / self.size) * np.matmul(self.X_train.T, dz)

self.W = self.W - self.alpha * dw

def setup(self, dtype):

np.random.seed(42)

self.size = 250

features = 16

self.X_train = np.random.rand(self.size, features).astype(dtype)

self.Y_train = np.random.choice(2, self.size).astype(dtype)

# Initialize weights

self.W = np.zeros((features, 1), dtype=dtype)

self.b = np.zeros((1, 1), dtype=dtype)

self.alpha = 0.1

def time_train(self, dtype):

self.train(1000)