CuDensityMat API Reference

CuDensityMat

Julia bindings for NVIDIA's cuDensityMat library — high-performance density matrix simulation for analog quantum dynamics (Lindblad master equation, time-dependent Hamiltonians, backward differentiation/gradients).

Part of the CuQuantum.jl package.

DenseMixedState{T}(ws, hilbert_space_dims; batch_size=1)

Mixed state in dense (density-matrix) representation.

Arguments

• ws::WorkStream: Execution context
• hilbert_space_dims::Vector{Int} or Tuple: Local Hilbert space dimensions
• batch_size::Int=1: Batch dimension

Example

ws = WorkStream()
rho = DenseMixedState{ComplexF64}(ws, (2, 2))
allocate_storage!(rho)

DensePureState{T}(ws, hilbert_space_dims; batch_size=1)

Pure state in dense (state-vector) representation.

Arguments

• ws::WorkStream: Execution context
• hilbert_space_dims::Vector{Int} or Tuple: Local Hilbert space dimensions
• batch_size::Int=1: Batch dimension

Example

ws = WorkStream()
psi = DensePureState{ComplexF64}(ws, (2, 2, 2))
allocate_storage!(psi)

ElementaryOperator

Wraps a cudensitymatElementaryOperator_t handle. Represents a single-mode operator that can be dense or multidiagonal.

Create via create_elementary_operator or create_elementary_operator_batch.

Expectation

Wraps a cudensitymatExpectation_t handle for computing expectation values of an operator with respect to a quantum state.

MatrixOperator

Wraps a cudensitymatMatrixOperator_t handle. Represents a full Hilbert-space dense matrix operator stored locally on a single GPU.

Operator

Wraps a cudensitymatOperator_t handle. A composite operator consisting of multiple OperatorTerms, each with a duality flag and scalar coefficient.

OperatorAction

Wraps a cudensitymatOperatorAction_t handle. Computes the action of one or more operators on input states, accumulating into an output state.

OperatorSpectrum

Wraps a cudensitymatOperatorSpectrum_t handle for computing eigenvalues and eigenstates of an operator.

OperatorTerm

Wraps a cudensitymatOperatorTerm_t handle. Represents one or more tensor products of elementary operators acting on specified modes.

WorkStream(; stream=nothing, memory_limit=nothing, device_id=nothing)

A workspace context bundling a CuDensityMat library handle, workspace descriptor, and CUDA stream. All CuDensityMat operations require a WorkStream.

Keyword Arguments

• stream::Union{Nothing, CUDA.CuStream}: CUDA stream. Uses the default stream if nothing.
• memory_limit::Union{Nothing, Int}: Maximum workspace memory in bytes. Default is unlimited.
• device_id::Union{Nothing, Int}: CUDA device ordinal. Uses current device if nothing.

Examples

ws = WorkStream()                          # default stream, current device
ws = WorkStream(stream=CUDA.CuStream())    # explicit stream
close(ws)                                  # release resources

close(ws::WorkStream)

Explicitly release all resources held by the WorkStream.

isopen(ws::WorkStream)

Check if the WorkStream is still valid (not yet closed/finalized).

_get_workspace_limit(ws, user_limit) -> Int

Return the workspace size limit in bytes. If user_limit is nothing, returns ws.memory_limit if set, or 80% of free GPU memory.

allocate_storage!(state)

Allocate an appropriately sized buffer and attach it to the state.

append_elementary_product!(term, elem_operators, modes_acted_on, mode_action_duality;
    coefficient=1.0+0im, kwargs...)

Append a tensor product of elementary operators to the term.

Arguments

• elem_operators: Vector of ElementaryOperators
• modes_acted_on: Flattened vector of mode indices (0-based) each operator acts on
• mode_action_duality: Per-mode duality flags (0=ket, nonzero=bra) — same length as modes_acted_on
• coefficient: Static complex scalar coefficient (default: 1.0+0im)

append_elementary_product_batch!(term, elem_operators, modes_acted_on,
    mode_action_duality, batch_size, static_coefficients; kwargs...)

Append a batched tensor product of elementary operators to the term.

Arguments

• static_coefficients: GPU array of ComplexF64 coefficients (length batch_size)
• total_coefficients: GPU storage for total coefficients, or nothing

