diff --git a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp index b475a6cd999..5374390722b 100644 --- a/backends/webgpu/runtime/WebGPUShaderRegistry.cpp +++ b/backends/webgpu/runtime/WebGPUShaderRegistry.cpp @@ -17,6 +17,10 @@ #include #include #include +#include +#include +#include +#include #include #include #include @@ -50,7 +54,6 @@ #include #include #include -#include #include #include #include @@ -65,20 +68,17 @@ #include #include #include -#include -#include +#include +#include #include -#include #include #include #include -#include #include #include #include #include #include -#include #include #include #include diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h index fe2315df30f..dcd2b2e46d8 100644 --- a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_floor_divide_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_floor_divide.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: baf71d277da79389315a6b96b439e7f0a55842e8288283f2af121f84536b3af3 inline constexpr const char* kBinaryFloorDivideWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h similarity index 97% rename from backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h index 88d9614ba59..c3f5f76c81f 100644 --- a/backends/webgpu/runtime/ops/minimum/binary_minimum_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_minimum_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_minimum.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: 929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972 inline constexpr const char* kBinaryMinimumWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h similarity index 98% rename from backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h index c9f60dbd200..68784e82eb2 100644 --- a/backends/webgpu/runtime/ops/mul/binary_mul_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_mul_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_mul.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: d248c0f1856b57115a5001a47f4936caa564dd3b787c02ceba504a13ab987812 inline constexpr const char* kBinaryMulWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl index 98076ed98b5..4b42665013f 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.wgsl @@ -16,11 +16,14 @@ override wg_size: u32 = 64u; $if USE_ALPHA: override alpha: f32 = 1.0; -fn op(a: f32, b: f32) -> f32 { - return ${OP_EXPR}; -} +$if INLINE: + @compute @workgroup_size(wg_size, 1, 1) +$else: + fn op(a: f32, b: f32) -> f32 { + return ${OP_EXPR}; + } -@compute @workgroup_size(wg_size, 1, 1) + @compute @workgroup_size(wg_size, 1, 1) fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { @@ -30,6 +33,8 @@ fn main( return; } + $if INLINE: + // Fast path: every input dim matches the output dim -> elementwise. var same = true; for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || @@ -38,10 +43,15 @@ fn main( } } if (same) { - output[idx] = op(input1[idx], input2[idx]); + $if INLINE: + output[idx] = ${SAME_EXPR}; + $else: + output[idx] = op(input1[idx], input2[idx]); return; } + $if INLINE: + // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. var rem = idx; var l1: u32 = 0u; var l2: u32 = 0u; @@ -51,5 +61,8 @@ fn main( l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; } - output[idx] = op(input1[l1], input2[l2]); + $if INLINE: + output[idx] = ${BROADCAST_EXPR}; + $else: + output[idx] = op(input1[l1], input2[l2]); } diff --git a/backends/webgpu/runtime/ops/binary_op/binary_op.yaml b/backends/webgpu/runtime/ops/binary_op/binary_op.yaml index 03040fcbcce..bd4cd452d35 100644 --- a/backends/webgpu/runtime/ops/binary_op/binary_op.yaml +++ b/backends/webgpu/runtime/ops/binary_op/binary_op.yaml @@ -2,6 +2,9 @@ binary_op: parameter_names_with_default_values: OP_EXPR: a + alpha * b USE_ALPHA: 1 + INLINE: 0 + SAME_EXPR: input1[idx] + input2[idx] + BROADCAST_EXPR: input1[l1] + input2[l2] shader_variants: - NAME: binary_div OP_EXPR: a / b @@ -9,3 +12,23 @@ binary_op: - NAME: binary_sub OP_EXPR: a - alpha * b USE_ALPHA: 1 + - NAME: binary_minimum + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: min(input1[idx], input2[idx]) + BROADCAST_EXPR: min(input1[l1], input2[l2]) + - NAME: binary_pow + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: pow(input1[idx], input2[idx]) + BROADCAST_EXPR: pow(input1[l1], input2[l2]) + - NAME: binary_floor_divide + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: floor(input1[idx] / input2[idx]) + BROADCAST_EXPR: floor(input1[l1] / input2[l2]) + - NAME: binary_mul + USE_ALPHA: 0 + INLINE: 1 + SAME_EXPR: input1[idx] * input2[idx] + BROADCAST_EXPR: input1[l1] * input2[l2] diff --git a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h b/backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h similarity index 98% rename from backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h rename to backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h index 3532c091160..776d4b6693f 100644 --- a/backends/webgpu/runtime/ops/pow/binary_pow_wgsl.h +++ b/backends/webgpu/runtime/ops/binary_op/binary_pow_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from binary_pow.wgsl - DO NOT EDIT. +// @generated from binary_op.wgsl - DO NOT EDIT. // wgsl-sha256: a88c161bd3f43d21a72ebd8ca6f8611b6b9b854e3572a8e6b820602091bc464c inline constexpr const char* kBinaryPowWGSL = R"( @group(0) @binding(0) var input1: array; diff --git a/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp b/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp index a35a7587311..98a76228003 100644 --- a/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/floor_divide/BinaryOp.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include #include diff --git a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl b/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl deleted file mode 100644 index 3b4edd788b0..00000000000 --- a/backends/webgpu/runtime/ops/floor_divide/binary_floor_divide.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = floor(input1[idx] / input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = floor(input1[l1] / input2[l2]); -} diff --git a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp index b6a7e12830f..358cb29fff0 100644 --- a/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp +++ b/backends/webgpu/runtime/ops/logical_and/LogicalAnd.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h b/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h similarity index 96% rename from backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h rename to backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h index 6a21a77a687..3c1f861f119 100644 --- a/backends/webgpu/runtime/ops/logical_and/logical_and_wgsl.h +++ b/backends/webgpu/runtime/ops/logical_binary/logical_and_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from logical_and.wgsl - DO NOT EDIT. +// @generated from logical_binary.wgsl - DO NOT EDIT. // wgsl-sha256: cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20 inline constexpr const char* kLogicalAndWGSL = R"( @group(0) @binding(0) var t_out: array; diff --git a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl b/backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl similarity index 73% rename from backends/webgpu/runtime/ops/logical_and/logical_and.wgsl rename to backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl index 9acb583f51c..fb5310cfde9 100644 --- a/backends/webgpu/runtime/ops/logical_and/logical_and.wgsl +++ b/backends/webgpu/runtime/ops/logical_binary/logical_binary.wgsl @@ -16,10 +16,13 @@ override wg_size: u32 = 64u; fn main( @builtin(global_invocation_id) gid: vec3, @builtin(num_workgroups) num_workgroups: vec3) { - // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + $if OP == "&": + // bool packed 4/word; canonical 0/1 bytes -> word-wise AND == per-byte AND. + $else: + // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. let w = gid.x + gid.y * (num_workgroups.x * wg_size); if (w >= params.num_words) { return; } - t_out[w] = t_a[w] & t_b[w]; + t_out[w] = t_a[w] ${OP} t_b[w]; } diff --git a/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml b/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml new file mode 100644 index 00000000000..65df4465b04 --- /dev/null +++ b/backends/webgpu/runtime/ops/logical_binary/logical_binary.yaml @@ -0,0 +1,13 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +logical_binary: + parameter_names_with_default_values: + OP: "&" + shader_variants: + - NAME: logical_and + - NAME: logical_or + OP: "|" diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h b/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h similarity index 96% rename from backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h rename to backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h index d64898cb523..e61317d2ff1 100644 --- a/backends/webgpu/runtime/ops/logical_or/logical_or_wgsl.h +++ b/backends/webgpu/runtime/ops/logical_binary/logical_or_wgsl.h @@ -12,7 +12,7 @@ namespace executorch::backends::webgpu { -// @generated from logical_or.wgsl - DO NOT EDIT. +// @generated from logical_binary.wgsl - DO NOT EDIT. // wgsl-sha256: 4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf inline constexpr const char* kLogicalOrWGSL = R"( @group(0) @binding(0) var t_out: array; diff --git a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp index d26f6b486d2..d750af11227 100644 --- a/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp +++ b/backends/webgpu/runtime/ops/logical_or/LogicalOr.cpp @@ -9,7 +9,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl b/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl deleted file mode 100644 index d7e6176ba32..00000000000 --- a/backends/webgpu/runtime/ops/logical_or/logical_or.wgsl +++ /dev/null @@ -1,25 +0,0 @@ -@group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_a: array; -@group(0) @binding(2) var t_b: array; - -struct Params { - num_words: u32, - pad0: u32, - pad1: u32, - pad2: u32, -} -@group(0) @binding(3) var params: Params; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // bool packed 4/word; canonical 0/1 bytes -> word-wise OR == per-byte OR. - let w = gid.x + gid.y * (num_workgroups.x * wg_size); - if (w >= params.num_words) { - return; - } - t_out[w] = t_a[w] | t_b[w]; -} diff --git a/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp b/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp index 3c150fc2580..f457bc9dc99 100644 --- a/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/minimum/BinaryOp.cpp @@ -8,7 +8,7 @@ #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl b/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl deleted file mode 100644 index e79cb2d2bcc..00000000000 --- a/backends/webgpu/runtime/ops/minimum/binary_minimum.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = min(input1[idx], input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = min(input1[l1], input2[l2]); -} diff --git a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp index fdb7984b9e1..27bff334dd3 100644 --- a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp @@ -10,7 +10,7 @@ #include #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl b/backends/webgpu/runtime/ops/mul/binary_mul.wgsl deleted file mode 100644 index f82a16e4b21..00000000000 --- a/backends/webgpu/runtime/ops/mul/binary_mul.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = input1[idx] * input2[idx]; - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = input1[l1] * input2[l2]; -} diff --git a/backends/webgpu/runtime/ops/pow/BinaryOp.cpp b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp index d2f923880bf..1ac90015ffe 100644 --- a/backends/webgpu/runtime/ops/pow/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/pow/BinaryOp.cpp @@ -8,7 +8,7 @@ #include #include -#include +#include #include diff --git a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl b/backends/webgpu/runtime/ops/pow/binary_pow.wgsl deleted file mode 100644 index 2114ef87fee..00000000000 --- a/backends/webgpu/runtime/ops/pow/binary_pow.wgsl +++ /dev/null @@ -1,51 +0,0 @@ -@group(0) @binding(0) var input1: array; -@group(0) @binding(1) var input2: array; -@group(0) @binding(2) var output: array; - -struct TensorMeta { - ndim: u32, - numel: u32, - sizes: array, 2>, - strides: array, 2>, -} -@group(0) @binding(3) var out_meta: TensorMeta; -@group(0) @binding(4) var in1_meta: TensorMeta; -@group(0) @binding(5) var in2_meta: TensorMeta; - -override wg_size: u32 = 64u; - -@compute @workgroup_size(wg_size, 1, 1) -fn main( - @builtin(global_invocation_id) gid: vec3, - @builtin(num_workgroups) num_workgroups: vec3) { - // 2D-folded flat index (lifts the 65535 1D-dispatch cap for large numel). - let idx = gid.x + gid.y * (num_workgroups.x * wg_size); - if (idx >= out_meta.numel) { - return; - } - - // Fast path: every input dim matches the output dim -> elementwise. - var same = true; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - if (in1_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u] || - in2_meta.sizes[d >> 2u][d & 3u] != out_meta.sizes[d >> 2u][d & 3u]) { - same = false; - } - } - if (same) { - output[idx] = pow(input1[idx], input2[idx]); - return; - } - - // Broadcast: out idx -> per-input coord (clamp size-1 dims), relinearize. - var rem = idx; - var l1: u32 = 0u; - var l2: u32 = 0u; - for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) { - let coord = rem / out_meta.strides[d >> 2u][d & 3u]; - rem = rem % out_meta.strides[d >> 2u][d & 3u]; - l1 = l1 + min(coord, in1_meta.sizes[d >> 2u][d & 3u] - 1u) * in1_meta.strides[d >> 2u][d & 3u]; - l2 = l2 + min(coord, in2_meta.sizes[d >> 2u][d & 3u] - 1u) * in2_meta.strides[d >> 2u][d & 3u]; - } - output[idx] = pow(input1[l1], input2[l2]); -} diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index aa8f000bceb..bbd6e13dc9d 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -38,10 +38,10 @@ ) from executorch.backends.webgpu.test.ops.test_avg_pool2d import AvgPool2dModule from executorch.backends.webgpu.test.ops.test_bitwise import ( + BITWISE_NOT_SHAPES, BitwiseAndModule, BitwiseNotModule, bw_gen_a, - bw_gen_b, ) from executorch.backends.webgpu.test.ops.test_cat import ( CatModule, @@ -71,14 +71,13 @@ make_qcs4w_linear_module, ) from executorch.backends.webgpu.test.ops.test_logical_and import ( - la_gen_a, - la_gen_b, + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, LogicalAndModule, ) from executorch.backends.webgpu.test.ops.test_logical_or import ( BitwiseOrModule, - lo_gen_a, - lo_gen_b, LogicalOrModule, ) from executorch.backends.webgpu.test.ops.test_minimum import MinimumModule @@ -253,7 +252,7 @@ def _rms_norm_suite() -> WebGPUTestSuite: @register_op_test("mul") def _mul_suite() -> WebGPUTestSuite: - # Full numeric coverage incl. broadcast (binary_mul.wgsl over a TensorMeta UBO); fp64 golden. + # Full binary_op-family numeric coverage, including broadcast. return WebGPUTestSuite( module_factory=lambda: MulModule(), cases=[ @@ -278,12 +277,19 @@ def _fn_config_suite(module_cls, configs) -> WebGPUTestSuite: @register_op_test("minimum") def _minimum_suite() -> WebGPUTestSuite: - # Same-shape numeric coverage (flat binary kernel; broadcast stays smoke). + # Same-shape and mixed-rank broadcast numeric coverage. return WebGPUTestSuite( module_factory=lambda: MinimumModule(), cases=[ Case(name="2d", inputs=((M1, M2), (M1, M2))), Case(name="3d", inputs=((S, S1, S2), (S, S1, S2))), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(-3.0, 3.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(-2.0, 4.0)), + ), + ), ], ) @@ -334,49 +340,33 @@ def _ge_suite() -> WebGPUTestSuite: return _compare_suite("ge") -@register_op_test("logical_and") -def _logical_and_suite() -> WebGPUTestSuite: - # out = (a>0) && (b>0): two bool masks derived on-GPU from float inputs via - # gt.Tensor (baked zeros), AND'd -> bool. Distinct a/b seeds so the masks - # differ (AND ~25% True, a real mix an OR mutant fails); all shapes numel % - # 4 == 0 (bool packs 4/word). float32 oracle (byte-exact bool golden). +def _logical_binary_suite(module_factory) -> WebGPUTestSuite: + # Every packed u32 receives all four Boolean pairs in byte order. def case(name, shape): return Case( name=name, construct={"shape": shape}, inputs=( - InputSpec(shape=shape, gen=la_gen_a), - InputSpec(shape=shape, gen=la_gen_b), + InputSpec(shape=shape, gen=logical_binary_gen_a), + InputSpec(shape=shape, gen=logical_binary_gen_b), ), ) return WebGPUTestSuite( - module_factory=lambda shape: LogicalAndModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + module_factory=module_factory, + cases=[case(name, shape) for name, shape in LOGICAL_BINARY_CASES], golden_dtype="float32", ) +@register_op_test("logical_and") +def _logical_and_suite() -> WebGPUTestSuite: + return _logical_binary_suite(lambda shape: LogicalAndModule(shape)) + + @register_op_test("bitwise_and") def _bitwise_and_suite() -> WebGPUTestSuite: - # bool bitwise AND == logical_and for canonical 0/1 (shares the handler). - # Two masks derived on-GPU from float inputs via gt.Tensor (baked zeros), - # distinct a/b seeds (AND ~25% True); all shapes numel % 4 == 0. - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=bw_gen_a), - InputSpec(shape=shape, gen=bw_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: BitwiseAndModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: BitwiseAndModule(shape)) @register_op_test("bitwise_not") @@ -392,52 +382,22 @@ def case(name, shape): return WebGPUTestSuite( module_factory=lambda shape: BitwiseNotModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], + cases=[ + case(name, shape) + for name, shape in zip(("2d", "3d", "sq"), BITWISE_NOT_SHAPES) + ], golden_dtype="float32", ) @register_op_test("logical_or") def _logical_or_suite() -> WebGPUTestSuite: - # out = (a>0) || (b>0): two bool masks derived on-GPU from float inputs via - # gt.Tensor (baked zeros), OR'd -> bool. Distinct a/b seeds (~50% each, - # independent -> OR ~75% True, a real mix an AND mutant fails); all shapes - # numel % 4 == 0. float32 oracle (byte-exact bool golden). - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=lo_gen_a), - InputSpec(shape=shape, gen=lo_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: LogicalOrModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: LogicalOrModule(shape)) @register_op_test("bitwise_or") def _bitwise_or_suite() -> WebGPUTestSuite: - # bool bitwise OR == logical_or for canonical 0/1 (shares the handler). - def case(name, shape): - return Case( - name=name, - construct={"shape": shape}, - inputs=( - InputSpec(shape=shape, gen=lo_gen_a), - InputSpec(shape=shape, gen=lo_gen_b), - ), - ) - - return WebGPUTestSuite( - module_factory=lambda shape: BitwiseOrModule(shape), - cases=[case("2d", (4, 8)), case("3d", (2, 3, 8)), case("sq", (16, 16))], - golden_dtype="float32", - ) + return _logical_binary_suite(lambda shape: BitwiseOrModule(shape)) @register_op_test("pow") @@ -460,6 +420,13 @@ def _pow_suite() -> WebGPUTestSuite: InputSpec(shape=(S, S1, S2), gen=_unary_lin(-2.0, 3.0)), ), ), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(0.1, 3.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(-2.0, 3.0)), + ), + ), ], ) @@ -493,6 +460,14 @@ def _floor_divide_suite() -> WebGPUTestSuite: ), golden_fn=_floor_div_golden, ), + Case( + name="broadcast_3d_2d", + inputs=( + InputSpec(shape=(2, 3, 8), gen=_unary_lin(-8.0, 8.0)), + InputSpec(shape=(3, 1), gen=_unary_lin(0.5, 4.0)), + ), + golden_fn=_floor_div_golden, + ), ], ) diff --git a/backends/webgpu/test/op_tests/test_generator.py b/backends/webgpu/test/op_tests/test_generator.py index fabf79171c0..65f765812be 100644 --- a/backends/webgpu/test/op_tests/test_generator.py +++ b/backends/webgpu/test/op_tests/test_generator.py @@ -10,7 +10,15 @@ import torch from executorch.backends.webgpu.test.op_tests import generate_op_tests as g -from executorch.backends.webgpu.test.op_tests.test_suite import op_test_registry +from executorch.backends.webgpu.test.op_tests.test_suite import ( + InputSpec, + op_test_registry, +) +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) def _add_regular_case(): @@ -112,3 +120,108 @@ def test_manifest_schema_roundtrip(tmp_path): gd = e["golden"] assert {"path", "shape", "dtype", "output_index"} <= set(gd) assert gd["output_index"] == 0 + + +def test_logical_binary_case_contract(): + expected_cases = ( + ("2d", (4, 8)), + ("3d", (2, 3, 8)), + ("sq", (16, 16)), + ("words63", (252,)), + ("words64", (256,)), + ("words65", (260,)), + ) + assert LOGICAL_BINARY_CASES == expected_cases + assert logical_binary_gen_a((8,)).tolist() == [ + -1.0, + -1.0, + 1.0, + 1.0, + -1.0, + -1.0, + 1.0, + 1.0, + ] + assert logical_binary_gen_b((8,)).tolist() == [ + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + -1.0, + 1.0, + ] + + for op in ("logical_and", "bitwise_and", "logical_or", "bitwise_or"): + suite = op_test_registry[op] + assert tuple((case.name, case.construct["shape"]) for case in suite.cases) == ( + expected_cases + ) + for case in suite.cases: + assert case.required is True + assert case.heavy is False + assert len(case.inputs) == 2 + assert case.inputs[0].gen is logical_binary_gen_a + assert case.inputs[1].gen is logical_binary_gen_b + + +def test_binary_shader_family_case_contract(): + expected = { + "minimum": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "pow": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "floor_divide": ( + ("2d", ((37, 41), (37, 41))), + ("3d", ((5, 7, 11), (5, 7, 11))), + ("broadcast_3d_2d", ((2, 3, 8), (3, 1))), + ), + "mul": ( + ("same", ((8, 32), (8, 32))), + ("bcast_lastdim", ((1, 1, 7, 896), (1, 1, 7, 1))), + ("bcast_firstdim", ((4, 4), (1, 4))), + ("bcast_4d_mixed", ((3, 5, 7, 11), (1, 5, 1, 11))), + ("mixedrank", ((4,), (3, 4))), + ), + } + + for op, cases in expected.items(): + suite = op_test_registry[op] + actual = tuple( + ( + case.name, + tuple( + spec.shape if isinstance(spec, InputSpec) else spec + for spec in case.inputs + ), + ) + for case in suite.cases + ) + assert actual == cases + + ranges = { + "minimum": ((-3.0, 3.0), (-2.0, 4.0)), + "pow": ((0.1, 3.0), (-2.0, 3.0)), + "floor_divide": ((-8.0, 8.0), (0.5, 4.0)), + } + for op, expected_ranges in ranges.items(): + case = next( + c for c in op_test_registry[op].cases if c.name == "broadcast_3d_2d" + ) + for spec, (start, end) in zip(case.inputs, expected_ranges): + assert isinstance(spec, InputSpec) and callable(spec.gen) + values = spec.gen(spec.shape).flatten() + assert torch.isclose(values[0], torch.tensor(start)) + assert torch.isclose(values[-1], torch.tensor(end)) + + assert all( + case.golden_fn is g.cases._floor_div_golden + for case in op_test_registry["floor_divide"].cases + ) diff --git a/backends/webgpu/test/ops/test_bitwise.py b/backends/webgpu/test/ops/test_bitwise.py index b61a37d9990..a8fee7a0095 100644 --- a/backends/webgpu/test/ops/test_bitwise.py +++ b/backends/webgpu/test/ops/test_bitwise.py @@ -19,6 +19,11 @@ import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) from executorch.exir import to_edge_transform_and_lower @@ -50,11 +55,10 @@ def g(shape): bw_gen_a = _bw_gen(0) -bw_gen_b = _bw_gen(1) # All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] +BITWISE_NOT_SHAPES = ((4, 8), (2, 3, 8), (16, 16)) class BitwiseTest(unittest.TestCase): @@ -75,13 +79,17 @@ def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ) def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = bw_gen_a(shape) - b = bw_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(op="bitwise_and", case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) self._assert_delegates( BitwiseAndModule(shape), (a, b), "bitwise_and", shape ) + + for shape in BITWISE_NOT_SHAPES: + with self.subTest(op="bitwise_not", shape=shape): + a = bw_gen_a(shape) self._assert_delegates( BitwiseNotModule(shape), (a,), "bitwise_not", shape ) diff --git a/backends/webgpu/test/ops/test_floor_divide.py b/backends/webgpu/test/ops/test_floor_divide.py index 6020480ff2c..141e8314cde 100644 --- a/backends/webgpu/test/ops/test_floor_divide.py +++ b/backends/webgpu/test/ops/test_floor_divide.py @@ -6,8 +6,8 @@ """`aten.div.Tensor_mode` module for the WebGPU op-test framework. -`FloorDivideModule` is imported by `cases.py`. Same-shape elementwise -`div(a, b, rounding_mode="floor")`. The kernel computes `floor(a/b)` mirroring +`FloorDivideModule` is imported by `cases.py`. Same-shape and broadcast +`div(a, b, rounding_mode="floor")` use `floor(a/b)`, mirroring the Vulkan `floor_divide` glsl (`floor(X/Y)`); this differs from torch's own fmod-corrected `div_floor` at rare fp boundaries, so the suite goldens against a `floor(a/b)` `golden_fn` (Vulkan-faithful), not this module's eager output. diff --git a/backends/webgpu/test/ops/test_logical_and.py b/backends/webgpu/test/ops/test_logical_and.py index 4c3939002d0..9cf474fd8de 100644 --- a/backends/webgpu/test/ops/test_logical_and.py +++ b/backends/webgpu/test/ops/test_logical_and.py @@ -4,17 +4,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.logical_and.default` module + configs for the WebGPU op-test framework. - -`LogicalAndModule` derives its two bool operands on-GPU from float inputs -(`a > 0`, `b > 0` via the delegated `gt.Tensor` against a baked zero buffer), so -the only runtime inputs are the two float tensors (the op-test framework is -float-input-only). `a`/`b` use distinct seeds so the two bool masks differ (each -~50% True, independent -> AND ~25% True), a real mix that a wrong op (e.g. OR) -would fail. Output is bool (byte-exact golden). `LogicalAndTest` is the -export-delegation smoke test. -""" +"""Delegation coverage for logical AND with packed truth-table inputs.""" +import math import unittest import torch @@ -32,29 +24,40 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.logical_and(a > self.z, b > self.z) -def _la_gen(seed): - # Distinct per-input seed so the two derived bool masks differ. - def g(shape): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=torch.float32) +LOGICAL_BINARY_CASES = ( + ("2d", (4, 8)), + ("3d", (2, 3, 8)), + ("sq", (16, 16)), + ("words63", (252,)), + ("words64", (256,)), + ("words65", (260,)), +) - return g +def _logical_binary_gen(pattern): + def generate(shape): + numel = math.prod(shape) + if numel == 0 or numel % len(pattern) != 0: + raise ValueError("logical-binary test shapes must have numel % 4 == 0") + return ( + torch.tensor(pattern, dtype=torch.float32) + .repeat(numel // len(pattern)) + .reshape(shape) + ) -la_gen_a = _la_gen(0) -la_gen_b = _la_gen(1) + return generate -# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] +logical_binary_gen_a = _logical_binary_gen((-1.0, -1.0, 1.0, 1.0)) +logical_binary_gen_b = _logical_binary_gen((-1.0, 1.0, -1.0, 1.0)) class LogicalAndTest(unittest.TestCase): def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = la_gen_a(shape) - b = la_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) ep = torch.export.export(LogicalAndModule(shape).eval(), (a, b)) edge = to_edge_transform_and_lower( ep, partitioner=[VulkanPartitioner()] diff --git a/backends/webgpu/test/ops/test_logical_or.py b/backends/webgpu/test/ops/test_logical_or.py index d71c1493ea9..420fee26336 100644 --- a/backends/webgpu/test/ops/test_logical_or.py +++ b/backends/webgpu/test/ops/test_logical_or.py @@ -4,23 +4,18 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -"""`aten.logical_or.default` / `aten.bitwise_or.Tensor` (bool) modules + configs. - -Mirrors the logical_and/bitwise_and tests: the modules derive their two bool -operands on-GPU from float inputs (`a > 0`, `b > 0` via the delegated `gt.Tensor` -against a baked zero buffer), so the only runtime inputs are the two float -tensors (the op-test framework is float-input-only). `a`/`b` use distinct seeds -so the two bool masks differ (each ~50% True, independent -> OR ~75% True), a -real mix that a wrong op (e.g. AND) would fail. `bitwise_or` on bool is identical -to `logical_or` (shares the handler). Output is bool (byte-exact golden). -`LogicalOrTest` is the export-delegation smoke test. -""" +"""Delegation coverage for logical and bitwise OR truth-table inputs.""" import unittest import torch from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.webgpu.test.ops.test_logical_and import ( + LOGICAL_BINARY_CASES, + logical_binary_gen_a, + logical_binary_gen_b, +) from executorch.exir import to_edge_transform_and_lower @@ -42,23 +37,6 @@ def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.bitwise_or(a > self.z, b > self.z) -def _lo_gen(seed): - # Distinct per-input seed so the two derived bool masks differ. - def g(shape): - gen = torch.Generator().manual_seed(seed) - return torch.randn(*shape, generator=gen, dtype=torch.float32) - - return g - - -lo_gen_a = _lo_gen(0) -lo_gen_b = _lo_gen(1) - - -# All shapes have numel % 4 == 0 (bool tensors pack 4 bytes/word). -SHAPES = [(4, 8), (2, 3, 8), (16, 16)] - - class LogicalOrTest(unittest.TestCase): def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ep = torch.export.export(mod.eval(), inputs) @@ -77,10 +55,10 @@ def _assert_delegates(self, mod, inputs, op_name, shape) -> None: ) def test_export_delegates(self) -> None: - for shape in SHAPES: - with self.subTest(shape=shape): - a = lo_gen_a(shape) - b = lo_gen_b(shape) + for case_name, shape in LOGICAL_BINARY_CASES: + with self.subTest(case=case_name, shape=shape): + a = logical_binary_gen_a(shape) + b = logical_binary_gen_b(shape) self._assert_delegates( LogicalOrModule(shape), (a, b), "logical_or", shape ) diff --git a/backends/webgpu/test/ops/test_minimum.py b/backends/webgpu/test/ops/test_minimum.py index f8f8088fef5..ba466317d5b 100644 --- a/backends/webgpu/test/ops/test_minimum.py +++ b/backends/webgpu/test/ops/test_minimum.py @@ -6,8 +6,8 @@ """`aten.minimum.default` module for the WebGPU op-test framework. -`MinimumModule` is imported by `cases.py`. minimum is a same-shape elementwise -binary op mirroring the landed `add`/`mul` pattern (flat 2D-dispatch kernel). +`MinimumModule` is imported by `cases.py`. minimum is an elementwise binary op +with the same-shape fast path and broadcast indexing used by `mul`. """ import torch diff --git a/backends/webgpu/test/ops/test_pow.py b/backends/webgpu/test/ops/test_pow.py index 93596dc5d9f..8498bd4f7aa 100644 --- a/backends/webgpu/test/ops/test_pow.py +++ b/backends/webgpu/test/ops/test_pow.py @@ -6,8 +6,8 @@ """`aten.pow.Tensor_Tensor` module for the WebGPU op-test framework. -`PowModule` is imported by `cases.py`. Same-shape elementwise `pow(a, b)`; the -suite uses a POSITIVE base so `pow(neg, frac)` (NaN) is never exercised. +`PowModule` is imported by `cases.py`. The suite covers same-shape and broadcast +`pow(a, b)` with POSITIVE bases so `pow(neg, frac)` never produces NaN. """ import torch diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index d437a0a42e8..574c3869864 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -213,7 +213,11 @@ def test_generated_output_manifest_digest(self) -> None: self.assertEqual(len(outputs), 134) self.assertEqual( digest.hexdigest(), - "ef97dca2336315ee2c8b0f9e896c6aa082834ae94948bda3ccb42b1145f2bd27", + "e502196846f0f8100f468e5d9f8f9c006b67e08df54e1e2e667daa2fc50d8844", + ) + self.assertEqual( + hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), + "492535b396833ad6ebfed29b093057e38f40b1182b5c7d6b3eb2f3577cab024e", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: @@ -950,10 +954,6 @@ def test_to_copy_convert_template_roundtrip_byte_identical(self) -> None: entries["to_copy_int_to_float"].include, "runtime/ops/to_copy/to_copy_int_to_float_wgsl.h", ) - self.assertEqual( - hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", - ) def test_extrema_template_roundtrip_byte_identical(self) -> None: extrema_dir = g.BACKEND_ROOT / "runtime/ops/extrema" @@ -1004,10 +1004,108 @@ def test_extrema_template_roundtrip_byte_identical(self) -> None: hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash ) - self.assertEqual( - hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "74f972fce4077f12a52dfcf67a0d20ebeea47748283ced9e5c0bcffd659fef74", - ) + def test_logical_binary_template_roundtrip_byte_identical(self) -> None: + logical_dir = g.BACKEND_ROOT / "runtime/ops/logical_binary" + template_path = logical_dir / "logical_binary.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "logical_and": ( + "&", + "cf7c1d1dbba94e429120796c9c25a6717786cca03c08f3bd1e291d5627089c20", + ), + "logical_or": ( + "|", + "4ad19ee04e2c7b396b4669cf44f95133d658c3ec2e6f37d7b271bedc0e582ecf", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + + for name, (op, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["OP"], op) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header_path = logical_dir / f"{name}_wgsl.h" + header = header_path.read_text() + body = header.split('R"(', 1)[1].split(')";', 1)[0][1:] + self.assertEqual(body, expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(body), (64, 1, 1)) + + entries = {entry.name: entry for entry in g.registry_entries()} + for name in expected: + self.assertEqual( + entries[name].include, + f"runtime/ops/logical_binary/{name}_wgsl.h", + ) + self.assertEqual(entries[name].symbol, g.symbol_base(name)) + + handler_hashes = { + "logical_and": "eb85a8f97ee7640298a661da49feb08aa79b8c24d3d4458b71d24d3f01bc388d", + "logical_or": "bda18617f7077fee5a812c21cdc495c89542a1688f7e1ef6739ed01da343a66b", + } + for name, expected_hash in handler_hashes.items(): + handler = ( + g.BACKEND_ROOT / f"runtime/ops/{name}/Logical{name[8:].title()}.cpp" + ) + self.assertEqual( + hashlib.sha256(handler.read_bytes()).hexdigest(), expected_hash + ) + + def test_binary_family_roundtrip_byte_identical(self) -> None: + binary_dir = g.BACKEND_ROOT / "runtime/ops/binary_op" + template_path = binary_dir / "binary_op.wgsl" + spec = g.parse_template_spec(template_path.with_suffix(".yaml")) + variants = {params["NAME"]: params for params in spec[template_path.stem]} + expected = { + "binary_div": ( + 0, + "e36b560fd623dd5337b9ae57acd8981c9c635b995d6021caf1331c182cd3f0cd", + ), + "binary_sub": ( + 0, + "63209ff70422a21fc340d9aadba0945bc259bba89bdf05db018a6507d01c7ae5", + ), + "binary_minimum": ( + 1, + "929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972", + ), + "binary_pow": ( + 1, + "a88c161bd3f43d21a72ebd8ca6f8611b6b9b854e3572a8e6b820602091bc464c", + ), + "binary_floor_divide": ( + 1, + "baf71d277da79389315a6b96b439e7f0a55842e8288283f2af121f84536b3af3", + ), + "binary_mul": ( + 1, + "d248c0f1856b57115a5001a47f4936caa564dd3b787c02ceba504a13ab987812", + ), + } + self.assertEqual(set(variants), set(expected)) + template = template_path.read_text() + entries = {entry.name: entry for entry in g.registry_entries()} + + for name, (inline, expected_hash) in expected.items(): + params = variants[name] + self.assertEqual(params["INLINE"], inline) + expanded = g.preprocess(template, {**g.WGSL_HELPERS, **params}) + self.assertEqual(g.wgsl_sha256(expanded), expected_hash) + + header = (binary_dir / f"{name}_wgsl.h").read_text() + literal = header.split('R"(', 1)[1].split(')";', 1)[0] + self.assertEqual(literal, "\n" + expanded) + self.assertEqual(g.embedded_sha256(header), expected_hash) + self.assertEqual(g.parse_workgroup_size(expanded), (64, 1, 1)) + self.assertIn(f"k{g.symbol_base(name)}WGSL", header) + self.assertEqual( + entries[name].include, + f"runtime/ops/binary_op/{name}_wgsl.h", + ) def test_unary_template_roundtrip_byte_identical(self) -> None: unary_dir = g.BACKEND_ROOT / "runtime/ops/unary"