Skip to content

fix(CMakeLists.txt): fix CUDA build for accelerated OP #53

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.

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

Merged
merged 2 commits into from
Aug 8, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

### Fixed

- Fix CUDA build for accelerated OP [@XuehaiPan](https://github.com/XuehaiPan) in [#53](https://github.com/metaopt/TorchOpt/pull/53).

### Removed

------
Expand Down
1 change: 1 addition & 0 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ find_package(CUDA)
if(CUDA_FOUND)
message(STATUS "Found CUDA, enabling CUDA support.")
enable_language(CUDA)
add_definitions(-D__CUDA_ENABLED__)

cuda_select_nvcc_arch_flags(CUDA_ARCH_FLAGS All)
list(APPEND CUDA_NVCC_FLAGS ${CUDA_ARCH_FLAGS})
Expand Down
16 changes: 8 additions & 8 deletions src/adam_op/adam_op.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
#include <pybind11/stl.h>

#include "include/adam_op/adam_op_impl_cpu.h"
#if defined(__CUDACC__)
#if defined(__CUDA_ENABLED__)
#include "include/adam_op/adam_op_impl_cuda.cuh"
#endif

Expand All @@ -30,7 +30,7 @@ TensorArray<3> adamForwardInplace(const torch::Tensor& updates,
const pyfloat_t b2, const pyfloat_t eps,
const pyfloat_t eps_root,
const pyuint_t count) {
#if defined(__CUDACC__)
#if defined(__CUDA_ENABLED__)
if (updates.device().is_cuda()) {
return adamForwardInplaceCUDA(updates, mu, nu, b1, b2, eps, eps_root,
count);
Expand All @@ -44,7 +44,7 @@ TensorArray<3> adamForwardInplace(const torch::Tensor& updates,
}
torch::Tensor adamForwardMu(const torch::Tensor& updates,
const torch::Tensor& mu, const pyfloat_t b1) {
#if defined(__CUDACC__)
#if defined(__CUDA_ENABLED__)
if (updates.device().is_cuda()) {
return adamForwardMuCUDA(updates, mu, b1);
}
Expand All @@ -58,7 +58,7 @@ torch::Tensor adamForwardMu(const torch::Tensor& updates,

torch::Tensor adamForwardNu(const torch::Tensor& updates,
const torch::Tensor& nu, const pyfloat_t b2) {
#if defined(__CUDACC__)
#if defined(__CUDA_ENABLED__)
if (updates.device().is_cuda()) {
return adamForwardNuCUDA(updates, nu, b2);
}
Expand All @@ -75,7 +75,7 @@ torch::Tensor adamForwardUpdates(const torch::Tensor& new_mu,
const pyfloat_t b1, const pyfloat_t b2,
const pyfloat_t eps, const pyfloat_t eps_root,
const pyuint_t count) {
#if defined(__CUDACC__)
#if defined(__CUDA_ENABLED__)
if (new_mu.device().is_cuda()) {
return adamForwardUpdatesCUDA(new_mu, new_nu, b1, b2, eps, eps_root, count);
}
Expand All @@ -90,7 +90,7 @@ torch::Tensor adamForwardUpdates(const torch::Tensor& new_mu,
TensorArray<2> adamBackwardMu(const torch::Tensor& dmu,
const torch::Tensor& updates,
const torch::Tensor& mu, const pyfloat_t b1) {
#if defined(__CUDACC__)
#if defined(__CUDA_ENABLED__)
if (dmu.device().is_cuda()) {
return adamBackwardMuCUDA(dmu, updates, mu, b1);
}
Expand All @@ -105,7 +105,7 @@ TensorArray<2> adamBackwardMu(const torch::Tensor& dmu,
TensorArray<2> adamBackwardNu(const torch::Tensor& dnu,
const torch::Tensor& updates,
const torch::Tensor& nu, const pyfloat_t b2) {
#if defined(__CUDACC__)
#if defined(__CUDA_ENABLED__)
if (dnu.device().is_cuda()) {
return adamBackwardNuCUDA(dnu, updates, nu, b2);
}
Expand All @@ -123,7 +123,7 @@ TensorArray<2> adamBackwardUpdates(const torch::Tensor& dupdates,
const torch::Tensor& new_nu,
const pyfloat_t b1, const pyfloat_t b2,
const pyuint_t count) {
#if defined(__CUDACC__)
#if defined(__CUDA_ENABLED__)
if (dupdates.device().is_cuda()) {
return adamBackwardUpdatesCUDA(dupdates, updates, new_mu, new_nu, b1, b2,
count);
Expand Down
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy