xref: /llvm-project/llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUMCExpr.cpp (revision bd9145c8c21334e099d51b3e66f49d51d24931ee) 
1 //===- AMDGPUMCExpr.cpp - AMDGPU specific MC expression classes -----------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "AMDGPUMCExpr.h"
10 #include "GCNSubtarget.h"
11 #include "Utils/AMDGPUBaseInfo.h"
12 #include "llvm/IR/Function.h"
13 #include "llvm/MC/MCAsmInfo.h"
14 #include "llvm/MC/MCAssembler.h"
15 #include "llvm/MC/MCContext.h"
16 #include "llvm/MC/MCStreamer.h"
17 #include "llvm/MC/MCSymbol.h"
18 #include "llvm/MC/MCValue.h"
19 #include "llvm/Support/KnownBits.h"
20 #include "llvm/Support/raw_ostream.h"
21 #include <optional>
22 
23 using namespace llvm;
24 using namespace llvm::AMDGPU;
25 
26 AMDGPUMCExpr::AMDGPUMCExpr(VariantKind Kind, ArrayRef<const MCExpr *> Args,
27                            MCContext &Ctx)
28     : Kind(Kind), Ctx(Ctx) {
29   assert(Args.size() >= 1 && "Needs a minimum of one expression.");
30   assert(Kind != AGVK_None && "Cannot construct AMDGPUMCExpr of kind none.");
31 
32   // Allocating the variadic arguments through the same allocation mechanism
33   // that the object itself is allocated with so they end up in the same memory.
34   //
35   // Will result in an asan failure if allocated on the heap through standard
36   // allocation (e.g., through SmallVector's grow).
37   RawArgs = static_cast<const MCExpr **>(
38       Ctx.allocate(sizeof(const MCExpr *) * Args.size()));
39   std::uninitialized_copy(Args.begin(), Args.end(), RawArgs);
40   this->Args = ArrayRef<const MCExpr *>(RawArgs, Args.size());
41 }
42 
43 AMDGPUMCExpr::~AMDGPUMCExpr() { Ctx.deallocate(RawArgs); }
44 
45 const AMDGPUMCExpr *AMDGPUMCExpr::create(VariantKind Kind,
46                                          ArrayRef<const MCExpr *> Args,
47                                          MCContext &Ctx) {
48   return new (Ctx) AMDGPUMCExpr(Kind, Args, Ctx);
49 }
50 
51 const MCExpr *AMDGPUMCExpr::getSubExpr(size_t Index) const {
52   assert(Index < Args.size() && "Indexing out of bounds AMDGPUMCExpr sub-expr");
53   return Args[Index];
54 }
55 
56 void AMDGPUMCExpr::printImpl(raw_ostream &OS, const MCAsmInfo *MAI) const {
57   switch (Kind) {
58   default:
59     llvm_unreachable("Unknown AMDGPUMCExpr kind.");
60   case AGVK_Or:
61     OS << "or(";
62     break;
63   case AGVK_Max:
64     OS << "max(";
65     break;
66   case AGVK_ExtraSGPRs:
67     OS << "extrasgprs(";
68     break;
69   case AGVK_TotalNumVGPRs:
70     OS << "totalnumvgprs(";
71     break;
72   case AGVK_AlignTo:
73     OS << "alignto(";
74     break;
75   case AGVK_Occupancy:
76     OS << "occupancy(";
77     break;
78   }
79   for (const auto *It = Args.begin(); It != Args.end(); ++It) {
80     (*It)->print(OS, MAI, /*InParens=*/false);
81     if ((It + 1) != Args.end())
82       OS << ", ";
83   }
84   OS << ')';
85 }
86 
87 static int64_t op(AMDGPUMCExpr::VariantKind Kind, int64_t Arg1, int64_t Arg2) {
88   switch (Kind) {
89   default:
90     llvm_unreachable("Unknown AMDGPUMCExpr kind.");
91   case AMDGPUMCExpr::AGVK_Max:
92     return std::max(Arg1, Arg2);
93   case AMDGPUMCExpr::AGVK_Or:
94     return Arg1 | Arg2;
95   }
96 }
97 
98 bool AMDGPUMCExpr::evaluateExtraSGPRs(MCValue &Res, const MCAssembler *Asm,
99                                       const MCFixup *Fixup) const {
100   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
101     MCValue MCVal;
102     if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) || !MCVal.isAbsolute())
103       return false;
104 
105     ConstantValue = MCVal.getConstant();
106     return true;
107   };
108 
109   assert(Args.size() == 3 &&
110          "AMDGPUMCExpr Argument count incorrect for ExtraSGPRs");
111   const MCSubtargetInfo *STI = Ctx.getSubtargetInfo();
112   uint64_t VCCUsed = 0, FlatScrUsed = 0, XNACKUsed = 0;
113 
114   bool Success = TryGetMCExprValue(Args[2], XNACKUsed);
115 
116   assert(Success && "Arguments 3 for ExtraSGPRs should be a known constant");
117   if (!Success || !TryGetMCExprValue(Args[0], VCCUsed) ||
118       !TryGetMCExprValue(Args[1], FlatScrUsed))
119     return false;
120 
121   uint64_t ExtraSGPRs = IsaInfo::getNumExtraSGPRs(
122       STI, (bool)VCCUsed, (bool)FlatScrUsed, (bool)XNACKUsed);
123   Res = MCValue::get(ExtraSGPRs);
124   return true;
125 }
126 
127 bool AMDGPUMCExpr::evaluateTotalNumVGPR(MCValue &Res, const MCAssembler *Asm,
128                                         const MCFixup *Fixup) const {
129   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
130     MCValue MCVal;
131     if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) || !MCVal.isAbsolute())
132       return false;
133 
134     ConstantValue = MCVal.getConstant();
135     return true;
136   };
137   assert(Args.size() == 2 &&
138          "AMDGPUMCExpr Argument count incorrect for TotalNumVGPRs");
139   const MCSubtargetInfo *STI = Ctx.getSubtargetInfo();
140   uint64_t NumAGPR = 0, NumVGPR = 0;
141 
142   bool Has90AInsts = AMDGPU::isGFX90A(*STI);
143 
144   if (!TryGetMCExprValue(Args[0], NumAGPR) ||
145       !TryGetMCExprValue(Args[1], NumVGPR))
146     return false;
147 
148   uint64_t TotalNum = Has90AInsts && NumAGPR ? alignTo(NumVGPR, 4) + NumAGPR
149                                              : std::max(NumVGPR, NumAGPR);
150   Res = MCValue::get(TotalNum);
151   return true;
152 }
153 
154 bool AMDGPUMCExpr::evaluateAlignTo(MCValue &Res, const MCAssembler *Asm,
155                                    const MCFixup *Fixup) const {
156   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
157     MCValue MCVal;
158     if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) || !MCVal.isAbsolute())
159       return false;
160 
161     ConstantValue = MCVal.getConstant();
162     return true;
163   };
164 
165   assert(Args.size() == 2 &&
166          "AMDGPUMCExpr Argument count incorrect for AlignTo");
167   uint64_t Value = 0, Align = 0;
168   if (!TryGetMCExprValue(Args[0], Value) || !TryGetMCExprValue(Args[1], Align))
169     return false;
170 
171   Res = MCValue::get(alignTo(Value, Align));
172   return true;
173 }
174 
175 bool AMDGPUMCExpr::evaluateOccupancy(MCValue &Res, const MCAssembler *Asm,
176                                      const MCFixup *Fixup) const {
177   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
178     MCValue MCVal;
179     if (!Arg->evaluateAsRelocatable(MCVal, Asm, Fixup) || !MCVal.isAbsolute())
180       return false;
181 
182     ConstantValue = MCVal.getConstant();
183     return true;
184   };
185   assert(Args.size() == 7 &&
186          "AMDGPUMCExpr Argument count incorrect for Occupancy");
187   uint64_t InitOccupancy, MaxWaves, Granule, TargetTotalNumVGPRs, Generation,
188       NumSGPRs, NumVGPRs;
189 
190   bool Success = true;
191   Success &= TryGetMCExprValue(Args[0], MaxWaves);
192   Success &= TryGetMCExprValue(Args[1], Granule);
193   Success &= TryGetMCExprValue(Args[2], TargetTotalNumVGPRs);
194   Success &= TryGetMCExprValue(Args[3], Generation);
195   Success &= TryGetMCExprValue(Args[4], InitOccupancy);
196 
197   assert(Success && "Arguments 1 to 5 for Occupancy should be known constants");
198 
199   if (!Success || !TryGetMCExprValue(Args[5], NumSGPRs) ||
200       !TryGetMCExprValue(Args[6], NumVGPRs))
201     return false;
202 
203   unsigned Occupancy = InitOccupancy;
204   if (NumSGPRs)
205     Occupancy = std::min(
206         Occupancy, IsaInfo::getOccupancyWithNumSGPRs(
207                        NumSGPRs, MaxWaves,
208                        static_cast<AMDGPUSubtarget::Generation>(Generation)));
209   if (NumVGPRs)
210     Occupancy = std::min(Occupancy,
211                          IsaInfo::getNumWavesPerEUWithNumVGPRs(
212                              NumVGPRs, Granule, MaxWaves, TargetTotalNumVGPRs));
213 
214   Res = MCValue::get(Occupancy);
215   return true;
216 }
217 
218 bool AMDGPUMCExpr::evaluateAsRelocatableImpl(MCValue &Res,
219                                              const MCAssembler *Asm,
220                                              const MCFixup *Fixup) const {
221   std::optional<int64_t> Total;
222   switch (Kind) {
223   default:
224     break;
225   case AGVK_ExtraSGPRs:
226     return evaluateExtraSGPRs(Res, Asm, Fixup);
227   case AGVK_AlignTo:
228     return evaluateAlignTo(Res, Asm, Fixup);
229   case AGVK_TotalNumVGPRs:
230     return evaluateTotalNumVGPR(Res, Asm, Fixup);
231   case AGVK_Occupancy:
232     return evaluateOccupancy(Res, Asm, Fixup);
233   }
234 
235   for (const MCExpr *Arg : Args) {
236     MCValue ArgRes;
237     if (!Arg->evaluateAsRelocatable(ArgRes, Asm, Fixup) || !ArgRes.isAbsolute())
238       return false;
239 
240     if (!Total.has_value())
241       Total = ArgRes.getConstant();
242     Total = op(Kind, *Total, ArgRes.getConstant());
243   }
244 
245   Res = MCValue::get(*Total);
246   return true;
247 }
248 
249 void AMDGPUMCExpr::visitUsedExpr(MCStreamer &Streamer) const {
250   for (const MCExpr *Arg : Args)
251     Streamer.visitUsedExpr(*Arg);
252 }
253 
254 MCFragment *AMDGPUMCExpr::findAssociatedFragment() const {
255   for (const MCExpr *Arg : Args) {
256     if (Arg->findAssociatedFragment())
257       return Arg->findAssociatedFragment();
258   }
259   return nullptr;
260 }
261 
262 /// Allow delayed MCExpr resolve of ExtraSGPRs (in case VCCUsed or FlatScrUsed
263 /// are unresolvable but needed for further MCExprs). Derived from
264 /// implementation of IsaInfo::getNumExtraSGPRs in AMDGPUBaseInfo.cpp.
265 ///
266 const AMDGPUMCExpr *AMDGPUMCExpr::createExtraSGPRs(const MCExpr *VCCUsed,
267                                                    const MCExpr *FlatScrUsed,
268                                                    bool XNACKUsed,
269                                                    MCContext &Ctx) {
270 
271   return create(AGVK_ExtraSGPRs,
272                 {VCCUsed, FlatScrUsed, MCConstantExpr::create(XNACKUsed, Ctx)},
273                 Ctx);
274 }
275 
276 const AMDGPUMCExpr *AMDGPUMCExpr::createTotalNumVGPR(const MCExpr *NumAGPR,
277                                                      const MCExpr *NumVGPR,
278                                                      MCContext &Ctx) {
279   return create(AGVK_TotalNumVGPRs, {NumAGPR, NumVGPR}, Ctx);
280 }
281 
282 /// Mimics GCNSubtarget::computeOccupancy for MCExpr.
283 ///
284 /// Remove dependency on GCNSubtarget and depend only only the necessary values
285 /// for said occupancy computation. Should match computeOccupancy implementation
286 /// without passing \p STM on.
287 const AMDGPUMCExpr *AMDGPUMCExpr::createOccupancy(unsigned InitOcc,
288                                                   const MCExpr *NumSGPRs,
289                                                   const MCExpr *NumVGPRs,
290                                                   const GCNSubtarget &STM,
291                                                   MCContext &Ctx) {
292   unsigned MaxWaves = IsaInfo::getMaxWavesPerEU(&STM);
293   unsigned Granule = IsaInfo::getVGPRAllocGranule(&STM);
294   unsigned TargetTotalNumVGPRs = IsaInfo::getTotalNumVGPRs(&STM);
295   unsigned Generation = STM.getGeneration();
296 
297   auto CreateExpr = [&Ctx](unsigned Value) {
298     return MCConstantExpr::create(Value, Ctx);
299   };
300 
301   return create(AGVK_Occupancy,
302                 {CreateExpr(MaxWaves), CreateExpr(Granule),
303                  CreateExpr(TargetTotalNumVGPRs), CreateExpr(Generation),
304                  CreateExpr(InitOcc), NumSGPRs, NumVGPRs},
305                 Ctx);
306 }
307 
308 bool AMDGPUMCExpr::isSymbolUsedInExpression(const MCSymbol *Sym) const {
309   for (const MCExpr *E : getArgs()) {
310     if (E->isSymbolUsedInExpression(Sym))
311       return true;
312   }
313   return false;
314 }
315 
316 static KnownBits fromOptionalToKnownBits(std::optional<bool> CompareResult) {
317   static constexpr unsigned BitWidth = 64;
318   const APInt True(BitWidth, 1);
319   const APInt False(BitWidth, 0);
320   if (CompareResult) {
321     return *CompareResult ? KnownBits::makeConstant(True)
322                           : KnownBits::makeConstant(False);
323   }
324 
325   KnownBits UnknownBool(/*BitWidth=*/1);
326   return UnknownBool.zext(BitWidth);
327 }
328 
329 using KnownBitsMap = DenseMap<const MCExpr *, KnownBits>;
330 static void knownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
331                                unsigned Depth = 0);
332 
333 static void binaryOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
334                                        unsigned Depth) {
335   static constexpr unsigned BitWidth = 64;
336   const MCBinaryExpr *BExpr = cast<MCBinaryExpr>(Expr);
337   const MCExpr *LHS = BExpr->getLHS();
338   const MCExpr *RHS = BExpr->getRHS();
339 
340   knownBitsMapHelper(LHS, KBM, Depth + 1);
341   knownBitsMapHelper(RHS, KBM, Depth + 1);
342   KnownBits LHSKnown = KBM[LHS];
343   KnownBits RHSKnown = KBM[RHS];
344 
345   switch (BExpr->getOpcode()) {
346   default:
347     KBM[Expr] = KnownBits(BitWidth);
348     return;
349   case MCBinaryExpr::Opcode::Add:
350     KBM[Expr] = KnownBits::add(LHSKnown, RHSKnown);
351     return;
352   case MCBinaryExpr::Opcode::And:
353     KBM[Expr] = LHSKnown & RHSKnown;
354     return;
355   case MCBinaryExpr::Opcode::Div:
356     KBM[Expr] = KnownBits::sdiv(LHSKnown, RHSKnown);
357     return;
358   case MCBinaryExpr::Opcode::EQ: {
359     std::optional<bool> CompareRes = KnownBits::eq(LHSKnown, RHSKnown);
360     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
361     return;
362   }
363   case MCBinaryExpr::Opcode::NE: {
364     std::optional<bool> CompareRes = KnownBits::ne(LHSKnown, RHSKnown);
365     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
366     return;
367   }
368   case MCBinaryExpr::Opcode::GT: {
369     std::optional<bool> CompareRes = KnownBits::sgt(LHSKnown, RHSKnown);
370     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
371     return;
372   }
373   case MCBinaryExpr::Opcode::GTE: {
374     std::optional<bool> CompareRes = KnownBits::sge(LHSKnown, RHSKnown);
375     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
376     return;
377   }
378   case MCBinaryExpr::Opcode::LAnd: {
379     std::optional<bool> CompareRes;
380     const APInt False(BitWidth, 0);
381     std::optional<bool> LHSBool =
382         KnownBits::ne(LHSKnown, KnownBits::makeConstant(False));
383     std::optional<bool> RHSBool =
384         KnownBits::ne(RHSKnown, KnownBits::makeConstant(False));
385     if (LHSBool && RHSBool)
386       CompareRes = *LHSBool && *RHSBool;
387     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
388     return;
389   }
390   case MCBinaryExpr::Opcode::LOr: {
391     const APInt False(BitWidth, 0);
392     KnownBits Bits = LHSKnown | RHSKnown;
393     std::optional<bool> CompareRes =
394         KnownBits::ne(Bits, KnownBits::makeConstant(False));
395     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
396     return;
397   }
398   case MCBinaryExpr::Opcode::LT: {
399     std::optional<bool> CompareRes = KnownBits::slt(LHSKnown, RHSKnown);
400     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
401     return;
402   }
403   case MCBinaryExpr::Opcode::LTE: {
404     std::optional<bool> CompareRes = KnownBits::sle(LHSKnown, RHSKnown);
405     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
406     return;
407   }
408   case MCBinaryExpr::Opcode::Mod:
409     KBM[Expr] = KnownBits::srem(LHSKnown, RHSKnown);
410     return;
411   case MCBinaryExpr::Opcode::Mul:
412     KBM[Expr] = KnownBits::mul(LHSKnown, RHSKnown);
413     return;
414   case MCBinaryExpr::Opcode::Or:
415     KBM[Expr] = LHSKnown | RHSKnown;
416     return;
417   case MCBinaryExpr::Opcode::Shl:
418     KBM[Expr] = KnownBits::shl(LHSKnown, RHSKnown);
419     return;
420   case MCBinaryExpr::Opcode::AShr:
421     KBM[Expr] = KnownBits::ashr(LHSKnown, RHSKnown);
422     return;
423   case MCBinaryExpr::Opcode::LShr:
424     KBM[Expr] = KnownBits::lshr(LHSKnown, RHSKnown);
425     return;
426   case MCBinaryExpr::Opcode::Sub:
427     KBM[Expr] = KnownBits::sub(LHSKnown, RHSKnown);
428     return;
429   case MCBinaryExpr::Opcode::Xor:
430     KBM[Expr] = LHSKnown ^ RHSKnown;
431     return;
432   }
433 }
434 
435 static void unaryOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
436                                       unsigned Depth) {
437   static constexpr unsigned BitWidth = 64;
438   const MCUnaryExpr *UExpr = cast<MCUnaryExpr>(Expr);
439   knownBitsMapHelper(UExpr->getSubExpr(), KBM, Depth + 1);
440   KnownBits KB = KBM[UExpr->getSubExpr()];
441 
442   switch (UExpr->getOpcode()) {
443   default:
444     KBM[Expr] = KnownBits(BitWidth);
445     return;
446   case MCUnaryExpr::Opcode::Minus: {
447     KB.makeNegative();
448     KBM[Expr] = KB;
449     return;
450   }
451   case MCUnaryExpr::Opcode::Not: {
452     KnownBits AllOnes(BitWidth);
453     AllOnes.setAllOnes();
454     KBM[Expr] = KB ^ AllOnes;
455     return;
456   }
457   case MCUnaryExpr::Opcode::Plus: {
458     KB.makeNonNegative();
459     KBM[Expr] = KB;
460     return;
461   }
462   }
463 }
464 
465 static void targetOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
466                                        unsigned Depth) {
467   static constexpr unsigned BitWidth = 64;
468   const AMDGPUMCExpr *AGVK = cast<AMDGPUMCExpr>(Expr);
469 
470   switch (AGVK->getKind()) {
471   default:
472     KBM[Expr] = KnownBits(BitWidth);
473     return;
474   case AMDGPUMCExpr::VariantKind::AGVK_Or: {
475     knownBitsMapHelper(AGVK->getSubExpr(0), KBM, Depth + 1);
476     KnownBits KB = KBM[AGVK->getSubExpr(0)];
477     for (const MCExpr *Arg : AGVK->getArgs()) {
478       knownBitsMapHelper(Arg, KBM, Depth + 1);
479       KB |= KBM[Arg];
480     }
481     KBM[Expr] = KB;
482     return;
483   }
484   case AMDGPUMCExpr::VariantKind::AGVK_Max: {
485     knownBitsMapHelper(AGVK->getSubExpr(0), KBM, Depth + 1);
486     KnownBits KB = KBM[AGVK->getSubExpr(0)];
487     for (const MCExpr *Arg : AGVK->getArgs()) {
488       knownBitsMapHelper(Arg, KBM, Depth + 1);
489       KB = KnownBits::umax(KB, KBM[Arg]);
490     }
491     KBM[Expr] = KB;
492     return;
493   }
494   case AMDGPUMCExpr::VariantKind::AGVK_ExtraSGPRs:
495   case AMDGPUMCExpr::VariantKind::AGVK_TotalNumVGPRs:
496   case AMDGPUMCExpr::VariantKind::AGVK_AlignTo:
497   case AMDGPUMCExpr::VariantKind::AGVK_Occupancy: {
498     int64_t Val;
499     if (AGVK->evaluateAsAbsolute(Val)) {
500       APInt APValue(BitWidth, Val);
501       KBM[Expr] = KnownBits::makeConstant(APValue);
502       return;
503     }
504     KBM[Expr] = KnownBits(BitWidth);
505     return;
506   }
507   }
508 }
509 
510 static void knownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
511                                unsigned Depth) {
512   static constexpr unsigned BitWidth = 64;
513 
514   int64_t Val;
515   if (Expr->evaluateAsAbsolute(Val)) {
516     APInt APValue(BitWidth, Val, /*isSigned=*/true);
517     KBM[Expr] = KnownBits::makeConstant(APValue);
518     return;
519   }
520 
521   if (Depth == 16) {
522     KBM[Expr] = KnownBits(BitWidth);
523     return;
524   }
525 
526   switch (Expr->getKind()) {
527   case MCExpr::ExprKind::Binary: {
528     binaryOpKnownBitsMapHelper(Expr, KBM, Depth);
529     return;
530   }
531   case MCExpr::ExprKind::Constant: {
532     const MCConstantExpr *CE = cast<MCConstantExpr>(Expr);
533     APInt APValue(BitWidth, CE->getValue(), /*isSigned=*/true);
534     KBM[Expr] = KnownBits::makeConstant(APValue);
535     return;
536   }
537   case MCExpr::ExprKind::SymbolRef: {
538     const MCSymbolRefExpr *RExpr = cast<MCSymbolRefExpr>(Expr);
539     const MCSymbol &Sym = RExpr->getSymbol();
540     if (!Sym.isVariable()) {
541       KBM[Expr] = KnownBits(BitWidth);
542       return;
543     }
544 
545     // Variable value retrieval is not for actual use but only for knownbits
546     // analysis.
547     const MCExpr *SymVal = Sym.getVariableValue(/*setUsed=*/false);
548     knownBitsMapHelper(SymVal, KBM, Depth + 1);
549 
550     // Explicitly copy-construct so that there exists a local KnownBits in case
551     // KBM[SymVal] gets invalidated after a potential growth through KBM[Expr].
552     KBM[Expr] = KnownBits(KBM[SymVal]);
553     return;
554   }
555   case MCExpr::ExprKind::Unary: {
556     unaryOpKnownBitsMapHelper(Expr, KBM, Depth);
557     return;
558   }
559   case MCExpr::ExprKind::Target: {
560     targetOpKnownBitsMapHelper(Expr, KBM, Depth);
561     return;
562   }
563   }
564 }
565 
566 static const MCExpr *tryFoldHelper(const MCExpr *Expr, KnownBitsMap &KBM,
567                                    MCContext &Ctx) {
568   if (!KBM.count(Expr))
569     return Expr;
570 
571   auto ValueCheckKnownBits = [](KnownBits &KB, unsigned Value) -> bool {
572     if (!KB.isConstant())
573       return false;
574 
575     return Value == KB.getConstant();
576   };
577 
578   if (Expr->getKind() == MCExpr::ExprKind::Constant)
579     return Expr;
580 
581   // Resolving unary operations to constants may make the value more ambiguous.
582   // For example, `~62` becomes `-63`; however, to me it's more ambiguous if a
583   // bit mask value is represented through a negative number.
584   if (Expr->getKind() != MCExpr::ExprKind::Unary) {
585     if (KBM[Expr].isConstant()) {
586       APInt ConstVal = KBM[Expr].getConstant();
587       return MCConstantExpr::create(ConstVal.getSExtValue(), Ctx);
588     }
589 
590     int64_t EvalValue;
591     if (Expr->evaluateAsAbsolute(EvalValue))
592       return MCConstantExpr::create(EvalValue, Ctx);
593   }
594 
595   switch (Expr->getKind()) {
596   default:
597     return Expr;
598   case MCExpr::ExprKind::Binary: {
599     const MCBinaryExpr *BExpr = cast<MCBinaryExpr>(Expr);
600     const MCExpr *LHS = BExpr->getLHS();
601     const MCExpr *RHS = BExpr->getRHS();
602 
603     switch (BExpr->getOpcode()) {
604     default:
605       return Expr;
606     case MCBinaryExpr::Opcode::Sub: {
607       if (ValueCheckKnownBits(KBM[RHS], 0))
608         return tryFoldHelper(LHS, KBM, Ctx);
609       break;
610     }
611     case MCBinaryExpr::Opcode::Add:
612     case MCBinaryExpr::Opcode::Or: {
613       if (ValueCheckKnownBits(KBM[LHS], 0))
614         return tryFoldHelper(RHS, KBM, Ctx);
615       if (ValueCheckKnownBits(KBM[RHS], 0))
616         return tryFoldHelper(LHS, KBM, Ctx);
617       break;
618     }
619     case MCBinaryExpr::Opcode::Mul: {
620       if (ValueCheckKnownBits(KBM[LHS], 1))
621         return tryFoldHelper(RHS, KBM, Ctx);
622       if (ValueCheckKnownBits(KBM[RHS], 1))
623         return tryFoldHelper(LHS, KBM, Ctx);
624       break;
625     }
626     case MCBinaryExpr::Opcode::Shl:
627     case MCBinaryExpr::Opcode::AShr:
628     case MCBinaryExpr::Opcode::LShr: {
629       if (ValueCheckKnownBits(KBM[RHS], 0))
630         return tryFoldHelper(LHS, KBM, Ctx);
631       if (ValueCheckKnownBits(KBM[LHS], 0))
632         return MCConstantExpr::create(0, Ctx);
633       break;
634     }
635     case MCBinaryExpr::Opcode::And: {
636       if (ValueCheckKnownBits(KBM[LHS], 0) || ValueCheckKnownBits(KBM[RHS], 0))
637         return MCConstantExpr::create(0, Ctx);
638       break;
639     }
640     }
641     const MCExpr *NewLHS = tryFoldHelper(LHS, KBM, Ctx);
642     const MCExpr *NewRHS = tryFoldHelper(RHS, KBM, Ctx);
643     if (NewLHS != LHS || NewRHS != RHS)
644       return MCBinaryExpr::create(BExpr->getOpcode(), NewLHS, NewRHS, Ctx,
645                                   BExpr->getLoc());
646     return Expr;
647   }
648   case MCExpr::ExprKind::Unary: {
649     const MCUnaryExpr *UExpr = cast<MCUnaryExpr>(Expr);
650     const MCExpr *SubExpr = UExpr->getSubExpr();
651     const MCExpr *NewSubExpr = tryFoldHelper(SubExpr, KBM, Ctx);
652     if (SubExpr != NewSubExpr)
653       return MCUnaryExpr::create(UExpr->getOpcode(), NewSubExpr, Ctx,
654                                  UExpr->getLoc());
655     return Expr;
656   }
657   case MCExpr::ExprKind::Target: {
658     const AMDGPUMCExpr *AGVK = cast<AMDGPUMCExpr>(Expr);
659     SmallVector<const MCExpr *, 8> NewArgs;
660     bool Changed = false;
661     for (const MCExpr *Arg : AGVK->getArgs()) {
662       const MCExpr *NewArg = tryFoldHelper(Arg, KBM, Ctx);
663       NewArgs.push_back(NewArg);
664       Changed |= Arg != NewArg;
665     }
666     return Changed ? AMDGPUMCExpr::create(AGVK->getKind(), NewArgs, Ctx) : Expr;
667   }
668   }
669   return Expr;
670 }
671 
672 const MCExpr *llvm::AMDGPU::foldAMDGPUMCExpr(const MCExpr *Expr,
673                                              MCContext &Ctx) {
674   KnownBitsMap KBM;
675   knownBitsMapHelper(Expr, KBM);
676   const MCExpr *NewExpr = tryFoldHelper(Expr, KBM, Ctx);
677 
678   return Expr != NewExpr ? NewExpr : Expr;
679 }
680 
681 void llvm::AMDGPU::printAMDGPUMCExpr(const MCExpr *Expr, raw_ostream &OS,
682                                      const MCAsmInfo *MAI) {
683   int64_t Val;
684   if (Expr->evaluateAsAbsolute(Val)) {
685     OS << Val;
686     return;
687   }
688 
689   Expr->print(OS, MAI);
690 }
691