xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/serialization/source_range_serialization_impl.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <torch/csrc/jit/serialization/source_range_serialization.h>
4 
5 namespace torch::jit {
6 
7 // Do this clownyness with virtual functions because of the split
8 // between ATen core and torch
9 
10 class ConcreteSourceRangeUnpickler : public SourceRangeUnpickler {
11  public:
12   ConcreteSourceRangeUnpickler(at::DataPtr&& data, size_t size);
13 
14   std::optional<SourceRange> findSourceRangeThatGenerated(
15       const SourceRange& range) override;
16 
17  private:
18   at::DataPtr data;
19   size_t size;
20 
21   void unpickle();
22 
23   std::mutex mutex;
24   std::shared_ptr<SourceRangeDeserializer> deserializer;
25   std::shared_ptr<SourceRangeRecords> unpickled_records;
26 };
27 
28 } // namespace torch::jit
29