blueqatSDKからcuStateVecを実行

Yuichiro Minato

2022/11/07 10:55

NVIDIAから量子コンピュータシミュレータのcuStateVecがでました。マルチGPUでその威力を発揮するのですが、blueqatでは今のところマルチGPU対応は考えておらず、googleのcirqやqsimを入れて使うのがおすすめですが、単体GPUでサクッと使いたいというご要望の場合にはcirqから使うのもあれだし、cuStateVec直接使うのも難しいという感じかと思います。

ちょっとcuStateVecをあまり深く追求できてないので、パフォーマンスは落ちるかもしれませんが、とりあえず使えるという意味でblueqatにつないでみます。

１，初期状態で状態ベクトルを準備
２，あとは各ゲート段階でゲートを定義してかける

のみです。簡単にできます。


import numpy as np
from blueqat.backends.backendbase import Backend
from blueqat import BlueqatGlobalSetting
from blueqat import Circuit
import cupy as cp
import cuquantum
from cuquantum import custatevec as cusv

class CuSv(Backend):

    
    def _preprocess_run(self, gates, n_qubits, args, kwargs):
        
        h_sv = np.asarray(np.full(2**n_qubits,0.0+0.0j), dtype=np.complex64)
        h_sv[0] = 1.0+0.0j
        d_sv = cp.asarray(h_sv)
        return gates, (d_sv, n_qubits)

    def _postprocess_run(self, ctx):
        print(ctx[0])
        return 

    def _one_qubit_gate_noargs(self, gate, ctx):
        for idx in gate.target_iter(ctx[1]):
            
            #欲しいゲートはこちらに記述
            if gate.lowername == "x":
                matrix = cp.asarray([0.0+0.0j, 1.0+0.0j, 1.0+0.0j, 0.0+0.0j], dtype=np.complex64)                
            else:
                matrix = cp.asarray([1.0+0.0j, 0.0+0.0j, 0.0+0.0j, 1.0+0.0j], dtype=np.complex64)                

            nIndexBits = ctx[1]
            nSvSize    = (1 << nIndexBits)
            nTargets   = 1
            nControls  = 0
            adjoint    = 0

            targets    = np.asarray([idx], dtype=np.int32)
            controls   = np.asarray([], dtype=np.int32)

            if isinstance(matrix, cp.ndarray):
                matrix_ptr = matrix.data.ptr
            elif isinstance(matrix, np.ndarray):
                matrix_ptr = matrix.ctypes.data
            else:
                raise ValueError

            # cuStateVec handle initialization
            handle = cusv.create()
            workspaceSize = cusv.apply_matrix_get_workspace_size(
                handle, cuquantum.cudaDataType.CUDA_C_32F, nIndexBits, matrix_ptr, cuquantum.cudaDataType.CUDA_C_32F,
                cusv.MatrixLayout.ROW, adjoint, nTargets, nControls, cuquantum.ComputeType.COMPUTE_32F)

            # check the size of external workspace
            if workspaceSize > 0:
                workspace = cp.cuda.memory.alloc(workspaceSize)
                workspace_ptr = workspace.ptr
            else:
                workspace_ptr = 0

            # apply gate
            cusv.apply_matrix(
                handle, ctx[0].data.ptr, cuquantum.cudaDataType.CUDA_C_32F, nIndexBits, matrix_ptr, cuquantum.cudaDataType.CUDA_C_32F,
                cusv.MatrixLayout.ROW, adjoint, targets.ctypes.data, nTargets, controls.ctypes.data, 0, nControls,
                cuquantum.ComputeType.COMPUTE_32F, workspace_ptr, workspaceSize)

            # destroy handle
            cusv.destroy(handle)

        return ctx

    gate_x = _one_qubit_gate_noargs
    gate_y = _one_qubit_gate_noargs
    gate_z = _one_qubit_gate_noargs
    gate_h = _one_qubit_gate_noargs
    gate_t = _one_qubit_gate_noargs
    gate_s = _one_qubit_gate_noargs
    
    def _one_qubit_gate_args_theta(self, gate, ctx):
        return ctx

    gate_rx = _one_qubit_gate_args_theta
    gate_ry = _one_qubit_gate_args_theta
    gate_rz = _one_qubit_gate_args_theta
    gate_phase = _one_qubit_gate_args_theta
    
    def _two_qubit_gate_noargs(self, gate, ctx):
        for control, target in gate.control_target_iter(ctx[1]):
            
            #欲しいゲートはこちらに記述
            if gate.lowername == "cx":
                matrix = cp.asarray([0.0+0.0j, 1.0+0.0j, 1.0+0.0j, 0.0+0.0j], dtype=np.complex64)                
            else:
                matrix = cp.asarray([1.0+0.0j, 0.0+0.0j, 0.0+0.0j, 1.0+0.0j], dtype=np.complex64)                

            nIndexBits = ctx[1]
            nSvSize    = (1 << nIndexBits)
            nTargets   = 1
            nControls  = 1
            adjoint    = 0

            targets    = np.asarray([target], dtype=np.int32)
            controls   = np.asarray([control], dtype=np.int32)

            if isinstance(matrix, cp.ndarray):
                matrix_ptr = matrix.data.ptr
            elif isinstance(matrix, np.ndarray):
                matrix_ptr = matrix.ctypes.data
            else:
                raise ValueError

            # cuStateVec handle initialization
            handle = cusv.create()
            workspaceSize = cusv.apply_matrix_get_workspace_size(
                handle, cuquantum.cudaDataType.CUDA_C_32F, nIndexBits, matrix_ptr, cuquantum.cudaDataType.CUDA_C_32F,
                cusv.MatrixLayout.ROW, adjoint, nTargets, nControls, cuquantum.ComputeType.COMPUTE_32F)

            # check the size of external workspace
            if workspaceSize > 0:
                workspace = cp.cuda.memory.alloc(workspaceSize)
                workspace_ptr = workspace.ptr
            else:
                workspace_ptr = 0

            # apply gate
            cusv.apply_matrix(
                handle, ctx[0].data.ptr, cuquantum.cudaDataType.CUDA_C_32F, nIndexBits, matrix_ptr, cuquantum.cudaDataType.CUDA_C_32F,
                cusv.MatrixLayout.ROW, adjoint, targets.ctypes.data, nTargets, controls.ctypes.data, 0, nControls,
                cuquantum.ComputeType.COMPUTE_32F, workspace_ptr, workspaceSize)

            # destroy handle
            cusv.destroy(handle)
        return ctx
    
    gate_cx = gate_cy = gate_cz = _two_qubit_gate_noargs

    def _three_qubit_gate_noargs(self, gate, ctx):
        return ctx

    gate_ccx = _three_qubit_gate_noargs
    gate_cswap = _three_qubit_gate_noargs

    def gate_measure(self, gate, ctx):
        return ctx

    gate_reset = _one_qubit_gate_noargs
    
BlueqatGlobalSetting.register_backend('cusv', CuSv)


from blueqat import Circuit

Circuit(3).x[0,1,2].cx[0,1].run(backend="cusv")

[0.+0.j 0.+0.j 0.+0.j 0.+0.j 0.+0.j 1.+0.j 0.+0.j 0.+0.j]

動きました。サンプリング機能をつけたり、期待値計算機能をつけたりすれば普通にNISQ向けに使えます。これまでのblueqatのアプリケーション資産もそのまま使えますので、ぜひ試してみてください。要望あれば作りこみます。