| {{ 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; |
| // IDs must index valid embedding rows. Defensive clamps prevent out-of-bounds |
| // reads if that precondition is violated. |
| 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(local_invocation_id) lid: vec3<u32>) { |
| // 2D-folded row index: wg.y carries the high bits past the |
| // per-axis dispatch fold width. |
| let token = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u; |
| 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; |
| } |
| } |
| |
