MM.c

#include "MM.h"


#include <stdio.h>
#include <immintrin.h>
#include <stddef.h>

#include <string.h>
#include <stdlib.h>

#ifdef _WIN32
#include <malloc.h>
#endif

static void *aligned_malloc(size_t alignment,size_t size)
{
#ifdef _WIN32
    return _aligned_malloc(size, alignment);
#else
    return aligned_alloc(alignment, size);
#endif
}

static void aligned_free(void *ptr)
{
#ifdef _WIN32
    _aligned_free(ptr);
#else
    free(ptr);
#endif
}


void MM(size_t n,double* a,double* b,double* c)
{
    size_t size = n*n*sizeof(double);
    memset(c, 0, size);
    for(size_t i=0;i<n;i++)
    {
        for(size_t j=0;j<n;j++)
        {
            double sum = 0.0;
            for(size_t k=0;k<n;k++)
            {
                   sum+=a[i*n+k]*b[k*n+j];
            }
            c[i*n+j]=sum;
        }
    }
}
void MMIKJ(size_t n,double* a,double* b,double* c)
{
    size_t size=n*n*sizeof(double);
    
    memset(c,0,size); //make sure c are 0.

    for(size_t i=0;i<n;i++)
    {
        for(size_t k=0;k<n;k++)
        {
            double a_ik = a[i*n+k];
            
            for(size_t j=0;j<n;j++)
            {
                c[i*n + j]+=a_ik*b[k*n+j];                 
            }
        }
    }       
}

void MMBlocked(size_t n,double* a,double* b,double* c)
{
    size_t size=n*n*sizeof(double);
    
    memset(c,0,size); //make sure c are 0.

    for(size_t ii=0;ii<n;ii+=32)
    {
        for(size_t kk=0;kk<n;kk+=32)
        {
            for(size_t jj=0;jj<n;jj+=32)
            {
                for(size_t i=ii;i<ii+32&&i<n;i++)
                {
                    for(size_t j=jj;j<jj+32&&j<n;j++)
                    {
                        double sum = 0.0;
                        for(size_t k=kk;k<kk+32&&k<n;k++)
                        {
                            sum+=a[i*n+k]*b[k*n+j];
                        }
                        c[i*n+j]+=sum;
                    }
                }
            }
        }
    }
}

void optimizedMM(size_t n,double* a,double* b,double* c)
{
    size_t size=n*n*sizeof(double);
    
    memset(c,0,size); //make sure c are 0.

    for(size_t ii=0;ii<n;ii+=32)
    {
        for(size_t kk=0;kk<n;kk+=32)
        {
            for(size_t jj=0;jj<n;jj+=32)
            {
                for(size_t i=ii;i<ii+32&&i<n;i++)
                {
                    for(size_t k=kk;k<kk+32&&k<n;k++)
                    {
                        double a_ik = a[i*n+k];
                        size_t j=jj;
                        size_t limit = jj+32<n?jj+32:n;
                        for(;j+3<limit;j+=4)
                        {
                            c[i*n + j]+=a_ik*b[k*n+j];
                            c[i*n + j+1]+=a_ik*b[k*n+j+1];
                            c[i*n + j+2]+=a_ik*b[k*n+j+2];
                            c[i*n + j+3]+=a_ik*b[k*n+j+3];
                        }
                        for(; j<limit; j++)
                        {
                            c[i*n + j] += a_ik * b[k*n + j];
                        }
                    }
                }
            }
        }
    }
}
void optimizedMM2(size_t n,double* a,double* b,double* c)
{   
    
    size_t paddingSize = ((n+31)/32)*32;
    
    
    double* paddingA = (double*)aligned_malloc(64,paddingSize*paddingSize*sizeof(double));
    double* paddingB = (double*)aligned_malloc(64,paddingSize*paddingSize*sizeof(double));
    double* paddingC = (double*)aligned_malloc(64,paddingSize*paddingSize*sizeof(double));
    
    if (paddingA == NULL || paddingB == NULL || paddingC == NULL) 
    {
    printf("Memory allocation failed!\n");
    return;
    }

    size_t size = paddingSize*paddingSize*sizeof(double);

    memset(paddingA,0,size);
    memset(paddingB,0,size);
    memset(paddingC,0,size); 

    double* pa = paddingA;
    double* aRow = a;
    double* pb = paddingB;
    double* bRow = b;
    for(size_t i=0;i<n;i++)
    {
        memcpy(pa,aRow,n*sizeof(double));
        pa+=paddingSize;
        aRow+=n;

        memcpy(pb,bRow,n*sizeof(double));
        pb+=paddingSize;
        bRow+=n;
    }


    for(size_t ii=0;ii<paddingSize;ii+=32)
    {
        for(size_t kk=0;kk<paddingSize;kk+=32)
        {
            for(size_t jj=0;jj<paddingSize;jj+=32)
            {
                for(size_t i=ii;i<ii+32;i++)
                {
                    for(size_t k=kk;k<kk+32;k++)
                    {
                        double a_ik = paddingA[i*paddingSize+k];
                        
                        for(size_t j=jj;j<jj+32;j+=4)
                        {
                            paddingC[i*paddingSize + j]+=a_ik*paddingB[k*paddingSize+j];
                            paddingC[i*paddingSize + j+1]+=a_ik*paddingB[k*paddingSize+j+1];
                            paddingC[i*paddingSize + j+2]+=a_ik*paddingB[k*paddingSize+j+2];
                            paddingC[i*paddingSize + j+3]+=a_ik*paddingB[k*paddingSize+j+3];
                        }
                    }
                }
            }
        }
    }   
    double* pc = paddingC;
    double* cRow = c;
    for (size_t i = 0; i < n; i++)
    {
        memcpy(cRow,pc,n*sizeof(double));
        cRow+=n;
        pc+=paddingSize;
    }

    aligned_free(paddingA);
    aligned_free(paddingB);
    aligned_free(paddingC);
    
}

