nn_f32.h

Go to the documentation of this file.

 
#ifndef __NN_F32_H
#define __NN_F32_H
 
#include <float.h>
#include <stdint.h>
 
#ifdef CONFIG_BACKEND_RISCV_V
  #include "riscv_vector.h"
#endif
 
typedef struct {
  float data;
} Tensor0D_F32;
 
 
typedef struct {
  size_t shape[1];
  float *data;
} Tensor1D_F32;
 
 
typedef struct {
  size_t shape[2];
  float *data;
} Tensor2D_F32;
 
typedef struct {
  size_t shape[3];
  float *data;
} Tensor3D_F32;
 
typedef struct {
  size_t shape[4];
  float *data;
} Tensor4D_F32;
 
 
static inline uint8_t nn_equal_f32(float golden, float actual, float rel_err) {
  return (fabs(actual - golden) < rel_err) || (fabs((actual - golden) / actual) < rel_err);
}
 
 
/* ======================================================================================================== */
/*                                           Tensor Creation                                                */
/* ======================================================================================================== */
 
Tensor0D_F32 *nn_tensor0d_f32(float data) {
  Tensor0D_F32 *tensor = (Tensor0D_F32 *)malloc(sizeof(Tensor0D_F32));
  tensor->data = data;
  return tensor;
}
 
Tensor1D_F32 *nn_tensor1d_f32(size_t shape[1], const float *data) {
  Tensor1D_F32 *tensor = (Tensor1D_F32 *)malloc(sizeof(Tensor1D_F32));
  tensor->shape[0] = shape[0];
 
  size_t n_bytes = shape[0] * sizeof(float);
  tensor->data = (float *)malloc(n_bytes);
  if (data != NULL) {
    memcpy(tensor->data, data, n_bytes);
  }
  return tensor;
}
 
Tensor2D_F32 *nn_tensor2d_f32(size_t shape[2], const float *data) {
  Tensor2D_F32 *tensor = (Tensor2D_F32 *)malloc(sizeof(Tensor2D_F32));
  tensor->shape[0] = shape[0];
  tensor->shape[1] = shape[1];
 
  size_t n_bytes = shape[0] * shape[1] * sizeof(float);
  tensor->data = (float *)malloc(n_bytes);
  if (data != NULL) {
    memcpy(tensor->data, data, n_bytes);
  }
  return tensor;
}
 
Tensor3D_F32 *nn_tensor3d_f32(size_t shape[3], const float *data) {
  Tensor3D_F32 *tensor = (Tensor3D_F32 *)malloc(sizeof(Tensor3D_F32));
  tensor->shape[0] = shape[0];
  tensor->shape[1] = shape[1];
  tensor->shape[2] = shape[2];
 
  size_t n_bytes = shape[0] * shape[1] * shape[2] * sizeof(float);
  tensor->data = (float *)malloc(n_bytes);
  if (data != NULL) {
    memcpy(tensor->data, data, n_bytes);
  }
  return tensor;
}
 
Tensor4D_F32 *nn_tensor4d_f32(size_t shape[4], const float *data) {
  Tensor4D_F32 *tensor = (Tensor4D_F32 *)malloc(sizeof(Tensor4D_F32));
  tensor->shape[0] = shape[0];
  tensor->shape[1] = shape[1];
  tensor->shape[2] = shape[2];
  tensor->shape[3] = shape[3];
 
  size_t n_bytes = shape[0] * shape[1] * shape[2] * shape[3] * sizeof(float);
  tensor->data = (float *)malloc(n_bytes);
  if (data != NULL) {
    memcpy(tensor->data, data, n_bytes);
  }
  return tensor;
}
 
Tensor1D_F32 *nn_as_tensor1d_f32(size_t shape[1], float *data) {
  Tensor1D_F32 *tensor = (Tensor1D_F32 *)malloc(sizeof(Tensor1D_F32));
  tensor->shape[0] = shape[0];
  tensor->data = data;
  return tensor;
}
 
Tensor2D_F32 *nn_as_tensor2d_f32(size_t shape[2], float *data) {
  Tensor2D_F32 *tensor = (Tensor2D_F32 *)malloc(sizeof(Tensor2D_F32));
  tensor->shape[0] = shape[0];
  tensor->shape[1] = shape[1];
  tensor->data = data;
  return tensor;
}
 
Tensor3D_F32 *nn_as_tensor3d_f32(size_t shape[3], float *data) {
  Tensor3D_F32 *tensor = (Tensor3D_F32 *)malloc(sizeof(Tensor3D_F32));
  tensor->shape[0] = shape[0];
  tensor->shape[1] = shape[1];
  tensor->shape[2] = shape[2];
  tensor->data = data;
  return tensor;
}
 
Tensor4D_F32 *nn_as_tensor4d_f32(size_t shape[4], float *data) {
  Tensor4D_F32 *tensor = (Tensor4D_F32 *)malloc(sizeof(Tensor4D_F32));
  tensor->shape[0] = shape[0];
  tensor->shape[1] = shape[1];
  tensor->shape[2] = shape[2];
  tensor->shape[3] = shape[3];
  tensor->data = data;
  return tensor;
}
 
 
Tensor0D_F32 *nn_zeros0d_f32() {
  Tensor0D_F32 *tensor = nn_tensor0d_f32(0);
  return tensor;
}
 
Tensor1D_F32 *nn_zeros1d_f32(size_t shape[1]) {
  Tensor1D_F32 *tensor = nn_tensor1d_f32(shape, NULL);
  size_t n = shape[0];
  for (size_t i = 0; i < n; i += 1) {
    tensor->data[i] = 0;
  }
  return tensor;
}
 
Tensor2D_F32 *nn_zeros2d_f32(size_t shape[2]) {
  Tensor2D_F32 *tensor = nn_tensor2d_f32(shape, NULL);
  size_t n = shape[0] * shape[1];
  for (size_t i = 0; i < n; i += 1) {
    tensor->data[i] = 0;
  }
  return tensor;
}
 
Tensor3D_F32 *nn_zeros3d_f32(size_t shape[3]) {
  Tensor3D_F32 *tensor = nn_tensor3d_f32(shape, NULL);
  size_t n = shape[0] * shape[1] * shape[2];
  for (size_t i = 0; i < n; i += 1) {
    tensor->data[i] = 0;
  }
  return tensor;
}
 
Tensor4D_F32 *nn_zeros4d_f32(size_t shape[4]) {
  Tensor4D_F32 *tensor = nn_tensor4d_f32(shape, NULL);
  size_t n = shape[0] * shape[1] * shape[2] * shape[3];
  for (size_t i = 0; i < n; i += 1) {
    tensor->data[i] = 0;
  }
  return tensor;
}
 
Tensor0D_F32 *nn_ones0d_f32() {
  Tensor0D_F32 *tensor = nn_tensor0d_f32(1);
  return tensor;
}
 
Tensor1D_F32 *nn_ones1d_f32(size_t shape[1]) {
  Tensor1D_F32 *tensor = nn_tensor1d_f32(shape, NULL);
  size_t n = shape[0];
  for (size_t i = 0; i < n; i += 1) {
    tensor->data[i] = 1;
  }
  return tensor;
}
 
