/*!
 *  \file lower_shared_tmem.cc
 *  \brief Convert shared.tmem buffers to plain shared + ptx init, and do
 *         coordinate translation (from logical address to physical address)
 */
#include "cuda/target_utils.h"
#include "op/builtin.h"
#include "support/check.h"
#include "tvm/ir/type.h"
#include <tvm/arith/analyzer.h>
#include <tvm/ir/cast.h>
#include <tvm/tirx/analysis.h>
#include <tvm/tirx/builtin.h>
#include <tvm/tirx/expr.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/stmt.h>
#include <tvm/tirx/stmt_functor.h>
#include <tvm/tirx/transform.h>
#include <unordered_set>

namespace tvm {
namespace tl {

using namespace tirx;
using namespace ffi;

using VarSet = std::unordered_set<Var, ObjectPtrHash, ObjectPtrEqual>;

/*!
 * \brief Collect TMEM buffers explicitly deallocated on fallthrough paths.
 *
 * A "fallthrough path" is one that reaches the end of the statement without
 * hitting thread_return().  Buffers deallocated on every such path already
 * have an explicit dealloc, so we can skip the auto-dealloc at block end.
 *
 * \return {buffers deallocated on fallthrough, whether the stmt can
 * fallthrough}
 */
static std::pair<VarSet, bool> CollectFallthroughDeallocs(const Stmt &stmt) {
  if (!stmt.defined())
    return {{}, true};

  // Unwrap transparent wrapper nodes
  if (stmt.as<BindNode>())
    return {{}, true};
  if (auto *n = stmt.as<AttrStmtNode>())
    return CollectFallthroughDeallocs(n->body);
  if (auto *n = stmt.as<SBlockNode>())
    return CollectFallthroughDeallocs(n->body);
  if (auto *n = stmt.as<SBlockRealizeNode>())
    return CollectFallthroughDeallocs(n->block->body);
  if (auto *n = stmt.as<ForNode>())
    return CollectFallthroughDeallocs(n->body);

  // Sequential: accumulate deallocs; stop if any child doesn't fallthrough
  if (auto *seq = stmt.as<SeqStmtNode>()) {
    VarSet deallocs;
    for (const auto &child : seq->seq) {
      auto [d, ft] = CollectFallthroughDeallocs(child);
      if (!ft)
        return {{}, false};
      deallocs.insert(d.begin(), d.end());
    }
    return {std::move(deallocs), true};
  }

  // Branch: collect deallocs only from branches that can fallthrough
  if (auto *iff = stmt.as<IfThenElseNode>()) {
    auto [then_d, then_ft] = CollectFallthroughDeallocs(iff->then_case);
    auto [else_d, else_ft] =
        iff->else_case.defined()
            ? CollectFallthroughDeallocs(iff->else_case.value())
            : std::pair<VarSet, bool>{{}, true};
    VarSet deallocs;
    if (then_ft)
      deallocs.insert(then_d.begin(), then_d.end());
    if (else_ft)
      deallocs.insert(else_d.begin(), else_d.end());
    return {std::move(deallocs), then_ft || else_ft};
  }

  // Leaf: detect deallocate_tmem and thread_return
  if (auto *eval = stmt.as<EvaluateNode>()) {
    if (auto *call = eval->value.as<CallNode>()) {
      if (call->op.same_as(tl::deallocate_tmem())) {
        ICHECK_EQ(call->args.size(), 1U);
        auto *buf = call->args[0].as<VarNode>();
        ICHECK(buf) << "tl.deallocate_tmem expects a buffer data Var";
        return {{GetRef<Var>(buf)}, true};
      }
      if (call->op.same_as(builtin::thread_return())) {
        return {{}, false};
      }
    }
  }

  return {{}, true};
}

class SharedTmemRewriter : public StmtExprMutator {
public:
  static Stmt Rewrite(Stmt body, Target target) {
    SharedTmemRewriter rewriter;
    rewriter.target_ = std::move(target);
    return rewriter(body);
  }

private:
  int GetNumColsAllocated(const Buffer &buffer) const {
    ICHECK_EQ(buffer->shape.size(), 2U);

    auto analyzer = std::make_shared<arith::Analyzer>();
    arith::ConstIntBound phy_col_bounds =
        analyzer->const_int_bound(buffer->shape[1]);
    int num_cols_required = phy_col_bounds->max_value;
    ICHECK(num_cols_required <= 512)
        << "The number of columns required for tmem buffer " << buffer->name
        << " is " << num_cols_required
        << ", which exceeds the maximum of 512 columns";

    int num_cols_allocated = 32; // Align num_cols_allocated to power of 2
    for (; num_cols_allocated < num_cols_required; num_cols_allocated *= 2) {
    }
    return num_cols_allocated;
  }