void optimizedMM3(size_t n,double* a,double* b,double* c)
{   //size(length or wide..) after add padding will be divisible by 48(use 48 step to unroaling)
    //also can be divisible by 4 and 6, (6x4 micro kernal)  

   
    size_t paddingSize = ((n+47)/48)*48; 
    
    
    double* paddingA = (double*)aligned_malloc(64,paddingSize*paddingSize*sizeof(double));
    double* paddingB = (double*)aligned_malloc(64,paddingSize*paddingSize*sizeof(double));
    double* paddingC = (double*)aligned_malloc(64,paddingSize*paddingSize*sizeof(double));
    
    if (!paddingA||!paddingB||!paddingC) 
    {
        printf("Memory allocation failed!\n");
        return;
    }

    size_t size = paddingSize*paddingSize*sizeof(double);

    memset(paddingA,0,size);
    memset(paddingB,0,size);
    memset(paddingC,0,size); 

    double* pa = paddingA;
    double* aRow = a;
    double* pb = paddingB;
    double* bRow = b;

    for(size_t i=0;i<n;i++)
    {
        memcpy(pa,aRow,n*sizeof(double));
        pa+=paddingSize;
        aRow+=n;

        memcpy(pb,bRow,n*sizeof(double));
        pb+=paddingSize;
        bRow+=n;
    }


    for(size_t ii=0;ii<paddingSize;ii+=48)
    {
        for(size_t kk=0;kk<paddingSize;kk+=48)
        {
            for(size_t jj=0;jj<paddingSize;jj+=48)
            {
                for(size_t i=ii;i<ii+48;i+=6)
                {   
                    for(size_t j=jj;j<jj+48;j+=4)
                    {   
                        __m256d c1 = _mm256_load_pd(paddingC+(i)*paddingSize+j);
                        __m256d c2 = _mm256_load_pd(paddingC+(i+1)*paddingSize+j);
                        __m256d c3 = _mm256_load_pd(paddingC+(i+2)*paddingSize+j);
                        __m256d c4 = _mm256_load_pd(paddingC+(i+3)*paddingSize+j);
                        __m256d c5 = _mm256_load_pd(paddingC+(i+4)*paddingSize+j);
                        __m256d c6 = _mm256_load_pd(paddingC+(i+5)*paddingSize+j);
                        for(size_t k=kk;k<kk+48;k++)
                        {
                            __m256d m_b = _mm256_load_pd(paddingB+k*paddingSize+j);

                            __m256d a1 =  _mm256_set1_pd(paddingA[i*paddingSize+k]);//use array index to shift
                            __m256d a2 =  _mm256_set1_pd(paddingA[(i+1)*paddingSize+k]);
                            __m256d a3 =  _mm256_set1_pd(paddingA[(i+2)*paddingSize+k]);
                            __m256d a4 =  _mm256_set1_pd(paddingA[(i+3)*paddingSize+k]);//use array index to shift
                            __m256d a5 =  _mm256_set1_pd(paddingA[(i+4)*paddingSize+k]);
                            __m256d a6 =  _mm256_set1_pd(paddingA[(i+5)*paddingSize+k]);

                            c1 = _mm256_fmadd_pd(a1,m_b,c1);
                            c2 = _mm256_fmadd_pd(a2,m_b,c2);
                            c3 = _mm256_fmadd_pd(a3,m_b,c3);
                            c4 = _mm256_fmadd_pd(a4,m_b,c4);
                            c5 = _mm256_fmadd_pd(a5,m_b,c5);
                            c6 = _mm256_fmadd_pd(a6,m_b,c6);

                           
                        }
                        _mm256_store_pd(paddingC+(i)*paddingSize+j,c1);
                        _mm256_store_pd(paddingC+(i+1)*paddingSize+j,c2);
                        _mm256_store_pd(paddingC+(i+2)*paddingSize+j,c3);
                        _mm256_store_pd(paddingC+(i+3)*paddingSize+j,c4);
                        _mm256_store_pd(paddingC+(i+4)*paddingSize+j,c5);
                        _mm256_store_pd(paddingC+(i+5)*paddingSize+j,c6);
                    }
                }
            }
        }
    }   
    double* pc = paddingC;
    double* cRow = c;
    for (size_t i = 0; i < n; i++)
    {
        memcpy(cRow,pc,n*sizeof(double));
        cRow+=n;
        pc+=paddingSize;
    }

    aligned_free(paddingA);
    aligned_free(paddingB);
    aligned_free(paddingC);
    
}