Tensor2D_F32 *nn_ones2d_f32(size_t shape[2]) {
  Tensor2D_F32 *tensor = nn_tensor2d_f32(shape, NULL);
  size_t n = shape[0] * shape[1];
  for (size_t i = 0; i < n; i += 1) {
    tensor->data[i] = 1;
  }
  return tensor;
}
 
Tensor0D_F32 *nn_full0d_f32(float data) {
  Tensor0D_F32 *tensor = nn_tensor0d_f32(data);
  return tensor;
}
 
Tensor1D_F32 *nn_full1d_f32(size_t shape[1], float data) {
  Tensor1D_F32 *tensor = nn_tensor1d_f32(shape, NULL);
  size_t n = shape[0];
  for (size_t i = 0; i < n; i += 1) {
    tensor->data[i] = data;
  }
  return tensor;
}
 
Tensor2D_F32 *nn_full2d_f32(size_t shape[2], float data) {
  Tensor2D_F32 *tensor = nn_tensor2d_f32(shape, NULL);
  size_t n = shape[0] * shape[1];
  for (size_t i = 0; i < n; i += 1) {
    tensor->data[i] = data;
  }
  return tensor;
}
 
Tensor0D_F32 *nn_rand0d_f32() {
  Tensor0D_F32 *tensor = nn_tensor0d_f32(rand());
  return tensor;
}
 
Tensor1D_F32 *nn_rand1d_f32(size_t shape[1]) {
  Tensor1D_F32 *tensor = nn_tensor1d_f32(shape, NULL);
  size_t n = shape[0];
  for (size_t i = 0; i < n; i += 1) {
    tensor->data[i] = rand();
  }
  return tensor;
}
 
Tensor2D_F32 *nn_rand2d_f32(size_t shape[2]) {
  Tensor2D_F32 *tensor = nn_tensor2d_f32(shape, NULL);
  size_t n = shape[0] * shape[1];
  for (size_t i = 0; i < n; i += 1) {
    tensor->data[i] = rand();
  }
  return tensor;
}
 
 
/* ======================================================================================================== */
/*                                           Tensor Prints                                                  */
/* ======================================================================================================== */
 
void nn_print_tensor1d_f32(const Tensor1D_F32 *tensor) {
  printf("[");
  for (size_t i=0; i<tensor->shape[0]; i+=1) {
    nn_print_f32(*((float *)tensor->data + i), 3);
    if (i < tensor->shape[0]-1) {
      printf(" ");
    }
  }
  printf("]\n");
}
 
void nn_print_tensor2d_f32(const Tensor2D_F32 *tensor) {
  printf("[");
  for (size_t i=0; i<tensor->shape[0]; i+=1) {
    if (i == 0) {
      printf("[");
    }
    else {
      printf(" [");
    }
    for (size_t j=0; j<tensor->shape[1]; j+=1) {
      nn_print_f32(*((float *)tensor->data + i*tensor->shape[1] + j), 3);
      if (j < tensor->shape[1]-1) {
        printf(" ");
      }
    }
    printf(" ]");
    if (i < tensor->shape[0]-1) {
      printf("\n");
    }
  }
  printf("]\n");
}
 
void nn_print_tensor3d_f32(const Tensor3D_F32 *tensor) {
  printf("[");
  for (size_t i=0; i<tensor->shape[0]; i+=1) {
    if (i == 0) {
      printf("[");
    }
    else {
      printf("\n [");
    }
    for (size_t j=0; j<tensor->shape[1]; j+=1) {
      if (j == 0) {
        printf("[");
      }
      else {
        printf("  [");
      }
      for (size_t k=0; k<tensor->shape[2]; k+=1) {
        nn_print_f32(*((float *)tensor->data + i*tensor->shape[1]*tensor->shape[2] + j*tensor->shape[2] + k), 3);
        if (k < tensor->shape[2]-1) {
          printf(" ");
        }
      }
      printf(" ]");
    }
    printf("]");
    if (i < tensor->shape[0]-1) {
      printf("\n");
    }
  }
  printf("]\n");
}
 
void nn_print_tensor4d_f32(const Tensor4D_F32 *tensor) {
  printf("[");
  for (size_t i=0; i<tensor->shape[0]; i+=1) {
    if (i == 0) {
      printf("[");
    }
    else {
      printf("\n [");
    }
    for (size_t j=0; j<tensor->shape[1]; j+=1) {
      if (j == 0) {
        printf("[");
      }
      else {
        printf("\n  [");
      }
      for (size_t k=0; k<tensor->shape[2]; k+=1) {
        if (k == 0) {
          printf("[");
        }
        else {
          printf("   [");
        }
        for (size_t l=0; l<tensor->shape[3]; l+=1) {
          nn_print_f32(*((float *)tensor->data + i*tensor->shape[1]*tensor->shape[2]*tensor->shape[3] + j*tensor->shape[2]*tensor->shape[3] + k*tensor->shape[3] + l), 3);
          if (l < tensor->shape[3]-1) {
            printf(" ");
          }
        }
        printf(" ]");
        if (k < tensor->shape[2]-1) {
          printf("\n");
        }
      }
      printf("]");
      if (j < tensor->shape[1]-1) {
        printf("\n");
      }
    }
    printf("]");
    if (i < tensor->shape[0]-1) {
      printf("\n");
    }
  }
  printf("]\n");
}
 
 
/* ======================================================================================================== */
/*                                           Comparision                                                    */
/* ======================================================================================================== */
uint8_t nn_equals0d_f32(const Tensor0D_F32 *a, const Tensor0D_F32 *b, float rel_err) {
  return nn_equal_f32(a->data, b->data, rel_err);
}
 
uint8_t nn_equals1d_f32(const Tensor1D_F32 *a, const Tensor1D_F32 *b, float rel_err) {
  nn_assert(a->shape[0] == b->shape[0], "Cannot compare tensors of different shapes");
 
  size_t n = a->shape[0];
  for (size_t i = 0; i < n; i += 1) {
    if (!nn_equal_f32(a->data[i], b->data[i], rel_err)) {
      return 0;
    }
  }
  return 1;
}
 
uint8_t nn_equals2d_f32(const Tensor2D_F32 *a, const Tensor2D_F32 *b, float rel_err) {
  nn_assert(a->shape[0] == b->shape[0] && a->shape[1] == b->shape[1], "Cannot compare tensors of different shapes");
 
  size_t n = a->shape[0] * a->shape[1];
  for (size_t i = 0; i < n; i += 1) {
    if (!nn_equal_f32(a->data[i], b->data[i], rel_err)) {
      return 0;
    }
  }
  return 1;
}
 
uint8_t nn_equals3d_f32(const Tensor3D_F32 *a, const Tensor3D_F32 *b, float rel_err) {
  nn_assert(a->shape[0] == b->shape[0] && a->shape[1] == b->shape[1] && a->shape[2] == b->shape[2], "Cannot compare tensors of different shapes");
 
  size_t n = a->shape[0] * a->shape[1] * a->shape[2];
  for (size_t i = 0; i < n; i += 1) {
    if (!nn_equal_f32(a->data[i], b->data[i], rel_err)) {
      return 0;
    }
  }
  return 1;
}
 
uint8_t nn_equals4d_f32(const Tensor4D_F32 *a, const Tensor4D_F32 *b, float rel_err) {
  nn_assert(a->shape[0] == b->shape[0] && a->shape[1] == b->shape[1] && a->shape[2] == b->shape[2] && a->shape[3] == b->shape[3], "Cannot compare tensors of different shapes");
 
  size_t n = a->shape[0] * a->shape[1] * a->shape[2] * a->shape[3];
  for (size_t i = 0; i < n; i += 1) {
    if (!nn_equal_f32(a->data[i], b->data[i], rel_err)) {
      return 0;
    }
  }
  return 1;
}
 
 
/* ======================================================================================================== */
/*                                           Unary                                                          */
/* ======================================================================================================== */
 
void nn_max1d_f32(Tensor0D_F32 *y, const Tensor1D_F32 *x) {
  size_t n = x->shape[0];
  float *x_data = x->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    vfloat32m1_t vec_max = __riscv_vfmv_s_f_f32m1(-FLT_MAX, 1);
 
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x = __riscv_vle32_v_f32m1(x_data, vl);
      vec_max = __riscv_vfredmax_vs_f32m1_f32m1(vec_x, vec_max, vl);
      x_data += vl;
      n -= vl;
    }
    y->data = __riscv_vfmv_f_s_f32m1_f32(vec_max);
  #else  /* scalar implementation */
    y->data = -FLT_MAX;
    for (size_t i = 0; i < n; i += 1) {
      float val = x_data[i];
      y->data = val > y->data ? val : y->data;
    }
  #endif
}
 
void nn_max2d_f32(Tensor0D_F32 *y, const Tensor2D_F32 *x) {
  size_t n = x->shape[0] * x->shape[1];
  float *x_data = x->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    vfloat32m1_t vec_max = __riscv_vfmv_s_f_f32m1(-FLT_MAX, 1);
 
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x = __riscv_vle32_v_f32m1(x_data, vl);
      vec_max = __riscv_vfredmax_vs_f32m1_f32m1(vec_x, vec_max, vl);
      x_data += vl;
      n -= vl;
    }
    y->data = __riscv_vfmv_f_s_f32m1_f32(vec_max);
  #else  /* scalar implementation */
    y->data = -FLT_MAX;
    for (size_t i = 0; i < n; i += 1) {
      float val = x_data[i];
      y->data = val > y->data ? val : y->data;
    }
  #endif
}
 
void nn_min1d_f32(Tensor0D_F32 *y, const Tensor1D_F32 *x) {
  size_t n = x->shape[0];
  float *x_data = x->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    vfloat32m1_t vec_min = __riscv_vfmv_s_f_f32m1(FLT_MAX, 1);
 
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x = __riscv_vle32_v_f32m1(x_data, vl);
      vec_min = __riscv_vfredmin_vs_f32m1_f32m1(vec_x, vec_min, vl);
      x_data += vl;
      n -= vl;
    }
    y->data = __riscv_vfmv_f_s_f32m1_f32(vec_min);
  #else  /* scalar implementation */
    y->data = FLT_MAX;
    for (size_t i = 0; i < n; i += 1) {
      float val = x_data[i];
      y->data = val < y->data ? val : y->data;
    }
  #endif
}
 
void nn_min2d_f32(Tensor0D_F32 *y, const Tensor2D_F32 *x) {
  size_t n = x->shape[0] * x->shape[1];
  float *x_data = x->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    vfloat32m1_t vec_min = __riscv_vfmv_s_f_f32m1(FLT_MAX, 1);
 
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x = __riscv_vle32_v_f32m1(x_data, vl);
      vec_min = __riscv_vfredmin_vs_f32m1_f32m1(vec_x, vec_min, vl);
      x_data += vl;
      n -= vl;
    }
    y->data = __riscv_vfmv_f_s_f32m1_f32(vec_min);
  #else  /* scalar implementation */
    y->data = FLT_MAX;
    for (size_t i = 0; i < n; i += 1) {
      float val = x_data[i];
      y->data = val < y->data ? val : y->data;
    }
  #endif
}
 
/* ======================================================================================================== */
/*                                           Addition                                                       */
/* ======================================================================================================== */
void nn_add1d_f32(Tensor1D_F32 *y, const Tensor1D_F32 *x1, const Tensor1D_F32 *x2) {
  nn_assert(x1->shape[0] == x2->shape[0], "Cannot add tensors of different shapes");
  nn_assert(y->shape[0] == x1->shape[0], "Cannot add tensors of different shapes");
 
  size_t n = y->shape[0];
  float *x1_data = x1->data;
  float *x2_data = x2->data;
  float *y_data = y->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x1 = __riscv_vle32_v_f32m1(x1_data, vl);
      vfloat32m1_t vec_x2 = __riscv_vle32_v_f32m1(x2_data, vl);
      vfloat32m1_t vec_y = __riscv_vfadd_vv_f32m1(vec_x1, vec_x2, vl);
      __riscv_vse32_v_f32m1(y_data, vec_y, vl);
      x1_data += vl;
      x2_data += vl;
      y_data += vl;
      n -= vl;
    }
  #else  /* scalar implementation */
    for (size_t i = 0; i < n; i += 1) {
      y_data[i] = x1_data[i] + x2_data[i];
    }
  #endif
}
 
void nn_add2d_f32(Tensor2D_F32 *y, const Tensor2D_F32 *x1, const Tensor2D_F32 *x2) {
  nn_assert(x1->shape[0] == x2->shape[0] && x1->shape[1] == x2->shape[1], "Cannot add tensors of different shapes");
  nn_assert(y->shape[0] == x1->shape[0] && y->shape[1] == x1->shape[1], "Cannot add tensors of different shapes");
 
  size_t n = y->shape[0] * y->shape[1];
  float *x1_data = x1->data;
  float *x2_data = x2->data;
  float *y_data = y->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x1 = __riscv_vle32_v_f32m1(x1_data, vl);
      vfloat32m1_t vec_x2 = __riscv_vle32_v_f32m1(x2_data, vl);
      vfloat32m1_t vec_y = __riscv_vfadd_vv_f32m1(vec_x1, vec_x2, vl);
      __riscv_vse32_v_f32m1(y_data, vec_y, vl);
      x1_data += vl;
      x2_data += vl;
      y_data += vl;
      n -= vl;
    }
  #else  /* scalar implementation */
    for (size_t i = 0; i < n; i += 1) {
      y_data[i] = x1_data[i] + x2_data[i];
    }
  #endif
}
 
