main.cpp

//------------------------------------------------------------------------------------
// Matrix multiplication code adapted from an SME programming example provided
// by Arm Ltd.
//
// Assumptions :
//  - Input and Output matricies are in row-major
//  - nbr in matLeft (M): any
//  - nbc in matLeft, nbr in matRight (K): any K > 2
//  - nbc in matRight (N): any
//------------------------------------------------------------------------------------


#include <iostream>
#include <random>
#include <chrono>
#include <cstdint>
#include <stdlib.h>
#include <Accelerate/Accelerate.h>
#define ERROR_TOLERANCE 0.0002f

void matmul_ref(const uint64_t rows_l, const uint64_t cols_l,const uint64_t cols_r, const float * input_left,
    const float * input_right, float * output) {
    for (uint64_t x = 0; x < rows_l; ++x) {
      for (uint64_t y = 0; y < cols_r; ++y) {
        float acc = 0.0f;
        for (uint64_t z = 0; z < cols_l; ++z) {
          acc += input_left[(x * cols_l) + z] * input_right[(z * cols_r) + y];
        }
        output[(x * cols_r) + y] = acc;
      }
    }
}

void random_matrix( int m, int n, float *a)
{
  double drand48();

  //#pragma omp parallel for num_threads(num)
  for(int i =0 ; i < m * n; i ++)
  	a[i]= 2.0 * (float)drand48( ) - 1.0 + 0.0000001 * (i/m);

}


uint64_t sme_cntw() {
  uint64_t cnt;
  asm volatile("  smstart sm         \n" // smstart sm
               "  cntw %[res]\n"
               "  smstop sm          \n" // smstop sm
               : [res] "=r"(cnt)
               :
               : "p0", "p1", "p2", "p3", "p4", "p5", "p6", "p7", "p8", "p9",
                 "p10", "p11", "p12", "p13", "p14", "p15", "z0", "z1", "z2",
                 "z3", "z4", "z5", "z6", "z7", "z8", "z9", "z10", "z11", "z12",
                 "z13", "z14", "z15", "z16", "z17", "z18", "z19", "z20", "z21",
                 "z22", "z23", "z24", "z25", "z26", "z27", "z28", "z29", "z30",
                 "z31");
  //cnt = 16;
  return cnt;
}

extern "C"{
  extern void sgemm_direct_sme1_2VLx2VL(const uint64_t rows_l, const uint64_t cols_l,
  const uint64_t cols_r, const float * input_left,
  const float * input_right, float * output);

  extern void sgemm_direct_sme1_preprocess(const uint64_t rows, const uint64_t cols,
      const float * a, float * a_mod);

  extern void sgemm_direct_sme2_preprocess(const uint64_t rows, const uint64_t cols,
      const float * a, float * a_mod);

  extern void sgemm_direct_sme2_2VLx2VL(const uint64_t rows_l, const uint64_t cols_l,
    const uint64_t cols_r, const float * input_left,
    const float * input_right, float * output);
}


#define L1_CACHE_SIZE 131072   // 128 KB
#define LOAD_WIDTH    (16 * 2) // 每次加载 32 个 float（128B）
#define FLOAT_SIZE    4

const int cacheK = (L1_CACHE_SIZE / FLOAT_SIZE / LOAD_WIDTH) * 8;


float cacheN[1][LOAD_WIDTH];
float cacheM[cacheK][LOAD_WIDTH];
float *cacheC;

void init_C() {
    for (size_t i = 0; i < LOAD_WIDTH; i++) {
        cacheN[0][i] = 2.0 * (float)drand48( ) - 1.0 + 0.0000001 * (i/LOAD_WIDTH);
    }

    for (size_t i = 0; i < cacheK; i++) {
        for (size_t j = 0; j < LOAD_WIDTH; j++) {
            cacheM[i][j] = 2.0 * (float)drand48( ) - 1.0 + 0.0000001 * (i/LOAD_WIDTH);
        }
    }

    posix_memalign( (void**) &cacheC, 128, cacheK * LOAD_WIDTH * sizeof(float));
    for (size_t i = 0; i < cacheK * LOAD_WIDTH; i++) {
        cacheC[i] = (float)i/16.52525;
    }
}

void func1() {
    float *ptr  = &cacheN[0][0];
    asm volatile("smstart" :::);
    for (int i = 0; i < cacheK; i++) {
        asm volatile(
            "ld1w {z0.s}, p0/z, [%[src]]\n"
            :
            : [src] "r"(ptr)
            : "p0", "z0"
        );
    }
    asm volatile("smstop" :::);
}



