#include <stdio.h>
#include <stdlib.h>
#include <omp.h>

int main(void)
{
    const int N     = 16384;
    const int total = N * N;   /* 256M floats per array */

    float *A = malloc(total * sizeof(float));
    float *B = malloc(total * sizeof(float));
    float *C = malloc(total * sizeof(float));

    if (!A || !B || !C) {
        fprintf(stderr, "allocation failed\n");
        return 1;
    }

    /* ------------------------------------------------------------------
     * Keep all three arrays resident on the GPU for both kernels.
     * map(alloc:) reserves device memory without copying anything to or
     * from the host - the initialization kernel below writes the values
     * directly on the device, so there is nothing to transfer in.
     * ------------------------------------------------------------------ */
    #pragma omp target data map(alloc: A[0:total], B[0:total], C[0:total])
    {
        /* Initialization kernel
         * collapse(2) merges the i and j loops into a single 256M-iteration
         * range, giving the runtime more parallelism to distribute across
         * the GPU thread grid. */
        #pragma omp target teams distribute parallel for collapse(2)
        for (int i = 0; i < N; ++i)
            for (int j = 0; j < N; ++j) {
                A[i*N+j] = (float)((i + j) %  97 + 1) * 0.01f;
                B[i*N+j] = (float)((i * j) % 101 + 1) * 0.01f;
                C[i*N+j] = 0.0f;
            }

        double t0 = omp_get_wtime();

        /* ------------------------------------------------------------------
         * Dense matrix multiply  C = A * B
         *
         * teams distribute  - partitions the collapsed (i,j) iteration space
         *                     across thread blocks (teams) on the GPU.
         * parallel for      - distributes iterations across threads within
         *                     each block.
         * The inner k-loop is sequential inside each thread, identical in
         * structure to the Fortran DO CONCURRENT and C++ par_unseq versions.
         * ------------------------------------------------------------------ */
        #pragma omp target teams distribute parallel for collapse(2)
        for (int i = 0; i < N; ++i)
            for (int j = 0; j < N; ++j) {
                float acc = 0.0f;
                for (int k = 0; k < N; ++k)
                    acc += A[i*N+k] * B[k*N+j];
                C[i*N+j] = acc;
            }

        double t1 = omp_get_wtime();

        double elapsed = t1 - t0;
        double gflops  = 2.0 * N * N * N / (elapsed * 1.0e9);

        /* Copy C back to host for the checksum; A and B stay on the device
         * and are discarded when the target data region closes. */
        #pragma omp target update from(C[0:total])

        float checksum = 0.0f;
        for (int i = 0; i < 16; ++i)
            for (int j = 0; j < 16; ++j)
                checksum += C[i*N+j];

        printf("Matrix size : %d x %d\n", N, N);
        printf("Elapsed     : %.3f s\n",   elapsed);
        printf("Performance : %.3f GFLOP/s\n", gflops);
        printf("Checksum    : %e\n",        checksum);
    }

    free(A);
    free(B);
    free(C);
    return 0;
}