Hello all 
.
I have written the below C++ extension (CPU) to fill nodes into fixed-size buckets in a FIFO manner. Each node goes into a single bucket in each of the L tables. I parallelize over L to avoid collisions.
Dimensions Used:- num_nodes=524288, num_buckets=4096, bucket_size=128, num_threads=48
(I use a system with about 1000G RAM, supporting up to 96 parallel threads.)
When using values of L < num_threads, I expect latency to be relatively constant w.r.t L. However, I find that latency varies significantly when I change L. What could be the reason for this? Does it have to do with caching? Thanks in advance.
L = 12 => 5.3450 ms
L = 24 => 12.8900 ms
L = 36 => 22.2119 ms
L = 48 => 29.6163 ms
void fill_buckets_FIFO(
const torch::Tensor& indices,
torch::Tensor& buckets,
torch::Tensor& bucket_counts) {
int32_t num_nodes = indices.size(0);
int32_t L = buckets.size(0);
// int32_t num_buckets = buckets.size(1);
int32_t bucket_size = buckets.size(2);
auto buckets_0 = buckets.accessor<int32_t, 3>(); // L x num_buckets x bucket_size
auto bucket_counts_0 = bucket_counts.accessor<int32_t, 2>(); // L x num_buckets
auto indices_0 = indices.accessor<int32_t, 2>(); // num_nodes x L
at::parallel_for(0, L, 0, [&](int32_t start, int32_t end) {
for (int32_t l = start; l < end; l++) {
auto buckets_1 = buckets_0[l];
auto bucket_counts_1 = bucket_counts_0[l];
for(int32_t i = 0; i < num_nodes; i++) {
int32_t bucket_index = indices_0[i][l];
int32_t &bucket_count = bucket_counts_1[bucket_index];
buckets_1[bucket_index][bucket_count % bucket_size] = i;
bucket_count++;
}
}
});
}