[EVM][DAGCombine] Expand SELECT of constant to arithmetic

This patch implements following optimization:

| Source pattern           | Rewritten form     | Proof (IR equivalent)               |
|--------------------------|--------------------|-------------------------------------|
| select C, X, C2          | C2 + (X - C2) * C  | https://alive2.llvm.org/ce/z/V2Cd9h |
| select C, C1, Y          | Y + (C1 - Y) * C   | https://alive2.llvm.org/ce/z/g3ki-b |
| select C, C1, C2 C1 > C2 | C2 + (C1 - C2) * C | https://alive2.llvm.org/ce/z/cOCHY_ |
| select C, C1, C2 C1 < C2 | C2 - (C2 - C1) * C | https://alive2.llvm.org/ce/z/P8jkjd |

When we have select of two constants, we are doing two
different transformations to avoid generating negative
constants, which are expensive in EVM.

Signed-off-by: Vladimir Radosavljevic <[email protected]>

Author: vladimirradosavljevic

llvm/lib/Target/EVM/EVMISelLowering.cpp

@@ -876,6 +876,49 @@ SDValue EVMTargetLowering::combineSELECT(SDNode *N,
     return DAG.getNode(ISD::OR, DL, VT, Neg, DAG.getFreeze(FalseV));
   }
 
+  // fold (Cond ? C1 : C2) -> (C2 + (C1 - C2) * Cond) if (C1 > C2)
+  // fold (Cond ? C1 : C2) -> (C2 - (C2 - C1) * Cond) if (C1 < C2)
+  //
+  // We distinguish between C1 > C2 and C1 < C2 to avoid generating
+  // negative constants, which are expensive in EVM.
+  if (isa<ConstantSDNode>(TrueV) && isa<ConstantSDNode>(FalseV)) {
+    APInt C1Val = cast<ConstantSDNode>(TrueV)->getAPIntValue();
+    APInt C2Val = cast<ConstantSDNode>(FalseV)->getAPIntValue();
+
+    // This should never happen, but just in case.
+    if (C1Val == C2Val)
+      return TrueV;
+
+    unsigned Opc = 0;
+    SDValue Scale;
+    if (C1Val.sgt(C2Val)) {
+      Opc = ISD::ADD;
+      Scale = DAG.getConstant(C1Val - C2Val, DL, VT);
+    } else {
+      // Use sub to avoid generating negative constants.
+      Opc = ISD::SUB;
+      Scale = DAG.getConstant(C2Val - C1Val, DL, VT);
+    }
+
+    SDValue Mul = DAG.getNode(ISD::MUL, DL, VT, CondV, Scale);
+    return DAG.getNode(Opc, DL, VT, FalseV, Mul);
+  }
+
+  // fold (Cond ? X : C) -> (C + (X* - C) * Cond)
+  // fold (Cond ? C : Y) -> (Y* + (C - Y*) * Cond)
+  if (isa<ConstantSDNode>(TrueV) || isa<ConstantSDNode>(FalseV)) {
+    if (!isa<ConstantSDNode>(TrueV))
+      TrueV = freeze(TrueV);
+    else if (!isa<ConstantSDNode>(FalseV))
+      FalseV = freeze(FalseV);
+    else
+      llvm_unreachable("Unexpected select operands");
+
+    SDValue Sub = DAG.getNode(ISD::SUB, DL, VT, TrueV, FalseV);
+    SDValue Mul = DAG.getNode(ISD::MUL, DL, VT, CondV, Sub);
+    return DAG.getNode(ISD::ADD, DL, VT, FalseV, Mul);
+  }
+
   if (SDValue V = tryFoldSelectIntoOp(N, DAG, TrueV, FalseV, /*Swapped=*/false))
     return V;
   // NOLINTNEXTLINE(readability-suspicious-call-argument)

llvm/test/CodeGen/EVM/select-const.ll

