Skip to content

[训练营]新添加算子clip,where, gather, reduce #164

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions include/infini_operators.h
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,23 @@
#include "ops/causal_softmax/causal_softmax.h"
#include "ops/global_avg_pool/global_avg_pool.h"
#include "ops/expand/expand.h"
#include "ops/gather_elements/gather_elements.h"
#include "ops/gather/gather.h"
#include "ops/gemm/gemm.h"
#include "ops/clip/clip.h"
#include "ops/conv/conv.h"
#include "ops/matmul/matmul.h"
#include "ops/max_pool/max_pool.h"
#include "ops/mlp/mlp.h"
#include "ops/random_sample/random_sample.h"
#include "ops/rearrange/rearrange.h"
#include "ops/reduceMax/reduceMax.h"
#include "ops/reduceMean/reduceMean.h"
#include "ops/reduceMin/reduceMin.h"
#include "ops/reduceSum/reduceSum.h"
#include "ops/relu/relu.h"
#include "ops/rms_norm/rms_norm.h"
#include "ops/rotary_embedding/rotary_embedding.h"
#include "ops/swiglu/swiglu.h"
#include "ops/where/where.h"
#include "tensor/tensor_descriptor.h"
28 changes: 28 additions & 0 deletions include/ops/clip/clip.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
#ifndef CLIP_H
#define CLIP_H

#include "../../export.h"
#include "../../operators.h"
#include <limits>

typedef struct{
Device device;
}ClipDescriptor;

typedef ClipDescriptor *infiniopClipDescriptor_t;

__C __export infiniopStatus_t infiniopCreateClipDescriptor(infiniopHandle_t handle,
infiniopClipDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t dst,
infiniopTensorDescriptor_t src,
float max,
float min);

__C __export infiniopStatus_t infiniopClip(infiniopClipDescriptor_t desc,
void *dst,
void const *src,
void *stream);

__C __export infiniopStatus_t infiniopDestroyClipDescriptor(infiniopClipDescriptor_t desc);

#endif
28 changes: 28 additions & 0 deletions include/ops/gather/gather.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
#ifndef GATHER_H
#define GATHER_H

#include "../../export.h"
#include "../../operators.h"

typedef struct{
Device device;
}GatherDescriptor;

typedef GatherDescriptor *infiniopGatherDescriptor_t;

__C __export infiniopStatus_t infiniopCreateGatherDescriptor(infiniopHandle_t handle,
infiniopGatherDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t dst,
infiniopTensorDescriptor_t data,
infiniopTensorDescriptor_t indices,
int axis);

__C __export infiniopStatus_t infiniopGather(infiniopGatherDescriptor_t desc,
void *dst,
void const *data,
void const *indices,
void *stream);

__C __export infiniopStatus_t infiniopDestroyGatherDescriptor(infiniopGatherDescriptor_t desc);

#endif
28 changes: 28 additions & 0 deletions include/ops/gather_elements/gather_elements.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
#ifndef GATHER_ELEMENTS_H
#define GATHER_ELEMENTS_H

#include "../../export.h"
#include "../../operators.h"

typedef struct{
Device device;
}GatherElementsDescriptor;

typedef GatherElementsDescriptor *infiniopGatherElementsDescriptor_t;

__C __export infiniopStatus_t infiniopCreateGatherElementsDescriptor(infiniopHandle_t handle,
infiniopGatherElementsDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t dst,
infiniopTensorDescriptor_t data,
infiniopTensorDescriptor_t indices,
int axis);

__C __export infiniopStatus_t infiniopGatherElements(infiniopGatherElementsDescriptor_t desc,
void *dst,
void const *data,
void const *indices,
void *stream);

__C __export infiniopStatus_t infiniopDestroyGatherElementsDescriptor(infiniopGatherElementsDescriptor_t desc);

#endif
25 changes: 25 additions & 0 deletions include/ops/reduceMax/reduceMax.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
#ifndef REDUCEMAX_H
#define REDUCEMAX_H

#include "../../export.h"
#include "../../operators.h"

typedef struct ReduceMaxDescriptor {
Device device;
} ReduceMaxDescriptor;
typedef ReduceMaxDescriptor *infiniopReduceMaxDescriptor_t;

__C __export infiniopStatus_t infiniopCreateReduceMaxDescriptor(infiniopHandle_t handle,
infiniopReduceMaxDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t dst,
infiniopTensorDescriptor_t src,
int* axis,
const int num_axis,
int const keepdims);

__C __export infiniopStatus_t infiniopGetReduceMaxWorkspaceSize(infiniopReduceMaxDescriptor_t desc, uint64_t *size);

