Commit Graph

31 Commits

Author SHA1 Message Date
Andy Ly 4c8fead3e0 Add support for token operands to mhlo.tuple.
mhlo.get_tuple_element supports extracting a mhlo.token type from a tuple. This updates the creation of tuples to allow for mhlo.token typed operands.

PiperOrigin-RevId: 324628663
2020-08-03 19:31:15 -07:00
Smit Hinsu 577a81a66d Sink standard dialect constants in sink_constants_to_control_flow pass
This is required before exporting HLO dialect ops with standard dialect constant to XLA.

Also, sink constants for sort op as well. Added a TODO to generalize this pass to handle more ops and non-const values defined outside.

PiperOrigin-RevId: 324301911
2020-08-03 19:30:29 -07:00
Stephan Herhut 9cbe5f2285 Constrain mhlo.const to static shaped tensors.
Constants of unknown shape cannot be materialized. In most cases, one likely wants to use a scalar constant and rely on broadcasting instead.

PiperOrigin-RevId: 324252475
2020-08-03 19:29:51 -07:00
Thomas Joerg 735ae2838c Integrate LLVM at llvm/llvm-project@cd4e8d7f6f
Updates LLVM usage to match
[cd4e8d7f6f5e](https://github.com/llvm/llvm-project/commit/cd4e8d7f6f5e)

PiperOrigin-RevId: 324173542
2020-08-03 19:29:29 -07:00
Tres Popp cce4bddf4b Remove unnecessary conversions between Shape and ExtentTensor.
PiperOrigin-RevId: 323981215
2020-08-03 19:29:04 -07:00
Tres Popp ffef8d6593 Support CHLO->LHLO lowering for broadcasting operations with both inputs unranked.
PiperOrigin-RevId: 323960733
2020-08-03 19:28:52 -07:00
Stephan Herhut b09bf2a4dc Move hlo end to end tests to the hlo directory tree.
PiperOrigin-RevId: 323955773
2020-08-03 19:28:37 -07:00
Stephan Herhut 1b0eb4baa7 Do not mandate the result type of shape computations but have it be inferred from context.
The computation of a broadcasted shape forced the use of the shape type unnecessarily, which blocked further canonicalizations.

PiperOrigin-RevId: 323783998
2020-08-03 19:28:12 -07:00
Benjamin Kramer b7c4314e7f Allow running tests with FILECHECK_OPTS=-enable-var-scope
PiperOrigin-RevId: 323420636
2020-08-03 19:27:38 -07:00
Thomas Joerg 739758f9cc Integrate LLVM at llvm/llvm-project@eed333149d
Updates LLVM usage to match
[eed333149d17](https://github.com/llvm/llvm-project/commit/eed333149d17)

PiperOrigin-RevId: 323354988
2020-08-03 19:27:25 -07:00
Robert Suderman 8023baa959 Modified HLOAbsOp lowering for differing types.
PiperOrigin-RevId: 323082107
2020-08-03 19:27:12 -07:00
Hanhan Wang 8f262ae8f5 Add support for lowering mhlo.iota to Linalg.
PiperOrigin-RevId: 322799853
2020-07-30 22:34:45 +00:00
Tres Popp 4251630426 Support CHLO broadcasting operations between scalar and unranked tensors.
This is done through reshaping the unranked tensor into a 1D ranked tensor which will result in a safe broadcast/indexing logic when the other operand is a scalar.

PiperOrigin-RevId: 322553661
2020-07-30 22:34:40 +00:00
Robert Suderman c23ad602c8 Add a transform for Gathers to torch_index_select.
Some gathers can be interpreted as torch index selects. Transforming these
cases allow torch_index_select lowerings to be used for certain gathers.

PiperOrigin-RevId: 322255835
2020-07-30 22:34:32 +00:00
Stephan Herhut c44e08351d Fix mhlo to lmhlo conversion for ReduceOp.
The existing conversion no longer worked and was not save to undo. Furthermore, the pattern for mhlo.return had been removed.

Also adds some tests to ensure this does not degrade again.

PiperOrigin-RevId: 321542071
2020-07-30 22:34:19 +00:00
A. Unique TensorFlower c8bb0ff54d Integrate LLVM at https://github.com/llvm/llvm-project/commit/f233b92f92a6
PiperOrigin-RevId: 321454533
2020-07-30 22:34:14 +00:00
Robert Suderman 98a1e3b108 Add an optimization that converts some Gathers to Slices.
Some Gathers can be represented as slices. This lowering transforms
these gathers into slices.

PiperOrigin-RevId: 321394868
2020-07-30 22:34:10 +00:00
Stephan Herhut 7a6adc6a84 Add canonicalization patterns for dynamic_broadcast_in_dim where the target shape is the shape of the operand.
PiperOrigin-RevId: 321312182
2020-07-30 22:34:06 +00:00
Uday Bondhugula d166b66cba PR #40925: [MLIR] Update lhlo.const to linalg lowering to use affine.store inste…
Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/40925

…ad of std.store

The xla_lhlo.const lowering uses std.store to store a constant to
0-d memrefs. Update it to affine.store since such an access is trivially
affine (no indices). An affine.store can always be lowered to std.store.
Copybara import of the project:

--
9e18ede72fbbca107177bd742921e4cbf77adc82 by Uday Bondhugula <uday@polymagelabs.com>:

[MLIR] Update lhlo.const to linalg lowering to use affine.store instead of std.store

The xla_lhlo.const lowering uses std.store to store a constant to
0-d memrefs. Update it to affine.store since such an access is trivially
affine (no indices). An affine.store can always be lowered to std.store.

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/40925 from polymage-labs:lhlo_to_linalg_affine_store 9e18ede72fbbca107177bd742921e4cbf77adc82
PiperOrigin-RevId: 320623152
2020-07-30 22:33:51 +00:00
Robert Suderman 06ae59074f Fold xla iota across a 1-length dimension into a zero value
Iota across length-1 is just a constant. Fold into it.

PiperOrigin-RevId: 320443468
2020-07-30 22:33:43 +00:00
Stephan Herhut c3be2474dd Integrate LLVM at https://github.com/llvm/llvm-project/commit/305b500eaf8c
PiperOrigin-RevId: 320395881
2020-07-30 22:33:38 +00:00
Uday Bondhugula de0578b4f9 PR #40745: [MLIR] Add constant folder for xla_hlo.broadcast_in_dim op
Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/40745

Fold broadcast_in_dim op if the operand is the result of a tensor splat.
Copybara import of the project:

--
26c9f631448b8d6ffd20ece39ea8d4132b5550c7 by Uday Bondhugula <uday@polymagelabs.com>:

[MLIR] Add constant folder for xla_hlo.broadcast_in_dim op

Fold broadcast_in_dim op if the operand is the result of a tensor
splat.

COPYBARA_INTEGRATE_REVIEW=https://github.com/tensorflow/tensorflow/pull/40745 from polymage-labs:broadcast_in_dim_fold 26c9f631448b8d6ffd20ece39ea8d4132b5550c7
PiperOrigin-RevId: 320365164
2020-07-30 22:33:34 +00:00
Mehdi Amini 506ddd9c4a Cleanup build rule names in compiler/mlir/hlo to remove the redundant/obsolete xla_ prefix
PiperOrigin-RevId: 320320140
2020-07-30 22:33:29 +00:00
Mehdi Amini 94dcb90d38 Rename xla_chlo dialect into chlo
Following on the plan of isolating the compiler/mlir/hlo directory.

PiperOrigin-RevId: 320212018
2020-07-30 22:33:16 +00:00
Mehdi Amini 7c4a5d62b5 Rename xla_lhlo dialect into lmhlo
Following on the plan of isolating the compiler/mlir/hlo directory.
Another xla_lhlo dialect will be created under compiler/mlir/xla/ later.

PiperOrigin-RevId: 320210326
2020-07-30 22:33:11 +00:00
Alexander Belyaev b076e018a8 [MLIR][LHLO] Legalize CallOp that call funcs with tensor args/results.
PiperOrigin-RevId: 320172723
2020-07-30 22:33:07 +00:00
Alexander Belyaev e8cfdee592 [MLIR][LHLO] Convert mhlo.dynamic_reshape -> lhlo.reshape_memref_cast.
PiperOrigin-RevId: 320149593
2020-07-30 22:33:02 +00:00
Alexander Belyaev 8692fde3f9 [MLIR] Convert FuncOp signature with unranked types in HLO->LHLO conversion.
PiperOrigin-RevId: 320146856
2020-07-30 22:32:58 +00:00
Robert Suderman e1651b6090 Canonicalize multidimensional iota to use broadcast
There is no reason to have a multidimensional iota for codegen.
This should be canonicalized to a single dimensional iota followed
by a broadcast. Changing iota to on a single dimension  and a broadcast
substantially simplifies implementing iota operations.

PiperOrigin-RevId: 320095470
2020-07-30 22:32:54 +00:00
Mehdi Amini 8900222fed Rename `xla_hlo` dialect to `mhlo`
This is part of the current refactoring of the HLO related dialect.
`xla_hlo` will be reintroduced in a new form later.

PiperOrigin-RevId: 319916753
2020-07-30 22:32:50 +00:00
Mehdi Amini fa057cc0bc Move HLO tests to `third_party/tensorflow/compiler/mlir/hlo/`
Also add a localized `mlir-hlo-opt` binary for the testing of
tensorflow/compiler/mlir/hlo/... ; this directory is intended to be self-contained
and depend only on MLIR.

PiperOrigin-RevId: 319878984
2020-07-30 22:32:45 +00:00