[AggressiveInstCombine] Recognize table-based ctz implementation and enable it for AARCH64 at -O3

[gsocshubham]

Summary

This patch is improved version of patch https://reviews.llvm.org/D113291 submitted by @djtodoro
Credit goes to @djtodoro for his submission.

I have added fix in AArch64 backend and I have followed an approach suggested by @craig.topper here -
https://github.com/llvm/llvm-project/issues/45779#issuecomment-1022851172

It generates below assembly -

rbit    w8, w0
clz     w8, w8
and     w0, w8, #0x1f
ret

which is similar assembly emitted by GCC AArch64 - https://gcc.gnu.org/bugzilla/show_bug.cgi?id=90838

Patch fixes - https://github.com/llvm/llvm-project/issues/45779

[xbolva00]

How hard is to implement support for x86?

[lebedev.ri]

Please split the lowering fix from the intrinsic recognition.

[craig.topper]

llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp

442

getPointerElementType() will be deleted when we move to OpaquePtr. Use LI->getType()?

450

GEP->getSourceElementType() I think.

468

Constant->ConstantInt? And isZeroValue->isZero

478

What does it take to address this FIXME so this will work on other targets?

llvm/test/Transforms/AggressiveInstCombine/AArch64/lower-table-based-ctz.ll

1

Why are there less tests than the other patch?

[dmgreen]

The changes in AArch64InstrInfo.td need to be their own patch in a separate commit (with their own tests). It's hard to see how they are correct as-is, but if you put up a new patch for them we can sort out what we need to do there.

Patches should also be uploaded with context, as per the llvm docs: https://llvm.org/docs/Phabricator.html#requesting-a-review-via-the-web-interface

[djtodoro]

Thanks for this.

This should be implemented as a patch on top of the https://reviews.llvm.org/D113291 (as a child of the D113291).

[Wilco1]

llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp

501

This doesn't check that ZeroTableElem is either zero or InputBits (eg. if table[0] = 1 in ctz2 below, it has to be rejected). Also does it avoid emitting the AND if it is InputBits?

int ctz2 (unsigned x)
{

const int u = 0;
static short table[64] =
  {
    32, 0, 1,12, 2, 6, u,13, 3, u, 7, u, u, u, u,14,
    10, 4, u, u, 8, u, u,25, u, u, u, u, u,21,27,15,
    31,11, 5, u, u, u, u, u, 9, u, u,24, u, u,20,26,
    30, u, u, u, u,23, u,19,29, u,22,18,28,17,16, u
  };

x = (x & -x) * 0x0450FBAF;
return table[x >> 26];

}

[gsocshubham]

In D119010#3297057, @xbolva00 wrote:

How hard is to implement support for x86?

It will be similar as lowering fix - https://reviews.llvm.org/D120462 once intrinsic recognition of cttz is merged.

[gsocshubham]

In D119010#3297125, @lebedev.ri wrote:

Please split the lowering fix from the intrinsic recognition.

Sure. I have created a different revision here - https://reviews.llvm.org/D120462
but it still needs some work.

[gsocshubham]

In D119010#3297623, @djtodoro wrote:

Thanks for this.

This should be implemented as a patch on top of the https://reviews.llvm.org/D113291 (as a child of the D113291).

Hello @djtodoro

What is the process to submit my patch based on https://reviews.llvm.org/D113291 ?
While submitting current patch, I could not find a option to make my patch as child patch of yours.

Can you please point to the process of doing it?

llvm/test/Transforms/AggressiveInstCombine/AArch64/lower-table-based-ctz.ll

1

This should go once review on this patch - https://reviews.llvm.org/D120462 is ready.

[djtodoro]

In D119010#3348354, @gsocshubham wrote:

In D119010#3297623, @djtodoro wrote:

Thanks for this.

This should be implemented as a patch on top of the https://reviews.llvm.org/D113291 (as a child of the D113291).

Hello @djtodoro

What is the process to submit my patch based on https://reviews.llvm.org/D113291 ?
While submitting current patch, I could not find a option to make my patch as child patch of yours.

