Wrap scal

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

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 <type_traits>
+
 namespace MDSPAN_IMPL_STANDARD_NAMESPACE {
 namespace MDSPAN_IMPL_PROPOSED_NAMESPACE {
 inline namespace __p1673_version_0 {
@@ -31,7 +34,7 @@ template<class ElementType,
          class Layout,
          class Accessor,
          class Scalar>
-void linalg_scale_rank_1(
+void generic_scale_rank_1(
   const Scalar alpha,
   mdspan<ElementType, extents<IndexType, ext0>, Layout, Accessor> x)
 {
@@ -40,6 +43,90 @@ void linalg_scale_rank_1(
   }
 }
 
+template<class Scalar, class MdspanType>
+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<Scalar, value_type> &&
+    impl::is_blas_value_type_v<value_type>;
+
+  constexpr bool blas_layout =
+    impl::is_blas_layout_type_v<typename MdspanType::layout_type>;
+
+  constexpr bool blas_accessor =
+    impl::is_blas_accessor_type_v<typename MdspanType::accessor_type>;
+
+  return blas_value_type && blas_layout && blas_accessor;
+} 
+
+// Return true if the BLAS call was successfull, false otherwise.
+template<class ElementType,
+	 class IndexType,
+         ::std::size_t ext0,
+         class Layout,
+         class Accessor,
+         class Scalar>
+bool try_blas_scale(
+  const Scalar alpha,
+  mdspan<ElementType, extents<IndexType, ext0>, Layout, Accessor> x)
+{
+  #ifdef KOKKOS_ENABLE_CBLAS
+  auto n = x.extent(0);
+  auto incx = x.stride(0);
+  if (x.is_strided() && n <= std::numeric_limits<CBLAS_INDEX>::max()) {
+    if constexpr (std::is_same_v<ElementType, float>) {
+      cblas_sscal(n, alpha, x.data_handle(), incx);
+      return true;
+    } else if constexpr (std::is_same_v<ElementType, double>) {
+      cblas_dscal(n, alpha, x.data_handle(), incx);
+      return true;
+    } else if constexpr (std::is_same_v<ElementType, std::complex<float>>) {
+      if constexpr(std::is_convertible_v<Scalar, float>) {
+        cblas_csscal(n, alpha, x.data_handle(), incx);
+        return true;
+      }
+      else if constexpr (std::is_convertible_v<Scalar, std::complex<float>>) {
+        auto converted_alpha = static_cast<std::complex<float>>(alpha);
+        cblas_cscal(n, &converted_alpha, x.data_handle(), incx);
+        return true;
+      }
+    } else if constexpr (std::is_same_v<ElementType, std::complex<double>>) {
+      if constexpr (std::is_convertible_v<Scalar, double>) {
+        cblas_zdscal(n, alpha, x.data_handle(), incx);
+        return true;
+      }
+      else if constexpr (std::is_convertible_v<Scalar, std::complex<double>>) {
+        auto converted_alpha = static_cast<std::complex<double>>(alpha);
+        cblas_zscal(n, &converted_alpha, x.data_handle(), incx);
+        return true;
+      }
+    }
+  }
+  #endif
+  return false;
+}
+
+template<class ElementType,
+	 class IndexType,
+         ::std::size_t ext0,
+         class Layout,
+         class Accessor,
+         class Scalar>
+void linalg_scale_rank_1(
+  const Scalar alpha,
+  mdspan<ElementType, extents<IndexType, ext0>, Layout, Accessor> x)
+{
+  bool done = false;
+  if constexpr (maybe_can_blas_scale<Scalar, decltype(x)>()) {
+    done = try_blas_scale(alpha, x);
+  }
+  if (!done) {
+    generic_scale_rank_1(alpha, x);
+  }
+}
+
 template<class ElementType,
 	 class IndexType,
          ::std::size_t numRows,

include/experimental/__p1673_bits/blas_helpers.hpp

@@ -0,0 +1,95 @@
+//@HEADER
+// ************************************************************************
+//
+//                        Kokkos v. 4.0
+//       Copyright (2022) National Technology & Engineering
+//               Solutions of Sandia, LLC (NTESS).
+//
+// Under the terms of Contract DE-NA0003525 with NTESS,
+// the U.S. Government retains certain rights in this software.
+//
+// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
+// See https://kokkos.org/LICENSE for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// ************************************************************************
+//@HEADER
+
+#ifndef LINALG_INCLUDE_EXPERIMENTAL___P1673_BITS_BLAS_HELPERS_HPP_
+#define LINALG_INCLUDE_EXPERIMENTAL___P1673_BITS_BLAS_HELPERS_HPP_
+
+#include <complex>
+#include <mdspan/mdspan.hpp>
+#include <type_traits>
+#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<class ValueType>
+constexpr bool is_blas_value_type_v =
+  std::is_same_v<ValueType, float> ||
+  std::is_same_v<ValueType, double> ||
+  std::is_same_v<ValueType, std::complex<float>> ||
+  std::is_same_v<ValueType, std::complex<double>>;
+
+template<class Layout>
+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<Layout, layout_left> ||
+  std::is_same_v<Layout, layout_right> ||
+  // AMK 5/20/26 - layout_left_padded and layout_right_padded were added in C++26.
+  // According to cppreference, padded layouts are covered by the same feature test as submdspan.
+  #ifdef __cpp_lib_submdspan
+  std::is_same_v<Layout, layout_left_padded> ||
+  std::is_same_v<Layout, layout_right_padded> ||
+  #endif
+  std::is_same_v<Layout, layout_stride>;
+
+// The BLAS accepts accessors that deal with pointers to memory:
+// default_accessor and aligned_accessor.
+// Those are both templated classes, 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<class Accessor>
+constexpr bool is_default_accessor_v = false;
+
+template<class ElementType>
+constexpr bool is_default_accessor_v<default_accessor<ElementType>> = true;
+
+template<class Accessor>
+constexpr bool is_aligned_accessor_v = false;
+
+#ifdef __cpp_lib_aligned_accessor
+template<class ElementType, std::size_t ByteAlignment>
+constexpr bool is_aligned_accessor_v<aligned_accessor<ElementType, ByteAlignment>> = true;
+#endif
+
+template<class Accessor>
+constexpr bool is_blas_accessor_type_v =
+  is_default_accessor_v<Accessor> ||
+  is_aligned_accessor_v<Accessor>;
+
+// 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_