Skip to content

Commit db3782f

Browse files
JCNTHJCNTH
committed
Update
[ghstack-poisoned]
1 parent f3385d0 commit db3782f

11 files changed

Lines changed: 934 additions & 0 deletions

File tree

‎backends/vulkan/op_registry.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1746,6 +1746,38 @@ def register_rms_norm():
17461746
)
17471747

17481748

1749+
1750+
1751+
@update_features(
1752+
[
1753+
exir_ops.edge.aten.ne.Scalar,
1754+
exir_ops.edge.aten.lt.Scalar,
1755+
exir_ops.edge.aten.le.Scalar,
1756+
exir_ops.edge.aten.ge.Scalar,
1757+
]
1758+
)
1759+
def register_compare_scalar_ops():
1760+
return OpFeatures(
1761+
inputs_storage=utils.ANY_STORAGE,
1762+
inputs_dtypes=utils.FP_INT_T,
1763+
outputs_dtypes=utils.BOOL_T,
1764+
supports_resize=True,
1765+
supports_highdim=True,
1766+
)
1767+
1768+
1769+
1770+
1771+
@update_features(exir_ops.edge.aten.logical_not.default)
1772+
def register_logical_not():
1773+
return OpFeatures(
1774+
inputs_storage=utils.ANY_STORAGE,
1775+
inputs_dtypes=utils.BOOL_T,
1776+
supports_resize=True,
1777+
supports_highdim=True,
1778+
)
1779+
1780+
17491781
#######################
17501782
## Utility functions ##
17511783
#######################

‎backends/webgpu/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,10 @@ set(WEBGPU_SRCS
5959
runtime/ops/reduce/Reduce.cpp
6060
runtime/ops/div/BinaryOp.cpp
6161
runtime/ops/sub/BinaryOp.cpp
62+
runtime/ops/where/Where.cpp
63+
runtime/ops/compare/Compare.cpp
6264
runtime/ops/linear/Linear.cpp
65+
runtime/ops/logical_not/LogicalNot.cpp
6366
)
6467

6568
add_library(webgpu_backend ${WEBGPU_SRCS})

