File: mem_file.h

package info (click to toggle)
pytorch-cuda 2.6.0%2Bdfsg-7
  • links: PTS, VCS
  • area: contrib
  • in suites: forky, sid, trixie
  • size: 161,620 kB
  • sloc: python: 1,278,832; cpp: 900,322; ansic: 82,710; asm: 7,754; java: 3,363; sh: 2,811; javascript: 2,443; makefile: 597; ruby: 195; xml: 84; objc: 68
file content (159 lines) | stat: -rw-r--r-- 4,591 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

#pragma once

#include <c10/util/error.h>
#include <elf.h>
#include <fcntl.h>
#include <fmt/format.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <torch/csrc/profiler/unwind/lexer.h>
#include <torch/csrc/profiler/unwind/unwind_error.h>
#include <unistd.h>
#include <cerrno>
#include <cstdio>
#include <cstring>

namespace torch::unwind {

struct Section {
  char* data = nullptr;
  size_t size = 0;
  const char* string(size_t offset) {
    return lexer(offset).readCString();
  }
  CheckedLexer lexer(size_t offset) {
    return CheckedLexer(data + offset, data, data + size);
  }
};

/// Memory maps a file into the address space read-only, and manages the
/// lifetime of the mapping. Here are a few use cases:
/// 1. Used in the loader to read in initial image, and to inspect
// ELF files for dependencies before callling dlopen.
///
/// 2. Used in unity to load the elf file.
struct MemFile {
  explicit MemFile(const char* filename_)
      : fd_(open(filename_, O_RDONLY)), name_(filename_) {
    UNWIND_CHECK(
        fd_ != -1,
        "failed to open {}: {}",
        filename_,
        c10::utils::str_error(errno));
    struct stat s {};
    if (-1 == fstat(fd_, &s)) {
      close(fd_); // destructors don't run during exceptions
      UNWIND_CHECK(
          false,
          "failed to stat {}: {}",
          filename_,
          c10::utils::str_error(errno));
    }
    n_bytes_ = s.st_size;
    UNWIND_CHECK(
        n_bytes_ > sizeof(Elf64_Ehdr), "empty shared library: {}", filename_);
    mem_ = (char*)mmap(nullptr, n_bytes_, PROT_READ, MAP_SHARED, fd_, 0);
    if (MAP_FAILED == mem_) {
      close(fd_);
      UNWIND_CHECK(
          false,
          "failed to mmap {}: {}",
          filename_,
          c10::utils::str_error(errno));
    }
    ehdr_ = (Elf64_Ehdr*)mem_;
#define ELF_CHECK(cond) UNWIND_CHECK(cond, "not an ELF file: {}", filename_)
    ELF_CHECK(ehdr_->e_ident[EI_MAG0] == ELFMAG0);
    ELF_CHECK(ehdr_->e_ident[EI_MAG1] == ELFMAG1);
    ELF_CHECK(ehdr_->e_ident[EI_MAG2] == ELFMAG2);
    ELF_CHECK(ehdr_->e_ident[EI_MAG3] == ELFMAG3);
    ELF_CHECK(ehdr_->e_ident[EI_CLASS] == ELFCLASS64);
    ELF_CHECK(ehdr_->e_ident[EI_VERSION] == EV_CURRENT);
    ELF_CHECK(ehdr_->e_version == EV_CURRENT);
    ELF_CHECK(ehdr_->e_machine == EM_X86_64);
#undef ELF_CHECK
    UNWIND_CHECK(
        ehdr_->e_shoff + sizeof(Elf64_Shdr) * ehdr_->e_shnum <= n_bytes_,
        "invalid section header table {} {} {}",
        ehdr_->e_shoff + sizeof(Elf64_Shdr) * ehdr_->e_shnum,
        n_bytes_,
        ehdr_->e_shnum);
    shdr_ = (Elf64_Shdr*)(mem_ + ehdr_->e_shoff);
    UNWIND_CHECK(
        ehdr_->e_shstrndx < ehdr_->e_shnum, "invalid strtab section offset");
    auto& strtab_hdr = shdr_[ehdr_->e_shstrndx];
    strtab_ = getSection(strtab_hdr);
  }

  MemFile(const MemFile&) = delete;
  MemFile(MemFile&&) = delete;
  MemFile& operator=(const MemFile&) = delete;
  MemFile& operator=(MemFile&&) = delete;
  [[nodiscard]] const char* data() const {
    return (const char*)mem_;
  }

  /// Returns whether or not the file descriptor
  /// of the underlying file is valid.
  int valid() {
    return fcntl(fd_, F_GETFD) != -1 || errno != EBADF;
  }

  ~MemFile() {
    if (mem_) {
      munmap((void*)mem_, n_bytes_);
    }
    if (fd_ >= 0) {
      close(fd_);
    }
  }

  /// Returns the size of the underlying file defined by the `MemFile`
  size_t size() {
    return n_bytes_;
  }
  [[nodiscard]] int fd() const {
    return fd_;
  }

  Section getSection(const Elf64_Shdr& shdr) {
    UNWIND_CHECK(shdr.sh_offset + shdr.sh_size <= n_bytes_, "invalid section");
    return Section{mem_ + shdr.sh_offset, shdr.sh_size};
  }

  Section getSection(const char* name, bool optional) {
    for (int i = 0; i < ehdr_->e_shnum; i++) {
      if (strcmp(strtab_.string(shdr_[i].sh_name), name) == 0) {
        return getSection(shdr_[i]);
      }
    }
    UNWIND_CHECK(optional, "{} has no section {}", name_, name);
    return Section{nullptr, 0};
  }

  Section strtab() {
    return strtab_;
  }

 private:
  template <typename T>
  T* load(size_t offset) {
    UNWIND_CHECK(offset < n_bytes_, "out of range");
    return (T*)(mem_ + offset);
  }
  int fd_;
  char* mem_{nullptr};
  size_t n_bytes_{0};
  std::string name_;
  Elf64_Ehdr* ehdr_;
  Elf64_Shdr* shdr_;
  Section strtab_ = {nullptr, 0};
};

} // namespace torch::unwind