When there are many options to a function, it's cleaner to use an options struct. This change is in preparation for adding another option.
Since ExecutableBuildOptions already has a DeviceAssignment field, we can use it rather than having our own. We can also sink the DeviceAssignment legality check into XLA proper, given it's now just testing that two parts of ExecutableBuildOptions agree.
PiperOrigin-RevId: 300407020
Change-Id: Idcd0acae6bb89eda34085ae2471b0eafd5cdef95
The XRT tuple allocation owns the device memory, which, in order for the lower level
aliasing to work, needs to be handed out as "owning" within the parameter's shape tree.
But if the parameter's shape tree get destroyed (for an intermediate error before execute)
the memory will get released and the tuple allocation will be pointing to free memory.
This CL introduces an ExecutionInput data structure which wraps a maybe-owning shape tree
together with the indices which should be released before the shape tree gets destroyed.
This allows the data structure to travel down until the point where the buffers lands
inside the ExecutionOutput, which uses a similar logic (until finally the result gets consumed).
Unfortunately the situation of the device memory data structures got a bit messy, with
Owning, MaybeOwning, ShapedBuffer, ScopedShapedBuffer, ... none of which can work nicely
with buffer sharing.
Ideally we should have something like std::shared_ptr<OwningDeviceMemory> and
ShapeTree<std::shared_ptr<OwningDeviceMemory>> and be done with it.
Unfortunately the change (I started that route first) towards that goal is pretty major.
PiperOrigin-RevId: 298498866
Change-Id: I2e27c11b7187fa2992ae3b606ea95c18f312cb5a
Move logic from ExecuteOnStreamWrapper, make it callback-based, and make ExecuteOnStreamWrapper a thin wrapper around ExecuteAsyncOnStreamWrapper than blocks.
Helps solve https://github.com/google/jax/issues/774 by making the XLA profiling code work on the async path used by JAX.
PiperOrigin-RevId: 262416495
Remove ExecuteOnStream virtual method, make ExecuteOnStream a non-virtual wrapper around ExecuteAsyncOnStream.
This means that backend authors have one method to implement (ExecuteAsyncOnStream) rather than two, and reduces the number of code paths to running an executable.
Comment that ExecuteAsyncOnStream may in fact not be async. While undesirable, this is a quality of implementation issue not a bug. Future changes can make implementations of ExecuteAsyncOnStream truly async.
PiperOrigin-RevId: 261922907
Currently a mutable execution_profile is attached to xla::Executable. This isn't thread safe, since the same Executable may be invoked concurrently. Instead, clients already have the ability to pass in their own ExecutionProfile via the ExecutableRunOptions; update that one instead.
Simplify APIs that accepted both an ExecutionProfile and an ExecutableRunOptions given the latter contains a pointer to an instance of the former.
PiperOrigin-RevId: 261570194
Currently the code keeps a mutable HloSnapshot attached to the xla::Executable object. This cannot work correctly in the presence of concurrent executions. Instead, keep only an immutable HloProto attached to xla::Executable and construct ephemeral HloSnapshots during dumping.
This has the minor downside that it requires copying the HloProto each time we dump, but presumably if you are dumping HLO snapshots you don't particularly care about performance.
PiperOrigin-RevId: 259934176
This change achieves three goals:
1. There are currently three different allocator abstractions in three different
places: XLA, stream executor, and tensorflow.
This change shrinks down the number of packages with allocator abstraction to
two.
2. Moving the allocator enables unifying ScopedDeviceMemory
and OwningDeviceMemory which both have "owning pointer" semantics,
but a slightly different API.
3. Moving the allocator enables moving RedzoneAllocator in stream executor,
which we would like to use in tensorflow to catch out-of-bound-writes in
CUDNN convolutions during the autotuning.
PiperOrigin-RevId: 247211996
[XLA:Python] Add support for asynchronous execution in the Python client.
Python isn't famous for being the world's fastest language, so for high performance eager-style dispatch it is helpful to be able to hide Python latency behind device computations by having the Python code dispatch device operations asynchronously.
The design here closely follows the design of asynchronous execution in TensorFlow and the TensorFlow/XLA client. We use three main streams:
* a compute stream, for running XLA computations,
* a host-to-device stream, for transferring data onto the device
* a device-to-host stream, for transferring data off the device.
Both host-to-device transfers and compute are asynchronous, that is, they return control to Python as soon as any necessary error checking is complete, but before the operation completes. This allows the Python code to enqueue any subsequent operations while the previously enqueued operations complete. Device-to-host transfers are still blocking, in the sense that they stall the host until the host-side data is ready.
[XLA] Add LocalExecutable::RunAsync() to obtain async execution on a stream. There is currently no way to achieve this via the LocalClient API, only by using internal XLA APIs.
[XLA:GPU] Implement ExecuteAsyncOnStream. It turns out that ExecuteOnStream is already more or less async anyway.
PiperOrigin-RevId: 246650968
The following flags (usually passed via the XLA_FLAGS envvar) are removed:
xla_dump_computations_to
xla_dump_executions_to
xla_dump_ir_to
xla_dump_optimized_hlo_proto_to
xla_dump_per_pass_hlo_proto_to
xla_dump_unoptimized_hlo_proto_to
xla_generate_hlo_graph
xla_generate_hlo_text_to
xla_hlo_dump_as_html
xla_hlo_graph_path
xla_log_hlo_text
The following new flags are added:
xla_dump_to
xla_dump_hlo_module_re
xla_dump_hlo_pass_re
xla_dump_hlo_as_text
xla_dump_hlo_as_proto
xla_dump_hlo_as_dot
xla_dump_hlo_as_url
xla_dump_hlo_as_html
xla_dump_ir
xla_dump_hlo_snapshots
The default is not to dump anything at all, but as soon as some dumping flag is
specified, we enable the following defaults (most of which can be overridden).
* dump to stdout (overridden by --xla_dump_to)
* dump HLO modules at the very beginning and end of the optimization pipeline
* don't dump between any HLO passes (overridden by --xla_dump_hlo_pass_re)
* dump all HLO modules (overridden by --xla_dump_hlo_module_re)
* dump in textual format (overridden by
--xla_dump_hlo_as_{text,proto,dot,url,html}).
For example, to dump optimized and unoptimized HLO text and protos to /tmp/foo,
pass
--xla_dump_to=/tmp/foo --xla_dump_hlo_as_text --xla_dump_hlo_as_proto
For details on these flags' meanings, see xla.proto.
The intent of this change is to make dumping both simpler to use and more
powerful.
For example:
* Previously there was no way to dump the HLO module during the pass pipeline
in HLO text format; the only option was --dump_per_pass_hlo_proto_to, which
dumped in proto format.
Now this is --xla_dump_pass_re=.* --xla_dump_hlo_as_text. (In fact, the
second flag is not necessary in this case, as dumping as text is the
default.)
* Previously there was no way to dump HLO as a graph before and after
compilation; the only option was --xla_generate_hlo_graph, which would dump
before/after every pass.
Now this is --xla_dump_hlo_as_{dot,url,html} (depending on what format you
want the graph in).
* Previously, there was no coordination between the filenames written by the
various flags, so info about one module might be dumped with various
filename prefixes. Now the filenames are consistent and all dumps from a
particular module are next to each other.
If you only specify some of these flags, we try to figure out what you wanted.
For example:
* --xla_dump_to implies --xla_dump_hlo_as_text unless you specify some
other --xla_dump_as_* flag.
* --xla_dump_hlo_as_text or --xla_dump_ir implies dumping to stdout unless you
specify a different --xla_dump_to directory. You can explicitly dump to
stdout with --xla_dump_to=-.
As part of this change, I simplified the debugging code in the HLO passes for
dumping HLO modules. Previously, many tests explicitly VLOG'ed the HLO module
before, after, and sometimes during the pass. I removed these VLOGs. If you
want dumps before/during/after an HLO pass, use --xla_dump_pass_re=<pass_name>.
PiperOrigin-RevId: 237920279
Unlike Printf, StrFormat does not require type-length qualifiers, e.g
%z, %ll. Nor does it require that you call c_str() to print strings.
So these are fixed up here as well.
PiperOrigin-RevId: 210435915
built for is equivalent to the device the it will run on.
Before this patch, if the device to run on was provided via a stream without
setting the device ordinal in the ExecutableRunOptions, we would check the
default device against the device the executable was built for.
PiperOrigin-RevId: 206892902
We have a Pool in XLA that maintains a freelist of Streams, to avoid
the overhead of repeatedly allocating new Streams. Streams have a
monotonic state machine; if a stream encounters any error, it will
remain in an error state forever.
The functional change in this CL is to ensure that streams which have
encountered an error are deleted, rather than being put back on the
pool. Without this change, a previously failed stream will be put back
on the pool, only to cause the next usage of the stream to trivially
fail.
I've chosen to replace the generic templatized Pool with a concrete
StreamPool, since this makes the logic more straightforward to reason
about. Also note that the only existing usage of Pool is to hold
streams.
The functional change is in stream_pool.cc; most of everything else is
mechanical updates.
PiperOrigin-RevId: 206100631
This lets replay_computation build an executable once and run it
multiple times. This is particularly important because in XLA:GPU, the
first run of an executable does some autotuning and therefore is
unrepresentative.
This change removes --xla_hlo_profile_last_run, because I don't see how
to support it in LocalClient -- LocalClient wants the do-profile bit to
be set when we *compile*. (There may not be an easy fix for this; it
worked with regular Client because we were recompiling every time we
ran.)
PiperOrigin-RevId: 198643577
Previously, only one layout was stored with an HLO module. This CL allows
HLO passes to modify the on-device layouts without affecting the on-host
layout (provided by the client)
PiperOrigin-RevId: 195014875
Previously, we returned a plain ShapedBuffer. But this doesn't capture
our semantics: It's up to the callee to free this ShapedBuffer.
PiperOrigin-RevId: 193854051
These are already notionally equivalent to T* and unique_ptr<T>, so
having a unique_ptr of a {Scoped,}ShapedBuffer is pretty redundant.
Also clean up the ScopedShapedBuffer API a bit.
PiperOrigin-RevId: 193599773
