/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/

#include "comm_overlap.h"
#include "torch_api/torch_ops_api.h"
namespace tmo {
namespace torch_api {

using namespace torch::indexing;

struct MatmulSplitM {
  const at::Tensor &a_;
  const at::Tensor &b_;
  const c10::optional<at::Tensor> &bias_;
  const c10::optional<at::Tensor> &c_;
  const c10::optional<at::Tensor> &d_;
  double alpha_;
  double beta_;
  int64_t block_m_;
  bool has_res_;
  bool has_output_;
  at::Tensor output_;

  MatmulSplitM(const at::Tensor &a,
               const at::Tensor &b,
               const c10::optional<at::Tensor> &bias,
               const c10::optional<at::Tensor> &c,
               const c10::optional<at::Tensor> &d,
               const double alpha,
               const double beta,
               const int64_t block_m)
      : a_(a), b_(b), bias_(bias), c_(c), d_(d), alpha_(alpha), beta_(beta), block_m_(block_m) {
    auto m = a_.size(0);
    if (block_m_ > m) {
      block_m_ = 1;
    } else if (block_m_ < 1) {
      block_m_ = m > 8192 ? 8 : (m < 2048 ? 1 : 4);
    }

    has_res_ = c_.has_value();
    has_output_ = d_.has_value();
    if (has_output_) {
      output_ = d_.value();
    } else {
      output_ = at::empty({a.size(0), b.size(0)}, a.options());
    }
  }

  auto split(const at::Tensor &a, const int64_t block_id) {
    auto m = a.size(0);
    auto m_per_blk = m / block_m_;
    auto remain = m % block_m_;
    auto start = block_id * m_per_blk + std::min(block_id, remain);
    auto end = (block_id + 1) * m_per_blk + std::min(block_id + 1, remain);
    return a.narrow(0, start, end - start);
  }

  at::Tensor forward(const int64_t block_id) {
    auto Di = split(output_, block_id);
    matmul(split(a_, block_id), b_, Di, bias_, has_res_ ? split(c_.value(), block_id) : c_, None,
           "none", alpha_, beta_, true, true, 1.0, 1.0, false, true);
    return Di;
  }

  int64_t getLoopNum() const { return block_m_; }
  at::Tensor getOutput() const { return output_; }
  at::Tensor getDSplit(const int64_t block_id) { return split(output_, block_id); }
};

at::Tensor matmul_allreduce(const int64_t cncl_comm,
                            const at::Tensor &a,
                            const at::Tensor &b,
                            const c10::optional<at::Tensor> &bias,
                            const c10::optional<at::Tensor> &c,
                            const c10::optional<at::Tensor> &d,
                            const double alpha,
                            const double beta,
                            const int64_t block_m) {
  MatmulSplitM mm(a, b, bias, c, d, alpha, beta, block_m);
  ParallelAllReduce<MatmulSplitM> parallel_rs(cncl_comm);
  return parallel_rs(mm);
}

}  //  namespace torch_api
}  //  namespace tmo
