GlobalISel: Add isKnownToBeAPowerOfTwo helper function

[paquette]

LGTM

[arsenm]

a427f15d6070fd50457c553a097e031139b40886

llvm/include/llvm/CodeGen/GlobalISel/Utils.h

 #include "llvm/CodeGen/Register.h"
 #include "llvm/Support/Alignment.h"
 #include "llvm/Support/LowLevelTypeImpl.h"
 #include <cstdint>
 
 namespace llvm {
 
 class AnalysisUsage;
+class GISelKnownBits;
 class MachineFunction;
 class MachineInstr;
 class MachineOperand;
 class MachineOptimizationRemarkEmitter;
 class MachineOptimizationRemarkMissed;
 struct MachinePointerInfo;
 class MachineRegisterInfo;
 class MCInstrDesc;

 
 Optional<APInt> ConstantFoldBinOp(unsigned Opcode, const Register Op1,
                                   const Register Op2,
                                   const MachineRegisterInfo &MRI);
 
 Optional<APInt> ConstantFoldExtOp(unsigned Opcode, const Register Op1,
                                   uint64_t Imm, const MachineRegisterInfo &MRI);
 
+/// Test if the given value is known to have exactly one bit set. This differs
+/// from computeKnownBits in that it doesn't necessarily determine which bit is
+/// set.
+bool isKnownToBeAPowerOfTwo(Register Val, const MachineRegisterInfo &MRI,
+                            GISelKnownBits *KnownBits = nullptr);
+
 /// Returns true if \p Val can be assumed to never be a NaN. If \p SNaN is true,
 /// this returns if \p Val can be assumed to never be a signaling NaN.
 bool isKnownNeverNaN(Register Val, const MachineRegisterInfo &MRI,
                      bool SNaN = false);
 
 /// Returns true if \p Val can be assumed to never be a signaling NaN.
 inline bool isKnownNeverSNaN(Register Val, const MachineRegisterInfo &MRI) {
   return isKnownNeverNaN(Val, MRI, true);

llvm/lib/CodeGen/GlobalISel/Utils.cpp

 //===- llvm/CodeGen/GlobalISel/Utils.cpp -------------------------*- C++ -*-==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 /// \file This file implements the utility functions used by the GlobalISel
 /// pipeline.
 //===----------------------------------------------------------------------===//
 
 #include "llvm/CodeGen/GlobalISel/Utils.h"
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/CodeGen/GlobalISel/GISelChangeObserver.h"
+#include "llvm/CodeGen/GlobalISel/GISelKnownBits.h"
 #include "llvm/CodeGen/GlobalISel/MIPatternMatch.h"
 #include "llvm/CodeGen/GlobalISel/RegisterBankInfo.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/StackProtector.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
case TargetOpcode::G_SEXT_INREG: {
       LLT Ty = MRI.getType(Op1);
       return MaybeOp1Cst->trunc(Imm).sext(Ty.getScalarSizeInBits());
     }
     }
   }
   return None;
 }
 
