From 8875fb5a02630f61481f9ab6669a7ff111292ff5 Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 10 Aug 2026 23:39:47 +0800 Subject: [PATCH 1/2] feat(nvidia): add Argmax provider --- src/native/cuda/nvidia/ops/argmax/kernel.cu | 91 ++++++++++++++++++++ src/native/cuda/nvidia/ops/argmax/kernel.cuh | 50 +++++++++++ src/native/cuda/nvidia/ops/argmax/kernel.h | 38 ++++++++ tests/test_argmax.py | 30 +++++++ 4 files changed, 209 insertions(+) create mode 100644 src/native/cuda/nvidia/ops/argmax/kernel.cu create mode 100644 src/native/cuda/nvidia/ops/argmax/kernel.cuh create mode 100644 src/native/cuda/nvidia/ops/argmax/kernel.h create mode 100644 tests/test_argmax.py diff --git a/src/native/cuda/nvidia/ops/argmax/kernel.cu b/src/native/cuda/nvidia/ops/argmax/kernel.cu new file mode 100644 index 000000000..5e5ef763e --- /dev/null +++ b/src/native/cuda/nvidia/ops/argmax/kernel.cu @@ -0,0 +1,91 @@ +#include "native/cuda/nvidia/ops/argmax/kernel.h" + +#include +#include +#include +#include + +#include "data_type.h" +#include "dispatcher.h" +#include "native/cuda/nvidia/caster.cuh" +#include "native/cuda/nvidia/ops/argmax/kernel.cuh" +#include "native/cuda/nvidia/runtime_.h" + +namespace infini::ops { + +Operator::Operator( + const Tensor input, const std::optional dim, const bool keepdim, + Tensor out) + : Argmax{input, dim, keepdim, out}, + numel_{input.numel()}, + workspace_size_{DispatchWorkspaceSize(input.dtype(), input.numel())} { + assert(!dim.has_value() && !keepdim && + "NVIDIA `Argmax` currently supports only flattened reduction"); + assert(input.IsContiguous() && + "NVIDIA `Argmax` requires contiguous input"); + assert(input.numel() > 0 && input.numel() <= INT_MAX && + "NVIDIA `Argmax` input size must fit in a positive `int`"); + assert(out.numel() == 1 && out.dtype() == DataType::kInt64 && + "NVIDIA `Argmax` requires one `int64` output value"); + assert(input.device() == out.device() && + "NVIDIA `Argmax` input and output must be on the same device"); + + auto error = Runtime::Malloc(&default_workspace_, + workspace_size_); + assert(error == cudaSuccess && + "NVIDIA `Argmax` failed to allocate workspace"); +} + +Operator::~Operator() { + auto error = Runtime::Free(default_workspace_); + assert(error == cudaSuccess && "NVIDIA `Argmax` failed to free workspace"); +} + +std::size_t +Operator::workspace_size_in_bytes() const { + return workspace_size_; +} + +void Operator::operator()( + const Tensor input, const std::optional dim, const bool keepdim, + Tensor out) const { + assert(!dim.has_value() && !keepdim); + assert(input.shape() == input_shape_ && input.strides() == input_strides_ + && input.dtype() == input_type_); + assert(out.shape() == out_shape_ && out.strides() == out_strides_ + && out.dtype() == out_type_); + + void* workspace = workspace_ ? workspace_ : default_workspace_; + auto workspace_size = workspace_ ? workspace_size_in_bytes_ : workspace_size_; + assert(workspace != nullptr && workspace_size >= workspace_size_ && + "NVIDIA `Argmax` received insufficient workspace"); + auto stream = reinterpret_cast(stream_); + + DispatchFunc>( + input.dtype(), + [&](auto dtype_tag) { + using T = typename decltype(dtype_tag)::type; + argmax_detail::Launch(workspace, workspace_size, + static_cast(input.data()), numel_, + static_cast(out.data()), stream); + }, + "NVIDIA Argmax"); +} + +std::size_t Operator::DispatchWorkspaceSize( + DataType dtype, std::size_t numel) { + std::size_t workspace_size = 0; + DispatchFunc>( + dtype, + [&](auto dtype_tag) { + using T = typename decltype(dtype_tag)::type; + workspace_size = argmax_detail::WorkspaceSize(numel); + }, + "NVIDIA Argmax workspace"); + + return workspace_size; +} + +} // namespace infini::ops diff --git a/src/native/cuda/nvidia/ops/argmax/kernel.cuh b/src/native/cuda/nvidia/ops/argmax/kernel.cuh new file mode 100644 index 000000000..fafaa247a --- /dev/null +++ b/src/native/cuda/nvidia/ops/argmax/kernel.cuh @@ -0,0 +1,50 @@ +#ifndef INFINI_OPS_NVIDIA_ARGMAX_KERNEL_CUH_ +#define INFINI_OPS_NVIDIA_ARGMAX_KERNEL_CUH_ + +#include +#include +#include +#include + +namespace infini::ops::argmax_detail { + +constexpr std::size_t Align256(std::size_t size) { + return (size + 255) & ~std::size_t{255}; +} + +template +std::size_t WorkspaceSize(std::size_t numel) { + std::size_t cub_workspace_size = 0; + auto error = cub::DeviceReduce::ArgMax( + nullptr, cub_workspace_size, static_cast(nullptr), + static_cast*>(nullptr), + static_cast(numel)); + assert(error == cudaSuccess && "`Argmax` failed to query CUB workspace"); + + return Align256(sizeof(cub::KeyValuePair)) + cub_workspace_size; +} + +template +__global__ void StoreIndex( + int64_t* out, const cub::KeyValuePair* result) { + *out = static_cast(result->key); +} + +template +void Launch(void* workspace, std::size_t workspace_size, const T* input, + std::size_t numel, int64_t* out, cudaStream_t stream) { + auto* result = static_cast*>(workspace); + auto* cub_workspace = static_cast(workspace) + + Align256(sizeof(cub::KeyValuePair)); + auto cub_workspace_size = workspace_size + - Align256(sizeof(cub::KeyValuePair)); + auto error = cub::DeviceReduce::ArgMax( + cub_workspace, cub_workspace_size, input, result, + static_cast(numel), stream); + assert(error == cudaSuccess && "`Argmax` CUB reduction failed"); + StoreIndex<<<1, 1, 0, stream>>>(out, result); +} + +} // namespace infini::ops::argmax_detail + +#endif // INFINI_OPS_NVIDIA_ARGMAX_KERNEL_CUH_ diff --git a/src/native/cuda/nvidia/ops/argmax/kernel.h b/src/native/cuda/nvidia/ops/argmax/kernel.h new file mode 100644 index 000000000..d94756ea0 --- /dev/null +++ b/src/native/cuda/nvidia/ops/argmax/kernel.h @@ -0,0 +1,38 @@ +#ifndef INFINI_OPS_NVIDIA_ARGMAX_KERNEL_H_ +#define INFINI_OPS_NVIDIA_ARGMAX_KERNEL_H_ + +#include +#include +#include + +#include "base/argmax.h" + +namespace infini::ops { + +template <> +class Operator : public Argmax { + public: + Operator(const Tensor input, const std::optional dim, + const bool keepdim, Tensor out); + + ~Operator() override; + + std::size_t workspace_size_in_bytes() const override; + + void operator()(const Tensor input, const std::optional dim, + const bool keepdim, Tensor out) const override; + + private: + static std::size_t DispatchWorkspaceSize(DataType dtype, + std::size_t numel); + + std::size_t numel_{0}; + + std::size_t workspace_size_{0}; + + void* default_workspace_{nullptr}; +}; + +} // namespace infini::ops + +#endif // INFINI_OPS_NVIDIA_ARGMAX_KERNEL_H_ diff --git a/tests/test_argmax.py b/tests/test_argmax.py new file mode 100644 index 000000000..dddce2e7f --- /dev/null +++ b/tests/test_argmax.py @@ -0,0 +1,30 @@ +import infini.ops +import pytest +import torch + +from tests.utils import get_stream + + +@pytest.mark.parametrize( + "dtype", + (torch.float32, torch.float16, torch.bfloat16), +) +def test_argmax_flattened(dtype, device, implementation_index): + if device != "cuda": + pytest.skip("argmax requires the NVIDIA backend") + + input = torch.randn(32_003, dtype=dtype, device=device) + input[17_391] = 100 + out = torch.full((), -1, dtype=torch.int64, device=device) + + result = infini.ops.argmax( + input, + None, + False, + out, + stream=get_stream(input.device), + implementation_index=implementation_index, + ) + + assert result is None + assert out.item() == torch.argmax(input).item() From c5ba52cbcfafccc7aa6ab41fcf46fa60dec6eece Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Tue, 11 Aug 2026 02:11:56 +0800 Subject: [PATCH 2/2] style(nvidia): format Argmax with clang-format 21 --- src/native/cuda/nvidia/ops/argmax/kernel.cuh | 18 +++++++++--------- src/native/cuda/nvidia/ops/argmax/kernel.h | 3 +-- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/src/native/cuda/nvidia/ops/argmax/kernel.cuh b/src/native/cuda/nvidia/ops/argmax/kernel.cuh index fafaa247a..2dd7e04da 100644 --- a/src/native/cuda/nvidia/ops/argmax/kernel.cuh +++ b/src/native/cuda/nvidia/ops/argmax/kernel.cuh @@ -25,8 +25,8 @@ std::size_t WorkspaceSize(std::size_t numel) { } template -__global__ void StoreIndex( - int64_t* out, const cub::KeyValuePair* result) { +__global__ void StoreIndex(int64_t* out, + const cub::KeyValuePair* result) { *out = static_cast(result->key); } @@ -34,13 +34,13 @@ template void Launch(void* workspace, std::size_t workspace_size, const T* input, std::size_t numel, int64_t* out, cudaStream_t stream) { auto* result = static_cast*>(workspace); - auto* cub_workspace = static_cast(workspace) - + Align256(sizeof(cub::KeyValuePair)); - auto cub_workspace_size = workspace_size - - Align256(sizeof(cub::KeyValuePair)); - auto error = cub::DeviceReduce::ArgMax( - cub_workspace, cub_workspace_size, input, result, - static_cast(numel), stream); + auto* cub_workspace = static_cast(workspace) + + Align256(sizeof(cub::KeyValuePair)); + auto cub_workspace_size = + workspace_size - Align256(sizeof(cub::KeyValuePair)); + auto error = + cub::DeviceReduce::ArgMax(cub_workspace, cub_workspace_size, input, + result, static_cast(numel), stream); assert(error == cudaSuccess && "`Argmax` CUB reduction failed"); StoreIndex<<<1, 1, 0, stream>>>(out, result); } diff --git a/src/native/cuda/nvidia/ops/argmax/kernel.h b/src/native/cuda/nvidia/ops/argmax/kernel.h index d94756ea0..006ac3d05 100644 --- a/src/native/cuda/nvidia/ops/argmax/kernel.h +++ b/src/native/cuda/nvidia/ops/argmax/kernel.h @@ -23,8 +23,7 @@ class Operator : public Argmax { const bool keepdim, Tensor out) const override; private: - static std::size_t DispatchWorkspaceSize(DataType dtype, - std::size_t numel); + static std::size_t DispatchWorkspaceSize(DataType dtype, std::size_t numel); std::size_t numel_{0};