impl.h

//
// Distributed Linear Algebra with Future (DLAF)
//
// Copyright (c) 2018-2021, ETH Zurich
// All rights reserved.
//
// Please, refer to the LICENSE file in the root directory.
// SPDX-License-Identifier: BSD-3-Clause
//
#pragma once

#include <hpx/include/util.hpp>
#include <hpx/local/execution.hpp>
#include <hpx/local/future.hpp>
#include <hpx/local/thread.hpp>

#include "dlaf/blas/tile.h"
#include "dlaf/common/index2d.h"
#include "dlaf/common/pipeline.h"
#include "dlaf/common/vector.h"
#include "dlaf/communication/communicator.h"
#include "dlaf/communication/communicator_grid.h"
#include "dlaf/communication/executor.h"
#include "dlaf/communication/kernels.h"
#include "dlaf/executors.h"
#include "dlaf/lapack/tile.h"
#include "dlaf/matrix/distribution.h"
#include "dlaf/matrix/matrix.h"
#include "dlaf/sender/when_all_lift.h"
#include "dlaf/solver/triangular/api.h"
#include "dlaf/util_matrix.h"

namespace dlaf {
namespace solver {
namespace internal {
template <Backend backend, typename T, typename InSender, typename OutSender>
void trsmBPanelTile(blas::Side side, blas::Uplo uplo, blas::Op op, blas::Diag diag, T alpha,
                    InSender&& in_tile, OutSender&& out_tile) {
  dlaf::internal::whenAllLift(side, uplo, op, diag, alpha, std::forward<InSender>(in_tile),
                              std::forward<OutSender>(out_tile)) |
      tile::trsm(dlaf::internal::Policy<backend>(hpx::threads::thread_priority::high)) |
      hpx::execution::experimental::detach();
}

template <Backend backend, typename T, typename ASender, typename BSender, typename CSender>
void gemmTrailingMatrixTile(hpx::threads::thread_priority priority, blas::Op op_a, blas::Op op_b, T beta,
                            ASender&& a_tile, BSender&& b_tile, CSender&& c_tile) {
  dlaf::internal::whenAllLift(op_a, op_b, beta, std::forward<ASender>(a_tile),
                              std::forward<BSender>(b_tile), T(1.0), std::forward<CSender>(c_tile)) |
      tile::gemm(dlaf::internal::Policy<backend>(priority)) | hpx::execution::experimental::detach();
}

template <Backend backend, Device device, class T>
void Triangular<backend, device, T>::call_LLN(blas::Diag diag, T alpha, Matrix<const T, device>& mat_a,
                                              Matrix<T, device>& mat_b) {
  using hpx::threads::thread_priority;

  constexpr auto Left = blas::Side::Left;
  constexpr auto Lower = blas::Uplo::Lower;
  constexpr auto NoTrans = blas::Op::NoTrans;

  SizeType m = mat_b.nrTiles().rows();
  SizeType n = mat_b.nrTiles().cols();

  for (SizeType k = 0; k < m; ++k) {
    for (SizeType j = 0; j < n; ++j) {
      auto kj = LocalTileIndex{k, j};

      // Triangular solve of k-th row Panel of B
      trsmBPanelTile<backend>(Left, Lower, blas::Op::NoTrans, diag, alpha,
                              mat_a.read_sender(LocalTileIndex{k, k}), mat_b.readwrite_sender(kj));

      for (SizeType i = k + 1; i < m; ++i) {
        // Choose queue priority
        const auto priority = (i == k - 1) ? thread_priority::high : thread_priority::normal;

        auto beta = static_cast<T>(-1.0) / alpha;
        // Update trailing matrix
        gemmTrailingMatrixTile<backend>(priority, blas::Op::NoTrans, blas::Op::NoTrans, beta,
                                        mat_a.read_sender(LocalTileIndex{i, k}), mat_b.read_sender(kj),
                                        mat_b.readwrite_sender(LocalTileIndex{i, j}));
      }
    }
  }
}

template <Backend backend, Device device, class T>
void Triangular<backend, device, T>::call_LLT(blas::Op op, blas::Diag diag, T alpha,
                                              Matrix<const T, device>& mat_a, Matrix<T, device>& mat_b) {
  using hpx::threads::thread_priority;

  constexpr auto Left = blas::Side::Left;
  constexpr auto Lower = blas::Uplo::Lower;
  constexpr auto NoTrans = blas::Op::NoTrans;

  SizeType m = mat_b.nrTiles().rows();
  SizeType n = mat_b.nrTiles().cols();

  for (SizeType k = m - 1; k > -1; --k) {
    for (SizeType j = n - 1; j > -1; --j) {
      auto kj = LocalTileIndex{k, j};
      // Triangular solve of k-th row Panel of B
      trsmBPanelTile<backend>(Left, Lower, op, diag, alpha, mat_a.read_sender(LocalTileIndex{k, k}),
                              mat_b.readwrite_sender(kj));

      for (SizeType i = k - 1; i > -1; --i) {
        // Choose queue priority
        const auto priority = (i == k - 1) ? thread_priority::high : thread_priority::normal;

        auto beta = static_cast<T>(-1.0) / alpha;

        // Update trailing matrix
        gemmTrailingMatrixTile<backend>(priority, op, blas::Op::NoTrans, beta,
                                        mat_a.read_sender(LocalTileIndex{k, i}), mat_b.read_sender(kj),
                                        mat_b.readwrite_sender(LocalTileIndex{i, j}));
      }
    }
  }
}

template <Backend backend, Device device, class T>
void Triangular<backend, device, T>::call_LUN(blas::Diag diag, T alpha, Matrix<const T, device>& mat_a,
                                              Matrix<T, device>& mat_b) {
  using hpx::threads::thread_priority;

  constexpr auto Left = blas::Side::Left;
  constexpr auto Upper = blas::Uplo::Upper;
  constexpr auto NoTrans = blas::Op::NoTrans;

  SizeType m = mat_b.nrTiles().rows();
  SizeType n = mat_b.nrTiles().cols();

  for (SizeType k = m - 1; k > -1; --k) {
    for (SizeType j = n - 1; j > -1; --j) {
      auto kj = LocalTileIndex{k, j};
      // Triangular solve of k-th row Panel of B
      trsmBPanelTile<backend>(Left, Upper, NoTrans, diag, alpha, mat_a.read_sender(LocalTileIndex{k, k}),
                              mat_b.readwrite_sender(kj));

      for (SizeType i = k - 1; i > -1; --i) {
        // Choose queue priority
        const auto priority = (i == k - 1) ? thread_priority::high : thread_priority::normal;

        auto beta = static_cast<T>(-1.0) / alpha;
        // Update trailing matrix
        gemmTrailingMatrixTile<backend>(priority, NoTrans, NoTrans, beta,
                                        mat_a.read_sender(LocalTileIndex{i, k}), mat_b.read_sender(kj),
                                        mat_b.readwrite_sender(LocalTileIndex{i, j}));
      }
    }
  }
}

template <Backend backend, Device device, class T>
void Triangular<backend, device, T>::call_LUT(blas::Op op, blas::Diag diag, T alpha,
                                              Matrix<const T, device>& mat_a, Matrix<T, device>& mat_b) {
  using hpx::threads::thread_priority;

  constexpr auto Left = blas::Side::Left;
  constexpr auto Upper = blas::Uplo::Upper;
  constexpr auto NoTrans = blas::Op::NoTrans;

  SizeType m = mat_b.nrTiles().rows();
  SizeType n = mat_b.nrTiles().cols();

  for (SizeType k = 0; k < m; ++k) {
    for (SizeType j = 0; j < n; ++j) {
      auto kj = LocalTileIndex{k, j};

      // Triangular solve of k-th row Panel of B
      trsmBPanelTile<backend>(Left, Upper, op, diag, alpha, mat_a.read_sender(LocalTileIndex{k, k}),
                              mat_b.readwrite_sender(kj));

      for (SizeType i = k + 1; i < m; ++i) {
        // Choose queue priority
        const auto priority = (i == k - 1) ? thread_priority::high : thread_priority::normal;

        auto beta = static_cast<T>(-1.0) / alpha;
        // Update trailing matrix
        gemmTrailingMatrixTile<backend>(priority, op, NoTrans, beta,
                                        mat_a.read_sender(LocalTileIndex{k, i}), mat_b.read_sender(kj),
                                        mat_b.readwrite_sender(LocalTileIndex{i, j}));
      }
    }
  }
}

template <Backend backend, Device device, class T>
void Triangular<backend, device, T>::call_RLN(blas::Diag diag, T alpha, Matrix<const T, device>& mat_a,
                                              Matrix<T, device>& mat_b) {
  using hpx::threads::thread_priority;

  constexpr auto Right = blas::Side::Right;
  constexpr auto Lower = blas::Uplo::Lower;
  constexpr auto NoTrans = blas::Op::NoTrans;

  SizeType m = mat_b.nrTiles().rows();
  SizeType n = mat_b.nrTiles().cols();

  for (SizeType k = n - 1; k > -1; --k) {
    for (SizeType i = m - 1; i > -1; --i) {
      auto ik = LocalTileIndex{i, k};

      // Triangular solve of k-th col Panel of B
      trsmBPanelTile<backend>(Right, Lower, NoTrans, diag, alpha,
                              mat_a.read_sender(LocalTileIndex{k, k}), mat_b.readwrite_sender(ik));

      for (SizeType j = k - 1; j > -1; --j) {
        // Choose queue priority
        const auto priority = (i == k - 1) ? thread_priority::high : thread_priority::normal;

        auto beta = static_cast<T>(-1.0) / alpha;
        // Update trailing matrix
        gemmTrailingMatrixTile<backend>(priority, NoTrans, NoTrans, beta, mat_b.read_sender(ik),
                                        mat_a.read_sender(LocalTileIndex{k, j}),
                                        mat_b.readwrite_sender(LocalTileIndex{i, j}));
      }
    }
  }
}

template <Backend backend, Device device, class T>
void Triangular<backend, device, T>::call_RLT(blas::Op op, blas::Diag diag, T alpha,
                                              Matrix<const T, device>& mat_a, Matrix<T, device>& mat_b) {
  using hpx::threads::thread_priority;

  constexpr auto Right = blas::Side::Right;
  constexpr auto Lower = blas::Uplo::Lower;
  constexpr auto NoTrans = blas::Op::NoTrans;

  SizeType m = mat_b.nrTiles().rows();
  SizeType n = mat_b.nrTiles().cols();

  for (SizeType k = 0; k < n; ++k) {
    for (SizeType i = 0; i < m; ++i) {
      auto ik = LocalTileIndex{i, k};

      // Triangular solve of k-th col Panel of B
      trsmBPanelTile<backend>(Right, Lower, op, diag, alpha, mat_a.read_sender(LocalTileIndex{k, k}),
                              mat_b.readwrite_sender(ik));

      for (SizeType j = k + 1; j < n; ++j) {
        // Choose queue priority
        const auto priority = (i == k - 1) ? thread_priority::high : thread_priority::normal;

        auto beta = static_cast<T>(-1.0) / alpha;
        // Update trailing matrix
        gemmTrailingMatrixTile<backend>(priority, NoTrans, op, beta, mat_b.read_sender(ik),
                                        mat_a.read_sender(LocalTileIndex{j, k}),
                                        mat_b.readwrite_sender(LocalTileIndex{i, j}));
      }
    }
  }
}

template <Backend backend, Device device, class T>
void Triangular<backend, device, T>::call_RUN(blas::Diag diag, T alpha, Matrix<const T, device>& mat_a,
                                              Matrix<T, device>& mat_b) {
  using hpx::threads::thread_priority;

  constexpr auto Right = blas::Side::Right;
  constexpr auto Upper = blas::Uplo::Upper;
  constexpr auto NoTrans = blas::Op::NoTrans;

  SizeType m = mat_b.nrTiles().rows();
  SizeType n = mat_b.nrTiles().cols();

  for (SizeType k = 0; k < n; ++k) {
    for (SizeType i = 0; i < m; ++i) {
      auto ik = LocalTileIndex{i, k};

      // Triangular solve of k-th col Panel of B
      trsmBPanelTile<backend>(Right, Upper, NoTrans, diag, alpha,
                              mat_a.read_sender(LocalTileIndex{k, k}), mat_b.readwrite_sender(ik));

      for (SizeType j = k + 1; j < n; ++j) {
        // Choose queue priority
        const auto priority = (i == k - 1) ? thread_priority::high : thread_priority::normal;

        auto beta = static_cast<T>(-1.0) / alpha;
        // Update trailing matrix
        gemmTrailingMatrixTile<backend>(priority, NoTrans, NoTrans, beta, mat_b.read_sender(ik),
                                        mat_a.read_sender(LocalTileIndex{k, j}),
                                        mat_b.readwrite_sender(LocalTileIndex{i, j}));
      }
    }
  }
}

template <Backend backend, Device device, class T>
void Triangular<backend, device, T>::call_RUT(blas::Op op, blas::Diag diag, T alpha,
                                              Matrix<const T, device>& mat_a, Matrix<T, device>& mat_b) {
  using hpx::threads::thread_priority;

  constexpr auto Right = blas::Side::Right;
  constexpr auto Upper = blas::Uplo::Upper;
  constexpr auto NoTrans = blas::Op::NoTrans;

  SizeType m = mat_b.nrTiles().rows();
  SizeType n = mat_b.nrTiles().cols();

  for (SizeType k = n - 1; k > -1; --k) {
    for (SizeType i = m - 1; i > -1; --i) {
      auto ik = LocalTileIndex{i, k};

      // Triangular solve of k-th col Panel of B
      trsmBPanelTile<backend>(Right, Upper, op, diag, alpha, mat_a.read_sender(LocalTileIndex{k, k}),
                              mat_b.readwrite_sender(ik));

      for (SizeType j = k - 1; j > -1; --j) {
        // Choose queue priority
        const auto priority = (i == k - 1) ? thread_priority::high : thread_priority::normal;

        auto beta = static_cast<T>(-1.0) / alpha;
        // Update trailing matrix
        gemmTrailingMatrixTile<backend>(priority, NoTrans, op, beta, mat_b.read_sender(ik),
                                        mat_a.read_sender(LocalTileIndex{j, k}),
                                        mat_b.readwrite_sender(LocalTileIndex{i, j}));
      }
    }
  }
}

template <Backend backend, Device device, class T>
void Triangular<backend, device, T>::call_LLN(comm::CommunicatorGrid grid, blas::Diag diag, T alpha,
                                              Matrix<const T, device>& mat_a, Matrix<T, device>& mat_b) {
  using hpx::execution::experimental::keep_future;
  using hpx::threads::thread_priority;

  constexpr auto Left = blas::Side::Left;
  constexpr auto Lower = blas::Uplo::Lower;
  constexpr auto NoTrans = blas::Op::NoTrans;

  using common::internal::vector;
  using ConstTileType = typename Matrix<T, device>::ConstTileType;

  auto executor_mpi = dlaf::getMPIExecutor<backend>();

  // Set up MPI
  common::Pipeline<comm::Communicator> mpi_col_task_chain(grid.colCommunicator());
  common::Pipeline<comm::Communicator> mpi_row_task_chain(grid.rowCommunicator());

  const matrix::Distribution& distr_a = mat_a.distribution();
  const matrix::Distribution& distr_b = mat_b.distribution();
  SizeType a_rows = mat_a.nrTiles().rows();
  auto a_local_rows = distr_a.localNrTiles().rows();
  auto b_local_cols = distr_b.localNrTiles().cols();

  for (SizeType k = 0; k < a_rows; ++k) {
    // Create a placeholder that will store the shared futures representing the panel
    vector<hpx::shared_future<ConstTileType>> panel(distr_b.localNrTiles().cols());

    auto k_rank_row = distr_a.rankGlobalTile<Coord::Row>(k);
    auto k_rank_col = distr_a.rankGlobalTile<Coord::Col>(k);

    hpx::shared_future<ConstTileType> kk_tile;

    if (mat_a.rankIndex().row() == k_rank_row) {
      auto k_local_row = distr_a.localTileFromGlobalTile<Coord::Row>(k);

      if (mat_a.rankIndex().col() == k_rank_col) {
        // Broadcast A(kk) row-wise
        auto k_local_col = distr_a.localTileFromGlobalTile<Coord::Col>(k);
        auto kk = LocalTileIndex{k_local_row, k_local_col};
        kk_tile = mat_a.read(kk);
        comm::scheduleSendBcast(executor_mpi, kk_tile, mpi_row_task_chain());
      }
      else {
        kk_tile =
            comm::scheduleRecvBcastAlloc<T, device>(executor_mpi, mat_a.tileSize(GlobalTileIndex(k, k)),
                                                    k_rank_col, mpi_row_task_chain());
      }
    }

    for (SizeType j_local = 0; j_local < b_local_cols; ++j_local) {
      auto j = distr_b.globalTileFromLocalTile<Coord::Col>(j_local);

      // Triangular solve B's k-th row panel and broadcast B(kj) column-wise
      if (mat_b.rankIndex().row() == k_rank_row) {
        auto k_local_row = distr_b.localTileFromGlobalTile<Coord::Row>(k);
        auto kj = LocalTileIndex{k_local_row, j_local};
        trsmBPanelTile<backend>(Left, Lower, blas::Op::NoTrans, diag, alpha, keep_future(kk_tile),
                                mat_b.readwrite_sender(kj));

        panel[j_local] = mat_b.read(kj);
        if (k != (mat_b.nrTiles().rows() - 1)) {
          comm::scheduleSendBcast(executor_mpi, panel[j_local], mpi_col_task_chain());
        }
      }
      else {
        if (k != (mat_b.nrTiles().rows() - 1)) {
          panel[j_local] = comm::scheduleRecvBcastAlloc<T, device>(executor_mpi,
                                                                   mat_b.tileSize(GlobalTileIndex(k, j)),
                                                                   k_rank_row, mpi_col_task_chain());
        }
      }
    }

    for (SizeType i_local = distr_a.nextLocalTileFromGlobalTile<Coord::Row>(k + 1);
         i_local < a_local_rows; ++i_local) {
      auto i = distr_a.globalTileFromLocalTile<Coord::Row>(i_local);

      // Choose queue priority
      const auto priority = (i == k - 1) ? thread_priority::high : thread_priority::normal;

      hpx::shared_future<ConstTileType> ik_tile;

      // Broadcast A(ik) row-wise
      if (mat_a.rankIndex().col() == k_rank_col) {
        auto k_local_col = distr_a.localTileFromGlobalTile<Coord::Col>(k);
        auto ik = LocalTileIndex{i_local, k_local_col};
        ik_tile = mat_a.read(ik);
        comm::scheduleSendBcast(executor_mpi, ik_tile, mpi_row_task_chain());
      }
      else {
        ik_tile =
            comm::scheduleRecvBcastAlloc<T, device>(executor_mpi, mat_a.tileSize(GlobalTileIndex(i, k)),
                                                    k_rank_col, mpi_row_task_chain());
      }

      // Update trailing matrix
      for (SizeType j_local = 0; j_local < b_local_cols; ++j_local) {
        T beta = T(-1.0) / alpha;
        gemmTrailingMatrixTile<backend>(priority, blas::Op::NoTrans, blas::Op::NoTrans, beta,
                                        keep_future(ik_tile), keep_future(panel[j_local]),
                                        mat_b.readwrite_sender(LocalTileIndex{i_local, j_local}));
      }
    }
  }
}

}
}
}