Skip to content

Commit 2aa007d

Browse files
DanieeelLiuDanieeelLiu
authored
[Feature](mluOpExecFFT): improve c2r 1d perf. (Cambricon#1286)
1 parent c2226e7 commit 2aa007d

2 files changed

Lines changed: 53 additions & 134 deletions

File tree

‎kernels/fft/common/fft_common_kernels.mlu‎

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -321,6 +321,7 @@ __mlu_func__ void generateIRFFTHalfDFTMatrixImpl(int n, void *output) {
321321
int32_t *offset_int_addr = (int32_t *)offset_addr;
322322
float *temp_addr = offset_addr + deal_size;
323323
float *row_addr = temp_addr;
324+
float *cos_sin = row_addr + deal_size;
324325

325326
// generate 0 to n indices
326327
__mluop_gen_stage_index(inc_addr, deal_size, 0.0f, 1.0f);
@@ -333,6 +334,8 @@ __mlu_func__ void generateIRFFTHalfDFTMatrixImpl(int n, void *output) {
333334
cos_coeff_addr[(n + 1) / 2] = 1.0;
334335
sin_coeff_addr[(n + 1) / 2] = -1.0;
335336

337+
// sin * -1 只需要实部结果即可，实部相减改为累加中减法部分用负数累加
338+
__bang_mul_scalar(sin_coeff_addr, sin_coeff_addr, (float)-1, deal_size);
336339
// generate sin and cos vectors
337340
const float scale = 2.0 * M_PI / n;
338341
__memcpy(offset_addr, inc_addr, deal_size * sizeof(float), NRAM2NRAM);
@@ -344,27 +347,35 @@ __mlu_func__ void generateIRFFTHalfDFTMatrixImpl(int n, void *output) {
344347
// generate offsets
345348
__memcpy(offset_addr, inc_addr, pad_col * sizeof(float), NRAM2NRAM);
346349
__bang_mul_scalar(offset_addr, offset_addr, (float)row_i, pad_col);
347-
__mluop_mod(offset_addr, temp_addr, (float)n, pad_col);
350+
__bang_rem(offset_addr, offset_addr, (float)n, pad_col);
348351

349352
genSelectOffsetVec(offset_addr, offset_int_addr, pad_col);
350353

351354
// select cos result
352-
selectVec(cos_addr, offset_int_addr, row_addr, col);
355+
__gather(row_addr, cos_addr, (uint32_t *)offset_int_addr, sizeof(float),
356+
NRAM2NRAM, sizeof(float), col);
353357
__bang_mul(row_addr, row_addr, cos_coeff_addr, pad_col);
354-
convert((DT *)row_addr, row_addr, pad_col);
355-
356-
// save cos result
357-
DT *dst_addr = (DT *)output + row_i * col;
358-
__memcpy(dst_addr, row_addr, col * sizeof(DT), NRAM2GDRAM);
358+
__bang_move(cos_sin, row_addr, sizeof(float), 2 * sizeof(float),
359+
sizeof(float), pad_col);
359360

360361
// select sin result
361-
selectVec(sin_addr, offset_int_addr, row_addr, col);
362+
__gather(row_addr, sin_addr, (uint32_t *)offset_int_addr, sizeof(float),
363+
NRAM2NRAM, sizeof(float), col);
362364
__bang_mul(row_addr, row_addr, sin_coeff_addr, pad_col);
363-
convert((DT *)row_addr, row_addr, pad_col);
364365

365-
// save sin result
366-
dst_addr = (DT *)output + (row_i + row) * col;
367-
__memcpy(dst_addr, row_addr, col * sizeof(DT), NRAM2GDRAM);
366+
__bang_move(cos_sin + 1, row_addr, sizeof(float), 2 * sizeof(float),
367+
sizeof(float), pad_col);
368+
convert((DT *)cos_sin, cos_sin, 2 * pad_col);
369+
370+
// irfft 结果需共轭反转
371+
int row_final = 0;
372+
if (row_i == 0) {
373+
row_final = 0;
374+
} else {
375+
row_final = row - row_i;
376+
}
377+
DT *dst_addr = (DT *)output + row_final * col * 2;
378+
__memcpy(dst_addr, cos_sin, 2 * col * sizeof(DT), NRAM2GDRAM);
368379
}
369380
}
370381

‎kernels/fft/irfft/irfft_host.cpp‎

Lines changed: 30 additions & 122 deletions
Original file line numberDiff line numberDiff line change
@@ -74,18 +74,17 @@ mluOpStatus_t makeIRFFT1dPolicy(mluOpHandle_t handle, mluOpFFTPlan_t fft_plan) {
7474
return MLUOP_STATUS_NOT_SUPPORTED;
7575
}
7676

77-
// Matmul Input : 2 * [batch, (n / 2 + 1)]
78-
// Matmul Matrix : 2 * [n, (n / 2 + 1)]
79-
// Matmul Result : 2 * [batch, n]
80-
int dft_mat_times = COMPLEX;
77+
// Matmul Input : [batch, (n / 2 + 1) * 2]
78+
// Matmul Matrix : [n, (n / 2 + 1) * 2]
79+
// Matmul Result : [batch, n * 2]
8180
int dim0 = n;
8281
int dim1 = FFT_HALF(n);
83-
int dft_mat_num = dft_mat_times * dim0 * dim1;
82+
int dft_mat_num = dim0 * dim1;
8483

8584
// reservespace size allocation
8685
fft_plan->reservespace_size = 0;
8786
fft_plan->reservespace_size +=
88-
dft_mat_num * mluOpDataTypeBytes(in_r_dtype);
87+
dft_mat_num * mluOpDataTypeBytes(in_c_dtype);
8988

9089
/* CNFFT_FUNC_MATMUL :
9190
-------------------------
@@ -104,22 +103,13 @@ mluOpStatus_t makeIRFFT1dPolicy(mluOpHandle_t handle, mluOpFFTPlan_t fft_plan) {
104103
| input_pad |
105104
-------------------------
106105
|
107-
| input trans: batch * (n / 2 + 1) * 2 --> 2 * batch * (n /
108-
2 + 1)
109-
\|/
110-
-------------------------
111-
| input_re |
112-
| input_im |
113-
-------------------------
114106
|
115-
| matmul
116107
\|/
117108
-------------------------
118-
| matmul_re_mul_re |
119-
| matmul_im_mul_im |(reuse output_contiguous)
109+
| input |
120110
-------------------------
121111
|
122-
| op_tensor
112+
| matmul
123113
\|/
124114
-------------------------
125115
| output_contiguous |
@@ -146,42 +136,14 @@ mluOpStatus_t makeIRFFT1dPolicy(mluOpHandle_t handle, mluOpFFTPlan_t fft_plan) {
146136
size_t padded_input_size = in_c_dtype_size * padded_input_num;
147137
fft_plan->workspace_size += need_pad ? padded_input_size : 0;
148138

149-
// input trans and workspace
150-
size_t transed_input_size = padded_input_size;
151-
fft_plan->workspace_size += transed_input_size;
152-
// input trans workspace: batch * (n / 2 + 1) * 2 --> 2 * batch * (n / 2 +
153-
// 1)
154-
const int trans_dim_num = 2;
155-
int64_t trans_input_dims[trans_dim_num] = {padded_input_num, COMPLEX};
156-
int trans_permute[trans_dim_num] = {1, 0};
157-
size_t trans_workspace_size = 0;
158-
status = fftGetTransposeWorkspaceSize(handle, trans_workspace_size,
159-
trans_dim_num, trans_input_dims,
160-
trans_permute, in_r_dtype, api);
161-
fft_plan->matmul_addrs.internal_workspace_size = std::max(
162-
fft_plan->matmul_addrs.internal_workspace_size, trans_workspace_size);
163-
164-
// matmul output(reuse output_coniguous)
165-
int matmul_times = COMPLEX;
166-
int per_matmul_output_num = batch * n;
167-
size_t per_matmul_output_size = in_r_dtype_size * per_matmul_output_num;
168-
fft_plan->workspace_size += (matmul_times - 1) * per_matmul_output_size;
169139
// matmul workspace
170140
size_t matmul_workspace_size = 0;
171-
status = fftGetMatMulWorkspaceSize(handle, matmul_workspace_size, batch,
172-
dim1, dim0, false, true, in_e_dtype,
173-
in_e_dtype, in_r_dtype, api);
141+
status = fftGetMatMulWorkspaceSize(
142+
handle, matmul_workspace_size, batch, dim1 * 2, dim0, false, true,
143+
in_e_dtype, in_e_dtype, in_r_dtype, api);
174144
fft_plan->matmul_addrs.internal_workspace_size =
175145
std::max(fft_plan->matmul_addrs.internal_workspace_size,
176146
matmul_workspace_size);
177-
// optensor workspace
178-
size_t optensor_workspace_size = 0;
179-
status =
180-
fftGetOptensorWorkspaceSize(handle, optensor_workspace_size,
181-
per_matmul_output_num, in_r_dtype, api);
182-
fft_plan->matmul_addrs.internal_workspace_size =
183-
std::max(fft_plan->matmul_addrs.internal_workspace_size,
184-
optensor_workspace_size);
185147

186148
// output contiguous
187149
size_t output_size =
@@ -413,15 +375,12 @@ static void configureIRFFT1dMatmulReserveAddrs(mluOpHandle_t handle,
413375

414376
switch (fft_plan->fft_strategy) {
415377
case CNFFT_FUNC_MATMUL: {
416-
// Matmul Matrix : 2 * [n, (n / 2 + 1)]
378+
// Matmul Matrix : [n, (n / 2 + 1)*2]
417379
int dim0 = n;
418380
int dim1 = FFT_HALF(n);
419-
size_t per_dft_mat_size = dim0 * dim1 * in_r_dtype_size;
381+
size_t per_dft_mat_size = dim0 * dim1 * mluOpDataTypeBytes(in_c_dtype);
420382
dft_mat_size = dft_mat_times * per_dft_mat_size;
421383
fft_plan->matmul_addrs.dft_matrix_addr = fft_plan->reservespace_addr;
422-
fft_plan->matmul_addrs.dft_re_matrix_addr = fft_plan->reservespace_addr;
423-
fft_plan->matmul_addrs.dft_im_matrix_addr =
424-
(uint8_t *)fft_plan->reservespace_addr + per_dft_mat_size;
425384
}; break;
426385
case CNFFT_FUNC_COOLEY_TUKEY:
427386
case CNFFT_FUNC_STOCKHAM: {
@@ -465,10 +424,10 @@ mluOpStatus_t setIRFFT1dReserveArea(mluOpHandle_t handle,
465424

466425
switch (fft_plan->fft_strategy) {
467426
case CNFFT_FUNC_MATMUL: {
468-
// Matmul Matrix : 2 * [n, (n / 2 + 1)]
427+
// Matmul Matrix : [n, (n / 2 + 1) * 2]
469428
int dim0 = n;
470-
int dim1 = (n / 2 + 1);
471-
int dft_mat_num = dft_mat_times * dim0 * dim1;
429+
int dim1 = (n / 2 + 1) * dft_mat_times;
430+
int dft_mat_num = dim0 * dim1;
472431
kernelGenerateIRFFTHalfDFTMatrix(k_dim, k_type, handle->queue, fft_plan,
473432
in_r_dtype, n);
474433
}; break;
@@ -566,17 +525,8 @@ static void configureIRFFT1dMatmulWorkspaceAddrs(mluOpHandle_t handle,
566525
fft_plan->matmul_addrs.input_contiguous_addr;
567526
}
568527

569-
if (fft_plan->fft_strategy == CNFFT_FUNC_MATMUL) {
570-
// input trans: batch * (n / 2 + 1) * 2 --> 2 * batch * (n / 2 + 1)
571-
size_t transed_input_size = padded_input_size;
572-
fft_plan->matmul_addrs.input_re_addr =
573-
(uint8_t *)workspace + workspace_cur_offset;
574-
fft_plan->matmul_addrs.input_im_addr =
575-
(uint8_t *)fft_plan->matmul_addrs.input_re_addr +
576-
transed_input_size / COMPLEX;
577-
workspace_cur_offset += transed_input_size;
578-
} else if (fft_plan->fft_strategy == CNFFT_FUNC_COOLEY_TUKEY ||
579-
fft_plan->fft_strategy == CNFFT_FUNC_STOCKHAM) {
528+
if (fft_plan->fft_strategy == CNFFT_FUNC_COOLEY_TUKEY ||
529+
fft_plan->fft_strategy == CNFFT_FUNC_STOCKHAM) {
580530
// input merge (transed_input and reversed_input reuse input_re)
581531
// 1st input trans: batch * (n / 2 + 1) * 2 --> 2 * batch * (n / 2 + 1)
582532
size_t transed_1st_input_size = padded_input_size;
@@ -625,17 +575,11 @@ static void configureIRFFT1dMatmulWorkspaceAddrs(mluOpHandle_t handle,
625575
}
626576

627577
// matmul output
578+
size_t per_matmul_output_size = 0;
628579
int per_matmul_output_num = batch * n;
629-
size_t per_matmul_output_size = in_r_dtype_size * per_matmul_output_num;
630-
if (fft_plan->fft_strategy == CNFFT_FUNC_MATMUL) {
631-
// matmut_im_mul_im reuse output_coniguous
632-
fft_plan->matmul_addrs.matmul_im_mul_im_addr =
633-
fft_plan->matmul_addrs.output_contiguous_addr;
634-
workspace_cur_offset_to_end += per_matmul_output_size;
635-
fft_plan->matmul_addrs.matmul_re_mul_re_addr =
636-
(uint8_t *)workspace_end - workspace_cur_offset_to_end;
637-
} else if (fft_plan->fft_strategy == CNFFT_FUNC_COOLEY_TUKEY ||
580+
if (fft_plan->fft_strategy == CNFFT_FUNC_COOLEY_TUKEY ||
638581
fft_plan->fft_strategy == CNFFT_FUNC_STOCKHAM) {
582+
per_matmul_output_size = in_r_dtype_size * per_matmul_output_num;
639583
workspace_cur_offset_to_end += per_matmul_output_size;
640584
fft_plan->matmul_addrs.matmul_im_mul_im_addr =
641585
(uint8_t *)workspace_end - workspace_cur_offset_to_end;
@@ -1073,23 +1017,7 @@ static mluOpStatus_t transposeIRFFT1dPaddedInput(mluOpHandle_t handle,
10731017
int batch = fft_plan->batch;
10741018
int n = fft_plan->n[0];
10751019

1076-
if (fft_plan->fft_strategy == CNFFT_FUNC_MATMUL) {
1077-
// transpose: batch * (n / 2 + 1) * 2 --> 2 * batch * (n / 2 + 1)
1078-
VLOG(5) << "launch mluOpTranspose for input MATMUL";
1079-
int padded_input_num = batch * FFT_HALF(n);
1080-
const int trans_dim_num = 2;
1081-
int64_t trans_input_dims[trans_dim_num] = {padded_input_num, COMPLEX};
1082-
int64_t trans_output_dims[trans_dim_num] = {COMPLEX, padded_input_num};
1083-
int trans_permute[trans_dim_num] = {1, 0};
1084-
1085-
status =
1086-
fftTranspose(handle, trans_dim_num, trans_input_dims, trans_output_dims,
1087-
trans_permute, fft_plan->matmul_addrs.input_pad_addr,
1088-
fft_plan->matmul_addrs.input_re_addr, in_r_dtype,
1089-
fft_plan->matmul_addrs.internal_workspace_addr,
1090-
fft_plan->matmul_addrs.internal_workspace_size, api);
1091-
CHECK_RETURN(api, status);
1092-
} else if (fft_plan->fft_strategy == CNFFT_FUNC_COOLEY_TUKEY) {
1020+
if (fft_plan->fft_strategy == CNFFT_FUNC_COOLEY_TUKEY) {
10931021
VLOG(5) << "launch mluOpTranspose for input COOLEY_TUKEY";
10941022
int L = fft_plan->L;
10951023
int m = (1 << fft_plan->m);
@@ -1162,35 +1090,15 @@ static mluOpStatus_t computeIRFFT1dMatmulResult(mluOpHandle_t handle,
11621090

11631091
if (fft_plan->fft_strategy == CNFFT_FUNC_MATMUL) {
11641092
VLOG(5) << "into computeIRFFT1dMatmulResult CNFFT_FUNC_MATMUL";
1165-
// input real matmul dft real
1166-
status = fftMatMul(
1167-
handle, batch, FFT_HALF(n), n, fft_plan->matmul_addrs.input_re_addr,
1168-
fft_plan->matmul_addrs.dft_re_matrix_addr,
1169-
fft_plan->matmul_addrs.matmul_re_mul_re_addr, false, true, scale_factor,
1170-
0.0, in_e_dtype, in_e_dtype, in_r_dtype,
1171-
fft_plan->matmul_addrs.internal_workspace_addr,
1172-
fft_plan->matmul_addrs.internal_workspace_size, api);
1173-
CHECK_RETURN(api, status);
1174-
1175-
// input imag matmul dft imag
1176-
status = fftMatMul(
1177-
handle, batch, FFT_HALF(n), n, fft_plan->matmul_addrs.input_im_addr,
1178-
fft_plan->matmul_addrs.dft_im_matrix_addr,
1179-
fft_plan->matmul_addrs.matmul_im_mul_im_addr, false, true, scale_factor,
1180-
0.0, in_e_dtype, in_e_dtype, in_r_dtype,
1181-
fft_plan->matmul_addrs.internal_workspace_addr,
1182-
fft_plan->matmul_addrs.internal_workspace_size, api);
1183-
CHECK_RETURN(api, status);
1184-
1185-
// real mul real add imag mul imag
1186-
int per_matmul_output_num = batch * n;
1187-
status = fftOptensor(handle, per_matmul_output_num,
1188-
fft_plan->matmul_addrs.matmul_re_mul_re_addr,
1189-
fft_plan->matmul_addrs.matmul_im_mul_im_addr,
1190-
fft_plan->matmul_addrs.output_contiguous_addr, 1.0,
1191-
1.0, 0.0, in_r_dtype, CNNL_OP_TENSOR_ADD,
1192-
fft_plan->matmul_addrs.internal_workspace_addr,
1193-
fft_plan->matmul_addrs.internal_workspace_size, api);
1093+
// input matmul dft
1094+
status =
1095+
fftMatMul(handle, batch, 2 * FFT_HALF(n), n,
1096+
fft_plan->matmul_addrs.input_pad_addr,
1097+
fft_plan->matmul_addrs.dft_matrix_addr,
1098+
fft_plan->matmul_addrs.output_contiguous_addr, false, true,
1099+
scale_factor, 0.0, in_e_dtype, in_e_dtype, in_r_dtype,
1100+
fft_plan->matmul_addrs.internal_workspace_addr,
1101+
fft_plan->matmul_addrs.internal_workspace_size, api);
11941102
CHECK_RETURN(api, status);
11951103
} else if (fft_plan->fft_strategy == CNFFT_FUNC_COOLEY_TUKEY) {
11961104
VLOG(5) << "into computeIRFFT1dMatmulResult CNFFT_FUNC_COOLEY_TUKEY";

0 commit comments

Comments
 (0)