{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
  // per-axis workgroup fold width.
  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
  if ({{ name }} >= {{ bound }}) {
    return;
  }{% endmacro %}
{{ env.wgsl.resourceDeclarations }}

// com.microsoft.EmbedLayerNormalization, mask_index pass. With a mask, return
// the first zero position or the sequence length when every position is set.
// Without a mask, initialize the optional output to zero.
{% if hasMask %}
const SEQUENCE: u32 = {{ sequenceLength }}u;
{% endif %}

@compute @workgroup_size({{ maskWorkgroupSize }}, 1, 1)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
{{ flat_index_2d(maskWorkgroupSize, "batch", "params.batch") }}
  var first_zero: i32 = 0;
{% if hasMask %}
  first_zero = i32(SEQUENCE);
  let base = batch * SEQUENCE;
  for (var s: u32 = 0u; s < SEQUENCE; s = s + 1u) {
    if (mask[base + s] == 0) {
      first_zero = i32(s);
      break;
    }
  }
{% endif %}
  mask_index[batch] = first_zero;
}