Skip to content

Commit 57bc896

Browse files
PanZezhong1725Gary828
PanZezhong1725
and
Gary828
committed
issue/1031 merge T1-1-4
Co-authored-by: Gary828 <garygali@csu.edu.cn>
1 parent 7f29544 commit 57bc896

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

98 files changed

+7009
-36
lines changed

‎include/infinicore/ops/all.hpp‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
#pragma once
2+
3+
#include "../device.hpp"
4+
#include "common/op.hpp"
5+
#include <optional>
6+
#include <vector>
7+
namespace infinicore::op {
8+
class All {
9+
public:
10+
using schema = void (*)(Tensor, Tensor, std::vector<size_t>, bool);
11+
static void execute(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim = false);
12+
static common::OpDispatcher<schema> &dispatcher();
13+
};
14+
15+
Tensor all(Tensor input, std::vector<size_t> dim, bool keepdim = false);
16+
void all_(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim = false);
17+
18+
} // namespace infinicore::op

‎include/infinicore/ops/sum.hpp‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
#pragma once
2+
3+
#include "../device.hpp"
4+
#include "common/op.hpp"
5+
#include <optional>
6+
#include <vector>
7+
8+
namespace infinicore::op {
9+
class Sum {
10+
public:
11+
using schema = void (*)(Tensor, Tensor, std::vector<size_t>, bool);
12+
static void execute(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim = false);
13+
static common::OpDispatcher<schema> &dispatcher();
14+
};
15+
16+
Tensor sum(Tensor input, std::vector<size_t> dim, bool keepdim = false);
17+
void sum_(Tensor output, Tensor input, std::vector<size_t> dim, bool keepdim = false);
18+
19+
} // namespace infinicore::op

‎include/infinicore/ops/topk.hpp‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
#pragma once
2+
3+
#include "../device.hpp"
4+
#include "common/op.hpp"
5+
namespace infinicore::op {
6+
class TopK {
7+
public:
8+
using schema = void (*)(Tensor, Tensor, Tensor, size_t, size_t, bool, bool);
9+
static void execute(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true);
10+
static common::OpDispatcher<schema> &dispatcher();
11+
};
12+
13+
std::pair<Tensor, Tensor> topk(Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true);
14+
void topk_(Tensor values_output, Tensor indices_output, Tensor input, size_t k, size_t dim, bool largest = true, bool sorted = true);
15+
16+
} // namespace infinicore::op

