ember_ir::cpu Namespace Reference

Functions
void	conv1d_causal_row (const float *x, const float *w, const float *bias, float *state, float *out, int C, int K)
void	selective_scan_row (const float *x, const float *dt, const float *B, const float *C, const float *A, const float *d, float *s, float *y, int D, int N, bool softplus)
void	sub (const float *a, const float *b, float *out, size_t n)
void	relu (const float *x, float *out, size_t n)
void	sigmoid (const float *x, float *out, size_t n)
void	affine (const float *x, float *out, size_t n, float scale, float shift)
void	mor_route (const float *l, int R, int &depth, float &gate)
void	mor_gate (const float *x_in, float *x, size_t n, float g)
void	scale (const float *x, const float *g, const float *b, float *out, int T, int w)
void	mul_col (const float *x, const float *c, float *out, int T, int w)
void	rms_row (const float *x, float *out, int T, int w, float eps)
void	depth_attend (const std::vector< const float * > &sources, const float *w, float *out, int T, int d, float eps)
void	mla_attention (const float *qn, const float *qr, const float *kv, float *out, int T, const std::vector< std::vector< int > > &rows, int H, int r, int dn, int dr, int dv, const float *Kb, const float *Vb, float scale)
void	split_groups (const float *x, int T, const std::vector< int > &sizes, int g, const std::vector< float * > &outs)

Function Documentation

affine()

void ember_ir::cpu::affine	(	const float *	x,
		float *	out,
		size_t	n,
		float	scale,
		float	shift
	)

Definition at line 26 of file graph_ir_ops.cpp.

                                                                            {
  for (size_t i = 0; i < n; ++i) {
    const float p = x[i] * scale;
    out[i] = p + shift;
  }
}

References j, n, out, p, and scale.

conv1d_causal_row()

void ember_ir::cpu::conv1d_causal_row	(	const float *	x,
		const float *	w,
		const float *	bias,
		float *	state,
		float *	out,
		int	C,
		int	K
	)

Definition at line 10 of file graph_ir_mamba.cpp.

                                                                                                                  {
  for (int c = 0; c < C; ++c) {
    float acc = 0.0f;
    for (int k = 0; k < K; ++k) {
      const float v = k < K - 1 ? state[static_cast<size_t>(k) * C + c] : x[c];
      acc = acc + w[static_cast<size_t>(c) * K + k] * v;
    }
    out[c] = bias ? acc + bias[c] : acc;
  }
  if (K > 1) {  // the state becomes the last K - 1 rows of [state; x]
    std::memmove(state, state + C, static_cast<size_t>(K - 2) * C * sizeof(float));
    std::memcpy(state + static_cast<size_t>(K - 2) * C, x, static_cast<size_t>(C) * sizeof(float));
  }
}

References c, j, k, out, state, and v.

depth_attend()

void ember_ir::cpu::depth_attend	(	const std::vector< const float * > &	sources,
		const float *	w,
		float *	out,
		int	T,
		int	d,
		float	eps
	)

Definition at line 81 of file graph_ir_ops.cpp.

                                                                                                               {
  const size_t S = sources.size();
  const size_t n = static_cast<size_t>(T) * d;
  if (S == 1) {  // softmax over one source is the identity mixture: s_0 bit for bit
    std::memcpy(out, sources[0], n * sizeof(float));
    return;
  }
  std::vector<float> a(S);
  for (int t = 0; t < T; ++t) {
    const size_t row = static_cast<size_t>(t) * d;
    // Logits: w . rmsnorm(s_i) (ops::rmsnorm with unit weight, ops::mul_row,
    // then the matmul with a ones column).
    for (size_t i = 0; i < S; ++i) {
      const float* s = sources[i] + row;
      float ss = 0.0f;
      for (int j = 0; j < d; ++j) {
        const float sq = s[j] * s[j];
        ss += sq;
      }
      ss /= static_cast<float>(d);
      const float r = 1.0f / std::sqrt(ss + eps);
      float l = 0.0f;
      for (int j = 0; j < d; ++j) {
        const float normed = s[j] * r;
        const float p = normed * w[j];
        l += p;
      }
      a[i] = l;
    }
    // Softmax over the sources (Matrix::apply_softmax): max-shifted, z in
    // ascending i, a_i = e_i / z.
    float mx = -std::numeric_limits<float>::infinity();
    for (size_t i = 0; i < S; ++i) mx = std::max(mx, a[i]);
    float z = 0.0f;
    for (size_t i = 0; i < S; ++i) {
      a[i] = std::exp(a[i] - mx);
      z += a[i];
    }
    for (size_t i = 0; i < S; ++i) a[i] /= z;
    // out = s_0 a_0 + s_1 a_1 + ... (ops::mul_col per source, ops::add in order).
    float* o = out + row;
    const float* s0 = sources[0] + row;
    for (int j = 0; j < d; ++j) o[j] = s0[j] * a[0];
    for (size_t i = 1; i < S; ++i) {
      const float* s = sources[i] + row;
      for (int j = 0; j < d; ++j) {
        const float p = s[j] * a[i];
        o[j] = o[j] + p;
      }
    }
  }
}

