/*!
 * \file tl/backend/common/op/finalize_reducer.h
 * \brief Shared tl.finalize_reducer lowering for GPU backends.
 */

#ifndef TVM_TL_BACKEND_COMMON_OP_FINALIZE_REDUCER_H_
#define TVM_TL_BACKEND_COMMON_OP_FINALIZE_REDUCER_H_

#include "op/finalize_reducer.h"
#include "support/check.h"

#include <tvm/tirx/builtin.h>

#include <array>
#include <cstdint>
#include <string>

namespace tvm {
namespace tl {
namespace backend {

using namespace tirx;
using namespace ffi;

template <typename Impl> struct FinalizeReducerLowerer {
  static Stmt Lower(const FinalizeReducerOpNode &op,
                    const LowerArgs &lower_args, arith::Analyzer *) {
    auto buffer = lower_args.buffer_remap[op.reducer];
    auto opt_layout = lower_args.layout_map.Get(op.reducer);
    ICHECK(opt_layout);
    ICHECK(opt_layout->as<Fragment>());
    auto layout = opt_layout->as<Fragment>().value();
    Array<PrimExpr> indices_0;
    indices_0.reserve(layout->OutputDim());
    for (int i = 0; i < layout->OutputDim(); ++i) {
      indices_0.push_back(Var("__finred_" + std::to_string(i)));
    }

    const int64_t *p_extent = as_const_int(layout->ReplicateExtent());
    ICHECK(p_extent);
    int extent = *p_extent;
    ICHECK(extent == 1 ||
           extent == *as_const_int(lower_args.thread_bounds->extent))
        << "Illegal finalize_reducer: extent=" << extent
        << "; T.thread_bounds=" << lower_args.thread_bounds;

    if (extent == 1) {
      return Evaluate(0);
    }

    std::array op_names{"tl::SumOp", "tl::MaxOp", "tl::MinOp"};
    auto op_str = op_names[static_cast<int>(op.op)];

    int reducing_threads = extent;
    auto thread_offset = lower_args.thread_bounds->min;

    int64_t layout_batch_size = 1;
    for (int i = 0; i < layout->OutputDim(); ++i) {
      const int64_t *p = as_const_int(layout->OutputShape()[i]);
      if (p == nullptr) {
        layout_batch_size = -1;
        break;
      }
      layout_batch_size *= *p;
    }

    int64_t effective_batch = static_cast<int64_t>(op.batch);

    if (effective_batch > 1 && layout_batch_size > 0) {
      ICHECK_LE(effective_batch, layout_batch_size)
          << "finalize_reducer: batch (" << effective_batch
          << ") exceeds total output elements (" << layout_batch_size << ")";
      ICHECK_EQ(layout_batch_size % effective_batch, 0)
          << "finalize_reducer: batch (" << effective_batch
          << ") must evenly divide total output elements (" << layout_batch_size
          << ")";
    }

    bool use_batch = effective_batch > 1 &&
                     reducing_threads > Impl::WarpSize(lower_args.target);

    if (use_batch) {
      int workspace_stride =
          static_cast<int>(*as_const_int(lower_args.thread_bounds->extent));
      std::string allreduce = Impl::MakeBatchAllReduce(
          op_str, reducing_threads, 1, thread_offset,
          lower_args.thread_bounds->extent, static_cast<int>(effective_batch),
          workspace_stride, lower_args.target);
      int ws_size = workspace_stride * static_cast<int>(effective_batch);
      PrimExpr workspace = lower_args.add_workspace(ws_size, buffer->dtype);
      Array<PrimExpr> args = {StringImm(allreduce), buffer->data, workspace};
      return Evaluate(Call(DataType::Handle(), builtin::call_extern(), args));
    }

    std::string allreduce = Impl::MakeScalarAllReduce(
        op_str, reducing_threads, 1, thread_offset,
        lower_args.thread_bounds->extent, lower_args.target);
    Array<PrimExpr> thread_reduce_args = {StringImm(allreduce),
                                          BufferLoad(buffer, indices_0)};
    if (reducing_threads > Impl::WarpSize(lower_args.target)) {
      PrimExpr workspace = lower_args.add_workspace(
          *as_const_int(lower_args.thread_bounds->extent), buffer->dtype);
      thread_reduce_args.push_back(workspace);
    }
    auto call = Call(buffer->dtype, builtin::call_extern(), thread_reduce_args);
    Stmt body = BufferStore(buffer, call, indices_0);

    for (int i = layout->OutputDim() - 1; i >= 0; i--) {
      body = For(indices_0[i].as<Var>().value(), 0, layout->OutputShape()[i],
                 ForKind::kParallel, body);
    }

    return body;
  }
};

} // namespace backend
} // namespace tl
} // namespace tvm

#endif // TVM_TL_BACKEND_COMMON_OP_FINALIZE_REDUCER_H_
