[LLVM][DAGCombiner] fold (shl (X * vscale(C0)), C1) -> (X * vscale(C0 << C1)).

Author: paulwalker-arm

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -10615,6 +10615,15 @@ SDValue DAGCombiner::visitSHL(SDNode *N) {
     return DAG.getVScale(DL, VT, C0 << C1);
   }
 
+  SDValue X;
+  APInt VS0;
+
+  // fold (shl (X * vscale(VS0)), C1) -> (X * vscale(VS0 << C1))
+  if (N1C && sd_match(N0, m_Mul(m_Value(X), m_VScale(m_ConstInt(VS0))))) {
+    SDValue VScale = DAG.getVScale(DL, VT, VS0 << N1C->getAPIntValue());
+    return DAG.getNode(ISD::MUL, DL, VT, X, VScale);
+  }
+
   // Fold (shl step_vector(C0), C1) to (step_vector(C0 << C1)).
   APInt ShlVal;
   if (N0.getOpcode() == ISD::STEP_VECTOR &&

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

@@ -120,10 +120,8 @@ define i32 @combine_shl_vscale_i32() nounwind {
 define i64 @combine_shl_mul_vscale(i64 %a) nounwind {
 ; CHECK-LABEL: combine_shl_mul_vscale:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    rdvl x8, #1
-; CHECK-NEXT:    lsr x8, x8, #4
-; CHECK-NEXT:    mul x8, x0, x8
-; CHECK-NEXT:    lsl x0, x8, #3
+; CHECK-NEXT:    cnth x8
+; CHECK-NEXT:    mul x0, x0, x8
 ; CHECK-NEXT:    ret
   %vscale = tail call i64 @llvm.vscale.i64()
   %mul = mul i64 %a, %vscale
@@ -134,10 +132,8 @@ define i64 @combine_shl_mul_vscale(i64 %a) nounwind {
 define i64 @combine_shl_mul_vscale_commuted(i64 %a) nounwind {
 ; CHECK-LABEL: combine_shl_mul_vscale_commuted:
 ; CHECK:       // %bb.0:
-; CHECK-NEXT:    rdvl x8, #1
-; CHECK-NEXT:    lsr x8, x8, #4
-; CHECK-NEXT:    mul x8, x8, x0
-; CHECK-NEXT:    lsl x0, x8, #3
+; CHECK-NEXT:    cnth x8
+; CHECK-NEXT:    mul x0, x0, x8
 ; CHECK-NEXT:    ret
   %vscale = tail call i64 @llvm.vscale.i64()
   %mul = mul i64 %vscale, %a

llvm/test/Transforms/LoopStrengthReduce/AArch64/vscale-fixups.ll

@@ -383,14 +383,14 @@ define void @vscale_squared_offset(ptr %alloc) #0 {
 ; COMMON-LABEL: vscale_squared_offset:
 ; COMMON:       // %bb.0: // %entry
 ; COMMON-NEXT:    rdvl x9, #1
+; COMMON-NEXT:    rdvl x10, #4
 ; COMMON-NEXT:    fmov z0.s, #4.00000000
-; COMMON-NEXT:    mov x8, xzr
 ; COMMON-NEXT:    lsr x9, x9, #4
 ; COMMON-NEXT:    fmov z1.s, #8.00000000
-; COMMON-NEXT:    cntw x10
+; COMMON-NEXT:    mov x8, xzr
 ; COMMON-NEXT:    ptrue p0.s, vl1
-; COMMON-NEXT:    umull x9, w9, w9
-; COMMON-NEXT:    lsl x9, x9, #6
+; COMMON-NEXT:    umull x9, w9, w10
+; COMMON-NEXT:    cntw x10
 ; COMMON-NEXT:    cmp x8, x10
 ; COMMON-NEXT:    b.ge .LBB6_2
 ; COMMON-NEXT:  .LBB6_1: // %for.body