/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */

/*!
 * \file codegen_metal.cc
 */
#include "codegen_metal.h"

#include <tvm/ffi/cast.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/map.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <tvm/tirx/transform.h>

#include <algorithm>
#include <cmath>
#include <sstream>
#include <string>
#include <unordered_map>
#include <utility>

#include "../../../runtime/thread_storage_scope.h"
#include "../../../target/build_common.h"
#include "metal_fallback_module.h"

namespace tvm {
namespace codegen {

namespace {

Var GetSimdgroupBufferVar(const Expr& data) {
  if (const auto* var = data.as<VarNode>()) {
    return ffi::GetRef<Var>(var);
  }
  if (const auto* call = data.as<CallNode>();
      call && call->op.same_as(tirx::builtin::buffer_data()) && call->args.size() == 1) {
    const auto* buffer = call->args[0].as<VarNode>();
    TVM_FFI_ICHECK(buffer && buffer->ty.as<BufferTypeNode>())
        << "Metal simdgroup data operands expect buffer_data to project a BufferVar";
    return ffi::GetRef<Var>(buffer);
  }
  TVM_FFI_THROW(InternalError)
      << "Metal simdgroup data operands must be a Var or buffer_data(BufferVar), but got " << data;
}

}  // namespace

void CodeGenMetal::InitFuncState(const PrimFunc& f) {
  CodeGenC::InitFuncState(f);
  // analyze the data;
  for (Var arg : f->params) {
    if (arg->ty.as<PointerTypeNode>()) {
      alloc_storage_scope_[arg.get()] = "global";
    }
  }
}

CodeGenMetal::CodeGenMetal(Target target) : target_(target) {
  decl_stream << "#include <metal_stdlib>\n";
  decl_stream << "using namespace metal;\n\n";
  decl_stream << "union __TVMArgUnion {\n"
              << " int v_int[2];\n"
              << "};\n\n";
}

void CodeGenMetal::AddFunction(const GlobalVar& gvar, const PrimFunc& func) {
  // NOTE: There is no inter-function calls among Metal kernels.
  // For now we keep the metal codegen without inter-function call
  // process.
  // We can switch to follow the flow with inter-function call process
  // after the Metal function declaration is properly printed.
  // In Metal, for PrimFuncs with signature
  //    def func(A: Buffer, B: Buffer, x: int, y: float) -> None
  // where there are trailing pod parameters, the codegen emits a struct
  //    struct func_params{ x: int; y: float; }
  // for the function. In the flow of inter-function call process,
  // the struct will be emitted for every time a function is declared.
  // So consequently there are duplicate appearances of a same struct,
  // which makes the Metal compiler unable to recognize.

  // clear previous generated state.
  this->InitFuncState(func);
  // skip the first underscore, so SSA variable starts from _1
  name_supply_->FreshName("v_");

  // add to alloc buffer type.
  auto global_symbol = func->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol);
  TVM_FFI_ICHECK(global_symbol.has_value())
      << "CodeGenC: Expect PrimFunc to have the global_symbol attribute";

  // Function header.
  this->stream << "kernel void " << static_cast<std::string>(global_symbol.value()) << "(";

