#pragma once
|
|
#include <c10/util/ArrayRef.h>
|
#include <c10/util/sparse_bitset.h>
|
#include <memory>
|
#include <unordered_map>
|
#include <unordered_set>
|
#include <vector>
|
|
#include <torch/csrc/WindowsTorchApiMacro.h>
|
|
// Uses a compressed index representation for faster comparisons
|
typedef c10::SparseBitVector<256> MemoryLocations;
|
namespace torch {
|
namespace jit {
|
|
struct Element;
|
struct Value;
|
|
// class MemoryDAG
|
//
|
// This class tracks the "A points to B" graph for all values. It is used by
|
// AliasDb to provide a higher-level API.
|
//
|
// We maintain a DAG where:
|
// - Vertices (called "elements") represent values and
|
// other aliasing entities (e.g. like the stuff inside a list)
|
// - Edges represent a "points-to" relationship.
|
//
|
// Leaves in this DAG are entities that don't point to anything, and thus
|
// correspond to unique "memory locations".
|
//
|
// So, by traversing the "points-to" graph to the leaves, you can determine
|
// which memory locations an element may point to.
|
class TORCH_API MemoryDAG {
|
public:
|
// explicitly delete copy constructor because otherwise windows build is
|
// confused for an exported class see
|
// https://stackoverflow.com/a/51033485/105137
|
MemoryDAG() {}
|
MemoryDAG(const MemoryDAG&) = delete;
|
MemoryDAG& operator=(const MemoryDAG&) = delete;
|
|
// Make `from` point at `to`.
|
void makePointerTo(Element* from, Element* to);
|
|
void addToContainedElements(Element* contained, Element* container);
|
|
// Make a fresh element (i.e. an element that doesn't point to anything) and
|
// return it.
|
Element* makeFreshValue(const Value* v);
|
|
// Do `a` and `b` potentially share a memory location?
|
bool mayAlias(const Element* a, const Element* b) const;
|
bool mayAlias(Element* a, Element* b) const;
|
|
// Does a hold reference to any memory that is stored in elem, or vice versa?
|
bool mayContainAlias(const Element* a, const Element* b) const;
|
bool mayContainAlias(Element* a, Element* b) const;
|
|
bool mayContainAlias(
|
const at::ArrayRef<Element*>& a,
|
const at::ArrayRef<Element*>& b) const;
|
|
// Converts from the compressed index representation
|
const Element* fromIndex(unsigned x) const;
|
Element* fromIndex(unsigned x);
|
|
private:
|
bool mayAliasImpl(const Element* a, const Element* b) const;
|
bool mayContainAliasImpl(const Element* contained, const Element* container)
|
const;
|
void collectAllContainedMemoryLocations(
|
const Element* elem, MemoryLocations& cont) const;
|
|
std::vector<std::unique_ptr<Element>> indexToElementMap_;
|
};
|
|
// `Element` represents the vertex in the points-to graph. It represents
|
// anything that could have an aliasing relationship, mostly IR `Value`s, but
|
// also the "inside of a list", or wildcards.
|
struct Element {
|
Element(MemoryDAG& dag_, const Value* value_, unsigned index_);
|
|
// Reference to the owning DAG.
|
MemoryDAG& dag;
|
// Index into the owning DAG's bit vector that represents this element.
|
unsigned index;
|
|
// All elements that this element *may* point to. It's possible to have
|
// multiple elements that you might point to due to control flow/complex ops
|
MemoryLocations pointsTo;
|
// Backreference for points-to.
|
MemoryLocations pointedFrom;
|
|
// Elements can contain other elements (e.g. List[Tensor])
|
MemoryLocations containedElements;
|
|
// Return the unique memory locations that `Element` might represent.
|
TORCH_API const MemoryLocations& getMemoryLocations() const;
|
|
// The value that this element corresponds to. May be null if this element
|
// doesn't represent a first-class value.
|
const Value* value = nullptr;
|
|
private:
|
// We do path compression to make repeated memory location queries faster.
|
// An empty cache means it is invalidated (it can never be empty otherwise,
|
// since every element must point to at least one memory location).
|
mutable MemoryLocations cachedMemoryLocations_;
|
|
enum class BfsDirection {
|
POINTS_TO,
|
POINTED_FROM,
|
};
|
// Do a breadth-first search over the graph, starting at `this` and
|
// traversing in the direction `dir`.`fn` will be run on each element.
|
void bfs(BfsDirection dir, MemoryLocations& res) const;
|
friend class MemoryDAG;
|
};
|
|
} // namespace jit
|
} // namespace torch
|