  Stmt VisitStmt_(const SBlockNode *op) final {
    SBlock block = GetRef<SBlock>(op);
    Array<Buffer> alloc_buffers = op->alloc_buffers;
    if (op->annotations.count(attr::kLayoutMap)) {
      auto layout_map = op->annotations.Get(attr::kLayoutMap);
      ICHECK(layout_map) << "layout map is not defined";
      layout_map_ = layout_map->as<Map<Buffer, Layout>>().value();
    }

    // Record the mapping from buffer data var to buffer for later lookup
    for (auto buffer : alloc_buffers) {
      buffer_map_.insert({buffer->data, buffer});
    }
    for (auto match_buffer : op->match_buffers) {
      buffer_map_.insert({match_buffer->buffer->data, match_buffer->buffer});
    }

    Array<Buffer> tmem_buffers;

    for (const auto &[data, buffer] : buffer_map_) {
      const auto *ptr_type =
          buffer->data->type_annotation.as<PointerTypeNode>();
      auto storage_scope = ptr_type->storage_scope;
      ICHECK(ptr_type) << "Buffer Var's type annotation must be of PointerType";
      if (storage_scope == "shared.tmem") {
        tmem_buffers.push_back(buffer);
      }
    }

    if (tmem_buffers.empty()) {
      return StmtExprMutator::VisitStmt_(op);
    }

    ICHECK(thread_var_.defined()) << "thread_var_ is not defined";

    auto [fallthrough_deallocs, _] = CollectFallthroughDeallocs(op->body);

    for (auto buffer : tmem_buffers) {
      buffer_data_to_buffer_.Set(buffer->data, buffer);
    }

    /*
    Transform the tmem buffers to new allocations
    transform:
        tmem_buf0 = T.alloc_buffer((128, 128,), "uint64",
    scope="shared.tmem")
        tmem_buf1 = T.alloc_buffer((128, 128,), "uint64",
    scope="shared.tmem")

    into:
        tmem_buf0 = T.alloc_buffer((1,), "uint64", scope="shared.tmem_addr")
        tmem_buf1 = T.alloc_buffer((1,), "uint64", scope="shared.tmem_addr")

        if tx == 0:
          T.ptx_init_tensor_memory(tmem_buf0[0], 128)
          T.ptx_init_tensor_memory(tmem_buf1[0], 128)
    */
    // 1. create new data vars
    Array<Var> new_data_vars;
    for (auto buffer : tmem_buffers) {
      auto data = buffer->data;
      if (var_remap_.count(data))
        continue;
      auto new_data =
          Var(data->name_hint, PointerType(PrimType(tmem_dtype_), "shared"));
      var_remap_.Set(data, new_data);
      new_data_vars.push_back(new_data);
    }

    // 2. create new buffers
    Array<Buffer> new_buffers;
    for (auto buffer : tmem_buffers) {
      auto data = buffer->data;
      ICHECK(var_remap_.find(data) != var_remap_.end())
          << "data not found in var_remap_";
      auto new_data = var_remap_.at(data);
      auto new_buffer = Buffer(new_data, tmem_dtype_, Array<PrimExpr>({1}),
                               Array<PrimExpr>({1}), PrimExpr(0), buffer->name,
                               buffer->data_alignment, buffer->offset_factor,
                               buffer->buffer_type);
      new_buffers.push_back(new_buffer);
      buffer_remap_.Set(buffer, new_buffer);
      buffer_data_to_buffer_.Set(new_data, new_buffer);
    }

    // remove the tmem buffers
    alloc_buffers.MutateByApply([this](Buffer buf) {
      if (buffer_remap_.find(buf) != buffer_remap_.end()) {
        return buffer_remap_.at(buf);
      }
      return buf;
    });
    if (!alloc_buffers.same_as(op->alloc_buffers)) {
      block.CopyOnWrite()->alloc_buffers = alloc_buffers;
    } else {
      return StmtExprMutator::VisitStmt_(op);
    }

    // If block has use_2cta attr, add use_2cta: 1 to tmem alloc/dealloc call
    // annotations.
    Map<String, ObjectRef> tmem_call_ann;
    if (op->annotations.count("use_2cta")) {
      PrimExpr val = Downcast<PrimExpr>(op->annotations["use_2cta"]);
      // Bool in TVM is a subclass of IntImm, so only check IntImm.
      if (const auto *i = val.as<IntImmNode>()) {
        if (i->value != 0) {
          tmem_call_ann.Set("use_2cta", IntImm(DataType::Int(32), 1));
        }
      }
    }

    // 3. create init & dealloc calls for new buffers
    std::vector<Stmt> init_mtmem_calls_;
    std::vector<Stmt> dealloc_tmem_calls_;
    for (auto buffer : tmem_buffers) {
      auto data = buffer->data;
      auto old_buffer = buffer_data_to_buffer_.at(data);
      auto new_buffer = buffer_remap_.at(old_buffer);
      int num_cols_allocated = GetNumColsAllocated(old_buffer);

      // Check that the number of rows doesn't exceed the tmem limit
      {
        auto analyzer = std::make_shared<arith::Analyzer>();
        arith::ConstIntBound phy_row_bounds =
            analyzer->const_int_bound(old_buffer->shape[0]);
        int num_rows_required = phy_row_bounds->max_value;
        ICHECK(num_rows_required <= 128)
            << "The number of rows required for tmem buffer "
            << old_buffer->name << " is " << num_rows_required
            << ", which exceeds the maximum of 128 rows";
      }

      tmem_num_cols_allocated_.insert({data, num_cols_allocated});
      tmem_call_annotations_.insert({data, tmem_call_ann});

      auto new_buffer_access = new_buffer.access_ptr(1, DataType::Handle(), 1,
                                                     PrimExpr(0), PrimExpr(1));
      auto alloc_call = Call(DataType::Handle(), tl::ptx_init_tensor_memory(),
                             {new_buffer_access, PrimExpr(num_cols_allocated)},
                             tmem_call_ann);
      init_mtmem_calls_.push_back(Evaluate(alloc_call));
      if (!fallthrough_deallocs.count(data)) {
        auto dealloc_call = Call(
            DataType::Handle(), tl::ptx_deallocate_tensor_memory(),
            {new_buffer_access, PrimExpr(num_cols_allocated)}, tmem_call_ann);
        dealloc_tmem_calls_.push_back(Evaluate(dealloc_call));
      }
    }
    auto compare_by_buffer_name = [&](const Stmt &a, const Stmt &b) {
      auto call_a = a.as<EvaluateNode>()->value.as<CallNode>();
      auto call_b = b.as<EvaluateNode>()->value.as<CallNode>();
      auto num_cols_a = call_a->args[1].as<IntImmNode>()->value;
      auto num_cols_b = call_b->args[1].as<IntImmNode>()->value;
      return num_cols_a > num_cols_b;
    };
    std::sort(init_mtmem_calls_.begin(), init_mtmem_calls_.end(),
              compare_by_buffer_name);

    Array<Stmt> new_body;
    ICHECK(target_.defined()) << "LowerSharedTmem requires a bound target";
    auto warp_size = TargetCudaGetWarpSize(target_);
    auto thread_var_div_warp_size =
        FloorDiv(thread_var_->var, IntImm(thread_var_->var->dtype, warp_size));
    new_body.push_back(IfThenElse(EQ(thread_var_div_warp_size, 0),
                                  init_mtmem_calls_.size() > 1
                                      ? SeqStmt(init_mtmem_calls_)
                                      : init_mtmem_calls_.back(),
                                  Stmt()));
    new_body.push_back(
        Evaluate(Call(DataType::Handle(), builtin::tvm_storage_sync(),
                      {StringImm("shared")})));
    new_body.push_back(block->body);
    if (!dealloc_tmem_calls_.empty()) {
      if (tmem_call_ann.find("use_2cta") != tmem_call_ann.end()) {
        new_body.push_back(
            Evaluate(Call(DataType::Handle(), tl::cluster_sync(), {})));
      }
      new_body.push_back(IfThenElse(EQ(thread_var_div_warp_size, 0),
                                    dealloc_tmem_calls_.size() > 1
                                        ? SeqStmt(dealloc_tmem_calls_)
                                        : dealloc_tmem_calls_.back(),
                                    Stmt()));
    }

    auto block_ptr = block.CopyOnWrite();
    block_ptr->annotations.erase(attr::kLayoutMap);
    block_ptr->body = SeqStmt(new_body);

    return StmtExprMutator::VisitStmt_(block.get());
  }

