compute_context.h#

Graph-level annotation context that combines value-tag inference, in-place reuse analysis, and per-node memory profiling.

:cpp:class:ComputeContext stores the results of all three analyses and exposes them in a single object, mirroring the way :cpp:class:core::shapes::ShapesContext stores inferred descriptors.

namespace onnx_light

Alias that makes onnx-light headers compatible with code that references ONNX_LIGHT_NAMESPACE (the macro used in the standard onnx package).

Set to ONNX_LIGHT_NAMESPACE so both names resolve to the same namespace.

Symbol-visibility attribute for the public onnx-light C++ API.

Maps the upstream compatibility macro to onnx-light’s explicit proto ABI annotation. This keeps declarations from vendored ONNX headers visible when lib_onnx_proto uses hidden visibility by default.

Namespace alias so that ONNX C++ code (and consumers such as onnxruntime) that refers to the literal onnx namespace — rather than the ONNX_NAMESPACE macro — resolves to the onnx-light namespace. The standard onnx package lives in namespace onnx; onnx-light uses onnx_light (via ONNX_LIGHT_NAMESPACE), so this alias keeps onnx-light a true drop-in. It is only introduced when the onnx-light namespace differs from onnx.

namespace core

namespace compute#

Typedefs

using ComputeEventLog = std::vector<ComputeEvent>#

using ShapeTag = std::string#

Represents a per-node memory snapshot computed by :cpp:class:ComputeContext.

The snapshot represents the memory footprint visible while one node runs:

already_allocated_bytes is the sum of the buffers already alive before the node starts (declared inputs, initializers and still-live intermediates), after shape inference and lifetime analysis;
output_allocation_bytes is the additional memory that must be allocated for the node’s outputs because no eligible in-place reuse opportunity covers them;
total_bytes is their sum.

The profile is stored as a std::map with seven well-known keys:

"total_bytes"
"already_allocated_bytes"
"output_allocation_bytes"
"inputs"
"initializers"
"intermediates"
"outputs"

The first three keys map to scalar :cpp:type:core::expressions::DimType values. The other four keys map to std::map<ShapeTag, DimType> buckets split by value tag ("shape", "axes", "weight", or the empty string for untagged values). "already_allocated_bytes" is the sum of the "inputs", "initializers" and "intermediates" maps; "output_allocation_bytes" is the sum of "outputs"; and "total_bytes" is the sum of those two scalar entries. Every amount is represented as a :cpp:type:core::expressions::DimType, so symbolic shapes retain their expression form instead of being dropped. The "outputs" map only counts the extra allocations performed at this node; outputs that reuse an existing input buffer in place contribute no additional bytes there.

using TaggedMemory = std::map<ShapeTag, expressions::DimType>#

using NodeMemoryProfileValue = std::variant<expressions::DimType, TaggedMemory>#

using NodeMemoryProfile = std::map<std::string, NodeMemoryProfileValue>#

Enums

enum class ComputeEventAction : int32_t#

Kind of decision recorded in the optional :cpp:class:ComputeContext decision log.

kInPlace — one output was matched to one input for in-place reuse.
kRelease — one value reached its last use at a node and can be released after that node.
kReleaseShapeTag — one released value was also classified as "shape" by value tagging.

Values:

enumerator kInPlace#

enumerator kRelease#

enumerator kReleaseShapeTag#

Functions

inline constexpr const char *ComputeEventActionName(ComputeEventAction action)#: Returns the short lowercase label for action.

Variables

constexpr const char *kNodeMemoryTotalBytesKey = "total_bytes"#

constexpr const char *kNodeMemoryAlreadyAllocatedBytesKey = "already_allocated_bytes"#

constexpr const char *kNodeMemoryOutputAllocationBytesKey = "output_allocation_bytes"#

constexpr const char *kNodeMemoryInputsKey = "inputs"#

constexpr const char *kNodeMemoryInitializersKey = "initializers"#

constexpr const char *kNodeMemoryIntermediatesKey = "intermediates"#

constexpr const char *kNodeMemoryOutputsKey = "outputs"#

class ComputeContext#

#include <compute_context.h>

Holds the in-place reuse opportunities computed for a graph, mirroring the way :cpp:class:core::shapes::ShapesContext holds the inferred descriptors.

The reuse guess is purely structural: it reports the opportunities implied by shape inference and value lifetimes, not whether a particular kernel actually performs the reuse. Populate the context with :cpp:func:ComputeInPlaceReuseGraph (consuming a :cpp:class:ShapesContext already filled by :cpp:func:ShapesContext::ComputeShapeGraph or :cpp:func:ShapesContext::ComputeShapeModel), then read the result through :cpp:func:Reuse / :cpp:func:NodeReuse or persist it into the graph with :cpp:func:WriteToMetadata.

Public Types

using CustomValueTagFn = std::function<void(ComputeContext&, const NodeProto&, std::size_t node_index)>#: Callback signature for custom value-tag behavior. Receives (ctx, node, node_index) where ctx can be mutated through :cpp:func:TrySetValueTag / :cpp:func:SetNodeTag.

using CustomValueTagMap = std::unordered_map<std::string, CustomValueTagFn>#

using InputShapes = std::unordered_map<std::string, SymTensor>#: Map from value name to the :cpp:class:SymTensor describing its shape and element type, used to seed shape inference for a :cpp:class:FunctionProto or a bare node list (neither of which carries declared input types).

Public Functions

ComputeContext() = default#

std::pair<std::unordered_map<std::string, std::string>, std::vector<std::string>> ComputeValueAndNodeTags(const GraphProto &graph)#

Infers semantic shape / axes / weight tags for the values and nodes in graph and stores the result in *this (replacing any previously computed tags).

Returns:: A pair (value_tags, node_tags) where value_tags maps value names to their inferred tag and node_tags follows the order of graph.node().

std::pair<std::unordered_map<std::string, std::string>, std::vector<std::string>> ComputeValueAndNodeTags(const FunctionProto &function)#: Same as :cpp:func:ComputeValueAndNodeTags(const GraphProto&) but for a function body.

std::pair<std::unordered_map<std::string, std::string>, std::vector<std::string>> ComputeValueAndNodeTags(const utils::RepeatedProtoField<NodeProto> &nodes)#: Same as :cpp:func:ComputeValueAndNodeTags(const GraphProto&) but for an arbitrary node list.

void SeedValueTag(const std::string &name, const std::string &tag)#: Seeds an initial value tag for a graph input, initializer, value_info or output so it participates in the incremental tag inference driven by :cpp:func:AppendNodeTags. Mirrors the seeding performed whole-graph by :cpp:func:CollectGraphSeedTags.

void AppendNodeTags(const utils::RepeatedProtoField<NodeProto> &nodes, std::size_t node_index)#: Incrementally updates the value/node tags after the node at node_index (the last node of nodes) has been appended. Only that node and the nodes whose values it changes are (re)processed through a monotone worklist — no whole-graph loop — so appending N nodes stays linear in the graph size instead of quadratic. The built-in inference rules converge to the same least fixed point as :cpp:func:ComputeValueAndNodeTags.

inline const std::unordered_map<std::string, std::string> &ValueTags() const noexcept#: Read-only access to the last value-tag map computed through :cpp:func:ComputeValueAndNodeTags.

inline const std::vector<std::string> &NodeTags() const noexcept#: Read-only access to the last per-node tag list computed through :cpp:func:ComputeValueAndNodeTags.

inline const std::string &NodeTag(std::size_t node_index) const#

Tag inferred for the node at node_index.

Throws:: std::out_of_range – when node_index is out of bounds.

void SeedConstant(const std::string &name)#: Seeds name as a constant value (e.g. an initializer) so it participates in the incremental constant analysis driven by :cpp:func:AppendNodeConstant.

void AppendNodeConstant(const NodeProto &node, std::size_t node_index)#: Incrementally updates the constant analysis after the node at node_index has been appended: records whether the node is constant and, when it is, marks its outputs as constant values. Appends exactly one entry to the per-node constant flag list.

inline bool IsConstantValue(const std::string &name) const noexcept#: Whether name is currently known to be a constant value.

