Add GPU implementation of SparseReshape #47251
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
sanjoy
suggested changes
Feb 25, 2021
| @@ -94,7 +94,7 @@ def testPropagatesFullyKnownDenseShapeWhenShapePartiallyKnown(self): | |||
| self.assertAllEqual((2, 3 * 4), sp_output.shape) | |||
|
|
|||
| def testSameShape(self): | |||
| with self.session(use_gpu=False) as sess: | |||
| with self.session(use_gpu=True) as sess: | |||
There was a problem hiding this comment.
You no longer need to set use_gpu=True explicitly, it default to True.
Comment on lines
36
to
49
| GPU_1D_KERNEL_LOOP(sparse_index, nnz) { | ||
| const Tindex* input_index = &input_indices[sparse_index * input_rank]; | ||
| Tindex* output_index = &output_indices[sparse_index * output_rank]; | ||
| Tindex dense_index = 0; | ||
| // Flatten input index from slowest- to fastest-changing dimension. | ||
| for (int i = 0; i < input_rank; ++i) { | ||
| dense_index = dense_index * input_shape[i] + input_index[i]; | ||
| } | ||
| // Compute output index from fastest- to slowest-changing dimension. | ||
| for (int i = output_rank; i-- > 0;) { | ||
| Tindex output_size = output_shape[i]; | ||
| output_index[i] = dense_index % output_size; | ||
| dense_index /= output_size; | ||
| } |
There was a problem hiding this comment.
Do you need to care about integer overflow in case the indices are in int32? Maybe always use int64 for dense_index?
| dense_index = dense_index * input_shape[i] + input_index[i]; | ||
| } | ||
| // Compute output index from fastest- to slowest-changing dimension. | ||
| for (int i = output_rank; i-- > 0;) { |
There was a problem hiding this comment.
Please use the more idiomatic i >= 0; i--. :)
| auto config = GetGpuLaunchConfig(nnz, device); | ||
| return GpuLaunchKernel(ReshapeSparseTensorKernel<int64>, config.block_count, | ||
| config.thread_per_block, 0, device.stream(), nnz, | ||
| /* input_rank = */ input_rank, |
There was a problem hiding this comment.
Please spell these as /*input_rank=*/; our internal tooling will then check that the param names match up.
- Remove use_gpu=True because it is already the default. - Use int64 for dense_index inside kernel to avoid integer overflow. - Change reverse for-loop style. - Reformat inline comments to ensure internal tooling picks them up.
sanjoy
approved these changes
Feb 25, 2021
google-ml-butler
bot
added
kokoro:force-run
gbaned
added
ready to pull
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This follows #46275.
cc @nluehr