Commit Graph

330 Commits

Author SHA1 Message Date
A. Unique TensorFlower 31536431e0 [MLIR][HLO] Eliminate duplicate broadcastable constraints
PiperOrigin-RevId: 376718433
2021-05-31 13:50:23 -07:00
A. Unique TensorFlower 0f341012c6 [MLIR][HLO] Eliminate duplicate broadcastable constraints
PiperOrigin-RevId: 376715240
2021-05-31 13:08:02 -07:00
A. Unique TensorFlower cc1b22e618 [HLO][Linalg] Support scalar broadcasts in point-wise converter
This is needed for operations that support this limited form of broadcasting,
namely `mhlo.select`.

PiperOrigin-RevId: 376655844
2021-05-31 03:50:23 -07:00
Hanhan Wang 402b74ed7f Fix type bug in mhlo.dynamic-update-slice lowering.
The operand type can be f32. We should not use operand type to do clamp
operations.

PiperOrigin-RevId: 376286524
2021-05-27 17:53:49 -07:00
Hanhan Wang 28c411606f Add support for lowering mhlo.dynamic-update-slice ops to Linalg and std ops.
PiperOrigin-RevId: 376042810
2021-05-26 15:31:05 -07:00
Robert Suderman 26a0053d7d Remove linalg.indexed_generic from mhlo lowerings to linalg
IndexedGeneric is going away. Transition to using linalg.Index instead.

PiperOrigin-RevId: 376002501
2021-05-26 12:24:23 -07:00
A. Unique TensorFlower 4ebcebf31c [MLIR][HLO] Exploit scalar properties in rank specialization lowering
Take advantage of the fact that scalars are already ranked and that they are
neutral elements to broadcasting. Do not reshape scalars, do not consider them
for broadcasting, and materialize ranked operations on scalars accordingly.

