Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

support GPU inference in torchscript model for v2.5 / v2.6 #188

Merged
merged 13 commits into from
May 18, 2021
Merged

support GPU inference in torchscript model for v2.5 / v2.6 #188

Show file tree
Hide file tree
Changes from 11 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
09f1c60
support GPU on torchscript
taoleicn May 12, 2021
69d28cd
support GPU on torchscript
taoleicn May 12, 2021
1be51bf
Merge branch 'torchscript_gpu_v2.5' of github.com:asappresearch/sru i…
taoleicn May 12, 2021
f2d8654
update version
taoleicn May 12, 2021
bcf6a0d
fix an issue
taoleicn May 13, 2021
72d9184
try to resolve a torchscript loading issue
taoleicn May 13, 2021
ed64427
update makefile
taoleicn May 13, 2021
d6daad6
update makefile
taoleicn May 13, 2021
85b4cd9
update version
taoleicn May 13, 2021
49eaa9a
rename registor op variables
taoleicn May 18, 2021
26e2579
update unit tests for CUDA/GPU
taoleicn May 18, 2021
1ca60b9
add backward tests
taoleicn May 18, 2021
32d9ebb
update version number
taoleicn May 18, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions sru/csrc/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,20 @@ target_compile_features(sru_cpu PRIVATE cxx_std_14)
# Link against LibTorch
target_link_libraries(sru_cpu "${TORCH_LIBRARIES}")

# Define our library target
add_library(sru_cuda SHARED sru_cuda_impl_dummy.cpp)
# Enable C++14
target_compile_features(sru_cuda PRIVATE cxx_std_14)
# Link against LibTorch
target_link_libraries(sru_cuda "${TORCH_LIBRARIES}")

add_executable(example_app main_test_cpp.cpp)
target_link_libraries(example_app "${TORCH_LIBRARIES}")
if (UNIX AND NOT APPLE)
target_link_libraries(example_app -Wl,--no-as-needed sru_cpu)
target_link_libraries(example_app -Wl,--no-as-needed sru_cuda)
else()
target_link_libraries(example_app -Wl,-all_load sru_cpu)
target_link_libraries(example_app -Wl,-all_load sru_cuda)
endif()
target_compile_features(example_app PRIVATE cxx_std_14)
8 changes: 4 additions & 4 deletions sru/csrc/sru_cpu_impl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,8 +90,8 @@ std::vector<at::Tensor> cpu_forward(
const float* reset_b_ptr = forget_b_ptr + hidden_size;
const float* U_ptr = U.data_ptr<float>();
const float* x_ptr = x.data_ptr<float>();
const float* pad_ptr = mask_pad.has_value() ?
mask_pad.value().data_ptr<float>() : NULL;
const bool* pad_ptr = mask_pad.has_value() ?
mask_pad.value().data_ptr<bool>() : NULL;

auto h = at::zeros({length, batch_size, hidden_size}, U.options());
auto c = c_init.clone();
Expand Down Expand Up @@ -175,8 +175,8 @@ std::vector<at::Tensor> cpu_bi_forward(
const float* reset_b_ptr = forget_b_ptr + hidden_size*2;
const float* U_ptr = U.data_ptr<float>();
const float* x_ptr = x.data_ptr<float>();
const float* pad_ptr = mask_pad.has_value() ?
mask_pad.value().data_ptr<float>() : NULL;
const bool* pad_ptr = mask_pad.has_value() ?
mask_pad.value().data_ptr<bool>() : NULL;

auto h = at::zeros({length, batch_size, hidden_size*2}, U.options());
auto c = c_init.clone();
Expand Down
Loading