View Javadoc
1   // ******************************************************************************
2   //
3   // Title:       Force Field X.
4   // Description: Force Field X - Software for Molecular Biophysics.
5   // Copyright:   Copyright (c) Michael J. Schnieders 2001-2025.
6   //
7   // This file is part of Force Field X.
8   //
9   // Force Field X is free software; you can redistribute it and/or modify it
10  // under the terms of the GNU General Public License version 3 as published by
11  // the Free Software Foundation.
12  //
13  // Force Field X is distributed in the hope that it will be useful, but WITHOUT
14  // ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
15  // FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
16  // details.
17  //
18  // You should have received a copy of the GNU General Public License along with
19  // Force Field X; if not, write to the Free Software Foundation, Inc., 59 Temple
20  // Place, Suite 330, Boston, MA 02111-1307 USA
21  //
22  // Linking this library statically or dynamically with other modules is making a
23  // combined work based on this library. Thus, the terms and conditions of the
24  // GNU General Public License cover the whole combination.
25  //
26  // As a special exception, the copyright holders of this library give you
27  // permission to link this library with independent modules to produce an
28  // executable, regardless of the license terms of these independent modules, and
29  // to copy and distribute the resulting executable under terms of your choice,
30  // provided that you also meet, for each linked independent module, the terms
31  // and conditions of the license of that module. An independent module is a
32  // module which is not derived from or based on this library. If you modify this
33  // library, you may extend this exception to your version of the library, but
34  // you are not obligated to do so. If you do not wish to do so, delete this
35  // exception statement from your version.
36  //
37  // ******************************************************************************
38  package ffx.numerics.fft;
39  
40  import jdk.incubator.vector.DoubleVector;
41  import jdk.incubator.vector.VectorShuffle;
42  import jdk.incubator.vector.VectorSpecies;
43  
44  import static java.lang.Math.fma;
45  
46  /**
47   * Mixed radix factor is extended by the pass classes to apply the mixed radix factor.
48   */
49  public abstract class MixedRadixFactor {
50  
51    private static final double[] negateReal = {-1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0};
52    private static final int[] shuffleMask = {1, 0, 3, 2, 5, 4, 7, 6};
53  
54    /**
55     * The preferred vector species for double precision.
56     */
57    protected static final VectorSpecies<Double> DOUBLE_SPECIES = DoubleVector.SPECIES_PREFERRED;
58    /**
59     * Vector used to change the sign of the imaginary members of the vector via multiplication.
60     */
61    protected static final DoubleVector NEGATE_IM;
62    /**
63     * Vector used to change the sign of the real members of the vector via multiplication.
64     */
65    protected static final DoubleVector NEGATE_RE;
66    /**
67     * Shuffle used to swap real and imaginary members of the vector.
68     */
69    protected static final VectorShuffle<Double> SHUFFLE_RE_IM;
70    /**
71     * The number of contiguous elements that will be read from the input data array.
72     */
73    protected static final int LENGTH = DOUBLE_SPECIES.length();
74    /**
75     * The number of complex elements that will be processed in each inner loop iteration.
76     * The number of elements to process in the inner loop must be evenly divisible by this loop increment.
77     */
78    protected static final int LOOP = LENGTH / 2;
79    /**
80     * The number of complex elements that will be processed in each inner loop iteration.
81     * The number of elements to process in the inner loop must be evenly divisible by this loop increment.
82     */
83    protected static final int BLOCK_LOOP = LENGTH;
84  
85    /**
86     * Vector used to change the sign of the imaginary members of the vector via multiplication.
87     */
88    protected static final DoubleVector NEGATE_IM_128;
89    /**
90     * Vector used to change the sign of the real members of the vector via multiplication.
91     */
92    protected static final DoubleVector NEGATE_RE_128;
93    /**
94     * Shuffle used to swap real and imaginary members of the vector.
95     */
96    protected static final VectorShuffle<Double> SHUFFLE_RE_IM_128;
97    /**
98     * The number of contiguous elements that will be read from the input data array.
99     */
100   protected static final int LENGTH_128 = DoubleVector.SPECIES_128.length();
101   /**
102    * The number of complex elements that will be processed in each inner loop iteration.
103    * The number of elements to process in the inner loop must be evenly divisible by this loop increment.
104    */
105   protected static final int LOOP_128 = LENGTH_128 / 2;
106   /**
107    * The number of complex elements that will be processed in each inner loop iteration.
108    * The number of elements to process in the inner loop must be evenly divisible by this loop increment.
109    */
110   protected static final int BLOCK_LOOP_128 = LENGTH_128;
111 
112   /**
113    * Vector used to change the sign of the imaginary members of the vector via multiplication.
114    */
115   protected static final DoubleVector NEGATE_IM_256;
116   /**
117    * Vector used to change the sign of the real members of the vector via multiplication.
118    */
119   protected static final DoubleVector NEGATE_RE_256;
120   /**
121    * Shuffle used to swap real and imaginary members of the vector.
122    */
123   protected static final VectorShuffle<Double> SHUFFLE_RE_IM_256;
124   /**
125    * The number of contiguous elements that will be read from the input data array.
126    */
127   protected static final int LENGTH_256 = DoubleVector.SPECIES_256.length();
128   /**
129    * The number of complex elements that will be processed in each inner loop iteration.
130    * The number of elements to process in the inner loop must be evenly divisible by this loop increment.
131    */
132   protected static final int LOOP_256 = LENGTH_256 / 2;
133   /**
134    * The number of complex elements that will be processed in each inner loop iteration.
135    * The number of elements to process in the inner loop must be evenly divisible by this loop increment.
136    */
137   protected static final int BLOCK_LOOP_256 = LENGTH_256;
138 
139   /**
140    * Vector used to change the sign of the imaginary members of the vector via multiplication.
141    */
142   protected static final DoubleVector NEGATE_IM_512;
143   /**
144    * Vector used to change the sign of the real members of the vector via multiplication.
145    */
146   protected static final DoubleVector NEGATE_RE_512;
147   /**
148    * Shuffle used to swap real and imaginary members of the vector.
149    */
150   protected static final VectorShuffle<Double> SHUFFLE_RE_IM_512;
151   /**
152    * The number of contiguous elements that will be read from the input data array.
153    */
154   protected static final int LENGTH_512 = DoubleVector.SPECIES_512.length();
155   /**
156    * The number of complex elements that will be processed in each inner loop iteration.
157    * The number of elements to process in the inner loop must be evenly divisible by this loop increment.
158    */
159   protected static final int LOOP_512 = LENGTH_512 / 2;
160   /**
161    * The number of complex elements that will be processed in each inner loop iteration.
162    * The number of elements to process in the inner loop must be evenly divisible by this loop increment.
163    */
164   protected static final int BLOCK_LOOP_512 = LENGTH_512;
165 
166   static {
167     // Assume that 512 is the largest vector size.
168     if (LENGTH > 8) {
169       throw new IllegalStateException("Unsupported SIMD vector size: " + LENGTH);
170     }
171 
172     NEGATE_RE_128 = DoubleVector.fromArray(DoubleVector.SPECIES_128, negateReal, 0);
173     NEGATE_IM_128 = NEGATE_RE_128.mul(-1.0);
174     SHUFFLE_RE_IM_128 = VectorShuffle.fromArray(DoubleVector.SPECIES_128, shuffleMask, 0);
175 
176     NEGATE_RE_256 = DoubleVector.fromArray(DoubleVector.SPECIES_256, negateReal, 0);
177     NEGATE_IM_256 = NEGATE_RE_256.mul(-1.0);
178     SHUFFLE_RE_IM_256 = VectorShuffle.fromArray(DoubleVector.SPECIES_256, shuffleMask, 0);
179 
180     NEGATE_RE_512 = DoubleVector.fromArray(DoubleVector.SPECIES_512, negateReal, 0);
181     NEGATE_IM_512 = NEGATE_RE_512.mul(-1.0);
182     SHUFFLE_RE_IM_512 = VectorShuffle.fromArray(DoubleVector.SPECIES_512, shuffleMask, 0);
183 
184     switch (LENGTH) {
185       case 2:
186         NEGATE_RE = NEGATE_RE_128;
187         NEGATE_IM = NEGATE_IM_128;
188         SHUFFLE_RE_IM = SHUFFLE_RE_IM_128;
189         break;
190       case 4:
191         NEGATE_RE = NEGATE_RE_256;
192         NEGATE_IM = NEGATE_IM_256;
193         SHUFFLE_RE_IM = SHUFFLE_RE_IM_256;
194         break;
195       case 8:
196         NEGATE_RE = NEGATE_RE_512;
197         NEGATE_IM = NEGATE_IM_512;
198         SHUFFLE_RE_IM = SHUFFLE_RE_IM_512;
199         break;
200       default:
201         throw new IllegalStateException("Unsupported SIMD DoubleVector size: " + LENGTH);
202     }
203   }
204 
205   /**
206    * The size of the input.
207    */
208   protected final int n;
209   /**
210    * The number of FFTs to process (default = 1).
211    */
212   protected final int nFFTs;
213   /**
214    * The imaginary offset.
215    */
216   protected final int im;
217   /**
218    * The mixed radix factor.
219    */
220   protected final int factor;
221   /**
222    * The product of all factors applied so far.
223    */
224   protected final int product;
225   /**
226    * The outer loop limit (n / product).
227    */
228   protected final int outerLoopLimit;
229   /**
230    * The inner loop limit (product / factor).
231    */
232   protected final int innerLoopLimit;
233   /**
234    * The next input (n / factor).
235    * This is the separation between the input data for each pass.
236    */
237   protected final int nextInput;
238   /**
239    * Equal to 2 * nextInput for interleaved complex data.
240    * Equal to nextInput for separate real and imaginary arrays.
241    */
242   protected final int di;
243   /**
244    * Equal to 2 * innerLoopLimit for interleaved complex data.
245    * Equal to innerLoopLimit for separate real and imaginary arrays.
246    */
247   protected final int dj;
248   /**
249    * The twiddle factors for this pass.
250    */
251   protected final double[][] twiddles;
252   /**
253    * The increment for input data within the inner loop.
254    * This is equal to 2 for interleaved complex data.
255    * This is equal to 1 for separate real and imaginary arrays.
256    */
257   protected final int ii;
258   /**
259    * Increment for the inner loop.
260    */
261   protected final int jstep;
262 
263   /**
264    * Constructor for the mixed radix factor.
265    *
266    * @param passConstants the pass constants.
267    */
268   public MixedRadixFactor(PassConstants passConstants) {
269     n = passConstants.n();
270     nFFTs = passConstants.nFFTs();
271     im = passConstants.im();
272     factor = passConstants.factor();
273     product = passConstants.product();
274     twiddles = passConstants.twiddles();
275     outerLoopLimit = n / product;
276     innerLoopLimit = (product / factor) * nFFTs;
277     nextInput = (n / factor) * nFFTs;
278     if (im == 1) {
279       ii = 2;
280       // For interleaved complex data, the di and dj offsets are doubled.
281       di = 2 * nextInput;
282       dj = 2 * innerLoopLimit;
283     } else {
284       ii = 1;
285       // For separate real and imaginary arrays, the di and dj offsets
286       // are the same as the next input and inner loop limit.
287       di = nextInput;
288       dj = innerLoopLimit;
289     }
290     jstep = (factor - 1) * dj;
291   }
292 
293   /**
294    * Return a string representation of the mixed radix factor.
295    *
296    * @return a string representation of the mixed radix factor.
297    */
298   public String toString() {
299     return "MixedRadixFactor {" +
300         "n=" + n +
301         ", nFFTs=" + nFFTs +
302         ", im=" + im +
303         ", factor=" + factor +
304         ", product=" + product +
305         ", outerLoopLimit=" + outerLoopLimit +
306         ", innerLoopLimit=" + innerLoopLimit +
307         ", nextInput=" + nextInput +
308         ", di=" + di +
309         ", dj=" + dj +
310         ", ii=" + ii +
311         ", jstep=" + jstep +
312         '}';
313   }
314 
315   /**
316    * Apply the mixed radix factor using scalar operations.
317    *
318    * @param passData the pass data.
319    */
320   protected abstract void passScalar(PassData passData);
321 
322   /**
323    * Apply the mixed radix factor using SIMD operations.
324    *
325    * @param passData the pass data.
326    */
327   protected abstract void passSIMD(PassData passData);
328 
329   /**
330    * Multiply two complex numbers [x_r, x_i] and [w_r, w_i] and store the result.
331    *
332    * @param x_r the real part of the complex number.
333    * @param x_i the imaginary part of the complex number.
334    * @param w_r the real part of the twiddle factor.
335    * @param w_i the imaginary part of the twiddle factor.
336    * @param ret the array to store the result.
337    * @param re  the real part index in the result array.
338    * @param im  the imaginary part index in the result array.
339    */
340   protected static void multiplyAndStore(double x_r, double x_i, double w_r, double w_i, double[] ret, int re, int im) {
341     ret[re] = fma(w_r, x_r, -w_i * x_i);
342     ret[im] = fma(w_r, x_i, w_i * x_r);
343   }
344 }