[T1-2-3]ZhaoShijie

include/infiniop.h

@@ -4,11 +4,18 @@
 #include "infiniop/handle.h"
 #include "infiniop/ops/add.h"
 #include "infiniop/ops/attention.h"
+#include "infiniop/ops/averagepool.h"
+#include "infiniop/ops/averagepool_backward.h"
 #include "infiniop/ops/causal_softmax.h"
 #include "infiniop/ops/clip.h"
 #include "infiniop/ops/conv.h"
+#include "infiniop/ops/conv_backward.h"
+#include "infiniop/ops/cross_entropy_loss.h"
 #include "infiniop/ops/dequantize.h"
 #include "infiniop/ops/gemm.h"
+#include "infiniop/ops/interpolate_nearest.h"
+#include "infiniop/ops/maxpool.h"
+#include "infiniop/ops/maxpool_backward.h"
 #include "infiniop/ops/mul.h"
 #include "infiniop/ops/random_sample.h"
 #include "infiniop/ops/rearrange.h"

include/infiniop/ops/averagepool.h

@@ -0,0 +1,29 @@
+#ifndef __INFINIOP_AVERAGEPOOL_H__
+#define __INFINIOP_AVERAGEPOOL_H__
+
+#include "../operator_descriptor.h"
+
+__C typedef struct InfiniopDescriptor *infiniopAvgPoolDescriptor_t;
+
+__C infiniStatus_t infiniopCreateAvgPoolDescriptor(infiniopHandle_t handle,
+                                                   infiniopAvgPoolDescriptor_t *desc_ptr,
+                                                   infiniopTensorDescriptor_t output_desc,
+                                                   infiniopTensorDescriptor_t input_desc,
+                                                   void *kernel_size,
+                                                   void *strides,
+                                                   void *pads,
+                                                   bool ceil_mode);
+
+__C infiniStatus_t infiniopGetAvgPoolWorkspaceSize(infiniopAvgPoolDescriptor_t desc,
+                                                   size_t *size);
+
+__C infiniStatus_t infiniopAvgPool(infiniopAvgPoolDescriptor_t desc,
+                                   void *workspace,
+                                   size_t workspace_size,
+                                   void *output,
+                                   const void *input,
+                                   void *stream);
+
+__C infiniStatus_t infiniopDestroyAvgPoolDescriptor(infiniopAvgPoolDescriptor_t desc);
+
+#endif // __INFINIOP_AVERAGEPOOL_H__

include/infiniop/ops/averagepool_backward.h

@@ -0,0 +1,31 @@
+#ifndef __INFINIOP_AVERAGEPOOL_BACKWARD_H__
+#define __INFINIOP_AVERAGEPOOL_BACKWARD_H__
+
+#include "../operator_descriptor.h"
+
+__C typedef struct InfiniopDescriptor *infiniopAvgPoolBackwardDescriptor_t;
+
+__C infiniStatus_t infiniopCreateAvgPoolBackwardDescriptor(infiniopHandle_t handle,
+                                                           infiniopAvgPoolBackwardDescriptor_t *desc_ptr,
+                                                           infiniopTensorDescriptor_t grad_input_desc,
+                                                           infiniopTensorDescriptor_t grad_output_desc,
+                                                           infiniopTensorDescriptor_t input_desc,
+                                                           void *kernel_size,
+                                                           void *strides,
+                                                           void *pads,
+                                                           bool ceil_mode);
+
+__C infiniStatus_t infiniopGetAvgPoolBackwardWorkspaceSize(infiniopAvgPoolBackwardDescriptor_t desc,
+                                                           size_t *size);
+
+__C infiniStatus_t infiniopAvgPoolBackward(infiniopAvgPoolBackwardDescriptor_t desc,
+                                           void *workspace,
+                                           size_t workspace_size,
+                                           void *grad_input,
+                                           const void *grad_output,
+                                           const void *input,
+                                           void *stream);
+
+__C infiniStatus_t infiniopDestroyAvgPoolBackwardDescriptor(infiniopAvgPoolBackwardDescriptor_t desc);
+
+#endif // __INFINIOP_AVERAGEPOOL_BACKWARD_H__

include/infiniop/ops/conv_backward.h

@@ -0,0 +1,34 @@
+#ifndef __INFINIOP_CONV_BACKWARD_API_H__
+#define __INFINIOP_CONV_BACKWARD_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopConvBackwardDescriptor_t;
+
+__C infiniStatus_t infiniopCreateConvBackwardDescriptor(infiniopHandle_t handle,
+                                                        infiniopConvBackwardDescriptor_t *desc_ptr,
+                                                        infiniopTensorDescriptor_t grad_output_desc,
+                                                        infiniopTensorDescriptor_t input_desc,
+                                                        infiniopTensorDescriptor_t weight_desc,
+                                                        infiniopTensorDescriptor_t bias_desc,
+                                                        void *pads,
+                                                        void *strides,
+                                                        void *dilations,
+                                                        size_t n);
+
+__C infiniStatus_t infiniopGetConvBackwardWorkspaceSize(infiniopConvBackwardDescriptor_t desc, size_t *size);
+
+__C infiniStatus_t infiniopConvBackward(infiniopConvBackwardDescriptor_t desc,
+                                        void *workspace,
+                                        size_t workspace_size,
+                                        void *grad_input,
+                                        void *grad_weight,
+                                        void *grad_bias,
+                                        const void *grad_output,
+                                        const void *input,
+                                        const void *weight,
+                                        void *stream);
+
+__C infiniStatus_t infiniopDestroyConvBackwardDescriptor(infiniopConvBackwardDescriptor_t desc);
+
+#endif

