xref: /aosp_15_r20/external/pytorch/torch/csrc/profiler/unwind/lexer.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 #include <cstdint>
3 #include <cstring>
4 #include <utility>
5 
6 #include <torch/csrc/profiler/unwind/dwarf_enums.h>
7 #include <torch/csrc/profiler/unwind/unwind_error.h>
8 
9 namespace torch::unwind {
10 
11 template <bool checked>
12 struct LexerImpl {
13   LexerImpl(void* data, void* base = nullptr, void* end = nullptr)
14       : next_((const char*)data),
15         base_((int64_t)base),
16         end_((const char*)end) {}
17 
18   template <typename T>
readLexerImpl19   T read() {
20     T result;
21     auto end = next_ + sizeof(T);
22     UNWIND_CHECK(
23         !checked || end <= end_,
24         "read out of bounds {} >= {}",
25         (void*)end,
26         (void*)end_);
27     memcpy(&result, next_, sizeof(T));
28     next_ = end;
29     return result;
30   }
31 
32   // SLEB/ULEB code adapted from LLVM equivalents
readSLEB128LexerImpl33   int64_t readSLEB128() {
34     int64_t Value = 0;
35     unsigned Shift = 0;
36     uint8_t Byte = 0;
37     do {
38       Byte = read<uint8_t>();
39       uint64_t Slice = Byte & 0x7f;
40       if ((Shift >= 64 && Slice != (Value < 0 ? 0x7f : 0x00)) ||
41           (Shift == 63 && Slice != 0 && Slice != 0x7f)) {
42         throw UnwindError("sleb128 too big for int64");
43       }
44       Value |= int64_t(Slice << Shift);
45       Shift += 7;
46     } while (Byte >= 128);
47     // Sign extend negative numbers if needed.
48     if (Shift < 64 && (Byte & 0x40)) {
49       Value |= int64_t((-1ULL) << Shift);
50     }
51     return Value;
52   }
53 
readULEB128LexerImpl54   uint64_t readULEB128() {
55     uint64_t Value = 0;
56     unsigned Shift = 0;
57     uint8_t p = 0;
58     do {
59       p = read<uint8_t>();
60       uint64_t Slice = p & 0x7f;
61       if ((Shift >= 64 && Slice != 0) || Slice << Shift >> Shift != Slice) {
62         throw UnwindError("uleb128 too big for uint64");
63       }
64       Value += Slice << Shift;
65       Shift += 7;
66     } while (p >= 128);
67     return Value;
68   }
readCStringLexerImpl69   const char* readCString() {
70     auto result = next_;
71     if (!checked) {
72       next_ += strlen(next_) + 1;
73       return result;
74     }
75     while (next_ < end_) {
76       if (*next_++ == '\0') {
77         return result;
78       }
79     }
80     UNWIND_CHECK(
81         false, "string is out of bounds {} >= {}", (void*)next_, (void*)end_);
82   }
readEncodedLexerImpl83   int64_t readEncoded(uint8_t enc) {
84     int64_t r = 0;
85     switch (enc & (~DW_EH_PE_indirect & 0xF0)) {
86       case DW_EH_PE_absptr:
87         break;
88       case DW_EH_PE_pcrel:
89         r = (int64_t)next_;
90         break;
91       case DW_EH_PE_datarel:
92         r = base_;
93         break;
94       default:
95         throw UnwindError("unknown encoding");
96     }
97     return r + readEncodedValue(enc);
98   }
readEncodedOrLexerImpl99   int64_t readEncodedOr(uint8_t enc, int64_t orelse) {
100     if (enc == DW_EH_PE_omit) {
101       return orelse;
102     }
103     return readEncoded(enc);
104   }
105 
read4or8LengthLexerImpl106   int64_t read4or8Length() {
107     return readSectionLength().first;
108   }
109 
readSectionLengthLexerImpl110   std::pair<int64_t, bool> readSectionLength() {
111     int64_t length = read<uint32_t>();
112     if (length == 0xFFFFFFFF) {
113       return std::make_pair(read<int64_t>(), true);
114     }
115     return std::make_pair(length, false);
116   }
117 
locLexerImpl118   void* loc() const {
119     return (void*)next_;
120   }
skipLexerImpl121   LexerImpl& skip(int64_t bytes) {
122     next_ += bytes;
123     return *this;
124   }
125 
readEncodedValueLexerImpl126   int64_t readEncodedValue(uint8_t enc) {
127     switch (enc & 0xF) {
128       case DW_EH_PE_udata2:
129         return read<uint16_t>();
130       case DW_EH_PE_sdata2:
131         return read<int16_t>();
132       case DW_EH_PE_udata4:
133         return read<uint32_t>();
134       case DW_EH_PE_sdata4:
135         return read<int32_t>();
136       case DW_EH_PE_udata8:
137         return read<uint64_t>();
138       case DW_EH_PE_sdata8:
139         return read<int64_t>();
140       case DW_EH_PE_uleb128:
141         return readULEB128();
142       case DW_EH_PE_sleb128:
143         return readSLEB128();
144       default:
145         throw UnwindError("not implemented");
146     }
147   }
148 
149  private:
150   const char* next_;
151   int64_t base_;
152   const char* end_;
153 };
154 
155 // using Lexer = LexerImpl<false>;
156 using CheckedLexer = LexerImpl<true>;
157 using Lexer = LexerImpl<false>;
158 
159 } // namespace torch::unwind
160