#include <stdio.h>
#include <stdlib.h>
#include <time.h> 
#include <arm_neon.h>

void matmul_optimized(float** A, float** B, float** C, int n) {
  int i,j,k;
  for (i = 0; i < n; i++) {
    for (j = 0; j < n; j++) {
      float32x4_t vecC = vdupq_n_f32(0.0f); 
      for (k = 0; k < n; k += 4) {
        float32x4_t vecA = vld1q_f32(&A[i][k]);
        float32x4_t vecB = vld1q_f32(&B[k][j]);
        vecC = vmlaq_f32(vecC, vecA, vecB); 
      }
      C[i][j] = vgetq_lane_f32(vecC, 0) + vgetq_lane_f32(vecC, 1) +
           vgetq_lane_f32(vecC, 2) + vgetq_lane_f32(vecC, 3);
    }
  }
}
int main() {
  int i,j;
  int n = 1024; 
  float** A = (float**)malloc(n * sizeof(float*));
  float** B = (float**)malloc(n * sizeof(float*));
  float** C = (float**)malloc(n * sizeof(float*));
  for (i = 0; i < n; i++) {
    A[i] = (float*)malloc(n * sizeof(float));
    B[i] = (float*)malloc(n * sizeof(float));
    C[i] = (float*)calloc(n, sizeof(float)); 
  }
  srand(time(NULL));
  for (i = 0; i < n; i++) {
    for (j = 0; j < n; j++) {
      A[i][j] = (float)(rand() % 100) / 100.0f;
      B[i][j] = (float)(rand() % 100) / 100.0f;
    }
  }

  clock_t start = clock();
  matmul_optimized(A, B, C, n);
  clock_t end = clock();
  double time_taken = (double)(end - start) / CLOCKS_PER_SEC;
  printf("Optimized matrix multiplication time: %f seconds\n", time_taken);
  for (int i = 0; i < n; i++) {
    free(A[i]);
    free(B[i]);
    free(C[i]);
  }
  free(A);
  free(B);
  free(C);

  return 0;
}