void nn_addscalar1d_f32(Tensor1D_F32 *y, const Tensor1D_F32 *x, float scalar) {
  nn_assert(y->shape[0] == x->shape[0], "Cannot add tensors of different shapes");
 
  size_t n = y->shape[0];
  float *x_data = x->data;
  float *y_data = y->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x = __riscv_vle32_v_f32m1(x_data, vl);
      vfloat32m1_t vec_y = __riscv_vfadd_vf_f32m1(vec_x, scalar, vl);
      __riscv_vse32_v_f32m1(y_data, vec_y, vl);
      x_data += vl;
      y_data += vl;
      n -= vl;
    }
  #else  /* scalar implementation */
    for (size_t i = 0; i < n; i += 1) {
      y_data[i] = x_data[i] + scalar;
    }
  #endif
}
 
void nn_addscalar2d_f32(Tensor2D_F32 *y, const Tensor2D_F32 *x, float scalar) {
  nn_assert(y->shape[0] == x->shape[0] && y->shape[1] == x->shape[1], "Cannot add tensors of different shapes");
 
  size_t n = y->shape[0] * y->shape[1];
  float *x_data = x->data;
  float *y_data = y->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x = __riscv_vle32_v_f32m1(x_data, vl);
      vfloat32m1_t vec_y = __riscv_vfadd_vf_f32m1(vec_x, scalar, vl);
      __riscv_vse32_v_f32m1(y_data, vec_y, vl);
      x_data += vl;
      y_data += vl;
      n -= vl;
    }
  #else  /* scalar implementation */
    for (size_t i = 0; i < n; i += 1) {
      y_data[i] = x_data[i] + scalar;
    }
  #endif
}
 
/* ======================================================================================================== */
/*                                           Multiplication                                                 */
/* ======================================================================================================== */
 
void nn_mul1d_f32(Tensor1D_F32 *y, const Tensor1D_F32 *x1, const Tensor1D_F32 *x2) {
  nn_assert(x1->shape[0] == x2->shape[0], "Cannot add tensors of different shapes");
  nn_assert(y->shape[0] == x1->shape[0], "Cannot add tensors of different shapes");
 
  size_t n = y->shape[0];
  float *x1_data = x1->data;
  float *x2_data = x2->data;
  float *y_data = y->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x1 = __riscv_vle32_v_f32m1(x1_data, vl);
      vfloat32m1_t vec_x2 = __riscv_vle32_v_f32m1(x2_data, vl);
      vfloat32m1_t vec_y = __riscv_vfmul_vv_f32m1(vec_x1, vec_x2, vl);
      __riscv_vse32_v_f32m1(y_data, vec_y, vl);
      x1_data += vl;
      x2_data += vl;
      y_data += vl;
      n -= vl;
    }
  #else  /* scalar implementation */
    for (size_t i = 0; i < n; i += 1) {
      y_data[i] = x1_data[i] * x2_data[i];
    }
  #endif
}
 
void nn_mul2d_f32(Tensor2D_F32 *y, const Tensor2D_F32 *x1, const Tensor2D_F32 *x2) {
  nn_assert(x1->shape[0] == x2->shape[0] && x1->shape[1] == x2->shape[1], "Cannot add tensors of different shapes");
  nn_assert(y->shape[0] == x1->shape[0] && y->shape[1] == x1->shape[1], "Cannot add tensors of different shapes");
 
  size_t n = y->shape[0] * y->shape[1];
  float *x1_data = x1->data;
  float *x2_data = x2->data;
  float *y_data = y->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x1 = __riscv_vle32_v_f32m1(x1_data, vl);
      vfloat32m1_t vec_x2 = __riscv_vle32_v_f32m1(x2_data, vl);
      vfloat32m1_t vec_y = __riscv_vfmul_vv_f32m1(vec_x1, vec_x2, vl);
      __riscv_vse32_v_f32m1(y_data, vec_y, vl);
      x1_data += vl;
      x2_data += vl;
      y_data += vl;
      n -= vl;
    }
  #else  /* scalar implementation */
    for (size_t i = 0; i < n; i += 1) {
      y_data[i] = x1_data[i] * x2_data[i];
    }
  #endif
}
 
void nn_mulscalar1d_f32(Tensor1D_F32 *y, const Tensor1D_F32 *x, float scalar) {
  nn_assert(y->shape[0] == x->shape[0], "Cannot add tensors of different shapes");
 
  size_t n = y->shape[0];
  float *x_data = x->data;
  float *y_data = y->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x = __riscv_vle32_v_f32m1(x_data, vl);
      vfloat32m1_t vec_y = __riscv_vfmul_vf_f32m1(vec_x, scalar, vl);
      __riscv_vse32_v_f32m1(y_data, vec_y, vl);
      x_data += vl;
      y_data += vl;
      n -= vl;
    }
  #else  /* scalar implementation */
    for (size_t i = 0; i < n; i += 1) {
      y_data[i] = x_data[i] * scalar;
    }
  #endif
}
 
void nn_mulscalar2d_f32(Tensor2D_F32 *y, const Tensor2D_F32 *x, float scalar) {
  nn_assert(y->shape[0] == x->shape[0] && y->shape[1] == x->shape[1], "Cannot add tensors of different shapes");
 
  size_t n = y->shape[0] * y->shape[1];
  float *x_data = x->data;
  float *y_data = y->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x = __riscv_vle32_v_f32m1(x_data, vl);
      vfloat32m1_t vec_y = __riscv_vfmul_vf_f32m1(vec_x, scalar, vl);
      __riscv_vse32_v_f32m1(y_data, vec_y, vl);
      x_data += vl;
      y_data += vl;
      n -= vl;
    }
  #else  /* scalar implementation */
    for (size_t i = 0; i < n; i += 1) {
      y_data[i] = x_data[i] * scalar;
    }
  #endif
}
 
 
/* ======================================================================================================== */
/*                                           Convolution                                                    */
/* ======================================================================================================== */
 
void nn_nchw_to_nhwc_f32(Tensor4D_F32 *out, const Tensor4D_F32 *in) {
  nn_assert(out->shape[0] == in->shape[0], "Cannot convert between tensors of different shapes");
  nn_assert(out->shape[1] == in->shape[2], "Cannot convert between tensors of different shapes");
  nn_assert(out->shape[2] == in->shape[3], "Cannot convert between tensors of different shapes");
  nn_assert(out->shape[3] == in->shape[1], "Cannot convert between tensors of different shapes");
 
  size_t batch_size = in->shape[0];
  size_t height = in->shape[2];
  size_t width = in->shape[3];
  size_t channels = in->shape[1];
 
  for (size_t n = 0; n < batch_size; n += 1) {
    for (size_t c = 0; c < channels; c += 1) {
      for (size_t h = 0; h < height; h += 1) {
        for (size_t w = 0; w < width; w += 1) {
          size_t nchw_index = n * channels * height * width + c * height * width + h * width + w;
          size_t nhwc_index = n * height * width * channels + h * width * channels + w * channels + c;
          ((float *)out->data)[nhwc_index] = ((float *)in->data)[nchw_index];
        }
      }
    }
  }
}
 
 
void nn_nhwc_to_nchw_f32(Tensor4D_F32 *out, const Tensor4D_F32 *in) {
  nn_assert(out->shape[0] == in->shape[0], "Cannot convert between tensors of different shapes");
  nn_assert(out->shape[1] == in->shape[3], "Cannot convert between tensors of different shapes");
  nn_assert(out->shape[2] == in->shape[1], "Cannot convert between tensors of different shapes");
  nn_assert(out->shape[3] == in->shape[2], "Cannot convert between tensors of different shapes");
 
  size_t batch_size = in->shape[0];
  size_t height = in->shape[1];
  size_t width = in->shape[2];
  size_t channels = in->shape[3];
 
  for (size_t n = 0; n < batch_size; n += 1) {
    for (size_t c = 0; c < channels; c += 1) {
      for (size_t h = 0; h < height; h += 1) {
        for (size_t w = 0; w < width; w += 1) {
          size_t nhwc_index = n * height * width * channels + h * width * channels + w * channels + c;
          size_t nchw_index = n * channels * height * width + c * height * width + h * width + w;
          ((float *)out->data)[nchw_index] = ((float *)in->data)[nhwc_index];
        }
      }
    }
  }
}
 
