queue-dispatch.hh

Go to the documentation of this file.

#pragma once
 
#include <complex>
#include <stdexcept>
#include <string>
#include <type_traits>
#include <utility>
 
#include <batchlas/backend_config.h>
#include <batchlas/blas/enums.hh>
#include <batchlas/settings.hh>
#include <batchlas/blas/matrix.hh>
#include <batchlas/util/sycl-device-queue.hh>
#include <batchlas/util/sycl-span.hh>
#include <batchlas/util/sycl-vector.hh>
 
/// @file
/// @brief Backend deduction from the Queue and the USM pointer check.
/// @ingroup api_dispatch
 
namespace batchlas {
 
/**
 * @addtogroup api_dispatch
 * @details
 * <b>Every entry point has two spellings.</b> `f<Backend::CUDA>(ctx, ...)` fixes the
 * backend at compile time; `f(ctx, ...)` takes it from `ctx.backend()` through
 * with_backend(). The deducing overload of a positional entry point is generated by
 * `BATCHLAS_DISPATCH_ON_QUEUE(f)`, and the option-struct overloads in
 * `blas/options.hh` deduce the same way. Library code templated on `Backend` must
 * call the `f<B>` spelling, never the deducing one.
 *
 * <b>Every entry point also accepts owning arguments.</b> A `Matrix` is a `MatrixView`
 * and a `Vector` is a `VectorView` (public inheritance), so an owning argument binds to
 * a view parameter directly and deduces `T` and `MatrixFormat` through derived-to-base
 * deduction, in any position and with any template arguments named.
 *
 * <b>Pointer checks.</b> The deducing spellings check that every non-empty pointer
 * argument is reachable from the queue's device and throw `std::invalid_argument`
 * naming the argument otherwise; the `f<Backend>` spellings do not check.
 * `BATCHLAS_SKIP_POINTER_CHECKS` turns the check off (latched once per process).
 *
 * The generating macros are omitted from this reference; see @ref design_api_conventions
 * for their design and traps.
 * @{
 */
 
namespace detail {
 
// Element type of a container argument (Matrix, MatrixView, Vector, VectorView, Span,
// UnifiedVector); void for anything else. Scalars are deliberately not containers: a
// float op may take a double knob (fill_graded_hermitian's condition number).
template <class A> struct container_elem { using type = void; };
// A one-parameter class template over the scalar (StedcParams<T>, GemmOptions<T>) names the
// op's type as well: stedc_buffer_size(ctx, n, batch, job, StedcParams<double>) has no container.
template <template <class> class P, class T> struct container_elem<P<T>> { using type = T; };
template <class T> struct container_elem<std::complex<T>> { using type = void; };
template <class T, MatrixFormat F> struct container_elem<Matrix<T, F>> { using type = T; };
template <class T, MatrixFormat F> struct container_elem<MatrixView<T, F>> { using type = T; };
template <class T> struct container_elem<Vector<T>> { using type = T; };
template <class T> struct container_elem<VectorView<T>> { using type = T; };
template <class T> struct container_elem<Span<T>> { using type = T; };
template <class T> struct container_elem<UnifiedVector<T>> { using type = T; };
 
template <class T>
inline constexpr bool is_fp64_scalar_v =
    std::is_same_v<std::remove_cv_t<T>, double> || std::is_same_v<std::remove_cv_t<T>, std::complex<double>>;
 
// True when any container argument holds double or complex<double>: Backend::SYCL
// instantiates those types only when BATCHLAS_HAS_DEVICE_FP64.
template <class... A>
inline constexpr bool needs_fp64_v =
    (is_fp64_scalar_v<typename container_elem<std::remove_cvref_t<A>>::type> || ...);
 
[[noreturn]] inline void throw_no_device_fp64(const Queue& ctx, const char* op) {
    const std::string what = std::string("BatchLAS: ") + (op ? op : "this call") +
                             " was called with double or complex<double> data on a queue whose device (" +
                             ctx.device().get_name() + ")";
    // A CPU queue pinned to Backend::SYCL has fp64; what it lacks is the build's SYCL double kernels.
    if (ctx.device().supports_fp64()) {
        throw batchlas::unsupported(
            what + " runs Backend::SYCL, and this build compiled no Backend::SYCL kernels for those "
            "types. Use Backend::NETLIB (the default backend of a CPU Queue), or convert the data "
            "to float or complex<float>.");
    }
    throw batchlas::unsupported(
        what + " has no fp64 support, so this build compiled no Backend::SYCL kernels for those "
        "types. Run double work on a CPU Queue (Queue(Device(\"cpu\"))), whose backend is "
        "NETLIB, or convert the data to float or complex<float>.");
}
 
/// Throws throw_no_device_fp64 for a double or complex<double> @p T on a GPU without fp64. For
/// kernels that are not Backend-templated (or NETLIB ones that launch on the queue's device),
/// whose double instantiations exist and would otherwise fail in the device JIT.
template <class T>
inline void require_device_fp64(const Queue& ctx, const char* op) {
    if constexpr (is_fp64_scalar_v<T> && !BATCHLAS_HAS_DEVICE_FP64) {
        if (ctx.device().type == DeviceType::GPU && !ctx.device().supports_fp64())
            throw_no_device_fp64(ctx, op);
    }
}
 
// with_backend with the fp64 gate: the SYCL branch is discarded (never instantiated) when
// NeedsFp64 and the build has no device fp64, so no <SYCL, double> symbol is referenced.
template <bool NeedsFp64, typename F>
inline auto with_backend_gated(Queue& ctx, const char* op, F&& f) {
    static_assert(BATCHLAS_HAS_CUDA_BACKEND || BATCHLAS_HAS_ROCM_BACKEND ||
                      BATCHLAS_HAS_SYCL_BACKEND || BATCHLAS_HAS_HOST_BACKEND,
                  "BatchLAS was built with no backends; nothing can be dispatched.");
 
    switch (ctx.backend()) {
#if BATCHLAS_HAS_CUDA_BACKEND
        case Backend::CUDA:
            return f(std::integral_constant<Backend, Backend::CUDA>{});
#endif
#if BATCHLAS_HAS_ROCM_BACKEND
        case Backend::ROCM:
            return f(std::integral_constant<Backend, Backend::ROCM>{});
#endif
#if BATCHLAS_HAS_SYCL_BACKEND
        case Backend::SYCL:
            if constexpr (NeedsFp64 && !BATCHLAS_HAS_DEVICE_FP64) {
                throw_no_device_fp64(ctx, op);
            } else {
                return f(std::integral_constant<Backend, Backend::SYCL>{});
            }
#endif
#if BATCHLAS_HAS_HOST_BACKEND
        case Backend::NETLIB:
            return f(std::integral_constant<Backend, Backend::NETLIB>{});
#endif
        default:
            break;
    }
    throw batchlas::unsupported(
        std::string("BatchLAS: backend ") + std::string(to_string(ctx.backend())) +
        " has no implementation in this build. "
        "Check Queue::backend_available() before pinning a backend.");
}
 
}  // namespace detail
 
/// @brief Call @p f with the queue's backend as a compile-time constant.
///
/// Converts `ctx.backend()` into `std::integral_constant<Backend, B>` and invokes
/// `f` with it, so the body can name `B.value` as a template argument:
/// @code
/// with_backend(ctx, [&](auto B) { return gemm<B.value, T>(ctx, ...); });
/// @endcode
/// This is a `switch` over existing instantiations, not a virtual call.
/// @tparam F     callable taking `std::integral_constant<Backend, B>`
/// @param ctx    queue whose backend is dispatched on
/// @param f      callable; every branch must return the same type
/// @return       whatever @p f returns
/// @throws batchlas::unsupported if `ctx.backend()` has no implementation in this
///         build (always the case for Backend::MAGMA)
/// @note In a build without device fp64 (`BATCHLAS_HAS_DEVICE_FP64` 0), Backend::SYCL
///       has no double or complex<double> instantiations: @p f must not name them for
///       Backend::SYCL. The entry points' deducing overloads check this and throw
///       batchlas::unsupported instead.
/// @see Queue::backend_available
template <typename F>
inline auto with_backend(Queue& ctx, F&& f) {
    return detail::with_backend_gated<false>(ctx, nullptr, std::forward<F>(f));
}
 
/// @brief with_backend() for a call on element type @p T: throws batchlas::unsupported
/// before dispatch when @p T is double or complex<double> and the queue's backend is
/// Backend::SYCL in a build without device fp64.
template <typename T, typename F>
inline auto with_backend_for(Queue& ctx, const char* op, F&& f) {
    return detail::with_backend_gated<detail::is_fp64_scalar_v<T>>(ctx, op, std::forward<F>(f));
}
 
/// @}
 
namespace detail {
 
// USM contract, enforced: host memory on a GPU queue throws invalid_argument here
// instead of aborting in the runtime. Latched once (it is on every call's path), so
// a settings reload does not reach it. Any value of BATCHLAS_SKIP_POINTER_CHECKS not
// starting with '0' disables it ("=off" too); do not tighten that to env_truthy.
// evidence: docs/design/api-conventions.md#api-conventions-the-usm-pointer-check
inline bool pointer_checks_enabled() {
    static const bool enabled = !batchlas::settings().unsafe.skip_pointer_checks;
    return enabled;
}
 
// Matrix/MatrixView spell it data_ptr(); Span/UnifiedVector spell it data().
template <typename A>
concept HasDataPtr = requires(const A& a) { { a.data_ptr() } -> std::convertible_to<const void*>; };
template <typename A>
concept HasData = requires(const A& a) {
    { a.data() } -> std::convertible_to<const void*>;
    { a.size() } -> std::convertible_to<size_t>;
};
 
// Extents: MatrixView rows()/cols()/batch_size(); VectorView size()/batch_size().
template <typename A>
concept HasMatrixExtent = requires(const A& a) {
    { a.rows() } -> std::convertible_to<long long>;
    { a.cols() } -> std::convertible_to<long long>;
    { a.batch_size() } -> std::convertible_to<long long>;
};
template <typename A>
concept HasVectorExtent = requires(const A& a) {
    { a.size() } -> std::convertible_to<long long>;
    { a.batch_size() } -> std::convertible_to<long long>;
};
 
// True when no kernel can dereference the argument, whatever its pointer.
template <typename A>
inline bool addresses_no_elements(const A& arg) {
    using U = std::remove_cvref_t<A>;
    if constexpr (HasMatrixExtent<U>) {
        return arg.rows() == 0 || arg.cols() == 0 || arg.batch_size() == 0;
    } else if constexpr (HasVectorExtent<U>) {
        return arg.size() == 0 || arg.batch_size() == 0;
    } else {
        return false;
    }
}
 
template <typename A>
inline void require_arg_accessible(const Queue& ctx, const A& arg, const std::string& what) {
    if constexpr (HasDataPtr<std::remove_cvref_t<A>>) {
        // A default-constructed view means "optional matrix not in use"; it must pass.
        if (addresses_no_elements(arg)) return;
        ctx.require_device_accessible(static_cast<const void*>(arg.data_ptr()), what.c_str());
    } else if constexpr (HasData<std::remove_cvref_t<A>>) {
        // BumpAllocator sizing passes hand out empty spans by design.
        if (arg.size() != 0) {
            ctx.require_device_accessible(static_cast<const void*>(arg.data()), what.c_str());
        }
    }
    // Anything else (option structs, scalars, enums) carries no pointer.
}
 
// Positional labelling: the dispatch macro forwards an unnamed pack.
template <typename... Args>
inline void require_pack_accessible(const Queue& ctx, const char* fn, const Args&... args) {
    if (!pointer_checks_enabled()) return;
    int pos = 0;
    (void)std::initializer_list<int>{
        (++pos,
         require_arg_accessible(ctx, args,
                                std::string(fn) + ": argument " + std::to_string(pos)),
         0)...};
}
 
// Any compiled-in backend; used only to probe "is the positional call well-formed?".
inline constexpr Backend kProbeBackend =
#if BATCHLAS_HAS_CUDA_BACKEND
    Backend::CUDA;
#elif BATCHLAS_HAS_ROCM_BACKEND
    Backend::ROCM;
#else
    Backend::NETLIB;
#endif
 
}  // namespace detail
 
}  // namespace batchlas
 
