[T1-2-2]wawahejun

include/infiniop.h

@@ -4,18 +4,32 @@
 #include "infiniop/handle.h"
 #include "infiniop/ops/add.h"
 #include "infiniop/ops/attention.h"
+#include "infiniop/ops/batch_norm.h"
+#include "infiniop/ops/batch_norm_backward.h"
 #include "infiniop/ops/causal_softmax.h"
 #include "infiniop/ops/clip.h"
 #include "infiniop/ops/conv.h"
+#include "infiniop/ops/gather.h"
 #include "infiniop/ops/gemm.h"
+#include "infiniop/ops/index_copy_inplace.h"
+#include "infiniop/ops/layer_norm.h"
+#include "infiniop/ops/layer_norm_backward.h"
+#include "infiniop/ops/linear.h"
+#include "infiniop/ops/linear_backward.h"
 #include "infiniop/ops/mul.h"
 #include "infiniop/ops/random_sample.h"
 #include "infiniop/ops/rearrange.h"
+#include "infiniop/ops/reduce_max.h"
+#include "infiniop/ops/reduce_mean.h"
 #include "infiniop/ops/relu.h"
 #include "infiniop/ops/rms_norm.h"
+#include "infiniop/ops/rms_norm_backward.h"
 #include "infiniop/ops/rope.h"
+#include "infiniop/ops/scatter.h"
 #include "infiniop/ops/sub.h"
 #include "infiniop/ops/swiglu.h"
+#include "infiniop/ops/tril.h"
+#include "infiniop/ops/triu.h"
 #include "infiniop/tensor_descriptor.h"
 
 #endif // __INFINIOP_API_H__

include/infiniop/ops/batch_norm.h

