template SafeLog

oneflow/core/common/util.h

@@ -146,13 +146,6 @@ inline uint32_t NewRandomSeed() {
   return gen();
 }
 
-// Work around the following issue on Windows
-// https://stackoverflow.com/questions/33218522/cuda-host-device-variables
-// const float LOG_THRESHOLD = 1e-20;
-#define LOG_THRESHOLD (1e-20)
-#define MAX_WITH_LOG_THRESHOLD(x) ((x) > LOG_THRESHOLD ? (x) : LOG_THRESHOLD)
-#define SAFE_LOG(x) logf(MAX_WITH_LOG_THRESHOLD(x))
-
 #if defined(WITH_CUDA)
 #define DEVICE_TYPE_SEQ                  \
   OF_PP_MAKE_TUPLE_SEQ(DeviceType::kCPU) \

oneflow/core/kernel/kernel_util.cuh

@@ -9,6 +9,17 @@ __device__ T gpu_atomic_add(T* address, const T val);
 template<typename T>
 __device__ T gpu_atomic_max(T* address, const T val);
 
+template<typename T>
+__host__ __device__ T MaxWithLogThreshold(T x) {
+  const T threshold = 1e-20;
+  return x > threshold ? x : threshold;
+}
+
+template<typename T>
+__host__ __device__ T SafeLog(T x) {
+  return logf(MaxWithLogThreshold(x));
+}
+
 }  // namespace oneflow
 
 #endif  // ONEFLOW_CORE_KERNEL_KERNEL_UTIL_CUH_

oneflow/core/kernel/sparse_cross_entropy_loss_kernel.cpp

 #include "oneflow/core/kernel/sparse_cross_entropy_loss_kernel.h"
+#include "oneflow/core/kernel/kernel_util.cuh"
 
 namespace oneflow {
 
@@ -39,7 +40,7 @@ struct SparseCrossEntropyLossKernelUtil<DeviceType::kCPU, PredType, LabelType> {
       int64_t label = static_cast<int64_t>(labels[i]);
       CHECK_GE(label, 0);
       CHECK_LT(label, num_of_classes);
-      loss[i] = -SAFE_LOG(prediction[i * num_of_classes + label]);
+      loss[i] = -SafeLog(prediction[i * num_of_classes + label]);
     }
   }
 
@@ -48,7 +49,7 @@ struct SparseCrossEntropyLossKernelUtil<DeviceType::kCPU, PredType, LabelType> {
                        PredType* prediction_diff) {
     for (int64_t i = 0; i < instance_num; ++i) {
       int64_t label = static_cast<int64_t>(labels[i]);
-      PredType prob = MAX_WITH_LOG_THRESHOLD(prediction[i * num_of_classes + label]);
+      PredType prob = MaxWithLogThreshold(prediction[i * num_of_classes + label]);
       prediction_diff[i * num_of_classes + label] = -1 / prob;
     }
   }

oneflow/core/kernel/sparse_cross_entropy_loss_kernel.cu

@@ -2,6 +2,7 @@
 #include "oneflow/core/common/util.h"
 #include "oneflow/core/device/cuda_util.h"
 #include "oneflow/core/kernel/kernel_util.h"
+#include "oneflow/core/kernel/kernel_util.cuh"
 #include "oneflow/core/kernel/sparse_cross_entropy_loss_kernel.h"
 
 namespace oneflow {
@@ -17,7 +18,7 @@ __global__ void SparseCrossEntropyLossForwardGpu(const int64_t instance_num,
     int64_t label = static_cast<int64_t>(labels[i]);
     assert(label >= 0);
     assert(label < num_of_classes);
-    loss[i] = -SAFE_LOG(prediction[i * num_of_classes + label]);
+    loss[i] = -SafeLog(prediction[i * num_of_classes + label]);
   }
 }
 
@@ -30,7 +31,7 @@ __global__ void SparseCrossEntropyLossBackwardGpu(const int64_t instance_num,
   CUDA_1D_KERNEL_LOOP(i, instance_num) {
     int64_t label = static_cast<int64_t>(labels[i]);
     PredType prob = prediction[i * num_of_classes + label];
-    prediction_diff[i * num_of_classes + label] = -1 / MAX_WITH_LOG_THRESHOLD(prob);
+    prediction_diff[i * num_of_classes + label] = -1 / MaxWithLogThreshold(prob);
   }
 }