xref: /freebsd/contrib/llvm-project/llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUMCExpr.cpp (revision 454322c08b8aa181939c8d920472f03cfd591032) 
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   llvm::uninitialized_copy(Args, 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     MAI->printExpr(OS, **It);
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,
99                                       const MCAssembler *Asm) const {
100   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
101     MCValue MCVal;
102     if (!Arg->evaluateAsRelocatable(MCVal, Asm) || !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,
128                                         const MCAssembler *Asm) const {
129   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
130     MCValue MCVal;
131     if (!Arg->evaluateAsRelocatable(MCVal, Asm) || !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) const {
155   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
156     MCValue MCVal;
157     if (!Arg->evaluateAsRelocatable(MCVal, Asm) || !MCVal.isAbsolute())
158       return false;
159 
160     ConstantValue = MCVal.getConstant();
161     return true;
162   };
163 
164   assert(Args.size() == 2 &&
165          "AMDGPUMCExpr Argument count incorrect for AlignTo");
166   uint64_t Value = 0, Align = 0;
167   if (!TryGetMCExprValue(Args[0], Value) || !TryGetMCExprValue(Args[1], Align))
168     return false;
169 
170   Res = MCValue::get(alignTo(Value, Align));
171   return true;
172 }
173 
174 bool AMDGPUMCExpr::evaluateOccupancy(MCValue &Res,
175                                      const MCAssembler *Asm) const {
176   auto TryGetMCExprValue = [&](const MCExpr *Arg, uint64_t &ConstantValue) {
177     MCValue MCVal;
178     if (!Arg->evaluateAsRelocatable(MCVal, Asm) || !MCVal.isAbsolute())
179       return false;
180 
181     ConstantValue = MCVal.getConstant();
182     return true;
183   };
184   assert(Args.size() == 7 &&
185          "AMDGPUMCExpr Argument count incorrect for Occupancy");
186   uint64_t InitOccupancy, MaxWaves, Granule, TargetTotalNumVGPRs, Generation,
187       NumSGPRs, NumVGPRs;
188 
189   bool Success = true;
190   Success &= TryGetMCExprValue(Args[0], MaxWaves);
191   Success &= TryGetMCExprValue(Args[1], Granule);
192   Success &= TryGetMCExprValue(Args[2], TargetTotalNumVGPRs);
193   Success &= TryGetMCExprValue(Args[3], Generation);
194   Success &= TryGetMCExprValue(Args[4], InitOccupancy);
195 
196   assert(Success && "Arguments 1 to 5 for Occupancy should be known constants");
197 
198   if (!Success || !TryGetMCExprValue(Args[5], NumSGPRs) ||
199       !TryGetMCExprValue(Args[6], NumVGPRs))
200     return false;
201 
202   unsigned Occupancy = InitOccupancy;
203   if (NumSGPRs)
204     Occupancy = std::min(
205         Occupancy, IsaInfo::getOccupancyWithNumSGPRs(
206                        NumSGPRs, MaxWaves,
207                        static_cast<AMDGPUSubtarget::Generation>(Generation)));
208   if (NumVGPRs)
209     Occupancy = std::min(Occupancy,
210                          IsaInfo::getNumWavesPerEUWithNumVGPRs(
211                              NumVGPRs, Granule, MaxWaves, TargetTotalNumVGPRs));
212 
213   Res = MCValue::get(Occupancy);
214   return true;
215 }
216 
217 bool AMDGPUMCExpr::isSymbolUsedInExpression(const MCSymbol *Sym,
218                                             const MCExpr *E) {
219   switch (E->getKind()) {
220   case MCExpr::Constant:
221     return false;
222   case MCExpr::Unary:
223     return isSymbolUsedInExpression(
224         Sym, static_cast<const MCUnaryExpr *>(E)->getSubExpr());
225   case MCExpr::Binary: {
226     const MCBinaryExpr *BE = static_cast<const MCBinaryExpr *>(E);
227     return isSymbolUsedInExpression(Sym, BE->getLHS()) ||
228            isSymbolUsedInExpression(Sym, BE->getRHS());
229   }
230   case MCExpr::SymbolRef: {
231     const MCSymbol &S = static_cast<const MCSymbolRefExpr *>(E)->getSymbol();
232     if (S.isVariable())
233       return isSymbolUsedInExpression(Sym, S.getVariableValue());
234     return &S == Sym;
235   }
236   case MCExpr::Specifier:
237   case MCExpr::Target: {
238     auto *TE = static_cast<const AMDGPUMCExpr *>(E);
239     for (const MCExpr *E : TE->getArgs())
240       if (isSymbolUsedInExpression(Sym, E))
241         return true;
242     return false;
243   }
244   }
245   llvm_unreachable("Unknown expr kind!");
246 }
247 
248 bool AMDGPUMCExpr::evaluateAsRelocatableImpl(MCValue &Res,
249                                              const MCAssembler *Asm) const {
250   std::optional<int64_t> Total;
251   switch (Kind) {
252   default:
253     break;
254   case AGVK_ExtraSGPRs:
255     return evaluateExtraSGPRs(Res, Asm);
256   case AGVK_AlignTo:
257     return evaluateAlignTo(Res, Asm);
258   case AGVK_TotalNumVGPRs:
259     return evaluateTotalNumVGPR(Res, Asm);
260   case AGVK_Occupancy:
261     return evaluateOccupancy(Res, Asm);
262   }
263 
264   for (const MCExpr *Arg : Args) {
265     MCValue ArgRes;
266     if (!Arg->evaluateAsRelocatable(ArgRes, Asm) || !ArgRes.isAbsolute())
267       return false;
268 
269     if (!Total.has_value())
270       Total = ArgRes.getConstant();
271     Total = op(Kind, *Total, ArgRes.getConstant());
272   }
273 
274   Res = MCValue::get(*Total);
275   return true;
276 }
277 
278 void AMDGPUMCExpr::visitUsedExpr(MCStreamer &Streamer) const {
279   for (const MCExpr *Arg : Args)
280     Streamer.visitUsedExpr(*Arg);
281 }
282 
283 MCFragment *AMDGPUMCExpr::findAssociatedFragment() const {
284   for (const MCExpr *Arg : Args) {
285     if (Arg->findAssociatedFragment())
286       return Arg->findAssociatedFragment();
287   }
288   return nullptr;
289 }
290 
291 /// Allow delayed MCExpr resolve of ExtraSGPRs (in case VCCUsed or FlatScrUsed
292 /// are unresolvable but needed for further MCExprs). Derived from
293 /// implementation of IsaInfo::getNumExtraSGPRs in AMDGPUBaseInfo.cpp.
294 ///
295 const AMDGPUMCExpr *AMDGPUMCExpr::createExtraSGPRs(const MCExpr *VCCUsed,
296                                                    const MCExpr *FlatScrUsed,
297                                                    bool XNACKUsed,
298                                                    MCContext &Ctx) {
299 
300   return create(AGVK_ExtraSGPRs,
301                 {VCCUsed, FlatScrUsed, MCConstantExpr::create(XNACKUsed, Ctx)},
302                 Ctx);
303 }
304 
305 const AMDGPUMCExpr *AMDGPUMCExpr::createTotalNumVGPR(const MCExpr *NumAGPR,
306                                                      const MCExpr *NumVGPR,
307                                                      MCContext &Ctx) {
308   return create(AGVK_TotalNumVGPRs, {NumAGPR, NumVGPR}, Ctx);
309 }
310 
311 /// Mimics GCNSubtarget::computeOccupancy for MCExpr.
312 ///
313 /// Remove dependency on GCNSubtarget and depend only only the necessary values
314 /// for said occupancy computation. Should match computeOccupancy implementation
315 /// without passing \p STM on.
316 const AMDGPUMCExpr *AMDGPUMCExpr::createOccupancy(
317     unsigned InitOcc, const MCExpr *NumSGPRs, const MCExpr *NumVGPRs,
318     unsigned DynamicVGPRBlockSize, const GCNSubtarget &STM, MCContext &Ctx) {
319   unsigned MaxWaves = IsaInfo::getMaxWavesPerEU(&STM);
320   unsigned Granule = IsaInfo::getVGPRAllocGranule(&STM, DynamicVGPRBlockSize);
321   unsigned TargetTotalNumVGPRs = IsaInfo::getTotalNumVGPRs(&STM);
322   unsigned Generation = STM.getGeneration();
323 
324   auto CreateExpr = [&Ctx](unsigned Value) {
325     return MCConstantExpr::create(Value, Ctx);
326   };
327 
328   return create(AGVK_Occupancy,
329                 {CreateExpr(MaxWaves), CreateExpr(Granule),
330                  CreateExpr(TargetTotalNumVGPRs), CreateExpr(Generation),
331                  CreateExpr(InitOcc), NumSGPRs, NumVGPRs},
332                 Ctx);
333 }
334 
335 static KnownBits fromOptionalToKnownBits(std::optional<bool> CompareResult) {
336   static constexpr unsigned BitWidth = 64;
337   const APInt True(BitWidth, 1);
338   const APInt False(BitWidth, 0);
339   if (CompareResult) {
340     return *CompareResult ? KnownBits::makeConstant(True)
341                           : KnownBits::makeConstant(False);
342   }
343 
344   KnownBits UnknownBool(/*BitWidth=*/1);
345   return UnknownBool.zext(BitWidth);
346 }
347 
348 using KnownBitsMap = DenseMap<const MCExpr *, KnownBits>;
349 static void knownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
350                                unsigned Depth = 0);
351 
352 static void binaryOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
353                                        unsigned Depth) {
354   static constexpr unsigned BitWidth = 64;
355   const MCBinaryExpr *BExpr = cast<MCBinaryExpr>(Expr);
356   const MCExpr *LHS = BExpr->getLHS();
357   const MCExpr *RHS = BExpr->getRHS();
358 
359   knownBitsMapHelper(LHS, KBM, Depth + 1);
360   knownBitsMapHelper(RHS, KBM, Depth + 1);
361   KnownBits LHSKnown = KBM[LHS];
362   KnownBits RHSKnown = KBM[RHS];
363 
364   switch (BExpr->getOpcode()) {
365   default:
366     KBM[Expr] = KnownBits(BitWidth);
367     return;
368   case MCBinaryExpr::Opcode::Add:
369     KBM[Expr] = KnownBits::add(LHSKnown, RHSKnown);
370     return;
371   case MCBinaryExpr::Opcode::And:
372     KBM[Expr] = LHSKnown & RHSKnown;
373     return;
374   case MCBinaryExpr::Opcode::Div:
375     KBM[Expr] = KnownBits::sdiv(LHSKnown, RHSKnown);
376     return;
377   case MCBinaryExpr::Opcode::EQ: {
378     std::optional<bool> CompareRes = KnownBits::eq(LHSKnown, RHSKnown);
379     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
380     return;
381   }
382   case MCBinaryExpr::Opcode::NE: {
383     std::optional<bool> CompareRes = KnownBits::ne(LHSKnown, RHSKnown);
384     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
385     return;
386   }
387   case MCBinaryExpr::Opcode::GT: {
388     std::optional<bool> CompareRes = KnownBits::sgt(LHSKnown, RHSKnown);
389     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
390     return;
391   }
392   case MCBinaryExpr::Opcode::GTE: {
393     std::optional<bool> CompareRes = KnownBits::sge(LHSKnown, RHSKnown);
394     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
395     return;
396   }
397   case MCBinaryExpr::Opcode::LAnd: {
398     std::optional<bool> CompareRes;
399     const APInt False(BitWidth, 0);
400     std::optional<bool> LHSBool =
401         KnownBits::ne(LHSKnown, KnownBits::makeConstant(False));
402     std::optional<bool> RHSBool =
403         KnownBits::ne(RHSKnown, KnownBits::makeConstant(False));
404     if (LHSBool && RHSBool)
405       CompareRes = *LHSBool && *RHSBool;
406     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
407     return;
408   }
409   case MCBinaryExpr::Opcode::LOr: {
410     const APInt False(BitWidth, 0);
411     KnownBits Bits = LHSKnown | RHSKnown;
412     std::optional<bool> CompareRes =
413         KnownBits::ne(Bits, KnownBits::makeConstant(False));
414     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
415     return;
416   }
417   case MCBinaryExpr::Opcode::LT: {
418     std::optional<bool> CompareRes = KnownBits::slt(LHSKnown, RHSKnown);
419     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
420     return;
421   }
422   case MCBinaryExpr::Opcode::LTE: {
423     std::optional<bool> CompareRes = KnownBits::sle(LHSKnown, RHSKnown);
424     KBM[Expr] = fromOptionalToKnownBits(CompareRes);
425     return;
426   }
427   case MCBinaryExpr::Opcode::Mod:
428     KBM[Expr] = KnownBits::srem(LHSKnown, RHSKnown);
429     return;
430   case MCBinaryExpr::Opcode::Mul:
431     KBM[Expr] = KnownBits::mul(LHSKnown, RHSKnown);
432     return;
433   case MCBinaryExpr::Opcode::Or:
434     KBM[Expr] = LHSKnown | RHSKnown;
435     return;
436   case MCBinaryExpr::Opcode::Shl:
437     KBM[Expr] = KnownBits::shl(LHSKnown, RHSKnown);
438     return;
439   case MCBinaryExpr::Opcode::AShr:
440     KBM[Expr] = KnownBits::ashr(LHSKnown, RHSKnown);
441     return;
442   case MCBinaryExpr::Opcode::LShr:
443     KBM[Expr] = KnownBits::lshr(LHSKnown, RHSKnown);
444     return;
445   case MCBinaryExpr::Opcode::Sub:
446     KBM[Expr] = KnownBits::sub(LHSKnown, RHSKnown);
447     return;
448   case MCBinaryExpr::Opcode::Xor:
449     KBM[Expr] = LHSKnown ^ RHSKnown;
450     return;
451   }
452 }
453 
454 static void unaryOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
455                                       unsigned Depth) {
456   static constexpr unsigned BitWidth = 64;
457   const MCUnaryExpr *UExpr = cast<MCUnaryExpr>(Expr);
458   knownBitsMapHelper(UExpr->getSubExpr(), KBM, Depth + 1);
459   KnownBits KB = KBM[UExpr->getSubExpr()];
460 
461   switch (UExpr->getOpcode()) {
462   default:
463     KBM[Expr] = KnownBits(BitWidth);
464     return;
465   case MCUnaryExpr::Opcode::Minus: {
466     KB.makeNegative();
467     KBM[Expr] = KB;
468     return;
469   }
470   case MCUnaryExpr::Opcode::Not: {
471     KnownBits AllOnes(BitWidth);
472     AllOnes.setAllOnes();
473     KBM[Expr] = KB ^ AllOnes;
474     return;
475   }
476   case MCUnaryExpr::Opcode::Plus: {
477     KB.makeNonNegative();
478     KBM[Expr] = KB;
479     return;
480   }
481   }
482 }
483 
484 static void targetOpKnownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
485                                        unsigned Depth) {
486   static constexpr unsigned BitWidth = 64;
487   const AMDGPUMCExpr *AGVK = cast<AMDGPUMCExpr>(Expr);
488 
489   switch (AGVK->getKind()) {
490   default:
491     KBM[Expr] = KnownBits(BitWidth);
492     return;
493   case AMDGPUMCExpr::VariantKind::AGVK_Or: {
494     knownBitsMapHelper(AGVK->getSubExpr(0), KBM, Depth + 1);
495     KnownBits KB = KBM[AGVK->getSubExpr(0)];
496     for (const MCExpr *Arg : AGVK->getArgs()) {
497       knownBitsMapHelper(Arg, KBM, Depth + 1);
498       KB |= KBM[Arg];
499     }
500     KBM[Expr] = KB;
501     return;
502   }
503   case AMDGPUMCExpr::VariantKind::AGVK_Max: {
504     knownBitsMapHelper(AGVK->getSubExpr(0), KBM, Depth + 1);
505     KnownBits KB = KBM[AGVK->getSubExpr(0)];
506     for (const MCExpr *Arg : AGVK->getArgs()) {
507       knownBitsMapHelper(Arg, KBM, Depth + 1);
508       KB = KnownBits::umax(KB, KBM[Arg]);
509     }
510     KBM[Expr] = KB;
511     return;
512   }
513   case AMDGPUMCExpr::VariantKind::AGVK_ExtraSGPRs:
514   case AMDGPUMCExpr::VariantKind::AGVK_TotalNumVGPRs:
515   case AMDGPUMCExpr::VariantKind::AGVK_AlignTo:
516   case AMDGPUMCExpr::VariantKind::AGVK_Occupancy: {
517     int64_t Val;
518     if (AGVK->evaluateAsAbsolute(Val)) {
519       APInt APValue(BitWidth, Val);
520       KBM[Expr] = KnownBits::makeConstant(APValue);
521       return;
522     }
523     KBM[Expr] = KnownBits(BitWidth);
524     return;
525   }
526   }
527 }
528 
529 static void knownBitsMapHelper(const MCExpr *Expr, KnownBitsMap &KBM,
530                                unsigned Depth) {
531   static constexpr unsigned BitWidth = 64;
532 
533   int64_t Val;
534   if (Expr->evaluateAsAbsolute(Val)) {
535     APInt APValue(BitWidth, Val, /*isSigned=*/true);
536     KBM[Expr] = KnownBits::makeConstant(APValue);
537     return;
538   }
539 
540   if (Depth == 16) {
541     KBM[Expr] = KnownBits(BitWidth);
542     return;
543   }
544 
545   switch (Expr->getKind()) {
546   case MCExpr::ExprKind::Binary: {
547     binaryOpKnownBitsMapHelper(Expr, KBM, Depth);
548     return;
549   }
550   case MCExpr::ExprKind::Constant: {
551     const MCConstantExpr *CE = cast<MCConstantExpr>(Expr);
552     APInt APValue(BitWidth, CE->getValue(), /*isSigned=*/true);
553     KBM[Expr] = KnownBits::makeConstant(APValue);
554     return;
555   }
556   case MCExpr::ExprKind::SymbolRef: {
557     const MCSymbolRefExpr *RExpr = cast<MCSymbolRefExpr>(Expr);
558     const MCSymbol &Sym = RExpr->getSymbol();
559     if (!Sym.isVariable()) {
560       KBM[Expr] = KnownBits(BitWidth);
561       return;
562     }
563 
564     // Variable value retrieval is not for actual use but only for knownbits
565     // analysis.
566     const MCExpr *SymVal = Sym.getVariableValue();
567     knownBitsMapHelper(SymVal, KBM, Depth + 1);
568 
569     // Explicitly copy-construct so that there exists a local KnownBits in case
570     // KBM[SymVal] gets invalidated after a potential growth through KBM[Expr].
571     KBM[Expr] = KnownBits(KBM[SymVal]);
572     return;
573   }
574   case MCExpr::ExprKind::Unary: {
575     unaryOpKnownBitsMapHelper(Expr, KBM, Depth);
576     return;
577   }
578   case MCExpr::ExprKind::Target: {
579     targetOpKnownBitsMapHelper(Expr, KBM, Depth);
580     return;
581   case MCExpr::Specifier:
582     llvm_unreachable("unused by this backend");
583   }
584   }
585 }
586 
587 static const MCExpr *tryFoldHelper(const MCExpr *Expr, KnownBitsMap &KBM,
588                                    MCContext &Ctx) {
589   if (!KBM.count(Expr))
590     return Expr;
591 
592   auto ValueCheckKnownBits = [](KnownBits &KB, unsigned Value) -> bool {
593     if (!KB.isConstant())
594       return false;
595 
596     return Value == KB.getConstant();
597   };
598 
599   if (Expr->getKind() == MCExpr::ExprKind::Constant)
600     return Expr;
601 
602   // Resolving unary operations to constants may make the value more ambiguous.
603   // For example, `~62` becomes `-63`; however, to me it's more ambiguous if a
604   // bit mask value is represented through a negative number.
605   if (Expr->getKind() != MCExpr::ExprKind::Unary) {
606     if (KBM[Expr].isConstant()) {
607       APInt ConstVal = KBM[Expr].getConstant();
608       return MCConstantExpr::create(ConstVal.getSExtValue(), Ctx);
609     }
610 
611     int64_t EvalValue;
612     if (Expr->evaluateAsAbsolute(EvalValue))
613       return MCConstantExpr::create(EvalValue, Ctx);
614   }
615 
616   switch (Expr->getKind()) {
617   default:
618     return Expr;
619   case MCExpr::ExprKind::Binary: {
620     const MCBinaryExpr *BExpr = cast<MCBinaryExpr>(Expr);
621     const MCExpr *LHS = BExpr->getLHS();
622     const MCExpr *RHS = BExpr->getRHS();
623 
624     switch (BExpr->getOpcode()) {
625     default:
626       return Expr;
627     case MCBinaryExpr::Opcode::Sub: {
628       if (ValueCheckKnownBits(KBM[RHS], 0))
629         return tryFoldHelper(LHS, KBM, Ctx);
630       break;
631     }
632     case MCBinaryExpr::Opcode::Add:
633     case MCBinaryExpr::Opcode::Or: {
634       if (ValueCheckKnownBits(KBM[LHS], 0))
635         return tryFoldHelper(RHS, KBM, Ctx);
636       if (ValueCheckKnownBits(KBM[RHS], 0))
637         return tryFoldHelper(LHS, KBM, Ctx);
638       break;
639     }
640     case MCBinaryExpr::Opcode::Mul: {
641       if (ValueCheckKnownBits(KBM[LHS], 1))
642         return tryFoldHelper(RHS, KBM, Ctx);
643       if (ValueCheckKnownBits(KBM[RHS], 1))
644         return tryFoldHelper(LHS, KBM, Ctx);
645       break;
646     }
647     case MCBinaryExpr::Opcode::Shl:
648     case MCBinaryExpr::Opcode::AShr:
649     case MCBinaryExpr::Opcode::LShr: {
650       if (ValueCheckKnownBits(KBM[RHS], 0))
651         return tryFoldHelper(LHS, KBM, Ctx);
652       if (ValueCheckKnownBits(KBM[LHS], 0))
653         return MCConstantExpr::create(0, Ctx);
654       break;
655     }
656     case MCBinaryExpr::Opcode::And: {
657       if (ValueCheckKnownBits(KBM[LHS], 0) || ValueCheckKnownBits(KBM[RHS], 0))
658         return MCConstantExpr::create(0, Ctx);
659       break;
660     }
661     }
662     const MCExpr *NewLHS = tryFoldHelper(LHS, KBM, Ctx);
663     const MCExpr *NewRHS = tryFoldHelper(RHS, KBM, Ctx);
664     if (NewLHS != LHS || NewRHS != RHS)
665       return MCBinaryExpr::create(BExpr->getOpcode(), NewLHS, NewRHS, Ctx,
666                                   BExpr->getLoc());
667     return Expr;
668   }
669   case MCExpr::ExprKind::Unary: {
670     const MCUnaryExpr *UExpr = cast<MCUnaryExpr>(Expr);
671     const MCExpr *SubExpr = UExpr->getSubExpr();
672     const MCExpr *NewSubExpr = tryFoldHelper(SubExpr, KBM, Ctx);
673     if (SubExpr != NewSubExpr)
674       return MCUnaryExpr::create(UExpr->getOpcode(), NewSubExpr, Ctx,
675                                  UExpr->getLoc());
676     return Expr;
677   }
678   case MCExpr::ExprKind::Target: {
679     const AMDGPUMCExpr *AGVK = cast<AMDGPUMCExpr>(Expr);
680     SmallVector<const MCExpr *, 8> NewArgs;
681     bool Changed = false;
682     for (const MCExpr *Arg : AGVK->getArgs()) {
683       const MCExpr *NewArg = tryFoldHelper(Arg, KBM, Ctx);
684       NewArgs.push_back(NewArg);
685       Changed |= Arg != NewArg;
686     }
687     return Changed ? AMDGPUMCExpr::create(AGVK->getKind(), NewArgs, Ctx) : Expr;
688   }
689   }
690   return Expr;
691 }
692 
693 const MCExpr *llvm::AMDGPU::foldAMDGPUMCExpr(const MCExpr *Expr,
694                                              MCContext &Ctx) {
695   KnownBitsMap KBM;
696   knownBitsMapHelper(Expr, KBM);
697   const MCExpr *NewExpr = tryFoldHelper(Expr, KBM, Ctx);
698 
699   return Expr != NewExpr ? NewExpr : Expr;
700 }
701 
702 void llvm::AMDGPU::printAMDGPUMCExpr(const MCExpr *Expr, raw_ostream &OS,
703                                      const MCAsmInfo *MAI) {
704   int64_t Val;
705   if (Expr->evaluateAsAbsolute(Val)) {
706     OS << Val;
707     return;
708   }
709 
710   MAI->printExpr(OS, *Expr);
711 }
712