More fixes

Author: peterjc123

aten/src/ATen/native/cuda/BinaryMulDivKernel.cu

@@ -66,6 +66,16 @@ static inline __host__ __device__ typename std::enable_if<!std::is_same<scalar_t
   floor_(scalar_t a) {
   return std::floor(a);
 }
+template <typename scalar_t>
+static inline __host__ __device__ typename std::enable_if<std::is_same<scalar_t, float>::value, scalar_t>::type
+  trunc_(scalar_t a) {
+  return std::truncf(a);
+}
+template <typename scalar_t>
+static inline __host__ __device__ typename std::enable_if<!std::is_same<scalar_t, float>::value, scalar_t>::type
+  trunc_(scalar_t a) {
+  return std::trunc(a);
+}
 template <typename scalar_t1, typename scalar_t2>
 static inline __host__ __device__ typename std::enable_if<std::is_same<scalar_t1, float>::value && std::is_same<scalar_t2, float>::value, scalar_t1>::type
   copysign_(scalar_t1 a, scalar_t2 b) {
@@ -80,6 +90,7 @@ static inline __host__ __device__ typename std::enable_if<!std::is_same<scalar_t
 #else
 #define ceil_ std::ceil
 #define floor_ std::floor
+#define trunc_ std::trunc
 #define copysign_ std::copysign
 #endif
 
@@ -121,13 +132,13 @@ void div_trunc_kernel_cuda(TensorIteratorBase& iter) {
       auto inv_b = accscalar_t(1.0) / iter.scalar_value<accscalar_t>(2);
       iter.remove_operand(2);
       gpu_kernel(iter, [inv_b] GPU_LAMBDA (scalar_t a) -> scalar_t {
-        return std::trunc(a * inv_b);
+        return trunc_(a * inv_b);
       });
     });
   } else {
     AT_DISPATCH_FLOATING_TYPES_AND2(kHalf, kBFloat16, dtype, "div_trunc_cuda", [&]() {
       gpu_kernel_with_scalars(iter, [] GPU_LAMBDA (scalar_t a, scalar_t b) -> scalar_t {
-        return std::trunc(a / b);
+        return trunc_(a / b);
       });
     });
   }

c10/util/BFloat16-math.h

@@ -17,14 +17,11 @@ inline c10::BFloat16 log(c10::BFloat16 a) { return std::log(float(a));}
 inline c10::BFloat16 log10(c10::BFloat16 a) { return std::log10(float(a));}
 inline c10::BFloat16 log1p(c10::BFloat16 a) { return std::log1p(float(a));}
 inline c10::BFloat16 log2(c10::BFloat16 a) { return std::log2(float(a));}
-inline c10::BFloat16 ceil(c10::BFloat16 a) { return std::ceil(float(a));}
 inline c10::BFloat16 cos(c10::BFloat16 a) { return std::cos(float(a));}
-inline c10::BFloat16 floor(c10::BFloat16 a) { return std::floor(float(a));}
 inline c10::BFloat16 nearbyint(c10::BFloat16 a) { return std::nearbyint(float(a));}
 inline c10::BFloat16 sin(c10::BFloat16 a) { return std::sin(float(a));}
 inline c10::BFloat16 tan(c10::BFloat16 a) { return std::tan(float(a));}
 inline c10::BFloat16 tanh(c10::BFloat16 a) { return std::tanh(float(a));}
-inline c10::BFloat16 trunc(c10::BFloat16 a) { return std::trunc(float(a));}
 inline c10::BFloat16 lgamma(c10::BFloat16 a) { return std::lgamma(float(a));}
 inline c10::BFloat16 sqrt(c10::BFloat16 a) { return std::sqrt(float(a));}
 inline c10::BFloat16 rsqrt(c10::BFloat16 a) { return 1.0 / std::sqrt(float(a));}
@@ -36,6 +33,15 @@ inline c10::BFloat16 pow(c10::BFloat16 a, double b) { return std::pow(float(a),
 #else
 inline c10::BFloat16 pow(c10::BFloat16 a, double b) { return std::pow(float(a), b);}
 #endif
+#if defined(_MSC_VER) && _MSC_VER >= 1928 && defined(__CUDACC__)
+inline c10::BFloat16 ceil(c10::BFloat16 a) { return std::ceilf(float(a));}
+inline c10::BFloat16 floor(c10::BFloat16 a) { return std::floorf(float(a));}
+inline c10::BFloat16 trunc(c10::BFloat16 a) { return std::truncf(float(a));}
+#else
+inline c10::BFloat16 ceil(c10::BFloat16 a) { return std::ceil(float(a));}
+inline c10::BFloat16 floor(c10::BFloat16 a) { return std::floor(float(a));}
+inline c10::BFloat16 trunc(c10::BFloat16 a) { return std::trunc(float(a));}
+#endif
 inline c10::BFloat16 pow(c10::BFloat16 a, c10::BFloat16 b) { return std::pow(float(a), float(b));}
 inline c10::BFloat16 fmod(c10::BFloat16 a, c10::BFloat16 b) { return std::fmod(float(a), float(b));}