{% if usesF16 %}
enable f16;
{% endif %}
{{ env.wgsl.resourceDeclarations }}

// com.microsoft.EmbedLayerNormalization, embedding-sum pass.
//   sum = word_embedding[input_ids[t]] + position_embedding[p(t)]
//         (+ segment_embedding[segment_ids[t] or 0])
// One workgroup writes one token's row of the summed embedding into `output`,
// where the normalization pass reads it back. The row bases -- which word,
// position and segment rows this token gathers -- are the same for every element
// of the row, so they are resolved once per workgroup.
// The sum is materialized in the tensor type -- it is also the optional
// `embedding_sum` output -- so the statistics the next pass computes describe
// the values this op emits rather than an unobservable wider intermediate.
const HIDDEN: u32 = {{ hidden }}u;
{% if not hasPositionIds or broadcastPositionIds %}
// Read only where the position row comes from the position within the sequence:
// a per-token position_ids indexes directly instead.
const SEQUENCE: u32 = {{ sequenceLength }}u;
{% endif %}
const WG: u32 = {{ workgroupSize }}u;
// Defensive row clamps keep an invalid id from reading outside an embedding
// table. Request validation rejects such ids before a valid dispatch.
const WORD_ROWS: u32 = {{ wordRows }}u;
const POSITION_ROWS: u32 = {{ positionRows }}u;
{% if hasSegment and hasSegmentIds %}
const SEGMENT_ROWS: u32 = {{ segmentRows }}u;
{% endif %}

@compute @workgroup_size(WG, 1, 1)
fn main(@builtin(workgroup_id) wg: vec3<u32>,
        @builtin(num_workgroups) nwg: vec3<u32>,
        @builtin(local_invocation_id) lid: vec3<u32>) {
  // 2D-folded row index: wg.y carries the high bits past the
  // maxComputeWorkgroupsPerDimension dispatch limit.
  let token = wg.x + wg.y * nwg.x;
  if (token >= params.tokens) {
    return;
  }
  let tid = lid.x;
  let base = token * HIDDEN;

  let word_row = min(u32(max(input_ids[token], 0)), WORD_ROWS - 1u) * HIDDEN;
{% if hasPositionIds %}
{% if broadcastPositionIds %}
  // A position_ids of shape (1, sequence_length) is shared by every batch, so
  // the row index drops back to the position within the sequence.
  let position_id = position_ids[token % SEQUENCE];
{% else %}
  let position_id = position_ids[token];
{% endif %}
  let position_row = min(u32(max(position_id, 0)), POSITION_ROWS - 1u) * HIDDEN;
{% else %}
  let position_row = min(token % SEQUENCE, POSITION_ROWS - 1u) * HIDDEN;
{% endif %}
{% if hasSegment %}
{% if hasSegmentIds %}
  let segment_row = min(u32(max(segment_ids[token], 0)), SEGMENT_ROWS - 1u) * HIDDEN;
{% else %}
  let segment_row = 0u;
{% endif %}
{% endif %}

  for (var i = tid; i < HIDDEN; i = i + WG) {
    // Materialize each addition in T; association is observable for f16.
{% if hasSegment and scalar == "f16" %}
    var stored = {{ scalar }}(f32(word_embedding[word_row + i]) + f32(segment_embedding[segment_row + i]));
    stored = {{ scalar }}(f32(stored) + f32(position_embedding[position_row + i]));
{% else %}
    var stored = {{ scalar }}(f32(word_embedding[word_row + i]) + f32(position_embedding[position_row + i]));
{% if hasSegment %}
    stored = {{ scalar }}(f32(stored) + f32(segment_embedding[segment_row + i]));
{% endif %}
{% endif %}
{% if writeEmbeddingSum %}
    embedding_sum[base + i] = stored;
{% endif %}
    output[base + i] = stored;
  }
}