MemoryCompaction/large_eval.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
125
126
#define CUB_STDERR

#include "compaction.cuh"
#include <iostream>
#include <random>
#include <thrust/device_vector.h>
#include <thrust/fill.h>
#include <thrust/host_vector.h>
#include <thrust/iterator/counting_iterator.h>

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

  constexpr int32_t block_threads    = 512;
  constexpr int32_t items_per_thread = 8; // Execution configuration will change based on chunk size
  constexpr int32_t tile_items       = block_threads * items_per_thread;
  constexpr uint64_t memory_pool_size = 1ULL << 28ULL; // # chunks -> 4GB of data

  /// Let's evaluate the performance of a highly fragmented memory pool
  std::default_random_engine eng{};
  std::uniform_int_distribution pool_dist(0,
                                          1024 * 1024); // For generating the size of an allocation
  std::uniform_int_distribution valid_dist(0, 1);       // For determining the allocation validity

  thrust::host_vector<uint64_t> allocation_indices_h{};
  thrust::host_vector<chunk_t> memory_pool_h(memory_pool_size);
  thrust::host_vector<chunk_t> expected_memory_pool_h{};

  /// Generate the test data
  std::cout << "Generating test data (this takes a minute)...\n";
  uint64_t prev_memory_pool_index    = 0;
  uint64_t current_memory_pool_index = 0;
  bool prev_valid                    = false;
  bool current_valid                 = valid_dist(eng);
  uint64_t iter                      = 0;
  while (prev_memory_pool_index < memory_pool_size)
  {
    uint64_t allocation_index =
      current_valid ? set_top_digit(current_memory_pool_index) : current_memory_pool_index;
    allocation_indices_h.push_back(allocation_index);

    // Backfill the memory pool
    for (auto i = 0; i < (current_memory_pool_index - prev_memory_pool_index); ++i)
    {
      memory_pool_h[prev_memory_pool_index + i] = iter;
      if (prev_valid)
      {
        expected_memory_pool_h.push_back(iter);
      }
    }

    // Update
    prev_memory_pool_index = current_memory_pool_index;
    current_memory_pool_index += pool_dist(eng);
    if (current_memory_pool_index > memory_pool_size)
    {
      // Ensure a full memory pool
      current_memory_pool_index = memory_pool_size;
    }
    prev_valid    = current_valid;
    current_valid = valid_dist(eng) && current_memory_pool_index < memory_pool_size;
    ++iter;
  }
  auto num_allocation_indices = allocation_indices_h.size() - 1;
  std::cout << "Test data generated.\n";

  /// 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 indices
      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(cudaGetLastError());
  CubDebugExit(cudaDeviceSynchronize());

  std::cout << "Kernel done. Checking results...\n";

  // Check the result count, otherwise assume here that it's correct (for now)
  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;
  }
  std::cout << "SUCCESS\n";

  return EXIT_SUCCESS;
}