PiperOrigin-RevId: 375968371
2021-05-26 09:59:13 -07:00
Benjamin Kramer edf5ec8084 Integrate LLVM at llvm/llvm-project@cb65419b1a
Updates LLVM usage to match
[cb65419b1ac0](https://github.com/llvm/llvm-project/commit/cb65419b1ac0)

PiperOrigin-RevId: 375915516
2021-05-26 04:47:24 -07:00
A. Unique TensorFlower cb46298a07 [MLIR][HLO] Support all smaller ranks in rank specialization cases
Rank specialization cases can be applied to all argument tensors of smaller
ranks than the expected maximum rank. This is crucial if all operands are
effectively scalars and the maximum reduced rank is 0.

PiperOrigin-RevId: 375712020
2021-05-25 08:38:53 -07:00
Adrian Kuegel a847109ac7 Support complex types when converting HLO multiply op.
We can lower it to the MulOp in the complex dialect.

PiperOrigin-RevId: 375675079
2021-05-25 04:35:34 -07:00
Adrian Kuegel 5816920258 Support complex types when converting HLO divide op.
We can lower it to the DivOp in the complex dialect.
Also add tests to hlo-legalize-to-linalg.mlir for CompareOp lowering of complex
types. These were forgotten in a previous commit.

PiperOrigin-RevId: 375669125
2021-05-25 03:43:46 -07:00
Adrian Kuegel 758ae7da6b Support complex types when converting HLO compare op (EQ/NE).
We can lower it to the EqualOp / NotEqualOp in the complex dialect.

PiperOrigin-RevId: 375655092
2021-05-25 01:54:27 -07:00
wyzhao b93e54d8a4 PR #49454: [MLIR][DISC] Upgrade to use the new `reifyReturnTypeShapes` interface.
Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/49454

The new interface is more safe to be used during dialect conversion
(e.g. converting from tensor world to buffer world).
Copybara import of the project:

--
a6968072d59bec3c3bbaef0121d297e807c37c91 by Wenyi Zhao <reyizero@gmail.com>:

[MLIR][DISC] Upgrade to use the new `reifyReturnTypeShapes` interface.

The new interface is more safe to be used during dialect conversion
(e.g. converting from tensor world to buffer world).

--
55e7c6b7f2f99b99e226645a57e2433fae3e90ed by Wenyi Zhao <reyizero@gmail.com>:

minor fix

PiperOrigin-RevId: 375500273
2021-05-24 10:11:55 -07:00
Hanhan Wang 1ba4c714c9 Add support for lowering mhlo.scatter ops to Linalg.
This only works for updating tensors, not add/min/max computations. It requires
the index depth to be 1 because of the limitation in Linalg. We can not compare
multiple indices without packing indices.

PiperOrigin-RevId: 375137721
2021-05-21 12:17:14 -07:00
A. Unique TensorFlower 97e6103933 [MLIR][HLO] Reshape to scalars in rank specialization
Scalars were incorrectly casted to scalar tensors when they have to be reshaped.

PiperOrigin-RevId: 375049088
2021-05-21 03:12:16 -07:00
A. Unique TensorFlower 3daf65578a [MLIR][HLO] Add scalar cases for binary rank specialization
For rank specialization clusters that have only two operands, we can materialize
two extra cases in which either of them is a scalar. This avoids redundant index
computations in these cases.

PiperOrigin-RevId: 375037390
2021-05-21 01:35:44 -07:00
Hanhan Wang cd8f585cf7 [MHLO:Linalg] Add support for lowering torch_index_select of unsigned tensors
Also fixes typos in tests.

PiperOrigin-RevId: 374979460
2021-05-20 17:03:05 -07:00
Rahul Joshi 41f663ce47 [HLO] Adopt custom syntax for convolution dimensions and window attributes (HLO)
PiperOrigin-RevId: 374923250
2021-05-20 12:13:50 -07:00
Rahul Joshi fc88cf1ff4 [HLO] Adopt custom syntax for convolution dims and window attributes for LMHLO_GPU
PiperOrigin-RevId: 374889917
2021-05-20 09:41:48 -07:00
A. Unique TensorFlower c62fd89663 [MLIR][HLO] Add equal shapes case to rank specialization
Also restructure lowering implementation to facilitate the addition or removal
of special cases.

PiperOrigin-RevId: 374626365
2021-05-19 05:38:42 -07:00
Stella Laurenzo 71394fb301 Properly handle if DynamicBroadcastInDimOp shape is not of index type.
* The op defines this to be index, any integer, or pred (i1).
* Many TensorFlow legalizations produce integers for the shape.

PiperOrigin-RevId: 374566113
2021-05-18 21:12:11 -07:00
Stella Laurenzo 0fe07e3814 Separate CHLO transforms for expanding compositions and lowering broadcasts.
* The former is typically invariant regardless of backend.
* The latter may need to be done differently depending on capabilities of the lowering target.

PiperOrigin-RevId: 374492924
2021-05-18 13:33:59 -07:00
A. Unique TensorFlower 6af3d2df91 [MLIR][HLO] Add rank specialization with multiple non-scalar operands
Add lowering pattern for rank specialization clusters with more than one
non-scalar operand. The lowering resembles that of the `TransformUnrankedHlo`
pass and switches cases for maximal ranks from 1 through 8.

PiperOrigin-RevId: 374377002
2021-05-18 03:02:45 -07:00
A. Unique TensorFlower 474e419729 [MLIR][HLO] Generalize rank specialization with single operand
The pattern can be generalized to also rank specialize operations with a single
non-scalar operand. Also extract helper functions that can be reused in
following specializations.

PiperOrigin-RevId: 374198381
2021-05-17 08:12:55 -07:00
A. Unique TensorFlower c514c73390 [MLIR][HLO] Extend rank specialization clustering pass
Also cluster operations that operate on same shape operands. These implicitly
satisfy the broadcasting semantics requirement. Also, add test cases for some
cases that appear in the current MLIR-generated kernels.

PiperOrigin-RevId: 374191950
2021-05-17 07:31:36 -07:00
A. Unique TensorFlower ccd70d5717 [MLIR][HLO] Add `rank-specialization-to-scf` pass
Currently the lowering is only implemented for the unary case. The n-ary case
will follow.

PiperOrigin-RevId: 374162772
2021-05-17 03:56:23 -07:00
Rahul Joshi a361253e4f [HLO] Add custom print/parse for window attributes of convolutions (in LMHLO)
PiperOrigin-RevId: 373807616
2021-05-14 09:47:25 -07:00
A. Unique TensorFlower 76341f3720 [MLIR][HLO] Add mixed test for `rank-specialization-cluster` pass
PiperOrigin-RevId: 373762814
2021-05-14 04:40:40 -07:00
A. Unique TensorFlower d2cc74317c Implement constant folding for mhlo.Sign.
PiperOrigin-RevId: 373550014
2021-05-13 03:54:04 -07:00
Hanhan Wang d764806c1e [MHLO:Linalg] Add support for lowering reshape of unsigned tensors
PiperOrigin-RevId: 373461627
2021-05-12 15:14:29 -07:00
A. Unique TensorFlower 420c42a0a1 [MLIR][HLO] Support CHLO unary operations in rank specialization clustering
PiperOrigin-RevId: 373397321
2021-05-12 10:20:43 -07:00
A. Unique TensorFlower 596918a6f1 [MLIR][HLO] Allow rank specialization clustering with `chlo.broadcast_select` op
PiperOrigin-RevId: 373379990
2021-05-12 08:56:49 -07:00
Rahul Joshi e260aa771c [HLO] Add custom print/parse for convolution dimension numbers (in LMHLO)
PiperOrigin-RevId: 373379227
2021-05-12 08:52:46 -07:00
A. Unique TensorFlower 875803e5e1 [MLIR][HLO] Add more tests for `rank-specialization-cluster` pass
PiperOrigin-RevId: 373343750
2021-05-12 04:46:30 -07:00
A. Unique TensorFlower 313d24bc8f [MLIR][HLO] Add `rank-specialization-cluster` pass
Add a pass to cluster unranked C/HLO operations in one
`chlo.rank_specialization_cluster` op. The C/HLO operations are moved to the
body of the operation. Later passes can use this to rank-specialize all these
operations together.

PiperOrigin-RevId: 373336725
2021-05-12 03:46:01 -07:00
A. Unique TensorFlower 7f7a86ad0d [MLIR][HLO] Implement `RegionBranchOpInterface` for rank specialization cluster
PiperOrigin-RevId: 373163196
2021-05-11 09:03:05 -07:00
A. Unique TensorFlower 96a47345cc [MLIR][HLO] Add `rank_specialization_cluster` op to CHLO
The operation will be used to cluster compatible operations that can be rank-
specialized collectively.

PiperOrigin-RevId: 373128557
2021-05-11 05:17:42 -07:00
Benjamin Kramer 86b7eb434c [MHLO] Don't crash trying to constant fold mhlo.convert on complex
MLIR still doesn't have a complex attribute so this can't be implemented, so
just bail out instead of trying to fold.

PiperOrigin-RevId: 373128307
2021-05-11 05:15:57 -07:00
A. Unique TensorFlower 7f86dd9f7e Constant fold compare EQ if one of the operands is true and compare NE if one of the operands is false.
PiperOrigin-RevId: 373058030
2021-05-10 18:53:49 -07:00
A. Unique TensorFlower 2af1796194 Integrate LLVM at llvm/llvm-project@5c7b43aa82
Updates LLVM usage to match
[5c7b43aa8298](https://github.com/llvm/llvm-project/commit/5c7b43aa8298)

PiperOrigin-RevId: 373028739
2021-05-10 15:46:34 -07:00
Rahul Joshi ce4c76314e [NFC] Remove all_gather_dimension from all-to-all in the unit test
PiperOrigin-RevId: 372463706
2021-05-06 18:14:52 -07:00
Rahul Joshi 8c854886cb [XLA:GPU] Allow all-gather operands to have different element types.
- XLA's all-gather combiner can create such all-gathers, so relax the same element type
  trait for all-gathers.

PiperOrigin-RevId: 372380446
2021-05-06 11:04:13 -07:00
Adrian Kuegel b2bc17c8b0 Integrate LLVM at llvm/llvm-project@632ebc4ab4
Updates LLVM usage to match
[632ebc4ab437](https://github.com/llvm/llvm-project/commit/632ebc4ab437)

PiperOrigin-RevId: 372330771
2021-05-06 06:37:39 -07:00
A. Unique TensorFlower d8c40b691c [MLIR][HLO] Add `shape.broadcast` canonicalization to unblock broadcast moving
PiperOrigin-RevId: 372120309
2021-05-05 07:16:49 -07:00
Geoffrey Martin-Noble ac68145565 [MHLO:Linalg] Add support for lowering concat of unsigned tensors
Nothing about concat here really. Just need to plumb through the type
conversion.

PiperOrigin-RevId: 372012957
2021-05-04 15:57:54 -07:00
Geoffrey Martin-Noble 5a60793b31 [MHLO:Linalg] Add support for lowering dynamic-slice on unsigned ints
PiperOrigin-RevId: 371979004
2021-05-04 13:08:36 -07:00
Adrian Kuegel 384b87fad0 Lower ReluGrad via chlo::BroadcastSelect.
This allows to get rid of the constraint that it needs to have a static shape.

PiperOrigin-RevId: 371862452
2021-05-04 01:03:02 -07:00
Benjamin Kramer f4414fcd66 [MHLO:Linalg] Add support for lowering unsigned ops
This strips away the signedness with a type converter, using unrealized
conversion casts. The rest is mostly mechanically pushing the original op down
the pipeline so lowerings can see the original types.

Signed types stay signless for now. This can be changed in the HLO bridge later.

I did a pass over all ops and added unsigned lowerings where they were missing.
There may be more.

Currently the lowering will die at a later stage because it doesn't understand
the unrealized casts.

PiperOrigin-RevId: 371077494
2021-04-29 02:27:35 -07:00
A. Unique TensorFlower 4d41b11f3b Integrate LLVM at llvm/llvm-project@671f0e2e18
Updates LLVM usage to match
[671f0e2e189c](https://github.com/llvm/llvm-project/commit/671f0e2e189c)

PiperOrigin-RevId: 371011125
2021-04-28 16:37:53 -07:00
Benjamin Kramer b2a23bf269 Integrate LLVM at llvm/llvm-project@4b13b7581d
Updates LLVM usage to match
[4b13b7581db5](https://github.com/llvm/llvm-project/commit/4b13b7581db5)

PiperOrigin-RevId: 370736351
2021-04-27 12:19:05 -07:00