// This program computes the dot products over a set of independent vectors

// and compares the runtime between baseline (sequential), OpenMP, C++ thread,

// and Taskflow implementations.

#include <taskflow/taskflow.hpp>

#include <random>

#include <numeric>

#include <fstream>

using matrix_t = std::vector<std::vector<float>>;

// ----------------------------------------------------------------------------

// Utility section

// ----------------------------------------------------------------------------

// Function: random_matrix

matrix_t random_matrix(size_t N) {

thread_local std::default_random_engine gen(0);

std::normal_distribution<float> d{0.0f, 1.0f};

std::ostringstream oss;

oss << "|----> generating " << N << "x" << N << " matrix by thread "

<< std::this_thread::get_id() << "\n";

std::cout << oss.str();

matrix_t mat(N);

for(auto& r : mat) {

r.resize(N);

for(auto& c : r) {

c = d(gen);

}

}

return mat;

}

// Operator: multiplication

matrix_t operator * (const matrix_t& A, const matrix_t& B) {

if(A.empty() || B.empty() || A[0].size() != B.size()) {

std::cout << A[0].size() << " " << B.size() << std::endl;

throw std::runtime_error("Dimension mismatched in matrix multiplication\n");

}

size_t M, K, N;

N = A.size();

K = A[0].size();

M = B[0].size();

printf("A[%zux%zu] * B[%zux%zu]\n", N, K, K, M);

// Initialize the matrix

matrix_t ret(N);

for(auto& r : ret) {

r.resize(M);

for(auto& c : r) {

c = 0.0f;

}

}

// Matrix multiplication

for(size_t i=0; i<N; ++i) {

for(size_t j=0; j<M; ++j) {

for(size_t k=0; k<K; ++k) {

ret[i][j] += A[i][k] * B[k][j];

}

}

}

return ret;

}

// ----------------------------------------------------------------------------

// Task section

// ----------------------------------------------------------------------------

// Procedure: baseline

void baseline(const std::vector<size_t>& D) {

std::cout << "========== baseline ==========\n";

auto tbeg = std::chrono::steady_clock::now();

std::cout << "Generating matrix As ...\n";

std::vector<matrix_t> As(D.size());

for(size_t j=0; j<D.size(); ++j) {

As[j] = random_matrix(D[j]);

}

std::cout << "Generating matrix Bs ...\n";

std::vector<matrix_t> Bs(D.size());

for(size_t j=0; j<D.size(); ++j) {

Bs[j] = random_matrix(D[j]);

}

std::cout << "Computing matrix product values Cs ...\n";

std::vector<matrix_t> Cs(D.size());

for(size_t j=0; j<D.size(); ++j) {

Cs[j] = As[j] * Bs[j];

}

auto tend = std::chrono::steady_clock::now();

std::cout << "Baseline takes "

<< std::chrono::duration_cast<std::chrono::milliseconds>(tend-tbeg).count()

<< " ms\n";

}

// Procedure: taskflow

void taskflow(const std::vector<size_t>& D) {

auto tbeg = std::chrono::steady_clock::now();

tf::Taskflow tf;

std::cout << "Generating task As ...\n";

std::vector<matrix_t> As(D.size());

std::vector<tf::Task> TaskAs;

for(size_t j=0; j<D.size(); ++j) {

TaskAs.push_back(tf.emplace([&, j] () {

As[j] = random_matrix(D[j]);

}));

}

std::cout << "Generating task Bs ...\n";

std::vector<matrix_t> Bs(D.size());

std::vector<tf::Task> TaskBs;

for(size_t j=0; j<D.size(); ++j) {

TaskBs.push_back(tf.emplace([&, j] () {

Bs[j] = random_matrix(D[j]);

}));

}

std::cout << "Generating task Cs ...\n";

std::vector<matrix_t> Cs(D.size());

std::vector<tf::Task> TaskCs;

for(size_t j=0; j<D.size(); ++j) {

TaskCs.push_back(tf.emplace([&, j] () {

Cs[j] = As[j] * Bs[j];

}));

}

// Build task dependency

for(size_t j=0; j<D.size(); ++j) {

TaskCs[j].gather({TaskAs[j], TaskBs[j]});

}

tf.wait_for_all();

auto tend = std::chrono::steady_clock::now();

std::cout << "Taskflow takes "

<< std::chrono::duration_cast<std::chrono::milliseconds>(tend-tbeg).count()

<< " ms\n";

}

// ------------------------------------------------------------------------------------------------

// Function: main

int main(int argc, char* argv[]) {

if(argc != 3) {

std::cerr << "usage: ./matrix [baseline|openmp|cppthread|taskflow] N\n";

std::exit(EXIT_FAILURE);

}

// Create a unbalanced dimension for vector products.

const auto N = std::stoul(argv[2]);

std::vector<size_t> dimensions(N);

std::default_random_engine engine(0);

std::uniform_int_distribution dis(1, 1000);

std::cout << "matrix sizes = [";

for(size_t i=0; i<dimensions.size(); ++i) {

dimensions[i] = dis(engine);

if(i) std::cout << ' ';

std::cout << dimensions[i];

}

std::cout << "]\n";

// Run methods

if(std::string_view method(argv[1]); method == "baseline") {

baseline(dimensions);

}

else if(method == "taskflow") {

taskflow(dimensions);

}

else {

std::cerr << "wrong method, shoud be [baseline|taskflow]\n";

}

return 0;

}