[MLIR][XLA] Add GatherOp to LHLO/HLO emitters #40578
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
|
@timshen91 can you help providing guidance about where we want the HLO->LHLO conversion be handled right now? |
pifon2a
reviewed
Jun 19, 2020
|
|
||
| // CHECK: func @gather(%[[ARG0:.*]]: [[TYPE0:.*]], %[[ARG1:.*]]: [[TYPE1:.*]], %[[RESULT:.*]]: [[RTYPE:.*]]) { | ||
| // CHECK: "xla_lhlo.gather"(%[[ARG0]], %[[ARG1]], %[[RESULT]]) | ||
| // CHECK: {collapsed_slice_dims = dense<0> : tensor<1xi64>, index_vector_dim = 2 : i64, offset_dims = dense<2> : tensor<1xi64>, slice_sizes = dense<[1, 10]> : tensor<2xi64>, start_index_map = dense<0> : tensor<1xi64>} : ([[TYPE0]], [[TYPE1]], [[RTYPE]]) -> () |
There was a problem hiding this comment.
nit: please use CHECK-SAME to wrap the line
google-ml-butler
bot
added
kokoro:force-run
joker-eph
reviewed
Jun 19, 2020
| input.offset_dims().end()); | ||
| std::vector<int64> gather_collapsed_slice_dims( | ||
| input.collapsed_slice_dims().begin(), input.collapsed_slice_dims().end()); | ||
| std::vector<int64> gather_start_index_map(input.start_index_map().begin(), |
There was a problem hiding this comment.
Nit: use SmallVector instead. I assume that these attributes aren't very large in general.
(see LhloDialectEmitter::HandleReduceWindow for example)
|
@xinan-jiang Can you please check @joker-eph's comments and keep us posted. Thanks! |
|
Unassigning myself since @pifon2a is more familiar with this code. |
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 is a PR from JIZHI, the AI platform in Tencent.
@sherhut @pifon2a
We work on TensorFlow/MLIR to make mlir_gpu enable. Could you tell me how to run the test? It seems need a new tool.