// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#include <algorithm>
#include <mutex>
#include <unordered_map>
#include "glog/logging.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/tensor_utils.h"

#include "paddle/common/enforce.h"
#include "paddle/phi/kernels/addmm_kernel.h"
#include "paddle/phi/kernels/elementwise_add_kernel.h"
#include "paddle/phi/kernels/impl/matmul_kernel_impl.h"
#include "paddle/phi/kernels/linear_v2_kernel.h"
#include "paddle/phi/kernels/reshape_kernel.h"
#include "paddle/phi/kernels/tile_kernel.h"

#ifdef PADDLE_WITH_HIP
#include <hip/hip_runtime.h>
#include <hip/hip_runtime_api.h>
#else
#include <cuda_runtime_api.h>  // NOLINT
#include "cuda.h"              // NOLINT
#endif

#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)

#include "paddle/common/flags.h"
#include "paddle/phi/backends/all_context.h"
#include "paddle/phi/common/amp_type_traits.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/enforce.h"
#include "paddle/utils/optional.h"
#if defined(PADDLE_WITH_CUDA)
#include "paddle/phi/backends/dynload/cublasLt.h"
#include "paddle/phi/backends/gpu/cuda/cuda_helper.h"
#include "paddle/phi/kernels/funcs/blas/blaslt_impl.cu.h"
#elif defined(PADDLE_WITH_HIP)
#include "paddle/phi/backends/dynload/hipblasLt.h"
#include "paddle/phi/backends/gpu/rocm/rocm_helper.h"
#include "paddle/phi/kernels/funcs/blas/blaslt_impl.hip.h"
#endif

#endif
COMMON_DECLARE_bool(use_legacy_linear);

namespace phi {

#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && \
    !defined(_WIN32) && CUDA_VERSION >= 11060

// Direct cublasLt matmul+bias, bypassing MatmulPlanner/DescriptorSetter/
// CublasLtBase. Uses persistent workspace from GPUContext.
template <typename T>
static void CublasLtMatmulBias(const GPUContext& ctx,
                               const T* x,
                               const T* w,
                               const T* bias,
                               T* out,
                               int64_t M,
                               int64_t N,
                               int64_t K,
                               bool trans_w) {
  using MT = typename MPTypeTrait<T>::Type;
  constexpr auto compute =
      std::is_same<T, double>::value ? CUBLAS_COMPUTE_64F : CUBLAS_COMPUTE_32F;
  const auto dtype = backends::gpu::ToCudaDataType<T>();
  const auto stype = backends::gpu::ToCudaDataType<MT>();

  MT alpha = static_cast<MT>(1), beta = static_cast<MT>(0);
  auto lt = ctx.cublaslt_handle();

  // op desc
  cublasLtMatmulDesc_t op = nullptr;
  PADDLE_ENFORCE_GPU_SUCCESS(
      dynload::cublasLtMatmulDescCreate(&op, compute, stype));

  // col-major: C(N,M) = A(weight) * B(input)
  cublasOperation_t ta = trans_w ? CUBLAS_OP_T : CUBLAS_OP_N;
  cublasOperation_t tb = CUBLAS_OP_N;
  PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulDescSetAttribute(
      op, CUBLASLT_MATMUL_DESC_TRANSA, &ta, sizeof(ta)));
  PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulDescSetAttribute(
      op, CUBLASLT_MATMUL_DESC_TRANSB, &tb, sizeof(tb)));

  cublasLtEpilogue_t epi = CUBLASLT_EPILOGUE_BIAS;
  PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulDescSetAttribute(
      op, CUBLASLT_MATMUL_DESC_EPILOGUE, &epi, sizeof(epi)));
  PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulDescSetAttribute(
      op, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias)));

  // matrix layouts (col-major)
  cublasLtMatrixLayout_t la = nullptr, lb = nullptr, lc = nullptr;
  if (trans_w) {
    PADDLE_ENFORCE_GPU_SUCCESS(
        dynload::cublasLtMatrixLayoutCreate(&la, dtype, K, N, K));
  } else {
    PADDLE_ENFORCE_GPU_SUCCESS(
        dynload::cublasLtMatrixLayoutCreate(&la, dtype, N, K, N));
  }
  PADDLE_ENFORCE_GPU_SUCCESS(
      dynload::cublasLtMatrixLayoutCreate(&lb, dtype, K, M, K));
  PADDLE_ENFORCE_GPU_SUCCESS(
      dynload::cublasLtMatrixLayoutCreate(&lc, dtype, N, M, N));

  // persistent workspace from context
  constexpr size_t kWsSize = 1024 * 1024;
  auto [ws, ws_sz] = ctx.cublaslt_workspace(kWsSize);

  // heuristic
  cublasLtMatmulPreference_t pref = nullptr;
  PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulPreferenceCreate(&pref));
  PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulPreferenceSetAttribute(
      pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws_sz, sizeof(ws_sz)));

  cublasLtMatmulHeuristicResult_t heur = {};
  int n_res = 0;
  PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmulAlgoGetHeuristic(
      lt, op, la, lb, lc, lc, pref, 1, &heur, &n_res));
  PADDLE_ENFORCE_GT(
      n_res,
      0,
      common::errors::Unavailable("No cublasLt GEMM algorithm available."));
  dynload::cublasLtMatmulPreferenceDestroy(pref);

  // execute
  PADDLE_ENFORCE_GPU_SUCCESS(dynload::cublasLtMatmul(lt,
                                                     op,
                                                     &alpha,
                                                     w,
                                                     la,
                                                     x,
                                                     lb,
                                                     &beta,
                                                     out,
                                                     lc,
                                                     out,
                                                     lc,
                                                     &heur.algo,
                                                     ws,
                                                     ws_sz,
                                                     ctx.stream()));

  // cleanup
  dynload::cublasLtMatmulDescDestroy(op);
  dynload::cublasLtMatrixLayoutDestroy(la);
  dynload::cublasLtMatrixLayoutDestroy(lb);
  dynload::cublasLtMatrixLayoutDestroy(lc);
}

