Reduce by Keys performance using linear index to rows index

Question

I have been checking the code rows_sum.cu from thrust examples and I could not understand what is exactly happening in linear_index_to_rows_index.

Can someone please explain me with some example what does convert a linear index to row index mean?

Reference: https://github.com/thrust/thrust/blob/master/examples/sum_rows.cu

Answer 1

In thrust, a common storage format is a vector container, (e.g. device_vector and host_vector). In typical usage, these containers are a one-dimensional storage format. The concept of "row" and "column" don't really apply to a 1D vector. We simply talk about the index of an element. In this case, it is a linear (1D) index.

In order to store a 2 dimensional item, such as a matrix where the concept of "row" and "column" indexing of elements is applicable, a common approach to using the 1D containers in thrust is to "flatten" or "linearize" the storage:

Matrix A:
        column:
        0    1 
row: 0  1    2
     1  3    4


Vector A:
index:   0    1    2    3
element: 1    2    3    4

At this point, we can conveniently use thrust operations. But what if we want to do operations on specific rows or columns of the original matrix? That information (row, column) is "lost" when we linearize or flatten it into a 1D vector. But if we know the dimensions of the original matrix, we can use that to convert a linear index:

  0     1     2     3

into a row/column index:

(0,0) (0,1) (1,0) (1,1)

and the reverse (row/column to linear index).

In the case of the thrust example you linked, the functor linear_index_to_row_index converts a given linear index to its associated row index. This allows the programmer to write a thrust operation that works on specific rows of data, such as summing the "rows" of the original matrix, even though it is now stored in a linear 1D vector.

Specifically, when given the following linear indicies:

  0    1    2    3

For my example, the functor would return:

  0    0    1    1

Because that represents the row of the original matrix that the specific element in the vector belongs to.

If I want to sum all the elements of each row together, producing one sum per row, then I can use the row index generated by that functor to identify the row of each element. At that point, reduce_by_key can easily sum together the elements that have the same row index, producing one result per row.