Skip to content

Commit f8e9cc2

Browse files
xu-zhengzhongxu-zhengzhong
committed
Use mcBLAS Ex APIs for axpy, blas_dot, nrm2, rot and scal operators
1 parent e64afb9 commit f8e9cc2

11 files changed

Lines changed: 202 additions & 70 deletions

File tree

‎src/infiniop/devices/metax/metax_ht2mc.h‎

Lines changed: 6 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
#ifdef ENABLE_METAX_MC_API
22
#define hpccDataType macaDataType
33
#define HPCC_R_32F MACA_R_32F
4+
#define HPCC_R_64F MACA_R_64F
45
#define HPCC_R_16F MACA_R_16F
56
#define HPCC_R_16BF MACA_R_16BF
67
#define hpcc_bfloat162 maca_bfloat162
@@ -125,24 +126,19 @@
125126
#define hcblasIdamin mcblasIdamin
126127
#define hcblasSasum mcblasSasum
127128
#define hcblasDasum mcblasDasum
128-
#define hcblasSaxpy mcblasSaxpy
129-
#define hcblasDaxpy mcblasDaxpy
129+
#define hcblasAxpyEx mcblasAxpyEx
130130
#define hcblasScopy mcblasScopy
131131
#define hcblasDcopy mcblasDcopy
132-
#define hcblasSdot mcblasSdot
133-
#define hcblasDdot mcblasDdot
134-
#define hcblasSnrm2 mcblasSnrm2
135-
#define hcblasDnrm2 mcblasDnrm2
136-
#define hcblasSrot mcblasSrot
137-
#define hcblasDrot mcblasDrot
132+
#define hcblasDotEx mcblasDotEx
133+
#define hcblasNrm2Ex mcblasNrm2Ex
134+
#define hcblasRotEx mcblasRotEx
138135
#define hcblasSrotg mcblasSrotg
139136
#define hcblasDrotg mcblasDrotg
140137
#define hcblasSrotm mcblasSrotm
141138
#define hcblasDrotm mcblasDrotm
142139
#define hcblasSrotmg mcblasSrotmg
143140
#define hcblasDrotmg mcblasDrotmg
144-
#define hcblasSscal mcblasSscal
145-
#define hcblasDscal mcblasDscal
141+
#define hcblasScalEx mcblasScalEx
146142
#define hcblasSswap mcblasSswap
147143
#define hcblasDswap mcblasDswap
148144
#define HCBLAS_STATUS_SUCCESS MCBLAS_STATUS_SUCCESS

‎src/infiniop/ops/axpy/metax/axpy_metax.cc‎

