softmax for fp16 with fp32 accumulator (apache#14098)

* softmax for fp16 with fp32 accumulator * return AType in kernel * add dtype * kernel * grad use in-out only when dtype override * simplify infer type * address comments
jessr92 · Feb 28, 2019 · 07a4bcf · 07a4bcf
1 parent 5cac999
commit 07a4bcf
Show file tree

Hide file tree

Showing 4 changed files with 326 additions and 71 deletions.
diff --git a/src/operator/mxnet_op.h b/src/operator/mxnet_op.h
@@ -249,6 +249,48 @@ inline int get_num_threads<cpu>(const int N) {
     LOG(FATAL) << "Unknown type enum " << type;            \
   }
 
+#define MXNET_REAL_ACC_TYPE_SWITCH(type, DType, AType, ...)\
+  switch (type) {                                          \
+  case mshadow::kFloat32:                                  \
+    {                                                      \
+      typedef float DType;                                 \
+      typedef double AType;                                \
+      {__VA_ARGS__}                                        \
+    }                                                      \
+    break;                                                 \
+  case mshadow::kFloat64:                                  \
+    {                                                      \
+      typedef double DType;                                \
+      typedef double AType;                                \
+      {__VA_ARGS__}                                        \
+    }                                                      \
+    break;                                                 \
+  case mshadow::kFloat16:                                  \
+    {                                                      \
+      typedef mshadow::half::half_t DType;                 \
+      typedef float AType;                                 \
+      {__VA_ARGS__}                                        \
+    }                                                      \
+    break;                                                 \
+  case mshadow::kUint8:                                    \
+    LOG(FATAL) << "This operation only support "           \
+                  "floating point types not uint8";        \
+    break;                                                 \
+  case mshadow::kInt8:                                     \
+    LOG(FATAL) << "This operation only support "           \
+                  "floating point types not int8";         \
+    break;                                                 \
+  case mshadow::kInt32:                                    \
+    LOG(FATAL) << "This operation only support "           \
+                  "floating point types, not int32";       \
+    break;                                                 \
+  case mshadow::kInt64:                                    \
+    LOG(FATAL) << "This operation only support "           \
+                  "floating point types, not int64";       \
+    break;                                                 \
+  default:                                                 \
+    LOG(FATAL) << "Unknown type enum " << type;            \
+  }
 
 /*!
  * \brief assign the val to out according