[InstCombine] Check if zext src and trunc dst have same types

[datta.nagraj]

Summary

Apply the opt only if zext src and trunc dst have same type and
do it only for datatypes of width more than 5, as for width less
than 6, this opt fails:

https://alive2.llvm.org/ce/z/5ThC64

The opt to be applied when width is less than 6 is shown here:

https://alive2.llvm.org/ce/z/ZtZm4w

llvm/lib/Transforms/InstCombine/InstCombineCasts.cpp

if (SrcWidth % DestWidth == 0) {
       auto *BitCastTo =
           VectorType::get(DestTy, BitCastNumElts, VecElts.isScalable());
       Value *BitCast = Builder.CreateBitCast(VecOp, BitCastTo);
       return ExtractElementInst::Create(BitCast, Builder.getInt32(NewIdx));
     }
   }
 
   // trunc (ctlz_i32(zext(A), B) --> add(ctlz_i16(A, B), C)
-  if (match(Src, m_OneUse(m_Intrinsic<Intrinsic::ctlz>(
-                     m_OneUse(m_ZExt(m_Value(A))), m_Value(B))))) {
-    Value *WidthDiff =
-        ConstantInt::get(A->getType(), Src->getType()->getScalarSizeInBits() -
-                                           A->getType()->getScalarSizeInBits());
+  if (match(Src, m_OneUse(m_Intrinsic<Intrinsic::ctlz>(m_ZExt(m_Value(A)),
+                                                       m_Value(B))))) {
+    unsigned AWidth = A->getType()->getScalarSizeInBits();
 
+    if (AWidth == DestWidth && AWidth > 5) {
+      Value *WidthDiff = ConstantInt::get(A->getType(), SrcWidth - AWidth);
       Value *NarrowCtlz =
           Builder.CreateIntrinsic(Intrinsic::ctlz, {Trunc.getType()}, {A, B});
       return BinaryOperator::CreateAdd(NarrowCtlz, WidthDiff);
     }
+  }
   return nullptr;
 }
 
 Instruction *InstCombinerImpl::transformZExtICmp(ICmpInst *Cmp, ZExtInst &Zext,
                                                  bool DoTransform) {
   // If we are just checking for a icmp eq of a single bit and zext'ing it
   // to an integer, then shift the bit to the appropriate place and then
   // cast to integer to avoid the comparison.

llvm/test/Transforms/InstCombine/zext-ctlz-trunc-to-ctlz-add.ll

;
   ret <2 x i17> %zz
 }
 
 ; Multiple uses of zext for which the opt is disabled
 
 define <vscale x 2 x i16> @trunc_ctlz_zext_nxv2i16_nxv2i63_multiple_uses(<vscale x 2 x i16> %x) {
 ; CHECK-LABEL: @trunc_ctlz_zext_nxv2i16_nxv2i63_multiple_uses(
 ; CHECK-NEXT:    [[Z:%.*]] = zext <vscale x 2 x i16> [[X:%.*]] to <vscale x 2 x i63>
-; CHECK-NEXT:    [[P:%.*]] = call <vscale x 2 x i63> @llvm.ctlz.nxv2i63(<vscale x 2 x i63> [[Z]], i1 true)
-; CHECK-NEXT:    [[ZZ:%.*]] = trunc <vscale x 2 x i63> [[P]] to <vscale x 2 x i16>
+; CHECK-NEXT:    [[TMP1:%.*]] = call <vscale x 2 x i16> @llvm.ctlz.nxv2i16(<vscale x 2 x i16> [[X]], i1 true)
+; CHECK-NEXT:    [[ZZ:%.*]] = add nuw nsw <vscale x 2 x i16> [[TMP1]], shufflevector (<vscale x 2 x i16> insertelement (<vscale x 2 x i16> undef, i16 47, i32 0), <vscale x 2 x i16> undef, <vscale x 2 x i32> zeroinitializer)
 ; CHECK-NEXT:    call void @use1(<vscale x 2 x i63> [[Z]])
 ; CHECK-NEXT:    ret <vscale x 2 x i16> [[ZZ]]
 ;
   %z = zext <vscale x 2 x i16> %x to <vscale x 2 x i63>
   %p = call <vscale x 2 x i63> @llvm.ctlz.nxv2i63(<vscale x 2 x i63> %z, i1 true)
   %zz = trunc <vscale x 2 x i63> %p to <vscale x 2 x i16>
   call void @use1(<vscale x 2 x i63> %z)
   ret <vscale x 2 x i16> %zz
 }
 
+; Negative case where types of x and zz don't match
+
+define i16 @trunc_ctlz_zext_i10_i32(i10 %x) {
+; CHECK-LABEL: @trunc_ctlz_zext_i10_i32(
+; CHECK-NEXT:    [[Z:%.*]] = zext i10 [[X:%.*]] to i32
+; CHECK-NEXT:    [[P:%.*]] = call i32 @llvm.ctlz.i32(i32 [[Z]], i1 false), !range [[RNG1:![0-9]+]]
+; CHECK-NEXT:    [[ZZ:%.*]] = trunc i32 [[P]] to i16
+; CHECK-NEXT:    ret i16 [[ZZ]]
+;
+  %z = zext i10 %x to i32
+  %p = call i32 @llvm.ctlz.i32(i32 %z, i1 false)
+  %zz = trunc i32 %p to i16
+  ret i16 %zz
+}
+
+; Negative case where size of x is less than 6
+
+define i5 @trunc_ctlz_zext_i5_i32(i5 %x) {
+; CHECK-LABEL: @trunc_ctlz_zext_i5_i32(
+; CHECK-NEXT:    [[Z:%.*]] = zext i5 [[X:%.*]] to i32
+; CHECK-NEXT:    [[P:%.*]] = call i32 @llvm.ctlz.i32(i32 [[Z]], i1 false), !range [[RNG2:![0-9]+]]
+; CHECK-NEXT:    [[ZZ:%.*]] = trunc i32 [[P]] to i5
+; CHECK-NEXT:    ret i5 [[ZZ]]
+;
+  %z = zext i5 %x to i32
+  %p = call i32 @llvm.ctlz.i32(i32 %z, i1 false)
+  %zz = trunc i32 %p to i5
+  ret i5 %zz
+}
+