@@ -0,0 +1,35 @@
+#ifndef __INFINIOP_BATCH_NORM_API_H__
+#define __INFINIOP_BATCH_NORM_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopBatchNormDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateBatchNormDescriptor(
+    infiniopHandle_t handle,
+    infiniopBatchNormDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t output_desc,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t weight_desc,
+    infiniopTensorDescriptor_t bias_desc,
+    infiniopTensorDescriptor_t running_mean_desc,
+    infiniopTensorDescriptor_t running_var_desc,
+    float momentum,
+    float eps);
+
+__C __export infiniStatus_t infiniopGetBatchNormWorkspaceSize(infiniopBatchNormDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopBatchNorm(infiniopBatchNormDescriptor_t desc,
+                                               void *workspace,
+                                               size_t workspace_size,
+                                               void *output,
+                                               const void *input,
+                                               const void *weight,
+                                               const void *bias,
+                                               void *running_mean,
+                                               void *running_var,
+                                               void *stream);
+
+__C __export infiniStatus_t infiniopDestroyBatchNormDescriptor(infiniopBatchNormDescriptor_t desc);
+
+#endif

include/infiniop/ops/batch_norm_backward.h

@@ -0,0 +1,39 @@
+#ifndef __INFINIOP_BATCH_NORM_BACKWARD_API_H__
+#define __INFINIOP_BATCH_NORM_BACKWARD_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopBatchNormBackwardDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateBatchNormBackwardDescriptor(
+    infiniopHandle_t handle,
+    infiniopBatchNormBackwardDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t grad_input_desc,
+    infiniopTensorDescriptor_t grad_weight_desc,
+    infiniopTensorDescriptor_t grad_bias_desc,
+    infiniopTensorDescriptor_t grad_output_desc,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t weight_desc,
+    infiniopTensorDescriptor_t running_mean_desc,
+    infiniopTensorDescriptor_t running_var_desc,
+    float momentum,
+    float eps);
+
+__C __export infiniStatus_t infiniopGetBatchNormBackwardWorkspaceSize(infiniopBatchNormBackwardDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopBatchNormBackward(infiniopBatchNormBackwardDescriptor_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,
+                                                      const void *running_mean,
+                                                      const void *running_var,
+                                                      void *stream);
+
+__C __export infiniStatus_t infiniopDestroyBatchNormBackwardDescriptor(infiniopBatchNormBackwardDescriptor_t desc);
+
+#endif

include/infiniop/ops/gather.h

@@ -0,0 +1,27 @@
+#ifndef __INFINIOP_GATHER_API_H__
+#define __INFINIOP_GATHER_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopGatherDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateGatherDescriptor(infiniopHandle_t handle,
+                                                            infiniopGatherDescriptor_t *desc_ptr,
+                                                            infiniopTensorDescriptor_t input,
+                                                            infiniopTensorDescriptor_t output,
+                                                            int dim,
+                                                            infiniopTensorDescriptor_t index);
+
+__C __export infiniStatus_t infiniopGetGatherWorkspaceSize(infiniopGatherDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopGather(infiniopGatherDescriptor_t desc,
+                                            void *workspace,
+                                            size_t workspace_size,
+                                            void *output,
+                                            void *input,
+                                            const void *index,
+                                            void *stream);
+
+__C __export infiniStatus_t infiniopDestroyGatherDescriptor(infiniopGatherDescriptor_t desc);
+
+#endif

include/infiniop/ops/index_copy_inplace.h

@@ -0,0 +1,27 @@
+#ifndef __INFINIOP_INDEX_COPY_INPLACE_API_H__
+#define __INFINIOP_INDEX_COPY_INPLACE_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopIndexCopyInplaceDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateIndexCopyInplaceDescriptor(infiniopHandle_t handle,
+                                                                      infiniopIndexCopyInplaceDescriptor_t *desc_ptr,
+                                                                      infiniopTensorDescriptor_t input,
+                                                                      infiniopTensorDescriptor_t output,
+                                                                      int dim,
+                                                                      infiniopTensorDescriptor_t index);
+
+__C __export infiniStatus_t infiniopGetIndexCopyInplaceWorkspaceSize(infiniopIndexCopyInplaceDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopIndexCopyInplace(infiniopIndexCopyInplaceDescriptor_t desc,
+                                                      void *workspace,
+                                                      size_t workspace_size,
+                                                      void *input,
+                                                      void *output,
+                                                      const void *index,
+                                                      void *stream);
+
+__C __export infiniStatus_t infiniopDestroyIndexCopyInplaceDescriptor(infiniopIndexCopyInplaceDescriptor_t desc);
+
+#endif

include/infiniop/ops/layer_norm.h

@@ -0,0 +1,27 @@
+#ifndef __INFINIOP_LAYER_NORM_API_H__
+#define __INFINIOP_LAYER_NORM_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLayerNormDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLayerNormDescriptor(
+    infiniopHandle_t handle,
+    infiniopLayerNormDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t output_desc,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t weight_desc,
+    infiniopTensorDescriptor_t bias_desc,
+    infiniopTensorDescriptor_t input_std_deviation_desc,
+    infiniopTensorDescriptor_t input_standardization_desc,
+    float epsilon);
+
+__C __export infiniStatus_t infiniopGetLayerNormWorkspaceSize(infiniopLayerNormDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLayerNorm(infiniopLayerNormDescriptor_t desc, void *workspace, size_t workspace_size,
+                                              void *output, const void *input, const void *weight, const void *bias,
+                                              void *input_std_deviation, void *input_standardization, void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLayerNormDescriptor(infiniopLayerNormDescriptor_t desc);
+
+#endif

include/infiniop/ops/layer_norm_backward.h

@@ -0,0 +1,30 @@
+#ifndef __INFINIOP_LAYER_NORM_BACKWARD_API_H__
+#define __INFINIOP_LAYER_NORM_BACKWARD_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLayerNormBackwardDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLayerNormBackwardDescriptor(
+    infiniopHandle_t handle,
+    infiniopLayerNormBackwardDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t grad_input_desc,
+    infiniopTensorDescriptor_t grad_weight_desc,
+    infiniopTensorDescriptor_t grad_bias_desc,
+    infiniopTensorDescriptor_t grad_output_desc,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t weight_desc,
+    infiniopTensorDescriptor_t input_std_deviation_desc,
+    infiniopTensorDescriptor_t input_standardization_desc,
+    float epsilon);
+
+__C __export infiniStatus_t infiniopGetLayerNormBackwardWorkspaceSize(infiniopLayerNormBackwardDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLayerNormBackward(infiniopLayerNormBackwardDescriptor_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,
+                                                      const void *input_std_deviation, const void *input_standardization, void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLayerNormBackwardDescriptor(infiniopLayerNormBackwardDescriptor_t desc);
+
+#endif

include/infiniop/ops/linear.h

@@ -0,0 +1,28 @@
+#ifndef __INFINIOP_LINEAR_API_H__
+#define __INFINIOP_LINEAR_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLinearDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLinearDescriptor(infiniopHandle_t handle,
+                                                            infiniopLinearDescriptor_t *desc_ptr,
+                                                            infiniopTensorDescriptor_t x,
+                                                            infiniopTensorDescriptor_t w,
+                                                            infiniopTensorDescriptor_t b,
+                                                            infiniopTensorDescriptor_t y);
+
+__C __export infiniStatus_t infiniopGetLinearWorkspaceSize(infiniopLinearDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLinear(infiniopLinearDescriptor_t desc,
+                                            void *workspace,
+                                            size_t workspace_size,
+                                            void *y,
+                                            const void *x,
+                                            const void *w,
+                                            const void *b,
+                                            void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLinearDescriptor(infiniopLinearDescriptor_t desc);
+
+#endif

include/infiniop/ops/linear_backward.h

@@ -0,0 +1,32 @@
+#ifndef __INFINIOP_LINEAR_BACKWARD_API_H__
+#define __INFINIOP_LINEAR_BACKWARD_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLinearBackwardDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLinearBackwardDescriptor(infiniopHandle_t handle,
+                                                                    infiniopLinearBackwardDescriptor_t *desc_ptr,
+                                                                    infiniopTensorDescriptor_t grad_y,
+                                                                    infiniopTensorDescriptor_t x,
+                                                                    infiniopTensorDescriptor_t w,
+                                                                    infiniopTensorDescriptor_t grad_x,
+                                                                    infiniopTensorDescriptor_t grad_w,
+                                                                    infiniopTensorDescriptor_t grad_b);
+
+__C __export infiniStatus_t infiniopGetLinearBackwardWorkspaceSize(infiniopLinearBackwardDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLinearBackward(infiniopLinearBackwardDescriptor_t desc,
+                                                    void *workspace,
+                                                    size_t workspace_size,
+                                                    void *grad_x,
+                                                    void *grad_w,
+                                                    void *grad_b,
+                                                    const void *grad_y,
+                                                    const void *x,
+                                                    const void *w,
+                                                    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLinearBackwardDescriptor(infiniopLinearBackwardDescriptor_t desc);
+
+#endif

include/infiniop/ops/reduce_max.h

@@ -0,0 +1,25 @@
+#ifndef __INFINIOP_REDUCE_MAX_API_H__
+#define __INFINIOP_REDUCE_MAX_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopReduceMaxDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateReduceMaxDescriptor(infiniopHandle_t handle,
+                                                              infiniopReduceMaxDescriptor_t *desc_ptr,
+                                                              infiniopTensorDescriptor_t output_desc,
+                                                              infiniopTensorDescriptor_t input_desc,
+                                                              size_t dim);
+
+__C __export infiniStatus_t infiniopGetReduceMaxWorkspaceSize(infiniopReduceMaxDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopReduceMax(infiniopReduceMaxDescriptor_t desc,
+                                              void *workspace,
+                                              size_t workspace_size,
+                                              void *output,
+                                              const void *input,
+                                              void *stream);
+
+__C __export infiniStatus_t infiniopDestroyReduceMaxDescriptor(infiniopReduceMaxDescriptor_t desc);
+
+#endif

include/infiniop/ops/reduce_mean.h

@@ -0,0 +1,25 @@
+#ifndef __INFINIOP_REDUCE_MEAN_API_H__
+#define __INFINIOP_REDUCE_MEAN_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopReduceMeanDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateReduceMeanDescriptor(infiniopHandle_t handle,
+                                                               infiniopReduceMeanDescriptor_t *desc_ptr,
+                                                               infiniopTensorDescriptor_t output_desc,
+                                                               infiniopTensorDescriptor_t input_desc,
+                                                               size_t dim);
+
+__C __export infiniStatus_t infiniopGetReduceMeanWorkspaceSize(infiniopReduceMeanDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopReduceMean(infiniopReduceMeanDescriptor_t desc,
+                                               void *workspace,
+                                               size_t workspace_size,
+                                               void *output,
+                                               const void *input,
+                                               void *stream);
+
+__C __export infiniStatus_t infiniopDestroyReduceMeanDescriptor(infiniopReduceMeanDescriptor_t desc);
+
+#endif

include/infiniop/ops/rms_norm_backward.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_RMS_NORM_BACKWARD_API_H__
+#define __INFINIOP_RMS_NORM_BACKWARD_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopRMSNormBackwardDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateRMSNormBackwardDescriptor(
+    infiniopHandle_t handle,
+    infiniopRMSNormBackwardDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t grad_x_desc,
+    infiniopTensorDescriptor_t grad_w_desc,
+    infiniopTensorDescriptor_t grad_y_desc,
+    infiniopTensorDescriptor_t x_desc,
+    infiniopTensorDescriptor_t w_desc,
+    float epsilon);
+
+__C __export infiniStatus_t infiniopGetRMSNormBackwardWorkspaceSize(infiniopRMSNormBackwardDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopRMSNormBackward(infiniopRMSNormBackwardDescriptor_t desc, void *workspace, size_t workspace_size,
+                                                    void *grad_x, void *grad_w,
+                                                    const void *grad_y, const void *x, const void *w, void *stream);
+
+__C __export infiniStatus_t infiniopDestroyRMSNormBackwardDescriptor(infiniopRMSNormBackwardDescriptor_t desc);
+
+#endif