void func2() {
    asm volatile("smstart" :::);
    for (size_t i = 0; i < cacheK; i++) {
        float *ptr = &cacheM[i][0];
        asm volatile(
            "ld1w {z0.s}, p0/z, [%[src]]\n"
            :
            : [src] "r"(ptr)
            : "p0","z0"
        );
    }
    asm volatile("smstop" :::);
}

void func3() {
    asm volatile("smstart" :::);
    for (size_t i = 0; i < cacheK; i++) {
        float *ptr = &cacheC[i * LOAD_WIDTH]; // 大跨度访问
        asm volatile(
            "ld1w {z0.s}, p0/z, [%[src]]\n"
            :
            : [src] "r"(ptr)
            : "p0","z0"
        );
    }
    asm volatile("smstop" :::);
}



void benchmark(const char *name, void (*func)()) {
    std::chrono::steady_clock::time_point l_start = std::chrono::steady_clock::now();
    func();
    std::chrono::steady_clock::time_point l_end = std::chrono::steady_clock::now();

    double elapsed = std::chrono::duration_cast< std::chrono::duration<double> >( l_end - l_start ).count();
    double total_bytes = cacheK * LOAD_WIDTH * FLOAT_SIZE;
    double bandwidth = total_bytes / elapsed / 1e9; // GB/s

    printf("%s:\n  Time: %.6f sec\n  Bandwidth: %.2f GB/s\n\n", name, elapsed, bandwidth);
}



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

  init_C();  // 分配并初始化大数组

  benchmark("func1 (same address load)", func1);
  benchmark("func2 (linear load)",       func2);
  benchmark("func3 (stride load)",       func3);

  free(cacheC);



  int M, N, K, M_mod,iterations,num_threads;

  // /* Size parameters */
