Commit Graph

154 Commits

Author SHA1 Message Date
Rahul Joshi 0651d58edf [MLIR:HLO] Change LHLO Outfeed to support multiple inputs.
PiperOrigin-RevId: 351447034
2021-01-12 14:13:47 -08:00
A. Unique TensorFlower 0b85d5c510 [MLIR][KernelGen] Add asin kernels and tests
PiperOrigin-RevId: 351381423
2021-01-12 09:02:46 -08:00
Alexander Belyaev ac38478350 Integrate LLVM at llvm/llvm-project@bcbdeafa9c
Updates LLVM usage to match
[bcbdeafa9cb3](https://github.com/llvm/llvm-project/commit/bcbdeafa9cb3)

PiperOrigin-RevId: 350763293
2021-01-08 07:57:20 -08:00
Alexander Belyaev 6c42f54298 [KERNEL_GEN] Restrict broadcast -> reshape canonicalization to identity dims.
This is needed to avoid the case, when the broadcast_in_dims also performs permutation.

PiperOrigin-RevId: 350650342
2021-01-07 15:30:28 -08:00
Rahul Joshi e3754d7b5c [MLIR:HLO] Change LHLO Infeed to support multiple outputs.
PiperOrigin-RevId: 348851612
2020-12-23 14:57:08 -08:00
Rahul Joshi bd8e768e23 [MLIR:HLO] Rename LMHLO Infeed and Outfeed to InfeedOp and OutfeedOp.
- To be consistent with other operation names.

PiperOrigin-RevId: 348824379
2020-12-23 11:38:05 -08:00
Rahul Joshi 8252eafa99 [NFC] Factor out repeated code out of InferFusibilityOpInterface.
PiperOrigin-RevId: 348671671
2020-12-22 12:04:29 -08:00
Rahul Joshi bc367971ec [MLIR:LHLO_GPU] Add additional constraints for batchnorm
- Constrain batchnorm inputs and outputs to be fp memrefs.

PiperOrigin-RevId: 348665747
2020-12-22 11:30:40 -08:00
Tres Popp a42213b870 Define lowering of [l]mhlo.pow.
For floating point operations, this uses std.pow.
For integer operations, this lowers to a loop.
This adds a dependency on scf.

PiperOrigin-RevId: 348537232
2020-12-21 15:27:40 -08:00
Smit Hinsu 8d051723c0 Use InferTypeOpInterface for HLO AbsOp and fix result shape inference
Shape inference in case of ops with complex element types need to use the element type of complex as the result element type and not the full operand type.

Before:
"mhlo.abs"(%arg0) : (tensor<4xcomplex<f32>>) -> tensor<4xtensor<4xcomplex<f32>>>
After:
"mhlo.abs"(%arg0) : (tensor<4xcomplex<f32>>) -> tensor<4xf32>
PiperOrigin-RevId: 348123967
2020-12-17 17:37:07 -08:00
Rahul Joshi 8134bff98d [XLA:GPU] Add layout attributes to LHLO_GPU Convolution operations.
- MLIR MemRefs do not preserve layout information correctly when unit dimensions
  are involved. Operations like convolution that use cuDNN however need the correct
  layout to be preserved so that we do not end up creating an incompatible combination
  of input/filter/output layout that is not supported by cuDNN.
- Add these layouts to convolution attributes in the form of I32ArrayAttr for representing
  the layout in "minor_to_major" form similar to XLA.

PiperOrigin-RevId: 348034757
2020-12-17 09:26:28 -08:00
Adrian Kuegel 1f244c3e2c Fix SignOp lowering for floating point values.
It didn't return 0 for 0.0 and -0.0.
Currently we emit -0.0 for -0.0 which is correct according to the HLO dialect.
For the TF_SignOp we should emit 0.0 in that case, we will leave that as a TODO.
Enable the tests which work now, and add another one for Int64.
Also improve the registration code, we should not register the Int32 kernel.

PiperOrigin-RevId: 347981124
2020-12-17 01:45:54 -08:00
Alexander Belyaev 65222893ae [KERNEL_GEN] Convert LHLO AddOp, SubOp (ComplexType) to complex ops.
PiperOrigin-RevId: 347805898
2020-12-16 05:45:06 -08:00
Alexander Belyaev e6e8920921 [KERNEL_GEN] Switch the pipeline to Linalg-on-Tensors.
PiperOrigin-RevId: 347781190
2020-12-16 01:51:15 -08:00
River Riddle 5ab94a00a7 [mlir][NFC] Replace usages or mlir/IR/StandardTypes.h with mlir/IR/BuiltinTypes.h
StandardTypes.h was moved to BuiltinTypes.h and is being removed.

PiperOrigin-RevId: 347751511
2020-12-15 21:28:18 -08:00
A. Unique TensorFlower f0c2695d31 Fix SignOp lowering for floating point values.
It didn't return 0 for 0.0 and -0.0.
Currently we emit -0.0 for -0.0 which is correct according to the HLO dialect.
For the TF_SignOp we should emit 0.0 in that case, we will leave that as a TODO.
Enable the tests which work now, and add another one for Int64.
Also improve the registration code, we should not register the Int32 kernel.

PiperOrigin-RevId: 347602378
2020-12-15 06:49:48 -08:00
Alexander Belyaev ddda2699fb [KERNEL_GEN] Switch the pipeline to Linalg-on-Tensors.
PiperOrigin-RevId: 347600145
2020-12-15 06:32:25 -08:00
Adrian Kuegel 79fa36bcbc Fix SignOp lowering for floating point values.
It didn't return 0 for 0.0 and -0.0.
Currently we emit -0.0 for -0.0 which is correct according to the HLO dialect.
For the TF_SignOp we should emit 0.0 in that case, we will leave that as a TODO.
Enable the tests which work now, and add another one for Int64.
Also improve the registration code, we should not register the Int32 kernel.

PiperOrigin-RevId: 347590340
2020-12-15 05:12:48 -08:00
Alexander Belyaev 8b35a75d4a [KERNEL_GEN] Switch the pipeline to Linalg-on-Tensors.
PiperOrigin-RevId: 347368063
2020-12-14 05:46:47 -08:00
River Riddle 6b439f7eee [mlir][NFC] Replace usages or mlir/IR/StandardTypes.h with mlir/IR/BuiltinTypes.h
StandardTypes.h was moved to BuiltinTypes.h and is being removed.

PiperOrigin-RevId: 347115952
2020-12-11 19:01:25 -08:00
Tim Shen 22b91d62ff [XLA/GPU] Migrate all unnested elementwise emitters.
Also fix GetHloOutputs/GetHloOperands to treat aliased operands correctly.

PiperOrigin-RevId: 347055290
2020-12-11 12:45:21 -08:00
Benjamin Kramer 9930c20c31 [mlir][hlo] Fix lowering of NE comparison. It should return true if either side is NaN
PiperOrigin-RevId: 346988987
2020-12-11 06:46:14 -08:00
Rahul Joshi f232da1f9d [MLIR:HLO] Add window_reversal attribute to convolution attributes.
- Add this attribute to match the corresponding XLA HLO attribute on convolution
  operations.
- A true value indicates a reversal of the corresponding kernel spatial dimension.
- Since XLA builder does not support this attribute, use a custom HLO converted to map
  from mlir::mhlo::ConvOp to XLA.

PiperOrigin-RevId: 346891737
2020-12-10 16:39:19 -08:00
A. Unique TensorFlower 6a05893169 [XLA/GPU] Migrate all unnested elementwise emitters.
PiperOrigin-RevId: 346725498
2020-12-10 00:21:36 -08:00
Tim Shen cfcf741932 [XLA/GPU] Migrate all unnested elementwise emitters.
PiperOrigin-RevId: 346716519
2020-12-09 23:08:13 -08:00
A. Unique TensorFlower 65ebd85563 [XLA/GPU] Migrate all unnested elementwise emitters.
PiperOrigin-RevId: 346658288
2020-12-09 16:05:43 -08:00
Tim Shen 1c10e1fec6 [XLA/GPU] Migrate all unnested elementwise emitters.
PiperOrigin-RevId: 346624905
2020-12-09 13:31:50 -08:00
A. Unique TensorFlower 1a5a1b5f41 [XLA/GPU] Migrate all unnested elementwise emitters.
PiperOrigin-RevId: 346574093
2020-12-09 09:54:16 -08:00
Tim Shen 3c33fe4b9e [XLA/GPU] Migrate all unnested elementwise emitters.
PiperOrigin-RevId: 346559170
2020-12-09 08:42:51 -08:00
Stephan Herhut c3790af758 Add plumbing for or and xor to hlo to lhlo and linalg lowerings.
PiperOrigin-RevId: 346311314
2020-12-08 06:39:02 -08:00
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