Skip to content

Make Tensor handle-backed with lazy materialization #86

@shinaoka

Description

@shinaoka

Motivation

The current Tensor{T,N} object is Julia-dense-data-first:

struct Tensor{T,N}
    data::Array{T,N}
    inds::Vector{Index}
    backend_handle::Union{Nothing,Ptr{Cvoid}}
    structured_storage::Union{Nothing,StructuredTensorStorage{T}}
end

The public constructor copies data, and backend operations such as tensor * tensor rebuild Rust tensor/index handles on every call. The contraction result is then copied back into Julia dense storage, reshaped, copied again by the constructor, and the backend handle is released.

This causes avoidable copies in Tensor4all.jl, especially in repeated tensor contraction / one-hot evaluation paths used by BubbleTeaCI.

Proposed direction

Make Tensor able to own or share a Rust tensor handle, with Julia dense data materialized lazily only when needed.

Possible shape:

mutable struct TensorHandle
    ptr::Ptr{Cvoid}
end

struct Tensor{T,N}
    data::Union{Nothing,Array{T,N}}
    inds::Vector{Index}
    backend_handle::Union{Nothing,TensorHandle}
    structured_storage::Union{Nothing,StructuredTensorStorage{T}}
end

The exact representation can differ, but the key property is that backend results can stay backend-resident.

Multi-index select wrapper

Expose a Tensor4all.jl wrapper around the tensor4all-rs multi-index select C API:

selectinds(t::Tensor, selections::Pair{Index,<:Integer}...)
selectinds(t::Tensor, indices::AbstractVector{Index}, values::AbstractVector{<:Integer})

Julia values should be 1-based and converted to 0-based at the C boundary.

This wrapper should be able to select multiple indices from the same tensor in one call. That matters for:

  • MPS-like TT tensors with one physical/site leg.
  • MPO-like TT tensors with multiple physical/site legs, for example input and output indices.
  • General TreeTN nodes with multiple site legs.
  • TreeTN/TT evaluation by selecting all local site legs of each node before contracting the remaining tensors.

Tensor * OneHotTensor can be implemented as a single-index convenience wrapper over selectinds, but the primary internal operation should support multiple selected indices.

Requirements

  • Tensor(data, inds) keeps existing public semantics.
  • A backend result can construct a Tensor that owns the returned Rust t4a_tensor handle without immediately copying dense data to Julia.
  • Array(tensor), indexing, and other Julia-data consumers materialize dense data on demand.
  • tensor * tensor uses existing handles when available and returns a handle-backed Tensor.
  • selectinds uses backend handles when available and returns a handle-backed Tensor.
  • Handle ownership uses a Julia object with a finalizer; do not store unmanaged raw Ptr{Cvoid} as an owned resource without lifetime management.
  • Index metadata remains exact (dim/id/tags/plev) and must not be weakened to id-only matching.

Acceptance criteria

  • Repeated Tensor contractions avoid rebuilding handles when inputs already have backend handles.
  • Backend contraction results do not eagerly copy dense data to Julia.
  • selectinds can select one or more indices from the same tensor without dense Julia materialization.
  • Existing tests for tensor arithmetic, contraction, SVD/QR, and ITensorCompat still pass.
  • New tests verify lazy materialization, handle lifetime behavior, and multi-index selection for MPS-like and MPO-like tensors.
  • Benchmarks or allocation tests show reduced allocations for tensor * tensor chains and tensor evaluation via local selections plus prod(tensors).

Related

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions