MCPcopy Create free account
hub / github.com/Oneflow-Inc/oneflow / TestCache

Function TestCache

oneflow/core/embedding/cache_test.cpp:36–212  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

34}
35
36void TestCache(Cache* cache, uint32_t line_size) {
37 std::unique_ptr<ep::DeviceManagerRegistry> device_manager_registry(
38 new ep::DeviceManagerRegistry());
39 auto device = device_manager_registry->GetDevice(DeviceType::kCUDA, 0);
40 ep::Stream* stream = device->CreateStream();
41
42 std::unordered_set<int64_t> in_cache;
43 const size_t n_iter = 32;
44 const uint32_t n_keys = 1024;
45 int64_t* d_keys;
46 int64_t* keys;
47 uint32_t* d_n_missing;
48 uint32_t* n_missing;
49 int64_t* d_missing_keys;
50 int64_t* missing_keys;
51 uint32_t* d_missing_indices;
52 uint32_t* missing_indices;
53 float* d_values;
54 float* values;
55 float* d_evicted_values;
56 float* evicted_values;
57 uint32_t* d_n_evicted;
58 uint32_t* n_evicted;
59 int64_t* d_evicted_keys;
60 int64_t* evicted_keys;
61 uint8_t* mask;
62 const size_t keys_size = n_keys * sizeof(int64_t);
63 OF_CUDA_CHECK(cudaMalloc(&d_keys, keys_size));
64 OF_CUDA_CHECK(cudaMallocHost(&keys, keys_size));
65 OF_CUDA_CHECK(cudaMalloc(&d_n_missing, sizeof(uint32_t)));
66 OF_CUDA_CHECK(cudaMallocHost(&n_missing, sizeof(uint32_t)));
67 OF_CUDA_CHECK(cudaMalloc(&d_missing_keys, keys_size));
68 OF_CUDA_CHECK(cudaMallocHost(&missing_keys, keys_size));
69 const size_t indices_size = n_keys * sizeof(uint32_t);
70 OF_CUDA_CHECK(cudaMalloc(&d_missing_indices, indices_size));
71 OF_CUDA_CHECK(cudaMallocHost(&missing_indices, indices_size));
72 const size_t values_size = n_keys * line_size * sizeof(float);
73 OF_CUDA_CHECK(cudaMalloc(&d_values, values_size));
74 OF_CUDA_CHECK(cudaMallocHost(&values, values_size));
75 OF_CUDA_CHECK(cudaMalloc(&d_evicted_values, values_size));
76 OF_CUDA_CHECK(cudaMallocHost(&evicted_values, values_size));
77 OF_CUDA_CHECK(cudaMalloc(&d_n_evicted, sizeof(uint32_t)));
78 OF_CUDA_CHECK(cudaMallocHost(&n_evicted, sizeof(uint32_t)));
79 OF_CUDA_CHECK(cudaMalloc(&d_evicted_keys, keys_size));
80 OF_CUDA_CHECK(cudaMallocHost(&evicted_keys, keys_size));
81 OF_CUDA_CHECK(cudaMalloc(&mask, n_keys));
82 std::vector<int64_t> random_keys(n_keys * 32);
83 std::iota(random_keys.begin(), random_keys.end(), 1);
84 std::random_device rd;
85 std::mt19937 g(rd());
86 for (size_t iter = 0; iter < n_iter; ++iter) {
87 std::shuffle(random_keys.begin(), random_keys.end(), g);
88 std::copy(random_keys.begin(), random_keys.begin() + n_keys, keys);
89 uint32_t expect_n_missing = 0;
90 std::unordered_set<int64_t> expect_missing_keys_set;
91 std::unordered_set<uint32_t> expect_missing_indices_set;
92 std::unordered_set<int64_t> keys_set;
93 for (size_t i = 0; i < n_keys; ++i) {

Callers 1

TESTFunction · 0.85

Calls 10

DumpCapacityMethod · 0.80
GetDeviceMethod · 0.45
CreateStreamMethod · 0.45
beginMethod · 0.45
endMethod · 0.45
GetMethod · 0.45
PutMethod · 0.45
DumpMethod · 0.45
sizeMethod · 0.45
DestroyStreamMethod · 0.45

Tested by

no test coverage detected