File: storage_context.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 (83 lines) | stat: -rw-r--r-- 2,487 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
#pragma once

#include <ATen/core/ivalue.h>

namespace torch::jit {

// Used in torch.package and TorchScript serialization to coordinate
// sharing of storages between models. Also used to create deterministic
// naming for storages.
class TORCH_API SerializationStorageContext {
 public:
  explicit SerializationStorageContext() = default;
  SerializationStorageContext operator=(const SerializationStorageContext&) =
      delete;
  SerializationStorageContext(const SerializationStorageContext&) = delete;

  uint64_t getOrAddStorage(const c10::Storage& storage) {
    if (!hasStorage(storage)) {
      uint64_t size = storage_id_map_.size();
      storage_id_map_[storage] = size;
    }
    return storage_id_map_[storage];
  }

  bool hasStorage(const c10::Storage& storage) {
    return storage_id_map_.find(storage) != storage_id_map_.end();
  }

  ~SerializationStorageContext() = default;

 private:
  class StorageSerializationHash {
   public:
    size_t operator()(const c10::Storage& storage) const {
      return std::hash<void*>()(
          reinterpret_cast<void*>(storage.unsafeGetStorageImpl()));
    }
  };

  class StorageSerializationEqual {
   public:
    bool operator()(const c10::Storage& lhs, const c10::Storage& rhs) const {
      return lhs.unsafeGetStorageImpl() == rhs.unsafeGetStorageImpl();
    }
  };

  std::unordered_map<
      c10::Storage,
      uint64_t,
      StorageSerializationHash,
      StorageSerializationEqual>
      storage_id_map_;
};

// Used in torch.package and TorchScript deserialization to coordinate
// sharing of storages between models.
class TORCH_API DeserializationStorageContext {
 public:
  explicit DeserializationStorageContext() = default;
  DeserializationStorageContext operator=(
      const DeserializationStorageContext&) = delete;
  DeserializationStorageContext(const DeserializationStorageContext&) = delete;

  void addStorage(std::string name, c10::Storage storage) {
    TORCH_INTERNAL_ASSERT(!hasStorage(name));
    name_storage_map_.emplace(std::move(name), std::move(storage));
  }

  bool hasStorage(const std::string& name) {
    return name_storage_map_.find(name) != name_storage_map_.end();
  }

  c10::Storage getStorage(const std::string& name) {
    TORCH_INTERNAL_ASSERT(hasStorage(name));
    return name_storage_map_.find(name)->second;
  }
  ~DeserializationStorageContext() = default;

 private:
  std::unordered_map<std::string, c10::Storage> name_storage_map_;
};

} // namespace torch::jit