References a, j, n, out, and p.

mla_attention()

void ember_ir::cpu::mla_attention	(	const float *	qn,
		const float *	qr,
		const float *	kv,
		float *	out,
		int	T,
		const std::vector< std::vector< int > > &	rows,
		int	H,
		int	r,
		int	dn,
		int	dr,
		int	dv,
		const float *	Kb,
		const float *	Vb,
		float	scale
	)

Definition at line 134 of file graph_ir_ops.cpp.

                                                 {
  const size_t w = static_cast<size_t>(r) + dr;
#pragma omp parallel for schedule(dynamic)
  for (int t = 0; t < T; ++t) {
    const std::vector<int>& vis = rows[t];
    std::vector<float> qa(static_cast<size_t>(r)), u(static_cast<size_t>(r)), s(vis.size());
    for (int h = 0; h < H; ++h) {
      const float* qnh = qn + (static_cast<size_t>(t) * H + h) * dn;
      const float* qrh = qr + (static_cast<size_t>(t) * H + h) * dr;
      const float* K = Kb + static_cast<size_t>(h) * r * dn;
      const float* V = Vb + static_cast<size_t>(h) * dv * r;
      for (int i = 0; i < r; ++i) {  // qa = K[h] qn_h
        float acc = 0.0f;
        for (int k = 0; k < dn; ++k) {
          const float p = K[static_cast<size_t>(i) * dn + k] * qnh[k];
          acc += p;
        }
        qa[i] = acc;
      }
      float mx = -std::numeric_limits<float>::infinity();
      for (size_t j = 0; j < vis.size(); ++j) {
        const float* c = kv + static_cast<size_t>(vis[j]) * w;
        float a = 0.0f, b = 0.0f;
        for (int i = 0; i < r; ++i) {
          const float p = qa[i] * c[i];
          a += p;
        }
        for (int e = 0; e < dr; ++e) {
          const float p = qrh[e] * c[r + e];
          b += p;
        }
        const float sum = a + b;
        s[j] = sum * scale;
        mx = std::max(mx, s[j]);
      }
      float z = 0.0f;
      for (size_t j = 0; j < vis.size(); ++j) {
        s[j] = std::exp(s[j] - mx);
        z += s[j];
      }
      std::fill(u.begin(), u.end(), 0.0f);
      for (size_t j = 0; j < vis.size(); ++j) {
        const float a = s[j] / z;
        const float* c = kv + static_cast<size_t>(vis[j]) * w;
        for (int i = 0; i < r; ++i) {
          const float p = a * c[i];
          u[i] += p;
        }
      }
      float* o = out + (static_cast<size_t>(t) * H + h) * dv;
      for (int k = 0; k < dv; ++k) {  // out = V[h] u
        float acc = 0.0f;
        for (int i = 0; i < r; ++i) {
          const float p = V[static_cast<size_t>(k) * r + i] * u[i];
          acc += p;
        }
        o[k] = acc;
      }
    }
  }
}

References a, c, j, k, out, p, and scale.

mor_gate()

void ember_ir::cpu::mor_gate	(	const float *	x_in,
		float *	x,
		size_t	n,
		float	g
	)

Definition at line 43 of file graph_ir_ops.cpp.

                                                              {
  for (size_t i = 0; i < n; ++i) {
    const float diff = x[i] - x_in[i];
    const float p = g * diff;
    x[i] = x_in[i] + p;
  }
}

References j, n, and p.

mor_route()

void ember_ir::cpu::mor_route	(	const float *	l,
		int	R,
		int &	depth,
		float &	gate
	)

Definition at line 33 of file graph_ir_ops.cpp.

                                                               {
  int j_max = 0;
  for (int j = 1; j < R; ++j)
    if (l[j] > l[j_max]) j_max = j;  // strict: a tie keeps the lower index
  float z = 0.0f;
  for (int j = 0; j < R; ++j) z += std::exp(l[j] - l[j_max]);
  depth = j_max + 1;
  gate = std::exp(l[j_max] - l[j_max]) / z;
}

References j.

mul_col()

