1 #pragma once 2 3 #include <c10/core/Allocator.h> 4 #include <torch/csrc/jit/frontend/source_range.h> 5 6 #include <ATen/core/ivalue.h> 7 8 #include <unordered_map> 9 #include <vector> 10 11 namespace c10 { 12 struct IValue; 13 } 14 15 namespace torch::jit { 16 17 class Pickler; 18 class SourceRangeSerializer; 19 static constexpr size_t kByteOffsetIndex = 0; 20 static constexpr size_t kSourceRangeIndex = 1; 21 static constexpr size_t kSourceRangeTagIndex = 2; 22 constexpr c10::string_view kFormatWithStringTable = "FORMAT_WITH_STRING_TABLE"; 23 24 class SourceRangePickler { 25 public: 26 SourceRangePickler(); 27 28 std::vector<char> pickle( 29 const SourceRangeRecords& ranges, 30 const SourceRangeTagMap& source_range_tags); 31 32 private: 33 std::shared_ptr<SourceRangeSerializer> srs; 34 }; 35 36 class SourceRangeDeserializer { 37 public: 38 SourceRangeDeserializer() = default; SourceRangeDeserializer(const c10::IValue & text_table)39 explicit SourceRangeDeserializer(const c10::IValue& text_table) { 40 for (const auto& x : text_table.toTuple()->elements()) { 41 text_table_.emplace_back(std::make_shared<std::string>(x.toStringRef())); 42 } 43 } 44 SourceRange deserialize(const c10::IValue& iv); 45 46 private: 47 std::shared_ptr<Source> deserialize_source(const c10::IValue& iv); 48 std::unordered_map< 49 c10::intrusive_ptr<c10::ivalue::Tuple>, 50 std::shared_ptr<Source>> 51 cached_sources; 52 std::vector<std::shared_ptr<std::string>> text_table_; 53 }; 54 55 class SourceRangeUnpickler { 56 public: 57 virtual std::optional<SourceRange> findSourceRangeThatGenerated( 58 const SourceRange& range) = 0; 59 60 virtual ~SourceRangeUnpickler() = default; 61 }; 62 63 TORCH_API void setShouldUseFormatWithStringTable( 64 bool should_use_format_with_string_table); 65 66 } // namespace torch::jit 67