  PrimExpr GetTmemOffset(const Buffer &buffer, const Array<PrimExpr> &indices) {
    ICHECK(buffer->shape.size() == 2);
    ICHECK(indices.size() == 2);
    ICHECK(layout_map_.defined());
    ICHECK(layout_map_.count(buffer))
        << "The layout of tmem buffer " << buffer->name
        << " is not defined in the layout map";
    auto layout = layout_map_[buffer];
    ICHECK(layout.defined());
    Array<PrimExpr> tmem_phy_coords = layout->Forward(indices);
    PrimExpr result =
        tmem_phy_coords[0] << 16 |
        tmem_phy_coords
            [1]; // https://docs.nvidia.com/cuda/parallel-thread-execution/#tensor-memory-addressing
    return result;
  }

  PrimExpr VisitExpr_(const BufferLoadNode *op) final {
    // Translate tmem[logical_row, logical_col] to tmem[0] + tmem_offset
    // Where
    // - (logical_row, logical_col) is the logical address in the tmem buffer
    // - tmem[0] is the base address allocated for the tmem buffer
    // - tmem_offset = tmem_phy_coords[0]<<16 | tmem_phy_coords[1]
    //   where tmem_phy_coords = layout.Forward(logical_row, logical_col)
    //   is the physical address in the tmem buffer
    auto load = Downcast<BufferLoad>(StmtExprMutator::VisitExpr_(op));
    auto buffer = load->buffer;
    auto indices = load->indices;

    if (buffer_remap_.count(buffer)) {
      auto new_buffer = buffer_remap_[load->buffer];
      return BufferLoad(new_buffer, {0}) + GetTmemOffset(buffer, indices);
    } else if (var_remap_.count(buffer->data)) {
      auto new_buffer = Buffer(
          var_remap_[buffer->data], tmem_dtype_, buffer->shape, buffer->strides,
          buffer->elem_offset, buffer->name, buffer->data_alignment,
          buffer->offset_factor, buffer->buffer_type);
      return BufferLoad(new_buffer, {0}) + GetTmemOffset(buffer, indices);
    }
    return load;
  }

