     auto grad = ctx.Input<framework::LoDTensor>("Grad");
     auto learning_rate = ctx.Input<framework::LoDTensor>("LearningRate");
 
+    int64_t numel = param->numel();
+    int grid = (numel + LARS_BLOCK_SIZE - 1) / LARS_BLOCK_SIZE;
     const framework::Tensor* master_param = nullptr;
     framework::Tensor* master_param_out = nullptr;
+    const MT* master_param_data = nullptr;
+    MT* master_param_out_data = nullptr;
+
     if (multi_precision) {
       bool has_master =
           ctx.HasInput("MasterParam") && ctx.HasOutput("MasterParamOut");


#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/operators/amp/fp16_type_traits.h"
#include "paddle/fluid/operators/optimizers/lars_momentum_op.h"

namespace paddle {
namespace operators {

                            "the attr `multi_precision` is true"));
      master_param = ctx.Input<framework::Tensor>("MasterParam");
      master_param_out = ctx.Output<framework::Tensor>("MasterParamOut");
      master_param_data = master_param->data<MT>();
      master_param_out_data =
          master_param_out->mutable_data<MT>(ctx.GetPlace());
    }
    MT mu = static_cast<MT>(ctx.Attr<float>("mu"));
    MT lars_coeff = static_cast<MT>(ctx.Attr<float>("lars_coeff"));
    MT lars_weight_decay =
        static_cast<MT>(ctx.Attr<float>("lars_weight_decay"));
    MT epsilon = static_cast<MT>(ctx.Attr<float>("epsilon"));
    MT rescale_grad = static_cast<MT>(ctx.Attr<float>("rescale_grad"));

    auto* param_data = param->data<T>();
    auto* grad_data = grad->data<T>();
    auto* velocity_data = velocity->data<MT>();
    auto* lr = learning_rate->data<MT>();
    auto& cuda_ctx = ctx.template device_context<platform::CUDADeviceContext>();
    T* param_out_data = param_out->mutable_data<T>(ctx.GetPlace());
    MT* velocity_out_data = velocity_out->mutable_data<MT>(ctx.GetPlace());

#if CUDA_VERSION >= 11000
    /*
    Once model trainning with lars optimizer, whose principal implementation
    is achieved by following two steps:
      1. Figure out the L2 norm statistic result of grad data and param data.
      2. Update param and velocity data with usage of L2 norm statistic result.

    Orignally, these two steps were fulfilled by respective eigen function and
    cuda kernel, however the overhead of eigen function occupied much ratio in
    total, consequently affect the performance of lars op, make it necessary
    to combine 2 steps into one cuda kernel.
    Since the step1 is l2 norm statistic, grid level reduce is needed. To
    achieve this and continuous calculation of step 2 in only one global
    lanuch, essential basis is to control all grid-threads while running. Apart
    from normal lanuch form, cuda9.0 provides `cudaLaunchCooperativeKernel`
    api :
      - The thread quantity shall less than pyhsical SM limited threads
      - Launches a device function where thread blocks can cooperate and
        synchronize as they execute.
    */
    // Figure out how many blocks can be active in each sm.
    int num_blocks_per_sm = 0;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks_per_sm,
                                                  MomentumLarsKernel<T, MT>,
                                                  LARS_BLOCK_SIZE, sizeof(MT));
    int sm_num = cuda_ctx.GetSMCount();
    int grid_real =
        std::min(std::min(sm_num * num_blocks_per_sm, grid), LARS_BLOCK_SIZE);
    framework::Tensor tmp_buffer_t =
        ctx.AllocateTmpTensor<MT, platform::CUDADeviceContext>(
            {LARS_BLOCK_SIZE << 1}, cuda_ctx);
    auto* p_buffer = tmp_buffer_t.mutable_data<MT>(ctx.GetPlace());
    auto* g_buffer = p_buffer + LARS_BLOCK_SIZE;
    int grid_stride = LARS_BLOCK_SIZE * grid;
    int repeat_times = (numel + grid_stride - 1) / grid_stride - 1;
    int thresh = 0;

    // Uniform kernel parameter for cudaLaunchCooperativeKernel
    void* cuda_param[] = {
        reinterpret_cast<void*>(&param_data),
        reinterpret_cast<void*>(&grad_data),
        reinterpret_cast<void*>(&velocity_data),
        reinterpret_cast<void*>(&param_out_data),
        reinterpret_cast<void*>(&velocity_out_data),
        reinterpret_cast<void*>(&master_param_data),
        reinterpret_cast<void*>(&master_param_out_data),
        reinterpret_cast<void*>(&lr),
        reinterpret_cast<void*>(&p_buffer),
        reinterpret_cast<void*>(&g_buffer),
        reinterpret_cast<void*>(&mu),
        reinterpret_cast<void*>(&lars_coeff),
        reinterpret_cast<void*>(&lars_weight_decay),
        reinterpret_cast<void*>(&epsilon),
        reinterpret_cast<void*>(&rescale_grad),
        reinterpret_cast<void*>(&repeat_times),
        reinterpret_cast<void*>(&thresh),  // Just a placeholder
        reinterpret_cast<void*>(&numel)};
    // Lanuch all sm theads.
    cudaLaunchCooperativeKernel(
        reinterpret_cast<void*>(MomentumLarsKernel<T, MT>), grid_real,
        LARS_BLOCK_SIZE, cuda_param, 0, cuda_ctx.stream());
#else
    // Determine to read 4 fp16 or float data once, but 2 double data once.
    int grid_lars =
        sizeof(T) < sizeof(double)
            ? (numel + (LARS_BLOCK_SIZE << 2) - 1) / (LARS_BLOCK_SIZE << 2)
            : (numel + (LARS_BLOCK_SIZE << 1) - 1) / (LARS_BLOCK_SIZE << 1);

    int grid_norm = std::min(grid, LARS_BLOCK_SIZE);
    framework::Tensor p_buffer_t =
        ctx.AllocateTmpTensor<MT, platform::CUDADeviceContext>(
            {LARS_BLOCK_SIZE << 1}, cuda_ctx);
    auto* p_buffer = p_buffer_t.mutable_data<MT>(ctx.GetPlace());
    auto* g_buffer = p_buffer + LARS_BLOCK_SIZE;

    const int grid_stride = LARS_BLOCK_SIZE * grid_norm;
    const int repeat_times = (numel + grid_stride - 1) / grid_stride - 1;

    L2NormKernel<T, MT><<<grid_norm, LARS_BLOCK_SIZE, 0, cuda_ctx.stream()>>>(
        param_data, grad_data, p_buffer, g_buffer, repeat_times, numel,
        rescale_grad);

    MomentumLarsKernel<
        T, MT><<<grid_lars, LARS_BLOCK_SIZE, 0, cuda_ctx.stream()>>>(
        param_data, grad_data, velocity_data, param_out_data, velocity_out_data,
        master_param_data, master_param_out_data, lr, p_buffer, g_buffer, mu,
        lars_coeff, lars_weight_decay, epsilon, rescale_grad, 0, grid_norm,
        numel);  // 0 is just a placeholder.
#endif
  }
};

