From baa08394ada22cf079bab7b3ba1f77d42ddcad6b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E5=AF=85?= <liyin@xiaomi.com>
Date: Fri, 1 Sep 2017 11:19:12 +0800
Subject: [PATCH] map neon allocator

---
 mace/core/allocator.cc | 10 ++++++----
 mace/core/operator.cc  |  2 +-
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/mace/core/allocator.cc b/mace/core/allocator.cc
index fd1f50c3..8b8a79ea 100644
--- a/mace/core/allocator.cc
+++ b/mace/core/allocator.cc
@@ -16,10 +16,12 @@ void SetCPUAllocator(CPUAllocator* alloc) {
 }
 
 Allocator* GetDeviceAllocator(DeviceType type) {
-  if (type == DeviceType::CPU) {
-    return cpu_allocator();
-  } else {
-    REQUIRE(false, "device type ", type, " is not supported.");
+  switch (type) {
+    case DeviceType::CPU:
+    case DeviceType::NEON:
+      return cpu_allocator();
+    default:
+      REQUIRE(false, "device type ", type, " is not supported.");
   }
   return nullptr;
 }
diff --git a/mace/core/operator.cc b/mace/core/operator.cc
index 078574a7..2e5086ac 100644
--- a/mace/core/operator.cc
+++ b/mace/core/operator.cc
@@ -23,7 +23,7 @@ MACE_DEFINE_REGISTRY(
     OperatorBase,
     const OperatorDef&,
     Workspace*);
-MACE_REGISTER_DEVICE_TYPE(DeviceType::NEON, CPUOperatorRegistry);
+MACE_REGISTER_DEVICE_TYPE(DeviceType::NEON, NEONOperatorRegistry);
 
 unique_ptr<OperatorBase> CreateOperator(
     const OperatorDef& operator_def,
-- 
GitLab