gras: jit factory api + unit tests

2 files changed, 151 insertions, 0 deletions

diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index 61d0146..63f6826 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -55,6 +55,7 @@ GR_ADD_TEST(sbuffer_test ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/sbuffe
 GR_ADD_TEST(query_test ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/query_test.py)
 GR_ADD_TEST(block_calls_test ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/block_calls_test.py)
 GR_ADD_TEST(time_tags_test ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/time_tags_test.py)
+GR_ADD_TEST(jit_factory_test ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/jit_factory_test.py)
 
 ########################################################################
 # Build an example loadable module
diff --git a/tests/jit_factory_test.py b/tests/jit_factory_test.py
new file mode 100644
index 0000000..326313c
--- /dev/null
+++ b/tests/jit_factory_test.py
@@ -0,0 +1,150 @@
+# Copyright (C) by Josh Blum. See LICENSE.txt for licensing information.
+
+import unittest
+import gras
+import numpy
+import time
+from gras import TestUtils
+
+#figure out the local include directories
+import os
+root_dir = os.path.join(os.path.dirname(__file__), '..')
+gras_inc = os.path.join(root_dir, 'include')
+pmc_inc = os.path.join(root_dir, 'PMC', 'include')
+
+ADD_F32_SOURCE = """
+#include <gras/block.hpp>
+#include <gras/factory.hpp>
+#include <iostream>
+
+struct MyAddFloat32 : gras::Block
+{
+    MyAddFloat32(void):
+        gras::Block("MyAddFloat32")
+    {
+        this->input_config(0).item_size = sizeof(float);
+        this->output_config(0).item_size = sizeof(float);
+    }
+
+    void work(const InputItems &ins, const OutputItems &outs)
+    {
+        const size_t n_nums = std::min(ins.min(), outs.min());
+        float *out = outs[0].cast<float *>();
+        const float *in0 = ins[0].cast<const float *>();
+        const float *in1 = ins[1].cast<const float *>();
+
+        for (size_t i = 0; i < n_nums; i++)
+        {
+            out[i] = in0[i] + in1[i];
+        }
+
+        this->consume(n_nums);
+        this->produce(n_nums);
+    }
+};
+
+GRAS_REGISTER_FACTORY0("/tests/my_add_f32", MyAddFloat32)
+"""
+
+ADD_CONST_F32_SOURCE = """
+#include <gras/block.hpp>
+#include <gras/factory.hpp>
+#include <iostream>
+
+struct MyAddConstFloat32 : gras::Block
+{
+    MyAddConstFloat32(void):
+        gras::Block("MyAddConstFloat32")
+    {
+        this->input_config(0).item_size = sizeof(float);
+        this->output_config(0).item_size = sizeof(float);
+        this->set_value(0); //initial state
+        this->register_call("set_value", &MyAddConstFloat32::set_value);
+        this->register_call("get_value", &MyAddConstFloat32::get_value);
+    }
+
+    void set_value(const float &value)
+    {
+        _value = value;
+    }
+
+    float get_value(void)
+    {
+        return _value;
+    }
+
+    void work(const InputItems &ins, const OutputItems &outs)
+    {
+        const size_t n_nums = std::min(ins.min(), outs.min());
+        float *out = outs[0].cast<float *>();
+        const float *in = ins[0].cast<const float *>();
+
+        for (size_t i = 0; i < n_nums; i++)
+        {
+            out[i] = in[i] + _value;
+        }
+
+        this->consume(n_nums);
+        this->produce(n_nums);
+    }
+
+    float _value;
+};
+
+GRAS_REGISTER_FACTORY0("/tests/my_add_const_f32", MyAddConstFloat32)
+"""
+
+class JITFactoryTest(unittest.TestCase):
+
+    def setUp(self):
+        self.tb = gras.TopBlock()
+
+    def tearDown(self):
+        self.tb = None
+
+    def test_add_float32(self):
+
+        gras.jit_factory(ADD_F32_SOURCE, ["-O3", "-I"+gras_inc, "-I"+pmc_inc])
+        op = gras.make("/tests/my_add_f32")
+
+        vec0 = numpy.array(numpy.random.randint(-150, +150, 10000), numpy.float32)
+        vec1 = numpy.array(numpy.random.randint(-150, +150, 10000), numpy.float32)
+
+        src0 = TestUtils.VectorSource(numpy.float32, vec0)
+        src1 = TestUtils.VectorSource(numpy.float32, vec1)
+        dst = TestUtils.VectorSink(numpy.float32)
+
+        self.tb.connect(src0, (op, 0))
+        self.tb.connect(src1, (op, 1))
+        self.tb.connect(op, dst)
+        self.tb.run()
+
+        expected_result = list(vec0 + vec1)
+        actual_result = list(dst.data())
+
+        self.assertEqual(expected_result, actual_result)
+
+    def test_add_const_float32(self):
+
+        gras.jit_factory(ADD_CONST_F32_SOURCE, ["-O3", "-I"+gras_inc, "-I"+pmc_inc])
+        op = gras.make("/tests/my_add_const_f32")
+
+        offset = 42.
+        op.set_value(offset) #set offset for test
+        self.assertAlmostEqual(op.get_value(), offset)
+
+        vec = numpy.array(numpy.random.randint(-150, +150, 10000), numpy.float32)
+
+        src = TestUtils.VectorSource(numpy.float32, vec)
+        dst = TestUtils.VectorSink(numpy.float32)
+
+        self.tb.connect(src, op, dst)
+        self.tb.run()
+
+        expected_result = list(vec + offset)
+        actual_result = list(dst.data())
+
+        self.assertEqual(expected_result, actual_result)
+
+if __name__ == '__main__':
+    unittest.main()