Repository navigation
[pull] master from tensorflow:master - #8847
Merged
Merged
Conversation
…uments and typed attributes
Writing a native custom-call handler requires hand-rolling a large amount of
boilerplate that XLA is in a better position to own, and several of those
hand-rolled pieces are easy to get subtly wrong. This change moves that work
into `NativeCustomCallEmitterContext`, which gains:
- `GetDeviceDescription()`. Handlers previously had to dig the device
description out of `GetTargetTopology()->gpu_target_config()`, guarding
against a target config that, in this code path, is always present.
- `GetResultShapedSlice()` / `GetOperandShapedSlice()`, which return the
`ShapedSlice` that thunks actually consume, rather than a bare slice the
handler then has to pair up with a shape by hand.
- `CreateKernelArguments()`, which builds the conventional kernel argument
list. Assembling `emitters::KernelArgument`s by hand leaves `slice_index_`
uninitialized, skips alignment and aliasing computation, and requires the
author to set `written` correctly - and `written` drives thunk scheduling,
so getting it wrong produces races rather than an obvious failure.
The implementation resolves buffers through `ThunkEmitter`'s allocation
overrides, which a handler could not have done itself.
- `GetFfiAttributes()` now returns a typed `ffi::Attributes` instead of an
`AttributesMap`, so reading an attribute is `attrs.Get<float>("epsilon")`
rather than two layers of hand-written `std::get_if` on nested variants.
The demonstrator handler is updated for the new signature.
PiperOrigin-RevId: 985413447
This frequently shows up in x86 CPU logs on the critical path, but we don't actually test this on CI. PiperOrigin-RevId: 985426107
This file contains hidden or 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
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
3 participants
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.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )