Expose linalg::dot in public API

cpp/include/raft/core/device_mdspan.hpp

@@ -266,14 +266,23 @@ auto make_device_matrix_view(ElementType* ptr, IndexType n_rows, IndexType n_col
  * @tparam LayoutPolicy policy for strides and layout ordering
  * @param[in] ptr on device to wrap
  * @param[in] n number of elements in pointer
+ * @param[in] stride the stride between consecutive elements in the vector. Setting to a value
+ *            other than 1 requires LayoutPolicy to be set to layout_stride
  * @return raft::device_vector_view
  */
 template <typename ElementType,
           typename IndexType    = std::uint32_t,
           typename LayoutPolicy = layout_c_contiguous>
-auto make_device_vector_view(ElementType* ptr, IndexType n)
+auto make_device_vector_view(ElementType* ptr, IndexType n, IndexType stride = 1)
 {
-  return device_vector_view<ElementType, IndexType, LayoutPolicy>{ptr, n};
+  if constexpr (std::is_same_v<layout_stride, LayoutPolicy>) {
+    vector_extent<IndexType> exts{n};
+    std::array<IndexType, 1> strides{stride};
+    auto layout = typename LayoutPolicy::template mapping<vector_extent<IndexType>>{exts, strides};
+    return device_vector_view<ElementType, IndexType, LayoutPolicy>{ptr, layout};
+  } else {
+    RAFT_EXPECTS(stride == 1, "Having a stride != 1 requires a layout_stride LayoutPolicy");
+    return device_vector_view<ElementType, IndexType, LayoutPolicy>{ptr, n};
+  }
 }
-
-}  // end namespace raft
+}  // end namespace raft

cpp/include/raft/linalg/axpy.cuh

@@ -62,66 +62,61 @@ void axpy(const raft::handle_t& handle,
  * @brief axpy function
  *  It computes the following equation: y = alpha * x + y
  *
- * @tparam InType  Type raft::device_mdspan
- * @tparam ScalarIdxType Index Type of scalar
  * @param [in] handle raft::handle_t
  * @param [in] alpha raft::device_scalar_view
  * @param [in] x Input vector
  * @param [inout] y Output vector
  */
-template <typename InType,
-          typename OutType,
-          typename ScalarIdxType,
-          typename = raft::enable_if_input_device_mdspan<InType>,
-          typename = raft::enable_if_output_device_mdspan<OutType>>
+template <typename ElementType,
+          typename IndexType       = std::uint32_t,
+          typename InLayoutPolicy  = layout_c_contiguous,
+          typename OutLayoutPolicy = layout_c_contiguous,
+          typename ScalarIdxType   = std::uint32_t>
 void axpy(const raft::handle_t& handle,
-          raft::device_scalar_view<const typename InType::value_type, ScalarIdxType> alpha,
-          InType x,
-          OutType y)
+          raft::device_scalar_view<const ElementType, ScalarIdxType> alpha,
+          raft::device_vector_view<const ElementType, IndexType, InLayoutPolicy> x,
+          raft::device_vector_view<ElementType, IndexType, OutLayoutPolicy> y)
 {
   RAFT_EXPECTS(y.size() == x.size(), "Size mismatch between Output and Input");
 
-  axpy<typename InType::value_type, true>(handle,
-                                          y.size(),
-                                          alpha.data_handle(),
-                                          x.data_handle(),
-                                          x.stride(0),
-                                          y.data_handle(),
-                                          y.stride(0),
-                                          handle.get_stream());
+  axpy<ElementType, true>(handle,
+                          y.size(),
+                          alpha.data_handle(),
+                          x.data_handle(),
+                          x.stride(0),
+                          y.data_handle(),
+                          y.stride(0),
+                          handle.get_stream());
 }
 
 /**
  * @brief axpy function
  *  It computes the following equation: y = alpha * x + y
- *
- * @tparam MdspanType  Type raft::device_mdspan
- * @tparam ScalarIdxType Index Type of scalar
  * @param [in] handle raft::handle_t
  * @param [in] alpha raft::device_scalar_view
  * @param [in] x Input vector
  * @param [inout] y Output vector
  */
-template <typename InType,
-          typename OutType,
-          typename ScalarIdxType,
-          typename = raft::enable_if_input_device_mdspan<InType>,
-          typename = raft::enable_if_output_device_mdspan<OutType>>
+template <typename ElementType,
+          typename IndexType       = std::uint32_t,
+          typename InLayoutPolicy  = layout_c_contiguous,
+          typename OutLayoutPolicy = layout_c_contiguous,
+          typename ScalarIdxType   = std::uint32_t>
 void axpy(const raft::handle_t& handle,
-          raft::host_scalar_view<const typename InType::value_type, ScalarIdxType> alpha,
-          InType x,
-          OutType y)
+          raft::host_scalar_view<const ElementType, ScalarIdxType> alpha,
+          raft::device_vector_view<const ElementType, IndexType, InLayoutPolicy> x,
+          raft::device_vector_view<ElementType, IndexType, OutLayoutPolicy> y)
 {
   RAFT_EXPECTS(y.size() == x.size(), "Size mismatch between Output and Input");
 
-  axpy<typename InType::value_type, false>(handle,
-                                           y.size(),
-                                           alpha.data_handle(),
-                                           x.data_handle(),
-                                           x.stride(0),
-                                           y.data_handle(),
-                                           y.stride(0),
-                                           handle.get_stream());
+  axpy<ElementType, false>(handle,
+                           y.size(),
+                           alpha.data_handle(),
+                           x.data_handle(),
+                           x.stride(0),
+                           y.data_handle(),
+                           y.stride(0),
+                           handle.get_stream());
 }
 
 /** @} */  // end of group axpy