void optimizedMM4(size_t n,double* a,double* b,double* c)
{   //size(length or wide..) after add padding will be divisible by 48(use 48 step to unroaling)
    //also can be divisible by 4 and 6, (6x4 micro kernal)  

   
    size_t paddingSize = ((n+47)/48)*48; 
    
    
    double* paddingA = (double*)aligned_malloc(64,paddingSize*paddingSize*sizeof(double));
    double* paddingB = (double*)aligned_malloc(64,paddingSize*paddingSize*sizeof(double));
    double* paddingC = (double*)aligned_malloc(64,paddingSize*paddingSize*sizeof(double));
    
    if (!paddingA||!paddingB||!paddingC) 
    {
        printf("Memory allocation failed!\n");
        return;
    }

    size_t size = paddingSize*paddingSize*sizeof(double);

    memset(paddingA,0,size);
    memset(paddingB,0,size);
    memset(paddingC,0,size); 

    double* pa = paddingA;
    double* aRow = a;
    double* pb = paddingB;
    double* bRow = b;

    for(size_t i=0;i<n;i++)
    {
        memcpy(pa,aRow,n*sizeof(double));
        pa+=paddingSize;
        aRow+=n;

        memcpy(pb,bRow,n*sizeof(double));
        pb+=paddingSize;
        bRow+=n;
    }


    for(size_t ii=0;ii<paddingSize;ii+=48)
    {
        for(size_t kk=0;kk<paddingSize;kk+=48)
        {
            for(size_t jj=0;jj<paddingSize;jj+=48)
            {
                for(size_t i=ii;i<ii+48;i++)
                {   
                    for(size_t k=kk;k<kk+48;k+=4)
                    {   
                       
                        
                        
                        // __m256d c1 = _mm256_load_pd(paddingC+(i)*paddingSize+j);
                        // __m256d c2 = _mm256_load_pd(paddingC+(i+1)*paddingSize+j);
                        // __m256d c3 = _mm256_load_pd(paddingC+(i+2)*paddingSize+j);
                        // __m256d c4 = _mm256_load_pd(paddingC+(i+3)*paddingSize+j);
                        // __m256d c5 = _mm256_load_pd(paddingC+(i+4)*paddingSize+j);
                        // __m256d c6 = _mm256_load_pd(paddingC+(i+5)*paddingSize+j);

                        //the naming number only represent the relative sequence, does not mean the actual row or col index
                        __m256d aRow1_Col1 = _mm256_set1_pd(paddingA[i*paddingSize+k]);
                        __m256d aRow1_Col2 = _mm256_set1_pd(paddingA[i*paddingSize+k+1]);
                        __m256d aRow1_Col3 = _mm256_set1_pd(paddingA[i*paddingSize+k+2]);
                        __m256d aRow1_Col4 = _mm256_set1_pd(paddingA[i*paddingSize+k+3]);
                        
                        for(size_t j=jj;j<jj+48;j+=4)
                        {
                           

                            //the naming number only represent the relative sequence, does not mean the actual row or col index
                            __m256d bRow1_Col1 = _mm256_load_pd(paddingB+k*paddingSize+j);
                            __m256d bRow2_Col1 = _mm256_load_pd(paddingB+(k+1)*paddingSize+j);
                            __m256d bRow3_Col1 = _mm256_load_pd(paddingB+(k+2)*paddingSize+j);
                            __m256d bRow4_Col1 = _mm256_load_pd(paddingB+(k+3)*paddingSize+j);
                           
                            __m256d cRow1_Col1 = _mm256_load_pd(paddingC+(i*paddingSize)+j);
                           
                           
                            cRow1_Col1 = _mm256_fmadd_pd(aRow1_Col2,bRow2_Col1,cRow1_Col1);
                            cRow1_Col1 = _mm256_fmadd_pd(aRow1_Col3,bRow3_Col1,cRow1_Col1);
                            cRow1_Col1 = _mm256_fmadd_pd(aRow1_Col4,bRow4_Col1,cRow1_Col1);
                            cRow1_Col1 = _mm256_fmadd_pd(aRow1_Col1,bRow1_Col1,cRow1_Col1);
                            
                           
                            _mm256_store_pd(paddingC+(i)*paddingSize+j,cRow1_Col1);

                        }
                        

                        
                    }
                }
            }
        }
    }   
    double* pc = paddingC;
    double* cRow = c;
    for (size_t i = 0; i < n; i++)
    {
        memcpy(cRow,pc,n*sizeof(double));
        cRow+=n;
        pc+=paddingSize;
    }

    aligned_free(paddingA);
    aligned_free(paddingB);
    aligned_free(paddingC);
    
}