1 /*
  2  * Copyright (c) 2026, Oracle and/or its affiliates. All rights reserved.
  3  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  4  *
  5  * This code is free software; you can redistribute it and/or modify it
  6  * under the terms of the GNU General Public License version 2 only, as
  7  * published by the Free Software Foundation.
  8  *
  9  * This code is distributed in the hope that it will be useful, but WITHOUT
 10  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 11  * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 12  * version 2 for more details (a copy is included in the LICENSE file that
 13  * accompanied this code).
 14  *
 15  * You should have received a copy of the GNU General Public License version
 16  * 2 along with this work; if not, write to the Free Software Foundation,
 17  * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 18  *
 19  * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 20  * or visit www.oracle.com if you need additional information or have any
 21  * questions.
 22  */
 23 
 24 /**
 25  * @test
 26  * @bug 8387073
 27  * @key randomness
 28  * @summary Narrower stores preceding masked vector stores must not be eliminated.
 29  * @modules jdk.incubator.vector
 30  * @library /test/lib /
 31  * @run driver ${test.main.class}
 32  */
 33 
 34 package compiler.igvn;
 35 
 36 import java.util.Arrays;
 37 import java.util.ArrayList;
 38 import java.util.Collections;
 39 import java.util.List;
 40 import java.util.Random;
 41 import java.util.Set;
 42 import java.util.stream.*;
 43 
 44 import jdk.incubator.vector.VectorShape;
 45 
 46 import jdk.test.lib.Utils;
 47 import compiler.lib.compile_framework.*;
 48 import compiler.lib.template_framework.*;
 49 import static compiler.lib.template_framework.Template.scope;
 50 import static compiler.lib.template_framework.Template.let;
 51 import compiler.lib.template_framework.library.*;
 52 
 53 public class TestMaskedStoreIdealization {
 54     private static final Random RANDOM = Utils.getRandomInstance();
 55     private static final String PACKAGE = "compiler.igvn.generated";
 56     private static final String CLASS_NAME = "TestMaskedStoreIdealizationGenerated";
 57 
 58     public static void main(String[] args) {
 59         final CompileFramework comp = new CompileFramework();
 60         comp.addJavaSourceCode(PACKAGE + "." + CLASS_NAME, generate(comp));
 61         comp.compile("--add-modules=jdk.incubator.vector");
 62 
 63         List<String> vmArgs = new ArrayList<>(List.of(
 64             "--add-modules=jdk.incubator.vector",
 65             "--add-opens", "jdk.incubator.vector/jdk.incubator.vector=ALL-UNNAMED"
 66         ));
 67         vmArgs.addAll(Arrays.asList(args)); // Forward args
 68         // Temporarily disable stress flag due to unrelated test failures.
 69         // TODO: Remove when JDK-8388490 is fixed.
 70         vmArgs.addAll(List.of("-XX:+IgnoreUnrecognizedVMOptions", "-XX:-StressReflectiveCode"));
 71         String[] vmArgsArray = vmArgs.toArray(new String[0]);
 72 
 73         comp.invoke(PACKAGE + "." + CLASS_NAME, "main", new Object[] { vmArgsArray });
 74     }
 75 
 76     private static String generate(CompileFramework comp) {
 77         final Set<String> imports = Set.of("java.util.Arrays",
 78                                            "java.util.Random",
 79                                            "jdk.incubator.vector.*",
 80                                            "jdk.test.lib.Utils",
 81                                            "compiler.lib.generators.*");
 82 
 83         // The preferred vector shape is the largest possible vector size.
 84         final int maxVecByteSize = VectorShape.preferredShape().vectorBitSize() / 8;
 85 
 86         final List<TemplateToken> tests = new ArrayList<>();
 87         // Add tests only for the vector shapes that
 88         tests.addAll(CodeGenerationDataNameType.VECTOR_VECTOR_TYPES
 89                         .stream()
 90                         .filter(vec -> vec.byteSize() <= maxVecByteSize && vec.elementType instanceof PrimitiveType)
 91                         .map(vec -> new TestPerShape(vec).generate())
 92                         .collect(Collectors.toList()));
 93         tests.add(PrimitiveType.generateLibraryRNG());
 94 
 95         return TestFrameworkClass.render(PACKAGE, CLASS_NAME, imports, comp.getEscapedClassPathOfCompiledClasses(), tests);
 96     }
 97 
 98     enum Operation {
 99         STORE_SCATTER,
100         STORE_MASK,
101         STORE_SCATTER_MASK,
102         STORE_VECTOR_AFTER_SCATTER,
103         RANDOM
104     }
105 
106     record TestPerShape(VectorType.Vector vec) {
107         TemplateToken generate() {
108             final String testName = vec.elementType.boxedTypeName() + vec.length;
109 
110             // Select the index where we set the mask to false. The index is biased to
111             // zero, as the original bug only triggered with the first element.
112             final int idx = RANDOM.nextBoolean() ? RANDOM.nextInt(0, vec.length) : 0;
113 
114             var irVerification = Template.make("op", "arraySize", (Operation op, Integer arraySize) -> {
115                 // No IR-verification for random test cases.
116                 if (op == Operation.RANDOM) {
117                     return scope("");
118                 }
119 
120                 final PrimitiveType pty = (PrimitiveType) vec.elementType;
121                 final String ptyIR = pty.abbrev().equals("S") ? "C" : pty.abbrev();
122 
123                 // Verify that the Scatter store for STORE_VECTOR_AFTER_SCATTER is not eliminated.
124                 var opVerification = Template.make(() -> {
125                     if (op != Operation.STORE_VECTOR_AFTER_SCATTER) {
126                         return scope("");
127                     }
128 
129                     if (vec.length <= 2) {
130                         return scope("    // No Vector nodes are emitted for vectors of length 2 or shorter.\n");
131                     }
132 
133                     if (vec.elementType.byteSize() < 4) {
134                         return scope("    // StoreVectorScatter is not emitted for vectors of subword types.\n");
135                     }
136 
137                     return scope(
138                         """
139                             @IR(counts = {IRNode.STORE_VECTOR_SCATTER, "=1"},
140                                 applyIfCPUFeatureOr = {"avx512", "true", "sve", "true"})
141                         """
142                     );
143                 });
144 
145                 return scope(
146                     let("pty", vec.elementType.name()),
147                     let("ptyIR", ptyIR),
148                     let("idx", idx),
149                     switch (op) {
150                         case STORE_MASK, STORE_SCATTER_MASK ->
151                         // For masked operations, depending on the generated mask and index map C2 does not manage to elide a branch from
152                         // if (mask.allTrue()) {
153                         //     intoArray(a, offset);
154                         // } else {
155                         //     intoArray(a, offset, mask);
156                         // }
157                         // leading to both stores of the diamond being live. This is highly profile dependent and cannot be predicted.
158                         """
159                             @IR(counts = {IRNode.START + "Store#{ptyIR}" + IRNode.MID + "(Memory: @aryptr:#{pty}\\\\[int:#{arraySize}\\\\]).*(:NotNull:exact\\\\[\\\\d+\\\\]).*" + IRNode.END, ">=1",
160                                           IRNode.START + "Store#{ptyIR}" + IRNode.MID + "(Memory: @aryptr:#{pty}\\\\[int:#{arraySize}\\\\]).*(:NotNull:exact\\\\[\\\\d+\\\\]).*" + IRNode.END, "<=2"},
161                                 applyIfCPUFeatureOr = {"avx512", "true", "sve", "true"},
162                                 phase = CompilePhase.BEFORE_MATCHING)
163                         """;
164                         case STORE_SCATTER ->
165                         """
166                             @IR(counts = {IRNode.START + "Store#{ptyIR}" + IRNode.MID + "(Memory: @aryptr:#{pty}\\\\[int:#{arraySize}\\\\]).*(:NotNull:exact\\\\[\\\\d+\\\\]).*" + IRNode.END, "=1"},
167                                 applyIfCPUFeatureOr = {"avx512", "true", "sve", "true"},
168                                 phase = CompilePhase.BEFORE_MATCHING)
169                         """;
170                         case STORE_VECTOR_AFTER_SCATTER -> opVerification.asToken();
171                         case  RANDOM -> "";
172                     }
173                 );
174             });
175 
176             var testBody = Template.make("testCaseRandom", "op", "arraySize", (Random testCaseRandom, Operation op, Integer arraySize) -> {
177                 if (op == Operation.RANDOM) {
178                     return scope(generateRandomTest(testCaseRandom).asToken());
179                 }
180 
181                 var maskGeneration = Template.make(() -> scope(
182                     let("idx", idx),
183                     let("boxedTy", vec.elementType.boxedTypeName()),
184                     let("species", vec.speciesName),
185                     "        VectorMask<#boxedTy> mask = VectorMask.fromLong(#species, ",
186                     testCaseRandom.nextInt(0, 10) == 0 ? testCaseRandom.nextLong() : "-1 - (1 << #idx)",
187                     ");\n"
188                 ));
189 
190                 var indexMapGeneration = Template.make(() -> {
191                     // For the scatter tests, the array is one larger than the number of lanes so we can
192                     // map indices starting at idx to the next index, omitting idx.
193                     int[] indexMap = IntStream.range(0, vec.length)
194                                               .map(i -> i >= idx ? i + 1 : i)
195                                               .toArray();
196                     String indexMapStr = Arrays.toString(indexMap)
197                                                .replace('[', '{')
198                                                .replace(']', '}');
199                     return scope(
200                         let("idx", idx),
201                         let("len", vec.length),
202                         let("idxMap", indexMapStr),
203                         """
204                                 final int[] indexMap = #{idxMap};
205                         """
206                     );
207                 });
208 
209                 var generation = switch (op) {
210                     case STORE_SCATTER, STORE_VECTOR_AFTER_SCATTER -> indexMapGeneration.asToken();
211                     case STORE_MASK                                -> maskGeneration.asToken();
212                     case STORE_SCATTER_MASK                        ->
213                         Template.make(() -> scope(indexMapGeneration.asToken(), maskGeneration.asToken())).asToken();
214                     case RANDOM                                    -> throw new RuntimeException("unreachable");
215                 };
216 
217                 var initStore  = switch (op) {
218                     case STORE_SCATTER, STORE_VECTOR_AFTER_SCATTER -> "        v.intoArray(a, 0, indexMap, 0);\n";
219                     case STORE_MASK                                -> "        v.intoArray(a, 0, mask);\n";
220                     case STORE_SCATTER_MASK                        -> "        v.intoArray(a, 0, indexMap, 0, mask);\n";
221                     case RANDOM                                    -> throw new RuntimeException("unreachable");
222                 };
223 
224                 var keepStore = switch (op) {
225                     case STORE_MASK, STORE_SCATTER, STORE_SCATTER_MASK -> "        a[#idx] = arrVal;\n";
226                     case STORE_VECTOR_AFTER_SCATTER                    -> "        v.intoArray(a, 0);\n";
227                     case RANDOM                                        -> throw new RuntimeException("unreachable");
228                 };
229 
230                 var holeStore  = switch (op) {
231                     case STORE_SCATTER              -> "        v.intoArray(a, 0, indexMap, 0);\n";
232                     case STORE_MASK                 -> "        v.intoArray(a, 0, mask);\n";
233                     case STORE_SCATTER_MASK         -> "        v.intoArray(a, 0, indexMap, 0, mask);\n";
234                     case STORE_VECTOR_AFTER_SCATTER -> "";
235                     case RANDOM                     -> throw new RuntimeException("unreachable");
236                 };
237 
238                 return scope(
239                     let("pty", vec.elementType.name()),
240                     let("vecTy", vec.name()),
241                     let("species", vec.speciesName),
242                     let("idx", idx),
243                     """
244                             #pty[] a = new #pty[#arraySize];
245                     """,
246                     generation,
247                     """
248 
249                             var v = #vecTy.broadcast(#species, broadcastVal);
250                     """,
251                     initStore,
252                     keepStore,
253                     holeStore,
254                     """
255                             return a;
256                     """
257                 );
258             });
259 
260             var testCase = Template.make("op", (Operation op) -> {
261                 // To get the same test body twice, use a new random instance for each generated test body with a seed fixed per test case.
262                 final int testCaseSeed = RANDOM.nextInt();
263 
264                 String testCaseName = testName + switch (op) {
265                     case STORE_SCATTER              -> "Scatter";
266                     case STORE_MASK                 -> "Mask";
267                     case STORE_SCATTER_MASK         -> "ScatterMask";
268                     case STORE_VECTOR_AFTER_SCATTER -> "VectorAfterScatter";
269                     case RANDOM                     -> "Random";
270                 };
271 
272                 // The array size needs to be one larger than the number of lanes for scatter tests, so we can not write to one element.
273                 final int arraySize = switch (op) {
274                     case STORE_SCATTER, STORE_SCATTER_MASK, STORE_VECTOR_AFTER_SCATTER -> vec.length + 1;
275                     case STORE_MASK, RANDOM                                            -> vec.length;
276                 };
277 
278                 return scope(
279                     let("pty", vec.elementType.name()),
280                     let("testCaseName", testCaseName),
281                     let("boxedTy", vec.elementType.boxedTypeName()),
282                     let("vecTy", vec.name()),
283                     let("lanes", vec.length),
284                     let("species", vec.speciesName),
285                     let("idx", idx),
286                     let("rngCall", vec.elementType.callLibraryRNG()),
287                     let("broadcastVal", vec.elementType.con()),
288                     let("arrVal", vec.elementType.con()),
289                 """
290                     @Run(test = "test#{testCaseName}")
291                     @Warmup(10_000)
292                     static void run#{testCaseName}(RunInfo info) {
293                         final #pty broadcastVal = #broadcastVal;
294                         final #pty arrVal = #arrVal;
295                         final #pty[] compiledResult = test#{testCaseName}(broadcastVal, arrVal);
296 
297                         if (!info.isWarmUp()) {
298                             final #pty[] interpreterResult = reference#{testCaseName}(broadcastVal, arrVal);
299                             if (!Arrays.equals(interpreterResult, compiledResult)) {
300                                 throw new RuntimeException("wrong result for test${testCaseName}:\\n" +
301                                                            "  interpreter result: " + Arrays.toString(interpreterResult) + "\\n" +
302                                                            "  compiled result: " + Arrays.toString(compiledResult));
303                             }
304                         }
305                     }
306 
307                     @Test
308                 """,
309                     irVerification.asToken(op, arraySize),
310                 """
311                     static #pty[] test#{testCaseName}(#pty broadcastVal, #pty arrVal) {
312                 """,
313                     testBody.asToken(new Random(testCaseSeed), op, arraySize),
314                 """
315                     }
316 
317                     @DontCompile
318                     static #pty[] reference#{testCaseName}(#pty broadcastVal, #pty arrVal) {
319                 """,
320                     testBody.asToken(new Random(testCaseSeed), op, arraySize),
321                 """
322                     }
323 
324                 """
325                 );
326             });
327 
328             return Template.make(() -> scope(
329                 Stream.of(Operation.class.getEnumConstants())
330                          .map(op -> testCase.asToken(op))
331                          .toList()
332             )).asToken();
333         }
334 
335         Template.ZeroArgs generateRandomTest(Random testCaseRandom) {
336             final int arraySize = testCaseRandom.nextInt(vec.length + 1, 5 * vec.length);
337             var maskHook = new Hook("MaskHook");
338             var genRandomMask = Template.make("maskName", "vecTy", (String maskName, VectorType.Vector vecTy) -> scope(
339                 let("boxedTy", vecTy.elementType.boxedTypeName()),
340                 let("species", vecTy.speciesName),
341                 let("maskVal", testCaseRandom.nextLong()),
342                 """
343                         VectorMask<#{boxedTy}> #maskName = VectorMask.fromLong(#species, #maskVal);
344                 """
345             ));
346             var genRandomIdxMap = Template.make("mapName", "maxIdx", "len", (String mapName, Integer maxIdx, Integer len) -> {
347                 ArrayList<Integer> possibleIndices = new ArrayList(IntStream.range(0, maxIdx).boxed().toList());
348                 Collections.shuffle(possibleIndices, testCaseRandom);
349                 return scope(
350                     let("map", String.join(", ", possibleIndices.stream().limit(vec.length).map(i -> i.toString()).toList())),
351                 """
352                         int[] #mapName = { #map };
353                 """
354                 );
355             });
356             var initialStore = Template.make(() -> {
357                 final int offset = testCaseRandom.nextInt(0, 4) == 0 ? testCaseRandom.nextInt(0, arraySize - vec.length) : 0;
358                 return scope(
359                     let("offset", offset),
360                     switch (testCaseRandom.nextInt(0,4)) {
361                         case 0 -> scope(
362                             """
363                                     v.intoArray(a, #offset);
364                             """
365                             );
366                         case 1 -> scope(
367                             maskHook.insert(genRandomMask.asToken("initMask", vec)),
368                             """
369                                     v.intoArray(a, #offset, initMask);
370                             """
371                             );
372                         case 2 -> scope(
373                             maskHook.insert(genRandomIdxMap.asToken("initMap", arraySize - offset, vec.length)),
374                             """
375                                     v.intoArray(a, #offset, initMap, 0);
376                             """
377                             );
378                         case 3 -> scope(
379                             maskHook.insert(scope(
380                                 genRandomIdxMap.asToken("initMap", arraySize - offset, vec.length),
381                                 genRandomMask.asToken("initMask", vec)
382                             )),
383                             """
384                                     v.intoArray(a, #offset, initMap, 0, initMask);
385                             """
386                             );
387                         default -> throw new RuntimeException("unreachable");
388                     }
389                 );
390             });
391 
392             var storeThatShouldNotBeLost = Template.make(() -> {
393                 final int selection = testCaseRandom.nextInt(0,6);
394                 final VectorType.Vector narrowerVector = CodeGenerationDataNameType.VECTOR_VECTOR_TYPES
395                                                             .stream()
396                                                             .filter(v -> v.elementType == vec.elementType && v.length <= vec.length)
397                                                             // Flip a coin on each reduction step to pick a random element.
398                                                             .reduce(null, (l, r) -> l == null ? r : (testCaseRandom.nextBoolean() ? l : r));
399                 final int offset = testCaseRandom.nextInt(0, 4) == 0 ? testCaseRandom.nextInt(0, arraySize - narrowerVector.length) : 0;
400                 return scope(
401                     let("offset", offset),
402                     switch (selection) {
403                         case 0, 1, 2, 3 -> scope(
404                             let("vecTy", narrowerVector),
405                             let("species", narrowerVector.speciesName),
406                             """
407                                     var vKeep = #vecTy.broadcast(#species, arrVal);
408                             """
409                             );
410                         default -> "";
411                     },
412                     switch (selection) {
413                         case 0 -> scope(
414                             """
415                                     vKeep.intoArray(a, #offset);
416                             """
417                             );
418                         case 1 -> scope(
419                             maskHook.insert(genRandomMask.asToken("keepMask", narrowerVector)),
420                             """
421                                     vKeep.intoArray(a, #offset, keepMask);
422                             """
423                             );
424                         case 2 -> scope(
425                             maskHook.insert(genRandomIdxMap.asToken("keepMap", arraySize - offset, narrowerVector.length)),
426                             """
427                                     vKeep.intoArray(a, #offset, keepMap, 0);
428                             """
429                             );
430                         case 3 -> scope(
431                             maskHook.insert(scope(
432                                 genRandomIdxMap.asToken("keepMap", arraySize - offset, narrowerVector.length),
433                                 genRandomMask.asToken("keepMask", narrowerVector)
434                             )),
435                             """
436                                     vKeep.intoArray(a, #offset, keepMap, 0, keepMask);
437                             """
438                             );
439                         case 4 -> scope(
440                             let("idx", testCaseRandom.nextInt(0, arraySize)),
441                             """
442                                     a[#idx] = arrVal;
443                             """
444                         );
445                         case 5 -> scope(
446                             let("idx", testCaseRandom.nextInt(0, arraySize - vec.length)),
447                             let("len", testCaseRandom.nextInt(1, vec.length + 1)),
448                             """
449                                     Arrays.fill(a, #idx, #idx + #len, arrVal);
450                             """
451                         );
452                         default -> throw new RuntimeException("unreachable");
453                     }
454                 );
455             });
456 
457             var storeWithHole = Template.make(() -> {
458                 final int offset = testCaseRandom.nextInt(0, 4) == 0 ? testCaseRandom.nextInt(0, arraySize - vec.length) : 0;
459                 return scope(
460                     let("offset", offset),
461                     switch (testCaseRandom.nextInt(0,3)) {
462                         case 0 -> scope(
463                             maskHook.insert(genRandomMask.asToken("holeMask", vec)),
464                             """
465                                     v.intoArray(a, #offset, holeMask);
466                             """
467                             );
468                         case 1 -> scope(
469                             maskHook.insert(genRandomIdxMap.asToken("holeMap", arraySize - offset, vec.length)),
470                             """
471                                     v.intoArray(a, #offset, holeMap, 0);
472                             """
473                             );
474                         case 2 -> scope(
475                             maskHook.insert(scope(
476                                 genRandomIdxMap.asToken("holeMap", arraySize - offset, vec.length),
477                                 genRandomMask.asToken("holeMask", vec)
478                             )),
479                             """
480                                     v.intoArray(a, #offset, holeMap, 0, holeMask);
481                             """
482                             );
483                         default -> throw new RuntimeException("unreachable");
484                     }
485                 );
486             });
487 
488             return Template.make(() -> scope(
489                 let("pty", vec.elementType.name()),
490                 let("arraySize", arraySize),
491                 """
492                         #pty[] a = new #pty[#arraySize];
493 
494                 """,
495                 maskHook.anchor(scope(
496                     // The code for generating masks and index maps goes here.
497                     let("vecTy", vec.name()),
498                     let("species", vec.speciesName),
499                 """
500                         var v = #vecTy.broadcast(#species, broadcastVal);
501                 """,
502                     initialStore.asToken(),
503                     storeThatShouldNotBeLost.asToken(),
504                     storeWithHole.asToken()
505                 )),
506                 """
507                         return a;
508                 """
509             ));
510         }
511     }
512 }