feat(nvidia): add Argmax provider

src/native/cuda/nvidia/ops/argmax/kernel.cu

@@ -0,0 +1,91 @@
+#include "native/cuda/nvidia/ops/argmax/kernel.h"
+
+#include <cassert>
+#include <climits>
+#include <cstddef>
+#include <cstdint>
+
+#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<Argmax, Device::Type::kNvidia>::Operator(
+    const Tensor input, const std::optional<int64_t> 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<Device::Type::kNvidia>::Malloc(&default_workspace_,
+                                                       workspace_size_);
+  assert(error == cudaSuccess &&
+         "NVIDIA `Argmax` failed to allocate workspace");
+}
+
+Operator<Argmax, Device::Type::kNvidia>::~Operator() {
+  auto error = Runtime<Device::Type::kNvidia>::Free(default_workspace_);
+  assert(error == cudaSuccess && "NVIDIA `Argmax` failed to free workspace");
+}
+
+std::size_t
+Operator<Argmax, Device::Type::kNvidia>::workspace_size_in_bytes() const {
+  return workspace_size_;
+}
+
+void Operator<Argmax, Device::Type::kNvidia>::operator()(
+    const Tensor input, const std::optional<int64_t> 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<cudaStream_t>(stream_);
+
+  DispatchFunc<Device::Type::kNvidia,
+               ConcatType<FloatTypes, ReducedFloatTypes>>(
+      input.dtype(),
+      [&](auto dtype_tag) {
+        using T = typename decltype(dtype_tag)::type;
+        argmax_detail::Launch(workspace, workspace_size,
+                              static_cast<const T*>(input.data()), numel_,
+                              static_cast<int64_t*>(out.data()), stream);
+      },
+      "NVIDIA Argmax");
+}
+
+std::size_t Operator<Argmax, Device::Type::kNvidia>::DispatchWorkspaceSize(
+    DataType dtype, std::size_t numel) {
+  std::size_t workspace_size = 0;
+  DispatchFunc<Device::Type::kNvidia,
+               ConcatType<FloatTypes, ReducedFloatTypes>>(
+      dtype,
+      [&](auto dtype_tag) {
+        using T = typename decltype(dtype_tag)::type;
+        workspace_size = argmax_detail::WorkspaceSize<T>(numel);
+      },
+      "NVIDIA Argmax workspace");
+
+  return workspace_size;
+}
+
+}  // namespace infini::ops

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 <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <cub/device/device_reduce.cuh>
+
+namespace infini::ops::argmax_detail {
+
+constexpr std::size_t Align256(std::size_t size) {
+  return (size + 255) & ~std::size_t{255};
+}
+
+template <typename T>
+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<const T*>(nullptr),
+      static_cast<cub::KeyValuePair<int, T>*>(nullptr),
+      static_cast<int>(numel));
+  assert(error == cudaSuccess && "`Argmax` failed to query CUB workspace");
+
+  return Align256(sizeof(cub::KeyValuePair<int, T>)) + cub_workspace_size;
+}
+
+template <typename T>
+__global__ void StoreIndex(int64_t* out,
+                           const cub::KeyValuePair<int, T>* result) {
+  *out = static_cast<int64_t>(result->key);
+}
+
+template <typename T>
+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<cub::KeyValuePair<int, T>*>(workspace);
+  auto* cub_workspace = static_cast<char*>(workspace) +
+                        Align256(sizeof(cub::KeyValuePair<int, T>));
+  auto cub_workspace_size =
+      workspace_size - Align256(sizeof(cub::KeyValuePair<int, T>));
+  auto error =
+      cub::DeviceReduce::ArgMax(cub_workspace, cub_workspace_size, input,
+                                result, static_cast<int>(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_

src/native/cuda/nvidia/ops/argmax/kernel.h

@@ -0,0 +1,37 @@
+#ifndef INFINI_OPS_NVIDIA_ARGMAX_KERNEL_H_
+#define INFINI_OPS_NVIDIA_ARGMAX_KERNEL_H_
+
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+
+#include "base/argmax.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Argmax, Device::Type::kNvidia> : public Argmax {
+ public:
+  Operator(const Tensor input, const std::optional<int64_t> 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<int64_t> 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_

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()