void ember_ir::cpu::mul_col	(	const float *	x,
		const float *	c,
		float *	out,
		int	T,
		int	w
	)

Definition at line 61 of file graph_ir_ops.cpp.

                                                                       {
  for (int t = 0; t < T; ++t) {
    const float ct = c[t];
    const size_t r = static_cast<size_t>(t) * w;
    for (int j = 0; j < w; ++j) out[r + j] = x[r + j] * ct;
  }
}

References c, j, and out.

relu()

void ember_ir::cpu::relu	(	const float *	x,
		float *	out,
		size_t	n
	)

Definition at line 18 of file graph_ir_ops.cpp.

                                                {
  for (size_t i = 0; i < n; ++i) out[i] = std::max(0.0f, x[i]);
}

References j, n, and out.

rms_row()

void ember_ir::cpu::rms_row	(	const float *	x,
		float *	out,
		int	T,
		int	w,
		float	eps
	)

Definition at line 69 of file graph_ir_ops.cpp.

                                                                  {
  for (int t = 0; t < T; ++t) {
    const float* r = x + static_cast<size_t>(t) * w;
    float ss = 0.0f;
    for (int j = 0; j < w; ++j) {
      const float sq = r[j] * r[j];
      ss += sq;
    }
    out[t] = std::sqrt(ss / static_cast<float>(w) + eps);
  }
}

References j, and out.

scale()

void ember_ir::cpu::scale	(	const float *	x,
		const float *	g,
		const float *	b,
		float *	out,
		int	T,
		int	w
	)

Definition at line 51 of file graph_ir_ops.cpp.

                                                                                     {
  for (int t = 0; t < T; ++t) {
    const size_t r = static_cast<size_t>(t) * w;
    if (b)
      for (int j = 0; j < w; ++j) out[r + j] = x[r + j] * g[j] + b[j];
    else
      for (int j = 0; j < w; ++j) out[r + j] = x[r + j] * g[j];
  }
}

References j, and out.

selective_scan_row()

void ember_ir::cpu::selective_scan_row	(	const float *	x,
		const float *	dt,
		const float *	B,
		const float *	C,
		const float *	A,
		const float *	d,
		float *	s,
		float *	y,
		int	D,
		int	N,
		bool	softplus
	)

Definition at line 25 of file graph_ir_mamba.cpp.

                                                                                         {
  for (int c = 0; c < D; ++c) {
    const float delta = softplus ? (dt[c] <= 20.0f ? std::log1p(std::exp(dt[c])) : dt[c]) : dt[c];
    const float xd = x[c] * delta;
    float* sc = s + static_cast<size_t>(c) * N;
    const float* ac = A + static_cast<size_t>(c) * N;
    float acc = 0.0f;
    for (int n = 0; n < N; ++n) {
      const float decay = sc[n] * std::exp(delta * ac[n]);
      sc[n] = decay + B[n] * xd;
      acc = acc + sc[n] * C[n];
    }
    y[c] = d ? acc + d[c] * x[c] : acc;
  }
}

References c, j, and n.

sigmoid()

void ember_ir::cpu::sigmoid	(	const float *	x,
		float *	out,
		size_t	n
	)

Definition at line 22 of file graph_ir_ops.cpp.

                                                   {
  for (size_t i = 0; i < n; ++i) out[i] = 1.0f / (1.0f + std::exp(-x[i]));
}

References j, n, and out.

split_groups()

void ember_ir::cpu::split_groups	(	const float *	x,
		int	T,
		const std::vector< int > &	sizes,
		int	g,
		const std::vector< float * > &	outs
	)

Definition at line 198 of file graph_ir_ops.cpp.

                                                                                                            {
  int group = 0;
  for (int s : sizes) group += s;
  const size_t width = static_cast<size_t>(group) * g;
  for (int t = 0; t < T; ++t) {
    const float* row = x + static_cast<size_t>(t) * width;
    for (int q = 0; q < g; ++q) {
      int col = 0;
      for (size_t k = 0; k < sizes.size(); ++k) {
        float* dst = outs[k] + (static_cast<size_t>(t) * g + q) * sizes[k];
        std::memcpy(dst, row + static_cast<size_t>(q) * group + col, static_cast<size_t>(sizes[k]) * sizeof(float));
        col += sizes[k];
      }
    }
  }
}

References col, j, k, and sizes.

sub()

void ember_ir::cpu::sub	(	const float *	a,
		const float *	b,
		float *	out,
		size_t	n
	)

Definition at line 14 of file graph_ir_ops.cpp.

                                                               {
  for (size_t i = 0; i < n; ++i) out[i] = a[i] - b[i];
}

References a, j, n, and out.