1 #pragma once 2 #include <torch/csrc/profiler/unwind/action.h> 3 #include <torch/csrc/profiler/unwind/unwind_error.h> 4 #include <cstdint> 5 #include <limits> 6 7 namespace torch::unwind { 8 9 struct UnwindState { 10 int64_t rip, rbp, rsp; 11 }; 12 13 struct Unwinder { UnwinderUnwinder14 Unwinder(Action rsp, Action rip, Action rbp) 15 : kind_(rip.kind == A_UNDEFINED ? END : STANDARD), 16 reg_(rsp.reg), 17 off_(rsp.data), 18 rip_off_(rip.data), 19 rbp_off_( 20 rbp.kind == A_UNDEFINED ? std::numeric_limits<int64_t>::max() 21 : rbp.data), 22 deref_(rsp.kind == A_REG_PLUS_DATA_DEREF) { 23 check(rsp.reg == D_RSP || rsp.reg == D_RBP); 24 check(rip.kind == A_UNDEFINED || rip.kind == A_LOAD_CFA_OFFSET); 25 if (rsp.kind == A_REG_PLUS_DATA) { 26 check(rbp.kind == A_LOAD_CFA_OFFSET || rbp.kind == A_UNDEFINED); 27 } else if (rsp.kind == A_REG_PLUS_DATA_DEREF) { 28 if (rbp.kind == A_REG_PLUS_DATA_DEREF) { 29 check(rbp.reg == rsp.reg); 30 rbp_off_ -= rsp.data; 31 } else { 32 check(rbp.kind == A_UNDEFINED); 33 } 34 } else { 35 check(false); 36 } 37 } checkUnwinder38 void check(bool cond) { 39 if (!cond) { 40 throw UnwindError("Unwinding actions do not follow supported patterns"); 41 } 42 } terminatorUnwinder43 bool terminator() const { 44 return kind_ != STANDARD; 45 } isUnknownUnwinder46 bool isUnknown() const { 47 return kind_ == UNKNOWN; 48 } 49 // unwinder representing some pattern unsupported in 50 // current implementation unknownUnwinder51 static Unwinder unknown() { 52 return Unwinder(); 53 } runUnwinder54 UnwindState run(const UnwindState& cur) const { 55 UnwindState r = cur; 56 r.rsp = (reg_ == D_RSP ? cur.rsp : cur.rbp) + off_; 57 r.rbp = rbp_off_ == std::numeric_limits<int64_t>::max() 58 ? cur.rbp 59 // NOLINTNEXTLINE(performance-no-int-to-ptr) 60 : *(int64_t*)(r.rsp + rbp_off_); 61 if (deref_) { 62 // NOLINTNEXTLINE(performance-no-int-to-ptr) 63 r.rsp = *(int64_t*)r.rsp; 64 } 65 // NOLINTNEXTLINE(performance-no-int-to-ptr) 66 r.rip = *(int64_t*)(r.rsp + rip_off_); 67 68 return r; 69 } 70 71 private: UnwinderUnwinder72 Unwinder() : kind_(UNKNOWN), reg_(0), off_(0), rip_off_(0), rbp_off_(0) {} 73 enum Kind { STANDARD, END, UNKNOWN } kind_; 74 uint32_t reg_; 75 int64_t off_; 76 int64_t rip_off_; 77 int64_t rbp_off_; 78 bool deref_{false}; 79 }; 80 81 } // namespace torch::unwind 82