path: root/src/core/runtime.c

#include "core/runtime.h"

#include "core/custom_gate.h"
#include "core/gates.h"
#include "core/kernel.h"
#include "maths/complex.h"
#include "maths/simd.h"

struct PsiRuntimeConfig psi_new_runtime_config(void)
{
    struct PsiRuntimeConfig config;
    config.parallel = false;
    config.simd = false;
    config.batched = false;
    config.structure_aware = false;
    config.parallel_threshold = PSI_PARALLEL_THRESHOLD;

    return config;
}

struct PsiRuntimeConfig psi_optimal_runtime_config(void)
{
    struct PsiRuntimeConfig config = psi_new_runtime_config();
    config.structure_aware = true;
    config.simd = true;
    config.parallel = true;

    return config;
}

struct PsiRuntimeConfig psi_runtime_to_config(enum PsiRuntime runtime)
{
    struct PsiRuntimeConfig config = psi_new_runtime_config();

    switch (runtime)
    {
        case PSI_RUNTIME_BASIC: break;
        case PSI_RUNTIME_BASIC_MT: config.parallel = true; break;
        case PSI_RUNTIME_BATCHED: config.batched = true; break;
        case PSI_RUNTIME_BATCHED_MT:
            config.batched = true;
            config.parallel = true;
            break;
        case PSI_RUNTIME_SIMD:
            config.batched = true;
            config.simd = true;
            break;
        case PSI_RUNTIME_SIMD_MT:
            config.batched = true;
            config.simd = true;
            config.parallel = true;
            break;
        case PSI_RUNTIME_STRUCTURE_AWARE:
            config.structure_aware = true;
            config.simd = true;
            break;
        case PSI_RUNTIME_STRUCTURE_AWARE_MT:
            config.structure_aware = true;
            config.simd = true;
            config.parallel = true;
            break;
    }

    return config;
}

static bool op_to_kernel(struct PsiGateOp op, struct PsiKernel* out)
{
    struct PsiMatrix matrix;
    const char* name;

    switch (op.kind)
    {
        case PSI_GATE_H:
            matrix = psi_hadamard_gate().matrix;
            name = "H";
            break;
        case PSI_GATE_X:
            matrix = psi_pauli_x_gate().matrix;
            name = "X";
            break;
        case PSI_GATE_Y:
            matrix = psi_pauli_y_gate().matrix;
            name = "Y";
            break;
        case PSI_GATE_Z:
            matrix = psi_pauli_z_gate().matrix;
            name = "Z";
            break;
        case PSI_GATE_S:
            matrix = psi_s_gate().matrix;
            name = "S";
            break;
        case PSI_GATE_T:
            matrix = psi_t_gate().matrix;
            name = "T";
            break;
        case PSI_GATE_SDG:
            matrix = psi_sdg_gate().matrix;
            name = "Sdg";
            break;
        case PSI_GATE_TDG:
            matrix = psi_tdg_gate().matrix;
            name = "Tdg";
            break;
        case PSI_GATE_SX:
            matrix = psi_sx_gate().matrix;
            name = "Sx";
            break;
        case PSI_GATE_SXDG:
            matrix = psi_sxdg_gate().matrix;
            name = "Sxdg";
            break;
        case PSI_GATE_RX:
            matrix = psi_rx_matrix(op.params[0]);
            name = "Rx";
            break;
        case PSI_GATE_RY:
            matrix = psi_ry_matrix(op.params[0]);
            name = "Ry";
            break;
        case PSI_GATE_RZ:
            matrix = psi_rz_matrix(op.params[0]);
            name = "Rz";
            break;
        case PSI_GATE_P:
            matrix = psi_p_matrix(op.params[0]);
            name = "P";
            break;
        case PSI_GATE_U1:
            matrix = psi_u1_matrix(op.params[0]);
            name = "U1";
            break;
        case PSI_GATE_U2:
            matrix = psi_u2_matrix(op.params[0], op.params[1]);
            name = "U2";
            break;
        case PSI_GATE_U3:
            matrix = psi_u3_matrix(op.params[0], op.params[1], op.params[2]);
            name = "U3";
            break;
        case PSI_GATE_CNOT:
            matrix = psi_cnot_gate().matrix;
            name = "CNOT";
            break;
        case PSI_GATE_CZ:
            matrix = psi_cz_gate().matrix;
            name = "CZ";
            break;
        case PSI_GATE_SWAP:
            matrix = psi_swap_gate().matrix;
            name = "SWAP";
            break;
        case PSI_GATE_CRX:
            matrix = psi_crx_matrix(op.params[0]);
            name = "CRx";
            break;
        case PSI_GATE_CRY:
            matrix = psi_cry_matrix(op.params[0]);
            name = "CRy";
            break;
        case PSI_GATE_CRZ:
            matrix = psi_crz_matrix(op.params[0]);
            name = "CRz";
            break;
        case PSI_GATE_CP:
            matrix = psi_cp_matrix(op.params[0]);
            name = "CP";
            break;
        case PSI_GATE_CCNOT:
            matrix = psi_toffoli_gate().matrix;
            name = "CCNOT";
            break;
        case PSI_GATE_CSWAP:
            matrix = psi_fredkin_gate().matrix;
            name = "CSWAP";
            break;
        case PSI_GATE_MEASURE: return false;
        case PSI_GATE_CUSTOM:
            matrix = psi_to_quantum_gate(*op.custom).matrix;
            name = "Custom";
            break;
    }

    size_t target_count;
    const size_t* targets = psi_gate_op_quantum_targets(&op, &target_count);
    *out = psi_new_kernel(name, matrix, targets, target_count);

    return true;
}

