Skip to content

Commit d3b0569

Browse files
gongchensugongchensu
committed
Issue/888 - Add gt,lt,ge,le,eq,ne,logical_and,logical_or,logical_xor,sin,bitwise_and, bitwise_or, bitwise_xor, bitwise_left_shift, bitwise_right_shift,floor_divide,atan2,exp2,log2,log10,rsqrt,square,hypot,copysign,remainder,isnan,isfinite,isinf,sinc,fmin,fmax,log1p binary operators.
1 parent 05096ea commit d3b0569

202 files changed

Lines changed: 4861 additions & 170 deletions

File tree

Some content is hidden

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

‎include/infiniop/ops/binary_ops_api.h‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,30 @@
1515

1616
// Declare all binary operator APIs
1717
BINARY_OP_API_DECLARE(div, Div)
18+
BINARY_OP_API_DECLARE(floor_divide, FloorDivide)
1819
BINARY_OP_API_DECLARE(pow, Pow)
20+
BINARY_OP_API_DECLARE(copysign, CopySign)
21+
BINARY_OP_API_DECLARE(hypot, Hypot)
22+
BINARY_OP_API_DECLARE(atan2, Atan2)
1923
BINARY_OP_API_DECLARE(mod, Mod)
24+
BINARY_OP_API_DECLARE(remainder, Remainder)
2025
BINARY_OP_API_DECLARE(max, Max)
2126
BINARY_OP_API_DECLARE(min, Min)
27+
BINARY_OP_API_DECLARE(fmax, Fmax)
28+
BINARY_OP_API_DECLARE(fmin, Fmin)
29+
BINARY_OP_API_DECLARE(gt, Gt)
30+
BINARY_OP_API_DECLARE(lt, Lt)
31+
BINARY_OP_API_DECLARE(ge, Ge)
32+
BINARY_OP_API_DECLARE(le, Le)
33+
BINARY_OP_API_DECLARE(eq, Eq)
34+
BINARY_OP_API_DECLARE(ne, Ne)
35+
BINARY_OP_API_DECLARE(logical_and, LogicalAnd)
36+
BINARY_OP_API_DECLARE(logical_or, LogicalOr)
37+
BINARY_OP_API_DECLARE(logical_xor, LogicalXor)
38+
BINARY_OP_API_DECLARE(bitwise_and, BitwiseAnd)
39+
BINARY_OP_API_DECLARE(bitwise_or, BitwiseOr)
40+
BINARY_OP_API_DECLARE(bitwise_xor, BitwiseXor)
41+
BINARY_OP_API_DECLARE(bitwise_left_shift, BitwiseLeftShift)
42+
BINARY_OP_API_DECLARE(bitwise_right_shift, BitwiseRightShift)
2243

2344
#endif // __INFINIOP_BINARY_OPS_API_H__

‎include/infiniop/ops/unary_ops_api.h‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,12 @@
1717
// Declare all unary operator APIs
1818
UNARY_OP_API_DECLARE(abs, Abs)
1919
UNARY_OP_API_DECLARE(log, Log)
20+
UNARY_OP_API_DECLARE(log2, Log2)
21+
UNARY_OP_API_DECLARE(log10, Log10)
22+
UNARY_OP_API_DECLARE(log1p, Log1p)
2023
UNARY_OP_API_DECLARE(sqrt, Sqrt)
24+
UNARY_OP_API_DECLARE(square, Square)
25+
UNARY_OP_API_DECLARE(rsqrt, Rsqrt)
2126
UNARY_OP_API_DECLARE(reciprocal, Reciprocal)
2227
UNARY_OP_API_DECLARE(neg, Neg)
2328
UNARY_OP_API_DECLARE(round, Round)
@@ -36,6 +41,12 @@ UNARY_OP_API_DECLARE(atan, Atan)
3641
UNARY_OP_API_DECLARE(acos, Acos)
3742
UNARY_OP_API_DECLARE(ceil, Ceil)
3843
UNARY_OP_API_DECLARE(exp, Exp)
44+
UNARY_OP_API_DECLARE(exp2, Exp2)
3945
UNARY_OP_API_DECLARE(hardswish, Hardswish)
46+
UNARY_OP_API_DECLARE(isnan, IsNan)
47+
UNARY_OP_API_DECLARE(isinf, IsInf)
48+
UNARY_OP_API_DECLARE(isfinite, IsFinite)
49+
UNARY_OP_API_DECLARE(sinc, Sinc)
50+
UNARY_OP_API_DECLARE(sin, Sin)
4051

4152
#endif // __INFINIOP_UNARY_OPS_API_H__

‎src/infiniop/elementwise/binary.h‎

Lines changed: 533 additions & 11 deletions
Large diffs are not rendered by default.

‎src/infiniop/elementwise/cpu/elementwise_cpu_impl.h‎

