Add unittest for linalg::axpy

cpp/include/raft/linalg/axpy.cuh

@@ -17,9 +17,11 @@
 #define __AXPY_H
 
 #pragma once
-
 #include "detail/axpy.cuh"
 
+#include <raft/core/device_mdspan.hpp>
+#include <raft/core/host_mdspan.hpp>
+
 namespace raft::linalg {
 
 /**
@@ -133,4 +135,4 @@ void axpy(const raft::handle_t& handle,
 
 }  // namespace raft::linalg
 
-#endif
+#endif

cpp/test/CMakeLists.txt

@@ -129,6 +129,7 @@ if(BUILD_TESTS)
     ConfigureTest(NAME LINALG_TEST
             PATH
             test/linalg/add.cu
+            test/linalg/axpy.cu
             test/linalg/binary_op.cu
             test/linalg/cholesky_r1.cu
             test/linalg/coalesced_reduction.cu

cpp/test/linalg/axpy.cu

@@ -0,0 +1,138 @@
+/*
+ * 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.
+ */
+#include <raft/linalg/axpy.cuh>
+
+#include "../test_utils.h"
+#include <gtest/gtest.h>
+#include <raft/random/rng.cuh>
+#include <raft/util/cuda_utils.cuh>
+
+namespace raft {
+namespace linalg {
+
+// Reference axpy implementation.
+template <typename T>
+__global__ void naiveAxpy(const int n, const T alpha, const T* x, T* y, int incx, int incy)
+{
+  int idx = threadIdx.x + blockIdx.x * blockDim.x;
+  if (idx < n) { y[idx * incy] += alpha * x[idx * incx]; }
+}
+
+template <typename InType, typename IdxType = int, typename OutType = InType>
+struct AxpyInputs {
+  OutType tolerance;
+  IdxType len;
+  InType alpha;
+  int incx;
+  int incy;
+  unsigned long long int seed;
+};
+
+template <typename T>
+class AxpyTest : public ::testing::TestWithParam<AxpyInputs<T>> {
+ protected:
+  raft::handle_t handle;
+  AxpyInputs<T> params;
+  rmm::device_uvector<T> refy;
+  rmm::device_uvector<T> y;
+
+ public:
+  AxpyTest()
+    : testing::TestWithParam<AxpyInputs<T>>(),
+      refy(0, handle.get_stream()),
+      y(0, handle.get_stream())
+  {
+    handle.sync_stream();
+  }
+
+ protected:
+  void SetUp() override
+  {
+    params = ::testing::TestWithParam<AxpyInputs<T>>::GetParam();
+
+    cudaStream_t stream = handle.get_stream();
+
+    raft::random::RngState r(params.seed);
+
+    int x_len = params.len * params.incx;
+    int y_len = params.len * params.incy;
+    rmm::device_uvector<T> x(x_len, stream);
+    y.resize(y_len, stream);
+
+    uniform(handle, r, x.data(), x_len, T(-1.0), T(1.0));
+    uniform(handle, r, y.data(), y_len, T(-1.0), T(1.0));
+
+    // Take a copy of the random generated values in y for the naive reference implementation
+    // this is necessary since axpy uses y for both input and output
+    refy = rmm::device_uvector<T>(y, stream);
+
+    int threads = 64;
+    int blocks  = raft::ceildiv<int>(params.len, threads);
+
+    naiveAxpy<T><<<blocks, threads, 0, stream>>>(
+      params.len, params.alpha, x.data(), refy.data(), params.incx, params.incy);
+
+    // TODO: passing incx/incy > 1 to axpy produces wrong results
+    // https://github.com/rapidsai/raft/issues/944
+    // We probably should be passing a strided mdspan here?
+    axpy(handle,
+         make_host_scalar_view<const T>(&params.alpha),
+         make_device_vector_view<const T>(x.data(), x_len),
+         make_device_vector_view<T>(y.data(), y_len),
+         params.incx,
+         params.incy);
+    handle.sync_stream();
+  }
+
+  void TearDown() override {}
+};
+
+const std::vector<AxpyInputs<float>> inputsf = {
+  {0.000001f, 1024 * 1024, 2.f, 1, 1, 1234ULL},
+  {0.000001f, 16 * 1024 * 1024, 128.f, 1, 1, 1234ULL},
+  {0.000001f, 98689, 4.f, 1, 1, 1234ULL},
+  {0.000001f, 4 * 1024 * 1024, -1, 1, 1, 1234ULL},
+};
+
+const std::vector<AxpyInputs<double>> inputsd = {
+  {0.000001f, 1024 * 1024, 2.f, 1, 1, 1234ULL},
+  {0.000001f, 16 * 1024 * 1024, 128.f, 1, 1, 1234ULL},
+  {0.000001f, 98689, 4.f, 1, 1, 1234ULL},
+  {0.000001f, 4 * 1024 * 1024, -1, 1, 1, 1234ULL},
+};
+
+typedef AxpyTest<float> AxpyTestF;
+TEST_P(AxpyTestF, Result)
+{
+  ASSERT_TRUE(raft::devArrMatch(
+    refy.data(), y.data(), params.len * params.incy, raft::CompareApprox<float>(params.tolerance)));
+}
+
+typedef AxpyTest<double> AxpyTestD;
+TEST_P(AxpyTestD, Result)
+{
+  ASSERT_TRUE(raft::devArrMatch(refy.data(),
+                                y.data(),
+                                params.len * params.incy,
+                                raft::CompareApprox<double>(params.tolerance)));
+}
+
+INSTANTIATE_TEST_SUITE_P(AxpyTests, AxpyTestF, ::testing::ValuesIn(inputsf));
+
+INSTANTIATE_TEST_SUITE_P(AxpyTests, AxpyTestD, ::testing::ValuesIn(inputsd));
+
+}  // end namespace linalg
+}  // end namespace raft