2021-07-05 15:54:15 +08:00
|
|
|
import numpy as np
|
|
|
|
|
import pytest
|
2020-06-26 19:32:21 -04:00
|
|
|
from pytest import approx
|
|
|
|
|
|
2021-06-23 17:12:25 +08:00
|
|
|
import taichi as ti
|
2022-02-10 12:37:36 +08:00
|
|
|
from tests import test_utils
|
2021-06-23 17:12:25 +08:00
|
|
|
|
2021-07-05 15:54:15 +08:00
|
|
|
OP_ADD = 0
|
|
|
|
|
OP_MIN = 1
|
|
|
|
|
OP_MAX = 2
|
|
|
|
|
OP_AND = 3
|
|
|
|
|
OP_OR = 4
|
|
|
|
|
OP_XOR = 5
|
|
|
|
|
|
|
|
|
|
ti_ops = {
|
|
|
|
|
OP_ADD: ti.atomic_add,
|
|
|
|
|
OP_MIN: ti.atomic_min,
|
|
|
|
|
OP_MAX: ti.atomic_max,
|
|
|
|
|
OP_AND: ti.atomic_and,
|
|
|
|
|
OP_OR: ti.atomic_or,
|
|
|
|
|
OP_XOR: ti.atomic_xor,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
np_ops = {
|
|
|
|
|
OP_ADD: np.sum,
|
|
|
|
|
OP_MIN: lambda a: a.min(),
|
|
|
|
|
OP_MAX: lambda a: a.max(),
|
|
|
|
|
OP_AND: np.bitwise_and.reduce,
|
|
|
|
|
OP_OR: np.bitwise_or.reduce,
|
|
|
|
|
OP_XOR: np.bitwise_xor.reduce,
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _test_reduction_single(dtype, criterion, op):
|
2020-06-26 19:32:21 -04:00
|
|
|
N = 1024 * 1024
|
2022-06-01 01:26:18 -07:00
|
|
|
if (
|
|
|
|
|
ti.lang.impl.current_cfg().arch == ti.opengl
|
|
|
|
|
or ti.lang.impl.current_cfg().arch == ti.vulkan
|
2023-01-09 20:47:57 +08:00
|
|
|
or ti.lang.impl.current_cfg().arch == ti.metal
|
2022-12-27 00:27:09 -08:00
|
|
|
or ti.lang.impl.current_cfg().arch == ti.gles
|
2022-06-01 01:26:18 -07:00
|
|
|
or ti.lang.impl.current_cfg().arch == ti.dx11
|
|
|
|
|
) and dtype == ti.f32:
|
2021-08-17 14:49:43 +08:00
|
|
|
# OpenGL/Vulkan are not capable of such large number in its float32...
|
2020-07-18 23:30:13 +08:00
|
|
|
N = 1024 * 16
|
2020-06-26 19:32:21 -04:00
|
|
|
|
2020-08-13 11:24:48 +08:00
|
|
|
a = ti.field(dtype, shape=N)
|
|
|
|
|
tot = ti.field(dtype, shape=())
|
2020-06-26 19:32:21 -04:00
|
|
|
|
2021-07-13 20:12:15 +08:00
|
|
|
if dtype in [ti.f32, ti.f64]:
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def fill():
|
|
|
|
|
for i in a:
|
|
|
|
|
a[i] = i + 0.5
|
2023-04-16 22:20:35 +08:00
|
|
|
|
2021-07-13 20:12:15 +08:00
|
|
|
else:
|
|
|
|
|
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def fill():
|
|
|
|
|
for i in a:
|
2021-09-22 17:36:02 +08:00
|
|
|
a[i] = i + 1
|
2020-06-26 19:32:21 -04:00
|
|
|
|
2021-07-05 15:54:15 +08:00
|
|
|
ti_op = ti_ops[op]
|
|
|
|
|
|
2020-06-26 19:32:21 -04:00
|
|
|
@ti.kernel
|
|
|
|
|
def reduce():
|
|
|
|
|
for i in a:
|
2021-07-05 15:54:15 +08:00
|
|
|
ti_op(tot[None], a[i])
|
2020-06-26 19:32:21 -04:00
|
|
|
|
2020-07-07 02:03:45 +09:00
|
|
|
@ti.kernel
|
|
|
|
|
def reduce_tmp() -> dtype:
|
2021-09-22 17:36:02 +08:00
|
|
|
s = ti.zero(tot[None]) if op == OP_ADD or op == OP_XOR else a[0]
|
2020-07-07 02:03:45 +09:00
|
|
|
for i in a:
|
2021-07-05 15:54:15 +08:00
|
|
|
ti_op(s, a[i])
|
2020-07-07 02:03:45 +09:00
|
|
|
return s
|
|
|
|
|
|
2020-06-26 19:32:21 -04:00
|
|
|
fill()
|
2021-09-22 17:36:02 +08:00
|
|
|
tot[None] = 0 if op in [OP_ADD, OP_XOR] else a[0]
|
2020-06-26 19:32:21 -04:00
|
|
|
reduce()
|
2020-07-07 02:03:45 +09:00
|
|
|
tot2 = reduce_tmp()
|
2020-06-26 19:32:21 -04:00
|
|
|
|
2021-09-22 17:36:02 +08:00
|
|
|
np_arr = a.to_numpy()
|
2021-07-13 20:12:15 +08:00
|
|
|
ground_truth = np_ops[op](np_arr)
|
2021-07-05 15:54:15 +08:00
|
|
|
|
2020-06-26 19:32:21 -04:00
|
|
|
assert criterion(tot[None], ground_truth)
|
2020-07-07 02:03:45 +09:00
|
|
|
assert criterion(tot2, ground_truth)
|
2020-06-26 19:32:21 -04:00
|
|
|
|
|
|
|
|
|
2021-07-05 15:54:15 +08:00
|
|
|
@pytest.mark.parametrize("op", [OP_ADD, OP_MIN, OP_MAX, OP_AND, OP_OR, OP_XOR])
|
2022-02-10 12:37:36 +08:00
|
|
|
@test_utils.test()
|
2021-07-05 15:54:15 +08:00
|
|
|
def test_reduction_single_i32(op):
|
2024-06-23 19:28:54 -07:00
|
|
|
_test_reduction_single(ti.i32, lambda x, y: int(x) % 2**32 == int(y) % 2**32, op)
|
2020-06-26 19:32:21 -04:00
|
|
|
|
|
|
|
|
|
2021-07-05 15:54:15 +08:00
|
|
|
@pytest.mark.parametrize("op", [OP_ADD])
|
2022-12-27 00:27:09 -08:00
|
|
|
@test_utils.test(exclude=[ti.opengl, ti.gles])
|
2021-07-05 15:54:15 +08:00
|
|
|
def test_reduction_single_u32(op):
|
2024-06-23 19:28:54 -07:00
|
|
|
_test_reduction_single(ti.u32, lambda x, y: int(x) % 2**32 == int(y) % 2**32, op)
|
2020-06-26 19:32:21 -04:00
|
|
|
|
|
|
|
|
|
2021-07-05 15:54:15 +08:00
|
|
|
@pytest.mark.parametrize("op", [OP_ADD, OP_MIN, OP_MAX])
|
2022-02-10 12:37:36 +08:00
|
|
|
@test_utils.test()
|
2021-07-05 15:54:15 +08:00
|
|
|
def test_reduction_single_f32(op):
|
|
|
|
|
_test_reduction_single(ti.f32, lambda x, y: x == approx(y, 3e-4), op)
|
2020-06-26 19:32:21 -04:00
|
|
|
|
|
|
|
|
|
2021-07-05 15:54:15 +08:00
|
|
|
@pytest.mark.parametrize("op", [OP_ADD])
|
2022-02-10 12:37:36 +08:00
|
|
|
@test_utils.test(require=ti.extension.data64)
|
2021-07-05 15:54:15 +08:00
|
|
|
def test_reduction_single_i64(op):
|
2024-06-23 19:28:54 -07:00
|
|
|
_test_reduction_single(ti.i64, lambda x, y: int(x) % 2**64 == int(y) % 2**64, op)
|
2020-06-26 19:32:21 -04:00
|
|
|
|
|
|
|
|
|
2021-07-05 15:54:15 +08:00
|
|
|
@pytest.mark.parametrize("op", [OP_ADD])
|
2022-12-27 00:27:09 -08:00
|
|
|
@test_utils.test(exclude=[ti.opengl, ti.gles], require=ti.extension.data64)
|
2021-07-05 15:54:15 +08:00
|
|
|
def test_reduction_single_u64(op):
|
2024-06-23 19:28:54 -07:00
|
|
|
_test_reduction_single(ti.u64, lambda x, y: int(x) % 2**64 == int(y) % 2**64, op)
|
2020-06-26 19:32:21 -04:00
|
|
|
|
|
|
|
|
|
2021-07-05 15:54:15 +08:00
|
|
|
@pytest.mark.parametrize("op", [OP_ADD])
|
2022-02-10 12:37:36 +08:00
|
|
|
@test_utils.test(require=ti.extension.data64)
|
2021-07-05 15:54:15 +08:00
|
|
|
def test_reduction_single_f64(op):
|
|
|
|
|
_test_reduction_single(ti.f64, lambda x, y: x == approx(y, 1e-12), op)
|
2020-07-27 14:59:36 +08:00
|
|
|
|
|
|
|
|
|
2022-02-10 12:37:36 +08:00
|
|
|
@test_utils.test()
|
2020-07-27 14:59:36 +08:00
|
|
|
def test_reduction_different_scale():
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def func(n: ti.template()) -> ti.i32:
|
|
|
|
|
x = 0
|
|
|
|
|
for i in range(n):
|
|
|
|
|
ti.atomic_add(x, 1)
|
|
|
|
|
return x
|
|
|
|
|
|
|
|
|
|
# 10 and 60 since OpenGL TLS stride size = 32
|
|
|
|
|
# 1024 and 100000 since OpenGL max threads per group ~= 1792
|
|
|
|
|
for n in [1, 10, 60, 1024, 100000]:
|
|
|
|
|
assert n == func(n)
|
2021-09-17 09:52:06 +08:00
|
|
|
|
|
|
|
|
|
2022-06-14 14:48:34 -07:00
|
|
|
@test_utils.test()
|
|
|
|
|
def test_reduction_non_full_warp():
|
|
|
|
|
@ti.kernel
|
|
|
|
|
def test() -> ti.i32:
|
|
|
|
|
hit_time = 1
|
|
|
|
|
ti.loop_config(block_dim=8)
|
|
|
|
|
for i in range(8):
|
|
|
|
|
ti.atomic_min(hit_time, 1)
|
|
|
|
|
return hit_time
|
|
|
|
|
|
|
|
|
|
assert test() == 1
|
|
|
|
|
|
|
|
|
|
|
2022-02-10 12:37:36 +08:00
|
|
|
@test_utils.test()
|
2022-04-01 17:53:06 +08:00
|
|
|
def test_reduction_ndarray():
|
2021-09-17 09:52:06 +08:00
|
|
|
@ti.kernel
|
2022-04-01 17:53:06 +08:00
|
|
|
def reduce(a: ti.types.ndarray()) -> ti.i32:
|
2021-09-17 09:52:06 +08:00
|
|
|
s = 0
|
|
|
|
|
for i in a:
|
|
|
|
|
ti.atomic_add(s, a[i])
|
2021-09-22 17:36:02 +08:00
|
|
|
ti.atomic_sub(s, 2)
|
2021-09-17 09:52:06 +08:00
|
|
|
return s
|
|
|
|
|
|
|
|
|
|
n = 1024
|
|
|
|
|
x = np.ones(n, dtype=np.int32)
|
2021-09-22 17:36:02 +08:00
|
|
|
assert reduce(x) == -n
|