diff --git a/backends/webgpu/runtime/ops/add/BinaryOp.cpp b/backends/webgpu/runtime/ops/add/BinaryOp.cpp index ca49e21f046..157317adcf0 100644 --- a/backends/webgpu/runtime/ops/add/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/add/BinaryOp.cpp @@ -76,90 +76,37 @@ void add_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_uniform_buffer_bytes(sizeof(AddParams)); - // Create shader module from built-in WGSL source - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kBinaryAddWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Create bind group layout: 3 storage buffers + 1 uniform - WGPUBindGroupLayoutEntry entries[4] = {}; - - // input1 - storage buffer, read-only - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - // input2 - storage buffer, read-only - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - // output - storage buffer, read-write - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - - // params - uniform buffer - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - // Create pipeline layout - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - // Create compute pipeline - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - // Create bind group with actual buffers const auto& in1_tensor = graph.get_tensor(in1_id); const auto& in2_tensor = graph.get_tensor(in2_id); - WGPUBindGroupEntry bg_entries[4] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = in1_tensor.buffer; - bg_entries[0].size = in1_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = in2_tensor.buffer; - bg_entries[1].size = in2_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_tensor.buffer; - bg_entries[2].size = out_tensor.nbytes; - - bg_entries[3].binding = 3; - bg_entries[3].buffer = uniform_buffer; - bg_entries[3].size = sizeof(AddParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kBinaryAddWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in1_tensor.buffer, + in1_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in2_tensor.buffer, + in2_tensor.nbytes}, + {2, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(AddParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "", + workgroup_count.y}); const size_t dispatch_idx = graph.num_dispatches() - 1; // Dynamic shapes: recompute numel/dispatch; out follows the larger operand. @@ -193,10 +140,6 @@ void add_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(in1_id, add_resize); graph.add_tensor_resize_hook(in2_id, add_resize); - // Release intermediate objects (pipeline and bind_group are kept by dispatch) - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so a resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp index 6f1febf77d8..ab2d7ce690b 100644 --- a/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp +++ b/backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp @@ -197,76 +197,38 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { wgpuBufferUnmap(uniform_buffer); graph.add_uniform_buffer_bytes(sizeof(EmbeddingParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kEmbeddingQ4gswWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: out (rw) + indices/weight/scales (ro storage) + uniform. - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - for (uint32_t i = 1; i <= 3; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - } - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = out.buffer; - bg_entries[0].size = out.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = indices.buffer; - bg_entries[1].size = indices.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = weight.buffer; - bg_entries[2].size = weight.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = scales.buffer; - bg_entries[3].size = scales.nbytes; - bg_entries[4].binding = 4; - bg_entries[4].buffer = uniform_buffer; - bg_entries[4].size = sizeof(EmbeddingParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + // Bindings: out (rw) + indices/weight/scales (ro storage) + uniform. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kEmbeddingQ4gswWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + indices.buffer, + indices.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight.buffer, + weight.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + scales.buffer, + scales.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(EmbeddingParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count, "embedding_q4gsw"}); + {bundle.pipeline, bundle.bind_group, workgroup_count, "embedding_q4gsw"}); // Dynamic shapes: recompute counts/dispatch; out = indices + [embed_dim]. const uint32_t gs_u = static_cast(group_size); @@ -297,9 +259,6 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector& args) { params_buf); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/index/Index.cpp b/backends/webgpu/runtime/ops/index/Index.cpp index 0c33d616c03..b122cc3af50 100644 --- a/backends/webgpu/runtime/ops/index/Index.cpp +++ b/backends/webgpu/runtime/ops/index/Index.cpp @@ -109,73 +109,33 @@ void index_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(IndexParams)); graph.add_uniform_buffer_bytes(sizeof(IndexParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kIndexWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // self (read), out (read_write), index (read i32), params (uniform). - WGPUBindGroupLayoutEntry entries[4] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[4] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = self_tensor.buffer; - bg_entries[0].size = self_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = index_tensor.buffer; - bg_entries[2].size = index_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = uniform_buffer; - bg_entries[3].size = sizeof(IndexParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - graph.add_dispatch({pipeline, bind_group, workgroup_count}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kIndexWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + self_tensor.buffer, + self_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + index_tensor.buffer, + index_tensor.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(IndexParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + // The bind group keeps the uniform buffer alive until release. wgpuBufferRelease(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp index a14e583037a..604955fc4ab 100644 --- a/backends/webgpu/runtime/ops/mul/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/mul/BinaryOp.cpp @@ -78,95 +78,36 @@ void mul_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, &in2_meta, sizeof(TensorMeta)); graph.add_uniform_buffer_bytes(3 * sizeof(TensorMeta)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kBinaryMulWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group: in1, in2, out (rw), out_meta, in1_meta, in2_meta (3 uniforms). - WGPUBindGroupLayoutEntry entries[6] = {}; - - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - entries[5].binding = 5; - entries[5].visibility = WGPUShaderStage_Compute; - entries[5].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 6; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[6] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = in1_tensor.buffer; - bg_entries[0].size = in1_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = in2_tensor.buffer; - bg_entries[1].size = in2_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_tensor.buffer; - bg_entries[2].size = out_tensor.nbytes; - - bg_entries[3].binding = 3; - bg_entries[3].buffer = out_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - - bg_entries[4].binding = 4; - bg_entries[4].buffer = in1_meta_buf; - bg_entries[4].size = sizeof(TensorMeta); - - bg_entries[5].binding = 5; - bg_entries[5].buffer = in2_meta_buf; - bg_entries[5].size = sizeof(TensorMeta); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kBinaryMulWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in1_tensor.buffer, + in1_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in2_tensor.buffer, + in2_tensor.nbytes}, + {2, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, in1_meta_buf, sizeof(TensorMeta)}, + {5, WGPUBufferBindingType_Uniform, in2_meta_buf, sizeof(TensorMeta)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "mul", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "mul", + workgroup_count.y}); // Dynamic shapes: rebuild all 3 broadcast TensorMeta UBOs + dispatch. WGPUBuffer o_buf = out_meta_buf, a_buf = in1_meta_buf, b_buf = in2_meta_buf; @@ -207,9 +148,6 @@ void mul_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(in1_id, mul_resize); graph.add_tensor_resize_hook(in2_id, mul_resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns them so a resize hook can rewrite them; freed in the dtor. graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(in1_meta_buf); diff --git a/backends/webgpu/runtime/ops/permute/Permute.cpp b/backends/webgpu/runtime/ops/permute/Permute.cpp index c3b9b5c70c6..fb23fce7019 100644 --- a/backends/webgpu/runtime/ops/permute/Permute.cpp +++ b/backends/webgpu/runtime/ops/permute/Permute.cpp @@ -108,80 +108,33 @@ void permute_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_uniform_buffer_bytes( 2 * sizeof(TensorMeta) + sizeof(PermuteParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kPermuteWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group: in, out (rw), out_meta, in_meta, params (3 uniforms). - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(PermuteParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kPermuteWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(PermuteParams)}, + }, + &wg_size_constant, + 1); graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "permute", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "permute", + workgroup_count.y}); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Drop our refs; the bind group keeps the uniforms alive until release. wgpuBufferRelease(out_meta_buf); wgpuBufferRelease(in_meta_buf); diff --git a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp index c07e58e8ec8..707971f5a38 100644 --- a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp +++ b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp @@ -134,85 +134,32 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { // group + dispatch. const bool use_vec4 = (row_width % 4u == 0u); - // Create shader module from built-in WGSL source - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {use_vec4 ? kRmsNormVec4WGSL : kRmsNormWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Create bind group layout: out (rw) + in/weight (ro storage) + params - WGPUBindGroupLayoutEntry entries[4] = {}; - - // t_out - storage buffer, read-write - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - - // t_in - storage buffer, read-only - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - // t_weight - storage buffer, read-only - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - - // params - uniform buffer - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 4; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - // Create pipeline layout - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - // Create compute pipeline - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - // Create bind group with actual buffers + // Bind group buffers: out (rw) + in/weight (ro storage) + params. const auto& out_tensor = graph.get_tensor(out_id); const auto& weight_tensor = graph.get_tensor(weight_id); - WGPUBindGroupEntry bg_entries[4] = {}; - - bg_entries[0].binding = 0; - bg_entries[0].buffer = out_tensor.buffer; - bg_entries[0].size = out_tensor.nbytes; - - bg_entries[1].binding = 1; - bg_entries[1].buffer = in_tensor.buffer; - bg_entries[1].size = in_tensor.nbytes; - - bg_entries[2].binding = 2; - bg_entries[2].buffer = weight_tensor.buffer; - bg_entries[2].size = weight_tensor.nbytes; - - bg_entries[3].binding = 3; - bg_entries[3].buffer = uniform_buffer; - bg_entries[3].size = sizeof(RmsNormParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 4; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + // Pass the selected vec4/scalar WGSL string; no override constant. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + use_vec4 ? kRmsNormVec4WGSL : kRmsNormWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + weight_tensor.buffer, + weight_tensor.nbytes}, + {3, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(RmsNormParams)}, + }); // One workgroup per row (kRmsNormWorkgroupSizeX threads cooperate per row) static_assert( @@ -222,7 +169,7 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { kRmsNormVec4WorkgroupSizeX == 64, "must match @workgroup_size and WG_SIZE in rms_norm_vec4.wgsl"); const size_t dispatch_idx = - graph.add_dispatch({pipeline, bind_group, num_rows}); + graph.add_dispatch({bundle.pipeline, bundle.bind_group, num_rows}); // Dynamic shapes: recompute num_rows + rewrite the UBO for the live input. WGPUBuffer params_buf = uniform_buffer; @@ -234,10 +181,6 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector& args) { g, in_id, out_id, row_width, epsilon, dispatch_idx, params_buf); }); - // Release intermediate objects (pipeline and bind_group are kept by dispatch) - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/select/Select.cpp b/backends/webgpu/runtime/ops/select/Select.cpp index 5d3e8103b7d..f2e69823974 100644 --- a/backends/webgpu/runtime/ops/select/Select.cpp +++ b/backends/webgpu/runtime/ops/select/Select.cpp @@ -108,76 +108,28 @@ void select_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(SelectParams)); graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(SelectParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kSelectWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group: in, out (rw), out_meta, in_meta, params (3 uniforms). - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(SelectParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kSelectWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(SelectParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); // Dynamic shapes: out = in minus `dim`; re-resolve index, meta, dispatch. graph.add_tensor_resize_hook( @@ -219,9 +171,6 @@ void select_impl(WebGPUGraph& graph, const std::vector& args) { g.device(), out_numel, wg_size, "select(resize)"); }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns them so the resize hook can rewrite them; freed in the dtor. graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(in_meta_buf); diff --git a/backends/webgpu/runtime/ops/slice/Slice.cpp b/backends/webgpu/runtime/ops/slice/Slice.cpp index 4fbbfccc5fe..34cfaa3c4d1 100644 --- a/backends/webgpu/runtime/ops/slice/Slice.cpp +++ b/backends/webgpu/runtime/ops/slice/Slice.cpp @@ -167,75 +167,27 @@ void slice_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(SliceParams)); graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(SliceParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kSliceWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group: in, out (rw), out_meta, in_meta, params (3 uniforms). - WGPUBindGroupLayoutEntry entries[5] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - entries[3].binding = 3; - entries[3].visibility = WGPUShaderStage_Compute; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - entries[4].binding = 4; - entries[4].visibility = WGPUShaderStage_Compute; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 5; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[5] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = in_tensor.buffer; - bg_entries[0].size = in_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = out_meta_buf; - bg_entries[2].size = sizeof(TensorMeta); - bg_entries[3].binding = 3; - bg_entries[3].buffer = in_meta_buf; - bg_entries[3].size = sizeof(TensorMeta); - bg_entries[4].binding = 4; - bg_entries[4].buffer = params_buf; - bg_entries[4].size = sizeof(SliceParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 5; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - graph.add_dispatch({pipeline, bind_group, workgroup_count}); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kSliceWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {3, WGPUBufferBindingType_Uniform, in_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, params_buf, sizeof(SliceParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); const size_t dispatch_idx = graph.num_dispatches() - 1; // Dynamic shapes: live start/end -> out[dim] len + meta/params/dispatch. @@ -289,9 +241,6 @@ void slice_impl(WebGPUGraph& graph, const std::vector& args) { } graph.add_tensor_resize_hook(in_id, recompute); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns the uniforms so the resize hook can rewrite them; freed in dtor. graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(in_meta_buf); diff --git a/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp b/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp index dc23a45eb91..b0137079b6d 100644 --- a/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp +++ b/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp @@ -119,72 +119,33 @@ void update_cache_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_uniform_buffer_bytes(sizeof(UpdateCacheParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kUpdateCacheWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - // Bind group layout: cache (rw storage) + value (ro storage) + params. - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].binding = 0; - entries[0].visibility = WGPUShaderStage_Compute; - entries[0].buffer.type = WGPUBufferBindingType_Storage; - entries[1].binding = 1; - entries[1].visibility = WGPUShaderStage_Compute; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].visibility = WGPUShaderStage_Compute; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 3; - bgl_desc.entries = entries; - WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); - - WGPUPipelineLayoutDescriptor pl_desc = {}; - pl_desc.bindGroupLayoutCount = 1; - pl_desc.bindGroupLayouts = &bgl; - WGPUPipelineLayout pipeline_layout = - wgpuDeviceCreatePipelineLayout(device, &pl_desc); - WGPUConstantEntry wg_size_constant = {}; wg_size_constant.key = {"wg_size", WGPU_STRLEN}; wg_size_constant.value = static_cast(wg_size); - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = 1; - pipeline_desc.compute.constants = &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - - WGPUBindGroupEntry bg_entries[3] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = cache_tensor.buffer; - bg_entries[0].size = cache_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = value_tensor.buffer; - bg_entries[1].size = value_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(UpdateCacheParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - graph.add_dispatch({pipeline, bind_group, workgroup_count_x}); - - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); + // Bindings: cache (rw storage) + value (ro storage) + params. + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kUpdateCacheWGSL, + { + {0, + WGPUBufferBindingType_Storage, + cache_tensor.buffer, + cache_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + value_tensor.buffer, + value_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(UpdateCacheParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count_x}); + // Drop our ref; the bind group keeps the uniform buffer alive until release. wgpuBufferRelease(uniform_buffer); }