Commit Graph

31 Commits

Author SHA1 Message Date
Jacques Pienaar 95ba03534f Allow variadic operands/result in MHLO while
This just adds support for it in the op, but keeps the production/uses as is (e.g., single tensor or tuple) matching what XLA export requires. In follow up here, would be to add pass for export to retuple and then the canonical form could be changed. Tuple'ing given control flow via regions & multi-result operations does not add representational power and all the get_tuple_element ops obscure the computation.

The old form allowed single tensor or tuple. The new variadic number of tensor or tuples as tuples may be nested, so the input could have (Tensor<..>, Tuple<Tensor<...>, Tuple<...>, ...>, Tensor<...>) and HLO_Tensor doesn't allow Tuples.

PiperOrigin-RevId: 378934388
2021-06-11 13:08:28 -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
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
Jacques Pienaar fdd75daed6 Add shape function for MHLO RngNormal and RngUniform
PiperOrigin-RevId: 368276963
2021-04-13 12:59:42 -07:00
Rahul Joshi ff2cbfa2ec [MLIR] Add support for representing variadic reduce-window in HLO/LMHLO dialect.
-  Fixed a subset of transformations to handle variadic reduce-window.

PiperOrigin-RevId: 366278650
2021-04-01 10:24:50 -07:00
Geoffrey Martin-Noble 2fb2a92c6e Verify mhlo.if region return types match op
This matches the behavior of mhlo.case. Additionally, fix the verification of CaseOp in the case of nested ops with mhlo.return-containing regions.

PiperOrigin-RevId: 365936672
2021-03-30 17:57:20 -07:00
Geoffrey Martin-Noble 7a9394dca5 Restrict MHLO control flow ops to single-block regions
PiperOrigin-RevId: 365935824
2021-03-30 17:51:03 -07:00
Jacques Pienaar 3de2024a9b Avoid creating tuple type only for verification
Make the error message a bit more verbose & it is cheaper to verify the elements rather than creating a (potentially) new type.

PiperOrigin-RevId: 363073909
2021-03-15 17:58:19 -07:00
Jacques Pienaar 329b1fd071 Verify compatible shapes in unpack verification rather than exact
Previously this would be too strict and fail if dynamic and static dims were
compared. Dynamic/unknown are treated as "maybe equal" to a static value without further info, so at this layer don't flag as invalid unless truly are.

PiperOrigin-RevId: 360189086
2021-03-01 08:00:16 -08:00
Richard Uhler b579bd5d9e Support dynamic-shaped operand in verification of BroadcastInDim.
Verification of HLO_BroadcastInDimOp was previously failing or crashing if the
operand had a dynamic shape or was unranked. Update the verification code to
allow the operand to be unranked or have dynamic shape.