Can you please point to the process of doing it?

Edit Related Revisions (Bellow "Download Raw DIff") -> Edit Child Revisions

[lebedev.ri]

This review may be stuck/dead, consider abandoning if no longer relevant.
Removing myself as reviewer in attempt to clean dashboard.

llvm/include/llvm/Analysis/TargetTransformInfo.h

   /// considered as "Slow".
   enum PopcntSupportKind { PSK_Software, PSK_SlowHardware, PSK_FastHardware };
 
+  /// Return true if the target would benefit from recognizing a table-based
+  /// implementation of cttz.
+  bool preferCTTZLowering() const;
+
   /// Return true if the specified immediate is legal add immediate, that
   /// is the target has add instructions which can add a register with the
   /// immediate without having to materialize the immediate into a register.

       APInt &UndefElts2, APInt &UndefElts3,
       std::function<void(Instruction *, unsigned, APInt, APInt &)>
           SimplifyAndSetOp) = 0;
+  virtual bool preferCTTZLowering() = 0;
   virtual bool isLegalAddImmediate(int64_t Imm) = 0;
   virtual bool isLegalICmpImmediate(int64_t Imm) = 0;
   virtual bool isLegalAddressingMode(Type *Ty, GlobalValue *BaseGV,

         IC, II, DemandedElts, UndefElts, UndefElts2, UndefElts3,
         SimplifyAndSetOp);
   }
+  bool preferCTTZLowering() override {
+    return Impl.preferCTTZLowering();
+  }
   bool isLegalAddImmediate(int64_t Imm) override {
     return Impl.isLegalAddImmediate(Imm);
   }

llvm/include/llvm/Analysis/TargetTransformInfoImpl.h

   void getPeelingPreferences(Loop *, ScalarEvolution &,
                              TTI::PeelingPreferences &) const {}
 
+  bool preferCTTZLowering() const { return false; }
+
   bool isLegalAddImmediate(int64_t Imm) const { return false; }
 
   bool isLegalICmpImmediate(int64_t Imm) const { return false; }

llvm/include/llvm/CodeGen/BasicTTIImpl.h

     return nullptr;
   }
 
+  bool preferCTTZLowering() {
+    return getTLI()->preferCTTZLowering();
+  }
+
   bool isLegalAddImmediate(int64_t imm) {
     return getTLI()->isLegalAddImmediate(imm);
   }

llvm/include/llvm/CodeGen/TargetLowering.h

     return true;
   }
 
+  /// Return true if the target would benefit from recognizing a table-based
+  /// implementation of cttz.
+  virtual bool preferCTTZLowering() const {
+    return false;
+  }
+
   /// Return true if the specified immediate is legal add immediate, that is the
   /// target has add instructions which can add a register with the immediate
   /// without having to materialize the immediate into a register.

llvm/lib/Analysis/TargetTransformInfo.cpp

   return TTIImpl->getPeelingPreferences(L, SE, PP);
 }
 
+bool TargetTransformInfo::preferCTTZLowering() const {
+  return TTIImpl->preferCTTZLowering();
+}
+
 bool TargetTransformInfo::isLegalAddImmediate(int64_t Imm) const {
   return TTIImpl->isLegalAddImmediate(Imm);
 }

llvm/lib/Target/AArch64/AArch64ISelLowering.h

   bool lowerInterleavedStore(StoreInst *SI, ShuffleVectorInst *SVI,
                              unsigned Factor) const override;
 
+  bool preferCTTZLowering() const override;
   bool isLegalAddImmediate(int64_t) const override;
   bool isLegalICmpImmediate(int64_t) const override;
 

llvm/lib/Target/AArch64/AArch64ISelLowering.cpp

   return LLT();
 }
 