‎backends/webgpu/runtime/ops/compare/Compare.cpp‎

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,190 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
10+
#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
11+
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
12+
#include <executorch/backends/webgpu/runtime/ops/compare/compare_wgsl.h>
13+
14+
#include <webgpu/webgpu.h>
15+
16+
#include <cstdint>
17+
#include <stdexcept>
18+
#include <string>
19+
#include <vector>
20+
21+
namespace executorch::backends::webgpu {
22+
23+
namespace {
24+
25+
struct CompareParams {
26+
uint32_t num_elements;
27+
uint32_t mode;
28+
float scalar;
29+
uint32_t _pad;
30+
};
31+
32+
float read_scalar(WebGPUGraph& graph, int id, const char* op_name) {
33+
if (graph.get_value_type(id) == WebGPUGraph::ValueType::Double) {
34+
return static_cast<float>(graph.get_double(id));
35+
}
36+
if (graph.get_value_type(id) == WebGPUGraph::ValueType::Int) {
37+
return static_cast<float>(graph.get_int(id));
38+
}
39+
throw std::runtime_error(std::string(op_name) + ": scalar is not int/double");
40+
}
41+
42+
// cmp(self[i], scalar) -> byte-packed bool; one u32 word packs 4 elems.
43+
void compare_impl(
44+
WebGPUGraph& graph,
45+
const std::vector<int>& args,
46+
uint32_t mode,
47+
const char* op_name) {
48+
const int self_id = args.at(0);
49+
const int out_id = args.at(args.size() - 1);
50+
const float scalar = read_scalar(graph, args.at(1), op_name);
51+
52+
WGPUDevice device = graph.device();
53+
const auto& self_tensor = graph.get_tensor(self_id);
54+
const auto& out_tensor = graph.get_tensor(out_id);
55+
56+
if (self_tensor.buffer == nullptr || out_tensor.buffer == nullptr) {
57+
throw std::runtime_error(std::string(op_name) + ": null buffer binding");
58+
}
59+
if (self_tensor.nbytes % sizeof(float) != 0) {
60+
throw std::runtime_error(std::string(op_name) + ": self is not fp32");
61+
}
62+
const uint32_t numel =
63+
static_cast<uint32_t>(self_tensor.nbytes / sizeof(float));
64+
if (out_tensor.nbytes != static_cast<size_t>(numel)) {
65+
throw std::runtime_error(
66+
std::string(op_name) + ": out is not a 1-byte (bool) tensor");
67+
}
68+
69+
const size_t out_bind_size = (out_tensor.nbytes + 3) & ~size_t(3);
70+
const uint32_t n_words = (numel + 3u) / 4u;
71+
72+
uint32_t wg_size = utils::clamp_workgroup_size(device, kCompareWorkgroupSizeX);
73+
uint32_t workgroup_count =
74+
utils::compute_1d_workgroup_count(device, n_words, wg_size, op_name);
75+
76+
WGPUConstantEntry wg_size_constant = {};
77+
wg_size_constant.key = {"wg_size", WGPU_STRLEN};
78+
wg_size_constant.value = static_cast<double>(wg_size);
79+
80+
CompareParams params = {numel, mode, scalar, 0u};
81+
WGPUBuffer params_buf =
82+
utils::make_uniform(device, &params, sizeof(CompareParams));
83+
graph.add_uniform_buffer_bytes(sizeof(CompareParams));
84+
85+
WGPUShaderSourceWGSL wgsl_desc = {};
86+
wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
87+
wgsl_desc.code = {kCompareWGSL, WGPU_STRLEN};
88+
WGPUShaderModuleDescriptor shader_desc = {};
89+
shader_desc.nextInChain = &wgsl_desc.chain;
90+
WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
91+
92+
WGPUBindGroupLayoutEntry entries[3] = {};
93+
entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
94+
entries[1].buffer.type = WGPUBufferBindingType_Storage;
95+
entries[2].buffer.type = WGPUBufferBindingType_Uniform;
96+
for (uint32_t i = 0; i < 3; i++) {
97+
entries[i].binding = i;
98+
entries[i].visibility = WGPUShaderStage_Compute;
99+
}
100+
101+
WGPUBindGroupLayoutDescriptor bgl_desc = {};
102+
bgl_desc.entryCount = 3;
103+
bgl_desc.entries = entries;
104+
WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
105+
106+
WGPUPipelineLayoutDescriptor pl_desc = {};
107+
pl_desc.bindGroupLayoutCount = 1;
108+
pl_desc.bindGroupLayouts = &bgl;
109+
WGPUPipelineLayout pipeline_layout =
110+
wgpuDeviceCreatePipelineLayout(device, &pl_desc);
111+
112+
WGPUComputePipelineDescriptor pipeline_desc = {};
113+
pipeline_desc.layout = pipeline_layout;
114+
pipeline_desc.compute.module = shader;
115+
pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
116+
pipeline_desc.compute.constantCount = 1;
117+
pipeline_desc.compute.constants = &wg_size_constant;
118+
WGPUComputePipeline pipeline =
119+
wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
120+
121+
WGPUBindGroupEntry bg_entries[3] = {};
122+
bg_entries[0].binding = 0;
123+
bg_entries[0].buffer = self_tensor.buffer;
124+
bg_entries[0].size = self_tensor.nbytes;
125+
bg_entries[1].binding = 1;
126+
bg_entries[1].buffer = out_tensor.buffer;
127+
bg_entries[1].size = out_bind_size;
128+
bg_entries[2].binding = 2;
129+
bg_entries[2].buffer = params_buf;
130+
bg_entries[2].size = sizeof(CompareParams);
131+
132+
WGPUBindGroupDescriptor bg_desc = {};
133+
bg_desc.layout = bgl;
134+
bg_desc.entryCount = 3;
135+
bg_desc.entries = bg_entries;
136+
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
137+
138+
const size_t dispatch_idx =
139+
graph.add_dispatch({pipeline, bind_group, workgroup_count});
140+
141+
WGPUBuffer p_buf = params_buf;
142+
auto cmp_resize = [self_id, out_id, mode, scalar, wg_size, dispatch_idx,
143+
p_buf, op_name](WebGPUGraph& g) {
144+
const auto& d = g.cur_dims(self_id);
145+
uint32_t n = 1u;
146+
for (auto x : d) {
147+
n *= static_cast<uint32_t>(x);
148+
}
149+
g.set_cur_dims(out_id, d);
150+
CompareParams p = {n, mode, scalar, 0u};
151+
wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p));
152+
const uint32_t nw = (n + 3u) / 4u;
153+
g.dispatch_at(dispatch_idx).workgroup_count_x =
154+
utils::compute_1d_workgroup_count(g.device(), nw, wg_size, op_name);
155+
};
156+
graph.add_tensor_resize_hook(self_id, cmp_resize);
157+
158+
wgpuShaderModuleRelease(shader);
159+
wgpuBindGroupLayoutRelease(bgl);
160+
wgpuPipelineLayoutRelease(pipeline_layout);
161+
graph.own_uniform_buffer(params_buf);
162+
}
163+
164+
void eq_scalar_impl(WebGPUGraph& graph, const std::vector<int>& args) {
165+
compare_impl(graph, args, 0u, "eq.Scalar");
166+
}
167+
void ne_scalar_impl(WebGPUGraph& graph, const std::vector<int>& args) {
168+
compare_impl(graph, args, 1u, "ne.Scalar");
169+
}
170+
void le_scalar_impl(WebGPUGraph& graph, const std::vector<int>& args) {
171+
compare_impl(graph, args, 2u, "le.Scalar");
172+
}
173+
void ge_scalar_impl(WebGPUGraph& graph, const std::vector<int>& args) {
174+
compare_impl(graph, args, 3u, "ge.Scalar");
175+
}
176+
void lt_scalar_impl(WebGPUGraph& graph, const std::vector<int>& args) {
177+
compare_impl(graph, args, 4u, "lt.Scalar");
178+
}
179+
180+
} // namespace
181+
182+
WEBGPU_REGISTER_OPERATORS {
183+
WEBGPU_REGISTER_OP(aten.eq.Scalar, eq_scalar_impl);
184+
WEBGPU_REGISTER_OP(aten.ne.Scalar, ne_scalar_impl);
185+
WEBGPU_REGISTER_OP(aten.le.Scalar, le_scalar_impl);
186+
WEBGPU_REGISTER_OP(aten.ge.Scalar, ge_scalar_impl);
187+
WEBGPU_REGISTER_OP(aten.lt.Scalar, lt_scalar_impl);
188+
}
189+
190+
} // namespace executorch::backends::webgpu

