[llvm] [GloalISel] Fold G_CTTZ if possible (PR #86224)

[Shilei Tian via llvm-commits]

https://github.com/shiltian created https://github.com/llvm/llvm-project/pull/86224

This patch tries to fold `G_CTTZ` if possible.


>From aa2d053057e5c07f223e0d4625e2b2a0df269222 Mon Sep 17 00:00:00 2001
From: Shilei Tian <i at tianshilei.me>
Date: Thu, 21 Mar 2024 20:47:33 -0400
Subject: [PATCH] [GloalISel] Fold G_CTTZ if possible

This patch tries to fold `G_CTTZ` if possible.
---
 llvm/include/llvm/CodeGen/GlobalISel/Utils.h  |  5 +++
 llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp | 16 ++++++++++
 llvm/lib/CodeGen/GlobalISel/Utils.cpp         | 31 +++++++++++++++++++
 3 files changed, 52 insertions(+)

diff --git a/llvm/include/llvm/CodeGen/GlobalISel/Utils.h b/llvm/include/llvm/CodeGen/GlobalISel/Utils.h
index f8900f3434ccaa..bcfc9337598223 100644
--- a/llvm/include/llvm/CodeGen/GlobalISel/Utils.h
+++ b/llvm/include/llvm/CodeGen/GlobalISel/Utils.h
@@ -313,6 +313,11 @@ std::optional<APFloat> ConstantFoldIntToFloat(unsigned Opcode, LLT DstTy,
 std::optional<SmallVector<unsigned>>
 ConstantFoldCTLZ(Register Src, const MachineRegisterInfo &MRI);
 
+/// Tries to constant fold a G_CTTZ operation on \p Src. If \p Src is a vector
+/// then it tries to do an element-wise constant fold.
+std::optional<SmallVector<unsigned>>
+ConstantFoldCTTZ(Register Src, 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.
diff --git a/llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp b/llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp
index 1869e0d41a51f6..4901b0e9d203dd 100644
--- a/llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/CSEMIRBuilder.cpp
@@ -272,6 +272,22 @@ MachineInstrBuilder CSEMIRBuilder::buildInstr(unsigned Opc,
           buildConstant(VecTy.getScalarType(), Cst).getReg(0));
     return buildBuildVector(DstOps[0], ConstantRegs);
   }
+  case TargetOpcode::G_CTTZ: {
+    assert(SrcOps.size() == 1 && "Expected one source");
+    assert(DstOps.size() == 1 && "Expected one dest");
+    auto MaybeCsts = ConstantFoldCTTZ(SrcOps[0].getReg(), *getMRI());
+    if (!MaybeCsts)
+      break;
+    if (MaybeCsts->size() == 1)
+      return buildConstant(DstOps[0], (*MaybeCsts)[0]);
+    // This was a vector constant. Build a G_BUILD_VECTOR for them.
+    SmallVector<Register> ConstantRegs;
+    LLT VecTy = DstOps[0].getLLTTy(*getMRI());
+    for (unsigned Cst : *MaybeCsts)
+      ConstantRegs.emplace_back(
+          buildConstant(VecTy.getScalarType(), Cst).getReg(0));
+    return buildBuildVector(DstOps[0], ConstantRegs);
+  }
   }
   bool CanCopy = checkCopyToDefsPossible(DstOps);
   if (!canPerformCSEForOpc(Opc))
diff --git a/llvm/lib/CodeGen/GlobalISel/Utils.cpp b/llvm/lib/CodeGen/GlobalISel/Utils.cpp
index a9fa73b60a097f..7a58b752f2ecc7 100644
--- a/llvm/lib/CodeGen/GlobalISel/Utils.cpp
+++ b/llvm/lib/CodeGen/GlobalISel/Utils.cpp
@@ -996,6 +996,37 @@ llvm::ConstantFoldCTLZ(Register Src, const MachineRegisterInfo &MRI) {
   return std::nullopt;
 }
 
+std::optional<SmallVector<unsigned>>
+llvm::ConstantFoldCTTZ(Register Src, const MachineRegisterInfo &MRI) {
+  LLT Ty = MRI.getType(Src);
+  SmallVector<unsigned> FoldedCTTZs;
+  auto tryFoldScalar = [&](Register R) -> std::optional<unsigned> {
+    auto MaybeCst = getIConstantVRegVal(R, MRI);
+    if (!MaybeCst)
+      return std::nullopt;
+    return MaybeCst->countTrailingZeros();
+  };
+  if (Ty.isVector()) {
+    // Try to constant fold each element.
+    auto *BV = getOpcodeDef<GBuildVector>(Src, MRI);
+    if (!BV)
+      return std::nullopt;
+    for (unsigned SrcIdx = 0; SrcIdx < BV->getNumSources(); ++SrcIdx) {
+      if (auto MaybeFold = tryFoldScalar(BV->getSourceReg(SrcIdx))) {
+        FoldedCTTZs.emplace_back(*MaybeFold);
+        continue;
+      }
+      return std::nullopt;
+    }
+    return FoldedCTTZs;
+  }
+  if (auto MaybeCst = tryFoldScalar(Src)) {
+    FoldedCTTZs.emplace_back(*MaybeCst);
+    return FoldedCTTZs;
+  }
+  return std::nullopt;
+}
+
 bool llvm::isKnownToBeAPowerOfTwo(Register Reg, const MachineRegisterInfo &MRI,
                                   GISelKnownBits *KB) {
   std::optional<DefinitionAndSourceRegister> DefSrcReg =