diff --git a/CMakeLists.txt b/CMakeLists.txt index 96977aea..2e699d2f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -116,10 +116,22 @@ if (NOT mdspan_FOUND) endif() endif() +# Unfortunately there doesn't seem to be a clean way to detect CBLAS support without trying to find a CBLAS header. find_package(BLAS) +find_path(CBLAS_INCLUDE_DIR cblas.h) +if(CBLAS_INCLUDE_DIR) + message(STATUS "Found CBLAS header at ${CBLAS_INCLUDE_DIR}, enabling BLAS support") + set(BLAS_FOUND TRUE) +else() + message(STATUS "Could not find CBLAS header, disabling BLAS support") + set(BLAS_FOUND FALSE) +endif() option(LINALG_ENABLE_BLAS "Assume that we are linking with a BLAS library." ${BLAS_FOUND}) +if(LINALG_ENABLE_BLAS AND NOT BLAS_FOUND) + message(FATAL_ERROR "Could not find CBLAS header, but LINALG_ENABLE_BLAS is ON. Please install a BLAS library and its CBLAS headers, or set LINALG_ENABLE_BLAS to OFF.") +endif() find_package(TBB) option(LINALG_ENABLE_TBB diff --git a/include/experimental/__p1673_bits/blas1_scale.hpp b/include/experimental/__p1673_bits/blas1_scale.hpp index dca5c2e2..0a918c73 100644 --- a/include/experimental/__p1673_bits/blas1_scale.hpp +++ b/include/experimental/__p1673_bits/blas1_scale.hpp @@ -18,6 +18,9 @@ #ifndef LINALG_INCLUDE_EXPERIMENTAL___P1673_BITS_BLAS1_SCALE_HPP_ #define LINALG_INCLUDE_EXPERIMENTAL___P1673_BITS_BLAS1_SCALE_HPP_ +#include "blas_helpers.hpp" +#include + namespace MDSPAN_IMPL_STANDARD_NAMESPACE { namespace MDSPAN_IMPL_PROPOSED_NAMESPACE { inline namespace __p1673_version_0 { @@ -31,7 +34,7 @@ template -void linalg_scale_rank_1( +void generic_scale_rank_1( const Scalar alpha, mdspan, Layout, Accessor> x) { @@ -40,6 +43,98 @@ void linalg_scale_rank_1( } } +template +constexpr bool maybe_can_blas_scale() { + // It's OK if Scalar and value_type aren't the same, + // as long as we can convert Scalar to value_type. + using value_type = typename MdspanType::value_type; + constexpr bool blas_value_type = + std::is_convertible_v && + impl::is_blas_value_type_v; + + constexpr bool blas_layout = + impl::is_blas_layout_type_v; + + constexpr bool blas_accessor = + impl::is_blas_accessor_type_v; + + return blas_value_type && blas_layout && blas_accessor; +} + +// Return true if a BLAS routine could be called to scale the vector, false otherwise. +template +bool try_blas_scale( + const Scalar alpha, + mdspan, Layout, Accessor> x) +{ + #ifdef KOKKOS_ENABLE_CBLAS + auto n = x.extent(0); + // We can't call x.stride(0) until we know that x.is_strided() is true. + if (x.is_strided() && n <= std::numeric_limits::max()) { + auto incx = x.stride(0); + if (incx > std::numeric_limits::max()) { + return false; + } + // Strided layout mappings can have stride zero; the BLAS cannot. + if (incx == 0) { + return false; + } + if constexpr (std::is_same_v) { + cblas_sscal(n, alpha, x.data_handle(), incx); + return true; + } else if constexpr (std::is_same_v) { + cblas_dscal(n, alpha, x.data_handle(), incx); + return true; + } else if constexpr (std::is_same_v>) { + if constexpr(std::is_convertible_v) { + cblas_csscal(n, alpha, x.data_handle(), incx); + return true; + } + else if constexpr (std::is_convertible_v>) { + auto converted_alpha = static_cast>(alpha); + cblas_cscal(n, &converted_alpha, x.data_handle(), incx); + return true; + } + } else if constexpr (std::is_same_v>) { + if constexpr (std::is_convertible_v) { + cblas_zdscal(n, alpha, x.data_handle(), incx); + return true; + } + else if constexpr (std::is_convertible_v>) { + auto converted_alpha = static_cast>(alpha); + cblas_zscal(n, &converted_alpha, x.data_handle(), incx); + return true; + } + } + } + #endif + return false; +} + +template +void linalg_scale_rank_1( + const Scalar alpha, + mdspan, Layout, Accessor> x) +{ + bool done = false; + if constexpr (maybe_can_blas_scale()) { + done = try_blas_scale(alpha, x); + } + if (!done) { + generic_scale_rank_1(alpha, x); + } +} + template +#include +#include +#ifdef KOKKOS_ENABLE_CBLAS +#include "cblas.h" +#endif + +namespace MDSPAN_IMPL_STANDARD_NAMESPACE { +namespace MDSPAN_IMPL_PROPOSED_NAMESPACE { +inline namespace __p1673_version_0 { +namespace linalg { +namespace impl { + +template +constexpr bool is_blas_value_type_v = + std::is_same_v || + std::is_same_v || + std::is_same_v> || + std::is_same_v>; + +// The padded layouts are class templates taking a size_t, so detecting them +// takes a partial specialization rather than is_same_v. +template +constexpr bool is_padded_layout_v = false; + +template +constexpr bool is_padded_layout_v> = true; + +template +constexpr bool is_padded_layout_v> = true; + +template +constexpr bool is_blas_layout_type_v = + // Assume that we have a C BLAS, which accepts + // both row-major and column-major layouts. + // + // This just means that the layouts COULD be valid. + // For layout_stride, we need to check the strides first. + std::is_same_v || + std::is_same_v || + is_padded_layout_v || + std::is_same_v; + +// The BLAS accepts accessors that deal with pointers to memory. +// default_accessor is a class template, so we can't just use is_same_v directly. +// +// scale doesn't accept conjugated_accessor or scaled_accessor +// because those are read-only accessors, and scale needs to +// write to the mdspan's elements. + +template +constexpr bool is_default_accessor_v = false; + +template +constexpr bool is_default_accessor_v> = true; + +template +constexpr bool is_blas_accessor_type_v = + is_default_accessor_v; + +#ifdef KOKKOS_ENABLE_CBLAS +// Deduce the BLAS integer index type from the first parameter +// (the length N) of cblas_dscal, as declared by the cblas.h in scope: +// int on an LP64 build, a 64-bit integer on an ILP64 build. +template A first_param(R (*)(A, Rest...)); +using cblas_index = decltype(first_param(cblas_dscal)); +#endif + +// We made the above queries traits, with their typical `_v` prefix. +// We make maybe_can_blas_scale() a function. +// That's a matter of taste; it could be a trait too. + +} // end namespace impl +} // end namespace linalg +} // end inline namespace __p1673_version_0 +} // end namespace MDSPAN_IMPL_PROPOSED_NAMESPACE +} // end namespace MDSPAN_IMPL_STANDARD_NAMESPACE + +#endif //LINALG_INCLUDE_EXPERIMENTAL___P1673_BITS_BLAS_HELPERS_HPP_ diff --git a/tests/native/scale.cpp b/tests/native/scale.cpp index 708e76dc..976171e4 100644 --- a/tests/native/scale.cpp +++ b/tests/native/scale.cpp @@ -77,6 +77,89 @@ namespace { } } } + + // Stride other than 1. Both halves matter: the referenced elements must be + // scaled and the ones between them must not, which is what a wrong incx gets + // wrong. + TEST(BLAS1_scale, mdspan_layout_stride) + { + using scalar_t = double; + using extents_t = extents; + using vector_t = mdspan; + + constexpr std::size_t vectorSize(5); + constexpr std::size_t stride(2); + std::vector storage(vectorSize * stride); + + layout_stride::mapping mapping{ + extents_t{vectorSize}, std::array{stride}}; + vector_t x(storage.data(), mapping); + + for (std::size_t k = 0; k < storage.size(); ++k) { + storage[k] = scalar_t (k) + 1.0; + } + const scalar_t scaleFactor = 5.0; + scale(scaleFactor, x); + for (std::size_t k = 0; k < storage.size(); ++k) { + const scalar_t storage_k = scalar_t (k) + 1.0; + const scalar_t expected = + (k % stride == 0) ? scaleFactor * storage_k : storage_k; + EXPECT_EQ( storage[k], expected ); + } + } + + TEST(BLAS1_scale, mdspan_layout_left_padded) + { + using scalar_t = double; + using extents_t = extents; + using layout_t = layout_left_padded<4>; + using vector_t = mdspan; + + constexpr std::size_t vectorSize(4); + constexpr std::size_t storageSize = vectorSize; + std::vector storage(storageSize); + + layout_t::mapping mapping{extents_t{}}; + vector_t x(storage.data(), mapping); + + for (std::size_t k = 0; k < vectorSize; ++k) { + const scalar_t x_k = scalar_t (k) + 1.0; + x(k) = x_k; + } + const scalar_t scaleFactor = 5.0; + scale(scaleFactor, x); + for (std::size_t k = 0; k < vectorSize; ++k) { + const scalar_t x_k = scalar_t (k) + 1.0; + EXPECT_EQ( x(k), scaleFactor * x_k ); + } + } + + TEST(BLAS1_scale, mdspan_layout_right_padded) + { + using scalar_t = double; + using extents_t = extents; + using layout_t = layout_right_padded; + using vector_t = mdspan; + + constexpr std::size_t vectorSize(4); + constexpr std::size_t paddingValue(4); + constexpr std::size_t storageSize = vectorSize; + std::vector storage(storageSize); + + layout_t::mapping mapping{extents_t{}, paddingValue}; + vector_t x(storage.data(), mapping); + + for (std::size_t k = 0; k < vectorSize; ++k) { + const scalar_t x_k = scalar_t (k) + 1.0; + x(k) = x_k; + } + const scalar_t scaleFactor = 5.0; + scale(scaleFactor, x); + for (std::size_t k = 0; k < vectorSize; ++k) { + const scalar_t x_k = scalar_t (k) + 1.0; + EXPECT_EQ( x(k), scaleFactor * x_k ); + } + } } // int main() {