Skip to content

Commit

Permalink
Fix #1786: check that work array is contiguous in SVD (#1793)
Browse files Browse the repository at this point in the history 
* Extend batched dense SVD test for #1786

* Batched SVD: check that work view is contiguous
  • Loading branch information
@brian-kelley
brian-kelley authored Apr 17, 2023
1 parent 03f48fa commit a176b93
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 30 deletions.
22 changes: 22 additions & 0 deletions batched/dense/impl/KokkosBatched_SVD_Serial_Impl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,19 @@ KOKKOS_INLINE_FUNCTION int SerialSVD::invoke(SVD_USV_Tag, const AViewType &A,
const SViewType &sigma,
const VViewType &Vt,
const WViewType &work) {
static_assert(Kokkos::is_view_v<AViewType> && AViewType::rank == 2,
"SVD: A must be a rank-2 view");
static_assert(Kokkos::is_view_v<UViewType> && UViewType::rank == 2,
"SVD: U must be a rank-2 view");
static_assert(Kokkos::is_view_v<SViewType> && SViewType::rank == 1,
"SVD: s must be a rank-1 view");
static_assert(Kokkos::is_view_v<VViewType> && VViewType::rank == 2,
"SVD: V must be a rank-2 view");
static_assert(Kokkos::is_view_v<WViewType> && WViewType::rank == 1,
"SVD: W must be a rank-1 view");
static_assert(
!std::is_same_v<typename WViewType::array_layout, Kokkos::LayoutStride>,
"SVD: W must be contiguous (not LayoutStride)");
using value_type = typename AViewType::non_const_value_type;
return KokkosBatched::SerialSVDInternal::invoke<value_type>(
A.extent(0), A.extent(1), A.data(), A.stride(0), A.stride(1), U.data(),
Expand All @@ -41,6 +54,15 @@ template <typename AViewType, typename SViewType, typename WViewType>
KOKKOS_INLINE_FUNCTION int SerialSVD::invoke(SVD_S_Tag, const AViewType &A,
const SViewType &sigma,
const WViewType &work) {
static_assert(Kokkos::is_view_v<AViewType> && AViewType::rank == 2,
"SVD: A must be a rank-2 view");
static_assert(Kokkos::is_view_v<SViewType> && SViewType::rank == 1,
"SVD: s must be a rank-1 view");
static_assert(Kokkos::is_view_v<WViewType> && WViewType::rank == 1,
"SVD: W must be a rank-1 view");
static_assert(
!std::is_same_v<typename WViewType::array_layout, Kokkos::LayoutStride>,
"SVD: W must be contiguous (not LayoutStride)");
using value_type = typename AViewType::non_const_value_type;
return KokkosBatched::SerialSVDInternal::invoke<value_type>(
A.extent(0), A.extent(1), A.data(), A.stride(0), A.stride(1), nullptr, 0,
Expand Down
127 changes: 97 additions & 30 deletions batched/dense/unit_test/Test_Batched_SerialSVD.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,54 +36,44 @@ float svdEpsilon() {
}
} // namespace Test

template <typename Vector>
double simpleNorm2(const Vector& v) {
using Scalar = typename Vector::non_const_value_type;
using KAT = Kokkos::ArithTraits<Scalar>;
auto vhost = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), v);
double d = 0;
for (size_t i = 0; i < v.extent(0); i++) {
double m = KAT::abs(vhost(i));
d += m * m;
}
return std::sqrt(d);
}

// NOTE: simpleDot and simpleNorm2 currently support only real scalars (OK since
// SVD does as well)
template <typename V1, typename V2>
typename V1::non_const_value_type simpleDot(const V1& v1, const V2& v2) {
using Scalar = typename V1::non_const_value_type;
using KAT = Kokkos::ArithTraits<Scalar>;
auto v1host = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), v1);
auto v2host = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), v2);
typename V1::non_const_value_type val = KAT::zero();
for (size_t i = 0; i < v1.extent(0); i++) {
val += v1host(i) * v2host(i);
}
return val;
Scalar d;
Kokkos::parallel_reduce(
Kokkos::RangePolicy<typename V1::execution_space>(0, v1.extent(0)),
KOKKOS_LAMBDA(int i, Scalar& ld) { ld += v1(i) * v2(i); }, d);
return d;
}
template <typename V>
typename V::non_const_value_type simpleNorm2(const V& v) {
return Kokkos::sqrt(simpleDot(v, v));
}