inline const std::vector<ConstantInfo> &NodeConstant() const noexcept#: Read-only access to the incremental per-node constant classifications. One entry per appended node, in graph order.

inline bool NodeConstant(std::size_t node_index) const#

Whether the node at node_index produces constant outputs.

Throws:: std::out_of_range – when node_index is out of bounds.

bool TrySetValueTag(const std::string &name, const std::string &tag)#: Sets or updates a value tag and returns true when the internal map changed. Returns false when name is empty, when tag is invalid/empty, or when setting it would not change the map.

bool SetNodeTag(std::size_t node_index, const std::string &tag)#

Sets or updates a per-node tag and returns true when the internal list changed. Returns false when tag is invalid/empty or does not change the current value.

Throws:: std::out_of_range – when node_index is out of bounds.

inline void ClearCustomValueTagChangedFlag() noexcept#: Internal flag helpers used by value-tag inference around custom callbacks.

inline bool ConsumeCustomValueTagChangedFlag() noexcept#

inline void SetCustomValueTagFunction(const std::string &domain, const std::string &op_type, CustomValueTagFn fn)#: Registers or replaces a custom value-tag callback for (domain, op_type). domain == "" is normalized to ai.onnx.

inline const CustomValueTagFn *GetCustomValueTagFunction(const std::string &domain, const std::string &op_type) const#: Returns a pointer to the custom value-tag callback registered for (domain, op_type), or nullptr if none is registered.

inline bool RemoveCustomValueTagFunction(const std::string &domain, const std::string &op_type)#: Removes the custom value-tag callback registered for (domain, op_type).

inline void ClearCustomValueTagFunctions()#: Removes every custom value-tag callback.

inline const CustomValueTagMap &CustomValueTagFunctions() const noexcept#: Read-only access to all registered custom value-tag callbacks.

void ComputeInPlaceReuseGraph(const GraphProto &graph, const ShapesContext &ctx, bool allow_input_overwrite = false, const std::unordered_map<std::string, std::string> &value_tags = {})#

Guesses, for every node of graph, which outputs may reuse which input buffers in place, using the shapes and element types already inferred into ctx, and stores the result in *this (replacing any previously computed result).

Parameters:

graph – Graph whose nodes are analysed, in topological order.
ctx – Shapes context already populated with the inferred descriptors for graph (graph inputs, initializers, intermediates and outputs).
allow_input_overwrite – When false (the default), declared graph inputs are never offered as reusable buffers, so a caller’s input is never overwritten in place. When true, a declared graph input may be reused like an intermediate (subject to the same lifetime and shape checks), allowing kernels to overwrite it.
value_tags – Optional map from value name to tag string ("shape", "axes", "weight"). When non-empty, values in the release list that carry the "shape" tag are also stored separately and exposed through :cpp:func:ReleaseAfterShapeTagged / :cpp:func:NodeReleaseAfterShapeTagged, and written to :cpp:var:kReleaseAfterShapeTagMetadataKey by :cpp:func:WriteToMetadata.

void SeedReuseInput(const std::string &name, bool is_graph_input, bool is_initializer, bool allow_input_overwrite)#: Seeds the incremental in-place-reuse lifetime state for a declared graph input (is_graph_input) or initializer (is_initializer). When allow_input_overwrite is false the value is protected (kept) from reuse; when true a graph input is instead made available before the first node (producer index -1) so it can be reused at its last use. Mirrors the seeding performed whole-graph by :cpp:func:ComputeResultLifetimeInfo.

void SeedReuseOutput(const std::string &name)#: Seeds the incremental in-place-reuse lifetime state for a declared graph output: the value is kept alive and removed from the release / not-used lists of any earlier node that had treated it as releasable.

void AppendNodeReuse(const NodeProto &node, std::size_t node_index, const ShapesContext &ctx)#: Incrementally updates the in-place reuse and release-after annotations after the node node at node_index has been appended, using the shapes already inferred into ctx. Only this node and the previous last-users of its inputs are touched — no whole-graph loop — via :cpp:func:ComputeSingleNodeReuse. Appends exactly one entry to each of the per-node result vectors so they stay aligned with :cpp:func:Size.

inline std::size_t Size() const noexcept#: Number of nodes for which reuse has been computed (one entry per node of the analysed graph, in graph.node() order). Zero before :cpp:func:ComputeInPlaceReuseGraph has been called.

inline bool Empty() const noexcept#: true when no reuse has been computed yet.

inline const std::vector<std::vector<InPlaceReuse>> &Reuse() const noexcept#: Read-only access to the per-node reuse opportunities. Entry i lists the opportunities discovered for graph.node()[i]; nodes without any opportunity carry an empty list.

inline const std::vector<InPlaceReuse> &NodeReuse(std::size_t node_index) const#

Reuse opportunities discovered for the node at node_index.

Throws:: std::out_of_range – when node_index is out of bounds.

inline const std::vector<std::vector<std::string>> &ReleaseAfterShapeTagged() const noexcept#: Read-only access to the per-node shape-tagged releasable values. When :cpp:func:ComputeInPlaceReuseGraph was called with a non-empty value_tags map, this vector has one entry per node (same order as graph.node()), and entry i lists the names from the release_after list that carry the "shape" value tag. When ComputeInPlaceReuseGraph was called without value_tags (or with an empty map), this vector is itself empty.

inline const std::vector<std::string> &NodeReleaseAfterShapeTagged(std::size_t node_index) const#

Shape-tagged releasable values for the node at node_index.

Throws:: std::out_of_range – when node_index is out of bounds, or when :cpp:func:ComputeInPlaceReuseGraph was called without value tags (in which case the vector is empty and every access is out of bounds).

inline const std::vector<NodeMemoryProfile> &Memory() const noexcept#: Read-only access to the per-node memory snapshots. Entry i describes the memory footprint observed while running graph.node()[i].

inline const NodeMemoryProfile &NodeMemory(std::size_t node_index) const#

Memory snapshot for the node at node_index.

Throws:: std::out_of_range – when node_index is out of bounds.

const ShapesContext &ComputeShapes(const GraphProto &graph)#

Runs shape inference on graph and stores the resulting descriptors in the :cpp:class:ShapesContext owned by *this.

Parameters:: graph – Graph whose nodes are analysed, in topological order.
Returns:: A reference to the owned :cpp:class:ShapesContext, now populated.

const ShapesContext &ComputeShapes(const ModelProto &model, bool prefill_with_value_info_output = false)#

Runs shape inference on model.graph() (also recording opset versions and local functions from model) and stores the resulting descriptors in the :cpp:class:ShapesContext owned by *this.

Parameters:

model – Model whose main graph is analysed.
prefill_with_value_info_output – Forwarded to :cpp:func:ShapesContext::ComputeShapeModel.

Returns:

A reference to the owned :cpp:class:ShapesContext, now populated.

const ShapesContext &ComputeShapes(const FunctionProto &function, const InputShapes &input_shapes = {})#

Runs shape inference on the body of function and stores the resulting descriptors in the :cpp:class:ShapesContext owned by *this.

A :cpp:class:FunctionProto only names its inputs, so their shapes and element types must be supplied through input_shapes. Any input consumed by a node but absent from input_shapes (and not produced by an earlier node) makes shape inference throw std::invalid_argument.

Parameters:

function – Function whose nodes are analysed, in topological order.
input_shapes – Shapes/types for the function inputs, seeded into the owned :cpp:class:ShapesContext before inference.

Returns:

A reference to the owned :cpp:class:ShapesContext, now populated.

const ShapesContext &ComputeShapes(const utils::RepeatedProtoField<NodeProto> &nodes, const InputShapes &input_shapes = {})#

Runs shape inference on the node list nodes and stores the resulting descriptors in the :cpp:class:ShapesContext owned by *this.

A bare node list has no declared inputs, so the shapes and element types of every value not produced by the list itself must be supplied through input_shapes. Any input consumed by a node but absent from input_shapes (and not produced by an earlier node) makes shape inference throw std::invalid_argument.

Parameters:

nodes – Nodes analysed in topological order.
input_shapes – Shapes/types for the values consumed by nodes but not produced by them, seeded into the owned :cpp:class:ShapesContext before inference.

Returns:

A reference to the owned :cpp:class:ShapesContext, now populated.

inline const ShapesContext &Shapes() const noexcept#: Read-only access to the :cpp:class:ShapesContext owned by *this. Empty until :cpp:func:ComputeShapes (or :cpp:func:Compute) has run.

inline ShapesContext &Shapes() noexcept#: Mutable access to the owned :cpp:class:ShapesContext, allowing callers to seed it (opset versions, custom inference functions, …) before :cpp:func:ComputeShapes.

const std::vector<int64_t> &ComputePeakMemory(const GraphProto &graph, Device device = Device::kUndefined)#

Computes the estimated peak scratch memory for every node of graph using the shapes already inferred into the owned :cpp:class:ShapesContext (via :cpp:func:ComputeShapes), storing one entry per node in graph.node() order. Nodes whose estimate is zero keep a 0 entry.

Parameters:

graph – Graph whose nodes are analysed, in topological order.
device – Logical device passed to the peak-memory dispatch function.

Returns:

A reference to the stored per-node peak-memory vector.

inline const std::vector<int64_t> &PeakMemory() const noexcept#: Read-only access to the per-node peak-memory estimates computed by :cpp:func:ComputePeakMemory. Empty before it has been called.

inline int64_t NodePeakMemory(std::size_t node_index) const#

Peak-memory estimate for the node at node_index.

Throws:: std::out_of_range – when node_index is out of bounds.

void Compute(const GraphProto &graph, Device device = Device::kUndefined, bool allow_input_overwrite = false)#

Runs every analysis on graph in order and stores all results in *this: shape inference (:cpp:func:ComputeShapes), value / node tagging (:cpp:func:ComputeValueAndNodeTags), in-place reuse together with the release-after and shape-tag classification (:cpp:func:ComputeInPlaceReuseGraph) and per-node peak memory (:cpp:func:ComputePeakMemory).

Parameters:

graph – Graph whose nodes are analysed, in topological order.
device – Logical device passed to the peak-memory dispatch function.
allow_input_overwrite – Forwarded to :cpp:func:ComputeInPlaceReuseGraph.

void Compute(const ModelProto &model, Device device = Device::kUndefined, bool allow_input_overwrite = false, bool prefill_with_value_info_output = false)#: Same as :cpp:func:Compute(const GraphProto&, ...) but seeds shape inference from model (opset versions and local functions) before analysing model.graph().

void WriteToGraph(GraphProto &graph) const#

Pushes every computed result into graph: the inferred shapes are written into graph (value_info / outputs, via :cpp:func:ShapesContext::ApplyInferredShapesToGraph), the in-place / release / shape-tag information into node metadata_props (via :cpp:func:WriteToMetadata) and the per-node peak-memory estimates under :cpp:var:kNodePeakMemoryMetadataKey.

Parameters:: graph – Graph whose value_info and node metadata are mutated in place; must be the same graph passed to :cpp:func:Compute / :cpp:func:ComputeInPlaceReuseGraph.

void WriteToModel(ModelProto &model) const#: Same as :cpp:func:WriteToGraph(GraphProto&) applied to model.graph().

runtime::ExecutionPlan BuildExecutionPlan(GraphProto &graph) const#: Creates the :cpp:class:runtime::ExecutionPlan derived from every result stored in *this. It first pushes the information into graph (see :cpp:func:WriteToGraph) and then builds the plan from the annotated graph, so the schedule is driven entirely by the analyses held by this context.

runtime::ExecutionPlan BuildExecutionPlan(ModelProto &model) const#: Same as :cpp:func:BuildExecutionPlan(GraphProto&, ...) applied to model.graph().

inline void set_events_enabled(bool enabled) noexcept#

inline bool events_enabled() const noexcept#

inline const ComputeEventLog &Events() const noexcept#: Append-only log of decisions made by :cpp:func:ComputeInPlaceReuseGraph.

