[T1-2-2] Vietlin

include/infiniop.h

@@ -17,5 +17,12 @@
 #include "infiniop/ops/sub.h"
 #include "infiniop/ops/swiglu.h"
 #include "infiniop/tensor_descriptor.h"
+#include "infiniop/ops/index_copy_inplace.h"
+#include "infiniop/ops/gather.h"
+#include "infiniop/ops/scatter.h"
+#include "infiniop/ops/tril.h"
+#include "infiniop/ops/triu.h"
+#include "infiniop/ops/linear.h"
+#include "infiniop/ops/linear_backward.h"
 
 #endif // __INFINIOP_API_H__

include/infiniop/ops/gather.h

@@ -0,0 +1,31 @@
+#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 output_desc,
+    infiniopTensorDescriptor_t input_desc,
+    infiniopTensorDescriptor_t index_desc,
+    size_t dim
+);
+
+__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,
+    const 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,30 @@
+#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 output_desc,
+	infiniopTensorDescriptor_t input_desc,
+	infiniopTensorDescriptor_t index_desc,
+	size_t dim
+);
+
+__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 * output,
+	const void * input,
+	const void * index,
+    void *stream
+);
+
+__C __export infiniStatus_t infiniopDestroyIndexCopyInplaceDescriptor(infiniopIndexCopyInplaceDescriptor_t desc);
+
+#endif

include/infiniop/ops/linear.h

@@ -0,0 +1,32 @@
+#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 y_desc,
+    infiniopTensorDescriptor_t x_desc,
+    infiniopTensorDescriptor_t w_desc,
+    infiniopTensorDescriptor_t b_desc
+);
+
+__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,35 @@
+#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_x_desc,
+	infiniopTensorDescriptor_t grad_w_desc,
+	infiniopTensorDescriptor_t grad_b_desc,
+	infiniopTensorDescriptor_t grad_y_desc,
+	infiniopTensorDescriptor_t x_desc,
+	infiniopTensorDescriptor_t w_desc
+);
+
+__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/scatter.h

@@ -0,0 +1,30 @@
+#ifndef __INFINIOP_SCATTER_API_H__
+#define __INFINIOP_SCATTER_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopScatterDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateScatterDescriptor(
+    infiniopHandle_t handle,
+    infiniopScatterDescriptor_t *desc_ptr,
+	infiniopTensorDescriptor_t output_desc,
+	infiniopTensorDescriptor_t input_desc,
+	infiniopTensorDescriptor_t index_desc,
+	size_t dim
+);
+
+__C __export infiniStatus_t infiniopGetScatterWorkspaceSize(infiniopScatterDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopScatter(infiniopScatterDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+	void * output,
+	const void * input,
+	const void * index,
+    void *stream
+);
+
+__C __export infiniStatus_t infiniopDestroyScatterDescriptor(infiniopScatterDescriptor_t desc);
+
+#endif

include/infiniop/ops/tril.h

@@ -0,0 +1,28 @@
+#ifndef __INFINIOP_TRIL_API_H__
+#define __INFINIOP_TRIL_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTrilDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateTrilDescriptor(
+    infiniopHandle_t handle,
+    infiniopTrilDescriptor_t *desc_ptr,
+	infiniopTensorDescriptor_t output_desc,
+	infiniopTensorDescriptor_t input_desc,
+	int diagonal
+);
+
+__C __export infiniStatus_t infiniopGetTrilWorkspaceSize(infiniopTrilDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopTril(infiniopTrilDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+	void * output,
+	const void * input,
+    void *stream
+);
+
+__C __export infiniStatus_t infiniopDestroyTrilDescriptor(infiniopTrilDescriptor_t desc);
+
+#endif

include/infiniop/ops/triu.h

@@ -0,0 +1,28 @@
+#ifndef __INFINIOP_TRIU_API_H__
+#define __INFINIOP_TRIU_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopTriuDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateTriuDescriptor(
+    infiniopHandle_t handle,
+    infiniopTriuDescriptor_t *desc_ptr,
+	infiniopTensorDescriptor_t output_desc,
+	infiniopTensorDescriptor_t input_desc,
+	int diagonal
+);
+
+__C __export infiniStatus_t infiniopGetTriuWorkspaceSize(infiniopTriuDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopTriu(infiniopTriuDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+	void * output,
+	const void * input,
+    void *stream
+);
+
+__C __export infiniStatus_t infiniopDestroyTriuDescriptor(infiniopTriuDescriptor_t desc);
+
+#endif

scripts/python_test.py

@@ -12,18 +12,25 @@
 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",
+        "index_copy_inplace.py",
+        "gather.py",
+        "scatter.py",
+        "tril.py",
+        "triu.py",
+        "linear.py",
+        "linear_backward.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(index_copy_inplace)
+DECLARE_INFINIOP_TEST(gather)
+DECLARE_INFINIOP_TEST(scatter)
+DECLARE_INFINIOP_TEST(tril)
+DECLARE_INFINIOP_TEST(triu)
+DECLARE_INFINIOP_TEST(linear)
+DECLARE_INFINIOP_TEST(linear_backward)
 
 #define REGISTER_INFINIOP_TEST(name)                      \
     {                                                     \
@@ -43,6 +50,13 @@ DECLARE_INFINIOP_TEST(sub)
         REGISTER_INFINIOP_TEST(causal_softmax) \
         REGISTER_INFINIOP_TEST(rearrange)      \
         REGISTER_INFINIOP_TEST(sub)            \
+        REGISTER_INFINIOP_TEST(index_copy_inplace)          \
+        REGISTER_INFINIOP_TEST(gather)                      \
+        REGISTER_INFINIOP_TEST(scatter)                     \
+        REGISTER_INFINIOP_TEST(tril)                        \
+        REGISTER_INFINIOP_TEST(triu)                        \
+        REGISTER_INFINIOP_TEST(linear)                      \
+        REGISTER_INFINIOP_TEST(linear_backward)             \
     }
 
 namespace infiniop_test {