diff --git a/tests/test_parallelreduce.py b/tests/test_parallelreduce.py index 384274b8..d470f328 100644 --- a/tests/test_parallelreduce.py +++ b/tests/test_parallelreduce.py @@ -45,7 +45,7 @@ def test_add_squares(self): @pk.workload -class SquareSumFloat: +class SquareSumDouble: def __init__(self, n): self.N: int = n self.total: pk.double = 0 @@ -55,12 +55,12 @@ def run(self): self.total = pk.parallel_reduce(self.N, self.squaresum) @pk.workunit - def squaresum(self, i: float, acc: pk.Acc[pk.double]): + def squaresum(self, i: pk.int64, acc: pk.Acc[pk.double]): acc += i * i @pk.workload -class SquareSumInt: +class SquareSumInt64: def __init__(self, n): self.N: int = n self.total: pk.int64 = 0 @@ -73,17 +73,74 @@ def run(self): def squaresum(self, i: pk.int64, acc: pk.Acc[pk.int64]): acc += i * i +@pk.workload +class SquareSumUInt32: + def __init__(self, n): + self.N: int = n + self.total: pk.uint32 = 0 + + @pk.main + def run(self): + self.total = pk.parallel_reduce(self.N, self.squaresum) + + @pk.workunit + def squaresum(self, i: pk.int32, acc: pk.Acc[pk.uint32]): + acc += i * i + +@pk.workload +class SquareSumInt16: + def __init__(self, n): + self.N: int = n + self.total: pk.int16 = 0 + + @pk.main + def run(self): + self.total = pk.parallel_reduce(self.N, self.squaresum) + + @pk.workunit + def squaresum(self, i: pk.int16, acc: pk.Acc[pk.int16]): + acc += i * i + +@pk.workload +class SquareSumUInt8: + def __init__(self, n): + self.N: int = n + self.total: pk.uint32 = 0 + + @pk.main + def run(self): + self.total = pk.parallel_reduce(self.N, self.squaresum) + + @pk.workunit + def squaresum(self, i: pk.uint8, acc: pk.Acc[pk.int32]): + acc += i * i @pytest.mark.parametrize("series_max", [10, 5000, 90000]) -@pytest.mark.parametrize("dtype", [np.float64, np.int64]) +@pytest.mark.parametrize("dtype", [np.float64, np.int64, np.uint32]) def test_squaresum_types(series_max, dtype): # check for the ability to match NumPy in # sum of squares reductions with various types expected = np.sum(np.arange(series_max, dtype=dtype) ** 2) if dtype == np.float64: - ss_instance = SquareSumFloat(series_max) + ss_instance = SquareSumDouble(series_max) elif dtype == np.int64: - ss_instance = SquareSumInt(series_max) + ss_instance = SquareSumInt64(series_max) + elif dtype == np.uint32: + ss_instance = SquareSumUInt32(series_max) + pk.execute(pk.ExecutionSpace.OpenMP, ss_instance) + actual = ss_instance.total + assert_allclose(actual, expected) + +@pytest.mark.parametrize("series_max", [10, 500]) +@pytest.mark.parametrize("dtype", [np.int16, np.uint8]) +def test_squaresum_types(series_max, dtype): + # check for the ability to match NumPy in + # sum of squares reductions with various types + expected = np.sum(np.arange(series_max, dtype=dtype) ** 2) + if dtype == np.int16: + ss_instance = SquareSumInt16(series_max) + elif dtype == np.uint8: + ss_instance = SquareSumUInt8(series_max) pk.execute(pk.ExecutionSpace.OpenMP, ss_instance) actual = ss_instance.total assert_allclose(actual, expected)