void nn_conv2d_f32(
  Tensor4D_F32 *out, const Tensor4D_F32 *in,
  const Tensor4D_F32 *weight, const Tensor1D_F32 *bias,
  const size_t *stride, const size_t *padding, const size_t *dilation, size_t groups) {
 
  size_t batch_size = in->shape[0];
  size_t in_height = in->shape[1];
  size_t in_width = in->shape[2];
  size_t in_channels = in->shape[3];
  
  size_t out_height = out->shape[1];
  size_t out_width = out->shape[2];
  size_t out_channels = out->shape[3];
 
  size_t kernel_height = weight->shape[0];
  size_t kernel_width = weight->shape[1];
  size_t stride_height = stride[0];
  size_t stride_width = stride[1];
  size_t padding_height = padding[0];
  size_t padding_width = padding[1];
  size_t dilation_height = dilation[0];
  size_t dilation_width = dilation[1];
 
  nn_assert(out->shape[0] == batch_size, "Cannot add tensors of different shapes");
  nn_assert(weight->shape[3] == out_channels, "Cannot add tensors of different shapes");
  nn_assert(weight->shape[2] * groups == in_channels, "Cannot add tensors of different shapes");
  nn_assert(out_height == (in_height + 2 * padding_height - dilation_height * (kernel_height - 1) - 1) / stride_height + 1, "Cannot add tensors of different shapes");
  nn_assert(out_width == (in_width + 2 * padding_width - dilation_width * (kernel_width - 1) - 1) / stride_width + 1, "Cannot add tensors of different shapes");
  nn_assert(groups > 0, "Cannot add tensors of different shapes");
  nn_assert(in_channels % groups == 0, "Cannot add tensors of different shapes");
  nn_assert(out_channels % groups == 0, "Cannot add tensors of different shapes");
 
 
  // Initialize output tensor to zeros
  memset(out->data, 0, batch_size * out_height * out_width * out_channels * sizeof(float));
 
  #ifdef CONFIG_BACKEND_GEMMINI
    if (!bias) {
      bias = nn_zeros1d_f32((size_t []){ out_channels });
    }
 
    if (groups == 1) {
      tiled_conv_auto(
          batch_size, in_height, in_width, in_channels,
          out_channels, out_height, out_width,
          stride_height, dilation_height, 1, padding_height, kernel_height, 
          0, 0, 0, 0, 0,
          in->data,
          weight->data,
          bias->data,
          out->data,
          NO_ACTIVATION, ACC_SCALE_IDENTITY,
          0, 0, 0,
          WS);
    }
    else if (groups == in_channels) {
      assert(weight->shape[2] == 1);
 
      Tensor *in_nchw = NN_tensor(4, (size_t[]){batch_size, in_channels, in_height, in_width}, DTYPE_F32, NULL);
      Tensor *out_nchw = NN_tensor(4, (size_t[]){batch_size, out_channels, out_height, out_width}, DTYPE_F32, NULL);
      
      Tensor *weight_1hwc = NN_tensor(4, (size_t[]){1, kernel_height, kernel_width, out_channels}, DTYPE_F32, weight->data);
      Tensor *weight_1chw = NN_tensor(4, (size_t[]){1, out_channels, kernel_height, kernel_width}, DTYPE_F32, NULL);
 
      NN_nhwc_to_nchw(in_nchw, in);
      NN_nhwc_to_nchw(weight_1chw, weight_1hwc);
 
      for (size_t g = 0; g < groups; g += 1) {
        tiled_conv_auto(
          batch_size, in_height, in_width, 1,
          1, out_height, out_width,
          stride_height, dilation_height, 1, padding_height, kernel_height, 
          0, 0, 0, 0, 0,
          ((float *)in_nchw->data) + g * in_height * in_width,
          ((float *)weight_1chw->data) + g * kernel_height * kernel_width,
          ((float *)bias->data) + g,
          ((float *)out_nchw->data) + g * out_height * out_width,
          NO_ACTIVATION, ACC_SCALE_IDENTITY,
          0, 0, 0,
          WS);
      }
 
      NN_nchw_to_nhwc(out, out_nchw);
 
    }
    else {
      printf("[ERROR] Unsupported conv2d operation for groups other than 1 or in_channels\n");
    }
 
 
  #else
 
    if (groups == 1) {
      // Standard convolution
      for (size_t b = 0; b < batch_size; b += 1) {
        for (size_t oc = 0; oc < out_channels; oc += 1) {
          for (size_t oh = 0; oh < out_height; oh += 1) {
            for (size_t ow = 0; ow < out_width; ow += 1) {
              float value = 0.0f;
              for (size_t kh = 0; kh < kernel_height; kh += 1) {
                for (size_t kw = 0; kw < kernel_width; kw += 1) {
                  for (size_t ic = 0; ic < in_channels; ic += 1) {
                    size_t ih = oh * stride_height + kh * dilation_height - padding_height;
                    size_t iw = ow * stride_width + kw * dilation_width - padding_width;
                    if (ih < in_height && iw < in_width) {
                      size_t in_idx = b * in_height * in_width * in_channels
                          + ih * in_width * in_channels
                          + iw * in_channels
                          + ic;
                      size_t weight_idx = kh * kernel_width * in_channels * out_channels
                          + kw * in_channels * out_channels
                          + ic * out_channels
                          + oc;
                      value += ((float*)in->data)[in_idx] * ((float*)weight->data)[weight_idx];
                    }
                  }
                }
              }
              if (bias != NULL) {
                value += ((float*)bias->data)[oc];
              }
              size_t out_idx = b * out_height * out_width * out_channels
                  + oh * out_width * out_channels
                  + ow * out_channels
                  + oc;
              ((float*)out->data)[out_idx] = value;
            }
          }
        }
      }
    }
    else if (groups == in_channels) {
      // Depthwise convolution
      for (size_t b = 0; b < batch_size; b += 1) {
        for (size_t oc = 0; oc < out_channels; oc += 1) {
          for (size_t oh = 0; oh < out_height; oh += 1) {
            for (size_t ow = 0; ow < out_width; ow += 1) {
              float value = 0.0f;
              for (size_t kh = 0; kh < kernel_height; kh += 1) {
                for (size_t kw = 0; kw < kernel_width; kw += 1) {
                  size_t ih = oh * stride_height + kh * dilation_height - padding_height;
                  size_t iw = ow * stride_width + kw * dilation_width - padding_width;
                  if (ih < in_height && iw < in_width) {
                    size_t in_idx = b * in_height * in_width * in_channels
                        + ih * in_width * in_channels
                        + iw * in_channels
                        + oc;
                    size_t weight_idx = kh * kernel_width * in_channels
                        + kw * in_channels
                        + oc;
                    value += ((float *)in->data)[in_idx] * ((float *)weight->data)[weight_idx];
                  }
                }
              }
              if (bias != NULL) {
                value += ((float *)bias->data)[oc];
              }
              size_t out_idx = b * out_height * out_width * out_channels
                  + oh * out_width * out_channels
                  + ow * out_channels
                  + oc;
              ((float *)out->data)[out_idx] = value;
            }
          }
        }
      }
    }
    else {
      printf("[ERROR] Unsupported conv2d operation for groups other than 1 or in_channels\n");
    }
  #endif
}
 
 
/* ======================================================================================================== */
/*                                           MatMul                                                         */
/* ======================================================================================================== */
 
