Commit Graph

347 Commits

Author SHA1 Message Date
A. Unique TensorFlower b580722041 [MLIR][KernelGen] Merge rank specialization clusters
Merge adjacent rank specialization clusters. Combine their operands, bodies, and
results.

PiperOrigin-RevId: 378433222
2021-06-09 10:07:47 -07:00
Adrian Kuegel b6d8160611 Add Broadcasting and BroadcastingElementwise traits to ConstantLikeOp.
This allows to include such ops in rank specialization clusters.

PiperOrigin-RevId: 378380915
2021-06-09 05:09:26 -07:00
A. Unique TensorFlower b9e45007d5 [MLIR][HLO] Extend broadcast propagation pass to enable more fusion
Move element-wise operations into assuming regions. This enables fusion
opportunities within the region.

PiperOrigin-RevId: 378362725
2021-06-09 03:03:37 -07:00
A. Unique TensorFlower d828b457b3 Handle empty tensors in SimplifyConcatSlice.
If the result of the slice is an empty tensor, do nothing.
This fixes a crash: we can't create a `concat` with an
empty operand range.

PiperOrigin-RevId: 378354956
2021-06-09 02:15:47 -07:00
Adrian Kuegel 9a8c254526 Support complex types for Sinh.
Because mhlo::ConstantLike doesn't support complex types, we need to use
GetScalarOfType and broadcast it to the needed shape.
Disable the tf2xla fallback, now that MLIR fully supports Sinh.

PiperOrigin-RevId: 378123151
2021-06-08 04:23:19 -07:00
A. Unique TensorFlower c47869f931 [MLIR][HLO] Rename `move-up-dynamic-broadcasts-for-fusion` to `broadcast-propagation`
PiperOrigin-RevId: 378102608
2021-06-08 01:51:10 -07:00
Benjamin Kramer d1c60df2fe [MHLO:linalg] Be more aggressive about turning mhlo.const into std.constant
On tensors the only difference between these ops is that mhlo.const supports unsigned types.

PiperOrigin-RevId: 377970948
2021-06-07 11:58:23 -07:00
Hanhan Wang 25b93c8d66 Add support for lowering mhlo.iota/dynamic_iota to Linalg on unsigned types.
PiperOrigin-RevId: 377956338
2021-06-07 10:59:33 -07:00
Adrian Kuegel 5315997402 Fix Sinh approximation for F16.
We should upcast F16 to F32 to prevent precision loss.
E.g. sinh(-9) would evaluate to -4042 previously instead of -4052.
This allows to enable the MLIR generated kernel for F16 type.

PiperOrigin-RevId: 377901896
2021-06-07 06:38:42 -07:00
Wenyi Zhao ade873a5e0 PR #49970: [MLIR][DISC] bufferize DynamicReshape and DynamicBroadcastInDim
Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/49970

1, add hlo-to-lhlo support for DynamicReshape and DynamicBroadcastInDim

2, add a flag `convert-to-lmhlo-only` to seperate following two case:
   - hlo-to-lhlo only. Simply lowers all mhlo ops to their lmhlo
     counterparts, do not apply any optimization (e.g. elide any
     buffer copy). Buffer optimization is not easy in dynamic
     shape world especially when involving control flow, thus we
     leave this to another dedicated pass.

   - hlo-to-lhlo-or-memref-directly. Lowers some metadata-only mhlo
     ops (e.g. reshape) to memref dialect directly and Lowers others
     to their lmhlo counterparts.
Copybara import of the project:

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

[MLIR][DISC] bufferize DynamicReshape and DynamicBroadcastInDim

1, add hlo-to-lhlo support for DynamicReshape and DynamicBroadcastInDim

2, add a flag `convert-to-lmhlo-only` to seperate following two case:
   - hlo-to-lhlo only. Simply lowers all mhlo ops to their lmhlo
     counterparts, do not apply any optimization (e.g. elide any
     buffer copy). Buffer optimization is not easy in dynamic
     shape world especially when involving control flow, thus we
     leave this to another dedicated pass.

   - hlo-to-lhlo-or-memref-directly. Lowers some metadata-only mhlo
     ops (e.g. reshape) to memref dialect directly and Lowers others
     to their lmhlo counterparts.

PiperOrigin-RevId: 377603395
2021-06-04 15:36:03 -07:00
A. Unique TensorFlower db05388a3c Integrate LLVM at llvm/llvm-project@da3ed58b97
Updates LLVM usage to match
[da3ed58b97c1](https://github.com/llvm/llvm-project/commit/da3ed58b97c1)

PiperOrigin-RevId: 377432380
2021-06-03 20:45:18 -07:00
A. Unique TensorFlower aba16adfa5 Add `mhlo.all_gather` op to MHLO dialect.
Adds import/export/verifier support as well.
Also makes `channel_handle` uniform across mhlo.all_reduce and mhlo.all-gather.

PiperOrigin-RevId: 377323468
2021-06-03 10:45:29 -07:00
A. Unique TensorFlower fe42a08fc9 Use channel_handle for ChannelHandles in MHLO ops. This makes the naming of these properties consistent across these ops.
PiperOrigin-RevId: 377309518
2021-06-03 09:49:47 -07:00
A. Unique TensorFlower 4620410f18 Integrate LLVM at llvm/llvm-project@b25546a4b4
Updates LLVM usage to match
[b25546a4b406](https://github.com/llvm/llvm-project/commit/b25546a4b406)

PiperOrigin-RevId: 377077163
2021-06-02 09:32:59 -07:00
A. Unique TensorFlower 557e56362e [MLIR][KernelGen] Simplify rank specialization tests with smaller target rank
For the tests rank specialize only up to rank 3. The remaining cases for higher
ranks are analogous.

PiperOrigin-RevId: 377024370
2021-06-02 03:48:07 -07:00
wyzhao 968d4b8709 PR #49598: [MLIR][DISC] legalize tensor_load inserted during hlo-to-lhlo conversion
Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/49598

This PR implements logic for lowering memref.tensor_load ops that are
inserted during `mhlo-legalize-to-lmhlo`
Copybara import of the project:

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

[MLIR][DISC] legalize tensor_load inserted during hlo-to-lhlo conversion

This PR implements logic for lowering memref.tensor_load ops that are
inserted during `mhlo-legalize-to-lmhlo`.

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

minor fix

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

minor refine

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

minor fix

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

fix

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

fix

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

fix

PiperOrigin-RevId: 376942547
2021-06-01 16:27:56 -07:00
A. Unique TensorFlower d1828625ab [MLIR][KernelGen] Make maximum supported rank in rank specialization configurable
The maximum supported target rank of 5 is sufficient for all operations but
`select`. Make the maximum target rank configurable in the rank specialization.
This reduces the number of generated kernels for operations that don't require
it.

PiperOrigin-RevId: 376822496
2021-06-01 06:54:31 -07:00
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