llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUMCExpr.cpp - rust-lang/llvm-project - Git at Google

 //===- AMDGPUMCExpr.cpp - AMDGPU specific MC expression classes -----------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//

 #include "AMDGPUMCExpr.h"
 #include "GCNSubtarget.h"
 #include "Utils/AMDGPUBaseInfo.h"
 #include "llvm/IR/Function.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCAssembler.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/MC/MCValue.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Support/raw_ostream.h"
 #include <optional>

 using namespace llvm;
 using namespace llvm::AMDGPU;

 AMDGPUMCExpr::AMDGPUMCExpr(VariantKind Kind, ArrayRef<const MCExpr *> Args,
                            MCContext &Ctx)
     : Kind(Kind), Ctx(Ctx) {
   assert(Args.size() >= 1 && "Needs a minimum of one expression.");
   assert(Kind != AGVK_None && "Cannot construct AMDGPUMCExpr of kind none.");

   // Allocating the variadic arguments through the same allocation mechanism
   // that the object itself is allocated with so they end up in the same memory.
   //
   // Will result in an asan failure if allocated on the heap through standard
   // allocation (e.g., through SmallVector's grow).
   RawArgs = static_cast<const MCExpr **>(
       Ctx.allocate(sizeof(const MCExpr *) * Args.size()));
   std::uninitialized_copy(Args.begin(), Args.end(), RawArgs);
   this->Args = ArrayRef<const MCExpr *>(RawArgs, Args.size());
 }

 AMDGPUMCExpr::~AMDGPUMCExpr() { Ctx.deallocate(RawArgs); }

 const AMDGPUMCExpr *AMDGPUMCExpr::create(VariantKind Kind,
                                          ArrayRef<const MCExpr *> Args,
                                          MCContext &Ctx) {
   return new (Ctx) AMDGPUMCExpr(Kind, Args, Ctx);
 }

 const MCExpr *AMDGPUMCExpr::getSubExpr(size_t Index) const {
   assert(Index < Args.size() && "Indexing out of bounds AMDGPUMCExpr sub-expr");
   return Args[Index];
 }

 void AMDGPUMCExpr::printImpl(raw_ostream &OS, const MCAsmInfo *MAI) const {
   switch (Kind) {
   default:
     llvm_unreachable("Unknown AMDGPUMCExpr kind.");
   case AGVK_Or:
     OS << "or(";
     break;
   case AGVK_Max:
     OS << "max(";
     break;
   case AGVK_ExtraSGPRs:
     OS << "extrasgprs(";
     break;
   case AGVK_TotalNumVGPRs:
     OS << "totalnumvgprs(";
     break;
   case AGVK_AlignTo:
     OS << "alignto(";
     break;
   case AGVK_Occupancy:
     OS << "occupancy(";
     break;
   }
   for (const auto *It = Args.begin(); It != Args.end(); ++It) {
     (*It)->print(OS, MAI, /*InParens=*/false);
     if ((It + 1) != Args.end())
       OS << ", ";
   }
   OS << ')';
 }

 static int64_t op(AMDGPUMCExpr::VariantKind Kind, int64_t Arg1, int64_t Arg2) {
   switch (Kind) {
   default:
     llvm_unreachable("Unknown AMDGPUMCExpr kind.");
   case AMDGPUMCExpr::AGVK_Max:
     return std::max(Arg1, Arg2);
   case AMDGPUMCExpr::AGVK_Or:
     return Arg1 | Arg2;
   }
 }

 bool AMDGPUMCExpr::evaluateExtraSGPRs(MCValue &Res, const MCAssembler *Asm,
                                       const MCFixup *Fixup) const {
   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
     MCValue MCVal;
     if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) || !MCVal.isAbsolute())
       return false;

     ConstantValue = MCVal.getConstant();
     return true;
   };

   assert(Args.size() == 3 &&
          "AMDGPUMCExpr Argument count incorrect for ExtraSGPRs");
   const MCSubtargetInfo *STI = Ctx.getSubtargetInfo();
   uint64_t VCCUsed = 0, FlatScrUsed = 0, XNACKUsed = 0;

   bool Success = TryGetMCExprValue(Args[2], XNACKUsed);

   assert(Success && "Arguments 3 for ExtraSGPRs should be a known constant");
   if (!Success || !TryGetMCExprValue(Args[0], VCCUsed) ||
       !TryGetMCExprValue(Args[1], FlatScrUsed))
     return false;

   uint64_t ExtraSGPRs = IsaInfo::getNumExtraSGPRs(
       STI, (bool)VCCUsed, (bool)FlatScrUsed, (bool)XNACKUsed);
   Res = MCValue::get(ExtraSGPRs);
   return true;
 }

 bool AMDGPUMCExpr::evaluateTotalNumVGPR(MCValue &Res, const MCAssembler *Asm,
                                         const MCFixup *Fixup) const {
   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
     MCValue MCVal;
     if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) || !MCVal.isAbsolute())
       return false;

     ConstantValue = MCVal.getConstant();
     return true;
   };
   assert(Args.size() == 2 &&
          "AMDGPUMCExpr Argument count incorrect for TotalNumVGPRs");
   const MCSubtargetInfo *STI = Ctx.getSubtargetInfo();
   uint64_t NumAGPR = 0, NumVGPR = 0;

   bool Has90AInsts = AMDGPU::isGFX90A(*STI);

   if (!TryGetMCExprValue(Args[0], NumAGPR) ||
       !TryGetMCExprValue(Args[1], NumVGPR))
     return false;

   uint64_t TotalNum = Has90AInsts && NumAGPR ? alignTo(NumVGPR, 4) + NumAGPR
                                              : std::max(NumVGPR, NumAGPR);
   Res = MCValue::get(TotalNum);
   return true;
 }

 bool AMDGPUMCExpr::evaluateAlignTo(MCValue &Res, const MCAssembler *Asm,
                                    const MCFixup *Fixup) const {
   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
     MCValue MCVal;
     if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) || !MCVal.isAbsolute())
       return false;

     ConstantValue = MCVal.getConstant();
     return true;
   };

   assert(Args.size() == 2 &&
          "AMDGPUMCExpr Argument count incorrect for AlignTo");
   uint64_t Value = 0, Align = 0;
   if (!TryGetMCExprValue(Args[0], Value) || !TryGetMCExprValue(Args[1], Align))
     return false;

   Res = MCValue::get(alignTo(Value, Align));
   return true;
 }

 bool AMDGPUMCExpr::evaluateOccupancy(MCValue &Res, const MCAssembler *Asm,
                                      const MCFixup *Fixup) const {
   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
     MCValue MCVal;
     if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) || !MCVal.isAbsolute())
       return false;

     ConstantValue = MCVal.getConstant();
     return true;
   };
   assert(Args.size() == 7 &&
          "AMDGPUMCExpr Argument count incorrect for Occupancy");
   uint64_t InitOccupancy, MaxWaves, Granule, TargetTotalNumVGPRs, Generation,
       NumSGPRs, NumVGPRs;

   bool Success = true;
   Success &= TryGetMCExprValue(Args[0], MaxWaves);
   Success &= TryGetMCExprValue(Args[1], Granule);
   Success &= TryGetMCExprValue(Args[2], TargetTotalNumVGPRs);
   Success &= TryGetMCExprValue(Args[3], Generation);
   Success &= TryGetMCExprValue(Args[4], InitOccupancy);

   assert(Success && "Arguments 1 to 5 for Occupancy should be known constants");

   if (!Success || !TryGetMCExprValue(Args[5], NumSGPRs) ||
       !TryGetMCExprValue(Args[6], NumVGPRs))
     return false;

   unsigned Occupancy = InitOccupancy;
   if (NumSGPRs)
     Occupancy = std::min(
         Occupancy, IsaInfo::getOccupancyWithNumSGPRs(
                        NumSGPRs, MaxWaves,
                        static_cast<AMDGPUSubtarget::Generation>(Generation)));
   if (NumVGPRs)
     Occupancy = std::min(Occupancy,
                          IsaInfo::getNumWavesPerEUWithNumVGPRs(
                              NumVGPRs, Granule, MaxWaves, TargetTotalNumVGPRs));

   Res = MCValue::get(Occupancy);
   return true;
 }

 bool AMDGPUMCExpr::evaluateAsRelocatableImpl(MCValue &Res,
                                              const MCAssembler *Asm,
                                              const MCFixup *Fixup) const {
   std::optional<int64_t> Total;
   switch (Kind) {
   default:
     break;
   case AGVK_ExtraSGPRs:
     return evaluateExtraSGPRs(Res, Asm, Fixup);
   case AGVK_AlignTo:
     return evaluateAlignTo(Res, Asm, Fixup);
   case AGVK_TotalNumVGPRs:
     return evaluateTotalNumVGPR(Res, Asm, Fixup);
   case AGVK_Occupancy:
     return evaluateOccupancy(Res, Asm, Fixup);
   }

   for (const MCExpr *Arg : Args) {
     MCValue ArgRes;
     if (!Arg->evaluateAsRelocatable(ArgRes, Asm, Fixup) || !ArgRes.isAbsolute())
       return false;

     if (!Total.has_value())
       Total = ArgRes.getConstant();
     Total = op(Kind, *Total, ArgRes.getConstant());
   }

   Res = MCValue::get(*Total);
   return true;
 }

 void AMDGPUMCExpr::visitUsedExpr(MCStreamer &Streamer) const {
   for (const MCExpr *Arg : Args)
     Streamer.visitUsedExpr(*Arg);
 }

 MCFragment *AMDGPUMCExpr::findAssociatedFragment() const {
   for (const MCExpr *Arg : Args) {
     if (Arg->findAssociatedFragment())
       return Arg->findAssociatedFragment();
   }
   return nullptr;
 }

 /// Allow delayed MCExpr resolve of ExtraSGPRs (in case VCCUsed or FlatScrUsed
 /// are unresolvable but needed for further MCExprs). Derived from
 /// implementation of IsaInfo::getNumExtraSGPRs in AMDGPUBaseInfo.cpp.
 ///
 const AMDGPUMCExpr *AMDGPUMCExpr::createExtraSGPRs(const MCExpr *VCCUsed,
                                                    const MCExpr *FlatScrUsed,
                                                    bool XNACKUsed,
                                                    MCContext &Ctx) {

   return create(AGVK_ExtraSGPRs,
                 {VCCUsed, FlatScrUsed, MCConstantExpr::create(XNACKUsed, Ctx)},
                 Ctx);
 }

 const AMDGPUMCExpr *AMDGPUMCExpr::createTotalNumVGPR(const MCExpr *NumAGPR,
                                                      const MCExpr *NumVGPR,
                                                      MCContext &Ctx) {
   return create(AGVK_TotalNumVGPRs, {NumAGPR, NumVGPR}, Ctx);
 }

 /// Mimics GCNSubtarget::computeOccupancy for MCExpr.
 ///
 /// Remove dependency on GCNSubtarget and depend only only the necessary values
 /// for said occupancy computation. Should match computeOccupancy implementation
 /// without passing \p STM on.
 const AMDGPUMCExpr *AMDGPUMCExpr::createOccupancy(unsigned InitOcc,
                                                   const MCExpr *NumSGPRs,
                                                   const MCExpr *NumVGPRs,
                                                   const GCNSubtarget &STM,
                                                   MCContext &Ctx) {
   unsigned MaxWaves = IsaInfo::getMaxWavesPerEU(&STM);
   unsigned Granule = IsaInfo::getVGPRAllocGranule(&STM);
   unsigned TargetTotalNumVGPRs = IsaInfo::getTotalNumVGPRs(&STM);
   unsigned Generation = STM.getGeneration();

   auto CreateExpr = [&Ctx](unsigned Value) {
     return MCConstantExpr::create(Value, Ctx);
   };

   return create(AGVK_Occupancy,
                 {CreateExpr(MaxWaves), CreateExpr(Granule),
                  CreateExpr(TargetTotalNumVGPRs), CreateExpr(Generation),
                  CreateExpr(InitOcc), NumSGPRs, NumVGPRs},
                 Ctx);
 }

 bool AMDGPUMCExpr::isSymbolUsedInExpression(const MCSymbol *Sym) const {
   for (const MCExpr *E : getArgs()) {
     if (E->isSymbolUsedInExpression(Sym))
       return true;
   }
   return false;
 }

 static KnownBits fromOptionalToKnownBits(std::optional<bool> CompareResult) {
   static constexpr unsigned BitWidth = 64;
   const APInt True(BitWidth, 1);
   const APInt False(BitWidth, 0);
   if (CompareResult) {
     return *CompareResult ? KnownBits::makeConstant(True)
                           : KnownBits::makeConstant(False);
   }

   KnownBits UnknownBool(/*BitWidth=*/1);
   return UnknownBool.zext(BitWidth);
 }

 using KnownBitsMap = DenseMap<const MCExpr *, KnownBits>;
 static void knownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
                                unsigned Depth = 0);

 static void binaryOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
                                        unsigned Depth) {
   static constexpr unsigned BitWidth = 64;
   const MCBinaryExpr *BExpr = cast<MCBinaryExpr>(Expr);
   const MCExpr *LHS = BExpr->getLHS();
   const MCExpr *RHS = BExpr->getRHS();

   knownBitsMapHelper(LHS, KBM, Depth + 1);
   knownBitsMapHelper(RHS, KBM, Depth + 1);
   KnownBits LHSKnown = KBM[LHS];
   KnownBits RHSKnown = KBM[RHS];

   switch (BExpr->getOpcode()) {
   default:
     KBM[Expr] = KnownBits(BitWidth);
     return;
   case MCBinaryExpr::Opcode::Add:
     KBM[Expr] = KnownBits::add(LHSKnown, RHSKnown);
     return;
   case MCBinaryExpr::Opcode::And:
     KBM[Expr] = LHSKnown & RHSKnown;
     return;
   case MCBinaryExpr::Opcode::Div:
     KBM[Expr] = KnownBits::sdiv(LHSKnown, RHSKnown);
     return;
   case MCBinaryExpr::Opcode::EQ: {
     std::optional<bool> CompareRes = KnownBits::eq(LHSKnown, RHSKnown);
     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
     return;
   }
   case MCBinaryExpr::Opcode::NE: {
     std::optional<bool> CompareRes = KnownBits::ne(LHSKnown, RHSKnown);
     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
     return;
   }
   case MCBinaryExpr::Opcode::GT: {
     std::optional<bool> CompareRes = KnownBits::sgt(LHSKnown, RHSKnown);
     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
     return;
   }
   case MCBinaryExpr::Opcode::GTE: {
     std::optional<bool> CompareRes = KnownBits::sge(LHSKnown, RHSKnown);
     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
     return;
   }
   case MCBinaryExpr::Opcode::LAnd: {
     std::optional<bool> CompareRes;
     const APInt False(BitWidth, 0);
     std::optional<bool> LHSBool =
         KnownBits::ne(LHSKnown, KnownBits::makeConstant(False));
     std::optional<bool> RHSBool =
         KnownBits::ne(RHSKnown, KnownBits::makeConstant(False));
     if (LHSBool && RHSBool)
       CompareRes = *LHSBool && *RHSBool;
     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
     return;
   }
   case MCBinaryExpr::Opcode::LOr: {
     const APInt False(BitWidth, 0);
     KnownBits Bits = LHSKnown | RHSKnown;
     std::optional<bool> CompareRes =
         KnownBits::ne(Bits, KnownBits::makeConstant(False));
     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
     return;
   }
   case MCBinaryExpr::Opcode::LT: {
     std::optional<bool> CompareRes = KnownBits::slt(LHSKnown, RHSKnown);
     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
     return;
   }
   case MCBinaryExpr::Opcode::LTE: {
     std::optional<bool> CompareRes = KnownBits::sle(LHSKnown, RHSKnown);
     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
     return;
   }
   case MCBinaryExpr::Opcode::Mod:
     KBM[Expr] = KnownBits::srem(LHSKnown, RHSKnown);
     return;
   case MCBinaryExpr::Opcode::Mul:
     KBM[Expr] = KnownBits::mul(LHSKnown, RHSKnown);
     return;
   case MCBinaryExpr::Opcode::Or:
     KBM[Expr] = LHSKnown | RHSKnown;
     return;
   case MCBinaryExpr::Opcode::Shl:
     KBM[Expr] = KnownBits::shl(LHSKnown, RHSKnown);
     return;
   case MCBinaryExpr::Opcode::AShr:
     KBM[Expr] = KnownBits::ashr(LHSKnown, RHSKnown);
     return;
   case MCBinaryExpr::Opcode::LShr:
     KBM[Expr] = KnownBits::lshr(LHSKnown, RHSKnown);
     return;
   case MCBinaryExpr::Opcode::Sub:
     KBM[Expr] = KnownBits::sub(LHSKnown, RHSKnown);
     return;
   case MCBinaryExpr::Opcode::Xor:
     KBM[Expr] = LHSKnown ^ RHSKnown;
     return;
   }
 }

 static void unaryOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
                                       unsigned Depth) {
   static constexpr unsigned BitWidth = 64;
   const MCUnaryExpr *UExpr = cast<MCUnaryExpr>(Expr);
   knownBitsMapHelper(UExpr->getSubExpr(), KBM, Depth + 1);
   KnownBits KB = KBM[UExpr->getSubExpr()];

   switch (UExpr->getOpcode()) {
   default:
     KBM[Expr] = KnownBits(BitWidth);
     return;
   case MCUnaryExpr::Opcode::Minus: {
     KB.makeNegative();
     KBM[Expr] = KB;
     return;
   }
   case MCUnaryExpr::Opcode::Not: {
     KnownBits AllOnes(BitWidth);
     AllOnes.setAllOnes();
     KBM[Expr] = KB ^ AllOnes;
     return;
   }
   case MCUnaryExpr::Opcode::Plus: {
     KB.makeNonNegative();
     KBM[Expr] = KB;
     return;
   }
   }
 }

 static void targetOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
                                        unsigned Depth) {
   static constexpr unsigned BitWidth = 64;
   const AMDGPUMCExpr *AGVK = cast<AMDGPUMCExpr>(Expr);

   switch (AGVK->getKind()) {
   default:
     KBM[Expr] = KnownBits(BitWidth);
     return;
   case AMDGPUMCExpr::VariantKind::AGVK_Or: {
     knownBitsMapHelper(AGVK->getSubExpr(0), KBM, Depth + 1);
     KnownBits KB = KBM[AGVK->getSubExpr(0)];
     for (const MCExpr *Arg : AGVK->getArgs()) {
       knownBitsMapHelper(Arg, KBM, Depth + 1);
       KB |= KBM[Arg];
     }
     KBM[Expr] = KB;
     return;
   }
   case AMDGPUMCExpr::VariantKind::AGVK_Max: {
     knownBitsMapHelper(AGVK->getSubExpr(0), KBM, Depth + 1);
     KnownBits KB = KBM[AGVK->getSubExpr(0)];
     for (const MCExpr *Arg : AGVK->getArgs()) {
       knownBitsMapHelper(Arg, KBM, Depth + 1);
       KB = KnownBits::umax(KB, KBM[Arg]);
     }
     KBM[Expr] = KB;
     return;
   }
   case AMDGPUMCExpr::VariantKind::AGVK_ExtraSGPRs:
   case AMDGPUMCExpr::VariantKind::AGVK_TotalNumVGPRs:
   case AMDGPUMCExpr::VariantKind::AGVK_AlignTo:
   case AMDGPUMCExpr::VariantKind::AGVK_Occupancy: {
     int64_t Val;
     if (AGVK->evaluateAsAbsolute(Val)) {
       APInt APValue(BitWidth, Val);
       KBM[Expr] = KnownBits::makeConstant(APValue);
       return;
     }
     KBM[Expr] = KnownBits(BitWidth);
     return;
   }
   }
 }

 static void knownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
                                unsigned Depth) {
   static constexpr unsigned BitWidth = 64;

   int64_t Val;
   if (Expr->evaluateAsAbsolute(Val)) {
     APInt APValue(BitWidth, Val, /*isSigned=*/true);
     KBM[Expr] = KnownBits::makeConstant(APValue);
     return;
   }

   if (Depth == 16) {
     KBM[Expr] = KnownBits(BitWidth);
     return;
   }

   switch (Expr->getKind()) {
   case MCExpr::ExprKind::Binary: {
     binaryOpKnownBitsMapHelper(Expr, KBM, Depth);
     return;
   }
   case MCExpr::ExprKind::Constant: {
     const MCConstantExpr *CE = cast<MCConstantExpr>(Expr);
     APInt APValue(BitWidth, CE->getValue(), /*isSigned=*/true);
     KBM[Expr] = KnownBits::makeConstant(APValue);
     return;
   }
   case MCExpr::ExprKind::SymbolRef: {
     const MCSymbolRefExpr *RExpr = cast<MCSymbolRefExpr>(Expr);
     const MCSymbol &Sym = RExpr->getSymbol();
     if (!Sym.isVariable()) {
       KBM[Expr] = KnownBits(BitWidth);
       return;
     }

     // Variable value retrieval is not for actual use but only for knownbits
     // analysis.
     const MCExpr *SymVal = Sym.getVariableValue(/*setUsed=*/false);
     knownBitsMapHelper(SymVal, KBM, Depth + 1);

     // Explicitly copy-construct so that there exists a local KnownBits in case
     // KBM[SymVal] gets invalidated after a potential growth through KBM[Expr].
     KBM[Expr] = KnownBits(KBM[SymVal]);
     return;
   }
   case MCExpr::ExprKind::Unary: {
     unaryOpKnownBitsMapHelper(Expr, KBM, Depth);
     return;
   }
   case MCExpr::ExprKind::Target: {
     targetOpKnownBitsMapHelper(Expr, KBM, Depth);
     return;
   }
   }
 }

 static const MCExpr *tryFoldHelper(const MCExpr *Expr, KnownBitsMap &KBM,
                                    MCContext &Ctx) {
   if (!KBM.count(Expr))
     return Expr;

   auto ValueCheckKnownBits = [](KnownBits &KB, unsigned Value) -> bool {
     if (!KB.isConstant())
       return false;

     return Value == KB.getConstant();
   };

   if (Expr->getKind() == MCExpr::ExprKind::Constant)
     return Expr;

   // Resolving unary operations to constants may make the value more ambiguous.
   // For example, `~62` becomes `-63`; however, to me it's more ambiguous if a
   // bit mask value is represented through a negative number.
   if (Expr->getKind() != MCExpr::ExprKind::Unary) {
     if (KBM[Expr].isConstant()) {
       APInt ConstVal = KBM[Expr].getConstant();
       return MCConstantExpr::create(ConstVal.getSExtValue(), Ctx);
     }

     int64_t EvalValue;
     if (Expr->evaluateAsAbsolute(EvalValue))
       return MCConstantExpr::create(EvalValue, Ctx);
   }

   switch (Expr->getKind()) {
   default:
     return Expr;
   case MCExpr::ExprKind::Binary: {
     const MCBinaryExpr *BExpr = cast<MCBinaryExpr>(Expr);
     const MCExpr *LHS = BExpr->getLHS();
     const MCExpr *RHS = BExpr->getRHS();

     switch (BExpr->getOpcode()) {
     default:
       return Expr;
     case MCBinaryExpr::Opcode::Sub: {
       if (ValueCheckKnownBits(KBM[RHS], 0))
         return tryFoldHelper(LHS, KBM, Ctx);
       break;
     }
     case MCBinaryExpr::Opcode::Add:
     case MCBinaryExpr::Opcode::Or: {
       if (ValueCheckKnownBits(KBM[LHS], 0))
         return tryFoldHelper(RHS, KBM, Ctx);
       if (ValueCheckKnownBits(KBM[RHS], 0))
         return tryFoldHelper(LHS, KBM, Ctx);
       break;
     }
     case MCBinaryExpr::Opcode::Mul: {
       if (ValueCheckKnownBits(KBM[LHS], 1))
         return tryFoldHelper(RHS, KBM, Ctx);
       if (ValueCheckKnownBits(KBM[RHS], 1))
         return tryFoldHelper(LHS, KBM, Ctx);
       break;
     }
     case MCBinaryExpr::Opcode::Shl:
     case MCBinaryExpr::Opcode::AShr:
     case MCBinaryExpr::Opcode::LShr: {
       if (ValueCheckKnownBits(KBM[RHS], 0))
         return tryFoldHelper(LHS, KBM, Ctx);
       if (ValueCheckKnownBits(KBM[LHS], 0))
         return MCConstantExpr::create(0, Ctx);
       break;
     }
     case MCBinaryExpr::Opcode::And: {
       if (ValueCheckKnownBits(KBM[LHS], 0) || ValueCheckKnownBits(KBM[RHS], 0))
         return MCConstantExpr::create(0, Ctx);
       break;
     }
     }
     const MCExpr *NewLHS = tryFoldHelper(LHS, KBM, Ctx);
     const MCExpr *NewRHS = tryFoldHelper(RHS, KBM, Ctx);
     if (NewLHS != LHS || NewRHS != RHS)
       return MCBinaryExpr::create(BExpr->getOpcode(), NewLHS, NewRHS, Ctx,
                                   BExpr->getLoc());
     return Expr;
   }
   case MCExpr::ExprKind::Unary: {
     const MCUnaryExpr *UExpr = cast<MCUnaryExpr>(Expr);
     const MCExpr *SubExpr = UExpr->getSubExpr();
     const MCExpr *NewSubExpr = tryFoldHelper(SubExpr, KBM, Ctx);
     if (SubExpr != NewSubExpr)
       return MCUnaryExpr::create(UExpr->getOpcode(), NewSubExpr, Ctx,
                                  UExpr->getLoc());
     return Expr;
   }
   case MCExpr::ExprKind::Target: {
     const AMDGPUMCExpr *AGVK = cast<AMDGPUMCExpr>(Expr);
     SmallVector<const MCExpr *, 8> NewArgs;
     bool Changed = false;
     for (const MCExpr *Arg : AGVK->getArgs()) {
       const MCExpr *NewArg = tryFoldHelper(Arg, KBM, Ctx);
       NewArgs.push_back(NewArg);
       Changed |= Arg != NewArg;
     }
     return Changed ? AMDGPUMCExpr::create(AGVK->getKind(), NewArgs, Ctx) : Expr;
   }
   }
   return Expr;
 }

 const MCExpr *llvm::AMDGPU::foldAMDGPUMCExpr(const MCExpr *Expr,
                                              MCContext &Ctx) {
   KnownBitsMap KBM;
   knownBitsMapHelper(Expr, KBM);
   const MCExpr *NewExpr = tryFoldHelper(Expr, KBM, Ctx);

   return Expr != NewExpr ? NewExpr : Expr;
 }

 void llvm::AMDGPU::printAMDGPUMCExpr(const MCExpr *Expr, raw_ostream &OS,
                                      const MCAsmInfo *MAI) {
   int64_t Val;
   if (Expr->evaluateAsAbsolute(Val)) {
     OS << Val;
     return;
   }

   Expr->print(OS, MAI);
 }
	//===- AMDGPUMCExpr.cpp - AMDGPU specific MC expression classes -----------===//
	//
	// 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
	//
	//===----------------------------------------------------------------------===//

	#include "AMDGPUMCExpr.h"
	#include "GCNSubtarget.h"
	#include "Utils/AMDGPUBaseInfo.h"
	#include "llvm/IR/Function.h"
	#include "llvm/MC/MCAsmInfo.h"
	#include "llvm/MC/MCAssembler.h"
	#include "llvm/MC/MCContext.h"
	#include "llvm/MC/MCStreamer.h"
	#include "llvm/MC/MCSymbol.h"
	#include "llvm/MC/MCValue.h"
	#include "llvm/Support/KnownBits.h"
	#include "llvm/Support/raw_ostream.h"
	#include <optional>

	using namespace llvm;
	using namespace llvm::AMDGPU;

	AMDGPUMCExpr::AMDGPUMCExpr(VariantKind Kind, ArrayRef<const MCExpr *> Args,
	MCContext &Ctx)
	: Kind(Kind), Ctx(Ctx) {
	assert(Args.size() >= 1 && "Needs a minimum of one expression.");
	assert(Kind != AGVK_None && "Cannot construct AMDGPUMCExpr of kind none.");

	// Allocating the variadic arguments through the same allocation mechanism
	// that the object itself is allocated with so they end up in the same memory.
	//
	// Will result in an asan failure if allocated on the heap through standard
	// allocation (e.g., through SmallVector's grow).
	RawArgs = static_cast<const MCExpr **>(
	Ctx.allocate(sizeof(const MCExpr ) Args.size()));
	std::uninitialized_copy(Args.begin(), Args.end(), RawArgs);
	this->Args = ArrayRef<const MCExpr *>(RawArgs, Args.size());
	}

	AMDGPUMCExpr::~AMDGPUMCExpr() { Ctx.deallocate(RawArgs); }

	const AMDGPUMCExpr *AMDGPUMCExpr::create(VariantKind Kind,
	ArrayRef<const MCExpr *> Args,
	MCContext &Ctx) {
	return new (Ctx) AMDGPUMCExpr(Kind, Args, Ctx);
	}

	const MCExpr *AMDGPUMCExpr::getSubExpr(size_t Index) const {
	assert(Index < Args.size() && "Indexing out of bounds AMDGPUMCExpr sub-expr");
	return Args[Index];
	}

	void AMDGPUMCExpr::printImpl(raw_ostream &OS, const MCAsmInfo *MAI) const {
	switch (Kind) {
	default:
	llvm_unreachable("Unknown AMDGPUMCExpr kind.");
	case AGVK_Or:
	OS << "or(";
	break;
	case AGVK_Max:
	OS << "max(";
	break;
	case AGVK_ExtraSGPRs:
	OS << "extrasgprs(";
	break;
	case AGVK_TotalNumVGPRs:
	OS << "totalnumvgprs(";
	break;
	case AGVK_AlignTo:
	OS << "alignto(";
	break;
	case AGVK_Occupancy:
	OS << "occupancy(";
	break;
	}
	for (const auto *It = Args.begin(); It != Args.end(); ++It) {
	(It)->print(OS, MAI, /InParens=*/false);
	if ((It + 1) != Args.end())
	OS << ", ";
	}
	OS << ')';
	}

	static int64_t op(AMDGPUMCExpr::VariantKind Kind, int64_t Arg1, int64_t Arg2) {
	switch (Kind) {
	default:
	llvm_unreachable("Unknown AMDGPUMCExpr kind.");
	case AMDGPUMCExpr::AGVK_Max:
	return std::max(Arg1, Arg2);
	case AMDGPUMCExpr::AGVK_Or:
	return Arg1 \| Arg2;
	}
	}

	bool AMDGPUMCExpr::evaluateExtraSGPRs(MCValue &Res, const MCAssembler *Asm,
	const MCFixup *Fixup) const {
	auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
	MCValue MCVal;
	if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) \|\| !MCVal.isAbsolute())
	return false;

	ConstantValue = MCVal.getConstant();
	return true;
	};

	assert(Args.size() == 3 &&
	"AMDGPUMCExpr Argument count incorrect for ExtraSGPRs");
	const MCSubtargetInfo *STI = Ctx.getSubtargetInfo();
	uint64_t VCCUsed = 0, FlatScrUsed = 0, XNACKUsed = 0;

	bool Success = TryGetMCExprValue(Args[2], XNACKUsed);

	assert(Success && "Arguments 3 for ExtraSGPRs should be a known constant");
	if (!Success \|\| !TryGetMCExprValue(Args[0], VCCUsed) \|\|
	!TryGetMCExprValue(Args[1], FlatScrUsed))
	return false;

	uint64_t ExtraSGPRs = IsaInfo::getNumExtraSGPRs(
	STI, (bool)VCCUsed, (bool)FlatScrUsed, (bool)XNACKUsed);
	Res = MCValue::get(ExtraSGPRs);
	return true;
	}

	bool AMDGPUMCExpr::evaluateTotalNumVGPR(MCValue &Res, const MCAssembler *Asm,
	const MCFixup *Fixup) const {
	auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
	MCValue MCVal;
	if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) \|\| !MCVal.isAbsolute())
	return false;

	ConstantValue = MCVal.getConstant();
	return true;
	};
	assert(Args.size() == 2 &&
	"AMDGPUMCExpr Argument count incorrect for TotalNumVGPRs");
	const MCSubtargetInfo *STI = Ctx.getSubtargetInfo();
	uint64_t NumAGPR = 0, NumVGPR = 0;

	bool Has90AInsts = AMDGPU::isGFX90A(*STI);

	if (!TryGetMCExprValue(Args[0], NumAGPR) \|\|
	!TryGetMCExprValue(Args[1], NumVGPR))
	return false;

	uint64_t TotalNum = Has90AInsts && NumAGPR ? alignTo(NumVGPR, 4) + NumAGPR
	: std::max(NumVGPR, NumAGPR);
	Res = MCValue::get(TotalNum);
	return true;
	}

	bool AMDGPUMCExpr::evaluateAlignTo(MCValue &Res, const MCAssembler *Asm,
	const MCFixup *Fixup) const {
	auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
	MCValue MCVal;
	if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) \|\| !MCVal.isAbsolute())
	return false;

	ConstantValue = MCVal.getConstant();
	return true;
	};

	assert(Args.size() == 2 &&
	"AMDGPUMCExpr Argument count incorrect for AlignTo");
	uint64_t Value = 0, Align = 0;
	if (!TryGetMCExprValue(Args[0], Value) \|\| !TryGetMCExprValue(Args[1], Align))
	return false;

	Res = MCValue::get(alignTo(Value, Align));
	return true;
	}

	bool AMDGPUMCExpr::evaluateOccupancy(MCValue &Res, const MCAssembler *Asm,
	const MCFixup *Fixup) const {
	auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
	MCValue MCVal;
	if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) \|\| !MCVal.isAbsolute())
	return false;

	ConstantValue = MCVal.getConstant();
	return true;
	};
	assert(Args.size() == 7 &&
	"AMDGPUMCExpr Argument count incorrect for Occupancy");
	uint64_t InitOccupancy, MaxWaves, Granule, TargetTotalNumVGPRs, Generation,
	NumSGPRs, NumVGPRs;

	bool Success = true;
	Success &= TryGetMCExprValue(Args[0], MaxWaves);
	Success &= TryGetMCExprValue(Args[1], Granule);
	Success &= TryGetMCExprValue(Args[2], TargetTotalNumVGPRs);
	Success &= TryGetMCExprValue(Args[3], Generation);
	Success &= TryGetMCExprValue(Args[4], InitOccupancy);

	assert(Success && "Arguments 1 to 5 for Occupancy should be known constants");

	if (!Success \|\| !TryGetMCExprValue(Args[5], NumSGPRs) \|\|
	!TryGetMCExprValue(Args[6], NumVGPRs))
	return false;

	unsigned Occupancy = InitOccupancy;
	if (NumSGPRs)
	Occupancy = std::min(
	Occupancy, IsaInfo::getOccupancyWithNumSGPRs(
	NumSGPRs, MaxWaves,
	static_cast<AMDGPUSubtarget::Generation>(Generation)));
	if (NumVGPRs)
	Occupancy = std::min(Occupancy,
	IsaInfo::getNumWavesPerEUWithNumVGPRs(
	NumVGPRs, Granule, MaxWaves, TargetTotalNumVGPRs));

	Res = MCValue::get(Occupancy);
	return true;
	}

	bool AMDGPUMCExpr::evaluateAsRelocatableImpl(MCValue &Res,
	const MCAssembler *Asm,
	const MCFixup *Fixup) const {
	std::optional<int64_t> Total;
	switch (Kind) {
	default:
	break;
	case AGVK_ExtraSGPRs:
	return evaluateExtraSGPRs(Res, Asm, Fixup);
	case AGVK_AlignTo:
	return evaluateAlignTo(Res, Asm, Fixup);
	case AGVK_TotalNumVGPRs:
	return evaluateTotalNumVGPR(Res, Asm, Fixup);
	case AGVK_Occupancy:
	return evaluateOccupancy(Res, Asm, Fixup);
	}

	for (const MCExpr *Arg : Args) {
	MCValue ArgRes;
	if (!Arg->evaluateAsRelocatable(ArgRes, Asm, Fixup) \|\| !ArgRes.isAbsolute())
	return false;

	if (!Total.has_value())
	Total = ArgRes.getConstant();
	Total = op(Kind, *Total, ArgRes.getConstant());
	}

	Res = MCValue::get(*Total);
	return true;
	}

	void AMDGPUMCExpr::visitUsedExpr(MCStreamer &Streamer) const {
	for (const MCExpr *Arg : Args)
	Streamer.visitUsedExpr(*Arg);
	}

	MCFragment *AMDGPUMCExpr::findAssociatedFragment() const {
	for (const MCExpr *Arg : Args) {
	if (Arg->findAssociatedFragment())
	return Arg->findAssociatedFragment();
	}
	return nullptr;
	}

	/// Allow delayed MCExpr resolve of ExtraSGPRs (in case VCCUsed or FlatScrUsed
	/// are unresolvable but needed for further MCExprs). Derived from
	/// implementation of IsaInfo::getNumExtraSGPRs in AMDGPUBaseInfo.cpp.
	///
	const AMDGPUMCExpr AMDGPUMCExpr::createExtraSGPRs(const MCExpr VCCUsed,
	const MCExpr *FlatScrUsed,
	bool XNACKUsed,
	MCContext &Ctx) {

	return create(AGVK_ExtraSGPRs,
	{VCCUsed, FlatScrUsed, MCConstantExpr::create(XNACKUsed, Ctx)},
	Ctx);
	}

	const AMDGPUMCExpr AMDGPUMCExpr::createTotalNumVGPR(const MCExpr NumAGPR,
	const MCExpr *NumVGPR,
	MCContext &Ctx) {
	return create(AGVK_TotalNumVGPRs, {NumAGPR, NumVGPR}, Ctx);
	}

	/// Mimics GCNSubtarget::computeOccupancy for MCExpr.
	///
	/// Remove dependency on GCNSubtarget and depend only only the necessary values
	/// for said occupancy computation. Should match computeOccupancy implementation
	/// without passing \p STM on.
	const AMDGPUMCExpr *AMDGPUMCExpr::createOccupancy(unsigned InitOcc,
	const MCExpr *NumSGPRs,
	const MCExpr *NumVGPRs,
	const GCNSubtarget &STM,
	MCContext &Ctx) {
	unsigned MaxWaves = IsaInfo::getMaxWavesPerEU(&STM);
	unsigned Granule = IsaInfo::getVGPRAllocGranule(&STM);
	unsigned TargetTotalNumVGPRs = IsaInfo::getTotalNumVGPRs(&STM);
	unsigned Generation = STM.getGeneration();

	auto CreateExpr = [&Ctx](unsigned Value) {
	return MCConstantExpr::create(Value, Ctx);
	};

	return create(AGVK_Occupancy,
	{CreateExpr(MaxWaves), CreateExpr(Granule),
	CreateExpr(TargetTotalNumVGPRs), CreateExpr(Generation),
	CreateExpr(InitOcc), NumSGPRs, NumVGPRs},
	Ctx);
	}

	bool AMDGPUMCExpr::isSymbolUsedInExpression(const MCSymbol *Sym) const {
	for (const MCExpr *E : getArgs()) {
	if (E->isSymbolUsedInExpression(Sym))
	return true;
	}
	return false;
	}

	static KnownBits fromOptionalToKnownBits(std::optional<bool> CompareResult) {
	static constexpr unsigned BitWidth = 64;
	const APInt True(BitWidth, 1);
	const APInt False(BitWidth, 0);
	if (CompareResult) {
	return *CompareResult ? KnownBits::makeConstant(True)
	: KnownBits::makeConstant(False);
	}

	KnownBits UnknownBool(/BitWidth=/1);
	return UnknownBool.zext(BitWidth);
	}

	using KnownBitsMap = DenseMap<const MCExpr *, KnownBits>;
	static void knownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
	unsigned Depth = 0);

	static void binaryOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
	unsigned Depth) {
	static constexpr unsigned BitWidth = 64;
	const MCBinaryExpr *BExpr = cast<MCBinaryExpr>(Expr);
	const MCExpr *LHS = BExpr->getLHS();
	const MCExpr *RHS = BExpr->getRHS();

	knownBitsMapHelper(LHS, KBM, Depth + 1);
	knownBitsMapHelper(RHS, KBM, Depth + 1);
	KnownBits LHSKnown = KBM[LHS];
	KnownBits RHSKnown = KBM[RHS];

	switch (BExpr->getOpcode()) {
	default:
	KBM[Expr] = KnownBits(BitWidth);
	return;
	case MCBinaryExpr::Opcode::Add:
	KBM[Expr] = KnownBits::add(LHSKnown, RHSKnown);
	return;
	case MCBinaryExpr::Opcode::And:
	KBM[Expr] = LHSKnown & RHSKnown;
	return;
	case MCBinaryExpr::Opcode::Div:
	KBM[Expr] = KnownBits::sdiv(LHSKnown, RHSKnown);
	return;
	case MCBinaryExpr::Opcode::EQ: {
	std::optional<bool> CompareRes = KnownBits::eq(LHSKnown, RHSKnown);
	KBM[Expr] = fromOptionalToKnownBits(CompareRes);
	return;
	}
	case MCBinaryExpr::Opcode::NE: {
	std::optional<bool> CompareRes = KnownBits::ne(LHSKnown, RHSKnown);
	KBM[Expr] = fromOptionalToKnownBits(CompareRes);
	return;
	}
	case MCBinaryExpr::Opcode::GT: {
	std::optional<bool> CompareRes = KnownBits::sgt(LHSKnown, RHSKnown);
	KBM[Expr] = fromOptionalToKnownBits(CompareRes);
	return;
	}
	case MCBinaryExpr::Opcode::GTE: {
	std::optional<bool> CompareRes = KnownBits::sge(LHSKnown, RHSKnown);
	KBM[Expr] = fromOptionalToKnownBits(CompareRes);
	return;
	}
	case MCBinaryExpr::Opcode::LAnd: {
	std::optional<bool> CompareRes;
	const APInt False(BitWidth, 0);
	std::optional<bool> LHSBool =
	KnownBits::ne(LHSKnown, KnownBits::makeConstant(False));
	std::optional<bool> RHSBool =
	KnownBits::ne(RHSKnown, KnownBits::makeConstant(False));
	if (LHSBool && RHSBool)
	CompareRes = LHSBool && RHSBool;
	KBM[Expr] = fromOptionalToKnownBits(CompareRes);
	return;
	}
	case MCBinaryExpr::Opcode::LOr: {
	const APInt False(BitWidth, 0);
	KnownBits Bits = LHSKnown \| RHSKnown;
	std::optional<bool> CompareRes =
	KnownBits::ne(Bits, KnownBits::makeConstant(False));
	KBM[Expr] = fromOptionalToKnownBits(CompareRes);
	return;
	}
	case MCBinaryExpr::Opcode::LT: {
	std::optional<bool> CompareRes = KnownBits::slt(LHSKnown, RHSKnown);
	KBM[Expr] = fromOptionalToKnownBits(CompareRes);
	return;
	}
	case MCBinaryExpr::Opcode::LTE: {
	std::optional<bool> CompareRes = KnownBits::sle(LHSKnown, RHSKnown);
	KBM[Expr] = fromOptionalToKnownBits(CompareRes);
	return;
	}
	case MCBinaryExpr::Opcode::Mod:
	KBM[Expr] = KnownBits::srem(LHSKnown, RHSKnown);
	return;
	case MCBinaryExpr::Opcode::Mul:
	KBM[Expr] = KnownBits::mul(LHSKnown, RHSKnown);
	return;
	case MCBinaryExpr::Opcode::Or:
	KBM[Expr] = LHSKnown \| RHSKnown;
	return;
	case MCBinaryExpr::Opcode::Shl:
	KBM[Expr] = KnownBits::shl(LHSKnown, RHSKnown);
	return;
	case MCBinaryExpr::Opcode::AShr:
	KBM[Expr] = KnownBits::ashr(LHSKnown, RHSKnown);
	return;
	case MCBinaryExpr::Opcode::LShr:
	KBM[Expr] = KnownBits::lshr(LHSKnown, RHSKnown);
	return;
	case MCBinaryExpr::Opcode::Sub:
	KBM[Expr] = KnownBits::sub(LHSKnown, RHSKnown);
	return;
	case MCBinaryExpr::Opcode::Xor:
	KBM[Expr] = LHSKnown ^ RHSKnown;
	return;
	}
	}

	static void unaryOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
	unsigned Depth) {
	static constexpr unsigned BitWidth = 64;
	const MCUnaryExpr *UExpr = cast<MCUnaryExpr>(Expr);
	knownBitsMapHelper(UExpr->getSubExpr(), KBM, Depth + 1);
	KnownBits KB = KBM[UExpr->getSubExpr()];

	switch (UExpr->getOpcode()) {
	default:
	KBM[Expr] = KnownBits(BitWidth);
	return;
	case MCUnaryExpr::Opcode::Minus: {
	KB.makeNegative();
	KBM[Expr] = KB;
	return;
	}
	case MCUnaryExpr::Opcode::Not: {
	KnownBits AllOnes(BitWidth);
	AllOnes.setAllOnes();
	KBM[Expr] = KB ^ AllOnes;
	return;
	}
	case MCUnaryExpr::Opcode::Plus: {
	KB.makeNonNegative();
	KBM[Expr] = KB;
	return;
	}
	}
	}

	static void targetOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
	unsigned Depth) {
	static constexpr unsigned BitWidth = 64;
	const AMDGPUMCExpr *AGVK = cast<AMDGPUMCExpr>(Expr);

	switch (AGVK->getKind()) {
	default:
	KBM[Expr] = KnownBits(BitWidth);
	return;
	case AMDGPUMCExpr::VariantKind::AGVK_Or: {
	knownBitsMapHelper(AGVK->getSubExpr(0), KBM, Depth + 1);
	KnownBits KB = KBM[AGVK->getSubExpr(0)];
	for (const MCExpr *Arg : AGVK->getArgs()) {
	knownBitsMapHelper(Arg, KBM, Depth + 1);
	KB \|= KBM[Arg];
	}
	KBM[Expr] = KB;
	return;
	}
	case AMDGPUMCExpr::VariantKind::AGVK_Max: {
	knownBitsMapHelper(AGVK->getSubExpr(0), KBM, Depth + 1);
	KnownBits KB = KBM[AGVK->getSubExpr(0)];
	for (const MCExpr *Arg : AGVK->getArgs()) {
	knownBitsMapHelper(Arg, KBM, Depth + 1);
	KB = KnownBits::umax(KB, KBM[Arg]);
	}
	KBM[Expr] = KB;
	return;
	}
	case AMDGPUMCExpr::VariantKind::AGVK_ExtraSGPRs:
	case AMDGPUMCExpr::VariantKind::AGVK_TotalNumVGPRs:
	case AMDGPUMCExpr::VariantKind::AGVK_AlignTo:
	case AMDGPUMCExpr::VariantKind::AGVK_Occupancy: {
	int64_t Val;
	if (AGVK->evaluateAsAbsolute(Val)) {
	APInt APValue(BitWidth, Val);
	KBM[Expr] = KnownBits::makeConstant(APValue);
	return;
	}
	KBM[Expr] = KnownBits(BitWidth);
	return;
	}
	}
	}

	static void knownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
	unsigned Depth) {
	static constexpr unsigned BitWidth = 64;

	int64_t Val;
	if (Expr->evaluateAsAbsolute(Val)) {
	APInt APValue(BitWidth, Val, /isSigned=/true);
	KBM[Expr] = KnownBits::makeConstant(APValue);
	return;
	}

	if (Depth == 16) {
	KBM[Expr] = KnownBits(BitWidth);
	return;
	}

	switch (Expr->getKind()) {
	case MCExpr::ExprKind::Binary: {
	binaryOpKnownBitsMapHelper(Expr, KBM, Depth);
	return;
	}
	case MCExpr::ExprKind::Constant: {
	const MCConstantExpr *CE = cast<MCConstantExpr>(Expr);
	APInt APValue(BitWidth, CE->getValue(), /isSigned=/true);
	KBM[Expr] = KnownBits::makeConstant(APValue);
	return;
	}
	case MCExpr::ExprKind::SymbolRef: {
	const MCSymbolRefExpr *RExpr = cast<MCSymbolRefExpr>(Expr);
	const MCSymbol &Sym = RExpr->getSymbol();
	if (!Sym.isVariable()) {
	KBM[Expr] = KnownBits(BitWidth);
	return;
	}

	// Variable value retrieval is not for actual use but only for knownbits
	// analysis.
	const MCExpr SymVal = Sym.getVariableValue(/setUsed=*/false);
	knownBitsMapHelper(SymVal, KBM, Depth + 1);

	// Explicitly copy-construct so that there exists a local KnownBits in case
	// KBM[SymVal] gets invalidated after a potential growth through KBM[Expr].
	KBM[Expr] = KnownBits(KBM[SymVal]);
	return;
	}
	case MCExpr::ExprKind::Unary: {
	unaryOpKnownBitsMapHelper(Expr, KBM, Depth);
	return;
	}
	case MCExpr::ExprKind::Target: {
	targetOpKnownBitsMapHelper(Expr, KBM, Depth);
	return;
	}
	}
	}

	static const MCExpr tryFoldHelper(const MCExpr Expr, KnownBitsMap &KBM,
	MCContext &Ctx) {
	if (!KBM.count(Expr))
	return Expr;

	auto ValueCheckKnownBits = [](KnownBits &KB, unsigned Value) -> bool {
	if (!KB.isConstant())
	return false;

	return Value == KB.getConstant();
	};

	if (Expr->getKind() == MCExpr::ExprKind::Constant)
	return Expr;

	// Resolving unary operations to constants may make the value more ambiguous.
	// For example, `~62` becomes `-63`; however, to me it's more ambiguous if a
	// bit mask value is represented through a negative number.
	if (Expr->getKind() != MCExpr::ExprKind::Unary) {
	if (KBM[Expr].isConstant()) {
	APInt ConstVal = KBM[Expr].getConstant();
	return MCConstantExpr::create(ConstVal.getSExtValue(), Ctx);
	}

	int64_t EvalValue;
	if (Expr->evaluateAsAbsolute(EvalValue))
	return MCConstantExpr::create(EvalValue, Ctx);
	}

	switch (Expr->getKind()) {
	default:
	return Expr;
	case MCExpr::ExprKind::Binary: {
	const MCBinaryExpr *BExpr = cast<MCBinaryExpr>(Expr);
	const MCExpr *LHS = BExpr->getLHS();
	const MCExpr *RHS = BExpr->getRHS();

	switch (BExpr->getOpcode()) {
	default:
	return Expr;
	case MCBinaryExpr::Opcode::Sub: {
	if (ValueCheckKnownBits(KBM[RHS], 0))
	return tryFoldHelper(LHS, KBM, Ctx);
	break;
	}
	case MCBinaryExpr::Opcode::Add:
	case MCBinaryExpr::Opcode::Or: {
	if (ValueCheckKnownBits(KBM[LHS], 0))
	return tryFoldHelper(RHS, KBM, Ctx);
	if (ValueCheckKnownBits(KBM[RHS], 0))
	return tryFoldHelper(LHS, KBM, Ctx);
	break;
	}
	case MCBinaryExpr::Opcode::Mul: {
	if (ValueCheckKnownBits(KBM[LHS], 1))
	return tryFoldHelper(RHS, KBM, Ctx);
	if (ValueCheckKnownBits(KBM[RHS], 1))
	return tryFoldHelper(LHS, KBM, Ctx);
	break;
	}
	case MCBinaryExpr::Opcode::Shl:
	case MCBinaryExpr::Opcode::AShr:
	case MCBinaryExpr::Opcode::LShr: {
	if (ValueCheckKnownBits(KBM[RHS], 0))
	return tryFoldHelper(LHS, KBM, Ctx);
	if (ValueCheckKnownBits(KBM[LHS], 0))
	return MCConstantExpr::create(0, Ctx);
	break;
	}
	case MCBinaryExpr::Opcode::And: {
	if (ValueCheckKnownBits(KBM[LHS], 0) \|\| ValueCheckKnownBits(KBM[RHS], 0))
	return MCConstantExpr::create(0, Ctx);
	break;
	}
	}
	const MCExpr *NewLHS = tryFoldHelper(LHS, KBM, Ctx);
	const MCExpr *NewRHS = tryFoldHelper(RHS, KBM, Ctx);
	if (NewLHS != LHS \|\| NewRHS != RHS)
	return MCBinaryExpr::create(BExpr->getOpcode(), NewLHS, NewRHS, Ctx,
	BExpr->getLoc());
	return Expr;
	}
	case MCExpr::ExprKind::Unary: {
	const MCUnaryExpr *UExpr = cast<MCUnaryExpr>(Expr);
	const MCExpr *SubExpr = UExpr->getSubExpr();
	const MCExpr *NewSubExpr = tryFoldHelper(SubExpr, KBM, Ctx);
	if (SubExpr != NewSubExpr)
	return MCUnaryExpr::create(UExpr->getOpcode(), NewSubExpr, Ctx,
	UExpr->getLoc());
	return Expr;
	}
	case MCExpr::ExprKind::Target: {
	const AMDGPUMCExpr *AGVK = cast<AMDGPUMCExpr>(Expr);
	SmallVector<const MCExpr *, 8> NewArgs;
	bool Changed = false;
	for (const MCExpr *Arg : AGVK->getArgs()) {
	const MCExpr *NewArg = tryFoldHelper(Arg, KBM, Ctx);
	NewArgs.push_back(NewArg);
	Changed \|= Arg != NewArg;
	}
	return Changed ? AMDGPUMCExpr::create(AGVK->getKind(), NewArgs, Ctx) : Expr;
	}
	}
	return Expr;
	}

	const MCExpr llvm::AMDGPU::foldAMDGPUMCExpr(const MCExpr Expr,
	MCContext &Ctx) {
	KnownBitsMap KBM;
	knownBitsMapHelper(Expr, KBM);
	const MCExpr *NewExpr = tryFoldHelper(Expr, KBM, Ctx);

	return Expr != NewExpr ? NewExpr : Expr;
	}

	void llvm::AMDGPU::printAMDGPUMCExpr(const MCExpr *Expr, raw_ostream &OS,
	const MCAsmInfo *MAI) {
	int64_t Val;
	if (Expr->evaluateAsAbsolute(Val)) {
	OS << Val;
	return;
	}

	Expr->print(OS, MAI);
	}