@@ -48,20 +48,13 @@ define i256 @select_const_int_harder(i1 %a) {
 ; CHECK-LABEL: select_const_int_harder:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH1 0x26
+; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    PUSH1 0x1
 ; CHECK-NEXT:    AND
-; CHECK-NEXT:    ISZERO
-; CHECK-NEXT:    PUSH4 @.BB3_2
-; CHECK-NEXT:    JUMPI
-; CHECK-NEXT:  ; %bb.1:
-; CHECK-NEXT:    PUSH1 0x6
-; CHECK-NEXT:    PUSH4 @.BB3_3
-; CHECK-NEXT:    JUMP
-; CHECK-NEXT:  .BB3_2:
-; CHECK-NEXT:    JUMPDEST
-; CHECK-NEXT:    PUSH1 0x26
-; CHECK-NEXT:  .BB3_3:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH1 0x5
+; CHECK-NEXT:    SHL
+; CHECK-NEXT:    XOR
 ; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    JUMP
   %1 = select i1 %a, i256 6, i256 38
@@ -306,20 +299,13 @@ define i256 @select_196_184(i1 %a) {
 ; CHECK-LABEL: select_196_184:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH1 0xC
+; CHECK-NEXT:    PUSH1 0xB8
+; CHECK-NEXT:    SWAP2
 ; CHECK-NEXT:    PUSH1 0x1
 ; CHECK-NEXT:    AND
-; CHECK-NEXT:    ISZERO
-; CHECK-NEXT:    PUSH4 @.BB20_2
-; CHECK-NEXT:    JUMPI
-; CHECK-NEXT:  ; %bb.1:
-; CHECK-NEXT:    PUSH1 0xC4
-; CHECK-NEXT:    PUSH4 @.BB20_3
-; CHECK-NEXT:    JUMP
-; CHECK-NEXT:  .BB20_2:
-; CHECK-NEXT:    JUMPDEST
-; CHECK-NEXT:    PUSH1 0xB8
-; CHECK-NEXT:  .BB20_3:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    MUL
+; CHECK-NEXT:    ADD
 ; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    JUMP
   %1 = select i1 %a, i256 196, i256 184
@@ -330,20 +316,13 @@ define i256 @select_184_196(i1 %a) {
 ; CHECK-LABEL: select_184_196:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH1 0xC
+; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    PUSH1 0x1
 ; CHECK-NEXT:    AND
-; CHECK-NEXT:    ISZERO
-; CHECK-NEXT:    PUSH4 @.BB21_2
-; CHECK-NEXT:    JUMPI
-; CHECK-NEXT:  ; %bb.1:
-; CHECK-NEXT:    PUSH1 0xB8
-; CHECK-NEXT:    PUSH4 @.BB21_3
-; CHECK-NEXT:    JUMP
-; CHECK-NEXT:  .BB21_2:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    MUL
 ; CHECK-NEXT:    PUSH1 0xC4
-; CHECK-NEXT:  .BB21_3:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    SUB
 ; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    JUMP
   %1 = select i1 %a, i256 184, i256 196
@@ -354,22 +333,14 @@ define i256 @select_n196_n184(i1 %a) {
 ; CHECK-LABEL: select_n196_n184:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH1 0xC
+; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    PUSH1 0x1
 ; CHECK-NEXT:    AND
-; CHECK-NEXT:    ISZERO
-; CHECK-NEXT:    PUSH4 @.BB22_2
-; CHECK-NEXT:    JUMPI
-; CHECK-NEXT:  ; %bb.1:
-; CHECK-NEXT:    PUSH1 0xC3
-; CHECK-NEXT:    NOT
-; CHECK-NEXT:    PUSH4 @.BB22_3
-; CHECK-NEXT:    JUMP
-; CHECK-NEXT:  .BB22_2:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    MUL
 ; CHECK-NEXT:    PUSH1 0xB7
 ; CHECK-NEXT:    NOT
-; CHECK-NEXT:  .BB22_3:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    SUB
 ; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    JUMP
   %1 = select i1 %a, i256 -196, i256 -184
@@ -380,22 +351,13 @@ define i256 @select_n184_n196(i1 %a) {
 ; CHECK-LABEL: select_n184_n196:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH1 0xC
+; CHECK-NEXT:    PUSH1 0xC4
+; CHECK-NEXT:    SWAP2
 ; CHECK-NEXT:    PUSH1 0x1
 ; CHECK-NEXT:    AND
-; CHECK-NEXT:    ISZERO
-; CHECK-NEXT:    PUSH4 @.BB23_2
-; CHECK-NEXT:    JUMPI
-; CHECK-NEXT:  ; %bb.1:
-; CHECK-NEXT:    PUSH1 0xB7
-; CHECK-NEXT:    NOT
-; CHECK-NEXT:    PUSH4 @.BB23_3
-; CHECK-NEXT:    JUMP
-; CHECK-NEXT:  .BB23_2:
-; CHECK-NEXT:    JUMPDEST
-; CHECK-NEXT:    PUSH1 0xC3
-; CHECK-NEXT:    NOT
-; CHECK-NEXT:  .BB23_3:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    MUL
+; CHECK-NEXT:    SUB
 ; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    JUMP
   %1 = select i1 %a, i256 -184, i256 -196
@@ -406,15 +368,15 @@ define i256 @select_var_12345(i1 %a, i256 %b) {
 ; CHECK-LABEL: select_var_12345:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH2 0x3039
+; CHECK-NEXT:    DUP1
+; CHECK-NEXT:    SWAP3
+; CHECK-NEXT:    SUB
+; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    PUSH1 0x1
 ; CHECK-NEXT:    AND
-; CHECK-NEXT:    PUSH4 @.BB24_2
-; CHECK-NEXT:    JUMPI
-; CHECK-NEXT:  ; %bb.1:
-; CHECK-NEXT:    POP
-; CHECK-NEXT:    PUSH2 0x3039
-; CHECK-NEXT:  .BB24_2:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    MUL
+; CHECK-NEXT:    ADD
 ; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    JUMP
   %1 = select i1 %a, i256 %b, i256 12345
@@ -425,16 +387,14 @@ define i256 @select_12345_var(i1 %a, i256 %b) {
 ; CHECK-LABEL: select_12345_var:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    DUP2
+; CHECK-NEXT:    PUSH2 0x3039
+; CHECK-NEXT:    SUB
+; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    PUSH1 0x1
 ; CHECK-NEXT:    AND
-; CHECK-NEXT:    ISZERO
-; CHECK-NEXT:    PUSH4 @.BB25_2
-; CHECK-NEXT:    JUMPI
-; CHECK-NEXT:  ; %bb.1:
-; CHECK-NEXT:    POP
-; CHECK-NEXT:    PUSH2 0x3039
-; CHECK-NEXT:  .BB25_2:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    MUL
+; CHECK-NEXT:    ADD
 ; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    JUMP
   %1 = select i1 %a, i256 12345, i256 %b
@@ -445,16 +405,15 @@ define i256 @select_var_n12345(i1 %a, i256 %b) {
 ; CHECK-LABEL: select_var_n12345:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    PUSH2 0x3039
+; CHECK-NEXT:    SWAP2
+; CHECK-NEXT:    DUP3
+; CHECK-NEXT:    ADD
+; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    PUSH1 0x1
 ; CHECK-NEXT:    AND
-; CHECK-NEXT:    PUSH4 @.BB26_2
-; CHECK-NEXT:    JUMPI
-; CHECK-NEXT:  ; %bb.1:
-; CHECK-NEXT:    POP
-; CHECK-NEXT:    PUSH2 0x3038
-; CHECK-NEXT:    NOT
-; CHECK-NEXT:  .BB26_2:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    MUL
+; CHECK-NEXT:    SUB
 ; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    JUMP
   %1 = select i1 %a, i256 %b, i256 -12345
@@ -465,17 +424,15 @@ define i256 @select_n12345_var(i1 %a, i256 %b) {
 ; CHECK-LABEL: select_n12345_var:
 ; CHECK:       ; %bb.0:
 ; CHECK-NEXT:    JUMPDEST
-; CHECK-NEXT:    PUSH1 0x1
-; CHECK-NEXT:    AND
-; CHECK-NEXT:    ISZERO
-; CHECK-NEXT:    PUSH4 @.BB27_2
-; CHECK-NEXT:    JUMPI
-; CHECK-NEXT:  ; %bb.1:
-; CHECK-NEXT:    POP
+; CHECK-NEXT:    DUP2
 ; CHECK-NEXT:    PUSH2 0x3038
 ; CHECK-NEXT:    NOT
-; CHECK-NEXT:  .BB27_2:
-; CHECK-NEXT:    JUMPDEST
+; CHECK-NEXT:    SUB
+; CHECK-NEXT:    SWAP1
+; CHECK-NEXT:    PUSH1 0x1
+; CHECK-NEXT:    AND
+; CHECK-NEXT:    MUL
+; CHECK-NEXT:    ADD
 ; CHECK-NEXT:    SWAP1
 ; CHECK-NEXT:    JUMP
   %1 = select i1 %a, i256 -12345, i256 %b