  // Buffer arguments
  size_t num_buffer = 0;
  size_t limit = target_->GetAttr<int64_t>("max_function_args").value();
  if (func->params.size() > limit) {
    LOG(WARNING) << "Probably you won't be able to execute your kernel due to high number of "
                    "buffers in the kernel";
  }
  for (size_t i = 0; i < func->params.size(); ++i, ++num_buffer) {
    Var v = func->params[i];
    if (!v->ty.as<PointerTypeNode>()) break;
    this->stream << "  ";
    std::string vid = AllocVarID(v.get());
    auto it = alloc_storage_scope_.find(v.get());
    if (it != alloc_storage_scope_.end()) {
      PrintStorageScope(it->second, this->stream);
    }
    PrintType(v->ty, this->stream);
    // Register handle data type
    // TODO(tvm-team): consider simply keep type info in the
    // type annotation(via a normalizing rewriting).
    if (auto* ptr = v->ty.as<PointerTypeNode>()) {
      if (auto* prim = ptr->element_type.as<PrimTypeNode>()) {
        RegisterHandleType(v.get(), ffi::GetRef<PrimType>(prim));
      }
    }
    this->stream << ' ' << vid << " [[ buffer(" << i << ") ]],\n";
  }
  // Setup normal arguments.
  size_t nargs = func->params.size() - num_buffer;
  std::string varg = name_supply_->FreshName("arg");
  if (nargs != 0) {
    std::string arg_buf_type = static_cast<std::string>(global_symbol.value()) + "_args_t";
    this->stream << "  constant " << arg_buf_type << "& " << varg << " [[ buffer(" << num_buffer
                 << ") ]],\n";
    // declare the struct
    decl_stream << "struct " << arg_buf_type << " {\n";
    for (size_t i = num_buffer; i < func->params.size(); ++i) {
      Var v = func->params[i];
      PrimType value_type = v->ty.as_or_throw<PrimType>();
      TVM_FFI_ICHECK(!value_type.IsVoid());
      std::string vid = AllocVarID(v.get());
      std::ostringstream vref;
      if (value_type.bits() == 32) {
        decl_stream << "  ";
        PrintType(value_type, decl_stream);
        decl_stream << " " << vid << "[2];\n";
        vref << varg << "." << vid << "[0]";
      } else if (value_type.bits() == 64) {
        decl_stream << "  ";
        PrintType(value_type, decl_stream);
        decl_stream << " " << vid << ";\n";
        vref << varg << "." << vid;
      } else {
        // For non 32bit type, ref through arg union.
        decl_stream << "  __TVMArgUnion " << vid << ";\n";
        vref << varg << "." << vid << ".v_";
        PrintType(value_type, vref);
      }
      var_idmap_[v.get()] = vref.str();
    }
    decl_stream << "};\n\n";
  }
  // Setup the thread group info.
  TVM_FFI_ICHECK_EQ(name_supply_->FreshName("threadIdx"), "threadIdx");
  TVM_FFI_ICHECK_EQ(name_supply_->FreshName("blockIdx"), "blockIdx");
  int work_dim = 0;
  auto launch_params =
      func->GetAttr<ffi::Array<ffi::String>>(tirx::attr::kKernelLaunchParams).value();
  for (const auto& tag : launch_params) {
    if (tag != runtime::launch_param::kUseDynamicSharedMemoryTag) {
      runtime::ThreadScope scope = runtime::ThreadScope::Create(tag);
      work_dim = std::max(work_dim, scope.dim_index + 1);
    }
  }

  if (work_dim != 0) {
    // use ushort by default for now
    stream << "  ";
    PrintType(PrimType::UInt(thread_index_bits_, work_dim), stream);
    stream << " blockIdx [[threadgroup_position_in_grid]],\n";
    stream << "  ";
    PrintType(PrimType::UInt(thread_index_bits_, work_dim), stream);
    stream << " threadIdx [[thread_position_in_threadgroup]]\n";
  }
  thread_work_dim_ = work_dim;

