chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Multiple updates and refactorings (#231)

Browse Source
This commit is contained in:
Ray Wang
2025-11-21 17:49:47 +08:00
committed by GitHub
parent bb4424aad4
commit 38f8ef73a4
80 changed files with 3767 additions and 2103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/jit/compiler.hpp
+14 -5
View File
@@ -26,12 +26,17 @@ public:
static std::string library_version;
static std::string get_library_version() {
std::stringstream ss;
std::vector<char> buffer;
for (const auto& f: collect_files(library_include_path / "deep_gemm")) {
std::ifstream in(f, std::ios::binary);
ss << in.rdbuf();
DG_HOST_ASSERT(in.is_open());
// Append into the buffer
buffer.insert(buffer.end(),
std::istreambuf_iterator<char>(in),
std::istreambuf_iterator<char>());
}
return get_hex_digest(ss.str());
return get_hex_digest(buffer);
}
static void prepare_init(const std::string& library_root_path,
@@ -62,8 +67,8 @@ public:
flags = fmt::format("-std=c++{} --diag-suppress=39,161,174,177,186,940 "
"--ptxas-options=--register-usage-level=10",
get_env<int>("DG_JIT_CPP_STANDARD", 20));
if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0))
flags += " --ptxas-options=--verbose";
if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0) or get_env("DG_JIT_PTXAS_CHECK", 0))
flags += " --ptxas-options=--verbose,--warn-on-local-memory-usage";
if (get_env("DG_JIT_WITH_LINEINFO", 0))
flags += " -Xcompiler -rdynamic -lineinfo";
}
@@ -178,6 +183,10 @@ public:
DG_HOST_ASSERT(false and "NVCC compilation failed");
}
// Check local memory usage
if (get_env("DG_JIT_PTXAS_CHECK", 0))
DG_HOST_ASSERT(not std::regex_search(output, std::regex(R"(Local memory used)")));
// Print PTXAS log
if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0))
printf("%s", output.c_str());
csrc/jit/device_runtime.hpp
+21 -1
View File
@@ -1,6 +1,7 @@
#pragma once
#include <cublasLt.h>
#include <torch/version.h>
#include <ATen/cuda/CUDAContext.h>
#include "../utils/exception.hpp"
@@ -14,10 +15,24 @@ class DeviceRuntime {
// cuBLASLt utils
static constexpr size_t kCublasLtWorkspaceSize = 32 * 1024 * 1024;
public:
#if TORCH_VERSION_MAJOR > 2 or (TORCH_VERSION_MAJOR == 2 and TORCH_VERSION_MINOR >= 3)
// For PyTorch 2.3+, share the PyTorch cuBLASLt handle
DeviceRuntime() = default;
static cublasLtHandle_t get_cublaslt_handle() {
return at::cuda::getCurrentCUDABlasLtHandle();
}
static torch::Tensor get_cublaslt_workspace() {
return torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA));
}
#else
// Otherwise, create the cuBLASLt handle ourselves
cublasLtHandle_t cublaslt_handle{};
std::shared_ptr<torch::Tensor> cublaslt_workspace;
public:
explicit DeviceRuntime() {
cublaslt_workspace = std::make_shared<torch::Tensor>(torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA)));
DG_CUBLASLT_CHECK(cublasLtCreate(&cublaslt_handle));
@@ -34,6 +49,7 @@ public:
torch::Tensor get_cublaslt_workspace() const {
return *cublaslt_workspace;
}
#endif
std::shared_ptr<cudaDeviceProp> get_prop() {
if (cached_prop == nullptr) {
@@ -77,6 +93,10 @@ public:
return num_sms;
}
int get_l2_cache_size() {
return get_prop()->l2CacheSize;
}
void set_tc_util(const int& new_tc_util) {
DG_HOST_ASSERT(0 <= new_tc_util and new_tc_util <= 100);
tc_util = new_tc_util;
csrc/jit/handle.hpp
+42 -5
View File
@@ -2,12 +2,49 @@
#include <cuda.h>
#include <cuda_runtime.h>
#include <dlfcn.h>
#include <filesystem>
#include "../utils/exception.hpp"
#include "../utils/compatibility.hpp"
namespace deep_gemm {
// Lazy loading all driver symbols
static void* get_driver_handle() {
static void* handle = nullptr;
if (handle == nullptr) {
handle = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL);
DG_HOST_ASSERT(handle != nullptr and "Failed to load CUDA driver `libcuda.so.1`");
}
return handle;
}
// Macro to define wrapper functions named `lazy_cu{API name}`
#define DECL_LAZY_CUDA_DRIVER_FUNCTION(name) \
template <typename... Args> \
static auto lazy_##name(Args&&... args) -> decltype(name(args...)) { \
using FuncType = decltype(&name); \
static FuncType func = nullptr; \
if (func == nullptr) { \
func = reinterpret_cast<FuncType>(dlsym(get_driver_handle(), #name)); \
DG_HOST_ASSERT(func != nullptr and "Failed to load CUDA driver API"); \
} \
return func(std::forward<decltype(args)>(args)...); \
}
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorName);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorString);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuFuncSetAttribute);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleLoad);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleUnload);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleGetFunction);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLaunchKernelEx);
#if DG_TENSORMAP_COMPATIBLE
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuTensorMapEncodeTiled);
#endif
#if CUDART_VERSION >= 12080 and defined(DG_JIT_USE_RUNTIME_API)
// Use CUDA runtime API
@@ -80,8 +117,8 @@ static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const s
LibraryHandle *library_opt = nullptr) {
LibraryHandle library;
KernelHandle kernel;
DG_CUDA_DRIVER_CHECK(cuModuleLoad(&library, cubin_path.c_str()));
DG_CUDA_DRIVER_CHECK(cuModuleGetFunction(&kernel, library, func_name.c_str()));
DG_CUDA_DRIVER_CHECK(lazy_cuModuleLoad(&library, cubin_path.c_str()));
DG_CUDA_DRIVER_CHECK(lazy_cuModuleGetFunction(&kernel, library, func_name.c_str()));
if (library_opt != nullptr)
*library_opt = library;
@@ -89,7 +126,7 @@ static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const s
}
static void unload_library(const LibraryHandle& library) {
const auto& error = cuModuleUnload(library);
const auto& error = lazy_cuModuleUnload(library);
DG_HOST_ASSERT(error == CUDA_SUCCESS or error == CUDA_ERROR_DEINITIALIZED);
}
@@ -97,7 +134,7 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
const cudaStream_t& stream, const int& smem_size,
const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim) {
if (smem_size > 0)
DG_CUDA_DRIVER_CHECK(cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size));
DG_CUDA_DRIVER_CHECK(lazy_cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size));
LaunchConfigHandle config;
config.gridDimX = grid_dim.x;
@@ -127,7 +164,7 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
template<typename... ActTypes>
static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) {
void *ptr_args[] = { &args... };
return cuLaunchKernelEx(&config, kernel, ptr_args, nullptr);
return lazy_cuLaunchKernelEx(&config, kernel, ptr_args, nullptr);
}
#endif