Implement T1-1-30: interpolate interpolate gaussian_nll_loss gaussia…

include/infiniop/ops/bitwise_right_shift.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_BITWISE_RIGHT_SHIFT_API_H__
+#define __INFINIOP_BITWISE_RIGHT_SHIFT_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopBitwiseRightShiftDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateBitwiseRightShiftDescriptor(infiniopHandle_t handle,
+                                                                      infiniopBitwiseRightShiftDescriptor_t *desc_ptr,
+                                                                      infiniopTensorDescriptor_t y,
+                                                                      infiniopTensorDescriptor_t x1,
+                                                                      infiniopTensorDescriptor_t x2);
+
+__C __export infiniStatus_t infiniopGetBitwiseRightShiftWorkspaceSize(infiniopBitwiseRightShiftDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopBitwiseRightShift(infiniopBitwiseRightShiftDescriptor_t desc,
+                                                      void *workspace,
+                                                      size_t workspace_size,
+                                                      void *y,
+                                                      const void *x1,
+                                                      const void *x2,
+                                                      void *stream);
+
+__C __export infiniStatus_t infiniopDestroyBitwiseRightShiftDescriptor(infiniopBitwiseRightShiftDescriptor_t desc);
+
+#endif

include/infiniop/ops/gaussian_nll_loss.h

@@ -0,0 +1,31 @@
+#ifndef __INFINIOP_GAUSSIAN_NLL_LOSS_API_H__
+#define __INFINIOP_GAUSSIAN_NLL_LOSS_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopGaussianNllLossDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateGaussianNllLossDescriptor(infiniopHandle_t handle,
+                                                                    infiniopGaussianNllLossDescriptor_t *desc_ptr,
+                                                                    infiniopTensorDescriptor_t y,
+                                                                    infiniopTensorDescriptor_t input,
+                                                                    infiniopTensorDescriptor_t target,
+                                                                    infiniopTensorDescriptor_t var,
+                                                                    int full,
+                                                                    double eps,
+                                                                    int reduction);
+
+__C __export infiniStatus_t infiniopGetGaussianNllLossWorkspaceSize(infiniopGaussianNllLossDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopGaussianNllLoss(infiniopGaussianNllLossDescriptor_t desc,
+                                                    void *workspace,
+                                                    size_t workspace_size,
+                                                    void *y,
+                                                    const void *input,
+                                                    const void *target,
+                                                    const void *var,
+                                                    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyGaussianNllLossDescriptor(infiniopGaussianNllLossDescriptor_t desc);
+
+#endif

include/infiniop/ops/interpolate.h

@@ -0,0 +1,28 @@
+#ifndef __INFINIOP_INTERPOLATE_API_H__
+#define __INFINIOP_INTERPOLATE_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopInterpolateDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateInterpolateDescriptor(infiniopHandle_t handle,
+                                                                infiniopInterpolateDescriptor_t *desc_ptr,
+                                                                infiniopTensorDescriptor_t y,
+                                                                infiniopTensorDescriptor_t x,
+                                                                const char *mode,
+                                                                void *size,
+                                                                void *scale_factor,
+                                                                int align_corners);
+
+__C __export infiniStatus_t infiniopGetInterpolateWorkspaceSize(infiniopInterpolateDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopInterpolate(infiniopInterpolateDescriptor_t desc,
+                                                void *workspace,
+                                                size_t workspace_size,
+                                                void *y,
+                                                const void *x,
+                                                void *stream);
+
+__C __export infiniStatus_t infiniopDestroyInterpolateDescriptor(infiniopInterpolateDescriptor_t desc);
+
+#endif

include/infiniop/ops/prelu.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_PRELU_API_H__
+#define __INFINIOP_PRELU_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopPreluDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreatePreluDescriptor(infiniopHandle_t handle,
+                                                          infiniopPreluDescriptor_t *desc_ptr,
+                                                          infiniopTensorDescriptor_t y,
+                                                          infiniopTensorDescriptor_t x,
+                                                          infiniopTensorDescriptor_t weight);
+
+__C __export infiniStatus_t infiniopGetPreluWorkspaceSize(infiniopPreluDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopPrelu(infiniopPreluDescriptor_t desc,
+                                          void *workspace,
+                                          size_t workspace_size,
+                                          void *y,
+                                          const void *x,
+                                          const void *weight,
+                                          void *stream);
+
+__C __export infiniStatus_t infiniopDestroyPreluDescriptor(infiniopPreluDescriptor_t desc);
+
+#endif

include/infiniop/ops/relu6.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_RELU6_API_H__
+#define __INFINIOP_RELU6_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopRelu6Descriptor_t;
+
+__C __export infiniStatus_t infiniopCreateRelu6Descriptor(infiniopHandle_t handle,
+                                                          infiniopRelu6Descriptor_t *desc_ptr,
+                                                          infiniopTensorDescriptor_t y,
+                                                          infiniopTensorDescriptor_t x);
+
+__C __export infiniStatus_t infiniopGetRelu6WorkspaceSize(infiniopRelu6Descriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopRelu6(infiniopRelu6Descriptor_t desc,
+                                          void *workspace,
+                                          size_t workspace_size,
+                                          void *y,
+                                          const void *x,
+                                          void *stream);
+
+__C __export infiniStatus_t infiniopDestroyRelu6Descriptor(infiniopRelu6Descriptor_t desc);
+
+#endif

src/infiniop/ops/bitwise_right_shift/bitwise_right_shift.h

@@ -0,0 +1,8 @@
+#ifndef __BITWISE_RIGHT_SHIFT_H__
+#define __BITWISE_RIGHT_SHIFT_H__
+
+#include "../../elementwise/elementwise.h"
+
+#define DESCRIPTOR(NAMESPACE) ELEMENTWISE_DESCRIPTOR(bitwise_right_shift, NAMESPACE)
+
+#endif // __BITWISE_RIGHT_SHIFT_H__

src/infiniop/ops/bitwise_right_shift/cpu/bitwise_right_shift_cpu.cc

@@ -0,0 +1,61 @@
+#include "bitwise_right_shift_cpu.h"
+
+namespace op::bitwise_right_shift::cpu {
+
+Descriptor::~Descriptor() = default;
+
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle_,
+    Descriptor **desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    std::vector<infiniopTensorDescriptor_t> input_desc_vec) {
+
+    auto handle = reinterpret_cast<device::cpu::Handle *>(handle_);
+    auto dtype = out_desc->dtype();
+
+    const auto &input_desc = input_desc_vec.at(0);
+    const auto &output_shape = out_desc->shape();
+    const auto &input_shape = input_desc->shape();
+
+    CHECK_DTYPE(dtype, INFINI_DTYPE_I8, INFINI_DTYPE_I16, INFINI_DTYPE_I32, INFINI_DTYPE_I64,
+                INFINI_DTYPE_U8, INFINI_DTYPE_U16, INFINI_DTYPE_U32, INFINI_DTYPE_U64);
+
+    CHECK_SAME_SHAPE(output_shape, input_shape);
+
+    CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec);
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Descriptor::calculate(
+    void *workspace,
+    size_t workspace_size,
+    void *output,
+    std::vector<const void *> inputs,
+    void *stream) const {
+
+    switch (_dtype) {
+    case INFINI_DTYPE_I8:
+        return _device_info->calculate<BitwiseRightShiftOp, int8_t, int8_t, int8_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_I16:
+        return _device_info->calculate<BitwiseRightShiftOp, int16_t, int16_t, int16_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_I32:
+        return _device_info->calculate<BitwiseRightShiftOp, int32_t, int32_t, int32_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_I64:
+        return _device_info->calculate<BitwiseRightShiftOp, int64_t, int64_t, int64_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_U8:
+        return _device_info->calculate<BitwiseRightShiftOp, uint8_t, uint8_t, uint8_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_U16:
+        return _device_info->calculate<BitwiseRightShiftOp, uint16_t, uint16_t, uint16_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_U32:
+        return _device_info->calculate<BitwiseRightShiftOp, uint32_t, uint32_t, uint32_t>(_info, output, inputs, stream);
+    case INFINI_DTYPE_U64:
+        return _device_info->calculate<BitwiseRightShiftOp, uint64_t, uint64_t, uint64_t>(_info, output, inputs, stream);
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+} // namespace op::bitwise_right_shift::cpu

src/infiniop/ops/bitwise_right_shift/cpu/bitwise_right_shift_cpu.h

@@ -0,0 +1,19 @@
+#ifndef __BITWISE_RIGHT_SHIFT_CPU_H__
+#define __BITWISE_RIGHT_SHIFT_CPU_H__
+
+#include "../../../elementwise/cpu/elementwise_cpu.h"
+
+ELEMENTWISE_DESCRIPTOR(bitwise_right_shift, cpu)
+
+namespace op::bitwise_right_shift::cpu {
+typedef struct BitwiseRightShiftOp {
+public:
+    static constexpr size_t num_inputs = 2;
+    template <typename T>
+    T operator()(const T &x, const T &shift) const {
+        return x >> shift;
+    }
+} BitwiseRightShiftOp;
+} // namespace op::bitwise_right_shift::cpu
+
+#endif // __BITWISE_RIGHT_SHIFT_CPU_H__

src/infiniop/ops/bitwise_right_shift/cuda/kernel.cuh

@@ -0,0 +1,14 @@
+#pragma once
+#include <cuda_runtime.h>
+#include <type_traits>
+
+namespace op::cuda {
+
+template <typename T>
+struct BitwiseRightShiftOp {
+    __device__ __forceinline__ T operator()(T x, T shift) const {
+        return x >> shift;
+    }
+};
+
+} // namespace op::cuda

src/infiniop/ops/bitwise_right_shift/metax/bitwise_right_shift_metax.h

@@ -0,0 +1,8 @@
+#ifndef __BITWISE_RIGHT_SHIFT_METAX_API_H__
+#define __BITWISE_RIGHT_SHIFT_METAX_API_H__
+
+#include "../../../elementwise/metax/elementwise_metax_api.h"
+
+ELEMENTWISE_DESCRIPTOR(bitwise_right_shift, metax)
+
+#endif // __BITWISE_RIGHT_SHIFT_METAX_API_H__

src/infiniop/ops/bitwise_right_shift/metax/bitwise_right_shift_metax.maca

@@ -0,0 +1,67 @@
+#include "bitwise_right_shift_metax.h"
+
+#include "../../../elementwise/metax/elementwise_metax.h"
+
+#include "../cuda/kernel.cuh"
+
+namespace op::bitwise_right_shift::metax {
+
+Descriptor::~Descriptor() = default;
+
+infiniStatus_t Descriptor::create(
+    infiniopHandle_t handle_,
+    Descriptor **desc_ptr,
+    infiniopTensorDescriptor_t out_desc,
+    std::vector<infiniopTensorDescriptor_t> input_desc_vec) {
+
+    auto handle = reinterpret_cast<device::metax::Handle *>(handle_);
+    auto dtype = out_desc->dtype();
+
+    const auto &input_desc = input_desc_vec.at(0);
+    const auto &output_shape = out_desc->shape();
+    const auto &input_shape = input_desc->shape();
+
+    CHECK_DTYPE(dtype, INFINI_DTYPE_I8, INFINI_DTYPE_I16, INFINI_DTYPE_I32, INFINI_DTYPE_I64,
+                INFINI_DTYPE_U8, INFINI_DTYPE_U16, INFINI_DTYPE_U32, INFINI_DTYPE_U64);
+
+    CHECK_SAME_SHAPE(output_shape, input_shape);
+
+    CREATE_ELEMENTWISE_METAX_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec)
+
+    return INFINI_STATUS_SUCCESS;
+}
+
+infiniStatus_t Descriptor::calculate(
+    void *workspace,
+    size_t workspace_size,
+    void *output,
+    std::vector<const void *> inputs,
+    void *stream) const {
+
+    if (workspace_size < _workspace_size) {
+        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
+    }
+
+    switch (_dtype) {
+    case INFINI_DTYPE_I8:
+        return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int8_t, int8_t, int8_t>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_I16:
+        return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int16_t, int16_t, int16_t>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_I32:
+        return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int32_t, int32_t, int32_t>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_I64:
+        return _device_info->calculate<256, cuda::BitwiseRightShiftOp, int64_t, int64_t, int64_t>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_U8:
+        return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint8_t, uint8_t, uint8_t>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_U16:
+        return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint16_t, uint16_t, uint16_t>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_U32:
+        return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint32_t, uint32_t, uint32_t>(_info, workspace, output, inputs, stream);
+    case INFINI_DTYPE_U64:
+        return _device_info->calculate<256, cuda::BitwiseRightShiftOp, uint64_t, uint64_t, uint64_t>(_info, workspace, output, inputs, stream);
+    default:
+        return INFINI_STATUS_BAD_TENSOR_DTYPE;
+    }
+
+    return INFINI_STATUS_SUCCESS;
+} // namespace op::bitwise_right_shift::metax

src/infiniop/ops/bitwise_right_shift/moore/bitwise_right_shift_moore.h

@@ -0,0 +1,8 @@
+#ifndef __BITWISE_RIGHT_SHIFT_MOORE_API_H__
+#define __BITWISE_RIGHT_SHIFT_MOORE_API_H__
+
+#include "../../../elementwise/moore/elementwise_moore_api.h"
+
+ELEMENTWISE_DESCRIPTOR(bitwise_right_shift, moore)
+
+#endif // __BITWISE_RIGHT_SHIFT_MOORE_API_H__