[GR-41901] Enforce injected branch probability invariant

PullRequest: graal/12963

Author: gergo-

compiler/src/org.graalvm.compiler.api.directives.test/src/org/graalvm/compiler/api/directives/test/ProbabilityDirectiveShortCircuitTest.java

@@ -0,0 +1,138 @@
+/*
+ * Copyright (c) 2022, 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
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.  Oracle designates this
+ * particular file as subject to the "Classpath" exception as provided
+ * by Oracle in the LICENSE file that accompanied this code.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package org.graalvm.compiler.api.directives.test;
+
+import static org.graalvm.compiler.api.directives.GraalDirectives.injectBranchProbability;
+import static org.graalvm.compiler.api.directives.GraalDirectives.sideEffect;
+
+import org.graalvm.compiler.core.test.GraalCompilerTest;
+import org.graalvm.compiler.debug.GraalError;
+import org.graalvm.compiler.nodes.IfNode;
+import org.graalvm.compiler.nodes.ProfileData.ProfileSource;
+import org.graalvm.compiler.nodes.ShortCircuitOrNode;
+import org.graalvm.compiler.nodes.StructuredGraph;
+import org.graalvm.compiler.phases.common.CanonicalizerPhase;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+public class ProbabilityDirectiveShortCircuitTest extends GraalCompilerTest {
+
+    private void checkProfiles(String snippetName) {
+        StructuredGraph graph = parseForCompile(getResolvedJavaMethod(snippetName));
+        CanonicalizerPhase canonicalizer = createCanonicalizerPhase();
+        createInliningPhase(canonicalizer).apply(graph, getDefaultHighTierContext());
+        canonicalizer.apply(graph, getDefaultHighTierContext());
+        for (IfNode ifNode : graph.getNodes(IfNode.TYPE)) {
+            Assert.assertEquals(ifNode + " profile source", ProfileSource.INJECTED, ifNode.profileSource());
+        }
+        for (ShortCircuitOrNode shortCircuit : graph.getNodes(ShortCircuitOrNode.TYPE)) {
+            Assert.assertEquals(shortCircuit + " profile source", ProfileSource.INJECTED, shortCircuit.getShortCircuitProbability().getProfileSource());
+        }
+    }
+
+    @Test
+    public void andIf() {
+        checkProfiles("andIfSnippet");
+    }
+
+    public static boolean andIfSnippet(boolean a, int b, double c) {
+        if (injectBranchProbability(0.125, a) && injectBranchProbability(0.125, b == 42) && injectBranchProbability(0.125, c > 0.0)) {
+            sideEffect();  // prevent folding to a conditional
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    @Test
+    public void andConditional() {
+        checkProfiles("andConditionalSnippet");
+    }
+
+    public static boolean andConditionalSnippet(boolean a, int b, double c) {
+        // This will fold to a conditional.
+        if (injectBranchProbability(0.125, a) && injectBranchProbability(0.125, b == 42) && injectBranchProbability(0.125, c > 0.0)) {
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    @Test(expected = GraalError.class)
+    public void andReturn() {
+        checkProfiles("andReturnSnippet");
+    }
+
+    @BytecodeParserForceInline
+    public static boolean andReturnSnippet(boolean a, int b, double c) {
+        /*
+         * This builds a graph shape with a BranchProbabilityNode used by a ValuePhi used by a
+         * Return. That is not accepted by the simplification in BranchProbabilityNode that
+         * propagates injected probabilities to the correct usage. Top-level snippets wanting to
+         * inject probabilities this way have to use an explicit if statement (which will then be
+         * folded correctly). However, inlining this pattern into a snippet works fine as long as
+         * the call site also injects a probability.
+         */
+        return injectBranchProbability(0.125, a) && injectBranchProbability(0.125, b == 42) && injectBranchProbability(0.125, c > 0.0);
+    }
+
+    @Test
+    public void andInlined() {
+        checkProfiles("andInlinedSnippet");
+    }
+
+    public static boolean andInlinedSnippet(boolean a, int b, double c) {
+        if (injectBranchProbability(0.25, andReturnSnippet(a, b, c))) {
+            sideEffect();
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    @Test
+    public void orInlined() {
+        checkProfiles("orInlinedSnippet");
+    }
+
+    @BytecodeParserForceInline
+    public static boolean orHelper(Integer i) {
+        if (injectBranchProbability(0.25, i == null) || injectBranchProbability(0.25, i < 42)) {
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    public static boolean orInlinedSnippet(Integer i) {
+        if (injectBranchProbability(0.25, orHelper(i))) {
+            sideEffect();
+            return true;
+        } else {
+            return false;
+        }
+    }
+}

compiler/src/org.graalvm.compiler.api.directives/src/org/graalvm/compiler/api/directives/GraalDirectives.java

@@ -162,7 +162,10 @@ public static long sideEffect(long a) {
 
     /**
      * Injects a probability for the given condition into the profiling information of a branch
-     * instruction. The probability must be a value between 0.0 and 1.0 (inclusive).
+     * instruction. The probability must be a value between 0.0 and 1.0 (inclusive). This directive
+     * should only be used for the condition of an if statement. The parameter condition should also
+     * only denote a simple condition and not a combined condition involving && or ||
+     * operators.
      *
      * Example usage (it specifies that the likelihood for a to be greater than b is 90%):
      *

compiler/src/org.graalvm.compiler.hotspot/src/org/graalvm/compiler/hotspot/replacements/HotSpotAllocationSnippets.java

@@ -210,7 +210,7 @@ private static Class<?> validateNewInstanceClass(Class<?> type, Class<?> classCl
         }
         Class<?> nonNullType = PiNode.piCastNonNullClass(type, SnippetAnchorNode.anchor());
         if (probability(DEOPT_PROBABILITY,
-                        DynamicNewInstanceNode.throwsInstantiationException(nonNullType, classClass))) {
+                        DynamicNewInstanceNode.throwsInstantiationExceptionInjectedProbability(DEOPT_PROBABILITY, nonNullType, classClass))) {
             DeoptimizeNode.deopt(None, RuntimeConstraint);
         }
         return nonNullType;

compiler/src/org.graalvm.compiler.nodes/src/org/graalvm/compiler/nodes/extended/BranchProbabilityNode.java

@@ -37,12 +37,15 @@
 import org.graalvm.compiler.nodeinfo.NodeInfo;
 import org.graalvm.compiler.nodes.ConstantNode;
 import org.graalvm.compiler.nodes.FixedGuardNode;
+import org.graalvm.compiler.nodes.FrameState;
 import org.graalvm.compiler.nodes.IfNode;
 import org.graalvm.compiler.nodes.NodeView;
 import org.graalvm.compiler.nodes.ProfileData.BranchProbabilityData;
 import org.graalvm.compiler.nodes.ProfileData.ProfileSource;
 import org.graalvm.compiler.nodes.ReturnNode;
+import org.graalvm.compiler.nodes.ShortCircuitOrNode;
 import org.graalvm.compiler.nodes.ValueNode;
+import org.graalvm.compiler.nodes.ValuePhiNode;
 import org.graalvm.compiler.nodes.calc.ConditionalNode;
 import org.graalvm.compiler.nodes.calc.FloatingNode;
 import org.graalvm.compiler.nodes.calc.IntegerEqualsNode;
@@ -113,6 +116,9 @@ public BranchProbabilityNode(ValueNode probability, ValueNode condition) {
         super(TYPE, StampFactory.forKind(JavaKind.Boolean));
         this.probability = probability;
         this.condition = condition;
+
+        GraalError.guarantee(!(condition instanceof ShortCircuitOrNode),
+                        "Branch probabilities must be injected on simple conditions, not short-circuiting && or ||: %s", condition);
     }
 
     public BranchProbabilityNode(ValueNode condition) {
@@ -180,6 +186,9 @@ public void simplify(SimplifierTool tool) {
                     }
                 }
             }
+            if (!usageFound) {
+                usageFound = hasValidPhiUsage();
+            }
             if (usageFound) {
                 ValueNode currentCondition = condition;
                 IntegerStamp currentStamp = (IntegerStamp) currentCondition.stamp(NodeView.DEFAULT);
@@ -241,6 +250,53 @@ private boolean isSubstitutionGraph() {
         return hasExactlyOneUsage() && usages().first() instanceof ReturnNode;
     }
 
+    /**
+     * Normally a branch probability should be consumed directly as a condition, but in some cases
+     * it can be used as a value itself. For example:
+     *
+     * <pre>
+     * boolean helper() {
+     *     if (probability(a, ...) || probability(b, ...) || probability(c, condition)) {
+     *         return true;
+     *     } else {
+     *         return false;
+     *     }
+     * }
+     *
+     * ...
+     * if (probability(d, helper()) {
+     *     ...
+     * }
+     * </pre>
+     *
+     * After inlining the helper, {@code probability(c, condition)} can be represented as a branch
+     * probability node that feeds into a phi which is then used in an {@code if} condition. This is
+     * benign if that {@code if} has an injected branch probability itself.
+     */
+    private boolean hasValidPhiUsage() {
+        for (Node usage : this.usages()) {
+            if (usage instanceof ValuePhiNode && !((ValuePhiNode) usage).isLoopPhi()) {
+                Node phi = usage;
+                // We want exactly one non-state usage, and it must be a branch probability node.
+                Node uniquePhiUsage = null;
+                for (Node phiUsage : phi.usages()) {
+                    if (phiUsage instanceof FrameState) {
+                        continue;
+                    } else if (uniquePhiUsage == null) {
+                        uniquePhiUsage = phiUsage;
+                    } else if (phiUsage != uniquePhiUsage) {
+                        uniquePhiUsage = null;
+                        break;
+                    }
+                }
+                if (uniquePhiUsage instanceof BranchProbabilityNode) {
+                    return true;
+                }
+            }
+        }
+        return false;
+    }
+
     /**
      * This intrinsic should only be used for the condition of an if statement. The parameter
      * condition should also only denote a simple condition and not a combined condition involving

compiler/src/org.graalvm.compiler.nodes/src/org/graalvm/compiler/nodes/java/DynamicNewInstanceNode.java

@@ -24,6 +24,8 @@
  */
 package org.graalvm.compiler.nodes.java;
 
+import static org.graalvm.compiler.nodes.extended.BranchProbabilityNode.probability;
+
 import java.lang.reflect.Modifier;
 
 import org.graalvm.compiler.core.common.type.ObjectStamp;
@@ -87,8 +89,16 @@ public Node canonical(CanonicalizerTool tool) {
         return this;
     }
 
-    public static boolean throwsInstantiationException(Class<?> type, Class<?> classClass) {
-        return type.isPrimitive() || type.isArray() || type.isInterface() || Modifier.isAbstract(type.getModifiers()) || type == classClass;
+    public static boolean throwsInstantiationExceptionInjectedProbability(double probability, Class<?> type, Class<?> classClass) {
+        /*
+         * This method is for use in a snippet and therefore injects probabilities for each
+         * disjunct.
+         */
+        return probability(probability, type.isPrimitive()) ||
+                        probability(probability, type.isArray()) ||
+                        probability(probability, type.isInterface()) ||
+                        probability(probability, Modifier.isAbstract(type.getModifiers())) ||
+                        probability(probability, type == classClass);
     }
 
     public static boolean throwsInstantiationException(ResolvedJavaType type, MetaAccessProvider metaAccess) {

compiler/src/org.graalvm.compiler.replacements/src/org/graalvm/compiler/replacements/arraycopy/ArrayCopySnippets.java

@@ -27,6 +27,7 @@
 import static jdk.vm.ci.services.Services.IS_BUILDING_NATIVE_IMAGE;
 import static org.graalvm.compiler.nodes.extended.BranchProbabilityNode.DEOPT_PROBABILITY;
 import static org.graalvm.compiler.nodes.extended.BranchProbabilityNode.FAST_PATH_PROBABILITY;
+import static org.graalvm.compiler.nodes.extended.BranchProbabilityNode.FREQUENT_PROBABILITY;
 import static org.graalvm.compiler.nodes.extended.BranchProbabilityNode.NOT_FREQUENT_PROBABILITY;
 import static org.graalvm.compiler.nodes.extended.BranchProbabilityNode.probability;
 
@@ -306,7 +307,7 @@ protected void doExactArraycopyWithExpandedLoopSnippet(Object src, int srcPos, O
         long destOffset = arrayBaseOffset + (long) destPos * scale;
 
         GuardingNode anchor = SnippetAnchorNode.anchor();
-        if (probability(NOT_FREQUENT_PROBABILITY, src == dest && srcPos < destPos)) {
+        if (probability(FREQUENT_PROBABILITY, src == dest) && probability(NOT_FREQUENT_PROBABILITY, srcPos < destPos)) {
             // bad aliased case so we need to copy the array from back to front
             for (int position = length - 1; probability(FAST_PATH_PROBABILITY, position >= 0); position--) {
                 Object value = GuardedUnsafeLoadNode.guardedLoad(src, sourceOffset + ((long) position) * scale, elementKind, arrayLocation, anchor);

compiler/src/org.graalvm.compiler.replacements/src/org/graalvm/compiler/replacements/gc/G1WriteBarrierSnippets.java

@@ -261,7 +261,7 @@ public void g1ArrayRangePreWriteBarrier(Address address, long length, @ConstantP
         Word thread = getThread();
         byte markingValue = thread.readByte(satbQueueMarkingActiveOffset(), SATB_QUEUE_MARKING_ACTIVE_LOCATION);
         // If the concurrent marker is not enabled or the vector length is zero, return.
-        if (probability(FREQUENT_PROBABILITY, markingValue == (byte) 0 || length == 0)) {
+        if (probability(FREQUENT_PROBABILITY, markingValue == (byte) 0) || probability(NOT_FREQUENT_PROBABILITY, length == 0)) {
             return;
         }