[llvm][CodeGen] DAG Combiner folds for vscale.

[fpetrogalli]

Summary

This patch simplifies the DAGs generated when using the intrinsic @llvm.vscale.* as follows:

• Fold (add (vscale * C0), (vscale * C1)) to (vscale * (C0 + C1)).
• Canonicalize (sub X, (vscale * C)) to (add X, (vscale * -C)).
• Fold (mul (vscale * C0), C1) to (vscale * (C0 * C1)).
• Fold (shl (vscale * C0), C1) to (vscale * (C0 << C1)).

The test sve-gep-ll have been updated to reflect the folding introduced by this patch.

[andwar]

@fpetrogalli Thank you for working on this! A few nits, but nothing blocking. LGTM!

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2330

[nit] Missing opening bracket: (vscale * C0 + C1)) --> (vscale * (C0 + C1))

3601

[nit] Brackets: Fold (mul (vscale * C0), C1) to (vscale * (C0 * C1))) --> Fold (mul (vscale * C0), C1)) to (vscale * (C0 * C1))`

7782

[nit] (shl (vscale * C0), C1) to (vscale * (C0 << C1))) --> (shl (vscale * C0), C1)) to (vscale * (C0 << C1)) :)

llvm/test/CodeGen/AArch64/sve-vscale-combine.ll

73

[nit] Extra line

[sdesmalen]

LGTM!

[fpetrogalli]

Hi @andwar, I have updated the comments. I have also fixed the sve-gep.ll test, as the changes in the combiner are now folding the multiplying constant inside the vscale node.

Thank you,

Francesco

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

SDValue DAGCombiner::visitADD(SDNode *N) {
   if (SDValue V = foldAddSubOfSignBit(N, DAG))
     return V;
 
   // fold (a+b) -> (a|b) iff a and b share no bits.
   if ((!LegalOperations || TLI.isOperationLegal(ISD::OR, VT)) &&
       DAG.haveNoCommonBitsSet(N0, N1))
     return DAG.getNode(ISD::OR, DL, VT, N0, N1);
 
+  // Fold (add (vscale * C0), (vscale * C1)) to (vscale * C0 + C1))
+  if (N0.getOpcode() == ISD::VSCALE && N1.getOpcode() == ISD::VSCALE) {
+    APInt C0 = N0->getConstantOperandAPInt(0);
+    APInt C1 = N1->getConstantOperandAPInt(0);
+    return DAG.getVScale(DL, VT, C0 + C1);
+  }
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitADDSAT(SDNode *N) {
   unsigned Opcode = N->getOpcode();
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   EVT VT = N0.getValueType();
if (N1.getOpcode() == ISD::SIGN_EXTEND_INREG) {
     VTSDNode *TN = cast<VTSDNode>(N1.getOperand(1));
     if (TN->getVT() == MVT::i1) {
       SDValue ZExt = DAG.getNode(ISD::AND, DL, VT, N1.getOperand(0),
                                  DAG.getConstant(1, DL, VT));
       return DAG.getNode(ISD::ADD, DL, VT, N0, ZExt);
     }
   }
 
+  // canonicalize (sub X, (vscale * C)) to (add X,  (vscale * -C))
+  if (N1.getOpcode() == ISD::VSCALE) {
+    APInt IntVal = N1.getConstantOperandAPInt(0);
+    return DAG.getNode(ISD::ADD, DL, VT, N0, DAG.getVScale(DL, VT, -IntVal));
+  }
+
   // Prefer an add for more folding potential and possibly better codegen:
   // sub N0, (lshr N10, width-1) --> add N0, (ashr N10, width-1)
   if (!LegalOperations && N1.getOpcode() == ISD::SRL && N1.hasOneUse()) {
     SDValue ShAmt = N1.getOperand(1);
     ConstantSDNode *ShAmtC = isConstOrConstSplat(ShAmt);
     if (ShAmtC &&
         ShAmtC->getAPIntValue() == (N1.getScalarValueSizeInBits() - 1)) {
       SDValue SRA = DAG.getNode(ISD::SRA, DL, VT, N1.getOperand(0), ShAmt);
if (DAG.isConstantIntBuildVectorOrConstantInt(N1) &&
       DAG.isConstantIntBuildVectorOrConstantInt(N0.getOperand(1)) &&
       isMulAddWithConstProfitable(N, N0, N1))
       return DAG.getNode(ISD::ADD, SDLoc(N), VT,
                          DAG.getNode(ISD::MUL, SDLoc(N0), VT,
                                      N0.getOperand(0), N1),
                          DAG.getNode(ISD::MUL, SDLoc(N1), VT,
                                      N0.getOperand(1), N1));
 
+  // Fold (mul (vscale * C0), C1) to (vscale * (C0 * C1)))
+  if (N0.getOpcode() == ISD::VSCALE)
+    if (ConstantSDNode *NC1 = isConstOrConstSplat(N1)) {
+      APInt C0 = N0.getConstantOperandAPInt(0);
+      APInt C1 = NC1->getAPIntValue();
+      return DAG.getVScale(SDLoc(N), VT, C0 * C1);
+    }
+
   // reassociate mul
   if (SDValue RMUL = reassociateOps(ISD::MUL, SDLoc(N), N0, N1, N->getFlags()))
     return RMUL;
 
   return SDValue();
 }
 
 /// Return true if divmod libcall is available.
if (N0.getOpcode() == ISD::MUL && N0.getNode()->hasOneUse() &&
     if (isConstantOrConstantVector(Shl))
       return DAG.getNode(ISD::MUL, SDLoc(N), VT, N0.getOperand(0), Shl);
   }
 
   if (N1C && !N1C->isOpaque())
     if (SDValue NewSHL = visitShiftByConstant(N))
       return NewSHL;
 
+  // Fold (shl (vscale * C0), C1) to (vscale * (C0 << C1)))
+  if (N0.getOpcode() == ISD::VSCALE)
+    if (ConstantSDNode *NC1 = isConstOrConstSplat(N->getOperand(1))) {
+      auto DL = SDLoc(N);
+      APInt C0 = N0.getConstantOperandAPInt(0);
+      APInt C1 = NC1->getAPIntValue();
+      return DAG.getVScale(DL, VT, C0 << C1);
+    }
+
   return SDValue();
 }
 
 SDValue DAGCombiner::visitSRA(SDNode *N) {
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   if (SDValue V = DAG.simplifyShift(N0, N1))
     return V;

llvm/test/CodeGen/AArch64/sve-vscale-combine.ll

+; RUN: llc -mtriple=aarch64--linux-gnu -mattr=+sve --asm-verbose=false < %s |FileCheck %s
+
+declare i32 @llvm.vscale.i32()
+declare i64 @llvm.vscale.i64()
+
+; Fold (add (vscale * C0), (vscale * C1)) to (vscale * (C0 + C1)))
+define i64 @combine_add_vscale_i64() nounwind {
+; CHECK-LABEL: combine_add_vscale_i64:
+; CHECK-NOT:   add
+; CHECK-NEXT:  cntd  x0
+; CHECK-NEXT:  ret
+ %vscale = call i64 @llvm.vscale.i64()
+ %add = add i64 %vscale, %vscale
+ ret i64 %add
+}
+
+define i32 @combine_add_vscale_i32() nounwind {
+; CHECK-LABEL: combine_add_vscale_i32:
+; CHECK-NOT:   add
+; CHECK-NEXT:  cntd  x0
+; CHECK-NEXT:  ret
+ %vscale = call i32 @llvm.vscale.i32()
+ %add = add i32 %vscale, %vscale
+ ret i32 %add
+}
+
+; Fold (mul (vscale * C0), C1) to (vscale * C0 * C1)))
+; In this test, C0 = 1, C1 = 32.
+define i64 @combine_mul_vscale_i64() nounwind {
+; CHECK-LABEL: combine_mul_vscale_i64:
+; CHECK-NOT:   mul
+; CHECK-NEXT:  rdvl  x0, #2
+; CHECK-NEXT:  ret
+ %vscale = call i64 @llvm.vscale.i64()
+ %mul = mul i64 %vscale, 32
+ ret i64 %mul
+}
+
+define i32 @combine_mul_vscale_i32() nounwind {
+; CHECK-LABEL: combine_mul_vscale_i32:
+; CHECK-NOT:   mul
+; CHECK-NEXT:  rdvl  x0, #3
+; CHECK-NEXT:  ret
+ %vscale = call i32 @llvm.vscale.i32()
+ %mul = mul i32 %vscale, 48
+ ret i32 %mul
+}
+
+; Canonicalize (sub X, (vscale * C)) to (add X,  (vscale * -C))
+define i64 @combine_sub_vscale_i64(i64 %in) nounwind {
+; CHECK-LABEL: combine_sub_vscale_i64:
+; CHECK-NOT:   sub
+; CHECK-NEXT:  rdvl  x8, #-1
+; CHECK-NEXT:  asr   x8, x8, #4
+; CHECK-NEXT:  add   x0, x0, x8
+; CHECK-NEXT:  ret
+ %vscale = call i64 @llvm.vscale.i64()
+ %sub = sub i64 %in,  %vscale
+ ret i64 %sub
+}
+
+define i32 @combine_sub_vscale_i32(i32 %in) nounwind {
+; CHECK-LABEL: combine_sub_vscale_i32:
+; CHECK-NOT:   sub
+; CHECK-NEXT:  rdvl  x8, #-1
+; CHECK-NEXT:  asr   x8, x8, #4
+; CHECK-NEXT:  add   w0, w0, w8
+; CHECK-NEXT:  ret
+ %vscale = call i32 @llvm.vscale.i32()
+ %sub = sub i32 %in, %vscale
+ ret i32 %sub
+}
+
+
+; Fold (shl (vscale * C0), C1) to (vscale * (C0 << C1)))
+; C0 = 1 , C1 = 4
+; At IR level,  %shl = 2^4 * VSCALE.
+; At Assembly level, the output of RDVL is also 2^4 * VSCALE.
+; Hence, the immediate for RDVL is #1.
+define i64 @combine_shl_vscale_i64() nounwind {
+; CHECK-LABEL: combine_shl_vscale_i64:
+; CHECK-NOT:   shl
+; CHECK-NEXT:  rdvl  x0, #1
+; CHECK-NEXT:  ret
+ %vscale = call i64 @llvm.vscale.i64()
+ %shl = shl i64 %vscale, 4
+ ret i64 %shl
+}
+
+define i32 @combine_shl_vscale_i32() nounwind {
+; CHECK-LABEL: combine_shl_vscale_i32:
+; CHECK-NOT:   shl
+; CHECK-NEXT:  rdvl  x0, #1
+; CHECK-NEXT:  ret
+ %vscale = call i32 @llvm.vscale.i32()
+ %shl = shl i32 %vscale, 4
+ ret i32 %shl
+}