[InstCombine] Detect different vscales in div by shift combine.

This attempts to fix a regression in code that performs `svcntb() / svcntw()`
(which is just a constant). https://godbolt.org/z/4o3a67s6n. We would previous
expand the svcnt into two different vscale intrinsics, CSE them in a later pass
and then fold udiv of shifts into a constant in a second instcombine.

After llvm#121386 we now introduce a cttz. This patch just adds an additional check
for vscale to the div of shift fold, allowing it to happen earlier and avoiding
the need to look through the awkward (but probably not impossible) cttz that
was introduced.

Author: davemgreen

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

@@ -1205,14 +1205,16 @@ static Value *foldIDivShl(BinaryOperator &I, InstCombiner::BuilderTy &Builder) {
 
   // If X << Y and X << Z does not overflow, then:
   // (X << Y) / (X << Z) -> (1 << Y) / (1 << Z) -> 1 << Y >> Z
-  if (match(Op0, m_Shl(m_Value(X), m_Value(Y))) &&
-      match(Op1, m_Shl(m_Specific(X), m_Value(Z)))) {
+  if ((match(Op0, m_Shl(m_Value(X), m_Value(Y))) &&
+       match(Op1, m_Shl(m_Specific(X), m_Value(Z)))) ||
+      (match(Op0, m_Shl(m_VScale(), m_Value(Y))) &&
+       match(Op1, m_Shl(m_VScale(), m_Value(Z))))) {
     auto *Shl0 = cast<OverflowingBinaryOperator>(Op0);
     auto *Shl1 = cast<OverflowingBinaryOperator>(Op1);
 
     if (IsSigned ? (Shl0->hasNoSignedWrap() && Shl1->hasNoSignedWrap())
                  : (Shl0->hasNoUnsignedWrap() && Shl1->hasNoUnsignedWrap())) {
-      Constant *One = ConstantInt::get(X->getType(), 1);
+      Constant *One = ConstantInt::get(Op0->getType(), 1);
       // Only preserve the nsw flag if dividend has nsw
       // or divisor has nsw and operator is sdiv.
       Value *Dividend = Builder.CreateShl(

llvm/test/Transforms/InstCombine/AArch64/sve-intrinsic-opts-counting-elems.ll

@@ -242,13 +242,7 @@ define i64 @cntd_all() {
 
 define i64 @udiv() vscale_range(1, 16) {
 ; CHECK-LABEL: @udiv(
-; CHECK-NEXT:    [[TMP2:%.*]] = call i64 @llvm.vscale.i64()
-; CHECK-NEXT:    [[B:%.*]] = shl nuw nsw i64 [[TMP2]], 4
-; CHECK-NEXT:    [[TMP3:%.*]] = call i64 @llvm.vscale.i64()
-; CHECK-NEXT:    [[B1:%.*]] = shl nuw nsw i64 [[TMP3]], 2
-; CHECK-NEXT:    [[TMP4:%.*]] = call range(i64 2, 65) i64 @llvm.cttz.i64(i64 [[B1]], i1 true)
-; CHECK-NEXT:    [[C:%.*]] = lshr i64 [[B]], [[TMP4]]
-; CHECK-NEXT:    ret i64 [[C]]
+; CHECK-NEXT:    ret i64 4
 ;
   %a = call i64 @llvm.aarch64.sve.cntb(i32 31)
   %b = call i64 @llvm.aarch64.sve.cntw(i32 31)

llvm/test/Transforms/InstCombine/div-shift.ll

@@ -1403,13 +1403,7 @@ start:
 define i32 @udiv_shl_pair_const_vscale() vscale_range(1, 16) {
 ; CHECK-LABEL: @udiv_shl_pair_const_vscale(
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[A:%.*]] = call i32 @llvm.vscale.i32()
-; CHECK-NEXT:    [[B:%.*]] = call i32 @llvm.vscale.i32()
-; CHECK-NEXT:    [[LHS:%.*]] = shl nuw nsw i32 [[A]], 2
-; CHECK-NEXT:    [[RHS:%.*]] = shl nuw nsw i32 [[B]], 1
-; CHECK-NEXT:    [[TMP0:%.*]] = call range(i32 1, 33) i32 @llvm.cttz.i32(i32 [[RHS]], i1 true)
-; CHECK-NEXT:    [[DIV:%.*]] = lshr i32 [[LHS]], [[TMP0]]
-; CHECK-NEXT:    ret i32 [[DIV]]
+; CHECK-NEXT:    ret i32 2
 ;
 entry:
   %a = call i32 @llvm.vscale()