void nn_dot_f32(Tensor1D_F32 *y, const Tensor1D_F32 *x1, const Tensor1D_F32 *x2) {
  nn_assert(x1->shape[0] == x2->shape[0], "Cannot dot tensors of different shapes");
  nn_assert(y->shape[0] == x1->shape[0], "Cannot dot tensors of different shapes");
 
  size_t n = y->shape[0];
  float *x1_data = x1->data;
  float *x2_data = x2->data;
  float *y_data = y->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x1 = __riscv_vle32_v_f32m1(x1_data, vl);
      vfloat32m1_t vec_x2 = __riscv_vle32_v_f32m1(x2_data, vl);
      vfloat32m1_t vec_y = __riscv_vfmul_vv_f32m1(vec_x1, vec_x2, vl);
      __riscv_vse32_v_f32m1(y_data, vec_y, vl);
      x1_data += vl;
      x2_data += vl;
      y_data += vl;
      n -= vl;
    }
  #else  /* scalar implementation */
    float sum = 0.0f;
    for (size_t i = 0; i < n; i += 1) {
      sum += x1_data[i] * x2_data[i];
    }
    y_data[0] = sum;
  #endif
}
 
void nn_mv_f32(Tensor1D_F32 *y, const Tensor2D_F32 *x1, const Tensor1D_F32 *x2) {
  nn_assert(x1->shape[1] == x2->shape[0], "Cannot perform MV on tensors of different shapes");
  nn_assert(y->shape[0] == x1->shape[0], "Cannot perform MV on tensors of different shapes");
 
  const size_t n = x1->shape[0]; // rows in matrix
  const size_t m = x1->shape[1]; // columns in matrix
  float *x1_data = x1->data;
  float *x2_data = x2->data;
  float *y_data = y->data;
 
  for (size_t i = 0; i < y->shape[0]; i += 1) {
    float sum = 0.0;
    for (size_t j = 0; j < m; j += 1) {
      sum += x1_data[i * m + j] * x2_data[j];
    }
    y_data[i] = sum;
  }
}
  
 
void nn_mm_f32(Tensor2D_F32 *y, const Tensor2D_F32 *x1, const Tensor2D_F32 *x2) {
  nn_assert(x1->shape[1] == x2->shape[0], "Cannot perform MatMul on tensors of different shapes");
  nn_assert(y->shape[0] == x1->shape[0] && y->shape[1] == x2->shape[1], "Cannot perform MatMul on tensors of different shapes");
 
  const size_t n = x1->shape[0];
  const size_t m = x1->shape[1];
  const size_t p = x2->shape[1];
 
  for (size_t i = 0; i < n; i += 1) {
    float *x1_row = x1->data + i * m;
    float *y_row = y->data + i * p;
 
    #ifdef CONFIG_BACKEND_RISCV_V
 
      size_t vlmax = __riscv_vsetvlmax_e32m1();
      for (size_t j = 0; j < p; j += 1) {
        vfloat32m1_t vec_zero = __riscv_vfmv_v_f_f32m1(0, vlmax);
        vfloat32m1_t vec_sum = __riscv_vfmv_v_f_f32m1(0, vlmax);
 
        float *x2_col = x2->data + j;
        size_t k = m;
 
        while (k > 0) {
          size_t vl = __riscv_vsetvl_e32m1(k);
          vfloat32m1_t vec_x1 = __riscv_vle32_v_f32m1(x1_row, vl);
          vfloat32m1_t vec_x2 = __riscv_vlse32_v_f32m1(x2_col, p * sizeof(float), vl);
          vec_sum = __riscv_vfmacc_vv_f32m1(vec_sum, vec_x1, vec_x2, vl);
 
          x1_row += vl;
          x2_col += vl * p;
          k -= vl;
        }
 
        #ifdef CONFIG_DEBUG_RISCV_V_USE_REDOSUM
          vec_sum = __riscv_vfredosum_vs_f32m1_f32m1(vec_sum, vec_zero, vlmax);
        #else
          vec_sum = __riscv_vfredusum_vs_f32m1_f32m1(vec_sum, vec_zero, vlmax);
        #endif
        y_row[j] = __riscv_vfmv_f_s_f32m1_f32(vec_sum);
      }
    #else
      for (size_t j = 0; j < p; j += 1) {
        float *x2_row = x2->data + j;
 
        float sum = 0.f;
        for (size_t k = 0; k < m; k += 1) {
          sum += x1_row[k] * x2_row[k * p];
        }
        y_row[j] = sum;
      }
    #endif
  }
}
 
void nn_addmm_f32(Tensor2D_F32 *y, const Tensor2D_F32 *c, const Tensor2D_F32 *x1, const Tensor2D_F32 *x2) {
  nn_assert(x1->shape[1] == x2->shape[0], "Cannot perform MatMulAdd on tensors of different shapes");
  nn_assert(y->shape[0] == c->shape[0] && y->shape[1] == x2->shape[1], "Cannot perform MatMulAdd on tensors of different shapes");
 
  const size_t n = x1->shape[0];
  const size_t m = x1->shape[1];
  const size_t p = x2->shape[1];
 
  for (size_t i = 0; i < n; i += 1) {
    float *x1_row = x1->data + i * m;
    float *c_row = c->data + i * p;
    float *y_row = y->data + i * p;
    
    #ifdef CONFIG_BACKEND_RISCV_V
 
      size_t vlmax = __riscv_vsetvlmax_e32m1();
      for (size_t j = 0; j < p; j += 1) {
        vfloat32m1_t vec_zero = __riscv_vfmv_v_f_f32m1(0, vlmax);
        vfloat32m1_t vec_sum = __riscv_vfmv_v_f_f32m1(0, vlmax);
 
        float *x2_col = x2->data + j;
        size_t k = m;
 
        while (k > 0) {
          size_t vl = __riscv_vsetvl_e32m1(k);
          vfloat32m1_t vec_x1 = __riscv_vle32_v_f32m1(x1_row, vl);
          vfloat32m1_t vec_x2 = __riscv_vlse32_v_f32m1(x2_col, p * sizeof(float), vl);
          vec_sum = __riscv_vfmacc_vv_f32m1(vec_sum, vec_x1, vec_x2, vl);
 
          x1_row += vl;
          x2_col += vl * p;
          k -= vl;
        }
 
        #ifdef CONFIG_DEBUG_RISCV_V_USE_REDOSUM
          vec_sum = __riscv_vfredosum_vs_f32m1_f32m1(vec_sum, vec_zero, vlmax);
        #else
          vec_sum = __riscv_vfredusum_vs_f32m1_f32m1(vec_sum, vec_zero, vlmax);
        #endif
        y_row[j] = __riscv_vfmv_f_s_f32m1_f32(vec_sum) + c_row[j];
      }
 
    #else
      for (size_t j = 0; j < p; j += 1) {
        float *x2_col = x2->data + j;
 
        float sum = 0.f;
        for (size_t k = 0; k < m; k += 1) {
          sum += x1_row[k] * x2_col[k * p];
        }
        y_row[j] = sum + c_row[j];
      }
    #endif
    x1_row += m;
    y_row += p;
  }
}
 
