Skip to content

Commit 6120325

Browse files
ZiminliZiminli
committed
Rebase to the latest dev
1 parent 4f15232 commit 6120325

File tree

31 files changed

+711
-3
lines changed

31 files changed

+711
-3
lines changed

‎include/include/data_type.h‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
#ifndef __DATA_TYPE_H__
2+
#define __DATA_TYPE_H__
3+
4+
typedef struct DataLayout {
5+
unsigned short
6+
packed : 8,
7+
sign : 1,
8+
size : 7,
9+
mantissa : 8,
10+
exponent : 8;
11+
12+
bool operator==(const DataLayout &other) const {
13+
union TypePun {
14+
DataLayout layout;
15+
unsigned int i;
16+
} pun;
17+
pun.layout = *this;
18+
auto a_ = pun.i;
19+
pun.layout = other;
20+
auto b_ = pun.i;
21+
return a_ == b_;
22+
}
23+
24+
bool operator!=(const DataLayout &other) const {
25+
return !(*this == other);
26+
}
27+
} DataLayout;
28+
29+
typedef struct DataLayout DT;
30+
31+
// clang-format off
32+
constexpr static struct DataLayout
33+
I8 = {1, 1, 1, 7, 0},
34+
I16 = {1, 1, 2, 15, 0},
35+
I32 = {1, 1, 4, 31, 0},
36+
I64 = {1, 1, 8, 63, 0},
37+
U8 = {1, 0, 1, 8, 0},
38+
U16 = {1, 0, 2, 16, 0},
39+
U32 = {1, 0, 4, 32, 0},
40+
U64 = {1, 0, 8, 64, 0},
41+
F16 = {1, 1, 2, 10, 5},
42+
BF16 = {1, 1, 2, 7, 8},
43+
F32 = {1, 1, 4, 23, 8},
44+
F64 = {1, 1, 8, 52, 11};
45+
// clang-format on
46+
47+
#endif// __DATA_TYPE_H__

‎include/include/device.h‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
#ifndef __DEVICE_H__
2+
#define __DEVICE_H__
3+
4+
enum DeviceEnum {
5+
DevCpu,
6+
DevNvGpu,
7+
DevCambriconMlu,
8+
DevAscendNpu,
9+
};
10+
11+
typedef enum DeviceEnum Device;
12+
13+
#endif// __DEVICE_H__

‎include/include/export.h‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
#ifndef __EXPORT_H__
2+
#define __EXPORT_H__
3+
4+
#if defined(_WIN32)
5+
#define __export __declspec(dllexport)
6+
#elif defined(__GNUC__) && ((__GNUC__ >= 4) || (__GNUC__ == 3 && __GNUC_MINOR__ >= 3))
7+
#define __export __attribute__((visibility("default")))
8+
#else
9+
#define __export
10+
#endif
11+
12+
#ifdef __cplusplus
13+
#define __C extern "C"
14+
#else
15+
#define __C
16+
#endif
17+
18+
#endif// __EXPORT_H__

‎include/include/handle.h‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
#ifndef INFINIOP_HANDLE_H
2+
#define INFINIOP_HANDLE_H
3+
4+
#include "device.h"
5+
6+
typedef struct HandleStruct {
7+
Device device;
8+
} HandleStruct;
9+
10+
typedef HandleStruct *infiniopHandle_t;
11+
12+
#endif

‎include/include/handle/handle_export.h‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
#ifndef INFINIOP_HANDLE_EXPORT_H
2+
#define INFINIOP_HANDLE_EXPORT_H
3+
#include "../status.h"
4+
#include "../handle.h"
5+
#include "../export.h"
6+
#include "../device.h"
7+
8+
__C __export infiniopStatus_t infiniopCreateHandle(infiniopHandle_t *handle_ptr, Device device, int device_id);
9+
10+
__C __export infiniopStatus_t infiniopDestroyHandle(infiniopHandle_t handle);
11+
12+
#endif // INFINIOP_HANDLE_EXPORT_H

‎include/include/infini_operators.h‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
#include "handle/handle_export.h"
2+
#include "ops/add/add.h"
3+
#include "ops/attention/attention.h"
4+
#include "ops/avg_pool/avg_pool.h"
5+
#include "ops/causal_softmax/causal_softmax.h"
6+
#include "ops/global_avg_pool/global_avg_pool.h"
7+
#include "ops/expand/expand.h"
8+
#include "ops/gemm/gemm.h"
9+
#include "ops/conv/conv.h"
10+
#include "ops/matmul/matmul.h"
11+
#include "ops/max_pool/max_pool.h"
12+
#include "ops/mlp/mlp.h"
13+
#include "ops/random_sample/random_sample.h"
14+
#include "ops/rearrange/rearrange.h"
15+
#include "ops/relu/relu.h"
16+
#include "ops/rms_norm/rms_norm.h"
17+
#include "ops/rotary_embedding/rotary_embedding.h"
18+
#include "ops/swiglu/swiglu.h"
19+
#include "tensor/tensor_descriptor.h"