cpp/include/raft/linalg/dot.cuh

@@ -0,0 +1,59 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * 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.
+ */
+#ifndef __DOT_H
+#define __DOT_H
+
+#pragma once
+
+#include <raft/linalg/detail/cublas_wrappers.hpp>
+
+#include <raft/core/device_mdspan.hpp>
+#include <raft/core/handle.hpp>
+#include <raft/core/host_mdspan.hpp>
+
+namespace raft::linalg {
+/**
+ * @brief Computes the dot product of two vectors.
+ * @param[in] handle   raft::handle_t
+ * @param[in] x        First input vector
+ * @param[in] y        Second input vector
+ * @param[out] out     The output dot product between the x and y vectors.
+ * @note The out parameter can be either a host_scalar_view or device_scalar_view
+ */
+template <typename ElementType,
+          typename IndexType       = std::uint32_t,
+          typename ScalarIndexType = std::uint32_t,
+          typename LayoutPolicy1   = layout_c_contiguous,
+          typename LayoutPolicy2   = layout_c_contiguous>
+void dot(const raft::handle_t& handle,
+         raft::device_vector_view<const ElementType, IndexType, LayoutPolicy1> x,
+         raft::device_vector_view<const ElementType, IndexType, LayoutPolicy2> y,
+         raft::device_scalar_view<ElementType, ScalarIndexType> out)
+{
+  RAFT_EXPECTS(x.size() == y.size(),
+               "Size mismatch between x and y input vectors in raft::linalg::dot");
+
+  RAFT_CUBLAS_TRY(detail::cublasdot(handle.get_cublas_handle(),
+                                    x.size(),
+                                    x.data_handle(),
+                                    x.stride(0),
+                                    y.data_handle(),
+                                    y.stride(0),
+                                    out.data_handle(),
+                                    handle.get_stream()));
+}
+}  // namespace raft::linalg
+#endif

cpp/test/CMakeLists.txt

@@ -138,6 +138,7 @@ if(BUILD_TESTS)
             test/linalg/cholesky_r1.cu
             test/linalg/coalesced_reduction.cu
             test/linalg/divide.cu
+            test/linalg/dot.cu
             test/linalg/eig.cu
             test/linalg/eig_sel.cu
             test/linalg/gemm_layout.cu

cpp/test/linalg/axpy.cu

@@ -31,17 +31,6 @@ __global__ void naiveAxpy(const int n, const T alpha, const T* x, T* y, int incx
   if (idx < n) { y[idx * incy] += alpha * x[idx * incx]; }
 }
 
-template <typename ElementType,
-          typename IndexType    = std::uint32_t,
-          typename LayoutPolicy = layout_stride>
-auto make_strided_device_vector_view(ElementType* ptr, IndexType n, IndexType stride)
-{
-  vector_extent<IndexType> exts{n};
-  std::array<IndexType, 1> strides{stride};
-  auto layout = typename LayoutPolicy::mapping<vector_extent<IndexType>>{exts, strides};
-  return device_vector_view<ElementType, IndexType, LayoutPolicy>{ptr, layout};
-}
-
 template <typename InType, typename IndexType = int, typename OutType = InType>
 struct AxpyInputs {
   OutType tolerance;
@@ -52,11 +41,11 @@ struct AxpyInputs {
   unsigned long long int seed;
 };
 
-template <typename T>
+template <typename T, typename IndexType = int>
 class AxpyTest : public ::testing::TestWithParam<AxpyInputs<T>> {
  protected:
   raft::handle_t handle;
-  AxpyInputs<T> params;
+  AxpyInputs<T, IndexType> params;
   rmm::device_uvector<T> refy;
   rmm::device_uvector<T> y;
 
@@ -78,8 +67,8 @@ class AxpyTest : public ::testing::TestWithParam<AxpyInputs<T>> {
 
     raft::random::RngState r(params.seed);
 
-    int x_len = params.len * params.incx;
-    int y_len = params.len * params.incy;
+    IndexType x_len = params.len * params.incx;
+    IndexType y_len = params.len * params.incy;
     rmm::device_uvector<T> x(x_len, stream);
     y.resize(y_len, stream);
     refy.resize(y_len, stream);
@@ -100,18 +89,21 @@ class AxpyTest : public ::testing::TestWithParam<AxpyInputs<T>> {
     if ((params.incx > 1) && (params.incy > 1)) {
       axpy(handle,
            make_host_scalar_view<const T>(&params.alpha),
-           make_strided_device_vector_view<const T>(x.data(), params.len, params.incx),
-           make_strided_device_vector_view<T>(y.data(), params.len, params.incy));
+           make_device_vector_view<const T, IndexType, layout_stride>(
+             x.data(), params.len, params.incx),
+           make_device_vector_view<T, IndexType, layout_stride>(y.data(), params.len, params.incy));
+
     } else if (params.incx > 1) {
       axpy(handle,
            make_host_scalar_view<const T>(&params.alpha),
-           make_strided_device_vector_view<const T>(x.data(), params.len, params.incx),
+           make_device_vector_view<const T, IndexType, layout_stride>(
+             x.data(), params.len, params.incx),
            make_device_vector_view<T>(y.data(), params.len));
     } else if (params.incy > 1) {
       axpy(handle,
            make_host_scalar_view<const T>(&params.alpha),
            make_device_vector_view<const T>(x.data(), params.len),
-           make_strided_device_vector_view<T>(y.data(), params.len, params.incy));
+           make_device_vector_view<T, IndexType, layout_stride>(y.data(), params.len, params.incy));
     } else {
       axpy(handle,
            make_host_scalar_view<const T>(&params.alpha),