Program Listing for File RegisterKernel.hpp

Return to documentation for file (pennylane_lightning/core/src/simulators/lightning_qubit/gates/RegisterKernel.hpp)

// Copyright 2018-2023 Xanadu Quantum Technologies Inc.

// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at

//     http://www.apache.org/licenses/LICENSE-2.0

// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include <tuple>

#include "Constant.hpp"
#include "ConstantUtil.hpp"
#include "DynamicDispatcher.hpp"
#include "Error.hpp"
#include "GateIndices.hpp"
#include "GateOperation.hpp"
#include "OpToMemberFuncPtr.hpp"

namespace {
using Pennylane::Util::lookup;
using Pennylane::Util::prepend_to_tuple;
} // namespace

namespace Pennylane::LightningQubit {
template <class PrecisionT, class ParamT, class GateImplementation,
          Pennylane::Gates::GateOperation gate_op>
constexpr auto gateOpToFunctor() {
    return [](std::complex<PrecisionT> *data, std::size_t num_qubits,
              const std::vector<std::size_t> &wires, bool inverse,
              const std::vector<PrecisionT> &params) {
        constexpr auto func_ptr =
            Gates::GateOpToMemberFuncPtr<PrecisionT, ParamT, GateImplementation,
                                         gate_op>::value;
        PL_ASSERT(params.size() ==
                  lookup(Pennylane::Gates::Constant::gate_num_params, gate_op));
        Gates::callGateOps(func_ptr, data, num_qubits, wires, inverse, params);
    };
}

template <class PrecisionT, class ParamT, class GateImplementation,
          Pennylane::Gates::ControlledGateOperation gate_op>
constexpr auto controlledGateOpToFunctor() {
    return [](std::complex<PrecisionT> *data, std::size_t num_qubits,
              const std::vector<std::size_t> &controlled_wires,
              const std::vector<bool> &controlled_values,
              const std::vector<std::size_t> &wires, bool inverse,
              const std::vector<PrecisionT> &params) {
        constexpr auto func_ptr = Gates::ControlledGateOpToMemberFuncPtr<
            PrecisionT, ParamT, GateImplementation, gate_op>::value;
        PL_ASSERT(params.size() ==
                  lookup(Pennylane::Gates::Constant::controlled_gate_num_params,
                         gate_op));
        Gates::callControlledGateOps(func_ptr, data, num_qubits,
                                     controlled_wires, controlled_values, wires,
                                     inverse, params);
    };
}

template <class PrecisionT, class ParamT, class GateImplementation,
          std::size_t gate_idx>
constexpr auto constructGateOpsFunctorTupleIter() {
    if constexpr (gate_idx == GateImplementation::implemented_gates.size()) {
        return std::tuple{};
    } else if (gate_idx < GateImplementation::implemented_gates.size()) {
        constexpr auto gate_op =
            GateImplementation::implemented_gates[gate_idx];
        return prepend_to_tuple(
            std::pair{gate_op, gateOpToFunctor<PrecisionT, ParamT,
                                               GateImplementation, gate_op>()},
            constructGateOpsFunctorTupleIter<
                PrecisionT, ParamT, GateImplementation, gate_idx + 1>());
    }
}
template <class PrecisionT, class ParamT, class GateImplementation,
          std::size_t gate_idx>
constexpr auto constructControlledGateOpsFunctorTupleIter() {
    if constexpr (gate_idx ==
                  GateImplementation::implemented_controlled_gates.size()) {
        return std::tuple{};
    } else if (gate_idx <
               GateImplementation::implemented_controlled_gates.size()) {
        constexpr auto gate_op =
            GateImplementation::implemented_controlled_gates[gate_idx];
        return prepend_to_tuple(
            std::pair{gate_op,
                      controlledGateOpToFunctor<PrecisionT, ParamT,
                                                GateImplementation, gate_op>()},
            constructControlledGateOpsFunctorTupleIter<
                PrecisionT, ParamT, GateImplementation, gate_idx + 1>());
    }
}

template <class PrecisionT, class GateImplementation, std::size_t gntr_idx>
constexpr auto constructGeneratorOpsFunctorTupleIter() {
    if constexpr (gntr_idx ==
                  GateImplementation::implemented_generators.size()) {
        return std::tuple{};
    } else if (gntr_idx < GateImplementation::implemented_generators.size()) {
        constexpr auto gntr_op =
            GateImplementation::implemented_generators[gntr_idx];
        return prepend_to_tuple(
            std::pair{gntr_op,
                      Gates::GeneratorOpToMemberFuncPtr<
                          PrecisionT, GateImplementation, gntr_op>::value},
            constructGeneratorOpsFunctorTupleIter<
                PrecisionT, GateImplementation, gntr_idx + 1>());
    }
}

template <class PrecisionT, class GateImplementation, std::size_t gntr_idx>
constexpr auto constructControlledGeneratorOpsFunctorTupleIter() {
    if constexpr (gntr_idx ==
                  GateImplementation::implemented_controlled_generators
                      .size()) {
        return std::tuple{};
    } else if (gntr_idx <
               GateImplementation::implemented_controlled_generators.size()) {
        constexpr auto gntr_op =
            GateImplementation::implemented_controlled_generators[gntr_idx];
        return prepend_to_tuple(
            std::pair{gntr_op,
                      Gates::ControlledGeneratorOpToMemberFuncPtr<
                          PrecisionT, GateImplementation, gntr_op>::value},
            constructControlledGeneratorOpsFunctorTupleIter<
                PrecisionT, GateImplementation, gntr_idx + 1>());
    }
}

template <class PrecisionT, class GateImplementation, std::size_t mat_idx>
constexpr auto constructMatrixOpsFunctorTupleIter() {
    if constexpr (mat_idx == GateImplementation::implemented_matrices.size()) {
        return std::tuple{};
    } else if (mat_idx < GateImplementation::implemented_matrices.size()) {
        constexpr auto mat_op =
            GateImplementation::implemented_matrices[mat_idx];
        return prepend_to_tuple(
            std::pair{
                mat_op,
                Gates::MatrixOpToMemberFuncPtr<PrecisionT, GateImplementation,
                                               mat_op>::value},
            constructMatrixOpsFunctorTupleIter<PrecisionT, GateImplementation,
                                               mat_idx + 1>());
    }
}

template <class PrecisionT, class GateImplementation, std::size_t mat_idx>
constexpr auto constructControlledMatrixOpsFunctorTupleIter() {
    if constexpr (mat_idx ==
                  GateImplementation::implemented_controlled_matrices.size()) {
        return std::tuple{};
    } else if (mat_idx <
               GateImplementation::implemented_controlled_matrices.size()) {
        constexpr auto mat_op =
            GateImplementation::implemented_controlled_matrices[mat_idx];
        return prepend_to_tuple(
            std::pair{mat_op,
                      Gates::ControlledMatrixOpToMemberFuncPtr<
                          PrecisionT, GateImplementation, mat_op>::value},
            constructControlledMatrixOpsFunctorTupleIter<
                PrecisionT, GateImplementation, mat_idx + 1>());
    }
}

template <class PrecisionT, class ParamT, class GateImplementation>
constexpr auto gate_op_functor_tuple =
    constructGateOpsFunctorTupleIter<PrecisionT, ParamT, GateImplementation,
                                     0>();

template <class PrecisionT, class ParamT, class GateImplementation>
constexpr auto controlled_gate_op_functor_tuple =
    constructControlledGateOpsFunctorTupleIter<PrecisionT, ParamT,
                                               GateImplementation, 0>();

template <class PrecisionT, class GateImplementation>
constexpr auto generator_op_functor_tuple =
    constructGeneratorOpsFunctorTupleIter<PrecisionT, GateImplementation, 0>();

template <class PrecisionT, class GateImplementation>
constexpr auto controlled_generator_op_functor_tuple =
    constructControlledGeneratorOpsFunctorTupleIter<PrecisionT,
                                                    GateImplementation, 0>();

template <class PrecisionT, class GateImplementation>
constexpr auto matrix_op_functor_tuple =
    constructMatrixOpsFunctorTupleIter<PrecisionT, GateImplementation, 0>();

template <class PrecisionT, class GateImplementation>
constexpr auto controlled_matrix_op_functor_tuple =
    constructControlledMatrixOpsFunctorTupleIter<PrecisionT, GateImplementation,
                                                 0>();

template <class PrecisionT, class ParamT, class GateImplementation>
void registerAllImplementedGateOps() {
    auto &dispatcher = DynamicDispatcher<PrecisionT>::getInstance();

    auto registerGateToDispatcher = [&dispatcher](
                                        const auto &gate_op_func_pair) {
        const auto &[gate_op, func] = gate_op_func_pair;
        dispatcher.registerGateOperation(gate_op, GateImplementation::kernel_id,
                                         func);
        return gate_op;
    };

    [[maybe_unused]] const auto registered_gate_ops = std::apply(
        [&registerGateToDispatcher](auto... elem) {
            return std::make_tuple(registerGateToDispatcher(elem)...);
        },
        gate_op_functor_tuple<PrecisionT, ParamT, GateImplementation>);
}

template <class PrecisionT, class ParamT, class GateImplementation>
void registerAllImplementedControlledGateOps() {
    auto &dispatcher = DynamicDispatcher<PrecisionT>::getInstance();

    auto registerControlledGateToDispatcher =
        [&dispatcher](const auto &gate_op_func_pair) {
            const auto &[gate_op, func] = gate_op_func_pair;
            dispatcher.registerControlledGateOperation(
                gate_op, GateImplementation::kernel_id, func);
            return gate_op;
        };

    [[maybe_unused]] const auto registered_gate_ops = std::apply(
        [&registerControlledGateToDispatcher](auto... elem) {
            return std::make_tuple(registerControlledGateToDispatcher(elem)...);
        },
        controlled_gate_op_functor_tuple<PrecisionT, ParamT,
                                         GateImplementation>);
}

template <class PrecisionT, class GateImplementation>
void registerAllImplementedGeneratorOps() {
    auto &dispatcher = DynamicDispatcher<PrecisionT>::getInstance();

    auto registerGeneratorToDispatcher =
        [&dispatcher](const auto &gntr_op_func_pair) {
            const auto &[gntr_op, func] = gntr_op_func_pair;
            dispatcher.registerGeneratorOperation(
                gntr_op, GateImplementation::kernel_id, func);
            return gntr_op;
        };

    [[maybe_unused]] const auto registered_gntr_ops = std::apply(
        [&registerGeneratorToDispatcher](auto... elem) {
            return std::make_tuple(registerGeneratorToDispatcher(elem)...);
        },
        generator_op_functor_tuple<PrecisionT, GateImplementation>);
}

template <class PrecisionT, class GateImplementation>
void registerAllImplementedControlledGeneratorOps() {
    auto &dispatcher = DynamicDispatcher<PrecisionT>::getInstance();

    auto registerControlledGeneratorToDispatcher =
        [&dispatcher](const auto &generator_op_func_pair) {
            const auto &[generator_op, func] = generator_op_func_pair;
            dispatcher.registerControlledGeneratorOperation(
                generator_op, GateImplementation::kernel_id, func);
            return generator_op;
        };

    [[maybe_unused]] const auto registered_generator_ops = std::apply(
        [&registerControlledGeneratorToDispatcher](auto... elem) {
            return std::make_tuple(
                registerControlledGeneratorToDispatcher(elem)...);
        },
        controlled_generator_op_functor_tuple<PrecisionT, GateImplementation>);
}

template <class PrecisionT, class GateImplementation>
void registerAllImplementedMatrixOps() {
    auto &dispatcher = DynamicDispatcher<PrecisionT>::getInstance();

    auto registerMatrixToDispatcher = [&dispatcher](
                                          const auto &mat_op_func_pair) {
        const auto &[mat_op, func] = mat_op_func_pair;
        dispatcher.registerMatrixOperation(mat_op,
                                           GateImplementation::kernel_id, func);
        return mat_op;
    };

    [[maybe_unused]] const auto registered_mat_ops = std::apply(
        [&registerMatrixToDispatcher](auto... elem) {
            return std::make_tuple(registerMatrixToDispatcher(elem)...);
        },
        matrix_op_functor_tuple<PrecisionT, GateImplementation>);
}

template <class PrecisionT, class GateImplementation>
void registerAllImplementedControlledMatrixOps() {
    auto &dispatcher = DynamicDispatcher<PrecisionT>::getInstance();

    auto registerControlledMatrixToDispatcher =
        [&dispatcher](const auto &mat_op_func_pair) {
            const auto &[mat_op, func] = mat_op_func_pair;
            dispatcher.registerControlledMatrixOperation(
                mat_op, GateImplementation::kernel_id, func);
            return mat_op;
        };

    [[maybe_unused]] const auto registered_mat_ops = std::apply(
        [&registerControlledMatrixToDispatcher](auto... elem) {
            return std::make_tuple(
                registerControlledMatrixToDispatcher(elem)...);
        },
        controlled_matrix_op_functor_tuple<PrecisionT, GateImplementation>);
}

template <class PrecisionT, class ParamT, class GateImplementation>
void registerKernel() {
    registerAllImplementedGateOps<PrecisionT, ParamT, GateImplementation>();
    registerAllImplementedGeneratorOps<PrecisionT, GateImplementation>();
    registerAllImplementedMatrixOps<PrecisionT, GateImplementation>();
    registerAllImplementedControlledGateOps<PrecisionT, ParamT,
                                            GateImplementation>();
    registerAllImplementedControlledGeneratorOps<PrecisionT,
                                                 GateImplementation>();
    registerAllImplementedControlledMatrixOps<PrecisionT, GateImplementation>();

    DynamicDispatcher<PrecisionT>::getInstance().registerKernelName(
        GateImplementation::kernel_id, std::string{GateImplementation::name});
}
} // namespace Pennylane::LightningQubit