Optimize LayerNorm Forward (#6842)

* layer_norm forward

* test case

* rm useless

* int count to T count

* fix

* fix T mask to int mask, refine code

* refine

* refine

* test case

* format

* fix

Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

CMakeLists.txt

@@ -35,6 +35,7 @@ option(WITH_COCOAPI "Option to build with COCO API" ON)
 option(BUILD_GIT_VERSION "" ON)
 option(BUILD_PROFILER "" OFF)
 option(OF_SOFTMAX_USE_FAST_MATH "" ON)
+option(OF_LAYER_NORM_USE_FAST_MATH "" ON)
 option(TREAT_WARNINGS_AS_ERRORS "" ON)
 # Reference:
 # https://medium.com/@alasher/colored-c-compiler-output-with-ninja-clang-gcc-10bfe7f2b949
@@ -125,6 +126,9 @@ endif()
 if (OF_SOFTMAX_USE_FAST_MATH)
   add_definitions(-DOF_SOFTMAX_USE_FAST_MATH)
 endif()
+if (OF_LAYER_NORM_USE_FAST_MATH)
+  add_definitions(-DOF_LAYER_NORM_USE_FAST_MATH)
+endif()
 if (OF_FORCE_COLORED_DIAGNOSTICS)
   add_compile_options(
     $<$<COMPILE_LANGUAGE:CXX>:$<$<CXX_COMPILER_ID:GNU>:-fdiagnostics-color=always>>