Vectorize index identification for faster GPU processing

Question

My input is a list of values, data_idx. In the example, the values range from [0, 5].

data_idx = [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0]

My desired output, filled_matrix is a tensor of shape max(value) by len(data_idx) where each row, r, of the tensor contains all of indices where data_idx == r and -1 for the rest of the row if the number of matched indices is fewer than len(data_idx)

For example, in the first row, r=0, data_idx==0 at indices [3, 11, 14]. The full output would look like:

filled_matrix = tensor([[ 3, 11, 14, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1], 
    [ 5, 10, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1],
    [ 0,  9, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1],
    [ 8, 12, 13, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1],
    [ 4,  6, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1],
    [ 1,  2,  7, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1]],
   dtype=torch.int8)

I have working for-loop code that accomplishes my goal.

import torch

max_idx = 6
data_idx = torch.tensor([2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0]).cuda()

max_number_data_idx = data_idx.shape[0]

filled_matrix = torch.zeros([max_idx, max_number_data_idx], dtype=torch.int8, device='cuda')
filled_matrix.fill_(-1)
for i in range(max_idx):
   same_idx = (data_idx == i).nonzero().flatten()
   filled_matrix[i][:same_idx.shape[0]] = same_idx

Now, I want to speed up this code. Specifically, I want it to be faster on the GPU. In the real scenario, the input, data_idx, can be a list containing millions of values. In that case, e.g 1 M of different values, the GPU will be call 1 M of time which make it very slow. My code is sequential and GPU hate sequential code.

Is there a function which will produce the same result more efficiently? Or a way to vectorize this for-loop?

Can you explain a little more what your goal is and what isn't working? For example, what is your desired output? What do you get instead? — Cecilia, May 07 '20 at 23:03
Thanks for you answer, my output is the "filled matrix". I want to find a way to optimize this code on GPU and have the same output — tani, May 08 '20 at 10:58
So just to clarify, this code does what you want on the gpu and now you want to move it to the cpu? — Cecilia, May 08 '20 at 14:32
It is working on both. I want to optimize it only for gpu. The code bellow is just an example. In the real scenario the input, data_idx, can be a list containing millions of values. In that case, e.g 1 M of different values, the gpu will be call 1 M of time which make it very slow. My code is sequential and gpu hate sequential code. So for me, the first option is either try to find a way to avoid the use of a loop (may be there is a hidden pytorch function that I do not know) or find a way to parallelize the loop. Thanks you again — tani, May 08 '20 at 16:47
based on your comments, I tried to re-write your question to be a little clearer. Hopefully, this will help you get the answer you are looking for. — Cecilia, May 08 '20 at 17:13

score 0 · Accepted Answer · answered May 08 '20 at 17:59

Disclaimer: I have not profiled this code too see if it is actually faster on a GPU.

One vectorized solution is to use tensor views to broadcast the comparison. Tensor views do not use additional memory. You can see more details in the documentation

First, make a matrix that contains the values that you want to compare to for each row. In this case, it's just the row indices.

comparison = torch.tensor(range(max_idx))

Now we are going to use expand and unsqueeze to make views of data_idx and comparison which are the same shape as filled_matrix.

comparison_view = comparison.unsqueeze(1).expand(max_idx, max_number_data_idx)
print(comparison_view)

# Each row is the index you want to compare to
# tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
    [2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
    [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3],
    [4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4],
    [5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5]])

data_idx_view = data_idx.expand(max_idx, max_number_data_idx)
print(data_idx_view)

# Each row is a copy of data_idx
# tensor([[2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0],
    [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0],
    [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0],
    [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0],
    [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0],
    [2, 5, 5, 0, 4, 1, 4, 5, 3, 2, 1, 0, 3, 3, 0]])

We can compare their equality and use nonzero to find the indices

mask = comparison_view == data_idx_view
mask_indices = mask.nonzero()

print(mask_indices)
# tensor([[ 0,  3],
    [ 0, 11],
    [ 0, 14],
    [ 1,  5],
    [ 1, 10],
    [ 2,  0],
    [ 2,  9],
    [ 3,  8],
    [ 3, 12],
    [ 3, 13],
    [ 4,  4],
    [ 4,  6],
    [ 5,  1],
    [ 5,  2],
    [ 5,  7]])

Now, you just need to manipulate these results into the format that you want for your output.

filled_matrix = torch.zeros([max_idx, max_number_data_idx], dtype=torch.int8)
filled_matrix.fill_(-1)
col_indices = [0, 1, 2, 0, 1, 0, 1, 0, 1, 2, 0, 1, 0, 1, 2]
filled_matrix[mask_indices[:, 0], col_indices] = mask_indices[:, 1].type(torch.int8)

I thought about several options for generating the col_indices list, but I couldn't come up with anything without a for loop.

col_indices = torch.zeros(mask_indices.shape[0])
for i in range(1, mask_indices.shape[0]):
    if mask_indices[i,0] == mask_indices[i-1,0]:
        col_indices[i] = col_indices[i-1]+1

You would need to do some profiling to see which code is actually faster.

Sorry for the late answer. Thank you very much for your help, it works very well. I tried it on my real code and it is 3-4 times faster. For generating the col_indices I used something else, explained here: https://stackoverflow.com/questions/55916932/multiple-ranges-np-arange . Thank you again it really helps me a lot! — tani, May 13 '20 at 07:20

Vectorize index identification for faster GPU processing

1 Answers1