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 workspace_pool.h |
22 | * \brief Workspace pool utility. |
23 | */ |
24 | #ifndef TVM_RUNTIME_WORKSPACE_POOL_H_ |
25 | #define TVM_RUNTIME_WORKSPACE_POOL_H_ |
26 | |
27 | #include <tvm/runtime/device_api.h> |
28 | |
29 | #include <memory> |
30 | #include <vector> |
31 | |
32 | namespace tvm { |
33 | namespace runtime { |
34 | /*! |
35 | * \brief A workspace pool to manage |
36 | * |
37 | * \note We have the following assumption about backend temporal |
38 | * workspace allocation, and will optimize for such assumption, |
39 | * some of these assumptions can be enforced by the compiler. |
40 | * |
41 | * - Only a few allocation will happen, and space will be released after use. |
42 | * - The release order is usually in reverse order of allocate |
43 | * - Repeative pattern of same allocations over different runs. |
44 | */ |
45 | class TVM_DLL WorkspacePool { |
46 | public: |
47 | /*! |
48 | * \brief Create pool with specific device type and device. |
49 | * \param device_type The device type. |
50 | * \param device_api The device API. |
51 | */ |
52 | WorkspacePool(DLDeviceType device_type, DeviceAPI* device_api); |
53 | /*! \brief destructor */ |
54 | ~WorkspacePool(); |
55 | /*! |
56 | * \brief Allocate temporal workspace. |
57 | * \param dev The device of allocation. |
58 | * \param size The size to be allocated. |
59 | */ |
60 | void* AllocWorkspace(Device dev, size_t size); |
61 | /*! |
62 | * \brief Free temporal workspace in backend execution. |
63 | * |
64 | * \param dev The device of allocation. |
65 | * \param ptr The pointer to be freed. |
66 | */ |
67 | void FreeWorkspace(Device dev, void* ptr); |
68 | |
69 | private: |
70 | class Pool; |
71 | /*! \brief pool of device local array */ |
72 | std::vector<Pool*> array_; |
73 | /*! \brief device type this pool support */ |
74 | DLDeviceType device_type_; |
75 | /*! \brief The device API */ |
76 | DeviceAPI* device_; |
77 | }; |
78 | |
79 | } // namespace runtime |
80 | } // namespace tvm |
81 | #endif // TVM_RUNTIME_WORKSPACE_POOL_H_ |
82 | |