diff --git a/ci/scripts/requirements.py b/ci/scripts/requirements.py index a591c9594..06fa8167a 100644 --- a/ci/scripts/requirements.py +++ b/ci/scripts/requirements.py @@ -1,13 +1,16 @@ #!/bin/env python import os + def deps(): import toml + c = toml.load("pyproject.toml") requires = c["build-system"]["requires"] develop = c["project"]["optional-dependencies"]["develop"] return requires + develop + def main(): ret = os.system("python -m pip install --prefer-binary toml") if ret != 0: diff --git a/conda/dev-environment-unix.yml b/conda/dev-environment-unix.yml index e060c6bce..9a0051b78 100644 --- a/conda/dev-environment-unix.yml +++ b/conda/dev-environment-unix.yml @@ -13,6 +13,7 @@ dependencies: - cyrus-sasl - deprecated - docutils<0.22.1 + - numba-cfunc-compiler>=1.0.0 - exprtk - flex - graphviz diff --git a/conda/dev-environment-win.yml b/conda/dev-environment-win.yml index 12391e9d2..a213c108b 100644 --- a/conda/dev-environment-win.yml +++ b/conda/dev-environment-win.yml @@ -13,6 +13,7 @@ dependencies: - cyrus-sasl - deprecated - docutils<0.22.1 + - numba-cfunc-compiler>=1.0.0 - exprtk # - flex # not available on windows - graphviz diff --git a/cpp/csp/engine/CppNode.h b/cpp/csp/engine/CppNode.h index f4d33f58e..634bd9406 100644 --- a/cpp/csp/engine/CppNode.h +++ b/cpp/csp/engine/CppNode.h @@ -92,7 +92,7 @@ class CppNode : public csp::Node if( nodedef.outputs.size() > OutputId::maxOutputs() ) CSP_THROW( ValueError, "number of outputs exceeds limit of " << OutputId::maxOutputs() << " on node " << name() ); - return csp::NodeDef{ INOUT_ID_TYPE( nodedef.inputs.size() ), INOUT_ID_TYPE( nodedef.outputs.size() ) }; + return csp::NodeDef{ nodedef.inputs.size(), nodedef.outputs.size() }; } void validateNodeDef() const diff --git a/cpp/csp/engine/Node.h b/cpp/csp/engine/Node.h index 23ab37a4d..3f64860f2 100644 --- a/cpp/csp/engine/Node.h +++ b/cpp/csp/engine/Node.h @@ -17,8 +17,16 @@ class InputAdapter; struct NodeDef { - NodeDef( INOUT_ID_TYPE numIn, INOUT_ID_TYPE numOut ) : numInputs( numIn ), numOutputs( numOut ) - {} + NodeDef( size_t numIn, size_t numOut ) + { + if( numIn > InputId::maxId() ) + CSP_THROW( ValueError, "number of inputs exceeds limit of " << InputId::maxId() ); + if( numOut > OutputId::maxId() ) + CSP_THROW( ValueError, "number of outputs exceeds limit of " << OutputId::maxId() ); + + numInputs = static_cast( numIn ); + numOutputs = static_cast( numOut ); + } INOUT_ID_TYPE numInputs; INOUT_ID_TYPE numOutputs; diff --git a/cpp/csp/python/CMakeLists.txt b/cpp/csp/python/CMakeLists.txt index bf7770005..e7c0488ae 100644 --- a/cpp/csp/python/CMakeLists.txt +++ b/cpp/csp/python/CMakeLists.txt @@ -68,6 +68,7 @@ add_library(cspimpl SHARED PyInputProxy.cpp PyNode.cpp PyNodeWrapper.cpp + NumbaStructFFI.cpp PyNumbaNode.cpp PyNumpyAdapter.cpp PyOutputAdapter.cpp diff --git a/cpp/csp/python/NumbaStructFFI.cpp b/cpp/csp/python/NumbaStructFFI.cpp new file mode 100644 index 000000000..7adec22ad --- /dev/null +++ b/cpp/csp/python/NumbaStructFFI.cpp @@ -0,0 +1,45 @@ +#include +#include +#include +#include + +#include +#include + +namespace csp::python +{ + +extern "C" CSPIMPL_EXPORT int64_t csp_numba_struct_enum_field_value( const void * struct_ptr, int64_t field_offset ) noexcept +{ + const char * bytes = static_cast( struct_ptr ); + return reinterpret_cast( bytes + field_offset ) -> value(); +} + +extern "C" CSPIMPL_EXPORT int8_t csp_numba_struct_enum_field_set( + void * struct_ptr, int64_t field_offset, int64_t value, const void * struct_type ) noexcept +{ + const auto * pyStructMeta = static_cast( struct_type ); + if( !struct_ptr || !pyStructMeta || !pyStructMeta -> structMeta || field_offset < 0 ) + return 0; + + const auto & fields = pyStructMeta -> structMeta -> fields(); + auto fieldIt = std::find_if( fields.begin(), fields.end(), [field_offset]( const auto & field ) + { + return field -> offset() == static_cast( field_offset ); + } ); + if( fieldIt == fields.end() || ( *fieldIt ) -> type() -> type() != CspType::Type::ENUM ) + return 0; + + try + { + const auto * enumType = static_cast( ( *fieldIt ) -> type().get() ); + ( *fieldIt ) -> setValue( static_cast( struct_ptr ), enumType -> meta() -> create( value ) ); + return 1; + } + catch( ... ) + { + return 0; + } +} + +} // namespace csp::python diff --git a/cpp/csp/python/PyNumbaNode.cpp b/cpp/csp/python/PyNumbaNode.cpp index 554a61f26..a0f507871 100644 --- a/cpp/csp/python/PyNumbaNode.cpp +++ b/cpp/csp/python/PyNumbaNode.cpp @@ -1,237 +1,323 @@ -#include #include +#include +#include +#include #include #include #include #include -#include -#include -#include -#include -#include -#include #include #include +#include +#include -#include - - -#ifdef DEBUG_NUMBA_NODE -#include -#define DEBUG_ONLY(v) do {v;} while(false) -#else -#define DEBUG_ONLY(v) do {} while(false) -#endif - -#define DEBUG_PRINT(v) DEBUG_ONLY(std::cout << v << std::endl) -#define DEBUG_PRINT_WITH_LOC(v) DEBUG_PRINT(__FILE__ << ':' << __LINE__ << ": " << v) - +#include +#include +#include namespace csp::python { -PyNumbaNode::PyNumbaNode(csp::Engine *engine, - void *stateObject, CallbackType numbaInitCallback, CallbackType numbaImplCallback, - PyObjectPtr inputs, - PyObjectPtr outputs, - NodeDef def, - PyObject *dataReference) : csp::Node(def, engine), - m_stateObject(stateObject), - m_numbaInitCallback(numbaInitCallback), - m_numbaImplCallback(numbaImplCallback), - m_dataReference(PyObjectPtr::incref(dataReference)) + +namespace { - DEBUG_PRINT_WITH_LOC("Created numba node: " << (int64_t) (this) << " stateObject=" << (int64_t) (stateObject) - << " numbaImplCallback=" - << (int64_t) numbaImplCallback); - init(inputs, outputs); +constexpr size_t OUTPUT_VALUE_SLOT_BYTES = sizeof( int64_t ); + +using NumbaTypeSwitch = PartialSwitchCspType< + CspType::Type::BOOL, + CspType::Type::INT64, + CspType::Type::DOUBLE, + CspType::Type::DATETIME, + CspType::Type::TIMEDELTA, + CspType::Type::ENUM, + CspType::Type::STRUCT +>; } -PyNumbaNode::~PyNumbaNode() +PyNumbaNode::PyNumbaNode( + csp::Engine * engine, + CompiledFuncT compiledFunc, + PyObjectPtr stateVariables, + NodeDef def, + PyObject * dataReference +) : csp::Node( def, engine ), + m_compiledFunc( compiledFunc ), + m_inputArgs( def.numInputs, nullptr ), + m_inputValid( def.numInputs, 0 ), + m_inputTicked( def.numInputs, 0 ), + m_inputEnumStorage( def.numInputs ), + m_outputArgs( def.numOutputs, nullptr ), + m_outputTicked( def.numOutputs, 0 ), + m_outputValueSlots( def.numOutputs ), + m_outputStructValues( def.numOutputs ), + m_dataReference( PyObjectPtr::incref( dataReference ) ) { + for( size_t i = 0; i < m_outputArgs.size(); ++i ) + { + // Numba writes primitive-compatible outputs by value into these fixed-size slots. + m_outputValueSlots[ i ] = std::make_unique( OUTPUT_VALUE_SLOT_BYTES ); + m_outputArgs[ i ] = m_outputValueSlots[ i ].get(); + } + initStateArrays( stateVariables ); } -void PyNumbaNode::init(PyObjectPtr inputs, PyObjectPtr outputs) +void PyNumbaNode::initStateArrays( PyObjectPtr stateVariables ) { + if( !stateVariables.ptr() || stateVariables.ptr() == Py_None || !PyTuple_Check( stateVariables.ptr() ) ) + return; - for (int input_idx = 0; input_idx < numInputs(); ++input_idx) + const size_t stateCount = PyTuple_Size( stateVariables.ptr() ); + + // Keep only buffers allocated here; the compiled callback owns container pointers. + auto allocateStateBuffer = [this]( size_t size ) { - PyObject * inputType = PyTuple_GET_ITEM(inputs.ptr(), input_idx); + m_ownedStateBuffers.emplace_back( new char[ size ]() ); + return m_ownedStateBuffers.back().get(); + }; - //if basket - if (PyTuple_Check(inputType)) + // Initialize state slots + for( size_t i = 0; i < stateCount; ++i ) + { + PyObject * stateVal = PyTuple_GET_ITEM( stateVariables.ptr(), i ); + void * stateArg = nullptr; + + if( PyObject_TypeCheck( stateVal, &PyStruct::PyType ) ) { - PyObject * shape = PyTuple_GET_ITEM(inputType, 0); - std::uint64_t basketSize; - if (PyLong_Check(shape)) - basketSize = fromPython(shape); - else - { - if (!PyList_Check(shape)) - CSP_THROW(TypeError, "Expected basket type as int or list, got " << Py_TYPE(shape)->tp_name); - basketSize = PyList_GET_SIZE(shape); - } - - initInputBasket(input_idx, basketSize, false); + auto * pyStruct = reinterpret_cast( stateVal ); + size_t structSize = pyStruct -> structMeta() -> size(); + char * buf = allocateStateBuffer( structSize ); + std::memcpy( buf, pyStruct -> struct_.get(), structSize ); + stateArg = buf; } + else if( PyBool_Check( stateVal ) ) + { + char * buf = allocateStateBuffer( 1 ); + *buf = stateVal == Py_True; + stateArg = buf; + } + else if( PyLong_Check( stateVal ) ) + { + // CSP enums are IntEnum instances, so their Numba representation + // is initialized through the same int64 path as Python integers. + int64_t stateValue = fromPython( stateVal ); + char * buf = allocateStateBuffer( sizeof( int64_t ) ); + *reinterpret_cast( buf ) = stateValue; + stateArg = buf; + } + else if( PyFloat_Check( stateVal ) ) + { + char * buf = allocateStateBuffer( sizeof( double ) ); + *reinterpret_cast( buf ) = PyFloat_AsDouble( stateVal ); + stateArg = buf; + } + m_stateArgs.push_back( stateArg ); } } -void PyNumbaNode::start() +void PyNumbaNode::invokeCompiled( LifecyclePhase phase ) { - m_numbaInitCallback(this, m_stateObject); + const int8_t status = m_compiledFunc( + m_outputArgs.data(), + m_outputTicked.data(), + m_stateArgs.data(), + static_cast( phase ), + m_inputArgs.data(), + m_inputTicked.data(), + m_inputValid.data() + ); + if( status != 1 ) + CSP_THROW( RuntimeException, "Compiled Numba callback failed" ); } -void PyNumbaNode::stop() +void PyNumbaNode::start() { -// PyObjectPtr rv = PyObjectPtr::own(PyObject_CallMethod(m_gen.ptr(), "close", nullptr)); -// if (!rv.ptr()) -// CSP_THROW(PythonPassthrough, ""); + invokeCompiled( LifecyclePhase::Start ); } -PyNumbaNode *PyNumbaNode::create(PyEngine *pyengine, PyObject *inputs, PyObject *outputs, - PyObject *stateObject, PyObject *numbaInitCallback, PyObject *numbaImplCallback, - PyObject *dataReference) +void PyNumbaNode::stop() { - //parse inputs/outputs, create outputs time series, etc - Py_ssize_t numInputs = PyTuple_GET_SIZE(inputs); - Py_ssize_t numOutputs = PyTuple_GET_SIZE(outputs); - - static_assert(sizeof(void *) == sizeof(int64_t)); - - void *stateObjectPtr = reinterpret_cast(fromPython(stateObject)); - CallbackType numbaInitCallbackPtr = reinterpret_cast(fromPython(numbaInitCallback)); - CallbackType numbaImplCallbackPtr = reinterpret_cast(fromPython(numbaImplCallback)); - - - if( size_t( numInputs ) > InputId::maxBasketElements() ) - CSP_THROW( ValueError, "number of inputs exceeds limit of " << InputId::maxBasketElements() ); - - if( size_t( numOutputs ) > OutputId::maxBasketElements() ) - CSP_THROW( ValueError, "number of outputs exceeds limit of " << OutputId::maxBasketElements() ); - - - return pyengine->engine()->createOwnedObject(stateObjectPtr, numbaInitCallbackPtr, - numbaImplCallbackPtr, PyObjectPtr::incref(inputs), - PyObjectPtr::incref(outputs), - NodeDef(numInputs, numOutputs), dataReference); + invokeCompiled( LifecyclePhase::Stop ); } -const char *PyNumbaNode::name() const +const char * PyNumbaNode::name() const { - // TODO: propagate name from python here? return "PyNumbaNode"; } -void PyNumbaNode::call_callback() -{ - m_numbaImplCallback(this, m_stateObject); -// if (!m_gen->ob_type->tp_iternext(m_gen.ptr())) -// CSP_THROW(PythonPassthrough, ""); -} - void PyNumbaNode::executeImpl() { - call_callback(); -} - -static PyObject *PyNumbaNode_create(PyObject * module, PyObject * args) -{ - CSP_BEGIN_METHOD ; - - PyEngine *engine; - PyObject * inputs; - PyObject * outputs; - PyObject * state; - PyObject * initCallback; - PyObject * implCallback; - PyObject * dataReference; - - if (!PyArg_ParseTuple(args, "O!O!O!O!O!O!O!", - &PyEngine::PyType, &engine, - &PyTuple_Type, &inputs, - &PyTuple_Type, &outputs, - &PyLong_Type, &state, - &PyLong_Type, &initCallback, - &PyLong_Type, &implCallback, - &PyTuple_Type, &dataReference)) - CSP_THROW(PythonPassthrough, ""); - - auto node = PyNumbaNode::create(engine, inputs, outputs, state, initCallback, implCallback, dataReference); - return PyNodeWrapper::create(node); - CSP_RETURN_NULL; -} - -REGISTER_MODULE_METHOD("PyNumbaNode", PyNumbaNode_create, METH_VARARGS, "PyNumbaNode"); - -} - -#define IMPLEMENT_TYPE_GETTERS_AND_SETTERS(typ_name, c_typ_name) \ -c_typ_name __csp_get_node_ ## typ_name ## _value__(int64_t node_ptr, int32_t input_index) \ -{ \ -auto node = reinterpret_cast(node_ptr); \ -auto res = node->tsinput(input_index)->lastValueTyped(); \ -return res; \ -} \ -void __csp_return_ ## typ_name ## _value__(int64_t node_ptr, int32_t output_index, c_typ_name value) \ -{ \ - auto node = reinterpret_cast(node_ptr); \ - node->tsoutput(output_index)->outputTickTyped(node->cycleCount(), node->now(), value); \ -} - -namespace -{ -inline csp::python::PyNumbaNode *getNodeFromInt(int64_t node_ptr) -{ - return reinterpret_cast(node_ptr); -} + const uint64_t currentCycleCount = rootEngine() -> cycleCount(); + // Refresh pointers: buffered values can move, and struct targets can change on each tick. + for( size_t i = 0; i < m_inputArgs.size(); ++i ) + { + const TimeSeriesProvider * ts = tsinput( i ); + if( ts ) + { + bool valid = ts -> valid(); + m_inputValid[ i ] = valid; + m_inputTicked[ i ] = ts -> lastCycleCount() == currentCycleCount; + + if( valid ) + { + const CspType * cspType = ts -> type(); + CSP_ASSERT( cspType != nullptr ); + NumbaTypeSwitch::invoke( cspType, [this, ts, i]( auto tag ) + { + using CType = typename decltype( tag )::type; + if constexpr( std::is_same_v ) + { + // CSP stores an enum pointer; Numba needs an int64 value that stays alive through the callback. + m_inputEnumStorage[ i ] = ts -> lastValueTyped().value(); + m_inputArgs[ i ] = &m_inputEnumStorage[ i ]; + } + else if constexpr( std::is_same_v ) + { + m_inputArgs[ i ] = ts -> lastValueTyped().get(); + } + else + m_inputArgs[ i ] = &( ts -> lastValueTyped() ); + } ); + } + } + } -inline csp::TimeSeriesProvider *getNonConstInput(csp::python::PyNumbaNode *node, int32_t input_index) -{ - return const_cast(node->tsinput(input_index)); -} - -} + // Clear output ticked flags before calling compiled function + std::fill( m_outputTicked.begin(), m_outputTicked.end(), static_cast( 0 ) ); + for( size_t i = 0; i < m_outputArgs.size(); ++i ) + { + const CspType * type = tsoutput( i ) -> type(); + if( type -> type() == CspType::Type::STRUCT ) + { + m_outputStructValues[ i ] = static_cast( type ) -> meta() -> create(); + m_outputArgs[ i ] = m_outputStructValues[ i ].get(); + } + } -extern "C" { + // Call the compiled Numba function + invokeCompiled( LifecyclePhase::Execute ); -int64_t __csp_create_datetime_nanoseconds__(int32_t year, int32_t month, int32_t day, - int32_t hour, int32_t minute, int32_t second, - int32_t nanosecond) -{ - return csp::DateTime(year, month, day, - hour, minute, second, - nanosecond).asNanoseconds(); -} + const DateTime currentTime = now(); -bool __csp_numba_node_ticked__(int64_t node_ptr, int32_t input_index) -{ - auto node = getNodeFromInt(node_ptr); - return node->inputTicked(input_index); + // Process output ticks + for( size_t i = 0; i < m_outputArgs.size(); ++i ) + { + if( m_outputTicked[ i ] ) + { + TimeSeriesProvider * ts = tsoutput( i ); + if( ts ) + { + const CspType * cspType = ts -> type(); + CSP_ASSERT( cspType != nullptr ); + NumbaTypeSwitch::invoke( cspType, [this, ts, cspType, i, currentCycleCount, currentTime]( auto tag ) + { + using CType = typename decltype( tag )::type; + if constexpr( std::is_same_v ) + { + ts -> outputTickTyped( + currentCycleCount, currentTime, + DateTime::fromNanoseconds( *static_cast( m_outputArgs[ i ] ) ) + ); + } + else if constexpr( std::is_same_v ) + { + ts -> outputTickTyped( + currentCycleCount, currentTime, + TimeDelta::fromNanoseconds( *static_cast( m_outputArgs[ i ] ) ) + ); + } + else if constexpr( std::is_same_v ) + { + auto enumType = static_cast( cspType ); + int64_t enumValue = *static_cast( m_outputArgs[ i ] ); + ts -> outputTickTyped( + currentCycleCount, currentTime, + enumType -> meta() -> create( enumValue ) + ); + } + else if constexpr( std::is_same_v ) + ts -> outputTickTyped( currentCycleCount, currentTime, m_outputStructValues[ i ] ); + else + ts -> outputTickTyped( + currentCycleCount, currentTime, + *static_cast( m_outputArgs[ i ] ) + ); + } ); + } + } + } } -bool __csp_numba_node_valid__(int64_t node_ptr, int32_t input_index) +PyNumbaNode * PyNumbaNode::create( + PyEngine * pyengine, + PyObject * compiledFuncPtr, + PyObject * inputs, + PyObject * outputs, + PyObject * stateVariables, + PyObject * dataReference +) { - auto node = getNodeFromInt(node_ptr); - return node->tsinput(input_index)->valid(); + // Parse compiled function pointer + if( !PyLong_Check( compiledFuncPtr ) ) + CSP_THROW( TypeError, "compiledFuncPtr must be an integer (function pointer)" ); + + CompiledFuncT funcPtr = reinterpret_cast( PyLong_AsVoidPtr( compiledFuncPtr ) ); + if( !funcPtr ) + CSP_THROW( ValueError, "compiledFuncPtr is null" ); + + // Count inputs/outputs + Py_ssize_t numInputs = 0; + Py_ssize_t numOutputs = 0; + + if( inputs && PyTuple_Check( inputs ) ) + numInputs = PyTuple_Size( inputs ); + + if( outputs && PyTuple_Check( outputs ) ) + numOutputs = PyTuple_Size( outputs ); + + return pyengine -> engine() -> createOwnedObject( + funcPtr, + PyObjectPtr::incref( stateVariables ), + csp::NodeDef( numInputs, numOutputs ), + dataReference + ); } -bool __csp_numba_node_make_passive__(int64_t node_ptr, int32_t input_index) +// Python module method to create PyNumbaNode +static PyObject * PyNumbaNode_create( PyObject * module, PyObject * args ) { - auto node = getNodeFromInt(node_ptr); - return getNonConstInput(node, input_index)->removeConsumer(node, csp::InputId(input_index)); -} + CSP_BEGIN_METHOD; + + PyEngine * engine; + PyObject * compiledFuncPtr; + PyObject * inputs; + PyObject * outputs; + PyObject * stateVariables; + PyObject * dataReference; + + if( !PyArg_ParseTuple( args, "O!OOOOO", + &PyEngine::PyType, &engine, + &compiledFuncPtr, + &inputs, + &outputs, + &stateVariables, + &dataReference ) ) + CSP_THROW( PythonPassthrough, "" ); + + + auto * node = PyNumbaNode::create( + engine, + compiledFuncPtr, + inputs, + outputs, + stateVariables, + dataReference + ); + return PyNodeWrapper::create( node ); -bool __csp_numba_node_make_active__(int64_t node_ptr, int32_t input_index) -{ - auto node = getNodeFromInt(node_ptr); - return getNonConstInput(node, input_index)->addConsumer(node, csp::InputId(input_index), true); + CSP_RETURN_NULL; } +REGISTER_MODULE_METHOD( "PyNumbaNode", PyNumbaNode_create, METH_VARARGS, "Create a Numba-compiled CSP node" ); -IMPLEMENT_TYPE_GETTERS_AND_SETTERS(int64, int64_t) -IMPLEMENT_TYPE_GETTERS_AND_SETTERS(double, double) -IMPLEMENT_TYPE_GETTERS_AND_SETTERS(bool, bool) - -} +} // namespace csp::python diff --git a/cpp/csp/python/PyNumbaNode.h b/cpp/csp/python/PyNumbaNode.h index cc71a05fe..8a4639fca 100644 --- a/cpp/csp/python/PyNumbaNode.h +++ b/cpp/csp/python/PyNumbaNode.h @@ -1,50 +1,105 @@ -#ifndef _IN_PYTHON_PYNUMBANODE_H -#define _IN_PYTHON_PYNUMBANODE_H +#ifndef _IN_CSP_PYTHON_PYNUMBANODE_H +#define _IN_CSP_PYTHON_PYNUMBANODE_H #include #include +#include #include #include +#include +#include +#include namespace csp::python { + class PyEngine; -typedef void (*CallbackType)(void *node, void *state); +// Type used for ticked/valid flags - must match int8_t used in Numba generated code +using ValidType = int8_t; + +// Values must match numba_cfunc_compiler.core.names. +enum class LifecyclePhase : ValidType +{ + Execute = 0, + Start = 1, + Stop = 2, +}; +/** + * The compiled function signature matches the source categories registered + * by CSP (Constant, Output, State, Lifecycle, Signal): + * + * int8_t(*)( void ** outputs, + * ValidType * outputs_ticked, + * void ** state_variables, + * ValidType lifecycle_phase, + * const void * const * inputs, + * const ValidType * inputs_ticked, + * const ValidType * inputs_valid ); + */ class PyNumbaNode final : public csp::Node { public: - // TODO: Add suppot for initialization callback as well - PyNumbaNode(csp::Engine *engine, - void *stateObject, CallbackType numbaInitCallback, CallbackType numbaImplCallback, PyObjectPtr inputs, - PyObjectPtr outputs, - NodeDef def, PyObject *dataReference); + using CompiledFuncT = int8_t(*)( + void ** outputs, + ValidType * outputs_ticked, + void ** state_variables, + ValidType lifecycle_phase, + const void * const * inputs, + const ValidType * inputs_ticked, + const ValidType * inputs_valid + ); - ~PyNumbaNode(); + PyNumbaNode( + csp::Engine * engine, + CompiledFuncT compiledFunc, + PyObjectPtr stateVariables, + csp::NodeDef def, + PyObject * dataReference + ); void executeImpl() override; - void start() override; - void stop() override; + const char * name() const override; - const char *name() const override; - - static PyNumbaNode *create(PyEngine *engine, PyObject *inputs, PyObject *outputs, - PyObject *stateObject, PyObject *numbaInitCallback, PyObject *numbaImplCallback, - PyObject *dataReference); + static PyNumbaNode * create( + PyEngine * engine, + PyObject * compiledFuncPtr, + PyObject * inputs, + PyObject * outputs, + PyObject * stateVariables, + PyObject * dataReference + ); private: - void init(PyObjectPtr inputs, PyObjectPtr outputs); + void initStateArrays( PyObjectPtr stateVariables ); + void invokeCompiled( LifecyclePhase phase ); + + // The compiled Numba function + CompiledFuncT m_compiledFunc = nullptr; + + // Active input arrays + std::vector m_inputArgs; + std::vector m_inputValid; + std::vector m_inputTicked; + std::vector m_inputEnumStorage; // Stable int64 values passed to Numba - void call_callback(); + // Output arrays + std::vector m_outputArgs; + std::vector m_outputTicked; + std::vector> m_outputValueSlots; // Primitive output storage + std::vector m_outputStructValues; // Fresh native storage for each callback - void *m_stateObject; - CallbackType m_numbaInitCallback; - CallbackType m_numbaImplCallback; + // State arrays + std::vector m_stateArgs; + std::vector> m_ownedStateBuffers; + + // Keep Python objects alive PyObjectPtr m_dataReference; }; -}; -#endif +} // namespace csp::python + +#endif // _IN_CSP_PYTHON_PYNUMBANODE_H diff --git a/csp/impl/wiring/__init__.py b/csp/impl/wiring/__init__.py index 8c8095b01..6d271205f 100644 --- a/csp/impl/wiring/__init__.py +++ b/csp/impl/wiring/__init__.py @@ -1,5 +1,6 @@ from csp.impl.wiring.adapters import * from csp.impl.wiring.context import * +from csp.impl.wiring.csp_numba.numba_node import * from csp.impl.wiring.delayed_edge import DelayedEdge from csp.impl.wiring.dynamic import dynamic from csp.impl.wiring.edge import * @@ -7,7 +8,6 @@ from csp.impl.wiring.graph import * from csp.impl.wiring.node import * from csp.impl.wiring.node_parser import * -from csp.impl.wiring.numba_node import * from csp.impl.wiring.outputs import * from csp.impl.wiring.runtime import * from csp.impl.wiring.signature import * diff --git a/csp/impl/wiring/csp_numba/__init__.py b/csp/impl/wiring/csp_numba/__init__.py new file mode 100644 index 000000000..5a6c046cd --- /dev/null +++ b/csp/impl/wiring/csp_numba/__init__.py @@ -0,0 +1,3 @@ +from .numba_node import NumbaNodeDef, numba_node + +__all__ = ["numba_node", "NumbaNodeDef"] diff --git a/csp/impl/wiring/csp_numba/csp_node_transformer.py b/csp/impl/wiring/csp_numba/csp_node_transformer.py new file mode 100644 index 000000000..110b65270 --- /dev/null +++ b/csp/impl/wiring/csp_numba/csp_node_transformer.py @@ -0,0 +1,222 @@ +import ast +from dataclasses import dataclass +from typing import TYPE_CHECKING, List, Union + +if TYPE_CHECKING: + from csp.impl.wiring.node_parser import ParsedNodeDefinition + + +@dataclass +class TransformedNode: + transformed_ast: ast.FunctionDef + start_body: List[ast.AST] + stop_body: List[ast.AST] + + +class CspNodeTransformer(ast.NodeTransformer): + """ + Transforms CSP node AST for Numba compilation. + + Transformations: + - csp.ticked(a, b, ...) -> a.ticked() or b.ticked() or ... + - csp.valid(a, b, ...) -> a.valid() and b.valid() and ... + - with csp.state(): variables -> State[type] annotations + - csp.output(a=x, b=y, ...) -> set_output('a', x), set_output('b', y), ... + + NOT transformed (kept as CSP native): + - ts[type] annotations (inputs and outputs) + """ + + def __init__(self): + self.csp_call_transformers = { + "ticked": lambda node: self._transform_signal_check( + node=node, + attr_name="ticked", + op=ast.Or(), + ), + "valid": lambda node: self._transform_signal_check( + node=node, + attr_name="valid", + op=ast.And(), + ), + } + + def _infer_type(self, name: str, value: ast.AST) -> ast.AST: + if isinstance(value, ast.Constant): + val = value.value + if isinstance(val, bool): + return ast.Name(id="bool", ctx=ast.Load()) + elif isinstance(val, int): + return ast.Name(id="int", ctx=ast.Load()) + elif isinstance(val, float): + return ast.Name(id="float", ctx=ast.Load()) + elif isinstance(val, str): + return ast.Name(id="str", ctx=ast.Load()) + elif isinstance(value, ast.List): + return ast.Name(id="list", ctx=ast.Load()) + elif isinstance(value, ast.Dict): + return ast.Name(id="dict", ctx=ast.Load()) + elif isinstance(value, ast.Call): + if isinstance(value.func, ast.Name): + return ast.Name(id=value.func.id, ctx=ast.Load()) + + raise ValueError( + f"Unable to infer type for state variable '{name}' from expression " + f"'{ast.unparse(value)}'; add an explicit annotation" + ) + + def _is_csp_call(self, node: ast.Call, method_name: str) -> bool: + if isinstance(node.func, ast.Attribute): + if isinstance(node.func.value, ast.Name) and node.func.value.id == "csp": + return node.func.attr == method_name + return False + + def _transform_signal_check(self, *, node: ast.Call, attr_name: str, op: ast.boolop) -> ast.AST: + """ + 1. csp.ticked(a, b, ...) -> a.ticked() or b.ticked() or ... + 2. csp.valid(a, b, ...) -> a.valid() and b.valid() and ... + """ + if not node.args: + raise ValueError(f"csp.{attr_name} requires at least one argument") + + signal_checks = [] + for arg in node.args: + signal_check = ast.Call( + func=ast.Attribute(value=self.visit(arg), attr=attr_name, ctx=ast.Load()), args=[], keywords=[] + ) + signal_checks.append(signal_check) + + if len(signal_checks) == 1: + return signal_checks[0] + + return ast.BoolOp(op=op, values=signal_checks) + + def _transform_output(self, node: ast.Call) -> List[ast.AST]: + """Transform csp.output(...) calls into set_output(...) helper calls.""" + if node.args: + if node.keywords or len(node.args) > 2: + raise ValueError("csp.output accepts one value, (name, value), or named output values") + if len(node.args) == 2: + if not isinstance(node.args[0], ast.Name): + raise TypeError("csp.output(name, value) requires a declared output name") + outputs = ((node.args[0].id, node.args[1]),) + else: + outputs = (("output_0", node.args[0]),) + else: + outputs = ((kw.arg, kw.value) for kw in node.keywords) + + statements = [] + for name, value in outputs: + set_output_call = ast.Call( + func=ast.Name(id="set_output", ctx=ast.Load()), + args=[ast.Constant(value=name), self.visit(value)], + keywords=[], + ) + statements.append(ast.Expr(value=set_output_call)) + + return statements + + def visit_Expr(self, node: ast.Expr) -> Union[ast.AST, List[ast.AST]]: + """Lower standalone ``csp.output(...)`` calls into output statements.""" + if isinstance(node.value, ast.Call) and self._is_csp_call(node.value, "output"): + return self._transform_output(node.value) + return self.generic_visit(node) + + def visit_Return(self, node: ast.Return) -> Union[ast.AST, List[ast.AST]]: + """Lower ``return csp.output(...)`` while preserving the early return.""" + if isinstance(node.value, ast.Call) and self._is_csp_call(node.value, "output"): + return [*self._transform_output(node.value), ast.Return(value=None)] + return self.generic_visit(node) + + def visit_Call(self, node: ast.Call) -> ast.AST: + for method_name, transformer in self.csp_call_transformers.items(): + if self._is_csp_call(node, method_name): + return transformer(node) + + if self._is_csp_call(node, "output"): + raise ValueError("csp.output(...) must be used as a standalone statement or returned") + + # Recursively transform arguments + node.args = [self.visit(arg) for arg in node.args] + node.keywords = [ast.keyword(arg=kw.arg, value=self.visit(kw.value)) for kw in node.keywords] + + return node + + @staticmethod + def _create_state_assignment( + *, + target: ast.expr, + type_annotation: ast.AST, + value: ast.AST, + simple: int, + ) -> ast.AnnAssign: + """Create a state assignment with State[type] annotation.""" + state_type = ast.Subscript(value=ast.Name(id="State", ctx=ast.Load()), slice=type_annotation, ctx=ast.Load()) + return ast.AnnAssign(target=target, annotation=state_type, value=value, simple=simple) + + def _transform_state_statements(self, statements) -> List[ast.AST]: + transformed = [] + + for stmt in statements: + if isinstance(stmt, ast.AnnAssign): + # Already annotated: y: float = 1.0 + if not isinstance(stmt.target, ast.Name): + raise TypeError("numba_node state declarations require a simple variable name") + transformed.append( + self._create_state_assignment( + target=stmt.target, + type_annotation=stmt.annotation, + value=stmt.value, + simple=stmt.simple, + ) + ) + + elif isinstance(stmt, ast.Assign): + # Unannotated: x = 0, need to infer type + for target in stmt.targets: + if not isinstance(target, ast.Name): + raise TypeError("numba_node state assignments require a simple variable name") + inferred_type = self._infer_type(target.id, stmt.value) + transformed.append( + self._create_state_assignment( + target=target, + type_annotation=inferred_type, + value=stmt.value, + simple=1, + ) + ) + + elif isinstance(stmt, ast.Expr): + transformed.extend(self._flatten_transformed_statements((stmt,))) + + elif not isinstance(stmt, ast.Pass): + raise TypeError(f"numba_node state blocks do not support {type(stmt).__name__} statements") + + return transformed + + def _flatten_transformed_statements(self, statements: List[ast.AST]) -> List[ast.AST]: + flattened = [] + for stmt in statements: + transformed_stmt = self.visit(stmt) + if isinstance(transformed_stmt, list): + flattened.extend(transformed_stmt) + else: + flattened.append(transformed_stmt) + return flattened + + def transform_csp_node(self, definition: "ParsedNodeDefinition") -> TransformedNode: + start_body = self._flatten_transformed_statements(definition.blocks.start) + stop_body = self._flatten_transformed_statements(definition.blocks.stop) + transformed_body = self._transform_state_statements(definition.blocks.state) + transformed_body.extend(self._flatten_transformed_statements(definition.blocks.body)) + func_def = definition.funcdef + new_func_def = ast.FunctionDef( + name=func_def.name, + args=func_def.args, # Keep original args with ts[type] annotations + body=transformed_body if transformed_body else [ast.Pass()], + decorator_list=[], # Remove decorators + returns=func_def.returns, # Keep original return annotation + ) + + ast.fix_missing_locations(new_func_def) + return TransformedNode(new_func_def, start_body, stop_body) diff --git a/csp/impl/wiring/csp_numba/enum_support.py b/csp/impl/wiring/csp_numba/enum_support.py new file mode 100644 index 000000000..29de81f8b --- /dev/null +++ b/csp/impl/wiring/csp_numba/enum_support.py @@ -0,0 +1,66 @@ +"""CSP IntEnum storage and literal binding for compiled nodes.""" + +import ast +import inspect +from dataclasses import dataclass +from enum import IntEnum + +from numba_cfunc_compiler.core.analysis import ParameterInfo +from numba_cfunc_compiler.types.base import HostType +from numba_cfunc_compiler.types.binding import ConstantPlan +from numba_cfunc_compiler.types.builtin.enum.host import register_enum_family +from numba_cfunc_compiler.types.factory import HostTypeFactory +from numba_cfunc_compiler.utils.ast import AST + + +def register_family(enum_type): + return register_enum_family( + enum_type, + abi_values={member.name: int(member.value) for member in enum_type}, + scalar_storage="int64", + set_bits=None, + ) + + +@dataclass(frozen=True) +class CspEnumHostType(HostType): + @classmethod + def is_type_supported(cls, var_type): + return isinstance(var_type, type) and issubclass(var_type, IntEnum) + + def get_numba_type_name(self): + return "int64" + + def get_state_payload(self): + return register_family(self.value).state_payload() + + def constant_plan(self, name, value, call_globals, slot_index): + binding = register_family(self.value) + literal_name = f"_csp_enum_constant_{name}" + call_globals[literal_name] = binding.literal(value.name) + expression = AST.function_call(literal_name) + return ConstantPlan((AST.assignment(name, expression),), ast.Name(id=name, ctx=ast.Load())) + + @classmethod + def try_parse_input(cls, param: inspect.Parameter, ann): + if cls.is_type_supported(ann): + return ParameterInfo(expected_type=ann) + return None + + @classmethod + def try_parse_state(cls, node, var_name, globalns): + try: + enum_type = eval(ast.unparse(node.annotation.slice), globalns) + except (NameError, AttributeError): + return None + if not cls.is_type_supported(enum_type): + return None + value = eval(ast.unparse(node.value), globalns) + if not isinstance(value, enum_type): + raise TypeError(f"State '{var_name}' requires a {enum_type.__name__} initializer") + register_family(enum_type) + return value, enum_type + + +def register(): + HostTypeFactory.register(CspEnumHostType, priority=0) diff --git a/csp/impl/wiring/csp_numba/input_handlers.py b/csp/impl/wiring/csp_numba/input_handlers.py new file mode 100644 index 000000000..58a18b35e --- /dev/null +++ b/csp/impl/wiring/csp_numba/input_handlers.py @@ -0,0 +1,34 @@ +"""Pass CSP-parsed timeseries input metadata to numba_cfunc_compiler.""" + +import inspect +from dataclasses import dataclass +from enum import Enum, auto +from typing import Any, Optional + +from numba_cfunc_compiler.core.analysis import InputTypeHandler, ParameterInfo + + +class CspInputCategory(Enum): + SIGNAL = auto() + SIGNAL_SET = auto() + + +@dataclass(frozen=True) +class CspInputMetadata: + category: CspInputCategory + + +class CspInputHandler(InputTypeHandler): + def try_parse(self, param: inspect.Parameter, ann: Any) -> Optional[ParameterInfo]: + if not isinstance(ann, CspInputMetadata): + return None + return ParameterInfo(expected_type=ann, category=ann.category) + + def validate_value(self, param_name: str, value: Any, expected_type: Any) -> Any: + return value + + +def register(): + from numba_cfunc_compiler.core.analysis import FunctionAnalyzer + + FunctionAnalyzer.register_input_handler(CspInputHandler()) diff --git a/csp/impl/wiring/csp_numba/numba_node.py b/csp/impl/wiring/csp_numba/numba_node.py new file mode 100644 index 000000000..c7e4006ca --- /dev/null +++ b/csp/impl/wiring/csp_numba/numba_node.py @@ -0,0 +1,224 @@ +import inspect +from enum import IntEnum +from typing import Callable, Optional, get_args, get_origin + +from numba_cfunc_compiler.api import NumbaDict, NumbaList +from numba_cfunc_compiler.core.compile import create_compiled_func +from numba_cfunc_compiler.core.context import CompilationContext +from numba_cfunc_compiler.core.defaults import register_all +from numba_cfunc_compiler.core.postprocess import CompilationOptions + +from csp.impl.__cspimpl import _cspimpl +from csp.impl.wiring.csp_numba.csp_node_transformer import CspNodeTransformer +from csp.impl.wiring.csp_numba.enum_support import register as register_enum_support, register_family +from csp.impl.wiring.csp_numba.input_handlers import ( + CspInputCategory, + CspInputMetadata, + register as register_input_handlers, +) +from csp.impl.wiring.csp_numba.output_handlers import register as register_output_handlers +from csp.impl.wiring.csp_numba.signal_support import register as register_signals +from csp.impl.wiring.csp_numba.struct_support import register as register_struct_support +from csp.impl.wiring.node import NodeDef, NodeDefMeta +from csp.impl.wiring.node_parser import NodeDefinitionParser + +__all__ = ( + "numba_node", + "NumbaNodeDef", +) + +# Lazily-initialized CompilationContext with all CSP-specific registrations. +# Created once on first use; reused for all subsequent compilations. +_csp_context: CompilationContext | None = None + + +def _get_csp_context() -> CompilationContext: + global _csp_context + if _csp_context is None: + ctx = CompilationContext() + with ctx: + register_all() + register_signals() + register_input_handlers() + register_output_handlers() + register_enum_support() + register_struct_support() + _csp_context = ctx + return _csp_context + + +class _NumbaNodeDefinitionParser(NodeDefinitionParser): + def parse_func_signature(self, funcdef): + inputs, defaults, outputs = super().parse_func_signature(funcdef) + normalized_inputs = [] + for input_def in inputs: + origin = get_origin(input_def.typ) + if origin is NumbaList: + input_def = input_def._replace(typ=list[get_args(input_def.typ)[0]]) + elif origin is NumbaDict: + input_def = input_def._replace(typ=dict[get_args(input_def.typ)]) + normalized_inputs.append(input_def) + return normalized_inputs, defaults, outputs + + +class NumbaNodeDefMeta(NodeDefMeta): + def _instantiate_impl(self, __forced_tvars, name, args, kwargs): + inputs, scalars, tvars = self._signature.parse_inputs(__forced_tvars, *args, **kwargs) + input_values = iter(inputs) + scalar_values = iter(scalars) + compiler_kwargs = {} + for input_def in self._signature.inputs: + if input_def.kind.is_any_ts(): + value = next(input_values) + compiler_kwargs[input_def.name] = dict(enumerate(value)) if isinstance(value, (list, tuple)) else value + else: + compiler_kwargs[input_def.name] = next(scalar_values) + + compilation = self._compile_call(**compiler_kwargs) + flattened_inputs = tuple(compilation.ordered_input_signals) + nodedef = type.__call__(self, flattened_inputs, scalars, tvars, self._impl, self._pre_create_hook, compilation) + return self._finalize_nodedef(nodedef, inputs, scalars, tvars, name) + + +class NumbaNodeDef(NodeDef): + def __init__(self, inputs, scalars, tvars, impl, pre_create_hook, compilation): + super().__init__(inputs, scalars, tvars, impl, pre_create_hook) + self._compilation = compilation + + def _create(self, engine, memo): + if self._pre_create_hook: + self._pre_create_hook(engine, memo) + + compiled_func = self._compilation.compiled_func + ordered_inputs = tuple(edge for _, edge in self.ts_inputs()) + + node = _cspimpl.PyNumbaNode( + engine, + compiled_func.address, + ordered_inputs, + self._output_types, + self._compilation.state_values, + compiled_func, + ) + + self._create_outputs(node) + return node + + +def numba_node( + func: Optional[Callable] = None, + *, + name: str = None, + globals: dict = None, +): + """ + Decorator that transforms a CSP node function into a Numba-compiled graph node. + + Args: + func: The function to transform + name: Custom name for the node type + globals: Additional globals to make available during compilation + (e.g., enum types defined in enclosing scopes) + + Returns: + A wrapper function that creates the numba node and returns Edge(s) + + Example: + @numba_node + def my_node(x: ts[float], y: ts[int]) -> ts[float]: + with csp.state(): + total = 0.0 + if csp.ticked(x, y): + total = total + x + y + return total + + """ + + func_frame = inspect.currentframe().f_back + + def _impl(func: Callable) -> Callable: + node_name = name or func.__name__ + parser = _NumbaNodeDefinitionParser(node_name, func, func_frame) + parser.parse() + definition = parser.definition + if definition.signature.alarms(): + raise NotImplementedError("numba_node does not support alarms") + + transformer = CspNodeTransformer() + transformed = transformer.transform_csp_node(definition) + transformed_ast = transformed.transformed_ast + + transformed_globals = func.__globals__.copy() + + if func.__code__.co_freevars and func.__closure__: + for var_name, cell in zip(func.__code__.co_freevars, func.__closure__): + transformed_globals[var_name] = cell.cell_contents + + if globals: + transformed_globals.update(globals) + + def get_edge_type(edge) -> type: + return edge.tstype.typ + + original_signature = inspect.signature( + func, + globals=transformed_globals, + locals=func_frame.f_locals, + eval_str=True, + ) + compiler_signature = original_signature.replace( + parameters=[ + param.replace( + annotation=CspInputMetadata( + category=CspInputCategory.SIGNAL_SET + if definition.signature.input(param.name).kind.is_basket() + else CspInputCategory.SIGNAL + ) + ) + if definition.signature.input(param.name).kind.is_any_ts() + else param + for param in original_signature.parameters.values() + ] + ) + + def compile_call(*args, **kwargs): + with _get_csp_context(): + for value in transformed_globals.values(): + if isinstance(value, type) and issubclass(value, IntEnum): + register_family(value) + return create_compiled_func( + transformed_ast, + *args, + extract_python_type_fn=get_edge_type, + decorator_name="@numba_node", + func_globals=transformed_globals, + signature=compiler_signature, + call_globals=transformed_globals, + start_body=transformed.start_body, + stop_body=transformed.stop_body, + options=CompilationOptions(return_status=True), + **kwargs, + ) + + nodetype = NumbaNodeDefMeta( + node_name, + (NumbaNodeDef,), + { + "_signature": definition.signature, + "_impl": func, + "_compile_call": staticmethod(compile_call), + "memoize": True, + "force_memoize": False, + "_cppimpl": None, + "_pre_create_hook": None, + "__wrapped__": func, + "__module__": func.__module__, + "__doc__": func.__doc__, + }, + ) + return nodetype + + if func is None: + return _impl + else: + return _impl(func) diff --git a/csp/impl/wiring/csp_numba/output_handlers.py b/csp/impl/wiring/csp_numba/output_handlers.py new file mode 100644 index 000000000..9e41248cc --- /dev/null +++ b/csp/impl/wiring/csp_numba/output_handlers.py @@ -0,0 +1,86 @@ +""" +CSP output handlers for numba_node. + +Provides: +- SingleTsOutputHandler for ts[type] return annotation +- CspOutputsHandler for csp.Outputs(name=ts[type], ...) return annotation +""" + +import ast +from typing import Any, Optional, get_args, get_origin + +from numba_cfunc_compiler.core.analysis import OutputAnalysis, OutputTypeHandler +from numba_cfunc_compiler.types.factory import HostTypeFactory + +from csp.impl.types.tstype import TsType + + +def validate_output_type(t: Any, context: str) -> None: + try: + HostTypeFactory.resolve(t) + except TypeError as exc: + raise TypeError(f"Unsupported type '{t}' for {context}") from exc + + +def _extract_ts_inner_type(ann: Any) -> Optional[type]: + if get_origin(ann) is not TsType: + return None + args = get_args(ann) + return args[0] if args else getattr(ann, "typ", None) + + +class SingleTsOutputHandler(OutputTypeHandler): + """Handles single ts[type] output annotations.""" + + def try_parse(self, return_annotation: Any, ast_tree: ast.AST) -> Optional[OutputAnalysis]: + if get_origin(return_annotation) is not TsType: + return None + + output_type = _extract_ts_inner_type(return_annotation) + if output_type is None: + raise TypeError("Return type ts[type] is missing type argument.") + + validate_output_type(output_type, "output[0]") + + return OutputAnalysis(outputs={None: output_type}) + + +class CspOutputsHandler(OutputTypeHandler): + """Handles csp.Outputs(name=ts[type], ...) for multiple named outputs.""" + + def try_parse(self, return_annotation: Any, ast_tree: ast.AST) -> Optional[OutputAnalysis]: + if not isinstance(return_annotation, type): + return None + if not hasattr(return_annotation, "__annotations__"): + return None + + annotations = return_annotation.__annotations__ + if not annotations: + return None + + named_outputs = {} + + for name, spec in annotations.items(): + if name == "__annotations__": + continue + + output_type = _extract_ts_inner_type(spec) + if output_type is None and hasattr(spec, "typ"): + output_type = spec.typ + if output_type is None: + return None + + validate_output_type(output_type, f"output[{name}]") + named_outputs[name] = output_type + + if not named_outputs: + return None + + return OutputAnalysis(outputs=named_outputs) + + +def register(): + from numba_cfunc_compiler.core.analysis import FunctionAnalyzer + + FunctionAnalyzer.register_output_handler(SingleTsOutputHandler()) + FunctionAnalyzer.register_output_handler(CspOutputsHandler()) diff --git a/csp/impl/wiring/csp_numba/signal_support.py b/csp/impl/wiring/csp_numba/signal_support.py new file mode 100644 index 000000000..25cdb6d9c --- /dev/null +++ b/csp/impl/wiring/csp_numba/signal_support.py @@ -0,0 +1,123 @@ +"""CSP callback inputs. Basket iterators yield slot positions.""" + +from numba import types +from numba_cfunc_compiler.core.variable_access import KeyedInputAccess, SourceAccess +from numba_cfunc_compiler.extension.callback_components import ( + CallbackComponent, + CfuncParam, + ComponentRegistry, + MaterializationPhase, +) +from numba_cfunc_compiler.extension.input_sources import ( + SLOT_INDEX, + KeyFilter, + SourceField, + register_input_source, + register_keyed_input, +) +from numba_cfunc_compiler.types.factory import HostTypeFactory + +from csp.impl.wiring.csp_numba.input_handlers import CspInputCategory + +ACTIVE_SOURCE = register_input_source( + identity=("csp", "signal.active.v1"), + fields=( + SourceField("valid", types.CPointer(types.int8)), + SourceField("ticked", types.CPointer(types.int8)), + SourceField("index", types.intp), + ), +) + + +@ACTIVE_SOURCE.method("valid") +def signal_valid(signal): + return signal.source.valid[signal.source.index] != 0 + + +@ACTIVE_SOURCE.method("ticked") +def signal_ticked(signal): + return signal.source.ticked[signal.source.index] != 0 + + +ACTIVE_BASKET = register_keyed_input( + identity=("csp", "basket.active.v1"), + element_source=ACTIVE_SOURCE, + fields=( + SourceField("slots", types.CPointer(types.voidptr)), + SourceField("valid_flags", types.CPointer(types.int8)), + SourceField("ticked_flags", types.CPointer(types.int8)), + SourceField("start", types.intp), + SourceField("length", types.intp), + ), + element_projection={"valid": "valid_flags", "ticked": "ticked_flags", "index": SLOT_INDEX}, +) +ACTIVE_BASKET.iter_method("keys", filter=KeyFilter.all()) +ACTIVE_BASKET.iter_method("validkeys", filter=KeyFilter.nonzero("valid_flags")) +ACTIVE_BASKET.iter_method("tickedkeys", filter=KeyFilter.nonzero("ticked_flags")) + + +@ACTIVE_BASKET.method("valid") +def basket_valid(basket): + for position in range(basket.length): + if basket.valid_flags[basket.start + position] == 0: + return False + return True + + +@ACTIVE_BASKET.method("ticked") +def basket_ticked(basket): + for position in range(basket.length): + if basket.ticked_flags[basket.start + position] != 0: + return True + return False + + +class SignalComponent(CallbackComponent): + id = "csp.signal" + order = 0 + materialization_phase = MaterializationPhase.EXECUTE + + @property + def cfunc_params(self): + return [ + CfuncParam("inputs", "CPointer(voidptr)"), + CfuncParam("input_ticked", "CPointer(int8)"), + CfuncParam("input_valid", "CPointer(int8)"), + ] + + def create_variables(self, inputs, factory): + ordered = [] + for name, edge in inputs.input_analysis.get_by_category(CspInputCategory.SIGNAL).items(): + index = len(ordered) + factory.add_variable( + SourceAccess( + index, + HostTypeFactory.resolve(inputs.extract_python_type_fn(edge)), + name, + "inputs", + ACTIVE_SOURCE, + {"valid": "input_valid", "ticked": "input_ticked", "index": index}, + ), + component=self.id, + ) + ordered.append(edge) + for name, basket in inputs.input_analysis.get_by_category(CspInputCategory.SIGNAL_SET).items(): + children = {} + for key, edge in basket.items(): + index = len(ordered) + children[key] = (index, HostTypeFactory.resolve(inputs.extract_python_type_fn(edge))) + ordered.append(edge) + factory.add_variable( + KeyedInputAccess( + name, + children, + ACTIVE_BASKET, + {"slots": "inputs", "valid_flags": "input_valid", "ticked_flags": "input_ticked"}, + ), + component=self.id, + ) + return {"ordered_input_signals": tuple(ordered)} + + +def register(): + ComponentRegistry.register(SignalComponent()) diff --git a/csp/impl/wiring/csp_numba/struct_native.py b/csp/impl/wiring/csp_numba/struct_native.py new file mode 100644 index 000000000..ee7ff37dd --- /dev/null +++ b/csp/impl/wiring/csp_numba/struct_native.py @@ -0,0 +1,289 @@ +"""CSP-owned Numba view of native structs, including presence and enum fields.""" + +import ctypes +import hashlib +from dataclasses import dataclass +from functools import cache + +from llvmlite import binding as llvm, ir +from numba import types +from numba.core.errors import TypingError +from numba.core.typing.templates import AttributeTemplate +from numba.extending import ( + infer_getattr, + intrinsic, + lower_getattr_generic, + lower_setattr_generic, + models, + register_model, +) +from numba_cfunc_compiler.types.native.ffi import declare_symbol +from numba_cfunc_compiler.types.native.state import StatePayload +from numba_cfunc_compiler.types.policy import ValueSemantics + +from csp.impl.__cspimpl import _cspimpl +from csp.impl.types.typing_utils import CspTypingUtils +from csp.impl.wiring.csp_numba.enum_support import register_family + +_PRIMITIVES = { + "BOOL": types.uint8, + "INT8": types.int8, + "UINT8": types.uint8, + "INT16": types.int16, + "UINT16": types.uint16, + "INT32": types.int32, + "UINT32": types.uint32, + "INT64": types.int64, + "UINT64": types.uint64, + "DOUBLE": types.float64, + "FLOAT": types.float32, +} + +_LIBRARY = None +_SYMBOLS = ( + "csp_numba_struct_enum_field_value", + "csp_numba_struct_enum_field_set", +) + + +def register_symbols(): + """Make CSP's exported native helpers available to Numba's JIT linker.""" + global _LIBRARY + if _LIBRARY is None: + _LIBRARY = ctypes.CDLL(_cspimpl.__file__) + for name in _SYMBOLS: + address = ctypes.cast(getattr(_LIBRARY, name), ctypes.c_void_p).value + llvm.add_symbol(name, address) + + +@dataclass(frozen=True) +class CspField: + name: str + offset: int + size: int + mask_offset: int + mask_bitmask: int + numba_type: types.Type + is_enum: bool + + @property + def key(self): + return ( + self.name, + self.offset, + self.size, + self.mask_offset, + self.mask_bitmask, + self.numba_type.name, + self.is_enum, + ) + + +def _fields(py_type, metadata): + fields = [] + optional_fields = ( + { + name + for name, annotation in py_type.metadata(typed=True).items() + if CspTypingUtils.is_optional_type(annotation) + } + if py_type.is_strict() + else set() + ) + for item in metadata.get("fields", ()): + name = item["fieldname"] + if name in optional_fields: + raise TypeError( + f"CSP struct field '{name}' is optional; optional native fields are not supported by numba_node" + ) + kind = item["type"].get("type", "STRUCT") + is_enum = kind == "ENUM" + if is_enum: + numba_type = register_family(item["type"]["pytype"]).value_type + else: + numba_type = _PRIMITIVES.get(kind) + if numba_type is None: + raise TypeError(f"CSP struct field '{name}' has unsupported native type '{kind}'") + if item["size"] != numba_type.bitwidth // 8: + raise ValueError(f"CSP struct field '{name}' has unexpected size {item['size']}") + fields.append( + CspField( + name, + item["offset"], + item["size"], + item["mask_offset"], + item["mask_bitmask"], + numba_type, + is_enum, + ) + ) + return tuple(fields) + + +class CspStructViewType(types.Type): + """Nominal pointer to a CSP-owned native struct.""" + + def __init__(self, py_type, metadata): + self.python_type = py_type # Retain the PyStructMeta used by native helpers. + self.fields = _fields(py_type, metadata) + self.host_size = metadata["size"] + self.alignment = max((ctypes.alignment(ctypes.c_void_p), *(item["alignment"] for item in metadata["fields"]))) + self.schema_key = ("csp-struct-view-v1", id(py_type), self.host_size, tuple(field.key for field in self.fields)) + fingerprint = hashlib.sha256(repr(self.schema_key).encode()).hexdigest()[:16] + super().__init__(f"CspStructView[{py_type.__module__}.{py_type.__qualname__}:{fingerprint}]") + + @property + def key(self): + return self.python_type, self.schema_key + + def get_field(self, name): + return next((field for field in self.fields if field.name == name), None) + + +@register_model(CspStructViewType) +class CspStructViewModel(models.PrimitiveModel): + def __init__(self, dmm, fe_type): + super().__init__(dmm, fe_type, ir.IntType(8).as_pointer()) + + +@cache +def csp_struct_view_type(py_type): + metadata = py_type._metadata_info() + if not metadata.get("is_native", False): + raise TypeError(f"numba_node only supports native csp.Struct types; '{py_type.__name__}' is non-native") + return CspStructViewType(py_type, metadata) + + +@cache +def csp_struct_view(py_type): + view_type = csp_struct_view_type(py_type) + + @intrinsic + def view(typingctx, raw_ptr): + if raw_ptr != types.voidptr: + return None + sig = view_type(raw_ptr) + + def codegen(context, builder, signature, args): + return builder.bitcast(args[0], context.get_value_type(view_type)) + + return sig, codegen + + return view + + +@cache +def _reject_replacement(py_type, name): + view_type = csp_struct_view_type(py_type) + + @intrinsic + def store(typingctx, slot, rhs): + raise TypingError(f"State[{view_type.python_type.__name__}] '{name}' does not support whole-struct replacement") + + return store + + +@cache +def csp_struct_payload(py_type): + view_type = csp_struct_view_type(py_type) + return StatePayload( + label=f"{py_type.__module__}.{py_type.__qualname__}", + semantics=ValueSemantics.BORROWED_VIEW, + native_type=view_type, + host_size=view_type.host_size, + alignment=view_type.alignment, + layout_key=view_type.schema_key, + replacement=False, + bind_function=csp_struct_view(py_type), + store_factory=lambda name: _reject_replacement(py_type, name), + mutable=True, + copy_output_bytes=True, + ) + + +@infer_getattr +class CspStructAttributes(AttributeTemplate): + key = CspStructViewType + + def generic_resolve(self, typ, attr): + field = typ.get_field(attr) + return field.numba_type if field is not None else None + + +def _type_pointer(context, builder, typ): + return builder.inttoptr(context.get_constant(types.intp, id(typ.python_type)), ir.IntType(8).as_pointer()) + + +def _field_offset(context, field): + return context.get_constant(types.int64, field.offset) + + +def _byte_pointer(context, builder, value, offset): + return builder.gep(value, [context.get_constant(types.intp, offset)]) + + +def _return_failure(context, builder): + # CSP compiles its callbacks with an int8 status result: 1 succeeds, 0 fails. + context.call_conv.return_value(builder, ir.Constant(ir.IntType(8), 0)) + + +def _guard_struct_pointer(context, builder, value): + null = builder.icmp_unsigned("==", value, ir.Constant(value.type, None)) + with builder.if_then(null, likely=False): + _return_failure(context, builder) + + +def _guard_field_read(context, builder, value, field): + _guard_struct_pointer(context, builder, value) + mask = builder.load(_byte_pointer(context, builder, value, field.mask_offset), align=1) + unset = builder.icmp_unsigned( + "==", builder.and_(mask, ir.Constant(ir.IntType(8), field.mask_bitmask)), ir.Constant(ir.IntType(8), 0) + ) + with builder.if_then(unset, likely=False): + _return_failure(context, builder) + + +@lower_getattr_generic(CspStructViewType) +def csp_struct_getattr(context, builder, typ, value, attr): + field = typ.get_field(attr) + if field is None: + raise AttributeError(attr) + _guard_field_read(context, builder, value, field) + if field.is_enum: + native_sig = ir.FunctionType(ir.IntType(64), [ir.IntType(8).as_pointer(), ir.IntType(64)]) + function = declare_symbol(builder, "csp_numba_struct_enum_field_value", native_sig) + return builder.call(function, [value, _field_offset(context, field)]) + + result_type = context.get_value_type(field.numba_type) + pointer = builder.bitcast(_byte_pointer(context, builder, value, field.offset), result_type.as_pointer()) + return builder.load(pointer, align=1) + + +@lower_setattr_generic(CspStructViewType) +def csp_struct_setattr(context, builder, sig, args, attr): + typ, value_type = sig.args + field = typ.get_field(attr) + if field is None: + raise AttributeError(attr) + value = context.cast(builder, args[1], value_type, field.numba_type) + _guard_struct_pointer(context, builder, args[0]) + if field.is_enum: + native_sig = ir.FunctionType( + ir.IntType(8), [ir.IntType(8).as_pointer(), ir.IntType(64), ir.IntType(64), ir.IntType(8).as_pointer()] + ) + function = declare_symbol(builder, "csp_numba_struct_enum_field_set", native_sig) + success = builder.call( + function, [args[0], _field_offset(context, field), value, _type_pointer(context, builder, typ)] + ) + with builder.if_then(builder.icmp_unsigned("==", success, ir.Constant(ir.IntType(8), 0)), likely=False): + _return_failure(context, builder) + return + + pointer = builder.bitcast( + _byte_pointer(context, builder, args[0], field.offset), + context.get_value_type(field.numba_type).as_pointer(), + ) + builder.store(value, pointer, align=1) + mask_pointer = _byte_pointer(context, builder, args[0], field.mask_offset) + mask = builder.load(mask_pointer, align=1) + builder.store(builder.or_(mask, ir.Constant(ir.IntType(8), field.mask_bitmask)), mask_pointer, align=1) diff --git a/csp/impl/wiring/csp_numba/struct_support.py b/csp/impl/wiring/csp_numba/struct_support.py new file mode 100644 index 000000000..d52c5daa8 --- /dev/null +++ b/csp/impl/wiring/csp_numba/struct_support.py @@ -0,0 +1,44 @@ +"""Host registration for CSP-owned native struct views.""" + +import ast + +from numba_cfunc_compiler.types.base import HostType +from numba_cfunc_compiler.types.factory import HostTypeFactory + +from csp.impl.struct import Struct as CspStruct +from csp.impl.wiring.csp_numba.struct_native import csp_struct_payload, register_symbols + + +class CspStructHostType(HostType): + @classmethod + def is_type_supported(cls, var_type): + return isinstance(var_type, type) and issubclass(var_type, CspStruct) and var_type is not CspStruct + + def get_state_payload(self): + return csp_struct_payload(self.value) + + def is_opaque_pointer(self): + return True + + @classmethod + def try_parse_state(cls, node, var_name, globalns): + try: + struct_type = eval(ast.unparse(node.annotation.slice), globalns) + except (NameError, AttributeError): + return None + if not cls.is_type_supported(struct_type): + return None + csp_struct_payload(struct_type) + value = ( + struct_type() + if node.value is None or isinstance(node.value, ast.Constant) and node.value.value is None + else eval(ast.unparse(node.value), globalns) + ) + if not isinstance(value, struct_type): + raise TypeError(f"State '{var_name}' requires a {struct_type.__name__} initializer") + return value, struct_type + + +def register(): + register_symbols() + HostTypeFactory.register(CspStructHostType, priority=0) diff --git a/csp/impl/wiring/node.py b/csp/impl/wiring/node.py index a905bee36..1d83546cb 100644 --- a/csp/impl/wiring/node.py +++ b/csp/impl/wiring/node.py @@ -29,11 +29,7 @@ def __init__(self, *args, **kwargs): python_kwargs["_cppimpl"] = None self.python = NodeDefMeta(self.__name__, self.__bases__, python_kwargs) - def _instantiate_impl(self, __forced_tvars, name, args, kwargs): - ## Parse inputs - inputs, scalars, tvars = self._signature.parse_inputs(__forced_tvars, *args, **kwargs) - nodedef = super().__call__(inputs, scalars, tvars, self._impl, self._pre_create_hook) - + def _finalize_nodedef(self, nodedef, inputs, scalars, tvars, name): basket_shape_eval_inputs = list(scalars) for input in inputs: if isinstance(input, list) or isinstance(input, dict): @@ -44,6 +40,12 @@ def _instantiate_impl(self, __forced_tvars, name, args, kwargs): nodedef.__name__ = name if name else self.__name__ return outputs + def _instantiate_impl(self, __forced_tvars, name, args, kwargs): + ## Parse inputs + inputs, scalars, tvars = self._signature.parse_inputs(__forced_tvars, *args, **kwargs) + nodedef = super().__call__(inputs, scalars, tvars, self._impl, self._pre_create_hook) + return self._finalize_nodedef(nodedef, inputs, scalars, tvars, name) + def _instantiate(self, __forced_tvars, name=None, *args, **kwargs): return self._instantiate_func(__forced_tvars, name=name, args=args, kwargs=kwargs) @@ -146,6 +148,10 @@ def ts_inputs(self): ts_idx += 1 + def _create_outputs(self, node): + for idx, output_type in enumerate(self._output_types): + node.create_output(idx, output_type) + def _create(self, engine, memo): if self._pre_create_hook: self._pre_create_hook(engine, memo) @@ -178,8 +184,7 @@ def _create(self, engine, memo): gen = self._impl(*self._scalars) node = _cspimpl.PyNode(engine, inputs, self._output_types, gen) - for idx, output_type in enumerate(self._output_types): - node.create_output(idx, output_type) + self._create_outputs(node) for idx, alarm in enumerate(alarms): node.create_alarm(idx, ContainerTypeNormalizer.normalized_type_to_actual_python_type(alarm.typ.typ)) diff --git a/csp/impl/wiring/node_parser.py b/csp/impl/wiring/node_parser.py index f1204bd2c..9d064395f 100644 --- a/csp/impl/wiring/node_parser.py +++ b/csp/impl/wiring/node_parser.py @@ -2,6 +2,7 @@ import copy import inspect import textwrap +from dataclasses import dataclass from warnings import warn from csp.impl import builtin_functions @@ -68,34 +69,27 @@ def __call__(self, proxy, node): ) -class NodeParser(BaseParser): - _CSP_NOW_FUNC = "_csp_now" - _CSP_ENGINE_START_TIME_FUNC = "_engine_start_time" - _CSP_ENGINE_END_TIME_FUNC = "_engine_end_time" - _CSP_ENGINE_STATS_FUNC = "_csp_engine_stats" +@dataclass(frozen=True) +class ParsedNodeBlocks: + state: tuple[ast.stmt, ...] + start: tuple[ast.stmt, ...] + stop: tuple[ast.stmt, ...] + body: tuple[ast.stmt, ...] - _CSP_STOP_ENGINE_FUNC = "_csp_stop_engine" - _CSP_IN_REALTIME_FUNC = "_csp_in_realtime" - _LOCAL_METHODS = { - _CSP_NOW_FUNC: _cspimpl._csp_now, - _CSP_ENGINE_START_TIME_FUNC: _cspimpl._csp_engine_start_time, - _CSP_ENGINE_END_TIME_FUNC: _cspimpl._csp_engine_end_time, - _CSP_STOP_ENGINE_FUNC: _cspimpl._csp_stop_engine, - _CSP_ENGINE_STATS_FUNC: _cspimpl._csp_engine_stats, - _CSP_IN_REALTIME_FUNC: _cspimpl._csp_in_realtime, - } +@dataclass(frozen=True) +class ParsedNodeDefinition: + funcdef: ast.FunctionDef + signature: Signature + blocks: ParsedNodeBlocks + + +class NodeDefinitionParser(BaseParser): _SPECIAL_BLOCKS_METH = {"alarms", "state", "start", "stop", "outputs"} _SPECIAL_BLOCKS_DEPR = {f"__{method}__" for method in _SPECIAL_BLOCKS_METH} _SPECIAL_BLOCKS_PYTHONIC = {f"{method}" for method in _SPECIAL_BLOCKS_METH if method != "outputs"} _SPECIAL_BLOCKS = _SPECIAL_BLOCKS_DEPR.union(_SPECIAL_BLOCKS_PYTHONIC) - # These match up with const names defined in PyNode.cpp - _NODE_P_VARNAME = "node_p" - _INPUT_VAR_VARNAME = "input_var" - _INPUT_PROXY_VARNAME = "input_proxy" - _OUTPUT_PROXY_VARNAME = "output_proxy" - def __init__(self, name, raw_func, func_frame, debug_print=False): super().__init__( name=name, @@ -106,13 +100,8 @@ def __init__(self, name, raw_func, func_frame, debug_print=False): self._stateblock = [] self._startblock = [] self._stopblock = [] + self._definition = None self._func_globals_modified.update(builtin_functions.CSP_BUILTIN_CONTEXT_DICT) - self._func_globals_modified.update(self._LOCAL_METHODS) - self._gen = None - - # To catch returning from within a for or while loop, which wouldnt work as it seems - self._inner_loop_count = 0 - self._returned_outputs = set() @_pythonic_depr_warning def _parse_alarms(self, node): @@ -233,6 +222,124 @@ def _parse_single_output_definition(self, name, arg_type_node, ts_idx, typ=None) def _extract_outputs_from_return_annotation(self, returns, *args, **kwargs): return super()._extract_outputs_from_return_annotation(returns=returns, enforce_shape=True) + def _parse_special_blocks(self, body): + # skip doc string + for index, node in enumerate(body): + if not isinstance(node, ast.Expr) or not ( + isinstance(node.value, ast.Constant) and isinstance(node.value.value, str) + ): + break + + last_special_block = None + + def consume_block(*args, **kwargs): + cur_index, cur_block_name = self._consume_special_block(*args, **kwargs) + if cur_block_name: + return cur_index, cur_block_name + return cur_index, last_special_block + + # Ensure blocks are processed in proper order i.e. start before stop + index, last_special_block = consume_block( + body, index, "alarms", self._parse_alarms, allow_with_block=True, allow_flat_call=False + ) + index, last_special_block = consume_block( + body, index, "__alarms__", self._parse_alarms, allow_with_block=True, allow_flat_call=True + ) + + index, last_special_block = consume_block( + body, index, "state", self._parse_state, allow_with_block=True, allow_flat_call=False + ) + index, last_special_block = consume_block( + body, index, "__state__", self._parse_state, allow_with_block=True, allow_flat_call=True + ) + + index, last_special_block = consume_block( + body, index, "start", self._parse_start, allow_with_block=True, allow_flat_call=False + ) + index, last_special_block = consume_block( + body, index, "__start__", self._parse_start, allow_with_block=True, allow_flat_call=False + ) + + index, last_special_block = consume_block( + body, index, "stop", self._parse_stop, allow_with_block=True, allow_flat_call=False + ) + index, last_special_block = consume_block( + body, index, "__stop__", self._parse_stop, allow_with_block=True, allow_flat_call=False + ) + + if index < len(body): + first_body_elem = body[index] + if self._is_special_with_or_call(first_body_elem, self._SPECIAL_BLOCKS): + for special_block_name in self._SPECIAL_BLOCKS: + if self._is_special_with_or_call(first_body_elem, special_block_name): + # For sure we must have previous block, it's a matter of out of order. Otherwise we can't get there + raise CspParseError( + f"{special_block_name} must be declared before {last_special_block}", + lineno=first_body_elem.lineno, + file=self._func_filename, + ) + # We should never get here + raise AssertionError("Special block error wasn't handled correctly") + return index + + def _is_ts_args_removed_from_signature(self): + return True + + def _parse_impl(self): + self._inputs, input_defaults, self._outputs = self.parse_func_signature(self._funcdef) + body_index = self._parse_special_blocks(self._funcdef.body) + self._signature = Signature( + self._name, + self._inputs, + self._outputs + self._special_outputs, + input_defaults, + special_outputs=self._special_outputs, + ) + self._definition = ParsedNodeDefinition( + funcdef=self._funcdef, + signature=self._signature, + blocks=ParsedNodeBlocks( + state=tuple(self._stateblock), + start=tuple(self._startblock), + stop=tuple(self._stopblock), + body=tuple(self._funcdef.body[body_index:]), + ), + ) + + @property + def definition(self): + return self._definition + + +class NodeParser(NodeDefinitionParser): + # These match up with const names defined in PyNode.cpp + _NODE_P_VARNAME = "node_p" + _INPUT_VAR_VARNAME = "input_var" + _INPUT_PROXY_VARNAME = "input_proxy" + _OUTPUT_PROXY_VARNAME = "output_proxy" + + _CSP_NOW_FUNC = "_csp_now" + _CSP_ENGINE_START_TIME_FUNC = "_engine_start_time" + _CSP_ENGINE_END_TIME_FUNC = "_engine_end_time" + _CSP_ENGINE_STATS_FUNC = "_csp_engine_stats" + _CSP_STOP_ENGINE_FUNC = "_csp_stop_engine" + _CSP_IN_REALTIME_FUNC = "_csp_in_realtime" + _LOCAL_METHODS = { + _CSP_NOW_FUNC: _cspimpl._csp_now, + _CSP_ENGINE_START_TIME_FUNC: _cspimpl._csp_engine_start_time, + _CSP_ENGINE_END_TIME_FUNC: _cspimpl._csp_engine_end_time, + _CSP_STOP_ENGINE_FUNC: _cspimpl._csp_stop_engine, + _CSP_ENGINE_STATS_FUNC: _cspimpl._csp_engine_stats, + _CSP_IN_REALTIME_FUNC: _cspimpl._csp_in_realtime, + } + + def __init__(self, name, raw_func, func_frame, debug_print=False): + super().__init__(name=name, raw_func=raw_func, func_frame=func_frame, debug_print=debug_print) + self._func_globals_modified.update(self._LOCAL_METHODS) + self._gen = None + self._inner_loop_count = 0 + self._returned_outputs = set() + def _node_proxy_expr(self, ctx=ast.Load()): return ast.Name(id="#node_p", ctx=ctx) @@ -684,66 +791,6 @@ def _validate_return_statements(self): if output.name not in self._returned_outputs: raise CspParseError(f"output '{output.name}' is never returned") - def _parse_special_blocks(self, body): - # skip doc string - for index, node in enumerate(body): - if not isinstance(node, ast.Expr) or not ( - isinstance(node.value, ast.Constant) and isinstance(node.value.value, str) - ): - break - - last_special_block = None - - def consume_block(*args, **kwargs): - cur_index, cur_block_name = self._consume_special_block(*args, **kwargs) - if cur_block_name: - return cur_index, cur_block_name - return cur_index, last_special_block - - # Ensure blocks are processed in proper order i.e. start before stop - index, last_special_block = consume_block( - body, index, "alarms", self._parse_alarms, allow_with_block=True, allow_flat_call=False - ) - index, last_special_block = consume_block( - body, index, "__alarms__", self._parse_alarms, allow_with_block=True, allow_flat_call=True - ) - - index, last_special_block = consume_block( - body, index, "state", self._parse_state, allow_with_block=True, allow_flat_call=False - ) - index, last_special_block = consume_block( - body, index, "__state__", self._parse_state, allow_with_block=True, allow_flat_call=True - ) - - index, last_special_block = consume_block( - body, index, "start", self._parse_start, allow_with_block=True, allow_flat_call=False - ) - index, last_special_block = consume_block( - body, index, "__start__", self._parse_start, allow_with_block=True, allow_flat_call=False - ) - - index, last_special_block = consume_block( - body, index, "stop", self._parse_stop, allow_with_block=True, allow_flat_call=False - ) - index, last_special_block = consume_block( - body, index, "__stop__", self._parse_stop, allow_with_block=True, allow_flat_call=False - ) - - if index < len(body): - first_body_elem = body[index] - if self._is_special_with_or_call(first_body_elem, self._SPECIAL_BLOCKS): - for special_block_name in self._SPECIAL_BLOCKS: - if self._is_special_with_or_call(first_body_elem, special_block_name): - # For sure we must have previous block, it's a matter of out of order. Otherwise we can't get there - raise CspParseError( - f"{special_block_name} must be declared before {last_special_block}", - lineno=first_body_elem.lineno, - file=self._func_filename, - ) - # We should never get here - raise AssertionError("Special block error wasn't handled correctly") - return index - @classmethod def _create_ast_args( cls, posonlyargs=[], args=[], kwonlyargs=[], defaults=[], vararg=None, kwarg=None, kw_defaults=[] @@ -758,19 +805,8 @@ def _create_ast_args( kw_defaults=kw_defaults, ) - def _is_ts_args_removed_from_signature(self): - return True - def _parse_impl(self): - self._inputs, input_defaults, self._outputs = self.parse_func_signature(self._funcdef) - idx = self._parse_special_blocks(self._funcdef.body) - self._signature = Signature( - self._name, - self._inputs, - self._outputs + self._special_outputs, - input_defaults, - special_outputs=self._special_outputs, - ) + super()._parse_impl() # Up front initialize timeseries inputs as local vars. Regular TS need two vars, the proxy and # the python value instance. Baskets will only need the proxy. We also provide a noderef for node # specific methods. Values are assigned to tuple( name, index ) which is replaced in PyNode.cpp init @@ -820,8 +856,8 @@ def _parse_impl(self): # innerbody is the while loop. Start off with a yield, will get called when something ticks innerbody = [ast.Expr(value=ast.Yield(value=None))] - for x in range(idx, len(self._funcdef.body)): - func_body_transformed = self.visit(self._funcdef.body[x]) + for body_node in self._definition.blocks.body: + func_body_transformed = self.visit(body_node) if isinstance(func_body_transformed, list): innerbody.extend(func_body_transformed) else: diff --git a/csp/impl/wiring/numba_node.py b/csp/impl/wiring/numba_node.py deleted file mode 100644 index a6ad5bfe6..000000000 --- a/csp/impl/wiring/numba_node.py +++ /dev/null @@ -1,217 +0,0 @@ -import inspect -import types -import typing - -from csp.impl.__cspimpl import _cspimpl -from csp.impl.error_handling import ExceptionContext -from csp.impl.mem_cache import csp_memoized_graph_object, function_full_name -from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer -from csp.impl.types.generic_values_resolver import GenericValuesResolver -from csp.impl.wiring.edge import Edge -from csp.impl.wiring.node import NodeDef, _create_node -from csp.impl.wiring.numba_node_parser import NumbaNodeParser - - -class NumbaNumbaNodeDefMeta(type): - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) - self._instantiate_func = self._instantiate_impl - if self.memoize or self.force_memoize: - self._instantiate_func = csp_memoized_graph_object( - self._instantiate_impl, force_memoize=self.force_memoize, function_name=function_full_name(self._impl) - ) - - def _instantiate_impl(self, __forced_tvars, args, kwargs): - ## Parse inputs - inputs, scalars, tvars = self._signature.parse_inputs(__forced_tvars, *args, **kwargs) - NumbaNodeDef = super().__call__( - inputs, scalars, self._state_types, tvars, self._impl, pre_create_hook=self._pre_create_hook - ) - outputs, output_types = self._signature.create_outputs(NumbaNodeDef, tvars) - NumbaNodeDef._outputs = output_types - return outputs - - def _instantiate(self, __forced_tvars, *args, **kwargs): - return self._instantiate_func(__forced_tvars, args=args, kwargs=kwargs) - - def __call__(cls, *args, **kwargs): - return cls._instantiate(None, *args, **kwargs) - - def using(cls, **__forced_tvars): - return lambda *args, **kwargs: cls._instantiate(__forced_tvars, *args, **kwargs) - - def __get__(self, instance, owner): - # This little bit of magic allows nodes to be defined as class members, this intercepts method calls - # and converts them to bound methods - if instance is not None: - return types.MethodType(self, instance) - return self - - -##Every NumbaNodeDef instance represents an instance of a wiring-time node -class NumbaNodeDef(NodeDef): - _COMPILED_NODES_CACHE = {} - - def __init__(self, inputs, scalars, state_types, tvars, impl, pre_create_hook): - super().__init__(inputs=inputs, scalars=scalars, tvars=tvars, impl=impl, pre_create_hook=pre_create_hook) - self._state_types = state_types - self._state_class = None - self._state_instance = None - self._init_func = None - self._next_func = None - - def _create(self, engine, memo): - from csp.impl.wiring.numba_utils.csp_cpp_numba_interface import NumbaTSTypedFunctionResolver, ffi_ptr_to_int - from csp.impl.wiring.numba_utils.numba_type_resolver import NumbaTypeResolver - - exposed_utility_values = {} - cache_key = [self._impl] - for (ts_idx, basket_idx), input in self.ts_inputs(): - exposed_utility_values[NumbaNodeParser.get_ts_input_value_getter_name(ts_idx)] = ( - NumbaTSTypedFunctionResolver.get_value_getter_function(input.tstype.typ) - ) - cache_key.append(input.tstype) - - for output in self._outputs: - assert output.kind.is_single_ts() - ts_idx = output.ts_idx - exposed_utility_values[NumbaNodeParser.get_ts_out_value_return_name(ts_idx)] = ( - NumbaTSTypedFunctionResolver.get_value_returner_function(output.typ.typ) - ) - - numba_scalar_types = [] - for i in self._signature.raw_inputs(): - if i.kind.is_scalar(): - resolved_python_type = GenericValuesResolver.resolve_generic_values(i.typ, self._tvars) - numba_type = NumbaTypeResolver.resolve_numba_type(resolved_python_type) - exposed_utility_values[NumbaNodeParser.get_arg_type_var_name(i.name)] = numba_type - cache_key.append(resolved_python_type) - numba_scalar_types.append(numba_type) - - for var_name, var_type in self._state_types.items(): - resolved_python_type = GenericValuesResolver.resolve_generic_values(var_type, self._tvars) - numba_type = NumbaTypeResolver.resolve_numba_type(resolved_python_type) - exposed_utility_values[NumbaNodeParser.get_state_type_var_name(var_name)] = numba_type - cache_key.append(resolved_python_type) - - cache_key = tuple(cache_key) - - compiled_node_tuple = self._COMPILED_NODES_CACHE.get(cache_key) - if compiled_node_tuple is None: - compiled_node_tuple = self._impl(*self._scalars, **exposed_utility_values) - self._COMPILED_NODES_CACHE[cache_key] = compiled_node_tuple - - self._state_class, self._init_func, self._next_func = compiled_node_tuple - - scalars = [NumbaTypeResolver.transform_scalar(v, t) for v, t in zip(self._scalars, numba_scalar_types)] - - self._state_instance = self._state_class(*scalars) if self._state_class is not None else None - # TODO: handle the following here: - # alarms - # baskets - - if len(self._signature.alarms()) != 0: - raise NotImplementedError("Alarms are not supported in numba nodes") - - # alarms = self._signature._create_alarms(self._tvars) - # inputs = [alarm.typ.typ for alarm in alarms] - inputs = [] - for input in self._inputs: - if isinstance(input, Edge): - inputs.append(input.tstype.typ) - else: - raise NotImplementedError(f"input of type {type(input)} is not supported") - # elif isinstance(input, list): - # inputs.append((len(input), input[0].tstype.typ)) - # else: # dict - # assert (isinstance(input, dict)) - # inputs.append((list(input.keys()), next(iter(input.values())).tstype.typ)) - - inputs = tuple(inputs) # tuple(x.tstype.typ if isinstance( x, Edge ) else ( len( x ), for x in self._inputs) - outputs = self._outputs - - if not isinstance(outputs, tuple): - outputs = tuple([outputs]) - - outputs = tuple(ContainerTypeNormalizer.normalized_type_to_actual_python_type(odef.typ.typ) for odef in outputs) - - init_callback_ptr_as_int = ffi_ptr_to_int(self._init_func.cffi) - callback_ptr_as_int = ffi_ptr_to_int(self._next_func.cffi) - state_ptr_as_int = self._state_instance.get_ptr() if self._state_instance is not None else 0 - - # We need the cpp to hold reference to all the relevant python objects so that they don't get garbage - # collected - data_reference = (self._init_func, self._next_func, self._state_instance, self._state_class) - - node = _cspimpl.PyNumbaNode( - engine, inputs, outputs, state_ptr_as_int, init_callback_ptr_as_int, callback_ptr_as_int, data_reference - ) - - for idx, output_type in enumerate(outputs): - node.create_output(idx, output_type) - - # for idx, alarm in enumerate(alarms): - # node.create_alarm(idx, alarm.typ.typ) - return node - - -# The decorator -def numba_node( - func=None, - *, - memoize=True, - force_memoize=False, - debug_print=False, - debug_print_python=False, - state_types: typing.Union[typing.Dict[str, typing.Any], None] = None, - pre_create_hook=None, -): - """ - :param func: - :param memoize: Specify whether the node should be memoized (default True) - :param force_memoize: If True, the node will be memoized even if csp.memoize(False) was called. Usually it should not be set, set - this to True ONLY if memoization required to guarantee correctness of the function (i.e the function must be called at most once with the for each - set of parameters). - :param debug_print: - :param debug_print_python: - :param state_types: - :return: - """ - raise NotImplementedError("numba nodes not yet supported") - - func_frame = inspect.currentframe().f_back - - def _impl(func): - with ExceptionContext(): - parser = NumbaNodeParser(func.__name__, func, func_frame, debug_print=debug_print, state_types=state_types) - parser.parse() - - nodetype = NumbaNumbaNodeDefMeta( - func.__name__, - (NumbaNodeDef,), - { - "_signature": parser._signature, - "_impl": parser._impl, - "_state_types": parser._state_types, - "memoize": memoize, - "force_memoize": force_memoize, - "python": _create_node( - func=func, - func_frame=func_frame, - debug_print=debug_print_python, - memoize=memoize, - force_memoize=force_memoize, - cppimpl=None, - pre_create_hook=pre_create_hook, - ), - "_pre_create_hook": pre_create_hook, - "__wrapped__": func, - "__doc__": parser._docstring, - }, - ) - return nodetype - - if func is None: - return _impl - else: - return _impl(func) diff --git a/csp/impl/wiring/numba_node_parser.py b/csp/impl/wiring/numba_node_parser.py deleted file mode 100644 index 90d0f1824..000000000 --- a/csp/impl/wiring/numba_node_parser.py +++ /dev/null @@ -1,487 +0,0 @@ -import ast -import copy -import typing -from ast import Attribute, Call, Name - -from csp.impl.types.common_definitions import ArgKind -from csp.impl.wiring.node_parser import CspParseError, NodeParser -from csp.impl.wiring.signature import Signature - - -class _StateVarTransformer(ast.NodeTransformer): - """Transform all usages of vars that are stored in state to be used from the state - - Example: - Assume that x, is a state variable, y is a function argument (function arguments treated as state variable since they should persist) and z is local variable. - x = x + y + z will be transformed to __csp_state__.x = __csp_state__.x + __csp_state__.y + z - """ - - def __init__(self, state_obj_name): - self._state_obj_name = state_obj_name - self.state_vars = set() - - def visit_Name(self, node: Name): - if node.id in self.state_vars: - return Attribute(value=Name(id=self._state_obj_name, ctx=ast.Load()), attr=node.id, ctx=node.ctx) - else: - return node - - -class NumbaNodeParser(NodeParser): - _STATE_ARG_NAME = "__csp_state_ptr__" - _STATE_INST_NAME = "__csp_state_inst__" - - def __init__(self, name, raw_func, func_frame, debug_print=False, state_types=None): - from csp.impl.wiring.numba_utils.csp_cpp_numba_interface import CSP_NUMBA_CPP_FUNCTIONS - from csp.impl.wiring.numba_utils.csp_numba_functions import ( - csp_print, - csp_unsafe_cast_class_ptr_to_int, - csp_unsafe_cast_int_to_class_ptr, - ) - - super().__init__(name=name, raw_func=raw_func, func_frame=func_frame, debug_print=debug_print) - self._func_globals_modified.update(CSP_NUMBA_CPP_FUNCTIONS) - self._func_globals_modified["csp_unsafe_cast_int_to_class_ptr"] = csp_unsafe_cast_int_to_class_ptr - self._func_globals_modified["csp_unsafe_cast_class_ptr_to_int"] = csp_unsafe_cast_class_ptr_to_int - self._func_globals_modified["__csp__print__"] = csp_print - - self._state_types = {} if state_types is None else state_types - self._state_vars_transformer = _StateVarTransformer(self._STATE_INST_NAME) - self._start_state_block_body = [] - self._state_class_name = f"__CSP__STATE__{raw_func.__name__}__" - for k in self._state_types.keys(): - self._state_vars_transformer.state_vars.add(k) - - def _parse_state(self, node): - if isinstance(node, ast.With): - assert len(node.items) == 1 - call = node.items[0].context_expr - body = node.body - else: - call = node - body = None - - for kwd in call.keywords: - if kwd.arg not in self._state_types: - if isinstance(kwd.value, ast.Constant) and isinstance(kwd.value.value, (str, bool, int, float)): - if isinstance(kwd.value.value, str): - self._state_types[kwd.arg] = str - elif isinstance(kwd.value.value, bool): - self._state_types[kwd.arg] = bool - else: - self._state_types[kwd.arg] = type(kwd.value.value) - else: - raise CspParseError( - f"Unable to resolve type of state variable {kwd.arg}, explicit type must be provided in numba_node decorator", - node.lineno, - ) - self._start_state_block_body.append( - ast.Assign(targets=[ast.Name(id=kwd.arg, ctx=ast.Store())], value=kwd.value) - ) - self._state_vars_transformer.state_vars.add(kwd.arg) - if body: - self._start_state_block_body += [self.visit(n) for n in body] - - def visit_Return(self, node): - if len(self._outputs) > 1: - raise CspParseError("returning from node with multiple outputs, use return csp.output(...)", node.lineno) - - res = [] - if node.value is not None: - return_value = self.visit(node.value) - res.append( - ast.Expr( - value=ast.Call( - func=ast.Name(id=self.get_ts_out_value_return_name(0), ctx=ast.Load()), - args=[ast.Name(id="__csp_node__", ctx=ast.Load()), ast.Constant(0), return_value], - keywords=[], - ) - ) - ) - - res.append(ast.Return(value=None)) - for i in res: - i.lineno = node.lineno - i.col_offset = node.col_offset - return res - - def visit_Name(self, node): - if self._signature is None or type(node.ctx) is not ast.Load: - return node - inp = self._signature.input(node.id, allow_missing=True) - if inp is None or inp.kind == ArgKind.SCALAR: - return node - - ts_idx = inp.ts_idx - if ts_idx is not None: - return ast.Call( - func=ast.Name(id=self.get_ts_input_value_getter_name(ts_idx), ctx=ast.Load()), - args=[ast.Name(id="__csp_node__", ctx=ast.Load()), ast.Constant(ts_idx)], - keywords=[], - ) - else: - return node - - def visit_Call(self, node): - # Hack to workaround numba print bug - if isinstance(node.func, ast.Name) and node.func.id == "print": - node.func.id = "__csp__print__" - - return super().visit_Call(node) - - @classmethod - def get_ts_input_value_getter_name(cls, ts_idx): - return f"__csp__value_getter_inp_{ts_idx}__" - - @classmethod - def get_ts_out_value_return_name(cls, ts_idx): - return f"__csp__value_return_out_{ts_idx}__" - - @classmethod - def get_arg_type_var_name(cls, arg_name): - return f"__csp__{arg_name}_type__" - - @classmethod - def get_state_type_var_name(cls, arg_name): - return f"__csp__state__{arg_name}_type__" - - def _create_single_input_builtin_call(self, arg, func): - input = self._signature.input(arg.id) - ts_idx = input.ts_idx - return ast.Call( - func=ast.Name(id=func, ctx=ast.Load()), - args=[ast.Name(id="__csp_node__", ctx=ast.Load()), ast.Constant(ts_idx)], - keywords=[], - ) - - def _create_single_input_ticked_expression(self, arg): - return self._create_single_input_builtin_call(arg, "__csp_numba_node_ticked__") - - def _create_single_input_valid_expression(self, arg): - return self._create_single_input_builtin_call(arg, "__csp_numba_node_valid__") - - @classmethod - def _csp_single_input_builtin_function_transformer(cls, cpp_function_name): - """A generic function that is used to transform multiple builtin functions - :param node: Call node - :param cpp_function_name: The name of the cpp function to call - """ - - def _transform_csp_call(self, node): - if len(node.args) != 1 or node.keywords: - raise CspParseError("invalid csp call csp.%s" % node.func.attr, node.lineno) - - return self._create_single_input_builtin_call(node.args[0], cpp_function_name) - - return _transform_csp_call - - def _create_state_class( - self, state_body: typing.List[typing.Any], args: typing.List[str], state_vars: typing.List[str] - ): - l_ctx = ast.Load() - - new_args = [] - for arg in args: - new_arg = copy.copy(arg) - new_arg.annotation = None - new_args.append(new_arg) - - cls_field_types = [ - ast.Tuple( - elts=[ast.Constant(arg.arg), ast.Name(id=self.get_arg_type_var_name(arg.arg), ctx=l_ctx)], ctx=l_ctx - ) - for arg in self.get_non_ts_args() - ] - cls_field_types += [ - ast.Tuple( - elts=[ast.Constant(state_var), ast.Name(id=self.get_state_type_var_name(state_var), ctx=l_ctx)], - ctx=l_ctx, - ) - for state_var in state_vars - ] - res = ast.ClassDef( - name=self._state_class_name, - bases=[], - keywords=[], - body=[ - ast.FunctionDef( - name="__init__", - args=self._create_ast_args( - posonlyargs=[], - args=[ast.arg(arg="self", annotation=None)] + new_args, - vararg=None, - kwonlyargs=[], - kw_defaults=[], - kwarg=None, - defaults=[], - ), - body=state_body, - decorator_list=[], - returns=None, - ), - ast.FunctionDef( - name="get_ptr", - args=self._create_ast_args( - posonlyargs=[], - args=[ast.arg(arg="self", annotation=None)], - vararg=None, - kwonlyargs=[], - kw_defaults=[], - kwarg=None, - defaults=[], - ), - body=[ - ast.Return( - value=ast.Call( - func=ast.Name(id="csp_unsafe_cast_class_ptr_to_int", ctx=l_ctx), - args=[ast.Name(id="self", ctx=l_ctx)], - keywords=[], - ) - ) - ], - decorator_list=[], - returns=None, - ), - ], - decorator_list=[ - ast.Call( - func=ast.Attribute(value=ast.Name(id="__csp_numba__", ctx=l_ctx), attr="jitclass", ctx=l_ctx), - args=[ast.List(elts=cls_field_types, ctx=l_ctx)], - keywords=[], - ) - ], - ) - res = ast.fix_missing_locations(res) - return res - - def _transform_state_block(self): - self._stateblock = [self.visit(node) for node in self._stateblock] - - if self.has_state(): - if not self._stateblock: - state_init_body = [ast.Pass()] - else: - state_init_body = self._stateblock - - self._stateblock = [ - self._create_state_class( - state_body=state_init_body, args=self.get_non_ts_args(), state_vars=list(self._state_types.keys()) - ) - ] - - def has_state(self): - return self._state_types or self._stateblock - - def _create_init_block(self): - init_block = [] - init_block.append(ast.Import(names=[ast.alias(name="numba", asname="__csp_numba__")])) - return init_block - - @classmethod - def _create_function_def(cls, name, body, args=None, decorator_list=None): - res = ast.FunctionDef(name=name, body=body, returns=None) - if args is None: - res.args = cls._create_ast_args( - posonlyargs=[], args=[], kwonlyargs=[], defaults=[], vararg=None, kwarg=None, kw_defaults=[] - ) - else: - res.args = args - if decorator_list is None: - res.decorator_list = [] - else: - res.decorator_list = decorator_list - return res - - def _create_jit_func(self, func_name, inner_body): - args = self._create_ast_args( - posonlyargs=[], - args=[ast.arg(arg="__csp_node__", annotation=None), ast.arg(arg=self._STATE_ARG_NAME, annotation=None)], - kwonlyargs=[], - defaults=[], - vararg=None, - kwarg=None, - kw_defaults=[], - ) - if self.has_state(): - body = [ - ast.Assign( - targets=[Name(id="__csp_state_inst__", ctx=ast.Store())], - value=Call( - func=Name(id="csp_unsafe_cast_int_to_class_ptr", ctx=ast.Load()), - args=[ - Name(id="__csp_state_ptr__", ctx=ast.Load()), - Name(id=self._state_class_name, ctx=ast.Load()), - ], - keywords=[], - ), - ) - ] + inner_body - else: - body = inner_body - - decorator_list = [ - Call( - func=Attribute(value=Name(id="__csp_numba__", ctx=ast.Load()), attr="cfunc", ctx=ast.Load()), - args=[ - Call( - func=Attribute( - value=Attribute( - value=Name(id="__csp_numba__", ctx=ast.Load()), attr="types", ctx=ast.Load() - ), - attr="void", - ctx=ast.Load(), - ), - args=[ - Attribute( - value=Attribute( - value=Name(id="__csp_numba__", ctx=ast.Load()), attr="types", ctx=ast.Load() - ), - attr="int64", - ctx=ast.Load(), - ), - Attribute( - value=Attribute( - value=Name(id="__csp_numba__", ctx=ast.Load()), attr="types", ctx=ast.Load() - ), - attr="int64", - ctx=ast.Load(), - ), - ], - keywords=[], - ) - ], - keywords=[ - ast.keyword(arg="cache", value=ast.Constant(value=False)), - ast.keyword(arg="debug", value=ast.Constant(value=False)), - ], - ) - ] - inner_body_func = self._create_function_def(name=func_name, args=args, body=body, decorator_list=decorator_list) - return inner_body_func - - def _create_wrapped_inner_body(self, inner_body): - if self._start_state_block_body: - state_block_body = [self._state_vars_transformer.visit(n) for n in self._start_state_block_body] - else: - state_block_body = [ast.Pass()] - inner_body = [self._state_vars_transformer.visit(n) for n in inner_body] - - init_state_func = self._create_jit_func("__csp__init__state__", state_block_body) - inner_body_func = self._create_jit_func("__csp__impl__", inner_body) - return [init_state_func, inner_body_func] - - def _finalize_func_body(self, newbody): - if self.has_state(): - state_cls = ast.Name(id=self._state_class_name, ctx=ast.Load()) - else: - state_cls = ast.Constant(value=None) - - newbody.append( - ast.Return( - value=ast.Tuple( - elts=[ - state_cls, - ast.Name(id="__csp__init__state__", ctx=ast.Load()), - ast.Name(id="__csp__impl__", ctx=ast.Load()), - ], - ctx=ast.Load(), - ) - ) - ) - return newbody - - def _create_top_level_func_def(self, body): - res = self._create_function_def(name=self._funcdef.name, body=body) - res.args.args += self.get_non_ts_args() - - for inp in self._inputs: - if inp.kind.is_single_ts(): - res.args.args.append(ast.arg(arg=self.get_ts_input_value_getter_name(inp.ts_idx), annotation=None)) - for out in self._outputs: - if out.kind.is_single_ts(): - res.args.args.append(ast.arg(arg=self.get_ts_out_value_return_name(out.ts_idx), annotation=None)) - - for arg in self.get_non_ts_args(): - res.args.args.append(ast.arg(arg=self.get_arg_type_var_name(arg.arg), annotation=None)) - - for state_var in self._state_types.keys(): - res.args.args.append(ast.arg(arg=self.get_state_type_var_name(state_var), annotation=None)) - - return res - - def _is_ts_args_removed_from_signature(self): - return True - - def _parse_impl(self): - self._inputs, input_defaults, self._outputs = self.parse_func_signature(self._funcdef) - # Should have inputs and outputs at this point - inner_block_start_idx = self._parse_special_blocks(self._funcdef.body) - self._signature = Signature(self._name, self._inputs, self._outputs, input_defaults) - for arg in self.get_non_ts_args(): - self._stateblock.append( - ast.Assign( - targets=[ast.Attribute(value=ast.Name(id="self", ctx=ast.Load()), attr=arg.arg, ctx=ast.Store())], - value=ast.Name(id=arg.arg, ctx=ast.Load()), - ) - ) - self._state_vars_transformer.state_vars.add(arg.arg) - # apply transform to stateblock after signature is ready for alarms and passive calls - self._transform_state_block() - - init_block = self._create_init_block() - - # inner_body is the while loop. Start off with a yield, will get called when something ticks - inner_body = [] - for x in range(inner_block_start_idx, len(self._funcdef.body)): - res = self.visit(self._funcdef.body[x]) - if isinstance(res, list): - inner_body.extend(res) - else: - inner_body.append(res) - - wrapped_inner_block = self._create_wrapped_inner_body(inner_body) - - if isinstance(wrapped_inner_block, list): - start_and_body = self._stateblock + wrapped_inner_block - else: - start_and_body = self._stateblock + [wrapped_inner_block] - - if self._stopblock: - # For stop we wrap start and body in a try / finally ( not the init block, if that fails it's unrecoverable ) - start_and_body = [ast.Try(body=start_and_body, finalbody=self._stopblock, handlers=[], orelse=[])] - - if init_block is None: - newbody = start_and_body - else: - newbody = init_block + start_and_body - - newbody = self._finalize_func_body(newbody) - - newfuncdef = self._create_top_level_func_def(newbody) - - if self._DEBUG_PARSE or self._debug_print: - import astor - - print(astor.to_source(newfuncdef)) - - newfuncdef = ast.fix_missing_locations(newfuncdef) - ast.increment_lineno(newfuncdef, self._raw_func.__code__.co_firstlineno - 1) - modast = self._create_ast_module([newfuncdef]) - - comp = compile(modast, filename=self._func_filename, mode="exec") - _locals = {} - eval(comp, self._func_globals_modified, _locals) - self._impl = _locals[newfuncdef.name] - - @classmethod - def _init_internal_maps(cls): - super()._init_internal_maps() - cls.METHOD_MAP["csp.make_passive"] = cls._csp_single_input_builtin_function_transformer( - "__csp_numba_node_make_passive__" - ) - cls.METHOD_MAP["csp.make_active"] = cls._csp_single_input_builtin_function_transformer( - "__csp_numba_node_make_active__" - ) - - -NumbaNodeParser._init_internal_maps() diff --git a/csp/impl/wiring/numba_utils/__init__.py b/csp/impl/wiring/numba_utils/__init__.py deleted file mode 100644 index dbf3290dc..000000000 --- a/csp/impl/wiring/numba_utils/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -# We need to import this to register the support for datetime in numba -from . import datetime_extension diff --git a/csp/impl/wiring/numba_utils/csp_cpp_numba_interface.py b/csp/impl/wiring/numba_utils/csp_cpp_numba_interface.py deleted file mode 100644 index 7eabe9b4b..000000000 --- a/csp/impl/wiring/numba_utils/csp_cpp_numba_interface.py +++ /dev/null @@ -1,69 +0,0 @@ -from cffi import FFI - -from csp.impl.__cspimpl import _cspimpl - -__all__ = [] - -ffi = FFI() -ffi.cdef( - """ - bool __csp_numba_node_ticked__(int64_t node_ptr, int32_t input_index); - bool __csp_numba_node_valid__(int64_t node_ptr, int32_t input_index); - bool __csp_numba_node_make_passive__(int64_t node_ptr, int32_t input_index); - bool __csp_numba_node_make_active__(int64_t node_ptr, int32_t input_index); - bool __csp_get_node_bool_value__(int64_t node_ptr, int32_t input_index); - int64_t __csp_get_node_int64_value__(int64_t node_ptr, int32_t input_index); - double __csp_get_node_double_value__(int64_t node_ptr, int32_t input_index); - void __csp_return_bool_value__(int64_t node_ptr, int32_t input_index, bool value); - void __csp_return_int64_value__(int64_t node_ptr, int32_t input_index, int64_t value); - void __csp_return_double_value__(int64_t node_ptr, int32_t output_index, double value); - int64_t __csp_create_datetime_nanoseconds__(int32_t year, int32_t month, int32_t day, - int32_t hour, int32_t minute, int32_t second, - int32_t nanosecond); -""" -) - -C = ffi.dlopen(_cspimpl.__file__) - -CSP_NUMBA_CPP_FUNCTIONS = {} - -for s in dir(C): - if s.startswith("__csp_"): - CSP_NUMBA_CPP_FUNCTIONS[s] = getattr(C, s) - - -class NumbaTSTypedFunctionResolver: - _SUPPORTED_TYPE_C_NAMES = {bool: "bool", int: "int64", float: "double"} - - __VALUE_GETTERS__ = {} - __VALUE_RETURNERS__ = {} - - def __init__(self): - pass - - @classmethod - def init_getters_and_setters_maps(cls): - for typ, typ_name in cls._SUPPORTED_TYPE_C_NAMES.items(): - cls.__VALUE_GETTERS__[typ] = CSP_NUMBA_CPP_FUNCTIONS[f"__csp_get_node_{typ_name}_value__"] - cls.__VALUE_RETURNERS__[typ] = CSP_NUMBA_CPP_FUNCTIONS[f"__csp_return_{typ_name}_value__"] - - @classmethod - def get_value_getter_function(cls, typ): - res = cls.__VALUE_GETTERS__.get(typ) - if res is None: - raise RuntimeError(f"Unable to resolve getter function for type {typ}") - return res - - @classmethod - def get_value_returner_function(cls, typ): - res = cls.__VALUE_RETURNERS__.get(typ) - if res is None: - raise RuntimeError(f"Unable to resolve return function for type {typ}") - return res - - -NumbaTSTypedFunctionResolver.init_getters_and_setters_maps() - - -def ffi_ptr_to_int(ffi_ptr_obj): - return int(ffi.cast("int64_t", ffi_ptr_obj)) diff --git a/csp/impl/wiring/numba_utils/csp_numba_functions.py b/csp/impl/wiring/numba_utils/csp_numba_functions.py deleted file mode 100644 index 356d8cd9b..000000000 --- a/csp/impl/wiring/numba_utils/csp_numba_functions.py +++ /dev/null @@ -1,97 +0,0 @@ -""" -Implementation of numba extensions and few wrappers to work with csp -""" - -import numba - - -@numba.extending.intrinsic -def csp_unsafe_cast_int_to_class_ptr(typingctx, src, dst): - """Numba extension function that allows casting from int64 (pointer) to class reference - - Provides an LLVM implementation of the casting operation. Can be used in any jit function. - The int pointer is assumed to be generated by csp_unsafe_cast_class_ptr_to_int. - Example usage: - @numba.jitclass([("val", numba.types.unicode_type)]) - class MyClass: - def __init__(self, val): - self.val = val - - def get_ptr(self): - return csp_unsafe_cast_class_ptr_to_int(self) - - def print_hello(self): - print("Hello my val is ", self.val, ", my address is:", self.get_ptr()) - - - @numba.njit - def print_from_ptr(ptr): - print("Ptr is ", ptr) - my_o: MyClass = csp_unsafe_cast_int_to_class_ptr(ptr, MyClass) - my_o.print_hello() - - - o = MyClass("Value1") - o2 = MyClass("Value2") - - print_from_ptr(o.get_ptr()) - print_from_ptr(o2.get_ptr()) - :param typingctx: - :param src: - :param dst: - :return: - """ - - def _csp_unsafe_cast_int_to_class_ptr(int_type, class_type): - inst_type = class_type.instance_type - sig = inst_type(numba.types.voidptr, class_type) - - def codegen(context, builder, signature, args): - ptr, _ = args - alloc_type = context.get_data_type(inst_type.get_data_type()) - inst_struct = context.make_helper(builder, inst_type) - inst_struct.meminfo = numba.cgutils.get_null_value(inst_struct.meminfo.type) - inst_struct.data = builder.bitcast(ptr, alloc_type.as_pointer()) - return inst_struct._getvalue() - - return sig, codegen - - if isinstance(src, numba.types.Integer) and isinstance(dst, numba.types.ClassType): - return _csp_unsafe_cast_int_to_class_ptr(src, dst) - return TypeError - - -@numba.extending.intrinsic -def csp_unsafe_cast_class_ptr_to_int(typingctx, src): - """Numba extension function that allows casting from class reference to int64 (ptr) - The returned pointer only valid usage is to store it and then convert back using csp_unsafe_cast_int_to_class_ptr - For usage example, see documentation of csp_unsafe_cast_int_to_class_ptr - - :param typingctx: - :param src: - """ - - def _csp_unsafe_cast_class_ptr_to_int(inst_type): - sig = numba.types.int64(inst_type) - - def codegen(context, builder, signature, args): - [src] = args - llrtype = context.get_value_type(numba.int64) - inst_struct = context.make_helper(builder, inst_type, src) - return builder.ptrtoint(inst_struct.data, llrtype) - - return sig, codegen - - if isinstance(src, numba.types.ClassInstanceType): - return _csp_unsafe_cast_class_ptr_to_int(src) - return TypeError - - -@numba.jit -def csp_print(*args): - """ - A wrapper around the print function, there is some bug when trying to print directly from some numba funcitons, - this wrapper seems to fix this problem - :param args: Values to print - """ - print(*args) diff --git a/csp/impl/wiring/numba_utils/csp_string_extension.py b/csp/impl/wiring/numba_utils/csp_string_extension.py deleted file mode 100644 index 35758c868..000000000 --- a/csp/impl/wiring/numba_utils/csp_string_extension.py +++ /dev/null @@ -1,53 +0,0 @@ -""" -An implementation of utility functions to work with numba strings -""" - -import numba -from numba import cgutils - -_BYTES_TYPE = numba.types.Bytes(numba.types.uint8, 1, "C", readonly=True) - - -def _csp_make_constant_bytes(context, builder, nbytes): - bstr_ctor = cgutils.create_struct_proxy(_BYTES_TYPE) - bstr = bstr_ctor(context, builder) - - if isinstance(nbytes, int): - nbytes = numba.llvmlite.ir.Constant(bstr.nitems.type, nbytes) - - bstr.meminfo = context.nrt.meminfo_alloc(builder, nbytes) - bstr.nitems = nbytes - bstr.itemsize = numba.llvmlite.ir.Constant(bstr.itemsize.type, 1) - bstr.data = context.nrt.meminfo_data(builder, bstr.meminfo) - bstr.parent = cgutils.get_null_value(bstr.parent.type) - # bstr.shape and bstr.strides are not used - bstr.shape = cgutils.get_null_value(bstr.shape.type) - bstr.strides = cgutils.get_null_value(bstr.strides.type) - return bstr - - -@numba.extending.lower_cast(numba.types.UnicodeType, numba.types.Bytes) -def csp_unicode_to_bytes_cast(context, builder, fromty, toty, val): - uni_str = cgutils.create_struct_proxy(fromty)(context, builder, value=val) - src1 = builder.bitcast(uni_str.data, numba.llvmlite.ir.IntType(8).as_pointer()) - notkind1 = builder.icmp_unsigned("!=", uni_str.kind, numba.llvmlite.ir.Constant(uni_str.kind.type, 1)) - src_length = uni_str.length - - with builder.if_then(notkind1): - context.call_conv.return_user_exc(builder, ValueError, ("cannot cast higher than 8-bit unicode_type to bytes",)) - - bstr = _csp_make_constant_bytes(context, builder, src_length) - cgutils.memcpy(builder, bstr.data, src1, bstr.nitems) - return bstr - - -@numba.extending.intrinsic -def csp_unicode_to_bytes(typingctx, s): - # used in _to_bytes method - assert s == numba.types.unicode_type - sig = _BYTES_TYPE(s) - - def codegen(context, builder, signature, args): - return csp_unicode_to_bytes_cast(context, builder, s, _BYTES_TYPE, args[0])._getvalue() - - return sig, codegen diff --git a/csp/impl/wiring/numba_utils/datetime_extension.py b/csp/impl/wiring/numba_utils/datetime_extension.py deleted file mode 100644 index 739cd1b0b..000000000 --- a/csp/impl/wiring/numba_utils/datetime_extension.py +++ /dev/null @@ -1,153 +0,0 @@ -""" -An extension module that enables construction of datetime and timedelta in numba functions. -The module needs to be imported to enable the support. -It enables construction that is compatible with datetime.datetime and datetime.timedelta, the -returned types are actually numpy.datetime64 and numpy.timedelta64 in units of nanoseconds. -Regular datetime and timedelta arithmetic can be used. -""" - -import datetime - -import numba - -from csp.impl.wiring.numba_utils.csp_cpp_numba_interface import C as csp_c - -csp_numba_datetime_type = numba.types.NPDatetime("ns") -csp_numba_timedelta_type = numba.types.NPTimedelta("ns") - - -def _datetime_cast_implementation(ret_type, inst_type): - """The actual implementation of the llvm casting to/from datetime and timedelta - :param ret_type: The return type of the cast - :param inst_type: The type of the casted value - :return: - """ - sig = ret_type(inst_type) - - def codegen(context, builder, signature, args): - [src] = args - llrtype = context.get_value_type(ret_type) - return builder.bitcast(src, llrtype) - - return sig, codegen - - -@numba.extending.intrinsic -def csp_datetime_from_int(typingctx, src): - """Numba extension function that implementes reinterpret_cast(int64_t) - :param typingctx: - :param src: - :return: - """ - if isinstance(src, numba.types.IntegerLiteral) or src is numba.types.int64: - return _datetime_cast_implementation(csp_numba_datetime_type, src) - return TypeError - - -@numba.extending.intrinsic -def csp_datetime_to_int(typingctx, src): - """Numba extension function that implementes reinterpret_cast(datetime64) - :param typingctx: - :param src: - :return: - """ - - if src is csp_numba_datetime_type: - return _datetime_cast_implementation(numba.types.int64, src) - return TypeError - - -@numba.extending.intrinsic -def csp_timedelta_from_int(typingctx, src): - """Numba extension function that implementes reinterpret_cast(int64_t) - :param typingctx: - :param src: - :return: - """ - if isinstance(src, numba.types.IntegerLiteral) or src is numba.types.int64: - return _datetime_cast_implementation(csp_numba_timedelta_type, src) - return TypeError - - -@numba.extending.intrinsic -def csp_timedelta_to_int(typingctx, src): - """Numba extension function that implementes reinterpret_cast(timedelta64) - :param typingctx: - :param src: - :return: - """ - if src is csp_numba_timedelta_type: - return _datetime_cast_implementation(numba.types.int64, src) - return TypeError - - -__csp_create_datetime_nanoseconds__ = csp_c.__csp_create_datetime_nanoseconds__ - - -@numba.extending.overload(datetime.datetime, inline="always") -def create_datetime(year, month, day, hour=0, minute=0, second=0, microsecond=0, nanosecond=0): - """An implementation of datetime __init_ call that will produce numpy.datetime64 object instead - :param year: - :param month: - :param day: - :param hour: - :param minute: - :param second: - :param microsecond: - :param nanosecond: - :return: - """ - - def create_datetime_impl(year, month, day, hour=0, minute=0, second=0, microsecond=0, nanosecond=0): - return csp_datetime_from_int( - __csp_create_datetime_nanoseconds__(year, month, day, hour, minute, second, microsecond * 1000 + nanosecond) - ) - - return create_datetime_impl - - -MICROS_MULT = 1000 -MILLIS_MULT = MICROS_MULT * 1000 -SECONDS_MULT = MILLIS_MULT * 1000 -MINUTES_MULT = SECONDS_MULT * 60 -HOURS_MULT = MINUTES_MULT * 60 -DAYS_MULT = HOURS_MULT * 24 -WEEKS_MULT = DAYS_MULT * 7 - - -@numba.extending.overload(datetime.timedelta, inline="always") -def create_timedelta(days=0, seconds=0, microseconds=0, milliseconds=0, minutes=0, hours=0, weeks=0, nanoseconds=0): - """An implementation of timedelta __init_ call that will produce numpy.timedelta64 object instead - :param year: - :param month: - :param day: - :param hour: - :param minute: - :param second: - :param microsecond: - :param nanosecond: - :return: - """ - - def create_datetime_impl( - days=0, seconds=0, microseconds=0, milliseconds=0, minutes=0, hours=0, weeks=0, nanoseconds=0 - ): - res = nanoseconds if nanoseconds is not None else 0 - if days != 0: - res += DAYS_MULT * days - if seconds != 0: - res += SECONDS_MULT * seconds - if microseconds != 0: - res += MICROS_MULT * microseconds - if milliseconds != 0: - res += MILLIS_MULT * milliseconds - if minutes != 0: - res += MINUTES_MULT * minutes - if hours != 0: - res += HOURS_MULT * hours - if weeks != 0: - res += WEEKS_MULT * weeks - - return csp_timedelta_from_int(res) - - return create_datetime_impl diff --git a/csp/impl/wiring/numba_utils/numba_type_resolver.py b/csp/impl/wiring/numba_utils/numba_type_resolver.py deleted file mode 100644 index 670642200..000000000 --- a/csp/impl/wiring/numba_utils/numba_type_resolver.py +++ /dev/null @@ -1,77 +0,0 @@ -import datetime -import typing - -import numba - -from csp.impl.types.typing_utils import CspTypingUtils -from csp.impl.wiring.numba_utils import datetime_extension - - -class NumbaTypeResolver(object): - _EPOCH = datetime.datetime.utcfromtimestamp(0) - - _PRIMITIVE_TYPE_MAPPING = { - bool: numba.types.boolean, - int: numba.types.int64, - float: numba.types.double, - str: numba.types.string, - datetime.datetime: datetime_extension.csp_numba_datetime_type, - datetime.timedelta: datetime_extension.csp_numba_timedelta_type, - } - - @classmethod - def _resolve_container_type(cls, python_type): - if CspTypingUtils.get_origin(python_type) is typing.Dict: - k, v = python_type.__args__ - numba_k = cls.resolve_numba_type(k) - numba_v = cls.resolve_numba_type(v) - return numba.types.DictType(numba_k, numba_v) - elif CspTypingUtils.get_origin(python_type) is typing.List: - (t,) = python_type.__args__ - numba_t = cls.resolve_numba_type(t) - return numba.types.ListType(numba_t) - elif CspTypingUtils.get_origin(python_type) is typing.Set: - raise NotImplementedError("Set scalars are not currently supported in numba nodes") - else: - raise RuntimeError(f"Unable to resolve numba type for {python_type}") - - @classmethod - def resolve_numba_type(cls, python_type): - res = cls._PRIMITIVE_TYPE_MAPPING.get(python_type) - if res is None: - if CspTypingUtils.is_generic_container(python_type): - return cls._resolve_container_type(python_type) - else: - raise RuntimeError(f"Unable to resolve numba type for {python_type}") - return res - - @classmethod - def _instantiate_container(cls, numba_type): - if isinstance(numba_type, numba.types.DictType): - return - # elif - raise TypeError(f"Don't know how to isntantiate {numba_type}") - - @classmethod - def transform_scalar(cls, scalar, scalar_type): - if isinstance(scalar_type, numba.types.DictType): - res = numba.typed.Dict.empty( - key_type=scalar_type.key_type, - value_type=scalar_type.value_type, - ) - - for k, v in scalar.items(): - res[cls.transform_scalar(k, scalar_type.key_type)] = cls.transform_scalar(v, scalar_type.value_type) - return res - elif isinstance(scalar_type, numba.types.ListType): - res = numba.typed.List.empty_list(scalar_type.item_type) - for v in scalar: - res.append(v) - return res - - elif isinstance(scalar, datetime.timedelta): - return datetime_extension.csp_numba_timedelta_type(int(scalar.total_seconds() * 1e9)) - elif isinstance(scalar, datetime.datetime): - datetime_extension.csp_numba_datetime_type(int((scalar - cls._EPOCH).total_seconds() * 1e9)) - else: - return scalar diff --git a/csp/tests/test_numba_features.py b/csp/tests/test_numba_features.py deleted file mode 100644 index 33a2577a1..000000000 --- a/csp/tests/test_numba_features.py +++ /dev/null @@ -1,189 +0,0 @@ -import typing -import unittest -from datetime import datetime - -import csp -from csp import ts - - -class TestTypeChecking(unittest.TestCase): - class Dummy: - pass - - class Dummy2(Dummy): - pass - - class Dummy3: - pass - - @unittest.skip("numba not yet used, tests fail on newer numba we get in our 3.8 build") - def test_graph_build_type_checking(self): - @csp.numba_node - def typed_ts(x: ts[int]): - if csp.ticked(x): - pass - - @csp.numba_node - def typed_scalar(x: ts[int], y: str): - if csp.ticked(x): - pass - - def graph(): - i = csp.const(5) - typed_ts(i) - - typed_scalar(i, "xyz") - - with self.assertRaisesRegex(TypeError, "Expected ts\\[int\\] for argument 'x', got ts\\[str\\]"): - s = csp.const("xyz") - ## THIS SHOULD RAISE, passing ts[str] but typed takes ts[int] - typed_ts(s) - - with self.assertRaisesRegex(TypeError, "Expected str for argument 'y', got 123 \\(int\\)"): - ## THIS SHOULD RAISE, passing int instead of str - typed_scalar(i, 123) - - csp.run(graph, starttime=datetime(2020, 2, 7, 9), endtime=datetime(2020, 2, 7, 9, 1)) - - @unittest.skip("numba not yet used, tests fail on newer numba we get in our 3.8 build") - def test_runtime_type_check(self): - ## native output type - @csp.numba_node - def typed_int(x: ts["T"]) -> ts[int]: - if csp.ticked(x): - return x - - # TODO: Uncomment - # @csp.numba_node - # def typed_alarm(v: '~T', alarm_type: 'V') -> outputs(ts['V']): - # with csp.alarms(): - # alarm = csp.alarm( 'V' ) - # with csp.start(): - # csp.schedule_alarm(alarm, timedelta(), v) - # - # if csp.ticked(alarm): - # return alarm - - # Valid - csp.run(typed_int, csp.const(5), starttime=datetime(2020, 2, 7)) - - # Invalid - with self.assertRaisesRegex(RuntimeError, "Unable to resolve getter function for type.*"): - csp.run(typed_int, csp.const("5"), starttime=datetime(2020, 2, 7)) - - # TODO: uncomment - # # valid - # csp.run(typed_alarm, 5, int, starttime=datetime(2020, 2, 7)) - # csp.run(typed_alarm, 5, object, starttime=datetime(2020, 2, 7)) - # csp.run(typed_alarm, [1, 2, 3], [int], starttime=datetime(2020, 2, 7)) - # - # # Invalid - # with self.assertRaisesRegex(TypeError, - # '"typed_alarm" node expected output type on output #0 to be of type "str" got type "int"'): - # csp.run(typed_alarm, 5, str, starttime=datetime(2020, 2, 7)) - # - # with self.assertRaisesRegex(TypeError, - # '"typed_alarm" node expected output type on output #0 to be of type "bool" got type "int"'): - # csp.run(typed_alarm, 5, bool, starttime=datetime(2020, 2, 7)) - # - # with self.assertRaisesRegex(TypeError, - # '"typed_alarm" node expected output type on output #0 to be of type "str" got type "list"'): - # csp.run(typed_alarm, [1, 2, 3], str, starttime=datetime(2020, 2, 7)) - - @unittest.skip("numba not yet used, tests fail on newer numba we get in our 3.8 build") - def test_dict_type_resolutions(self): - @csp.numba_node - def typed_dict_int_int(x: {int: int}): - pass - - @csp.numba_node - def typed_dict_int_int2(x: typing.Dict[int, int]): - pass - - @csp.numba_node - def typed_dict_int_float(x: {int: int}): - pass - - @csp.numba_node - def typed_dict_float_float(x: {float: float}): - pass - - @csp.numba_node - def typed_dict(x: {"T": "V"}): - pass - - @csp.numba_node - def deep_nested_generic_resolution(x: "T1", y: "T2", z: {"T1": {"T2": [{"T1"}]}}): - pass - - def graph(): - d_i_i = csp.const({1: 2, 3: 4}) - - # Ok int dict expected - typed_dict_int_int({1: 2, 3: 4}) - - # Ok int dict expected - typed_dict_int_int2({1: 2, 3: 4}) - - typed_dict_float_float({1: 2}) - typed_dict_float_float({1.0: 2}) - typed_dict_float_float({}) - - with self.assertRaisesRegex(TypeError, r"Expected typing.Dict\[int, int\] for argument 'x', got .*"): - # Passing a float value instead of expected ints - typed_dict_int_int2({1: 2, 3: 4.0}) - - l_good = csp.const.using(T={int: float})({}) - l_good = csp.const.using(T={int: float})({2: 1}) - l_good = csp.const.using(T={int: float})({2: 1.0}) - with self.assertRaises(TypeError): - # passing float to int - l_bad = csp.const.using(T={int: float})({2.0: 1}) - - csp.run(graph, starttime=datetime(2020, 2, 7, 9), endtime=datetime(2020, 2, 7, 9, 1)) - - @unittest.skip("numba not yet used, tests fail on newer numba we get in our 3.8 build") - def test_list_type_resolutions(self): - @csp.numba_node - def typed_list_int(x: [int]): - pass - - @csp.numba_node - def typed_list_int2(x: typing.List[int]): - pass - - @csp.numba_node - def typed_list_float(x: [float]): - pass - - def graph(): - l_i = csp.const([1, 2, 3, 4]) - - typed_list_int([]) - typed_list_int([1, 2, 3]) - typed_list_int2([1, 2, 3]) - typed_list_float([1, 2, 3]) - typed_list_float([1, 2, 3.0]) - - with self.assertRaisesRegex(TypeError, r"Expected typing.List\[int\] for argument 'x', got .*"): - # Passing a float value instead of expected ints - typed_list_int([1, 2, 3.0]) - - with self.assertRaisesRegex(TypeError, r"Expected typing.List\[float\] for argument 'x', got .*"): - # Passing a Dummy value instead of expected float - typed_list_float([TestTypeChecking.Dummy()]) - - l_good = csp.const.using(T=[int])([]) - with self.assertRaisesRegex( - TypeError, - "Unable to resolve container type for type variable T explicit value must have" - + " uniform values and be non empty.*", - ): - # Passing a Dummy value instead of expected float - l_bad = csp.const([]) - - csp.run(graph, starttime=datetime(2020, 2, 7, 9), endtime=datetime(2020, 2, 7, 9, 1)) - - -if __name__ == "__main__": - unittest.main() diff --git a/csp/tests/test_numba_node.py b/csp/tests/test_numba_node.py new file mode 100644 index 000000000..e81387426 --- /dev/null +++ b/csp/tests/test_numba_node.py @@ -0,0 +1,1436 @@ +import unittest +from datetime import datetime, timedelta, timezone +from enum import IntEnum + +from numba_cfunc_compiler.api import NumbaDict, NumbaList, create_new_dict, create_new_list + +import csp +from csp import ts +from csp.impl.wiring.csp_numba import numba_node + + +class TestBasicTypes(unittest.TestCase): + def test_primitives(self): + @numba_node + def primitives_node(x: ts[int]) -> ts[float]: + with csp.state(): + s_count: int = 0 + s_sum: float = 0.0 + s_flag: bool = False + + s_count = s_count + 1 + s_sum = s_sum + x + s_flag = not s_flag + # Add 10 on odd ticks (when s_flag is True) + if s_flag: + return s_count + s_sum + 10.0 + else: + return s_count + s_sum + + @csp.graph + def g(): + # x = 1, 2, 3 + values = csp.curve(int, [(timedelta(seconds=i), i + 1) for i in range(3)]) + result = primitives_node(values) + csp.add_graph_output("result", result) + + # tick 1: s_count=1, s_sum=1.0, s_flag=True -> output = 1 + 1.0 + 10 = 12.0 + # tick 2: s_count=2, s_sum=3.0, s_flag=False -> output = 2 + 3.0 = 5.0 + # tick 3: s_count=3, s_sum=6.0, s_flag=True -> output = 3 + 6.0 + 10 = 19.0 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [12.0, 5.0, 19.0]) + + def test_constant_inputs(self): + @numba_node + def with_constants(x: ts[int], multiplier: int, offset: float) -> ts[float]: + return x * multiplier + offset + + @csp.graph + def g(): + values = csp.curve(int, [(timedelta(seconds=i), i + 1) for i in range(3)]) + result = with_constants(values, 10, 0.5) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [10.5, 20.5, 30.5]) + + def test_timedelta(self): + @numba_node + def timedelta_node(td: ts[timedelta], scale: timedelta) -> ts[timedelta]: + with csp.state(): + s_total: timedelta = timedelta(seconds=10) + + local_incr = timedelta(seconds=1) + s_total = s_total + td + local_incr + return s_total + scale + + @csp.graph + def g(): + deltas = csp.curve( + timedelta, + [ + (timedelta(seconds=0), timedelta(seconds=5)), + (timedelta(seconds=1), timedelta(seconds=10)), + ], + ) + result = timedelta_node(deltas, timedelta(seconds=100)) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual[0], timedelta(seconds=116)) + self.assertEqual(actual[1], timedelta(seconds=127)) + + def test_datetime(self): + @numba_node + def datetime_node(dt: ts[datetime], base: datetime) -> ts[datetime]: + with csp.state(): + s_last: datetime = 0 + + with csp.start(): + s_last = base + + # If dt is after threshold, return dt; otherwise return s_last + local_threshold = datetime(2024, 6, 1, tzinfo=timezone.utc) + if dt > local_threshold: + result = dt + else: + result = s_last + s_last = dt + return result + + @csp.graph + def g(): + dates = csp.curve( + datetime, + [ + (timedelta(seconds=0), datetime(2024, 1, 1, tzinfo=timezone.utc)), # before threshold + (timedelta(seconds=1), datetime(2024, 7, 1, tzinfo=timezone.utc)), # after threshold + ], + ) + result = datetime_node(dates, datetime(2020, 1, 1, tzinfo=timezone.utc)) + csp.add_graph_output("result", result) + + # tick 1: dt=2024-01-01 < threshold, return s_last=2020-01-01, then s_last=2024-01-01 + # tick 2: dt=2024-07-01 > threshold, return dt=2024-07-01, then s_last=2024-07-01 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual( + actual, + [ + datetime(2020, 1, 1), + datetime(2024, 7, 1), + ], + ) + + def test_datetime_timedelta_combined(self): + @numba_node + def time_calc(x: ts[int], base: datetime, interval: timedelta) -> ts[datetime]: + with csp.state(): + s_current: datetime = 0 + + with csp.start(): + s_current = base + + s_current = s_current + interval + return s_current + + @csp.graph + def g(): + values = csp.curve(int, [(timedelta(seconds=i), i) for i in range(3)]) + result = time_calc(values, datetime(2024, 1, 1, tzinfo=timezone.utc), timedelta(hours=2)) + csp.add_graph_output("result", result) + + # tick 1: s_current = base + 2h = 2024-01-01 02:00 + # tick 2: s_current = base + 4h = 2024-01-01 04:00 + # tick 3: s_current = base + 6h = 2024-01-01 06:00 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual( + actual, + [ + datetime(2024, 1, 1, 2), + datetime(2024, 1, 1, 4), + datetime(2024, 1, 1, 6), + ], + ) + + def test_enum(self): + class Mode(csp.Enum): + IDLE = 0 + RUNNING = 1 + PAUSED = 2 + + @numba_node + def enum_node(cmd: ts[int], default_mode: Mode) -> ts[Mode]: + with csp.state(): + s_mode: Mode = Mode.IDLE + + local_running = Mode.RUNNING + local_paused = Mode.PAUSED + if cmd == 1: + s_mode = local_running + elif cmd == 2: + s_mode = local_paused + else: + s_mode = default_mode + return s_mode + + @csp.graph + def g(): + commands = csp.curve( + int, + [ + (timedelta(seconds=0), 0), + (timedelta(seconds=1), 1), + (timedelta(seconds=2), 2), + ], + ) + result = enum_node(commands, Mode.IDLE) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [Mode.IDLE, Mode.RUNNING, Mode.PAUSED]) + + def test_int_enum(self): + class Mode(IntEnum): + IDLE = 0 + RUNNING = 1 + PAUSED = 2 + + @numba_node + def enum_node(cmd: ts[Mode], default_mode: Mode) -> ts[Mode]: + with csp.state(): + s_mode: Mode = Mode.IDLE + + if cmd == Mode.RUNNING: + s_mode = cmd + else: + s_mode = default_mode + return s_mode + + @csp.graph + def g(): + commands = csp.curve( + Mode, + [ + (timedelta(seconds=0), Mode.RUNNING), + (timedelta(seconds=1), Mode.PAUSED), + ], + ) + result = enum_node(commands, Mode.IDLE) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [Mode.RUNNING, Mode.IDLE]) + + def test_unsupported_state_expression(self): + with self.assertRaisesRegex(ValueError, "Unable to infer type for state variable 'total'"): + + @numba_node + def bad_state_inference(x: ts[int]) -> ts[int]: + with csp.state(): + total = x + 1 + return x + + +class TestContainerTypes(unittest.TestCase): + def test_list_in_start_phase(self): + @numba_node + def scale_with_start_weights(x: ts[int], weights: NumbaList[int]) -> ts[int]: + with csp.state(): + s_scale = 0 + + with csp.start(): + s_scale = weights[0] + weights[1] + + return x * s_scale + + @csp.graph + def g(): + values = csp.curve(int, [(timedelta(seconds=i), i + 1) for i in range(3)]) + result = scale_with_start_weights(values, [2, 3]) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [5, 10, 15]) + + def test_dict_in_stop_phase(self): + @numba_node + def read_constant_mapping_on_stop(x: ts[int], adjustments: NumbaDict[int, int]) -> ts[int]: + with csp.state(): + s_stop_value = 0 + + with csp.stop(): + s_stop_value = adjustments[1] + + return x + 1 + + @csp.graph + def g(): + values = csp.curve(int, [(timedelta(seconds=i), i + 1) for i in range(3)]) + result = read_constant_mapping_on_stop(values, {1: 7, 2: 11}) + csp.add_graph_output("result", result) + + # The execute output verifies normal operation; successful graph shutdown + # verifies the constant dict is available in the stop lifecycle phase. + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [2, 3, 4]) + + def test_list_basic_operations(self): + @numba_node + def list_node(x: ts[int]) -> csp.Outputs( + length=ts[int], first=ts[int], last=ts[int], popped=ts[int], total=ts[int] + ): + with csp.state(): + s_values: NumbaList[int] = create_new_list(int) + + s_values.append(x) + s_values.append(x * 2) + csp.output(length=len(s_values)) + csp.output(first=s_values[0]) + csp.output(last=s_values[len(s_values) - 1]) + csp.output(popped=s_values.pop()) + total = 0 + for v in s_values: + total = total + v + csp.output(total=total) + + @csp.graph + def g(): + inputs = csp.curve(int, [(timedelta(seconds=i), i + 1) for i in range(3)]) + result = list_node(inputs) + csp.add_graph_output("length", result.length) + csp.add_graph_output("first", result.first) + csp.add_graph_output("last", result.last) + csp.add_graph_output("popped", result.popped) + csp.add_graph_output("total", result.total) + + # tick 1: x=1 -> [1,2], len=2, first=1, last=2, pop->2, total=1 + # tick 2: x=2 -> [1,2,4], len=3, first=1, last=4, pop->4, total=3 + # tick 3: x=3 -> [1,2,3,6], len=4, first=1, last=6, pop->6, total=6 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["length"]], [2, 3, 4]) + self.assertEqual([v for _, v in results["first"]], [1, 1, 1]) + self.assertEqual([v for _, v in results["last"]], [2, 4, 6]) + self.assertEqual([v for _, v in results["popped"]], [2, 4, 6]) + self.assertEqual([v for _, v in results["total"]], [1, 3, 6]) + + def test_list_float(self): + @numba_node + def running_avg(x: ts[float], window: int) -> ts[float]: + with csp.state(): + s_values: NumbaList[float] = create_new_list(float) + + s_values.append(x) + # Keep only last 'window' values + while len(s_values) > window: + # Shift elements left (poor man's deque) + for i in range(len(s_values) - 1): + s_values[i] = s_values[i + 1] + s_values.pop() + + total = 0.0 + for v in s_values: + total = total + v + return total / len(s_values) + + @csp.graph + def g(): + inputs = csp.curve(float, [(timedelta(seconds=i), float(i + 1)) for i in range(5)]) + result = running_avg(inputs, 3) + csp.add_graph_output("result", result) + + # window=3: [1]->1, [1,2]->1.5, [1,2,3]->2, [2,3,4]->3, [3,4,5]->4 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [1.0, 1.5, 2.0, 3.0, 4.0]) + + def test_list_index_assignment(self): + @numba_node + def modify_list(x: ts[int]) -> ts[int]: + with csp.state(): + s_values: NumbaList[int] = create_new_list(int) + + s_values.append(x) + if len(s_values) >= 2: + # Double the first element + s_values[0] = s_values[0] * 2 + return s_values[0] + + @csp.graph + def g(): + inputs = csp.curve(int, [(timedelta(seconds=i), 10) for i in range(4)]) + result = modify_list(inputs) + csp.add_graph_output("result", result) + + # tick 1: [10] -> first=10 + # tick 2: [10,10], first*=2 -> [20,10] -> first=20 + # tick 3: [20,10,10], first*=2 -> [40,10,10] -> first=40 + # tick 4: [40,10,10,10], first*=2 -> [80,10,10,10] -> first=80 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [10, 20, 40, 80]) + + def test_list_multiple_lists(self): + @numba_node + def dual_lists(x: ts[int]) -> csp.Outputs(evens=ts[int], odds=ts[int]): + with csp.state(): + s_evens: NumbaList[int] = create_new_list(int) + s_odds: NumbaList[int] = create_new_list(int) + + if x % 2 == 0: + s_evens.append(x) + csp.output(evens=len(s_evens)) + else: + s_odds.append(x) + csp.output(odds=len(s_odds)) + + @csp.graph + def g(): + inputs = csp.curve(int, [(timedelta(seconds=i), i + 1) for i in range(6)]) + result = dual_lists(inputs) + csp.add_graph_output("evens", result.evens) + csp.add_graph_output("odds", result.odds) + + # 1(odd), 2(even), 3(odd), 4(even), 5(odd), 6(even) + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["evens"]], [1, 2, 3]) + self.assertEqual([v for _, v in results["odds"]], [1, 2, 3]) + + def test_dict_count_occurrences(self): + @numba_node + def count_node(x: ts[int]) -> ts[int]: + with csp.state(): + s_counts: NumbaDict[int, int] = create_new_dict(int, int) + + s_counts[x] += 1 + return s_counts[x] + + @csp.graph + def g(): + inputs = csp.curve(int, [(timedelta(seconds=i), v) for i, v in enumerate([1, 2, 1, 1, 2])]) + result = count_node(inputs) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [1, 1, 2, 3, 2]) + + def test_dict_contains(self): + @numba_node + def contains_node(x: ts[int]) -> csp.Outputs(found=ts[int], not_found=ts[int]): + with csp.state(): + s_seen: NumbaDict[int, int] = create_new_dict(int, int) + + if x in s_seen: + csp.output(found=s_seen[x]) + else: + csp.output(not_found=x) + s_seen[x] = x * 10 + + @csp.graph + def g(): + # 1, 2, 1, 3, 2 + inputs = csp.curve(int, [(timedelta(seconds=i), v) for i, v in enumerate([1, 2, 1, 3, 2])]) + result = contains_node(inputs) + csp.add_graph_output("found", result.found) + csp.add_graph_output("not_found", result.not_found) + + # t=0: x=1 not in {} -> not_found=1, then s_seen={1:10} + # t=1: x=2 not in {1:10} -> not_found=2, then s_seen={1:10,2:20} + # t=2: x=1 in {1:10,2:20} -> found=10, then s_seen={1:10,2:20} + # t=3: x=3 not in {1:10,2:20} -> not_found=3, then s_seen={1:10,2:20,3:30} + # t=4: x=2 in {1:10,2:20,3:30} -> found=20 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["not_found"]], [1, 2, 3]) + self.assertEqual([v for _, v in results["found"]], [10, 20]) + + def test_dict_float_values(self): + @numba_node + def sum_by_key(key: ts[int], val: ts[float]) -> ts[float]: + with csp.state(): + s_sums: NumbaDict[int, float] = create_new_dict(int, float) + + if csp.ticked(key) and csp.ticked(val): + s_sums[key] += val + return s_sums[key] + + @csp.graph + def g(): + keys = csp.curve(int, [(timedelta(seconds=i), k) for i, k in enumerate([1, 2, 1, 2, 1])]) + vals = csp.curve(float, [(timedelta(seconds=i), float(i + 1) * 10) for i in range(5)]) + result = sum_by_key(keys, vals) + csp.add_graph_output("result", result) + + # key=1,val=10 -> {1:10} -> 10 + # key=2,val=20 -> {1:10,2:20} -> 20 + # key=1,val=30 -> {1:40,2:20} -> 40 + # key=2,val=40 -> {1:40,2:60} -> 60 + # key=1,val=50 -> {1:90,2:60} -> 90 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [10.0, 20.0, 40.0, 60.0, 90.0]) + + def test_list_state_with_basket_inputs(self): + @numba_node + def basket_aggregator(xs: list[ts[float]]) -> csp.Outputs(avg=ts[float], delta=ts[float]): + with csp.state(): + s_history: NumbaList[float] = create_new_list(float) + + total = 0.0 + count = 0 + for i in range(len(xs)): + if xs.at(i).ticked(): + total = total + xs.at(i) + count = count + 1 + + if count > 0: + avg = total / count + s_history.append(avg) + csp.output(avg=avg) + if len(s_history) >= 2: + prev = s_history[len(s_history) - 2] + csp.output(delta=avg - prev) + + @csp.graph + def g(): + a = csp.curve(float, [(timedelta(seconds=0), 10.0), (timedelta(seconds=1), 20.0)]) + b = csp.curve(float, [(timedelta(seconds=0), 100.0), (timedelta(seconds=2), 300.0)]) + result = basket_aggregator([a, b]) + csp.add_graph_output("avg", result.avg) + csp.add_graph_output("delta", result.delta) + + # t=0: a=10, b=100 tick -> avg=(10+100)/2=55 + # t=1: a=20 ticks -> avg=20, delta=20-55=-35 + # t=2: b=300 ticks -> avg=300, delta=300-20=280 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["avg"]], [55.0, 20.0, 300.0]) + self.assertEqual([v for _, v in results["delta"]], [-35.0, 280.0]) + + def test_list_with_enum_mode_state(self): + """Test list state coexisting with enum state and windowed processing.""" + + class Mode(csp.Enum): + SUM = 0 + AVG = 1 + + @numba_node + def windowed_processor(x: ts[float], mode: ts[Mode], window_size: int) -> ts[float]: + with csp.state(): + s_values: NumbaList[float] = create_new_list(float) + s_mode: Mode = Mode.SUM + + if csp.ticked(mode): + s_mode = mode + + if csp.ticked(x): + s_values.append(x) + + # Calculate over last window_size values + start_idx = 0 + if len(s_values) > window_size: + start_idx = len(s_values) - window_size + + total = 0.0 + count = 0 + for i in range(start_idx, len(s_values)): + total = total + s_values[i] + count = count + 1 + + if s_mode == Mode.SUM: + return total + return total / count + + @csp.graph + def g(): + values = csp.curve(float, [(timedelta(seconds=i), float(i + 1)) for i in range(6)]) + # Start with SUM, switch to AVG at t=3 + mode = csp.curve( + Mode, + [ + (timedelta(seconds=0), Mode.SUM), + (timedelta(seconds=3), Mode.AVG), + ], + ) + result = windowed_processor(values, mode, 3) + csp.add_graph_output("result", result) + + # window_size=3 + # t=0: x=1, window=[1], mode=SUM -> 1.0 + # t=1: x=2, window=[1,2], mode=SUM -> 3.0 + # t=2: x=3, window=[1,2,3], mode=SUM -> 6.0 + # t=3: x=4, window=[2,3,4], mode=AVG -> 9/3 = 3.0 + # t=4: x=5, window=[3,4,5], mode=AVG -> 12/3 = 4.0 + # t=5: x=6, window=[4,5,6], mode=AVG -> 15/3 = 5.0 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [1.0, 3.0, 6.0, 3.0, 4.0, 5.0]) + + +class TestFeatures(unittest.TestCase): + def test_rejects_mismatched_signal_type(self): + @numba_node + def typed_signal(x: ts[int]) -> ts[int]: + return 1 + + @csp.graph + def g(): + csp.add_graph_output("result", typed_signal(csp.const(1.25))) + + with self.assertRaisesRegex( + TypeError, r"cannot validate ts\[float\] as ts\[int\]|Expected ts\[int\].*got ts\[float\]" + ): + csp.build_graph(g) + + def test_allows_int_to_float_signal_upcast(self): + @numba_node + def typed_signal(x: ts[float]) -> ts[float]: + return x + 0.5 + + @csp.graph + def g(): + csp.add_graph_output("result", typed_signal(csp.const(1))) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1)) + self.assertEqual([v for _, v in results["result"]], [1.5]) + + def test_rejects_mismatched_basket_element_type(self): + @numba_node + def typed_basket(xs: list[ts[int]]) -> ts[int]: + return 1 + + @csp.graph + def g(): + csp.add_graph_output("result", typed_basket([csp.const(1), csp.const(2.5)])) + + with self.assertRaisesRegex( + TypeError, r"cannot validate ts\[float\] as ts\[int\]|Expected list\[.*TsType\[int\]\]" + ): + csp.build_graph(g) + + def test_rejects_mismatched_dict_basket_key_type(self): + @numba_node + def typed_basket(xs: dict[int, ts[int]]) -> ts[int]: + return 1 + + @csp.graph + def g(): + csp.add_graph_output("result", typed_basket({"bad": csp.const(1)})) + + with self.assertRaisesRegex( + TypeError, r"Input should be a valid integer|Expected dict\[int, .*TsType\[int\]\]" + ): + csp.build_graph(g) + + def test_lifecycle_start(self): + @numba_node + def with_start(x: ts[int]) -> ts[int]: + with csp.state(): + s_mult = 0 + s_base = 100 + + with csp.start(): + s_mult = 10 + s_base = 500 + + s_base = s_base + x * s_mult + return s_base + + @csp.graph + def g(): + values = csp.curve(int, [(timedelta(seconds=i), i + 1) for i in range(3)]) + result = with_start(values) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [510, 530, 560]) + + def test_ticked_and_valid_methods(self): + @numba_node + def check_signals(a: ts[int], b: ts[int]) -> csp.Outputs( + a_count=ts[int], + b_count=ts[int], + any_ticked=ts[int], + a_valid=ts[int], + b_valid=ts[int], + both_valid=ts[int], + ): + with csp.state(): + s_a_count: int = 0 + s_b_count: int = 0 + + if csp.ticked(a): + s_a_count = s_a_count + 1 + csp.output(a_count=s_a_count) + if csp.ticked(b): + s_b_count = s_b_count + 1 + csp.output(b_count=s_b_count) + + if csp.ticked(a, b): + csp.output(any_ticked=s_a_count * 10 + s_b_count) + + if csp.valid(a): + csp.output(a_valid=a) + if csp.valid(b): + csp.output(b_valid=b) + if csp.valid(a, b): + csp.output(both_valid=a + b) + + @csp.graph + def g(): + a = csp.curve(int, [(timedelta(seconds=0), 1), (timedelta(seconds=2), 2)]) + b = csp.curve(int, [(timedelta(seconds=1), 5), (timedelta(seconds=2), 15)]) + result = check_signals(a, b) + csp.add_graph_output("a_count", result.a_count) + csp.add_graph_output("b_count", result.b_count) + csp.add_graph_output("any_ticked", result.any_ticked) + csp.add_graph_output("a_valid", result.a_valid) + csp.add_graph_output("b_valid", result.b_valid) + csp.add_graph_output("both_valid", result.both_valid) + + # t=0: a ticks, only a is valid + # t=1: b ticks, both a and b are valid + # t=2: a and b both tick, both are valid + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["a_count"]], [1, 2]) + self.assertEqual([v for _, v in results["b_count"]], [1, 2]) + self.assertEqual([v for _, v in results["any_ticked"]], [10, 11, 22]) + self.assertEqual([v for _, v in results["a_valid"]], [1, 1, 2]) + self.assertEqual([v for _, v in results["b_valid"]], [5, 15]) + self.assertEqual([v for _, v in results["both_valid"]], [6, 17]) + + def test_list_basket_input(self): + @numba_node + def basket_stats(xs: [ts[int]]) -> csp.Outputs(ticked_count=ts[int], ticked_sum=ts[int]): + ticked_count = 0 + ticked_sum = 0 + for i in range(len(xs)): + if xs.at(i).ticked(): + ticked_count = ticked_count + 1 + ticked_sum = ticked_sum + xs.at(i) + csp.output(ticked_count=ticked_count) + csp.output(ticked_sum=ticked_sum) + + @csp.graph + def g(): + a = csp.curve(int, [(timedelta(seconds=0), 10), (timedelta(seconds=2), 20)]) + b = csp.curve(int, [(timedelta(seconds=1), 100)]) + c = csp.const(1000) + result = basket_stats([a, b, c]) + csp.add_graph_output("ticked_count", result.ticked_count) + csp.add_graph_output("ticked_sum", result.ticked_sum) + + # t=0: a=10, c=1000 tick -> ticked: a,c + # t=1: b=100 ticks -> ticked: b + # t=2: a=20 ticks -> ticked: a + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["ticked_count"]], [2, 1, 1]) + self.assertEqual([v for _, v in results["ticked_sum"]], [1010, 100, 20]) + + def test_list_basket_input_modern_annotation(self): + @numba_node + def basket_stats(xs: list[ts[int]]) -> csp.Outputs(ticked_count=ts[int], ticked_sum=ts[int]): + ticked_count = 0 + ticked_sum = 0 + for i in range(len(xs)): + if xs.at(i).ticked(): + ticked_count = ticked_count + 1 + ticked_sum = ticked_sum + xs.at(i) + csp.output(ticked_count=ticked_count) + csp.output(ticked_sum=ticked_sum) + + @csp.graph + def g(): + a = csp.curve(int, [(timedelta(seconds=0), 10), (timedelta(seconds=2), 20)]) + b = csp.curve(int, [(timedelta(seconds=1), 100)]) + c = csp.const(1000) + result = basket_stats([a, b, c]) + csp.add_graph_output("ticked_count", result.ticked_count) + csp.add_graph_output("ticked_sum", result.ticked_sum) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["ticked_count"]], [2, 1, 1]) + self.assertEqual([v for _, v in results["ticked_sum"]], [1010, 100, 20]) + + def test_dict_basket_input(self): + @numba_node + def basket_stats(prices: {str: ts[float]}) -> csp.Outputs(ticked_sum=ts[float], valid_sum=ts[float]): + ticked_sum = 0.0 + valid_sum = 0.0 + for key in prices.tickedkeys(): + ticked_sum = ticked_sum + prices.at(key) + for key in prices.validkeys(): + valid_sum = valid_sum + prices.at(key) + csp.output(ticked_sum=ticked_sum) + csp.output(valid_sum=valid_sum) + + @csp.graph + def g(): + a = csp.curve(float, [(timedelta(seconds=0), 100.0), (timedelta(seconds=2), 200.0)]) + b = csp.curve(float, [(timedelta(seconds=1), 50.0)]) + result = basket_stats({"a": a, "b": b}) + csp.add_graph_output("ticked_sum", result.ticked_sum) + csp.add_graph_output("valid_sum", result.valid_sum) + + # t=0: a=100 ticks -> ticked: a; valid: a + # t=1: b=50 ticks -> ticked: b; valid: a,b + # t=2: a=200 ticks -> ticked: a; valid: a,b + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["ticked_sum"]], [100.0, 50.0, 200.0]) + self.assertEqual([v for _, v in results["valid_sum"]], [100.0, 150.0, 250.0]) + + def test_dict_basket_input_modern_annotation(self): + @numba_node + def basket_stats(prices: dict[str, ts[float]]) -> csp.Outputs(ticked_sum=ts[float], valid_sum=ts[float]): + ticked_sum = 0.0 + valid_sum = 0.0 + for key in prices.tickedkeys(): + ticked_sum = ticked_sum + prices.at(key) + for key in prices.validkeys(): + valid_sum = valid_sum + prices.at(key) + csp.output(ticked_sum=ticked_sum) + csp.output(valid_sum=valid_sum) + + @csp.graph + def g(): + a = csp.curve(float, [(timedelta(seconds=0), 100.0), (timedelta(seconds=2), 200.0)]) + b = csp.curve(float, [(timedelta(seconds=1), 50.0)]) + result = basket_stats({"a": a, "b": b}) + csp.add_graph_output("ticked_sum", result.ticked_sum) + csp.add_graph_output("valid_sum", result.valid_sum) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["ticked_sum"]], [100.0, 50.0, 200.0]) + self.assertEqual([v for _, v in results["valid_sum"]], [100.0, 150.0, 250.0]) + + def test_dict_basket_string_keys(self): + @numba_node + def basket_stats(prices: dict[str, ts[float]]) -> csp.Outputs(ticked_sum=ts[float], valid_sum=ts[float]): + ticked_sum = 0.0 + valid_sum = 0.0 + for key in prices.tickedkeys(): + ticked_sum = ticked_sum + prices.at(key) + for key in prices.validkeys(): + valid_sum = valid_sum + prices.at(key) + csp.output(ticked_sum=ticked_sum) + csp.output(valid_sum=valid_sum) + + @csp.graph + def g(): + a = csp.curve(float, [(timedelta(seconds=0), 100.0), (timedelta(seconds=2), 200.0)]) + b = csp.curve(float, [(timedelta(seconds=1), 50.0)]) + result = basket_stats({"b": b, "a": a}) + csp.add_graph_output("ticked_sum", result.ticked_sum) + csp.add_graph_output("valid_sum", result.valid_sum) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["ticked_sum"]], [100.0, 50.0, 200.0]) + self.assertEqual([v for _, v in results["valid_sum"]], [100.0, 150.0, 250.0]) + + def test_dict_basket_int_keys(self): + @numba_node + def weighted_sum(prices: dict[int, ts[float]]) -> ts[float]: + total = 0.0 + if prices[10].valid(): + total = total + 10 * prices[10] + if prices[20].valid(): + total = total + 20 * prices[20] + return total + + @csp.graph + def g(): + ten = csp.curve(float, [(timedelta(seconds=0), 1.5), (timedelta(seconds=2), 3.0)]) + twenty = csp.curve(float, [(timedelta(seconds=1), 2.0)]) + result = weighted_sum({20: twenty, 10: ten}) + csp.add_graph_output("result", result) + + # t=0: 10 * 1.5 = 15 + # t=1: 10 * 1.5 + 20 * 2.0 = 55 + # t=2: 10 * 3.0 + 20 * 2.0 = 70 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [15.0, 55.0, 70.0]) + + def test_basket_ticked_and_valid(self): + @numba_node + def check_basket(xs: list[ts[int]]) -> csp.Outputs(tick_count=ts[int], all_valid=ts[int]): + with csp.state(): + s_count: int = 0 + + s_count = s_count + 1 + csp.output(tick_count=s_count) + if csp.valid(xs): + csp.output(all_valid=1) + else: + csp.output(all_valid=0) + + @csp.graph + def g(): + a = csp.curve(int, [(timedelta(seconds=0), 1)]) + b = csp.curve(int, [(timedelta(seconds=1), 2)]) + result = check_basket([a, b]) + csp.add_graph_output("tick_count", result.tick_count) + csp.add_graph_output("all_valid", result.all_valid) + + # t=0: a ticks, b not valid -> tick_count=1, all_valid=0 + # t=1: b ticks, both valid -> tick_count=2, all_valid=1 + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["tick_count"]], [1, 2]) + self.assertEqual([v for _, v in results["all_valid"]], [0, 1]) + + def test_nested_iteration_over_same_basket(self): + @numba_node + def cross_product(xs: list[ts[int]]) -> ts[int]: + total = 0 + for i in xs.keys(): + for j in xs.keys(): + total = total + xs.at(i) * xs.at(j) + return total + + @csp.graph + def g(): + result = cross_product([csp.const(2), csp.const(3)]) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1)) + self.assertEqual([v for _, v in results["result"]], [25]) + + def test_multiple_outputs(self): + @numba_node + def multi_out(x: ts[int]) -> csp.Outputs(doubled=ts[int], squared=ts[int], positive=ts[int]): + csp.output(doubled=x * 2) + csp.output(squared=x * x) + if x > 0: + csp.output(positive=x) + + @csp.graph + def g(): + values = csp.curve( + int, + [ + (timedelta(seconds=0), -2), + (timedelta(seconds=1), 3), + (timedelta(seconds=2), 0), + ], + ) + outputs = multi_out(values) + csp.add_graph_output("doubled", outputs.doubled) + csp.add_graph_output("squared", outputs.squared) + csp.add_graph_output("positive", outputs.positive) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["doubled"]], [-4, 6, 0]) + self.assertEqual([v for _, v in results["squared"]], [4, 9, 0]) + self.assertEqual([v for _, v in results["positive"]], [3]) + + def test_multiple_outputs_in_single_call(self): + @numba_node + def multi_out(x: ts[int]) -> csp.Outputs(doubled=ts[int], squared=ts[int]): + csp.output(doubled=x * 2, squared=x * x) + + @csp.graph + def g(): + outputs = multi_out(csp.const(3)) + csp.add_graph_output("doubled", outputs.doubled) + csp.add_graph_output("squared", outputs.squared) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1)) + self.assertEqual([v for _, v in results["doubled"]], [6]) + self.assertEqual([v for _, v in results["squared"]], [9]) + + def test_positional_named_outputs(self): + @numba_node + def multi_out(x: ts[int]) -> csp.Outputs(doubled=ts[int], squared=ts[int]): + csp.output(doubled, x * 2) + return csp.output(squared, x * x) + + @csp.graph + def g(): + outputs = multi_out(csp.const(3)) + csp.add_graph_output("doubled", outputs.doubled) + csp.add_graph_output("squared", outputs.squared) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1)) + self.assertEqual([v for _, v in results["doubled"]], [6]) + self.assertEqual([v for _, v in results["squared"]], [9]) + + def test_return_multiple_outputs(self): + @numba_node + def multi_out(x: ts[int]) -> csp.Outputs(doubled=ts[int], squared=ts[int]): + if x > 0: + return csp.output(doubled=x * 2, squared=x * x) + + @csp.graph + def g(): + values = csp.curve(int, [(timedelta(seconds=0), -2), (timedelta(seconds=1), 3)]) + outputs = multi_out(values) + csp.add_graph_output("doubled", outputs.doubled) + csp.add_graph_output("squared", outputs.squared) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=2)) + self.assertEqual([v for _, v in results["doubled"]], [6]) + self.assertEqual([v for _, v in results["squared"]], [9]) + + def test_single_set_output(self): + @numba_node + def emit_scaled(x: ts[int]) -> ts[int]: + csp.output(x * 10) + + @csp.graph + def g(): + values = csp.curve(int, [(timedelta(seconds=i), i + 1) for i in range(3)]) + result = emit_scaled(values) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [10, 20, 30]) + + def test_conditional_output(self): + @numba_node + def emit_every_n(x: ts[int]) -> ts[int]: + with csp.state(): + s_count = 0 + + s_count = s_count + 1 + if s_count % 3 == 0: + return x + + @csp.graph + def g(): + values = csp.curve(int, [(timedelta(seconds=i), (i + 1) * 10) for i in range(9)]) + result = emit_every_n(values) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [30, 60, 90]) + + def test_dict_basket_with_enum_output(self): + class Trend(csp.Enum): + DOWN = -1 + FLAT = 0 + UP = 1 + + @numba_node + def trend_analyzer(prices: {str: ts[float]}) -> ts[Trend]: + with csp.state(): + s_prev = 0.0 + s_has_prev = 0 + + total = 0.0 + for key in prices.validkeys(): + total = total + prices.at(key) + + if s_has_prev == 0: + s_prev = total + s_has_prev = 1 + return Trend.FLAT + + diff = total - s_prev + s_prev = total + if diff > 0.01: + return Trend.UP + elif diff < -0.01: + return Trend.DOWN + return Trend.FLAT + + @csp.graph + def g(): + a = csp.curve(float, [(timedelta(seconds=0), 100.0), (timedelta(seconds=1), 110.0)]) + b = csp.curve(float, [(timedelta(seconds=0), 200.0), (timedelta(seconds=1), 180.0)]) + result = trend_analyzer({"a": a, "b": b}) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + self.assertEqual(actual, [Trend.FLAT, Trend.DOWN]) + + def test_running_statistics(self): + @numba_node + def running_stats(x: ts[float]) -> csp.Outputs(mean=ts[float], variance=ts[float]): + with csp.state(): + s_count = 0 + s_mean = 0.0 + s_m2 = 0.0 + + s_count = s_count + 1 + delta = x - s_mean + s_mean = s_mean + delta / s_count + delta2 = x - s_mean + s_m2 = s_m2 + delta * delta2 + + csp.output(mean=s_mean) + if s_count > 1: + var = s_m2 / (s_count - 1) + csp.output(variance=var) + + @csp.graph + def g(): + values = csp.curve( + float, [(timedelta(milliseconds=i), v) for i, v in enumerate([2.0, 4.0, 4.0, 4.0, 5.0, 5.0, 7.0, 9.0])] + ) + outputs = running_stats(values) + csp.add_graph_output("mean", outputs.mean) + csp.add_graph_output("variance", outputs.variance) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1)) + means = [v for _, v in results["mean"]] + variances = [v for _, v in results["variance"]] + self.assertAlmostEqual(means[-1], 5.0, places=5) + self.assertAlmostEqual(variances[-1], 32 / 7, places=5) + + +class TestStructSupport(unittest.TestCase): + def test_invalid_struct_input_fails_callback(self): + class Point(csp.Struct): + x: int + + @numba_node + def read_x(trigger: ts[int], value: ts[Point]) -> ts[int]: + return trigger + value.x + + @csp.graph + def g(): + value = csp.curve(Point, [(timedelta(seconds=1), Point(x=2))]) + csp.add_graph_output("result", read_x(csp.const(1), value)) + + with self.assertRaisesRegex(RuntimeError, "Compiled Numba callback failed"): + csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=2)) + + def test_unset_primitive_field_fails_callback(self): + class Point(csp.Struct): + x: int + + @numba_node + def read_x(value: ts[Point]) -> ts[int]: + return value.x + + @csp.graph + def g(): + csp.add_graph_output("result", read_x(csp.const(Point()))) + + with self.assertRaisesRegex(RuntimeError, "Compiled Numba callback failed"): + csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1)) + + def test_unset_enum_field_fails_callback(self): + class Direction(IntEnum): + UP = 1 + + class Velocity(csp.Struct): + direction: Direction + + @numba_node + def read_direction(value: ts[Velocity]) -> ts[int]: + if value.direction == Direction.UP: + return 1 + return 0 + + @csp.graph + def g(): + csp.add_graph_output("result", read_direction(csp.const(Velocity()))) + + with self.assertRaisesRegex(RuntimeError, "Compiled Numba callback failed"): + csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1)) + + def test_assign_unset_struct_enum_field(self): + class Direction(IntEnum): + UP = 1 + DOWN = -1 + + class Velocity(csp.Struct): + direction: Direction + + @numba_node + def set_direction(x: ts[int]) -> ts[int]: + with csp.state(): + s_velocity: Velocity = Velocity() + + s_velocity.direction = Direction.UP + if s_velocity.direction == Direction.UP: + return x + return 0 + + @csp.graph + def g(): + csp.add_graph_output("result", set_direction(csp.const(7))) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=1)) + self.assertEqual([value for _, value in results["result"]], [7]) + + def test_struct_enum_input_uses_declared_family(self): + class Direction(csp.Enum): + DOWN = -1 + UP = 7 + + class Velocity(csp.Struct): + direction: Direction + + @numba_node + def direction_code(value: ts[Velocity]) -> ts[int]: + if value.direction == Direction.DOWN: + return -1 + if value.direction == Direction.UP: + return 7 + return 0 + + @csp.graph + def g(): + values = csp.curve( + Velocity, + [ + (timedelta(seconds=0), Velocity(direction=Direction.DOWN)), + (timedelta(seconds=1), Velocity(direction=Direction.UP)), + ], + ) + csp.add_graph_output("result", direction_code(values)) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=2)) + self.assertEqual([value for _, value in results["result"]], [-1, 7]) + + def test_struct_output_from_state_has_independent_ticks(self): + class Point(csp.Struct): + x: float + y: float + + @numba_node + def emit_struct(x: ts[float]) -> ts[Point]: + with csp.state(): + s_point: Point = Point(x=1.0, y=2.0) + + s_point.x = s_point.x + x + return s_point + + @csp.graph + def g(): + values = csp.curve(float, [(timedelta(seconds=i), float(i + 1)) for i in range(3)]) + result = emit_struct(values) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=3)) + self.assertEqual([(value.x, value.y) for _, value in results["result"]], [(2.0, 2.0), (4.0, 2.0), (7.0, 2.0)]) + + def test_struct_input_to_named_output(self): + class Point(csp.Struct): + x: float + y: float + + @numba_node + def forward(value: ts[Point]) -> csp.Outputs(point=ts[Point], x=ts[float]): + csp.output(point=value) + csp.output(x=value.x) + + @csp.graph + def g(): + values = csp.curve( + Point, + [ + (timedelta(seconds=0), Point(x=1.0, y=2.0)), + (timedelta(seconds=1), Point(x=3.0, y=4.0)), + ], + ) + result = forward(values) + csp.add_graph_output("point", result.point) + csp.add_graph_output("x", result.x) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=2)) + self.assertEqual([(value.x, value.y) for _, value in results["point"]], [(1.0, 2.0), (3.0, 4.0)]) + self.assertEqual([value for _, value in results["x"]], [1.0, 3.0]) + + def test_struct_output_preserves_enum_field(self): + class Direction(csp.Enum): + UP = 1 + DOWN = 2 + + class Velocity(csp.Struct): + direction: Direction + speed: float + + @numba_node + def forward(value: ts[Velocity]) -> ts[Velocity]: + return value + + @csp.graph + def g(): + values = csp.curve( + Velocity, + [ + (timedelta(seconds=0), Velocity(direction=Direction.UP, speed=1.0)), + (timedelta(seconds=1), Velocity(direction=Direction.DOWN, speed=2.0)), + ], + ) + csp.add_graph_output("result", forward(values)) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=2)) + self.assertEqual( + [(value.direction, value.speed) for _, value in results["result"]], + [(Direction.UP, 1.0), (Direction.DOWN, 2.0)], + ) + + def test_invalid_struct_var(self): + class NonNativePoint(csp.Struct): + x: float + label: str + + @numba_node + def non_native_struct_state(x: ts[float]) -> ts[float]: + with csp.state(): + s_point: NonNativePoint = NonNativePoint(x=0.0, label="start") + + s_point.x = s_point.x + x + return s_point.x + + @csp.graph + def g(): + result = non_native_struct_state(csp.const(1.0)) + csp.add_graph_output("result", result) + + with self.assertRaisesRegex( + TypeError, r"numba_node only supports native csp\.Struct types; 'NonNativePoint' is non-native" + ): + csp.build_graph(g) + + def test_struct_state_multiple_structs(self): + """Test multiple struct state variables.""" + + class Direction(csp.Enum): + UP = 1 + DOWN = -1 + + class Position(csp.Struct): + x: float + steps: int + active: bool + + class Velocity(csp.Struct): + vx: float + step_mult: int + direction: Direction + + @numba_node + def physics_node(dt: ts[float]) -> ts[float]: + with csp.state(): + s_pos: Position = Position(x=0.0, steps=0, active=False) + s_vel: Velocity = Velocity(vx=1.0, step_mult=2, direction=Direction.UP) + + s_pos.steps = s_pos.steps + 1 + s_pos.active = not s_pos.active + + if s_pos.active: + delta = s_vel.vx * dt * s_vel.step_mult + if s_vel.direction == Direction.UP: + s_pos.x = s_pos.x + delta + else: + s_pos.x = s_pos.x - delta + + return s_pos.x + s_pos.steps + + @csp.graph + def g(): + deltas = csp.curve(float, [(timedelta(seconds=i), 0.5) for i in range(4)]) + result = physics_node(deltas) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + # steps: 1, 2, 3, 4 + # active toggles: True, False, True, False + # x updates only on active ticks: +1.0, +0.0, +1.0, +0.0 + # result = x + steps -> 2.0, 3.0, 5.0, 6.0 + self.assertEqual(actual, [2.0, 3.0, 5.0, 6.0]) + + def test_struct_state_conditional_field_update(self): + class Accumulator(csp.Struct): + positive_sum: float + negative_sum: float + + @numba_node + def accumulate_node(x: ts[float]) -> csp.Outputs(positive=ts[float], negative=ts[float]): + with csp.state(): + s_acc: Accumulator = Accumulator(positive_sum=0.0, negative_sum=0.0) + + if x > 0: + s_acc.positive_sum = s_acc.positive_sum + x + csp.output(positive=s_acc.positive_sum) + else: + s_acc.negative_sum = s_acc.negative_sum + x + csp.output(negative=s_acc.negative_sum) + + @csp.graph + def g(): + values = csp.curve( + float, + [ + (timedelta(seconds=0), 10.0), + (timedelta(seconds=1), -5.0), + (timedelta(seconds=2), 20.0), + (timedelta(seconds=3), -10.0), + ], + ) + result = accumulate_node(values) + csp.add_graph_output("positive", result.positive) + csp.add_graph_output("negative", result.negative) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + self.assertEqual([v for _, v in results["positive"]], [10.0, 30.0]) + self.assertEqual([v for _, v in results["negative"]], [-5.0, -15.0]) + + def test_iterate_basket_of_structs(self): + class Point(csp.Struct): + x: float + y: float + + @numba_node + def sum_x_from_basket(points: list[ts[Point]]) -> ts[float]: + total = 0.0 + for i in range(len(points)): + if points.at(i).ticked(): + total = total + points.at(i).x + return total + + @csp.graph + def g(): + # Create individual Point signals (basket of 3 struct signals) + p1 = csp.curve( + Point, + [ + (timedelta(seconds=0), Point(x=1.0, y=10.0)), + (timedelta(seconds=2), Point(x=10.0, y=100.0)), + ], + ) + p2 = csp.curve( + Point, + [ + (timedelta(seconds=0), Point(x=2.0, y=20.0)), + (timedelta(seconds=1), Point(x=5.0, y=50.0)), + ], + ) + p3 = csp.curve( + Point, + [ + (timedelta(seconds=0), Point(x=3.0, y=30.0)), + ], + ) + + # Pass as a basket + result = sum_x_from_basket([p1, p2, p3]) + csp.add_graph_output("result", result) + + results = csp.run(g, starttime=datetime(2024, 1, 1), endtime=timedelta(seconds=10)) + actual = [v for _, v in results["result"]] + # t=0: all 3 tick: 1.0 + 2.0 + 3.0 = 6.0 + # t=1: only p2 ticks: 5.0 + # t=2: only p1 ticks: 10.0 + self.assertEqual(actual, [6.0, 5.0, 10.0]) + + +if __name__ == "__main__": + unittest.main() diff --git a/csp/tests/test_numba_type_checking.py b/csp/tests/test_numba_type_checking.py deleted file mode 100644 index 7486fc7bc..000000000 --- a/csp/tests/test_numba_type_checking.py +++ /dev/null @@ -1,170 +0,0 @@ -import typing -import unittest -from datetime import datetime - -import csp -from csp import ts - - -class TestNumbaTypeChecking(unittest.TestCase): - @unittest.skip("numba not yet used, tests fail on newer numba we get in our 3.8 build") - def test_graph_build_type_checking(self): - @csp.numba_node - def typed_ts(x: ts[int]): - if csp.ticked(x): - pass - - @csp.numba_node - def typed_scalar(x: ts[int], y: str): - if csp.ticked(x): - pass - - @csp.graph - def graph(): - i = csp.const(5) - typed_ts(i) - - typed_scalar(i, "xyz") - - with self.assertRaisesRegex(TypeError, "Expected ts\\[int\\] for argument 'x', got ts\\[str\\]"): - s = csp.const("xyz") - ## THIS SHOULD RAISE, passing ts[str] but typed takes ts[int] - typed_ts(s) - - with self.assertRaisesRegex(TypeError, "Expected str for argument 'y', got 123 \\(int\\)"): - ## THIS SHOULD RAISE, passing int instead of str - typed_scalar(i, 123) - - csp.run(graph, starttime=datetime(2020, 2, 7, 9), endtime=datetime(2020, 2, 7, 9, 1)) - - @unittest.skip("numba not yet used, tests fail on newer numba we get in our 3.8 build") - def test_runtime_type_check(self): - ## native output type - @csp.numba_node - def typed_int(x: ts["T"]) -> ts[int]: - if csp.ticked(x): - return x - - # TODO: Uncomment - # @csp.numba_node - # def typed_alarm(v: '~T', alarm_type: 'V') -> outputs(ts['V']): - # with csp.alarms(): - # alarm = csp.alarm( 'V' ) - # with csp.start(): - # csp.schedule_alarm(alarm, timedelta(), v) - # - # if csp.ticked(alarm): - # return alarm - - # Valid - csp.run(typed_int, csp.const(5), starttime=datetime(2020, 2, 7)) - - # Invalid - with self.assertRaisesRegex(RuntimeError, "Unable to resolve getter function for type.*"): - csp.run(typed_int, csp.const("5"), starttime=datetime(2020, 2, 7)) - - # TODO: uncomment - # # valid - # csp.run(typed_alarm, 5, int, starttime=datetime(2020, 2, 7)) - # csp.run(typed_alarm, 5, object, starttime=datetime(2020, 2, 7)) - # csp.run(typed_alarm, [1, 2, 3], [int], starttime=datetime(2020, 2, 7)) - # - # # Invalid - # with self.assertRaisesRegex(TypeError, - # '"typed_alarm" node expected output type on output #0 to be of type "str" got type "int"'): - # csp.run(typed_alarm, 5, str, starttime=datetime(2020, 2, 7)) - # - # with self.assertRaisesRegex(TypeError, - # '"typed_alarm" node expected output type on output #0 to be of type "bool" got type "int"'): - # csp.run(typed_alarm, 5, bool, starttime=datetime(2020, 2, 7)) - # - # with self.assertRaisesRegex(TypeError, - # '"typed_alarm" node expected output type on output #0 to be of type "str" got type "list"'): - # csp.run(typed_alarm, [1, 2, 3], str, starttime=datetime(2020, 2, 7)) - - @unittest.skip("numba not yet used, tests fail on newer numba we get in our 3.8 build") - def test_dict_type_resolutions(self): - @csp.numba_node - def typed_dict_int_int(x: {int: int}): - pass - - @csp.numba_node - def typed_dict_int_int2(x: typing.Dict[int, int]): - pass - - @csp.numba_node - def typed_dict_int_float(x: {int: int}): - pass - - @csp.numba_node - def typed_dict_float_float(x: {float: float}): - pass - - @csp.numba_node - def typed_dict(x: {"T": "V"}): - pass - - @csp.numba_node - def deep_nested_generic_resolution(x: "T1", y: "T2", z: {"T1": {"T2": [{"T1"}]}}): - pass - - @csp.graph - def graph(): - d_i_i = csp.const({1: 2, 3: 4}) - csp.add_graph_output("o1", d_i_i) - - # Ok int dict expected - typed_dict_int_int({1: 2, 3: 4}) - - # Ok int dict expected - typed_dict_int_int2({1: 2, 3: 4}) - - typed_dict_float_float({1: 2}) - typed_dict_float_float({1.0: 2}) - typed_dict_float_float({}) - - with self.assertRaisesRegex(TypeError, r"Expected typing.Dict\[int, int\] for argument 'x', got .*"): - # Passing a float value instead of expected ints - typed_dict_int_int2({1: 2, 3: 4.0}) - - l_good = csp.const.using(T={int: float})({}) - csp.add_graph_output("o2", l_good) - l_good = csp.const.using(T={int: float})({2: 1}) - csp.add_graph_output("o3", l_good) - l_good = csp.const.using(T={int: float})({2: 1.0}) - csp.add_graph_output("o4", l_good) - - csp.run(graph, starttime=datetime(2020, 2, 7, 9), endtime=datetime(2020, 2, 7, 9, 1)) - - @unittest.skip("numba not yet used, tests fail on newer numba we get in our 3.8 build") - def test_list_type_resolutions(self): - @csp.numba_node - def typed_list_int(x: [int]): - pass - - @csp.numba_node - def typed_list_int2(x: typing.List[int]): - pass - - @csp.numba_node - def typed_list_float(x: [float]): - pass - - def graph(): - l_i = csp.const([1, 2, 3, 4]) - - typed_list_int([]) - typed_list_int([1, 2, 3]) - typed_list_int2([1, 2, 3]) - typed_list_float([1, 2, 3]) - typed_list_float([1, 2, 3.0]) - - with self.assertRaisesRegex(TypeError, r"Expected typing.List\[int\] for argument 'x', got .*"): - # Passing a float value instead of expected ints - typed_list_int([1, 2, 3.0]) - - csp.run(graph, starttime=datetime(2020, 2, 7, 9), endtime=datetime(2020, 2, 7, 9, 1)) - - -if __name__ == "__main__": - unittest.main() diff --git a/examples/01_basics/e3_show_graph.py b/examples/01_basics/e3_show_graph.py index d6cd0a2e5..c219ac257 100644 --- a/examples/01_basics/e3_show_graph.py +++ b/examples/01_basics/e3_show_graph.py @@ -25,13 +25,14 @@ def my_graph(): csp.print("spread2", s2) -def main(): - # open in graphviz viewer - # csp.show_graph(my_graph) - # or output to file - csp.show_graph(my_graph, graph_filename="tmp.png") +def main(show_graph: bool = False): + if show_graph: + # open in graphviz viewer + # csp.show_graph(my_graph) + # or output to file + csp.show_graph(my_graph, graph_filename="tmp.png") csp.run(my_graph, starttime=datetime(2020, 3, 1), endtime=timedelta(seconds=10)) if __name__ == "__main__": - main() + main(show_graph=True) diff --git a/examples/98_just_for_fun/e1_csp_nand_computer.py b/examples/98_just_for_fun/e1_csp_nand_computer.py index 378663f9f..f95d20db1 100644 --- a/examples/98_just_for_fun/e1_csp_nand_computer.py +++ b/examples/98_just_for_fun/e1_csp_nand_computer.py @@ -104,11 +104,12 @@ def my_graph(bits: int = 16): csp.print("x+y_bits", basket_to_bitstring(add)) -def main(): - # Show graph with 4-bit ints to limit size - csp.showgraph.show_graph(my_graph, 4) +def main(show_graph: bool = False): + if show_graph: + # Show graph with 4-bit ints to limit size + csp.showgraph.show_graph(my_graph, 4) csp.run(my_graph, starttime=datetime(2022, 6, 24)) if __name__ == "__main__": - main() + main(show_graph=True) diff --git a/pyproject.toml b/pyproject.toml index 1793b3825..7e90a06e8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ requires-python = ">=3.10" dependencies = [ "deprecated", + "numba_cfunc_compiler>=1.0.0", "numpy>=1", "packaging", "pandas",