#include <immintrin.h>
#include <omp.h>
#include <algorithm>
#include <iostream>

extern "C" {
    // MAX LIMIT मैट्रिक्स मल्टीप्लिकेशन इंजन
    void vlpu_max_matmul(const float* __restrict__ A, const float* __restrict__ B, float* __restrict__ C, int M, int N, int K) {
        
        const int BM = 64; 
        const int BN = 128; // चौड़ाई दोगुनी कर दी गई है
        const int BK = 64;

        // थ्रेड्स को स्थिर रखने के लिए static शेड्यूलिंग
        #pragma omp parallel for collapse(2) schedule(static)
        for (int i = 0; i < M; i += BM) {
            for (int j = 0; j < N; j += BN) {
                for (int k = 0; k < K; k += BK) {
                    
                    for (int ii = i; ii < std::min(i + BM, M); ++ii) {
                        for (int kk = k; kk < std::min(k + BK, K); ++kk) {
                            
                            __m256 a_val = _mm256_set1_ps(A[ii * K + kk]);
                            int jj = j;
                            
                            // एक्सट्रीम अनरोलिंग: एक झटके में 64 कैलकुलेशन (8x8 रजिस्टर्स)
                            for (; jj <= std::min(j + BN, N) - 64; jj += 64) {
                                
                                // पहले 32 के लिए
                                __m256 b0 = _mm256_loadu_ps(&B[kk * N + jj]);
                                __m256 c0 = _mm256_loadu_ps(&C[ii * N + jj]);
                                c0 = _mm256_fmadd_ps(a_val, b0, c0);
                                _mm256_storeu_ps(&C[ii * N + jj], c0);

                                __m256 b1 = _mm256_loadu_ps(&B[kk * N + jj + 8]);
                                __m256 c1 = _mm256_loadu_ps(&C[ii * N + jj + 8]);
                                c1 = _mm256_fmadd_ps(a_val, b1, c1);
                                _mm256_storeu_ps(&C[ii * N + jj + 8], c1);

                                __m256 b2 = _mm256_loadu_ps(&B[kk * N + jj + 16]);
                                __m256 c2 = _mm256_loadu_ps(&C[ii * N + jj + 16]);
                                c2 = _mm256_fmadd_ps(a_val, b2, c2);
                                _mm256_storeu_ps(&C[ii * N + jj + 16], c2);

                                __m256 b3 = _mm256_loadu_ps(&B[kk * N + jj + 24]);
                                __m256 c3 = _mm256_loadu_ps(&C[ii * N + jj + 24]);
                                c3 = _mm256_fmadd_ps(a_val, b3, c3);
                                _mm256_storeu_ps(&C[ii * N + jj + 24], c3);

                                // अगले 32 के लिए
                                __m256 b4 = _mm256_loadu_ps(&B[kk * N + jj + 32]);
                                __m256 c4 = _mm256_loadu_ps(&C[ii * N + jj + 32]);
                                c4 = _mm256_fmadd_ps(a_val, b4, c4);
                                _mm256_storeu_ps(&C[ii * N + jj + 32], c4);

                                __m256 b5 = _mm256_loadu_ps(&B[kk * N + jj + 40]);
                                __m256 c5 = _mm256_loadu_ps(&C[ii * N + jj + 40]);
                                c5 = _mm256_fmadd_ps(a_val, b5, c5);
                                _mm256_storeu_ps(&C[ii * N + jj + 40], c5);

                                __m256 b6 = _mm256_loadu_ps(&B[kk * N + jj + 48]);
                                __m256 c6 = _mm256_loadu_ps(&C[ii * N + jj + 48]);
                                c6 = _mm256_fmadd_ps(a_val, b6, c6);
                                _mm256_storeu_ps(&C[ii * N + jj + 48], c6);

                                __m256 b7 = _mm256_loadu_ps(&B[kk * N + jj + 56]);
                                __m256 c7 = _mm256_loadu_ps(&C[ii * N + jj + 56]);
                                c7 = _mm256_fmadd_ps(a_val, b7, c7);
                                _mm256_storeu_ps(&C[ii * N + jj + 56], c7);
                            }
                            
                            // बचे हुए छोटे डेटा के लिए
                            for (; jj < std::min(j + BN, N); ++jj) {
                                C[ii * N + jj] += A[ii * K + kk] * B[kk * N + jj];
                            }
                        }
                    }
                }
            }
        }
    }
}