Skip to content

Commit 343aa70

Browse files
JCNTHJCNTH
committed
Update (base update)
[ghstack-poisoned]
1 parent f16809e commit 343aa70

13 files changed

Lines changed: 62 additions & 63 deletions

File tree

‎backends/vulkan/custom_ops_lib.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1256,25 +1256,25 @@ def adamw_step_meta(
12561256

12571257

12581258
# STE weight gradient d_out^T @ x through the frozen 4-bit linear_q4gsw base.
1259-
def linear_q4gsw_dW_impl(
1259+
def linear_dW_impl(
12601260
d_out: torch.Tensor,
12611261
x: torch.Tensor,
12621262
) -> torch.Tensor:
12631263
return d_out.reshape(-1, d_out.shape[-1]).t() @ x.reshape(-1, x.shape[-1])
12641264

12651265

1266-
def linear_q4gsw_dW_meta(
1266+
def linear_dW_meta(
12671267
d_out: torch.Tensor,
12681268
x: torch.Tensor,
12691269
) -> torch.Tensor:
12701270
return d_out.new_empty((d_out.shape[-1], x.shape[-1]))
12711271

12721272

1273-
name = "linear_q4gsw_dW"
1273+
name = "linear_dW"
12741274
lib.define(f"{name}(Tensor d_out, Tensor x) -> Tensor")
1275-
lib.impl(name, linear_q4gsw_dW_impl, "CompositeExplicitAutograd")
1276-
lib.impl(name, linear_q4gsw_dW_meta, "Meta")
1277-
linear_q4gsw_dW_op = getattr(getattr(torch.ops, namespace), name)
1275+
lib.impl(name, linear_dW_impl, "CompositeExplicitAutograd")
1276+
lib.impl(name, linear_dW_meta, "Meta")
1277+
linear_dW_op = getattr(getattr(torch.ops, namespace), name)
12781278

12791279

12801280
##################

‎backends/vulkan/op_registry.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1806,8 +1806,8 @@ def register_adamw_step():
18061806
)
18071807

18081808

1809-
@update_features(exir_ops.edge.et_vk.linear_q4gsw_dW.default)
1810-
def register_linear_q4gsw_dW():
1809+
@update_features(exir_ops.edge.et_vk.linear_dW.default)
1810+
def register_linear_dW():
18111811
return OpFeatures(
18121812
inputs_storage=utils.CONTIGUOUS_ANY,
18131813
inputs_dtypes=utils.FP_T,

backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_dW.glsl renamed to backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl

File renamed without changes.

backends/vulkan/runtime/graph/ops/glsl/linear_q4gsw_dW.yaml renamed to backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7-
linear_q4gsw_dW:
7+
linear_dW:
88
parameter_names_with_default_values:
99
DTYPE: float
1010
STORAGE: buffer
@@ -14,4 +14,4 @@ linear_q4gsw_dW:
1414
DTYPE:
1515
- VALUE: float
1616
shader_variants:
17-
- NAME: linear_q4gsw_dW
17+
- NAME: linear_dW

backends/vulkan/runtime/graph/ops/impl/QuantizedLinearDw.cpp renamed to backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
namespace vkcompute {
1616

17-
void resize_linear_q4gsw_dW_node(
17+
void resize_linear_dW_node(
1818
ComputeGraph* graph,
1919
const std::vector<ArgGroup>& args,
2020
const std::vector<ValueRef>& resize_args) {
@@ -27,7 +27,7 @@ void resize_linear_q4gsw_dW_node(
2727
graph->virtual_resize(dW, {N, K});
2828
}
2929

30-
utils::uvec3 linear_q4gsw_dW_global_wg_size(
30+
utils::uvec3 linear_dW_global_wg_size(
3131
ComputeGraph* graph,
3232
const vkapi::ShaderInfo& shader,
3333
const std::vector<ArgGroup>& args,
@@ -41,7 +41,7 @@ utils::uvec3 linear_q4gsw_dW_global_wg_size(
4141
return {tiles, 1u, 1u};
4242
}
4343

44-
utils::uvec3 linear_q4gsw_dW_local_wg_size(
44+
utils::uvec3 linear_dW_local_wg_size(
4545
ComputeGraph* graph,
4646
const vkapi::ShaderInfo& shader,
4747
const utils::uvec3& global_workgroup_size,
@@ -55,7 +55,7 @@ utils::uvec3 linear_q4gsw_dW_local_wg_size(
5555
return {64u, 1u, 1u};
5656
}
5757

58-
void linear_q4gsw_dW(ComputeGraph& graph, const std::vector<ValueRef>& args) {
58+
void linear_dW(ComputeGraph& graph, const std::vector<ValueRef>& args) {
5959
int32_t i = 0;
6060
const ValueRef d_out = args.at(i++);
6161
const ValueRef x = args.at(i++);
@@ -79,18 +79,18 @@ void linear_q4gsw_dW(ComputeGraph& graph, const std::vector<ValueRef>& args) {
7979
utils::div_up_4(static_cast<uint32_t>(K));
8080
VK_CHECK_COND(
8181
(tiles + 63u) / 64u <= 65535u,
82-
"linear_q4gsw_dW: tile count exceeds max workgroup count");
82+
"linear_dW: tile count exceeds max workgroup count");
8383

84-
std::string kernel_name = "linear_q4gsw_dW";
84+
std::string kernel_name = "linear_dW";
8585
kernel_name.reserve(kShaderNameReserve);
8686
add_storage_type_suffix(kernel_name, graph.storage_type_of(dW));
8787
add_dtype_suffix(kernel_name, graph.dtype_of(dW));
8888

8989
graph.execute_nodes().emplace_back(new DynamicDispatchNode(
9090
graph,
9191
VK_KERNEL_FROM_STR(kernel_name),
92-
linear_q4gsw_dW_global_wg_size,
93-
linear_q4gsw_dW_local_wg_size,
92+
linear_dW_global_wg_size,
93+
linear_dW_local_wg_size,
9494
// Inputs and Outputs
9595
{{dW, vkapi::kWrite}, {{d_out, x}, vkapi::kRead}},
9696
// Shader params buffers
@@ -102,11 +102,11 @@ void linear_q4gsw_dW(ComputeGraph& graph, const std::vector<ValueRef>& args) {
102102
// Resize Args
103103
{},
104104
// Resizing Logic
105-
resize_linear_q4gsw_dW_node));
105+
resize_linear_dW_node));
106106
}
107107

108108
REGISTER_OPERATORS {
109-
VK_REGISTER_OP(et_vk.linear_q4gsw_dW.default, linear_q4gsw_dW);
109+
VK_REGISTER_OP(et_vk.linear_dW.default, linear_dW);
110110
}
111111

112112
} // namespace vkcompute

‎backends/vulkan/test/op_tests/CMakeLists.txt‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -120,8 +120,7 @@ if(TARGET vulkan_backend AND LIB_TORCH)
120120
adamw_step_test ${CMAKE_CURRENT_SOURCE_DIR}/adamw_step_test.cpp test_utils
121121
)
122122
vulkan_op_test(
123-
linear_q4gsw_dW_test ${CMAKE_CURRENT_SOURCE_DIR}/linear_q4gsw_dW_test.cpp
124-
test_utils
123+
linear_dW_test ${CMAKE_CURRENT_SOURCE_DIR}/linear_dW_test.cpp test_utils
125124
)
126125

127126
# Only build generated op tests if a path to tags.yaml and

backends/vulkan/test/op_tests/linear_q4gsw_dW_test.cpp renamed to backends/vulkan/test/op_tests/linear_dW_test.cpp

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@
2121
//
2222

2323
// Golden: dW[N, K] = d_out^T @ x, contracting over the flattened leading dims.
24-
// Mirrors the CPU-eager linear_q4gsw_dW_impl in custom_ops_lib.py.
25-
at::Tensor linear_q4gsw_dW_reference_impl(
24+
// Mirrors the CPU-eager linear_dW_impl in custom_ops_lib.py.
25+
at::Tensor linear_dW_reference_impl(
2626
const at::Tensor& d_out,
2727
const at::Tensor& x) {
2828
const int64_t N = d_out.size(-1);
@@ -34,15 +34,15 @@ at::Tensor linear_q4gsw_dW_reference_impl(
3434
// Test function
3535
//
3636

37-
void test_vulkan_linear_q4gsw_dW_impl(
37+
void test_vulkan_linear_dW_impl(
3838
const std::vector<int64_t>& d_out_sizes,
3939
const std::vector<int64_t>& x_sizes,
4040
const vkcompute::utils::StorageType storage = vkcompute::utils::kBuffer) {
4141
at::Tensor d_out =
4242
at::rand(d_out_sizes, at::device(at::kCPU).dtype(at::kFloat));
4343
at::Tensor x = at::rand(x_sizes, at::device(at::kCPU).dtype(at::kFloat));
4444

45-
at::Tensor dW_ref = linear_q4gsw_dW_reference_impl(d_out, x);
45+
at::Tensor dW_ref = linear_dW_reference_impl(d_out, x);
4646

4747
// Build Vulkan graph
4848
using namespace vkcompute;
@@ -58,7 +58,7 @@ void test_vulkan_linear_q4gsw_dW_impl(
5858
const ValueRef r_dW = graph.add_tensor(
5959
dW_ref.sizes().vec(), from_at_scalartype(dW_ref.scalar_type()), storage);
6060

61-
VK_GET_OP_FN("et_vk.linear_q4gsw_dW.default")
61+
VK_GET_OP_FN("et_vk.linear_dW.default")
6262
(graph, {r_d_out.value, r_x.value, r_dW});
6363

6464
ValueRef staging_out = graph.set_output_tensor(r_dW);
@@ -91,20 +91,20 @@ void test_vulkan_linear_q4gsw_dW_impl(
9191
}
9292

9393
// Tile-aligned 2D shapes (M, N, K all multiples of 4).
94-
TEST(VulkanLinearQ4gswDwTest, test_tile_aligned) {
95-
test_vulkan_linear_q4gsw_dW_impl(
94+
TEST(VulkanLinearDwTest, test_tile_aligned) {
95+
test_vulkan_linear_dW_impl(
9696
/*d_out_sizes=*/{8, 16}, /*x_sizes=*/{8, 32});
9797
}
9898

9999
// Non-tile-multiple shapes (M, N, K each not a multiple of 4) to exercise the
100100
// partial-tile min()-clamp paths in the shader.
101-
TEST(VulkanLinearQ4gswDwTest, test_non_tile_multiple) {
102-
test_vulkan_linear_q4gsw_dW_impl(
101+
TEST(VulkanLinearDwTest, test_non_tile_multiple) {
102+
test_vulkan_linear_dW_impl(
103103
/*d_out_sizes=*/{5, 6}, /*x_sizes=*/{5, 10});
104104
}
105105

106106
// Leading dims > 2D: M is the flattened product of all leading dims.
107-
TEST(VulkanLinearQ4gswDwTest, test_leading_dims_flatten) {
108-
test_vulkan_linear_q4gsw_dW_impl(
107+
TEST(VulkanLinearDwTest, test_leading_dims_flatten) {
108+
test_vulkan_linear_dW_impl(
109109
/*d_out_sizes=*/{2, 3, 16}, /*x_sizes=*/{2, 3, 32});
110110
}

‎backends/vulkan/test/op_tests/targets.bzl‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,7 @@ def define_common_targets(is_fbcode = False):
187187
]
188188
)
189189
define_test_targets(
190-
"linear_q4gsw_dW_test",
190+
"linear_dW_test",
191191
extra_deps = [
192192
":test_utils",
193193
]

‎backends/webgpu/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ set(WEBGPU_SRCS
7070
runtime/ops/linear/Linear.cpp
7171
runtime/ops/embedding/Embedding.cpp
7272
runtime/ops/adamw/AdamwStep.cpp
73-
runtime/ops/quantized_linear/QuantizedLinearDw.cpp
73+
runtime/ops/quantized_linear/LinearDw.cpp
7474
runtime/ops/quantized_linear/QuantizedLinearRequant.cpp
7575
)
7676

backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearDw.cpp renamed to backends/webgpu/runtime/ops/quantized_linear/LinearDw.cpp

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
1010
#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
1111
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
12-
#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_dW_wgsl.h>
12+
#include <executorch/backends/webgpu/runtime/ops/quantized_linear/linear_dW_wgsl.h>
1313

1414
#include <webgpu/webgpu.h>
1515

@@ -21,16 +21,16 @@ namespace executorch::backends::webgpu {
2121

2222
namespace {
2323

24-
struct Q4gswDwParams {
24+
struct LinearDwParams {
2525
uint32_t M;
2626
uint32_t N;
2727
uint32_t K;
2828
uint32_t _pad;
2929
};
30-
static_assert(sizeof(Q4gswDwParams) == 16, "params must be 16 bytes");
30+
static_assert(sizeof(LinearDwParams) == 16, "params must be 16 bytes");
3131

3232
// STE weight gradient d_W[N,K] = d_out^T @ x.
33-
void q4gsw_dW_impl(WebGPUGraph& graph, const std::vector<int>& args) {
33+
void linear_dW_impl(WebGPUGraph& graph, const std::vector<int>& args) {
3434
const int dout_id = args.at(0);
3535
const int x_id = args.at(1);
3636
const int dw_id = args.at(2);
@@ -41,12 +41,12 @@ void q4gsw_dW_impl(WebGPUGraph& graph, const std::vector<int>& args) {
4141
const auto& dw = graph.get_tensor(dw_id);
4242

4343
if (dw.dims.size() != 2 || dout.dims.empty() || x.dims.empty()) {
44-
throw std::runtime_error("q4gsw_dW: bad tensor ranks");
44+
throw std::runtime_error("linear_dW: bad tensor ranks");
4545
}
4646
const uint32_t N = static_cast<uint32_t>(dw.dims[0]);
4747
const uint32_t K = static_cast<uint32_t>(dw.dims[1]);
4848
if (N == 0 || K == 0) {
49-
throw std::runtime_error("q4gsw_dW: N or K == 0");
49+
throw std::runtime_error("linear_dW: N or K == 0");
5050
}
5151

5252
uint64_t dout_numel = 1;
@@ -59,42 +59,42 @@ void q4gsw_dW_impl(WebGPUGraph& graph, const std::vector<int>& args) {
5959
}
6060
if (static_cast<uint32_t>(dout.dims.back()) != N ||
6161
static_cast<uint32_t>(x.dims.back()) != K) {
62-
throw std::runtime_error("q4gsw_dW: d_out/x last dim mismatch");
62+
throw std::runtime_error("linear_dW: d_out/x last dim mismatch");
6363
}
6464
const uint32_t M = static_cast<uint32_t>(dout_numel / N);
6565
if (dout_numel % N != 0 || x_numel % K != 0 || x_numel / K != M) {
66-
throw std::runtime_error("q4gsw_dW: M mismatch across d_out/x");
66+
throw std::runtime_error("linear_dW: M mismatch across d_out/x");
6767
}
6868

6969
// fp32-only byte-size guards (mirror the forward's byte checks).
7070
if (dw.nbytes != static_cast<uint64_t>(N) * K * sizeof(float) ||
7171
dout.nbytes != dout_numel * sizeof(float) ||
7272
x.nbytes != x_numel * sizeof(float)) {
73-
throw std::runtime_error("q4gsw_dW: fp32-only (byte-size mismatch)");
73+
throw std::runtime_error("linear_dW: fp32-only (byte-size mismatch)");
7474
}
7575

76-
Q4gswDwParams params = {};
76+
LinearDwParams params = {};
7777
params.M = M;
7878
params.N = N;
7979
params.K = K;
8080

8181
const uint32_t wg_size =
82-
utils::clamp_workgroup_size(device, kQ4gswDwWorkgroupSizeX);
82+
utils::clamp_workgroup_size(device, kLinearDwWorkgroupSizeX);
8383
const uint64_t tiles =
8484
utils::div_up<uint64_t>(N, 4u) * utils::div_up<uint64_t>(K, 4u);
8585
if (tiles > UINT32_MAX) {
86-
throw std::runtime_error("q4gsw_dW: tile count exceeds u32");
86+
throw std::runtime_error("linear_dW: tile count exceeds u32");
8787
}
8888
const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
89-
device, static_cast<uint32_t>(tiles), wg_size, "q4gsw_dW");
89+
device, static_cast<uint32_t>(tiles), wg_size, "linear_dW");
9090

9191
WGPUBuffer uniform_buffer =
9292
utils::make_uniform(device, &params, sizeof(params));
9393
graph.add_uniform_buffer_bytes(sizeof(params));
9494

9595
WGPUShaderSourceWGSL wgsl_desc = {};
9696
wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
97-
wgsl_desc.code = {kQ4gswDwWGSL, WGPU_STRLEN};
97+
wgsl_desc.code = {kLinearDwWGSL, WGPU_STRLEN};
9898
WGPUShaderModuleDescriptor shader_desc = {};
9999
shader_desc.nextInChain = &wgsl_desc.chain;
100100
WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
@@ -156,7 +156,7 @@ void q4gsw_dW_impl(WebGPUGraph& graph, const std::vector<int>& args) {
156156
bg_desc.entries = bg_entries;
157157
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
158158

159-
graph.add_dispatch({pipeline, bind_group, workgroup_count, "q4gsw_dW"});
159+
graph.add_dispatch({pipeline, bind_group, workgroup_count, "linear_dW"});
160160

161161
wgpuShaderModuleRelease(shader);
162162
wgpuBindGroupLayoutRelease(bgl);
@@ -167,7 +167,7 @@ void q4gsw_dW_impl(WebGPUGraph& graph, const std::vector<int>& args) {
167167
} // namespace
168168

169169
WEBGPU_REGISTER_OPERATORS {
170-
WEBGPU_REGISTER_OP(et_vk.linear_q4gsw_dW.default, q4gsw_dW_impl);
170+
WEBGPU_REGISTER_OP(et_vk.linear_dW.default, linear_dW_impl);
171171
}
172172

173173
} // namespace executorch::backends::webgpu

0 commit comments

Comments
 (0)