[T1-1-2] tianyuxbear

include/infiniop.h

@@ -3,17 +3,27 @@
 
 #include "infiniop/handle.h"
 #include "infiniop/ops/add.h"
+#include "infiniop/ops/all_equal.h"
 #include "infiniop/ops/attention.h"
 #include "infiniop/ops/causal_softmax.h"
 #include "infiniop/ops/clip.h"
 #include "infiniop/ops/conv.h"
+#include "infiniop/ops/cross_entropy_loss_backward.h"
+#include "infiniop/ops/div.h"
+#include "infiniop/ops/equal.h"
+#include "infiniop/ops/gelu.h"
+#include "infiniop/ops/gelu_backward.h"
 #include "infiniop/ops/gemm.h"
+#include "infiniop/ops/logical_and.h"
+#include "infiniop/ops/logical_or.h"
 #include "infiniop/ops/mul.h"
 #include "infiniop/ops/random_sample.h"
 #include "infiniop/ops/rearrange.h"
 #include "infiniop/ops/relu.h"
+#include "infiniop/ops/relu_backward.h"
 #include "infiniop/ops/rms_norm.h"
 #include "infiniop/ops/rope.h"
+#include "infiniop/ops/silu.h"
 #include "infiniop/ops/sub.h"
 #include "infiniop/ops/swiglu.h"
 #include "infiniop/tensor_descriptor.h"

