1/*
2 * Licensed to the Apache Software Foundation (ASF) under one
3 * or more contributor license agreements. See the NOTICE file
4 * distributed with this work for additional information
5 * regarding copyright ownership. The ASF licenses this file
6 * to you under the Apache License, Version 2.0 (the
7 * "License"); you may not use this file except in compliance
8 * with the License. You may obtain a copy of the License at
9 *
10 * http://www.apache.org/licenses/LICENSE-2.0
11 *
12 * Unless required by applicable law or agreed to in writing,
13 * software distributed under the License is distributed on an
14 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15 * KIND, either express or implied. See the License for the
16 * specific language governing permissions and limitations
17 * under the License.
18 */
19
20/*!
21 * \file cpu_device_api.cc
22 */
23#include <dmlc/thread_local.h>
24#include <tvm/runtime/device_api.h>
25#include <tvm/runtime/logging.h>
26#include <tvm/runtime/registry.h>
27
28#include <cstdlib>
29#include <cstring>
30
31#include "workspace_pool.h"
32
33#ifdef __ANDROID__
34#include <android/api-level.h>
35#endif
36
37namespace tvm {
38namespace runtime {
39class CPUDeviceAPI final : public DeviceAPI {
40 public:
41 void SetDevice(Device dev) final {}
42 void GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* rv) final {
43 if (kind == kExist) {
44 *rv = 1;
45 }
46 }
47 void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint) final {
48 void* ptr;
49#if _MSC_VER
50 ptr = _aligned_malloc(nbytes, alignment);
51 if (ptr == nullptr) throw std::bad_alloc();
52#elif defined(__ANDROID__) && __ANDROID_API__ < 17
53 ptr = memalign(alignment, nbytes);
54 if (ptr == nullptr) throw std::bad_alloc();
55#else
56 // posix_memalign is available in android ndk since __ANDROID_API__ >= 17
57 int ret = posix_memalign(&ptr, alignment, nbytes);
58 if (ret != 0) throw std::bad_alloc();
59#endif
60 return ptr;
61 }
62
63 void FreeDataSpace(Device dev, void* ptr) final {
64#if _MSC_VER
65 _aligned_free(ptr);
66#else
67 free(ptr);
68#endif
69 }
70
71 void StreamSync(Device dev, TVMStreamHandle stream) final {}
72
73 void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final;
74 void FreeWorkspace(Device dev, void* data) final;
75
76 static CPUDeviceAPI* Global() {
77 // NOTE: explicitly use new to avoid exit-time destruction of global state
78 // Global state will be recycled by OS as the process exits.
79 static auto* inst = new CPUDeviceAPI();
80 return inst;
81 }
82
83 protected:
84 void CopyDataFromTo(const void* from, size_t from_offset, void* to, size_t to_offset, size_t size,
85 Device dev_from, Device dev_to, DLDataType type_hint,
86 TVMStreamHandle stream) final {
87 memcpy(static_cast<char*>(to) + to_offset, static_cast<const char*>(from) + from_offset, size);
88 }
89};
90
91struct CPUWorkspacePool : public WorkspacePool {
92 CPUWorkspacePool() : WorkspacePool(kDLCPU, CPUDeviceAPI::Global()) {}
93};
94
95void* CPUDeviceAPI::AllocWorkspace(Device dev, size_t size, DLDataType type_hint) {
96 return dmlc::ThreadLocalStore<CPUWorkspacePool>::Get()->AllocWorkspace(dev, size);
97}
98
99void CPUDeviceAPI::FreeWorkspace(Device dev, void* data) {
100 dmlc::ThreadLocalStore<CPUWorkspacePool>::Get()->FreeWorkspace(dev, data);
101}
102
103TVM_REGISTER_GLOBAL("device_api.cpu").set_body([](TVMArgs args, TVMRetValue* rv) {
104 DeviceAPI* ptr = CPUDeviceAPI::Global();
105 *rv = static_cast<void*>(ptr);
106});
107} // namespace runtime
108} // namespace tvm
109