Commit Graph

124 Commits

Author SHA1 Message Date
Benjamin Kramer 5235eceea0 Lower mhlo shifts to linalg
PiperOrigin-RevId: 346161253
2020-12-07 13:02:32 -08:00
Smit Hinsu 9bd1995f90 Legalize XlaReplicaId to HLO replica-id op
Also, define shape inference function for HLO replica-id op.

PiperOrigin-RevId: 345714342
2020-12-04 11:04:40 -08:00
Rahul Joshi e48881af81 [MLIR:LHLO_GPU] Add fused convolution operation without any side inputs.
- Add a variant of the fused convolution that does not need a side input and side input scale.
- Rename the existing one to `ConvForwardFusedSideInputOp`.
- Update tests to exercise all variants of the convolution ops in the GPU dialect.
- Eliminate unused `LHLO_ExtentBuffer` and changed LHLO_Buffer to allow any integer element
  type to match what XLA can generate sometimes for scratch buffers.

PiperOrigin-RevId: 345701569
2020-12-04 10:09:27 -08:00
A. Unique TensorFlower 3691e39f62 Integrate LLVM at llvm/llvm-project@f5d52916ce
Updates LLVM usage to match
[f5d52916ce34](https://github.com/llvm/llvm-project/commit/f5d52916ce34)

PiperOrigin-RevId: 345672828
2020-12-04 07:24:17 -08:00
Tim Shen 73cb134cca [MLIR] Add XLA HLO -> LMHLO conversion for all elementwise ops.
PiperOrigin-RevId: 345557248
2020-12-03 16:12:02 -08:00
Rahul Joshi dbbdfea95b [MLIR:HLO] Generate enum decls for HLO and LHLO GPU dialects.
- Split out enum definitions in hlo dialect into a separate .td file (similar to structs)
  and generate enum decl/defs for these enums.
- Also split out the LHLO GPU enums into a separate .td file and generate enum
  decl/defs for these enums as well.
- Remove unused dialect from ConvolutionAttributes and generate lhlo_gpu enums.
- Add appropriate namespace for all the enums.

PiperOrigin-RevId: 345277240
2020-12-02 11:39:23 -08:00
Rahul Joshi d7bd5233ab [XLA:GPU] Migrate GEMM Thunk emission to MLIR.
- Map Custom call for GEMM in XLA HLO to Gemm/Gemm bias operations in LHLO GPU
  dialect.
- Make 'algorithm' an optional attribute to better match with XLA HLO backend config.
- Replace 'alpha' with 'alpha_real' and 'alpha_complex' to support complex GEMM correctly.
- Generate GemmThunk off of LHLO GPU Gemm operations.

PiperOrigin-RevId: 345250840
2020-12-02 09:43:12 -08:00
Rahul Joshi 5350b8080d [NFC] Add default value for attributes in lmhlo::FusionOp custom builder.
PiperOrigin-RevId: 344856882
2020-11-30 11:47:39 -08:00
Alexander Belyaev 5583c63cab [KERNEL_GEN] Add unranked Conj kernel.
PiperOrigin-RevId: 344243271
2020-11-25 06:37:26 -08:00
Tim Shen 3e01448481 [MLIR] Move some walk()ing functions to the lmhlo::FusionOp API.
PiperOrigin-RevId: 344109366
2020-11-24 12:23:57 -08:00
Lucy Fox 85f92a1651 [KernelGen] Lower tf.Erf and tf.Erfc ops to CHLO.
This does not include the lowerings from CHLO to LMHLO.

PiperOrigin-RevId: 344091604
2020-11-24 10:55:43 -08:00
Rahul Joshi ac54c5ccfa [XLA:GPU] Convert Cholesky custom call in XLA HLO to LHLO GPU Dialect.
- Restructured LHLO GPU Cholesky to better match XLA HLO by eliminating the
  untyped buffer and changing is_upper attribute to is_lower.
- Change LhloDialectEmitter to emit LHLO GPU Cholesky operation.

PiperOrigin-RevId: 343873516
2020-11-23 10:06:21 -08:00
Tim Shen 5a6edaa588 Roll-forward with fix:
[XLA/GPU] Migrate nested reduce emitter to take LMHLO.

PiperOrigin-RevId: 343582798
2020-11-20 16:06:58 -08:00
A. Unique TensorFlower bae7c7bd2b [XLA/GPU] Migrate nested reduce emitter to take LMHLO.
PiperOrigin-RevId: 343388053
2020-11-19 16:00:58 -08:00
Tim Shen 303181924c [XLA/GPU] Migrate nested reduce emitter to take LMHLO.
PiperOrigin-RevId: 343373088
2020-11-19 14:41:55 -08:00
A. Unique TensorFlower 7f239c7ba2 Add canonicalizer for Reshape(Broadcast(X)) pattern when it is an identity sequence
PiperOrigin-RevId: 343251257
2020-11-19 02:32:45 -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
Alexander Belyaev 3d930d08c2 [HLO] Delete LHLO memref cast ops and migrate to STD ones.
PiperOrigin-RevId: 340663578
2020-11-04 09:26:34 -08:00
Sean Silva d3ea3abdec Remove `results_escape_functions` from HloLegalizeToLhlo
PiperOrigin-RevId: 340464958
2020-11-03 09:49:56 -08:00
Tres Popp 76b30fd426 Move unranked chlo lowering to transform_unranked_hlo.
Additionally:
- Forward listeners through new if/else op builders.
This corrects an error that led to incomplete legalization of broadcasted op
lowering.
- Use OpConversionPattern to ensure up to date operand values are used.
PiperOrigin-RevId: 339838833
2020-10-30 02:56:44 -07:00
Alex Zinenko e77b5ae534 Update MLIR ODS syntax to use OpBuidlerDAG instead of OpBuilder
The latter has been deprecated in favor of the former and will be removed.

PiperOrigin-RevId: 339523966
2020-10-28 13:16:46 -07:00
Kazuaki Ishizaki 5f2b7a74be PR #44277: NFC - minor spelling tweaks of documents under compiler directory
Imported from GitHub PR https://github.com/tensorflow/tensorflow/pull/44277

This PR addresses minor spelling tweaks of md/td files under compiler directory
Copybara import of the project:

--
4fedebde8f7d48ce2917642ebaab966c9ce49f3e by Kazuaki Ishizaki <ishizaki@jp.ibm.com>:

minor spelling tweaks

PiperOrigin-RevId: 339260830
2020-10-27 09:08:22 -07:00
Smit Hinsu 6eda9ed273 Add compare_type optional attribute to CompareOp in HLO dialects
If unspecified, `compare_type` is FLOAT for float element types, SIGNED for signed element types and UNSIGNED for unsigned element types. compare_type can be TOTALORDER for float element types.

- Added import and export support the attribute.
- Restricted legalization from HLO to TF to the default compare types.
- Updated existing usage of the CompareOp

PiperOrigin-RevId: 339099219
2020-10-26 12:58:29 -07: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
Benjamin Kramer 31c1c3aa1f Integrate LLVM at llvm/llvm-project@c89447b659
Updates LLVM usage to match
[c89447b65984](https://github.com/llvm/llvm-project/commit/c89447b65984)

PiperOrigin-RevId: 338560059
2020-10-22 15:23:24 -07:00
A. Unique TensorFlower 33c450e4cb Fix the MHLO to LMHLO lowering of 'gather'
The lowering assumes that the 'gather' op attributes are identical in both MHLO and LMHLO. But that's not true; some time ago the MHLO version was changed to pack 4 of its attributes into a struct. By doing the same for the LMHLO version we both fix the lowering for this op and resolve a longstanding TODO.

PiperOrigin-RevId: 337943946
2020-10-19 15:14:05 -07:00
A. Unique TensorFlower 7a983ea389 Add folder for mhlo::pad
PiperOrigin-RevId: 337827560
2020-10-19 04:21:44 -07:00
A. Unique TensorFlower 4a18aa41ee Add folder to mhlo::round_nearest_afz
PiperOrigin-RevId: 337823786
2020-10-19 03:45:15 -07:00
A. Unique TensorFlower 706718b4fb Permit vector types in lmhlo to std lowering.
PiperOrigin-RevId: 337523303
2020-10-16 09:47:02 -07:00
A. Unique TensorFlower 51cd4200b6 Make LMHLO's Dot have the same power as MHLO's DotGeneral.
PiperOrigin-RevId: 337391565
2020-10-15 15:09:06 -07:00
A. Unique TensorFlower 05ee41baf8 Add folder for mhlo::scatter
PiperOrigin-RevId: 337274351
2020-10-15 03:26:05 -07:00
Adrian Kuegel 0f36979e2c Integrate LLVM at llvm/llvm-project@220de1f32a
Updates LLVM usage to match
[220de1f32add](https://github.com/llvm/llvm-project/commit/220de1f32add)

PiperOrigin-RevId: 337270545
2020-10-15 02:54:24 -07:00
Rahul Joshi f6b4e6758a Add GPU specific LMHLO level ops
- Introduce operations in a new lmhlo_gpu dialect that map to GPU library function calls
  in the XLA:GPU backend.
- Add basic unit tests as well.

PiperOrigin-RevId: 337132166
2020-10-14 11:23:55 -07:00
Adrian Kuegel 53d16d9f8e Integrate LLVM at llvm/llvm-project@93377888ae
Updates LLVM usage to match
[93377888ae89](https://github.com/llvm/llvm-project/commit/93377888ae89)

PiperOrigin-RevId: 336669842
2020-10-12 08:42:31 -07:00
Jacques Pienaar ab4c03b84b Replace operands now injected in OpBuilder
These are now injected automatically.

PiperOrigin-RevId: 336356284
2020-10-09 13:39:33 -07:00
Rahul Joshi 41436ea0d9 [NFC] Extract common attributes of MHLO and LMHLO ConvOp
- Create a common class that hold ConvOp attributes that are common to MHLO and
  LHLO Dialects and use Tablegen DAG concat to append these common attributes to
  dialect specific arguments in ODS for ConvOp.

PiperOrigin-RevId: 336114003
2020-10-08 10:29:32 -07:00
Pavel Krajcevski ac483f4091 Remove superfluous quote from lhlo_ops.td
PiperOrigin-RevId: 335921873
2020-10-07 12:06:36 -07:00
A. Unique TensorFlower 3736c5542f [MLIR][KernelGen] Fix unranked codegeneration in kernel generator
PiperOrigin-RevId: 335847086
2020-10-07 05:39:55 -07:00
A. Unique TensorFlower bae0815ef0 [MLIR][KernelGen] Legalize `atan` to approximation
PiperOrigin-RevId: 335417836
2020-10-05 08:05:52 -07:00
A. Unique TensorFlower 7f84a86cf5 [MLIR][KerneGen] Lower `tf.Atan` all the way to LLVM
PiperOrigin-RevId: 335394668
2020-10-05 05:07:13 -07:00
A. Unique TensorFlower 7367eac074 Add folder for mhlo::remainder
PiperOrigin-RevId: 335372628
2020-10-05 02:20:01 -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
Chao Xie 5f303440da [MLIR][KerneGen] Lower `tf.Atan` all the way to LLVM
PiperOrigin-RevId: 334843070
2020-10-01 10:25:51 -07:00
A. Unique TensorFlower 458e861254 [MLIR][KerneGen] Lower `tf.Atan` all the way to LLVM
PiperOrigin-RevId: 334810730
2020-10-01 07:38:40 -07:00
Rahul Joshi bce128b070 Introduce CustomCall operation in LHLO Dialect
- And add conversion from MHLO CustomCall to LHLO CustomCall
- According to XLA documentation, the called function should not be side effecting,
  so marking the argument MemRefs as MemRead.

PiperOrigin-RevId: 334737196
2020-09-30 20:56:18 -07:00
Benjamin Kramer dfe64d3958 Implement InferShapedTypeOpInterface for mhlo.complex
Binary companion for 8bcd33e4b7

PiperOrigin-RevId: 334651523
2020-09-30 12:14:15 -07:00
Benjamin Kramer c8919f8419 Implement InferShapedTypeOpInterface and use inferReturnTypes for mhlo.imag and mhlo.real
This makes the lhlo lowering work with dynamic shapes.

PiperOrigin-RevId: 334553472
2020-09-30 02:02:26 -07:00
Ahmed S. Taei 39389587d2 {mhlo.is_finite, lmhlo.is_finite} -> linalg.generic conversion
PiperOrigin-RevId: 334414295
2020-09-29 10:50:42 -07:00
Benjamin Kramer 336ee14538 Make lmhlo.sign to linalg lowering work for more floating point and integer types
PiperOrigin-RevId: 334400341
2020-09-29 09:52:13 -07:00