append_matrix_product!(term, matrix_operators, conjugations, action_duality;
    coefficient=1.0+0im, kwargs...)

Append a product of matrix operators to the term.

Arguments

• matrix_operators: Vector of MatrixOperators
• conjugations: Per-operator conjugation flags (0=normal, nonzero=conjugate-transpose)
• action_duality: Per-operator duality flags (0=ket, nonzero=bra)
• coefficient: Static complex scalar coefficient (default: 1.0+0im)

append_term!(operator, term; duality=0, coefficient=1.0+0im, kwargs...)

Append an operator term to the composite operator.

Arguments

• duality: 0 for ket-side (default), nonzero for bra-side
• coefficient: Static complex scalar coefficient (default: 1.0+0im)

append_term_batch!(operator, term, batch_size, static_coefficients;
    duality=0, kwargs...)

Append a batched operator term.

Arguments

• static_coefficients: GPU array of ComplexF64 coefficients (length batch_size)
• total_coefficients: GPU storage for total coefficients, or nothing

attach_storage!(state, data::CuVector{T})

Attach a GPU buffer to the state. The buffer must be F-contiguous and match the required storage size.

clone(state, buf::CuVector{T}) -> State

Clone a state with a new storage buffer.

component_storage_size(state) -> Vector{Int}

Storage size in bytes for each component.

compute_action!(ws, action, states_in, state_out;
    time=0.0, batch_size=0, params=nothing)

Execute the operator action: state_out = sum_i operator_i(states_in[i]).

compute_expectation!(ws, expectation, state, result;
    time=0.0, batch_size=1, num_params=0, params=nothing)

Compute expectation value: result = Tr(operator * state).

Arguments

• result: GPU array to store the expectation value(s). For complex operators, use CuVector{ComplexF64} of length batch_size.

compute_operator_action!(ws, operator, state_in, state_out;
    time=0.0, batch_size=0, params=nothing)

Execute single-operator action: state_out = operator * state_in.

compute_operator_action_backward!(ws, operator, state_in, state_out_adj,
    state_in_adj, params_grad; time=0.0, batch_size=0,
    num_params=0, params=nothing)

Compute backward differentiation of a single-operator action.

Computes ∂L/∂statein (into `stateinadj) and ∂L/∂params (intoparamsgrad) given ∂L/∂state_out (stateoutadj`).

Arguments

• state_in: Forward-pass input state
• state_out_adj: Adjoint of forward-pass output (∂L/∂state_out)
• state_in_adj: Output — receives ∂L/∂state_in (accumulated, += semantics)
• params_grad: GPU vector — receives ∂L/∂params (accumulated, += semantics)

compute_spectrum!(ws, spectrum, num_eigenstates, eigenstates, eigenvalues;
    time=0.0, batch_size=1, num_params=0, params=nothing) -> tolerances

Compute eigenvalues and eigenstates.

Arguments

• num_eigenstates: Number of eigenvalues/eigenstates to compute (≤ max from prepare)
• eigenstates: Vector of pre-allocated AbstractStates to receive eigenstates
• eigenvalues: GPU array to receive eigenvalues. For Hermitian operators, use CuVector{Float64} of length num_eigenstates * batch_size. For non-Hermitian, use CuVector{ComplexF64}.

Returns

• tolerances: Vector{Float64} of convergence tolerances for each eigenvalue

configure_spectrum!(ws, spectrum, attribute, value)

Configure a spectrum solver parameter.

Arguments

• attribute: One of:
  • CUDENSITYMAT_OPERATOR_SPECTRUM_CONFIG_MAX_EXPANSION — max Krylov expansion ratio (Int32, default 5)
  • CUDENSITYMAT_OPERATOR_SPECTRUM_CONFIG_MAX_RESTARTS — max restarts (Int32, default 20)
  • CUDENSITYMAT_OPERATOR_SPECTRUM_CONFIG_MIN_BLOCK_SIZE — min block size (Int32, default 1)
• value: Int32 configuration value

create() -> cudensitymatHandle_t

Create a CuDensityMat library context handle.

create_elementary_operator(ws, space_mode_extents, data::CuArray{T};
    sparsity=:none, diagonal_offsets=Int32[],
    tensor_callback=nothing, tensor_gradient_callback=nothing) -> ElementaryOperator

