boxing snippets: assign primitive cache profiles

Author: davleopo

compiler/src/jdk.graal.compiler.test/src/jdk/graal/compiler/core/test/BoxingTest.java

@@ -56,4 +56,58 @@ public void test300() {
     public void testMinus300() {
         test("boxSnippet", -300);
     }
+
+    public static Object boxInt(int arg) {
+        return arg;
+    }
+
+    public static Object boxLong(long arg) {
+        return arg;
+    }
+
+    public static Object boxByte(byte arg) {
+        return arg;
+    }
+
+    public static Object boxChar(char arg) {
+        return arg;
+    }
+
+    public static Object boxBoolean(boolean arg) {
+        return arg;
+    }
+
+    public static Object boxShort(short arg) {
+        return arg;
+    }
+
+    @Test
+    public void boxInt0() {
+        test("boxInt", 0);
+    }
+
+    @Test
+    public void boxByte0() {
+        test("boxByte", (byte) 0);
+    }
+
+    @Test
+    public void boxLong0() {
+        test("boxLong", (long) 0);
+    }
+
+    @Test
+    public void boxChar0() {
+        test("boxChar", (char) 0);
+    }
+
+    @Test
+    public void boxBoolean0() {
+        test("boxBoolean", false);
+    }
+
+    @Test
+    public void boxShort0() {
+        test("boxShort", (short) 0);
+    }
 }

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/replacements/BoxingSnippets.java

@@ -26,21 +26,33 @@
 
 import java.util.EnumMap;
 
+import org.graalvm.collections.UnmodifiableEconomicMap;
 import org.graalvm.word.LocationIdentity;
 
 import jdk.graal.compiler.api.replacements.Snippet;
 import jdk.graal.compiler.api.replacements.Snippet.ConstantParameter;
+import jdk.graal.compiler.debug.Assertions;
+import jdk.graal.compiler.debug.DebugContext;
 import jdk.graal.compiler.debug.GraalError;
+import jdk.graal.compiler.graph.Node;
+import jdk.graal.compiler.graph.NodeBitMap;
+import jdk.graal.compiler.nodes.AbstractBeginNode;
+import jdk.graal.compiler.nodes.DeoptimizeNode;
 import jdk.graal.compiler.nodes.FieldLocationIdentity;
+import jdk.graal.compiler.nodes.IfNode;
 import jdk.graal.compiler.nodes.NamedLocationIdentity;
 import jdk.graal.compiler.nodes.PiNode;
+import jdk.graal.compiler.nodes.ProfileData;
+import jdk.graal.compiler.nodes.calc.CompareNode;
 import jdk.graal.compiler.nodes.extended.AbstractBoxingNode;
 import jdk.graal.compiler.nodes.extended.BoxNode;
+import jdk.graal.compiler.nodes.extended.BranchProbabilityNode;
 import jdk.graal.compiler.nodes.extended.UnboxNode;
 import jdk.graal.compiler.nodes.spi.CoreProviders;
 import jdk.graal.compiler.nodes.spi.LoweringTool;
 import jdk.graal.compiler.options.OptionValues;
 import jdk.graal.compiler.phases.util.Providers;
+import jdk.vm.ci.meta.DeoptimizationReason;
 import jdk.vm.ci.meta.JavaKind;
 import jdk.vm.ci.meta.ResolvedJavaField;
 
@@ -254,7 +266,9 @@ public void lower(BoxNode box, LoweringTool tool) {
             args.add("valueOfCounter", valueOfCounter);
             SnippetTemplate template = template(tool, box, args);
             box.getDebug().log("Lowering integerValueOf in %s: node=%s, template=%s, arguments=%s", box.graph(), box, template, args);
-            template.instantiate(tool.getMetaAccess(), box, SnippetTemplate.DEFAULT_REPLACER, args);
+
+            UnmodifiableEconomicMap<Node, Node> duplicates = template.instantiate(tool.getMetaAccess(), box, SnippetTemplate.DEFAULT_REPLACER, args);
+            assignPrimitiveCacheProfiles(box.getBoxingKind(), duplicates);
         }
 
         public void lower(UnboxNode unbox, LoweringTool tool) {
@@ -267,4 +281,63 @@ public void lower(UnboxNode unbox, LoweringTool tool) {
             template.instantiate(tool.getMetaAccess(), unbox, SnippetTemplate.DEFAULT_REPLACER, args);
         }
     }
+
+    /**
+     * Boxing primitive values typically utilizes a primitive cache, but profiling information for
+     * this cache is absent within snippets. This function injects the necessary profiles to address
+     * that gap.
+     */
+    private static void assignPrimitiveCacheProfiles(JavaKind kind, UnmodifiableEconomicMap<Node, Node> duplicates) {
+        NodeBitMap controlFlow = null;
+        for (Node originalNode : duplicates.getKeys()) {
+            if (originalNode instanceof IfNode originalIf) {
+                if (isBoundsCheck(originalIf)) {
+                    // Ignore the bounds check of the cache, that is not relevant
+                    continue;
+                }
+                if (controlFlow == null) {
+                    controlFlow = originalIf.graph().createNodeBitMap();
+                }
+                controlFlow.mark(originalIf);
+                ProfileData.ProfileSource source = originalIf.getProfileData().getProfileSource();
+                assert source.isUnknown() || source.isInjected() : Assertions.errorMessage(originalIf, originalIf.getProfileData(), "Profile should be unknown inside the snippet");
+            }
+        }
+        if (controlFlow == null) {
+            return;
+        }
+        GraalError.guarantee(controlFlow.count() == 1, "Must only have a single control flow element - the branch into the cache but found more %s", controlFlow);
+
+        ProfileData.BranchProbabilityData b = null;
+        switch (kind) {
+            case Byte:
+                // cache contains all byte values, should not see any control flow and thus never
+                // enter this branch
+                throw GraalError.shouldNotReachHere("Byte.valueOf should all go to cache and never contain control flow " + controlFlow);
+            case Boolean:
+                // only two cases, truly 50 / 50
+                b = ProfileData.BranchProbabilityData.injected(0.5);
+                break;
+            case Char:
+            case Short:
+            case Int:
+            case Long:
+                b = ProfileData.BranchProbabilityData.injected(BranchProbabilityNode.FREQUENT_PROBABILITY);
+                break;
+            default:
+                throw GraalError.shouldNotReachHere("Unknown control flow in boxing code, did a JDK change trigger this error? Consider adding new logic to set the profile of " + controlFlow);
+        }
+        IfNode cacheIf = (IfNode) controlFlow.first();
+        IfNode inlinedNode = (IfNode) duplicates.get(cacheIf);
+        inlinedNode.setTrueSuccessorProbability(b);
+        inlinedNode.graph().getDebug().dump(DebugContext.VERY_DETAILED_LEVEL, inlinedNode.graph(), "After updating profile of %s", inlinedNode);
+    }
+
+    private static boolean isBoundsCheck(IfNode ifNode) {
+        if (ifNode.condition() instanceof CompareNode) {
+            AbstractBeginNode fs = ifNode.falseSuccessor();
+            return fs.next() instanceof DeoptimizeNode deopt && deopt.getReason() == DeoptimizationReason.BoundsCheckException;
+        }
+        return false;
+    }
 }