Vectorize index identification for faster GPU processing
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?

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.