Create a dense or multidiagonal elementary operator for a single mode.

Arguments

• ws::WorkStream: Execution context
• space_mode_extents::Vector{Int}: Dimensions of the space modes (e.g., [d] for single qudit)
• data::CuArray{T}: Operator data on GPU (column-major)
• sparsity: :none for dense, :multidiagonal for sparse
• diagonal_offsets: Required for multidiagonal sparsity

create_elementary_operator_batch(ws, space_mode_extents, data::CuArray{T}, batch_size;
    kwargs...) -> ElementaryOperator

Create a batched elementary operator.

create_expectation(ws, operator) -> Expectation

Create an expectation value descriptor for the given operator.

create_matrix_operator(ws, space_mode_extents, data::CuArray{T};
    tensor_callback=nothing, tensor_gradient_callback=nothing) -> MatrixOperator

Create a dense local matrix operator covering the full Hilbert space.

create_matrix_operator_batch(ws, space_mode_extents, data::CuArray{T}, batch_size;
    kwargs...) -> MatrixOperator

Create a batched dense local matrix operator.

create_operator(ws, hilbert_space_dims) -> Operator

Create an empty composite operator for the given Hilbert space.

create_operator_action(ws, operators::Vector{Operator}) -> OperatorAction

Create an operator action descriptor for multiple operators.

create_operator_spectrum(ws, operator;
    is_hermitian=true,
    spectrum_kind=CUDENSITYMAT_OPERATOR_SPECTRUM_LARGEST) -> OperatorSpectrum

Create an eigenspectrum solver for the given operator.

Arguments

• is_hermitian: Whether the operator is Hermitian (enables optimized solver)
• spectrum_kind: Which eigenvalues to compute:
  • CUDENSITYMAT_OPERATOR_SPECTRUM_LARGEST — largest by magnitude
  • CUDENSITYMAT_OPERATOR_SPECTRUM_SMALLEST — smallest by magnitude
  • CUDENSITYMAT_OPERATOR_SPECTRUM_LARGEST_REAL — largest real part
  • CUDENSITYMAT_OPERATOR_SPECTRUM_SMALLEST_REAL — smallest real part

create_operator_term(ws, hilbert_space_dims) -> OperatorTerm

Create an empty operator term for the given Hilbert space dimensions.

create_workspace(handle) -> cudensitymatWorkspaceDescriptor_t

Create a workspace descriptor.

destroy(handle::cudensitymatHandle_t)

Destroy a CuDensityMat library context handle.

destroy_workspace(workspace)

Destroy a workspace descriptor.

get_num_ranks(ws::WorkStream) -> Int

Return the number of distributed processes.

get_proc_rank(ws::WorkStream) -> Int

Return the rank of the current process.

handle() -> cudensitymatHandle_t

Get a cuDensityMat library handle for the current CUDA context. Handles are cached per-context and reused.

initialize_zero!(state)

Set all elements of the state to zero using the C API.

inner_product(left, right) -> Vector{T}

Compute inner product(s) ⟨left|right⟩. Returns a CPU vector of length batch_size.

inplace_accumulate!(dest, src, factors=1)

Accumulate: dest += factors * src. Both states must be compatible.

inplace_scale!(state, factors)

Scale the state by scalar factor(s).

factors can be a scalar, a Vector, or a CuVector of length batch_size.

local_info(state) -> (shape::Tuple, offsets::Tuple)

Local storage buffer dimensions and mode offsets. The last dimension is always the batch dimension.

norm(state) -> Vector{real(T)}

Compute the squared Frobenius norm(s). Returns a CPU vector of length batch_size.

num_components(state) -> Int

Number of storage components (1 for single-GPU, >1 for distributed).

prepare_action!(ws, action, states_in, state_out;
    compute_type=CUDENSITYMAT_COMPUTE_64F, workspace_limit=nothing)

Prepare the operator action computation (one-time setup).

If workspace_limit is nothing (default), uses 80% of free GPU memory.

prepare_expectation!(ws, expectation, state;
    compute_type=CUDENSITYMAT_COMPUTE_64F, workspace_limit=nothing)

Prepare expectation value computation (one-time setup).

prepare_operator_action!(ws, operator, state_in, state_out;
    compute_type=CUDENSITYMAT_COMPUTE_64F, workspace_limit=nothing)