PiperOrigin-RevId: 358056793
2021-02-17 16:18:09 -08:00
A. Unique TensorFlower 4a29ca3b1d Add layout to mhlo::InfeedOp td.
PiperOrigin-RevId: 356286875
2021-02-08 09:48:14 -08:00
A. Unique TensorFlower 04110a4b1c Integrate LLVM at llvm/llvm-project@67dfe9c8d7
Updates LLVM usage to match
[67dfe9c8d70c](https://github.com/llvm/llvm-project/commit/67dfe9c8d70c)

PiperOrigin-RevId: 355235205
2021-02-02 13:09:20 -08:00
Rahul Joshi bbaad416a1 [MLIR] Update tests to eliminate public function declarations.
- Also fixed kernel_gen embed_memref_prints to mark the declarations inserted private.
- This is in prep for proposed MLIR change to disallow public declarations.

PiperOrigin-RevId: 342081252
2020-11-12 10:42:59 -08:00
Rahul Joshi 1958f228ec [MLIR:HLO] Extend CustomCall to support multiple outputs.
- Extend MHLO CustomCall to have multiple tensors as results.
- Extend LHLO CustomCall to have multiple memrefs for output operands.
- Fix HLO->LHLO and XLA HLO->LHLO mapping for CustomCall to setup the
  operand_segment_sizes attribute correctly.

PiperOrigin-RevId: 342067762
2020-11-12 09:46:25 -08:00
Smit Hinsu 4ef12aa000 Update GetDimensionSize and SetDimensionSize ops to use I64 attribute for dimension
This is to match with HLO semantics and general dimension semantics in MLIR.

Also,

* Define minimal verifier for these ops.
* Add folder for SetDimensionSize op on static shaped dimension.
* Fix assumption of ranked shape in GetDimensionSize op.

PiperOrigin-RevId: 341150923
2020-11-06 18:03:04 -08:00
Dmitry Volodin 1821c69910 PR #44405: Fix typos in compiler directory
Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/44405

Splitting #43857 by top-level directories.
Copybara import of the project:

--
fa5da7d5478649d11321dcac9f867b0a57e4798a by Dmitry Volodin <mr.molkree@gmail.com>:

fix typos in compiler dir

--
4d3c9f047f7ecb8ab299f1bf28a86fd39096eee7 by Dmitry Volodin <mr.molkree@gmail.com>:

fix one test as "atleast" in it comes from Bazel

--
9440ebaaa9fc4a735f7f72f0c8f0de4ec58afbd6 by Dmitry Volodin <mr.molkree@gmail.com>:

a bit more

PiperOrigin-RevId: 340819994
2020-11-05 03:31:54 -08:00
Richard Uhler 82031b356c Improve error message for improperly shaped slice indices.
The slice indices must be rank-1 and have the same number of elements of the
rank of the operand. Give reasonable error messages for violations of these
requirements instead of a misleading error message that the types of the
indices don't all match.

PiperOrigin-RevId: 340660822
2020-11-04 09:10:51 -08:00
Richard Uhler f9843fabe1 Use InferTypeOpInterface for HLO_SliceOp.
Instead of having a custom builder to construct a slice op without an explicit
return type.

PiperOrigin-RevId: 339058864
2020-10-26 09:54:13 -07:00
Jacques Pienaar 27968619b7 Verify non-scalar inputs for HLO concat
XLA HLO concat does not accept scalars, so fail verification if this occurs. Avoids segfault when accessing an empty output shape.

PiperOrigin-RevId: 337618167
2020-10-16 19:39:31 -07:00
Tim Shen c708bfd6d0 [MLIR] Add cbrt, reduce-precision, and bitcast ops to MHLO.
PiperOrigin-RevId: 335109804
2020-10-02 15:13:18 -07:00
Robert Suderman 26ac5baae4 Make mhlo.sort return variadic results instead of a tuple
Tuple is only used on XLA's sort to return multiple inputs. MLIR supports
multiple inputs, switch to a tuple return.

PiperOrigin-RevId: 334226937
2020-09-28 13:32:23 -07:00
Rahul Joshi 87d9e6951e [MLIR] Use MLIR provided 0DTensorOf/1DTensorOf instead of custom constraints.
- Use MLIR provided constraints for HLO_ScalarIntTensor and HLO_DimensionTensor.
- Update unit tests to expect new error messages.

PiperOrigin-RevId: 333313131
2020-09-23 09:57:12 -07:00
Prakalp Srivastava b7248424ae Fix return type of replica-id to unsigned int 32 tensor.
PiperOrigin-RevId: 330012181
2020-09-03 16:08:38 -07:00
Andy Ly 53fdda7f3e Update mhlo.constant to use a custom assembly format instead of a custom printer and parser (NFC).
PiperOrigin-RevId: 325560779
2020-08-07 22:23:06 -07:00
Lucy Fox d742477c02 Verify that MHLO DynamicUpdateSlice start indices have matching element types.
HLO requires that the element types match for all start index parameters. Right now we don't catch this invalid case until export, so adding a check in the verifier so that we catch this sooner.

This also requires a small tweak to the TF InplaceUpdate op lowering.

PiperOrigin-RevId: 325463796
2020-08-07 22:21:40 -07:00
Lucy Fox cd22ecd136 Relax DynamicBroadcastInDim verifier when dimensions are dynamic.
For input and output dimensions which must match, we shouldn't fail in the case where one dim is dynamic and the other is static. This is insufficient information to conclude a dimension mismatch.

PiperOrigin-RevId: 325344738
2020-08-07 22:18:38 -07:00
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
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
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