profile.cpp

#include <iostream>
#include <mpi.h>
#include <cassert>
#include <cmath>
#include <nvToolsExt.h>
#include <functional>
#include "matrix.hpp"
#include "distributedmatrix.hpp"

// Fonction utilitaire pour comparer deux matrices
bool matricesEqual(const Matrix& a, const Matrix& b, double tol = 1e-8) {
    if (a.numRows() != b.numRows() || a.numCols() != b.numCols()) return false;
    for (int i = 0; i < a.numRows(); ++i) {
        for (int j = 0; j < a.numCols(); ++j) {
            if (std::abs(a.get(i, j) - b.get(i, j)) > tol) return false;
        }
    }
    return true;
}

void testMultiplyTransposed(int rowsA, int colsA, int rowsB) {
    int rank, numProcs;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);
    MPI_Comm_size(MPI_COMM_WORLD, &numProcs);

    // Générer matrices pleines
    Matrix matrix1Full(rowsA, colsA); 
    Matrix matrix2Full(rowsB, colsA); 

    for (int i = 0; i < rowsA; i++) {
        for (int j = 0; j < colsA; j++) {
            matrix1Full.set(i, j, sin(i * colsA + j));
        }
    }

    for (int i = 0; i < rowsB; i++) {
        for (int j = 0; j < colsA; j++) {
            matrix2Full.set(i, j, cos(i * colsA + j));
        }
    }

    DistributedMatrix distA(matrix1Full, numProcs);
    DistributedMatrix distB(matrix2Full, numProcs);

    // Calcule de référence
    Matrix expected = matrix1Full * matrix2Full.transpose();

    if (rank == 0) {
        std::cout << "Expected result [0][0]: " << expected.get(0, 0) << std::endl;
    }

    // Test de multiplyTransposed avec marquage NVTX
    nvtxRangePushA("DistributedMatrix::multiplyTransposed");
    Matrix result = distA.multiplyTransposed(distB);
    nvtxRangePop();

    if (rank == 0) {
        std::cout << "Result [0][0]: " << result.get(0, 0) << std::endl;
    }

    // Vérification du résultat
    assert(matricesEqual(result, expected, 1e-8));
    if (rank == 0) {
        std::cout << "Test passed! Matrices are equal." << std::endl;
    }
}

int main(int argc, char** argv) {
    int initialized;
    MPI_Initialized(&initialized);
    if (!initialized) {
        MPI_Init(&argc, &argv);
    }
    
    int rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    if (rank == 0) {
        std::cout << "Starting MultiplyTransposed tests..." << std::endl;
    }

    try {
        testMultiplyTransposed(512, 256, 512);  // Taille paramétrable
        if (rank == 0) {
            std::cout << "Tests passed successfully!" << std::endl;
        }
    } 
    catch (std::exception& e) {
        if (rank == 0) {
            std::cerr << "Test failed with exception: " << e.what() << std::endl;
        }
        MPI_Abort(MPI_COMM_WORLD, 1);
    }
    
    // Finalize MPI if we initialized it
    // int finalized;
    // MPI_Finalized(&finalized);
    // if (!finalized && initialized) {
    //     MPI_Finalize();
    // }

    MPI_Finalize();
    return 0;
}