Prepare single-operator action computation.

If workspace_limit is nothing (default), uses 80% of free GPU memory.

prepare_operator_action_backward!(ws, operator, state_in, state_out_adj;
    compute_type=CUDENSITYMAT_COMPUTE_64F, workspace_limit=nothing)

Prepare backward differentiation of a single-operator action.

Arguments

• state_in: The forward-pass input state
• state_out_adj: The adjoint of the forward-pass output state (∂L/∂state_out)

prepare_spectrum!(ws, spectrum, max_eigenstates, state;
    compute_type=CUDENSITYMAT_COMPUTE_64F, workspace_limit=nothing)

Prepare eigenspectrum computation.

Arguments

• max_eigenstates: Maximum number of eigenvalues/eigenstates to compute
• state: A template state (defines the Hilbert space structure)

release_handle(h::cudensitymatHandle_t)

Return a handle to the cache for reuse.

release_workspace!(ws::WorkStream)

Release the workspace buffer and create a fresh workspace descriptor.

set_communicator!(ws::WorkStream, provider::Symbol; comm_ptr=nothing, comm_size=nothing)

Configure distributed execution on the WorkStream.

Arguments

• provider: :none, :mpi, or :nccl
• comm_ptr: Pointer to the communicator (MPIComm* or ncclCommt*)
• comm_size: Size of the communicator object in bytes

MPI Example

using MPI
MPI.Init()
comm = MPI.COMM_WORLD
set_communicator!(ws, :mpi;
    comm_ptr=MPI.API.MPI_Comm_c2f(comm),
    comm_size=sizeof(MPI.MPI_Comm))

set_random_seed!(ws::WorkStream, seed::Integer)

Reset the random seed used by the library's internal random number generator.

state_view(state) -> CuArray

Return a multidimensional view of the state's storage buffer.

storage_size(state) -> Int

Storage buffer size in number of elements of the state's data type.

trace(state) -> Vector{T}

Compute the trace(s). Returns a CPU vector of length batch_size.

unregister_callback!(refs)

Release a callback from the global registry, allowing GC. Call this after destroying the operator that uses the callback.

version()

Return the CuDensityMat library version as a VersionNumber.

workspace_allocate!(ws::WorkStream, size::Int; memspace=:device, kind=:scratch)

Allocate and attach a workspace buffer of the given size.

workspace_query_size(ws::WorkStream; memspace=:device, kind=:scratch) -> Int

Query the required workspace buffer size in bytes after a prepare call.

wrap_scalar_callback(f; gradient=nothing) -> (cb, gcb, refs)

Wrap a Julia function as a scalar callback for time-dependent coefficients.

Arguments

• f: Function with signature f(time::Float64, params::Matrix{Float64}, storage::Vector{T}) where T matches the operator's data type. params is (numParams, batchSize), storage is (batchSize,). Write the coefficient values into storage.
• gradient: Optional gradient function with signature g(time::Float64, params::Matrix{Float64}, scalar_grad::Vector{T}, params_grad::Matrix{Float64})

Returns

• cb::cudensitymatWrappedScalarCallback_t — the callback struct
• gcb::cudensitymatWrappedScalarGradientCallback_t — the gradient callback struct
• refs — opaque reference holder (keep alive while callback is in use)

Example

# Time-dependent coefficient: f(t) = exp(iΩt) where Ω = params[1,:]
function my_coeff(time, params, storage)
    for i in axes(storage, 1)
        ω = params[1, i]
        storage[i] = complex(cos(ω * time), sin(ω * time))
    end
end

cb, gcb, refs = wrap_scalar_callback(my_coeff)

wrap_tensor_callback(f; gradient=nothing) -> (cb, gcb, refs)

Wrap a Julia function as a tensor callback for time-dependent operator elements.

Arguments

• f: Function with signature f(time::Float64, params::Matrix{Float64}, storage::Array{T}) where storage shape is (d1, d2, ..., d1, d2, ..., batchSize) for dense operators. Write the tensor elements into storage.
• gradient: Optional gradient function.

Returns

• cb::cudensitymatWrappedTensorCallback_t
• gcb::cudensitymatWrappedTensorGradientCallback_t
• refs — opaque reference holder