Skip to content

Commit

Permalink
Change to CUDNN_CALL (apache#6048)
Browse files Browse the repository at this point in the history 
* Change to CUDNN_CALL

* Fix lint
  • Loading branch information
@ZihengJiang @piiswrong
ZihengJiang authored and piiswrong committed May 1, 2017
1 parent 1460b16 commit f0bfbfa
Show file tree
Hide file tree
Showing 11 changed files with 836 additions and 847 deletions.
2 changes: 1 addition & 1 deletion include/mxnet/operator.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ struct OpContext {

/*!
* \brief Operator interface.
* Operator defins basic operation unit of optimized computation graph in mxnet.
* Operator defines basic operation unit of optimized computation graph in mxnet.
* This interface relies on pre-allocated memory in TBlob, the caller need to set
* the memory region in TBlob correctly before calling Forward and Backward.
*
Expand Down
104 changes: 51 additions & 53 deletions src/operator/cudnn_activation-inl.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,18 +36,16 @@ class CuDNNActivationOp : public Operator {
}
#if CUDNN_MAJOR >= 5
nan_prop_ = CUDNN_NOT_PROPAGATE_NAN;
CHECK_EQ(cudnnCreateActivationDescriptor(&desc_),
CUDNN_STATUS_SUCCESS);
CHECK_EQ(cudnnSetActivationDescriptor(desc_, mode_, nan_prop_, relu_ceil_),
CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnCreateActivationDescriptor(&desc_));
CUDNN_CALL(cudnnSetActivationDescriptor(desc_, mode_, nan_prop_, relu_ceil_));
#endif
}

~CuDNNActivationOp() {
if (init_cudnn_) {
CHECK_EQ(cudnnDestroyTensorDescriptor(shape_desc_), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnDestroyTensorDescriptor(shape_desc_));
#if CUDNN_MAJOR >= 5
CHECK_EQ(cudnnDestroyActivationDescriptor(desc_), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnDestroyActivationDescriptor(desc_));
#endif
}
}
Expand Down Expand Up @@ -89,33 +87,33 @@ class CuDNNActivationOp : public Operator {
CHECK_EQ(s->dnn_handle_ownership_, mshadow::Stream<gpu>::OwnHandle);
if (!init_cudnn_) {
init_cudnn_ = true;
CHECK_EQ(cudnnCreateTensorDescriptor(&shape_desc_), CUDNN_STATUS_SUCCESS);
CHECK_EQ(cudnnSetTensor4dDescriptor(shape_desc_,
CUDNN_TENSOR_NCHW,
dtype_,
data.shape_[0],
data.shape_[1],
data.shape_[2],
data.shape_[3]), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnCreateTensorDescriptor(&shape_desc_));
CUDNN_CALL(cudnnSetTensor4dDescriptor(shape_desc_,
CUDNN_TENSOR_NCHW,
dtype_,
data.shape_[0],
data.shape_[1],
data.shape_[2],
data.shape_[3]));
}
#if CUDNN_MAJOR <= 4
CHECK_EQ(cudnnActivationForward(s->dnn_handle_,
mode_,
&alpha,
shape_desc_,
data.dptr_,
&beta,
shape_desc_,
out.dptr_), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnActivationForward(s->dnn_handle_,
mode_,
&alpha,
shape_desc_,
data.dptr_,
&beta,
shape_desc_,
out.dptr_));
#elif CUDNN_MAJOR >= 5
CHECK_EQ(cudnnActivationForward(s->dnn_handle_,
CUDNN_CALL(cudnnActivationForward(s->dnn_handle_,
desc_,
&alpha,
shape_desc_,
data.dptr_,
&beta,
shape_desc_,
out.dptr_), CUDNN_STATUS_SUCCESS);
&alpha,
shape_desc_,
data.dptr_,
&beta,
shape_desc_,
out.dptr_));
#endif
}

Expand Down Expand Up @@ -166,31 +164,31 @@ class CuDNNActivationOp : public Operator {
}
CHECK_EQ(s->dnn_handle_ownership_, mshadow::Stream<gpu>::OwnHandle);
#if CUDNN_MAJOR <= 4
CHECK_EQ(cudnnActivationBackward(s->dnn_handle_,
mode_,
&alpha,
shape_desc_,
output_data.dptr_,
shape_desc_,
grad.dptr_,
shape_desc_,
data.dptr_,
&beta,
shape_desc_,
input_grad.dptr_), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnActivationBackward(s->dnn_handle_,
mode_,
&alpha,
shape_desc_,
output_data.dptr_,
shape_desc_,
grad.dptr_,
shape_desc_,
data.dptr_,
&beta,
shape_desc_,
input_grad.dptr_));
#elif CUDNN_MAJOR >= 5
CHECK_EQ(cudnnActivationBackward(s->dnn_handle_,
desc_,
&alpha,
shape_desc_,
output_data.dptr_,
shape_desc_,
grad.dptr_,
shape_desc_,
data.dptr_,
&beta,
shape_desc_,
input_grad.dptr_), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnActivationBackward(s->dnn_handle_,
desc_,
&alpha,
shape_desc_,
output_data.dptr_,
shape_desc_,
grad.dptr_,
shape_desc_,
data.dptr_,
&beta,
shape_desc_,
input_grad.dptr_));
#endif
}

Expand Down
163 changes: 82 additions & 81 deletions src/operator/cudnn_batch_norm-inl.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,8 @@ class CuDNNBatchNormOp : public Operator {

~CuDNNBatchNormOp() {
if (init_cudnn_) {
CHECK_EQ(cudnnDestroyTensorDescriptor(io_desc_), CUDNN_STATUS_SUCCESS);
CHECK_EQ(cudnnDestroyTensorDescriptor(mean_desc_), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnDestroyTensorDescriptor(io_desc_));
CUDNN_CALL(cudnnDestroyTensorDescriptor(mean_desc_));
}
}

Expand Down Expand Up @@ -73,18 +73,18 @@ class CuDNNBatchNormOp : public Operator {
shape_[i] = 1;
}
}
CHECK_EQ(cudnnCreateTensorDescriptor(&io_desc_), CUDNN_STATUS_SUCCESS);
CHECK_EQ(cudnnCreateTensorDescriptor(&mean_desc_), CUDNN_STATUS_SUCCESS);
CHECK_EQ(cudnnSetTensor4dDescriptor(io_desc_,
CUDNN_TENSOR_NCHW,
dtype_,
shape_[0],
shape_[1],
shape_[2],
shape_[3]), CUDNN_STATUS_SUCCESS);
CHECK_EQ(cudnnDeriveBNTensorDescriptor(mean_desc_,
io_desc_,
CUDNN_BATCHNORM_SPATIAL), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnCreateTensorDescriptor(&io_desc_));
CUDNN_CALL(cudnnCreateTensorDescriptor(&mean_desc_));
CUDNN_CALL(cudnnSetTensor4dDescriptor(io_desc_,
CUDNN_TENSOR_NCHW,
dtype_,
shape_[0],
shape_[1],
shape_[2],
shape_[3]));
CUDNN_CALL(cudnnDeriveBNTensorDescriptor(mean_desc_,
io_desc_,
CUDNN_BATCHNORM_SPATIAL));
init_cudnn_ = true;
}

Expand Down Expand Up @@ -117,38 +117,38 @@ class CuDNNBatchNormOp : public Operator {
Tensor<gpu, 1, DTypeParam> save_inv_var =
out_data[cudnnbatchnorm::kInvVar]
.get_with_shape<gpu, 1, DTypeParam>(Shape1(shape_[1]), s);
CHECK_EQ(cudnnBatchNormalizationForwardTraining(s->dnn_handle_,
CUDNN_BATCHNORM_SPATIAL,
&a,
&b,
io_desc_,
x.dptr_,
io_desc_,
y.dptr_,
mean_desc_,
gamma.dptr_,
beta.dptr_,
1 - param_.momentum,
moving_mean.dptr_,
moving_inv_var.dptr_,
param_.eps,
save_mean.dptr_,
save_inv_var.dptr_), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnBatchNormalizationForwardTraining(s->dnn_handle_,
CUDNN_BATCHNORM_SPATIAL,
&a,
&b,
io_desc_,
x.dptr_,
io_desc_,
y.dptr_,
mean_desc_,
gamma.dptr_,
beta.dptr_,
1 - param_.momentum,
moving_mean.dptr_,
moving_inv_var.dptr_,
param_.eps,
save_mean.dptr_,
save_inv_var.dptr_));
} else {
CHECK_EQ(cudnnBatchNormalizationForwardInference(s->dnn_handle_,
CUDNN_BATCHNORM_SPATIAL,
&a,
&b,
io_desc_,
x.dptr_,
io_desc_,
y.dptr_,
mean_desc_,
gamma.dptr_,
beta.dptr_,
moving_mean.dptr_,
moving_inv_var.dptr_,
param_.eps), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnBatchNormalizationForwardInference(s->dnn_handle_,
CUDNN_BATCHNORM_SPATIAL,
&a,
&b,
io_desc_,
x.dptr_,
io_desc_,
y.dptr_,
mean_desc_,
gamma.dptr_,
beta.dptr_,
moving_mean.dptr_,
moving_inv_var.dptr_,
param_.eps));
}
})
}
Expand Down Expand Up @@ -197,25 +197,26 @@ class CuDNNBatchNormOp : public Operator {

if (param_.fix_gamma) gamma = 1.f;

CHECK_EQ(cudnnBatchNormalizationBackward(s->dnn_handle_,
CUDNN_BATCHNORM_SPATIAL,
&a,
&b,
&a,
req[cudnnbatchnorm::kGamma] == kWriteTo ? &b: &b_add,
io_desc_,
x.dptr_,
io_desc_,
dy.dptr_,
io_desc_,
dx.dptr_,
mean_desc_,
gamma.dptr_,
dgamma.dptr_,
dbeta.dptr_,
param_.eps,
save_mean.dptr_,
save_inv_var.dptr_), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnBatchNormalizationBackward(
s->dnn_handle_,
CUDNN_BATCHNORM_SPATIAL,
&a,
&b,
&a,
req[cudnnbatchnorm::kGamma] == kWriteTo ? &b: &b_add,
io_desc_,
x.dptr_,
io_desc_,
dy.dptr_,
io_desc_,
dx.dptr_,
mean_desc_,
gamma.dptr_,
dgamma.dptr_,
dbeta.dptr_,
param_.eps,
save_mean.dptr_,
save_inv_var.dptr_));
if (param_.fix_gamma) dgamma = 0.f;
})
#else // CUDNN_VERSION < 4007
Expand All @@ -237,23 +238,23 @@ class CuDNNBatchNormOp : public Operator {
CHECK_EQ(s->dnn_handle_ownership_, mshadow::Stream<gpu>::OwnHandle);

if (param_.fix_gamma) gamma = 1.f;
CHECK_EQ(cudnnBatchNormalizationBackward(s->dnn_handle_,
CUDNN_BATCHNORM_SPATIAL,
&a,
&b,
io_desc_,
x.dptr_,
io_desc_,
dy.dptr_,
io_desc_,
dx.dptr_,
mean_desc_,
gamma.dptr_,
dgamma.dptr_,
dbeta.dptr_,
param_.eps,
save_mean.dptr_,
save_inv_var.dptr_), CUDNN_STATUS_SUCCESS);
CUDNN_CALL(cudnnBatchNormalizationBackward(s->dnn_handle_,
CUDNN_BATCHNORM_SPATIAL,
&a,
&b,
io_desc_,
x.dptr_,
io_desc_,
dy.dptr_,
io_desc_,
dx.dptr_,
mean_desc_,
gamma.dptr_,
dgamma.dptr_,
dbeta.dptr_,
param_.eps,
save_mean.dptr_,
save_inv_var.dptr_));
if (param_.fix_gamma) dgamma = 0.f;
})
#endif
Expand Down
Loading

0 comments on commit f0bfbfa

Please sign in to comment.