include/infiniop/ops/cross_entropy_loss.h

@@ -0,0 +1,27 @@
+#ifndef __INFINIOP_CROSS_ENTROPY_LOSS_API_H__
+#define __INFINIOP_CROSS_ENTROPY_LOSS_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopCrossEntropyLossDescriptor_t;
+
+__C infiniStatus_t infiniopCreateCrossEntropyLossDescriptor(infiniopHandle_t handle,
+                                                            infiniopCrossEntropyLossDescriptor_t *desc_ptr,
+                                                            infiniopTensorDescriptor_t loss_desc,
+                                                            infiniopTensorDescriptor_t logits_desc,
+                                                            infiniopTensorDescriptor_t target_desc);
+
+__C infiniStatus_t infiniopGetCrossEntropyLossWorkspaceSize(infiniopCrossEntropyLossDescriptor_t desc,
+                                                            size_t *size);
+
+__C infiniStatus_t infiniopCrossEntropyLoss(infiniopCrossEntropyLossDescriptor_t desc,
+                                            void *workspace,
+                                            size_t workspace_size,
+                                            void *loss,
+                                            const void *logits,
+                                            const void *target,
+                                            void *stream);
+
+__C infiniStatus_t infiniopDestroyCrossEntropyLossDescriptor(infiniopCrossEntropyLossDescriptor_t desc);
+
+#endif // __INFINIOP_CROSS_ENTROPY_LOSS_API_H__

include/infiniop/ops/interpolate_nearest.h

@@ -0,0 +1,25 @@
+#ifndef __INFINIOP_INTERPOLATE_NEAREST_H__
+#define __INFINIOP_INTERPOLATE_NEAREST_H__
+
+#include "../operator_descriptor.h"
+
+__C typedef struct InfiniopDescriptor *infiniopInterpolateNearestDescriptor_t;
+
+__C infiniStatus_t infiniopCreateInterpolateNearestDescriptor(infiniopHandle_t handle,
+                                                              infiniopInterpolateNearestDescriptor_t *desc_ptr,
+                                                              infiniopTensorDescriptor_t output_desc,
+                                                              infiniopTensorDescriptor_t input_desc);
+
+__C infiniStatus_t infiniopGetInterpolateNearestWorkspaceSize(infiniopInterpolateNearestDescriptor_t desc,
+                                                              size_t *size);
+
+__C infiniStatus_t infiniopInterpolateNearest(infiniopInterpolateNearestDescriptor_t desc,
+                                              void *workspace,
+                                              size_t workspace_size,
+                                              void *output,
+                                              const void *input,
+                                              void *stream);
+
+__C infiniStatus_t infiniopDestroyInterpolateNearestDescriptor(infiniopInterpolateNearestDescriptor_t desc);
+
+#endif // __INFINIOP_INTERPOLATE_NEAREST_H__

include/infiniop/ops/maxpool.h

@@ -0,0 +1,29 @@
+#ifndef __INFINIOP_MAX_POOL_H__
+#define __INFINIOP_MAX_POOL_H__
+
+#include "../operator_descriptor.h"
+
+__C typedef struct InfiniopDescriptor *infiniopMaxPoolDescriptor_t;
+
+__C infiniStatus_t infiniopCreateMaxPoolDescriptor(infiniopHandle_t handle,
+                                                   infiniopMaxPoolDescriptor_t *desc_ptr,
+                                                   infiniopTensorDescriptor_t output_desc,
+                                                   infiniopTensorDescriptor_t input_desc,
+                                                   void *kernel_size,
+                                                   void *strides,
+                                                   void *pads,
+                                                   bool ceil_mode);
+
+__C infiniStatus_t infiniopGetMaxPoolWorkspaceSize(infiniopMaxPoolDescriptor_t desc,
+                                                   size_t *size);
+
+__C infiniStatus_t infiniopMaxPool(infiniopMaxPoolDescriptor_t desc,
+                                   void *workspace,
+                                   size_t workspace_size,
+                                   void *output,
+                                   const void *input,
+                                   void *stream);
+
+__C infiniStatus_t infiniopDestroyMaxPoolDescriptor(infiniopMaxPoolDescriptor_t desc);
+
+#endif // __INFINIOP_MAX_POOL_H__

include/infiniop/ops/maxpool_backward.h

