removed additional helper file
This commit is contained in:
tankya2
@@ -1,270 +0,0 @@
|
|||||||
# Copyright (c) 2021-2023, NVIDIA CORPORATION & AFFILIATES
|
|
||||||
#
|
|
||||||
# SPDX-License-Identifier: BSD-3-Clause
|
|
||||||
|
|
||||||
import itertools
|
|
||||||
|
|
||||||
import cupy as cp
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
import cuquantum
|
|
||||||
from cuquantum import cutensornet as cutn
|
|
||||||
|
|
||||||
class MPSHelper:
|
|
||||||
|
|
||||||
"""
|
|
||||||
MPSHelper(num_sites, phys_extent, max_virtual_extent, initial_state, data_type, compute_type)
|
|
||||||
|
|
||||||
Create an MPSHelper object for gate splitting algorithm.
|
|
||||||
i j
|
|
||||||
-------A-------B------- i j k
|
|
||||||
p| |q -------> -------A`-------B`-------
|
|
||||||
GGGGGGGGG r| |s
|
|
||||||
r| |s
|
|
||||||
|
|
||||||
Args:
|
|
||||||
num_sites: The number of sites in the MPS.
|
|
||||||
phys_extents: The extent for the physical mode where the gate tensors are acted on.
|
|
||||||
max_virtual_extent: The maximal extent allowed for the virtual mode shared between adjacent MPS tensors.
|
|
||||||
initial_state: A sequence of :class:`cupy.ndarray` representing the initial state of the MPS.
|
|
||||||
data_type (cuquantum.cudaDataType): The data type for all tensors and gates.
|
|
||||||
compute_type (cuquantum.ComputeType): The compute type for all gate splitting.
|
|
||||||
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(self, num_sites, phys_extent, max_virtual_extent, initial_state, data_type, compute_type):
|
|
||||||
self.num_sites = num_sites
|
|
||||||
self.phys_extent = phys_extent
|
|
||||||
self.data_type = data_type
|
|
||||||
self.compute_type = compute_type
|
|
||||||
|
|
||||||
self.phys_modes = []
|
|
||||||
self.virtual_modes = []
|
|
||||||
self.new_mode = itertools.count(start=0, step=1)
|
|
||||||
|
|
||||||
for i in range(num_sites+1):
|
|
||||||
self.virtual_modes.append(next(self.new_mode))
|
|
||||||
if i != num_sites:
|
|
||||||
self.phys_modes.append(next(self.new_mode))
|
|
||||||
|
|
||||||
untruncated_max_extent = phys_extent ** (num_sites // 2)
|
|
||||||
if max_virtual_extent == 0:
|
|
||||||
self.max_virtual_extent = untruncated_max_extent
|
|
||||||
else:
|
|
||||||
self.max_virtual_extent = min(max_virtual_extent, untruncated_max_extent)
|
|
||||||
|
|
||||||
self.handle = cutn.create()
|
|
||||||
self.work_desc = cutn.create_workspace_descriptor(self.handle)
|
|
||||||
self.svd_config = cutn.create_tensor_svd_config(self.handle)
|
|
||||||
self.svd_info = cutn.create_tensor_svd_info(self.handle)
|
|
||||||
self.gate_algo = cutn.GateSplitAlgo.DIRECT
|
|
||||||
|
|
||||||
self.desc_tensors = []
|
|
||||||
self.state_tensors = []
|
|
||||||
|
|
||||||
# create tensor descriptors
|
|
||||||
for i in range(self.num_sites):
|
|
||||||
temp = initial_state[i]
|
|
||||||
self.state_tensors.append(initial_state[i].astype(temp.dtype, order="F"))
|
|
||||||
extent = self.get_tensor_extent(i)
|
|
||||||
modes = self.get_tensor_modes(i)
|
|
||||||
desc_tensor = cutn.create_tensor_descriptor(self.handle, 3, extent, 0, modes, self.data_type)
|
|
||||||
self.desc_tensors.append(desc_tensor)
|
|
||||||
|
|
||||||
def get_tensor(self, site):
|
|
||||||
"""Get the tensor operands for a specific site."""
|
|
||||||
return self.state_tensors[site]
|
|
||||||
|
|
||||||
def get_tensor_extent(self, site):
|
|
||||||
"""Get the extent of the MPS tensor at a specific site."""
|
|
||||||
return self.state_tensors[site].shape
|
|
||||||
|
|
||||||
def get_tensor_modes(self, site):
|
|
||||||
"""Get the current modes of the MPS tensor at a specific site."""
|
|
||||||
return (self.virtual_modes[site], self.phys_modes[site], self.virtual_modes[site+1])
|
|
||||||
|
|
||||||
def set_svd_config(self, abs_cutoff, rel_cutoff, renorm, partition):
|
|
||||||
"""Update the SVD truncation setting.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
abs_cutoff: The cutoff value for absolute singular value truncation.
|
|
||||||
rel_cutoff: The cutoff value for relative singular value truncation.
|
|
||||||
renorm (cuquantum.cutensornet.TensorSVDNormalization): The option for renormalization of the truncated singular values.
|
|
||||||
partition (cuquantum.cutensornet.TensorSVDPartition): The option for partitioning of the singular values.
|
|
||||||
"""
|
|
||||||
|
|
||||||
if partition != cutn.TensorSVDPartition.UV_EQUAL:
|
|
||||||
raise NotImplementedError("this basic example expects partition to be cutensornet.TensorSVDPartition.UV_EQUAL")
|
|
||||||
|
|
||||||
svd_config_attributes = [cutn.TensorSVDConfigAttribute.ABS_CUTOFF,
|
|
||||||
cutn.TensorSVDConfigAttribute.REL_CUTOFF,
|
|
||||||
cutn.TensorSVDConfigAttribute.S_NORMALIZATION,
|
|
||||||
cutn.TensorSVDConfigAttribute.S_PARTITION]
|
|
||||||
|
|
||||||
for (attr, value) in zip(svd_config_attributes, [abs_cutoff, rel_cutoff, renorm, partition]):
|
|
||||||
dtype = cutn.tensor_svd_config_get_attribute_dtype(attr)
|
|
||||||
value = np.array([value], dtype=dtype)
|
|
||||||
cutn.tensor_svd_config_set_attribute(self.handle,
|
|
||||||
self.svd_config, attr, value.ctypes.data, value.dtype.itemsize)
|
|
||||||
|
|
||||||
def set_gate_algorithm(self, gate_algo):
|
|
||||||
"""Set the algorithm to use for all gate split operations.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
gate_algo (cuquantum.cutensornet.GateSplitAlgo): The gate splitting algorithm to use.
|
|
||||||
"""
|
|
||||||
|
|
||||||
self.gate_algo = gate_algo
|
|
||||||
|
|
||||||
def compute_max_workspace_sizes(self):
|
|
||||||
"""Compute the maximal workspace needed for MPS gating algorithm."""
|
|
||||||
modes_in_A = [ord(c) for c in ('i', 'p', 'j')]
|
|
||||||
modes_in_B = [ord(c) for c in ('j', 'q', 'k')]
|
|
||||||
modes_in_G = [ord(c) for c in ('p', 'q', 'r', 's')]
|
|
||||||
modes_out_A = [ord(c) for c in ('i', 'r', 'j')]
|
|
||||||
modes_out_B = [ord(c) for c in ('j', 's', 'k')]
|
|
||||||
|
|
||||||
max_extents_AB = (self.max_virtual_extent, self.phys_extent, self.max_virtual_extent)
|
|
||||||
extents_in_G = (self.phys_extent, self.phys_extent, self.phys_extent, self.phys_extent)
|
|
||||||
|
|
||||||
desc_tensor_in_A = cutn.create_tensor_descriptor(self.handle, 3, max_extents_AB, 0, modes_in_A, self.data_type)
|
|
||||||
desc_tensor_in_B = cutn.create_tensor_descriptor(self.handle, 3, max_extents_AB, 0, modes_in_B, self.data_type)
|
|
||||||
desc_tensor_in_G = cutn.create_tensor_descriptor(self.handle, 4, extents_in_G, 0, modes_in_G, self.data_type)
|
|
||||||
desc_tensor_out_A = cutn.create_tensor_descriptor(self.handle, 3, max_extents_AB, 0, modes_out_A, self.data_type)
|
|
||||||
desc_tensor_out_B = cutn.create_tensor_descriptor(self.handle, 3, max_extents_AB, 0, modes_out_B, self.data_type)
|
|
||||||
|
|
||||||
cutn.workspace_compute_gate_split_sizes(self.handle,
|
|
||||||
desc_tensor_in_A, desc_tensor_in_B, desc_tensor_in_G,
|
|
||||||
desc_tensor_out_A, desc_tensor_out_B,
|
|
||||||
self.gate_algo, self.svd_config, self.compute_type, self.work_desc)
|
|
||||||
|
|
||||||
workspace_size = cutn.workspace_get_memory_size(self.handle, self.work_desc, cutn.WorksizePref.MIN, cutn.Memspace.DEVICE, cutn.WorkspaceKind.SCRATCH)
|
|
||||||
|
|
||||||
# free resources
|
|
||||||
cutn.destroy_tensor_descriptor(desc_tensor_in_A)
|
|
||||||
cutn.destroy_tensor_descriptor(desc_tensor_in_B)
|
|
||||||
cutn.destroy_tensor_descriptor(desc_tensor_in_G)
|
|
||||||
cutn.destroy_tensor_descriptor(desc_tensor_out_A)
|
|
||||||
cutn.destroy_tensor_descriptor(desc_tensor_out_B)
|
|
||||||
return workspace_size
|
|
||||||
|
|
||||||
def set_workspace(self, work, workspace_size):
|
|
||||||
"""Compute the maximal workspace needed for MPS gating algorithm.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
work: Pointer to the allocated workspace.
|
|
||||||
workspace_size: The required workspace size on the device.
|
|
||||||
"""
|
|
||||||
cutn.workspace_set_memory(self.handle, self.work_desc, cutn.Memspace.DEVICE, cutn.WorkspaceKind.SCRATCH, work.ptr, workspace_size)
|
|
||||||
|
|
||||||
def apply_gate(self, site_A, site_B, gate, verbose, stream):
|
|
||||||
"""Inplace execution of the apply gate algoritm on site A and site B.
|
|
||||||
|
|
||||||
Args:
|
|
||||||
site_A: The first site on which the gate is applied to.
|
|
||||||
site_B: The second site on which the gate is applied to.
|
|
||||||
gate (cupy.ndarray): The input data for the gate tensor.
|
|
||||||
verbose: Whether to print out the runtime information during truncation.
|
|
||||||
stream (cupy.cuda.Stream): The CUDA stream on which the computation is performed.
|
|
||||||
"""
|
|
||||||
if site_B - site_A != 1:
|
|
||||||
raise ValueError("Site B must be the right site of site A")
|
|
||||||
if site_B >= self.num_sites:
|
|
||||||
raise ValueError("Site index cannot exceed maximum number of sites")
|
|
||||||
|
|
||||||
desc_tensor_in_A = self.desc_tensors[site_A]
|
|
||||||
desc_tensor_in_B = self.desc_tensors[site_B]
|
|
||||||
|
|
||||||
phys_mode_in_A = self.phys_modes[site_A]
|
|
||||||
phys_mode_in_B = self.phys_modes[site_B]
|
|
||||||
phys_mode_out_A = next(self.new_mode)
|
|
||||||
phys_mode_out_B = next(self.new_mode)
|
|
||||||
modes_G = (phys_mode_in_A, phys_mode_in_B, phys_mode_out_A, phys_mode_out_B)
|
|
||||||
extent_G = (self.phys_extent, self.phys_extent, self.phys_extent, self.phys_extent)
|
|
||||||
desc_tensor_in_G = cutn.create_tensor_descriptor(self.handle, 4, extent_G, 0, modes_G, self.data_type)
|
|
||||||
|
|
||||||
# construct and initialize the expected output A and B
|
|
||||||
tensor_in_A = self.state_tensors[site_A]
|
|
||||||
tensor_in_B = self.state_tensors[site_B]
|
|
||||||
left_extent_A = tensor_in_A.shape[0]
|
|
||||||
extent_AB_in = tensor_in_A.shape[2]
|
|
||||||
right_extent_B = tensor_in_B.shape[2]
|
|
||||||
combined_extent_left = min(left_extent_A, extent_AB_in * self.phys_extent) * self.phys_extent
|
|
||||||
combined_extent_right = min(right_extent_B, extent_AB_in * self.phys_extent) * self.phys_extent
|
|
||||||
extent_Aout_B = min(combined_extent_left, combined_extent_right, self.max_virtual_extent)
|
|
||||||
|
|
||||||
extent_out_A = (left_extent_A, self.phys_extent, extent_Aout_B)
|
|
||||||
extent_out_B = (extent_Aout_B, self.phys_extent, right_extent_B)
|
|
||||||
|
|
||||||
tensor_out_A = cp.zeros(extent_out_A, dtype=tensor_in_A.dtype, order="F")
|
|
||||||
tensor_out_B = cp.zeros(extent_out_B, dtype=tensor_in_B.dtype, order="F")
|
|
||||||
|
|
||||||
# create tensor descriptors for output A and B
|
|
||||||
modes_out_A = (self.virtual_modes[site_A], phys_mode_out_A, self.virtual_modes[site_A+1])
|
|
||||||
modes_out_B = (self.virtual_modes[site_B], phys_mode_out_B, self.virtual_modes[site_B+1])
|
|
||||||
|
|
||||||
desc_tensor_out_A = cutn.create_tensor_descriptor(self.handle, 3, extent_out_A, 0, modes_out_A, self.data_type)
|
|
||||||
desc_tensor_out_B = cutn.create_tensor_descriptor(self.handle, 3, extent_out_B, 0, modes_out_B, self.data_type)
|
|
||||||
|
|
||||||
cutn.gate_split(self.handle,
|
|
||||||
desc_tensor_in_A, tensor_in_A.data.ptr,
|
|
||||||
desc_tensor_in_B, tensor_in_B.data.ptr,
|
|
||||||
desc_tensor_in_G, gate.data.ptr,
|
|
||||||
desc_tensor_out_A, tensor_out_A.data.ptr,
|
|
||||||
0, # we factorize singular values equally onto output A and B.
|
|
||||||
desc_tensor_out_B, tensor_out_B.data.ptr,
|
|
||||||
self.gate_algo, self.svd_config, self.compute_type,
|
|
||||||
self.svd_info, self.work_desc, stream.ptr)
|
|
||||||
|
|
||||||
if verbose:
|
|
||||||
full_extent = np.array([0], dtype=cutn.tensor_svd_info_get_attribute_dtype(cutn.TensorSVDInfoAttribute.FULL_EXTENT))
|
|
||||||
reduced_extent = np.array([0], dtype=cutn.tensor_svd_info_get_attribute_dtype(cutn.TensorSVDInfoAttribute.REDUCED_EXTENT))
|
|
||||||
discarded_weight = np.array([0], dtype=cutn.tensor_svd_info_get_attribute_dtype(cutn.TensorSVDInfoAttribute.DISCARDED_WEIGHT))
|
|
||||||
|
|
||||||
cutn.tensor_svd_info_get_attribute(
|
|
||||||
self.handle, self.svd_info, cutn.TensorSVDInfoAttribute.FULL_EXTENT,
|
|
||||||
full_extent.ctypes.data, full_extent.dtype.itemsize)
|
|
||||||
cutn.tensor_svd_info_get_attribute(
|
|
||||||
self.handle, self.svd_info, cutn.TensorSVDInfoAttribute.REDUCED_EXTENT,
|
|
||||||
reduced_extent.ctypes.data, reduced_extent.dtype.itemsize)
|
|
||||||
cutn.tensor_svd_info_get_attribute(
|
|
||||||
self.handle, self.svd_info, cutn.TensorSVDInfoAttribute.DISCARDED_WEIGHT,
|
|
||||||
discarded_weight.ctypes.data, discarded_weight.dtype.itemsize)
|
|
||||||
|
|
||||||
print("Virtual bond truncated from {0} to {1} with a discarded weight of {2:.6f}".format(full_extent[0], reduced_extent[0], discarded_weight[0]))
|
|
||||||
|
|
||||||
self.phys_modes[site_A] = phys_mode_out_A
|
|
||||||
self.phys_modes[site_B] = phys_mode_out_B
|
|
||||||
self.desc_tensors[site_A] = desc_tensor_out_A
|
|
||||||
self.desc_tensors[site_B] = desc_tensor_out_B
|
|
||||||
|
|
||||||
extent_out_A = np.zeros((3,), dtype=np.int64)
|
|
||||||
extent_out_B = np.zeros((3,), dtype=np.int64)
|
|
||||||
extent_out_A, strides_out_A = cutn.get_tensor_details(self.handle, desc_tensor_out_A)[2:]
|
|
||||||
extent_out_B, strides_out_B = cutn.get_tensor_details(self.handle, desc_tensor_out_B)[2:]
|
|
||||||
|
|
||||||
# Recall that `cutensornet.gate_split` can potentially find reduced extent during SVD truncation when value-based truncation is used.
|
|
||||||
# Therefore we here update the container for output tensor A and B.
|
|
||||||
if extent_out_A[2] != extent_Aout_B:
|
|
||||||
# note strides in cutensornet are in the unit of count and strides in cupy/numpy are in the unit of nbytes
|
|
||||||
strides_out_A = [i * tensor_out_A.itemsize for i in strides_out_A]
|
|
||||||
strides_out_B = [i * tensor_out_B.itemsize for i in strides_out_B]
|
|
||||||
tensor_out_A = cp.ndarray(extent_out_A, dtype=tensor_out_A.dtype, memptr=tensor_out_A.data, strides=strides_out_A)
|
|
||||||
tensor_out_B = cp.ndarray(extent_out_B, dtype=tensor_out_B.dtype, memptr=tensor_out_B.data, strides=strides_out_B)
|
|
||||||
|
|
||||||
self.state_tensors[site_A] = tensor_out_A
|
|
||||||
self.state_tensors[site_B] = tensor_out_B
|
|
||||||
|
|
||||||
cutn.destroy_tensor_descriptor(desc_tensor_in_A)
|
|
||||||
cutn.destroy_tensor_descriptor(desc_tensor_in_B)
|
|
||||||
cutn.destroy_tensor_descriptor(desc_tensor_in_G)
|
|
||||||
|
|
||||||
def __del__(self):
|
|
||||||
"""Free all resources owned by the object."""
|
|
||||||
for desc_tensor in self.desc_tensors:
|
|
||||||
cutn.destroy_tensor_descriptor(desc_tensor)
|
|
||||||
cutn.destroy(self.handle)
|
|
||||||
cutn.destroy_workspace_descriptor(self.work_desc)
|
|
||||||
cutn.destroy_tensor_svd_config(self.svd_config)
|
|
||||||
cutn.destroy_tensor_svd_info(self.svd_info)
|
|
||||||
|
|
||||||
Reference in New Issue
Block a user