void nn_linear_f32(Tensor2D_F32 *y, const Tensor2D_F32 *x, const Tensor2D_F32 *weight, const Tensor1D_F32 *bias) {
  nn_assert(x->shape[1] == weight->shape[1], "Cannot perform Linear on tensors of different shapes");
  nn_assert(!bias || bias->shape[0] == weight->shape[0], "Cannot perform Linear on tensors of different shapes");
  nn_assert(y->shape[0] == x->shape[0] && y->shape[1] == weight->shape[0], "Cannot perform Linear on tensors of different shapes");
 
  const size_t batch_size = x->shape[0];
  const size_t in_features = x->shape[1];
  const size_t out_features = weight->shape[0];
 
  float *x_batch_data = x->data;
  float *y_batch_data = y->data;
 
  for (size_t i = 0; i < batch_size; i += 1) {
    float *x_data = x_batch_data;
    float *y_data = y_batch_data;
 
    #ifdef CONFIG_BACKEND_RISCV_V
      size_t vlmax = __riscv_vsetvlmax_e32m1();
 
      for (size_t j = 0; j < out_features; j += 1) {
        vfloat32m1_t vec_zero = __riscv_vfmv_v_f_f32m1(0, vlmax);
        vfloat32m1_t vec_sum = __riscv_vfmv_v_f_f32m1(0, vlmax);
 
        float *weight_row = weight->data + j * in_features;
        size_t n = in_features;
 
        while (n > 0) {
          size_t vl = __riscv_vsetvl_e32m1(n);
          vfloat32m1_t vec_x = __riscv_vle32_v_f32m1(x_data, vl);
          vfloat32m1_t vec_w = __riscv_vle32_v_f32m1(weight_row, vl);
          vec_sum = __riscv_vfmacc_vv_f32m1(vec_sum, vec_x, vec_w, vl);
 
          x_data += vl;
          weight_row += vl;
          n -= vl;
        }
 
        #ifdef CONFIG_DEBUG_RISCV_V_USE_REDOSUM
          vec_sum = __riscv_vfredosum_vs_f32m1_f32m1(vec_sum, vec_zero, vlmax);
        #else
          vec_sum = __riscv_vfredusum_vs_f32m1_f32m1(vec_sum, vec_zero, vlmax);
        #endif
 
        float sum = __riscv_vfmv_f_s_f32m1_f32(vec_sum);
        if (bias) {
          sum += bias->data[j];
        }
        y_data[j] = sum;
        x_data = x_batch_data; // reset x_data pointer for next output feature
      }
    #else  /* scalar implementation */
      for (size_t j = 0; j < out_features; j += 1) {
        float *weight_row = weight->data + j * in_features;
 
        float sum = 0.f;
        for (size_t k = 0; k < in_features; k += 1) {
          sum += x_data[k] * weight_row[k];
        }
        if (bias) {
          sum += bias->data[j];
        }
        y_data[j] = sum;
      }
    #endif
 
    x_batch_data += in_features;
    y_batch_data += out_features;
  }
}
 
 
/* ======================================================================================================== */
/*                                           Non-linear                                                     */
/* ======================================================================================================== */
 
void nn_elu2d_f32(Tensor2D_F32 *y, const Tensor2D_F32 *x, float alpha) {
  nn_assert(x->shape[0] == y->shape[0] && x->shape[1] == y->shape[1], "Cannot perform ELU on tensors of different shapes");
 
  const size_t n = y->shape[0] * y->shape[1];
  float *x_data = x->data;
  float *y_data = y->data;
  
  for (size_t i = 0; i < n; i += 1) {
    if (x_data[i] > 0) {
      y_data[i] = x_data[i];
    }
    else {
      y_data[i] = alpha * (expf(x_data[i]) - 1.f);
    }
  }
}
 
void nn_relu2d_f32(Tensor2D_F32 *y, const Tensor2D_F32 *x) {
  nn_assert(x->shape[0] == y->shape[0] && x->shape[1] == y->shape[1], "Cannot perform ReLU on tensors of different shapes");
 
  size_t n = y->shape[0] * y->shape[1];
  float *x_data = x->data;
  float *y_data = y->data;
 
  #ifdef CONFIG_BACKEND_RISCV_V
    float zero = 0.0f;
 
    while (n > 0) {
      size_t vl = __riscv_vsetvl_e32m1(n);
      vfloat32m1_t vec_x = __riscv_vle32_v_f32m1(x_data, vl);
      vfloat32m1_t vec_y = __riscv_vfmax_vf_f32m1(vec_x, zero, vl);
      __riscv_vse32_v_f32m1(y_data, vec_y, vl);
      x_data += vl;
      y_data += vl;
      n -= vl;
    }
  #else  /* scalar implementation */
    for (size_t i = 0; i < n; i += 1) {
      float x_val = x_data[i];
      y_data[i] = x_val > 0 ? x_val : 0;
    }
  #endif
}
 
void nn_silu1d_f32(Tensor1D_F32 *y, const Tensor1D_F32 *x) {
  nn_assert(x->shape[0] == y->shape[0], "Cannot perform SiLU on tensors of different shapes");
 
  const size_t n = y->shape[0];
  float *x_data = x->data;
  float *y_data = y->data;
 
  for (size_t i = 0; i < n; i++) {
    float x_i = x_data[i];
    float sigmoid_x = 1.0f / (1.0f + expf(-x_i));
    y_data[i] = x_i * sigmoid_x;
  }
}
 
void nn_softmax1d_f32(Tensor1D_F32 *y, const Tensor1D_F32 *x) {
  nn_assert(y->shape[0] == x->shape[0], "Cannot add tensors of different shapes");
 
  const size_t n = y->shape[0];
  float *x_data = x->data;
  float *y_data = y->data;
 
  float sum = 0.0f;
  for (size_t i = 0; i < n; i += 1) {
    y_data[i] = expf(x_data[i]);
    sum += y_data[i];
  }
  // normalize
  for (size_t i = 0; i < n; i += 1) {
    y_data[i] /= sum;
  }
}
 