// Backend-deducing overload `NAME(ctx, args...)` of `NAME<Backend, T, ...>(ctx, ...)`.
// The pack is forwarded, not restated, so the primary's defaults still apply. The
// requires-clause must stay: unconstrained, this beats option-struct overloads and
// fails inside its body. Even constrained, it beats `const MA&` for prvalues, so
// checks must also live on the <Backend> option overloads.
// evidence: docs/design/api-conventions.md#api-conventions-the-variadic-dispatch-overload-and-its-requires-clause
#define BATCHLAS_DISPATCH_ON_QUEUE(NAME)                                        \
    template <typename... Args>                                                 \
        requires requires(Queue& probe_ctx, Args&&... probe_args) {             \
            NAME<::batchlas::detail::kProbeBackend>(probe_ctx,                  \
                                                    std::forward<Args>(probe_args)...); \
        }                                                                       \
    inline auto NAME(Queue& ctx, Args&&... args) {                              \
        ::batchlas::detail::require_pack_accessible(ctx, #NAME, args...);       \
        return ::batchlas::detail::with_backend_gated<                          \
            ::batchlas::detail::needs_fp64_v<Args...>>(ctx, #NAME, [&](auto Back) { \
            return NAME<Back.value>(ctx, std::forward<Args>(args)...);          \
        });                                                                     \
    }