  Stmt VisitStmt_(const BufferStoreNode *op) final {
    auto store = Downcast<BufferStore>(StmtExprMutator::VisitStmt_(op));
    auto buffer = store->buffer;
    ICHECK(buffer.scope() != "shared.tmem")
        << "We should never directly store data into tmem!";
    return store;
  }

  PrimExpr VisitExpr_(const CallNode *op) final {
    if (op->op.same_as(tl::deallocate_tmem())) {
      ICHECK_EQ(op->args.size(), 1U);
      Var buffer_data = Downcast<Var>(op->args[0]);
      auto num_cols_it = tmem_num_cols_allocated_.find(buffer_data);
      ICHECK(num_cols_it != tmem_num_cols_allocated_.end())
          << "tl.deallocate_tmem expects a TMEM buffer allocated in the same "
             "or an enclosing block";
      ICHECK(buffer_data_to_buffer_.count(buffer_data))
          << "TMEM buffer for tl.deallocate_tmem is not tracked";
      Buffer old_buffer = buffer_data_to_buffer_.at(buffer_data);
      ICHECK(buffer_remap_.count(old_buffer))
          << "TMEM buffer for tl.deallocate_tmem has not been remapped";
      Buffer new_buffer = buffer_remap_[old_buffer];
      auto new_buffer_access = new_buffer.access_ptr(1, DataType::Handle(), 1,
                                                     PrimExpr(0), PrimExpr(1));

      Map<String, ObjectRef> ann;
      auto ann_it = tmem_call_annotations_.find(buffer_data);
      if (ann_it != tmem_call_annotations_.end()) {
        ann = ann_it->second;
      }
      return Call(DataType::Handle(), tl::ptx_deallocate_tensor_memory(),
                  {new_buffer_access, PrimExpr(num_cols_it->second)}, ann);
    }
    if (op->op.same_as(builtin::tvm_access_ptr())) {
      ICHECK_EQ(op->args.size(), 5U);
      Var buffer_data = Downcast<Var>(op->args[1]);
      if (!var_remap_.count(buffer_data)) {
        return StmtExprMutator::VisitExpr_(op);
      }
      Var new_data = var_remap_[buffer_data];
      return Call(
          op->dtype, op->op,
          {op->args[0], new_data, op->args[2], op->args[3], op->args[4]});
    }
    auto expr = StmtExprMutator::VisitExpr_(op);
    return expr;
  }
  PrimExpr VisitExpr_(const VarNode *op) final {
    Var var = GetRef<Var>(op);
    if (var_remap_.count(var)) {
      return var_remap_[var];
    }
    return var;
  }

