| { |
| "op": "ai.onnx.Gather", |
| "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] }, |
| "cases": [ |
| { |
| "name": "gather-axis0-f32-4096x128-idx512", |
| "preset": "smoke", |
| "vars": { "rows": 4096, "cols": 128, "indexCount": 512 }, |
| "attrs": { "axis": 0 }, |
| "inputs": { |
| "data": { "shape": [4096, 128], "dtype": "float32", "dist": "normal", "seed": 121, "scale": 0.2 }, |
| "indices": { "shape": [512], "dtype": "int32", "dist": "linearMod", "seed": 121, "step": 17, "mod": 4096 } |
| }, |
| "outputs": { "output": { "shape": [512, 128], "dtype": "float32" } }, |
| "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "args.indexCount * args.cols * 4 * 2" }] } |
| }, |
| { |
| "name": "gather-axis1-f32-1024x512-idx128", |
| "preset": "smoke", |
| "vars": { "rows": 1024, "cols": 512, "indexCount": 128 }, |
| "attrs": { "axis": 1 }, |
| "inputs": { |
| "data": { "shape": [1024, 512], "dtype": "float32", "dist": "normal", "seed": 122, "scale": 0.2 }, |
| "indices": { "shape": [128], "dtype": "int32", "dist": "linearMod", "seed": 122, "step": 17, "mod": 512 } |
| }, |
| "outputs": { "output": { "shape": [1024, 128], "dtype": "float32" } }, |
| "bench": { "primary": true, "metrics": [{ "type": "bandwidth", "value": "args.rows * args.indexCount * 4 * 2" }] } |
| }, |
| { |
| "name": "gather-axis0-u32-bool-like-4096x128-idx512", |
| "preset": "edge", |
| "vars": { "rows": 4096, "cols": 128, "indexCount": 512, "dtype": "uint32" }, |
| "attrs": { "axis": 0 }, |
| "inputs": { |
| "data": { "shape": [4096, 128], "dtype": "uint32", "dist": "linearMod", "mod": 2 }, |
| "indices": { "shape": [512], "dtype": "int32", "dist": "linearMod", "seed": 123, "step": 17, "mod": 4096 } |
| }, |
| "outputs": { "output": { "shape": [512, 128], "dtype": "uint32" } }, |
| "bench": { |
| "metrics": [{ "type": "bandwidth", "value": "args.indexCount * args.cols * dtypeBytes(args.dtype) * 2" }] |
| } |
| }, |
| { |
| "name": "gather-embedding-f16-32000x512-idx1024", |
| "preset": "smoke", |
| "vars": { "rows": 32000, "cols": 512, "indexCount": 1024, "dtype": "float16" }, |
| "attrs": { "axis": 0 }, |
| "inputs": { |
| "data": { "shape": [32000, 512], "dtype": "float16", "dist": "normal", "seed": 124, "scale": 0.2 }, |
| "indices": { "shape": [1024], "dtype": "int32", "dist": "linearMod", "seed": 124, "step": 29, "mod": 32000 } |
| }, |
| "outputs": { "output": { "shape": [1024, 512], "dtype": "float16" } }, |
| "bench": { |
| "metrics": [{ "type": "bandwidth", "value": "args.indexCount * args.cols * dtypeBytes(args.dtype) * 2" }] |
| } |
| }, |
| { |
| "name": "gather-embedding-f16-32000x513-idx1024-unaligned", |
| "preset": "edge", |
| "vars": { "rows": 32000, "cols": 513, "indexCount": 1024, "dtype": "float16" }, |
| "attrs": { "axis": 0 }, |
| "inputs": { |
| "data": { "shape": [32000, 513], "dtype": "float16", "dist": "normal", "seed": 131, "scale": 0.2 }, |
| "indices": { "shape": [1024], "dtype": "int32", "dist": "linearMod", "seed": 131, "step": 29, "mod": 32000 } |
| }, |
| "outputs": { "output": { "shape": [1024, 513], "dtype": "float16" } }, |
| "bench": { |
| "metrics": [{ "type": "bandwidth", "value": "args.indexCount * args.cols * dtypeBytes(args.dtype) * 2" }] |
| } |
| }, |
| { |
| "name": "gather-axis0-f32-4096x127-idx512-unaligned", |
| "preset": "edge", |
| "vars": { "rows": 4096, "cols": 127, "indexCount": 512, "dtype": "float32" }, |
| "attrs": { "axis": 0 }, |
| "inputs": { |
| "data": { "shape": [4096, 127], "dtype": "float32", "dist": "normal", "seed": 132, "scale": 0.2 }, |
| "indices": { "shape": [512], "dtype": "int32", "dist": "linearMod", "seed": 132, "step": 17, "mod": 4096 } |
| }, |
| "outputs": { "output": { "shape": [512, 127], "dtype": "float32" } }, |
| "bench": { |
| "metrics": [{ "type": "bandwidth", "value": "args.indexCount * args.cols * dtypeBytes(args.dtype) * 2" }] |
| } |
| }, |
| { |
| "name": "gather-axis0-f32-50000x16-idx2048-narrow-rows", |
| "preset": "edge", |
| "vars": { "rows": 50000, "cols": 16, "indexCount": 2048, "dtype": "float32" }, |
| "attrs": { "axis": 0 }, |
| "inputs": { |
| "data": { "shape": [50000, 16], "dtype": "float32", "dist": "normal", "seed": 133, "scale": 0.2 }, |
| "indices": { "shape": [2048], "dtype": "int32", "dist": "linearMod", "seed": 133, "step": 29, "mod": 50000 } |
| }, |
| "outputs": { "output": { "shape": [2048, 16], "dtype": "float32" } }, |
| "bench": { |
| "metrics": [{ "type": "bandwidth", "value": "args.indexCount * args.cols * dtypeBytes(args.dtype) * 2" }] |
| } |
| }, |
| { |
| "name": "gather-axis0-f32-1000000x128-idx131072-saturated-grid", |
| "preset": "edge", |
| "vars": { "rows": 1000000, "cols": 128, "indexCount": 131072, "dtype": "float32" }, |
| "attrs": { "axis": 0 }, |
| "inputs": { |
| "data": { "shape": [1000000, 128], "dtype": "float32", "dist": "normal", "seed": 134, "scale": 0.2 }, |
| "indices": { "shape": [131072], "dtype": "int32", "dist": "linearMod", "seed": 134, "step": 7, "mod": 1000000 } |
| }, |
| "outputs": { "output": { "shape": [131072, 128], "dtype": "float32" } }, |
| "bench": { |
| "metrics": [{ "type": "bandwidth", "value": "args.indexCount * args.cols * dtypeBytes(args.dtype) * 2" }] |
| } |
| } |
| ] |
| } |
|
|
