Unverified Commit 191cf2d1 authored by Jinze Xue's avatar Jinze Xue Committed by GitHub
Browse files

nvcc<11 does not have (if constexpr) c++17 feature (#574)

parent 4d10aceb
...@@ -96,7 +96,7 @@ def cuda_extension(build_all=False): ...@@ -96,7 +96,7 @@ def cuda_extension(build_all=False):
pkg='torchani.cuaev', pkg='torchani.cuaev',
sources=glob.glob('torchani/cuaev/*.cu'), sources=glob.glob('torchani/cuaev/*.cu'),
include_dirs=maybe_download_cub(), include_dirs=maybe_download_cub(),
extra_compile_args={'cxx': ['-std=c++17'], 'nvcc': nvcc_args}) extra_compile_args={'cxx': ['-std=c++14'], 'nvcc': nvcc_args})
def cuaev_kwargs(): def cuaev_kwargs():
......
...@@ -306,7 +306,7 @@ __global__ void pairwiseDistance_backward_or_doublebackward( ...@@ -306,7 +306,7 @@ __global__ void pairwiseDistance_backward_or_doublebackward(
const DataT dely = pos_t[mol_idx][j][1] - pos_t[mol_idx][i][1]; const DataT dely = pos_t[mol_idx][j][1] - pos_t[mol_idx][i][1];
const DataT delz = pos_t[mol_idx][j][2] - pos_t[mol_idx][i][2]; const DataT delz = pos_t[mol_idx][j][2] - pos_t[mol_idx][i][2];
if constexpr (is_double_backward) { if (is_double_backward) {
auto& grad_force = grad_coord_or_force; auto& grad_force = grad_coord_or_force;
DataT grad_force_coord_Rij_item = (grad_force[mol_idx][j][0] - grad_force[mol_idx][i][0]) * delx / Rij + DataT grad_force_coord_Rij_item = (grad_force[mol_idx][j][0] - grad_force[mol_idx][i][0]) * delx / Rij +
(grad_force[mol_idx][j][1] - grad_force[mol_idx][i][1]) * dely / Rij + (grad_force[mol_idx][j][1] - grad_force[mol_idx][i][1]) * dely / Rij +
...@@ -689,7 +689,7 @@ __global__ void cuAngularAEVs_backward_or_doublebackward( ...@@ -689,7 +689,7 @@ __global__ void cuAngularAEVs_backward_or_doublebackward(
factor1 * grad_factor2_dist * sdz[kk] / Rik * fc_ijk + factor1 * grad_factor2_dist * sdz[kk] / Rik * fc_ijk +
factor1 * factor2 * fc_ij * grad_fc_ik * sdz[kk] / Rik); factor1 * factor2 * fc_ij * grad_fc_ik * sdz[kk] / Rik);
if constexpr (is_double_backward) { if (is_double_backward) {
int atomj_idx = d_Rij[start_idx + jj].j; int atomj_idx = d_Rij[start_idx + jj].j;
int atomk_idx = d_Rij[start_idx + kk].j; int atomk_idx = d_Rij[start_idx + kk].j;
auto& grad_force = grad_output; auto& grad_force = grad_output;
...@@ -741,7 +741,7 @@ __global__ void cuAngularAEVs_backward_or_doublebackward( ...@@ -741,7 +741,7 @@ __global__ void cuAngularAEVs_backward_or_doublebackward(
} }
} }
if constexpr (!is_double_backward) { if (!is_double_backward) {
auto& grad_coord = grad_input; auto& grad_coord = grad_input;
int atomi_idx = i; int atomi_idx = i;
atomicAdd(&grad_coord[mol_idx][atomi_idx][0], sdix_grad); atomicAdd(&grad_coord[mol_idx][atomi_idx][0], sdix_grad);
...@@ -833,7 +833,7 @@ __global__ void cuRadialAEVs_backward_or_doublebackward( ...@@ -833,7 +833,7 @@ __global__ void cuRadialAEVs_backward_or_doublebackward(
DataT fc_grad = -0.5 * (PI / aev_params.Rcr) * sin(PI * Rij / aev_params.Rcr); DataT fc_grad = -0.5 * (PI / aev_params.Rcr) * sin(PI * Rij / aev_params.Rcr);
DataT upstream_grad; DataT upstream_grad;
if constexpr (is_double_backward) { if (is_double_backward) {
upstream_grad = grad_dist[idx]; upstream_grad = grad_dist[idx];
} }
...@@ -844,7 +844,7 @@ __global__ void cuRadialAEVs_backward_or_doublebackward( ...@@ -844,7 +844,7 @@ __global__ void cuRadialAEVs_backward_or_doublebackward(
DataT GmR_grad = -EtaR * (-2 * ShfR + 2 * Rij) * GmR; DataT GmR_grad = -EtaR * (-2 * ShfR + 2 * Rij) * GmR;
DataT jacobian = GmR_grad * fc + GmR * fc_grad; DataT jacobian = GmR_grad * fc + GmR * fc_grad;
if constexpr (is_double_backward) { if (is_double_backward) {
atomicAdd(&grad_aev[mol_idx][i][type_j * aev_params.radial_sublength + ishfr], upstream_grad * jacobian); atomicAdd(&grad_aev[mol_idx][i][type_j * aev_params.radial_sublength + ishfr], upstream_grad * jacobian);
} else { } else {
upstream_grad = grad_aev[mol_idx][i][type_j * aev_params.radial_sublength + ishfr]; upstream_grad = grad_aev[mol_idx][i][type_j * aev_params.radial_sublength + ishfr];
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment