diff options
author | Mateusz Chudyk <mateuszchudyk@gmail.com> | 2019-07-19 19:20:18 +0300 |
---|---|---|
committer | Mateusz Chudyk <mateuszchudyk@gmail.com> | 2019-07-19 19:20:18 +0300 |
commit | 87e51cd18a05f503d4f04709ac4121388b206c48 (patch) | |
tree | e7bce31f87e0bf7abbad76a5b83a3492cb64bf58 | |
parent | a80efb933528ffbad8d17f7f7a915aeee1c3e0e7 (diff) |
Add Writer callback
-rw-r--r-- | callbacks/configs.h | 7 | ||||
-rw-r--r-- | callbacks/implementations.inl | 16 |
2 files changed, 23 insertions, 0 deletions
diff --git a/callbacks/configs.h b/callbacks/configs.h index 78986c5..4a01e2e 100644 --- a/callbacks/configs.h +++ b/callbacks/configs.h @@ -6,6 +6,13 @@ namespace callbacks { struct Dummy { }; +template <typename OutputBufferType> +struct Write { + OutputBufferType* addr; + + Write(OutputBufferType* addr) : addr(addr) {} +}; + struct UnquantizeAndWrite { float unquant_mult; float* addr; diff --git a/callbacks/implementations.inl b/callbacks/implementations.inl index 6e3a10f..773a5f8 100644 --- a/callbacks/implementations.inl +++ b/callbacks/implementations.inl @@ -53,6 +53,22 @@ public: }; /* + * Write + */ +template <typename OutputBufferType> +class CallbackImpl<CPUType::CPU_NAME, Write<OutputBufferType>> { +public: + CPU_ATTR CallbackImpl(const Write<OutputBufferType>& config) : config(config) {} + + CPU_ATTR void operator()(vi input, const OutputBufferInfo& info) { + kernels::write(input, config.addr, info.row_idx * info.cols + info.col_idx); + } + +private: + Write<OutputBufferType> config; +}; + +/* * UnquantizeAndWrite */ template <> class CallbackImpl<CPUType::CPU_NAME, UnquantizeAndWrite> { |