Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.Max

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
f544dde
·
verified ·
1 Parent(s): 743d5e7

sync c928d21e6cc1

Browse files
Files changed (3) hide show
  1. build/webgpu/manifest.json +12 -8
  2. build/webgpu/metadata.json +4 -4
  3. build/webgpu/minmax-broadcast.wgsl.jinja +18 -8
build/webgpu/manifest.json CHANGED
@@ -368,11 +368,12 @@
368
  "bRank": "ranks.B",
369
  "yRank": "ranks.y",
370
  "hasC": "false",
371
- "op": "\"max\""
 
372
  }
373
  },
374
  "bindings": "scalarTwo",
375
- "dispatch": { "threads": "numel(shapes.y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
376
  }
377
  ]
378
  },
@@ -397,11 +398,12 @@
397
  "cRank": "ranks.C",
398
  "yRank": "ranks.y",
399
  "hasC": "true",
400
- "op": "\"max\""
 
401
  }
402
  },
403
  "bindings": "scalarThree",
404
- "dispatch": { "threads": "numel(shapes.y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
405
  }
406
  ]
407
  },
@@ -429,11 +431,12 @@
429
  "op": "\"max\"",
430
  "dShape": "shapes.D",
431
  "dRank": "ranks.D",
432
- "hasD": "true"
 
433
  }
434
  },
435
  "bindings": "scalarFour",
436
- "dispatch": { "threads": "numel(shapes.y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
437
  }
438
  ]
439
  },
@@ -465,11 +468,12 @@
465
  "hasE": "true",
466
  "extraInputs": "[\"c\", \"d\", \"e\"]",
467
  "eRank": "ranks.E",
468
- "eShape": "shapes.E"
 
469
  }
470
  },
471
  "bindings": "scalarFive",
472
- "dispatch": { "threads": "numel(shapes.y)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
473
  }
474
  ]
475
  }
 
368
  "bRank": "ranks.B",
369
  "yRank": "ranks.y",
370
  "hasC": "false",
371
+ "op": "\"max\"",
372
+ "itemsPerInvocation": 4
373
  }
374
  },
375
  "bindings": "scalarTwo",
376
+ "dispatch": { "threads": "ceilDiv(numel(shapes.y), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
377
  }
378
  ]
379
  },
 
398
  "cRank": "ranks.C",
399
  "yRank": "ranks.y",
400
  "hasC": "true",
401
+ "op": "\"max\"",
402
+ "itemsPerInvocation": 4
403
  }
404
  },
405
  "bindings": "scalarThree",
406
+ "dispatch": { "threads": "ceilDiv(numel(shapes.y), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
407
  }
408
  ]
409
  },
 
431
  "op": "\"max\"",
432
  "dShape": "shapes.D",
433
  "dRank": "ranks.D",
434
+ "hasD": "true",
435
+ "itemsPerInvocation": 4
436
  }
437
  },
438
  "bindings": "scalarFour",
439
+ "dispatch": { "threads": "ceilDiv(numel(shapes.y), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
440
  }
441
  ]
442
  },
 
468
  "hasE": "true",
469
  "extraInputs": "[\"c\", \"d\", \"e\"]",
470
  "eRank": "ranks.E",
471
+ "eShape": "shapes.E",
472
+ "itemsPerInvocation": 4
473
  }
474
  },
475
  "bindings": "scalarFive",
476
+ "dispatch": { "threads": "ceilDiv(numel(shapes.y), 4)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
477
  }
478
  ]
479
  }
build/webgpu/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "ai.onnx.Max",
3
- "id": "_ai_onnx_max_webgpu_07f8f49",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
@@ -9,12 +9,12 @@
9
  "files": {
10
  "bench.json": "RNgdFMBNRNyK6rosPqaYVu6EZMBRczN8MaT1PXgx5qs=",
11
  "datamove-elementwise-copy.wgsl.jinja": "J5yC2bAddPiP+odXLgVGS3TJ9jeNsfRTvedKrj/fhZg=",
12
- "manifest.json": "5ogrMJ0Z1/7/SjIm5WqLTPzgTFa4TQ1h0VsARPggnLI=",
13
- "minmax-broadcast.wgsl.jinja": "PfO5v9vE6fpTpa/V/4d6sfxmLxTnjx65wMuLfCLBxQI=",
14
  "minmax-vec4.wgsl.jinja": "8EEl0UubHzIfW8gQD2Ra6SzRpRL4uv8ao7U3FgQxm+g=",
15
  "test.json": "3Te1JTq0Cf0qFrtb5eAvL3A5SvDrmI+XmzWbln4BJ5s="
16
  }