‎include/include/operators.h‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
#ifndef __OPERATORS_H__
2+
#define __OPERATORS_H__
3+
4+
#include "device.h"
5+
#include "tensor.h"
6+
#include "handle.h"
7+
#include "status.h"
8+
9+
#endif// __OPERATORS_H__

‎include/include/ops/add/add.h‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
#ifndef ADD_H
2+
#define ADD_H
3+
4+
#include "../../export.h"
5+
#include "../../operators.h"
6+
7+
typedef struct AddDescriptor {
8+
Device device;
9+
} AddDescriptor;
10+
11+
typedef AddDescriptor *infiniopAddDescriptor_t;
12+
13+
__C __export infiniopStatus_t infiniopCreateAddDescriptor(infiniopHandle_t handle,
14+
infiniopAddDescriptor_t *desc_ptr,
15+
infiniopTensorDescriptor_t c,
16+
infiniopTensorDescriptor_t a,
17+
infiniopTensorDescriptor_t b);
18+
19+
__C __export infiniopStatus_t infiniopAdd(infiniopAddDescriptor_t desc,
20+
void *c,
21+
void const *a,
22+
void const *b,
23+
void *stream);
24+
25+
__C __export infiniopStatus_t infiniopDestroyAddDescriptor(infiniopAddDescriptor_t desc);
26+
27+
#endif

‎include/include/ops/attention/attention.h‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
#ifndef ATTENTION_H
2+
#define ATTENTION_H
3+
4+
#include "../../export.h"
5+
#include "../../operators.h"
6+
#include "../matmul/matmul.h"
7+
#include "../swiglu/swiglu.h"
8+
9+
typedef struct AttentionDescriptor {
10+
Device device;
11+
} AttentionDescriptor;
12+
13+
typedef AttentionDescriptor *infiniopAttentionDescriptor_t;
14+
15+
__C __export infiniopStatus_t infiniopCreateAttentionDescriptor(infiniopHandle_t handle,
16+
infiniopAttentionDescriptor_t *desc_ptr,
17+
infiniopTensorDescriptor_t out_desc,
18+
infiniopTensorDescriptor_t q_desc,
19+
infiniopTensorDescriptor_t k_desc,
20+
infiniopTensorDescriptor_t v_desc,
21+
infiniopTensorDescriptor_t k_cache_desc,
22+
infiniopTensorDescriptor_t v_cache_desc,
23+
uint64_t pos);
24+
25+
__C __export infiniopStatus_t infiniopGetAttentionWorkspaceSize(infiniopAttentionDescriptor_t desc, uint64_t *size);
26+
27+
__C __export infiniopStatus_t infiniopAttention(infiniopAttentionDescriptor_t desc,
28+
void *workspace,
29+
uint64_t workspace_size,
30+
void *out,
31+
void const *q,
32+
void const *k,
33+
void const *v,
34+
void *k_cache,
35+
void *v_cache,
36+
void *stream);
37+
38+
__C __export infiniopStatus_t infiniopDestroyAttentionDescriptor(infiniopAttentionDescriptor_t desc);
39+
#endif

‎include/include/ops/avg_pool/avg_pool.h‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
#ifndef AVG_POOL_H
2+
#define AVG_POOL_H
3+
4+
#include "../../export.h"
5+
#include "../../operators.h"
6+
7+
typedef struct AvgPoolDescriptor {
8+
Device device;
9+
} AvgPoolDescriptor;
10+
typedef AvgPoolDescriptor *infiniopAvgPoolDescriptor_t;
11+
12+
__C __export infiniopStatus_t infiniopCreateAvgPoolDescriptor(infiniopHandle_t handle,
13+
infiniopAvgPoolDescriptor_t *desc_ptr,
14+
infiniopTensorDescriptor_t y,
15+
infiniopTensorDescriptor_t x,
16+
uint64_t const *kernel_shape,
17+
uint64_t const *pads,
18+
int64_t const *strides,
19+
uint64_t n);
20+
21+
__C __export infiniopStatus_t infiniopGetAvgPoolWorkspaceSize(infiniopAvgPoolDescriptor_t desc, uint64_t *size);
22+
23+
__C __export infiniopStatus_t infiniopAvgPool(infiniopAvgPoolDescriptor_t desc,
24+
void *workspace, uint64_t workspace_size,
25+
void *y, void const *x, void *stream);
26+
27+
__C __export infiniopStatus_t infiniopDestroyAvgPoolDescriptor(infiniopAvgPoolDescriptor_t desc);
28+
#endif

0 commit comments

Comments
 (0)