Tensor Registry#

class TensorRegistry#

Central registry of all tensor specs needed by a TRT engine.

TensorRegistry holds tensor specs with symbolic shape templates. bindAll() resolves all shapes from InferenceDims and calls setTensorAddress + setInputShape in a single pass, replacing many manual binding calls. Per-layer tensors (e.g. KV cache) are auto-expanded from d templates.

Public Functions

void addTensor(TensorSpec spec)#

Add a tensor specification to the registry.

Parameters:

spec – The tensor specification to add

bool bindAll(
nvinfer1::IExecutionContext *ctx,
TensorMap const &tensorMap,
InferenceDims const &inferenceDims
) const#

Bind all registered tensors to a TRT execution context.

Resolves symbolic shapes from inferenceDims, then calls setTensorAddress and (for inputs) setInputShape on each tensor.

Parameters:

  • ctx – TRT execution context

  • tensorMap – Name-to-tensor mapping

  • inferenceDims – Current symbolic dimension values

Returns:

True if all tensors were bound successfully

std::vector<int64_t InferenceDims::*> const &referencedMembers(
) const#

Return the set of InferenceDims members referenced by any symbolic ShapeDim in the registry.

Used by EngineExecutor::prepare to restrict its field-completeness assertion to dims the engine actually consumes. Computed once (lazily) when the expanded spec set is materialized; cached across calls.

Returns:

Vector of int64_t InferenceDims::* — deduplicated

std::vector<std::string> allTensorNames() const#

Return all tensor names after per-layer expansion.

Returns:

Vector of concrete tensor names

bool contains(std::string_view name) const#

Test whether a tensor name is registered (post-expansion).

Used by EngineExecutor::prepare to skip the registry-managed bindings in its fall-back loop: those are already covered by bindAll, so re-binding them via the tensor map’s actual TRTDims would override the registry-computed shape.

Parameters:

name – Concrete tensor name (after d expansion if applicable)

Returns:

True if the registry has a spec for name

std::vector<TensorSpec> const &allExpandedSpecs() const#

Return all tensor specs after per-layer expansion.

Returns:

Vector of expanded TensorSpec objects

nvinfer1::Dims resolveShape(
std::vector<ShapeDim> const &shape,
InferenceDims const &inferenceDims
) const#

Resolve a shape template into concrete TRT Dims.

Fixed dimensions are used as-is. Symbolic dimensions are looked up from inferenceDims.

Parameters:

  • shape – Shape template with fixed and/or symbolic dims

  • inferenceDims – Current symbolic dimension values

Returns:

Resolved TRT Dims

struct TensorSpec#

Declarative specification for a single tensor binding.

Per-layer tensors use a printf-style d placeholder in the name and a positive perLayer count; the registry expands these into concrete specs at query time.

Public Members

std::string name#

Tensor name (may contain “%d” for per-layer)

TensorIO io = {TensorIO::kInput}#

Binding direction.

nvinfer1::DataType dtype = {nvinfer1::DataType::kHALF}#

Data type.

std::vector<ShapeDim> shape#

Shape template.

int32_t perLayer = {0}#

>0: expand “%d” for 0..perLayer-1