17
  },
18
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
19
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Max" }
20
  }
 
1
  {
2
  "name": "ai.onnx.Max",
3
+ "id": "_ai_onnx_max_webgpu_119ddc2",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
 
9
  "files": {
10
  "bench.json": "RNgdFMBNRNyK6rosPqaYVu6EZMBRczN8MaT1PXgx5qs=",
11
  "datamove-elementwise-copy.wgsl.jinja": "J5yC2bAddPiP+odXLgVGS3TJ9jeNsfRTvedKrj/fhZg=",
12
+ "manifest.json": "8zQFEWWGskKzn+PwXWpjtiuiivgDxLObzk65ZMx6oh8=",
13
+ "minmax-broadcast.wgsl.jinja": "gmj/AG7mXiBYYsC9j8N1ZoNVUBuzM9M56hvIAR1qpw4=",
14
  "minmax-vec4.wgsl.jinja": "8EEl0UubHzIfW8gQD2Ra6SzRpRL4uv8ao7U3FgQxm+g=",
15
  "test.json": "3Te1JTq0Cf0qFrtb5eAvL3A5SvDrmI+XmzWbln4BJ5s="
16
  }
17
  },
18
+ "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
19
  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.Max" }
20
  }
build/webgpu/minmax-broadcast.wgsl.jinja CHANGED
@@ -1,3 +1,19 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  {% set extraInputs = source.extraInputs if source.extraInputs is defined else (["c"] if source.hasC else []) + (["d"] if source.hasD else []) %}
2
  {% if usesF16 %}
3
  enable f16;
@@ -85,14 +101,7 @@ fn {{ fn_name }}({% if out_numel != 0 and op_numel != 1 %}out_index: u32{% endif
85
  {% endfor %}
86
 
87
 
88
- @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
89
- fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
90
- // 2D-folded flat index: gid.y carries the high bits past the
91
- // maxComputeWorkgroupsPerDimension limit.
92
- let i = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
93
- if (i >= params.count) {
94
- return;
95
- }
96
  {% if scalar == "i32" or scalar == "u32" %}
97
  let av = a[{{ broadcast_offset_call("a_offset", source.aShape, source.yShape, "i") }}];
98
  let bv = b[{{ broadcast_offset_call("b_offset", source.bShape, source.yShape, "i") }}];
@@ -119,4 +128,5 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
119
  {% endfor %}
120
  y[i] = {{ scalar }}(out);
121
  {% endif %}
 
122
  }
 
1
+ {% macro flat_tail_open() %}
2
+ @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
3
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
4
+ // 2D-folded flat index: gid.y carries the high bits when the element count exceeds the
5
+ // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills the rest into y).
6
+ let invocation = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
7
+ // Tail-safe scalar x4 keeps vector-like dispatch density without requiring
8
+ // the logical tensor length (or its storage binding) to be vec4 aligned.
9
+ let begin = invocation * {{ source.itemsPerInvocation }}u;
10
+ let end = min(begin + {{ source.itemsPerInvocation }}u, params.count);
11
+ for (var i = begin; i < end; i = i + 1u) {
12
+ {%- endmacro %}
13
+ {% macro flat_tail_close() %}
14
+ }
15
+ {% endmacro %}
16
+
17
  {% set extraInputs = source.extraInputs if source.extraInputs is defined else (["c"] if source.hasC else []) + (["d"] if source.hasD else []) %}
18
  {% if usesF16 %}
19
  enable f16;
 
101
  {% endfor %}
102
 
103
 
104
+ {{ flat_tail_open() }}
 
 
 
 
 
 
 
105
  {% if scalar == "i32" or scalar == "u32" %}
106
  let av = a[{{ broadcast_offset_call("a_offset", source.aShape, source.yShape, "i") }}];
107
  let bv = b[{{ broadcast_offset_call("b_offset", source.bShape, source.yShape, "i") }}];
 
128
  {% endfor %}
129
  y[i] = {{ scalar }}(out);
130
  {% endif %}
131
+ {{ flat_tail_close() -}}
132
  }