Skip to content

Commit 9da0b53

Browse files
xu-zhengzhongxu-zhengzhong
committed
Split BLAS op info into local headers
1 parent f8e9cc2 commit 9da0b53

65 files changed

Lines changed: 875 additions & 902 deletions

Some content is hidden

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

‎src/infiniop/ops/asum/asum.h‎

Lines changed: 1 addition & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,8 @@
11
#ifndef __ASUM_H__
22
#define __ASUM_H__
33

4-
#include "../../../utils.h"
54
#include "../../operator.h"
6-
#include "../../tensor.h"
7-
#include "infiniop/ops/asum.h"
5+
#include "info.h"
86

97
#define DESCRIPTOR(NAMESPACE) \
108
\
@@ -46,45 +44,4 @@
4644
}; \
4745
}
4846

49-
class AsumInfo {
50-
private:
51-
size_t _size;
52-
ptrdiff_t _incx;
53-
infiniDtype_t _dtype;
54-
55-
AsumInfo(size_t size,
56-
ptrdiff_t incx,
57-
infiniDtype_t dtype)
58-
: _size(size), _incx(incx), _dtype(dtype) {}
59-
60-
public:
61-
inline size_t getSize() const { return _size; }
62-
inline ptrdiff_t getIncx() const { return _incx; }
63-
inline infiniDtype_t getDtype() const { return _dtype; }
64-
65-
using ResultType = utils::Result<AsumInfo>;
66-
67-
static ResultType createAsumInfo(
68-
infiniopTensorDescriptor_t x_desc,
69-
infiniopTensorDescriptor_t result_desc) {
70-
71-
CHECK_OR_RETURN(x_desc != nullptr, INFINI_STATUS_NULL_POINTER);
72-
CHECK_OR_RETURN(result_desc != nullptr, INFINI_STATUS_NULL_POINTER);
73-
74-
auto dtype = x_desc->dtype();
75-
76-
CHECK_OR_RETURN(result_desc->dtype() == dtype, INFINI_STATUS_BAD_TENSOR_DTYPE);
77-
CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64);
78-
79-
CHECK_OR_RETURN(x_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE);
80-
CHECK_OR_RETURN(result_desc->numel() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE);
81-
82-
auto size = x_desc->numel();
83-
auto incx = x_desc->stride(0);
84-
85-
AsumInfo info(size, incx, dtype);
86-
return ResultType(std::move(info));
87-
}
88-
};
89-
9047
#endif // __ASUM_H__