static struct PsiVector new_zero_state(size_t num_qubits)
{
    size_t dim = (size_t)1 << num_qubits;
    struct PsiVector state = psi_new_vector(dim, PSI_COLUMN_VECTOR);
    state.data[0] = psi_new_complex(1.0, 0.0);

    return state;
}

static void execute_kernels(struct PsiVector* state, const struct PsiKernel* kernels, size_t count,
                            size_t num_qubits, struct PsiRuntimeConfig config)
{
    bool use_parallel = config.parallel && num_qubits >= config.parallel_threshold;

    for (size_t i = 0; i < count; i++)
    {
        struct PsiKernel kernel = kernels[i];

        if (config.simd && kernel.target_count == 1)
        {
            struct PsiComplex gate[2][2] = {
                { kernel.matrix.data[0], kernel.matrix.data[1] },
                { kernel.matrix.data[2], kernel.matrix.data[3] },
            };

            if (use_parallel)
                psi_apply_single_qubit_gate_simd_parallel(state->data, gate, kernel.targets[0],
                                                          num_qubits);
            else
                psi_apply_single_qubit_gate_simd(state->data, gate, kernel.targets[0], num_qubits);
        }
        else if (use_parallel)
            psi_apply_kernel_parallel(state, kernel, num_qubits);
        else
            psi_apply_kernel(state, kernel, num_qubits);
    }
}

struct PsiVector psi_compute_runtime_config(struct PsiRuntimeConfig config, size_t num_qubits,
                                            const struct PsiGateOp* operations, size_t op_count)
{
    struct PsiVector state = new_zero_state(num_qubits);

    if (config.structure_aware)
    {
        struct PsiStructureAwareBatch batch = psi_new_structure_aware_batch(num_qubits);
        for (size_t i = 0; i < op_count; i++)
        {
            struct PsiKernel kernel;
            if (op_to_kernel(operations[i], &kernel))
                psi_add_structure_aware_kernel(&batch, kernel);
        }

        psi_optimize_structure_aware_batch(&batch);
        execute_kernels(&state, batch.kernels, batch.count, num_qubits, config);
        psi_free_structure_aware_batch(&batch);

        return state;
    }

    struct PsiKernelBatch batch = psi_new_kernel_batch(num_qubits);
    for (size_t i = 0; i < op_count; i++)
    {
        struct PsiKernel kernel;
        if (op_to_kernel(operations[i], &kernel))
            psi_add_kernel(&batch, kernel);
    }

    if (config.batched)
        psi_optimize_kernel_batch(&batch);

    execute_kernels(&state, batch.kernels, batch.count, num_qubits, config);
    psi_free_kernel_batch(&batch);

    return state;
}

struct PsiVector psi_compute_runtime(enum PsiRuntime runtime, size_t num_qubits,
                                     const struct PsiGateOp* operations, size_t op_count)
{
    return psi_compute_runtime_config(psi_runtime_to_config(runtime), num_qubits, operations,
                                      op_count);
}

const struct PsiVector* psi_compute_circuit_with_config(struct PsiQuantumCircuit* circuit,
                                                        struct PsiRuntimeConfig config)
{
    if (!circuit->is_computed)
    {
        psi_free_vector(&circuit->computed_state);
        circuit->computed_state = psi_compute_runtime_config(
                config, circuit->num_qubits, circuit->operations, circuit->operation_count);
        circuit->is_computed = true;
    }

    return &circuit->computed_state;
}

const struct PsiVector* psi_compute_circuit_with(struct PsiQuantumCircuit* circuit,
                                                 enum PsiRuntime runtime)
{
    return psi_compute_circuit_with_config(circuit, psi_runtime_to_config(runtime));
}

const struct PsiVector* psi_compute_circuit(struct PsiQuantumCircuit* circuit)
{
    return psi_compute_circuit_with(circuit, PSI_RUNTIME_BASIC);
}

double psi_circuit_probability(struct PsiQuantumCircuit* circuit, size_t state_index)
{
    const struct PsiVector* state = psi_compute_circuit(circuit);
    return psi_norm2_complex(state->data[state_index]);
}

void psi_circuit_probabilities(struct PsiQuantumCircuit* circuit, double* out)
{
    const struct PsiVector* state = psi_compute_circuit(circuit);
    size_t dim = (size_t)1 << circuit->num_qubits;
    for (size_t i = 0; i < dim; i++)
        out[i] = psi_norm2_complex(state->data[i]);
}