__C __export infiniopStatus_t infiniopReduceMax(infiniopReduceMaxDescriptor_t desc, void *workspace, uint64_t workspace_size, void *dst, void const *src, void *stream);

__C __export infiniopStatus_t infiniopDestroyReduceMaxDescriptor(infiniopReduceMaxDescriptor_t desc);
#endif
25 changes: 25 additions & 0 deletions include/ops/reduceMean/reduceMean.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
#ifndef REDUCEMEAN_H
#define REDUCEMEAN_H

#include "../../export.h"
#include "../../operators.h"

typedef struct ReduceMeanDescriptor {
Device device;
} ReduceMeanDescriptor;
typedef ReduceMeanDescriptor *infiniopReduceMeanDescriptor_t;

__C __export infiniopStatus_t infiniopCreateReduceMeanDescriptor(infiniopHandle_t handle,
infiniopReduceMeanDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t dst,
infiniopTensorDescriptor_t src,
int* axis,
const int num_axis,
int const keepdims);

__C __export infiniopStatus_t infiniopGetReduceMeanWorkspaceSize(infiniopReduceMeanDescriptor_t desc, uint64_t *size);

__C __export infiniopStatus_t infiniopReduceMean(infiniopReduceMeanDescriptor_t desc, void *workspace, uint64_t workspace_size, void *dst, void const *src, void *stream);

__C __export infiniopStatus_t infiniopDestroyReduceMeanDescriptor(infiniopReduceMeanDescriptor_t desc);
#endif
25 changes: 25 additions & 0 deletions include/ops/reduceMin/reduceMin.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
#ifndef REDUCEMIN_H
#define REDUCEMIN_H

#include "../../export.h"
#include "../../operators.h"

typedef struct ReduceMinDescriptor {
Device device;
} ReduceMinDescriptor;
typedef ReduceMinDescriptor *infiniopReduceMinDescriptor_t;

__C __export infiniopStatus_t infiniopCreateReduceMinDescriptor(infiniopHandle_t handle,
infiniopReduceMinDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t dst,
infiniopTensorDescriptor_t src,
int* axis,
const int num_axis,
int const keepdims);

__C __export infiniopStatus_t infiniopGetReduceMinWorkspaceSize(infiniopReduceMinDescriptor_t desc, uint64_t *size);

__C __export infiniopStatus_t infiniopReduceMin(infiniopReduceMinDescriptor_t desc, void *workspace, uint64_t workspace_size, void *dst, void const *src, void *stream);

__C __export infiniopStatus_t infiniopDestroyReduceMinDescriptor(infiniopReduceMinDescriptor_t desc);
#endif
25 changes: 25 additions & 0 deletions include/ops/reduceSum/reduceSum.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
#ifndef REDUCESUM_H
#define REDUCESUM_H

#include "../../export.h"
#include "../../operators.h"

typedef struct ReduceSumDescriptor {
Device device;
} ReduceSumDescriptor;
typedef ReduceSumDescriptor *infiniopReduceSumDescriptor_t;

__C __export infiniopStatus_t infiniopCreateReduceSumDescriptor(infiniopHandle_t handle,
infiniopReduceSumDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t dst,
infiniopTensorDescriptor_t src,
int* axis,
const int num_axis,
int const keepdims);

__C __export infiniopStatus_t infiniopGetReduceSumWorkspaceSize(infiniopReduceSumDescriptor_t desc, uint64_t *size);

__C __export infiniopStatus_t infiniopReduceSum(infiniopReduceSumDescriptor_t desc, void *workspace, uint64_t workspace_size, void *dst, void const *src, void *stream);

__C __export infiniopStatus_t infiniopDestroyReduceSumDescriptor(infiniopReduceSumDescriptor_t desc);
#endif
29 changes: 29 additions & 0 deletions include/ops/where/where.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
#ifndef WHERE_H
#define WHERE_H

#include "../../export.h"
#include "../../operators.h"

typedef struct WhereDescriptor{
Device device;
} WhereDescriptor;

typedef WhereDescriptor *infiniopWhereDescriptor_t;

__C __export infiniopStatus_t infiniopCreateWhereDescriptor(infiniopHandle_t handle,
infiniopWhereDescriptor_t *desc_ptr,
infiniopTensorDescriptor_t dst,
infiniopTensorDescriptor_t x,
infiniopTensorDescriptor_t y,
infiniopTensorDescriptor_t condition);

__C __export infiniopStatus_t infiniopWhere(infiniopWhereDescriptor_t desc,
void *dst,
void const *x,
void const *y,
void const *condition,
void *stream);

__C __export infiniopStatus_t infiniopDestroyWhereDescriptor(infiniopWhereDescriptor_t desc);

#endif
Loading