‎src/infiniop/ops/asum/bang/asum_bang.mlu‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,11 @@ infiniStatus_t Descriptor::create(
1313
infiniopTensorDescriptor_t result_desc) {
1414

1515
auto handle = reinterpret_cast<device::bang::Handle *>(handle_);
16-
auto info = AsumInfo::createAsumInfo(x_desc, result_desc);
17-
CHECK_RESULT(info);
16+
auto result = AsumInfo::createAsumInfo(x_desc, result_desc);
17+
CHECK_RESULT(result);
1818

1919
*desc_ptr = new Descriptor(
20-
info.take(),
20+
result.take(),
2121
0,
2222
nullptr,
2323
handle->device,
@@ -33,8 +33,8 @@ infiniStatus_t calculateAsum(
3333
Tdata *result,
3434
cnrtQueue_t queue) {
3535

36-
const size_t size = info.getSize();
37-
const ptrdiff_t incx = info.getIncx();
36+
const size_t n = info.n;
37+
const ptrdiff_t incx = info.incx;
3838

3939
cnrtDim3_t k_dim;
4040
cnrtFunctionType_t k_type;
@@ -46,12 +46,12 @@ infiniStatus_t calculateAsum(
4646

4747
if (incx == 1) {
4848
asumKernelContiguous<Tdata><<<k_dim, k_type, queue>>>(
49-
size,
49+
n,
5050
x,
5151
result);
5252
} else {
5353
asumKernelStrided<Tdata><<<k_dim, k_type, queue>>>(
54-
size,
54+
n,
5555
x,
5656
incx,
5757
result);
@@ -78,7 +78,7 @@ infiniStatus_t Descriptor::calculate(
7878
(void)workspace;
7979
(void)workspace_size;
8080

81-
switch (_info.getDtype()) {
81+
switch (_info.data_type) {
8282
case INFINI_DTYPE_F16:
8383
return CALCULATE_ASUM(half);
8484
case INFINI_DTYPE_F32:

‎src/infiniop/ops/asum/cpu/asum_cpu.cc‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,11 +12,11 @@ infiniStatus_t Descriptor::create(
1212
infiniopTensorDescriptor_t result_desc) {
1313

1414
auto handle = reinterpret_cast<device::cpu::Handle *>(handle_);
15-
auto info = AsumInfo::createAsumInfo(x_desc, result_desc);
16-
CHECK_RESULT(info);
15+
auto result = AsumInfo::createAsumInfo(x_desc, result_desc);
16+
CHECK_RESULT(result);
1717

1818
*desc_ptr = new Descriptor(
19-
info.take(),
19+
result.take(),
2020
0,
2121
nullptr,
2222
handle->device,
@@ -31,21 +31,21 @@ infiniStatus_t calculateAsum(
3131
const Tdata *x,
3232
Tdata *result) {
3333

34-
const ptrdiff_t size = info.getSize();
35-
const ptrdiff_t incx = info.getIncx();
34+
const ptrdiff_t n = info.n;
35+
const ptrdiff_t incx = info.incx;
3636

3737
if constexpr (std::is_same<Tdata, fp16_t>::value || std::is_same<Tdata, bf16_t>::value) {
3838
float total_sum = 0.0;
3939

40-
for (ptrdiff_t i = 0; i < size; ++i) {
40+
for (ptrdiff_t i = 0; i < n; ++i) {
4141
total_sum += std::abs(utils::cast<float>(x[i * incx]));
4242
}
4343

4444
result[0] = utils::cast<Tdata>(total_sum);
4545
} else {
4646
Tdata total_sum = 0.0;
4747

48-
for (ptrdiff_t i = 0; i < size; ++i) {
48+
for (ptrdiff_t i = 0; i < n; ++i) {
4949
total_sum += std::abs(x[i * incx]);
5050
}
5151

@@ -70,7 +70,7 @@ infiniStatus_t Descriptor::calculate(
7070
(void)workspace;
7171
(void)workspace_size;
7272

73-
switch (_info.getDtype()) {
73+
switch (_info.data_type) {
7474
case INFINI_DTYPE_F16:
7575
return CALCULATE_ASUM(fp16_t);
7676
case INFINI_DTYPE_BF16:

‎src/infiniop/ops/asum/info.h‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
#ifndef __ASUM_INFO_H__
2+
#define __ASUM_INFO_H__
3+
4+
#include "../../../utils.h"
5+
#include "../../tensor.h"
6+
7+
class AsumInfo {
8+
private:
9+
AsumInfo() = default;
10+
11+
public:
12+
size_t n;
13+
ptrdiff_t incx;
14+
infiniDtype_t data_type;
15+
16+
static utils::Result<AsumInfo> createAsumInfo(
17+
infiniopTensorDescriptor_t x_desc,
18+
infiniopTensorDescriptor_t result_desc) {
19+
20+
CHECK_OR_RETURN(x_desc != nullptr, INFINI_STATUS_NULL_POINTER);
21+
CHECK_OR_RETURN(result_desc != nullptr, INFINI_STATUS_NULL_POINTER);
22+
23+
auto data_type = x_desc->dtype();
24+
25+
CHECK_OR_RETURN(result_desc->dtype() == data_type, INFINI_STATUS_BAD_TENSOR_DTYPE);
26+
CHECK_DTYPE(data_type, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64);
27+
28+
CHECK_OR_RETURN(x_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE);
29+
CHECK_OR_RETURN(result_desc->numel() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE);
30+
31+
auto n = x_desc->numel();
32+
auto incx = x_desc->stride(0);
33+
34+
return utils::Result<AsumInfo>(AsumInfo{
35+
n,
36+
incx,
37+
data_type});
38+
}
39+
};
40+
41+
#endif // __ASUM_INFO_H__

‎src/infiniop/ops/asum/metax/asum_metax.cc‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,11 @@ infiniStatus_t Descriptor::create(
1919
infiniopTensorDescriptor_t result_desc) {
2020

2121
auto handle = reinterpret_cast<device::metax::Handle *>(handle_);
22-
auto info = AsumInfo::createAsumInfo(x_desc, result_desc);
23-
CHECK_RESULT(info);
22+
auto result = AsumInfo::createAsumInfo(x_desc, result_desc);
23+
CHECK_RESULT(result);
2424

2525
*desc_ptr = new Descriptor(
26-
info.take(),
26+
result.take(),
2727
0,
2828
new Opaque{handle->internal()},
2929
handle->device,
@@ -42,9 +42,9 @@ infiniStatus_t Descriptor::calculate(
4242
(void)workspace;
4343
(void)workspace_size;
4444

45-
const size_t size = _info.getSize();
46-
const ptrdiff_t incx = _info.getIncx();
47-
const infiniDtype_t data_type = _info.getDtype();
45+
const size_t n = _info.n;
46+
const ptrdiff_t incx = _info.incx;
47+
const infiniDtype_t data_type = _info.data_type;
4848

4949
CHECK_STATUS(_opaque->internal->useMcblas(
5050
(hcStream_t)stream,
@@ -53,10 +53,10 @@ infiniStatus_t Descriptor::calculate(
5353

5454
switch (data_type) {
5555
case INFINI_DTYPE_F32:
56-
CHECK_MCBLAS(hcblasSasum(handle, size, (const float *)x, incx, (float *)result));
56+
CHECK_MCBLAS(hcblasSasum(handle, n, (const float *)x, incx, (float *)result));
5757
break;
5858
case INFINI_DTYPE_F64:
59-
CHECK_MCBLAS(hcblasDasum(handle, size, (const double *)x, incx, (double *)result));
59+
CHECK_MCBLAS(hcblasDasum(handle, n, (const double *)x, incx, (double *)result));
6060
break;
6161
default:
6262
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;

‎src/infiniop/ops/axpy/axpy.h‎

Lines changed: 1 addition & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,8 @@
11
#ifndef __AXPY_H__
22
#define __AXPY_H__
33

4-
#include "../../../utils.h"
54
#include "../../operator.h"
6-
#include "../../tensor.h"
7-
#include "infiniop/ops/axpy.h"
5+
#include "info.h"
86

97
#define DESCRIPTOR(NAMESPACE) \
108
\
@@ -48,54 +46,4 @@
4846
}; \
4947
}
5048

51-
class AxpyInfo {
52-
private:
53-
size_t _size;
54-
ptrdiff_t _incx;
55-
ptrdiff_t _incy;
56-
infiniDtype_t _dtype;
57-
58-
AxpyInfo(size_t size,
59-
ptrdiff_t incx,
60-
ptrdiff_t incy,
61-
infiniDtype_t dtype)
62-
: _size(size), _incx(incx), _incy(incy), _dtype(dtype) {}
63-
64-
public:
65-
inline size_t getSize() const { return _size; }
66-
inline ptrdiff_t getIncx() const { return _incx; }
67-
inline ptrdiff_t getIncy() const { return _incy; }
68-
inline infiniDtype_t getDtype() const { return _dtype; }
69-
70-
using ResultType = utils::Result<AxpyInfo>;
71-
72-
static ResultType createAxpyInfo(
73-
infiniopTensorDescriptor_t alpha_desc,
74-
infiniopTensorDescriptor_t x_desc,
75-
infiniopTensorDescriptor_t y_desc) {
76-
77-
CHECK_OR_RETURN(alpha_desc != nullptr, INFINI_STATUS_NULL_POINTER);
78-
CHECK_OR_RETURN(x_desc != nullptr, INFINI_STATUS_NULL_POINTER);
79-
CHECK_OR_RETURN(y_desc != nullptr, INFINI_STATUS_NULL_POINTER);
80-
81-
auto dtype = x_desc->dtype();
82-
83-
CHECK_OR_RETURN(alpha_desc->dtype() == dtype, INFINI_STATUS_BAD_TENSOR_DTYPE);
84-
CHECK_OR_RETURN(x_desc->dtype() == dtype, INFINI_STATUS_BAD_TENSOR_DTYPE);
85-
CHECK_DTYPE(dtype, INFINI_DTYPE_F16, INFINI_DTYPE_BF16, INFINI_DTYPE_F32, INFINI_DTYPE_F64);
86-
87-
CHECK_OR_RETURN(alpha_desc->numel() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE);
88-
CHECK_OR_RETURN(x_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE);
89-
CHECK_OR_RETURN(y_desc->ndim() == 1, INFINI_STATUS_BAD_TENSOR_SHAPE);
90-
CHECK_OR_RETURN(x_desc->numel() == y_desc->numel(), INFINI_STATUS_BAD_TENSOR_SHAPE);
91-
92-
auto size = x_desc->numel();
93-
auto incx = x_desc->stride(0);
94-
auto incy = y_desc->stride(0);
95-
96-
AxpyInfo info(size, incx, incy, dtype);
97-
return ResultType(std::move(info));
98-
}
99-
};
100-
10149
#endif // __AXPY_H__

‎src/infiniop/ops/axpy/bang/axpy_bang.mlu‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,12 @@ infiniStatus_t Descriptor::create(
1414
infiniopTensorDescriptor_t y_desc) {
1515

1616
auto handle = reinterpret_cast<device::bang::Handle *>(handle_);
17-
auto info = AxpyInfo::createAxpyInfo(alpha_desc, x_desc, y_desc);
18-
CHECK_RESULT(info);
17+
auto result = AxpyInfo::createAxpyInfo(alpha_desc, x_desc, y_desc);
18+
CHECK_RESULT(result);
1919

2020
// Create descriptor
2121
*desc_ptr = new Descriptor(
22-
info.take(),
22+
result.take(),
2323
0,
2424
nullptr,
2525
handle->device,
@@ -36,9 +36,9 @@ infiniStatus_t calculateAxpy(
3636
Tdata *y,
3737
cnrtQueue_t queue) {
3838

39-
const size_t size = info.getSize();
40-
const ptrdiff_t incx = info.getIncx();
41-
const ptrdiff_t incy = info.getIncy();
39+
const size_t size = info.n;
40+
const ptrdiff_t incx = info.incx;
41+
const ptrdiff_t incy = info.incy;
4242

4343
cnrtDim3_t k_dim;
4444
cnrtFunctionType_t k_type;
@@ -87,7 +87,7 @@ infiniStatus_t Descriptor::calculate(
8787
(void)workspace;
8888
(void)workspace_size;
8989

90-
switch (_info.getDtype()) {
90+
switch (_info.data_type) {
9191
case INFINI_DTYPE_F16:
9292
return CALCULATE_AXPY(half);
9393
case INFINI_DTYPE_F32:

‎src/infiniop/ops/axpy/cpu/axpy_cpu.cc‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -13,11 +13,11 @@ infiniStatus_t Descriptor::create(
1313
infiniopTensorDescriptor_t y_desc) {
1414

1515
auto handle = reinterpret_cast<device::cpu::Handle *>(handle_);
16-
auto info = AxpyInfo::createAxpyInfo(alpha_desc, x_desc, y_desc);
17-
CHECK_RESULT(info);
16+
auto result = AxpyInfo::createAxpyInfo(alpha_desc, x_desc, y_desc);
17+
CHECK_RESULT(result);
1818

1919
*desc_ptr = new Descriptor(
20-
info.take(),
20+
result.take(),
2121
0,
2222
nullptr,
2323
handle->device,
@@ -33,9 +33,9 @@ infiniStatus_t calculateAxpy(
3333
const Tdata *x,
3434
Tdata *y) {
3535

36-
const ptrdiff_t size = info.getSize();
37-
const ptrdiff_t incx = info.getIncx();
38-
const ptrdiff_t incy = info.getIncy();
36+
const ptrdiff_t size = info.n;
37+
const ptrdiff_t incx = info.incx;
38+
const ptrdiff_t incy = info.incy;
3939

4040
if constexpr (std::is_same<Tdata, fp16_t>::value || std::is_same<Tdata, bf16_t>::value) {
4141
const float alpha_f = utils::cast<float>(alpha[0]);
@@ -72,7 +72,7 @@ infiniStatus_t Descriptor::calculate(
7272
(void)workspace_size;
7373
(void)stream;
7474

75-
switch (_info.getDtype()) {
75+
switch (_info.data_type) {
7676
case INFINI_DTYPE_F16:
7777
return CALCULATE_AXPY(fp16_t);
7878
case INFINI_DTYPE_F32:

0 commit comments

Comments
 (0)