@@ -0,0 +1,31 @@
+#ifndef __INFINIOP_MAXPOOL_BACKWARD_H__
+#define __INFINIOP_MAXPOOL_BACKWARD_H__
+
+#include "../operator_descriptor.h"
+
+__C typedef struct InfiniopDescriptor *infiniopMaxPoolBackwardDescriptor_t;
+
+__C infiniStatus_t infiniopCreateMaxPoolBackwardDescriptor(infiniopHandle_t handle,
+                                                           infiniopMaxPoolBackwardDescriptor_t *desc_ptr,
+                                                           infiniopTensorDescriptor_t grad_input_desc,
+                                                           infiniopTensorDescriptor_t grad_output_desc,
+                                                           infiniopTensorDescriptor_t input_desc,
+                                                           void *kernel_size,
+                                                           void *strides,
+                                                           void *pads,
+                                                           bool ceil_mode);
+
+__C infiniStatus_t infiniopGetMaxPoolBackwardWorkspaceSize(infiniopMaxPoolBackwardDescriptor_t desc,
+                                                           size_t *size);
+
+__C infiniStatus_t infiniopMaxPoolBackward(infiniopMaxPoolBackwardDescriptor_t desc,
+                                           void *workspace,
+                                           size_t workspace_size,
+                                           void *grad_input,
+                                           const void *grad_output,
+                                           const void *input,
+                                           void *stream);
+
+__C infiniStatus_t infiniopDestroyMaxPoolBackwardDescriptor(infiniopMaxPoolBackwardDescriptor_t desc);
+
+#endif // __INFINIOP_MAXPOOL_BACKWARD_H__

scripts/python_test.py

@@ -25,6 +25,13 @@ def run_tests(args):
         "sub.py",
         "swiglu.py",
         "softplus.py",
+        "averagepool_backward.py",
+        "averagepool.py",
+        "maxpool_backward.py",
+        "maxpool.py",
+        "interpolate_nearest.py",
+        "conv_backward.py",
+        "cross_entropy_loss.py",
     ]:
         result = subprocess.run(
             f"python {test} {args} --debug", text=True, encoding="utf-8", shell=True

src/infiniop-test/include/ops.hpp

@@ -16,6 +16,13 @@ DECLARE_INFINIOP_TEST(add)
 DECLARE_INFINIOP_TEST(causal_softmax)
 DECLARE_INFINIOP_TEST(rearrange)
 DECLARE_INFINIOP_TEST(sub)
+DECLARE_INFINIOP_TEST(cross_entropy_loss)
+DECLARE_INFINIOP_TEST(averagepool)
+DECLARE_INFINIOP_TEST(averagepool_backward)
+DECLARE_INFINIOP_TEST(interpolate_nearest)
+DECLARE_INFINIOP_TEST(conv_backward)
+DECLARE_INFINIOP_TEST(maxpool)
+DECLARE_INFINIOP_TEST(maxpool_backward)
 
 #define REGISTER_INFINIOP_TEST(name)                      \
     {                                                     \
@@ -30,19 +37,26 @@ DECLARE_INFINIOP_TEST(sub)
 /*
  * Register all the tests here
  */
-#define TEST_BUILDER_MAPPINGS                  \
-    {                                          \
-        REGISTER_INFINIOP_TEST(gemm)           \
-        REGISTER_INFINIOP_TEST(random_sample)  \
-        REGISTER_INFINIOP_TEST(add)            \
-        REGISTER_INFINIOP_TEST(mul)            \
-        REGISTER_INFINIOP_TEST(clip)           \
-        REGISTER_INFINIOP_TEST(swiglu)         \
-        REGISTER_INFINIOP_TEST(rope)           \
-        REGISTER_INFINIOP_TEST(rms_norm)       \
-        REGISTER_INFINIOP_TEST(causal_softmax) \
-        REGISTER_INFINIOP_TEST(rearrange)      \
-        REGISTER_INFINIOP_TEST(sub)            \
+#define TEST_BUILDER_MAPPINGS                        \
+    {                                                \
+        REGISTER_INFINIOP_TEST(gemm)                 \
+        REGISTER_INFINIOP_TEST(random_sample)        \
+        REGISTER_INFINIOP_TEST(add)                  \
+        REGISTER_INFINIOP_TEST(mul)                  \
+        REGISTER_INFINIOP_TEST(clip)                 \
+        REGISTER_INFINIOP_TEST(swiglu)               \
+        REGISTER_INFINIOP_TEST(rope)                 \
+        REGISTER_INFINIOP_TEST(rms_norm)             \
+        REGISTER_INFINIOP_TEST(causal_softmax)       \
+        REGISTER_INFINIOP_TEST(rearrange)            \
+        REGISTER_INFINIOP_TEST(sub)                  \
+        REGISTER_INFINIOP_TEST(cross_entropy_loss)   \
+        REGISTER_INFINIOP_TEST(averagepool)          \
+        REGISTER_INFINIOP_TEST(averagepool_backward) \
+        REGISTER_INFINIOP_TEST(interpolate_nearest)  \
+        REGISTER_INFINIOP_TEST(conv_backward)        \
+        REGISTER_INFINIOP_TEST(maxpool)              \
+        REGISTER_INFINIOP_TEST(maxpool_backward)     \
     }
 
 namespace infiniop_test {