Fix stride(0) ordering, bound incx, reject incx == 0 and deduce the CBLAS index type

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,98 @@ 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 a BLAS routine could be called to scale the vector, 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);
+  // We can't call x.stride(0) until we know that x.is_strided() is true.
+  if (x.is_strided() && n <= std::numeric_limits<impl::cblas_index>::max()) {
+    auto incx = x.stride(0);
+    if (incx > std::numeric_limits<impl::cblas_index>::max()) {
+      return false;
+    }
+    // Strided layout mappings can have stride zero; the BLAS cannot.
+    if (incx == 0) {
+      return false;
+    }
+    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,99 @@
+//@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>>;
+
+// The padded layouts are class templates taking a size_t, so detecting them
+// takes a partial specialization rather than is_same_v.
+template<class Layout>
+constexpr bool is_padded_layout_v = false;
+
+template<std::size_t PaddingValue>
+constexpr bool is_padded_layout_v<layout_left_padded<PaddingValue>> = true;
+
+template<std::size_t PaddingValue>
+constexpr bool is_padded_layout_v<layout_right_padded<PaddingValue>> = true;
+
+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> ||
+  is_padded_layout_v<Layout> ||
+  std::is_same_v<Layout, layout_stride>;
+
+// 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<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_blas_accessor_type_v =
+  is_default_accessor_v<Accessor>;
+
+#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<class R, class A, class... Rest> 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_

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<std::size_t, dynamic_extent>;
+    using vector_t = mdspan<scalar_t, extents_t, layout_stride>;
+
+    constexpr std::size_t vectorSize(5);
+    constexpr std::size_t stride(2);
+    std::vector<scalar_t> storage(vectorSize * stride);
+
+    layout_stride::mapping<extents_t> mapping{
+      extents_t{vectorSize}, std::array<std::size_t, 1>{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<std::size_t, 4>;
+    using layout_t = layout_left_padded<4>;
+    using vector_t = mdspan<scalar_t, extents_t, layout_t>;
+
+    constexpr std::size_t vectorSize(4);
+    constexpr std::size_t storageSize = vectorSize;
+    std::vector<scalar_t> storage(storageSize);
+
+    layout_t::mapping<extents_t> 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<std::size_t, 4>;
+    using layout_t = layout_right_padded<dynamic_extent>;
+    using vector_t = mdspan<scalar_t, extents_t, layout_t>;
+
+    constexpr std::size_t vectorSize(4);
+    constexpr std::size_t paddingValue(4);
+    constexpr std::size_t storageSize = vectorSize;
+    std::vector<scalar_t> storage(storageSize);
+
+    layout_t::mapping<extents_t> 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() {