fix column para_batch bug

Cambricon · Jul 10, 2024 · ada7ea1 · ada7ea1
1 parent 4fb5ed6
commit ada7ea1
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/kernels/fft/fft.cpp b/kernels/fft/fft.cpp
@@ -1049,12 +1049,26 @@ mluOpStatus_t MLUOP_WIN_API fftTwoStepFactor(mluOpFFTPlan_t fft_plan,
       }
     } else {
       // column major
-      for (int cur_r = 64; cur_r > 1; cur_r--) {
+      switch (_n)
+      {
+      case 2048:
+        if (n % 16 == 0) {
+          r = 16;
+        } else if ((n % 8) == 0) {
+          r = 8;
+        }
+        break;
+
+      default:
+        for (int cur_r = 64; cur_r > 1; cur_r--) {
         if (n % cur_r == 0) {
           r = cur_r;
           break;
         }
       }
+        break;
+      }
+
     }
     n /= r;
     switch (factor_type) {

diff --git a/kernels/fft/fft_optm_device/fft_c2c_stockham_gdram.h b/kernels/fft/fft_optm_device/fft_c2c_stockham_gdram.h
@@ -325,7 +325,8 @@ __mlu_func__ void computeMutiStageOnchipColumn(DT *input, DT *output,
       int tw_offset = factors[small_factors_offset + 1];
       int small_twiddles_size = factors[small_factors_offset + 2];
       const DT *small_twiddles = _twiddles + tw_offset * 2;
-      max_para_batch = (6144 / radix) > batch ? batch : (6144 / radix);
+      max_para_batch = factors[small_factors_offset + 3] > (t_end - t_start)?
+       batch : (t_end - t_start);
       for (int t = t_start; t < t_end; t += max_para_batch) {
         int para_batch =
             (max_para_batch < (t_end - t)) ? max_para_batch : (t_end - t);