+bool llvm::isKnownToBeAPowerOfTwo(Register Reg, const MachineRegisterInfo &MRI,
+                                  GISelKnownBits *KB) {
+  Optional<DefinitionAndSourceRegister> DefSrcReg =
+      getDefSrcRegIgnoringCopies(Reg, MRI);
+  if (!DefSrcReg)
+    return false;
+
+  const MachineInstr &MI = *DefSrcReg->MI;
+  const LLT Ty = MRI.getType(Reg);
+
+  switch (MI.getOpcode()) {
+  case TargetOpcode::G_CONSTANT: {
+    unsigned BitWidth = Ty.getScalarSizeInBits();
+    const ConstantInt *CI = MI.getOperand(1).getCImm();
+    return CI->getValue().zextOrTrunc(BitWidth).isPowerOf2();
+  }
+  case TargetOpcode::G_SHL: {
+    // A left-shift of a constant one will have exactly one bit set because
+    // shifting the bit off the end is undefined.
+
+    // TODO: Constant splat
+    if (auto ConstLHS = getConstantVRegVal(MI.getOperand(1).getReg(), MRI)) {
+      if (*ConstLHS == 1)
+        return true;
+    }
+
+    break;
+  }
+  case TargetOpcode::G_LSHR: {
+    if (auto ConstLHS = getConstantVRegVal(MI.getOperand(1).getReg(), MRI)) {
+      if (ConstLHS->isSignMask())
+        return true;
+    }
+
+    break;
+  }
+  default:
+    break;
+  }
+
+  // TODO: Are all operands of a build vector constant powers of two?
+  if (!KB)
+    return false;
+
+  // More could be done here, though the above checks are enough
+  // to handle some common cases.
+
+  // Fall back to computeKnownBits to catch other known cases.
+  KnownBits Known = KB->getKnownBits(Reg);
+  return (Known.countMaxPopulation() == 1) && (Known.countMinPopulation() == 1);
+}
+
 void llvm::getSelectionDAGFallbackAnalysisUsage(AnalysisUsage &AU) {
   AU.addPreserved<StackProtector>();
 }
 
 static unsigned getLCMSize(unsigned OrigSize, unsigned TargetSize) {
   unsigned Mul = OrigSize * TargetSize;
   unsigned GCDSize = greatestCommonDivisor(OrigSize, TargetSize);
   return Mul / GCDSize;

llvm/unittests/CodeGen/GlobalISel/KnownBitsTest.cpp

TEST_F(AMDGPUGISelMITest, TestTargetKnownAlign) {
 
   EXPECT_EQ(Align(4), Info.computeKnownAlignment(CopyDispatchPtr));
   EXPECT_EQ(Align(4), Info.computeKnownAlignment(CopyQueuePtr));
   EXPECT_EQ(Align(4), Info.computeKnownAlignment(CopyKernargSegmentPtr));
   EXPECT_EQ(Align(4), Info.computeKnownAlignment(CopyImplicitArgPtr));
   EXPECT_EQ(Align(4), Info.computeKnownAlignment(CopyImplicitBufferPtr));
 }
 
+TEST_F(AMDGPUGISelMITest, TestIsKnownToBeAPowerOfTwo) {
+
+  StringRef MIRString = R"MIR(
+  %zero:_(s32) = G_CONSTANT i32 0
+  %one:_(s32) = G_CONSTANT i32 1
+  %two:_(s32) = G_CONSTANT i32 2
+  %three:_(s32) = G_CONSTANT i32 3
+  %five:_(s32) = G_CONSTANT i32 5
+  %copy_zero:_(s32) = COPY %zero
+  %copy_one:_(s32) = COPY %one
+  %copy_two:_(s32) = COPY %two
+  %copy_three:_(s32) = COPY %three
+
+  %trunc_two:_(s1) = G_TRUNC %two
+  %trunc_three:_(s1) = G_TRUNC %three
+  %trunc_five:_(s1) = G_TRUNC %five
+
+  %copy_trunc_two:_(s1) = COPY %trunc_two
+  %copy_trunc_three:_(s1) = COPY %trunc_three
+  %copy_trunc_five:_(s1) = COPY %trunc_five
+
+  %ptr:_(p1) = G_IMPLICIT_DEF
+  %shift_amt:_(s32) = G_LOAD %ptr :: (load 4, addrspace 1)
+
+  %shl_1:_(s32) = G_SHL %one, %shift_amt
+  %copy_shl_1:_(s32) = COPY %shl_1
+
+  %shl_2:_(s32) = G_SHL %two, %shift_amt
+  %copy_shl_2:_(s32) = COPY %shl_2
+
+  %not_sign_mask:_(s32) = G_LOAD %ptr :: (load 4, addrspace 1)
+  %sign_mask:_(s32) = G_CONSTANT i32 -2147483648
+
+  %lshr_not_sign_mask:_(s32) = G_LSHR %not_sign_mask, %shift_amt
+  %copy_lshr_not_sign_mask:_(s32) = COPY %lshr_not_sign_mask
+
+  %lshr_sign_mask:_(s32) = G_LSHR %sign_mask, %shift_amt
+  %copy_lshr_sign_mask:_(s32) = COPY %lshr_sign_mask
+
+  %or_pow2:_(s32) = G_OR %zero, %two
+  %copy_or_pow2:_(s32) = COPY %or_pow2
+
+)MIR";
+  setUp(MIRString);
+  if (!TM)
+    return;
+
+  GISelKnownBits KB(*MF);
+
+  Register CopyZero = Copies[Copies.size() - 12];
+  Register CopyOne = Copies[Copies.size() - 11];
+  Register CopyTwo = Copies[Copies.size() - 10];
+  Register CopyThree = Copies[Copies.size() - 9];
+  Register CopyTruncTwo = Copies[Copies.size() - 8];
+  Register CopyTruncThree = Copies[Copies.size() - 7];
+  Register CopyTruncFive = Copies[Copies.size() - 6];
+
+  Register CopyShl1 = Copies[Copies.size() - 5];
+  Register CopyShl2 = Copies[Copies.size() - 4];
+
+  Register CopyLShrNotSignMask = Copies[Copies.size() - 3];
+  Register CopyLShrSignMask = Copies[Copies.size() - 2];
+  Register CopyOrPow2 = Copies[Copies.size() - 1];
+
+  EXPECT_FALSE(isKnownToBeAPowerOfTwo(CopyZero, *MRI, &KB));
+  EXPECT_TRUE(isKnownToBeAPowerOfTwo(CopyOne, *MRI, &KB));
+  EXPECT_TRUE(isKnownToBeAPowerOfTwo(CopyTwo, *MRI, &KB));
+  EXPECT_FALSE(isKnownToBeAPowerOfTwo(CopyThree, *MRI, &KB));
+
+  EXPECT_FALSE(isKnownToBeAPowerOfTwo(CopyTruncTwo, *MRI, &KB));
+  EXPECT_TRUE(isKnownToBeAPowerOfTwo(CopyTruncThree, *MRI, &KB));
+  EXPECT_TRUE(isKnownToBeAPowerOfTwo(CopyTruncFive, *MRI, &KB));
+
+  EXPECT_TRUE(isKnownToBeAPowerOfTwo(CopyShl1, *MRI, &KB));
+  EXPECT_FALSE(isKnownToBeAPowerOfTwo(CopyShl2, *MRI, &KB));
+
+  EXPECT_FALSE(isKnownToBeAPowerOfTwo(CopyLShrNotSignMask, *MRI, &KB));
+  EXPECT_TRUE(isKnownToBeAPowerOfTwo(CopyLShrSignMask, *MRI, &KB));
+  EXPECT_TRUE(isKnownToBeAPowerOfTwo(CopyOrPow2, *MRI, &KB));
+}
+
 TEST_F(AArch64GISelMITest, TestMetadata) {
   StringRef MIRString = "  %imp:_(p0) = G_IMPLICIT_DEF\n"
                         "  %load:_(s8) = G_LOAD %imp(p0) :: (load 1)\n"
                         "  %ext:_(s32) = G_ZEXT %load(s8)\n"
                         "  %cst:_(s32) = G_CONSTANT i32 1\n"
                         "  %and:_(s32) = G_AND %ext, %cst\n"
                         "  %copy:_(s32) = COPY %and(s32)\n";
   setUp(MIRString);