MemoryCompaction/small_test.cu at main · kevkrist/MemoryCompaction · GitHub

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
#include "compaction.cuh"
#include <iostream>
#include <thrust/device_vector.h>
#include <thrust/host_vector.h>
#include <thrust/iterator/counting_iterator.h>

/*
 * @author Kevin Kristensen
 */

int main()
{
  using chunk_t           = uint64_t;
  using scan_tile_state_t = cub::ScanTileState<uint64_t>;

  constexpr int32_t block_threads    = 32;
  constexpr int32_t items_per_thread = 2;
  constexpr int32_t tile_items       = block_threads * items_per_thread;

  /*
   * We have 5 allocations. The first is invalid, after which their validity alternates between
   * allocations. Thus, after compaction, we expect only the 1, 3 allocations in the memory pool.
   * We mark the validity of each allocation index entry in the MSB of the entry. This has lead to
   * slightly more complexity than I anticipated, but it is obviously more memory efficient.
   * We will use 4 thread blocks, each consisting of 32 threads with 2 items per thread.
   * Unfortunately, we cannot work on a smaller problem size because decoupled look-back requires
   * 32 threads per thread block.
   * The problem is designed to produce an interesting distribution of indices across the tiles.
   * The last allocation index points to the end of the memory pool.
   */
  thrust::host_vector<uint64_t> allocation_indices_h{0,
                                                     set_top_digit<uint64_t>(64),
                                                     125,
                                                     set_top_digit<uint64_t>(187),
                                                     193,
                                                     195}; // Pointer to end of memory pool
  auto num_allocation_indices = allocation_indices_h.size() - 1;

  /// Initialize the memory pool and the expected result
  thrust::host_vector<chunk_t> memory_pool_h{};
  thrust::host_vector<chunk_t> expected_memory_pool_h{};
  for (uint64_t i = 0; i < num_allocation_indices; ++i)
  {
    auto length =
      remove_top_digit(allocation_indices_h[i + 1]) - remove_top_digit(allocation_indices_h[i]);
    auto is_valid = evaluate_top_digit(allocation_indices_h[i]);
    for (auto j = 0; j < length; ++j)
    {
      memory_pool_h.push_back(i);
      if (is_valid)
      {
        expected_memory_pool_h.push_back(i);
      }
    }
  }

  /// Copy test data to device
  thrust::device_vector<uint64_t> allocation_indices = allocation_indices_h;
  thrust::device_vector<chunk_t> memory_pool         = memory_pool_h;

  /// Initialize the merge partitions
  auto num_tiles =
    static_cast<int32_t>(cuda::ceil_div(memory_pool.size() + num_allocation_indices, tile_items));
  auto num_merge_partitions = num_tiles + 1;
  thrust::device_vector<uint64_t> merge_partitions(num_merge_partitions);
  partition_merge_path_kernel<tile_items>
    <<<cuda::ceil_div(num_merge_partitions, block_threads), block_threads>>>(
      allocation_indices.begin(),
      num_allocation_indices,
      thrust::make_counting_iterator<uint64_t>(0),
      memory_pool.size(),
      thrust::raw_pointer_cast(merge_partitions.data()),
      num_merge_partitions);

  /// Determine the temporary storage needed for decoupled look-back
  std::size_t scan_tile_state_storage_bytes = 0;
  scan_tile_state_t::AllocationSize(num_tiles, scan_tile_state_storage_bytes);
  thrust::device_vector<uint8_t> scan_tile_state_storage(scan_tile_state_storage_bytes);

  /// Initialize the temporary storage for decoupled look-back
  scan_tile_state_t scan_tile_state{};
  scan_tile_state.Init(num_tiles,
                       thrust::raw_pointer_cast(scan_tile_state_storage.data()),
                       scan_tile_state_storage_bytes);
  scan_tile_state_init_kernel<<<cuda::ceil_div(num_tiles, block_threads), block_threads>>>(
    scan_tile_state,
    num_tiles);

  /// Launch the compaction kernel
  thrust::device_vector<uint64_t> compaction_counter(1);
  compact_kernel<block_threads, items_per_thread, chunk_t>
    <<<num_tiles, block_threads>>>(thrust::raw_pointer_cast(allocation_indices.data()),
                                   num_allocation_indices,
                                   thrust::raw_pointer_cast(memory_pool.data()),
                                   memory_pool.size(),
                                   thrust::raw_pointer_cast(merge_partitions.data()),
                                   scan_tile_state,
                                   thrust::raw_pointer_cast(compaction_counter.data()));
  CubDebugExit(cudaDeviceSynchronize());

  // Check the results
  if (expected_memory_pool_h.size() != compaction_counter[0])
  {
    std::cerr << "Expected compaction size (" << expected_memory_pool_h.size()
              << ") != actual compaction size (" << compaction_counter[0] << ")\n";
    return EXIT_FAILURE;
  }
  for (auto i = 0; i < expected_memory_pool_h.size(); ++i)
  {
    // Print the result, since it's small
    std::cout << memory_pool[i] << ", ";

    // Validate against expected result
    if (expected_memory_pool_h[i] != memory_pool[i])
    {
      std::cerr << "Expected value (" << expected_memory_pool_h[i] << ") at index " << i
                << " != actual value (" << memory_pool[i] << ")\n";
      return EXIT_FAILURE;
    }
  }
  std::cout << "\nSUCCESS\n";

  return EXIT_SUCCESS;
}