PiperOrigin-RevId: 554932893

It's been over three months since the new PJRT C API client was
enabled by default
(https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-8-march-29-2023).

PiperOrigin-RevId: 554935166

PiperOrigin-RevId: 554937981

Since all the subfolders in tensorflow/core/lib have BUILD files, no files are picked up by this glob.

PiperOrigin-RevId: 554943490

These tests seem to use different number of GPUs across multiple unit test cases. This causes issues with maintaining the PjRtClient and DeviceCompiler lifetimes.

PiperOrigin-RevId: 554953351

PiperOrigin-RevId: 554958089

PiperOrigin-RevId: 554958217

PiperOrigin-RevId: 554959602

PiperOrigin-RevId: 554965356

PiperOrigin-RevId: 554967781

Only move TpuCopyWithDynamicShape out of launches if nothing in the
launch depends on them.

PiperOrigin-RevId: 554971170

PiperOrigin-RevId: 554972724

PiperOrigin-RevId: 554973779

PiperOrigin-RevId: 554976689

… not empty.

PiperOrigin-RevId: 554976987

PiperOrigin-RevId: 554977992

…to `WARNING`.

PiperOrigin-RevId: 554982965

… HIDDEN.

Add tf_export'ed raw_ops to modules_with_exports.

PiperOrigin-RevId: 554985433

PiperOrigin-RevId: 555002989

…n MLIR & Flatbuffers

PiperOrigin-RevId: 555004085

…able_xla_auto_cluster_for_pjrt_device

PiperOrigin-RevId: 5550338

- psimd was used by XNNPack, but is no longer needed

PiperOrigin-RevId: 555038715

PiperOrigin-RevId: 555043552

Added a test case for checking the behaviour wrt numpy for the changes done in tf.experimental.numpy.vander.

Resolved indentation error in vander function

For now, we just disable the in-place DUS/memcpy detection logic.

PiperOrigin-RevId: 555067189

PiperOrigin-RevId: 555069429

PiperOrigin-RevId: 555077671

PiperOrigin-RevId: 555091363

PiperOrigin-RevId: 555091467

This is a small efficiency gain, as we can potentially already detect a
difference before calling Identical().

PiperOrigin-RevId: 555093049

Update comments, rename variables.

PiperOrigin-RevId: 555093353

The logic assumed that we just need to check the opcode, but for several ops we
also need to check additional attributes. IdenticalSlowPath() is the method to
check this. Expose this through another SameOp method in HloInstruction.

PiperOrigin-RevId: 555094486

Fixes crash when using unsupported F8 types.

PiperOrigin-RevId: 555116700

Change the data structure describing tensors in the rewriter and its handling to make it more flexible and support more kinds transformations.

PiperOrigin-RevId: 555125358

Committing again with a fix.

PiperOrigin-RevId: 555134874

PiperOrigin-RevId: 555138718

Now contracting dimensions with split-K are represented as two fragments of one dimension instead of two separate dimensions.

PiperOrigin-RevId: 555139630

PiperOrigin-RevId: 555145026

PiperOrigin-RevId: 555154633

…nd layouts. This was triggering for semantically correct ReduceScatter HLO.

PiperOrigin-RevId: 555169308

PiperOrigin-RevId: 555172036

PiperOrigin-RevId: 555177148

PiperOrigin-RevId: 555184894

… have strict dependencies.

PiperOrigin-RevId: 555194981

PiperOrigin-RevId: 555202253

This change include modifying 3 passes: ConvertTFQuantTypes, ConvertTFQuantOpsToMHLO and ConvertMHLOQuantToInt.

* In ConvertTFQuantTypes pass, we add qint <-> int tf.cast ops to surround TF UQ ops. We added some special op legality check in ConversionTarget to make sure the resulting TF UQ and tf.cast ops are legal.
* In ConvertTFQuantOpsToMHLO pass, we add uq <-> int mhlo.convert ops to surround TF UQ ops and convert qint <-> int tf.cast ops to int <-> int mhlo.convert ops.
* In ConvertMHLOQuantToInt pass, we convert uq <-> int mhlo.convert ops to int <-> int mhlo.convert ops

The above passes may introduce int <-> int mhlo.convert ops. They are no-ops and can be removed in a Canonicalizer pass afterwards.

PiperOrigin-RevId: 555206832

…ency

hiding scheduler for the GPU.

PiperOrigin-RevId: 555212403

PiperOrigin-RevId: 555223964

…t examines elements along the diagonal only.

PiperOrigin-RevId: 555226962

Imported from GitHub PR openxla/xla#4864

amdgpu_compiler build error due to openxla/xla@6a68378

@akuegel @ddunl Thanks in advance!
Copybara import of the project:

--
2479f31ce053a2335e532772afb47d250fe3739f by Chao Chen <cchen104@amd.com>:

add thread_pool to amdgpu_compiler

Merging this change closes tensorflow#4864

PiperOrigin-RevId: 555227394

…ocinfo.

PiperOrigin-RevId: 555232216

PiperOrigin-RevId: 555232359

…tch-1

PiperOrigin-RevId: 555233600

…ater

PiperOrigin-RevId: 555235182

… remove custom-call to sharding.

PiperOrigin-RevId: 555242148

… directory.

PiperOrigin-RevId: 555245780

- Assume that in a multi-controller setting, XLA/GPU should see identical HLOs
  when multiple HLOs are compiled and communicate through collectives.
