#include "ccv.h"

#include "ccv_internal.h"

#include "nnc/ccv_nnc.h"

#include "nnc/ccv_nnc_easy.h"

#include "nnc/ccv_nnc_internal.h"

#ifdef USE_OPENMP

#include <omp.h>

#endif

#ifdef USE_DISPATCH

#include <dispatch/dispatch.h>

#endif

// Shared methods.

#include "../_ccv_nnc_cpu_ref.h"

static int _ccv_nnc_rmsnorm_forw(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)

{

  assert(input_size == 2);

  ccv_nnc_tensor_view_t* const a = (ccv_nnc_tensor_view_t*)inputs[0];

  ccv_nnc_tensor_view_t* const scale = (ccv_nnc_tensor_view_t*)inputs[1];

  ccv_nnc_tensor_view_t* const b = (ccv_nnc_tensor_view_t*)outputs[0];

  ccv_nnc_tensor_view_t* const saved_inv_std = (ccv_nnc_tensor_view_t*)outputs[1];

  assert(ccv_nnc_tensor_nd(a->info.dim) <= CCV_NNC_MAX_DIM + 2);

  assert(ccv_nnc_tensor_nd(b->info.dim) <= CCV_NNC_MAX_DIM + 2);

  // Assuming this is float 32.

  int adim[CCV_NNC_MAX_DIM_ALLOC];

  int rdim[CCV_NNC_MAX_DIM_ALLOC];

  ccv_nnc_tensor_view_get_dim(a, adim);

  ccv_nnc_tensor_view_get_dim(saved_inv_std, rdim);

  assert(ccv_nnc_tensor_view_check_dim(b, adim));

  assert(CCV_NNC_MAX_DIM == 2); // Need to change this logic for CCV_NNC_MAX_DIM == other number.

  int astride[CCV_NNC_MAX_DIM_ALLOC];

  int bstride[CCV_NNC_MAX_DIM_ALLOC];

  int scale_stride[CCV_NNC_MAX_DIM_ALLOC];

  ccv_nnc_tensor_view_get_stride(a, astride);

  ccv_nnc_tensor_view_get_stride(scale, scale_stride);

  ccv_nnc_tensor_view_get_stride(b, bstride);

  // The epsilon is used a little bit differently from batch norm, it is outside of the sqrt in this case.

  const float epsilon = cmd.info.rmsnorm.epsilon;

  int saved_inv_std_stride[CCV_NNC_MAX_DIM_ALLOC];

  ccv_nnc_tensor_view_get_stride(saved_inv_std, saved_inv_std_stride);

  int x;

  int n = 1;

  for (x = 0; x < CCV_NNC_MAX_DIM + 2; 
x++20
)

    n *= adim[x];

  for (x = 0; x < CCV_NNC_MAX_DIM + 2; 
x++20
)

    n /= rdim[x];

  const float inv_n = 1. / n;

  ccv_nnc_tensor_zero(saved_inv_std);

  float* const ap = a->data.f32;

  float* const varp = saved_inv_std->data.f32;

  int i[CCV_NNC_MAX_DIM + 2];

  for (i[0] = 0; i[0] < adim[0]; 
i[0]++22
)

  {

    float* const ap0 = ap + i[0] * astride[0];

    float* const varp0 = rdim[0] == 1 ? 
varp0
 : varp + i[0] * saved_inv_std_stride[0];

    for (i[1] = 0; i[1] < adim[1]; 
i[1]++76
)

    {

      float* ap1 = ap0 + i[1] * astride[1];

      float* const varp1 = rdim[1] == 1 ? varp0 : 
varp0 + i[1] * saved_inv_std_stride[1]0
;

      for (i[2] = 0; i[2] < adim[2]; 
i[2]++280
)

      {

        float* const varp2 = rdim[2] == 1 ? varp1 : 
varp1 + i[2] * saved_inv_std_stride[2]0
;

        if (rdim[3] == 1)

          
for (x = 0; 280
x < adim[3]; 
x++2.80k
)

          {

            float w = ap1[x * astride[3]];

            varp2[0] += w * w;

          }

        else

          for (x = 0; x < adim[3]; x++)

          {

            float w = ap1[x * astride[3]];

            varp2[x] += w * w;

          }

        ap1 += astride[2];

      }

    }

  }

  for (i[0] = 0; i[0] < rdim[0]; 
i[0]++22
)

  {

    float* const varp0 = varp + i[0] * saved_inv_std_stride[0];

    for (i[1] = 0; i[1] < rdim[1]; 
i[1]++22
)

    {

      float* const varp1 = varp0 + i[1] * saved_inv_std_stride[1];

      for (i[2] = 0; i[2] < rdim[2]; 
i[2]++22
)

      {

        float* const varp2 = varp1 + i[2] * saved_inv_std_stride[2];

        for (x = 0; x < rdim[3]; 
x++22
)

          varp2[x] = 1. / sqrtf(varp2[x] * inv_n + epsilon);

      }

    }

  }

  float* const scalep = scale->data.f32;

  int sdim[CCV_NNC_MAX_DIM_ALLOC];

  ccv_nnc_tensor_view_get_dim(scale, sdim);

  // Do the straight-forward one, y = x * inv_std * scale + bias, we cannot allocate extra memory to help.

  // There is no need for precompute since scale / bias is per element.

  float* const bp = b->data.f32;

  for (i[0] = 0; i[0] < adim[0]; 
i[0]++22
)

  {

    float* const ap0 = ap + i[0] * astride[0];

    float* const bp0 = bp + i[0] * bstride[0];

    float* const varp0 = rdim[0] == 1 ? 
varp0
 : varp + i[0] * saved_inv_std_stride[0];

    float* const scalep0 = sdim[0] == 1 ? scalep : 
scalep + i[0] * scale_stride[0]0
;

    for (i[1] = 0; i[1] < adim[1]; 
i[1]++76
)

    {

      float* ap1 = ap0 + i[1] * astride[1];

      float* bp1 = bp0 + i[1] * bstride[1];

      float* const varp1 = rdim[1] == 1 ? varp0 : 
varp0 + i[1] * saved_inv_std_stride[1]0
;

      float* const scalep1 = sdim[1] == 1 ? 
scalep00
 : scalep0 + i[1] * scale_stride[1];

      for (i[2] = 0; i[2] < adim[2]; 
i[2]++280
)

      {

        float* const varp2 = rdim[2] == 1 ? varp1 : 
varp1 + i[2] * saved_inv_std_stride[2]0
;

        float* const scalep2 = sdim[2] == 1 ? 
scalep10
 : scalep1 + i[2] * scale_stride[2];

        if (rdim[3] == 1)

          
for (x = 0; 280
x < adim[3]; 
x++2.80k
)

            bp1[x] = ap1[x * astride[3]] * varp2[0] * scalep2[sdim[3] == 1 ? 
00
 : x];

        else

          for (x = 0; x < adim[3]; x++)

            bp1[x] = ap1[x * astride[3]] * varp2[x] * scalep2[sdim[3] == 1 ? 0 : x];

        ap1 += astride[2];

        bp1 += bstride[2];

      }

    }

  }

  return CCV_NNC_EXEC_SUCCESS;

}