Lines changed: 86 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -32,59 +32,71 @@
3232
/**
3333
* @brief Common Calculate Switch Cases (F16 & F32)
3434
*/
35-
#define _IMPL_CALC_CASES_COMMON \
36-
case INFINI_DTYPE_F16: \
35+
#define _IMPL_CALC_CASES_COMMON \
36+
case INFINI_DTYPE_F16: \
3737
return _device_info->template calculate<Op, fp16_t>(_info, output, inputs, stream); \
38-
case INFINI_DTYPE_F32: \
38+
case INFINI_DTYPE_F32: \
3939
return _device_info->template calculate<Op, float>(_info, output, inputs, stream);
4040

4141
/**
4242
* @brief Extended Calculate Switch Cases (Adds F64 & BF16)
4343
*/
44-
#define _IMPL_CALC_CASES_EXTENDED \
45-
_IMPL_CALC_CASES_COMMON \
46-
case INFINI_DTYPE_F64: \
44+
#define _IMPL_CALC_CASES_EXTENDED \
45+
_IMPL_CALC_CASES_COMMON \
46+
case INFINI_DTYPE_F64: \
4747
return _device_info->template calculate<Op, double>(_info, output, inputs, stream); \
48-
case INFINI_DTYPE_BF16: \
48+
case INFINI_DTYPE_BF16: \
4949
return _device_info->template calculate<Op, bf16_t>(_info, output, inputs, stream);
5050