‎include/infinicore/ops/var.hpp‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
#pragma once
2+
3+
#include "../device.hpp"
4+
#include "common/op.hpp"
5+
#include <optional>
6+
#include <utility>
7+
#include <vector>
8+
namespace infinicore::op {
9+
class Var {
10+
public:
11+
using schema = void (*)(Tensor, Tensor, std::vector<size_t>, bool, bool); // var_output, input, dim, unbiased, keepdim
12+
static void execute(Tensor var_output, Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
13+
static common::OpDispatcher<schema> &dispatcher();
14+
};
15+
16+
Tensor var(Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
17+
void var_(Tensor var_output, Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
18+
19+
} // namespace infinicore::op

‎include/infinicore/ops/var_mean.hpp‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
#pragma once
2+
3+
#include "../device.hpp"
4+
#include "common/op.hpp"
5+
#include <optional>
6+
#include <utility>
7+
#include <vector>
8+
namespace infinicore::op {
9+
class Var_Mean {
10+
public:
11+
using schema = void (*)(Tensor, Tensor, Tensor, std::vector<size_t>, bool, bool); // var_output, mean_output, input, dim, unbiased, keepdim
12+
static void execute(Tensor var_output, Tensor mean_output, Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
13+
static common::OpDispatcher<schema> &dispatcher();
14+
};
15+
16+
std::pair<Tensor, Tensor> var_mean(Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
17+
void var_mean_(Tensor var_output, Tensor mean_output, Tensor input, std::vector<size_t> dim, bool unbiased = true, bool keepdim = false);
18+
19+
} // namespace infinicore::op

‎include/infiniop.h‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
#include "infiniop/ops/adaptive_max_pool1d.h"
66
#include "infiniop/ops/add.h"
77
#include "infiniop/ops/add_rms_norm.h"
8+
#include "infiniop/ops/all.h"
89
#include "infiniop/ops/asinh.h"
910
#include "infiniop/ops/attention.h"
1011
#include "infiniop/ops/causal_softmax.h"
@@ -38,17 +39,21 @@
3839
#include "infiniop/ops/softmax.h"
3940
#include "infiniop/ops/softplus.h"
4041
#include "infiniop/ops/sub.h"
42+
#include "infiniop/ops/sum.h"
4143
#include "infiniop/ops/swiglu.h"
4244
#include "infiniop/ops/tanh.h"
45+
#include "infiniop/ops/topk.h"
4346
#include "infiniop/ops/topkrouter.h"
4447
#include "infiniop/ops/topksoftmax.h"
48+
#include "infiniop/ops/var.h"
49+
#include "infiniop/ops/var_mean.h"
4550
#include "infiniop/ops/zeros.h"
4651
#include "infiniop/tensor_descriptor.h"
4752

48-
#include "infiniop/ops/cross_entropy.h"
49-
#include "infiniop/ops/hardswish.h"
5053
#include "infiniop/ops/avg_pool1d.h"
54+
#include "infiniop/ops/cross_entropy.h"
5155
#include "infiniop/ops/equal.h"
56+
#include "infiniop/ops/hardswish.h"
5257
#include "infiniop/ops/hardtanh.h"
5358

5459
#endif // __INFINIOP_API_H__

‎include/infiniop/ops/all.h‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
#ifndef __INFINIOP_ALL_API_H__
2+
#define __INFINIOP_ALL_API_H__
3+
4+
#include "../operator_descriptor.h"
5+
#include <cstddef>
6+
#include <vector>
7+
typedef struct InfiniopDescriptor *infiniopAllDescriptor_t;
8+
9+
__C __export infiniStatus_t infiniopCreateAllDescriptor(infiniopHandle_t handle,
10+
infiniopAllDescriptor_t *desc_ptr,
11+
infiniopTensorDescriptor_t output_desc,
12+
infiniopTensorDescriptor_t input_desc,
13+
size_t *dim,
14+
size_t dim_size,
15+
bool keepdim);
16+
17+
__C __export infiniStatus_t infiniopGetAllWorkspaceSize(infiniopAllDescriptor_t desc, size_t *size);
18+
19+
__C __export infiniStatus_t infiniopAll(infiniopAllDescriptor_t desc,
20+
void *workspace,
21+
size_t workspace_size,
22+
void *output,
23+
const void *input,
24+
size_t *dim,
25+
size_t dim_size,
26+
bool keepdim,
27+
void *stream);
28+
29+
__C __export infiniStatus_t infiniopDestroyAllDescriptor(infiniopAllDescriptor_t desc);
30+
31+
#endif

‎include/infiniop/ops/sum.h‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
#ifndef __INFINIOP_SUM_API_H__
2+
#define __INFINIOP_SUM_API_H__
3+
4+
#include "../operator_descriptor.h"
5+
#include <cstddef>
6+
#include <vector>
7+
typedef struct InfiniopDescriptor *infiniopSumDescriptor_t;
8+
9+
__C __export infiniStatus_t infiniopCreateSumDescriptor(infiniopHandle_t handle,
10+
infiniopSumDescriptor_t *desc_ptr,
11+
infiniopTensorDescriptor_t output_desc,
12+
infiniopTensorDescriptor_t input_desc,
13+
size_t *dim,
14+
size_t dim_size,
15+
bool keepdim);
16+
17+
__C __export infiniStatus_t infiniopGetSumWorkspaceSize(infiniopSumDescriptor_t desc, size_t *size);
18+
19+
__C __export infiniStatus_t infiniopSum(infiniopSumDescriptor_t desc,
20+
void *workspace,
21+
size_t workspace_size,
22+
void *output,
23+
const void *input,
24+
size_t *dim,
25+
size_t dim_size,
26+
bool keepdim,
27+
void *stream);
28+
29+
__C __export infiniStatus_t infiniopDestroySumDescriptor(infiniopSumDescriptor_t desc);
30+
31+
#endif

‎include/infiniop/ops/topk.h‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
#ifndef __INFINIOP_TOPK_API_H__
2+
#define __INFINIOP_TOPK_API_H__
3+
4+
#include "../operator_descriptor.h"
5+
#include <cstddef>
6+
#include <vector>
7+
typedef struct InfiniopDescriptor *infiniopTopKDescriptor_t;
8+
9+
__C __export infiniStatus_t infiniopCreateTopKDescriptor(infiniopHandle_t handle,
10+
infiniopTopKDescriptor_t *desc_ptr,
11+
infiniopTensorDescriptor_t values_output_desc,
12+
infiniopTensorDescriptor_t indices_output_desc,
13+
infiniopTensorDescriptor_t input_desc,
14+
size_t k,
15+
size_t dim,
16+
bool largest,
17+
bool sorted);
18+
19+
__C __export infiniStatus_t infiniopGetTopKWorkspaceSize(infiniopTopKDescriptor_t desc, size_t *size);
20+
21+
__C __export infiniStatus_t infiniopTopK(infiniopTopKDescriptor_t desc,
22+
void *workspace,
23+
size_t workspace_size,
24+
void *values_output,
25+
void *indices_output,
26+
const void *input,
27+
size_t k,
28+
size_t dim,
29+
bool largest,
30+
bool sorted,
31+
void *stream);
32+
33+
__C __export infiniStatus_t infiniopDestroyTopKDescriptor(infiniopTopKDescriptor_t desc);
34+
35+
#endif

‎include/infiniop/ops/var.h‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
#ifndef __INFINIOP_VAR_API_H__
2+
#define __INFINIOP_VAR_API_H__
3+
4+
#include "../operator_descriptor.h"
5+
#include <cstddef>
6+
#include <vector>
7+
typedef struct InfiniopDescriptor *infiniopVarDescriptor_t;
8+
9+
__C __export infiniStatus_t infiniopCreateVarDescriptor(infiniopHandle_t handle,
10+
infiniopVarDescriptor_t *desc_ptr,
11+
infiniopTensorDescriptor_t var_output_desc,
12+
infiniopTensorDescriptor_t input_desc,
13+
size_t *dim,
14+
size_t dim_size,
15+
bool unbiased,
16+
bool keepdim);
17+
18+
__C __export infiniStatus_t infiniopGetVarWorkspaceSize(infiniopVarDescriptor_t desc, size_t *size);
19+
20+
__C __export infiniStatus_t infiniopVar(infiniopVarDescriptor_t desc,
21+
void *workspace,
22+
size_t workspace_size,
23+
void *var_output,
24+
const void *input,
25+
size_t *dim,
26+
size_t dim_size,
27+
bool unbiased,
28+
bool keepdim,
29+
void *stream);
30+
31+
__C __export infiniStatus_t infiniopDestroyVarDescriptor(infiniopVarDescriptor_t desc);
32+
33+
#endif

0 commit comments

Comments
 (0)