- Given that, restrict CollectiveScheduleLinearizer pass to only handle cases
  where compiler auto tuning can case divergence in schedules across different
  compilations of the same HLO.

PiperOrigin-RevId: 555248173

PiperOrigin-RevId: 555249702

PiperOrigin-RevId: 555255267

This makes the MLIR side effect modelling of the op consistent with ACD, and improves run-time performance of some GPU graphs with these ops.

PiperOrigin-RevId: 555261171

Added missing xnnpack_plugin.cc. This PR resolves tensorflow#60261 issue.

PiperOrigin-RevId: 555266459

The distutils module has been removed under Python 3.12, and the code here fails with a ModuleNotFoundError. This occurs because:

```
$ ipython
Python 3.12.0rc1 (main, Aug  7 2023, 13:07:47) [GCC 12.2.0]
Type 'copyright', 'credits' or 'license' for more information
IPython 8.14.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: import importlib.util

In [2]: importlib.util.find_spec('distutils.sysconfig')
---------------------------------------------------------------------------
ModuleNotFoundError                       Traceback (most recent call last)
Cell In[2], line 1
----> 1 importlib.util.find_spec('distutils.sysconfig')

File <frozen importlib.util>:91, in find_spec(name, package)

ModuleNotFoundError: No module named 'distutils'
```

However, as far as I can tell, we can just unconditionally use sysconfig since Python 3.2, which is well below the minimum Python version requirement of this code.

This change effectively reverts: tensorflow#54292
That change will no longer be tenable once Python 3.12 is released.

PiperOrigin-RevId: 555273752

PiperOrigin-RevId: 555274419

…rgets.

PiperOrigin-RevId: 555277315

PiperOrigin-RevId: 555277542

…le element in the parameter instructions that have no users.

PiperOrigin-RevId: 555286907

…e ordinal

as an input.

PiperOrigin-RevId: 555289711

PiperOrigin-RevId: 555292150

PiperOrigin-RevId: 555296787

PiperOrigin-RevId: 555298522

PiperOrigin-RevId: 555302357

… and maximum op

PiperOrigin-RevId: 555305030

PiperOrigin-RevId: 555305766

If the device flag is specified, then the cache will only be used for that specific device. By default, this change does impact how the compilation cache works.

This flag is being added to prevent errors from occurring when running on hardware with multiple device types.

PiperOrigin-RevId: 555311811

This change adds `xla_client.DeviceList` that is implemented in C++
`jax::PyDeviceList`. `jax::PyDeviceList` implements the features of
`pxla._DeviceAssignment` as a functional drop-in replacement.
`jax::PyDeviceList` internally has `xla::ifrt::DeviceList`, which will be used
when using IFRT APIs without having to construct a new copy of a potentially
large device list.

`pxla._DeviceAssignment`'s interface is changed slightly to encourage avoiding
conversion to tuple.

Note that for the backward compatibility (and fast `xla_client.Device`
conversion), `jax::PyDeviceList` still uses a Python tuple whose element can be
any Python object matches `xla_client.Device` interface with duck typing. This
duck typing support will be removed when such use case is deprecated.
Eventually, we can try to avoid any type conversion to remove a shadow copy of
device list in JAX.

PiperOrigin-RevId: 555317152

…cluded instruction with a constant broadcasted from a random number.

PiperOrigin-RevId: 555317504

The parsed key solution won't have a correct buf_ and could break the object lifespan and ParsedKey::FullKey()

PiperOrigin-RevId: 555327807

Updates LLVM usage to match
[6556e2902570](llvm/llvm-project@6556e2902570)

PiperOrigin-RevId: 555348867

http://github.com/tensorflow/runtime/commit/8293af247eb4703428bdf857dfcb846205b712cf.

PiperOrigin-RevId: 555363067

PiperOrigin-RevId: 555370639

PiperOrigin-RevId: 555379542

PiperOrigin-RevId: 555392263

…nected` when the filter is per-axis quantized.

This change adds a new pattern in `UniformQuantizedStablehloToTflPass` that converts quantized `stablehlo.dot_general` to `tfl.fully_connected`. This pattern is applied when the filter is per-axis quantized.

PiperOrigin-RevId: 555405376

…A graph integration

Similar to other Gpu libraries move the platform specific implementation to GpuDriver.

PiperOrigin-RevId: 555408360

Derive block pointer parameters separately for each input tensor so that one scope of a dot can access tensors with different shapes and layouts. This is ~NFC for now, it is required to enable more fusions but they will only come with other updates to the rewriter and the analysis.

PiperOrigin-RevId: 555414634

PiperOrigin-RevId: 555416852

PiperOrigin-RevId: 555417053

Also fix IsInputFusibleReduction(), it was using HasRealReductionHero()
instead of IsReductionFromOrToContiguousDimensions()

PiperOrigin-RevId: 555458297

PiperOrigin-RevId: 555458586

PiperOrigin-RevId: 555461404

… POD and string type of

tensor match on list of tensors.

PiperOrigin-RevId: 555508096

Adds detection of partial replicated input compatible with fully sharded output
for contracting dimension matching case.

PiperOrigin-RevId: 555512892

Updates LLVM usage to match
[6448d5ba581a](llvm/llvm-project@6448d5ba581a)

PiperOrigin-RevId: 555515803