From 879b4bda69842c65482fc7c35c00369971ad092d Mon Sep 17 00:00:00 2001
From: =?utf8?q?=EB=B0=95=EC=B2=9C=EA=B5=90/On-Device=20Lab=28SR=29/Enginee?=
 =?utf8?q?r/=EC=82=BC=EC=84=B1=EC=A0=84=EC=9E=90?= <ch.bahk@samsung.com>
Date: Wed, 6 Nov 2019 13:30:08 +0900
Subject: [PATCH] [exo] Convert TensorReduce (#8741)

* [exo] Convert TensorReduce

This commit introduces convert stage for canonical TensorReduce node

Signed-off-by: Cheongyo Bahk <ch.bahk@samsung.com>

* Strict check
---
 .../exo/src/Conversion/TensorReduceConverter.cpp   | 93 ++++++++++++++++++++++
 .../exo/src/Conversion/TensorReduceConverter.h     | 46 +++++++++++
 2 files changed, 139 insertions(+)
 create mode 100644 compiler/exo/src/Conversion/TensorReduceConverter.cpp
 create mode 100644 compiler/exo/src/Conversion/TensorReduceConverter.h

diff --git a/compiler/exo/src/Conversion/TensorReduceConverter.cpp b/compiler/exo/src/Conversion/TensorReduceConverter.cpp
new file mode 100644
index 0000000..56d992f
--- /dev/null
+++ b/compiler/exo/src/Conversion/TensorReduceConverter.cpp
@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "TensorReduceConverter.h"
+
+#include "Dialect/IR/TFLNodes.h"
+#include "Check.h"
+
+#include <loco.h>
+#include <loco/Service/ShapeInference.h>
+
+namespace
+{
+
+/**
+ * @brief  Convert given TensorReduce as TFLMean
+ *
+ * <Before>
+ *    In --- loco::TensorReduce --- Out(s)
+ *
+ * <After>
+ *    In -------- locoex::TFLMean --- Out(s)
+ *                /
+ *    TFLConst ---
+ *    (reduction indices)
+ */
+bool convert_as_mean(loco::TensorReduce *origin)
+{
+  EXO_ASSERT(origin->func() == loco::ReduceFunc::Mean, "func should be Mean for this helper");
+  EXO_ASSERT(origin->input(), "TensorReduce has no input");
+
+  auto *graph = origin->graph();
+
+  // Make reduction indicies TFLConst node
+  auto reduction = graph->nodes()->create<locoex::TFLConst>();
+  {
+    auto input_rank = loco::shape_get(origin->input()).as<loco::TensorShape>().rank();
+
+    std::vector<int32_t> red_vec;
+    for (uint32_t axis = 0; axis < input_rank; ++axis)
+      if (origin->axes()->defined(axis))
+        red_vec.push_back(static_cast<int32_t>(axis));
+
+    const loco::DataType S32 = loco::DataType::S32;
+
+    reduction->dtype(S32);
+    reduction->rank(1);
+    reduction->dim(0) = red_vec.size();
+    reduction->size<S32>(red_vec.size());
+    for (uint32_t i = 0; i < red_vec.size(); ++i)
+      reduction->at<S32>(i) = red_vec.at(i);
+  }
+
+  // Make TFLMean node to replace
+  auto mean = graph->nodes()->create<locoex::TFLMean>();
+  mean->input(origin->input());
+  mean->reduction_indices(reduction);
+  mean->keep_dims(true); // Canonical TensorReduce always keep dimensions
+
+  // replace canonical node
+  loco::replace(origin).with(mean);
+  origin->input(nullptr);
+
+  return true;
+}
+
+} // namespace
+
+namespace exo
+{
+
+bool TensorReduceConverter::convert(loco::TensorReduce *origin)
+{
+  if (origin->func() == loco::ReduceFunc::Mean)
+    return convert_as_mean(origin);
+  else
+    throw std::runtime_error("NYI ReduceFunc");
+}
+
+} // namespace exo
diff --git a/compiler/exo/src/Conversion/TensorReduceConverter.h b/compiler/exo/src/Conversion/TensorReduceConverter.h
new file mode 100644
index 0000000..dfd65ad
--- /dev/null
+++ b/compiler/exo/src/Conversion/TensorReduceConverter.h
@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2019 Samsung Electronics Co., Ltd. All Rights Reserved
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef __TENSOR_REDUCE_CONVERTER__
+#define __TENSOR_REDUCE_CONVERTER__
+
+#include "CanonicalNodeConverter.h"
+
+#include <loco.h>
+
+namespace exo
+{
+
+/**
+ * @brief  Convert loco::TensorReduce to appropriate TFL reduce operation
+ * @note   loco::TensorReduce always keep dimensions
+ *
+ * Currently support:
+ *   - When loco::TensorReduce::func() == Mean, convert to TFLMean + TFLConst
+ *   - TODO Support other cases
+ */
+class TensorReduceConverter : public CanonicalNodeConverter<loco::TensorReduce>
+{
+public:
+  const char *name(void) const final { return "exo::TensorReduceConverter"; }
+
+public:
+  bool convert(loco::TensorReduce *origin) final;
+};
+
+} // namespace exo
+
+#endif // __TENSOR_REDUCE_CONVERTER__
-- 
2.7.4