Fixed builds.

aschi87 · Feb 12, 2016 · 958f634 · 958f634
1 parent 820cfc2
commit 958f634
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 7 deletions.
diff --git a/Source/Math/CuDnnConvolutionEngine.cuh b/Source/Math/CuDnnConvolutionEngine.cuh
@@ -5,7 +5,6 @@
 
 #pragma once
 
-#include <device_launch_parameters.h>
 #ifdef _MSC_VER
 #pragma warning(push)
 #pragma warning(disable : 4100)
@@ -108,6 +107,18 @@ namespace Microsoft { namespace MSR { namespace CNTK {
         *(reinterpret_cast<float4*>(dst)) = v;
     }
 
+    template <typename T>
+    __device__ __forceinline__ T Shuffle(T input, int srcLane)
+    {
+        // shfl is supported only on Kepler+. We really don't care about Fermi anymore but our build still has sm_20.
+#if __CUDA_ARCH__ >= 300
+        return cub::ShuffleIndex(input, srcLane);
+#else
+        assert(false);
+        return input;
+#endif
+    }
+
     struct Operations
     {
         template <typename T>
@@ -217,16 +228,16 @@ namespace Microsoft { namespace MSR { namespace CNTK {
             for (int i = 1; i < CUB_PTX_WARP_THREADS / BlockDimX; i *= 2)
             {
                 int srcLane = laneId + BlockDimX * i;
-                int n2 = cub::ShuffleIndex(n, srcLane);
+                int n2 = Shuffle(n, srcLane);
                 int nsum = n + n2;
                 T d[U];
 #pragma unroll
                 for (int k = 0; k < U; k++)
                 {
-                    d[k] = cub::ShuffleIndex(mean[k], srcLane) - mean[k];
+                    d[k] = Shuffle(mean[k], srcLane) - mean[k];
                     T dScaled = d[k] * n2 / nsum;
                     mean[k] += dScaled;
-                    m2[k] += cub::ShuffleIndex(m2[k], srcLane) + d[k] * n * dScaled;
+                    m2[k] += Shuffle(m2[k], srcLane) + d[k] * n * dScaled;
                 }
                 n = nsum;
             }
@@ -357,16 +368,16 @@ namespace Microsoft { namespace MSR { namespace CNTK {
             for (int i = 1; i < CUB_PTX_WARP_THREADS; i *= 2)
             {
                 int srcLane = laneId + i;
-                int n2 = cub::ShuffleIndex(n, srcLane);
+                int n2 = Shuffle(n, srcLane);
                 int nsum = n + n2;
                 T d[U];
 #pragma unroll
                 for (int k = 0; k < U; k++)
                 {
-                    d[k] = cub::ShuffleIndex(mean[k], srcLane) - mean[k];
+                    d[k] = Shuffle(mean[k], srcLane) - mean[k];
                     T dScaled = d[k] * n2 / nsum;
                     mean[k] += dScaled;
-                    m2[k] += cub::ShuffleIndex(m2[k], srcLane) + d[k] * n * dScaled;
+                    m2[k] += Shuffle(m2[k], srcLane) + d[k] * n * dScaled;
                 }
                 n = nsum;
             }

diff --git a/Tests/UnitTests/MathTests/ConvolutionEngineTests.cpp b/Tests/UnitTests/MathTests/ConvolutionEngineTests.cpp
@@ -658,6 +658,9 @@ BOOST_AUTO_TEST_SUITE(BatchNormalizationSuite)
 
 BOOST_AUTO_TEST_CASE(BatchNormalizationForwardTrain)
 {
+    if (!IsCuDnnSupported())
+        return;
+
     std::mt19937 rng(0);
     std::normal_distribution<float> nd;
 
@@ -779,6 +782,9 @@ BOOST_AUTO_TEST_CASE(BatchNormalizationForwardTrain)
 
 BOOST_AUTO_TEST_CASE(BatchNormalizationBackward)
 {
+    if (!IsCuDnnSupported())
+        return;
+
     std::mt19937 rng(0);
     std::normal_distribution<float> nd;