File size: 6,325 Bytes

// Pass 1 of parallel DynamicQuantizeLinear. Regular mode reduces one contiguous
// WG * EPT block per workgroup. Grid-stride mode caps the partial count and has
// every lane revisit the tensor at grid-sized strides. Both modes use the same
// subgroup/tree combine and write one min/max pair per workgroup. The subgroup
// combine publishes one shared-memory slot per invocation (its subgroup pair
// from the elected lane, the neutral 0.0 from every other lane) and every
// subgroup folds all of them.
//
// The local min/max start at 0.0 because the ONNX DynamicQuantizeLinear range
// always includes zero; out-of-range lanes contribute the same neutral value.
// f32 min/max is order-independent, so subgroup and workgroup-tree combinations
// produce the same extrema. Without subgroups, a full workgroup tree combines
// the per-invocation partials.
{% set gridStride = gridStride if gridStride is defined else false %}
{% if useSubgroups %}
enable subgroups;
{% endif %}
{{ env.wgsl.resourceDeclarations }}

const WG: u32 = {{ workgroupSize }}u;
{% if not gridStride and not vec4 %}
const EPT: u32 = {{ elemsPerThread }}u;

{% endif %}
var<workgroup> wgMin: array<f32, WG>;
var<workgroup> wgMax: array<f32, WG>;

@compute @workgroup_size(WG, 1, 1)
fn main(@builtin(workgroup_id) wg: vec3<u32>,
        @builtin(local_invocation_id) lid: vec3<u32>{{ "," if gridStride or useSubgroups else "" }}
{% if gridStride %}
        @builtin(num_workgroups) nwg: vec3<u32>{{ "," if useSubgroups else "" }}
{% endif %}
{% if useSubgroups %}
        @builtin(subgroup_invocation_id) sgLane: u32,
        @builtin(subgroup_size) sgSize: u32
{% endif %}
) {
  let tid = lid.x;
{% if gridStride %}
  let blk = wg.x;
  // The capped x-only dispatch is the full logical grid. Each workgroup writes
  // exactly one partial, matching finalize's numPartials uniform.
  let gridThreads = nwg.x * WG;
{% else %}
  // 2D-folded block index: wg.y carries the high bits past the device's
  // per-dimension dispatch limit. The whole over-dispatched workgroup returns
  // before barriers or scratch stores.
  // blk/numBlocks are workgroup-uniform, so the return can't split the barrier.
  let blk = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
{% if vec4 %}
  let numBlocks = (params.count / 4u + WG - 1u) / WG;
{% else %}
  let numBlocks = (params.count + WG * EPT - 1u) / (WG * EPT);
{% endif %}
  if (blk >= numBlocks) { return; }
{% endif %}
  var localMin = 0.0;
  var localMax = 0.0;
{% if gridStride %}
{% if vec4 %}
  // Adjacent lanes read adjacent vec4s on each grid-stride iteration.
  let count4 = params.count / 4u;
  for (var idx4 = blk * WG + tid; idx4 < count4; idx4 = idx4 + gridThreads) {
    let v = x[idx4];
    localMin = min(localMin, min(min(v.x, v.y), min(v.z, v.w)));
    localMax = max(localMax, max(max(v.x, v.y), max(v.z, v.w)));
  }
{% else %}
  for (var idx = blk * WG + tid; idx < params.count; idx = idx + gridThreads) {
    let value = x[idx];
    localMin = min(localMin, value);
    localMax = max(localMax, value);
  }
{% endif %}
{% else %}
{% if vec4 %}
  // Each thread reads one contiguous vec4; this path assumes params.count is
  // divisible by 4.
  let count4 = params.count / 4u;
  let idx4 = blk * WG + tid;
  if (idx4 < count4) {
    let v = x[idx4];
    localMin = min(localMin, min(min(v.x, v.y), min(v.z, v.w)));
    localMax = max(localMax, max(max(v.x, v.y), max(v.z, v.w)));
  }
{% else %}
  let base = blk * WG * EPT;
  for (var e = 0u; e < EPT; e = e + 1u) {
    let idx = base + e * WG + tid;
    if (idx < params.count) {
      let v = x[idx];
      localMin = min(localMin, v);
      localMax = max(localMax, v);
    }
  }
{% endif %}
{% endif %}
{% if useSubgroups %}
  let sgMin = subgroupMin(localMin);
  let sgMax = subgroupMax(localMax);
  // Cross-subgroup fold that assumes nothing about which invocations share a
  // subgroup or how many subgroups there are (it does require the uniform
  // control flow this entry point already has): every invocation owns the slot
  // at its own index, the elected lane
  // publishes its subgroup pair there and every other lane publishes 0.0, which
  // is an exact identity here because every lane's local range already includes
  // zero (so every published minimum is <= 0 and every maximum >= 0). Each
  // subgroup then folds all WG slots — lane `sgLane` walks slots sgLane,
  // sgLane + sgSize, ... — and one more collective merges the lane partials, so
  // every slot is merged exactly once at any legal width. min/max is commutative
  // and associative, so the fold order does not change the result.
  var totalMin = sgMin;
  var totalMax = sgMax;
  // A one-subgroup workgroup is already fully reduced by the collectives above.
  // The test reads the `subgroup_size` builtin, which is uniform; a collective's
  // result is not uniform to WGSL's analysis and may not guard a barrier.
  if (sgSize != WG) {
    // The coordinates are the BUILTINS, never `subgroupExclusiveAdd(1u)` /
    // `subgroupAdd(1u)`: those agree with them under this contract, but a driver
    // in the wild answers a claim derived from them with zero and leaves most of
    // the workgroup's slots unclaimed. Every lane of this entry point is active
    // here, so the full subgroup width is the right stride.
    let rank = sgLane;
    let count = sgSize;
    let leader = rank == 0u;
    wgMin[tid] = select(0.0, sgMin, leader);
    wgMax[tid] = select(0.0, sgMax, leader);
    workgroupBarrier();
    var foldMin = 0.0;
    var foldMax = 0.0;
    for (var i = rank; i < WG; i = i + count) {
      foldMin = min(foldMin, wgMin[i]);
      foldMax = max(foldMax, wgMax[i]);
    }
    totalMin = subgroupMin(foldMin);
    totalMax = subgroupMax(foldMax);
  }
  if (tid == 0u) {
    partial_min[blk] = totalMin;
    partial_max[blk] = totalMax;
  }
{% else %}
  // No-subgroups fallback: full workgroup-tree min/max over per-thread partials.
  wgMin[tid] = localMin;
  wgMax[tid] = localMax;
  workgroupBarrier();
  for (var step = WG >> 1u; step > 0u; step = step >> 1u) {
    if (tid < step) {
      wgMin[tid] = min(wgMin[tid], wgMin[tid + step]);
      wgMax[tid] = max(wgMax[tid], wgMax[tid + step]);
    }
    workgroupBarrier();
  }
  if (tid == 0u) {
    partial_min[blk] = wgMin[0];
    partial_max[blk] = wgMax[0];
  }
{% endif %}
}