#endif

template <typename T, typename Context>
void LinearV2Kernel(const Context& dev_ctx,
                    const DenseTensor& input,
                    const DenseTensor& weight,
                    const DenseTensor& bias,
                    const bool transpose_weight,
                    DenseTensor* out) {
  dev_ctx.template Alloc<T>(out);
  if (out->numel() == 0) {
    return;
  }

#if defined(PADDLE_WITH_CUDA) && !defined(PADDLE_WITH_HIP) && \
    !defined(_WIN32) && CUDA_VERSION >= 11060
  if (!FLAGS_use_legacy_linear) {
    const auto out_dim_original = out->dims();
    const auto [M, N, K] = canonicalize_dims(input, weight, transpose_weight);

    DenseTensor input_processed = input;
    DenseTensor weight_processed = weight;
    input_processed.Resize({M, K});
    if (transpose_weight) {
      weight_processed.Resize({N, K});
    } else {
      weight_processed.Resize({K, N});
    }
    out->Resize({M, N});

    if (N > 1 && K > 1) {
      DenseTensor bias_processed;
      if (bias.numel() != N) {
        TileKernel<T, Context>(dev_ctx, bias, {N}, &bias_processed);
      } else {
        bias_processed = bias;
      }
      CublasLtMatmulBias<T>(dev_ctx,
                            input_processed.data<T>(),
                            weight_processed.data<T>(),
                            bias_processed.data<T>(),
                            out->data<T>(),
                            M,
                            N,
                            K,
                            transpose_weight);
    } else {
      // When N=1 or K=1, {N,K} and {K,N} have identical memory layout,
      // so just reshape to {K,N} which is what AddmmKernel expects.
      weight_processed.Resize({K, N});
      DenseTensor bias_processed = bias;
      if (bias.numel() != (M * N)) {
        bias_processed.Resize({1, bias.numel()});
        TileKernel<T, Context>(
            dev_ctx, bias_processed, {M, 1}, &bias_processed);
      }
      AddmmKernel<T>(dev_ctx,
                     bias_processed,
                     input_processed,
                     weight_processed,
                     1.0f,
                     1.0f,
                     out);
    }
    out->Resize(out_dim_original);
  } else  // NOLINT
#endif
  {  // NOLINT
    std::vector<std::int64_t> input_dims_vec = vectorize(input.dims());
    std::vector<std::int64_t> weight_dims_vec = vectorize(weight.dims());

    MatMulFunction<Context, T>(dev_ctx,
                               input,
                               weight,
                               input_dims_vec,
                               weight_dims_vec,
                               out,
                               false,
                               transpose_weight);
    AddKernel<T, Context>(dev_ctx, *out, bias, out);
  }
}

}  // namespace phi

PD_REGISTER_KERNEL(linear_v2,
                   GPU,
                   ALL_LAYOUT,
                   phi::LinearV2Kernel,
                   float,
                   double,
                   phi::float16,
                   phi::bfloat16) {}