[GlobalISel] Implement computeKnownBits for G_ASSERT_ZEXT

[paquette]

Summary

It's the same as the ZEXT/TRUNC case, except SrcBitWidth is given by the immediate operand.

Update KnownBitsTest.cpp and a MIR test for a concrete example.

[arsenm]

llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp

416

This should get braces

llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp

924

I think testing the boundary conditions 1, 63 are more useful than a bunch of powers of 2

[paquette]

• Add missing braces
• Test boundary values

[aemerson]

LGTM.

llvm/lib/CodeGen/GlobalISel/GISelKnownBits.cpp

computeKnownBitsImpl(MI.getOperand(2).getReg(), RHSKnown, DemandedElts,
                          Depth + 1);
     Known = KnownBits::shl(LHSKnown, RHSKnown);
     break;
   }
   case TargetOpcode::G_INTTOPTR:
   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()
     assert(SrcBitWidth && "SrcBitWidth can't be zero");
     Known = Known.zextOrTrunc(SrcBitWidth);
     computeKnownBitsImpl(SrcReg, Known, DemandedElts, Depth + 1);
     Known = Known.zextOrTrunc(BitWidth);
     if (BitWidth > SrcBitWidth)
       Known.Zero.setBitsFrom(SrcBitWidth);
     break;
   }

llvm/test/CodeGen/AArch64/GlobalISel/postlegalizer-combiner-and-trivial-mask.mir

bb.0:
     %x:_(s32) = COPY $w0
     %y:_(s32) = COPY $w1
     %z:_(s32) = COPY $w2
     %one:_(s32) = G_CONSTANT i32 32
     %and:_(s32) = G_AND %z(s32), %one
     $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

llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp

)";
 
   // We don't know what the result of the shift is, but we should not crash
   EXPECT_TRUE(EqSizeRes.One.isNullValue());
   EXPECT_TRUE(EqSizeRes.Zero.isNullValue());
 
   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
+
+   %assert16:_(s64) = G_ASSERT_ZEXT %copy, 16
+   %copy_assert16:_(s64) = COPY %assert16
+
+   %assert32:_(s64) = G_ASSERT_ZEXT %copy, 32
+   %copy_assert32:_(s64) = COPY %assert32
+
+   %assert3:_(s64) = G_ASSERT_ZEXT %copy, 3
+   %copy_assert3:_(s64) = COPY %assert3
+)";
+
+  setUp(MIRString);
+  if (!TM)
+    return;
+
+  Register CopyAssert8 = Copies[Copies.size() - 4];
+  Register CopyAssert16 = Copies[Copies.size() - 3];
+  Register CopyAssert32 = 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 16-bit value.
+  Copy = MRI->getVRegDef(CopyAssert16);
+  SrcReg = Copy->getOperand(1).getReg();
+  Res = Info.getKnownBits(SrcReg);
+  EXPECT_EQ(64u, Res.getBitWidth());
+  EXPECT_EQ(0u, Res.One.getZExtValue());
+  EXPECT_EQ(0xFFFFFFFFFFFF0000u, Res.Zero.getZExtValue());
+
+  // Assert zero-extension from a 32-bit value.
+  Copy = MRI->getVRegDef(CopyAssert32);
+  SrcReg = Copy->getOperand(1).getReg();
+  Res = Info.getKnownBits(SrcReg);
+  EXPECT_EQ(64u, Res.getBitWidth());
+  EXPECT_EQ(0u, Res.One.getZExtValue());
+  EXPECT_EQ(0xFFFFFFFF00000000u, 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());
+}