static int _ccv_nnc_rmsnorm_back(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_t hint, const int flags, ccv_nnc_tensor_t* const* const inputs, const int input_size, ccv_nnc_tensor_t* const* const outputs, const int output_size, ccv_nnc_stream_context_t* const stream_context)

{

  assert(input_size == 6);

  assert(output_size >= 1);

  ccv_nnc_tensor_view_t* const g = (ccv_nnc_tensor_view_t*)inputs[0];

  ccv_nnc_tensor_view_t* const a = (ccv_nnc_tensor_view_t*)inputs[2];

  ccv_nnc_tensor_view_t* const scale = (ccv_nnc_tensor_view_t*)inputs[3];

  ccv_nnc_tensor_view_t* const saved_inv_std = (ccv_nnc_tensor_view_t*)inputs[5];

  ccv_nnc_tensor_view_t* const h = (ccv_nnc_tensor_view_t*)outputs[0];

  ccv_nnc_tensor_view_t* const dscale = output_size > 1 ? (ccv_nnc_tensor_view_t*)outputs[1] : 
00
;

  assert(ccv_nnc_tensor_nd(g->info.dim) <= CCV_NNC_MAX_DIM + 2);

  assert(ccv_nnc_tensor_nd(a->info.dim) <= CCV_NNC_MAX_DIM + 2);

  assert(ccv_nnc_tensor_nd(h->info.dim) <= CCV_NNC_MAX_DIM + 2);

  // Assuming this is float 32.

  int gdim[CCV_NNC_MAX_DIM_ALLOC];

  int rdim[CCV_NNC_MAX_DIM_ALLOC];

  ccv_nnc_tensor_view_get_dim(g, gdim);

  ccv_nnc_tensor_view_get_dim(saved_inv_std, rdim);

  int sdim[CCV_NNC_MAX_DIM_ALLOC];

  ccv_nnc_tensor_view_get_dim(scale, sdim);

  if (dscale)

    { assert(ccv_nnc_tensor_view_check_dim(dscale, sdim)); }

  assert(ccv_nnc_tensor_view_check_dim(a, gdim));

  assert(ccv_nnc_tensor_view_check_dim(h, gdim));

  int astride[CCV_NNC_MAX_DIM_ALLOC];

  int gstride[CCV_NNC_MAX_DIM_ALLOC];

  int hstride[CCV_NNC_MAX_DIM_ALLOC];

  int scale_stride[CCV_NNC_MAX_DIM_ALLOC];

  int inv_std_stride[CCV_NNC_MAX_DIM_ALLOC];

  int dscale_stride[CCV_NNC_MAX_DIM_ALLOC];

  ccv_nnc_tensor_view_get_stride(a, astride);

  ccv_nnc_tensor_view_get_stride(g, gstride);

  ccv_nnc_tensor_view_get_stride(h, hstride);

  ccv_nnc_tensor_view_get_stride(scale, scale_stride);

  ccv_nnc_tensor_view_get_stride(saved_inv_std, inv_std_stride);

  if (dscale)

    ccv_nnc_tensor_view_get_stride(dscale, dscale_stride);

  // Need to allocate two additional memory:

  // 1. normalized a;

  // 2. scale * inv_std / n;

  assert(!(flags & CCV_NNC_ZERO_MEMORY_ALLOC));

  int x;

  int n = 1;

  for (x = 0; x < CCV_NNC_MAX_DIM + 2; 
x++12
)

    n *= gdim[x];

  for (x = 0; x < CCV_NNC_MAX_DIM + 2; 
x++12
)

    n /= rdim[x];

  int gcount = 1, rcount = 1;

  for (x = 0; x < CCV_NNC_MAX_DIM + 2; 
x++12
)

    gcount *= gdim[x], rcount *= rdim[x];

  float* const ah = (float*)ccv_nnc_stream_context_get_workspace(stream_context, sizeof(float) * gcount * 2 + sizeof(float) * rcount, CCV_TENSOR_CPU_MEMORY);

  float* const gss = ah + gcount; // g * scale * inv_std

  float* const ahgssr = gss + gcount; // ah * gss then reduced to inv_std dimension.

  int i[CCV_NNC_MAX_DIM + 2];

  float* ahp = ah;

  const float* const inv_stdp = saved_inv_std->data.f32;

  const float* const ap = a->data.f32;

  for (i[0] = 0; i[0] < gdim[0]; 
i[0]++12
)

  {

    const float* const ap0 = ap + i[0] * astride[0];

    const float* const inv_stdp0 = rdim[0] == 1 ? 
inv_stdp0
 : inv_stdp + i[0] * inv_std_stride[0];

    for (i[1] = 0; i[1] < gdim[1]; 
i[1]++40
)

    {

      const float* ap1 = ap0 + i[1] * astride[1];

      const float* const inv_stdp1 = rdim[1] == 1 ? inv_stdp0 : 
inv_stdp0 + i[1] * inv_std_stride[1]0
;

      for (i[2] = 0; i[2] < gdim[2]; 
i[2]++144
)

      {

        const float* const inv_stdp2 = rdim[2] == 1 ? inv_stdp1 : 
inv_stdp1 + i[2] * inv_std_stride[2]0
;

        if (rdim[3] == 1)

          
for (x = 0; 144
x < gdim[3]; 
x++1.44k
)

            ahp[x] = ap1[x] * inv_stdp2[0];

        else

          for (x = 0; x < gdim[3]; x++)

            ahp[x] = ap1[x] * inv_stdp2[x];

        ap1 += astride[2];

        ahp += gdim[3];

      }

    }

  }

  if (dscale)

  {

    ccv_nnc_tensor_zero(dscale);

    ahp = ah;

    float* gssp = gss;

    const float* const gp = g->data.f32;

    const float* const scalep = scale->data.f32;

    float* const dscalep = dscale->data.f32;

    for (i[0] = 0; i[0] < gdim[0]; 
i[0]++10
)

    {

      const float* const gp0 = gp + i[0] * gstride[0];

      const float* const inv_stdp0 = rdim[0] == 1 ? 
inv_stdp0
 : inv_stdp + i[0] * inv_std_stride[0];

      const float* const scalep0 = sdim[0] == 1 ? scalep : 
scalep + i[0] * scale_stride[0]0
;

      float* const dscalep0 = sdim[0] == 1 ? dscalep : 
dscalep + i[0] * dscale_stride[0]0
;

      for (i[1] = 0; i[1] < gdim[1]; 
i[1]++36
)

      {

        const float* gp1 = gp0 + i[1] * gstride[1];

        const float* const inv_stdp1 = rdim[1] == 1 ? inv_stdp0 : 
inv_stdp0 + i[1] * inv_std_stride[1]0
;

        const float* const scalep1 = sdim[1] == 1 ? 
scalep00
 : scalep0 + i[1] * scale_stride[1];

        float* const dscalep1 = sdim[1] == 1 ? 
dscalep00
 : dscalep0 + i[1] * dscale_stride[1];

        for (i[2] = 0; i[2] < gdim[2]; 
i[2]++136
)

        {

          const float* const inv_stdp2 = rdim[2] == 1 ? inv_stdp1 : 
inv_stdp1 + i[2] * inv_std_stride[2]0
;

          const float* const scalep2 = sdim[2] == 1 ? 
scalep10
 : scalep1 + i[2] * scale_stride[2];

          float* const dscalep2 = sdim[2] == 1 ? 
dscalep10
 : dscalep1 + i[2] * dscale_stride[2];

          if (sdim[3] == 1)

            for (x = 0; x < gdim[3]; x++)

            {

              gssp[x] = gp1[x] * scalep2[0] * inv_stdp2[rdim[3] == 1 ? 0 : x];

              dscalep2[0] += ahp[x] * gp1[x];

            }

          else

            
for (x = 0; 136
x < gdim[3]; 
x++1.36k
)

            {

              gssp[x] = gp1[x] * scalep2[x] * inv_stdp2[rdim[3] == 1 ? 0 : 
x0
];

              dscalep2[x] += ahp[x] * gp1[x];

            }

          gp1 += gstride[2];

          ahp += gdim[3];

          gssp += gdim[3];

        }

      }

    }

  } else {

    float* gssp = gss;

    const float* const gp = g->data.f32;

    const float* const scalep = scale->data.f32;

    for (i[0] = 0; i[0] < gdim[0]; 
i[0]++2
)

    {

      const float* const gp0 = gp + i[0] * gstride[0];

      const float* const inv_stdp0 = rdim[0] == 1 ? 
inv_stdp0
 : inv_stdp + i[0] * inv_std_stride[0];

      const float* const scalep0 = sdim[0] == 1 ? scalep : 
scalep + i[0] * scale_stride[0]0
;

      for (i[1] = 0; i[1] < gdim[1]; 
i[1]++4
)

      {

        const float* gp1 = gp0 + i[1] * gstride[1];

        const float* const inv_stdp1 = rdim[1] == 1 ? inv_stdp0 : 
inv_stdp0 + i[1] * inv_std_stride[1]0
;

        const float* const scalep1 = sdim[1] == 1 ? 
scalep00
 : scalep0 + i[1] * scale_stride[1];

        for (i[2] = 0; i[2] < gdim[2]; 
i[2]++8
)

        {

          const float* const inv_stdp2 = rdim[2] == 1 ? inv_stdp1 : 
inv_stdp1 + i[2] * inv_std_stride[2]0
;

          const float* const scalep2 = sdim[2] == 1 ? 
scalep10
 : scalep1 + i[2] * scale_stride[2];

          if (sdim[3] == 1)

            for (x = 0; x < gdim[3]; x++)

              gssp[x] = gp1[x] * scalep2[0] * inv_stdp2[rdim[3] == 1 ? 0 : x];

          else

            
for (x = 0; 8
x < gdim[3]; 
x++80
)

              gssp[x] = gp1[x] * scalep2[x] * inv_stdp2[rdim[3] == 1 ? 0 : 
x0
];

          gp1 += gstride[2];

          gssp += gdim[3];

        }

      }

    }

  }

  ahp = ah;

  float* gssp = gss;

  ccv_nnc_tensor_t ahgssrt = ccv_nnc_tensor(ahgssr, saved_inv_std->info, 0);

  ccv_nnc_tensor_zero(&ahgssrt);

  float* const ahgssrp = ahgssr;

  for (i[0] = 0; i[0] < gdim[0]; 
i[0]++12
)

  {

    float* const ahgssrp0 = rdim[0] == 1 ? 
ahgssrp0
 : ahgssrp + i[0] * rdim[1] * rdim[2] * rdim[3];

    for (i[1] = 0; i[1] < gdim[1]; 
i[1]++40
)

    {

      float* const ahgssrp1 = rdim[1] == 1 ? ahgssrp0 : 
ahgssrp0 + i[1] * rdim[2] * rdim[3]0
;

      for (i[2] = 0; i[2] < gdim[2]; 
i[2]++144
)

      {

        float* const ahgssrp2 = rdim[2] == 1 ? ahgssrp1 : 
ahgssrp1 + i[2] * rdim[3]0
;

        if (rdim[3] == 1)

          
for (x = 0; 144
x < gdim[3]; 
x++1.44k
)

            ahgssrp2[0] += ahp[x] * gssp[x];

        else

          for (x = 0; x < gdim[3]; x++)

            ahgssrp2[x] += ahp[x] * gssp[x];

        ahp += gdim[3];

        gssp += gdim[3];

      }

    }

  }

  // Now the part to compute dx (h).

  float* const hp = h->data.f32;

  ahp = ah;

  const float inv_n = 1. / n;

  gssp = gss;

  for (i[0] = 0; i[0] < gdim[0]; 
i[0]++12
)

  {

    float* const hp0 = hp + i[0] * hstride[0];

    const float* const ahgssrp0 = rdim[0] == 1 ? 
ahgssrp0
 : ahgssrp + i[0] * rdim[1] * rdim[2] * rdim[3];

    for (i[1] = 0; i[1] < gdim[1]; 
i[1]++40
)

    {

      float* hp1 = hp0 + i[1] * hstride[1];

      const float* const ahgssrp1 = rdim[1] == 1 ? ahgssrp0 : 
ahgssrp0 + i[1] * rdim[2] * rdim[3]0
;

      for (i[2] = 0; i[2] < gdim[2]; 
i[2]++144
)

      {

        const float* const ahgssrp2 = rdim[2] == 1 ? ahgssrp1 : 
ahgssrp1 + i[2] * rdim[3]0
;

        if (rdim[3] == 1)

          
for (x = 0; 144
x < gdim[3]; 
x++1.44k
)

            hp1[x] = gssp[x] - inv_n * ahp[x] * ahgssrp2[0];

        else

          for (x = 0; x < gdim[3]; x++)

            hp1[x] = gssp[x] - inv_n * ahp[x] * ahgssrp2[x];

        hp1 += hstride[2];

        ahp += gdim[3];

        gssp += gdim[3];

      }

    }

  }

  return CCV_NNC_EXEC_SUCCESS;

}

