From 87e51cd18a05f503d4f04709ac4121388b206c48 Mon Sep 17 00:00:00 2001
From: Mateusz Chudyk <mateuszchudyk@gmail.com>
Date: Fri, 19 Jul 2019 17:20:18 +0100
Subject: Add Writer callback

---
 callbacks/configs.h           |  7 +++++++
 callbacks/implementations.inl | 16 ++++++++++++++++
 2 files changed, 23 insertions(+)

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
@@ -52,6 +52,22 @@ public:
   CPU_ATTR void operator()(vi, const OutputBufferInfo&) {}
 };
 
+/*
+ * 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
  */
-- 
cgit v1.2.3