void nn_softmax2d_f32(Tensor2D_F32 *y, const Tensor2D_F32 *x, size_t dim) {
  nn_assert(y->shape[0] == x->shape[0] && y->shape[1] == x->shape[1], "Cannot add tensors of different shapes");
 
  float *y_data = y->data;
  float *x_data = x->data;
 
  if (dim == 0) {
    for (size_t i = 0; i < y->shape[1]; i += 1) {
      size_t n = y->shape[0];
      size_t m = y->shape[1];
      float sum = 0.0f;
      for (size_t j = 0; j < n; j += 1) {
        sum += expf(x_data[j * m]);
      }
 
      for (size_t j = 0; j < n; j += 1) {
        y_data[j * m] = expf(x_data[j * m]) / sum;
      }
 
      x_data += 1;
      y_data += 1;
    }
  }
  else if (dim == 1) {
    // HACK: fix batch size
    for (size_t i = 0; i < y->shape[0]; i += 1) {
      size_t n = y->shape[1];
      float sum = 0.0f;
      for (size_t j = 0; j < n; j += 1) {
        sum += expf(x_data[j]);
      }
 
      for (size_t j = 0; j < n; j += 1) {
        y_data[j] = expf(x_data[j]) / sum;
      }
 
      x_data += n;
      y_data += n;
    }
  }
  else {
    nn_assert(0, "Invalid dimension for softmax");
  }
}
 
void nn_tanh2d_f32(Tensor2D_F32 *y, const Tensor2D_F32 *x) {
  nn_assert(x->shape[0] == y->shape[0] && x->shape[1] == y->shape[1], "Cannot perform ReLU on tensors of different shapes");
 
  const size_t n = y->shape[0] * y->shape[1];
  float *x_data = x->data;
  float *y_data = y->data;
 
  for (size_t i = 0; i < n; i += 1) {
    float x_val = x_data[i];
    y_data[i] = tanhf(x_val);
  }
}
 
void nn_rms_norm1d_f32(Tensor1D_F32 *y, const Tensor1D_F32 *x, const Tensor1D_F32 *weight, float eps) {
  nn_assert(x->shape[0] == y->shape[0], "Cannot perform RMSNorm on tensors of different shapes");
 
  const size_t n = y->shape[0];
  float *x_data = x->data;
  float *y_data = y->data;
  float *w_data = weight->data;
  
  float ss = 0.0f;
  for (size_t i = 0; i < n; i += 1) {
    ss += x_data[i] * x_data[i];
  }
  ss /= n;
  ss += eps;
 
  // normalize and scale
  // y = (x / ss) * w
  nn_mulscalar1d_f32(y, x, 1.0f / sqrtf(ss));
  nn_mul1d_f32(y, y, weight);
}
 
 
/* ======================================================================================================== */
/*                                           Attention                                                      */
/* ======================================================================================================== */
 
void nn_scaled_dot_product_attention_f32(Tensor4D_F32 *y, const Tensor4D_F32 *query, const Tensor4D_F32 *key, const Tensor4D_F32 *value) {
  nn_assert(query->shape[0] == key->shape[0] && query->shape[0] == value->shape[0], "Query, key, and value must have the same batch size");
  nn_assert(query->shape[1] == key->shape[1] && query->shape[1] == value->shape[1], "Query, key, and value must have the same number of heads");
  nn_assert(key->shape[2] == value->shape[2], "Key and value must have the same sequence length");
  nn_assert(query->shape[3] == key->shape[3], "Query and key must have the same embedding dimension");
 
  size_t n = query->shape[0]; // batch size
  size_t h = query->shape[1]; // head count
  size_t l = query->shape[2]; // target sequence length (query)
  size_t s = key->shape[2];   // source sequence length (key/value)
  size_t e = query->shape[3]; // embedding dimension
  size_t ev = value->shape[3]; // value embedding dimension
 
  // scale_factor = 1 / math.sqrt(query.size(-1))
  float scale_factor = 1.0f / sqrt(e);
 
  // Process each batch
  for (size_t batch = 0; batch < n; batch += 1) {
    // Process each head
    for (size_t head = 0; head < h; head += 1) {
      // Set up tensor views for the current batch and head
      size_t query_head_dims[2] = {l, e};
      size_t key_head_dims[2] = {s, e};  // Corrected: should be s, not l
      size_t key_transposed_dims[2] = {e, s};  // Transposed key dimensions
      size_t attn_weight_head_dims[2] = {l, s};
      size_t value_head_dims[2] = {s, ev};
      size_t y_head_dims[2] = {l, ev};
 
      // Get the data pointers for the current batch and head
      float *query_data = (float *)query->data + (batch * h * l * e) + (head * l * e);
      float *key_data = (float *)key->data + (batch * h * s * e) + (head * s * e);
      float *value_data = (float *)value->data + (batch * h * s * ev) + (head * s * ev);
      float *y_data = (float *)y->data + (batch * h * l * ev) + (head * l * ev);
 
      // Create tensor views
      Tensor2D_F32 *query_head = nn_as_tensor2d_f32(query_head_dims, query_data);
      Tensor2D_F32 *key_head = nn_as_tensor2d_f32(key_head_dims, key_data);
      Tensor2D_F32 *value_head = nn_as_tensor2d_f32(value_head_dims, value_data);
      Tensor2D_F32 *y_head = nn_as_tensor2d_f32(y_head_dims, y_data);
 
      // Create and transpose the key matrix manually (key.transpose(-2, -1))
      Tensor2D_F32 *key_transposed = nn_tensor2d_f32(key_transposed_dims, NULL);
      for (size_t i = 0; i < s; i += 1) {
        for (size_t j = 0; j < e; j += 1) {
          key_transposed->data[j * s + i] = key_head->data[i * e + j];
        }
      }
 
      // Calculate attention weights: attn_weight = query @ key.transpose(-2, -1)
      Tensor2D_F32 *attn_weight_head = nn_tensor2d_f32(attn_weight_head_dims, NULL);
      nn_mm_f32(attn_weight_head, query_head, key_transposed);
 
      // Apply scaling: attn_weight = attn_weight * scale_factor
      nn_mulscalar2d_f32(attn_weight_head, attn_weight_head, scale_factor);
 
      // attn_weight = torch.softmax(attn_weight, dim=-1)
      nn_softmax2d_f32(attn_weight_head, attn_weight_head, 1);
 
      // (n, h, l, ev) = (n, h, l, s) @ (n, h, s, ev)
      // output = attn_weight @ value
      nn_mm_f32(y_head, attn_weight_head, value_head);
 
      // Free the temporary tensors we created
      free(query_head);
      free(key_head);
      free(key_transposed->data);
      free(key_transposed);
      free(attn_weight_head->data);
      free(attn_weight_head);
      free(value_head);
      free(y_head);
    }
  }
}
 
 
#endif // __NN_F32_H