diff --git a/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp b/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp index 1da013f77a3..e62f387d258 100644 --- a/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp +++ b/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp @@ -84,61 +84,22 @@ void dispatch_bool_op( utils::make_uniform(device, ¶ms, sizeof(BoolOpParams)); graph.add_uniform_buffer_bytes(sizeof(BoolOpParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {wgsl, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].buffer.type = WGPUBufferBindingType_Storage; - entries[2].buffer.type = WGPUBufferBindingType_Uniform; - for (uint32_t i = 0; i < 3; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - } - - 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); - - 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 = self_tensor.buffer; - bg_entries[0].size = self_bind_size; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out_tensor.buffer; - bg_entries[1].size = out_bind_size; - bg_entries[2].binding = 2; - bg_entries[2].buffer = params_buf; - bg_entries[2].size = sizeof(BoolOpParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + wgsl, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + self_tensor.buffer, + self_bind_size}, + {1, WGPUBufferBindingType_Storage, out_tensor.buffer, out_bind_size}, + {2, WGPUBufferBindingType_Uniform, params_buf, sizeof(BoolOpParams)}, + }, + &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}); WGPUBuffer p_buf = params_buf; auto resize = @@ -158,9 +119,6 @@ void dispatch_bool_op( }; graph.add_tensor_resize_hook(self_id, resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(params_buf); } diff --git a/backends/webgpu/runtime/ops/div/BinaryOp.cpp b/backends/webgpu/runtime/ops/div/BinaryOp.cpp index ec359b6684d..279b50e5ec5 100644 --- a/backends/webgpu/runtime/ops/div/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/div/BinaryOp.cpp @@ -78,94 +78,35 @@ void div_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 = {kBinaryDivWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - 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, + kBinaryDivWGSL, + { + {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, "div", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "div", + 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; @@ -206,9 +147,6 @@ void div_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(in1_id, div_resize); graph.add_tensor_resize_hook(in2_id, div_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/expand_copy/ExpandCopy.cpp b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp index 507677c7bca..5bfc0fa3bf1 100644 --- a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp +++ b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp @@ -60,72 +60,26 @@ void expand_copy_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kExpandCopyWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - 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_Uniform; - 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 = 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); - - 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, + kExpandCopyWGSL, + { + {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)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); + wgpuBufferRelease(out_meta_buf); wgpuBufferRelease(in_meta_buf); } diff --git a/backends/webgpu/runtime/ops/fill/Fill.cpp b/backends/webgpu/runtime/ops/fill/Fill.cpp index 2ccd48cc8fa..8da3b31577d 100644 --- a/backends/webgpu/runtime/ops/fill/Fill.cpp +++ b/backends/webgpu/runtime/ops/fill/Fill.cpp @@ -67,57 +67,28 @@ void add_fill( utils::make_uniform(device, ¶ms, sizeof(FillParams)); graph.add_uniform_buffer_bytes(sizeof(FillParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kFillWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[2] = {}; - 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_Uniform; - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 2; - 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[2] = {}; - 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 = uniform_buffer; - bg_entries[1].size = sizeof(FillParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 2; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kFillWGSL, + { + {0, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {1, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(FillParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "", + workgroup_count.y}); // Dynamic shapes: recompute num_elements/dispatch from the live output dims. WGPUBuffer params_buf = uniform_buffer; @@ -136,9 +107,6 @@ void add_fill( g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(uniform_buffer); } diff --git a/backends/webgpu/runtime/ops/gather/Gather.cpp b/backends/webgpu/runtime/ops/gather/Gather.cpp index d7e170129d1..9236a6714a2 100644 --- a/backends/webgpu/runtime/ops/gather/Gather.cpp +++ b/backends/webgpu/runtime/ops/gather/Gather.cpp @@ -93,81 +93,31 @@ void gather_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(GatherParams)); graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(GatherParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kGatherWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[6] = {}; - entries[0].binding = 0; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].binding = 1; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].binding = 2; - entries[2].buffer.type = WGPUBufferBindingType_Storage; - entries[3].binding = 3; - entries[3].buffer.type = WGPUBufferBindingType_Uniform; - entries[4].binding = 4; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - entries[5].binding = 5; - entries[5].buffer.type = WGPUBufferBindingType_Uniform; - for (auto& e : entries) { - e.visibility = WGPUShaderStage_Compute; - } + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kGatherWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + self_tensor.buffer, + self_tensor.nbytes}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + index_tensor.buffer, + index_tensor.nbytes}, + {2, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {3, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {4, WGPUBufferBindingType_Uniform, self_meta_buf, sizeof(TensorMeta)}, + {5, WGPUBufferBindingType_Uniform, params_buf, sizeof(GatherParams)}, + }, + &wg_size_constant, + 1); + + graph.add_dispatch({bundle.pipeline, bundle.bind_group, workgroup_count}); - 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 = self_tensor.buffer; - bg_entries[0].size = self_tensor.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = index_tensor.buffer; - bg_entries[1].size = index_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 = self_meta_buf; - bg_entries[4].size = sizeof(TensorMeta); - bg_entries[5].binding = 5; - bg_entries[5].buffer = params_buf; - bg_entries[5].size = sizeof(GatherParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 6; - 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); wgpuBufferRelease(out_meta_buf); wgpuBufferRelease(self_meta_buf); wgpuBufferRelease(params_buf); 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/log_softmax/LogSoftmax.cpp b/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp index f4fb9952706..e88114e6516 100644 --- a/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp +++ b/backends/webgpu/runtime/ops/log_softmax/LogSoftmax.cpp @@ -120,67 +120,26 @@ void log_softmax_impl(WebGPUGraph& graph, const std::vector& args) { wgpuBufferUnmap(uniform_buffer); graph.add_uniform_buffer_bytes(sizeof(LogSoftmaxParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kLogSoftmaxWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - 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; - - 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 = in.buffer; - bg_entries[0].size = in.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out.buffer; - bg_entries[1].size = out.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(LogSoftmaxParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kLogSoftmaxWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(LogSoftmaxParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count, "log_softmax"}); + {bundle.pipeline, bundle.bind_group, workgroup_count, "log_softmax"}); // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch. WGPUBuffer params_buf = uniform_buffer; @@ -215,9 +174,6 @@ void log_softmax_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, std::vector(d.begin(), d.end())); }); - 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/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/reduce/Reduce.cpp b/backends/webgpu/runtime/ops/reduce/Reduce.cpp index b276354d2a8..7e185d38f82 100644 --- a/backends/webgpu/runtime/ops/reduce/Reduce.cpp +++ b/backends/webgpu/runtime/ops/reduce/Reduce.cpp @@ -132,67 +132,30 @@ void reduce_impl( wgpuBufferUnmap(uniform_buffer); graph.add_uniform_buffer_bytes(sizeof(ReduceParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kReduceWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - 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; - - 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 = in.buffer; - bg_entries[0].size = in.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out.buffer; - bg_entries[1].size = out.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(ReduceParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kReduceWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(ReduceParams)}, + }, + &wg_size_constant, + 1); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, op_name, workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + op_name, + workgroup_count.y}); // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch. WGPUBuffer params_buf = uniform_buffer; @@ -248,9 +211,6 @@ void reduce_impl( g.set_cur_dims(out_id, od); }); - 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..f4fcb6088a8 100644 --- a/backends/webgpu/runtime/ops/slice/Slice.cpp +++ b/backends/webgpu/runtime/ops/slice/Slice.cpp @@ -167,75 +167,26 @@ 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 +240,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/softmax/Softmax.cpp b/backends/webgpu/runtime/ops/softmax/Softmax.cpp index 72b48cd8bdd..e5f60f1d4c2 100644 --- a/backends/webgpu/runtime/ops/softmax/Softmax.cpp +++ b/backends/webgpu/runtime/ops/softmax/Softmax.cpp @@ -109,67 +109,26 @@ void softmax_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, ¶ms, sizeof(SoftmaxParams)); graph.add_uniform_buffer_bytes(sizeof(SoftmaxParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kSoftmaxWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[3] = {}; - 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; - - 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 = in.buffer; - bg_entries[0].size = in.nbytes; - bg_entries[1].binding = 1; - bg_entries[1].buffer = out.buffer; - bg_entries[1].size = out.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = uniform_buffer; - bg_entries[2].size = sizeof(SoftmaxParams); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 3; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - - const size_t dispatch_idx = - graph.add_dispatch({pipeline, bind_group, workgroup_count, "softmax"}); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kSoftmaxWGSL, + { + {0, WGPUBufferBindingType_ReadOnlyStorage, in.buffer, in.nbytes}, + {1, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(SoftmaxParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = graph.add_dispatch( + {bundle.pipeline, bundle.bind_group, workgroup_count, "softmax"}); // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch. WGPUBuffer params_buf = uniform_buffer; @@ -204,9 +163,6 @@ void softmax_impl(WebGPUGraph& graph, const std::vector& args) { g.set_cur_dims(out_id, d); }); - 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/split_with_sizes/SplitWithSizes.cpp b/backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp new file mode 100644 index 00000000000..75542377f69 --- /dev/null +++ b/backends/webgpu/runtime/ops/split_with_sizes/SplitWithSizes.cpp @@ -0,0 +1,147 @@ +/* + * 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. + */ + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +namespace executorch::backends::webgpu { + +namespace { + +struct SliceParams { + uint32_t dim; + uint32_t start; + uint32_t step; + uint32_t _pad; +}; + +// aten.split_with_sizes_copy: N contiguous chunks via per-output slice gather. +void split_with_sizes_impl(WebGPUGraph& graph, const std::vector& args) { + if (args.size() < 4) { + throw std::runtime_error("WebGPU split_with_sizes: expected >=4 args"); + } + const int in_id = args.at(0); + const std::vector& sizes = graph.get_int_list(args.at(1)); + int64_t dim = graph.get_int(args.at(2)); + const std::vector& outs = graph.get_value_list(args.at(args.size() - 1)); + + WGPUDevice device = graph.device(); + const auto& in_tensor = graph.get_tensor(in_id); + const int in_ndim = static_cast(in_tensor.dims.size()); + if (dim < 0) { + dim += in_ndim; + } + if (dim < 0 || dim >= in_ndim) { + throw std::runtime_error("split_with_sizes: dim out of range"); + } + if (outs.size() != sizes.size()) { + throw std::runtime_error("split_with_sizes: outputs != sizes count"); + } + + // Validate the split contract up front (before any dispatch): each size is + // non-negative and they sum to the split dim's extent -- a negative or + // oversize value would otherwise wrap when cast to u32 into an out-of-bounds + // gather offset. + int64_t sizes_sum = 0; + for (int64_t s : sizes) { + if (s < 0) { + throw std::runtime_error("split_with_sizes: negative split size"); + } + sizes_sum += s; + } + if (sizes_sum != in_tensor.dims[dim]) { + throw std::runtime_error( + "split_with_sizes: sizes must sum to the split dim extent"); + } + + TensorMeta in_meta; + fill_tensor_meta(in_tensor, &in_meta); + if (in_tensor.nbytes != static_cast(in_meta.numel) * sizeof(float)) { + throw std::runtime_error("split_with_sizes: non-fp32 input"); + } + + uint32_t start = 0; + for (size_t i = 0; i < outs.size(); i++) { + const auto& out_tensor = graph.get_tensor(outs[i]); + TensorMeta out_meta; + fill_tensor_meta(out_tensor, &out_meta); + if (out_tensor.nbytes != + static_cast(out_meta.numel) * sizeof(float)) { + throw std::runtime_error("split_with_sizes: non-fp32 output"); + } + + SliceParams params = {}; + params.dim = static_cast(dim); + params.start = start; + params.step = 1u; + start += static_cast(sizes[i]); + + uint32_t wg_size = + utils::clamp_workgroup_size(device, kSliceWorkgroupSizeX); + uint32_t workgroup_count = utils::compute_1d_workgroup_count( + device, out_meta.numel, wg_size, "split_with_sizes"); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUBuffer out_meta_buf = + utils::make_uniform(device, &out_meta, sizeof(TensorMeta)); + WGPUBuffer in_meta_buf = + utils::make_uniform(device, &in_meta, sizeof(TensorMeta)); + WGPUBuffer params_buf = + utils::make_uniform(device, ¶ms, sizeof(SliceParams)); + graph.add_uniform_buffer_bytes( + 2 * sizeof(TensorMeta) + sizeof(SliceParams)); + + 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}); + + graph.own_uniform_buffer(out_meta_buf); + graph.own_uniform_buffer(in_meta_buf); + graph.own_uniform_buffer(params_buf); + } +} + +} // namespace + +WEBGPU_REGISTER_OPERATORS { + WEBGPU_REGISTER_OP(aten.split_with_sizes_copy.default, split_with_sizes_impl); +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sub/BinaryOp.cpp b/backends/webgpu/runtime/ops/sub/BinaryOp.cpp index 6945feb8da2..13b609cc1e7 100644 --- a/backends/webgpu/runtime/ops/sub/BinaryOp.cpp +++ b/backends/webgpu/runtime/ops/sub/BinaryOp.cpp @@ -87,94 +87,35 @@ void sub_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 = {kBinarySubWGSL, WGPU_STRLEN}; - - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - 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 = 2; - pipeline_desc.compute.constants = constants; - 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, + kBinarySubWGSL, + { + {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)}, + }, + constants, + 2); const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count.x, "sub", workgroup_count.y}); + {bundle.pipeline, + bundle.bind_group, + workgroup_count.x, + "sub", + 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; @@ -215,9 +156,6 @@ void sub_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(in1_id, sub_resize); graph.add_tensor_resize_hook(in2_id, sub_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/update_cache/UpdateCache.cpp b/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp index dc23a45eb91..5bdba562633 100644 --- a/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp +++ b/backends/webgpu/runtime/ops/update_cache/UpdateCache.cpp @@ -119,72 +119,32 @@ 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); + 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); } diff --git a/backends/webgpu/runtime/ops/where/Where.cpp b/backends/webgpu/runtime/ops/where/Where.cpp index bde779498dd..f8aba4a5d01 100644 --- a/backends/webgpu/runtime/ops/where/Where.cpp +++ b/backends/webgpu/runtime/ops/where/Where.cpp @@ -92,81 +92,36 @@ void where_impl(WebGPUGraph& graph, const std::vector& args) { utils::make_uniform(device, &b_meta, sizeof(TensorMeta)); graph.add_uniform_buffer_bytes(4 * sizeof(TensorMeta)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {kWhereWGSL, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - - WGPUBindGroupLayoutEntry entries[8] = {}; - entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; - entries[3].buffer.type = WGPUBufferBindingType_Storage; - entries[4].buffer.type = WGPUBufferBindingType_Uniform; - entries[5].buffer.type = WGPUBufferBindingType_Uniform; - entries[6].buffer.type = WGPUBufferBindingType_Uniform; - entries[7].buffer.type = WGPUBufferBindingType_Uniform; - for (uint32_t i = 0; i < 8; i++) { - entries[i].binding = i; - entries[i].visibility = WGPUShaderStage_Compute; - } - - WGPUBindGroupLayoutDescriptor bgl_desc = {}; - bgl_desc.entryCount = 8; - 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[8] = {}; - bg_entries[0].binding = 0; - bg_entries[0].buffer = cond_tensor.buffer; - bg_entries[0].size = cond_bind_size; - bg_entries[1].binding = 1; - bg_entries[1].buffer = a_tensor.buffer; - bg_entries[1].size = a_tensor.nbytes; - bg_entries[2].binding = 2; - bg_entries[2].buffer = b_tensor.buffer; - bg_entries[2].size = b_tensor.nbytes; - bg_entries[3].binding = 3; - bg_entries[3].buffer = out_tensor.buffer; - bg_entries[3].size = out_tensor.nbytes; - bg_entries[4].binding = 4; - bg_entries[4].buffer = out_meta_buf; - bg_entries[4].size = sizeof(TensorMeta); - bg_entries[5].binding = 5; - bg_entries[5].buffer = cond_meta_buf; - bg_entries[5].size = sizeof(TensorMeta); - bg_entries[6].binding = 6; - bg_entries[6].buffer = a_meta_buf; - bg_entries[6].size = sizeof(TensorMeta); - bg_entries[7].binding = 7; - bg_entries[7].buffer = b_meta_buf; - bg_entries[7].size = sizeof(TensorMeta); - - WGPUBindGroupDescriptor bg_desc = {}; - bg_desc.layout = bgl; - bg_desc.entryCount = 8; - bg_desc.entries = bg_entries; - WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kWhereWGSL, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + cond_tensor.buffer, + cond_bind_size}, + {1, + WGPUBufferBindingType_ReadOnlyStorage, + a_tensor.buffer, + a_tensor.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + b_tensor.buffer, + b_tensor.nbytes}, + {3, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {4, WGPUBufferBindingType_Uniform, out_meta_buf, sizeof(TensorMeta)}, + {5, WGPUBufferBindingType_Uniform, cond_meta_buf, sizeof(TensorMeta)}, + {6, WGPUBufferBindingType_Uniform, a_meta_buf, sizeof(TensorMeta)}, + {7, WGPUBufferBindingType_Uniform, b_meta_buf, sizeof(TensorMeta)}, + }, + &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: rebuild the 4 broadcast TensorMeta UBOs + dispatch count. WGPUBuffer o_buf = out_meta_buf, c_buf = cond_meta_buf, a_buf = a_meta_buf, @@ -223,9 +178,6 @@ void where_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(a_id, where_resize); graph.add_tensor_resize_hook(b_id, where_resize); - wgpuShaderModuleRelease(shader); - wgpuBindGroupLayoutRelease(bgl); - wgpuPipelineLayoutRelease(pipeline_layout); graph.own_uniform_buffer(out_meta_buf); graph.own_uniform_buffer(cond_meta_buf); graph.own_uniform_buffer(a_meta_buf); diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 9ccb8ad1376..a73e8e8d04a 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -812,6 +812,52 @@ def _batch_norm_suite() -> WebGPUTestSuite: ) +from executorch.backends.webgpu.test.ops.test_split_with_sizes_copy import ( + _det_input as _split_det_input, + SplitWithSizesModule, +) + + +@register_op_test("split_with_sizes_copy") +def _split_with_sizes_copy_suite() -> WebGPUTestSuite: + # YOLO Detect-head split of concatenated predictions. Multi-output: the + # framework compares chunk 0 (out_index 0) while each case runs all N + # per-chunk dispatches. Covers a 3-way channel split, a dim-0 split, and + # a last-dim split. copy is bit-exact -> float32 golden. + return WebGPUTestSuite( + module_factory=lambda sizes, dim, out_order=None: SplitWithSizesModule( + sizes, dim, out_order + ), + cases=[ + Case( + name="three_dim1", + construct={"sizes": [2, 3, 3], "dim": 1}, + inputs=(InputSpec(shape=(1, 8, 4, 4), gen=_split_det_input),), + ), + Case( + name="two_dim0", + construct={"sizes": [3, 2], "dim": 0}, + inputs=(InputSpec(shape=(5, 4), gen=_split_det_input),), + ), + Case( + name="dim_last", + construct={"sizes": [4, 4], "dim": -1}, + inputs=(InputSpec(shape=(2, 8), gen=_split_det_input),), + ), + # Reorder so chunk 1 (running offset > 0) is output 0 -> verifies the + # per-chunk start accumulation (out_index 0 is all the framework checks). + Case( + name="offset_chunk1_first", + construct={"sizes": [2, 3, 3], "dim": 1, "out_order": [1, 0, 2]}, + inputs=(InputSpec(shape=(1, 8, 4, 4), gen=_split_det_input),), + ), + ], + golden_dtype="float32", + atol=1e-4, + rtol=1e-3, + ) + + from executorch.backends.webgpu.test.ops.test_max_pool2d import ( _det_input as _maxpool_det_input, MaxPool2dModule, diff --git a/backends/webgpu/test/ops/test_split_with_sizes_copy.py b/backends/webgpu/test/ops/test_split_with_sizes_copy.py new file mode 100644 index 00000000000..c1771039e41 --- /dev/null +++ b/backends/webgpu/test/ops/test_split_with_sizes_copy.py @@ -0,0 +1,39 @@ +# 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. + +"""`aten.split_with_sizes_copy.default` module for the WebGPU op-test framework. + +split_with_sizes is on YOLO's Detect head (splitting the concatenated +box/objectness/class predictions). `torch.split` by a size list lowers to +`split_with_sizes_copy` in the edge dialect. +""" + +import torch + + +class SplitWithSizesModule(torch.nn.Module): + def __init__(self, sizes, dim: int, out_order=None) -> None: + super().__init__() + self.sizes = sizes + self.dim = dim + # Reorder the returned chunks so a non-first chunk can be output 0 (the + # op-test framework compares out_index 0) -- exercises the running offset. + self.out_order = out_order + + def forward(self, x: torch.Tensor): + chunks = torch.split(x, self.sizes, self.dim) + if self.out_order is not None: + return tuple(chunks[i] for i in self.out_order) + return chunks + + +def _det_input(shape) -> torch.Tensor: + # ((i % 23) - 11) / 16: exact in fp32, spans negatives through positives. + n = 1 + for s in shape: + n *= s + idx = torch.arange(n, dtype=torch.int64) + return (((idx % 23) - 11).to(torch.float32) / 16.0).reshape(shape)