Lines changed: 47 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,21 +49,58 @@ infiniStatus_t Descriptor::calculate(
4949
const ptrdiff_t incy = _info.getIncy();
5050
const infiniDtype_t data_type = _info.getDtype();
5151

52+
hpccDataType alpha_type, x_type, y_type;
53+
hpccDataType execution_type;
54+
55+
switch (data_type) {
56+
case INFINI_DTYPE_F16:
57+
alpha_type = x_type = y_type = HPCC_R_16F;
58+
execution_type = HPCC_R_32F;
59+
break;
60+
case INFINI_DTYPE_BF16:
61+
alpha_type = x_type = y_type = HPCC_R_16BF;
62+
execution_type = HPCC_R_32F;
63+
break;
64+
case INFINI_DTYPE_F32:
65+
alpha_type = x_type = y_type = HPCC_R_32F;
66+
execution_type = HPCC_R_32F;
67+
break;
68+
case INFINI_DTYPE_F64:
69+
alpha_type = x_type = y_type = HPCC_R_64F;
70+
execution_type = HPCC_R_64F;
71+
break;
72+
default:
73+
return INFINI_STATUS_BAD_TENSOR_DTYPE;
74+
}
75+
5276
CHECK_STATUS(_opaque->internal->useMcblas(
5377
(hcStream_t)stream,
5478
[&](hcblasHandle_t handle) {
5579
CHECK_MCBLAS(hcblasSetPointerMode(handle, HCBLAS_POINTER_MODE_DEVICE));
5680

57-
switch (data_type) {
58-
case INFINI_DTYPE_F32:
59-
CHECK_MCBLAS(hcblasSaxpy(handle, size, (const float *)alpha, (const float *)x, incx, (float *)y, incy));
60-
break;
61-
case INFINI_DTYPE_F64:
62-
CHECK_MCBLAS(hcblasDaxpy(handle, size, (const double *)alpha, (const double *)x, incx, (double *)y, incy));
63-
break;
64-
default:
65-
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
66-
}
81+
// switch (data_type) {
82+
// case INFINI_DTYPE_F32:
83+
// CHECK_MCBLAS(hcblasSaxpy(handle, size, (const float *)alpha, (const float *)x, incx, (float *)y, incy));
84+
// break;
85+
// case INFINI_DTYPE_F64:
86+
// CHECK_MCBLAS(hcblasDaxpy(handle, size, (const double *)alpha, (const double *)x, incx, (double *)y, incy));
87+
// break;
88+
// default:
89+
// return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
90+
// }
91+
92+
hcblasAxpyEx(
93+
handle,
94+
size,
95+
alpha,
96+
alpha_type,
97+
x,
98+
x_type,
99+
incx,
100+
y,
101+
y_type,
102+
incy,
103+
execution_type);
67104

68105
return INFINI_STATUS_SUCCESS;
69106
}));

‎src/infiniop/ops/blas_dot/metax/blas_dot_metax.cc‎

Lines changed: 36 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -49,21 +49,47 @@ infiniStatus_t Descriptor::calculate(
4949
const ptrdiff_t incy = _info.getIncy();
5050
const infiniDtype_t data_type = _info.getDtype();
5151

52+
hpccDataType x_type, y_type, result_type;
53+
hpccDataType execution_type;
54+
55+
switch (data_type) {
56+
case INFINI_DTYPE_F16:
57+
x_type = y_type = result_type = HPCC_R_16F;
58+
execution_type = HPCC_R_32F;
59+
break;
60+
case INFINI_DTYPE_BF16:
61+
x_type = y_type = result_type = HPCC_R_16BF;
62+
execution_type = HPCC_R_32F;
63+
break;
64+
case INFINI_DTYPE_F32:
65+
x_type = y_type = result_type = HPCC_R_32F;
66+
execution_type = HPCC_R_32F;
67+
break;
68+
case INFINI_DTYPE_F64:
69+
x_type = y_type = result_type = HPCC_R_64F;
70+
execution_type = HPCC_R_64F;
71+
break;
72+
default:
73+
return INFINI_STATUS_BAD_TENSOR_DTYPE;
74+
}
75+
5276
CHECK_STATUS(_opaque->internal->useMcblas(
5377
(hcStream_t)stream,
5478
[&](hcblasHandle_t handle) {
5579
CHECK_MCBLAS(hcblasSetPointerMode(handle, HCBLAS_POINTER_MODE_DEVICE));
5680

57-
switch (data_type) {
58-
case INFINI_DTYPE_F32:
59-
CHECK_MCBLAS(hcblasSdot(handle, size, (const float *)x, incx, (const float *)y, incy, (float *)result));
60-
break;
61-
case INFINI_DTYPE_F64:
62-
CHECK_MCBLAS(hcblasDdot(handle, size, (const double *)x, incx, (const double *)y, incy, (double *)result));
63-
break;
64-
default:
65-
return INFINI_STATUS_BAD_TENSOR_DTYPE;
66-
}
81+
CHECK_MCBLAS(hcblasDotEx(
82+
handle,
83+
size,
84+
x,
85+
x_type,
86+
incx,
87+
y,
88+
y_type,
89+
incy,
90+
result,
91+
result_type,
92+
execution_type));
6793

6894
return INFINI_STATUS_SUCCESS;
6995
}));

‎src/infiniop/ops/nrm2/metax/nrm2_metax.cc‎

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -46,21 +46,44 @@ infiniStatus_t Descriptor::calculate(
4646
const ptrdiff_t incx = _info.getIncx();
4747
const infiniDtype_t data_type = _info.getDtype();
4848

49+
hpccDataType x_type, result_type;
50+
hpccDataType execution_type;
51+
52+
switch (data_type) {
53+
case INFINI_DTYPE_F16:
54+
x_type = result_type = HPCC_R_16F;
55+
execution_type = HPCC_R_32F;
56+
break;
57+
case INFINI_DTYPE_BF16:
58+
x_type = result_type = HPCC_R_16BF;
59+
execution_type = HPCC_R_32F;
60+
break;
61+
case INFINI_DTYPE_F32:
62+
x_type = result_type = HPCC_R_32F;
63+
execution_type = HPCC_R_32F;
64+
break;
65+
case INFINI_DTYPE_F64:
66+
x_type = result_type = HPCC_R_64F;
67+
execution_type = HPCC_R_64F;
68+
break;
69+
default:
70+
return INFINI_STATUS_BAD_TENSOR_DTYPE;
71+
}
72+
4973
CHECK_STATUS(_opaque->internal->useMcblas(
5074
(hcStream_t)stream,
5175
[&](hcblasHandle_t handle) {
5276
CHECK_MCBLAS(hcblasSetPointerMode(handle, HCBLAS_POINTER_MODE_DEVICE));
5377

54-
switch (data_type) {
55-
case INFINI_DTYPE_F32:
56-
CHECK_MCBLAS(hcblasSnrm2(handle, size, (const float *)x, incx, (float *)result));
57-
break;
58-
case INFINI_DTYPE_F64:
59-
CHECK_MCBLAS(hcblasDnrm2(handle, size, (const double *)x, incx, (double *)result));
60-
break;
61-
default:
62-
return INFINI_STATUS_BAD_TENSOR_DTYPE;
63-
}
78+
CHECK_MCBLAS(hcblasNrm2Ex(
79+
handle,
80+
size,
81+
x,
82+
x_type,
83+
incx,
84+
result,
85+
result_type,
86+
execution_type));
6487

6588
return INFINI_STATUS_SUCCESS;
6689
}));

‎src/infiniop/ops/rot/metax/rot_metax.cc‎

Lines changed: 37 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -51,21 +51,48 @@ infiniStatus_t Descriptor::calculate(
5151
const ptrdiff_t incy = _info.getIncy();
5252
const infiniDtype_t data_type = _info.getDtype();
5353

54+
hpccDataType x_type, y_type, cs_type;
55+
hpccDataType execution_type;
56+
57+
switch (data_type) {
58+
case INFINI_DTYPE_F16:
59+
x_type = y_type = cs_type = HPCC_R_16F;
60+
execution_type = HPCC_R_32F;
61+
break;
62+
case INFINI_DTYPE_BF16:
63+
x_type = y_type = cs_type = HPCC_R_16BF;
64+
execution_type = HPCC_R_32F;
65+
break;
66+
case INFINI_DTYPE_F32:
67+
x_type = y_type = cs_type = HPCC_R_32F;
68+
execution_type = HPCC_R_32F;
69+
break;
70+
case INFINI_DTYPE_F64:
71+
x_type = y_type = cs_type = HPCC_R_64F;
72+
execution_type = HPCC_R_64F;
73+
break;
74+
default:
75+
return INFINI_STATUS_BAD_TENSOR_DTYPE;
76+
}
77+
5478
CHECK_STATUS(_opaque->internal->useMcblas(
5579
(hcStream_t)stream,
5680
[&](hcblasHandle_t handle) {
5781
CHECK_MCBLAS(hcblasSetPointerMode(handle, HCBLAS_POINTER_MODE_DEVICE));
5882

59-
switch (data_type) {
60-
case INFINI_DTYPE_F32:
61-
CHECK_MCBLAS(hcblasSrot(handle, size, (float *)x, incx, (float *)y, incy, (const float *)c, (const float *)s));
62-
break;
63-
case INFINI_DTYPE_F64:
64-
CHECK_MCBLAS(hcblasDrot(handle, size, (double *)x, incx, (double *)y, incy, (const double *)c, (const double *)s));
65-
break;
66-
default:
67-
return INFINI_STATUS_BAD_TENSOR_DTYPE;
68-
}
83+
CHECK_MCBLAS(hcblasRotEx(
84+
handle,
85+
size,
86+
x,
87+
x_type,
88+
incx,
89+
y,
90+
y_type,
91+
incy,
92+
c,
93+
s,
94+
cs_type,
95+
execution_type));
6996

7097
return INFINI_STATUS_SUCCESS;
7198
}));

‎src/infiniop/ops/scal/metax/scal_metax.cc‎

Lines changed: 33 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -46,21 +46,44 @@ infiniStatus_t Descriptor::calculate(
4646
const ptrdiff_t incx = _info.getIncx();
4747
const infiniDtype_t data_type = _info.getDtype();
4848

49+
hpccDataType alpha_type, x_type;
50+
hpccDataType execution_type;
51+
52+
switch (data_type) {
53+
case INFINI_DTYPE_F16:
54+
alpha_type = x_type = HPCC_R_16F;
55+
execution_type = HPCC_R_32F;
56+
break;
57+
case INFINI_DTYPE_BF16:
58+
alpha_type = x_type = HPCC_R_16BF;
59+
execution_type = HPCC_R_32F;
60+
break;
61+
case INFINI_DTYPE_F32:
62+
alpha_type = x_type = HPCC_R_32F;
63+
execution_type = HPCC_R_32F;
64+
break;
65+
case INFINI_DTYPE_F64:
66+
alpha_type = x_type = HPCC_R_64F;
67+
execution_type = HPCC_R_64F;
68+
break;
69+
default:
70+
return INFINI_STATUS_BAD_TENSOR_DTYPE;
71+
}
72+
4973
CHECK_STATUS(_opaque->internal->useMcblas(
5074
(hcStream_t)stream,
5175
[&](hcblasHandle_t handle) {
5276
CHECK_MCBLAS(hcblasSetPointerMode(handle, HCBLAS_POINTER_MODE_DEVICE));
5377

54-
switch (data_type) {
55-
case INFINI_DTYPE_F32:
56-
CHECK_MCBLAS(hcblasSscal(handle, size, (const float *)alpha, (float *)x, incx));
57-
break;
58-
case INFINI_DTYPE_F64:
59-
CHECK_MCBLAS(hcblasDscal(handle, size, (const double *)alpha, (double *)x, incx));
60-
break;
61-
default:
62-
return INFINI_STATUS_DEVICE_TYPE_NOT_SUPPORTED;
63-
}
78+
CHECK_MCBLAS(hcblasScalEx(
79+
handle,
80+
size,
81+
alpha,
82+
alpha_type,
83+
x,
84+
x_type,
85+
incx,
86+
execution_type));
6487

6588
return INFINI_STATUS_SUCCESS;
6689
}));

‎test/infiniop/axpy.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,10 @@
3333
]
3434

3535
_TENSOR_DTYPES = [
36-
# InfiniDtype.F16,
36+
InfiniDtype.F16,
3737
InfiniDtype.F32,
3838
# InfiniDtype.F64,
39-
# InfiniDtype.BF16,
39+
InfiniDtype.BF16,
4040
]
4141

4242
_TOLERANCE_MAP = {

‎test/infiniop/blas_dot.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,10 @@
3333
]
3434

3535
_TENSOR_DTYPES = [
36-
# InfiniDtype.F16,
36+
InfiniDtype.F16,
3737
InfiniDtype.F32,
3838
# InfiniDtype.F64,
39-
# InfiniDtype.BF16,
39+
InfiniDtype.BF16,
4040
]
4141

4242
_TOLERANCE_MAP = {

‎test/infiniop/nrm2.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,10 @@
3333
]
3434

3535
_TENSOR_DTYPES = [
36-
# InfiniDtype.F16,
36+
InfiniDtype.F16,
3737
InfiniDtype.F32,
3838
# InfiniDtype.F64,
39-
# InfiniDtype.BF16,
39+
InfiniDtype.BF16,
4040
]
4141

4242
_TOLERANCE_MAP = {

‎test/infiniop/rot.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,10 +29,10 @@
2929
]
3030

3131
_TENSOR_DTYPES = [
32-
# InfiniDtype.F16,
32+
InfiniDtype.F16,
3333
InfiniDtype.F32,
3434
# InfiniDtype.F64,
35-
# InfiniDtype.BF16,
35+
InfiniDtype.BF16,
3636
]
3737

3838
_TOLERANCE_MAP = {

0 commit comments

Comments
 (0)