std::linalg::triangular_matrix_product(C++26)

#include <array>
#include <iostream>
#include <linalg>
#include <mdspan>
#include <vector>

template <class Matrix>
void print_mat(const Matrix& A) {
  for(int i = 0; i < A.extent(0); ++i) {
    for(int j = 0; j < A.extent(1) - 1; ++j) {
      std::cout << A[i, j] << ' ';
    }
    std::cout << A[i, A.extent(1) - 1] << '\n';
  }
}

template <class Matrix>
void init_mat(Matrix& A, typename Matrix::value_type geta = 1) {
  for(int i = 0; i < A.extent(0); ++i) {
    for(int j = 0; j < A.extent(1); ++j) {
      A[i, j] = i * A.extent(1) + j + geta;
    }
  }
}

template <class Matrix>
void init_tria_mat(Matrix& A) {
  for(int i = 0; i < A.extent(0); ++i) {
    for(int j = i + 1; j < A.extent(1); ++j) {
      A[i, j] = i * A.extent(1) + j;
    }
  }
}

int main()
{
  constexpr size_t N = 2;

  using Scalar = double;
  using Vector = std::vector<Scalar>;
  using TriangularMatrix = std::mdspan<
      Scalar,
      std::extents<size_t, N, N>,
      std::linalg::layout_blas_packed<
        std::linalg::upper_triangle_t,
        std::linalg::row_major_t>
    >;

  Vector A_vec(N * N);
  Vector B_vec(N * N);
  Vector C_vec(N * N);
  Vector E_vec(N * N);

  std::mdspan C(C_vec.data(), N, N);
  std::mdspan E(E_vec.data(), N, N);

  init_mat(E, N * N);

  {
    TriangularMatrix A(A_vec.data());
    std::mdspan      B(B_vec.data(), N, N);

    init_tria_mat(A);
    init_mat(B);

    // (1)
    std::cout << "(1)\n";
    std::linalg::triangular_matrix_product(
      A,
      std::linalg::upper_triangle,
      std::linalg::implicit_unit_diagonal,
      B,
      C);
    print_mat(C);

    // (2)
    std::cout << "(2)\n";
    std::linalg::triangular_matrix_product(
      std::execution::par,
      A,
      std::linalg::upper_triangle,
      std::linalg::implicit_unit_diagonal,
      B,
      C);
    print_mat(C);
  }

  {
    std::mdspan      A(A_vec.data(), N, N);
    TriangularMatrix B(B_vec.data());

    init_mat(A);
    init_tria_mat(B);

    // (3)
    std::cout << "(3)\n";
    std::linalg::triangular_matrix_product(
      A,
      B,
      std::linalg::upper_triangle,
      std::linalg::implicit_unit_diagonal,
      C);
    print_mat(C);

    // (4)
    std::cout << "(4)\n";
    std::linalg::triangular_matrix_product(
      std::execution::par,
      A,
      B,
      std::linalg::upper_triangle,
      std::linalg::implicit_unit_diagonal,
      C);
    print_mat(C);
  }

  {
    TriangularMatrix A(A_vec.data());
    std::mdspan      B(B_vec.data(), N, N);

    init_tria_mat(A);
    init_mat(B);

    // (5)
    std::cout << "(5)\n";
    std::linalg::triangular_matrix_product(
      A,
      std::linalg::upper_triangle,
      std::linalg::implicit_unit_diagonal,
      B,
      E,
      C);
    print_mat(C);

    // (6)
    std::cout << "(6)\n";
    std::linalg::triangular_matrix_product(
      std::execution::par,
      A,
      std::linalg::upper_triangle,
      std::linalg::implicit_unit_diagonal,
      B,
      E,
      C);
    print_mat(C);
  }

  {
    std::mdspan      A(A_vec.data(), N, N);
    TriangularMatrix B(B_vec.data());

    init_mat(A);
    init_tria_mat(B);

    // (7)
    std::cout << "(7)\n";
    std::linalg::triangular_matrix_product(
      A,
      B,
      std::linalg::upper_triangle,
      std::linalg::implicit_unit_diagonal,
      E,
      C);
    print_mat(C);

    // (8)
    std::cout << "(8)\n";
    std::linalg::triangular_matrix_product(
      std::execution::par,
      A,
      B,
      std::linalg::upper_triangle,
      std::linalg::implicit_unit_diagonal,
      E,
      C);
    print_mat(C);
  }

  return 0;
}