From ae47d4c0eca5f6125b04e5aeb20ce2602d957f90 Mon Sep 17 00:00:00 2001
From: David Nadlinger <code@klickverbot.at>
Date: Sun, 9 Aug 2020 17:45:20 +0100
Subject: [PATCH] test/coredevice: Add host/device consistency checks for NumPy
 math

---
 artiq/test/coredevice/test_numpy.py | 95 +++++++++++++++++++++++++++++
 1 file changed, 95 insertions(+)
 create mode 100644 artiq/test/coredevice/test_numpy.py

diff --git a/artiq/test/coredevice/test_numpy.py b/artiq/test/coredevice/test_numpy.py
new file mode 100644
index 000000000..7cc906087
--- /dev/null
+++ b/artiq/test/coredevice/test_numpy.py
@@ -0,0 +1,95 @@
+from artiq.experiment import *
+import numpy
+from artiq.test.hardware_testbench import ExperimentCase
+from artiq.compiler import math_fns
+
+
+class _RunOnDevice(EnvExperiment):
+    def build(self):
+        self.setattr_device("core")
+
+    @kernel
+    def run_on_kernel_unary(self, a, callback, numpy):
+        self.run(a, callback, numpy)
+
+    @kernel
+    def run_on_kernel_binary(self, a, b, callback, numpy):
+        self.run(a, b, callback, numpy)
+
+
+# Binary operations supported for scalars and arrays of any dimension, including
+# broadcasting.
+ELEM_WISE_BINOPS = ["+", "*", "//", "%", "**", "-", "/"]
+
+
+class CompareHostDeviceTest(ExperimentCase):
+    def _test_binop(self, op, a, b):
+        exp = self.create(_RunOnDevice)
+        exp.run = kernel_from_string(["a", "b", "callback", "numpy"],
+                                     "callback(a " + op + "b)",
+                                     decorator=portable)
+        checked = False
+
+        def with_host_result(host):
+            def with_both_results(device):
+                nonlocal checked
+                checked = True
+                self.assertTrue(
+                    numpy.allclose(host, device),
+                    "Discrepancy in binop test for '{}': Expexcted ({}, {}) -> {}, got {}"
+                    .format(op, a, b, host, device))
+
+            exp.run_on_kernel_binary(a, b, with_both_results, numpy)
+
+        exp.run(a, b, with_host_result, numpy)
+        self.assertTrue(checked, "Test did not run")
+
+    def _test_unaryop(self, op, a):
+        exp = self.create(_RunOnDevice)
+        exp.run = kernel_from_string(["a", "callback", "numpy"],
+                                     "callback(" + op + ")",
+                                     decorator=portable)
+        checked = False
+
+        def with_host_result(host):
+            def with_both_results(device):
+                nonlocal checked
+                checked = True
+                self.assertTrue(
+                    numpy.allclose(host, device),
+                    "Discrepancy in unaryop test for '{}': Expexcted {} -> {}, got {}"
+                    .format(op, a, host, device))
+
+            exp.run_on_kernel_unary(a, with_both_results, numpy)
+
+        exp.run(a, with_host_result, numpy)
+        self.assertTrue(checked, "Test did not run")
+
+    def test_scalar_scalar_binops(self):
+        # Some arbitrarily chosen arguments of different types. Could be turned into
+        # randomised tests instead.
+        # TODO: Provoke overflows, division by zero, etc., and compare results.
+        args = [(typ(a), typ(b)) for a, b in [(0, 1), (3, 2), (11, 6)]
+                for typ in [numpy.int32, numpy.int64, numpy.float]]
+        for op in ELEM_WISE_BINOPS:
+            for arg in args:
+                self._test_binop(op, *arg)
+
+    def test_scalar_matrix_binops(self):
+        for typ in [numpy.int32, numpy.int64, numpy.float]:
+            scalar = typ(3)
+            matrix = numpy.array([[4, 5, 6], [7, 8, 9]], dtype=typ)
+            for op in ELEM_WISE_BINOPS:
+                self._test_binop(op, scalar, matrix)
+                self._test_binop(op, matrix, scalar)
+                self._test_binop(op, matrix, matrix)
+
+    def test_unary_math_fns(self):
+        names = [
+            a for a, _ in math_fns.unary_fp_intrinsics + math_fns.unary_fp_runtime_calls
+        ]
+        for name in names:
+            op = "numpy.{}(a)".format(name)
+            print(op)
+            self._test_unaryop(op, 0.5)
+            self._test_unaryop(op, numpy.array([[0.3, 0.4], [0.5, 0.6]]))