Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

webgpu ep support for argmax/argmin #24089

Merged
merged 6 commits into from
Mar 19, 2025
Merged

webgpu ep support for argmax/argmin #24089

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
316de54
argmin/argmax
guschmue Mar 18, 2025
d51cfb7
no multiaxis
guschmue Mar 18, 2025
e4f62ea
support for select_last_index
guschmue Mar 18, 2025
4444e2e
support for select_last_index
guschmue Mar 18, 2025
b5550b8
Merge branch 'main' into gs/argminmax
guschmue Mar 18, 2025
f365965
reflect changes in mail
guschmue Mar 19, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 31 additions & 1 deletion onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,14 @@ REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ReduceLogSumExp, 11, 12);
REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ReduceLogSumExp, 13, 17);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceLogSumExp, 18);

REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ArgMax, 1, 10);
REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ArgMax, 11, 12);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ArgMax, 13);

REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ArgMin, 1, 10);
REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ArgMin, 11, 12);
REGISTER_UNARY_ELEMENTWISE_KERNEL(ArgMin, 13);

Status ReduceKernelProgram::GenerateShaderCode(ShaderHelper& shader) const {
const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias | ShaderUsage::UseValueTypeAlias);
if (is_input_empty_) {
Expand All @@ -105,6 +113,9 @@ Status ReduceKernelProgram::GenerateShaderCode(ShaderHelper& shader) const {
std::stringstream ss;
std::string index = "i" + std::to_string(i);
ss << "for (var " << index << " : u32 = 0; " << index << " < " << input.IndicesGet("uniforms.input_shape", i) << "; " << index << "++) {\n";
if (loop_body.find("last_index") != std::string::npos) {
ss << "let last_index = " + index + ";\n";
}
ss << input.IndicesSet("input_indices", i, index) << ";\n";
ss << loop_body << "\n";
ss << "}\n";
Expand Down Expand Up @@ -328,5 +339,24 @@ ReduceOpSpecificCode ReduceLogSumExp::GetOpSpecificCode(const Tensor* input_tens
return code;
}

ReduceOpSpecificCode ArgMin::GetOpSpecificCode(const Tensor* input_tensor) const {
ORT_UNUSED_PARAMETER(input_tensor);
std::string loop_header = "var best_element = first_element; var best_index = u32(0);";
std::string loop_body = "if (current_element < best_element) { best_element = current_element; best_index = last_index; };";
std::string loop_footer = "let output_value = output_value_t(best_index);";
ReduceOpSpecificCode code({loop_header, loop_body, loop_footer});
return code;
}

ReduceOpSpecificCode ArgMax::GetOpSpecificCode(const Tensor* input_tensor) const {
ORT_UNUSED_PARAMETER(input_tensor);
std::string loop_header = "var best_element = first_element; var best_index = u32(0);";
std::string loop_body = "if (current_element > best_element) { best_element = current_element; best_index = last_index; };";
std::string loop_footer = "let output_value = output_value_t(best_index);";
ReduceOpSpecificCode code({loop_header, loop_body, loop_footer});
return code;
}


} // namespace webgpu
} // namespace onnxruntime
} // namespace onnxruntime
12 changes: 12 additions & 0 deletions onnxruntime/core/providers/webgpu/reduction/reduction_ops.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,5 +119,17 @@ class ReduceLogSumExp final : public ReduceKernel<true> {
ReduceOpSpecificCode GetOpSpecificCode(const Tensor* input_tensor) const override;
};

class ArgMin final : public ReduceKernel<true> {
public:
ArgMin(const OpKernelInfo& info) : ReduceKernel<true>(info, "ArgMin", true) {}
ReduceOpSpecificCode GetOpSpecificCode(const Tensor* input_tensor) const override;
};

class ArgMax final : public ReduceKernel<true> {
public:
ArgMax(const OpKernelInfo& info) : ReduceKernel<true>(info, "ArgMax", true) {}
ReduceOpSpecificCode GetOpSpecificCode(const Tensor* input_tensor) const override;
};

} // namespace webgpu
} // namespace onnxruntime
24 changes: 12 additions & 12 deletions onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -297,12 +297,12 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13,
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 12, MatMul);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, MatMul);

class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, float, ArgMax);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, float, ArgMax);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, float, ArgMax);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, float, ArgMin);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, float, ArgMin);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, float, ArgMin);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, ArgMax);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, ArgMax);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, ArgMax);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, ArgMin);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, ArgMin);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, ArgMin);

class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, Softmax);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, Softmax);
Expand Down Expand Up @@ -624,13 +624,13 @@ std::unique_ptr<KernelRegistry> RegisterKernels() {
// BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 12, MatMul)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, MatMul)>,

// BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, float, ArgMax)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, float, ArgMax)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, float, ArgMax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, ArgMax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, ArgMax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, ArgMax)>,

// BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, float, ArgMin)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, float, ArgMin)>,
// BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, float, ArgMin)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, ArgMin)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, ArgMin)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 13, ArgMin)>,

BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 1, 10, Softmax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 11, 12, Softmax)>,
Expand Down