pytorch/torch/csrc/jit/frontend/source_ref.h
Zhengxu Chen 3a3959d253 [jit] Add a utility class SourceRef to represent Source as keys (#57396)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/57396

A new type SourceRef is introduced to represent a unique identifier to source
text. The type holds refcount to underlying source, and supports comparators
and hash functions, such that it can be used in C++ and Python maps. In later
diffs we will use this to aggregate and print profiling information.

Test Plan: Imported from OSS

Reviewed By: nikithamalgifb

Differential Revision: D28133578

fbshipit-source-id: c3d5199a8269c5006c85a145b281bcaaf3e2dc1c
2021-05-18 18:20:53 -07:00

44 lines
1.1 KiB
C++

#pragma once
#include <functional>
#include <memory>
#include <c10/macros/Export.h>
#include <torch/csrc/jit/frontend/source_range.h>
namespace torch {
namespace jit {
/**
* SourceRef does two things:
* 1. Owns a Source object.
* 2. Serves as lookup key to the owned Source in associative containers, for
* runtime data aggregation.
* We don't want to use std::shared_ptr<Source> directly because we want to
* support heteogeneous lookup, and also shared_ptr is an implementation detail
* which should be encapsulated.
*/
class TORCH_API SourceRef {
public:
explicit SourceRef(std::shared_ptr<Source> source)
: source_(std::move(source)) {}
bool operator==(const SourceRef& other) const {
return source_ == other.source_;
}
bool operator<(const Source& other) const {
return source_.get() < &other;
}
friend bool operator<(const Source& other, const SourceRef& self) {
return &other < self.source_.get();
}
bool operator<(const SourceRef& other) const {
return *this < *other.source_.get();
}
private:
std::shared_ptr<Source> source_;
};
} // namespace jit
} // namespace torch