feat(cuda): add canonical GELU and sigmoid providers

src/base/gelu.h

@@ -1,6 +1,7 @@
 #ifndef INFINI_OPS_BASE_GELU_H_
 #define INFINI_OPS_BASE_GELU_H_
 
+#include <cassert>
 #include <string>
 
 #include "operator.h"
@@ -17,7 +18,14 @@ class Gelu : public Operator<Gelu> {
         out_strides_{out.strides()},
         out_type_{out.dtype()},
         approximate_{approximate},
-        device_index_{out.device().index()} {}
+        output_size_{out.numel()},
+        ndim_{out.ndim()},
+        is_input_contiguous_{input.IsContiguous()},
+        is_out_contiguous_{out.IsContiguous()},
+        device_index_{out.device().index()} {
+    assert((approximate == "none" || approximate == "tanh") &&
+           "`Gelu` requires `approximate` to be `none` or `tanh`");
+  }
 
   virtual void operator()(const Tensor input, const std::string approximate,
                           Tensor out) const = 0;
@@ -37,6 +45,14 @@ class Gelu : public Operator<Gelu> {
 
   std::string approximate_{};
 
+  Tensor::Size output_size_{0};
+
+  Tensor::Size ndim_{0};
+
+  bool is_input_contiguous_{false};
+
+  bool is_out_contiguous_{false};
+
   int device_index_{0};
 };
 

src/base/sigmoid.h

@@ -14,6 +14,10 @@ class Sigmoid : public Operator<Sigmoid> {
         out_shape_{out.shape()},
         out_strides_{out.strides()},
         out_type_{out.dtype()},
+        output_size_{out.numel()},
+        ndim_{out.ndim()},
+        is_input_contiguous_{input.IsContiguous()},
+        is_out_contiguous_{out.IsContiguous()},
         device_index_{out.device().index()} {}
 
   virtual void operator()(const Tensor input, Tensor out) const = 0;
@@ -31,6 +35,14 @@ class Sigmoid : public Operator<Sigmoid> {
 
   DataType out_type_;
 
+  Tensor::Size output_size_{0};
+
+  Tensor::Size ndim_{0};
+
+  bool is_input_contiguous_{false};
+
+  bool is_out_contiguous_{false};
+
   int device_index_{0};
 };
 

src/native/cuda/iluvatar/ops/gelu/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_ILUVATAR_GELU_KERNEL_H_
+#define INFINI_OPS_ILUVATAR_GELU_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/iluvatar/caster.cuh"
+#include "native/cuda/iluvatar/runtime_.h"
+#include "native/cuda/ops/gelu/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Gelu, Device::Type::kIluvatar>
+    : public CudaGelu<Runtime<Device::Type::kIluvatar>> {
+ public:
+  using CudaGelu<Runtime<Device::Type::kIluvatar>>::CudaGelu;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/iluvatar/ops/sigmoid/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_ILUVATAR_SIGMOID_KERNEL_H_
+#define INFINI_OPS_ILUVATAR_SIGMOID_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/iluvatar/caster.cuh"
+#include "native/cuda/iluvatar/runtime_.h"
+#include "native/cuda/ops/sigmoid/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Sigmoid, Device::Type::kIluvatar>
+    : public CudaSigmoid<Runtime<Device::Type::kIluvatar>> {
+ public:
+  using CudaSigmoid<Runtime<Device::Type::kIluvatar>>::CudaSigmoid;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/metax/ops/gelu/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_METAX_GELU_KERNEL_H_
+#define INFINI_OPS_METAX_GELU_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/metax/caster.cuh"
+#include "native/cuda/metax/runtime_.h"
+#include "native/cuda/ops/gelu/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Gelu, Device::Type::kMetax>
+    : public CudaGelu<Runtime<Device::Type::kMetax>> {
+ public:
+  using CudaGelu<Runtime<Device::Type::kMetax>>::CudaGelu;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/metax/ops/sigmoid/kernel.h

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_METAX_SIGMOID_KERNEL_H_
+#define INFINI_OPS_METAX_SIGMOID_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/metax/caster.cuh"
+#include "native/cuda/metax/runtime_.h"
+#include "native/cuda/ops/sigmoid/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Sigmoid, Device::Type::kMetax>
+    : public CudaSigmoid<Runtime<Device::Type::kMetax>> {
+ public:
+  using CudaSigmoid<Runtime<Device::Type::kMetax>>::CudaSigmoid;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/moore/ops/gelu/kernel.h

@@ -0,0 +1,22 @@
+#ifndef INFINI_OPS_MOORE_GELU_KERNEL_H_
+#define INFINI_OPS_MOORE_GELU_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/moore/caster.cuh"
+#include "native/cuda/moore/polyfills.cuh"
+#include "native/cuda/moore/runtime_.h"
+#include "native/cuda/ops/gelu/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Gelu, Device::Type::kMoore>
+    : public CudaGelu<Runtime<Device::Type::kMoore>> {
+ public:
+  using CudaGelu<Runtime<Device::Type::kMoore>>::CudaGelu;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/moore/ops/sigmoid/kernel.h

@@ -0,0 +1,22 @@
+#ifndef INFINI_OPS_MOORE_SIGMOID_KERNEL_H_
+#define INFINI_OPS_MOORE_SIGMOID_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/moore/caster.cuh"
+#include "native/cuda/moore/polyfills.cuh"
+#include "native/cuda/moore/runtime_.h"
+#include "native/cuda/ops/sigmoid/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Sigmoid, Device::Type::kMoore>
+    : public CudaSigmoid<Runtime<Device::Type::kMoore>> {
+ public:
+  using CudaSigmoid<Runtime<Device::Type::kMoore>>::CudaSigmoid;
+};
+
+}  // namespace infini::ops
+
+#endif

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

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_NVIDIA_GELU_KERNEL_H_
+#define INFINI_OPS_NVIDIA_GELU_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/nvidia/caster.cuh"
+#include "native/cuda/nvidia/runtime_.h"
+#include "native/cuda/ops/gelu/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Gelu, Device::Type::kNvidia>
+    : public CudaGelu<Runtime<Device::Type::kNvidia>> {
+ public:
+  using CudaGelu<Runtime<Device::Type::kNvidia>>::CudaGelu;
+};
+
+}  // namespace infini::ops
+
+#endif

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

@@ -0,0 +1,21 @@
+#ifndef INFINI_OPS_NVIDIA_SIGMOID_KERNEL_H_
+#define INFINI_OPS_NVIDIA_SIGMOID_KERNEL_H_
+
+#include <utility>
+
+#include "native/cuda/nvidia/caster.cuh"
+#include "native/cuda/nvidia/runtime_.h"
+#include "native/cuda/ops/sigmoid/kernel.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<Sigmoid, Device::Type::kNvidia>
+    : public CudaSigmoid<Runtime<Device::Type::kNvidia>> {
+ public:
+  using CudaSigmoid<Runtime<Device::Type::kNvidia>>::CudaSigmoid;
+};
+
+}  // namespace infini::ops
+
+#endif

src/native/cuda/ops/gelu/kernel.h

@@ -0,0 +1,127 @@
+#ifndef INFINI_OPS_CUDA_GELU_KERNEL_H_
+#define INFINI_OPS_CUDA_GELU_KERNEL_H_
+
+#include <algorithm>
+#include <cassert>
+#include <cstddef>
+#include <cstring>
+#include <vector>
+
+#include "base/gelu.h"
+#include "common/generic_utils.h"
+#include "data_type.h"
+#include "dispatcher.h"
+#include "native/cuda/kernel_commons.cuh"
+#include "native/cuda/ops/gelu_infinilm/kernel.cuh"
+#include "native/cuda/ops/gelutanh_infinilm/kernel.cuh"
+#include "native/cuda/runtime_utils.h"
+
+namespace infini::ops {
+
+template <typename Backend>
+class CudaGelu : public Gelu {
+ public:
+  CudaGelu(const Tensor input, const std::string approximate, Tensor out)
+      : Gelu{input, approximate, out} {
+    assert(input_shape_ == out_shape_ &&
+           "`Gelu` input and output shapes must match");
+    assert(input_type_ == out_type_ &&
+           "`Gelu` input and output dtypes must match");
+    assert(input.device() == out.device() &&
+           "`Gelu` input and output devices must match");
+    assert(!out.HasBroadcastDim() &&
+           "`Gelu` output must not have broadcasted dimensions");
+
+    const size_t shape_size = ndim_ * sizeof(*d_input_shape_);
+    const size_t strides_size = ndim_ * sizeof(*d_input_strides_);
+    const size_t metadata_size = 2 * (shape_size + strides_size);
+    if (metadata_size == 0) {
+      return;
+    }
+
+    std::vector<std::byte> metadata(metadata_size);
+    Backend::Malloc(reinterpret_cast<void**>(&d_metadata_), metadata_size);
+
+    size_t offset = 0;
+    d_input_shape_ = reinterpret_cast<Tensor::Size*>(d_metadata_ + offset);
+    std::memcpy(metadata.data() + offset, input_shape_.data(), shape_size);
+    offset += shape_size;
+
+    d_out_shape_ = reinterpret_cast<Tensor::Size*>(d_metadata_ + offset);
+    std::memcpy(metadata.data() + offset, out_shape_.data(), shape_size);
+    offset += shape_size;
+
+    d_input_strides_ = reinterpret_cast<Tensor::Stride*>(d_metadata_ + offset);
+    std::memcpy(metadata.data() + offset, input_strides_.data(), strides_size);
+    offset += strides_size;
+
+    d_out_strides_ = reinterpret_cast<Tensor::Stride*>(d_metadata_ + offset);
+    std::memcpy(metadata.data() + offset, out_strides_.data(), strides_size);
+
+    Backend::Memcpy(d_metadata_, metadata.data(), metadata_size,
+                    Backend::kMemcpyHostToDevice);
+  }
+
+  ~CudaGelu() {
+    if (d_metadata_ != nullptr) {
+      Backend::Free(d_metadata_);
+    }
+  }
+
+  void operator()(const Tensor input, const std::string approximate,
+                  Tensor out) const override {
+    assert(approximate == approximate_ &&
+           "`Gelu` attributes changed after descriptor creation");
+    if (output_size_ == 0) {
+      return;
+    }
+
+    auto cuda_stream =
+        static_cast<typename Backend::Stream>(stream_ ? stream_ : 0);
+    int block_size = std::min(
+        RuntimeUtils<Backend::kDeviceType>::GetOptimalBlockSize(), 1024);
+    dim3 block(std::min(static_cast<Tensor::Size>(block_size), output_size_));
+    dim3 grid(utils::CeilDiv(output_size_, block.x));
+
+    DispatchFunc<AllFloatTypes, List<128, 256, 512, 1024>>(
+        {static_cast<int64_t>(out_type_), block_size},
+        [&](auto list_tag) {
+          using T = TypeMapType<Backend::kDeviceType, ListGet<0>(list_tag)>;
+          constexpr int kBlockSize = ListGet<1>(list_tag);
+
+          if (approximate_ == "tanh") {
+            GelutanhInfinilmKernel<Backend::kDeviceType, T, kBlockSize>
+                <<<grid, block, 0, cuda_stream>>>(
+                    reinterpret_cast<T*>(out.data()),
+                    reinterpret_cast<const T*>(input.data()), d_out_shape_,
+                    d_input_shape_, d_out_strides_, d_input_strides_,
+                    output_size_, ndim_, is_out_contiguous_,
+                    is_input_contiguous_);
+          } else {
+            GeluInfinilmKernel<Backend::kDeviceType, T, kBlockSize>
+                <<<grid, block, 0, cuda_stream>>>(
+                    reinterpret_cast<T*>(out.data()),
+                    reinterpret_cast<const T*>(input.data()), d_out_shape_,
+                    d_input_shape_, d_out_strides_, d_input_strides_,
+                    output_size_, ndim_, is_out_contiguous_,
+                    is_input_contiguous_);
+          }
+        },
+        "CudaGelu::operator()");
+  }
+
+ private:
+  std::byte* d_metadata_{nullptr};
+
+  Tensor::Size* d_input_shape_{nullptr};
+
+  Tensor::Size* d_out_shape_{nullptr};
+
+  Tensor::Stride* d_input_strides_{nullptr};
+
+  Tensor::Stride* d_out_strides_{nullptr};
+};
+
+}  // namespace infini::ops
+
+#endif