trending/tilelang
1
0
Fork 0
mirror of https://github.com/tile-ai/tilelang.git synced 2026-10-02 06:34:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Ascend] Lazy-load libascendcl via an ascendcl stub and gate the backend behind USE_ASCEND

Browse Source 
Formalize the ad-hoc dlopen loader in ascend_module.cc into a CUDA/ROCm-style
stub library (libstub_ascendcl.so): the runtime module now calls the ACL
entrypoints directly and the stub resolves them lazily, preferring the copy
already loaded by torch_npu (RTLD_DEFAULT/RTLD_NEXT), then LD_LIBRARY_PATH,
then well-known CANN install roots. aclGetRecentErrMsg degrades to nullptr so
error reporting never throws; a missing required symbol is reported by name.

Add USE_ASCEND to TILELANG_BACKENDS with env-var handling and CANN
auto-detection (probed before the pip-provided CUDA toolkit). Only
target_utils.cc and layout/ascend_layouts.cc stay unconditional: common
transforms call TargetIsAscend, and tilelang.layout registers
tl.AscendFractalLayout at import time. TILELANG_USE_ASCEND_STUBS=OFF links the
real libascendcl instead.

Verified on NPU hardware (testing/ascend, 1010 passed) and on a CANN-less
GB200 host where a single USE_CUDA+USE_ROCM+USE_ASCEND wheel imports cleanly
and runs CUDA kernels.
This commit is contained in:
LeiWang1999
2026-09-11 22:35:22 +08:00
parent 8aa0f76044
commit 9b1b879999
6 changed files with 450 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CMakeLists.txt
+33 -1
View File
@@ -274,12 +274,13 @@ else()
endif()
# Configs
set(TILELANG_BACKENDS CUDA ROCM METAL LLVM)
set(TILELANG_BACKENDS CUDA ROCM METAL LLVM ASCEND)
set(TILELANG_BACKEND_DOC_CUDA "Enable CUDA backend (ON/OFF/or CUDA SDK path)")
set(TILELANG_BACKEND_DOC_ROCM "Enable ROCm backend (ON/OFF/or ROCm SDK path)")
set(TILELANG_BACKEND_DOC_METAL "Enable Metal backend")
set(TILELANG_BACKEND_DOC_LLVM "Enable LLVM backend")
set(TILELANG_BACKEND_DOC_ASCEND "Enable Ascend backend")
# TVM's config.cmake redefines USE_* options later, so we cache the user's choice
# (including explicit -DUSE_XXX arguments) before we include TVM and restore it
@@ -354,6 +355,22 @@ unset(_TILELANG_USE_HIP_STUBS_DEFAULT)
# /opt/rocm/include -> 3rdparty/hip-headers/include (vendored default).
set(TILELANG_HIP_INCLUDE_DIR "" CACHE PATH
"Path to HIP headers when building with USE_ROCM=ON on a host without a ROCm runtime")
# The ascendcl stub library avoids a hard DT_NEEDED dependency on
# libascendcl.so, so Ascend-enabled wheels can be imported on machines without
# CANN installed.
#
# The stub is POSIX-only (dlopen/dlsym via <dlfcn.h>); the Ascend backend
# itself does not support Windows. Only meaningful when USE_ASCEND is enabled.
if(WIN32 AND NOT CYGWIN)
set(_TILELANG_USE_ASCEND_STUBS_DEFAULT OFF)
else()
set(_TILELANG_USE_ASCEND_STUBS_DEFAULT ON)
endif()
option(TILELANG_USE_ASCEND_STUBS
"Use POSIX dlopen-based ascendcl stub library for portable wheels"
${_TILELANG_USE_ASCEND_STUBS_DEFAULT})
unset(_TILELANG_USE_ASCEND_STUBS_DEFAULT)
# Configs end
include(cmake/load_tvm.cmake)
@@ -454,10 +471,25 @@ if(NOT TILELANG_BACKEND_USER_SELECTED)
endif()
endif()
if(DEFINED ENV{USE_ASCEND})
set(_tilelang_backend_env_selected ON)
if($ENV{USE_ASCEND})
set(USE_ASCEND ON)
else()
set(USE_ASCEND OFF)
endif()
endif()
if(NOT _tilelang_backend_env_selected)
if(APPLE)
message(STATUS "Enable Metal support by default.")
set(USE_METAL ON)
elseif(IS_DIRECTORY "$ENV{ASCEND_HOME_PATH}" OR
IS_DIRECTORY "/usr/local/Ascend/ascend-toolkit/latest")
# A physical CANN install is a stronger signal than a (possibly
# pip-provided) CUDA toolkit, so probe it first.
message(STATUS "CANN toolkit found; enable Ascend support by default.")
set(USE_ASCEND ON)
elseif(TILELANG_CUDA_TOOLKIT_AVAILABLE)
message(STATUS "Enable CUDA support by default.")
set(USE_CUDA ON)
docs/runtime_internals/stubs.md
+31 -5
View File
@@ -1,7 +1,7 @@
# CUDA and ROCm Stub Libraries
# CUDA, ROCm, and Ascend Stub Libraries
This document describes TileLang's stub mechanism for GPU driver/runtime
libraries (CUDA and ROCm/HIP).
This document describes TileLang's stub mechanism for GPU/NPU driver and
runtime libraries (CUDA, ROCm/HIP, and Ascend/CANN).
## Purpose
@@ -28,10 +28,20 @@ ROCm:
loads `libhiprtc.so` and exposes the minimal HIPRTC API subset used by
TileLang/TVM.
Ascend:
1. **CANN Runtime (`ascendcl_stub`, file `libstub_ascendcl.so`)**: Allows
TileLang to be imported on systems without CANN installed by lazy-loading
`libascendcl.so` only when needed. The stub prefers the copy already loaded
by torch_npu (via `RTLD_DEFAULT` / `RTLD_NEXT`) before searching the
filesystem. There is no RTC-style stub for Ascend: device kernels are
compiled by invoking the `bisheng` compiler as a subprocess, not through a
library API.
## Implementation
The CUDA stubs in `src/cuda/stubs/` and ROCm stubs in
`src/rocm/stubs/` implement a lazy-loading mechanism:
The CUDA stubs in `src/cuda/stubs/`, ROCm stubs in `src/rocm/stubs/`, and
Ascend stubs in `src/ascend/stubs/` implement a lazy-loading mechanism:
- **Lazy Loading**: Libraries are loaded via `dlopen` only upon the first API call.
- **Global Symbol Reuse**: For `cudart` and `nvrtc`, the stubs first check the global namespace (`RTLD_DEFAULT`) to use any already loaded symbols (e.g., from PyTorch).
@@ -39,6 +49,16 @@ The CUDA stubs in `src/cuda/stubs/` and ROCm stubs in
falls back to `dlopen("libamdhip64.so")`. It additionally provides wrappers
for `hsa_init` / `hsa_shut_down` so that ROCm-enabled wheels do not record a
hard dependency on `libhsa-runtime64` at import time.
- **Ascend Notes**: `ascendcl_stub` checks `RTLD_DEFAULT` / `RTLD_NEXT` first
(reusing the `libascendcl` copy loaded by torch_npu), then
`dlopen("libascendcl.so")` (which honors the `LD_LIBRARY_PATH` set by CANN's
`set_env.sh`), and finally well-known install roots: `$ASCEND_HOME_PATH`,
`$ASCEND_TOOLKIT_HOME`, and `/usr/local/Ascend/ascend-toolkit/latest`, each
with the `lib64/` and `runtime/lib64/` suffixes. The stub exposes only the
small entrypoint set used by the Ascend runtime module; `aclGetRecentErrMsg`
degrades gracefully (returns `nullptr` instead of throwing) so error
reporting never fails. No CANN headers are needed to build the stub — the
ABI is expressed with opaque pointer and fixed-width integer types.
- **Versioning Support**: Handles ABI differences between CUDA versions (e.g., `cudaGraphInstantiate` changes in CUDA 12).
## Build Option
@@ -49,3 +69,9 @@ The CUDA stubs in `src/cuda/stubs/` and ROCm stubs in
- `TILELANG_USE_HIP_STUBS` (Default: `ON`) controls ROCm stubs. When enabled
(and `USE_ROCM=ON`), TileLang/TVM link against `hip_stub` / `hiprtc_stub`
instead of the system ROCm libraries.
- `TILELANG_USE_ASCEND_STUBS` (Default: `ON` on non-Windows) controls the
Ascend stub. When enabled (and `USE_ASCEND=ON`), TileLang links against
`ascendcl_stub` instead of the system CANN library. When disabled, TileLang
links directly against the real `libascendcl` found under
`$ASCEND_HOME_PATH` / `$ASCEND_TOOLKIT_HOME` /
`/usr/local/Ascend/ascend-toolkit/latest`.
src/ascend/CMakeLists.txt
+63 -7
View File
@@ -1,11 +1,27 @@
# Ascend backend: source files and build configuration.
# Ascend backend: source files, stubs, and build configuration.
#
# NOTE: Unlike the CUDA/ROCm/Metal backends, Ascend has no `USE_ASCEND` guard
# yet, so its sources are compiled unconditionally. The codegen emits AscendC
# source text and carries no CANN build-time dependency, so it is safe to build
# on any host. A `USE_ASCEND` toggle can be added later without moving sources.
# Ascend target helpers and layout registrations are pure C++ registrations
# used by common code (TargetIsAscend is called from common transforms via
# backend/common/target_utils.h; tilelang.layout registers the
# tl.AscendFractalLayout object type at import time). Compile them regardless
# of USE_ASCEND.
file(GLOB TILE_LANG_ASCEND_ALWAYS_SRCS
src/ascend/target_utils.cc
src/ascend/layout/ascend_layouts.cc
)
list(APPEND TILE_LANG_SRCS ${TILE_LANG_ASCEND_ALWAYS_SRCS})
file(GLOB TILE_LANG_ASCEND_SRCS
if(NOT USE_ASCEND)
return()
endif()
message(STATUS "Ascend Backend is enabled")
# The codegen emits AscendC source text and carries no CANN build-time
# dependency; the runtime module links against the ascendcl stub below (or the
# real libascendcl when stubs are disabled), so USE_ASCEND=ON is safe on any
# non-Windows host.
file(GLOB TILE_LANG_ASCEND_ACTIVE_SRCS
src/ascend/*.cc
src/ascend/codegen/*.cc
src/ascend/layout/*.cc
@@ -14,4 +30,44 @@ file(GLOB TILE_LANG_ASCEND_SRCS
src/ascend/transform/*.cc
src/ascend/transform/auto_schedule/*.cc
)
list(APPEND TILE_LANG_SRCS ${TILE_LANG_ASCEND_SRCS})
list(REMOVE_ITEM TILE_LANG_ASCEND_ACTIVE_SRCS ${TILE_LANG_ASCEND_ALWAYS_SRCS})
list(APPEND TILE_LANG_SRCS ${TILE_LANG_ASCEND_ACTIVE_SRCS})
if(TILELANG_USE_ASCEND_STUBS)
# ==========================================================================
# ascendcl Stub Library (libstub_ascendcl.so)
# ==========================================================================
# Provides drop-in replacements for the CANN runtime entrypoints used by the
# Ascend runtime module by lazily loading libascendcl.so at runtime
# (preferring a copy already loaded by torch_npu). This keeps wheels free of
# a hard DT_NEEDED on libascendcl.so so they can be imported without CANN.
# ==========================================================================
add_library(ascendcl_stub SHARED src/ascend/stubs/ascendcl.cc)
target_compile_definitions(ascendcl_stub PRIVATE TILELANG_ASCENDCL_STUB_EXPORTS)
target_link_libraries(ascendcl_stub PRIVATE ${CMAKE_DL_LIBS})
set_target_properties(ascendcl_stub PROPERTIES
LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/lib"
# Keep the stub filename distinct from the real CANN library, matching the
# CUDA/ROCm stub naming convention.
OUTPUT_NAME "stub_ascendcl"
)
list(APPEND TILELANG_ACTIVE_BACKEND_STUB_LINK ascendcl_stub)
list(APPEND TILELANG_ACTIVE_BACKEND_STUB_TARGETS ascendcl_stub)
elseif(NOT WIN32)
# Stubs disabled: link the real CANN runtime directly (on-hardware builds).
find_library(TILELANG_ASCENDCL_LIBRARY ascendcl
HINTS $ENV{ASCEND_HOME_PATH} $ENV{ASCEND_TOOLKIT_HOME}
/usr/local/Ascend/ascend-toolkit/latest
PATH_SUFFIXES lib64 runtime/lib64)
if(NOT TILELANG_ASCENDCL_LIBRARY)
message(FATAL_ERROR
"TILELANG_USE_ASCEND_STUBS=OFF requires libascendcl, but it was not "
"found. Install the CANN toolkit and set ASCEND_HOME_PATH (or "
"ASCEND_TOOLKIT_HOME), or keep TILELANG_USE_ASCEND_STUBS=ON.")
endif()
message(STATUS "Linking Ascend runtime against ${TILELANG_ASCENDCL_LIBRARY}")
list(APPEND TILELANG_ACTIVE_BACKEND_STUB_LINK ${TILELANG_ASCENDCL_LIBRARY})
endif()
src/ascend/runtime/ascend_module.cc
+18 -104
View File
@@ -5,10 +5,12 @@
* The module mirrors CUDA's binary runtime module: codegen stores executable
* device bytes and launch metadata, while the runtime loads functions and
* launches them on the stream supplied by TVM-FFI's DLPack Exchange API.
* CANN symbols are resolved lazily so TileLang keeps no CANN build-time
* dependency. The Ascend backend is Linux-only.
* CANN symbols are resolved lazily by the ascendcl stub library
* (src/ascend/stubs/) so TileLang keeps no CANN build-time dependency.
* The Ascend backend is Linux-only.
*/
#include <dlfcn.h>
#include "ascend/stubs/ascendcl.h"
#include <tvm/ffi/extra/c_env_api.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/function.h>
@@ -329,98 +331,11 @@ struct AclLaunchKernelCfg {
size_t num_attrs;
};
class AscendDriver {
public:
static AscendDriver *Global() {
static auto *driver = new AscendDriver();
return driver;
}
AclError BinaryLoadFromData(const void *data, size_t size,
AclBinHandle *handle) const {
return binary_load_from_data_(data, size, nullptr, handle);
}
AclError BinaryGetFunction(AclBinHandle binary, const char *name,
AclFuncHandle *function) const {
return binary_get_function_(binary, name, function);
}
AclError BinaryUnload(AclBinHandle binary) const {
return binary_unload_(binary);
}
AclError GetDevice(int32_t *device_id) const {
return get_device_(device_id);
}
AclError LaunchKernelWithHostArgs(AclFuncHandle function, uint32_t num_blocks,
AclStream stream,
AclLaunchKernelCfg *config, void *args,
size_t args_size) const {
return launch_kernel_with_host_args_(function, num_blocks, stream, config,
args, args_size, nullptr, 0);
}
const char *GetRecentErrorMessage() const {
return get_recent_error_message_ == nullptr ? nullptr
: get_recent_error_message_();
}
private:
using BinaryLoadFromDataFn = AclError (*)(const void *, size_t, const void *,
AclBinHandle *);
using BinaryGetFunctionFn = AclError (*)(AclBinHandle, const char *,
AclFuncHandle *);
using BinaryUnloadFn = AclError (*)(AclBinHandle);
using GetDeviceFn = AclError (*)(int32_t *);
using LaunchKernelWithHostArgsFn = AclError (*)(AclFuncHandle, uint32_t,
AclStream,
AclLaunchKernelCfg *, void *,
size_t, void *, size_t);
using GetRecentErrorMessageFn = const char *(*)();
AscendDriver() {
library_ = dlopen("libascendcl.so", RTLD_LAZY | RTLD_LOCAL);
TVM_FFI_CHECK(library_ != nullptr, RuntimeError)
<< "Ascend runtime could not load libascendcl.so: " << dlerror();
binary_load_from_data_ =
LoadSymbol<BinaryLoadFromDataFn>("aclrtBinaryLoadFromData");
binary_get_function_ =
LoadSymbol<BinaryGetFunctionFn>("aclrtBinaryGetFunction");
binary_unload_ = LoadSymbol<BinaryUnloadFn>("aclrtBinaryUnLoad");
get_device_ = LoadSymbol<GetDeviceFn>("aclrtGetDevice");
launch_kernel_with_host_args_ =
LoadSymbol<LaunchKernelWithHostArgsFn>("aclrtLaunchKernelWithHostArgs");
get_recent_error_message_ =
LoadSymbol<GetRecentErrorMessageFn>("aclGetRecentErrMsg");
}
template <typename FunctionType> FunctionType LoadSymbol(const char *name) {
dlerror();
void *symbol = dlsym(library_, name);
const char *error = dlerror();
TVM_FFI_CHECK(symbol != nullptr && error == nullptr, RuntimeError)
<< "Ascend runtime could not resolve " << name
<< " from libascendcl.so: "
<< (error == nullptr ? "symbol not found" : error);
return reinterpret_cast<FunctionType>(symbol);
}
void *library_{nullptr};
BinaryLoadFromDataFn binary_load_from_data_{nullptr};
BinaryGetFunctionFn binary_get_function_{nullptr};
BinaryUnloadFn binary_unload_{nullptr};
GetDeviceFn get_device_{nullptr};
LaunchKernelWithHostArgsFn launch_kernel_with_host_args_{nullptr};
GetRecentErrorMessageFn get_recent_error_message_{nullptr};
};
void CheckAcl(AclError result, const char *operation) {
if (result == kAclSuccess) {
return;
}
const char *message = AscendDriver::Global()->GetRecentErrorMessage();
const char *message = aclGetRecentErrMsg();
TVM_FFI_THROW(RuntimeError)
<< operation << " failed with ACL error " << result
<< (message == nullptr ? "" : std::string(": ") + message);
@@ -440,10 +355,11 @@ public:
if (device_modules_.empty()) {
return;
}
AscendDriver *driver = AscendDriver::Global();
for (const auto &[device_id, device_module] : device_modules_) {
if (device_module.binary != nullptr) {
AclError result = driver->BinaryUnload(device_module.binary);
// binary != nullptr implies a prior successful stub call, so the
// lazy-loaded library is guaranteed to be present here.
AclError result = aclrtBinaryUnLoad(device_module.binary);
if (result != kAclSuccess) {
LOG(WARNING) << "aclrtBinaryUnLoad failed for Ascend device "
<< device_id << " with error " << result;
@@ -487,11 +403,10 @@ public:
AclFuncHandle GetFunctionHandle(int32_t device_id, const std::string &name) {
std::lock_guard<std::mutex> lock(mutex_);
AscendDriver *driver = AscendDriver::Global();
DeviceModule &device_module = device_modules_[device_id];
if (device_module.binary == nullptr) {
CheckAcl(driver->BinaryLoadFromData(code_.data(), code_.size(),
&device_module.binary),
CheckAcl(aclrtBinaryLoadFromData(code_.data(), code_.size(), nullptr,
&device_module.binary),
"aclrtBinaryLoadFromData");
}
@@ -501,9 +416,9 @@ public:
}
AclFuncHandle function{nullptr};
CheckAcl(driver->BinaryGetFunction(device_module.binary, name.c_str(),
&function),
"aclrtBinaryGetFunction");
CheckAcl(
aclrtBinaryGetFunction(device_module.binary, name.c_str(), &function),
"aclrtBinaryGetFunction");
device_module.functions.emplace(name, function);
return function;
}
@@ -555,9 +470,8 @@ public:
<< "Ascend dynamic UBUF size exceeds uint32 range for kernel "
<< function_name_;
AscendDriver *driver = AscendDriver::Global();
int32_t device_id = 0;
CheckAcl(driver->GetDevice(&device_id), "aclrtGetDevice");
CheckAcl(aclrtGetDevice(&device_id), "aclrtGetDevice");
AclFuncHandle function =
module_->GetFunctionHandle(device_id, function_name_);
AclStream stream = TVMFFIEnvGetStream(kDLExtDev, device_id);
@@ -574,11 +488,11 @@ public:
config_ptr = &config;
}
AclError result = driver->LaunchKernelWithHostArgs(
AclError result = aclrtLaunchKernelWithHostArgs(
function, static_cast<uint32_t>(num_blocks), stream, config_ptr,
packed_args, packed_args_size);
packed_args, packed_args_size, nullptr, 0);
if (result != kAclSuccess) {
const char *message = driver->GetRecentErrorMessage();
const char *message = aclGetRecentErrMsg();
std::ostringstream error;
error << "aclrtLaunchKernelWithHostArgs failed for " << function_name_
<< " with ACL error " << result << ", grid=" << num_blocks
src/ascend/stubs/ascendcl.cc
+208
View File
@@ -0,0 +1,208 @@
/**
* \file ascendcl.cc
* \brief Implementation of the CANN runtime (libascendcl) stub library.
*
* This implements lazy loading of libascendcl.so and provides exported global
* wrapper functions that serve as drop-in replacements for the ACL runtime
* entrypoints used by TileLang's Ascend runtime module.
*
* The implementation mirrors src/rocm/stubs/hip.cc:
* - Resolve symbols via dlopen/dlsym on first use.
* - Prefer RTLD_DEFAULT/RTLD_NEXT when CANN is already loaded by another
* framework (e.g. torch_npu).
* - Fall back to well-known CANN install roots when LD_LIBRARY_PATH is not
* set up (no set_env.sh sourced).
*/
#include "ascendcl.h"
#if defined(_WIN32) && !defined(__CYGWIN__)
#error "ascendcl_stub is POSIX-only (requires <dlfcn.h> / dlopen). " \
"The Ascend backend does not support Windows; configure with " \
"-DTILELANG_USE_ASCEND_STUBS=OFF."
#endif
#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif
#include <dlfcn.h>
#include <cstdlib>
#include <stdexcept>
#include <string>
#include <vector>
namespace tvm::tl::ascendcl {
namespace {
// First missing required symbol, recorded so get() can name it.
std::string missing_symbol;
template <typename T> T GetSymbol(void *handle, const char *name) {
(void)dlerror();
void *sym = dlsym(handle, name);
const char *error = dlerror();
if (error != nullptr) {
return nullptr;
}
return reinterpret_cast<T>(sym);
}
std::vector<std::string> LibAscendCLPathCandidates() {
std::vector<std::string> candidates;
// Honors LD_LIBRARY_PATH (set by CANN's set_env.sh).
candidates.emplace_back("libascendcl.so");
std::vector<std::string> roots;
for (const char *env : {"ASCEND_HOME_PATH", "ASCEND_TOOLKIT_HOME"}) {
const char *value = std::getenv(env);
if (value != nullptr && value[0] != '\0') {
roots.emplace_back(value);
}
}
roots.emplace_back("/usr/local/Ascend/ascend-toolkit/latest");
for (const std::string &root : roots) {
candidates.push_back(root + "/lib64/libascendcl.so");
candidates.push_back(root + "/runtime/lib64/libascendcl.so");
}
return candidates;
}
void *TryLoadLibAscendCL() {
// Prefer already-loaded symbols (e.g. if torch_npu is imported first).
// We use a representative symbol and ensure we don't just find ourselves.
void *sym = dlsym(RTLD_DEFAULT, "aclrtGetDevice");
if (sym != nullptr && sym != reinterpret_cast<void *>(&aclrtGetDevice)) {
return RTLD_DEFAULT;
}
sym = dlsym(RTLD_NEXT, "aclrtGetDevice");
if (sym != nullptr && sym != reinterpret_cast<void *>(&aclrtGetDevice)) {
return RTLD_NEXT;
}
// Otherwise, attempt to dlopen the library directly. If a copy is already
// mapped (e.g. under a different path), the dynamic loader dedups by SONAME
// and returns the existing mapping.
for (const std::string &path : LibAscendCLPathCandidates()) {
void *handle = dlopen(path.c_str(), RTLD_LAZY | RTLD_LOCAL);
if (handle != nullptr) {
return handle;
}
}
return nullptr;
}
AscendCLAPI CreateAscendCLAPI() {
AscendCLAPI api{};
void *handle = AscendCLAPI::get_handle();
if (handle == nullptr) {
return api;
}
// Required symbols: record the first missing name (older CANN releases may
// lack e.g. aclrtLaunchKernelWithHostArgs) so get() can report it instead of
// leaving a null pointer behind.
#define LOOKUP(symbol) \
api.symbol##_ = GetSymbol<decltype(api.symbol##_)>(handle, #symbol); \
if (api.symbol##_ == nullptr) { \
if (missing_symbol.empty()) { \
missing_symbol = #symbol; \
} \
return AscendCLAPI{}; \
}
LOOKUP(aclrtBinaryLoadFromData)
LOOKUP(aclrtBinaryGetFunction)
LOOKUP(aclrtBinaryUnLoad)
LOOKUP(aclrtGetDevice)
LOOKUP(aclrtLaunchKernelWithHostArgs)
#undef LOOKUP
// Optional: wrappers return nullptr when absent.
api.aclGetRecentErrMsg_ = GetSymbol<decltype(api.aclGetRecentErrMsg_)>(
handle, "aclGetRecentErrMsg");
return api;
}
} // namespace
void *AscendCLAPI::get_handle() {
static void *handle = TryLoadLibAscendCL();
return handle;
}
bool AscendCLAPI::is_available() { return get_handle() != nullptr; }
AscendCLAPI *AscendCLAPI::get_or_null() {
static AscendCLAPI singleton = CreateAscendCLAPI();
return &singleton;
}
AscendCLAPI *AscendCLAPI::get() {
AscendCLAPI *api = get_or_null();
if (!is_available()) {
throw std::runtime_error(
"CANN runtime library (libascendcl.so) not found. Install the Ascend "
"CANN toolkit and source its set_env.sh (or import torch_npu first) "
"before using TileLang's Ascend backend.");
}
if (!missing_symbol.empty()) {
throw std::runtime_error(
"libascendcl.so was found but does not provide the required symbol `" +
missing_symbol +
"`. TileLang's Ascend backend requires a newer CANN toolkit.");
}
return api;
}
} // namespace tvm::tl::ascendcl
// ============================================================================
// Global wrapper function implementations
// ============================================================================
using tvm::tl::ascendcl::AscendCLAPI;
extern "C" {
int32_t aclrtBinaryLoadFromData(const void *data, size_t size,
const void *options, void **binHandle) {
return AscendCLAPI::get()->aclrtBinaryLoadFromData_(data, size, options,
binHandle);
}
int32_t aclrtBinaryGetFunction(void *binHandle, const char *name,
void **funcHandle) {
return AscendCLAPI::get()->aclrtBinaryGetFunction_(binHandle, name,
funcHandle);
}
int32_t aclrtBinaryUnLoad(void *binHandle) {
return AscendCLAPI::get()->aclrtBinaryUnLoad_(binHandle);
}
int32_t aclrtGetDevice(int32_t *deviceId) {
return AscendCLAPI::get()->aclrtGetDevice_(deviceId);
}
int32_t aclrtLaunchKernelWithHostArgs(void *func, uint32_t numBlocks,
void *stream, void *config, void *args,
size_t argsSize, void *reserved,
size_t reservedSize) {
return AscendCLAPI::get()->aclrtLaunchKernelWithHostArgs_(
func, numBlocks, stream, config, args, argsSize, reserved, reservedSize);
}
const char *aclGetRecentErrMsg(void) {
// Error-message path: must never throw, callers handle nullptr.
AscendCLAPI *api = AscendCLAPI::get_or_null();
if (!AscendCLAPI::is_available() || api->aclGetRecentErrMsg_ == nullptr) {
return nullptr;
}
return api->aclGetRecentErrMsg_();
}
} // extern "C"
src/ascend/stubs/ascendcl.h
+97
View File
@@ -0,0 +1,97 @@
/**
* \file ascendcl.h
* \brief Stub library header for lazy loading the CANN runtime (libascendcl).
*
* This mirrors the CUDA/ROCm stubs in src/cuda/stubs/ and src/rocm/stubs/:
* - Instead of linking against libascendcl.so at build time, TileLang links
* against a small stub library (libstub_ascendcl.so) that resolves the ACL
* symbols via dlopen()/dlsym() on first use.
*
* This enables:
* 1. Importing TileLang on machines without CANN installed.
* 2. Reusing the libascendcl copy already loaded by torch_npu.
* 3. Building a single wheel that can run across environments.
*
* Unlike the CUDA/ROCm stubs, no CANN headers are required (or vendored): the
* TileLang runtime only uses a handful of entrypoints whose ABI is expressed
* here with opaque pointer and fixed-width integer types.
*
* Usage:
* #include "ascend/stubs/ascendcl.h"
* int32_t error = aclrtGetDevice(&device_id);
*/
#pragma once
#include <cstddef>
#include <cstdint>
// Symbol visibility macros for shared library export.
#if defined(_WIN32) || defined(__CYGWIN__)
#ifdef TILELANG_ASCENDCL_STUB_EXPORTS
#define TILELANG_ASCENDCL_STUB_API __declspec(dllexport)
#else
#define TILELANG_ASCENDCL_STUB_API __declspec(dllimport)
#endif
#else
#define TILELANG_ASCENDCL_STUB_API __attribute__((visibility("default")))
#endif
namespace tvm::tl::ascendcl {
/**
* \brief ACL API accessor struct with lazy loading support.
*
* Similar to tvm::tl::hip::HIPDriverAPI, this struct resolves libascendcl.so
* symbols lazily on first use.
*
* Function pointer members have a trailing underscore to avoid collisions with
* the exported global wrapper functions.
*/
struct TILELANG_ASCENDCL_STUB_API AscendCLAPI {
int32_t (*aclrtBinaryLoadFromData_)(const void *, size_t, const void *,
void **);
int32_t (*aclrtBinaryGetFunction_)(void *, const char *, void **);
int32_t (*aclrtBinaryUnLoad_)(void *);
int32_t (*aclrtGetDevice_)(int32_t *);
int32_t (*aclrtLaunchKernelWithHostArgs_)(void *, uint32_t, void *, void *,
void *, size_t, void *, size_t);
// Optional: absent on CANN versions that predate the API.
const char *(*aclGetRecentErrMsg_)(void);
/// Throws std::runtime_error when libascendcl.so (or a required symbol
/// within it) is unavailable.
static AscendCLAPI *get();
/// Never throws; returns nullptr when the library could not be loaded.
static AscendCLAPI *get_or_null();
static bool is_available();
static void *get_handle();
};
} // namespace tvm::tl::ascendcl
// ============================================================================
// Global wrapper functions for the lazy-loaded ACL API
// ============================================================================
// These functions provide drop-in replacements for the CANN runtime
// entrypoints used by TileLang. The implementations are in ascendcl.cc.
extern "C" {
TILELANG_ASCENDCL_STUB_API int32_t aclrtBinaryLoadFromData(const void *data,
size_t size,
const void *options,
void **binHandle);
TILELANG_ASCENDCL_STUB_API int32_t aclrtBinaryGetFunction(void *binHandle,
const char *name,
void **funcHandle);
TILELANG_ASCENDCL_STUB_API int32_t aclrtBinaryUnLoad(void *binHandle);
TILELANG_ASCENDCL_STUB_API int32_t aclrtGetDevice(int32_t *deviceId);
TILELANG_ASCENDCL_STUB_API int32_t aclrtLaunchKernelWithHostArgs(
void *func, uint32_t numBlocks, void *stream, void *config, void *args,
size_t argsSize, void *reserved, size_t reservedSize);
/// Returns nullptr (instead of throwing) when libascendcl.so or the symbol is
/// unavailable, so error-reporting paths stay usable.
TILELANG_ASCENDCL_STUB_API const char *aclGetRecentErrMsg(void);
} // extern "C"