map neon allocator

baa08394 · 李寅 · 8f3f7d9c · baa08394 · baa08394
隐藏空白更改
内联并排

Showing with 7 addition and 5 deletion

mace/core/allocator.cc mace/core/allocator.cc +6 -4

mace/core/operator.cc mace/core/operator.cc +1 -1

未找到文件。
--- 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) {
+  switch (type) {
-    return cpu_allocator();
+    case DeviceType::CPU:
-  } else {
+    case DeviceType::NEON:
-    REQUIRE(false, "device type ", type, " is not supported.");
+      return cpu_allocator();
+    default:
+      REQUIRE(false, "device type ", type, " is not supported.");
  }
  return nullptr;
 }

--- 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,