[NFC] polish colossalai/kernel/cuda_native/csrc/multi_tensor_sgd_kernel.cu code style (#978)

colossalai/kernel/cuda_native/csrc/multi_tensor_sgd_kernel.cu

-// modified from https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_sgd_kernel.cu
+// modified from
+// https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_sgd_kernel.cu
 #include <ATen/ATen.h>
 #include <ATen/AccumulateType.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <ATen/cuda/Exceptions.h>
-#include "multi_tensor_apply.cuh"
-#include "compat.h"
-
 #include <assert.h>
 #include <cuda_runtime.h>
 
+#include "compat.h"
+#include "multi_tensor_apply.cuh"
+
 #define BLOCK_SIZE 512
 #define ILP 4
 
@@ -28,24 +29,13 @@
  * wd_after_momentum : apply weight decay _after_ momentum instead of before
  **/
 template <int N, typename T_grad, typename T_weight>
-struct SGDFunctor
-{
+struct SGDFunctor {
   __device__ __forceinline__ void operator()(
-        int chunk_size,
-        volatile int *noop_gmem,
-        TensorListMetadata<N> &tl,
-        float wd,
-        float momentum,
-        float dampening,
-        float lr,
-        bool nesterov,
-        bool first_run,
-        bool wd_after_momentum,
-        float scale)
-    {
+      int chunk_size, volatile int *noop_gmem, TensorListMetadata<N> &tl,
+      float wd, float momentum, float dampening, float lr, bool nesterov,
+      bool first_run, bool wd_after_momentum, float scale) {
     // Early exit if we don't need to do anything
-        if (*noop_gmem)
-            return;
+    if (*noop_gmem) return;
 
     int tensor_loc = tl.block_to_tensor[blockIdx.x];
     int chunk_idx = tl.block_to_chunk[blockIdx.x];
@@ -61,8 +51,7 @@ struct SGDFunctor
     mom_in += chunk_idx * chunk_size;
 
     at::Half *model_weights_out = nullptr;
-        if (N == 4)
-        {
+    if (N == 4) {
       model_weights_out = (at::Half *)tl.addresses[3][tensor_loc];
       model_weights_out += chunk_idx * chunk_size;
     }
@@ -73,19 +62,15 @@ struct SGDFunctor
     float incoming_grads[ILP];
     float incoming_weights[ILP];
     float incoming_moms[ILP];
-        for (int i_start = 0;
-             i_start < n && i_start < chunk_size;
-             i_start += blockDim.x * ILP)
-        {
+    for (int i_start = 0; i_start < n && i_start < chunk_size;
+         i_start += blockDim.x * ILP) {
 #pragma unroll
-            for (int ii = 0; ii < ILP; ii++)
-            {
+      for (int ii = 0; ii < ILP; ii++) {
         incoming_grads[ii] = 0;
         incoming_weights[ii] = 0;
         incoming_moms[ii] = 0;
         int i = i_start + threadIdx.x + ii * blockDim.x;
-                if (i < n && i < chunk_size)
-                {
+        if (i < n && i < chunk_size) {
           incoming_grads[ii] = static_cast<float>(grad_in[i]) * scale;
           incoming_weights[ii] = static_cast<float>(weight_in[i]);
           incoming_moms[ii] = static_cast<float>(mom_in[i]);
@@ -98,19 +83,17 @@ struct SGDFunctor
 // Put another way, the STGs are dependent on the LDGs, but not on each other.
 // There is still compute ILP benefit from unrolling the loop though.
 #pragma unroll
-            for (int ii = 0; ii < ILP; ii++)
-            {
+      for (int ii = 0; ii < ILP; ii++) {
         int i = i_start + threadIdx.x + ii * blockDim.x;
-                if (i < n && i < chunk_size)
-                {
+        if (i < n && i < chunk_size) {
           // apply weight decay before momentum if necessary
           if (wd != 0.f && !wd_after_momentum)
             incoming_grads[ii] += wd * incoming_weights[ii];
 
-                    if (momentum != 0.f)
-                    {
+          if (momentum != 0.f) {
             if (!first_run)
-                            incoming_moms[ii] = incoming_moms[ii] * momentum + (1.f - dampening) * incoming_grads[ii];
+              incoming_moms[ii] = incoming_moms[ii] * momentum +
+                                  (1.f - dampening) * incoming_grads[ii];
             else  // initialize momentums to current incoming grads
               incoming_moms[ii] = incoming_grads[ii];
 
@@ -132,27 +115,18 @@ struct SGDFunctor
             model_weights_out[i] = static_cast<at::Half>(weight_in[i]);
 
           // also write out the new momentum
-                    if (momentum != 0.f)
-                        mom_in[i] = incoming_moms[ii];
+          if (momentum != 0.f) mom_in[i] = incoming_moms[ii];
         }
       }
     }
   }
 };
 
-void multi_tensor_sgd_cuda(
-    int chunk_size,
-    at::Tensor noop_flag,
+void multi_tensor_sgd_cuda(int chunk_size, at::Tensor noop_flag,
                            std::vector<std::vector<at::Tensor>> tensor_lists,
-    float wd,
-    float momentum,
-    float dampening,
-    float lr,
-    bool nesterov,
-    bool first_run,
-    bool wd_after_momentum,
-    float scale)
-{
+                           float wd, float momentum, float dampening, float lr,
+                           bool nesterov, bool first_run,
+                           bool wd_after_momentum, float scale) {
   auto num_tensors = tensor_lists.size();
   auto grad_type = tensor_lists[0][0].scalar_type();
   auto weight_type = tensor_lists[1][0].scalar_type();
@@ -162,7 +136,8 @@ void multi_tensor_sgd_cuda(
       TORCH_CHECK(tensor_lists[3][i].scalar_type() == at::ScalarType::Half,
                   "Additional output tensors should always be fp16.");
 
-    TORCH_CHECK(noop_flag.device() == tensor_lists[0][0].device(), "expected noop flag to be on the same device as tensors");
+  TORCH_CHECK(noop_flag.device() == tensor_lists[0][0].device(),
+              "expected noop flag to be on the same device as tensors");
 
   // We have 3 possibilities to handle here, in terms of
   // grad_type, param_type, momentum_type, requires_fp16_copy
@@ -176,22 +151,10 @@ void multi_tensor_sgd_cuda(
 
   // Case 1. fp16, fp16, fp16, No
   if (grad_type == at::ScalarType::Half &&
-        weight_type == at::ScalarType::Half &&
-        num_tensors == 3)
-    {
-        multi_tensor_apply<3>(
-            BLOCK_SIZE,
-            chunk_size,
-            noop_flag,
-            tensor_lists,
-            SGDFunctor<3, at::Half, at::Half>(),
-            wd,
-            momentum,
-            dampening,
-            lr,
-            nesterov,
-            first_run,
-            wd_after_momentum,
+      weight_type == at::ScalarType::Half && num_tensors == 3) {
+    multi_tensor_apply<3>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
+                          SGDFunctor<3, at::Half, at::Half>(), wd, momentum,
+                          dampening, lr, nesterov, first_run, wd_after_momentum,
                           scale);
   }
   // Case 2. fp16, fp32, fp32, No
@@ -214,68 +177,33 @@ void multi_tensor_sgd_cuda(
   // }
   // Case 2. fp32, fp32, fp32, No
   else if (grad_type == at::ScalarType::Float &&
-             weight_type == at::ScalarType::Float &&
-             num_tensors == 3)
-    {
-        multi_tensor_apply<3>(
-            BLOCK_SIZE,
-            chunk_size,
-            noop_flag,
-            tensor_lists,
-            SGDFunctor<3, float, float>(),
-            wd,
-            momentum,
-            dampening,
-            lr,
-            nesterov,
-            first_run,
-            wd_after_momentum,
+           weight_type == at::ScalarType::Float && num_tensors == 3) {
+    multi_tensor_apply<3>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
+                          SGDFunctor<3, float, float>(), wd, momentum,
+                          dampening, lr, nesterov, first_run, wd_after_momentum,
                           scale);
   }
   // Case 3. fp16, fp32, fp32, Yes
   else if (grad_type == at::ScalarType::Half &&
-             weight_type == at::ScalarType::Float &&
-             num_tensors == 4)
-    {
-        multi_tensor_apply<4>(
-            BLOCK_SIZE,
-            chunk_size,
-            noop_flag,
-            tensor_lists,
-            SGDFunctor<4, at::Half, float>(),
-            wd,
-            momentum,
-            dampening,
-            lr,
-            nesterov,
-            first_run,
-            wd_after_momentum,
+           weight_type == at::ScalarType::Float && num_tensors == 4) {
+    multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
+                          SGDFunctor<4, at::Half, float>(), wd, momentum,
+                          dampening, lr, nesterov, first_run, wd_after_momentum,
                           scale);
   }
   // Case 4. fp32, fp32, fp32, Yes
   else if (grad_type == at::ScalarType::Float &&
-             weight_type == at::ScalarType::Float &&
-             num_tensors == 4)
-    {
-        multi_tensor_apply<4>(
-            BLOCK_SIZE,
-            chunk_size,
-            noop_flag,
-            tensor_lists,
-            SGDFunctor<4, float, float>(),
-            wd,
-            momentum,
-            dampening,
-            lr,
-            nesterov,
-            first_run,
-            wd_after_momentum,
+           weight_type == at::ScalarType::Float && num_tensors == 4) {
+    multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
+                          SGDFunctor<4, float, float>(), wd, momentum,
+                          dampening, lr, nesterov, first_run, wd_after_momentum,
                           scale);
-    }
-    else
-    {
-        AT_ERROR("multi_tensor_sgd only supports some combinations of gradient & weight types. Given: ",
-                 "gradient: ", grad_type, ", weight: ", weight_type, ", num_lists: ", num_tensors);
+  } else {
+    AT_ERROR(
+        "multi_tensor_sgd only supports some combinations of gradient & weight "
+        "types. Given: ",
+        "gradient: ", grad_type, ", weight: ", weight_type,
+        ", num_lists: ", num_tensors);
   }
 
   AT_CUDA_CHECK(cudaGetLastError());