The QuantumFramework paclet integrates quantum computation capabilities into the Wolfram Language, providing symbolic representations for quantum bases, states, and operators to model and simulate quantum circuits and algorithms. Currently, we are focused on optimizing tensor contraction operations using Einstein Summation, which involves updating tensor indices and mapping superscripts to subscripts. The existing implementation, reliant on functionalities like EinsteinSummation, Transpose, and FindPermutation, constructs and models network graphs but does not leverage the Wolfram Compiler, leading to inefficiencies. We plan to develop compiled versions of these operations, to optimize for speed and memory.

Project Description

Background

Motivation

The QuantumFramework paclet enhances Wolfram Language with quantum computation capabilities, including discrete quantum mechanics. It provides symbolic representations of quantum bases, states, and operators, enabling the modeling and simulation of quantum circuits and algorithms. The paclet also allows inter operation with external platforms and facilitates queries to quantum processing units from Wolfram notebooks.
Our current focus is on optimizing tensor contraction operations using Einstein Summation, which involves updating tensor indices and mapping superscripts to subscripts. The existing implementation utilizes various system functionalities such as EinsteinSummation, Transpose, and FindPermutation.
The current process involves constructing, translating, and modeling a network graph by creating edge lists, mapping vertices, connecting edges, and ranking them through annotation. However, this approach doesn’t leverage the Wolfram Compiler and lacks a compiled version of the function, resulting in slow and inefficient performance.
Our goal is to use FunctionCompile to compile the function body for the network graph, producing machine-independent code for the computer to execute. This should significantly improve the efficiency and speed of the operations.

Objectives

