diff --git a/python/infinicore/_preload.py b/python/infinicore/_preload.py index 69121ff2f..bcdfd34b8 100644 --- a/python/infinicore/_preload.py +++ b/python/infinicore/_preload.py @@ -6,7 +6,6 @@ import sys from typing import Iterable, List - _PRELOADED_HANDLES: List[ctypes.CDLL] = [] @@ -98,9 +97,7 @@ def _import_extension_global(module_name: str, path: str) -> None: raise module_path = getattr(module, "__file__", None) or path - _PRELOADED_HANDLES.append( - ctypes.CDLL(module_path, mode=ctypes.RTLD_GLOBAL) - ) + _PRELOADED_HANDLES.append(ctypes.CDLL(module_path, mode=ctypes.RTLD_GLOBAL)) def preload_flash_attn() -> None: diff --git a/src/bridge/infini/rt.hpp b/src/bridge/infini/rt.hpp index c5bec7ecd..c75cb1335 100644 --- a/src/bridge/infini/rt.hpp +++ b/src/bridge/infini/rt.hpp @@ -33,6 +33,8 @@ inline constexpr ::infini::rt::Device::Type translate_to(infiniDevice_t device) return ::infini::rt::Device::Type::kIluvatar; case INFINI_DEVICE_HYGON: return ::infini::rt::Device::Type::kHygon; + case INFINI_DEVICE_ALI: + return ::infini::rt::Device::Type::kThead; default: return ::infini::rt::Device::Type::kCount; } @@ -56,6 +58,8 @@ inline constexpr infiniDevice_t translate_from(::infini::rt::Device::Type device return INFINI_DEVICE_ILUVATAR; case ::infini::rt::Device::Type::kHygon: return INFINI_DEVICE_HYGON; + case ::infini::rt::Device::Type::kThead: + return INFINI_DEVICE_ALI; default: return INFINI_DEVICE_TYPE_COUNT; } diff --git a/src/infinicore/context/context_impl.cc b/src/infinicore/context/context_impl.cc index 29a470679..37fda83a3 100644 --- a/src/infinicore/context/context_impl.cc +++ b/src/infinicore/context/context_impl.cc @@ -90,15 +90,18 @@ void ContextImpl::initializeRuntime() { ContextImpl::ContextImpl() { initializeRuntime(); initializeRuntime(); + initializeRuntime(); + initializeRuntime(); + initializeRuntime(); + initializeRuntime(); initializeRuntime(); initializeRuntime(); - initializeRuntime(); + initializeRuntime(); if (current_runtime_ == nullptr && !runtime_table_[static_cast(Device::Type::CPU)].empty()) { current_runtime_ = runtime_table_[static_cast(Device::Type::CPU)][0].get(); } } - namespace context { void setDevice(Device device) { diff --git a/xmake.lua b/xmake.lua index 9ec14e9ca..2c9b3c69f 100644 --- a/xmake.lua +++ b/xmake.lua @@ -306,6 +306,71 @@ local function get_infinirt_root() return nil end +local function add_moore_runtime_sdk_dirs() + if not has_config("moore-gpu") then + return + end + local musa_root = os.getenv("MUSA_ROOT") or os.getenv("MUSA_HOME") or os.getenv("MUSA_PATH") or "/usr/local/musa" + add_includedirs(path.join(musa_root, "include")) + add_linkdirs( + path.join(musa_root, "lib"), + path.join(musa_root, "lib64")) + add_rpathdirs(path.join(musa_root, "lib"), path.join(musa_root, "lib64")) +end + +local function add_metax_runtime_sdk_dirs() + if not has_config("metax-gpu") then + return + end + local maca_root = os.getenv("MACA_PATH") or os.getenv("MACA_HOME") or os.getenv("MACA_ROOT") or "/opt/maca" + add_includedirs(path.join(maca_root, "include")) + add_linkdirs(path.join(maca_root, "lib"), path.join(maca_root, "lib64")) + add_rpathdirs(path.join(maca_root, "lib"), path.join(maca_root, "lib64")) +end + +local function add_ascend_runtime_sdk_dirs() + if not has_config("ascend-npu") then + return + end + local ascend_home = os.getenv("ASCEND_HOME") or os.getenv("ASCEND_TOOLKIT_HOME") or "/usr/local/Ascend/ascend-toolkit/latest" + add_includedirs(path.join(ascend_home, "include"), path.join(ascend_home, "include", "aclnn")) + add_linkdirs(path.join(ascend_home, "lib64")) + add_rpathdirs(path.join(ascend_home, "lib64")) +end + +local function get_ali_cuda_root() + local ali_cuda_root = os.getenv("ALI_CUDA_ROOT") or os.getenv("PPU_CUDA_ROOT") + local ppu_sdk_root = os.getenv("PPU_SDK_ROOT") + if (not ali_cuda_root or ali_cuda_root == "") and ppu_sdk_root and ppu_sdk_root ~= "" then + ali_cuda_root = path.join(ppu_sdk_root, "CUDA_SDK") + end + return ali_cuda_root or get_config("cuda") or "/usr/local/PPU_SDK/CUDA_SDK" +end + +local function add_ali_runtime_sdk_dirs() + if not has_config("ali-ppu") then + return + end + local ali_cuda_root = get_ali_cuda_root() + for _, include_dir in ipairs({ + path.join(ali_cuda_root, "include"), + path.join(ali_cuda_root, "targets", "x86_64-linux", "include"), + }) do + if os.isdir(include_dir) then + add_includedirs(include_dir) + end + end + for _, link_dir in ipairs({ + path.join(ali_cuda_root, "lib64"), + path.join(ali_cuda_root, "targets", "x86_64-linux", "lib"), + }) do + if os.isdir(link_dir) then + add_linkdirs(link_dir) + add_rpathdirs(link_dir) + end + end +end + local function get_infiniops_cuda_architectures() local arch_opt = get_config("cuda_arch") if not arch_opt or arch_opt == "" then @@ -508,6 +573,10 @@ target_end() target("infiniop") set_kind("shared") add_external_infinirt() + add_moore_runtime_sdk_dirs() + add_metax_runtime_sdk_dirs() + add_ascend_runtime_sdk_dirs() + add_ali_runtime_sdk_dirs() local public_cuda_root = get_config("cuda") or os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") if public_cuda_root and public_cuda_root ~= "" then @@ -529,6 +598,9 @@ target("infiniop") end elseif has_config("nv-gpu") then add_includedirs(path.join("/usr/local/cuda", "include")) + elseif has_config("iluvatar-gpu") then + add_includedirs(path.join("/usr/local/corex", "include")) + add_linkdirs(path.join("/usr/local/corex", "lib64")) end if has_config("cpu") then @@ -625,6 +697,7 @@ target("infiniccl") add_files("src/infiniccl/*.cc") add_installfiles("include/infiniccl.h", {prefixdir = "include"}) + add_installfiles("include/infinirt.h", {prefixdir = "include"}) set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini")) target_end() @@ -666,8 +739,13 @@ target("infinicore_cpp_api") end add_includedirs(INFINI_ROOT.."/include", { public = true }) add_external_infinirt() - if has_config("nv-gpu") then - local cuda_root = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or get_config("cuda") or "/usr/local/cuda" + add_moore_runtime_sdk_dirs() + add_metax_runtime_sdk_dirs() + add_ascend_runtime_sdk_dirs() + add_ali_runtime_sdk_dirs() + if has_config("nv-gpu") or has_config("iluvatar-gpu") then + local default_cuda_root = has_config("iluvatar-gpu") and "/usr/local/corex" or "/usr/local/cuda" + local cuda_root = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or get_config("cuda") or default_cuda_root add_includedirs(cuda_root .. "/include") end if has_config("infiniops") then @@ -910,6 +988,7 @@ target("infinicore_cpp_api") add_installfiles("include/infinicore/(**/*.hpp)",{prefixdir = "include/infinicore"}) add_installfiles("include/infinicore.h", {prefixdir = "include"}) add_installfiles("include/infinicore.hpp", {prefixdir = "include"}) + add_installfiles("include/infinirt.h", {prefixdir = "include"}) after_build(function (target) print(YELLOW .. "[Congratulations!] Now you can install the libraries with \"xmake install\"" .. NC) end) target_end() @@ -932,6 +1011,14 @@ target("_infinicore") local INFINI_ROOT = os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini") add_includedirs(INFINI_ROOT.."/include", { public = true }) add_external_infinirt() + add_moore_runtime_sdk_dirs() + add_metax_runtime_sdk_dirs() + add_ascend_runtime_sdk_dirs() + add_ali_runtime_sdk_dirs() + if has_config("iluvatar-gpu") then + local cuda_root = os.getenv("CUDA_HOME") or os.getenv("CUDA_PATH") or get_config("cuda") or "/usr/local/corex" + add_cxxflags("-idirafter", cuda_root .. "/include", {force = true}) + end add_linkdirs(INFINI_ROOT.."/lib") add_links("infiniop", "infiniccl") diff --git a/xmake/ali.lua b/xmake/ali.lua index 550f7e0fc..ccec64a41 100644 --- a/xmake/ali.lua +++ b/xmake/ali.lua @@ -80,6 +80,7 @@ target_end() target("infinirt-ali") set_kind("static") + set_default(false) add_deps("infini-utils") on_install(function (target) end) diff --git a/xmake/ascend.lua b/xmake/ascend.lua index f8b2fe9ac..3ef451e7b 100644 --- a/xmake/ascend.lua +++ b/xmake/ascend.lua @@ -80,6 +80,7 @@ target_end() target("infinirt-ascend") set_kind("static") + set_default(false) set_languages("cxx17") on_install(function (target) end) add_deps("infini-utils") diff --git a/xmake/bang.lua b/xmake/bang.lua index 0be267c94..0f9cb9479 100644 --- a/xmake/bang.lua +++ b/xmake/bang.lua @@ -54,6 +54,7 @@ target_end() target("infinirt-cambricon") set_kind("static") + set_default(false) add_deps("infini-utils") set_languages("cxx17") on_install(function (target) end) diff --git a/xmake/iluvatar.lua b/xmake/iluvatar.lua index fea75a5d8..7f04b62e7 100644 --- a/xmake/iluvatar.lua +++ b/xmake/iluvatar.lua @@ -81,6 +81,7 @@ target_end() target("infinirt-iluvatar") set_kind("static") + set_default(false) add_deps("infini-utils") on_install(function (target) end) diff --git a/xmake/metax.lua b/xmake/metax.lua index 3e700fff2..de4c39c70 100644 --- a/xmake/metax.lua +++ b/xmake/metax.lua @@ -147,6 +147,7 @@ target_end() target("infinirt-metax") set_kind("static") + set_default(false) set_languages("cxx17") on_install(function (target) end) add_deps("infini-utils") diff --git a/xmake/moore.lua b/xmake/moore.lua index 6e8800be8..be4403b0c 100644 --- a/xmake/moore.lua +++ b/xmake/moore.lua @@ -70,6 +70,7 @@ target_end() target("infinirt-moore") set_kind("static") + set_default(false) set_languages("cxx17") on_install(function (target) end) add_deps("infini-utils")