/*
 * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

// This example performs the same computation as the one in the cuBLAS gemm example
// https://github.com/NVIDIA/CUDALibrarySamples/blob/master/cuBLAS/Level-3/gemm/cublas_gemm_example.cu

#include <cstdio>
#include <cstdlib>
#include <vector>
#include "../utils/stdblas_utils.h"
#include <linalg>

using data_type    = double;
using extents_type = std::extents<int, std::dynamic_extent, std::dynamic_extent>;

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

    const int m = 2;
    const int n = 2;
    const int k = 2;
    const int lda = 2;
    const int ldb = 2;
    const int ldc = 2;
    /*
     *   A = | 1.0 | 2.0 |
     *       | 3.0 | 4.0 |
     *
     *   B = | 5.0 | 6.0 |
     *       | 7.0 | 8.0 |
     */

    const std::vector<data_type> A = {1.0, 2.0, 3.0, 4.0};
    const std::vector<data_type> B = {5.0, 6.0, 7.0, 8.0};
    std::vector<data_type> C(m * n);

    printf("A\n");
    print_matrix(m, k, A.data(), lda);
    printf("=====\n");

    printf("B\n");
    print_matrix(k, n, B.data(), ldb);
    printf("=====\n");

    std::mdspan<const data_type, extents_type, std::layout_left > mds_A( A.data(), m, k );
    std::mdspan<const data_type, extents_type, std::layout_left > mds_B( B.data(), k, n );
    std::mdspan<      data_type, extents_type, std::layout_left > mds_C( C.data(), m, n );

    std::linalg::matrix_product( std::execution::par, mds_A, mds_B, mds_C );

    /*
     *   C = | 23.0 | 31.0 |
     *       | 34.0 | 46.0 |
     */

    printf("C\n");
    print_matrix(m, n, C.data(), ldc);
    printf("=====\n");

    return EXIT_SUCCESS;
}
