wehub/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensoradapter/include/tensoradapter.h
+113
View File
@@ -0,0 +1,113 @@
/**
* Copyright (c) 2020-2022 by Contributors
* @file tensoradapter.h
* @brief Header file for functions exposed by the adapter library.
*
* Functions in this library must be exported with extern "C" so that DGL can
* locate them with dlsym(3) (or GetProcAddress on Windows).
*/
#ifndef TENSORADAPTER_H_
#define TENSORADAPTER_H_
#ifdef DGL_USE_CUDA
#include <cuda_runtime.h>
#endif // DGL_USE_CUDA
namespace tensoradapter {
extern "C" {
/**
* @brief Allocate a piece of CPU memory via
* PyTorch's CPUAllocator
*
* @param nbytes The size to be allocated.
* @return Pointer to the allocated memory.
*/
void* CPURawAlloc(size_t nbytes);
/**
* @brief Free the CPU memory.
*
* @param ptr Pointer to the memory to be freed.
*/
void CPURawDelete(void* ptr);
#ifdef DGL_USE_CUDA
/**
* @brief Allocate a piece of GPU memory via
* PyTorch's THCCachingAllocator.
*
* @param nbytes The size to be allocated.
* @param stream The stream to be allocated on.
* @return Pointer to the allocated memory.
*/
void* CUDARawAlloc(size_t nbytes, cudaStream_t stream);
/**
* @brief Free the GPU memory.
*
* @param ptr Pointer to the memory to be freed.
*/
void CUDARawDelete(void* ptr);
/**
* @brief Get the current CUDA stream.
*/
cudaStream_t CUDACurrentStream();
/**
* @brief Let the caching allocator know which streams are using this tensor.
*
* @param ptr Pointer of the tensor to be recorded.
* @param stream The stream that is using this tensor.
* @param device_id Device of the tensor.
*/
void RecordStream(void* ptr, cudaStream_t stream, int device_id);
/**
* @brief Allocate a piece of pinned CPU memory via
* PyTorch's CachingHostAllocator.
*
* @param nbytes The size to be allocated.
* @param ctx Pointer to the PyTorch storage ctx ptr returned from the
* allocator.
* @param deleter Pointer to the delete function ptr returned from the
* allocator.
* @return Raw pointer to the allocated memory.
*/
void* CUDARawHostAlloc(size_t nbytes, void** ctx, void** raw_deleter);
/**
* @brief 'Free' the pinned CPU memory via
* inserting the memory block back to the free list.
*
* @param deleter Pointer to the delete function ptr returned from the
* allocator.
*/
void CUDARawHostDelete(void** raw_deleter);
/**
* @brief 'Record' a CUDA stream (usually from a copy kernel) for the pinned
* memory via PyTorch's CachingHostAllocator.
*
* @param data Pointer of the tensor to be recorded.
* @param ctx PyTorch storage ctx ptr returned from the allocator.
* @param stream The stream that currently consumes this tensor.
* @param device_id Device of the tensor.
*/
void CUDARecordHostAlloc(
void* data, void* ctx, cudaStream_t stream, int device_id);
/**
* @brief Release cached pinned memory allocations via cudaHostFree.
*/
void CUDAHostAllocatorEmptyCache();
#endif // DGL_USE_CUDA
}
}; // namespace tensoradapter
#endif // TENSORADAPTER_H_
tensoradapter/include/tensoradapter_exports.h
+16
View File
@@ -0,0 +1,16 @@
/**
* Copyright (c) 2020 by Contributors
* @file tensoradapter_exports.h
* @brief Header file for functions exposed by the adapter library.
*/
#ifndef TENSORADAPTER_EXPORTS_H_
#define TENSORADAPTER_EXPORTS_H_
#if defined(WIN32) || defined(_WIN32)
#define TA_EXPORTS __declspec(dllexport)
#else
#define TA_EXPORTS
#endif
#endif // TENSORADAPTER_EXPORTS_H_
tensoradapter/pytorch/CMakeLists.txt
+48
View File
@@ -0,0 +1,48 @@
cmake_minimum_required(VERSION 3.5)
project(tensoradapter_pytorch C CXX)
# Find PyTorch cmake files and PyTorch versions with the python interpreter $PYTHON_INTERP
# ("python3" or "python" if empty)
if(NOT PYTHON_INTERP)
find_program(PYTHON_INTERP NAMES python3 python)
endif()
message(STATUS "Using Python interpreter: ${PYTHON_INTERP}")
file(TO_NATIVE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/find_cmake.py FIND_CMAKE_PY)
execute_process(
COMMAND ${PYTHON_INTERP} ${FIND_CMAKE_PY}
OUTPUT_VARIABLE TORCH_PREFIX_VER
OUTPUT_STRIP_TRAILING_WHITESPACE)
message(STATUS "find_cmake.py output: ${TORCH_PREFIX_VER}")
list(GET TORCH_PREFIX_VER 0 TORCH_PREFIX)
list(GET TORCH_PREFIX_VER 1 TORCH_VER)
message(STATUS "Configuring for PyTorch ${TORCH_VER}")
if(USE_CUDA)
add_definitions(-DDGL_USE_CUDA)
endif()
set(Torch_DIR "${TORCH_PREFIX}/Torch")
message(STATUS "Setting directory to ${Torch_DIR}")
find_package(Torch REQUIRED)
set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${TORCH_C_FLAGS}")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -O0 -g3 -ggdb")
set(TORCH_TARGET_NAME "tensoradapter_pytorch_${TORCH_VER}")
file(GLOB TA_TORCH_SRC *.cpp)
add_library(${TORCH_TARGET_NAME} SHARED "${TA_TORCH_SRC}")
# use the library name rather than the path
set(TENSORADAPTER_TORCH_LIBS torch)
message(STATUS "tensoradapter found PyTorch includes: ${TORCH_INCLUDE_DIRS}")
message(STATUS "tensoradapter found PyTorch lib: ${TENSORADAPTER_TORCH_LIBS}")
target_include_directories(
${TORCH_TARGET_NAME} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/../include")
target_include_directories(
${TORCH_TARGET_NAME} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/../../third_party/dlpack/include")
target_include_directories(
${TORCH_TARGET_NAME} PRIVATE "${TORCH_INCLUDE_DIRS}")
target_link_libraries(${TORCH_TARGET_NAME} PRIVATE "${TENSORADAPTER_TORCH_LIBS}")
set_property(TARGET ${TORCH_TARGET_NAME} PROPERTY CXX_STANDARD 17)
message(STATUS "Configured target ${TORCH_TARGET_NAME}")
tensoradapter/pytorch/build.bat
+31
View File
@@ -0,0 +1,31 @@
REM Helper script to build tensor adapter libraries for PyTorch
@ECHO OFF
SETLOCAL EnableDelayedExpansion
MD "%BINDIR%\tensoradapter\pytorch"
DEL /S /Q build
MD build
PUSHD build
IF x%1x == xx GOTO single
FOR %%X IN (%*) DO (
DEL /S /Q *
"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% -DUSE_CUDA=%USE_CUDA% -DPYTHON_INTERP=%%X .. -G "Visual Studio 16 2019" || EXIT /B 1
msbuild tensoradapter_pytorch.sln /m /nr:false || EXIT /B 1
COPY /Y Release\*.dll "%BINDIR%\tensoradapter\pytorch" || EXIT /B 1
)
GOTO end
:single
DEL /S /Q *
"%CMAKE_COMMAND%" -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% -DUSE_CUDA=%USE_CUDA% .. -G "Visual Studio 16 2019" || EXIT /B 1
msbuild tensoradapter_pytorch.sln /m /nr:false || EXIT /B 1
COPY /Y Release\*.dll "%BINDIR%\tensoradapter\pytorch" || EXIT /B 1
:end
POPD
ENDLOCAL
tensoradapter/pytorch/build.sh
+31
View File
@@ -0,0 +1,31 @@
#!/bin/bash
# Helper script to build tensor adapter libraries for PyTorch
set -e
mkdir -p build
mkdir -p $BINDIR/tensoradapter/pytorch
cd build
if [ $(uname) = 'Darwin' ]; then
CPSOURCE=*.dylib
else
CPSOURCE=*.so
fi
CMAKE_FLAGS="-DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR -DTORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST -DUSE_CUDA=$USE_CUDA"
if [ $# -eq 0 ]; then
$CMAKE_COMMAND $CMAKE_FLAGS ..
make -j
cp -v $CPSOURCE $BINDIR/tensoradapter/pytorch
else
for PYTHON_INTERP in $@; do
TORCH_VER=$($PYTHON_INTERP -c 'import torch; print(torch.__version__.split("+")[0])')
mkdir -p $TORCH_VER
cd $TORCH_VER
$CMAKE_COMMAND $CMAKE_FLAGS -DPYTHON_INTERP=$PYTHON_INTERP ../..
make -j
cp -v $CPSOURCE $BINDIR/tensoradapter/pytorch
cd ..
done
fi
tensoradapter/pytorch/find_cmake.py
+11
View File
@@ -0,0 +1,11 @@
import os
import torch
cmake_prefix_path = getattr(
torch.utils,
"cmake_prefix_path",
os.path.join(os.path.dirname(torch.__file__), "share", "cmake"),
)
version = torch.__version__.split("+")[0]
print(";".join([cmake_prefix_path, version]))
tensoradapter/pytorch/torch.cpp
+106
View File
@@ -0,0 +1,106 @@
/**
* Copyright (c) 2020-2022 by Contributors
* @file torch/torch.cpp
* @brief Implementation of PyTorch adapter library.
*/
#include <c10/core/CPUAllocator.h>
#include <tensoradapter_exports.h>
#ifdef DGL_USE_CUDA
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CachingHostAllocator.h>
#include <c10/cuda/CUDACachingAllocator.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>
#endif // DGL_USE_CUDA
namespace tensoradapter {
extern "C" {
TA_EXPORTS void* CPURawAlloc(size_t nbytes) {
return c10::GetCPUAllocator()->raw_allocate(nbytes);
}
TA_EXPORTS void CPURawDelete(void* ptr) {
c10::GetCPUAllocator()->raw_deallocate(ptr);
}
#ifdef DGL_USE_CUDA
TA_EXPORTS void* CUDARawAlloc(size_t nbytes, cudaStream_t stream) {
at::globalContext().lazyInitDevice(at::kCUDA);
return c10::cuda::CUDACachingAllocator::raw_alloc_with_stream(nbytes, stream);
}
TA_EXPORTS void CUDARawDelete(void* ptr) {
c10::cuda::CUDACachingAllocator::raw_delete(ptr);
}
TA_EXPORTS cudaStream_t CUDACurrentStream() {
return at::cuda::getCurrentCUDAStream();
}
TA_EXPORTS void RecordStream(void* ptr, cudaStream_t stream, int device_id) {
c10::DataPtr data_ptr{
ptr, ptr, c10::cuda::CUDACachingAllocator::get()->raw_deleter(),
c10::Device(c10::DeviceType::CUDA, device_id)};
c10::cuda::CUDACachingAllocator::recordStream(
data_ptr,
// getStreamFromExternal doesn't exist before PyTorch 1.10, just copy it
// here
c10::cuda::CUDAStream(
c10::cuda::CUDAStream::UNCHECKED,
c10::Stream(
c10::Stream::UNSAFE,
c10::Device(c10::DeviceType::CUDA, device_id),
reinterpret_cast<int64_t>(stream))));
data_ptr.release_context();
}
class CUDAHostDeleter {
public:
explicit CUDAHostDeleter(std::unique_ptr<void, c10::DeleterFnPtr> ptr)
: ptr_(std::move(ptr)) {}
private:
std::unique_ptr<void, c10::DeleterFnPtr> ptr_;
};
TA_EXPORTS void* CUDARawHostAlloc(
size_t nbytes, void** ctx, void** raw_deleter) {
auto data_ptr = at::cuda::getCachingHostAllocator()->allocate(nbytes);
auto raw = data_ptr.get();
// Return the raw ctx ptr for recording event.
*ctx = data_ptr.get_context();
// Transfer ownership to raw_deleter.
auto* data_deleter = new CUDAHostDeleter(data_ptr.move_context());
*raw_deleter = static_cast<void*>(data_deleter);
return raw;
}
// Designated CUDAHostDeleter for CUDARawHostAlloc.
TA_EXPORTS void CUDARawHostDelete(void** raw_deleter) {
delete static_cast<CUDAHostDeleter*>(*raw_deleter);
*raw_deleter = nullptr;
}
TA_EXPORTS void CUDARecordHostAlloc(
void* ptr, void* ctx, cudaStream_t stream, int device_id) {
at::cuda::CachingHostAllocator_recordEvent(
ptr, ctx,
c10::cuda::CUDAStream(
c10::cuda::CUDAStream::UNCHECKED,
c10::Stream(
c10::Stream::UNSAFE,
c10::Device(c10::DeviceType::CUDA, device_id),
reinterpret_cast<int64_t>(stream))));
}
TA_EXPORTS void CUDAHostAllocatorEmptyCache() {
at::cuda::CachingHostAllocator_emptyCache();
}
#endif // DGL_USE_CUDA
};
}; // namespace tensoradapter