Update unit test for parallel_reduce.

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)