//   if (argc == 5) {
//     iterations = strtoul(argv[1], NULL, 0);
//     M = strtoul(argv[2], NULL, 0);
//     N = strtoul(argv[3], NULL, 0);
//     K = strtoul(argv[4], NULL, 0);
//     num_threads = 1;
//   } else {
//     /* Default: 128x128x128 */
//     M = 128;
//     N = 128;
//     K = 128;
//     iterations = 1;
//     num_threads = 1;
//   }

  M = 32;
  N = 32;
  K = 8192;
  iterations = 1000;

  if (!(K > 2)) {
    printf("Invalid Matrix Dimensions. Please ensure that :\n\t- K dimension "
           "is Greater-Than 2 \n\nRuntime args ==> ./BINARY_NAME iterations M "
           "K N\n");
    exit(1);
  }

  printf("Left Matrix Dimensions = %d x %d\nRight Matrix Dimensions = %d x "
         "%d\nSME Loop iterations = %d\n\n",
         M, K, K, N, iterations);

  /* Calculate M of transformed matLeft.  */
  uint64_t vl_elms = sme_cntw();
  printf("%lu\n", vl_elms);

  M_mod = ceil((double)M / (double)vl_elms) * vl_elms; // ceil(M/cntw)*cntw
  printf("M_mod = %d\n", M_mod);

  pthread_set_qos_class_self_np( QOS_CLASS_USER_INTERACTIVE, 0 );


  float *matLeft = NULL;
  posix_memalign( (void**) &matLeft, 128, M * K * sizeof(float) );
  float *matLeft_mod = NULL;
  posix_memalign( (void**) &matLeft_mod, 128, M_mod * K * sizeof(float) );
  float *matRight = NULL;
  posix_memalign( (void**) &matRight, 128,K * N * sizeof(float) );
  float *matResult_ref = NULL;
  posix_memalign( (void**) &matResult_ref, 128, M * N * sizeof(float) );
  float *matResult_opt = NULL;
  posix_memalign( (void**) &matResult_opt, 128, M * N * sizeof(float) );

  random_matrix(M, K, matLeft);
  random_matrix(K, N, matRight);


  // for(int i=0;i<M*K; i++){
  //   printf(" %.2f ",matLeft[i]);
  // }
  // printf("\n");
  // printf("\n");
  // for(int j=0;j< M;j++){
  //     for(int z=0;z< K;z++){
  //         printf(" %.2f ",matLeft[z + j * K]);
  //     }
  //     printf("\n");
  // }

  // printf("\n");
  // printf("\n");

  // sgemm_direct_sme1_preprocess(M, K, matLeft, matLeft_mod);

  // for(int i=0;i<M*K; i++){
  //   printf(" %.2f ",matLeft_mod[i]);
  // }
  // printf("\n");
  // printf("\n");


  // for(int j=0;j< M;j++){
  //     for(int z=0;z< K;z++){
  //         printf(" %.2f ",matLeft_mod[z + j * K]);
  //     }
  //     printf("\n");
  // }

  // set up dispatch
  // dispatch_queue_attr_t l_attr = dispatch_queue_attr_make_with_qos_class( DISPATCH_QUEUE_CONCURRENT,
  //   QOS_CLASS_USER_INTERACTIVE,
  //   0 );

  // dispatch_queue_t l_queue = dispatch_queue_create( "bench_queue", l_attr );

  // dispatch_group_t l_group = dispatch_group_create();

  double ops = (double )M *  N  * K * 1.0e-09 * 2;

  //Calculate matrix multiply naively
  
  printf("Calculating Reference Matrix Multiply...\n");
  std::chrono::steady_clock::time_point l_start = std::chrono::steady_clock::now();
  // for( int l_td = 0; l_td < num_threads; l_td++ ) {
  //   dispatch_group_async( l_group,
  //                         l_queue,
  //                         ^{  for (int i = 0; i < iterations; i++) {
  //                             matmul_ref(M, K, N, matLeft, matRight, matResult_ref); }
  //                           } );
  // }
  // dispatch_group_wait( l_group, DISPATCH_TIME_FOREVER );
  for (int i = 0; i < iterations; i++) {
    //matmul_ref(M, K, N, matLeft, matRight, matResult_ref);
    cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans, M, N, K, 1.0, matLeft, K, matRight, N, 0.0, matResult_ref, N);
  }
  std::chrono::steady_clock::time_point l_end = std::chrono::steady_clock::now();
  double cost = std::chrono::duration_cast< std::chrono::duration<double> >( l_end - l_start ).count() / iterations;
  printf("SGEMM:  M= %d N=%d K=%d cost=%lf Gflops = %lf \n", M, N, K, cost,  ops/cost);
  printf("Done\n");

  // for (int i = 0; i < 3; i++) {
  //   sgemm_direct_sme1_preprocess(M, K, matLeft, matLeft_mod);
  //   sgemm_direct_sme1_2VLx2VL(M, K, N, matLeft_mod, matRight, matResult_opt);
  // }
  
  
  // Looping multiple times to isolate performance to this function call
  printf("Calculating %d iterations of SME Matrix Multiply...\n", iterations);
  l_start = std::chrono::steady_clock::now();
  //#pragma omp parallel for num_threads(10)
  // for( int l_td = 0; l_td < num_threads; l_td++ ) {
  //   dispatch_group_async( l_group,
  //                         l_queue,
  //                         ^{  
  //                             for (int i = 0; i < iterations; i++) {
  //                             sgemm_direct_sme1_preprocess(M, K, matLeft, matLeft_mod);
  //                             sgemm_direct_sme1_2VLx2VL(M, K, N, matLeft_mod, matRight, matResult_opt);
  //                           } 
  //                         } );
  // }
  // dispatch_group_wait( l_group, DISPATCH_TIME_FOREVER );
  for (int i = 0; i < iterations; i++) {
      sgemm_direct_sme2_preprocess(M, K, matLeft, matLeft_mod);
      sgemm_direct_sme2_2VLx2VL(M, K, N, matLeft_mod, matRight, matResult_opt);
  }
  l_end = std::chrono::steady_clock::now();
  cost = std::chrono::duration_cast< std::chrono::duration<double> >( l_end - l_start ).count() / iterations;
  
  printf("SGEMMxsme:  M= %d N=%d K=%d cost=%lf ops=%f Gflops = %lf \n", M, N, K, cost, ops, ops/cost);
  printf("Done\n\n");

  //clang -march=armv9-a+sme+sme2 -O0 -c matmul_opt.s preprocess.s

  //clang -O3 -c main.c -o main.o
  //clang -march=armv9-a+sme+sme2 -O0 -c matmul_opt.s preprocess.s
  //clang main.o matmul_opt.o preprocess.o -o matmul
  

  printf("ERROR TOLERANCE = %.4f%%\n", (float)ERROR_TOLERANCE);
  int error = 0;
  for (uint64_t i = 0; i < M * N; i++) {
    if (fabsf(matResult_ref[i] - matResult_opt[i]) > fabsf((float)ERROR_TOLERANCE * matResult_ref[i])) {
        error = 1;
        printf("%lu, %f, %f\n", i, matResult_ref[i],matResult_opt[i]);
      }
  }

  // for (uint64_t i = 0; i < M*N; i++) {
  //   printf("%lu, %f, %f\n", i, matResult_ref[i],matResult_opt[i]);
  // }

  // if (error)
  // {
  //   printf("FAILED\n");
  //   printf("100 , %f, %f\n", matResult_ref[100],matResult_opt[100]);
  // }
    
  // else
  // {
  //   printf("PASS\n");
  //   printf("100 , %f, %f\n", matResult_ref[100],matResult_opt[100]);
  // }
    
  free(matLeft);
  free(matLeft_mod);
  free(matRight);
  free(matResult_ref);
  free(matResult_opt);

  return 0;
}