‎backends/webgpu/runtime/ops/compare/compare.wgsl‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
@group(0) @binding(0) var<storage, read> input: array<f32>;
2+
@group(0) @binding(1) var<storage, read_write> output: array<u32>;
3+
4+
struct Params {
5+
num_elements: u32,
6+
mode: u32,
7+
scalar: f32,
8+
_pad: u32,
9+
}
10+
@group(0) @binding(2) var<uniform> params: Params;
11+
12+
override wg_size: u32 = 64u;
13+
14+
// One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
15+
@compute @workgroup_size(wg_size, 1, 1)
16+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
17+
let word_idx = gid.x;
18+
let n_words = (params.num_elements + 3u) / 4u;
19+
if (word_idx >= n_words) {
20+
return;
21+
}
22+
var packed: u32 = 0u;
23+
for (var j: u32 = 0u; j < 4u; j = j + 1u) {
24+
let i = word_idx * 4u + j;
25+
if (i < params.num_elements) {
26+
let v = input[i];
27+
var r: bool;
28+
if (params.mode == 0u) {
29+
r = v == params.scalar;
30+
} else if (params.mode == 1u) {
31+
r = v != params.scalar;
32+
} else if (params.mode == 2u) {
33+
r = v <= params.scalar;
34+
} else if (params.mode == 3u) {
35+
r = v >= params.scalar;
36+
} else {
37+
r = v < params.scalar;
38+
}
39+
if (r) {
40+
packed = packed | (1u << (j * 8u));
41+
}
42+
}
43+
}
44+
output[word_idx] = packed;
45+
}

‎backends/webgpu/runtime/ops/compare/compare_wgsl.h‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <cstdint>
12+
13+
namespace executorch::backends::webgpu {
14+
15+
// @generated from compare.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: f13da085195696aa6975cae62b4d0b2f5837fc02584d95b0a46fd06dc418c4a4
17+
inline constexpr const char* kCompareWGSL = R"(
18+
@group(0) @binding(0) var<storage, read> input: array<f32>;
19+
@group(0) @binding(1) var<storage, read_write> output: array<u32>;
20+
21+
struct Params {
22+
num_elements: u32,
23+
mode: u32,
24+
scalar: f32,
25+
_pad: u32,
26+
}
27+
@group(0) @binding(2) var<uniform> params: Params;
28+
29+
override wg_size: u32 = 64u;
30+
31+
// One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
32+
@compute @workgroup_size(wg_size, 1, 1)
33+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
34+
let word_idx = gid.x;
35+
let n_words = (params.num_elements + 3u) / 4u;
36+
if (word_idx >= n_words) {
37+
return;
38+
}
39+
var packed: u32 = 0u;
40+
for (var j: u32 = 0u; j < 4u; j = j + 1u) {
41+
let i = word_idx * 4u + j;
42+
if (i < params.num_elements) {
43+
let v = input[i];
44+
var r: bool;
45+
if (params.mode == 0u) {
46+
r = v == params.scalar;
47+
} else if (params.mode == 1u) {
48+
r = v != params.scalar;
49+
} else if (params.mode == 2u) {
50+
r = v <= params.scalar;
51+
} else if (params.mode == 3u) {
52+
r = v >= params.scalar;
53+
} else {
54+
r = v < params.scalar;
55+
}
56+
if (r) {
57+
packed = packed | (1u << (j * 8u));
58+
}
59+
}
60+
}
61+
output[word_idx] = packed;
62+
}
63+
)";
64+
65+
inline constexpr uint32_t kCompareWorkgroupSizeX = 64;
66+
inline constexpr uint32_t kCompareWorkgroupSizeY = 1;
67+
inline constexpr uint32_t kCompareWorkgroupSizeZ = 1;
68+
69+
} // namespace executorch::backends::webgpu

0 commit comments

Comments
 (0)