File: Test04_ParallelFor_RangePolicy.hpp

package info (click to toggle)
kokkos 4.7.01-2
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 16,636 kB
  • sloc: cpp: 223,676; sh: 2,446; makefile: 2,437; python: 91; fortran: 4; ansic: 2
file content (145 lines) | stat: -rw-r--r-- 4,593 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
//@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

#include <Kokkos_Core.hpp>
#include <gtest/gtest.h>

/// @Kokkos_Feature_Level_Required:4
// parallel-for unit test.
// In this test, different elements of an array are updated by different
// threads.

namespace Test {

using value_type = double;
int num_elements = 10;

struct ParallelForFunctor {
  value_type *_data;
  const value_type _value;

  ParallelForFunctor(value_type *data, const value_type value)
      : _data(data), _value(value) {}

  KOKKOS_INLINE_FUNCTION
  void operator()(const int i) const { _data[i] = (i + 1) * _value; }
};

template <class ExecSpace>
struct TestParallel_For {
  // Memory space type for Device and Host data
  using d_memspace_type = typename ExecSpace::memory_space;
  using h_memspace_type = Kokkos::HostSpace;

  value_type *deviceData, *hostData;
  const value_type value = 0.5;

  // Check if the array values are updated correctly.
  void correctness_check(value_type *data) {
    for (int i = 0; i < num_elements; ++i) {
      ASSERT_EQ((i + 1) * value, data[i])
          << "Values in index " << i << " are incorrect";
    }
  }

  // Routine to allocate memory in a specific memory space.
  template <class MemSpace>
  value_type *allocate_mem(int N) {
    return (static_cast<value_type *>(
        Kokkos::kokkos_malloc<MemSpace>("deviceData", N * sizeof(value_type))));
  }

  // Routine to free the memory from a specific memory space.
  template <class MemSpace>
  void free_mem(value_type *data) {
    Kokkos::kokkos_free<MemSpace>(data);
  }

  void init() {
    // Allocate memory on Device space.
    deviceData = allocate_mem<d_memspace_type>(num_elements);
    ASSERT_NE(deviceData, nullptr);

    // Allocate memory on Host space.
    hostData = allocate_mem<h_memspace_type>(num_elements);
    ASSERT_NE(hostData, nullptr);
  }

  void check_correctness_and_cleanup() {
    // Copy the data back to Host memory space
    Kokkos::Impl::DeepCopy<h_memspace_type, d_memspace_type>(
        hostData, deviceData, num_elements * sizeof(value_type));
    Kokkos::fence("Fence after copying data to host memory space");

    // Check if all data has been update correctly
    correctness_check(hostData);

    // free the allocated memory
    free_mem<d_memspace_type>(deviceData);
    free_mem<h_memspace_type>(hostData);
  }

  // A simple parallel for test with functors
  void simple_test() {
    // Allocates memory for num_elements number of value_type elements in the
    // host and device memory spaces.
    init();

    // parallel-for functor called for num_elements number of iterations.
    Kokkos::parallel_for("parallel_for",
                         Kokkos::RangePolicy<ExecSpace>(0, num_elements),
                         ParallelForFunctor(deviceData, value));

    Kokkos::fence();
    // Checks if parallel_for gave the correct results.
    // Frees the allocated memory in init().
    check_correctness_and_cleanup();
  }

  // A parallel_for test with user defined RangePolicy
  void range_policy() {
    // Allocates memory for num_elements number of value_type elements in the
    // host and device memory spaces.
    init();

    // Creates a range policy that uses dynamic scheduling.
    using range_policy_t =
        Kokkos::RangePolicy<ExecSpace, Kokkos::Schedule<Kokkos::Dynamic> >;

    // parallel-for functor with range-policy from 0 to num_elements iterations.
    Kokkos::parallel_for("RangePolicy_ParallelFor",
                         range_policy_t(0, num_elements),
                         ParallelForFunctor(deviceData, value));

    // Checks if parallel_for gave the correct results.
    // Free the allocated memory in init().
    check_correctness_and_cleanup();
  }
};

TEST(TEST_CATEGORY, IncrTest_04_simple_parallelFor) {
  if (std::is_same_v<Kokkos::DefaultExecutionSpace, TEST_EXECSPACE>) {
    TestParallel_For<TEST_EXECSPACE> test;
    test.simple_test();
  }
}

TEST(TEST_CATEGORY, IncrTest_04_RangePolicy_parallelFor) {
  TestParallel_For<TEST_EXECSPACE> test;
  test.range_policy();
}

}  // namespace Test