// Check that all columns of X are unit length and pairwise orthogonal
template <typename Mat>
void verifyOrthogonal(const Mat& X) {
using value_type = typename Mat::non_const_value_type;
int k = X.extent(1);
using Scalar = typename Mat::non_const_value_type;
int k = X.extent(1);
for (int i = 0; i < k; i++) {
auto col1 = Kokkos::subview(X, Kokkos::ALL(), i);
double len = simpleNorm2(col1);
Test::EXPECT_NEAR_KK(len, 1.0, Test::svdEpsilon<value_type>());
Test::EXPECT_NEAR_KK(len, 1.0, Test::svdEpsilon<Scalar>());
for (int j = 0; j < i; j++) {
auto col2 = Kokkos::subview(X, Kokkos::ALL(), j);
double d = Kokkos::ArithTraits<value_type>::abs(simpleDot(col1, col2));
Test::EXPECT_NEAR_KK(d, 0.0, Test::svdEpsilon<value_type>());
double d = Kokkos::ArithTraits<Scalar>::abs(simpleDot(col1, col2));
Test::EXPECT_NEAR_KK(d, 0.0, Test::svdEpsilon<Scalar>());
}
}
}

template <typename AView, typename UView, typename VtView, typename SigmaView>
void verifySVD(const AView& A, const UView& U, const VtView& Vt,
const SigmaView& sigma) {
using value_type = typename AView::non_const_value_type;
using KAT = Kokkos::ArithTraits<value_type>;
using Scalar = typename AView::non_const_value_type;
using KAT = Kokkos::ArithTraits<Scalar>;
// Check that U/V columns are unit length and orthogonal, and that U *
// diag(sigma) * V^T == A
int m = A.extent(0);
Expand All @@ -93,7 +83,7 @@ void verifySVD(const AView& A, const UView& U, const VtView& Vt,
// NOTE: V^T being square and orthonormal implies that V is, so we don't have
// to transpose it here.
verifyOrthogonal(Vt);
AView usvt("USV^T", m, n);
Kokkos::View<Scalar**, typename AView::device_type> usvt("USV^T", m, n);
for (int i = 0; i < maxrank; i++) {
auto Ucol =
Kokkos::subview(U, Kokkos::ALL(), Kokkos::make_pair<int>(i, i + 1));
Expand All @@ -103,7 +93,7 @@ void verifySVD(const AView& A, const UView& U, const VtView& Vt,
}
for (int i = 0; i < m; i++) {
for (int j = 0; j < n; j++) {
Test::EXPECT_NEAR_KK(usvt(i, j), A(i, j), Test::svdEpsilon<value_type>());
Test::EXPECT_NEAR_KK(usvt(i, j), A(i, j), Test::svdEpsilon<Scalar>());
}
}
// Make sure all singular values are positive
Expand Down Expand Up @@ -389,11 +379,86 @@ void testSVD() {
testSerialSVDSingularValuesOnly<Scalar, Layout, Device>(10, 8);
}

template <typename ViewT>
KOKKOS_INLINE_FUNCTION constexpr auto Determinant(ViewT F)
-> std::enable_if_t<Kokkos::is_view<ViewT>::value && ViewT::rank == 2,
double> {
return (F(0, 0) * F(1, 1) * F(2, 2) + F(0, 1) * F(1, 2) * F(2, 0) +
F(0, 2) * F(1, 0) * F(2, 1) -
(F(0, 2) * F(1, 1) * F(2, 0) + F(0, 1) * F(1, 0) * F(2, 2) +
F(0, 0) * F(1, 2) * F(2, 1)));
}

template <typename ExeSpace, typename ViewT>
void GenerateTestData(ViewT data) {
using memory_space = typename ExeSpace::memory_space;
// finite difference should return dPK2dU. So, we can analyze two cases.
Kokkos::Random_XorShift64_Pool<memory_space> random(13718);
Kokkos::fill_random(data, random, 1.0);
Kokkos::parallel_for(
Kokkos::RangePolicy<ExeSpace>(0, data.extent(0)), KOKKOS_LAMBDA(int i) {
auto data_i = Kokkos::subview(data, i, Kokkos::ALL(), Kokkos::ALL());
while (Determinant(data_i) < 0.5) {
data_i(0, 0) += 1.0;
data_i(1, 1) += 1.0;
data_i(2, 2) += 1.0;
}
});
}

template <typename Scalar, typename Layout, typename ExeSpace, int N = 3>
void testIssue1786() {
using memory_space = typename ExeSpace::memory_space;
constexpr int num_tests = 4;
Kokkos::View<Scalar * [3][3], Layout, memory_space> matrices("data",
num_tests);
GenerateTestData<ExeSpace>(matrices);
Kokkos::View<Scalar * [N][N], Layout, memory_space> Us("Us",
matrices.extent(0));
Kokkos::View<Scalar * [N], Layout, memory_space> Ss("Ss", matrices.extent(0));
Kokkos::View<Scalar * [N][N], Layout, memory_space> Vts("Vts",
matrices.extent(0));
// Make sure the 2nd dimension of works is contiguous
Kokkos::View<Scalar * [N], Kokkos::LayoutRight, memory_space> works(
"works", matrices.extent(0));
Kokkos::View<Scalar * [N][N], Layout, memory_space> matrices_copy(
"matrices_copy", matrices.extent(0));
// make a copy of the input data to avoid overwriting it
Kokkos::deep_copy(matrices_copy, matrices);
auto policy = Kokkos::RangePolicy<ExeSpace>(0, matrices.extent(0));
Kokkos::parallel_for(
"polar decomposition", policy, KOKKOS_LAMBDA(int i) {
auto matrix_copy =
Kokkos::subview(matrices_copy, i, Kokkos::ALL(), Kokkos::ALL());
auto U = Kokkos::subview(Us, i, Kokkos::ALL(), Kokkos::ALL());
auto S = Kokkos::subview(Ss, i, Kokkos::ALL());
auto Vt = Kokkos::subview(Vts, i, Kokkos::ALL(), Kokkos::ALL());
auto work = Kokkos::subview(works, i, Kokkos::ALL());
KokkosBatched::SerialSVD::invoke(KokkosBatched::SVD_USV_Tag{},
matrix_copy, U, S, Vt, work);
});

auto Us_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, Us);
auto Ss_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, Ss);
auto Vts_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, Vts);
auto matrices_h =
Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, matrices);
for (int i = 0; i < num_tests; i++) {
auto A = Kokkos::subview(matrices_h, i, Kokkos::ALL(), Kokkos::ALL());
auto U = Kokkos::subview(Us_h, i, Kokkos::ALL(), Kokkos::ALL());
auto S = Kokkos::subview(Ss_h, i, Kokkos::ALL());
auto Vt = Kokkos::subview(Vts_h, i, Kokkos::ALL(), Kokkos::ALL());
verifySVD(A, U, Vt, S);
}
}

#if defined(KOKKOSKERNELS_INST_DOUBLE)
TEST_F(TestCategory, batched_scalar_serial_svd_double) {
// Test general SVD on a few different input sizes (full rank randomized)
testSVD<double, Kokkos::LayoutLeft, TestExecSpace>();
testSVD<double, Kokkos::LayoutRight, TestExecSpace>();
testIssue1786<double, Kokkos::LayoutLeft, TestExecSpace>();
testIssue1786<double, Kokkos::LayoutRight, TestExecSpace>();
}
#endif

Expand All @@ -402,5 +467,7 @@ TEST_F(TestCategory, batched_scalar_serial_svd_float) {
// Test general SVD on a few different input sizes (full rank randomized)
testSVD<float, Kokkos::LayoutLeft, TestExecSpace>();
testSVD<float, Kokkos::LayoutRight, TestExecSpace>();
testIssue1786<float, Kokkos::LayoutLeft, TestExecSpace>();
testIssue1786<float, Kokkos::LayoutRight, TestExecSpace>();
}
#endif

0 comments on commit a176b93

Please sign in to comment.