  // the function scope.
  stream << ") {\n";
  int func_scope = this->BeginScope();
  this->PrintStmt(func->body);
  this->EndScope(func_scope);
  this->PrintIndent();
  this->stream << "}\n\n";
}

void CodeGenMetal::BindThreadIndex(const IterVar& iv) {
  TVM_FFI_ICHECK(!var_idmap_.count(iv->var.get()));
  // if we only have threadIdx.x
  // metal will directly print as threadIdx
  std::string vname = iv->thread_tag;
  if (thread_work_dim_ <= 1) {
    vname = vname.substr(0, iv->thread_tag.length() - 2);
  }
  var_idmap_[iv->var.get()] = CastFromTo(vname, PrimType::UInt(thread_index_bits_), iv->var.ty());
}

void CodeGenMetal::PrintType(const PrimType& t, std::ostream& os) {  // NOLINT(*)
  int lanes = t.lanes();
  if (t.IsVoid()) {
    os << "void";
    return;
  }
  if (t == PrimType::Bool()) {
    os << "bool";
    return;
  }
  bool fail = false;
  if (t.code() == DLDataTypeCode::kDLFloat) {
    // Need to care about sizes and alignment of half3/float3 because tirx representation might not
    // be aware of Metal half3/float3 details and can treat them as just three elements,
    // while sizes and alignmnents of half3/float3 are one element more (half3-8 bytes/
    // float13 - 16bytes).
    // Example of problematic pattern: filling of threadgroup packed array using float3 elements
    // by threads concurrently can lead to datarace and wrong data in threadgroup shared array.
    // packed_(half3/float3) are exactly datatypes dealing with 3 elements and per-element
    // alignment
    if (lanes == 3) {
      os << "packed_";
    }
    switch (t.bits()) {
      case 16:
        os << "half";
        break;
      case 32:
        os << "float";
        break;
      default:
        fail = true;
        break;
    }
    if (!fail && lanes == 1) return;
    if (!fail && (lanes >= 2 && lanes <= 4)) {
      os << lanes;
      return;
    }
  } else if (t.MatchesCode(DLDataTypeCode::kDLUInt, DLDataTypeCode::kDLInt)) {
    if (t.MatchesCode(DLDataTypeCode::kDLUInt)) {
      os << 'u';
    }
    switch (t.bits()) {
      case 8:
        os << "char";
        break;
      case 16:
        os << "short";
        break;
      case 32:
        os << "int";
        break;
      case 64:
        os << "long";
        break;
      case 1:
        os << "bool";
        break;
      default:
        fail = true;
        break;
    }
    if (!fail && lanes == 1) return;
    if (!fail && (lanes >= 2 && lanes <= 4)) {
      os << lanes;
      return;
    }
  } else if (t.MatchesElementType(DLDataTypeCode::kDLBfloat, 16)) {
    os << "bfloat";
    return;
  }
  TVM_FFI_THROW(InternalError) << "Cannot convert type " << ffi::DLDataTypeToString(t->dtype)
                               << " to Metal type";
}

void CodeGenMetal::PrintStorageSync(const CallNode* op) {
  const std::string& sync = op->args[0].as<StringImmNode>()->value;
  if (sync == "warp") {
    this->PrintIndent();
    this->stream << "simdgroup_barrier(mem_flags::mem_threadgroup);\n";
  } else if (sync == "shared") {
    this->PrintIndent();
    this->stream << "threadgroup_barrier(mem_flags::mem_threadgroup);\n";
  } else if (sync == "global") {
    TVM_FFI_THROW(InternalError) << "global barrier not supported";
  }
}

void CodeGenMetal::PrintVecElemLoad(const std::string& vec, const PrimType& t, int i,
                                    std::ostream& os) {  // NOLINT(*)
  os << vec << "[" << i << "]";
}

void CodeGenMetal::PrintVecElemStore(const std::string& vec, const PrimType& t, int i,
                                     const std::string& value) {
  this->PrintIndent();
  stream << vec << "[" << i << "]"
         << " = " << value << ";\n";
}

void CodeGenMetal::PrintStorageScope(const std::string& scope, std::ostream& os) {  // NOLINT(*)
  if (scope == "global") {
    os << "device ";
  } else if (scope == "shared") {
    os << "threadgroup ";
  } else if (scope == "local") {
    os << "thread ";
  } else {
    TVM_FFI_THROW(InternalError) << "Unknown storage scope `" << scope << "`";
  }
}

void CodeGenMetal::VisitStmt_(const AllocBufferNode* op) {
  TVM_FFI_ICHECK(op->buffer.defined());
  std::string vid = AllocVarID(op->buffer.get());

  this->PrintIndent();
  // Compute constant_size from buffer shape
  size_t constant_size = 1;
  for (const auto& dim : op->buffer->shape) {
    const IntImmNode* dim_imm = dim.as<IntImmNode>();
    TVM_FFI_ICHECK(dim_imm) << "Can only handle constant size stack allocation for now";
    constant_size *= dim_imm->value;
  }
  TVM_FFI_ICHECK_GT(constant_size, 0) << "Can only handle constant size stack allocation for now";

  auto scope = op->buffer.scope();
  alloc_storage_scope_[op->buffer.get()] = scope;
  const PrimType& dtype = op->buffer->dtype;
  if (scope == "metal.simdgroup") {
    bool supported_simdgroup_dtype = dtype == PrimType::Float(16) || dtype == PrimType::Float(32) ||
                                     dtype == PrimType::BFloat(16);
    TVM_FFI_ICHECK(supported_simdgroup_dtype)
        << "Only float16, float32, and bfloat16 are supported, but got "
        << ffi::DLDataTypeToString(dtype->dtype);
    TVM_FFI_ICHECK(constant_size % 64 == 0)
        << "Only 8x8 matrix is supported, but got " << constant_size << " bytes\n";

    std::ostringstream dtype_os;
    PrintType(dtype, dtype_os);
    std::string dtype_str = dtype_os.str();
    simdgroup_dtype_[op->buffer.get()] = dtype_str;
    stream << "simdgroup_" << dtype_str << "8x8 " << vid << '[' << constant_size / 64 << "];\n";
  } else {
    PrintStorageScope(scope, stream);
    PrintType(dtype, stream);
    stream << ' ' << vid << '[' << constant_size << "];\n";
  }

  RegisterHandleType(op->buffer.get(), op->buffer->dtype);
  if (op->annotations.count(tirx::attr::kVolatile)) {
    MarkVolatile(op->buffer.get());
  }
}

void CodeGenMetal::VisitExpr_(const SelectNode* op, std::ostream& os) {  // NOLINT(*)
  os << "select(" << PrintExpr(op->false_value) << ", " << PrintExpr(op->true_value) << ", "
     << PrintExpr(op->condition) << ")";
}

void CodeGenMetal::VisitExpr_(const BroadcastNode* op, std::ostream& os) {  // NOLINT(*)
  std::string v = PrintExpr(op->value);
  int lanes = op->ty.as_or_throw<PrimType>().lanes();
  PrintType(op->ty.as_or_throw<PrimType>(), os);
  os << "(";
  for (int i = 0; i < lanes; ++i) {
    if (i != 0) os << ", ";
    os << v;
  }
  os << ')';
}

void CodeGenMetal::VisitExpr_(const CallNode* op, std::ostream& os) {  // NOLINT(*)
  TVM_FFI_ICHECK(!op->op.as<GlobalVarNode>())
      << "CodegenMetal does not support inter-function calls, "
      << "but expression " << ffi::GetRef<Call>(op) << " calls PrimFunc " << op->op;
  auto f_check_simdgroup_shape = [](PrimExpr col, PrimExpr row) {
    TVM_FFI_ICHECK(col->IsInstance<IntImmNode>() && row->IsInstance<IntImmNode>())
        << "Only constant shape is supported for simdgroup matrix, but got " << col << "x" << row;
    int col_val = col.as<IntImmNode>()->value;
    int row_val = row.as<IntImmNode>()->value;
    TVM_FFI_ICHECK(col_val == 8 && row_val == 8)
        << "Only 8x8 matrix is supported, but got " << col_val << "x" << row_val;
  };

  static const Op& make_filled_simdgroup_matrix_op = Op::Get("tirx.make_filled_simdgroup_matrix");
  static const Op& simdgroup_load_op = Op::Get("tirx.simdgroup_load");
  static const Op& simdgroup_store_op = Op::Get("tirx.simdgroup_store");
  static const Op& simdgroup_multiply_accumulate_op = Op::Get("tirx.simdgroup_multiply_accumulate");

  if (op->op.same_as(make_filled_simdgroup_matrix_op)) {
    TVM_FFI_ICHECK_EQ(op->args.size(), 5);
    Var var = GetSimdgroupBufferVar(op->args[0]);
    // Get the data type of the simdgroup matrix
    auto it = simdgroup_dtype_.find(var.get());
    TVM_FFI_ICHECK(it != simdgroup_dtype_.end())
        << "Cannot find variable allocation for simdgroup: " << var;
    const std::string& dtype_str = it->second;
    f_check_simdgroup_shape(op->args[3].as_or_throw<PrimExpr>(),
                            op->args[4].as_or_throw<PrimExpr>());
    os << PrintExpr(var) << "[" << PrintExpr(op->args[1]) << "] = make_filled_simdgroup_matrix<"
       << dtype_str << ", " << PrintExpr(op->args[3]) << ", " << PrintExpr(op->args[4]) << ">("
       << PrintExpr(op->args[2]) << ")";
  } else if (op->op.same_as(simdgroup_load_op)) {
    TVM_FFI_ICHECK_EQ(op->args.size(), 7);
    Var var = GetSimdgroupBufferVar(op->args[0]);
    f_check_simdgroup_shape(op->args[4].as_or_throw<PrimExpr>(),
                            op->args[5].as_or_throw<PrimExpr>());
    os << "simdgroup_load(" << PrintExpr(var) << "[" << PrintExpr(op->args[1]) << "], "
       << PrintExpr(op->args[2]) << ", " << PrintExpr(op->args[3]) << ", 0, "
       << PrintExpr(op->args[6]) << ")";
  } else if (op->op.same_as(simdgroup_store_op)) {
    TVM_FFI_ICHECK_EQ(op->args.size(), 7);
    Var var = GetSimdgroupBufferVar(op->args[0]);
    f_check_simdgroup_shape(op->args[4].as_or_throw<PrimExpr>(),
                            op->args[5].as_or_throw<PrimExpr>());
    os << "simdgroup_store(" << PrintExpr(var) << "[" << PrintExpr(op->args[1]) << "], "
       << PrintExpr(op->args[2]) << ", " << PrintExpr(op->args[3]) << ", 0, "
       << PrintExpr(op->args[6]) << ")";
  } else if (op->op.same_as(simdgroup_multiply_accumulate_op)) {
    TVM_FFI_ICHECK_EQ(op->args.size(), 8);
    Var d = GetSimdgroupBufferVar(op->args[0]);
    Var a = GetSimdgroupBufferVar(op->args[2]);
    Var b = GetSimdgroupBufferVar(op->args[4]);
    Var c = GetSimdgroupBufferVar(op->args[6]);
    os << "simdgroup_multiply_accumulate("                        //
       << PrintExpr(d) << "[" << PrintExpr(op->args[1]) << "], "  //
       << PrintExpr(a) << "[" << PrintExpr(op->args[3]) << "], "  //
       << PrintExpr(b) << "[" << PrintExpr(op->args[5]) << "], "  //
       << PrintExpr(c) << "[" << PrintExpr(op->args[7]) << "])";
  } else if (op->op.same_as(builtin::reinterpret())) {
    if (!op->ty.as<PrimTypeNode>() || !op->args[0]->ty.as<PrimTypeNode>()) {
      return CodeGenC::VisitExpr_(op, os);
    }
    // generate as_type<TYPE>(ARG)
    os << "(as_type<";
    this->PrintType(op->ty.as_or_throw<PrimType>(), os);
    os << ">(";
    this->PrintExpr(op->args[0], os);
    os << "))";
  } else {
    CodeGenC::VisitExpr_(op, os);
  }
}

void CodeGenMetal::VisitExpr_(const FloatImmNode* op, std::ostream& os) {  // NOLINT(*)
  std::ostringstream temp;
  if (std::isinf(op->value)) {
    if (op->value < 0) {
      temp << "-";
    }
    temp << "INFINITY";
  } else if (std::isnan(op->value)) {
    temp << "NAN";
  } else {
    temp << std::scientific << op->value;
    if (op->ty.as_or_throw<PrimType>().bits() == 32)
      temp << 'f';
    else if (op->ty.as_or_throw<PrimType>().bits() == 16)
      temp << 'h';
  }
  MarkConst(temp.str());
  os << temp.str();
}

ffi::Module BuildMetal(IRModule mod, Target target) {
  bool output_ssa = false;
  mod = tirx::transform::PointerValueTypeRewrite()(std::move(mod));

  std::ostringstream source_maker;
  // Per-kernel payload: kernel-name -> bytes (MSL source for fmt="metal" /
  // compiled metallib bytes for fmt="metallib").  See
  // tasks/...-tvm-unify-device-module... "Per-kernel smap shape" — uniform
  // Map<String, Bytes> across all multi-shader backends.
  ffi::Map<ffi::String, ffi::Bytes> smap;
  const auto fmetal_compile = tvm::ffi::Function::GetGlobal("tvm_callback_metal_compile");
  // Default payload format. A callback may override it per the contract below.
  std::string fmt = fmetal_compile ? "metallib" : "metal";
  bool fmt_locked = false;

  for (auto kv : mod->functions) {
    TVM_FFI_ICHECK(kv.second->IsInstance<PrimFuncNode>()) << "CodeGenMetal: Can only take PrimFunc";
    auto global_symbol = kv.second->GetAttr<ffi::String>(tvm::attr::kGlobalSymbol);
    TVM_FFI_ICHECK(global_symbol.has_value());
    std::string func_name = global_symbol.value();

    source_maker << "// Function: " << func_name << "\n";
    CodeGenMetal cg(target);
    cg.Init(output_ssa);
    auto f = kv.second.as_or_throw<PrimFunc>();
    auto calling_conv = f->GetAttr<CallingConv>(tvm::attr::kCallingConv);
    TVM_FFI_ICHECK(calling_conv.has_value())
        << "CodeGenMetal: expected kCallingConv attribute to be set.";
    TVM_FFI_ICHECK(calling_conv.value() == CallingConv::kDeviceKernelLaunch)
        << "CodeGenMetal: expect calling_conv equals CallingConv::kDeviceKernelLaunch, but got "
        << static_cast<int>(calling_conv.value());

    cg.AddFunction(kv.first, f);

    std::string fsource = cg.Finish();
    source_maker << fsource << "\n";
    if (fmetal_compile) {
      ffi::Any ret = (*fmetal_compile)(fsource, target);
      // Backward-compatible contract for tvm_callback_metal_compile:
      //   * returning a str/bytes    -> treated as a compiled metallib payload
      //   * returning (payload, fmt) -> the callback declares the payload format
      std::string kernel_fmt;
      if (auto ret_tuple = ret.try_cast<ffi::Array<ffi::Any>>()) {
        TVM_FFI_ICHECK_EQ(ret_tuple->size(), 2)
            << "tvm_callback_metal_compile must return either a payload or a "
               "(payload, format) pair, but got a tuple of size "
            << ret_tuple->size();
        fsource = (*ret_tuple)[0].cast<std::string>();
        kernel_fmt = (*ret_tuple)[1].cast<std::string>();
        TVM_FFI_ICHECK(kernel_fmt == "metal" || kernel_fmt == "metallib")
            << "tvm_callback_metal_compile returned unsupported format \"" << kernel_fmt
            << "\"; expected \"metal\" or \"metallib\"";
      } else {
        // Backward-compatible behavior
        fsource = ret.cast<std::string>();
        kernel_fmt = "metallib";
      }
      // All kernels of a module share a single declared format
      TVM_FFI_ICHECK(!fmt_locked || fmt == kernel_fmt)
          << "tvm_callback_metal_compile returned inconsistent formats across kernels: \"" << fmt
          << "\" vs \"" << kernel_fmt << "\"";
      fmt = kernel_fmt;
      fmt_locked = true;
    }
    smap.Set(func_name, ffi::Bytes(std::move(fsource)));
  }

  // The aggregated MSL source dump is preserved in the in-memory source
  // map keyed by "metal" — only used by InspectSource and never serialized.
  ffi::Map<ffi::String, ffi::String> source;
  source.Set("metal", source_maker.str());
  return target::MetalModuleCreateWithFallback(std::move(smap), ffi::String(fmt),
                                               ExtractFuncInfo(mod), std::move(source));
}

void RegisterMetalCodegen() {
  static bool registered = false;
  if (registered) return;
  registered = true;

  namespace refl = tvm::ffi::reflection;
  refl::GlobalDef().def("target.build.metal", BuildMetal);
}
}  // namespace codegen
}  // namespace tvm
