[NPU] Support NPU kernel for TopKV2 op #34599
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
|
|
|
✅ This PR's description meets the template requirements! |
|
Thanks for your contribution! |
zhiqiu
reviewed
Aug 10, 2021
| @@ -0,0 +1,250 @@ | |||
| /* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved. | |||
| template <typename T> | ||
| class TopkV2NPUKernel : public framework::OpKernel<T> { | ||
| public: | ||
| // Use Ascend TopKV2 operator to implement paddle TopKV2Op |
There was a problem hiding this comment.
I think this comment is not literally necessary since it is easy for the reader to understand.
| public: | ||
| // Use Ascend TopKV2 operator to implement paddle TopKV2Op | ||
| void Compute(const framework::ExecutionContext& context) const override { | ||
| // Read message from context |
| indices->mutable_data<int64_t>(context.GetPlace()); | ||
|
|
||
| // Allocate space for output indices of Ascend topkV2 operator | ||
| framework::Tensor* indices_int32 = |
| new Tensor(framework::proto::VarType::INT32); | ||
| indices_int32->Resize(output_dims); | ||
| indices_int32->mutable_data<int32_t>(context.GetPlace()); | ||
| VLOG(4) << "input:" << *input; |
There was a problem hiding this comment.
LOG the tensor may cost much time.
| .AddAttr("dim", axis) | ||
| .AddAttr("largest", largest) | ||
| .Run(npu_stream_topkv2); | ||
| VLOG(4) << "output:" << *out; |
Comment on lines
89
to
91
| auto npu_stream_cast = | ||
| context.template device_context<paddle::platform::NPUDeviceContext>() | ||
| .stream(); |
There was a problem hiding this comment.
It is exactly the same stream with npu_stream_topkv2
… add_npu_op_tok_k_v2
…into add_npu_op_tok_k_v2
…into add_npu_op_tok_k_v2
…into add_npu_op_tok_k_v2
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.
PR types
New features
PR changes
OPs
Describe
Develop the NPU kernel for TopKV2 op and reuse the Python unit test for CPU kernel to test the new kernel.
Unit test result
NPU op call result