◼
  • Implemented ArrayReshape, Transpose and Dot with arrays as input
    • ◼
  • Annotated each of these functions with appropriate types (primarily `NumericArray`, `PackedArray`for structures and elements with `ComplexReal32` and `MachineInteger` for dimensions
    • ◼
  • Utilized `FunctionCompile` to generate compiled versions of these respective functions and added them to the compiler

    • Tensor Networks: An Introduction

    Reference: https://www.wolframcloud.com/obj/nikm/Published/TensorNetworkContractionCompilation.nb
    In[]:=
    <<Wolfram`QuantumFramework`

    Example: Quantum Circuit Operator and Tensor Networks

    We represent a quantum circuit using the `QuantumCircuitOperator` function in the QuantumFramework paclet
    Reference: https://resources.wolframcloud.com/PacletRepository/resources/Wolfram/QuantumFramework/
    In the call below, we have `CHSH` inequality which is used in the Bell’s circuit that states that hidden variables cannot account for some consequences of quantum entanglement in quantum circuits
    In[]:=
    circuit = QuantumCircuitOperator["CHSH"[p]][[;;-5]]
    Out[]=
    QuantumCircuitOperator
    ​
    

    Generating the TensorNetwork

    The `TensorNetwork`can be generated from the `circuit` object defined above:
    In[]:=
    net = circuit["TensorNetwork", EdgeLabels  "EdgeTag"]
    Out[]=
    The tensor network is a graph annotated with tensors and contraction indices.
    Each edge represents (i.e., is tagged by) a contraction

    Getting Indices

    TensorNetworkIndexGraph[...] transforms a tensor network into a new graph with indices as vertices:
    In[]:=
    TensorNetworkIndices[net]
    Out[]=
    {
    1
    1
    ,
    4
    1
    },{
    2
    2
    },{
    3
    3
    },{
    1
    4
    ,
    4
    1
    },{
    1
    5
    ,
    2
    5
    ,
    5
    1
    ,
    5
    2
    },{
    4
    6
    ,
    6
    4
    },
    3
    7
    ,
    4
    7
    ,
    7
    3
    ,
    7
    4
    

    Explicit Representation

    The first tensor in the network has the following contents:
    In[]:=
    TensorNetworkTensors[net][[1]] // Normal
    Out[]=
    
    1
    2
    ,0,0,
    1
    2
    
    The second tensor in the network is rank-2 matrix:
    In[]:=
    TensorNetworkTensors[net][[2]] // Dimensions
    Out[]=
    {2}
    The fifth tensor in the network is rank-4 matrix:
    In[]:=
    TensorNetworkTensors[net][[5]] // Dimensions
    Out[]=
    {2,2,2,2}

    Tensor Network Contraction

    Goal: To optimize the current implementation of the operation in terms of computation time and memory efficiency
    In[]:=
    ContractTensorNetwork[net]//AbsoluteTiming
    Out[]=
    0.002258,SparseArray
    Specified elements: 16
    Dimensions: {2,2,2,2}
    

    Ground work

    In order to optimize the above operation we need to compile parts of it.
    These following two functions need to be adapted for this end, i.e. we want to modify these functions so that they return a Function that we can then feed to FunctionCompile:

    Einstein Summation
    

    Tensor Network Contract

    Reference: https://github.com/sw1sh/QuantumFramework/blob/8d9234277b25cb74c316247692b3e37052a1153a/QuantumFramework/Kernel/QuantumCircuitOperator/TensorNetwork.m#L14

    Previous Version: TensorNetworkContractPath

    TensorNetworkContractPath performs a series of tensor contractions on a tensor network following a specified sequence of operations, with the goal of reducing the network to a single tensor while preserving the free indices.
    It allows the user to define a specific sequence (path) of contraction steps, specifying which tensors and indices to contract at each step, ensuring that the resulting single tensor maintains the correct order of free indices as specified by the network.
    Steps:
    1
    .
    Input: Tensor Network
    2
    .
    Initialization
    2
    .
    1
    .
    Extracting tensors, indices and free indices from network
    2
    .
    2
    .
    Updating each index with edge tags
    3
    .
    Iteration over Optimal Path
    3
    .
    2
    .
    Looping over each element in path
    3
    .
    3
    .
    For each tensor index:
    3
    .
    3
    .
    1
    .
    Moving the tensor to the end of the list
    3
    .
    4
    .
    For a pair of tensor indices:
    3
    .
    4
    .
    1
    .
    Using `SymmetricDifference` calculating output indices
    3
    .
    4
    .
    2
    .
    Using `EinsteinSummation` to contract both tensors
    3
    .
    4
    .
    3
    .
    Updating tensors and indices
    4
    .
    Transposing the final tensor to match order of free indices
    In[]:=
    TensorNetworkContractPath2[net_ ? TensorNetworkQ, path_] := Enclose @ Block[{tensors, indices, freeIndices},​​ tensors = TensorNetworkTensors[net];​​ indices = TensorNetworkIndices[net];​​ freeIndices = TensorNetworkFreeIndices[net];​​ indices = Replace[indices, Rule @@@ EdgeTags[net], {2}];​​ Do[​​ Replace[p, {​​ {i_} :> (​​ {tensors, indices} = Append[Delete[#, {i}], #[[i]]] & /@ {tensors, indices}​​ ),​​ {i_, j_} :> Block[{out, tensor},​​ out = SymmetricDifference @@ Extract[indices, {{i}, {j}}];​​ tensor = einsum[indices[[{i, j}]] -> out, tensors[[i]], tensors[[j]]];​​ tensors = Append[Delete[tensors, {{i}, {j}}], tensor];​​ indices = Append[Delete[indices, {{i}, {j}}], out]​​ ]​​ }],​​ {p, path}​​ ];​​ ConfirmAssert[Length[tensors] == Length[indices] == 1];​​ ConfirmAssert[ContainsAll[indices[[1]], freeIndices]];​​ Transpose[tensors[[1]], FindPermutation[indices[[1]], freeIndices]]​​];
    Contract the net object defined in the previous section alongside an specified path:

    Modified Version: TensorNetworkContractPathInstructions

    ◼
  • Input: net (tensor network), path (contraction path)
    • ◼
  • Steps:
    • ◼
  • Iterating through each pair or single index specified in `path`
    • ◼
  • Single Tensor Contraction:
    • ◼
  • For each index `i`, we update the dimensions and indices parameters and move the specific tensor to the end of the list
    • ◼
  • Pair of Tensors Contraction:
    • ◼
  • For each pair of indices `{i, j}`, we determine the output indices post contraction by computing the symmetric difference
    • ◼
  • Next, we append the new instruction to our list specifying the contraction operation
    • ◼
  • We utilize the EinsteinSummation (defined above) to denote the contraction between tensors `i` and `j
    • ◼
  • `i`: index of the first tensor in the pair
    • ◼
  • `j`: index of second tensor in the pair
    • ◼
  • Updating dimensions and indices to reflect new tensor resulting from contraction
    • ◼
  • Appending the final transpose instruction to match the order of free indices

    • Implementation of ArrayReshape and Transpose in the Wolfram Compiler

    Compiling Tensor Network Contraction

    Compiling the output of ToEinsteinSummationFunction

    Benchmarks

    TODO: Make some benchmarks

    Preliminary Efforts and Observations

    The following subsections below highlight the initial attempts towards creating manually implementing the built-in functions for ArrayReshape and Transpose. In addition, we also attempted to declare a new component within the compiler for the type `RanklessNumericArray` which gives us the ability to reshape without depending on the rank of each matrix, however this approach does not yield appropriate results are internally the rank is required. One positive effort lead to faster compiled versions of both ArrayReshape (for 1-dimensions) and Transpose (for 1-Dimension) than the paclet versions of these respective functions. Through this process, we believe if the `FunctionCompile` operation would be less restrictive in terms of typing and subtyping, in addition, to restricting the format of declaring a pure, anonymous function to be accepted by the `FunctionCompile` might be restrictive as well.

    Initial Work on Compiling ArrayReshape and Transpose: A review of different cases

    Next Steps (Future Deliverables)

    ◼
  • Optimizing ArrayReshape using BLAS operation
    • ◼
  • The current implementation of ArrayReshape assumes that the number of elements in the output is the same as in the input, which involves simply adjusting the lengths. However, this approach might not be sufficient for more complex cases where the number of elements in the output differs from the input
    • ◼
  • BLAS would be set of subprograms which have levels of definitions, implementing in higher-order levels can affect time complexity (Level 1 is O(n) however Level 3 is O(n^3), but higher levels help us compute complex dot products and matrix-vector operations, this would allow us to produce highly optimized code
    • ◼
  • Compiling Transpose and ArrayReshape for the following overloaded versions, by adding cases to the existing functionality (current progress):
    • ◼
  • Comparing against existing benchmark for computational efficiency and time complexity

    • Acknowledgments

    I want to thank my Mentors (Daniel Sanchez and Christian Pasquel) for the constant support, assistance and guidance throughout the period working on this project and special thanks to Nikolay Murzin and Tom Wickham-Jones for their feedback . I hope to work with them in future endeavours and have a chance to collaborate on future work.

    References

    1
    .
    Abdul Dakkak, Tom Wickham-Jones, and Wen-mei Hwu. 2020. The design and implementation of the wolfram language compiler. In Proceedings of the 18th ACM/IEEE International Symposium on Code Generation and Optimization (CGO 2020). Association for Computing Machinery, New York, NY, USA, 212–228. https://doi.org/10.1145/3368826.3377913
    2
    .
    QuantumFramework: https://resources.wolframcloud.com/PacletRepository/resources/Wolfram/QuantumFramework

    Cite This Notebook

    “Optimizing tensor contraction using Wolfram compiler”
    by Vedant Tewari​
    ​https://community.wolfram.com/groups/-/m/t/3209893