inline ComputeEventLog &Events() noexcept#

inline void ClearEvents() noexcept#: Empties the decision log without touching computed results.

void WriteToMetadata(GraphProto &graph) const#

Records the computed opportunities into each node’s metadata_props of graph under :cpp:var:kInPlaceReuseMetadataKey, :cpp:var:kReleaseAfterMetadataKey, :cpp:var:kNotUsedAfterMetadataKey, and (when :cpp:func:ComputeInPlaceReuseGraph was called with value tags) :cpp:var:kReleaseAfterShapeTagMetadataKey.

For every node that has at least one in-place opportunity, a single metadata entry is added (or updated in place if the key already exists) whose value lists the opportunities as output_index:input_index:kind triplets separated by ; (kind being equal or greater).

For every node that has releasable last-use inputs, one metadata entry is added (or updated in place) under :cpp:var:kReleaseAfterMetadataKey; the value is a ;-separated list of releasable names.

For every node that has declared graph inputs / initializers reaching their last use, one metadata entry is added (or updated in place) under :cpp:var:kNotUsedAfterMetadataKey; the value is a ;-separated list of those names.

When shape-tag information was provided to :cpp:func:ComputeInPlaceReuseGraph, a further metadata entry is added under :cpp:var:kReleaseAfterShapeTagMetadataKey for every node that has at least one shape-tagged releasable value; the value is a ;-separated list of those names.

Nodes without in-place opportunities, without releasable names, and without last-use input/initializer names are left untouched.

graph must be the same graph passed to :cpp:func:ComputeInPlaceReuseGraph, so that node indices line up with the stored result.

Parameters:: graph – Graph whose nodes are mutated in place.
Throws:: std::invalid_argument – when graph has a different number of nodes than the result stored in *this.

inline void Clear() noexcept#: Empties the stored result.

Private Members

std::unordered_map<std::string, std::string> value_tags_#

std::vector<std::string> node_tags_#

std::unordered_set<std::string> constant_values_#

std::vector<ConstantInfo> node_constant_#

bool custom_value_tags_changed_ = false#

CustomValueTagMap custom_value_tags_#

std::vector<std::vector<InPlaceReuse>> reuse_#

std::vector<std::vector<std::string>> release_after_#

std::vector<std::vector<std::string>> not_used_after_#

std::vector<std::vector<std::string>> release_after_shape_tagged_#

std::vector<NodeMemoryProfile> memory_#

ShapesContext shapes_#

std::vector<int64_t> peak_memory_#

ComputeEventLog events_#

bool events_enabled_ = false#

std::vector<char> node_tag_custom_override_#

std::unordered_map<std::string, int> tag_producer_node_#

std::unordered_map<std::string, std::vector<int>> tag_consumers_#

std::unordered_map<std::string, int> incr_producer_#

std::unordered_map<std::string, int> incr_last_use_#

std::unordered_set<std::string> incr_keep_#

std::unordered_set<std::string> incr_graph_inputs_#

std::unordered_set<std::string> incr_graph_initializers_#

std::unordered_set<std::string> incr_graph_outputs_#

std::unordered_map<std::string, std::optional<expressions::DimType>> incr_byte_size_expr_cache_#

expressions::SimplifiedExpressionCache incr_simplified_dim_cache_#

Private Static Functions

static inline std::string NormalizeDomain(const std::string &domain)#

static inline std::string MakeCustomValueTagKey(const std::string &domain, const std::string &op_type)#

struct ComputeEvent#

#include <compute_context.h>

One entry of the optional :cpp:class:ComputeContext decision log.

Public Members

ComputeEventAction action = ComputeEventAction::kInPlace #: Decision kind.

int64_t node_index = -1#: Node index in graph.node() where the decision was made.

std::string name#: Value name for kRelease / kReleaseShapeTag decisions.

int64_t output_index = -1#: Output index for kInPlace decisions; -1 otherwise.

int64_t input_index = -1#: Input index for kInPlace decisions; -1 otherwise.

InPlaceReuseKind kind = InPlaceReuseKind::kEqual #: Match kind for kInPlace decisions.