[GR-61244] Implement direct jump table fast path

PullRequest: graal/19687

Author: kittyoracle

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/core/aarch64/AArch64LIRGenerator.java

@@ -423,7 +423,7 @@ private static ConditionFlag toFloatConditionFlag(Condition cond, boolean unorde
     /**
      * Takes a Condition and returns the correct AArch64 specific ConditionFlag.
      */
-    private static ConditionFlag toIntConditionFlag(Condition cond) {
+    public static ConditionFlag toIntConditionFlag(Condition cond) {
         switch (cond) {
             case EQ:
                 return ConditionFlag.EQ;
@@ -566,8 +566,8 @@ protected StrategySwitchOp createStrategySwitchOp(SwitchStrategy strategy, Label
     }
 
     @Override
-    protected void emitRangeTableSwitch(int lowKey, LabelRef defaultTarget, LabelRef[] targets, AllocatableValue key) {
-        append(new RangeTableSwitchOp(lowKey, defaultTarget, targets, key));
+    protected void emitRangeTableSwitch(int lowKey, LabelRef defaultTarget, LabelRef[] targets, SwitchStrategy remainingStrategy, LabelRef[] remainingTargets, AllocatableValue key) {
+        append(new RangeTableSwitchOp(lowKey, defaultTarget, targets, remainingStrategy, remainingTargets, key));
     }
 
     @Override

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/core/amd64/AMD64LIRGenerator.java

@@ -1116,10 +1116,8 @@ public void emitStrategySwitch(SwitchStrategy strategy, AllocatableValue key, La
     }
 
     @Override
-    protected void emitRangeTableSwitch(int lowKey, LabelRef defaultTarget, LabelRef[] targets, AllocatableValue key) {
-        Variable scratch = newVariable(LIRKind.value(target().arch.getWordKind()));
-        Variable idxScratch = newVariable(key.getValueKind());
-        append(new RangeTableSwitchOp(lowKey, defaultTarget, targets, key, scratch, idxScratch));
+    protected void emitRangeTableSwitch(int lowKey, LabelRef defaultTarget, LabelRef[] targets, SwitchStrategy remainingStrategy, LabelRef[] remainingTargets, AllocatableValue key) {
+        append(new RangeTableSwitchOp(this, lowKey, defaultTarget, targets, remainingStrategy, remainingTargets, key));
     }
 
     @Override

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/lir/SwitchStrategy.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013, 2018, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2013, 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
@@ -31,6 +31,7 @@
 import jdk.graal.compiler.asm.Label;
 import jdk.graal.compiler.core.common.calc.Condition;
 import jdk.graal.compiler.debug.Assertions;
+import jdk.graal.compiler.debug.GraalError;
 import jdk.graal.compiler.lir.asm.CompilationResultBuilder;
 
 import jdk.vm.ci.meta.Constant;
@@ -212,7 +213,6 @@ public double getAverageEffort() {
     private EffortClosure effortClosure;
 
     public SwitchStrategy(double[] keyProbabilities) {
-        assert keyProbabilities.length >= 2 : Assertions.errorMessage(keyProbabilities);
         this.keyProbabilities = keyProbabilities;
     }
 
@@ -334,8 +334,9 @@ public static class RangesStrategy extends PrimitiveStrategy {
 
         public RangesStrategy(final double[] keyProbabilities, JavaConstant[] keyConstants) {
             super(keyProbabilities, keyConstants);
-
             int keyCount = keyConstants.length;
+            GraalError.guarantee(keyCount > 1, "%s", Arrays.toString(keyConstants));
+
             indexes = new Integer[keyCount];
             for (int i = 0; i < keyCount; i++) {
                 indexes[i] = i;
@@ -398,6 +399,7 @@ public static class BinaryStrategy extends PrimitiveStrategy {
 
         public BinaryStrategy(double[] keyProbabilities, JavaConstant[] keyConstants) {
             super(keyProbabilities, keyConstants);
+            GraalError.guarantee(keyProbabilities.length > 1, "%s", Arrays.toString(keyConstants));
             probabilitySums = new double[keyProbabilities.length + 1];
             double sum = 0;
             for (int i = 0; i < keyConstants.length; i++) {
@@ -501,8 +503,14 @@ private void recurseBinarySwitch(SwitchClosure closure, int left, int right, int
     public abstract void run(SwitchClosure closure);
 
     private static SwitchStrategy[] getStrategies(double[] keyProbabilities, JavaConstant[] keyConstants, LabelRef[] keyTargets) {
-        SwitchStrategy[] strategies = new SwitchStrategy[]{new SequentialStrategy(keyProbabilities, keyConstants), new RangesStrategy(keyProbabilities, keyConstants),
-                        new BinaryStrategy(keyProbabilities, keyConstants)};
+        SwitchStrategy[] strategies;
+        if (keyProbabilities.length == 1) {
+            strategies = new SwitchStrategy[]{new SequentialStrategy(keyProbabilities, keyConstants)};
+        } else {
+            strategies = new SwitchStrategy[]{new SequentialStrategy(keyProbabilities, keyConstants), new RangesStrategy(keyProbabilities, keyConstants),
+                            new BinaryStrategy(keyProbabilities, keyConstants)};
+        }
+
         for (SwitchStrategy strategy : strategies) {
             strategy.effortClosure = strategy.new EffortClosure(keyTargets);
             strategy.run(strategy.effortClosure);

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/lir/aarch64/AArch64ControlFlow.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2013, 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
@@ -44,6 +44,7 @@
 import jdk.graal.compiler.asm.aarch64.AArch64MacroAssembler.ScratchRegister;
 import jdk.graal.compiler.code.CompilationResult.JumpTable;
 import jdk.graal.compiler.code.CompilationResult.JumpTable.EntryFormat;
+import jdk.graal.compiler.core.aarch64.AArch64LIRGenerator;
 import jdk.graal.compiler.core.common.NumUtil;
 import jdk.graal.compiler.core.common.calc.Condition;
 import jdk.graal.compiler.debug.Assertions;
@@ -385,7 +386,8 @@ private static void emitCompareHelper(CompilationResultBuilder crb, AArch64Macro
 
     /**
      * This operation jumps to the appropriate destination as specified within a JumpTable, or to
-     * the default condition if there is no match within the JumpTable.
+     * another destination according to {@code remainingStrategy} if there is no match within the
+     * JumpTable.
      *
      * <p>
      * The JumpTable contains a series of target offsets, relative to the start of the jump table,
@@ -395,8 +397,12 @@ private static void emitCompareHelper(CompilationResultBuilder crb, AArch64Macro
      * <ol>
      * <li>Determine whether the index is within the JumpTable. This is accomplished by first
      * normalizing the index (normalizedIdx == index - lowKey), and then checking whether
-     * <code>(unsigned(normalizedIdx) &lt;= highKey - lowKey</code>). If not, then one must jump to
-     * the defaultTarget.</li>
+     * <code>(unsigned(normalizedIdx) &lt;= highKey - lowKey</code>).
+     *
+     * <li>If normalizedIdx is not in the JumpTable, the destination is decided by {@code
+     * remainingStrategy}. If {@code remainingStrategy == null}, then the destination is {@code
+     * defaultTarget}. Otherwise, the destination is {@code defaultTarget} or one of {@code
+     * remainingTargets} based on the value of {@code key}.</li>
      *
      * <li>If normalizedIdx is within the JumpTable, then jump to JumpTableStart +
      * JumpTable[normalizedIdx].</li>
@@ -407,53 +413,80 @@ public static final class RangeTableSwitchOp extends AArch64BlockEndOp {
         private final int lowKey;
         private final LabelRef defaultTarget;
         private final LabelRef[] targets;
-        @Use({REG}) protected AllocatableValue index;
+        private final SwitchStrategy remainingStrategy;
+        private final LabelRef[] remainingTargets;
+        @Alive({REG}) protected AllocatableValue key;
 
-        public RangeTableSwitchOp(final int lowKey, final LabelRef defaultTarget, final LabelRef[] targets, AllocatableValue index) {
+        public RangeTableSwitchOp(int lowKey, LabelRef defaultTarget, LabelRef[] targets, SwitchStrategy remainingStrategy, LabelRef[] remainingTargets, AllocatableValue key) {
             super(TYPE);
             this.lowKey = lowKey;
             assert defaultTarget != null;
             this.defaultTarget = defaultTarget;
             this.targets = targets;
-            this.index = index;
+            this.remainingStrategy = remainingStrategy;
+            this.remainingTargets = remainingTargets;
+            this.key = key;
         }
 
         @Override
         public void emitCode(CompilationResultBuilder crb, AArch64MacroAssembler masm) {
             try (ScratchRegister sc1 = masm.getScratchRegister(); ScratchRegister sc2 = masm.getScratchRegister()) {
+                Register keyReg = asRegister(key);
                 Register scratch1 = sc1.getRegister();
                 Register scratch2 = sc2.getRegister();
+                GraalError.guarantee(!keyReg.equals(scratch1) && !keyReg.equals(scratch2), "must not alias");
                 /* Compare index against jump table bounds */
                 int highKey = lowKey + targets.length - 1;
-                masm.sub(32, scratch2, asRegister(index), lowKey);
-                int keyDiff = highKey - lowKey; // equivalent to targets.length - 1
-                if (AArch64MacroAssembler.isComparisonImmediate(keyDiff)) {
-                    masm.compare(32, scratch2, keyDiff);
+                Register keyOffsetReg = keyReg;
+                if (lowKey != 0) {
+                    masm.sub(32, scratch2, keyReg, lowKey);
+                    keyOffsetReg = scratch2;
+                }
+
+                int interval = highKey - lowKey;
+                if (AArch64MacroAssembler.isComparisonImmediate(interval)) {
+                    masm.compare(32, keyOffsetReg, interval);
                 } else {
-                    masm.mov(scratch1, keyDiff);
-                    masm.cmp(32, scratch2, scratch1);
+                    masm.mov(scratch1, interval);
+                    masm.cmp(32, keyOffsetReg, scratch1);
                 }
 
-                // Jump to default target if index is not within the jump table
-                masm.branchConditionally(ConditionFlag.HI, defaultTarget.label());
+                Label outOfRangeLabel = defaultTarget.label();
+                if (remainingStrategy != null) {
+                    Label remainingLabel = new Label();
+                    outOfRangeLabel = remainingLabel;
+
+                    crb.getLIR().addSlowPath(this, () -> {
+                        masm.bind(remainingLabel);
+                        new StrategySwitchOp(remainingStrategy, remainingTargets, defaultTarget, key, AArch64LIRGenerator::toIntConditionFlag).emitCode(crb, masm);
+                    });
+                }
+                // Jump to outOfRangeLabel if index is not within the jump table
+                masm.branchConditionally(ConditionFlag.HI, outOfRangeLabel);
 
-                emitJumpTable(crb, masm, scratch1, scratch2, lowKey, highKey, Arrays.stream(targets).map(LabelRef::label));
+                emitJumpTable(crb, masm, keyOffsetReg, scratch1, scratch2, lowKey, highKey, Arrays.stream(targets).map(LabelRef::label));
             }
         }
 
-        public static void emitJumpTable(CompilationResultBuilder crb, AArch64MacroAssembler masm, Register scratch, Register index, int lowKey, int highKey, Stream<Label> targets) {
+        public static void emitJumpTable(CompilationResultBuilder crb, AArch64MacroAssembler masm, Register scratch, Register keyScratch, int lowKey, int highKey, Stream<Label> targets) {
+            emitJumpTable(crb, masm, keyScratch, scratch, keyScratch, lowKey, highKey, targets);
+        }
+
+        private static void emitJumpTable(CompilationResultBuilder crb, AArch64MacroAssembler masm, Register key, Register scratch, Register idxScratch, int lowKey, int highKey,
+                        Stream<Label> targets) {
+            GraalError.guarantee(!key.equals(scratch), "must not alias");
             Label jumpTable = new Label();
             // load start of jump table
             masm.adr(scratch, jumpTable);
             /*
              * Note scratch holds the start of the jump table and index stores the normalized index.
              * Because each jumpTable index is 4 bytes large, index should be scaled.
              */
-            AArch64Address jumpTableEntryAddr = AArch64Address.createExtendedRegisterOffsetAddress(32, scratch, index, true, ExtendType.UXTW);
+            AArch64Address jumpTableEntryAddr = AArch64Address.createExtendedRegisterOffsetAddress(32, scratch, key, true, ExtendType.UXTW);
             // load relative target offset
-            masm.ldrs(64, 32, index, jumpTableEntryAddr);
+            masm.ldrs(64, 32, idxScratch, jumpTableEntryAddr);
             // compute target address (jumpTableStart + target offset)
-            masm.add(64, scratch, scratch, index);
+            masm.add(64, scratch, scratch, idxScratch);
             // jump to target
             masm.jmp(scratch);
 

compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/lir/amd64/AMD64ControlFlow.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2011, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2011, 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
@@ -43,6 +43,7 @@
 import jdk.graal.compiler.asm.amd64.AMD64MacroAssembler;
 import jdk.graal.compiler.code.CompilationResult.JumpTable;
 import jdk.graal.compiler.code.CompilationResult.JumpTable.EntryFormat;
+import jdk.graal.compiler.core.common.LIRKind;
 import jdk.graal.compiler.core.common.NumUtil;
 import jdk.graal.compiler.core.common.Stride;
 import jdk.graal.compiler.core.common.calc.Condition;
@@ -57,6 +58,7 @@
 import jdk.graal.compiler.lir.SwitchStrategy;
 import jdk.graal.compiler.lir.Variable;
 import jdk.graal.compiler.lir.asm.CompilationResultBuilder;
+import jdk.graal.compiler.lir.gen.LIRGenerator;
 import jdk.vm.ci.amd64.AMD64;
 import jdk.vm.ci.amd64.AMD64Kind;
 import jdk.vm.ci.code.Register;
@@ -655,59 +657,98 @@ protected void conditionalJump(int index, Condition condition, Label target) {
         }
     }
 
+    /**
+     * See {@code LIRGenerator::emitRangeTableSwitch}.
+     */
     public static final class RangeTableSwitchOp extends AMD64BlockEndOp {
         public static final LIRInstructionClass<RangeTableSwitchOp> TYPE = LIRInstructionClass.create(RangeTableSwitchOp.class);
         private final int lowKey;
         private final LabelRef defaultTarget;
         private final LabelRef[] targets;
-        @LIRInstruction.Use protected Value index;
-        @LIRInstruction.Temp({LIRInstruction.OperandFlag.REG, LIRInstruction.OperandFlag.HINT}) protected Value idxScratch;
-        @LIRInstruction.Temp protected Value scratch;
+        private final SwitchStrategy remainingStrategy;
+        private final LabelRef[] remainingTargets;
+        @LIRInstruction.Use(OperandFlag.REG) protected AllocatableValue key;
+        @LIRInstruction.Temp(OperandFlag.REG) protected AllocatableValue scratch1;
+        @LIRInstruction.Temp(OperandFlag.REG) protected AllocatableValue scratch2;
 
-        public RangeTableSwitchOp(final int lowKey, final LabelRef defaultTarget, final LabelRef[] targets, Value index, Variable scratch, Variable idxScratch) {
+        public RangeTableSwitchOp(LIRGenerator gen, int lowKey, LabelRef defaultTarget, LabelRef[] targets, SwitchStrategy remainingStrategy, LabelRef[] remainingTargets, AllocatableValue key) {
             super(TYPE);
             this.lowKey = lowKey;
             assert defaultTarget != null;
             this.defaultTarget = defaultTarget;
             this.targets = targets;
-            this.index = index;
-            this.scratch = scratch;
-            this.idxScratch = idxScratch;
+            this.remainingStrategy = remainingStrategy;
+            this.remainingTargets = remainingTargets;
+            this.key = key;
+            this.scratch1 = gen.newVariable(LIRKind.value(AMD64Kind.DWORD));
+            this.scratch2 = gen.newVariable(LIRKind.value(AMD64Kind.DWORD));
         }
 
         @Override
         public void emitCode(CompilationResultBuilder crb, AMD64MacroAssembler masm) {
-            Register indexReg = asRegister(index, AMD64Kind.DWORD);
-            Register idxScratchReg = asRegister(idxScratch, AMD64Kind.DWORD);
-            Register scratchReg = asRegister(scratch, AMD64Kind.QWORD);
-
-            if (!indexReg.equals(idxScratchReg)) {
-                masm.movl(idxScratchReg, indexReg);
+            Register keyReg = asRegister(key);
+            AllocatableValue scratch;
+            AllocatableValue idxScratch;
+            if (asRegister(scratch1).equals(keyReg)) {
+                // keyReg cannot alias with scratchReg but may alias with idxScratchReg
+                scratch = scratch2;
+                idxScratch = key;
+            } else {
+                scratch = scratch1;
+                idxScratch = scratch2;
             }
+            Register scratchReg = asRegister(scratch);
+            Register idxScratchReg = asRegister(idxScratch);
 
             // Compare index against jump table bounds
             int highKey = lowKey + targets.length - 1;
+            Register keyOffsetReg;
             if (lowKey != 0) {
                 // subtract the low value from the switch value
-                masm.subl(idxScratchReg, lowKey);
+                if (keyReg.equals(idxScratchReg)) {
+                    masm.addl(idxScratchReg, -lowKey);
+                } else {
+                    masm.lead(idxScratchReg, new AMD64Address(keyReg, -lowKey));
+                }
                 masm.cmpl(idxScratchReg, highKey - lowKey);
+                keyOffsetReg = idxScratchReg;
             } else {
-                masm.cmpl(idxScratchReg, highKey);
+                masm.cmpl(keyReg, highKey);
+                keyOffsetReg = keyReg;
             }
 
-            // Jump to default target if index is not within the jump table
-            masm.jcc(ConditionFlag.Above, defaultTarget.label());
+            Label outOfRangeLabel = defaultTarget.label();
+            if (remainingStrategy != null) {
+                Label remainingLabel = new Label();
+                outOfRangeLabel = remainingLabel;
+                boolean needsKeyRecover = lowKey != 0 && keyReg.equals(idxScratchReg);
 
-            emitJumpTable(crb, masm, scratchReg, idxScratchReg, lowKey, highKey, Arrays.stream(targets).map(LabelRef::label));
+                crb.getLIR().addSlowPath(this, () -> {
+                    masm.bind(remainingLabel);
+                    if (needsKeyRecover) {
+                        masm.addl(keyReg, lowKey);
+                    }
+                    new StrategySwitchOp(remainingStrategy, remainingTargets, defaultTarget, key, scratch).emitCode(crb, masm);
+                });
+            }
+            masm.jcc(ConditionFlag.Above, outOfRangeLabel);
+
+            emitJumpTable(crb, masm, keyOffsetReg, scratchReg, idxScratchReg, lowKey, highKey, Arrays.stream(targets).map(LabelRef::label));
         }
 
         public static void emitJumpTable(CompilationResultBuilder crb, AMD64MacroAssembler masm, Register scratchReg, Register idxScratchReg, int lowKey, int highKey, Stream<Label> targets) {
+            emitJumpTable(crb, masm, idxScratchReg, scratchReg, idxScratchReg, lowKey, highKey, targets);
+        }
+
+        private static void emitJumpTable(CompilationResultBuilder crb, AMD64MacroAssembler masm, Register keyReg, Register scratchReg, Register idxScratchReg, int lowKey, int highKey,
+                        Stream<Label> targets) {
+            GraalError.guarantee(!keyReg.equals(scratchReg), "must not alias");
             // Set scratch to address of jump table
             masm.leaq(scratchReg, new AMD64Address(AMD64.rip, 0));
             final int afterLea = masm.position();
 
             // Load jump table entry into scratch and jump to it
-            masm.movslq(idxScratchReg, new AMD64Address(scratchReg, idxScratchReg, Stride.S4, 0));
+            masm.movslq(idxScratchReg, new AMD64Address(scratchReg, keyReg, Stride.S4, 0));
             masm.addq(scratchReg, idxScratchReg);
             masm.jmp(scratchReg);