tokenspeed-deepjit 0.1.0__tar.gz

tokenspeed_deepjit-0.1.0/MANIFEST.in

@@ -0,0 +1,4 @@
+include README.md
+recursive-include src/tokenspeed_deepjit/cmake *.cmake
+recursive-include src/tokenspeed_deepjit/include/deep_jit *
+recursive-include tests CMakeLists.txt *.cc

tokenspeed_deepjit-0.1.0/PKG-INFO

@@ -0,0 +1,79 @@
+Metadata-Version: 2.4
+Name: tokenspeed-deepjit
+Version: 0.1.0
+Summary: Python packaging for the header-only DeepJIT C++ library.
+Author: TokenSpeed maintainers
+Project-URL: Homepage, https://github.com/deepseek-ai/DeepJIT
+Project-URL: Source, https://github.com/deepseek-ai/DeepJIT
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: C++
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+
+# tokenspeed-deepjit
+
+Python packaging for the header-only
+[deepseek-ai/DeepJIT](https://github.com/deepseek-ai/DeepJIT) C++20 library.
+The `tokenspeed-deepjit` distribution vendors the headers from upstream commit
+[`8b3ef868705a3792cc1a14ab570b539f72fd3d94`](https://github.com/deepseek-ai/DeepJIT/commit/8b3ef868705a3792cc1a14ab570b539f72fd3d94)
+and provides Python and CMake helpers for locating them.
+
+## Installation
+
+```bash
+pip install tokenspeed-deepjit
+```
+
+The wheel is platform-independent because DeepJIT is header-only. Applications
+still need the device toolchain and libraries required by the selected backend.
+CUDA consumers need CUDA headers 12.4 or newer, NVCC 12.9 or newer, and PyTorch
+with CUDA support. Ascend consumers need the CANN compiler and headers, ACL, and
+`torch_npu`. Pybind11 is also required when exposing DeepJIT through a Python
+extension.
+
+## Python helpers
+
+```python
+import tokenspeed_deepjit
+
+print(tokenspeed_deepjit.include_dir())
+print(tokenspeed_deepjit.cmake_prefix_path())
+```
+
+The distribution does not install a `deep_jit` Python import package. The
+upstream import package is empty; omitting it avoids conflicts with extensions
+that embed DeepJIT.
+
+## CMake usage
+
+```cmake
+find_package(Python COMPONENTS Interpreter REQUIRED)
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -c
+          "import tokenspeed_deepjit; print(tokenspeed_deepjit.cmake_prefix_path())"
+  OUTPUT_VARIABLE tokenspeed_deepjit_ROOT
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  COMMAND_ERROR_IS_FATAL ANY
+)
+
+list(PREPEND CMAKE_PREFIX_PATH "${tokenspeed_deepjit_ROOT}")
+find_package(deep_jit CONFIG REQUIRED)
+target_link_libraries(my_target PRIVATE deep_jit::deep_jit)
+```
+
+The imported target requests C++20 for C++ sources. CUDA projects must also set
+`CUDA_STANDARD 20` or pass the equivalent standard flag to NVCC.
+
+## Local build
+
+Populate the upstream headers before building:
+
+```bash
+git clone https://github.com/deepseek-ai/DeepJIT.git /tmp/DeepJIT
+git -C /tmp/DeepJIT checkout 8b3ef868705a3792cc1a14ab570b539f72fd3d94
+./prepare_headers.sh /tmp/DeepJIT
+python -m build
+```
+
+See the [upstream documentation](https://github.com/deepseek-ai/DeepJIT) for
+the runtime API and backend integration details.

tokenspeed_deepjit-0.1.0/README.md

@@ -0,0 +1,67 @@
+# tokenspeed-deepjit
+
+Python packaging for the header-only
+[deepseek-ai/DeepJIT](https://github.com/deepseek-ai/DeepJIT) C++20 library.
+The `tokenspeed-deepjit` distribution vendors the headers from upstream commit
+[`8b3ef868705a3792cc1a14ab570b539f72fd3d94`](https://github.com/deepseek-ai/DeepJIT/commit/8b3ef868705a3792cc1a14ab570b539f72fd3d94)
+and provides Python and CMake helpers for locating them.
+
+## Installation
+
+```bash
+pip install tokenspeed-deepjit
+```
+
+The wheel is platform-independent because DeepJIT is header-only. Applications
+still need the device toolchain and libraries required by the selected backend.
+CUDA consumers need CUDA headers 12.4 or newer, NVCC 12.9 or newer, and PyTorch
+with CUDA support. Ascend consumers need the CANN compiler and headers, ACL, and
+`torch_npu`. Pybind11 is also required when exposing DeepJIT through a Python
+extension.
+
+## Python helpers
+
+```python
+import tokenspeed_deepjit
+
+print(tokenspeed_deepjit.include_dir())
+print(tokenspeed_deepjit.cmake_prefix_path())
+```
+
+The distribution does not install a `deep_jit` Python import package. The
+upstream import package is empty; omitting it avoids conflicts with extensions
+that embed DeepJIT.
+
+## CMake usage
+
+```cmake
+find_package(Python COMPONENTS Interpreter REQUIRED)
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -c
+          "import tokenspeed_deepjit; print(tokenspeed_deepjit.cmake_prefix_path())"
+  OUTPUT_VARIABLE tokenspeed_deepjit_ROOT
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  COMMAND_ERROR_IS_FATAL ANY
+)
+
+list(PREPEND CMAKE_PREFIX_PATH "${tokenspeed_deepjit_ROOT}")
+find_package(deep_jit CONFIG REQUIRED)
+target_link_libraries(my_target PRIVATE deep_jit::deep_jit)
+```
+
+The imported target requests C++20 for C++ sources. CUDA projects must also set
+`CUDA_STANDARD 20` or pass the equivalent standard flag to NVCC.
+
+## Local build
+
+Populate the upstream headers before building:
+
+```bash
+git clone https://github.com/deepseek-ai/DeepJIT.git /tmp/DeepJIT
+git -C /tmp/DeepJIT checkout 8b3ef868705a3792cc1a14ab570b539f72fd3d94
+./prepare_headers.sh /tmp/DeepJIT
+python -m build
+```
+
+See the [upstream documentation](https://github.com/deepseek-ai/DeepJIT) for
+the runtime API and backend integration details.

tokenspeed_deepjit-0.1.0/pyproject.toml

@@ -0,0 +1,30 @@
+[build-system]
+requires = ["setuptools>=77", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "tokenspeed-deepjit"
+version = "0.1.0"
+description = "Python packaging for the header-only DeepJIT C++ library."
+readme = "README.md"
+requires-python = ">=3.9"
+authors = [
+  {name = "TokenSpeed maintainers"}
+]
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "Programming Language :: C++",
+]
+
+[project.urls]
+Homepage = "https://github.com/deepseek-ai/DeepJIT"
+Source = "https://github.com/deepseek-ai/DeepJIT"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+tokenspeed_deepjit = [
+  "cmake/*.cmake",
+  "include/deep_jit/**/*",
+]

tokenspeed_deepjit-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tokenspeed_deepjit-0.1.0/src/tokenspeed_deepjit/__init__.py

@@ -0,0 +1,27 @@
+"""Helpers for locating the packaged DeepJIT headers and CMake config."""
+
+from pathlib import Path
+
+__all__ = [
+    "__upstream_commit__",
+    "__version__",
+    "cmake_prefix_path",
+    "include_dir",
+]
+
+__version__ = "0.1.0"
+__upstream_commit__ = "8b3ef868705a3792cc1a14ab570b539f72fd3d94"
+
+
+def _package_root() -> Path:
+    return Path(__file__).resolve().parent
+
+
+def cmake_prefix_path() -> str:
+    """Return the directory containing the DeepJIT CMake package config."""
+    return str(_package_root())
+
+
+def include_dir() -> str:
+    """Return the directory containing the packaged ``deep_jit`` headers."""
+    return str(_package_root() / "include")

tokenspeed_deepjit-0.1.0/src/tokenspeed_deepjit/cmake/deep_jitConfig.cmake

@@ -0,0 +1,17 @@
+get_filename_component(_TOKENSPEED_DEEPJIT_PREFIX "${CMAKE_CURRENT_LIST_DIR}/.." ABSOLUTE)
+set(_TOKENSPEED_DEEPJIT_INCLUDE_DIR "${_TOKENSPEED_DEEPJIT_PREFIX}/include")
+
+if(NOT EXISTS "${_TOKENSPEED_DEEPJIT_INCLUDE_DIR}/deep_jit/runtime/runtime.hpp")
+  message(FATAL_ERROR "tokenspeed-deepjit is missing packaged DeepJIT headers")
+endif()
+
+if(NOT TARGET deep_jit::deep_jit)
+  add_library(deep_jit::deep_jit INTERFACE IMPORTED)
+  set_target_properties(deep_jit::deep_jit PROPERTIES
+    INTERFACE_COMPILE_FEATURES "cxx_std_20"
+    INTERFACE_INCLUDE_DIRECTORIES "${_TOKENSPEED_DEEPJIT_INCLUDE_DIR}"
+  )
+endif()
+
+set(deep_jit_VERSION "0.1.0")
+set(deep_jit_FOUND TRUE)

tokenspeed_deepjit-0.1.0/src/tokenspeed_deepjit/cmake/deep_jitConfigVersion.cmake

@@ -0,0 +1,11 @@
+set(PACKAGE_VERSION "0.1.0")
+
+if(PACKAGE_FIND_VERSION STREQUAL PACKAGE_VERSION)
+  set(PACKAGE_VERSION_EXACT TRUE)
+endif()
+
+if(PACKAGE_FIND_VERSION VERSION_LESS_EQUAL PACKAGE_VERSION)
+  set(PACKAGE_VERSION_COMPATIBLE TRUE)
+else()
+  set(PACKAGE_VERSION_COMPATIBLE FALSE)
+endif()

tokenspeed_deepjit-0.1.0/src/tokenspeed_deepjit/include/deep_jit/backend/ascend/backend.hpp

@@ -0,0 +1,185 @@
+#pragma once
+
+#include <filesystem>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <deep_jit/backend/ascend/device.hpp>
+#include <deep_jit/backend/ascend/kernel.hpp>
+#include <deep_jit/backend/ascend/options.hpp>
+#include <deep_jit/runtime/config.hpp>
+#include <deep_jit/runtime/runtime.hpp>
+#include <deep_jit/utils/command.hpp>
+#include <deep_jit/utils/env.hpp>
+#include <deep_jit/utils/exception.hpp>
+#include <deep_jit/utils/filesystem.hpp>
+#include <deep_jit/utils/gil.hpp>
+#include <deep_jit/utils/hash.hpp>
+#include <deep_jit/utils/json.hpp>
+#include <deep_jit/utils/str.hpp>
+
+namespace deep_jit {
+
+class Ascend {
+public:
+    using Device = ascend::Device;
+    using Kernel = ascend::Kernel;
+    using CompilerOptions = ascend::CompilerOptions;
+    using LaunchOptions = ascend::LaunchOptions;
+
+    struct CompilerInfo {
+        std::filesystem::path path;
+        std::string version;
+        std::filesystem::path linker_path;
+        std::string linker_version;
+
+        [[nodiscard]] std::string get_hash() const {
+            return hash::FNV1a().update(version).update(linker_version).get_hex_digest();
+        }
+
+        [[nodiscard]] json to_json() const {
+            return json::object_t {
+                {"path", path.string()},
+                {"version", version},
+                {"linker_path", linker_path.string()},
+                {"linker_version", linker_version},
+            };
+        }
+    };
+
+    struct Toolkit {
+        std::filesystem::path bisheng;
+        std::filesystem::path ld_lld;
+        std::filesystem::path asc_devkit_include;
+    };
+
+    Toolkit toolkit;
+    CompilerInfo compiler_info;
+
+    explicit Ascend(const Env&)
+        : toolkit(find_ascend_toolkit()),
+          compiler_info(get_compiler_info()) {}
+
+    [[nodiscard]] CompilerInfo get_compiler_info() const {
+        return {
+            .path = toolkit.bisheng,
+            .version = call_external_command(toolkit.bisheng.string() + " --version"),
+            .linker_path = toolkit.ld_lld,
+            .linker_version = call_external_command(toolkit.ld_lld.string() + " --version"),
+        };
+    }
+
+    void compile(const std::string& source,
+                 std::filesystem::path dir,
+                 const Env& env,
+                 const Config& config,
+                 const CompilerOptions& options) const {
+        // Release GIL to let other Python threads run
+        GilScopedRelease gil_release;
+
+        // Paths
+        dir = std::filesystem::absolute(dir).lexically_normal();
+        const auto source_path = dir / "kernel.asc";
+        const auto relocatable_path = dir / "kernel.rel.o";
+        const auto binary_path = dir / "kernel.o";
+        const bool print_compiler_command = env.get<bool>("JIT_DEBUG", false) or
+                                            env.get<bool>("JIT_PRINT_COMPILER_COMMAND", false);
+
+        // Write source code
+        write_file_sync(source_path, source);
+
+        // Build commands
+        std::vector<std::string> base_bisheng_args = {toolkit.bisheng.string()};
+        const auto bisheng_flags = options.get_bisheng_flags();
+        base_bisheng_args.insert(base_bisheng_args.end(), bisheng_flags.begin(), bisheng_flags.end());
+        for (const auto& include_dir: config.include_dirs)
+            base_bisheng_args.emplace_back("-I" + include_dir.string());
+        base_bisheng_args.emplace_back("-I" + toolkit.asc_devkit_include.string());
+
+        auto bisheng_args = base_bisheng_args;
+        bisheng_args.insert(bisheng_args.end(), {
+            source_path.string(),
+            "-o",
+            relocatable_path.string(),
+        });
+        const auto bisheng_command = str::join(bisheng_args);
+
+        std::vector<std::string> linker_args = {toolkit.ld_lld.string()};
+        const auto linker_flags = options.get_linker_flags();
+        linker_args.insert(linker_args.end(), linker_flags.begin(), linker_flags.end());
+        linker_args.insert(linker_args.end(), {
+            relocatable_path.string(),
+            "-o",
+            binary_path.string(),
+        });
+        const auto linker_command = str::join(linker_args);
+
+        // Compile and link
+        const auto cd_command = "cd " + dir.string() + " && ";
+        call_external_command(cd_command + bisheng_command, print_compiler_command);
+        call_external_command(cd_command + linker_command, print_compiler_command);
+        safe_remove_all(relocatable_path);
+        DJ_HOST_ASSERT(std::filesystem::is_regular_file(binary_path) and std::filesystem::file_size(binary_path) != 0,
+                       "Bisheng did not produce a valid Ascend binary: {}", binary_path.string());
+
+        // Dump assembly
+        if (options.dump_asm.value_or(false)) {
+            const auto asm_dir = dir / "asm";
+            make_dirs(asm_dir);
+            auto asm_args = base_bisheng_args;
+            asm_args.insert(asm_args.end(), {
+                "-save-temps",
+                source_path.string(),
+                "-o",
+                (asm_dir / "kernel.rel.o").string(),
+            });
+            call_external_command("cd " + asm_dir.string() + " && " + str::join(asm_args),
+                                  print_compiler_command);
+        }
+
+        // Write metadata
+        const json metadata = json::object_t {
+            {"command", bisheng_command + " && " + linker_command},
+            {"config", config.to_json()},
+            {"compiler_info", compiler_info.to_json()},
+            {"compiler_options", options.to_json()},
+        };
+        write_file_sync(dir / "meta.json", metadata.dump());
+    }
+
+    [[nodiscard]] static std::shared_ptr<Kernel> load(const std::filesystem::path& dir, const Env& env) {
+        return Kernel::load(dir, env);
+    }
+
+    static Toolkit find_ascend_toolkit() {
+        // Find Ascend home
+        std::filesystem::path home_path = get_env<std::string>("ASCEND_HOME_PATH");
+        if (home_path.empty())
+            home_path = get_env<std::string>("ASCEND_TOOLKIT_HOME");
+        if (home_path.empty() and std::filesystem::exists("/usr/local/Ascend/ascend-toolkit/latest"))
+            home_path = "/usr/local/Ascend/ascend-toolkit/latest";
+        if (home_path.empty() and std::filesystem::exists("/usr/local/Ascend/cann"))
+            home_path = "/usr/local/Ascend/cann";
+
+        DJ_HOST_ASSERT(not home_path.empty() and std::filesystem::is_directory(home_path),
+                       "Ascend toolkit home was not found");
+        home_path = std::filesystem::absolute(home_path).lexically_normal();
+
+        const auto bisheng = home_path / "bin/bisheng";
+        const auto ld_lld = home_path / "bin/ld.lld";
+        DJ_HOST_ASSERT(is_executable(bisheng), "Bisheng compiler is not executable: {}", bisheng.string());
+        DJ_HOST_ASSERT(is_executable(ld_lld), "Ascend linker is not executable: {}", ld_lld.string());
+
+        const auto asc_devkit_include = home_path / "aarch64-linux/asc/include/adv_api";
+        DJ_HOST_ASSERT(std::filesystem::is_directory(asc_devkit_include),
+                       "Ascend adv_api include directory was not found: {}", asc_devkit_include.string());
+        return {
+            .bisheng = bisheng,
+            .ld_lld = ld_lld,
+            .asc_devkit_include = asc_devkit_include,
+        };
+    }
+};
+
+}  // namespace deep_jit

tokenspeed_deepjit-0.1.0/src/tokenspeed_deepjit/include/deep_jit/backend/ascend/device.hpp

@@ -0,0 +1,79 @@
+#pragma once
+
+#include <cstdint>
+#include <string>
+
+#include <deep_jit/backend/ascend/driver.hpp>
+#include <deep_jit/utils/exception.hpp>
+
+namespace deep_jit::ascend {
+
+struct DeviceProperties {
+    std::string soc_name;
+    int64_t num_aicpu_cores = 0;
+    int64_t num_aicore_cores = 0;
+    int64_t num_cube_cores = 0;
+    int64_t num_vec_cores = 0;
+    int64_t num_lanes_per_warp = 0;
+    int64_t num_max_threads_per_vec_core = 0;
+    int64_t num_ubuf_bytes_per_vec_core = 0;
+    int64_t num_total_global_mem_bytes = 0;
+    int64_t num_l2_cache_bytes = 0;
+    int64_t npu_arch = 0;
+};
+
+class Device {
+    DeviceProperties prop{};
+    bool initialized = false;
+
+public:
+    const DeviceProperties& get_prop() {
+        if (not initialized) {
+            int device_index = 0;
+            DJ_ACL_CHECK(driver::lazy_aclrtGetDevice(&device_index));
+            const char* soc_name = driver::lazy_aclrtGetSocName();
+            DJ_HOST_ASSERT(soc_name != nullptr, "aclrtGetSocName returned null");
+            prop.soc_name = soc_name;
+
+#define DJ_ASCEND_GET_PROP(name, attribute) \
+            DJ_ACL_CHECK(driver::lazy_aclrtGetDeviceInfo(device_index, attribute, &prop.name))
+            DJ_ASCEND_GET_PROP(num_aicpu_cores, ACL_DEV_ATTR_AICPU_CORE_NUM);
+            DJ_ASCEND_GET_PROP(num_aicore_cores, ACL_DEV_ATTR_AICORE_CORE_NUM);
+            DJ_ASCEND_GET_PROP(num_cube_cores, ACL_DEV_ATTR_CUBE_CORE_NUM);
+            DJ_ASCEND_GET_PROP(num_vec_cores, ACL_DEV_ATTR_VECTOR_CORE_NUM);
+            DJ_ASCEND_GET_PROP(num_lanes_per_warp, ACL_DEV_ATTR_WARP_SIZE);
+            DJ_ASCEND_GET_PROP(num_max_threads_per_vec_core, ACL_DEV_ATTR_MAX_THREAD_PER_VECTOR_CORE);
+            DJ_ASCEND_GET_PROP(num_ubuf_bytes_per_vec_core, ACL_DEV_ATTR_UBUF_PER_VECTOR_CORE);
+            DJ_ASCEND_GET_PROP(num_total_global_mem_bytes, ACL_DEV_ATTR_TOTAL_GLOBAL_MEM_SIZE);
+            DJ_ASCEND_GET_PROP(num_l2_cache_bytes, ACL_DEV_ATTR_L2_CACHE_SIZE);
+            DJ_ASCEND_GET_PROP(npu_arch, ACL_DEV_ATTR_NPU_ARCH);
+#undef DJ_ASCEND_GET_PROP
+            initialized = true;
+        }
+        return prop;
+    }
+
+    int get_npu_arch() { return static_cast<int>(get_prop().npu_arch); }
+
+    int get_num_aicore_cores() { return static_cast<int>(get_prop().num_aicore_cores); }
+
+    int get_num_sms() { return get_num_aicore_cores(); }
+
+    int get_num_vec_cores() { return static_cast<int>(get_prop().num_vec_cores); }
+
+    int get_num_cube_cores() { return static_cast<int>(get_prop().num_cube_cores); }
+
+    int get_num_vec_cores_per_ai_core() {
+        const auto num_ai_cores = get_num_aicore_cores();
+        const auto num_vector_cores = get_num_vec_cores();
+        DJ_HOST_ASSERT(num_ai_cores > 0 and num_vector_cores % num_ai_cores == 0,
+                       "Ascend vector core count must be divisible by AI core count");
+        return num_vector_cores / num_ai_cores;
+    }
+
+    int64_t get_num_l2_cache_bytes() { return get_prop().num_l2_cache_bytes; }
+
+    int64_t get_num_ubuf_bytes_per_vec_core() { return get_prop().num_ubuf_bytes_per_vec_core; }
+};
+
+}  // namespace deep_jit::ascend

tokenspeed_deepjit-0.1.0/src/tokenspeed_deepjit/include/deep_jit/backend/ascend/driver.hpp

@@ -0,0 +1,37 @@
+#pragma once
+
+#include <acl/acl.h>
+
+#include <deep_jit/utils/exception.hpp>
+#include <deep_jit/utils/lazy.hpp>
+
+namespace deep_jit::ascend::driver {
+
+DJ_DECL_LAZY_DL_HANDLE(get_acl_handle, "libascendcl.so");
+
+DJ_DECL_LAZY_DL_FUNCTION(get_acl_handle, aclGetRecentErrMsg);
+DJ_DECL_LAZY_DL_FUNCTION(get_acl_handle, aclrtBinaryLoadFromFile);
+DJ_DECL_LAZY_DL_FUNCTION(get_acl_handle, aclrtBinaryUnLoad);
+DJ_DECL_LAZY_DL_FUNCTION(get_acl_handle, aclrtBinaryGetFunction);
+DJ_DECL_LAZY_DL_FUNCTION(get_acl_handle, aclrtLaunchKernelWithHostArgs);
+DJ_DECL_LAZY_DL_FUNCTION(get_acl_handle, aclrtGetDevice);
+DJ_DECL_LAZY_DL_FUNCTION(get_acl_handle, aclrtGetSocName);
+DJ_DECL_LAZY_DL_FUNCTION(get_acl_handle, aclrtGetDeviceInfo);
+DJ_DECL_LAZY_DL_FUNCTION(get_acl_handle, aclrtSynchronizeDevice);
+
+inline void check_acl(const aclError error, const char* expression) {
+    if (error == ACL_SUCCESS)
+        return;
+
+    const char* message = lazy_aclGetRecentErrMsg();
+    DJ_PANIC("{} failed with ACL error {}: {}",
+             expression,
+             static_cast<int>(error),
+             message == nullptr ? "unknown" : message);
+}
+
+#ifndef DJ_ACL_CHECK
+#define DJ_ACL_CHECK(expr) ::deep_jit::ascend::driver::check_acl((expr), #expr)
+#endif
+
+}  // namespace deep_jit::ascend::driver