[TypePromotion] Replace Zext to Truncate for the case src bitwidth is larger

Fix: llvm/llvm-project#58843

Reviewed By: samtebbs

Differential Revision: https://reviews.llvm.org/D137613

Author: bcl5980

llvm/lib/CodeGen/TypePromotion.cpp

@@ -578,7 +578,8 @@ void IRPromoter::TruncateSinks() {
 void IRPromoter::Cleanup() {
   LLVM_DEBUG(dbgs() << "IR Promotion: Cleanup..\n");
   // Some zexts will now have become redundant, along with their trunc
-  // operands, so remove them
+  // operands, so remove them.
+  // Some zexts need to be replaced with truncate if src bitwidth is larger.
   for (auto *V : Visited) {
     if (!isa<ZExtInst>(V))
       continue;
@@ -593,6 +594,11 @@ void IRPromoter::Cleanup() {
                         << "\n");
       ReplaceAllUsersOfWith(ZExt, Src);
       continue;
+    } else if (ZExt->getSrcTy()->getScalarSizeInBits() > PromotedWidth) {
+      IRBuilder<> Builder{ZExt};
+      Value *Trunc = Builder.CreateTrunc(Src, ZExt->getDestTy());
+      ReplaceAllUsersOfWith(ZExt, Trunc);
+      continue;
     }
 
     // We've inserted a trunc for a zext sink, but we already know that the

llvm/test/Transforms/TypePromotion/AArch64/pr58843.ll

@@ -0,0 +1,20 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -mtriple=aarch64 -type-promotion -verify -S %s -o - | FileCheck %s
+target datalayout = "e-m:e-i8:8:32-i16:16:32-i64:64-i128:128-n32:64-S128"
+
+; Check the case don't crash due to zext source type bitwidth
+; larger than dest type bitwidth.
+define i1 @test(i8 %arg) {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:    [[EXT1:%.*]] = zext i8 [[ARG:%.*]] to i64
+; CHECK-NEXT:    [[TMP1:%.*]] = and i64 [[EXT1]], 7
+; CHECK-NEXT:    [[TMP2:%.*]] = trunc i64 [[TMP1]] to i32
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ne i32 [[TMP2]], 0
+; CHECK-NEXT:    ret i1 [[CMP]]
+;
+  %ext1 = zext i8 %arg to i64
+  %trunc = trunc i64 %ext1 to i3
+  %ext2 = zext i3 %trunc to i8
+  %cmp = icmp ne i8 %ext2, 0
+  ret i1 %cmp
+}