REGISTER_COMMAND_BACKEND(CCV_NNC_RMSNORM_FORWARD, CCV_NNC_BACKEND_CPU_REF)(ccv_nnc_cmd_backend_registry_t* const registry)

{

  registry->tensor_formats = CCV_TENSOR_FORMAT_NHWC | CCV_TENSOR_FORMAT_NCHW | CCV_TENSOR_FORMAT_CHWN;

  registry->tensor_datatypes = CCV_32F;

  registry->tensor_memory = CCV_TENSOR_CPU_MEMORY;

  registry->algorithms = 1;

  registry->exec = _ccv_nnc_rmsnorm_forw;

}

REGISTER_COMMAND_BACKEND(CCV_NNC_RMSNORM_BACKWARD, CCV_NNC_BACKEND_CPU_REF)(ccv_nnc_cmd_backend_registry_t* const registry)

{

  registry->tensor_formats = CCV_TENSOR_FORMAT_NHWC | CCV_TENSOR_FORMAT_NCHW | CCV_TENSOR_FORMAT_CHWN;

  registry->tensor_datatypes = CCV_32F;

  registry->tensor_memory = CCV_TENSOR_CPU_MEMORY;

  registry->algorithms = 1;

  registry->exec = _ccv_nnc_rmsnorm_back;

}