include/infiniop/ops/all_equal.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_ALL_EQUAL_API_H__
+#define __INFINIOP_ALL_EQUAL_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopAllEqualDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateAllEqualDescriptor(infiniopHandle_t handle,
+                                                             infiniopAllEqualDescriptor_t *desc_ptr,
+                                                             infiniopTensorDescriptor_t c,
+                                                             infiniopTensorDescriptor_t a,
+                                                             infiniopTensorDescriptor_t b);
+
+__C __export infiniStatus_t infiniopGetAllEqualWorkspaceSize(infiniopAllEqualDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopAllEqual(infiniopAllEqualDescriptor_t desc,
+                                             void *workspace,
+                                             size_t workspace_size,
+                                             void *c,
+                                             const void *a,
+                                             const void *b,
+                                             void *stream);
+
+__C __export infiniStatus_t infiniopDestroyAllEqualDescriptor(infiniopAllEqualDescriptor_t desc);
+
+#endif

include/infiniop/ops/cross_entropy_loss_backward.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_CROSS_ENTROPY_LOSS_BACKWARD_API_H__
+#define __INFINIOP_CROSS_ENTROPY_LOSS_BACKWARD_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopCrossEntropyLossBackwardDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateCrossEntropyLossBackwardDescriptor(infiniopHandle_t handle,
+                                                                             infiniopCrossEntropyLossBackwardDescriptor_t *desc_ptr,
+                                                                             infiniopTensorDescriptor_t grad_logits,
+                                                                             infiniopTensorDescriptor_t probs,
+                                                                             infiniopTensorDescriptor_t target);
+
+__C __export infiniStatus_t infiniopGetCrossEntropyLossBackwardWorkspaceSize(infiniopCrossEntropyLossBackwardDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopCrossEntropyLossBackward(infiniopCrossEntropyLossBackwardDescriptor_t desc,
+                                                             void *workspace,
+                                                             size_t workspace_size,
+                                                             void *grad_logits,
+                                                             const void *probs,
+                                                             const void *target,
+                                                             void *stream);
+
+__C __export infiniStatus_t infiniopDestroyCrossEntropyLossBackwardDescriptor(infiniopCrossEntropyLossBackwardDescriptor_t desc);
+
+#endif

include/infiniop/ops/div.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_DIV_API_H__
+#define __INFINIOP_DIV_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopDivDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateDivDescriptor(infiniopHandle_t handle,
+                                                        infiniopDivDescriptor_t *desc_ptr,
+                                                        infiniopTensorDescriptor_t c,
+                                                        infiniopTensorDescriptor_t a,
+                                                        infiniopTensorDescriptor_t b);
+
+__C __export infiniStatus_t infiniopGetDivWorkspaceSize(infiniopDivDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopDiv(infiniopDivDescriptor_t desc,
+                                        void *workspace,
+                                        size_t workspace_size,
+                                        void *c,
+                                        const void *a,
+                                        const void *b,
+                                        void *stream);
+
+__C __export infiniStatus_t infiniopDestroyDivDescriptor(infiniopDivDescriptor_t desc);
+
+#endif

include/infiniop/ops/equal.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_EQUAL_API_H__
+#define __INFINIOP_EQUAL_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopEqualDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateEqualDescriptor(infiniopHandle_t handle,
+                                                          infiniopEqualDescriptor_t *desc_ptr,
+                                                          infiniopTensorDescriptor_t c,
+                                                          infiniopTensorDescriptor_t a,
+                                                          infiniopTensorDescriptor_t b);
+
+__C __export infiniStatus_t infiniopGetEqualWorkspaceSize(infiniopEqualDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopEqual(infiniopEqualDescriptor_t desc,
+                                          void *workspace,
+                                          size_t workspace_size,
+                                          void *c,
+                                          const void *a,
+                                          const void *b,
+                                          void *stream);
+
+__C __export infiniStatus_t infiniopDestroyEqualDescriptor(infiniopEqualDescriptor_t desc);
+
+#endif

include/infiniop/ops/gelu.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_GELU_API_H__
+#define __INFINIOP_GELU_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopGeluDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateGeluDescriptor(infiniopHandle_t handle,
+                                                         infiniopGeluDescriptor_t *desc_ptr,
+                                                         infiniopTensorDescriptor_t output,
+                                                         infiniopTensorDescriptor_t intput);
+
+__C __export infiniStatus_t infiniopGetGeluWorkspaceSize(infiniopGeluDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopGelu(infiniopGeluDescriptor_t desc,
+                                         void *workspace,
+                                         size_t workspace_size,
+                                         void *output,
+                                         const void *intput,
+                                         void *stream);
+
+__C __export infiniStatus_t infiniopDestroyGeluDescriptor(infiniopGeluDescriptor_t desc);
+
+#endif

include/infiniop/ops/gelu_backward.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_GELU_BACKWARD_API_H__
+#define __INFINIOP_GELU_BACKWARD_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopGeluBackwardDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateGeluBackwardDescriptor(infiniopHandle_t handle,
+                                                                 infiniopGeluBackwardDescriptor_t *desc_ptr,
+                                                                 infiniopTensorDescriptor_t grad_input,
+                                                                 infiniopTensorDescriptor_t intput,
+                                                                 infiniopTensorDescriptor_t grad_output);
+
+__C __export infiniStatus_t infiniopGetGeluBackwardWorkspaceSize(infiniopGeluBackwardDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopGeluBackward(infiniopGeluBackwardDescriptor_t desc,
+                                                 void *workspace,
+                                                 size_t workspace_size,
+                                                 void *grad_input,
+                                                 const void *input,
+                                                 const void *grad_output,
+                                                 void *stream);
+
+__C __export infiniStatus_t infiniopDestroyGeluBackwardDescriptor(infiniopGeluBackwardDescriptor_t desc);
+
+#endif

include/infiniop/ops/logical_and.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_LOGICAL_AND_API_H__
+#define __INFINIOP_LOGICAL_AND_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLogicalAndDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLogicalAndDescriptor(infiniopHandle_t handle,
+                                                               infiniopLogicalAndDescriptor_t *desc_ptr,
+                                                               infiniopTensorDescriptor_t c,
+                                                               infiniopTensorDescriptor_t a,
+                                                               infiniopTensorDescriptor_t b);
+
+__C __export infiniStatus_t infiniopGetLogicalAndWorkspaceSize(infiniopLogicalAndDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLogicalAnd(infiniopLogicalAndDescriptor_t desc,
+                                               void *workspace,
+                                               size_t workspace_size,
+                                               void *c,
+                                               const void *a,
+                                               const void *b,
+                                               void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLogicalAndDescriptor(infiniopLogicalAndDescriptor_t desc);
+
+#endif

include/infiniop/ops/logical_or.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_LOGICAL_OR_API_H__
+#define __INFINIOP_LOGICAL_OR_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopLogicalOrDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateLogicalOrDescriptor(infiniopHandle_t handle,
+                                                              infiniopLogicalOrDescriptor_t *desc_ptr,
+                                                              infiniopTensorDescriptor_t c,
+                                                              infiniopTensorDescriptor_t a,
+                                                              infiniopTensorDescriptor_t b);
+
+__C __export infiniStatus_t infiniopGetLogicalOrWorkspaceSize(infiniopLogicalOrDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopLogicalOr(infiniopLogicalOrDescriptor_t desc,
+                                              void *workspace,
+                                              size_t workspace_size,
+                                              void *c,
+                                              const void *a,
+                                              const void *b,
+                                              void *stream);
+
+__C __export infiniStatus_t infiniopDestroyLogicalOrDescriptor(infiniopLogicalOrDescriptor_t desc);
+
+#endif

include/infiniop/ops/relu_backward.h

@@ -0,0 +1,26 @@
+#ifndef __INFINIOP_RELU_BACKWARD_API_H__
+#define __INFINIOP_RELU_BACKWARD_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopReluBackwardDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateReluBackwardDescriptor(infiniopHandle_t handle,
+                                                                 infiniopReluBackwardDescriptor_t *desc_ptr,
+                                                                 infiniopTensorDescriptor_t grad_input,
+                                                                 infiniopTensorDescriptor_t intput,
+                                                                 infiniopTensorDescriptor_t grad_output);
+
+__C __export infiniStatus_t infiniopGetReluBackwardWorkspaceSize(infiniopReluBackwardDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopReluBackward(infiniopReluBackwardDescriptor_t desc,
+                                                 void *workspace,
+                                                 size_t workspace_size,
+                                                 void *grad_input,
+                                                 const void *input,
+                                                 const void *grad_output,
+                                                 void *stream);
+
+__C __export infiniStatus_t infiniopDestroyReluBackwardDescriptor(infiniopReluBackwardDescriptor_t desc);
+
+#endif

include/infiniop/ops/silu.h

@@ -0,0 +1,24 @@
+#ifndef __INFINIOP_SILU_API_H__
+#define __INFINIOP_SILU_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopSiluDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateSiluDescriptor(infiniopHandle_t handle,
+                                                         infiniopSiluDescriptor_t *desc_ptr,
+                                                         infiniopTensorDescriptor_t output,
+                                                         infiniopTensorDescriptor_t intput);
+
+__C __export infiniStatus_t infiniopGetSiluWorkspaceSize(infiniopSiluDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopSilu(infiniopSiluDescriptor_t desc,
+                                         void *workspace,
+                                         size_t workspace_size,
+                                         void *output,
+                                         const void *intput,
+                                         void *stream);
+
+__C __export infiniStatus_t infiniopDestroySiluDescriptor(infiniopSiluDescriptor_t desc);
+
+#endif

scripts/python_test.py

@@ -12,21 +12,31 @@
 def run_tests(args):
     failed = []
     for test in [
-        "add.py",
-        "attention.py",
-        "causal_softmax.py",
-        "clip.py",
-        "gemm.py",
-        "mul.py",
-        "random_sample.py",
-        "rearrange.py",
-        "rms_norm.py",
-        "rope.py",
-        "sub.py",
-        "swiglu.py",
+        # "add.py",
+        # "attention.py",
+        # "causal_softmax.py",
+        # "clip.py",
+        # "gemm.py",
+        # "mul.py",
+        # "random_sample.py",
+        # "rearrange.py",
+        # "rms_norm.py",
+        # "rope.py",
+        # "sub.py",
+        # "swiglu.py",
+        "silu.py",
+        "div.py",
+        "logical_and.py",
+        "logical_or.py",
+        "equal.py",
+        "all_equal.py",
+        "relu_backward.py",
+        "gelu.py",
+        "gelu_backward.py",
+        "cross_entropy_loss_backward.py"
     ]:
         result = subprocess.run(
-            f"python {test} {args} --debug", text=True, encoding="utf-8", shell=True
+            f"python {test} {args} --profile", text=True, encoding="utf-8", shell=True
         )
         if result.returncode != 0:
             failed.append(test)

src/infiniop-test/include/ops.hpp

@@ -16,6 +16,16 @@ DECLARE_INFINIOP_TEST(add)
 DECLARE_INFINIOP_TEST(causal_softmax)
 DECLARE_INFINIOP_TEST(rearrange)
 DECLARE_INFINIOP_TEST(sub)
+DECLARE_INFINIOP_TEST(silu)
+DECLARE_INFINIOP_TEST(div)
+DECLARE_INFINIOP_TEST(logical_and)
+DECLARE_INFINIOP_TEST(logical_or)
+DECLARE_INFINIOP_TEST(equal)
+DECLARE_INFINIOP_TEST(all_equal)
+DECLARE_INFINIOP_TEST(relu_backward)
+DECLARE_INFINIOP_TEST(gelu)
+DECLARE_INFINIOP_TEST(gelu_backward)
+DECLARE_INFINIOP_TEST(cross_entropy_loss_backward)
 
 #define REGISTER_INFINIOP_TEST(name)                      \
     {                                                     \
@@ -30,19 +40,29 @@ 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(silu)                        \
+        REGISTER_INFINIOP_TEST(div)                         \
+        REGISTER_INFINIOP_TEST(logical_and)                 \
+        REGISTER_INFINIOP_TEST(logical_or)                  \
+        REGISTER_INFINIOP_TEST(equal)                       \
+        REGISTER_INFINIOP_TEST(all_equal)                   \
+        REGISTER_INFINIOP_TEST(relu_backward)               \
+        REGISTER_INFINIOP_TEST(gelu)                        \
+        REGISTER_INFINIOP_TEST(gelu_backward)               \
+        REGISTER_INFINIOP_TEST(cross_entropy_loss_backward) \
     }
 
 namespace infiniop_test {