def test_pow():
# note hat a**2 is converted to a*a and a**0.5 to sqrt(a)
@python2shader_and_validate
def compute_shader(
index_xyz: ("input", "GlobalInvocationId", ivec3),
data1: ("buffer", 0, Array(f32)),
data2: ("buffer", 1, Array(vec4)),
):
index = index_xyz.x
a = data1[index]
data2[index] = vec4(a**2, a**0.5, a**3.0, a**3.1)
skip_if_no_wgpu()
values1 = [i - 5 for i in range(10)]
inp_arrays = {0: (ctypes.c_float * 10)(*values1)}
out_arrays = {1: ctypes.c_float * 40}
out = compute_with_buffers(inp_arrays, out_arrays, compute_shader)
res = list(out[1])
assert res[0::4] == [i**2 for i in values1]
assert iters_close(res[1::4], [i**0.5 for i in values1])
assert res[2::4] == [i**3 for i in values1]
assert iters_close(res[3::4], [i**3.1 for i in values1])
def test_mix():
@python2shader_and_validate
def compute_shader(
index_xyz: ("input", "GlobalInvocationId", ivec3),
data1: ("buffer", 0, Array(vec4)),
data2: ("buffer", 1, Array(vec4)),
):
index = index_xyz.x
v = data1[index]
v1 = mix(v.x, v.y, v.z)
v2 = mix(vec2(v.x, v.x), vec2(v.y, v.y), v.z)
data2[index] = vec4(v1, v2.x, v2.y, 0.0)
skip_if_no_wgpu()
values1 = [-4, -3, -2, -1, +0, +0, +1, +2, +3, +4]
values2 = [-2, -5, -5, +2, +2, -1, +3, +1, +1, -6]
weights = [0.1 * i for i in range(10)]
stubs = [0] * 10
values = sum(zip(values1, values2, weights, stubs), ())
inp_arrays = {0: (ctypes.c_float * 40)(*values)}
out_arrays = {1: ctypes.c_float * 40}
out = compute_with_buffers(inp_arrays, out_arrays, compute_shader, n=10)
res = list(out[1])
ref = [
values1[i] * (1 - w) + values2[i] * w for i, w in enumerate(weights)
]
assert iters_close(res[0::4], ref)
assert iters_close(res[1::4], ref)
assert iters_close(res[2::4], ref)
def test_abs():
@python2shader_and_validate
def compute_shader(
index_xyz: ("input", "GlobalInvocationId", ivec3),
data1: ("buffer", 0, Array(f32)),
data2: ("buffer", 1, Array(i32)),
data3: ("buffer", 2, Array(vec2)),
):
index = index_xyz.x
v1 = abs(data1[index]) # float
v2 = abs(data2[index]) # int
data3[index] = vec2(f32(v1), v2)
skip_if_no_wgpu()
values1 = [random.uniform(-2, 2) for i in range(10)]
values2 = [random.randint(-100, 100) for i in range(10)]
inp_arrays = {
0: (ctypes.c_float * 10)(*values1),
1: (ctypes.c_int * 10)(*values2)
}
out_arrays = {2: ctypes.c_float * 20}
out = compute_with_buffers(inp_arrays, out_arrays, compute_shader, n=10)
res = list(out[2])
assert iters_close(res[0::2], [abs(v) for v in values1])
assert res[1::2] == [abs(v) for v in values2]
def test_normalize():
@python2shader_and_validate
def compute_shader(
index_xyz: ("input", "GlobalInvocationId", ivec3),
data1: ("buffer", 0, Array(f32)),
data2: ("buffer", 1, Array(vec2)),
):
index = index_xyz.x
v = data1[index]
data2[index] = normalize(vec2(v, v))
skip_if_no_wgpu()
values1 = [i - 5 for i in range(10)]
inp_arrays = {0: (ctypes.c_float * 10)(*values1)}
out_arrays = {1: ctypes.c_float * 20}
out = compute_with_buffers(inp_arrays, out_arrays, compute_shader, n=10)
res = list(out[1])
assert iters_close(res[:10], [-(2**0.5) / 2 for i in range(10)])
assert iters_close(res[-8:], [+(2**0.5) / 2 for i in range(8)])
assert math.isnan(res[10]) and math.isnan(
res[11]) # or can this also be inf?
def test_sqrt():
@python2shader_and_validate
def compute_shader(
index_xyz: ("input", "GlobalInvocationId", ivec3),
data1: ("buffer", 0, Array(f32)),
data2: ("buffer", 1, Array(vec4)),
):
index = index_xyz.x
a = data1[index]
data2[index] = vec4(a**0.5, math.sqrt(a), stdlib.sqrt(a), 0.0)
skip_if_no_wgpu()
values1 = [i for i in range(10)]
inp_arrays = {0: (ctypes.c_float * 10)(*values1)}
out_arrays = {1: ctypes.c_float * 40}
out = compute_with_buffers(inp_arrays, out_arrays, compute_shader)
res = list(out[1])
ref = [i**0.5 for i in values1]
assert iters_close(res[0::4], ref)
assert iters_close(res[1::4], ref)
assert iters_close(res[2::4], ref)
def test_math_constants():
@python2shader_and_validate
def compute_shader(
index_xyz: ("input", "GlobalInvocationId", ivec3),
data2: ("buffer", 1, Array(f32)),
):
index = index_xyz.x
if index % 2 == 0:
data2[index] = math.pi
else:
data2[index] = math.e
skip_if_no_wgpu()
out = compute_with_buffers({}, {1: ctypes.c_float * 10},
compute_shader,
n=10)
res = list(out[1])
assert iters_close(res, [math.pi, math.e] * 5)
def test_length():
@python2shader_and_validate
def compute_shader(
index_xyz: ("input", "GlobalInvocationId", ivec3),
data1: ("buffer", 0, Array(vec2)),
data2: ("buffer", 1, Array(f32)),
):
index = index_xyz.x
data2[index] = length(data1[index])
skip_if_no_wgpu()
values1 = [random.uniform(-2, 2) for i in range(20)]
inp_arrays = {0: (ctypes.c_float * 20)(*values1)}
out_arrays = {1: ctypes.c_float * 10}
out = compute_with_buffers(inp_arrays, out_arrays, compute_shader, n=10)
res = list(out[1])
ref = [(values1[i * 2]**2 + values1[i * 2 + 1]**2)**0.5 for i in range(10)]
assert iters_close(res, ref)