Environment variables to use native add-bias and bias-grad

apache · Dec 8, 2021 · 464a07e · 464a07e
1 parent 9d1df0f
commit 464a07e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/src/operator/nn/convolution.cu b/src/operator/nn/convolution.cu
@@ -58,7 +58,8 @@ void ConvolutionCompute<gpu>(const nnvm::NodeAttrs& attrs,
       CHECK_EQ(inputs[conv::kBias].shape_.ndim(), 1);
       auto layout = static_cast<mshadow::LayoutFlag>(param.layout.value());
       auto li     = cudnn::GetLayoutInfo(layout);
-      if (!cudnn::LegacyAddBias(ctx, li, outputs[conv::kOut], inputs[conv::kBias])) {
+      if (dmlc::GetEnv("MXNET_NATIVE_ADD_BIAS", li.channel_last) ||
+          !cudnn::LegacyAddBias(ctx, li, outputs[conv::kOut], inputs[conv::kBias])) {
         int k  = inputs[conv::kBias].shape_.Size();
         auto b = inputs[conv::kBias].reshape(cudnn::ExpandChannelDims(layout, k));
         BinaryBroadcastRTCCompute{"add"}(  // NOLINT(whitespace/braces)
@@ -142,7 +143,8 @@ void ConvolutionGradCompute<gpu>(const nnvm::NodeAttrs& attrs,
     if (ok && !param.no_bias && req[conv::kBias] != kNullOp) {
       auto li     = cudnn::GetLayoutInfo(static_cast<mshadow::LayoutFlag>(param.layout.value()));
       auto add_to = req[conv::kBias] == kAddTo;
-      if (!cudnn::LegacyBiasGrad(ctx, li, add_to, outputs[conv::kBias], inputs[0])) {
+      if (dmlc::GetEnv("MXNET_NATIVE_BIAS_GRAD", li.channel_last) ||
+          !cudnn::LegacyBiasGrad(ctx, li, add_to, outputs[conv::kBias], inputs[0])) {
         if (li.channel_last) {
           // This kernel should be faster.
           auto y_grad = FlattenAs2DHead<gpu, DType>(inputs[0], ctx);

diff --git a/src/operator/nn/deconvolution.cu b/src/operator/nn/deconvolution.cu
@@ -57,7 +57,8 @@ void DeconvolutionCompute<gpu>(const nnvm::NodeAttrs& attrs,
       CHECK_EQ(inputs[deconv::kBias].shape_.ndim(), 1);
       auto layout = static_cast<mshadow::LayoutFlag>(param.layout.value());
       auto li     = cudnn::GetLayoutInfo(layout);
-      if (!cudnn::LegacyAddBias(ctx, li, outputs[deconv::kOut], inputs[deconv::kBias])) {
+      if (dmlc::GetEnv("MXNET_NATIVE_ADD_BIAS", li.channel_last) ||
+          !cudnn::LegacyAddBias(ctx, li, outputs[deconv::kOut], inputs[deconv::kBias])) {
         int k  = inputs[deconv::kBias].shape_.Size();
         auto b = inputs[deconv::kBias].reshape(cudnn::ExpandChannelDims(layout, k));
         BinaryBroadcastRTCCompute{"add"}(  // NOLINT(whitespace/braces)
@@ -120,7 +121,8 @@ void DeconvolutionGradCompute<gpu>(const nnvm::NodeAttrs& attrs,
     if (ok && !param.no_bias && req[deconv::kBias] != kNullOp) {
       auto li     = cudnn::GetLayoutInfo(static_cast<mshadow::LayoutFlag>(param.layout.value()));
       auto add_to = req[conv::kBias] == kAddTo;
-      if (!cudnn::LegacyBiasGrad(ctx, li, add_to, outputs[deconv::kBias], inputs[0])) {
+      if (dmlc::GetEnv("MXNET_NATIVE_BIAS_GRAD", li.channel_last) ||
+          !cudnn::LegacyBiasGrad(ctx, li, add_to, outputs[deconv::kBias], inputs[0])) {
         if (li.channel_last) {
           // This kernel should be faster.
           auto y_grad = FlattenAs2DHead<gpu, DType>(inputs[0], ctx);