[GR-67783] Correctly update vector architecture for SVM runtime compilations, delay Vector API stamp computation

PullRequest: graal/21546

Author: gergo-

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/nodes/GraphState.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2022, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -610,6 +610,7 @@ public enum StageFlag {
         LOOP_OVERFLOWS_CHECKED,
         PARTIAL_ESCAPE,
         FINAL_PARTIAL_ESCAPE,
+        VECTOR_API_EXPANSION,
         HIGH_TIER_LOWERING,
         /* Stages applied by mid tier. */
         FLOATING_READS,

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/vector/architecture/VectorArchitecture.java

@@ -50,6 +50,7 @@
 import jdk.graal.compiler.options.OptionValues;
 import jdk.graal.compiler.vector.nodes.simd.LogicValueStamp;
 import jdk.graal.compiler.vector.nodes.simd.SimdStamp;
+import jdk.graal.compiler.vector.replacements.vectorapi.VectorAPIType;
 import jdk.vm.ci.meta.JavaKind;
 
 /**
@@ -69,6 +70,9 @@ public abstract class VectorArchitecture {
     protected final Stamp oopMaskStamp;
     protected int cachedMaxVectorLength;
 
+    /** The table of Vector API types associated with this vector architecture. */
+    protected VectorAPIType.Table vectorAPITypeTable;
+
     @Override
     public boolean equals(Object o) {
         if (o == null || getClass() != o.getClass()) {
@@ -472,4 +476,13 @@ public Stamp maskStamp(Stamp elementStamp) {
             return LogicValueStamp.UNRESTRICTED;
         }
     }
+
+    public VectorAPIType.Table getVectorAPITypeTable() {
+        return vectorAPITypeTable;
+    }
+
+    public void setVectorAPITypeTable(VectorAPIType.Table table) {
+        GraalError.guarantee(this.vectorAPITypeTable == null, "Vector API type table must only be set once per vector architecture");
+        this.vectorAPITypeTable = table;
+    }
 }

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/vector/architecture/amd64/VectorAMD64.java

@@ -432,9 +432,15 @@ private static int maxVectorSizeForArchitecture(AMD64 arch) {
         return Math.min(maxPhysicalSize, maxDesiredSize);
     }
 
-    public void updateMaxVectorSizeForArchitecture(AMD64 newArch) {
+    /**
+     * To be called only when (re-)configuring the compiler for an SVM runtime compilation. Resets
+     * precomputed values stored in this vector architecture instance for the now known runtime
+     * target architecture.
+     */
+    public void updateForRuntimeArchitecture(AMD64 newArch) {
         this.cachedMaxVectorLength = 0;  // force recomputation
         this.maxVectorSize = maxVectorSizeForArchitecture(newArch);
+        this.vectorAPITypeTable = null;
     }
 
     /**

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/vector/replacements/vectorapi/VectorAPIExpansionPhase.java

@@ -211,10 +211,19 @@ public void checkInvariants() {
 
     @Override
     protected void run(StructuredGraph graph, HighTierContext context) {
+        graph.getGraphState().setDuringStage(GraphState.StageFlag.VECTOR_API_EXPANSION);
         if (!graph.hasNode(VectorAPIMacroNode.TYPE)) {
             return;
         }
 
+        /*
+         * Canonicalize first. Needed for computing SIMD stamps, since we delay their computation to
+         * compile time. We can't generally compute SIMD stamps at the time we build the macro nodes
+         * because the target architecture for SVM runtime compilations is not known at that time.
+         */
+        canonicalizer.applyIncremental(graph, context, graph.getNodes(VectorAPIMacroNode.TYPE));
+        graph.getDebug().dump(DebugContext.DETAILED_LEVEL, graph, "after Vector API macro canonicalization");
+
         VectorArchitecture vectorArch = ((VectorLoweringProvider) context.getLowerer()).getVectorArchitecture();
         /*
          * Nodes visited while discovering the equivalence classes. All nodes that are in one of the

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/vector/replacements/vectorapi/VectorAPIType.java

@@ -26,15 +26,14 @@
 
 import org.graalvm.collections.EconomicMap;
 
-import jdk.graal.compiler.vector.architecture.VectorArchitecture;
-import jdk.graal.compiler.vector.nodes.simd.SimdStamp;
-
 import jdk.graal.compiler.core.common.type.IntegerStamp;
 import jdk.graal.compiler.core.common.type.Stamp;
 import jdk.graal.compiler.core.common.type.StampFactory;
 import jdk.graal.compiler.debug.GraalError;
 import jdk.graal.compiler.nodes.ValueNode;
 import jdk.graal.compiler.nodes.spi.CoreProviders;
+import jdk.graal.compiler.vector.architecture.VectorArchitecture;
+import jdk.graal.compiler.vector.nodes.simd.SimdStamp;
 import jdk.vm.ci.meta.ConstantReflectionProvider;
 import jdk.vm.ci.meta.JavaKind;
 import jdk.vm.ci.meta.ResolvedJavaField;
@@ -112,11 +111,12 @@ public static VectorAPIType ofConstant(ValueNode value, CoreProviders providers)
     public static final String VECTOR_PACKAGE_NAME = "jdk.incubator.vector";
 
     /**
-     * Lookup table for Java Vector API types. An unmodifiable singleton instance of this table is
-     * built lazily on the first lookup request. Users of {@link VectorAPIType} should not interact
-     * with this class directly, only through {@link #ofType}.
+     * Lookup table for Java Vector API types. An unmodifiable instance of this table is built
+     * lazily on the first lookup request for each {@link VectorArchitecture} instance. The vector
+     * architecture is the owner of its associated table. Users of {@link VectorAPIType} should not
+     * interact with this class directly, only through {@link #ofType}.
      */
-    private static final class Table {
+    public static final class Table {
 
         /** The element types for which the Vector API provides vector types. */
         private static final JavaKind[] KINDS = new JavaKind[]{JavaKind.Byte, JavaKind.Short, JavaKind.Int, JavaKind.Long, JavaKind.Float, JavaKind.Double};
@@ -132,17 +132,14 @@ private static final class Table {
         /** The vector architecture for which this table was built. */
         private final VectorArchitecture vectorArch;
 
-        /** Singleton instance lazily initialized via double-checked locking. */
-        private static volatile Table instance;
-
         private static Table instance(VectorArchitecture vectorArch) {
-            Table result = instance;
+            Table result = vectorArch.getVectorAPITypeTable();
             if (result == null) {
-                synchronized (Table.class) {
-                    result = instance;
+                synchronized (vectorArch) {
+                    result = vectorArch.getVectorAPITypeTable();
                     if (result == null) {
                         result = buildLookupTable(vectorArch);
-                        instance = result;
+                        vectorArch.setVectorAPITypeTable(result);
                     }
                 }
             }
@@ -194,18 +191,6 @@ private static void recordEntries(EconomicMap<String, VectorAPIType> table, Java
         }
     }
 
-    /**
-     * Reset the singleton {@link Table} instance to {@code null}. The next lookup in the table will
-     * reinitialize it. This method is for testing only: Native image tests can build multiple
-     * images in the same process with different vector architectures, so we need to forget the type
-     * table between builds.
-     */
-    public static void resetTableForTestingOnly() {
-        synchronized (Table.class) {
-            Table.instance = null;
-        }
-    }
-
     /**
      * Verify that this data structure corresponds to the given {@code javaType}. Specifically, the
      * type must be initialized, and its {@code VLENGTH} and {@code ETYPE} fields must correspond to

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/vector/replacements/vectorapi/nodes/VectorAPIMacroNode.java

@@ -47,9 +47,11 @@
 import jdk.graal.compiler.nodes.ConstantNode;
 import jdk.graal.compiler.nodes.FixedNode;
 import jdk.graal.compiler.nodes.FrameState;
+import jdk.graal.compiler.nodes.GraphState;
 import jdk.graal.compiler.nodes.Invoke;
 import jdk.graal.compiler.nodes.NodeView;
 import jdk.graal.compiler.nodes.PiNode;
+import jdk.graal.compiler.nodes.StructuredGraph;
 import jdk.graal.compiler.nodes.ValueNode;
 import jdk.graal.compiler.nodes.ValuePhiNode;
 import jdk.graal.compiler.nodes.extended.GetClassNode;
@@ -264,6 +266,16 @@ protected static SimdStamp improveVectorStamp(SimdStamp oldVectorStamp, ValueNod
         ValueNode vmClass = arguments[vmClassIndex];
         ValueNode eClass = arguments[elementClassIndex];
         ValueNode length = arguments[lengthIndex];
+        StructuredGraph graph = vmClass.graph();
+        if (graph == null || graph.isBeforeStage(GraphState.StageFlag.VECTOR_API_EXPANSION)) {
+            /*
+             * Don't compute SIMD stamps before we actually start a compilation. The reason we want
+             * to delay until actual compilation time is that the stamp may differ between native
+             * image build time and runtime compilation time since the native image target
+             * architecture may have different CPU features from the runtime target architecture.
+             */
+            return null;
+        }
         return VectorAPIUtils.stampForVectorClass(vmClass, eClass, length, providers);
     }
 

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/CPUFeatureAccess.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, 2019, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2018, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -29,6 +29,7 @@
 import org.graalvm.nativeimage.ImageSingletons;
 
 import jdk.graal.compiler.api.replacements.Fold;
+import jdk.graal.compiler.nodes.spi.LoweringProvider;
 import jdk.vm.ci.code.Architecture;
 
 public interface CPUFeatureAccess {
@@ -43,7 +44,7 @@ static CPUFeatureAccess singleton() {
     @Uninterruptible(reason = "Thread state not set up yet.")
     void verifyHostSupportsArchitectureEarlyOrExit();
 
-    void enableFeatures(Architecture architecture);
+    void enableFeatures(Architecture architecture, LoweringProvider runtimeLowerer);
 
     /**
      * Compute the CPU features enabled at image run time.

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/aarch64/AArch64CPUFeatureAccess.java

@@ -38,6 +38,7 @@
 import com.oracle.svm.core.graal.stackvalue.UnsafeStackValue;
 import com.oracle.svm.core.util.VMError;
 
+import jdk.graal.compiler.nodes.spi.LoweringProvider;
 import jdk.vm.ci.aarch64.AArch64;
 import jdk.vm.ci.code.Architecture;
 
@@ -85,7 +86,7 @@ public void verifyHostSupportsArchitectureEarlyOrExit() {
     }
 
     @Override
-    public void enableFeatures(Architecture runtimeArchitecture) {
+    public void enableFeatures(Architecture runtimeArchitecture, LoweringProvider runtimeLowerer) {
         AArch64 architecture = (AArch64) runtimeArchitecture;
         EnumSet<AArch64.CPUFeature> features = determineHostCPUFeatures();
         architecture.getFeatures().addAll(features);

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/amd64/AMD64CPUFeatureAccess.java

@@ -41,6 +41,9 @@
 import com.oracle.svm.core.jdk.JVMCISubstitutions;
 import com.oracle.svm.core.util.VMError;
 
+import jdk.graal.compiler.nodes.spi.LoweringProvider;
+import jdk.graal.compiler.vector.architecture.VectorLoweringProvider;
+import jdk.graal.compiler.vector.architecture.amd64.VectorAMD64;
 import jdk.vm.ci.amd64.AMD64;
 import jdk.vm.ci.amd64.AMD64Kind;
 import jdk.vm.ci.code.Architecture;
@@ -103,7 +106,7 @@ public static boolean canUpdateCPUFeatures() {
     }
 
     @Override
-    public void enableFeatures(Architecture runtimeArchitecture) {
+    public void enableFeatures(Architecture runtimeArchitecture, LoweringProvider runtimeLowerer) {
         if (!canUpdateCPUFeatures()) {
             return;
         }
@@ -115,5 +118,11 @@ public void enableFeatures(Architecture runtimeArchitecture) {
         // update largest storable kind
         AMD64Kind largestStorableKind = new AMD64(features).getLargestStorableKind(AMD64.XMM);
         JVMCISubstitutions.updateLargestStorableKind(architecture, largestStorableKind);
+        AMD64Kind largestStorableMaskKind = new AMD64(features).getLargestStorableKind(AMD64.MASK);
+        JVMCISubstitutions.updateLargestStorableMaskKind(architecture, largestStorableMaskKind);
+
+        // update vector architecture
+        VectorAMD64 initialVectorArch = (VectorAMD64) ((VectorLoweringProvider) runtimeLowerer).getVectorArchitecture();
+        initialVectorArch.updateForRuntimeArchitecture((AMD64) runtimeArchitecture);
     }
 }

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/JVMCISubstitutions.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019, 2019, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2019, 2025, Oracle and/or its affiliates. All rights reserved.
  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
  *
  * This code is free software; you can redistribute it and/or modify it
@@ -70,7 +70,7 @@ public static Map<String, String> getSavedProperties() {
 @TargetClass(value = AMD64.class)
 final class Target_jdk_vm_ci_amd64_AMD64 {
     @Alias AMD64Kind largestKind;
-
+    @Alias AMD64Kind largestMaskKind;
 }
 
 /** Dummy class to have a class with the file's name. */
@@ -80,4 +80,10 @@ public static void updateLargestStorableKind(AMD64 architecture, AMD64Kind large
         Target_jdk_vm_ci_amd64_AMD64 arch = SubstrateUtil.cast(architecture, Target_jdk_vm_ci_amd64_AMD64.class);
         arch.largestKind = largestStorableKind;
     }
+
+    @Platforms(Platform.AMD64.class)
+    public static void updateLargestStorableMaskKind(AMD64 architecture, AMD64Kind largestStorableMaskKind) {
+        Target_jdk_vm_ci_amd64_AMD64 arch = SubstrateUtil.cast(architecture, Target_jdk_vm_ci_amd64_AMD64.class);
+        arch.largestMaskKind = largestStorableMaskKind;
+    }
 }