51+
/**
52+
* @brief Integral Calculate Switch Cases (I32, I64, U8)
53+
* For bitwise operations that only support integral types
54+
*/
55+
#define _IMPL_CALC_CASES_INTEGRAL \
56+
case INFINI_DTYPE_I32: \
57+
return _device_info->template calculate<Op, int32_t>(_info, output, inputs, stream); \
58+
case INFINI_DTYPE_I64: \
59+
return _device_info->template calculate<Op, int64_t>(_info, output, inputs, stream); \
60+
case INFINI_DTYPE_U8: \
61+
return _device_info->template calculate<Op, uint8_t>(_info, output, inputs, stream);
62+
5163
/**
5264
* @brief Generic Template for the Calculate method
5365
* @param CASES_MACRO The macro containing the switch cases to use
5466
*/
55-
#define _IMPL_CALCULATE_METHOD(CASES_MACRO) \
56-
infiniStatus_t Descriptor::calculate( \
57-
void *workspace, \
58-
size_t workspace_size, \
59-
void *output, \
60-
std::vector<const void *> inputs, \
61-
void *stream) const { \
62-
switch (_dtype) { \
63-
CASES_MACRO \
64-
default: \
65-
return INFINI_STATUS_BAD_TENSOR_DTYPE; \
66-
} \
67+
#define _IMPL_CALCULATE_METHOD(CASES_MACRO) \
68+
infiniStatus_t Descriptor::calculate( \
69+
void *workspace, \
70+
size_t workspace_size, \
71+
void *output, \
72+
std::vector<const void *> inputs, \
73+
void *stream) const { \
74+
switch (_dtype) { \
75+
CASES_MACRO \
76+
default: \
77+
return INFINI_STATUS_BAD_TENSOR_DTYPE; \
78+
} \
6779
}
6880

6981
/**
7082
* @brief Generic Template for the Create method
7183
* @param SHAPE_CHECK_BLOCK Code block to execute for shape checking
7284
* @param ... Variadic arguments for allowed data types in CHECK_DTYPE
7385
*/
74-
#define _IMPL_CREATE_METHOD(SHAPE_CHECK_BLOCK, ...) \
75-
Descriptor::~Descriptor() = default; \
76-
infiniStatus_t Descriptor::create( \
77-
infiniopHandle_t handle_, \
78-
Descriptor **desc_ptr, \
79-
infiniopTensorDescriptor_t out_desc, \
80-
std::vector<infiniopTensorDescriptor_t> input_desc_vec) { \
81-
auto handle = reinterpret_cast<device::cpu::Handle *>(handle_); \
82-
auto dtype = out_desc->dtype(); \
83-
const auto &out_shape = out_desc->shape(); \
84-
SHAPE_CHECK_BLOCK \
85-
CHECK_DTYPE(dtype, __VA_ARGS__); \
86+
#define _IMPL_CREATE_METHOD(SHAPE_CHECK_BLOCK, ...) \
87+
Descriptor::~Descriptor() = default; \
88+
infiniStatus_t Descriptor::create( \
89+
infiniopHandle_t handle_, \
90+
Descriptor **desc_ptr, \
91+
infiniopTensorDescriptor_t out_desc, \
92+
std::vector<infiniopTensorDescriptor_t> input_desc_vec) { \
93+
auto handle = reinterpret_cast<device::cpu::Handle *>(handle_); \
94+
auto dtype = out_desc->dtype(); \
95+
const auto &out_shape = out_desc->shape(); \
96+
SHAPE_CHECK_BLOCK \
97+
CHECK_DTYPE(dtype, __VA_ARGS__); \
8698
CREATE_ELEMENTWISE_CPU_DESCRIPTOR(handle, dtype, out_desc, input_desc_vec); \
87-
return INFINI_STATUS_SUCCESS; \
99+
return INFINI_STATUS_SUCCESS; \
88100
}
89101

90102
// =========================================================================
@@ -103,15 +115,14 @@
103115
* ELEMENTWISE_CPU_IMPL_BINARY(pow)
104116
* }
105117
*/
106-
#define ELEMENTWISE_CPU_IMPL_BINARY(OP) \
107-
_IMPL_CREATE_METHOD( \
108-
const auto &a_desc = input_desc_vec.at(0); \
109-
const auto &b_desc = input_desc_vec.at(1); \
110-
const auto &a_shape = a_desc->shape(); \
111-
const auto &b_shape = b_desc->shape(); \
112-
CHECK_SAME_SHAPE(out_shape, a_shape, b_shape);, \
113-
INFINI_DTYPE_F16, INFINI_DTYPE_F32 \
114-
) \
118+
#define ELEMENTWISE_CPU_IMPL_BINARY(OP) \
119+
_IMPL_CREATE_METHOD( \
120+
const auto &a_desc = input_desc_vec.at(0); \
121+
const auto &b_desc = input_desc_vec.at(1); \
122+
const auto &a_shape = a_desc->shape(); \
123+
const auto &b_shape = b_desc->shape(); \
124+
CHECK_SAME_SHAPE(out_shape, a_shape, b_shape);, \
125+
INFINI_DTYPE_F16, INFINI_DTYPE_F32) \
115126
_IMPL_CALCULATE_METHOD(_IMPL_CALC_CASES_COMMON)
116127

117128
/**
@@ -126,13 +137,12 @@
126137
* ELEMENTWISE_CPU_IMPL_UNARY(sqrt)
127138
* }
128139
*/
129-
#define ELEMENTWISE_CPU_IMPL_UNARY(OP) \
130-
_IMPL_CREATE_METHOD( \
131-
const auto &x_desc = input_desc_vec.at(0); \
132-
const auto &x_shape = x_desc->shape(); \
133-
CHECK_SAME_SHAPE(out_shape, x_shape);, \
134-
INFINI_DTYPE_F16, INFINI_DTYPE_F32 \
135-
) \
140+
#define ELEMENTWISE_CPU_IMPL_UNARY(OP) \
141+
_IMPL_CREATE_METHOD( \
142+
const auto &x_desc = input_desc_vec.at(0); \
143+
const auto &x_shape = x_desc->shape(); \
144+
CHECK_SAME_SHAPE(out_shape, x_shape);, \
145+
INFINI_DTYPE_F16, INFINI_DTYPE_F32) \
136146
_IMPL_CALCULATE_METHOD(_IMPL_CALC_CASES_COMMON)
137147

138148
/**
@@ -147,13 +157,34 @@
147157
* ELEMENTWISE_CPU_IMPL_UNARY_EXTENDED(exp)
148158
* }
149159
*/
150-
#define ELEMENTWISE_CPU_IMPL_UNARY_EXTENDED(OP) \
151-
_IMPL_CREATE_METHOD( \
152-
const auto &x_desc = input_desc_vec.at(0); \
153-
const auto &x_shape = x_desc->shape(); \
154-
CHECK_SAME_SHAPE(out_shape, x_shape);, \
155-
INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16 \
156-
) \
160+
#define ELEMENTWISE_CPU_IMPL_UNARY_EXTENDED(OP) \
161+
_IMPL_CREATE_METHOD( \
162+
const auto &x_desc = input_desc_vec.at(0); \
163+
const auto &x_shape = x_desc->shape(); \
164+
CHECK_SAME_SHAPE(out_shape, x_shape);, \
165+
INFINI_DTYPE_F16, INFINI_DTYPE_F32, INFINI_DTYPE_F64, INFINI_DTYPE_BF16) \
157166
_IMPL_CALCULATE_METHOD(_IMPL_CALC_CASES_EXTENDED)
158167

168+
/**
169+
* @brief Implementation for Binary Operators with Integral Types (I32, I64, U8)
170+
*
171+
* This macro generates the Descriptor destructor, create, and calculate methods
172+
* for binary operators that only support integral types (e.g., bitwise operations).
173+
*
174+
* Usage:
175+
* namespace op::bitwise_and::cpu {
176+
* using Op = op::elementwise::binary::BinaryOp<BinaryMode::BitwiseAnd>;
177+
* ELEMENTWISE_CPU_IMPL_BINARY_INTEGRAL(bitwise_and)
178+
* }
179+
*/
180+
#define ELEMENTWISE_CPU_IMPL_BINARY_INTEGRAL(OP) \
181+
_IMPL_CREATE_METHOD( \
182+
const auto &a_desc = input_desc_vec.at(0); \
183+
const auto &b_desc = input_desc_vec.at(1); \
184+
const auto &a_shape = a_desc->shape(); \
185+
const auto &b_shape = b_desc->shape(); \
186+
CHECK_SAME_SHAPE(out_shape, a_shape, b_shape);, \
187+
INFINI_DTYPE_I32, INFINI_DTYPE_I64, INFINI_DTYPE_U8) \
188+
_IMPL_CALCULATE_METHOD(_IMPL_CALC_CASES_INTEGRAL)
189+
159190
#endif // __INFINIOP_ELEMENTWISE_CPU_IMPL_H__

0 commit comments

Comments
 (0)