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-tirx.git
The following commit(s) were added to refs/heads/main by this push:
new 2fa19e5a Hide internal symbols in CUDA and macOS libraries (#16)
2fa19e5a is described below
commit 2fa19e5a4621e4c8bd98d103c8fcc4983db93d9e
Author: Tianqi Chen <[email protected]>
AuthorDate: Wed Oct 21 17:04:35 2026 -0400
Hide internal symbols in CUDA and macOS libraries (#16)
Keep backend implementation symbols private across CUDA and macOS
libraries. WorkspacePool becomes a private header-only utility with
inline methods, preserving allocation and cache behavior while removing
backend imports of its implementation. Internal CUDA memory/device query
helpers retain their FFI registrations without DLL exports.
An annotated TIRx CMake helper applies archive visibility at each final
shared-library link. On Darwin it uses native per-library overrides for
directly linked archive paths and static/unknown imported targets,
preserving shared/interface/object dependencies. Non-Darwin targets
retain the existing FFI policy. Static build targets remain available
internally, but the static runtime archive is omitted from installation
and wheels. The accepted FFI dependency pin is unchanged.
---
CMakeLists.txt | 4 +-
cmake/modules/CUDA.cmake | 4 +-
cmake/utils/Library.cmake | 28 +++++
src/backend/cuda/runtime/cuda_device_api.cc | 4 +-
src/runtime/workspace_pool.cc | 187 ----------------------------
src/runtime/workspace_pool.h | 157 ++++++++++++++++++++++-
6 files changed, 190 insertions(+), 194 deletions(-)
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 6e81ddcd..70d6f70b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -267,7 +267,7 @@ foreach(_target tvm_runtime tvm_tirx_compiler
${TVM_RUNTIME_BACKEND_LIBS})
elseif(UNIX)
target_link_options(${_target} PRIVATE "-Wl,--no-undefined")
endif()
- tvm_ffi_hide_static_linked_lib_symbols(${_target})
+ tvm_tirx_hide_static_linked_lib_symbols(${_target})
tvm_ffi_add_apple_dsymutil(${_target})
endforeach()
@@ -278,7 +278,7 @@ install(DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/include/"
DESTINATION include COMPONENT tvm_tirx FILES_MATCHING PATTERN "*.h")
# Libraries
-install(TARGETS tvm_tirx_header tvm_tirx_compiler tvm_runtime
tvm_runtime_static
+install(TARGETS tvm_tirx_header tvm_tirx_compiler tvm_runtime
ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT tvm_tirx
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT tvm_tirx
RUNTIME DESTINATION ${CMAKE_INSTALL_LIBDIR} COMPONENT tvm_tirx)
diff --git a/cmake/modules/CUDA.cmake b/cmake/modules/CUDA.cmake
index c9852709..b623e39b 100644
--- a/cmake/modules/CUDA.cmake
+++ b/cmake/modules/CUDA.cmake
@@ -80,8 +80,8 @@ if(TVM_TIRX_USE_CUDA)
add_library(tvm_runtime_cuda_objs OBJECT ${RUNTIME_CUDA_SRCS})
target_link_libraries(tvm_runtime_cuda_objs PUBLIC tvm_tirx_header)
# These sources compile into tvm_runtime_cuda.dll, so their TVM_RUNTIME_DLL /
- # TVM_FFI_DLL symbols must be dllexport on MSVC (e.g. GetCudaDeviceCount in
- # cuda_device_api.cc). Mirror tvm_runtime_objs; a no-op on non-MSVC
platforms.
+ # TVM_FFI_DLL symbols must be dllexport on MSVC. Mirror tvm_runtime_objs;
+ # a no-op on non-MSVC platforms.
target_compile_definitions(tvm_runtime_cuda_objs PRIVATE TVM_RUNTIME_EXPORTS
TVM_FFI_EXPORTS)
set_target_properties(tvm_runtime_cuda_objs PROPERTIES
POSITION_INDEPENDENT_CODE ON)
add_library(tvm_runtime_cuda SHARED $<TARGET_OBJECTS:tvm_runtime_cuda_objs>)
diff --git a/cmake/utils/Library.cmake b/cmake/utils/Library.cmake
index a37a8a97..506bf96f 100644
--- a/cmake/utils/Library.cmake
+++ b/cmake/utils/Library.cmake
@@ -33,3 +33,31 @@ function(tvm_tirx_configure_target target)
set_target_properties(${target} PROPERTIES BUILD_RPATH "\$ORIGIN"
INSTALL_RPATH "\$ORIGIN")
endif()
endfunction()
+
+# Hide archive dependencies in the final shared library, without changing the
archives themselves.
+function(tvm_tirx_hide_static_linked_lib_symbols target_name)
+ if(NOT APPLE)
+ tvm_ffi_hide_static_linked_lib_symbols(${target_name})
+ return()
+ endif()
+
+ # CMake evaluates link features in the caller's directory scope, after this
function returns.
+ set(CMAKE_LINK_LIBRARY_USING_tvm_hidden "LINKER:-load_hidden,<LIBRARY>"
PARENT_SCOPE)
+ set(CMAKE_LINK_LIBRARY_USING_tvm_hidden_SUPPORTED TRUE PARENT_SCOPE)
+ get_target_property(_libs ${target_name} LINK_LIBRARIES)
+ foreach(_lib IN LISTS _libs)
+ if(TARGET ${_lib})
+ # Shared, interface and object targets do not provide dependency
archives to hide.
+ get_target_property(_type ${_lib} TYPE)
+ if(_type STREQUAL "STATIC_LIBRARY" OR _type STREQUAL "UNKNOWN_LIBRARY")
+ # UNKNOWN imports may resolve to a shared library; inspect the
selected configuration's file.
+ set(_is_archive
"$<STREQUAL:$<PATH:GET_EXTENSION,LAST_ONLY,$<TARGET_FILE:${_lib}>>,.a>")
+ set_property(TARGET ${target_name} PROPERTY
LINK_LIBRARY_OVERRIDE_${_lib}
+ "$<IF:${_is_archive},tvm_hidden,DEFAULT>")
+ endif()
+ elseif(_lib MATCHES "[.]a$")
+ # llvm-config also supplies archive paths directly, without imported
CMake targets.
+ set_property(TARGET ${target_name} PROPERTY
LINK_LIBRARY_OVERRIDE_${_lib} tvm_hidden)
+ endif()
+ endforeach()
+endfunction()
diff --git a/src/backend/cuda/runtime/cuda_device_api.cc
b/src/backend/cuda/runtime/cuda_device_api.cc
index c72d56b7..745d10be 100644
--- a/src/backend/cuda/runtime/cuda_device_api.cc
+++ b/src/backend/cuda/runtime/cuda_device_api.cc
@@ -387,7 +387,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_RUNTIME_DLL ffi::String GetCudaFreeMemory() {
+ffi::String GetCudaFreeMemory() {
size_t free_mem, total_mem;
TVM_FFI_CHECK_CUDA_ERROR(cudaMemGetInfo(&free_mem, &total_mem));
std::stringstream ss;
@@ -409,7 +409,7 @@ TVM_FFI_STATIC_INIT_BLOCK() {
});
}
-TVM_RUNTIME_DLL int GetCudaDeviceCount() {
+int GetCudaDeviceCount() {
int count;
TVM_FFI_CHECK_CUDA_ERROR(cudaGetDeviceCount(&count));
return count;
diff --git a/src/runtime/workspace_pool.cc b/src/runtime/workspace_pool.cc
deleted file mode 100644
index 5a8ca217..00000000
--- a/src/runtime/workspace_pool.cc
+++ /dev/null
@@ -1,187 +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 workspace_pool.h
- * \brief Workspace pool utility.
- */
-#include "workspace_pool.h"
-
-#include <algorithm>
-#include <memory>
-
-namespace tvm {
-namespace runtime {
-
-// page size.
-constexpr size_t kWorkspacePageSize = 4 << 10;
-
-class WorkspacePool::Pool {
- public:
- // constructor
- Pool() {
- // safe guard header on each list.
- Entry e;
- e.data = nullptr;
- e.size = 0;
- free_list_.push_back(e);
- allocated_.push_back(e);
- }
- // allocate from pool
- 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 = [&]() {
- 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) {
- // resize the page
- e.owner = ffi::Tensor();
- e.owner = allocate();
- e.data = e.owner.data_ptr();
- e.size = nbytes;
- }
- } else if (free_list_.size() == 1) {
- e.owner = allocate();
- e.data = e.owner.data_ptr();
- e.size = nbytes;
- } else {
- if (free_list_.back().size >= nbytes) {
- // find smallest fit
- auto it = free_list_.end() - 2;
- for (; it->size >= nbytes; --it) {
- }
- e = *(it + 1);
- free_list_.erase(it + 1);
- } else {
- // resize the page
- e = free_list_.back();
- free_list_.pop_back();
- e.owner = ffi::Tensor();
- e.owner = allocate();
- e.data = e.owner.data_ptr();
- e.size = nbytes;
- }
- }
- allocated_.push_back(e);
- return e.data;
- }
- // free resource back to pool
- void Free(void* data) {
- Entry e;
- if (allocated_.back().data == data) {
- // quick path, last allocated.
- e = allocated_.back();
- allocated_.pop_back();
- } else {
- int index = static_cast<int>(allocated_.size()) - 2;
- for (; index > 0 && allocated_[index].data != data; --index) {
- }
- TVM_FFI_ICHECK_GT(index, 0) << "trying to free things that has not been
allocated";
- e = allocated_[index];
- allocated_.erase(allocated_.begin() + index);
- }
- // 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) {
- free_list_.push_back(free_list_.back());
- free_list_[1] = e;
- } else {
- size_t i = free_list_.size() - 1;
- free_list_.resize(free_list_.size() + 1);
- for (; e.size < free_list_[i].size; --i) {
- free_list_[i + 1] = free_list_[i];
- }
- free_list_[i + 1] = e;
- }
- }
- void Clear() {
- free_list_.resize(1);
- for (auto& entry : allocated_) entry.size = 0;
- }
-
- 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 {
- ffi::Tensor owner;
- void* data;
- size_t size;
- };
- /*! \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, DeviceAPI* device, Allocator alloc)
- : device_(device), alloc_(alloc) {}
-
-WorkspacePool::~WorkspacePool() {
- for (auto& streams : array_) {
- for (auto& [stream, pool] : streams) delete pool;
- }
-}
-
-void WorkspacePool::Clear() {
- for (auto& streams : array_) {
- for (auto& [stream, pool] : streams) pool->Clear();
- }
-}
-
-void* WorkspacePool::AllocWorkspace(Device dev, size_t size) {
- if (static_cast<size_t>(dev.device_id) >= array_.size()) {
- array_.resize(dev.device_id + 1);
- }
- // 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) {
- 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
-} // namespace tvm
diff --git a/src/runtime/workspace_pool.h b/src/runtime/workspace_pool.h
index e57d0897..0d63797e 100644
--- a/src/runtime/workspace_pool.h
+++ b/src/runtime/workspace_pool.h
@@ -26,6 +26,7 @@
#include <tvm/runtime/device_api.h>
+#include <algorithm>
#include <memory>
#include <unordered_map>
#include <vector>
@@ -43,7 +44,7 @@ namespace runtime {
* - The release order is usually in reverse order of allocate
* - Repeative pattern of same allocations over different runs.
*/
-class TVM_RUNTIME_DLL WorkspacePool {
+class WorkspacePool {
public:
/*!
* \brief Create pool with specific device type and device.
@@ -80,6 +81,160 @@ class TVM_RUNTIME_DLL WorkspacePool {
Allocator alloc_;
};
+class WorkspacePool::Pool {
+ public:
+ // constructor
+ Pool() {
+ // safe guard header on each list.
+ Entry e;
+ e.data = nullptr;
+ e.size = 0;
+ free_list_.push_back(e);
+ allocated_.push_back(e);
+ }
+ // allocate from pool
+ 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 = [&]() {
+ 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) {
+ // resize the page
+ e.owner = ffi::Tensor();
+ e.owner = allocate();
+ e.data = e.owner.data_ptr();
+ e.size = nbytes;
+ }
+ } else if (free_list_.size() == 1) {
+ e.owner = allocate();
+ e.data = e.owner.data_ptr();
+ e.size = nbytes;
+ } else {
+ if (free_list_.back().size >= nbytes) {
+ // find smallest fit
+ auto it = free_list_.end() - 2;
+ for (; it->size >= nbytes; --it) {
+ }
+ e = *(it + 1);
+ free_list_.erase(it + 1);
+ } else {
+ // resize the page
+ e = free_list_.back();
+ free_list_.pop_back();
+ e.owner = ffi::Tensor();
+ e.owner = allocate();
+ e.data = e.owner.data_ptr();
+ e.size = nbytes;
+ }
+ }
+ allocated_.push_back(e);
+ return e.data;
+ }
+ // free resource back to pool
+ void Free(void* data) {
+ Entry e;
+ if (allocated_.back().data == data) {
+ // quick path, last allocated.
+ e = allocated_.back();
+ allocated_.pop_back();
+ } else {
+ int index = static_cast<int>(allocated_.size()) - 2;
+ for (; index > 0 && allocated_[index].data != data; --index) {
+ }
+ TVM_FFI_ICHECK_GT(index, 0) << "trying to free things that has not been
allocated";
+ e = allocated_[index];
+ allocated_.erase(allocated_.begin() + index);
+ }
+ // 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) {
+ free_list_.push_back(free_list_.back());
+ free_list_[1] = e;
+ } else {
+ size_t i = free_list_.size() - 1;
+ free_list_.resize(free_list_.size() + 1);
+ for (; e.size < free_list_[i].size; --i) {
+ free_list_[i + 1] = free_list_[i];
+ }
+ free_list_[i + 1] = e;
+ }
+ }
+ void Clear() {
+ free_list_.resize(1);
+ for (auto& entry : allocated_) entry.size = 0;
+ }
+
+ bool Contains(void* data) const {
+ return std::any_of(allocated_.begin() + 1, allocated_.end(),
+ [data](const Entry& entry) { return entry.data == data;
});
+ }
+
+ private:
+ // Page size.
+ static constexpr size_t kWorkspacePageSize = 4 << 10;
+
+ /*! \brief a single entry in the pool */
+ struct Entry {
+ ffi::Tensor owner;
+ void* data;
+ size_t size;
+ };
+ /*! \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_;
+};
+
+inline WorkspacePool::WorkspacePool(DLDeviceType, DeviceAPI* device, Allocator
alloc)
+ : device_(device), alloc_(alloc) {}
+
+inline WorkspacePool::~WorkspacePool() {
+ for (auto& streams : array_) {
+ for (auto& [stream, pool] : streams) delete pool;
+ }
+}
+
+inline void WorkspacePool::Clear() {
+ for (auto& streams : array_) {
+ for (auto& [stream, pool] : streams) pool->Clear();
+ }
+}
+
+inline void* WorkspacePool::AllocWorkspace(Device dev, size_t size) {
+ if (static_cast<size_t>(dev.device_id) >= array_.size()) {
+ array_.resize(dev.device_id + 1);
+ }
+ // 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);
+}
+
+inline void WorkspacePool::FreeWorkspace(Device dev, void* 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
} // namespace tvm
#endif // TVM_RUNTIME_WORKSPACE_POOL_H_
