vLPU7 / vlpu_max_core.cpp
Vedika
Create vlpu_max_core.cpp
1e6870b verified
Raw
History Blame Contribute Delete
4.49 kB
#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];
}
}
}
}
}
}
}
}