Skip to content

Commit

Permalink
Don't assume the default memory space is used (#1969)
Browse files Browse the repository at this point in the history 
(Fixes #1910)
In the source and especially unit tests, assume the device type we
want can have any memory space (not just the default memory space of
the execution space).

Lets TestExecSpace either be a Kokkos::Device<E, M> type, or
just an execution space (the old behavior).

In all Cuda unit tests, use <Cuda, CudaUVMSpace> as the TestExecSpace
if CudaUVMSpace is instantiated but CudaSpace is not. Otherwise, use
<Cuda, CudaSpace>. All other backends are unchanged, but this
would also let us test devices like <HIP, HIPManagedSpace> and
<SYCL, SYCLSharedUSMSpace> in the future.
  • Loading branch information
@brian-kelley
brian-kelley authored Sep 18, 2023
1 parent 5cce74d commit f5ad8b8
Show file tree
Hide file tree
Showing 67 changed files with 355 additions and 290 deletions.
3 changes: 2 additions & 1 deletion batched/dense/unit_test/Test_Batched_SerialAxpy.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ namespace Axpy {

template <typename DeviceType, typename ViewType, typename alphaViewType>
struct Functor_TestBatchedSerialAxpy {
using execution_space = typename DeviceType::execution_space;
const alphaViewType _alpha;
const ViewType _X;
const ViewType _Y;
Expand All @@ -54,7 +55,7 @@ struct Functor_TestBatchedSerialAxpy {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType> policy(0, _X.extent(0));
Kokkos::RangePolicy<execution_space> policy(0, _X.extent(0));
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
3 changes: 2 additions & 1 deletion batched/dense/unit_test/Test_Batched_SerialEigendecomposition.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ namespace Test {
typename ViewRank2Type,
typename WorkViewType>
struct Functor_TestBatchedSerialEigendecomposition {
using execution_space = typename DeviceType::execution_space;
ViewRank3Type _A;
ViewRank2Type _Er, _Ei;
ViewRank3Type _UL, _UR;
Expand Down Expand Up @@ -70,7 +71,7 @@ namespace Test {
>::value ? "::ComplexFloat" : std::is_same<value_type,Kokkos::complex<double>
>::value ? "::ComplexDouble" : "::UnknownValueType" ); std::string name =
name_region + name_value_type; Kokkos::Profiling::pushRegion( name.c_str() );
Kokkos::RangePolicy<DeviceType> policy(0, _A.extent(0));
Kokkos::RangePolicy<execution_space> policy(0, _A.extent(0));
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
3 changes: 2 additions & 1 deletion batched/dense/unit_test/Test_Batched_SerialGemm.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ struct ParamTag {
template <typename DeviceType, typename ViewType, typename ScalarType,
typename ParamTagType, typename AlgoTagType>
struct Functor_TestBatchedSerialGemm {
using execution_space = typename DeviceType::execution_space;
ViewType _a, _b, _c;

ScalarType _alpha, _beta;
Expand All @@ -66,7 +67,7 @@ struct Functor_TestBatchedSerialGemm {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType, ParamTagType> policy(0, _c.extent(0));
Kokkos::RangePolicy<execution_space, ParamTagType> policy(0, _c.extent(0));
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
3 changes: 2 additions & 1 deletion batched/dense/unit_test/Test_Batched_SerialGesv.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ namespace Gesv {
template <typename DeviceType, typename MatrixType, typename VectorType,
typename AlgoTagType>
struct Functor_TestBatchedSerialGesv {
using execution_space = typename DeviceType::execution_space;
const MatrixType _A;
const MatrixType _tmp;
const VectorType _X;
Expand All @@ -61,7 +62,7 @@ struct Functor_TestBatchedSerialGesv {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType> policy(0, _X.extent(0));
Kokkos::RangePolicy<execution_space> policy(0, _X.extent(0));
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
9 changes: 6 additions & 3 deletions batched/dense/unit_test/Test_Batched_SerialInverseLU.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ struct ParamTag {
template <typename DeviceType, typename ViewType, typename ScalarType,
typename ParamTagType, typename AlgoTagType>
struct Functor_BatchedSerialGemm {
using execution_space = typename DeviceType::execution_space;
ViewType _a, _b, _c;

ScalarType _alpha, _beta;
Expand Down Expand Up @@ -72,14 +73,15 @@ struct Functor_BatchedSerialGemm {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType, ParamTagType> policy(0, _c.extent(0));
Kokkos::RangePolicy<execution_space, ParamTagType> policy(0, _c.extent(0));
Kokkos::parallel_for((name + "::GemmFunctor").c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
};

template <typename DeviceType, typename ViewType, typename AlgoTagType>
struct Functor_BatchedSerialLU {
using execution_space = typename DeviceType::execution_space;
ViewType _a;

KOKKOS_INLINE_FUNCTION
Expand All @@ -100,7 +102,7 @@ struct Functor_BatchedSerialLU {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType> policy(0, _a.extent(0));
Kokkos::RangePolicy<execution_space> policy(0, _a.extent(0));
Kokkos::parallel_for((name + "::LUFunctor").c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand All @@ -109,6 +111,7 @@ struct Functor_BatchedSerialLU {
template <typename DeviceType, typename AViewType, typename WViewType,
typename AlgoTagType>
struct Functor_TestBatchedSerialInverseLU {
using execution_space = typename DeviceType::execution_space;
AViewType _a;
WViewType _w;

Expand All @@ -130,7 +133,7 @@ struct Functor_TestBatchedSerialInverseLU {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType> policy(0, _a.extent(0));
Kokkos::RangePolicy<execution_space> policy(0, _a.extent(0));
Kokkos::parallel_for((name + "::InverseLUFunctor").c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
3 changes: 2 additions & 1 deletion batched/dense/unit_test/Test_Batched_SerialLU.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ namespace Test {

template <typename DeviceType, typename ViewType, typename AlgoTagType>
struct Functor_TestBatchedSerialLU {
using execution_space = typename DeviceType::execution_space;
ViewType _a;

KOKKOS_INLINE_FUNCTION
Expand All @@ -52,7 +53,7 @@ struct Functor_TestBatchedSerialLU {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType> policy(0, _a.extent(0));
Kokkos::RangePolicy<execution_space> policy(0, _a.extent(0));
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
9 changes: 5 additions & 4 deletions batched/dense/unit_test/Test_Batched_SerialSVD.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -406,13 +406,14 @@ void GenerateTestData(ViewT data) {
});
}

template <typename Scalar, typename Layout, typename ExeSpace, int N = 3>
template <typename Scalar, typename Layout, typename Device, int N = 3>
void testIssue1786() {
using memory_space = typename ExeSpace::memory_space;
using execution_space = typename Device::execution_space;
using memory_space = typename Device::memory_space;
constexpr int num_tests = 4;
Kokkos::View<Scalar * [3][3], Layout, memory_space> matrices("data",
num_tests);
GenerateTestData<ExeSpace>(matrices);
GenerateTestData<execution_space>(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));
Expand All @@ -425,7 +426,7 @@ void testIssue1786() {
"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));
auto policy = Kokkos::RangePolicy<execution_space>(0, matrices.extent(0));
Kokkos::parallel_for(
"polar decomposition", policy, KOKKOS_LAMBDA(int i) {
auto matrix_copy =
Expand Down
9 changes: 6 additions & 3 deletions batched/dense/unit_test/Test_Batched_SerialSolveLU.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ struct ParamTag {
template <typename DeviceType, typename ViewType, typename ScalarType,
typename ParamTagType, typename AlgoTagType>
struct Functor_BatchedSerialGemm {
using execution_space = typename DeviceType::execution_space;
ViewType _a, _b, _c;

ScalarType _alpha, _beta;
Expand Down Expand Up @@ -72,14 +73,15 @@ struct Functor_BatchedSerialGemm {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType, ParamTagType> policy(0, _c.extent(0));
Kokkos::RangePolicy<execution_space, ParamTagType> policy(0, _c.extent(0));
Kokkos::parallel_for((name + "::GemmFunctor").c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
};

template <typename DeviceType, typename ViewType, typename AlgoTagType>
struct Functor_BatchedSerialLU {
using execution_space = typename DeviceType::execution_space;
ViewType _a;

KOKKOS_INLINE_FUNCTION
Expand All @@ -100,7 +102,7 @@ struct Functor_BatchedSerialLU {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType> policy(0, _a.extent(0));
Kokkos::RangePolicy<execution_space> policy(0, _a.extent(0));
Kokkos::parallel_for((name + "::LUFunctor").c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand All @@ -109,6 +111,7 @@ struct Functor_BatchedSerialLU {
template <typename DeviceType, typename ViewType, typename TransType,
typename AlgoTagType>
struct Functor_TestBatchedSerialSolveLU {
using execution_space = typename DeviceType::execution_space;
ViewType _a;
ViewType _b;

Expand All @@ -130,7 +133,7 @@ struct Functor_TestBatchedSerialSolveLU {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType> policy(0, _a.extent(0));
Kokkos::RangePolicy<execution_space> policy(0, _a.extent(0));
Kokkos::parallel_for((name + "::SolveLUFunctor").c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
3 changes: 2 additions & 1 deletion batched/dense/unit_test/Test_Batched_SerialTrmm.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ struct ParamTag {
template <typename DeviceType, typename ViewType, typename ScalarType,
typename ParamTagType, typename AlgoTagType>
struct Functor_TestBatchedSerialTrmm {
using execution_space = typename DeviceType::execution_space;
ViewType _a, _b;

ScalarType _alpha;
Expand All @@ -138,7 +139,7 @@ struct Functor_TestBatchedSerialTrmm {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType, ParamTagType> policy(0, _a.extent(0));
Kokkos::RangePolicy<execution_space, ParamTagType> policy(0, _a.extent(0));
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
3 changes: 2 additions & 1 deletion batched/dense/unit_test/Test_Batched_SerialTrsm.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ struct ParamTag {
template <typename DeviceType, typename ViewType, typename ScalarType,
typename ParamTagType, typename AlgoTagType>
struct Functor_TestBatchedSerialTrsm {
using execution_space = typename DeviceType::execution_space;
ViewType _a, _b;

ScalarType _alpha;
Expand All @@ -65,7 +66,7 @@ struct Functor_TestBatchedSerialTrsm {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType, ParamTagType> policy(0, _b.extent(0));
Kokkos::RangePolicy<execution_space, ParamTagType> policy(0, _b.extent(0));
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
3 changes: 2 additions & 1 deletion batched/dense/unit_test/Test_Batched_SerialTrsv.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ struct ParamTag {
template <typename DeviceType, typename ViewType, typename ScalarType,
typename ParamTagType, typename AlgoTagType>
struct Functor_TestBatchedSerialTrsv {
using execution_space = typename DeviceType::execution_space;
ViewType _a, _b;

ScalarType _alpha;
Expand All @@ -64,7 +65,7 @@ struct Functor_TestBatchedSerialTrsv {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType, ParamTagType> policy(0, _b.extent(0));
Kokkos::RangePolicy<execution_space, ParamTagType> policy(0, _b.extent(0));
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
3 changes: 2 additions & 1 deletion batched/dense/unit_test/Test_Batched_SerialTrtri.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ struct ParamTag {
template <typename DeviceType, typename ViewType, typename ParamTagType,
typename AlgoTagType>
struct Functor_TestBatchedSerialTrtri {
using execution_space = typename DeviceType::execution_space;
ViewType _a;

KOKKOS_INLINE_FUNCTION
Expand All @@ -132,7 +133,7 @@ struct Functor_TestBatchedSerialTrtri {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::RangePolicy<DeviceType, ParamTagType> policy(0, _a.extent(0));
Kokkos::RangePolicy<execution_space, ParamTagType> policy(0, _a.extent(0));
Kokkos::parallel_for("Functor_TestBatchedSerialTrtri", policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
5 changes: 3 additions & 2 deletions batched/dense/unit_test/Test_Batched_TeamAxpy.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ namespace TeamAxpy {

template <typename DeviceType, typename ViewType, typename alphaViewType>
struct Functor_TestBatchedTeamAxpy {
using execution_space = typename DeviceType::execution_space;
const alphaViewType _alpha;
const ViewType _X;
const ViewType _Y;
Expand Down Expand Up @@ -65,8 +66,8 @@ struct Functor_TestBatchedTeamAxpy {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::TeamPolicy<DeviceType> policy(_X.extent(0) / _N_team,
Kokkos::AUTO(), Kokkos::AUTO());
Kokkos::TeamPolicy<execution_space> policy(_X.extent(0) / _N_team,
Kokkos::AUTO(), Kokkos::AUTO());
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
5 changes: 3 additions & 2 deletions batched/dense/unit_test/Test_Batched_TeamGemm.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ struct ParamTag {
template <typename DeviceType, typename ViewType, typename ScalarType,
typename ParamTagType, typename AlgoTagType>
struct Functor_TestBatchedTeamGemm {
using execution_space = typename DeviceType::execution_space;
ViewType _a, _b, _c;

ScalarType _alpha, _beta;
Expand Down Expand Up @@ -73,8 +74,8 @@ struct Functor_TestBatchedTeamGemm {
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
const int league_size = _c.extent(0);
Kokkos::TeamPolicy<DeviceType, ParamTagType> policy(league_size,
Kokkos::AUTO);
Kokkos::TeamPolicy<execution_space, ParamTagType> policy(league_size,
Kokkos::AUTO);
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
5 changes: 3 additions & 2 deletions batched/dense/unit_test/Test_Batched_TeamGesv.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ namespace TeamGesv {
template <typename DeviceType, typename MatrixType, typename VectorType,
typename AlgoTagType>
struct Functor_TestBatchedTeamGesv {
using execution_space = typename DeviceType::execution_space;
const MatrixType _A;
const VectorType _X;
const VectorType _B;
Expand Down Expand Up @@ -62,8 +63,8 @@ struct Functor_TestBatchedTeamGesv {
const std::string name_value_type = Test::value_type_name<value_type>();
std::string name = name_region + name_value_type;
Kokkos::Profiling::pushRegion(name.c_str());
Kokkos::TeamPolicy<DeviceType> policy(_X.extent(0), Kokkos::AUTO(),
Kokkos::AUTO());
Kokkos::TeamPolicy<execution_space> policy(_X.extent(0), Kokkos::AUTO(),
Kokkos::AUTO());

using MatrixViewType =
Kokkos::View<typename MatrixType::non_const_value_type **,
Expand Down
5 changes: 3 additions & 2 deletions batched/dense/unit_test/Test_Batched_TeamInverseLU.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ struct ParamTag {
template <typename DeviceType, typename ViewType, typename ScalarType,
typename ParamTagType, typename AlgoTagType>
struct Functor_BatchedTeamGemm {
using execution_space = typename DeviceType::execution_space;
ViewType _a, _b, _c;

ScalarType _alpha, _beta;
Expand Down Expand Up @@ -82,8 +83,8 @@ struct Functor_BatchedTeamGemm {
Kokkos::Profiling::pushRegion(name.c_str());

const int league_size = _c.extent(0);
Kokkos::TeamPolicy<DeviceType, ParamTagType> policy(league_size,
Kokkos::AUTO);
Kokkos::TeamPolicy<execution_space, ParamTagType> policy(league_size,
Kokkos::AUTO);
Kokkos::parallel_for((name + "::GemmFunctor").c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
4 changes: 3 additions & 1 deletion batched/dense/unit_test/Test_Batched_TeamLU.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@ namespace TeamLU {

template <typename DeviceType, typename ViewType, typename AlgoTagType>
struct Functor_TestBatchedTeamLU {
using execution_space = typename DeviceType::execution_space;

ViewType _a;

KOKKOS_INLINE_FUNCTION
Expand All @@ -60,7 +62,7 @@ struct Functor_TestBatchedTeamLU {
Kokkos::Profiling::pushRegion(name.c_str());

const int league_size = _a.extent(0);
Kokkos::TeamPolicy<DeviceType> policy(league_size, Kokkos::AUTO);
Kokkos::TeamPolicy<execution_space> policy(league_size, Kokkos::AUTO);
Kokkos::parallel_for(name.c_str(), policy, *this);
Kokkos::Profiling::popRegion();
}
Expand Down
Loading

0 comments on commit f5ad8b8

Please sign in to comment.