-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[Feature] Tvm integration (dmlc#2367)
Co-authored-by: Zihao Ye <[email protected]>
- Loading branch information
Showing
15 changed files
with
435 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,34 @@ | ||
| cmake_minimum_required(VERSION 3.5) | ||
|
|
||
| project(featgraph C CXX) | ||
| message(STATUS "Start configuring project ${PROJECT_NAME}") | ||
|
|
||
| # Find CUDA | ||
| include(../cmake/util/FindCUDA.cmake) | ||
| find_cuda(ON) | ||
| message(STATUS "${CUDA_INCLUDE_DIRS}") | ||
|
|
||
| add_custom_target( | ||
| featgraph_kernel | ||
| COMMAND python ${CMAKE_CURRENT_SOURCE_DIR}/pack_featgraph.py | ||
| COMMENT "Creating featgraph kernels..." | ||
| ) | ||
|
|
||
| set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -O2 -fPIC") | ||
| file(GLOB FEATGRAPH_SRC | ||
| src/featgraph.cc | ||
| src/tvm_runtime_pack.cc | ||
| ) | ||
| add_library(featgraph_runtime SHARED ${FEATGRAPH_SRC}) | ||
| target_include_directories(featgraph_runtime PRIVATE ${CUDA_INCLUDE_DIRS}) | ||
| target_include_directories(featgraph_runtime PRIVATE "./include") | ||
| target_include_directories(featgraph_runtime PRIVATE "../third_party/tvm/include") | ||
| target_include_directories(featgraph_runtime PRIVATE "../third_party/tvm/3rdparty/dmlc-core/include") | ||
| target_include_directories(featgraph_runtime PRIVATE "../third_party/tvm/3rdparty/dlpack/include") | ||
| target_link_libraries(featgraph_runtime "dl" # dynamic linking | ||
| ${CUDA_CUDART_LIBRARY} | ||
| ${CUDA_CUDA_LIBRARY} | ||
| ${CUDA_NVRTC_LIBRARY}) | ||
| add_dependencies(featgraph_runtime featgraph_kernel) | ||
|
|
||
| install(TARGETS featgraph_runtime LIBRARY DESTINATION lib) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,21 @@ | ||
| # FeatGraph-DGL | ||
|
|
||
| FeatGraph is an efficient backend for Graph Neural Networks based on TVM. | ||
|
|
||
| - Original repo: https://github.com/amazon-research/FeatGraph | ||
| - SC2020 Paper: https://www.csl.cornell.edu/~zhiruz/pdfs/featgraph-sc2020.pdf | ||
|
|
||
| This folder contains the code for integrating featgraph kernels to DGL. | ||
|
|
||
| ## Usage | ||
|
|
||
| After building DGL with `USE_TVM=ON`, you should be able to run: | ||
| ```bash | ||
| python test.py | ||
| ``` | ||
| to verify correctness. | ||
|
|
||
| ## Reference | ||
|
|
||
| - [TVM Tutorial on Deploy TVM Module using C++ API](https://tvm.apache.org/docs/deploy/cpp_deploy.html). | ||
|
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,25 @@ | ||
| /*! | ||
| * Copyright (c) 2020 by Contributors | ||
| * \file featgraph/include/featgraph.h | ||
| * \brief FeatGraph kernel headers. | ||
| */ | ||
| #ifndef FEATGRAPH_H_ | ||
| #define FEATGRAPH_H_ | ||
|
|
||
| #include <dlpack/dlpack.h> | ||
|
|
||
| namespace dgl { | ||
| namespace featgraph { | ||
|
|
||
| /* \brief Load Featgraph module from given path. */ | ||
| void LoadFeatGraphModule(const std::string& path); | ||
|
|
||
| /* \brief Call Featgraph's SDDMM kernel. */ | ||
| void SDDMMTreeReduction(DLManagedTensor* row, DLManagedTensor* col, | ||
| DLManagedTensor* lhs, DLManagedTensor* rhs, | ||
| DLManagedTensor* out); | ||
|
|
||
| } // namespace featgraph | ||
| } // namespace dgl | ||
|
|
||
| #endif // FEATGRAPH_H_ |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,40 @@ | ||
| """ Export featgraph kernels to a shared library. """ | ||
| import tvm | ||
| from sddmm import sddmm_tree_reduction_gpu | ||
|
|
||
|
|
||
| def get_sddmm_kernels_gpu(idtypes, dtypes): | ||
| """ | ||
| Parameters | ||
| ---------- | ||
| idtypes: List[str] | ||
| Possible index types. | ||
| dtypes: List[str] | ||
| Possible data types. | ||
| Returns | ||
| ------- | ||
| List[IRModule]: | ||
| The list of IRModules. | ||
| """ | ||
| ret = [] | ||
| # SDDMM Tree Reduction | ||
| for dtype in dtypes: | ||
| for idtype in idtypes: | ||
| ret.append(sddmm_tree_reduction_gpu(idtype, dtype)) | ||
|
|
||
| return ret | ||
|
|
||
|
|
||
| if __name__ == '__main__': | ||
| binary_path = 'libfeatgraph_kernels.so' | ||
| kernels = [] | ||
| idtypes = ['int32', 'int64'] | ||
| dtypes = ['float16', 'float64', 'float32', 'int32', 'int64'] | ||
|
|
||
| kernels += get_sddmm_kernels_gpu(idtypes, dtypes) | ||
|
|
||
| # build kernels and export the module to libfeatgraph_kernels.so | ||
| module = tvm.build(kernels, target='cuda', target_host='llvm') | ||
| module.export_library(binary_path) | ||
|
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,53 @@ | ||
| """ The compute function and schedules for SDDMM kernels written in TVM. """ | ||
| import tvm | ||
| from tvm import te | ||
|
|
||
|
|
||
| def sddmm_tree_reduction_gpu(idx_type, feat_type): | ||
| """ SDDMM kernels on GPU optimized with Tree Reduction. | ||
| Parameters | ||
| ---------- | ||
| idx_type : str | ||
| The data type for indexing tensors. | ||
| feat_type : str | ||
| The data type of feature tensor. | ||
| Returns | ||
| ------- | ||
| IRModule | ||
| The result IRModule. | ||
| """ | ||
| # define vars and placeholders | ||
| nnz = te.var('nnz', idx_type) | ||
| num_rows = te.var('num_rows', idx_type) | ||
| num_cols = te.var('num_cols', idx_type) | ||
| H = te.var('num_heads', idx_type) | ||
| D = te.var('feat_len', idx_type) | ||
| row = te.placeholder((nnz,), idx_type, 'row') | ||
| col = te.placeholder((nnz,), idx_type, 'col') | ||
| ufeat = te.placeholder((num_rows, H, D), feat_type, 'ufeat') | ||
| vfeat = te.placeholder((num_cols, H, D), feat_type, 'vfeat') | ||
| # define edge computation function | ||
| def edge_func(eid, h, i): | ||
| k = te.reduce_axis((0, D), name='k') | ||
| return te.sum(ufeat[row[eid], h, k] * vfeat[col[eid], h, k], axis=k) | ||
| out = te.compute((nnz, H, tvm.tir.IntImm(idx_type, 1)), edge_func, name='out') | ||
| # define schedules | ||
| sched = te.create_schedule(out.op) | ||
| edge_axis, head_axis, _ = out.op.axis | ||
| reduce_axis = out.op.reduce_axis[0] | ||
| _, red_inner = sched[out].split(reduce_axis, factor=32) | ||
| edge_outer, edge_inner = sched[out].split(edge_axis, factor=32) | ||
| sched[out].bind(red_inner, te.thread_axis('threadIdx.x')) | ||
| sched[out].bind(edge_inner, te.thread_axis('threadIdx.y')) | ||
| sched[out].bind(edge_outer, te.thread_axis('blockIdx.x')) | ||
| sched[out].bind(head_axis, te.thread_axis('blockIdx.y')) | ||
| return tvm.lower(sched, [row, col, ufeat, vfeat, out], | ||
| name='SDDMMTreeReduction_{}_{}'.format(idx_type, feat_type)) | ||
|
|
||
|
|
||
| if __name__ == '__main__': | ||
| kernel0 = sddmm_tree_reduction_gpu('int32', 'float32') | ||
| print(kernel0) | ||
|
|
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,78 @@ | ||
| /*! | ||
| * Copyright (c) 2020 by Contributors | ||
| * \file featgraph/src/featgraph.cc | ||
| * \brief FeatGraph kernels. | ||
| */ | ||
| #include <tvm/runtime/module.h> | ||
| #include <tvm/runtime/packed_func.h> | ||
| #include <tvm/runtime/registry.h> | ||
| #include <dmlc/logging.h> | ||
| #include <featgraph.h> | ||
|
|
||
| namespace dgl { | ||
| namespace featgraph { | ||
|
|
||
| /* \brief Singleton that loads the featgraph module. */ | ||
| class FeatGraphModule { | ||
| public: | ||
| static FeatGraphModule* Global() { | ||
| static FeatGraphModule inst; | ||
| return &inst; | ||
| } | ||
|
|
||
| void Load(const std::string& path) { | ||
| mod = tvm::runtime::Module::LoadFromFile(path); | ||
| } | ||
|
|
||
| inline tvm::runtime::ModuleNode* Get() { | ||
| auto ret = mod.operator->(); | ||
| if (!ret) { | ||
| LOG(FATAL) << "FeatGraph module have not been loaded. " | ||
| << "Please set path of featgraph shared library."; | ||
| } | ||
| return ret; | ||
| } | ||
| private: | ||
| tvm::runtime::Module mod; | ||
| FeatGraphModule() {} | ||
| }; | ||
|
|
||
| /* \brief Load Featgraph module from given path. */ | ||
| void LoadFeatGraphModule(const std::string& path) { | ||
| FeatGraphModule::Global()->Load(path); | ||
| } | ||
|
|
||
| /* \brief Convert DLDataType to string. */ | ||
| inline std::string DTypeAsStr(const DLDataType& t) { | ||
| switch(t.code) { | ||
| case 0U: return "int" + std::to_string(t.bits); | ||
| case 1U: return "uint" + std::to_string(t.bits); | ||
| case 2U: return "float" + std::to_string(t.bits); | ||
| case 3U: return "bfloat" + std::to_string(t.bits); | ||
| default: LOG(FATAL) << "Type code " << t.code << " not recognized"; | ||
| } | ||
| } | ||
|
|
||
| /* \brief Get operator filename. */ | ||
| inline std::string GetOperatorName( | ||
| const std::string& base_name, | ||
| const DLDataType& dtype, | ||
| const DLDataType& idtype) { | ||
| return base_name + "_" + DTypeAsStr(dtype) + "_" + DTypeAsStr(idtype); | ||
| } | ||
|
|
||
| /* \brief Call FeatGraph's SDDMM kernel. */ | ||
| void SDDMMTreeReduction(DLManagedTensor* row, DLManagedTensor* col, | ||
| DLManagedTensor* lhs, DLManagedTensor* rhs, | ||
| DLManagedTensor* out) { | ||
| tvm::runtime::ModuleNode* mod = FeatGraphModule::Global()->Get(); | ||
| std::string f_name = GetOperatorName("SDDMMTreeReduction", | ||
| (row->dl_tensor).dtype, | ||
| (lhs->dl_tensor).dtype); | ||
| tvm::runtime::PackedFunc f = mod->GetFunction(f_name); | ||
| if (f != nullptr) | ||
| f(row, col, lhs, rhs, out); | ||
| } | ||
|
|
||
| } // namespace featgraph | ||
| } // namespace dgl |
Oops, something went wrong.