CSR row indexing is more parallel, but now there's an unfortunate list conversion.
@@ -61,20 +61,20 @@ def _get_csr_submatrix(Ap, Aj, Ax,
_csr_row_index_ker = core.ElementwiseKernel(
- '''raw I rows, raw I Ap, raw I Aj, raw T Ax, raw I Bp''',
+ '''raw I out_rows, raw I in_rows, raw I Ap, raw I Aj,
+ raw T Ax, raw I Bp''',
'''raw I Bj, raw T Bx''', '''
- const I row = rows[i];
- const I row_start = Ap[row];
- const I row_end = Ap[row+1];
- I out_row_idx = Bp[i];
+ const I out_row = out_rows[i];
+ const I in_row = in_rows[i];
- // Copy columns
- for(I j = row_start; j < row_end; j++) {
- Bj[out_row_idx] = Aj[j];
- Bx[out_row_idx] = Ax[j];
- out_row_idx++;
- }
+ // Look up starting offset
+ const I starting_output_offset = Bp[out_row];
+ const I output_offset = i - starting_output_offset;
+ const I starting_input_offset = Ap[in_row];
+
+ Bj[i] = Aj[starting_input_offset + output_offset];
+ Bx[i] = Ax[starting_input_offset + output_offset];
''', 'csr_row_index_ker', no_return=True)
@@ -100,7 +100,15 @@ def _csr_row_index(rows,
Bj = cupy.empty(nnz, dtype=Aj.dtype)
Bx = cupy.empty(nnz, dtype=Ax.dtype)
- _csr_row_index_ker(rows, Ap, Aj, Ax, Bp, Bj, Bx, size=rows.size)
+ Bp_counts = cupy.diff(Bp).tolist()
+
+ # TODO(cjnolet): Using a list for repeat is very slow.
+ out_rows = cupy.arange(rows.size, dtype=rows.dtype)
+ in_rows = cupy.repeat(rows, Bp_counts)
+ out_rows = cupy.repeat(out_rows, Bp_counts)
+
+ _csr_row_index_ker(out_rows, in_rows, Ap, Aj, Ax, Bp, Bj, Bx,
+ size=in_rows.size)
return Bj, Bx
