The following code example illustrates how to integrate cuTensorNet functionalities to perform basic MPS simulation.
The workflow is encapsulated in an MPSHelper
class. The full code can be found in
the NVIDIA/cuQuantum repository (here).
Define MPSHelper class¶
We first define an MPSHelper
class to keep track of the modes and extents of all physical and virtual bonds.
The simulation settings are also stored in this class. Once out of scope, all resource owned by this class will be freed.
75class MPSHelper
76{
77 public:
78 /**
79 * \brief Construct an MPSHelper object for gate splitting algorithm.
80 * i j k
81 * -------A-------B------- i j k
82 * p| |q -------> -------A`-------B`-------
83 * GGGGGGGGG r| |s
84 * r| |s
85 * \param[in] numSites The number of sites in the MPS
86 * \param[in] physExtent The extent for the physical mode where the gate tensors are acted on.
87 * \param[in] maxVirtualExtent The maximal extent allowed for the virtual mode shared between adjacent MPS tensors.
88 * \param[in] initialVirtualExtents A vector of size \p numSites-1 where the ith element denotes the extent of the shared mode for site i and site i+1 in the beginning of the simulation.
89 * \param[in] typeData The data type for all tensors and gates
90 * \param[in] typeCompute The compute type for all gate splitting process
91 */
92 MPSHelper(int32_t numSites,
93 int64_t physExtent,
94 int64_t maxVirtualExtent,
95 const std::vector<int64_t>& initialVirtualExtents,
96 cudaDataType_t typeData,
97 cutensornetComputeType_t typeCompute);
98
99 /**
100 * \brief Initialize the MPS metadata and cutensornet library.
101 */
102 cutensornetStatus_t initialize();
103
104 /**
105 * \brief Compute the maximal number of elements for each site.
106 */
107 std::vector<size_t> getMaxTensorElements() const;
108
109 /**
110 * \brief Update the SVD truncation setting.
111 * \param[in] absCutoff The cutoff value for absolute singular value truncation.
112 * \param[in] relCutoff The cutoff value for relative singular value truncation.
113 * \param[in] renorm The option for renormalization of the truncated singular values.
114 * \param[in] partition The option for partitioning of the singular values.
115 */
116 cutensornetStatus_t setSVDConfig(double absCutoff,
117 double relCutoff,
118 cutensornetTensorSVDNormalization_t renorm,
119 cutensornetTensorSVDPartition_t partition);
120
121 /**
122 * \brief Update the algorithm to use for the gating process.
123 * \param[in] gateAlgo The gate algorithm to use for MPS simulation.
124 */
125 void setGateAlgorithm(cutensornetGateSplitAlgo_t gateAlgo) {gateAlgo_ = gateAlgo;}
126
127 /**
128 * \brief Compute the maximal workspace needed for MPS gating algorithm.
129 * \param[out] workspaceSize The required workspace size on the device.
130 */
131 cutensornetStatus_t computeMaxWorkspaceSizes(int64_t* workspaceSize);
132
133 /**
134 * \brief Compute the maximal workspace needed for MPS gating algorithm.
135 * \param[in] work Pointer to the allocated workspace.
136 * \param[in] workspaceSize The required workspace size on the device.
137 */
138 cutensornetStatus_t setWorkspace(void* work, int64_t workspaceSize);
139
140 /**
141 * \brief In-place execution of the apply gate algorithm on \p siteA and \p siteB.
142 * \param[in] siteA The first site where the gate is applied to.
143 * \param[in] siteB The second site where the gate is applied to. Must be adjacent to \p siteA.
144 * \param[in,out] dataInA The data for the MPS tensor at \p siteA. The input will be overwritten with output mps tensor data.
145 * \param[in,out] dataInB The data for the MPS tensor at \p siteB. The input will be overwritten with output mps tensor data.
146 * \param[in] dataInG The input data for the gate tensor.
147 * \param[in] verbose Whether to print out the runtime information regarding truncation.
148 * \param[in] stream The CUDA stream on which the computation is performed.
149 */
150 cutensornetStatus_t applyGate(uint32_t siteA,
151 uint32_t siteB,
152 void* dataInA,
153 void* dataInB,
154 const void* dataInG,
155 bool verbose,
156 cudaStream_t stream);
157
158 /**
159 * \brief Free all the tensor descriptors in mpsHelper.
160 */
161 ~MPSHelper()
162 {
163 if (inited_)
164 {
165 for (auto& descTensor: descTensors_)
166 {
167 cutensornetDestroyTensorDescriptor(descTensor);
168 }
169 cutensornetDestroy(handle_);
170 cutensornetDestroyWorkspaceDescriptor(workDesc_);
171 }
172 if (svdConfig_ != nullptr)
173 {
174 cutensornetDestroyTensorSVDConfig(svdConfig_);
175 }
176 if (svdInfo_ != nullptr)
177 {
178 cutensornetDestroyTensorSVDInfo(svdInfo_);
179 }
180 }
181
182 private:
183 int32_t numSites_; ///< Number of sites in the MPS
184 int64_t physExtent_; ///< Extent for the physical index
185 int64_t maxVirtualExtent_{0}; ///< The maximal extent allowed for the virtual dimension
186 cudaDataType_t typeData_;
187 cutensornetComputeType_t typeCompute_;
188
189 bool inited_{false};
190 std::vector<int32_t> physModes_; ///< A vector of length \p numSites_ storing the physical mode of each site.
191 std::vector<int32_t> virtualModes_; ///< A vector of length \p numSites_+1; For site i, virtualModes_[i] and virtualModes_[i+1] represents the left and right virtual mode.
192 std::vector<int64_t> extentsPerSite_; ///< A vector of length \p numSites_+1; For site i, extentsPerSite_[i] and extentsPerSite_[i+1] represents the left and right virtual extent.
193
194 cutensornetHandle_t handle_;
195 std::vector<cutensornetTensorDescriptor_t> descTensors_; /// A vector of length \p numSites_ storing the cutensornetTensorDescriptor_t for each site
196 cutensornetWorkspaceDescriptor_t workDesc_{nullptr};
197 cutensornetTensorSVDConfig_t svdConfig_{nullptr};
198 cutensornetTensorSVDInfo_t svdInfo_{nullptr};
199 cutensornetGateSplitAlgo_t gateAlgo_{CUTENSORNET_GATE_SPLIT_ALGO_DIRECT};
200 int32_t nextMode_{0}; /// The next mode label to use for labelling site tensors and gates.
201};
Note
For full definition of all the methods, please refer to the sample here.
Setup MPS simulation setting¶
Next, in the main function, we need to choose the simulation setting for the MPS simulation (i.e., the number of sites, the initial extents, and the data type).
557 /***********************************
558 * Step 1: basic MPS setup
559 ************************************/
560
561 // setup the simulation setting for the MPS
562 typedef std::complex<double> complexType;
563 cudaDataType_t typeData = CUDA_C_64F;
564 cutensornetComputeType_t typeCompute = CUTENSORNET_COMPUTE_64F;
565 int32_t numSites = 16;
566 int64_t physExtent = 2;
567 int64_t maxVirtualExtent = 12;
568 const std::vector<int64_t> initialVirtualExtents(numSites-1, 1); // starting MPS with shared extent of 1;
569
570 // initialize an MPSHelper to dynamically update tensor metadats
571 MPSHelper mpsHelper(numSites, physExtent, maxVirtualExtent, initialVirtualExtents, typeData, typeCompute);
572 HANDLE_ERROR( mpsHelper.initialize() );
The MPS metadata and all cuTensorNet library objects will be managed by the MPSHelper
while the data pointers are explicitly
managed in the main function.
Allocate memory and initialize data¶
Next, we allocate memory for the MPS operands and four 2-qubit gate tensors.
The largest tensor size for each MPS tensor can be queried through the MPSHelper
class.
The MPS tensors are initialized to a state corresponding to |00..000>
and the gate tensors are filled with random values.
575 /***********************************
576 * Step 2: data allocation
577 ************************************/
578
579 // query largest tensor sizes for the MPS
580 const std::vector<size_t> maxElementsPerSite = mpsHelper.getMaxTensorElements();
581 std::vector<void*> tensors_h;
582 std::vector<void*> tensors_d;
583 for (int32_t i=0; i<numSites; i++)
584 {
585 size_t maxSize = sizeof(complexType) * maxElementsPerSite.at(i);
586 void* data_h = malloc(maxSize);
587 memset(data_h, 0, maxSize);
588 // initialize state to |0000..0000>
589 *(complexType*)(data_h) = complexType(1,0);
590 void* data_d;
591 HANDLE_CUDA_ERROR( cudaMalloc(&data_d, maxSize) );
592 // data transfer from host to device
593 HANDLE_CUDA_ERROR( cudaMemcpy(data_d, data_h, maxSize, cudaMemcpyHostToDevice) );
594 tensors_h.push_back(data_h);
595 tensors_d.push_back(data_d);
596 }
597
598 // initialize 4 random gate tensors on host and copy them to device
599 const int32_t numRandomGates = 4;
600 const int64_t numGateElements = physExtent * physExtent * physExtent * physExtent; // shape (2, 2, 2, 2)
601 size_t gateSize = sizeof(complexType) * numGateElements;
602 complexType* gates_h[numRandomGates];
603 void* gates_d[numRandomGates];
604
605 for (int i=0; i<numRandomGates; i++)
606 {
607 gates_h[i] = (complexType*) malloc(gateSize);
608 HANDLE_CUDA_ERROR( cudaMalloc((void**) &gates_d[i], gateSize) );
609 for (int j=0; j<numGateElements; j++)
610 {
611 gates_h[i][j] = complexType(((float) rand())/RAND_MAX, ((float) rand())/RAND_MAX);
612 }
613 HANDLE_CUDA_ERROR( cudaMemcpy(gates_d[i], gates_h[i], gateSize, cudaMemcpyHostToDevice) );
614 }
615
Setup gate split options¶
Then we setup the SVD truncation parameters and the algorithm cutensornetGateSplitAlgo_t
to use for the gate split process.
617 /*****************************************
618 * Step 3: setup options for gate operation
619 ******************************************/
620
621 double absCutoff = 1e-2;
622 double relCutoff = 1e-2;
623 cutensornetTensorSVDNormalization_t renorm = CUTENSORNET_TENSOR_SVD_NORMALIZATION_L2; // renormalize the L2 norm of truncated singular values to 1.
624 cutensornetTensorSVDPartition_t partition = CUTENSORNET_TENSOR_SVD_PARTITION_UV_EQUAL; // equally partition the singular values onto U and V;
625 HANDLE_ERROR( mpsHelper.setSVDConfig(absCutoff, relCutoff, renorm, partition));
626
627 cutensornetGateSplitAlgo_t gateAlgo = CUTENSORNET_GATE_SPLIT_ALGO_REDUCED;
628 mpsHelper.setGateAlgorithm(gateAlgo);
Query and allocate required workspace¶
Once all simulation settings are set, we can query the required workspace size. Inside the MPSHelper, the required workspace size is estimated on the largest tensor sizes involved in the simulation.
631 /********************************************
632 * Step 4: workspace size query and allocation
633 *********************************************/
634
635 int64_t workspaceSize;
636 HANDLE_ERROR( mpsHelper.computeMaxWorkspaceSizes(&workspaceSize) );
637
638 void *work = nullptr;
639 std::cout << "Maximal workspace size required: " << workspaceSize << std::endl;
640 HANDLE_CUDA_ERROR( cudaMalloc(&work, workspaceSize) );
641
642 HANDLE_ERROR( mpsHelper.setWorkspace(work, workspaceSize));
643
Execution¶
At this stage, we can execute the simulation by iterating over all the gate tensors.
All the metadata of the MPS will be managed and updated inside the MPSHelper
.
645 /***********************************
646 * Step 5: execution
647 ************************************/
648
649 cudaStream_t stream;
650 HANDLE_CUDA_ERROR( cudaStreamCreate(&stream) );
651 uint32_t numLayers = 10; // 10 layers of gate
652 for (uint32_t i=0; i<numLayers; i++)
653 {
654 uint32_t start_site = i % 2;
655 std::cout << "Cycle " << i << ":" << std::endl;
656 bool verbose = (i == numLayers - 1);
657 for (uint32_t j=start_site; j<numSites-1; j=j+2)
658 {
659 uint32_t gateIdx = rand() % numRandomGates; // pick a random gate tensor
660 std::cout << "apply gate " << gateIdx << " on " << j << " and " << j+1<< std::endl;
661 void *dataA = tensors_d[j];
662 void *dataB = tensors_d[j+1];
663 void *dataG = gates_d[gateIdx];
664 HANDLE_ERROR( mpsHelper.applyGate(j, j+1, dataA, dataB, dataG, verbose, stream) );
665 }
666 }
667
668 HANDLE_CUDA_ERROR( cudaStreamSynchronize(stream) );
Free resources¶
After the simulation, we free up all the data pointers allocated in the main function.
671 /***********************************
672 * Step 6: free resources
673 ************************************/
674
675 std::cout << "Free all resources" << std::endl;
676
677 for (int i=0; i<numRandomGates; i++)
678 {
679 free(gates_h[i]);
680 HANDLE_CUDA_ERROR( cudaFree(gates_d[i]) );
681 }
682
683 for (int32_t i=0; i<numSites; i++)
684 {
685 free(tensors_h.at(i));
686 HANDLE_CUDA_ERROR( cudaFree(tensors_d.at(i)) );
687 }
688
689 HANDLE_CUDA_ERROR( cudaFree(work) );
690 // The MPSHelper destructor will free all internal resources when out of scope
691 return 0;
692}
All cuTensorNet library objects owned by the MPSHelper will be freed once out of scope.