This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new aca96b3358 [Refactor][Runtime] Unify tensor pools and VM ownership
(#20670)
aca96b3358 is described below
commit aca96b3358e9b7eb2b4568be8f2a58c1c4bf440e
Author: Tianqi Chen <[email protected]>
AuthorDate: Mon Oct 12 11:59:57 2026 -0400
[Refactor][Runtime] Unify tensor pools and VM ownership (#20670)
---
CMakeLists.txt | 1 -
apps/android_rpc/app/src/main/jni/tvm_runtime.h | 1 -
docs/arch/relax_vm.rst | 6 +-
include/tvm/runtime/device_api.h | 25 ++
include/tvm/runtime/disco/cuda_ipc_memory.h | 6 +-
include/tvm/runtime/memory/memory_manager.h | 199 --------------
include/tvm/runtime/tensor.h | 171 ++++++++++++
include/tvm/runtime/vm/vm.h | 18 +-
python/tvm/rpc/remote.py | 8 +-
python/tvm/runtime/vm.py | 33 +--
src/backend/cuda/runtime/cuda_device_api.cc | 21 ++
src/backend/hexagon/runtime/hexagon_device_api.cc | 18 ++
src/backend/hexagon/runtime/hexagon_device_api.h | 17 +-
src/backend/metal/runtime/metal_common.h | 7 +-
src/backend/metal/runtime/metal_device_api.mm | 8 +
src/backend/opencl/runtime/opencl_common.h | 24 +-
src/backend/opencl/runtime/opencl_device_api.cc | 294 +++++++++++----------
src/backend/rocm/runtime/rocm_common.h | 5 +
src/backend/rocm/runtime/rocm_device_api.cc | 15 +-
src/backend/vulkan/runtime/vulkan_device_api.cc | 12 +
src/backend/vulkan/runtime/vulkan_device_api.h | 11 +
src/relax/backend/vm/codegen_vm.cc | 1 +
src/runtime/cpu_device_api.cc | 22 ++
src/runtime/device_api.cc | 41 ++-
.../extra/contrib/nvshmem/memory_allocator.cc | 82 +++---
src/runtime/extra/disco/builtin.cc | 6 +-
.../extra/disco/cuda_ipc/cuda_ipc_memory.cc | 94 +++----
.../extra/disco/cuda_ipc/custom_allreduce.cc | 1 -
src/runtime/memory/memory_manager.cc | 278 -------------------
src/runtime/memory/naive_allocator.h | 143 ----------
src/runtime/memory/pooled_allocator.h | 176 ------------
src/runtime/tensor.cc | 11 +
src/runtime/vm/builtin.cc | 64 +++--
src/runtime/vm/lm_support.cc | 1 -
src/runtime/vm/paged_kv_cache.cc | 1 -
src/runtime/vm/vm.cc | 57 ++--
src/runtime/workspace_pool.cc | 121 +++------
src/runtime/workspace_pool.h | 19 +-
tests/python/disco/test_custom_allreduce.py | 4 +-
.../relax/test_vm_alloc_storage_with_scope.py | 5 +-
tests/python/relax/test_vm_execbuilder.py | 4 +-
web/emcc/wasm_runtime.cc | 1 -
web/src/runtime.ts | 7 -
43 files changed, 755 insertions(+), 1284 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 9ec883f229..6070108922 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -368,7 +368,6 @@ list(APPEND COMPILER_SRCS ${CODEGEN_SRCS})
tvm_file_glob(GLOB RUNTIME_SRCS
src/runtime/*.cc
src/runtime/vm/*.cc
- src/runtime/memory/*.cc
src/runtime/rpc/minrpc/*.cc
)
# Note: src/runtime/extra/disco/** moves to libtvm_runtime_extra.
diff --git a/apps/android_rpc/app/src/main/jni/tvm_runtime.h
b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
index ea9c7eb278..fa08f82e31 100644
--- a/apps/android_rpc/app/src/main/jni/tvm_runtime.h
+++ b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
@@ -47,7 +47,6 @@
#include "../src/runtime/device_api.cc"
#include "../src/runtime/file_utils.cc"
#include "../src/runtime/logging.cc"
-#include "../src/runtime/memory/memory_manager.cc"
#include "../src/runtime/registry.cc"
#include "../src/runtime/rpc/rpc_channel.cc"
#include "../src/runtime/rpc/rpc_endpoint.cc"
diff --git a/docs/arch/relax_vm.rst b/docs/arch/relax_vm.rst
index 91ca6d107a..26ced6469d 100644
--- a/docs/arch/relax_vm.rst
+++ b/docs/arch/relax_vm.rst
@@ -225,9 +225,9 @@ At runtime, a ``VirtualMachine`` is created and initialized:
Under the hood:
1. **LoadExecutable**: the bytecode and metadata are loaded from the
``VMExecutable``.
-2. **Init**: devices and memory allocators are set up. Each device gets an
``Allocator``
- (either ``NAIVE_ALLOCATOR`` or ``POOLED_ALLOCATOR``, defaulting to pooled).
A CPU device
- is always added for shape computations.
+2. **Init**: devices and owning tensor pools are set up. Each device gets a
+ ``TensorMemoryPool`` from the shared device factory. A CPU device is always
+ added for shape computations.
3. **InitFuncPool**: the function pool is populated — ``kPackedFunc`` entries
are resolved from
imports or the global registry; ``kVMFunc`` entries are wrapped in
``VMClosure`` objects.
diff --git a/include/tvm/runtime/device_api.h b/include/tvm/runtime/device_api.h
index 124dcdd14c..19d420d53a 100644
--- a/include/tvm/runtime/device_api.h
+++ b/include/tvm/runtime/device_api.h
@@ -203,6 +203,31 @@ class TVM_RUNTIME_DLL DeviceAPI {
* \param event_dst The destination stream to synchronize.
*/
virtual void SyncStreamFromTo(Device dev, TVMStreamHandle event_src,
TVMStreamHandle event_dst);
+ /*!
+ * \brief Allocate temporal workspace for backend execution.
+ *
+ * \note We have the following assumption about backend temporal
+ * workspace allocation, and backend will optimize for such assumption:
+ *
+ * - Only a few allocation will happen, and space will be released after
use.
+ * - The release order is usually in reverse order of allocate (stack
style).
+ * - Repeative pattern of same allocations over different runs.
+ * - Workspace should not overlap between different threads(i.e. be
threadlocal)
+ *
+ * \param dev The device of allocation.
+ * \param nbytes The size to be allocated.
+ * \param type_hint The type of elements. Only needed by certain backends
such
+ * as OpenGL, as nbytes is sufficient for most backends.
+ */
+ virtual void* AllocWorkspace(Device dev, size_t nbytes, DLDataType type_hint
= {});
+ /*!
+ * \brief Free temporal workspace in backend execution.
+ *
+ * \param dev The device of allocation.
+ * \param ptr The pointer to be freed.
+ */
+ virtual void FreeWorkspace(Device dev, void* ptr);
+
/*!
* \brief Get device API based on device.
* \param dev The device
diff --git a/include/tvm/runtime/disco/cuda_ipc_memory.h
b/include/tvm/runtime/disco/cuda_ipc_memory.h
index 38fe61a014..346fe09831 100644
--- a/include/tvm/runtime/disco/cuda_ipc_memory.h
+++ b/include/tvm/runtime/disco/cuda_ipc_memory.h
@@ -21,7 +21,7 @@
#define TVM_RUNTIME_DISCO_CUDA_IPC_MEMORY_H_
#include <tvm/runtime/base.h>
-#include <tvm/runtime/memory/memory_manager.h>
+#include <tvm/runtime/tensor.h>
#include <vector>
@@ -79,8 +79,8 @@ class CUDAIPCMemoryObj : public ffi::Object {
*/
class CUDAIPCMemory : public ffi::ObjectRef {
public:
- /*! \brief Get the global singleton CUDAIPCMemory allocator. */
- TVM_RUNTIME_DLL static memory::Allocator* GlobalAllocator();
+ /*! \brief Get the global singleton CUDAIPCMemory pool. */
+ TVM_RUNTIME_DLL static TensorMemoryPool* GlobalMemoryPool();
/*!
* \brief Given a local CUDA data pointer, return the CUDAIPCMemory object
of the pointer.
* \note The pointer's CUDAIPCMemory is expected to have been allocated
diff --git a/include/tvm/runtime/memory/memory_manager.h
b/include/tvm/runtime/memory/memory_manager.h
deleted file mode 100644
index cc499481d2..0000000000
--- a/include/tvm/runtime/memory/memory_manager.h
+++ /dev/null
@@ -1,199 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you 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.
- */
-
-/*!
- * \file tvm/runtime/memory/memory_manager.h
- * \brief Abstract device memory management API
- */
-#ifndef TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
-#define TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
-
-#include <tvm/runtime/base.h>
-#include <tvm/runtime/tensor.h>
-
-#include <functional>
-#include <memory>
-#include <mutex>
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-namespace tvm {
-namespace runtime {
-namespace memory {
-
-enum AllocatorType {
- kNaive = 1,
- kPooled,
-};
-
-struct Buffer {
- /*! \brief The pointer to the allocated block of memory. */
- void* data{nullptr};
- /*! \brief The size of the block. */
- size_t size{0};
- /*! \brief The context of the allocated buffers. */
- Device device;
- /*! \brief The allocator that created this buffer. */
- AllocatorType alloc_type;
-};
-
-class TVM_RUNTIME_DLL Allocator {
- public:
- explicit Allocator(AllocatorType type) : type_(type) {}
- virtual ~Allocator() = default;
- /*! \brief Allocate an empty Tensor using from the allocator.
- * \param shape The shape of the Tensor.
- * \param dtype The datatype of the Tensor.
- * \param dev The device where the array is allocated.
- * \param mem_scope The device memory scope hint.
- * \return The empty Tensor.
- */
- ffi::Tensor Empty(ffi::Shape shape, DLDataType dtype, Device dev,
- ffi::Optional<ffi::String> mem_scope = std::nullopt);
- /*! \brief Return the allocator type. */
- inline AllocatorType type() const { return type_; }
- /*! \brief Allocate a buffer given a size, alignment and type.
- * \param dev The device where the array is allocated.
- * \param nbytes The size of the buffer.
- * \param alignment The alignment of the buffer.
- * \param type_hint A type hint to the allocator.
- * \return A sized allocation in the form of a buffer.
- */
- virtual Buffer Alloc(Device dev, size_t nbytes, size_t alignment, DLDataType
type_hint) = 0;
- /*! \brief Allocate a buffer given a shape and type.
- * \param dev The device where the array is allocated.
- * \param shape The shape of the tensor.
- * \param type_hint A type hint to the allocator.
- * \param mem_scope A memory scope of the buffer.
- * \return A sized allocation in the form of a buffer.
- */
- virtual Buffer Alloc(Device dev, ffi::Shape shape, DLDataType type_hint,
- const std::string& mem_scope = "");
-
- /*! \brief Create a view for the buffer given a shape, type and scope.
- * \param buffer The existing buffer upon which we need to create a view.
- * \param shape The shape of the view.
- * \param type_hint A type hint to the view.
- * \param mem_scope A memory scope of the view.
- * \return A device pointer to the created view.
- */
- virtual void* CreateView(const Buffer& buffer, ffi::Shape shape, DLDataType
type_hint,
- const std::string& mem_scope = "global") {
- return buffer.data;
- }
-
- /*! \brief Release the view .
- * \param dev is the device where this view is created
- * \param data The view pointer to be freed.
- */
- virtual void FreeView(Device dev, void* data) {}
-
- /*! \brief Free a buffer allocated by the allocator.
- * \param buffer The buffer to free.
- */
- virtual void Free(const Buffer& buffer) = 0;
- /*! \brief Clear the allocated memory. */
- virtual void Clear();
- /*! \brief The amount of memory currently allocated.
- * \return The amount of memory currently allocated.
- */
- virtual size_t UsedMemory() const = 0;
-
- protected:
- /*! \brief Check if the given memory scope is allowed to allocate by the
allocator. */
- virtual bool AllowMemoryScope(const std::string& mem_scope) const;
-
- private:
- AllocatorType type_;
-};
-
-class MemoryManager {
- public:
- TVM_RUNTIME_DLL static MemoryManager* Global();
- /*!
- * \brief Get or create an allocator given the context and allocator type.
- * \param dev The TVM device
- * \param type The allocator type
- * \return The memory allocator.
- */
- TVM_RUNTIME_DLL static Allocator* GetOrCreateAllocator(Device dev,
AllocatorType type);
- /*!
- * \brief Get an allocator given the context.
- * \param dev The TVM device
- * \param type The allocator type
- * \return The memory allocator.
- */
- TVM_RUNTIME_DLL static Allocator* GetAllocator(Device dev, AllocatorType
type);
- /*! \brief Clear the allocators. */
- static void Clear();
-
- private:
- MemoryManager() {}
-
- protected:
- std::mutex mu_;
- std::unordered_map<Device, std::unordered_map<AllocatorType,
std::unique_ptr<Allocator>>>
- allocators_;
-};
-
-/*! \brief An object representing a storage allocation. */
-class StorageObj : public ffi::Object {
- public:
- /*! \brief The index into the VM function table. */
- Buffer buffer;
- /*! \brief The allocator where the storage buffer is allocated from. */
- Allocator* allocator = nullptr;
-
- /*! \brief Allocate an Tensor from a given piece of storage. */
- TVM_RUNTIME_DLL ffi::Tensor AllocTensor(int64_t offset, ffi::Shape shape,
DLDataType dtype);
-
- /*! \brief Allocate an Tensor with memory scope from a given piece of
storage. */
- TVM_RUNTIME_DLL ffi::Tensor AllocTensorScoped(int64_t offset, ffi::Shape
shape, DLDataType dtype,
- ffi::String scope = "global");
-
- ~StorageObj() {
- if (allocator) {
- allocator->Free(buffer);
- }
- }
-
- static constexpr const bool _type_mutable = true;
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("vm.Storage", StorageObj, ffi::Object);
-};
-
-/*! \brief reference to storage. */
-class Storage : public ffi::ObjectRef {
- public:
- TVM_RUNTIME_DLL explicit Storage(Buffer buffer, Allocator* allocator);
-
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Storage, ffi::ObjectRef,
StorageObj);
-};
-
-} // namespace memory
-
-using memory::Allocator;
-using memory::AllocatorType;
-using memory::MemoryManager;
-using memory::StorageObj;
-
-} // namespace runtime
-} // namespace tvm
-
-#endif // TVM_RUNTIME_MEMORY_MEMORY_MANAGER_H_
diff --git a/include/tvm/runtime/tensor.h b/include/tvm/runtime/tensor.h
index ac14fec090..c729bef0d7 100644
--- a/include/tvm/runtime/tensor.h
+++ b/include/tvm/runtime/tensor.h
@@ -24,6 +24,7 @@
#ifndef TVM_RUNTIME_TENSOR_H_
#define TVM_RUNTIME_TENSOR_H_
+#include <tvm/ffi/cast.h>
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/dtype.h>
@@ -34,8 +35,12 @@
#include <tvm/support/io.h>
#include <tvm/support/serializer.h>
+#include <algorithm>
#include <atomic>
#include <functional>
+#include <limits>
+#include <mutex>
+#include <unordered_map>
#include <utility>
#include <vector>
@@ -194,6 +199,172 @@ inline Device GetPreferredHostDevice(Device device) {
}
}
+class TensorMemoryPool;
+
+/*! \brief Create the default pool for a device. */
+TVM_RUNTIME_DLL ffi::ObjectPtr<TensorMemoryPool>
CreateTensorMemoryPool(DLDevice device);
+
+/*!
+ * \brief A cache of owning device tensors.
+ *
+ * Construct pools with ffi::make_object. Live tensors retain the pool
directly;
+ * only unpooled backing tensors enter its cache. A backend supplies
independent
+ * backing owners, never a deleter that calls or retains this pool.
+ *
+ * This class protects host ownership, not device execution. An adapter must
keep
+ * separate pools for incompatible contexts/streams, or establish device
ordering
+ * before reuse. Callers must finish using tensors before releasing their
owners.
+ */
+class TensorMemoryPool : public ffi::Object {
+ public:
+ explicit TensorMemoryPool(size_t page_size = 4096) : page_size_(page_size),
used_memory_(0) {
+ TVM_FFI_CHECK_GT(page_size_, 0, ValueError) << "Pool page size must be
positive";
+ TVM_FFI_CHECK_LE(page_size_,
static_cast<size_t>(std::numeric_limits<int64_t>::max()),
+ ValueError);
+ }
+
+ virtual ~TensorMemoryPool() { Clear(); }
+
+ /*! \brief Return an owning compact tensor in ordinary global memory. */
+ virtual ffi::Tensor Empty(ffi::ShapeView shape, DLDataType dtype, DLDevice
device,
+ ffi::Optional<ffi::String> mem_scope =
std::nullopt) {
+ std::lock_guard<std::recursive_mutex> lock(mu_);
+ if (mem_scope.has_value() && !(*mem_scope).empty() && *mem_scope !=
"global") {
+ return DeviceAPI::Get(device)->Empty(shape, dtype, device, mem_scope);
+ }
+ size_t nbytes = ffi::GetDataSize(shape.Product(), dtype);
+ TVM_FFI_CHECK_LE(std::max<size_t>(nbytes, 1),
+ static_cast<size_t>(std::numeric_limits<int64_t>::max())
- (page_size_ - 1),
+ ValueError)
+ << "Rounded tensor allocation size overflows";
+ size_t size = ((std::max<size_t>(nbytes, 1) + page_size_ - 1) /
page_size_) * page_size_;
+ size_t alignment = kAllocAlignment;
+ size_t element_bytes = (static_cast<size_t>(dtype.bits) * dtype.lanes + 7)
/ 8;
+ while (alignment < element_bytes) alignment *= 2;
+
+ ffi::Tensor backing;
+ auto it = memory_pool_.find(size);
+ if (it != memory_pool_.end()) {
+ auto& pool = it->second;
+ bool addressable =
DeviceAPI::Get(device)->SupportsDevicePointerArithmeticsOnHost();
+ auto candidate = std::find_if(pool.begin(), pool.end(), [&](const
ffi::Tensor& tensor) {
+ return tensor.device().device_type == device.device_type &&
+ tensor.device().device_id == device.device_id &&
+ (!addressable || reinterpret_cast<uintptr_t>(tensor.data_ptr())
% alignment == 0);
+ });
+ if (candidate != pool.end()) {
+ backing = std::move(*candidate);
+ pool.erase(candidate);
+ }
+ }
+ if (!backing.defined()) {
+ // Collective backends cannot retry by unilaterally releasing cached
buffers.
+ // Any allocation-pressure retry belongs to the non-collective adapter.
+ backing = AllocBacking(size, alignment, device);
+ used_memory_.fetch_add(size, std::memory_order_relaxed);
+ }
+ Lease lease{ffi::GetObjectPtr<TensorMemoryPool>(this), std::move(backing)};
+ try {
+ return ffi::Tensor::FromNDAlloc(lease, shape, dtype, device);
+ } catch (...) {
+ // Tensor metadata allocation may fail before the owning tensor exists.
+ Recycle(std::move(lease.backing));
+ throw;
+ }
+ }
+
+ /*! \brief Create a bounded owning view of a backing tensor. */
+ virtual ffi::Tensor CreateView(ffi::Tensor backing, ffi::Shape shape,
DLDataType dtype,
+ uint64_t byte_offset = 0,
+ ffi::Optional<ffi::String> mem_scope =
std::nullopt) {
+ TVM_FFI_CHECK(!mem_scope.has_value() || (*mem_scope).empty() || *mem_scope
== "global" ||
+ (backing.device().device_type == kDLHexagon &&
+ (*mem_scope == "global.vtcm" || *mem_scope ==
"global.ddr")),
+ ValueError)
+ << "Unsupported tensor view memory scope";
+ auto view = runtime::CreateView(std::move(backing), std::move(shape),
dtype, byte_offset);
+ if (view.device().device_type == kDLHexagon &&
+ (!mem_scope.has_value() || (*mem_scope).empty() || *mem_scope ==
"global")) {
+ // Hexagon kernels receive the adjusted pointer rather than a byte
offset.
+ auto* tensor = const_cast<DLTensor*>(view.GetDLTensorPtr());
+ if (tensor->byte_offset != 0) {
+ tensor->data = static_cast<char*>(tensor->data) + tensor->byte_offset;
+ }
+ tensor->byte_offset = 0;
+ }
+ return view;
+ }
+
+ /*! \brief Release cached backing; live tensors remain valid and may return
later. */
+ void Clear() {
+ std::lock_guard<std::recursive_mutex> lock(mu_);
+ for (const auto& [size, pool] : memory_pool_) {
+ used_memory_.fetch_sub(size * pool.size(), std::memory_order_relaxed);
+ }
+ memory_pool_.clear();
+ }
+
+ /*! \brief Bytes held by this pool, including checked-out backing. */
+ size_t UsedMemory() const { return
used_memory_.load(std::memory_order_relaxed); }
+
+ static constexpr bool _type_mutable = true;
+ TVM_FFI_DECLARE_OBJECT_INFO("runtime.TensorMemoryPool", TensorMemoryPool,
ffi::Object);
+
+ protected:
+ /*!
+ * \brief Allocate fresh owning uint8[nbytes] backing at the requested
alignment.
+ * \note Opaque handles must satisfy backend alignment requirements. The
backing
+ * owner releases its resources independently of the pool and must not
allocate
+ * through a DeviceAPI::Empty implementation that routes back to this pool.
+ */
+ virtual ffi::Tensor AllocBacking(size_t nbytes, size_t alignment, DLDevice
device) {
+ auto allocate = [&]() {
+ auto* api = DeviceAPI::Get(device);
+ DLDataType bytes{kDLUInt, 8, 1};
+ if (alignment <= kAllocAlignment) {
+ return api->Empty({static_cast<int64_t>(nbytes)}, bytes, device);
+ }
+ // Preserve the backend's dtype alignment contract for wide vector
elements.
+ TVM_FFI_CHECK_LE(alignment, std::numeric_limits<uint16_t>::max(),
ValueError);
+ auto owner = api->Empty({static_cast<int64_t>((nbytes + alignment - 1) /
alignment)},
+ DLDataType{kDLUInt, 8,
static_cast<uint16_t>(alignment)}, device);
+ return runtime::CreateView(std::move(owner),
{static_cast<int64_t>(nbytes)}, bytes);
+ };
+ try {
+ return allocate();
+ } catch (const ffi::Error&) {
+ // This default allocator is non-collective; release cached memory and
retry.
+ Clear();
+ return allocate();
+ }
+ }
+
+ private:
+ struct Lease {
+ ffi::ObjectPtr<TensorMemoryPool> pool;
+ ffi::Tensor backing;
+
+ void AllocData(DLTensor* tensor) { tensor->data = backing.data_ptr(); }
+ void FreeData(DLTensor*) { pool->Recycle(std::move(backing)); }
+ };
+
+ void Recycle(ffi::Tensor backing) {
+ std::lock_guard<std::recursive_mutex> lock(mu_);
+ size_t size = ffi::GetDataSize(backing);
+ try {
+ memory_pool_[size].push_back(std::move(backing));
+ } catch (const std::bad_alloc&) {
+ // A cache insertion failure must not escape tensor destruction.
+ used_memory_.fetch_sub(size, std::memory_order_relaxed);
+ }
+ }
+
+ size_t page_size_;
+ std::atomic<size_t> used_memory_;
+ std::unordered_map<size_t, std::vector<ffi::Tensor>> memory_pool_;
+ std::recursive_mutex mu_;
+};
+
} // namespace runtime
} // namespace tvm
diff --git a/include/tvm/runtime/vm/vm.h b/include/tvm/runtime/vm/vm.h
index 3ebab66fab..bc240043e7 100644
--- a/include/tvm/runtime/vm/vm.h
+++ b/include/tvm/runtime/vm/vm.h
@@ -24,25 +24,19 @@
#define TVM_RUNTIME_VM_VM_H_
#include <tvm/ffi/extra/module.h>
+#include <tvm/runtime/device_api.h>
+#include <tvm/runtime/tensor.h>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
-#include "../memory/memory_manager.h"
#include "./bytecode.h"
#include "./executable.h"
namespace tvm {
namespace runtime {
-
-using memory::Allocator;
-using memory::AllocatorType;
-using memory::MemoryManager;
-using memory::Storage;
-using memory::StorageObj;
-
namespace vm {
/*!
@@ -128,10 +122,8 @@ class VirtualMachine : public ffi::ModuleObj {
/*!
* \brief Initialize the virtual machine for a set of devices.
* \param devices The set of TVM devices.
- * \param alloc_types The allocator types for each device.
*/
- virtual void Init(const std::vector<Device>& devices,
- const std::vector<AllocatorType>& alloc_types) = 0;
+ virtual void Init(const std::vector<Device>& devices) = 0;
/*!
* \brief Load the executable for the virtual machine.
* \param exec The executable.
@@ -210,8 +202,8 @@ class VirtualMachine : public ffi::ModuleObj {
//--------------------------------------------------------------------------
// The following section contains states that other builtin can depend on
//--------------------------------------------------------------------------
- /*! \brief The memory allocators. */
- std::vector<Allocator*> allocators;
+ /*! \brief Owning memory pools for the physical devices. */
+ std::vector<ffi::ObjectPtr<TensorMemoryPool>> memory_pools;
/*! \brief Runtime physical device list. */
std::vector<Device> devices;
/*! \brief The VM extensions. Mapping from the type index of the extension
to the extension
diff --git a/python/tvm/rpc/remote.py b/python/tvm/rpc/remote.py
index 03a3a03433..f4994a519b 100644
--- a/python/tvm/rpc/remote.py
+++ b/python/tvm/rpc/remote.py
@@ -250,14 +250,10 @@ def time_evaluator(
)
-def virtual_machine(module, device, memory_cfg=None):
+def virtual_machine(module, device):
"""Initialize a remote VM with session-checked endpoint device
descriptors."""
from tvm.runtime.vm import VirtualMachine
devices = list(device) if isinstance(device, list | tuple) else [device]
local = [_ffi_api.ValidateDevice(module, dev) for dev in devices]
- if isinstance(memory_cfg, dict):
- memory_cfg = {
- _ffi_api.ValidateDevice(module, dev): value for dev, value in
memory_cfg.items()
- }
- return VirtualMachine(module, local, memory_cfg)
+ return VirtualMachine(module, local)
diff --git a/python/tvm/runtime/vm.py b/python/tvm/runtime/vm.py
index 0e7eaee280..3dcf8f4f40 100644
--- a/python/tvm/runtime/vm.py
+++ b/python/tvm/runtime/vm.py
@@ -39,14 +39,10 @@ class VMInstrumentReturnKind(IntEnum):
class VirtualMachine:
"""Relax VM runtime."""
- NAIVE_ALLOCATOR = 1
- POOLED_ALLOCATOR = 2
-
def __init__(
self,
rt_mod: tvm.runtime.Module | tvm.runtime.Executable,
device: Device | list[Device],
- memory_cfg: str | dict[Device, str] | None = None,
) -> None:
"""
Construct a VirtualMachine wrapper object.
@@ -58,14 +54,6 @@ class VirtualMachine:
device : Union[Device, List[Device]]
The device to deploy the module.
-
- memory_cfg : Optional[Union[str, Dict[Device, str]]]
- Config the type of memory allocator. The allocator type can be
["naive",
- "pooled"]. If memory_cfg is None, all devices will use pooled
allocator
- by default. If memory_cfg is string, all devices will use the
specified
- allocator type. If memory_cfg is a dict, each device uses the
allocator
- type specified in the dict, or pooled allocator if not specified
in the
- dict.
"""
if not isinstance(rt_mod, tvm.runtime.Module):
if isinstance(rt_mod, tvm.runtime.Executable):
@@ -83,10 +71,10 @@ class VirtualMachine:
self._get_function_arity = self.module["get_function_arity"]
self._get_function_param_name = self.module["get_function_param_name"]
self._set_instrument = self.module["set_instrument"]
- self._setup_device(device, memory_cfg)
+ self._setup_device(device)
- def _setup_device(self, dev: Device, memory_cfg: str | dict[Device, str])
-> None:
- """init devices and allocators."""
+ def _setup_device(self, dev: Device | list[Device]) -> None:
+ """Initialize devices and owning tensor pools."""
devs = dev
if not isinstance(dev, list | tuple):
if not isinstance(dev, tvm.runtime.Device):
@@ -97,25 +85,10 @@ class VirtualMachine:
if devs[-1].dlpack_device_type() != tvm.cpu().dlpack_device_type():
devs.append(tvm.cpu())
- default_alloc_type = VirtualMachine.POOLED_ALLOCATOR
- if memory_cfg is None:
- memory_cfg = {}
- elif isinstance(memory_cfg, str):
- assert memory_cfg in ["naive", "pooled"]
- if memory_cfg == "naive":
- default_alloc_type = VirtualMachine.NAIVE_ALLOCATOR
- memory_cfg = {}
- elif not isinstance(memory_cfg, dict):
- raise TypeError(
- "memory_cfg is expected be string or dictionary, "
- + f"but received {type(memory_cfg)}"
- )
init_args = []
for device in devs:
init_args.append(device.dlpack_device_type())
init_args.append(device.index)
- alloc_type = memory_cfg[device] if device in memory_cfg else
default_alloc_type
- init_args.append(alloc_type)
self.module["vm_initialization"](*init_args)
def __getitem__(self, key: str) -> Function:
diff --git a/src/backend/cuda/runtime/cuda_device_api.cc
b/src/backend/cuda/runtime/cuda_device_api.cc
index 03f627d981..bec8e0c680 100644
--- a/src/backend/cuda/runtime/cuda_device_api.cc
+++ b/src/backend/cuda/runtime/cuda_device_api.cc
@@ -34,6 +34,8 @@
#include <algorithm>
#include <cstring>
+#include "../../../runtime/workspace_pool.h"
+
namespace tvm {
namespace runtime {
@@ -291,6 +293,14 @@ class CUDADeviceAPI final : public DeviceAPI {
TVM_FFI_CHECK_CUDA_ERROR(cudaStreamSynchronize(static_cast<cudaStream_t>(stream)));
}
+ void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final {
+ return ThreadLocalWorkspacePool(dev.device_type)->AllocWorkspace(dev,
size);
+ }
+
+ void FreeWorkspace(Device dev, void* data) final {
+ ThreadLocalWorkspacePool(dev.device_type)->FreeWorkspace(dev, data);
+ }
+
bool SupportsDevicePointerArithmeticsOnHost() final { return true; }
static CUDADeviceAPI* Global() {
@@ -301,12 +311,23 @@ class CUDADeviceAPI final : public DeviceAPI {
}
private:
+ static WorkspacePool* ThreadLocalWorkspacePool(DLDeviceType device_type);
+
static void GPUCopy(const void* from, void* to, size_t size, cudaMemcpyKind
kind,
cudaStream_t stream) {
TVM_FFI_CHECK_CUDA_ERROR(cudaMemcpyAsync(to, from, size, kind, stream));
}
};
+WorkspacePool* CUDADeviceAPI::ThreadLocalWorkspacePool(DLDeviceType
device_type) {
+ if (device_type == kDLCUDAHost) {
+ static thread_local WorkspacePool pool(kDLCUDAHost,
CUDADeviceAPI::Global());
+ return &pool;
+ }
+ static thread_local WorkspacePool pool(kDLCUDA, CUDADeviceAPI::Global());
+ return &pool;
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
diff --git a/src/backend/hexagon/runtime/hexagon_device_api.cc
b/src/backend/hexagon/runtime/hexagon_device_api.cc
index 763cb82c1d..ed6491ba18 100644
--- a/src/backend/hexagon/runtime/hexagon_device_api.cc
+++ b/src/backend/hexagon/runtime/hexagon_device_api.cc
@@ -166,6 +166,24 @@ void HexagonDeviceAPI::FreeDataSpace(Device dev, void*
ptr) {
}
}
+void* HexagonDeviceAPI::AllocWorkspace(Device dev, size_t size, DLDataType
type_hint) {
+ TVM_FFI_ICHECK(IsValidDevice(dev)) << "dev.device_type: " << dev.device_type;
+ TVM_FFI_ICHECK(runtime_hexbuffs) << "Call AcquireResources before allocating
Hexagon workspace";
+ std::lock_guard<std::mutex> lock(workspace_mutex_);
+ auto& pool = workspace_pools_[std::this_thread::get_id()];
+ if (!pool) pool = std::make_unique<WorkspacePool>(kDLHexagon, this);
+ return pool->AllocWorkspace(dev, size);
+}
+
+void HexagonDeviceAPI::FreeWorkspace(Device dev, void* data) {
+ TVM_FFI_ICHECK(IsValidDevice(dev)) << "dev.device_type: " << dev.device_type;
+ std::lock_guard<std::mutex> lock(workspace_mutex_);
+ auto it = workspace_pools_.find(std::this_thread::get_id());
+ TVM_FFI_ICHECK(it != workspace_pools_.end())
+ << "Attempt made to free unknown workspace allocation";
+ it->second->FreeWorkspace(dev, data);
+}
+
void HexagonDeviceAPI::CopyDataFromTo(DLTensor* from, DLTensor* to,
TVMStreamHandle stream) {
TVM_FFI_ICHECK_EQ(from->byte_offset, 0);
TVM_FFI_ICHECK_EQ(to->byte_offset, 0);
diff --git a/src/backend/hexagon/runtime/hexagon_device_api.h
b/src/backend/hexagon/runtime/hexagon_device_api.h
index 016a00b4ba..319517c69b 100644
--- a/src/backend/hexagon/runtime/hexagon_device_api.h
+++ b/src/backend/hexagon/runtime/hexagon_device_api.h
@@ -27,6 +27,7 @@
#include <memory>
#include <mutex>
#include <string>
+#include <thread>
#include <unordered_map>
#include <utility>
#include <vector>
@@ -84,7 +85,10 @@ class HexagonDeviceAPI final : public DeviceAPI {
TVM_FFI_ICHECK(runtime_threads) << "runtime_threads was not created in
AcquireResources";
runtime_threads.reset();
- ResetWorkspacePool(kDLHexagon);
+ {
+ std::lock_guard<std::mutex> lock(workspace_mutex_);
+ for (auto& [thread, pool] : workspace_pools_) pool->Clear();
+ }
TVM_FFI_ICHECK(runtime_hexbuffs) << "runtime_hexbuffs was not created in
AcquireResources";
runtime_hexbuffs.reset();
@@ -114,6 +118,12 @@ class HexagonDeviceAPI final : public DeviceAPI {
//! \brief Allocate a flat allocation of global memory wrapped in a
HexagonBuffer.
void* AllocDataSpace(Device dev, size_t nbytes, size_t alignment, DLDataType
type_hint);
+ //! \brief Request a dynamically allocated buffer from the thread's
workspace pool.
+ void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final;
+
+ //! \brief Return a workspace allocation to its owning pool.
+ void FreeWorkspace(Device dev, void* data) final;
+
//! \brief Free the allocated HexagonBuffer.
void FreeDataSpace(Device dev, void* ptr);
@@ -233,6 +243,11 @@ class HexagonDeviceAPI final : public DeviceAPI {
// so a late deleter never frees an allocation through a replacement session.
std::shared_ptr<HexagonBufferManager> runtime_hexbuffs;
+ // Backend ownership permits session teardown to clear every thread's cached
workspace.
+ // Outstanding owners retain their original buffer manager until
FreeWorkspace.
+ std::unordered_map<std::thread::id, std::unique_ptr<WorkspacePool>>
workspace_pools_;
+ std::mutex workspace_mutex_;
+
//! \brief Thread manager
std::unique_ptr<HexagonThreadManager> runtime_threads;
const unsigned threads{6};
diff --git a/src/backend/metal/runtime/metal_common.h
b/src/backend/metal/runtime/metal_common.h
index 27f639e57f..e308264f46 100644
--- a/src/backend/metal/runtime/metal_common.h
+++ b/src/backend/metal/runtime/metal_common.h
@@ -41,6 +41,8 @@
#include <utility>
#include <vector>
+#include "../../../runtime/workspace_pool.h"
+
/* Macro for convenience in using AutoReleasePoolWrapper.
* With this macro we can add AutoReleasePoolWrapper to our ObjC code in more
* native way.
@@ -317,6 +319,8 @@ class MetalWorkspace final : public DeviceAPI {
TVMStreamHandle CreateStream(Device dev) final;
void FreeStream(Device dev, TVMStreamHandle stream) final;
void StreamSync(Device dev, TVMStreamHandle stream) final;
+ void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final;
+ void FreeWorkspace(Device dev, void* data) final;
void ReinitializeDefaultStreams();
/**
@@ -413,7 +417,8 @@ class MetalThreadEntry {
};
std::vector<StagingBufferPool> staging_pools_; // per device
// constructor
- MetalThreadEntry() {
+ WorkspacePool pool;
+ MetalThreadEntry() : pool(kDLMetal, MetalWorkspace::Global()) {
device.device_id = 0;
device.device_type = static_cast<DLDeviceType>(kDLMetal);
MetalWorkspace* global_ws = MetalWorkspace::Global();
diff --git a/src/backend/metal/runtime/metal_device_api.mm
b/src/backend/metal/runtime/metal_device_api.mm
index aabee12bcb..357a568667 100644
--- a/src/backend/metal/runtime/metal_device_api.mm
+++ b/src/backend/metal/runtime/metal_device_api.mm
@@ -319,6 +319,14 @@ void MetalWorkspace::StreamSync(Device dev,
TVMStreamHandle stream) {
};
}
+void* MetalWorkspace::AllocWorkspace(Device dev, size_t size, DLDataType
type_hint) {
+ return MetalThreadEntry::ThreadLocal()->pool.AllocWorkspace(dev, size);
+}
+
+void MetalWorkspace::FreeWorkspace(Device dev, void* data) {
+ MetalThreadEntry::ThreadLocal()->pool.FreeWorkspace(dev, data);
+}
+
MetalThreadEntry::~MetalThreadEntry() {
for (auto x : temp_buffer_) {
if (x != nil) {
diff --git a/src/backend/opencl/runtime/opencl_common.h
b/src/backend/opencl/runtime/opencl_common.h
index 1cd54b8c46..e197ded776 100644
--- a/src/backend/opencl/runtime/opencl_common.h
+++ b/src/backend/opencl/runtime/opencl_common.h
@@ -28,7 +28,6 @@
#include <tvm/ffi/function.h>
#include <tvm/runtime/base.h>
#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/memory/memory_manager.h>
#include <tvm/runtime/tensor.h>
#include <tvm/runtime/timer.h>
@@ -74,14 +73,13 @@
#include "../../../runtime/metadata.h"
#include "../../../runtime/pack_args.h"
#include "../../../runtime/thread_storage_scope.h"
+#include "../../../runtime/workspace_pool.h"
#include "texture.h"
namespace tvm {
namespace runtime {
namespace cl {
-using tvm::runtime::memory::Buffer;
-
static_assert(sizeof(cl_mem) == sizeof(void*), "Required to store cl_mem
inside void*");
inline const char* CLGetErrorString(cl_int error) {
@@ -266,6 +264,8 @@ class OpenCLWorkspace : public DeviceAPI {
// the mutex for initialization
std::mutex mu;
+ OpenCLWorkspace();
+
// destructor
~OpenCLWorkspace() {
for (auto& it : contexts) {
@@ -356,13 +356,19 @@ class OpenCLWorkspace : public DeviceAPI {
ffi::Tensor Empty(ffi::ShapeView shape, DLDataType dtype, DLDevice device,
ffi::Optional<ffi::String> mem_scope = std::nullopt) final;
void* AllocDataSpace(Device dev, int ndim, const int64_t* shape, DLDataType
dtype,
- ffi::Optional<ffi::String> mem_scope = std::nullopt);
+ ffi::Optional<ffi::String> mem_scope = std::nullopt,
+ TensorMemoryPool* pool = nullptr);
void* AllocDataSpace(Device dev, size_t width, size_t height, size_t depth,
DLDataType type_hint,
- ffi::Optional<ffi::String> mem_scope = std::nullopt);
+ ffi::Optional<ffi::String> mem_scope = std::nullopt,
+ TensorMemoryPool* pool = nullptr);
void* GetNativePtr(const tvm::ffi::Tensor& narr);
void SetNativePtr(const tvm::ffi::Tensor& narr, void* host_ptr, size_t
buf_size);
void SetPerfHint(Device dev, cl_uint perf_hint);
void FreeDataSpace(Device dev, void* ptr);
+ void ClearMemoryPool();
+ void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final;
+ void FreeWorkspace(Device dev, void* data) final;
+ static ffi::Tensor AllocWorkspaceTensor(DeviceAPI* api, Device dev, size_t
size);
void StreamSync(Device dev, TVMStreamHandle stream) final;
// cl_mem alloc utils
void* AllocCLBuffer(Device dev, size_t size, size_t alignment, DLDataType
type_hint);
@@ -387,6 +393,7 @@ class OpenCLWorkspace : public DeviceAPI {
return "";
}
std::string noDevicesErrorMsg = "";
+ ffi::ObjectPtr<TensorMemoryPool> memory_pool_;
};
/*! \brief Thread local workspace */
@@ -401,10 +408,13 @@ class OpenCLThreadEntry {
};
/*! \brief The current device */
Device device;
+ /*! \brief Thread-local linear workspace allocation cache. */
+ WorkspacePool pool;
/*! \brief The thread-local kernel table */
std::vector<KTEntry> kernel_table;
// constructor
- OpenCLThreadEntry(DLDeviceType device_type, DeviceAPI* device_api) {
+ OpenCLThreadEntry(DLDeviceType device_type, DeviceAPI* device_api)
+ : pool(device_type, device_api, OpenCLWorkspace::AllocWorkspaceTensor) {
device.device_id = 0;
device.device_type = device_type;
}
@@ -449,7 +459,7 @@ struct BufferDescriptor {
cl::BufferDescriptor* back_buffer{nullptr};
cl_uchar* host_ptr{nullptr};
MemoryLayout layout{MemoryLayout::kBuffer1D};
- Buffer mbuf{nullptr}; // MemoryManager ref.
+ ffi::Tensor backing; // Owns the pooled buffer underlying an image.
bool is_compat_view{false};
};
} // namespace cl
diff --git a/src/backend/opencl/runtime/opencl_device_api.cc
b/src/backend/opencl/runtime/opencl_device_api.cc
index 510b8a0f82..04e4301bb0 100644
--- a/src/backend/opencl/runtime/opencl_device_api.cc
+++ b/src/backend/opencl/runtime/opencl_device_api.cc
@@ -27,7 +27,6 @@
#include <sstream>
-#include "../../../runtime/memory/pooled_allocator.h"
#include "../../../support/env.h"
#include "opencl_common.h"
@@ -244,29 +243,118 @@ void*
OpenCLWorkspace::CreateHostPtrIfEnabled(cl::BufferDescriptor* desc, Device
return desc;
}
-ffi::Tensor OpenCLWorkspace::Empty(ffi::ShapeView shape, DLDataType dtype,
DLDevice device,
- ffi::Optional<ffi::String> mem_scope) {
- struct NDAlloc {
- OpenCLWorkspace* api;
- ffi::Optional<ffi::String> scope;
- void AllocData(DLTensor* tensor) {
- tensor->data =
- api->AllocDataSpace(tensor->device, tensor->ndim, tensor->shape,
tensor->dtype, scope);
+// Reuse follows the workspace's existing in-order command queue per device.
+class OpenCLTensorMemoryPool final : public TensorMemoryPool {
+ public:
+ ffi::Tensor Empty(ffi::ShapeView shape, DLDataType dtype, DLDevice device,
+ ffi::Optional<ffi::String> mem_scope = std::nullopt) final
{
+ if (!mem_scope.has_value() || (*mem_scope).empty() || *mem_scope ==
"global") {
+ return TensorMemoryPool::Empty(shape, dtype, device);
}
- void FreeData(DLTensor* tensor) {
- auto* desc = static_cast<BufferDescriptor*>(tensor->data);
- if (desc->layout == BufferDescriptor::MemoryLayout::kBuffer1D) {
- MemoryManager::GetAllocator(tensor->device,
desc->mbuf.alloc_type)->Free(desc->mbuf);
- } else {
- api->FreeDataSpace(tensor->device, desc);
+ struct NDAlloc {
+ OpenCLTensorMemoryPool* pool;
+ ffi::Optional<ffi::String> scope;
+ void AllocData(DLTensor* tensor) {
+ tensor->data = OpenCLWorkspace::Global()->AllocDataSpace(
+ tensor->device, tensor->ndim, tensor->shape, tensor->dtype, scope,
pool);
+ }
+ void FreeData(DLTensor* tensor) {
+ OpenCLWorkspace::Global()->FreeDataSpace(tensor->device, tensor->data);
}
+ };
+ return ffi::Tensor::FromNDAlloc(NDAlloc{this, mem_scope}, shape, dtype,
device);
+ }
+
+ ffi::Tensor CreateView(ffi::Tensor backing, ffi::Shape shape, DLDataType
dtype,
+ uint64_t byte_offset = 0,
+ ffi::Optional<ffi::String> mem_scope = std::nullopt)
final {
+ TVM_FFI_CHECK_EQ(backing.device().device_type, kDLOpenCL, ValueError);
+ TVM_FFI_CHECK(backing.IsContiguous(), ValueError) << "OpenCL views require
compact backing";
+ auto* desc = static_cast<BufferDescriptor*>(backing.data_ptr());
+ if ((!mem_scope.has_value() || (*mem_scope).empty() || *mem_scope ==
"global") &&
+ desc->layout == BufferDescriptor::MemoryLayout::kBuffer1D) {
+ return runtime::CreateView(std::move(backing), std::move(shape), dtype,
byte_offset);
}
- };
- return ffi::Tensor::FromNDAlloc(NDAlloc{this, mem_scope}, shape, dtype,
device);
+ TVM_FFI_CHECK(byte_offset == 0 && backing.GetDLTensorPtr()->byte_offset ==
0, ValueError)
+ << "OpenCL physical views require a zero byte offset";
+ TVM_FFI_CHECK_LE(ffi::GetDataSize(shape.Product(), dtype),
ffi::GetDataSize(backing),
+ ValueError)
+ << "OpenCL view exceeds its backing tensor";
+ auto* api = OpenCLWorkspace::Global();
+ auto device = backing.device();
+ if (mem_scope.has_value() && !(*mem_scope).empty() && *mem_scope !=
"global") {
+ TVM_FFI_CHECK_GE(shape.size(), 2, ValueError);
+ TVM_FFI_CHECK_EQ(shape.back(), 4, ValueError) << "OpenCL textures
require four channels";
+ // Physical image rows include backend padding, which must also fit the
buffer.
+ if (api->IsBufferToImageSupported(device.device_id)) {
+ auto scope = BufferDescriptor::ScopeFromMemoryLayout(desc->layout);
+ size_t capacity =
+ desc->layout == BufferDescriptor::MemoryLayout::kBuffer1D
+ ? ffi::GetDataSize(backing)
+ : GetTextureMemorySize(backing.shape(), backing.dtype().bits,
backing.dtype().lanes,
+ scope,
api->GetImageAlignment(device.device_id));
+ size_t needed = GetTextureMemorySize(shape, dtype.bits, dtype.lanes,
*mem_scope,
+
api->GetImageAlignment(device.device_id));
+ TVM_FFI_CHECK_LE(needed, capacity, ValueError)
+ << "OpenCL image view exceeds its backing buffer";
+ }
+ }
+ struct NDAlloc {
+ ffi::Tensor backing;
+ ffi::Optional<ffi::String> scope;
+ void AllocData(DLTensor* tensor) {
+ tensor->data = OpenCLWorkspace::Global()->AllocDataSpaceView(
+ tensor->device, backing.data_ptr(),
+ ffi::Shape(tensor->shape, tensor->shape + tensor->ndim),
tensor->dtype, scope);
+ }
+ void FreeData(DLTensor* tensor) {
+ if (tensor->data != backing.data_ptr()) {
+ OpenCLWorkspace::Global()->FreeDataSpaceView(tensor->device,
tensor->data);
+ }
+ }
+ };
+ return ffi::Tensor::FromNDAlloc(NDAlloc{std::move(backing), mem_scope},
shape, dtype, device);
+ }
+
+ protected:
+ ffi::Tensor AllocBacking(size_t nbytes, size_t alignment, DLDevice device)
final {
+ struct NDAlloc {
+ size_t alignment;
+ void AllocData(DLTensor* tensor) {
+ tensor->data = OpenCLWorkspace::Global()->AllocCLBuffer(
+ tensor->device, ffi::GetDataSize(*tensor), alignment,
tensor->dtype);
+ }
+ void FreeData(DLTensor* tensor) {
+ OpenCLWorkspace::Global()->FreeDataSpace(tensor->device, tensor->data);
+ }
+ };
+ auto allocate = [&]() {
+ return ffi::Tensor::FromNDAlloc(NDAlloc{alignment},
{static_cast<int64_t>(nbytes)},
+ DLDataType{kDLUInt, 8, 1}, device);
+ };
+ try {
+ return allocate();
+ } catch (const ffi::Error& err) {
+ LOG(WARNING) << "OpenCL pool got an error during allocation: " <<
err.what();
+ LOG(WARNING) << "Trying to release all unused memory and reallocate...";
+ Clear();
+ return allocate();
+ }
+ }
+};
+
+OpenCLWorkspace::OpenCLWorkspace() :
memory_pool_(ffi::make_object<OpenCLTensorMemoryPool>()) {}
+
+void OpenCLWorkspace::ClearMemoryPool() { memory_pool_->Clear(); }
+
+ffi::Tensor OpenCLWorkspace::Empty(ffi::ShapeView shape, DLDataType dtype,
DLDevice device,
+ ffi::Optional<ffi::String> mem_scope) {
+ return memory_pool_->Empty(shape, dtype, device, mem_scope);
}
void* OpenCLWorkspace::AllocDataSpace(Device dev, size_t width, size_t height,
size_t depth,
- DLDataType type_hint,
ffi::Optional<ffi::String> mem_scope) {
+ DLDataType type_hint,
ffi::Optional<ffi::String> mem_scope,
+ TensorMemoryPool* pool) {
// Texture allocation given width and height
cl_uint row_align = GetImageAlignment(dev.device_id);
size_t pixel_size = (type_hint.bits * type_hint.lanes + 7) / 8;
@@ -274,43 +362,35 @@ void* OpenCLWorkspace::AllocDataSpace(Device dev, size_t
width, size_t height, s
size_t mem_size = row_pitch * height * depth;
// Alloc back buffer from pool
+ ffi::Tensor backing;
cl::BufferDescriptor* back_buffer = nullptr;
if (IsBufferToImageSupported(dev.device_id)) {
- auto buf = MemoryManager::GetOrCreateAllocator(dev, AllocatorType::kPooled)
- ->Alloc(dev, mem_size, kTempAllocaAlignment, type_hint);
- back_buffer = static_cast<cl::BufferDescriptor*>(buf.data);
- back_buffer->mbuf = buf;
+ backing = (pool ? pool : memory_pool_.get())
+ ->Empty({static_cast<int64_t>(mem_size)},
DLDataType{kDLUInt, 8, 1}, dev);
+ back_buffer = static_cast<cl::BufferDescriptor*>(backing.data_ptr());
}
if (!mem_scope.has_value()) {
mem_scope = ffi::String("global.texture");
}
- try {
- return AllocCLImage(dev, back_buffer, width, height, depth, row_pitch,
type_hint, mem_scope);
- } catch (...) {
- if (back_buffer) {
- MemoryManager::GetAllocator(dev,
back_buffer->mbuf.alloc_type)->Free(back_buffer->mbuf);
- }
- throw;
- }
+ auto* image = static_cast<cl::BufferDescriptor*>(
+ AllocCLImage(dev, back_buffer, width, height, depth, row_pitch,
type_hint, mem_scope));
+ image->backing = std::move(backing);
+ return image;
}
void* OpenCLWorkspace::AllocDataSpace(Device dev, int ndim, const int64_t*
shape, DLDataType dtype,
- ffi::Optional<ffi::String> mem_scope) {
+ ffi::Optional<ffi::String> mem_scope,
+ TensorMemoryPool* pool) {
this->Init();
if (!mem_scope.has_value() || (*mem_scope).empty() || (*mem_scope) ==
"global") {
size_t size = GetMemObjectSize(dev, ndim, shape, dtype);
- cl::BufferDescriptor* ret_buffer = nullptr;
- auto buf = MemoryManager::GetOrCreateAllocator(dev, AllocatorType::kPooled)
- ->Alloc(dev, size, kTempAllocaAlignment, dtype);
- ret_buffer = static_cast<cl::BufferDescriptor*>(buf.data);
- ret_buffer->mbuf = buf;
- return ret_buffer;
+ return AllocCLBuffer(dev, size, kTempAllocaAlignment, dtype);
}
size_t axis = DefaultTextureLayoutSeparator(ndim, mem_scope.value());
auto texture = ApplyTexture2DFlattening<int64_t>(shape, ndim, axis);
- return AllocDataSpace(dev, texture.width, texture.height, texture.depth,
dtype, mem_scope);
+ return AllocDataSpace(dev, texture.width, texture.height, texture.depth,
dtype, mem_scope, pool);
}
void* OpenCLWorkspace::AllocCLBuffer(Device dev, size_t size, size_t alignment,
@@ -485,8 +565,6 @@ void OpenCLWorkspace::FreeDataSpace(Device dev, void* ptr) {
if (desc->back_buffer) {
// 2D Image w/ back buffer allocated from pool
OPENCL_CALL(clReleaseMemObject(desc->buffer));
- MemoryManager::GetAllocator(dev, desc->back_buffer->mbuf.alloc_type)
- ->Free(desc->back_buffer->mbuf);
delete desc;
} else {
if (desc->layout == cl::BufferDescriptor::MemoryLayout::kBuffer1D) {
@@ -586,6 +664,27 @@ void OpenCLWorkspace::CopyDataFromTo(DLTensor* from,
DLTensor* to, TVMStreamHand
}
}
+ffi::Tensor OpenCLWorkspace::AllocWorkspaceTensor(DeviceAPI* api, Device dev,
size_t size) {
+ struct NDAlloc {
+ OpenCLWorkspace* api;
+ void AllocData(DLTensor* tensor) {
+ tensor->data = api->AllocCLBuffer(tensor->device,
ffi::GetDataSize(*tensor),
+ kTempAllocaAlignment, tensor->dtype);
+ }
+ void FreeData(DLTensor* tensor) { api->FreeDataSpace(tensor->device,
tensor->data); }
+ };
+ return ffi::Tensor::FromNDAlloc(NDAlloc{static_cast<OpenCLWorkspace*>(api)},
+ {static_cast<int64_t>(size)},
DLDataType{kDLUInt, 8, 1}, dev);
+}
+
+void* OpenCLWorkspace::AllocWorkspace(Device dev, size_t size, DLDataType) {
+ return GetThreadEntry()->pool.AllocWorkspace(dev, size);
+}
+
+void OpenCLWorkspace::FreeWorkspace(Device dev, void* data) {
+ GetThreadEntry()->pool.FreeWorkspace(dev, data);
+}
+
void OpenCLWorkspace::StreamSync(Device dev, TVMStreamHandle stream) {
this->Init();
TVM_FFI_ICHECK(stream == nullptr);
@@ -806,107 +905,24 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-class OpenCLPooledAllocator final : public memory::PooledAllocator {
- public:
- explicit OpenCLPooledAllocator() : PooledAllocator() {}
- ~OpenCLPooledAllocator() override { ReleaseAll(); }
-
- bool AllowMemoryScope(const std::string& mem_scope) const final {
- return ((mem_scope.find("texture") != std::string::npos) ||
mem_scope.empty() ||
- ("global" == mem_scope));
- }
-
- Buffer Alloc(Device dev, size_t nbytes, size_t alignment, DLDataType
type_hint) override {
- std::lock_guard<std::recursive_mutex> lock(mu_);
- size_t size = ((nbytes + page_size_ - 1) / page_size_) * page_size_;
- auto&& it = memory_pool_.find(size);
- if (it != memory_pool_.end() && !it->second.empty()) {
- auto&& pool = it->second;
- auto ret = pool.back();
- pool.pop_back();
- return ret;
- }
- Buffer buf;
- buf.device = dev;
- buf.size = size;
- buf.alloc_type = AllocatorType::kPooled;
- try {
- buf.data = DeviceAllocDataSpace(dev, size, alignment, type_hint);
- } catch (tvm::ffi::Error& err) {
- LOG(WARNING) << "PooledAllocator got InternalError during allocation: "
<< err.what();
- LOG(WARNING) << "Trying to release all unused memory and reallocate...";
- ReleaseAll();
- buf.data = DeviceAllocDataSpace(dev, size, alignment, type_hint);
- }
-
- used_memory_.fetch_add(size, std::memory_order_relaxed);
- VLOG(1) << "allocate " << size << " B, used memory " << used_memory_ << "
B";
- return buf;
- }
-
- Buffer Alloc(Device dev, ffi::Shape shape, DLDataType type_hint,
- const std::string& mem_scope) override {
- if (AllowMemoryScope(mem_scope)) {
- size_t size = GetMemObjectSize(dev, shape.size(), shape.data(),
type_hint);
- if (mem_scope.empty() || mem_scope == "global") {
- return Alloc(dev, size, kTempAllocaAlignment, type_hint);
- }
- Buffer buf;
- buf.device = dev;
- buf.size = size;
- buf.alloc_type = AllocatorType::kPooled;
- buf.data = OpenCLWorkspace::Global()->AllocDataSpace(dev, shape.size(),
shape.data(),
- type_hint,
ffi::String(mem_scope));
- return buf;
- }
- TVM_FFI_THROW(InternalError) << "Unsupported memory scope for this
Allocator:" << mem_scope;
- return {};
- }
-
- void Free(const Buffer& buffer) override {
- std::lock_guard<std::recursive_mutex> lock(mu_);
- auto* desc = static_cast<BufferDescriptor*>(buffer.data);
- if (desc->layout != BufferDescriptor::MemoryLayout::kBuffer1D) {
- OpenCLWorkspace::Global()->FreeDataSpace(buffer.device, buffer.data);
- return;
- }
- if (memory_pool_.find(buffer.size) == memory_pool_.end()) {
- memory_pool_.emplace(buffer.size, std::vector<Buffer>{});
- }
- memory_pool_.at(buffer.size).push_back(buffer);
- VLOG(1) << "reclaim buffer " << buffer.size;
- }
-
- void* CreateView(const Buffer& buffer, ffi::Shape shape, DLDataType
type_hint,
- const std::string& mem_scope) final {
- OpenCLWorkspace* ws_ = OpenCLWorkspace::Global();
- return ws_->AllocDataSpaceView(buffer.device, buffer.data, shape,
type_hint,
- ffi::String(mem_scope));
- }
-
- void FreeView(Device dev, void* data) final {
- OpenCLWorkspace* ws_ = OpenCLWorkspace::Global();
- return ws_->FreeDataSpaceView(dev, data);
- }
-
- protected:
- // Empty uses this pool for its backing storage. Keep these hooks below Empty
- // so a cache miss does not recursively enter owning tensor allocation.
- void* DeviceAllocDataSpace(Device dev, size_t size, size_t alignment,
- DLDataType type_hint) override {
- return OpenCLWorkspace::Global()->AllocCLBuffer(dev, size, alignment,
type_hint);
- }
- void DeviceFreeDataSpace(Device dev, void* data) override {
- OpenCLWorkspace::Global()->FreeDataSpace(dev, data);
- }
-};
-
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def_packed("DeviceAllocator.opencl", [](ffi::PackedArgs
args, ffi::Any* rv) {
- Allocator* alloc = new OpenCLPooledAllocator();
- *rv = static_cast<void*>(alloc);
- });
+ refl::GlobalDef()
+ .def("runtime.opencl.clear_pool", []() {
OpenCLWorkspace::Global()->ClearMemoryPool(); })
+ .def("runtime.opencl.alloc_tensor",
+ [](ffi::Shape shape, DLDataType dtype, DLDevice device,
+ ffi::Optional<ffi::String> scope) {
+ return OpenCLWorkspace::Global()->Empty(shape, dtype, device,
scope);
+ })
+ .def("runtime.opencl.create_view",
+ [](ffi::Tensor backing, ffi::Shape shape, DLDataType dtype,
+ ffi::Optional<ffi::String> scope) {
+ return ffi::make_object<OpenCLTensorMemoryPool>()->CreateView(
+ std::move(backing), std::move(shape), dtype, 0, scope);
+ })
+ .def("runtime.create_tensor_memory_pool.opencl", [](DLDevice) {
+ return
ffi::ObjectPtr<TensorMemoryPool>(ffi::make_object<OpenCLTensorMemoryPool>());
+ });
}
} // namespace cl
diff --git a/src/backend/rocm/runtime/rocm_common.h
b/src/backend/rocm/runtime/rocm_common.h
index 96b41f71e8..efd67c7a21 100644
--- a/src/backend/rocm/runtime/rocm_common.h
+++ b/src/backend/rocm/runtime/rocm_common.h
@@ -30,6 +30,8 @@
#include <string>
+#include "../../../runtime/workspace_pool.h"
+
namespace tvm {
namespace runtime {
@@ -53,6 +55,9 @@ class ROCMThreadEntry {
public:
/*! \brief The hip stream */
hipStream_t stream{nullptr};
+ /*! \brief thread local pool */
+ WorkspacePool pool;
+ WorkspacePool host_pool;
/*! \brief constructor */
ROCMThreadEntry();
// get the threadlocal workspace
diff --git a/src/backend/rocm/runtime/rocm_device_api.cc
b/src/backend/rocm/runtime/rocm_device_api.cc
index 6c5efcc947..8addf188c3 100644
--- a/src/backend/rocm/runtime/rocm_device_api.cc
+++ b/src/backend/rocm/runtime/rocm_device_api.cc
@@ -233,6 +233,18 @@ class ROCMDeviceAPI final : public DeviceAPI {
ROCM_CALL(hipStreamSynchronize(static_cast<hipStream_t>(stream)));
}
+ void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final {
+ auto* entry = ROCMThreadEntry::ThreadLocal();
+ auto& pool = dev.device_type == kDLROCMHost ? entry->host_pool :
entry->pool;
+ return pool.AllocWorkspace(dev, size);
+ }
+
+ void FreeWorkspace(Device dev, void* data) final {
+ auto* entry = ROCMThreadEntry::ThreadLocal();
+ auto& pool = dev.device_type == kDLROCMHost ? entry->host_pool :
entry->pool;
+ pool.FreeWorkspace(dev, data);
+ }
+
static ROCMDeviceAPI* Global() {
static ROCMDeviceAPI* inst = new ROCMDeviceAPI();
return inst;
@@ -249,7 +261,8 @@ class ROCMDeviceAPI final : public DeviceAPI {
}
};
-ROCMThreadEntry::ROCMThreadEntry() {}
+ROCMThreadEntry::ROCMThreadEntry()
+ : pool(kDLROCM, ROCMDeviceAPI::Global()), host_pool(kDLROCMHost,
ROCMDeviceAPI::Global()) {}
ROCMThreadEntry* ROCMThreadEntry::ThreadLocal() {
static thread_local ROCMThreadEntry inst;
diff --git a/src/backend/vulkan/runtime/vulkan_device_api.cc
b/src/backend/vulkan/runtime/vulkan_device_api.cc
index c5174e56ed..bff473d868 100644
--- a/src/backend/vulkan/runtime/vulkan_device_api.cc
+++ b/src/backend/vulkan/runtime/vulkan_device_api.cc
@@ -286,6 +286,18 @@ void VulkanDeviceAPI::FreeDataSpace(Device dev, void* ptr)
{
delete pbuf;
}
+void* VulkanDeviceAPI::AllocWorkspace(Device dev, size_t size, DLDataType
type_hint) {
+ auto& pool = pool_per_thread.GetOrMake(kDLVulkan, this);
+ return pool.AllocWorkspace(dev, size);
+}
+
+void VulkanDeviceAPI::FreeWorkspace(Device dev, void* data) {
+ auto* pool = pool_per_thread.Get();
+ TVM_FFI_ICHECK(pool) << "Attempted to free a vulkan workspace on a
CPU-thread "
+ << "that has never allocated a workspace";
+ pool->FreeWorkspace(dev, data);
+}
+
TVMStreamHandle VulkanDeviceAPI::CreateStream(Device dev) { return nullptr; }
void VulkanDeviceAPI::FreeStream(Device dev, TVMStreamHandle stream) {
diff --git a/src/backend/vulkan/runtime/vulkan_device_api.h
b/src/backend/vulkan/runtime/vulkan_device_api.h
index 7cc41e759a..7e2ee41e79 100644
--- a/src/backend/vulkan/runtime/vulkan_device_api.h
+++ b/src/backend/vulkan/runtime/vulkan_device_api.h
@@ -26,6 +26,7 @@
#include <string>
#include <vector>
+#include "../../../runtime/workspace_pool.h"
#include "thread_map.h"
#include "vulkan/vulkan_core.h"
#include "vulkan_device.h"
@@ -47,6 +48,9 @@ class VulkanDeviceAPI final : public DeviceAPI {
ffi::Tensor Empty(ffi::ShapeView shape, DLDataType dtype, DLDevice device,
ffi::Optional<ffi::String> mem_scope = std::nullopt) final;
+ void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final;
+ void FreeWorkspace(Device dev, void* data) final;
+
// Current vulkan implementation has one "stream" per CPU thread,
// with all commands writing into a single command buffer that is
// submitted on a call to StreamSync. Therefore, for now, these are
@@ -114,6 +118,13 @@ class VulkanDeviceAPI final : public DeviceAPI {
*/
std::vector<VulkanDevice> devices_;
+ /*! \brief One pool of device memory for each CPU thread.
+ *
+ * These allocate memory based on the devices stored in devices_.
+ * The memory pools must be destructed before devices_.
+ */
+ ThreadMap<WorkspacePool> pool_per_thread;
+
/*! \brief The index of the active device for each CPU thread.
*
* To mimic the semantics of cudaSetDevice, each CPU thread can set
diff --git a/src/relax/backend/vm/codegen_vm.cc
b/src/relax/backend/vm/codegen_vm.cc
index d2ad11c7d9..bf5388f534 100644
--- a/src/relax/backend/vm/codegen_vm.cc
+++ b/src/relax/backend/vm/codegen_vm.cc
@@ -364,6 +364,7 @@ class CodeGenVM : public ExprFunctor<Instruction::Arg(const
Expr&)> {
void EmitAllocTensor(const Call& call_node, RegName dst_reg) {
TVM_FFI_ICHECK_EQ(call_node->args.size(), 5);
std::vector<Instruction::Arg> args;
+ args.push_back(Instruction::Arg::Register(Instruction::kVMRegister));
for (int i = 0; i < 4; ++i) {
args.push_back(this->VisitExpr(call_node->args[i]));
}
diff --git a/src/runtime/cpu_device_api.cc b/src/runtime/cpu_device_api.cc
index a3b43e16c3..374e89f4c6 100644
--- a/src/runtime/cpu_device_api.cc
+++ b/src/runtime/cpu_device_api.cc
@@ -29,6 +29,8 @@
#include <cstdlib>
#include <cstring>
+#include "workspace_pool.h"
+
#ifdef __ANDROID__
#include <android/api-level.h>
#endif
@@ -132,6 +134,9 @@ class CPUDeviceAPI final : public DeviceAPI {
public:
void StreamSync(Device dev, TVMStreamHandle stream) final {}
+ void* AllocWorkspace(Device dev, size_t size, DLDataType type_hint) final;
+ void FreeWorkspace(Device dev, void* data) final;
+
bool SupportsDevicePointerArithmeticsOnHost() final { return true; }
static CPUDeviceAPI* Global() {
@@ -149,6 +154,23 @@ class CPUDeviceAPI final : public DeviceAPI {
}
};
+struct CPUWorkspacePool : public WorkspacePool {
+ CPUWorkspacePool() : WorkspacePool(kDLCPU, CPUDeviceAPI::Global()) {}
+};
+
+static CPUWorkspacePool* CPUWorkspacePoolThreadLocal() {
+ static thread_local CPUWorkspacePool inst;
+ return &inst;
+}
+
+void* CPUDeviceAPI::AllocWorkspace(Device dev, size_t size, DLDataType
type_hint) {
+ return CPUWorkspacePoolThreadLocal()->AllocWorkspace(dev, size);
+}
+
+void CPUDeviceAPI::FreeWorkspace(Device dev, void* data) {
+ CPUWorkspacePoolThreadLocal()->FreeWorkspace(dev, data);
+}
+
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef().def_packed("device_api.cpu", [](ffi::PackedArgs args,
ffi::Any* rv) {
diff --git a/src/runtime/device_api.cc b/src/runtime/device_api.cc
index bfeb959165..a94b7f28ee 100644
--- a/src/runtime/device_api.cc
+++ b/src/runtime/device_api.cc
@@ -37,7 +37,7 @@
#include <cctype>
#include <cstdlib>
#include <limits>
-#include <memory>
+#include <map>
#include <mutex>
#include <sstream>
#include <string>
@@ -45,8 +45,6 @@
#include <unordered_map>
#include <variant>
-#include "workspace_pool.h"
-
namespace tvm {
namespace runtime {
@@ -99,6 +97,28 @@ int DeviceAPI::GetDeviceIndex() {
TVM_FFI_THROW(NotImplementedError) << "Device does not support
GetDeviceIndex.";
}
+namespace {
+// Default workspace allocation has no cache. Backends with reuse own their
pools.
+std::unordered_map<void*, ffi::Tensor>& DefaultWorkspaces(Device dev) {
+ static thread_local std::map<std::pair<int, int>, std::unordered_map<void*,
ffi::Tensor>>
+ allocations;
+ return allocations[{dev.device_type, dev.device_id}];
+}
+} // namespace
+
+void* DeviceAPI::AllocWorkspace(Device dev, size_t size, DLDataType type_hint)
{
+ int64_t count =
+ static_cast<int64_t>(size / kTempAllocaAlignment + (size %
kTempAllocaAlignment != 0));
+ auto owner = Empty(ffi::ShapeView(&count, 1), DLDataType{kDLUInt, 8,
kTempAllocaAlignment}, dev);
+ void* ptr = owner.data_ptr();
+ DefaultWorkspaces(dev).emplace(ptr, std::move(owner));
+ return ptr;
+}
+
+void DeviceAPI::FreeWorkspace(Device dev, void* ptr) {
+ TVM_FFI_ICHECK_EQ(DefaultWorkspaces(dev).erase(ptr), 1);
+}
+
void DeviceAPI::CopyDataFromTo(DLTensor* from, DLTensor* to, TVMStreamHandle
stream) {
// by default, we can always redirect to the flat memory copy operation.
size_t nbytes = ffi::GetDataSize(*from);
@@ -200,25 +220,18 @@ int TVMBackendGetDeviceIndex(int device_type, int*
out_device_index) {
TVM_FFI_SAFE_CALL_END();
}
-namespace {
-WorkspacePool* BackendWorkspacePool(DLDevice dev) {
- static thread_local std::unordered_map<int, std::unique_ptr<WorkspacePool>>
pools;
- auto& pool = pools[static_cast<int>(dev.device_type)];
- if (!pool) pool = std::make_unique<WorkspacePool>(dev.device_type,
DeviceAPI::Get(dev));
- return pool.get();
-}
-} // namespace
-
void* TVMBackendAllocWorkspace(int device_type, int device_id, uint64_t size,
int dtype_code_hint,
int dtype_bits_hint) {
DLDevice dev{static_cast<DLDeviceType>(device_type), device_id};
TVM_FFI_CHECK_LE(size, std::numeric_limits<size_t>::max(), ValueError)
<< "Workspace size exceeds addressable memory";
- return BackendWorkspacePool(dev)->AllocWorkspace(dev,
static_cast<size_t>(size));
+ DLDataType type_hint{static_cast<uint8_t>(dtype_code_hint),
static_cast<uint8_t>(dtype_bits_hint),
+ 1};
+ return DeviceAPI::Get(dev)->AllocWorkspace(dev, static_cast<size_t>(size),
type_hint);
}
int TVMBackendFreeWorkspace(int device_type, int device_id, void* ptr) {
DLDevice dev{static_cast<DLDeviceType>(device_type), device_id};
- BackendWorkspacePool(dev)->FreeWorkspace(dev, ptr);
+ DeviceAPI::Get(dev)->FreeWorkspace(dev, ptr);
return 0;
}
diff --git a/src/runtime/extra/contrib/nvshmem/memory_allocator.cc
b/src/runtime/extra/contrib/nvshmem/memory_allocator.cc
index 2c2682a553..3c006e2fe4 100644
--- a/src/runtime/extra/contrib/nvshmem/memory_allocator.cc
+++ b/src/runtime/extra/contrib/nvshmem/memory_allocator.cc
@@ -21,76 +21,58 @@
#include <tvm/ffi/extra/cuda/base.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/memory/memory_manager.h>
+#include <tvm/runtime/tensor.h>
-#include <thread>
-
-#include "../../../memory/pooled_allocator.h"
#include "../../disco/utils.h"
namespace tvm {
namespace runtime {
-using tvm::runtime::memory::Buffer;
-using tvm::runtime::memory::PooledAllocator;
-
/*!
- * \brief The memory allocator of NVSHMEM.
- * Overriding PooledAllocator for efficient memory management.
+ * \brief Tensor pool for NVSHMEM symmetric memory.
+ * Allocations, releases and finalization retain NVSHMEM's collective ordering
+ * requirements. All tensor owners must be released before finalization.
*/
-class NVSHMEMAllocator final : public PooledAllocator {
+class NVSHMEMPool final : public TensorMemoryPool {
public:
- explicit NVSHMEMAllocator() : PooledAllocator() {}
-
- ~NVSHMEMAllocator() { PooledAllocator::ReleaseAll(); }
-
- void Clear() final { PooledAllocator::ReleaseAll(); }
+ explicit NVSHMEMPool() = default;
- bool AllowMemoryScope(const std::string& mem_scope) const final {
- // The allowed memory scope of NVSHMEM is "nvshmem";
- return mem_scope == "nvshmem";
+ /*! \brief Return the global NVSHMEM singleton pool. */
+ static NVSHMEMPool* Global() {
+ static auto* pool = new
ffi::ObjectPtr<NVSHMEMPool>(ffi::make_object<NVSHMEMPool>());
+ return pool->get();
}
- /*! \brief Return the global NVSHMEM singleton allocator. */
- static NVSHMEMAllocator* Global() {
- static NVSHMEMAllocator* allocator = new NVSHMEMAllocator();
- return allocator;
- }
-
- ffi::Tensor Empty(ffi::Shape shape, DLDataType dtype, Device device) {
- class NVSHMEMAlloc {
- public:
- explicit NVSHMEMAlloc(Buffer buffer) : buffer_(buffer) {}
- void AllocData(DLTensor* tensor) { tensor->data = buffer_.data; }
- void FreeData(DLTensor* tensor) {
NVSHMEMAllocator::Global()->Free(buffer_); }
-
- private:
- Buffer buffer_;
- };
-
- Buffer buffer = PooledAllocator::Alloc(device, shape, dtype,
ffi::String("nvshmem"));
- return ffi::Tensor::FromNDAlloc(NVSHMEMAlloc(buffer), shape, dtype,
device);
+ ffi::Tensor Empty(ffi::ShapeView shape, DLDataType dtype, DLDevice device,
+ ffi::Optional<ffi::String> mem_scope = std::nullopt) final
{
+ TVM_FFI_ICHECK_EQ(device.device_type, DLDeviceType::kDLCUDA)
+ << "nvshmem can only allocate cuda device memory space.";
+ TVM_FFI_ICHECK(dtype.code == kDLInt || dtype.code == kDLUInt || dtype.code
== kDLFloat ||
+ dtype.code == kDLBfloat)
+ << "nvshmem can only allocate tensor with int, unsigned int, float, or
bfloat data types.";
+ TVM_FFI_CHECK(!mem_scope.has_value() || mem_scope.value() == "nvshmem",
ValueError)
+ << "NVSHMEM memory requires the nvshmem scope";
+ return TensorMemoryPool::Empty(shape, dtype, device);
}
private:
- void* DeviceAllocDataSpace(Device dev, size_t size, size_t alignment,
- DLDataType type_hint) final {
- TVM_FFI_ICHECK_EQ(dev.device_type, DLDeviceType::kDLCUDA)
- << "nvshmem can only allocate cuda device memory space.";
- TVM_FFI_ICHECK(
- type_hint.code == DLDataTypeCode::kDLInt || type_hint.code ==
DLDataTypeCode::kDLUInt ||
- type_hint.code == DLDataTypeCode::kDLFloat || type_hint.code ==
DLDataTypeCode::kDLBfloat)
- << "nvshmem can only allocate tensor with int, usingned int, float, or
bfloat data types.";
- return nvshmem_align(alignment, size);
+ ffi::Tensor AllocBacking(size_t size, size_t alignment, DLDevice device)
final {
+ struct Backing {
+ size_t size;
+ size_t alignment;
+ void AllocData(DLTensor* tensor) { tensor->data =
nvshmem_align(alignment, size); }
+ void FreeData(DLTensor* tensor) { nvshmem_free(tensor->data); }
+ };
+ int64_t extent = static_cast<int64_t>(size);
+ return ffi::Tensor::FromNDAlloc(Backing{size, alignment},
ffi::ShapeView(&extent, 1),
+ DLDataType{kDLUInt, 8, 1}, device);
}
-
- void DeviceFreeDataSpace(Device dev, void* ptr) final { nvshmem_free(ptr); }
};
ffi::Tensor NVSHMEMEmpty(ffi::Shape shape, DLDataType dtype,
ffi::Optional<Device> device) {
Device resolved_device =
device.has_value() ? device.value() :
UseDefaultDeviceIfNone(std::nullopt);
- return NVSHMEMAllocator::Global()->Empty(shape, dtype, resolved_device);
+ return NVSHMEMPool::Global()->Empty(shape, dtype, resolved_device);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -99,7 +81,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
}
void NVSHMEMFinalize() {
- NVSHMEMAllocator::Global()->Clear();
+ NVSHMEMPool::Global()->Clear();
nvshmem_finalize();
}
diff --git a/src/runtime/extra/disco/builtin.cc
b/src/runtime/extra/disco/builtin.cc
index ce6f5e9729..1651441e24 100644
--- a/src/runtime/extra/disco/builtin.cc
+++ b/src/runtime/extra/disco/builtin.cc
@@ -23,7 +23,7 @@
#include <tvm/runtime/disco/builtin.h>
#include <tvm/runtime/disco/disco_worker.h>
#include <tvm/runtime/disco/session.h>
-#include <tvm/runtime/vm/vm.h>
+#include <tvm/runtime/tensor.h>
#include <sstream>
@@ -65,9 +65,7 @@ ffi::Module LoadVMModule(std::string path,
ffi::Optional<Device> device) {
<< "File `" << path
<< "` is not built by RelaxVM, because `vm_initialization` does not
exist";
}
- (*vm_initialization)(static_cast<int>(dev.device_type),
static_cast<int>(dev.device_id),
- static_cast<int>(AllocatorType::kPooled),
static_cast<int>(kDLCPU), 0,
- static_cast<int>(AllocatorType::kPooled));
+ ffi::Function::GetGlobalRequired("vm.builtin.initialize_default")(mod, dev);
return mod;
}
diff --git a/src/runtime/extra/disco/cuda_ipc/cuda_ipc_memory.cc
b/src/runtime/extra/disco/cuda_ipc/cuda_ipc_memory.cc
index 1d7cb73e6f..aea7665361 100644
--- a/src/runtime/extra/disco/cuda_ipc/cuda_ipc_memory.cc
+++ b/src/runtime/extra/disco/cuda_ipc/cuda_ipc_memory.cc
@@ -22,10 +22,11 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/disco/cuda_ipc_memory.h>
-#include <tvm/runtime/memory/memory_manager.h>
+#include <tvm/runtime/tensor.h>
+
+#include <unordered_map>
#include "../../../../../3rdparty/tensorrt_llm/custom_allreduce_kernels.h"
-#include "../../../memory/pooled_allocator.h"
#include "../nccl/nccl_context.h"
namespace tvm {
@@ -34,7 +35,6 @@ namespace cuda_ipc {
using tensorrt_llm::MAX_ALL_REDUCE_BLOCKS;
using tensorrt_llm::MAX_RANKS_PER_NODE;
-using tvm::runtime::memory::Buffer;
/*!
* \brief All-gather the IPC memory handles across all distributed workers.
@@ -66,15 +66,17 @@ std::vector<cudaIpcMemHandle_t>
AllGatherIPCHandles(nccl::CCLThreadLocalContext*
/*!
* \brief The memory allocator of CUDAIPCMemory.
- * Overriding PooledAllocator for efficient memory management.
+ * Overriding TensorMemoryPool for efficient memory management.
*/
-class CUDAIPCMemoryAllocator final : public memory::PooledAllocator {
+class CUDAIPCMemoryPool final : public TensorMemoryPool {
public:
- explicit CUDAIPCMemoryAllocator() : PooledAllocator() {}
+ explicit CUDAIPCMemoryPool() = default;
- bool AllowMemoryScope(const std::string& mem_scope) const final {
- // The allowed memory scope of CUDAIPCMemory is "ipc_memory";
- return mem_scope == "ipc_memory";
+ ffi::Tensor Empty(ffi::ShapeView shape, DLDataType dtype, DLDevice device,
+ ffi::Optional<ffi::String> mem_scope = std::nullopt) final
{
+ TVM_FFI_CHECK(!mem_scope.has_value() || mem_scope.value() == "ipc_memory",
ValueError)
+ << "CUDA IPC memory requires the ipc_memory scope";
+ return TensorMemoryPool::Empty(shape, dtype, device);
}
CUDAIPCMemory GetIPCMemoryFromDevicePtr(void* ptr) const {
@@ -86,16 +88,16 @@ class CUDAIPCMemoryAllocator final : public
memory::PooledAllocator {
}
/*! \brief Return the global CUDAIPCMemory singleton allocator. */
- static CUDAIPCMemoryAllocator* Global() {
- static CUDAIPCMemoryAllocator* allocator = new CUDAIPCMemoryAllocator();
- return allocator;
+ static CUDAIPCMemoryPool* Global() {
+ static auto* pool =
+ new
ffi::ObjectPtr<CUDAIPCMemoryPool>(ffi::make_object<CUDAIPCMemoryPool>());
+ return pool->get();
}
private:
- void* DeviceAllocDataSpace(Device dev, size_t size, size_t alignment,
- DLDataType type_hint) final {
- auto [data_ptr, data_comm_ptrs] =
- AllocIPCMemory(dev, size, alignment, type_hint,
/*reset_memory_to_zero=*/false);
+ ffi::Tensor AllocBacking(size_t size, size_t alignment, DLDevice dev) final {
+ auto [data_ptr, data_comm_ptrs] = AllocIPCMemory(
+ dev, size, alignment, DLDataType{kDLUInt, 8, 1},
/*reset_memory_to_zero=*/false);
int barrier_ptr_size = sizeof(uint32_t) * (MAX_ALL_REDUCE_BLOCKS + 2) *
MAX_RANKS_PER_NODE;
auto [barrier_in_ptr, barrier_in_comm_ptrs] =
AllocIPCMemory(dev, barrier_ptr_size, alignment, DLDataType{kDLUInt,
32, 1},
@@ -113,21 +115,24 @@ class CUDAIPCMemoryAllocator final : public
memory::PooledAllocator {
ipc_memory->barrier_flag = 1;
ipc_memory->num_workers = nccl_ctx->worker->num_workers;
ipc_memory->worker_id = nccl_ctx->worker->worker_id;
- ipc_memory_map_.insert_or_assign(data_ptr,
- CUDAIPCMemory(ffi::UnsafeInit{},
std::move(ipc_memory)));
- return data_ptr;
- }
-
- void DeviceFreeDataSpace(Device dev, void* ptr) final {
- TVM_FFI_ICHECK(dev.device_type == kDLCUDA);
- ffi::CUDADeviceGuard device_guard(dev.device_id);
- nccl::CCLThreadLocalContext* ctx = nccl::CCLThreadLocalContext::Get();
- auto it = ipc_memory_map_.find(ptr);
- TVM_FFI_ICHECK(it != ipc_memory_map_.end());
- FreeIPCMemory(it->second->remote_data, ctx->worker->worker_id);
- FreeIPCMemory(it->second->barrier_in, ctx->worker->worker_id);
- FreeIPCMemory(it->second->barrier_out, ctx->worker->worker_id);
- ipc_memory_map_.erase(it);
+ CUDAIPCMemory memory(ffi::UnsafeInit{}, std::move(ipc_memory));
+ ipc_memory_map_.insert_or_assign(data_ptr, memory);
+ struct Backing {
+ CUDAIPCMemory memory;
+ void AllocData(DLTensor* tensor) { tensor->data =
memory->remote_data[memory->worker_id]; }
+ void FreeData(DLTensor* tensor) {
+ ffi::CUDADeviceGuard device_guard(tensor->device.device_id);
+ FreeIPCMemory(memory->remote_data, memory->worker_id);
+ FreeIPCMemory(memory->barrier_in, memory->worker_id);
+ FreeIPCMemory(memory->barrier_out, memory->worker_id);
+ // The original global lookup map has process lifetime, independently
of
+ // the thread-local NCCL context used to exchange the handles.
+ Global()->ipc_memory_map_.erase(tensor->data);
+ }
+ };
+ int64_t extent = static_cast<int64_t>(size);
+ return ffi::Tensor::FromNDAlloc(Backing{std::move(memory)},
ffi::ShapeView(&extent, 1),
+ DLDataType{kDLUInt, 8, 1}, dev);
}
/*!
@@ -185,7 +190,7 @@ class CUDAIPCMemoryAllocator final : public
memory::PooledAllocator {
}
/*! \brief Free the IPC memory pointers. */
- void FreeIPCMemory(std::vector<void*> comm_ptrs, int worker_id) {
+ static void FreeIPCMemory(std::vector<void*> comm_ptrs, int worker_id) {
for (int i = 0; i < static_cast<int>(comm_ptrs.size()); ++i) {
if (i != worker_id) {
// Free ipc handle.
@@ -202,35 +207,32 @@ class CUDAIPCMemoryAllocator final : public
memory::PooledAllocator {
};
/*!
- * \brief Allocate a storage object with CUDA IPC memory.
- * \param buffer_shape The shape of the storage to allocate.
- * \param dtype_hint The dtype of the storage to allocate.
- * \return The allocated storage object with internal CUDA IPC memory buffer.
+ * \brief Allocate owning CUDA IPC backing for VM tensors.
*/
-memory::Storage IPCAllocStorage(ffi::Shape buffer_shape, DLDataType
dtype_hint) {
+ffi::Tensor IPCAllocTensor(ffi::Shape buffer_shape, DLDataType dtype_hint) {
nccl::CCLThreadLocalContext* nccl_ctx = nccl::CCLThreadLocalContext::Get();
Device device{DLDeviceType::kDLCUDA, nccl_ctx->device_id};
- CUDAIPCMemoryAllocator* allocator = CUDAIPCMemoryAllocator::Global();
- auto buffer = allocator->Alloc(device, std::move(buffer_shape), dtype_hint,
"ipc_memory");
- return memory::Storage(std::move(buffer), allocator);
+ return CUDAIPCMemoryPool::Global()->Empty(buffer_shape, dtype_hint, device,
+ ffi::String("ipc_memory"));
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
- .def("runtime.disco.cuda_ipc.alloc_storage", IPCAllocStorage)
+ .def("runtime.disco.cuda_ipc.empty", IPCAllocTensor)
+ .def("runtime.disco.cuda_ipc.alloc_storage", IPCAllocTensor)
.def("runtime.disco.cuda_ipc.cuda_ipc_memory_allocator_clear",
- []() { CUDAIPCMemoryAllocator::Global()->Clear(); });
+ []() { CUDAIPCMemoryPool::Global()->Clear(); });
}
/******************** CUDAIPCMemoryObj ********************/
-// Direct to CUDAIPCMemoryAllocator::Global.
-memory::Allocator* CUDAIPCMemory::GlobalAllocator() { return
CUDAIPCMemoryAllocator::Global(); }
+// Direct to CUDAIPCMemoryPool::Global.
+TensorMemoryPool* CUDAIPCMemory::GlobalMemoryPool() { return
CUDAIPCMemoryPool::Global(); }
-// Direct to CUDAIPCMemoryAllocator::GlobalGetIPCMemoryFromDevicePtr.
+// Direct to CUDAIPCMemoryPool::GlobalGetIPCMemoryFromDevicePtr.
CUDAIPCMemory CUDAIPCMemory::GetIPCMemoryFromDevicePtr(void* ptr) {
- return CUDAIPCMemoryAllocator::Global()->GetIPCMemoryFromDevicePtr(ptr);
+ return CUDAIPCMemoryPool::Global()->GetIPCMemoryFromDevicePtr(ptr);
}
} // namespace cuda_ipc
diff --git a/src/runtime/extra/disco/cuda_ipc/custom_allreduce.cc
b/src/runtime/extra/disco/cuda_ipc/custom_allreduce.cc
index 3eaca5ba98..4ee47ed312 100644
--- a/src/runtime/extra/disco/cuda_ipc/custom_allreduce.cc
+++ b/src/runtime/extra/disco/cuda_ipc/custom_allreduce.cc
@@ -21,7 +21,6 @@
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/disco/cuda_ipc_memory.h>
-#include <tvm/runtime/memory/memory_manager.h>
#include "../../../../../3rdparty/tensorrt_llm/custom_allreduce_kernels.h"
#include "../nccl/nccl_context.h"
diff --git a/src/runtime/memory/memory_manager.cc
b/src/runtime/memory/memory_manager.cc
deleted file mode 100644
index 04eaf8960a..0000000000
--- a/src/runtime/memory/memory_manager.cc
+++ /dev/null
@@ -1,278 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you 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.
- */
-
-/*!
- * \file tvm/runtime/memory/memory_manager.cc
- * \brief Allocate and manage memory for the runtime.
- */
-#include <tvm/ffi/cast.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/runtime/memory/memory_manager.h>
-
-#include <memory>
-#include <utility>
-
-#include "naive_allocator.h"
-#include "pooled_allocator.h"
-
-namespace tvm {
-namespace runtime {
-namespace memory {
-
-Storage::Storage(Buffer buffer, Allocator* allocator) {
- auto n = ffi::make_object<StorageObj>();
- n->buffer = std::move(buffer);
- n->allocator = allocator;
- data_ = std::move(n);
-}
-
-inline void VerifyDataType(DLDataType dtype) {
- TVM_FFI_ICHECK_GE(dtype.lanes, 1);
- if (dtype.code == kDLFloat) {
- TVM_FFI_ICHECK_EQ(dtype.bits % 8, 0);
- } else {
- // allow uint1 as a special flag for bool.
- if (dtype.bits == 1 && dtype.code == kDLUInt) return;
- TVM_FFI_ICHECK_EQ(dtype.bits % 8, 0);
- }
- TVM_FFI_ICHECK_EQ(dtype.bits & (dtype.bits - 1), 0);
-}
-
-inline size_t GetDataAlignment(const DLDataType& dtype) {
- size_t align = dtype.lanes * dtype.bits / 8;
- if (align < kAllocAlignment) return kAllocAlignment;
- return align;
-}
-
-ffi::Tensor StorageObj::AllocTensorScoped(int64_t offset, ffi::Shape shape,
DLDataType dtype,
- ffi::String scope) {
- if (scope == "global" || scope.empty()) {
- return AllocTensor(offset, shape, dtype);
- }
- VerifyDataType(dtype);
-
- struct StorageScopedAlloc {
- public:
- explicit StorageScopedAlloc(Storage storage) : storage_(storage) {}
-
- void AllocData(DLTensor* tensor, const ffi::Shape& shape, const
ffi::String& scope,
- int64_t byte_offset) {
- tensor->data = storage_->allocator->CreateView(storage_->buffer, shape,
tensor->dtype, scope);
- tensor->byte_offset = byte_offset;
- }
- void FreeData(DLTensor* tensor) {
storage_->allocator->FreeView(tensor->device, tensor->data); }
-
- private:
- Storage storage_;
- };
-
- size_t needed_size = ffi::GetDataSize(shape.Product(), dtype);
- TVM_FFI_ICHECK(offset + needed_size <= this->buffer.size)
- << "storage allocation failure, attempted to allocate " << needed_size
<< " at offset "
- << offset << " in region that is " << this->buffer.size << "bytes";
-
- return
ffi::Tensor::FromNDAlloc(StorageScopedAlloc(ffi::GetRef<Storage>(this)), shape,
dtype,
- this->buffer.device, shape, scope, offset);
-}
-
-ffi::Tensor StorageObj::AllocTensor(int64_t offset, ffi::Shape shape,
DLDataType dtype) {
- VerifyDataType(dtype);
-
- size_t needed_size = ffi::GetDataSize(shape.Product(), dtype);
- TVM_FFI_ICHECK(offset + needed_size <= this->buffer.size)
- << "storage allocation failure, attempted to allocate " << needed_size
<< " at offset "
- << offset << " in region that is " << this->buffer.size << "bytes";
- class StorageAlloc {
- public:
- explicit StorageAlloc(Storage storage) : storage_(storage) {}
-
- void AllocData(DLTensor* tensor, int64_t offset) {
- if (storage_->buffer.device.device_type == kDLHexagon) {
- // For Hexagon, non-zero offset support simply requires adjusting the
- // beginning of data pointer
- auto offset_ptr = reinterpret_cast<uint8_t*>(storage_->buffer.data) +
offset;
- tensor->data = reinterpret_cast<void*>(offset_ptr);
- tensor->byte_offset = 0;
- } else {
- tensor->data = storage_->buffer.data;
- tensor->byte_offset = offset;
- }
- }
- void FreeData(DLTensor* tensor) {}
-
- private:
- Storage storage_;
- };
-
- return ffi::Tensor::FromNDAlloc(StorageAlloc(ffi::GetRef<Storage>(this)),
shape, dtype,
- this->buffer.device, offset);
-}
-
-MemoryManager* MemoryManager::Global() {
- // NOTE: explicitly use new to avoid exit-time destruction of global state
- // Global state will be recycled by OS as the process exits.
- static auto* inst = new MemoryManager();
- return inst;
-}
-
-std::string DeviceTypeStr(DLDeviceType type) {
- switch (type) {
- case kDLOpenCL:
- return "opencl";
- break;
- case kDLVulkan:
- return "vulkan";
- break;
- default:
- return "";
- }
-}
-
-Allocator* GetDeviceSpecificAllocator(Device dev, AllocatorType type) {
- std::string dev_str = DeviceTypeStr(dev.device_type);
- auto device_alloc_helper = tvm::ffi::Function::GetGlobal("DeviceAllocator."
+ dev_str);
- void* valloc;
- Allocator* allocator = nullptr;
- if (device_alloc_helper) {
- valloc = (*device_alloc_helper)(dev, static_cast<int>(type)).cast<void*>();
- allocator = static_cast<Allocator*>(valloc);
- }
- if (nullptr == allocator) {
- switch (type) {
- case kNaive: {
- VLOG(1) << "New naive allocator for " << dev;
- allocator = new NaiveAllocator();
- break;
- }
- case kPooled: {
- VLOG(1) << "New pooled allocator for " << dev;
- allocator = new PooledAllocator();
- break;
- }
- default:
- TVM_FFI_THROW(InternalError) << "Unknown allocator type: " << type;
- }
- }
- return allocator;
-}
-
-Allocator* MemoryManager::GetOrCreateAllocator(Device dev, AllocatorType type)
{
- MemoryManager* m = MemoryManager::Global();
- std::lock_guard<std::mutex> lock(m->mu_);
- if (m->allocators_.find(dev) == m->allocators_.end()) {
- m->allocators_.emplace(dev, std::unordered_map<AllocatorType,
std::unique_ptr<Allocator>>());
- }
- if (m->allocators_.at(dev).find(type) == m->allocators_.at(dev).end()) {
- std::unique_ptr<Allocator> alloc;
- alloc.reset(GetDeviceSpecificAllocator(dev, type));
- auto ret = alloc.get();
- m->allocators_.at(dev).emplace(type, std::move(alloc));
- return ret;
- }
- auto alloc = m->allocators_.at(dev).at(type).get();
-
- return alloc;
-}
-
-Allocator* MemoryManager::GetAllocator(Device dev, AllocatorType type) {
- MemoryManager* m = MemoryManager::Global();
- std::lock_guard<std::mutex> lock(m->mu_);
- auto it = m->allocators_.find(dev);
- if (it == m->allocators_.end()) {
- TVM_FFI_THROW(InternalError) << "Allocator for " << dev << " has not been
created yet.";
- }
- if (it->second.find(type) == it->second.end()) {
- TVM_FFI_THROW(InternalError) << "Allocator for " << dev << " of type " <<
type
- << " has not been created yet.";
- }
- return it->second.at(type).get();
-}
-
-void MemoryManager::Clear() {
- MemoryManager* m = MemoryManager::Global();
- std::lock_guard<std::mutex> lock(m->mu_);
- for (const auto& [device, allocators] : m->allocators_) {
- for (const auto& [allocator_type, allocator] : allocators) {
- allocator->Clear();
- }
- }
-}
-
-ffi::Tensor Allocator::Empty(ffi::Shape shape, DLDataType dtype, DLDevice dev,
- ffi::Optional<ffi::String> mem_scope) {
- VerifyDataType(dtype);
-
- class BufferAlloc {
- public:
- explicit BufferAlloc(Buffer buffer) : buffer_(buffer) {}
-
- void AllocData(DLTensor* tensor) { tensor->data = buffer_.data; }
- void FreeData(DLTensor* tensor) {
- MemoryManager::GetAllocator(buffer_.device,
buffer_.alloc_type)->Free(buffer_);
- }
-
- private:
- Buffer buffer_;
- };
-
- size_t alignment = GetDataAlignment(dtype);
- size_t size = ffi::GetDataSize(shape.Product(), dtype);
-
- Buffer buffer;
- if (!mem_scope.has_value() || (*mem_scope).empty() || (*mem_scope) ==
"global") {
- buffer = this->Alloc(dev, size, alignment, dtype);
- } else {
- buffer = this->Alloc(dev, shape, dtype, *mem_scope);
- }
- return ffi::Tensor::FromNDAlloc(BufferAlloc(buffer), shape, dtype, dev);
-}
-
-bool Allocator::AllowMemoryScope(const std::string& mem_scope) const {
- return mem_scope.empty() || mem_scope == "global";
-}
-
-Buffer Allocator::Alloc(Device dev, ffi::Shape shape, DLDataType type_hint,
- const std::string& mem_scope) {
- if (AllowMemoryScope(mem_scope)) {
- // by default, we can always redirect to the flat memory allocations
- size_t alignment = GetDataAlignment(type_hint);
- size_t size = ffi::GetDataSize(shape.Product(), type_hint);
- return Alloc(dev, size, alignment, type_hint);
- }
- TVM_FFI_THROW(InternalError) << "Allocator cannot allocate data space with "
- << "specified memory scope: " << mem_scope;
- return {};
-}
-
-void Allocator::Clear() {
- // This function by default does nothing.
- // For naive allocator, no explicit manual clear is needed.
- // Pooled allocator will override this method.
-}
-
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef().def("vm.builtin.memory_manager.clear",
MemoryManager::Clear);
-}
-
-} // namespace memory
-} // namespace runtime
-} // namespace tvm
diff --git a/src/runtime/memory/naive_allocator.h
b/src/runtime/memory/naive_allocator.h
deleted file mode 100644
index 6b0f6d0b05..0000000000
--- a/src/runtime/memory/naive_allocator.h
+++ /dev/null
@@ -1,143 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you 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.
- */
-
-/*!
- * \file src/runtime/memory/naive_allocator.h
- */
-#ifndef TVM_RUNTIME_MEMORY_NAIVE_ALLOCATOR_H_
-#define TVM_RUNTIME_MEMORY_NAIVE_ALLOCATOR_H_
-
-#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/runtime/memory/memory_manager.h>
-
-#include <algorithm>
-#include <atomic>
-#include <limits>
-#include <memory>
-#include <mutex>
-#include <string>
-#include <unordered_map>
-
-namespace tvm {
-namespace runtime {
-namespace memory {
-
-class NaiveAllocator final : public Allocator {
- public:
- explicit NaiveAllocator() : Allocator(kNaive), used_memory_(0) {}
-
- Buffer Alloc(Device dev, size_t nbytes, size_t alignment, DLDataType
type_hint) final {
- Buffer buf;
- buf.device = dev;
- buf.size = nbytes;
- buf.alloc_type = kNaive;
- TVM_FFI_CHECK(alignment && (alignment & (alignment - 1)) == 0, ValueError)
- << "Allocation alignment must be a power of two";
- size_t natural_alignment = std::max(static_cast<size_t>(kAllocAlignment),
- static_cast<size_t>(type_hint.bits /
8) * type_hint.lanes);
- DLDataType storage_dtype = type_hint;
- size_t storage_bytes = nbytes;
- bool align_cpu = dev.device_type == kDLCPU && alignment >
natural_alignment;
- if (align_cpu) {
- TVM_FFI_CHECK_LE(nbytes, std::numeric_limits<size_t>::max() - (alignment
- 1), ValueError)
- << "Aligned allocation size overflow";
- storage_bytes += alignment - 1;
- } else if (alignment > natural_alignment) {
- TVM_FFI_CHECK_LE(alignment, std::numeric_limits<uint16_t>::max(),
ValueError)
- << "Backend allocation alignment cannot be represented";
- storage_dtype = DLDataType{kDLUInt, 8, static_cast<uint16_t>(alignment)};
- }
- size_t element_bits = static_cast<size_t>(storage_dtype.bits) *
storage_dtype.lanes;
- TVM_FFI_CHECK_GT(element_bits, 0, ValueError) << "Invalid allocation
datatype";
- TVM_FFI_CHECK_LE(storage_bytes, (std::numeric_limits<int64_t>::max() -
element_bits) / 8,
- ValueError)
- << "Allocation size overflow";
- int64_t length = static_cast<int64_t>(
- (static_cast<uint64_t>(storage_bytes) * 8 + element_bits - 1) /
element_bits);
- ffi::Tensor owner = DeviceAPI::Get(dev)->Empty(ffi::ShapeView(&length, 1),
storage_dtype, dev);
- void* data = owner.data_ptr();
- if (align_cpu) {
- size_t space = ffi::GetDataSize(owner);
- TVM_FFI_ICHECK(std::align(alignment, nbytes, data, space) != nullptr);
- }
- buf.data = data;
- {
- std::lock_guard<std::mutex> lock(mu_);
- owners_.emplace(buf.data, std::move(owner));
- }
- used_memory_.fetch_add(nbytes, std::memory_order_relaxed);
- DLOG(INFO) << "allocate " << nbytes << " B, used memory " << used_memory_
<< " B";
- return buf;
- }
-
- Buffer Alloc(Device dev, ffi::Shape shape, DLDataType type_hint,
- const std::string& mem_scope) final {
- Buffer buf;
- size_t nbytes = 1;
- for (int i = 0; i < static_cast<int>(shape.size()); ++i) {
- nbytes *= static_cast<size_t>(shape[i]);
- }
- nbytes *= (type_hint.bits * type_hint.lanes + 7) / 8;
- buf.device = dev;
- if (AllowMemoryScope(mem_scope)) {
- auto tmp_buf = Allocator::Alloc(dev, shape, type_hint, mem_scope);
- buf.size = tmp_buf.size;
- buf.data = tmp_buf.data;
- buf.alloc_type = kNaive;
- return buf;
- }
-
- buf.size = nbytes;
- auto owner = DeviceAPI::Get(dev)->Empty(shape, type_hint, dev,
ffi::String(mem_scope));
- buf.data = owner.data_ptr();
- {
- std::lock_guard<std::mutex> lock(mu_);
- owners_.emplace(buf.data, std::move(owner));
- }
- used_memory_.fetch_add(nbytes, std::memory_order_relaxed);
- DLOG(INFO) << "allocate " << nbytes << " B, used memory " << used_memory_
<< " B";
- buf.alloc_type = kNaive;
- return buf;
- }
-
- void Free(const Buffer& buffer) override {
- {
- std::lock_guard<std::mutex> lock(mu_);
- auto it = owners_.find(buffer.data);
- TVM_FFI_ICHECK(it != owners_.end()) << "Unknown allocation owner";
- owners_.erase(it);
- }
- used_memory_.fetch_sub(buffer.size, std::memory_order_relaxed);
- DLOG(INFO) << "free " << buffer.size << " B, used memory " << used_memory_
<< " B";
- }
-
- size_t UsedMemory() const override { return
used_memory_.load(std::memory_order_relaxed); }
-
- private:
- std::atomic<size_t> used_memory_;
- std::mutex mu_;
- std::unordered_multimap<void*, ffi::Tensor> owners_;
-};
-
-} // namespace memory
-} // namespace runtime
-} // namespace tvm
-
-#endif // TVM_RUNTIME_MEMORY_NAIVE_ALLOCATOR_H_
diff --git a/src/runtime/memory/pooled_allocator.h
b/src/runtime/memory/pooled_allocator.h
deleted file mode 100644
index 2d20039fa0..0000000000
--- a/src/runtime/memory/pooled_allocator.h
+++ /dev/null
@@ -1,176 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one
- * or more contributor license agreements. See the NOTICE file
- * distributed with this work for additional information
- * regarding copyright ownership. The ASF licenses this file
- * to you 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.
- */
-
-/*!
- * \file src/runtime/memory/pooled_allocator.h
- */
-#ifndef TVM_RUNTIME_MEMORY_POOLED_ALLOCATOR_H_
-#define TVM_RUNTIME_MEMORY_POOLED_ALLOCATOR_H_
-
-#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/logging.h>
-#include <tvm/runtime/memory/memory_manager.h>
-
-#include <algorithm>
-#include <atomic>
-#include <limits>
-#include <memory>
-#include <mutex>
-#include <string>
-#include <unordered_map>
-#include <vector>
-
-namespace tvm {
-namespace runtime {
-namespace memory {
-
-class PooledAllocator : public Allocator {
- public:
- static constexpr size_t kDefaultPageSize = 4096;
-
- explicit PooledAllocator(size_t page_size = kDefaultPageSize)
- : Allocator(kPooled), page_size_(page_size), used_memory_(0) {}
-
- ~PooledAllocator() { ReleaseAll(); }
-
- Buffer Alloc(Device dev, size_t nbytes, size_t alignment, DLDataType
type_hint) override {
- std::lock_guard<std::recursive_mutex> lock(mu_);
- size_t size = ((nbytes + page_size_ - 1) / page_size_) * page_size_;
- auto&& it = memory_pool_.find(size);
- if (it != memory_pool_.end() && !it->second.empty()) {
- auto&& pool = it->second;
- bool addressable =
DeviceAPI::Get(dev)->SupportsDevicePointerArithmeticsOnHost();
- auto candidate = std::find_if(pool.begin(), pool.end(), [&](const
Buffer& buffer) {
- return !addressable ||
- (alignment != 0 && reinterpret_cast<uintptr_t>(buffer.data) %
alignment == 0);
- });
- if (candidate != pool.end()) {
- Buffer ret = *candidate;
- pool.erase(candidate);
- return ret;
- }
- }
- Buffer buf;
- buf.device = dev;
- buf.size = size;
- buf.alloc_type = kPooled;
- try {
- buf.data = DeviceAllocDataSpace(dev, size, alignment, type_hint);
- } catch (tvm::ffi::Error& err) {
- LOG(WARNING) << "PooledAllocator got InternalError during allocation: "
<< err.what();
- LOG(WARNING) << "Trying to release all unused memory and reallocate...";
- ReleaseAll();
- buf.data = DeviceAllocDataSpace(dev, size, alignment, type_hint);
- }
-
- used_memory_.fetch_add(size, std::memory_order_relaxed);
- VLOG(1) << "allocate " << size << " B, used memory " << used_memory_ << "
B";
- return buf;
- }
-
- Buffer Alloc(Device dev, ffi::Shape shape, DLDataType type_hint,
- const std::string& mem_scope) override {
- if (AllowMemoryScope(mem_scope)) {
- return Allocator::Alloc(dev, shape, type_hint, mem_scope);
- }
- TVM_FFI_THROW(InternalError) << "This alloc should be implemented";
- return {};
- }
-
- void Free(const Buffer& buffer) override {
- std::lock_guard<std::recursive_mutex> lock(mu_);
- if (memory_pool_.find(buffer.size) == memory_pool_.end()) {
- memory_pool_.emplace(buffer.size, std::vector<Buffer>{});
- }
- memory_pool_.at(buffer.size).push_back(buffer);
- VLOG(1) << "reclaim buffer " << buffer.size;
- }
-
- void Clear() override { ReleaseAll(); }
-
- size_t UsedMemory() const override { return
used_memory_.load(std::memory_order_relaxed); }
-
- protected:
- virtual void* DeviceAllocDataSpace(Device dev, size_t nbytes, size_t
alignment,
- DLDataType type_hint) {
- TVM_FFI_CHECK(alignment && (alignment & (alignment - 1)) == 0, ValueError)
- << "Allocation alignment must be a power of two";
- size_t natural_alignment = std::max(static_cast<size_t>(kAllocAlignment),
- static_cast<size_t>(type_hint.bits /
8) * type_hint.lanes);
- DLDataType storage_dtype = type_hint;
- size_t storage_bytes = nbytes;
- bool align_cpu = dev.device_type == kDLCPU && alignment >
natural_alignment;
- if (align_cpu) {
- TVM_FFI_CHECK_LE(nbytes, std::numeric_limits<size_t>::max() - (alignment
- 1), ValueError)
- << "Aligned allocation size overflow";
- storage_bytes += alignment - 1;
- } else if (alignment > natural_alignment) {
- TVM_FFI_CHECK_LE(alignment, std::numeric_limits<uint16_t>::max(),
ValueError)
- << "Backend allocation alignment cannot be represented";
- storage_dtype = DLDataType{kDLUInt, 8, static_cast<uint16_t>(alignment)};
- }
- size_t element_bits = static_cast<size_t>(storage_dtype.bits) *
storage_dtype.lanes;
- TVM_FFI_CHECK_GT(element_bits, 0, ValueError) << "Invalid allocation
datatype";
- TVM_FFI_CHECK_LE(storage_bytes, (std::numeric_limits<int64_t>::max() -
element_bits) / 8,
- ValueError)
- << "Allocation size overflow";
- int64_t length = static_cast<int64_t>(
- (static_cast<uint64_t>(storage_bytes) * 8 + element_bits - 1) /
element_bits);
- ffi::Tensor owner = DeviceAPI::Get(dev)->Empty(ffi::ShapeView(&length, 1),
storage_dtype, dev);
- void* data = owner.data_ptr();
- if (align_cpu) {
- size_t space = ffi::GetDataSize(owner);
- TVM_FFI_ICHECK(std::align(alignment, nbytes, data, space) != nullptr);
- }
- owners_.emplace(data, std::move(owner));
- return data;
- }
-
- virtual void DeviceFreeDataSpace(Device dev, void* ptr) {
- auto it = owners_.find(ptr);
- TVM_FFI_ICHECK(it != owners_.end()) << "Unknown allocation owner";
- owners_.erase(it);
- }
-
- virtual void ReleaseAll() {
- std::lock_guard<std::recursive_mutex> lock(mu_);
- for (auto const& it : memory_pool_) {
- auto const& pool = it.second;
- for (auto const& buf : pool) {
- DeviceFreeDataSpace(buf.device, buf.data);
- }
- }
- memory_pool_.clear();
- used_memory_ = 0;
- VLOG(1) << "release all buffers";
- }
-
- protected:
- std::unordered_multimap<void*, ffi::Tensor> owners_;
- size_t page_size_;
- std::atomic<size_t> used_memory_;
- std::unordered_map<size_t, std::vector<Buffer>> memory_pool_;
- std::recursive_mutex mu_;
-};
-
-} // namespace memory
-} // namespace runtime
-} // namespace tvm
-
-#endif // TVM_RUNTIME_MEMORY_POOLED_ALLOCATOR_H_
diff --git a/src/runtime/tensor.cc b/src/runtime/tensor.cc
index f6bea42368..b04c13b47c 100644
--- a/src/runtime/tensor.cc
+++ b/src/runtime/tensor.cc
@@ -89,6 +89,17 @@ ffi::Tensor Empty(ffi::Shape shape, DLDataType dtype, Device
dev,
return DeviceAPI::Get(dev)->Empty(shape, dtype, dev, mem_scope);
}
+ffi::ObjectPtr<TensorMemoryPool> CreateTensorMemoryPool(DLDevice device) {
+ auto factory =
ffi::Function::GetGlobal(std::string("runtime.create_tensor_memory_pool.") +
+
DLDeviceType2Str(device.device_type));
+ if (factory.has_value()) {
+ auto pool = (*factory)(device).cast<ffi::ObjectPtr<TensorMemoryPool>>();
+ TVM_FFI_CHECK(pool != nullptr, ValueError) << "Tensor memory pool factory
returned null";
+ return pool;
+ }
+ return ffi::make_object<TensorMemoryPool>();
+}
+
ffi::Tensor CreateView(ffi::Tensor source, ffi::Shape shape, DLDataType dtype,
uint64_t relative_byte_offset) {
TVM_FFI_ICHECK(source.defined());
diff --git a/src/runtime/vm/builtin.cc b/src/runtime/vm/builtin.cc
index f8662dddbf..85ea9c8a79 100644
--- a/src/runtime/vm/builtin.cc
+++ b/src/runtime/vm/builtin.cc
@@ -28,7 +28,6 @@
#include <tvm/ffi/memory.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/memory/memory_manager.h>
#include <tvm/runtime/tensor.h>
#include <tvm/runtime/vm/builtin.h>
#include <tvm/runtime/vm/bytecode.h>
@@ -43,6 +42,15 @@ namespace vm {
using tvm::ffi::Tensor;
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::GlobalDef().def("vm.builtin.initialize_default", [](ffi::Module mod,
DLDevice dev) {
+ auto init = mod->GetFunction("vm_initialization");
+ TVM_FFI_CHECK(init.has_value(), ValueError) << "Module has no VM
initialization function";
+ (*init)(static_cast<int>(dev.device_type), dev.device_id,
static_cast<int>(kDLCPU), 0);
+ });
+}
+
//-------------------------------------------------
// Shape/Type handling.
//-------------------------------------------------
@@ -54,7 +62,7 @@ using tvm::ffi::Tensor;
*/
ffi::Tensor AllocShapeHeap(void* ctx_ptr, int64_t size) {
VirtualMachine* vm = static_cast<VirtualMachine*>(ctx_ptr);
- // use host allocator, which is always last element.
+ // Use the host pool, which is always the last element.
size_t host_device_index = vm->devices.size() - 1;
// specially handle hexagon on-device RT.
// TODO(relax-team): visit and consider other possible choices.
@@ -63,8 +71,8 @@ ffi::Tensor AllocShapeHeap(void* ctx_ptr, int64_t size) {
} else {
TVM_FFI_ICHECK_EQ(vm->devices[host_device_index].device_type, kDLCPU);
}
- auto* alloc = vm->allocators[host_device_index];
- return alloc->Empty({size}, DLDataType{kDLInt, 64, 1},
vm->devices[host_device_index]);
+ auto& pool = vm->memory_pools[host_device_index];
+ return pool->Empty({size}, DLDataType{kDLInt, 64, 1},
vm->devices[host_device_index]);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -377,24 +385,19 @@ TVM_FFI_STATIC_INIT_BLOCK() {
//-------------------------------------------------
// Storage management.
//-------------------------------------------------
-Storage VMAllocStorage(void* ctx_ptr, ffi::Shape buffer_shape, Index
device_index,
- DLDataType dtype_hint, ffi::String mem_scope) {
+ffi::Tensor VMAllocStorage(void* ctx_ptr, ffi::Shape buffer_shape, Index
device_index,
+ DLDataType dtype_hint, ffi::String mem_scope) {
VirtualMachine* vm = static_cast<VirtualMachine*>(ctx_ptr);
-
- TVM_FFI_ICHECK_LT(device_index, vm->devices.size())
- << "The device index is out of VM physical devices list";
-
if (device_index == -1) {
// Allocate on host. Host is always the last element of vm->devices.
device_index = vm->devices.size() - 1;
}
-
- auto* alloc = vm->allocators[device_index];
- TVM_FFI_ICHECK(alloc) << "Did you forget to init the VirtualMachine with
devices?";
-
- auto buffer = alloc->Alloc(vm->devices[device_index], buffer_shape,
dtype_hint, mem_scope);
-
- return Storage(buffer, alloc);
+ TVM_FFI_ICHECK_GE(device_index, 0);
+ TVM_FFI_ICHECK_LT(device_index, vm->devices.size())
+ << "The device index is out of VM physical devices list";
+ auto& pool = vm->memory_pools[device_index];
+ TVM_FFI_ICHECK(pool) << "Did you forget to init the VirtualMachine with
devices?";
+ return pool->Empty(buffer_shape, dtype_hint, vm->devices[device_index],
mem_scope);
}
TVM_FFI_STATIC_INIT_BLOCK() {
@@ -402,16 +405,25 @@ TVM_FFI_STATIC_INIT_BLOCK() {
refl::GlobalDef()
.def("vm.builtin.alloc_storage", VMAllocStorage)
.def_packed("vm.builtin.alloc_tensor", [](ffi::PackedArgs args,
ffi::Any* rv) {
- Storage sobj = args[0].cast<Storage>();
- int64_t offset = args[1].cast<int64_t>();
- ffi::Shape shape = args[2].cast<ffi::Shape>();
- DLDataType dtype = args[3].cast<DLDataType>();
- if (args.size() == 5) {
- ffi::String scope = args[4].cast<ffi::String>();
- *rv = sobj->AllocTensorScoped(offset, shape, dtype, scope);
- } else {
- *rv = sobj->AllocTensor(offset, shape, dtype);
+ VirtualMachine* vm = VirtualMachine::GetContextPtr(args[0]);
+ ffi::Tensor backing = args[1].cast<ffi::Tensor>();
+ int64_t offset = args[2].cast<int64_t>();
+ ffi::Shape shape = args[3].cast<ffi::Shape>();
+ DLDataType dtype = args[4].cast<DLDataType>();
+ ffi::Optional<ffi::String> scope = std::nullopt;
+ if (args.size() == 6) {
+ scope = args[5].cast<ffi::String>();
+ }
+ TVM_FFI_CHECK_GE(offset, 0, ValueError) << "Tensor byte offset must be
nonnegative";
+ DLDevice device = backing.device();
+ for (size_t i = 0; i < vm->devices.size(); ++i) {
+ if (vm->devices[i].device_type == device.device_type &&
+ vm->devices[i].device_id == device.device_id) {
+ *rv = vm->memory_pools[i]->CreateView(std::move(backing), shape,
dtype, offset, scope);
+ return;
+ }
}
+ TVM_FFI_THROW(ValueError) << "Tensor backing device is not initialized
in the VM";
});
}
diff --git a/src/runtime/vm/lm_support.cc b/src/runtime/vm/lm_support.cc
index 79d81a61f7..38286f60a2 100644
--- a/src/runtime/vm/lm_support.cc
+++ b/src/runtime/vm/lm_support.cc
@@ -42,7 +42,6 @@
#include <tvm/ffi/memory.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
-#include <tvm/runtime/memory/memory_manager.h>
#include <tvm/runtime/tensor.h>
#include <tvm/runtime/vm/vm.h>
diff --git a/src/runtime/vm/paged_kv_cache.cc b/src/runtime/vm/paged_kv_cache.cc
index 59d6192bbf..836a90c948 100644
--- a/src/runtime/vm/paged_kv_cache.cc
+++ b/src/runtime/vm/paged_kv_cache.cc
@@ -28,7 +28,6 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/device_api.h>
#include <tvm/runtime/disco/disco_worker.h>
-#include <tvm/runtime/memory/memory_manager.h>
#include <tvm/runtime/tensor.h>
#include <tvm/support/cuda/nvtx.h>
diff --git a/src/runtime/vm/vm.cc b/src/runtime/vm/vm.cc
index 751762bc79..12aa482c18 100644
--- a/src/runtime/vm/vm.cc
+++ b/src/runtime/vm/vm.cc
@@ -25,7 +25,6 @@
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/function.h>
#include <tvm/runtime/logging.h>
-#include <tvm/runtime/memory/memory_manager.h>
#include <tvm/runtime/vm/vm.h>
#include <tvm/support/cuda/nvtx.h>
@@ -87,27 +86,27 @@ ffi::Any IndexIntoNestedObject(ffi::Any obj,
ffi::PackedArgs args, int starting_
return obj;
}
-ffi::Tensor ConvertTensorToDevice(ffi::Tensor src, const DLDevice& dev,
Allocator* alloc,
+ffi::Tensor ConvertTensorToDevice(ffi::Tensor src, const DLDevice& dev,
TensorMemoryPool* pool,
ffi::String scope = "global") {
TVM_FFI_ICHECK(src.defined());
if (src.device().device_type == dev.device_type && src.device().device_id ==
dev.device_id) {
return src;
} else {
- auto res = alloc->Empty(src.shape(), src.dtype(), dev, scope);
+ auto res = pool->Empty(src.shape(), src.dtype(), dev, scope);
tvm::runtime::CopyFromTo(src, res);
return res;
}
}
-Any ConvertObjectToDevice(Any src, const Device& dev, Allocator* alloc,
+Any ConvertObjectToDevice(Any src, const Device& dev, TensorMemoryPool* pool,
ffi::String scope = "global") {
if (src.as<ffi::Tensor::ContainerType>()) {
- return ConvertTensorToDevice(src.cast<ffi::Tensor>(), dev, alloc, scope);
+ return ConvertTensorToDevice(src.cast<ffi::Tensor>(), dev, pool, scope);
} else if (src.as<ffi::ArrayObj>()) {
std::vector<Any> ret;
auto arr = src.cast<ffi::Array<Any>>();
for (size_t i = 0; i < arr.size(); i++) {
- ret.push_back(ConvertObjectToDevice(arr[i], dev, alloc));
+ ret.push_back(ConvertObjectToDevice(arr[i], dev, pool));
}
return ffi::Array<Any>(ret.begin(), ret.end());
} else {
@@ -115,18 +114,18 @@ Any ConvertObjectToDevice(Any src, const Device& dev,
Allocator* alloc,
}
}
-ffi::Any ConvertArgToDevice(ffi::AnyView input, Device dev, Allocator* alloc) {
+ffi::Any ConvertArgToDevice(ffi::AnyView input, Device dev, TensorMemoryPool*
pool) {
// in terms of memory-behavior.
// To be extra careful, we copy DLTensor.
// The developer can still explicitly allocate Tensor
// in TVM Native API or Tensor::FromDLPack to regain zero copy behavior.
ffi::Any ret;
if (auto opt_obj = input.as<ffi::ObjectRef>()) {
- ret = ConvertObjectToDevice(opt_obj.value(), dev, alloc);
+ ret = ConvertObjectToDevice(opt_obj.value(), dev, pool);
} else if (auto opt_dltensor = input.as<DLTensor*>()) {
DLTensor* tensor = opt_dltensor.value();
- std::vector<int64_t> shape(tensor->shape, tensor->shape + tensor->ndim);
- auto dst = alloc->Empty(shape, tensor->dtype, dev);
+ ffi::ShapeView shape(tensor->shape, tensor->ndim);
+ auto dst = pool->Empty(shape, tensor->dtype, dev);
tvm::runtime::CopyFromTo(tensor, dst);
ret = dst;
} else {
@@ -135,11 +134,11 @@ ffi::Any ConvertArgToDevice(ffi::AnyView input, Device
dev, Allocator* alloc) {
return ret;
}
-ffi::Any ConvertRegToDevice(ffi::Any input, Device dev, Allocator* alloc,
+ffi::Any ConvertRegToDevice(ffi::Any input, Device dev, TensorMemoryPool* pool,
ffi::String scope = "global") {
ffi::Any ret;
if (auto opt_obj = input.as<ffi::ObjectRef>()) {
- ret = ConvertObjectToDevice(opt_obj.value(), dev, alloc, scope);
+ ret = ConvertObjectToDevice(opt_obj.value(), dev, pool, scope);
} else {
ret = input;
}
@@ -194,8 +193,7 @@ class VirtualMachineImpl : public VirtualMachine {
// Public facing functions overloading
//---------------------------------------------------
void LoadExecutable(ffi::ObjectPtr<VMExecutable> exec) final;
- void Init(const std::vector<Device>& devices,
- const std::vector<AllocatorType>& alloc_types) final;
+ void Init(const std::vector<Device>& devices) final;
VMClosure GetClosure(const ffi::String& func_name) final {
return this->GetClosureInternal(func_name, false).value();
}
@@ -463,23 +461,19 @@ void
VirtualMachineImpl::LoadExecutable(ffi::ObjectPtr<VMExecutable> exec) {
this->imports_ = exec->imports();
}
-void VirtualMachineImpl::Init(const std::vector<Device>& devices,
- const std::vector<AllocatorType>& alloc_types) {
- TVM_FFI_ICHECK_EQ(devices.size(), alloc_types.size());
-
+void VirtualMachineImpl::Init(const std::vector<Device>& devices) {
this->devices.reserve(devices.size());
- this->allocators.reserve(alloc_types.size());
- for (size_t i = 0; i < devices.size(); i++) {
- auto alloc = MemoryManager::GetOrCreateAllocator(devices[i],
alloc_types[i]);
- this->devices.push_back(devices[i]);
- this->allocators.push_back(alloc);
+ this->memory_pools.reserve(devices.size());
+ for (Device device : devices) {
+ this->devices.push_back(device);
+ this->memory_pools.push_back(CreateTensorMemoryPool(device));
}
// Setup constant sections.
this->const_pool_.reserve(exec_->constants.size());
for (size_t i = 0; i < exec_->constants.size(); ++i) {
if (auto opt_nd = exec_->constants[i].as<ffi::Tensor>(); opt_nd &&
opt_nd->defined()) {
- this->const_pool_.push_back(
- ConvertRegToDevice(opt_nd.value(), devices[0], allocators[0],
exec_->memory_scopes[i]));
+ this->const_pool_.push_back(ConvertRegToDevice(
+ opt_nd.value(), devices[0], memory_pools[0].get(),
exec_->memory_scopes[i]));
} else {
this->const_pool_.push_back(exec_->constants[i]);
}
@@ -519,7 +513,7 @@ void VirtualMachineImpl::SetInput(std::string func_name,
bool with_param_module,
// call param func to get the arguments(usually corresponds to param
pack.)
func_args[i] =
(args[i].cast<ffi::Module>())->GetFunction("get_params").value()();
} else {
- func_args[i] = ConvertArgToDevice(args[i], devices[0], allocators[0]);
+ func_args[i] = ConvertArgToDevice(args[i], devices[0],
memory_pools[0].get());
}
}
inputs_[func_name] = func_args;
@@ -585,7 +579,7 @@ void VirtualMachineImpl::SaveClosure(const ffi::String&
func_name, const ffi::St
VMClosure clo = this->GetClosure(func_name);
std::vector<RegType> inputs(args.size());
for (int i = 0; i < args.size(); ++i) {
- inputs[i] = ConvertArgToDevice(args[i], this->devices[0],
this->allocators[0]);
+ inputs[i] = ConvertArgToDevice(args[i], this->devices[0],
this->memory_pools[0].get());
}
ffi::Function impl = VMClosure::BindLastArgs(clo->impl, inputs);
if (!include_return) {
@@ -837,17 +831,14 @@ ffi::ObjectPtr<VirtualMachine> VirtualMachine::Create() {
//--------------------------------------------------------------------
void VirtualMachineImpl::_Init(ffi::PackedArgs args, ffi::Any* rv) {
- TVM_FFI_ICHECK_EQ(args.size() % 3, 0);
+ TVM_FFI_ICHECK_EQ(args.size() % 2, 0);
std::vector<Device> devices;
- std::vector<AllocatorType> alloc_types;
- for (int i = 0; i < args.size(); i += 3) {
+ for (int i = 0; i < args.size(); i += 2) {
int device_type = args[i].cast<int>();
int device_id = args[i + 1].cast<int>();
- int alloc_type = args[i + 2].cast<int>();
devices.push_back(Device{DLDeviceType(device_type), device_id});
- alloc_types.push_back(AllocatorType(alloc_type));
}
- this->Init(devices, alloc_types);
+ this->Init(devices);
}
void VirtualMachineImpl::_SaveClosure(ffi::PackedArgs args, ffi::Any* rv) {
diff --git a/src/runtime/workspace_pool.cc b/src/runtime/workspace_pool.cc
index c0ba44f647..1549d3cbef 100644
--- a/src/runtime/workspace_pool.cc
+++ b/src/runtime/workspace_pool.cc
@@ -24,33 +24,11 @@
#include "workspace_pool.h"
#include <algorithm>
-#include <array>
-#include <limits>
#include <memory>
namespace tvm {
namespace runtime {
-namespace {
-struct WorkspacePools {
- std::mutex mutex;
- std::vector<WorkspacePool*> pools;
-};
-WorkspacePools* AllWorkspacePools() {
- // Kept alive through thread-local pool destruction at process exit.
- static auto* registry = new WorkspacePools();
- return registry;
-}
-} // namespace
-
-void ResetWorkspacePool(DLDeviceType device_type) {
- auto* registry = AllWorkspacePools();
- std::lock_guard<std::mutex> lock(registry->mutex);
- for (auto* pool : registry->pools) {
- if (pool->device_type_ == device_type) pool->Reset();
- }
-}
-
// page size.
constexpr size_t kWorkspacePageSize = 4 << 10;
@@ -66,29 +44,25 @@ class WorkspacePool::Pool {
allocated_.push_back(e);
}
// allocate from pool
- void* Alloc(Device dev, DeviceAPI* device, size_t nbytes) {
- TVM_FFI_CHECK_LE(nbytes, std::numeric_limits<size_t>::max() -
kWorkspacePageSize, ValueError)
- << "Workspace size overflow";
+ void* Alloc(Device dev, DeviceAPI* device, Allocator alloc, size_t nbytes) {
// Allocate align to page.
nbytes = (nbytes + (kWorkspacePageSize - 1)) / kWorkspacePageSize *
kWorkspacePageSize;
if (nbytes == 0) nbytes = kWorkspacePageSize;
Entry e;
auto allocate = [&]() {
- TVM_FFI_CHECK_LE(nbytes / kTempAllocaAlignment,
std::numeric_limits<int64_t>::max(),
- ValueError)
- << "Workspace shape overflow";
- int64_t shape = static_cast<int64_t>(nbytes / kTempAllocaAlignment +
- (nbytes % kTempAllocaAlignment !=
0));
- DLDataType dtype{kDLUInt, 8, kTempAllocaAlignment};
- return device->Empty(ffi::ShapeView(&shape, 1), dtype, dev);
+ if (alloc) return alloc(device, dev, nbytes);
+ // Preserve the workspace alignment when adapting raw allocations to
Empty.
+ int64_t size = static_cast<int64_t>(nbytes / kTempAllocaAlignment +
+ (nbytes % kTempAllocaAlignment !=
0));
+ return device->Empty(ffi::ShapeView(&size, 1), DLDataType{kDLUInt, 8,
kTempAllocaAlignment},
+ dev);
};
if (free_list_.size() == 2) {
e = free_list_.back();
free_list_.pop_back();
if (e.size < nbytes) {
- // Release the cached page before growing it.
+ // resize the page
e.owner = ffi::Tensor();
- e.data = nullptr;
e.owner = allocate();
e.data = e.owner.data_ptr();
e.size = nbytes;
@@ -110,19 +84,16 @@ class WorkspacePool::Pool {
e = free_list_.back();
free_list_.pop_back();
e.owner = ffi::Tensor();
- e.data = nullptr;
e.owner = allocate();
e.data = e.owner.data_ptr();
e.size = nbytes;
}
}
- e.generation = generation_;
allocated_.push_back(e);
return e.data;
}
// free resource back to pool
void Free(void* data) {
- TVM_FFI_CHECK(data != nullptr, ValueError) << "Cannot free a null
workspace";
Entry e;
if (allocated_.back().data == data) {
// quick path, last allocated.
@@ -136,7 +107,8 @@ class WorkspacePool::Pool {
e = allocated_[index];
allocated_.erase(allocated_.begin() + index);
}
- if (e.generation != generation_) return;
+ // Clear marks outstanding allocations so they cannot enter a replacement
session's cache.
+ if (e.size == 0) return;
if (free_list_.back().size < e.size) {
free_list_.push_back(e);
} else if (free_list_.size() == 2) {
@@ -151,77 +123,64 @@ class WorkspacePool::Pool {
free_list_[i + 1] = e;
}
}
- void Reset() {
- // Live allocations keep their original owners and are discarded when
freed.
- ++generation_;
- free_list_.clear();
- free_list_.emplace_back();
+ void Clear() {
+ free_list_.resize(1);
+ for (auto& entry : allocated_) entry.size = 0;
}
- // Release all resources
- void Release() {
- free_list_.clear();
- allocated_.clear();
+
+ bool Contains(void* data) const {
+ return std::any_of(allocated_.begin() + 1, allocated_.end(),
+ [data](const Entry& entry) { return entry.data == data;
});
}
private:
/*! \brief a single entry in the pool */
struct Entry {
- uint64_t generation{0};
ffi::Tensor owner;
- void* data{nullptr};
- size_t size{0};
+ void* data;
+ size_t size;
};
- uint64_t generation_{0};
/*! \brief List of free items, sorted from small to big size */
std::vector<Entry> free_list_;
/*! \brief List of allocated items */
std::vector<Entry> allocated_;
};
-WorkspacePool::WorkspacePool(DLDeviceType device_type, DeviceAPI* device)
- : device_type_(device_type), device_(device) {
- auto* registry = AllWorkspacePools();
- std::lock_guard<std::mutex> lock(registry->mutex);
- registry->pools.push_back(this);
-}
+WorkspacePool::WorkspacePool(DLDeviceType device_type, DeviceAPI* device,
Allocator alloc)
+ : device_type_(device_type), device_(device), alloc_(alloc) {}
WorkspacePool::~WorkspacePool() {
- auto* registry = AllWorkspacePools();
- std::lock_guard<std::mutex> lock(registry->mutex);
- auto& pools = registry->pools;
- pools.erase(std::remove(pools.begin(), pools.end(), this), pools.end());
- for (size_t i = 0; i < array_.size(); ++i) {
- if (array_[i] != nullptr) {
- array_[i]->Release();
- delete array_[i];
- }
+ for (auto& streams : array_) {
+ for (auto& [stream, pool] : streams) delete pool;
}
}
-void WorkspacePool::Reset() {
- std::lock_guard<std::mutex> lock(mutex_);
- for (auto* pool : array_) {
- if (pool) pool->Reset();
+void WorkspacePool::Clear() {
+ for (auto& streams : array_) {
+ for (auto& [stream, pool] : streams) pool->Clear();
}
}
void* WorkspacePool::AllocWorkspace(Device dev, size_t size) {
- TVM_FFI_CHECK_GE(dev.device_id, 0, ValueError) << "Invalid workspace device
index";
- std::lock_guard<std::mutex> lock(mutex_);
if (static_cast<size_t>(dev.device_id) >= array_.size()) {
- array_.resize(dev.device_id + 1, nullptr);
+ array_.resize(dev.device_id + 1);
}
- if (array_[dev.device_id] == nullptr) {
- array_[dev.device_id] = new Pool();
- }
- return array_[dev.device_id]->Alloc(dev, device_, size);
+ // Keep the historical size-sorted cache separate for each execution stream.
+ auto& pool = array_[dev.device_id][device_->GetCurrentStream(dev)];
+ if (pool == nullptr) pool = new Pool();
+ return pool->Alloc(dev, device_, alloc_, size);
}
void WorkspacePool::FreeWorkspace(Device dev, void* ptr) {
- std::lock_guard<std::mutex> lock(mutex_);
- TVM_FFI_ICHECK(static_cast<size_t>(dev.device_id) < array_.size() &&
- array_[dev.device_id] != nullptr);
- array_[dev.device_id]->Free(ptr);
+ TVM_FFI_ICHECK(static_cast<size_t>(dev.device_id) < array_.size());
+ // A caller may change streams between allocation and release.
+ for (auto& [stream, pool] : array_[dev.device_id]) {
+ if (pool->Contains(ptr)) {
+ pool->Free(ptr);
+ return;
+ }
+ }
+ TVM_FFI_THROW(InternalError) << "trying to free things that has not been
allocated";
}
} // namespace runtime
diff --git a/src/runtime/workspace_pool.h b/src/runtime/workspace_pool.h
index 27e9005d8a..f72e0089af 100644
--- a/src/runtime/workspace_pool.h
+++ b/src/runtime/workspace_pool.h
@@ -27,14 +27,11 @@
#include <tvm/runtime/device_api.h>
#include <memory>
-#include <mutex>
+#include <unordered_map>
#include <vector>
namespace tvm {
namespace runtime {
-// Invalidate cached workspace owners when a backend tears down a resource
session.
-TVM_RUNTIME_DLL void ResetWorkspacePool(DLDeviceType device_type);
-
/*!
* \brief A workspace pool to manage
*
@@ -46,14 +43,16 @@ TVM_RUNTIME_DLL void ResetWorkspacePool(DLDeviceType
device_type);
* - The release order is usually in reverse order of allocate
* - Repeative pattern of same allocations over different runs.
*/
-class WorkspacePool {
+class TVM_RUNTIME_DLL WorkspacePool {
public:
/*!
* \brief Create pool with specific device type and device.
* \param device_type The device type.
* \param device_api The device API.
+ * \param alloc Optional backend allocation for unpooled owning workspace
tensors.
*/
- WorkspacePool(DLDeviceType device_type, DeviceAPI* device_api);
+ using Allocator = ffi::Tensor (*)(DeviceAPI*, Device, size_t);
+ WorkspacePool(DLDeviceType device_type, DeviceAPI* device_api, Allocator
alloc = nullptr);
/*! \brief destructor */
~WorkspacePool();
/*!
@@ -69,18 +68,18 @@ class WorkspacePool {
* \param ptr The pointer to be freed.
*/
void FreeWorkspace(Device dev, void* ptr);
+ /*! \brief Release cached allocations; outstanding allocations remain valid
until freed. */
+ void Clear();
private:
class Pool;
- friend void ResetWorkspacePool(DLDeviceType device_type);
- void Reset();
- std::mutex mutex_;
/*! \brief pool of device local array */
- std::vector<Pool*> array_;
+ std::vector<std::unordered_map<TVMStreamHandle, Pool*>> array_;
/*! \brief device type this pool support */
DLDeviceType device_type_;
/*! \brief The device API */
DeviceAPI* device_;
+ Allocator alloc_;
};
} // namespace runtime
diff --git a/tests/python/disco/test_custom_allreduce.py
b/tests/python/disco/test_custom_allreduce.py
index db23257703..255571e20f 100644
--- a/tests/python/disco/test_custom_allreduce.py
+++ b/tests/python/disco/test_custom_allreduce.py
@@ -64,10 +64,8 @@ def test_allreduce(shape, ccl, strategy):
num_elements = reduce(lambda x, y: x * y, shape)
dtype = "float32"
falloc_ipc_storage =
sess.get_global_func("runtime.disco.cuda_ipc.alloc_storage")
- falloc_tensor = sess.get_global_func("vm.builtin.alloc_tensor")
fallreduce =
sess.get_global_func("runtime.disco.cuda_ipc.custom_allreduce")
- d_storage = sess.call_packed(falloc_ipc_storage, Shape(shape),
DataType(dtype))
- d_input = sess.call_packed(falloc_tensor, d_storage, 0, Shape(shape),
DataType(dtype))
+ d_input = sess.call_packed(falloc_ipc_storage, Shape(shape),
DataType(dtype))
array_1 = np.arange(num_elements, dtype="float32").reshape(*shape)
array_2 = np.arange(start=1, stop=-(num_elements - 1), step=-1,
dtype="float32").reshape(*shape)
diff --git a/tests/python/relax/test_vm_alloc_storage_with_scope.py
b/tests/python/relax/test_vm_alloc_storage_with_scope.py
index 4b929c8fe1..a7ceb5a4bf 100644
--- a/tests/python/relax/test_vm_alloc_storage_with_scope.py
+++ b/tests/python/relax/test_vm_alloc_storage_with_scope.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Test Naive allocator with memory scope for Relax VM"""
+"""Test tensor allocation with memory scope for Relax VM"""
import numpy as np
@@ -66,8 +66,7 @@ def test_alloc_storage_with_scope_global():
lib = tvm.relax.build(mod, target=target)
dev = tvm.cpu()
- # This is the important line which tests nd allocator
- vm_rt = relax.VirtualMachine(lib, dev, memory_cfg="naive")
+ vm_rt = relax.VirtualMachine(lib, dev)
x = tvm.runtime.tensor(arg0, dev)
vm_rt.set_input("main", x)
vm_rt.invoke_stateful("main")
diff --git a/tests/python/relax/test_vm_execbuilder.py
b/tests/python/relax/test_vm_execbuilder.py
index 35ac921956..785aea568c 100644
--- a/tests/python/relax/test_vm_execbuilder.py
+++ b/tests/python/relax/test_vm_execbuilder.py
@@ -183,7 +183,9 @@ def test_vm_storage():
dst=ib.r(1),
)
ib.emit_call(
- "vm.builtin.alloc_tensor", args=[ib.r(1), ib.imm(0), shape,
dtype], dst=ib.r(2)
+ "vm.builtin.alloc_tensor",
+ args=[ib.vm_state(), ib.r(1), ib.imm(0), shape, dtype],
+ dst=ib.r(2),
)
ib.emit_ret(ib.r(2))
ex = ib.get()
diff --git a/web/emcc/wasm_runtime.cc b/web/emcc/wasm_runtime.cc
index b1eaf6c32f..6d3c1b7bb0 100644
--- a/web/emcc/wasm_runtime.cc
+++ b/web/emcc/wasm_runtime.cc
@@ -67,7 +67,6 @@
#include "3rdparty/tvm-ffi/src/ffi/object.cc"
#include "3rdparty/tvm-ffi/src/ffi/tensor.cc"
#include "3rdparty/tvm-ffi/src/ffi/testing/testing.cc"
-#include "src/runtime/memory/memory_manager.cc"
#include "src/runtime/vm/attn_backend.cc"
#include "src/runtime/vm/builtin.cc"
#include "src/runtime/vm/bytecode.cc"
diff --git a/web/src/runtime.ts b/web/src/runtime.ts
index d615f2bbbd..1f3b26a74b 100644
--- a/web/src/runtime.ts
+++ b/web/src/runtime.ts
@@ -833,11 +833,6 @@ export class TVMArray extends TVMObject {
}
}
-export enum VMAllocatorKind {
- NAIVE_ALLOCATOR = 1,
- POOLED_ALLOCATOR = 2,
-}
-
/**
* VirtualMachine Executor.
*
@@ -857,11 +852,9 @@ export class VirtualMachine implements Disposable {
this.mod.getFunction("vm_initialization")(
new Scalar(device.deviceType, "int"),
new Scalar(device.deviceId, "int"),
- new Scalar(VMAllocatorKind.POOLED_ALLOCATOR, "int"),
// explicitly specify host device type
new Scalar(DeviceStrToEnum.cpu, "int"),
new Scalar(0, "int"),
- new Scalar(VMAllocatorKind.POOLED_ALLOCATOR, "int"),
);
}