  Stmt VisitStmt_(const AttrStmtNode *op) final {
    if (op->attr_key == tirx::attr::thread_extent) {
      IterVar iv = Downcast<IterVar>(op->node);
      if (iv->thread_tag == "threadIdx.x") {
        ICHECK(iv->dom->extent.as<IntImmNode>());
        thread_var_ = iv;
      }
    }
    return StmtExprMutator::VisitStmt_(op);
  }

  // Datatypes for tmem
  const DataType tmem_dtype_ = DataType::UInt(32);
  // This is a workaround for cpu backend,
  // we need to define a thread_var for the serial loop.
  IterVar thread_var_;
  Target target_;
  Map<Var, Var> var_remap_;
  Map<Var, Buffer> buffer_data_to_buffer_;
  Map<Buffer, Buffer> buffer_remap_;
  // Mapping from data Var of a Buffer to Buffer, for lookup
  std::unordered_map<Var, Buffer, ObjectPtrHash, ObjectPtrEqual> buffer_map_;
  std::unordered_map<Var, int, ObjectPtrHash, ObjectPtrEqual>
      tmem_num_cols_allocated_;
  std::unordered_map<Var, Map<String, ObjectRef>, ObjectPtrHash, ObjectPtrEqual>
      tmem_call_annotations_;
  Map<Buffer, Layout> layout_map_;
};

PrimFunc LowerSharedTmem(PrimFunc f) {
  auto target = f->GetAttr<Target>(tvm::attr::kTarget);
  ICHECK(target.defined()) << "LowerSharedTmem: Require the target attribute";
  f.CopyOnWrite()->body = SharedTmemRewriter::Rewrite(f->body, target.value());
  return f;
}

namespace transform {
using namespace tirx::transform;

tvm::transform::Pass LowerSharedTmem() {
  auto pass_func = [=](PrimFunc f, IRModule m, PassContext ctx) {
    return tl::LowerSharedTmem(std::move(f));
  };
  return CreatePrimFuncPass(pass_func, 0, "tl.LowerSharedTmem", {});
}

TVM_FFI_STATIC_INIT_BLOCK() {
  namespace refl = reflection;
  refl::GlobalDef().def("tl.cuda.transform.LowerSharedTmem", LowerSharedTmem);
}

} // namespace transform
} // namespace tl
} // namespace tvm
