diff --git a/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp b/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp --- a/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp +++ b/llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp @@ -403,13 +403,21 @@ case TargetOpcode::G_PTRTOINT: // Fall through and handle them the same as zext/trunc. LLVM_FALLTHROUGH; + case TargetOpcode::G_ASSERT_ZEXT: case TargetOpcode::G_ZEXT: case TargetOpcode::G_TRUNC: { Register SrcReg = MI.getOperand(1).getReg(); LLT SrcTy = MRI.getType(SrcReg); - unsigned SrcBitWidth = SrcTy.isPointer() - ? DL.getIndexSizeInBits(SrcTy.getAddressSpace()) - : SrcTy.getSizeInBits(); + unsigned SrcBitWidth; + + // G_ASSERT_ZEXT stores the original bitwidth in the immediate operand. + if (Opcode == TargetOpcode::G_ASSERT_ZEXT) + SrcBitWidth = MI.getOperand(2).getImm(); + else { + SrcBitWidth = SrcTy.isPointer() + ? DL.getIndexSizeInBits(SrcTy.getAddressSpace()) + : SrcTy.getSizeInBits(); + } assert(SrcBitWidth && "SrcBitWidth can't be zero"); Known = Known.zextOrTrunc(SrcBitWidth); computeKnownBitsImpl(SrcReg, Known, DemandedElts, Depth + 1); diff --git a/llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-combiner-and-trivial-mask.mir b/llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-combiner-and-trivial-mask.mir --- a/llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-combiner-and-trivial-mask.mir +++ b/llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-combiner-and-trivial-mask.mir @@ -221,3 +221,49 @@ $w0 = COPY %and(s32) RET_ReallyLR implicit $w0 ... +--- +name: remove_and_assert_zext +legalized: true +tracksRegLiveness: true +body: | + bb.0: + liveins: $w0 + ; G_ASSERT_ZEXT communicates that only the bottom 8 bits of %x can be set. + ; So, the G_AND can be removed. + + ; CHECK-LABEL: name: remove_and_assert_zext + ; CHECK: liveins: $w0 + ; CHECK: %x:_(s32) = COPY $w0 + ; CHECK: %assert_zext:_(s32) = G_ASSERT_ZEXT %x, 8 + ; CHECK: $w0 = COPY %assert_zext(s32) + ; CHECK: RET_ReallyLR implicit $w0 + %x:_(s32) = COPY $w0 + %assert_zext:_(s32) = G_ASSERT_ZEXT %x(s32), 8 + %mask:_(s32) = G_CONSTANT i32 255 + %and:_(s32) = G_AND %assert_zext(s32), %mask + $w0 = COPY %and(s32) + RET_ReallyLR implicit $w0 +... +--- +name: dont_remove_and_assert_zext_wrong_mask +legalized: true +tracksRegLiveness: true +body: | + bb.0: + liveins: $w0 + ; The mask here is for 8 bits, not 16. + + ; CHECK-LABEL: name: dont_remove_and_assert_zext + ; CHECK: liveins: $w0 + ; CHECK: %x:_(s32) = COPY $w0 + ; CHECK: %assert_zext:_(s32) = G_ASSERT_ZEXT %x, 16 + ; CHECK: %mask:_(s32) = G_CONSTANT i32 255 + ; CHECK: %and:_(s32) = G_AND %assert_zext, %mask + ; CHECK: $w0 = COPY %and(s32) + ; CHECK: RET_ReallyLR implicit $w0 + %x:_(s32) = COPY $w0 + %assert_zext:_(s32) = G_ASSERT_ZEXT %x(s32), 16 + %mask:_(s32) = G_CONSTANT i32 255 + %and:_(s32) = G_AND %assert_zext(s32), %mask + $w0 = COPY %and(s32) + RET_ReallyLR implicit $w0 diff --git a/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp b/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp --- a/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp +++ b/llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp @@ -913,3 +913,67 @@ EXPECT_TRUE(BiggerSizeRes.One.isNullValue()); EXPECT_TRUE(BiggerSizeRes.Zero.isNullValue()); } + +TEST_F(AArch64GISelMITest, TestKnownBitsAssertZext) { + StringRef MIRString = R"( + %copy:_(s64) = COPY $x0 + + %assert8:_(s64) = G_ASSERT_ZEXT %copy, 8 + %copy_assert8:_(s64) = COPY %assert8 + + %assert1:_(s64) = G_ASSERT_ZEXT %copy, 1 + %copy_assert1:_(s64) = COPY %assert1 + + %assert63:_(s64) = G_ASSERT_ZEXT %copy, 63 + %copy_assert63:_(s64) = COPY %assert63 + + %assert3:_(s64) = G_ASSERT_ZEXT %copy, 3 + %copy_assert3:_(s64) = COPY %assert3 +)"; + + setUp(MIRString); + if (!TM) + return; + + Register CopyAssert8 = Copies[Copies.size() - 4]; + Register CopyAssert1 = Copies[Copies.size() - 3]; + Register CopyAssert63 = Copies[Copies.size() - 2]; + Register CopyAssert3 = Copies[Copies.size() - 1]; + + GISelKnownBits Info(*MF); + MachineInstr *Copy; + Register SrcReg; + KnownBits Res; + + // Assert zero-extension from an 8-bit value. + Copy = MRI->getVRegDef(CopyAssert8); + SrcReg = Copy->getOperand(1).getReg(); + Res = Info.getKnownBits(SrcReg); + EXPECT_EQ(64u, Res.getBitWidth()); + EXPECT_EQ(0u, Res.One.getZExtValue()); + EXPECT_EQ(0xFFFFFFFFFFFFFF00u, Res.Zero.getZExtValue()); + + // Assert zero-extension from a 1-bit value. + Copy = MRI->getVRegDef(CopyAssert1); + SrcReg = Copy->getOperand(1).getReg(); + Res = Info.getKnownBits(SrcReg); + EXPECT_EQ(64u, Res.getBitWidth()); + EXPECT_EQ(0u, Res.One.getZExtValue()); + EXPECT_EQ(0xFFFFFFFFFFFFFFFE, Res.Zero.getZExtValue()); + + // Assert zero-extension from a 63-bit value. + Copy = MRI->getVRegDef(CopyAssert63); + SrcReg = Copy->getOperand(1).getReg(); + Res = Info.getKnownBits(SrcReg); + EXPECT_EQ(64u, Res.getBitWidth()); + EXPECT_EQ(0u, Res.One.getZExtValue()); + EXPECT_EQ(0x8000000000000000u, Res.Zero.getZExtValue()); + + // Assert zero-extension from a 3-bit value. + Copy = MRI->getVRegDef(CopyAssert3); + SrcReg = Copy->getOperand(1).getReg(); + Res = Info.getKnownBits(SrcReg); + EXPECT_EQ(64u, Res.getBitWidth()); + EXPECT_EQ(0u, Res.One.getZExtValue()); + EXPECT_EQ(0xFFFFFFFFFFFFFFF8u, Res.Zero.getZExtValue()); +}