+bool AArch64TargetLowering::preferCTTZLowering() const {
+  return true;
+}
+
 // 12-bit optionally shifted immediates are legal for adds.
 bool AArch64TargetLowering::isLegalAddImmediate(int64_t Immed) const {
   if (Immed == std::numeric_limits<int64_t>::min()) {

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 def  REV16Xr : OneXRegData<0b001, "rev16", null_frag>;
 
 def : Pat<(cttz GPR32:$Rn),
-          (CLZWr (RBITWr GPR32:$Rn))>;
+          (ANDWri (CLZWr (RBITWr GPR32:$Rn)), (i32 4))>;
 def : Pat<(cttz GPR64:$Rn),
-          (CLZXr (RBITXr GPR64:$Rn))>;
+          (ANDXri (CLZXr (RBITXr GPR64:$Rn)), (i64 5))>;
 def : Pat<(ctlz (or (shl (xor (sra GPR32:$Rn, (i64 31)), GPR32:$Rn), (i64 1)),
                 (i32 1))),
           (CLSWr GPR32:$Rn)>;

llvm/lib/Transforms/AggressiveInstCombine/AggressiveInstCombine.cpp

 #include "llvm/Analysis/BasicAliasAnalysis.h"
 #include "llvm/Analysis/GlobalsModRef.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"

   return false;
 }
 
+static bool isCTTZTable(const ConstantDataArray &Table, uint64_t Mul,
+                        uint64_t Shift, uint64_t InputBits) {
+  unsigned Length = Table.getNumElements();
+  uint64_t IntWidth = InputBits / CHAR_BIT;
+  if (Length < InputBits || Length > InputBits * 2)
+    return false;
+
+  APInt Mask(InputBits, ((IntWidth << (InputBits - Shift)) - 1) << Shift);
+  unsigned Matched = 0;
+
+  for (unsigned i = 0; i < Length; i++) {
+    uint64_t Element = Table.getElementAsInteger(i);
+    if (Element < InputBits &&
+        (((Mul << Element) & Mask.getZExtValue()) >> Shift) == i)
+      Matched++;
+  }
+
+  return Matched == InputBits;
+}
+
+// Try to recognize table-based ctz implementation.
+// E.g., an example in C (for more cases please see the llvm/tests):
+// int f(unsigned x) {
+//    static const char table[32] =
+//      {0, 1, 28, 2, 29, 14, 24, 3, 30,
+//       22, 20, 15, 25, 17, 4, 8, 31, 27,
+//       13, 23, 21, 19, 16, 7, 26, 12, 18, 6, 11, 5, 10, 9};
+//    return table[((unsigned)((x & -x) * 0x077CB531U)) >> 27];
+// }
+// this can be lowered to `cttz` instruction.
+// There is also a special case when the element is 0.
+//
+// Here are some examples or IR for AARCH64 target:
+//
+// CASE 1:
+// %sub = sub i32 0, %x
+// %and = and i32 %sub, %x
+// %mul = mul i32 %and, 125613361
+// %shr = lshr i32 %mul, 27
+// %idxprom = zext i32 %shr to i64
+// %arrayidx = getelementptr inbounds [32 x i8], [32 x i8]* @ctz1.table, i64 0,
+// i64 %idxprom %0 = load i8, i8* %arrayidx, align 1, !tbaa !8
+//
+// CASE 2:
+// %sub = sub i32 0, %x
+// %and = and i32 %sub, %x
+// %mul = mul i32 %and, 72416175
+// %shr = lshr i32 %mul, 26
+// %idxprom = zext i32 %shr to i64
+// %arrayidx = getelementptr inbounds [64 x i16], [64 x i16]* @ctz2.table, i64
+// 0, i64 %idxprom %0 = load i16, i16* %arrayidx, align 2, !tbaa !8
+//
+// CASE 3:
+// %sub = sub i32 0, %x
+// %and = and i32 %sub, %x
+// %mul = mul i32 %and, 81224991
+// %shr = lshr i32 %mul, 27
+// %idxprom = zext i32 %shr to i64
+// %arrayidx = getelementptr inbounds [32 x i32], [32 x i32]* @ctz3.table, i64
+// 0, i64 %idxprom %0 = load i32, i32* %arrayidx, align 4, !tbaa !8
+//
+// CASE 4:
+// %sub = sub i64 0, %x
+// %and = and i64 %sub, %x
+// %mul = mul i64 %and, 283881067100198605
+// %shr = lshr i64 %mul, 58
+// %arrayidx = getelementptr inbounds [64 x i8], [64 x i8]* @table, i64 0, i64
+// %shr %0 = load i8, i8* %arrayidx, align 1, !tbaa !8
+//
+// All this can be lowered to @llvm.cttz.i32/64 intrinsic.
+static bool tryToRecognizeTableBasedCttz(Instruction &I) {
+  LoadInst *LI = dyn_cast<LoadInst>(&I);
+  if (!LI)
+    return false;
+
+  Type *ElType = LI->getPointerOperandType()->getPointerElementType();
+  if (!ElType->isIntegerTy())
+    return false;
+
+  GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(LI->getPointerOperand());
+  if (!GEP || !GEP->isInBounds() || GEP->getNumIndices() != 2)
+    return false;
+
+  Type *GEPPointeeType = GEP->getPointerOperandType()->getPointerElementType();
+  if (!GEPPointeeType->isArrayTy())
+    return false;
+
+  uint64_t ArraySize = GEPPointeeType->getArrayNumElements();
+  if (ArraySize != 32 && ArraySize != 64)
+    return false;
+
+  User *GEPUser = dyn_cast<User>(GEP->getPointerOperand());
+  if (!GEPUser)
+    return false;
+
+  ConstantDataArray *ConstData =
+      dyn_cast<ConstantDataArray>(GEPUser->getOperand(0));
+  if (!ConstData)
+    return false;
+
+  Value *Idx1 = GEP->idx_begin()->get();
+  Constant *Zero = dyn_cast<Constant>(Idx1);
+  if (!Zero || !Zero->isZeroValue())
+    return false;
+
+  Value *Idx2 = std::next(GEP->idx_begin())->get();
+
+  bool ConstIsWide = !match(Idx2, m_ZExt(m_Value()));
+
+  Value *X1;
+  uint64_t MulConst, ShiftConst;
+  // FIXME: AArch64 has i64 type for the GEP index, so this match will
+  // probably fail for other targets.
+  if (!match(Idx2,
+             m_ZExtOrSelf(m_LShr(
+                 m_ZExtOrSelf(m_Mul(m_c_And(m_Neg(m_Value(X1)), m_Deferred(X1)),
+                                    m_ConstantInt(MulConst))),
+                 m_ConstantInt(ShiftConst)))))
+    return false;
+
+  unsigned InputBits = ConstIsWide ? 64 : 32;
+
+  // Shift should extract top 5..7 bits.
+  if (ShiftConst < InputBits - 7 || ShiftConst > InputBits - 5)
+    return false;
+
+  Type *XType = X1->getType();
+  if (!XType->isIntegerTy(InputBits))
+    return false;
+
+  if (!isCTTZTable(*ConstData, MulConst, ShiftConst, InputBits))
+    return false;
+
+  auto ZeroTableElem = ConstData->getElementAsInteger(0);
+  bool DefinedForZero = ZeroTableElem == InputBits;
+
+  IRBuilder<> B(LI);
+  ConstantInt *BoolConst = B.getInt1(!DefinedForZero);
+  auto Cttz = B.CreateIntrinsic(Intrinsic::cttz, {XType}, {X1, BoolConst});
+  Value *ZExtOrTrunc = nullptr;
+  ZExtOrTrunc = B.CreateZExtOrTrunc(Cttz, ElType);
+  LI->replaceAllUsesWith(ZExtOrTrunc);
+
+  return true;
+}
+
 /// This is the entry point for folds that could be implemented in regular
 /// InstCombine, but they are separated because they are not expected to
 /// occur frequently and/or have more than a constant-length pattern match.
-static bool foldUnusualPatterns(Function &F, DominatorTree &DT) {
+static bool foldUnusualPatterns(Function &F, DominatorTree &DT,
+    TargetTransformInfo &TTI) {
   bool MadeChange = false;
   for (BasicBlock &BB : F) {
     // Ignore unreachable basic blocks.

       MadeChange |= foldAnyOrAllBitsSet(I);
       MadeChange |= foldGuardedFunnelShift(I, DT);
       MadeChange |= tryToRecognizePopCount(I);
+      if (TTI.preferCTTZLowering())
+        MadeChange |= tryToRecognizeTableBasedCttz(I);
     }
   }
 

 /// This is the entry point for all transforms. Pass manager differences are
 /// handled in the callers of this function.
 static bool runImpl(Function &F, AssumptionCache &AC, TargetLibraryInfo &TLI,
-                    DominatorTree &DT) {
+                    DominatorTree &DT, TargetTransformInfo &TTI) {
   bool MadeChange = false;
   const DataLayout &DL = F.getParent()->getDataLayout();
   TruncInstCombine TIC(AC, TLI, DL, DT);
   MadeChange |= TIC.run(F);
-  MadeChange |= foldUnusualPatterns(F, DT);
+  MadeChange |= foldUnusualPatterns(F, DT, TTI);
   return MadeChange;
 }
 

   AU.setPreservesCFG();
   AU.addRequired<AssumptionCacheTracker>();
   AU.addRequired<DominatorTreeWrapperPass>();
+  AU.addRequired<TargetTransformInfoWrapperPass>();
   AU.addRequired<TargetLibraryInfoWrapperPass>();
   AU.addPreserved<AAResultsWrapperPass>();
   AU.addPreserved<BasicAAWrapperPass>();

   auto &AC = getAnalysis<AssumptionCacheTracker>().getAssumptionCache(F);
   auto &TLI = getAnalysis<TargetLibraryInfoWrapperPass>().getTLI(F);
   auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
-  return runImpl(F, AC, TLI, DT);
+  auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
+  return runImpl(F, AC, TLI, DT, TTI);
 }
 
 PreservedAnalyses AggressiveInstCombinePass::run(Function &F,

   auto &AC = AM.getResult<AssumptionAnalysis>(F);
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
   auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
-  if (!runImpl(F, AC, TLI, DT)) {
+  auto &TTI = AM.getResult<TargetIRAnalysis>(F);
+  if (!runImpl(F, AC, TLI, DT, TTI)) {
     // No changes, all analyses are preserved.
     return PreservedAnalyses::all();
   }

                       "Combine pattern based expressions", false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)
 INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(TargetTransformInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_END(AggressiveInstCombinerLegacyPass, "aggressive-instcombine",
                     "Combine pattern based expressions", false, false)

llvm/test/Transforms/AggressiveInstCombine/AArch64/lit.local.cfg

+if not 'AArch64' in config.root.targets:
+    config.unsupported = True

llvm/test/Transforms/AggressiveInstCombine/AArch64/lower-table-based-ctz.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -aggressive-instcombine -mtriple aarch64-linux-gnu -S < %s | FileCheck %s
+
[e-mail] = internal unnamed_addr constant [32 x i8] c"\00\01\1C\02\1D\0E\18\03\1E\16\14\0F\19\11\04\08\1F\1B\0D\17\15\13\10\07\1A\0C\12\06\0B\05\0A\09", align 1
+
+define i32 @f(i32 %x) {
+; CHECK-LABEL: @f(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP2:%.*]] = call i32 @llvm.cttz.i32(i32 [[X:%.*]], i1 true)
+; CHECK-NEXT:    [[TMP3:%.*]] = trunc i32 [[TMP2]] to i8
+; CHECK-NEXT:    [[CONV:%.*]] = zext i8 [[TMP3]] to i32
+; CHECK-NEXT:    ret i32 [[CONV]]
+;
+entry:
+  %sub = sub i32 0, %x
+  %and = and i32 %sub, %x
+  %mul = mul i32 %and, 125613361
+  %shr = lshr i32 %mul, 27
+  %idxprom = zext i32 %shr to i64
+  %arrayidx = getelementptr inbounds [32 x i8], [32 x i8]* @f.table, i64 0, i64 %idxprom
+  %0 = load i8, i8* %arrayidx, align 1
+  %conv = zext i8 %0 to i32
+  ret i32 %conv
+}