Row-wise Argmax
For every row of a matrix, find the column that holds the largest value. If the largest value appears more than once, report the first (smallest) column.
Input
solve receives two pointers to GPU memory and two sizes:
x, the input:Nrows andMcolumns.X[i][j]isx[i * m + j].y, where the answer goes:Nentries.Y[i]isy[i].n, the number of rows, andm, the number of columns.
The sizes are:
1 <= N <= 2^201 <= M <= 4096N * M <= 2^24
Every entry of x is a signed 32-bit integer. x is an array of int, and
y is an array of unsigned int.
Output
Write the column index of the largest entry of every row into y:
Y[i] = smallest j such that X[i][j] = max_k X[i][k]
The answer must be exact. There is no tolerance.
Examples
Example 1
One row:
X = [ 1 3 2 ]
Output:
Y = [ 1 ]
Example 2
The largest value can appear more than once. Ties go to the first column:
X = [ 5 5 1 5 ]
[ -1 -2 -3 -4 ]
Output:
Y = [ 0 0 ]
Example 3
All entries can be negative, so starting the search at zero is wrong:
X = [ -5 -3 -4 ]
Output:
Y = [ 1 ]
Note
Use a strict comparison (
>), not>=. With>=, a row of equal entries reports its last column instead of its first.