clean up; add comment

example/21_gemm_layernorm/gemm_xdl_layernorm_single_kernel_fp16.cpp

@@ -2,7 +2,6 @@
 #include <numeric>
 #include <initializer_list>
 #include <cstdlib>
-#include <stdlib.h>
 #include <half.hpp>
 #include "check_err.hpp"
 #include "config.hpp"
@@ -17,6 +16,13 @@
 #include "reference_gemm_layernorm.hpp"
 #include "gemm_specialization.hpp"
 
+// This example demonstrate a single kernel that runs GEMM layer and laynorm in one fused kernel
+//
+// The GEMM + Layernorm implementation is a specialized kernel which allows fusing both layers
+// together given the condition GEMM extents N of MNK is spanned by a single workgroup. For example,
+// a kernel configured with NPerBlock = 128 allows to operate on all GEMM sizes if N <= 128
+//
+// D = Layernorm(acc_element_op(A * B + broadcast(bias)) + add) * broadcast(gamma) + broadcast(beta)
 template <ck::index_t... Is>
 using S = ck::Sequence<Is...>;
 

include/ck/tensor_operation/gpu/device/device_gemm_xdl_layernorm_cshuffle.hpp

@@ -22,6 +22,8 @@ namespace device {
 // Note: inter-wave loop scheduler is rolled out to c-shuffle version first. Becuase non c-shuffle
 // version currently has compiler issues with register spill which further causes validation
 // failures.
+//
+// D = Layernorm(acc_element_op(A * B + broadcast(bias)) + add) * broadcast(gamma) + broadcast(beta)
 template <typename ALayout,
           typename BLayout,
           typename CLayout,

include/ck/tensor_operation/gpu/grid/gridwise_gemm_xdl_layernorm_cshuffle_v1.hpp

@@ -15,7 +15,7 @@
 
 namespace ck {
 
-// D = Layernorm(A * B + broadcast(bias)) * broadcast(gamma) + broadcast(beta)
+// D = Layernorm(acc_element_op(A * B + broadcast(bias)) + add) * broadcast(gamma) + broadcast(beta)
 template <typename GridwiseGemm,
           typename FloatAB,
           typename FloatC,

library/include/ck/library/reference_tensor_operation/cpu/reference_gemm_layernorm.hpp

@@ -9,6 +9,7 @@ namespace ck {
 namespace tensor_operation {
 namespace host {
 
+// D = Layernorm(acc_element_op(A * B + broadcast(bias)) + add) * broadcast(gamma) + broadcast(beta)
 template <typename ADataType,
           typename BDataType,
           typename CDataType,
@@ -28,7 +29,6 @@ struct ReferenceGemmLayernorm : public device::BaseOperator
                                                 BElementwiseOperation,
                                                 element_wise::PassThrough>;
 
-    // D = Layernorm(acc + broadcast(bias)) * broadcast(gamma) + broadcast(beta)
     template <typename InDataType, typename OutDataType, typename ComputeDataType>
     static void RunLayernorm(Tensor<OutDataType>& result,
                              const Tensor<ComputeDataType>& acc, // MxN