colin
/
mlir-hlo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[MLIR][KernelGen] Lower `tf.Acos` to LMHLO. 

- Add ranked code generation for `mhlo.compare/select`
- Add bufferization for `tensor_cast`
- Add lowerings for `Atan2Op`

PiperOrigin-RevId: 332407734
This commit is contained in:
A. Unique TensorFlower 2020-09-18 01:39:48 -07:00 committed by TensorFlow MLIR Team
parent 794e275904
commit 2fbbbe9cf1
3 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
include/mlir-hlo/Dialect/mhlo/transforms/map_hlo_to_lhlo_op.h
View File

@ -40,6 +40,7 @@ using HloToLhloOp = typename HloToLhloOpImpl<HloOpTy>::Type;
MAP_HLO_TO_LHLO(AbsOp);
MAP_HLO_TO_LHLO(AddOp);
MAP_HLO_TO_LHLO(AndOp);
MAP_HLO_TO_LHLO(Atan2Op);
MAP_HLO_TO_LHLO(BroadcastInDimOp);
MAP_HLO_TO_LHLO(CeilOp);
MAP_HLO_TO_LHLO(ConstOp);

1
lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
View File

@ -488,6 +488,7 @@ void populateHLOToLHLOConversionPattern(
HloToLhloOpConverter<mhlo::AbsOp>,
HloToLhloOpConverter<mhlo::AddOp>,
HloToLhloOpConverter<mhlo::AndOp>,
HloToLhloOpConverter<mhlo::Atan2Op>,
HloToLhloOpConverter<mhlo::BroadcastInDimOp>,
HloToLhloOpConverter<mhlo::CeilOp>,
HloToLhloOpConverter<mhlo::CompareOp>,

6
lib/Dialect/mhlo/transforms/transform_unranked_hlo.cc
View File

@ -145,6 +145,8 @@ struct TransformUnrankedHloPass
MAP_CHLO_OPERATION_CWISE_UNARY(ADD_LEGAL_CHLO, ;);
#undef ADD_LEGAL_MHLO
#undef ADD_LEGAL_CHLO
AddLegalOpOnRankedTensor<mhlo::CompareOp>(&target);
AddLegalOpOnRankedTensor<mhlo::SelectOp>(&target);
// Populate rewrite patterns.
OwningRewritePatternList patterns;
@ -168,7 +170,9 @@ void PopulateTransformUnrankedHloPatterns(MLIRContext *context,
patterns->insert<
MAP_XLA_OPERATION_CWISE_UNARY(MAP_UNARY, COMMA),
MAP_XLA_OPERATION_CWISE_BINARY(MAP_BINARY, COMMA),
MAP_CHLO_OPERATION_CWISE_UNARY(MAP_CHLO_UNARY, COMMA)>(context);
MAP_CHLO_OPERATION_CWISE_UNARY(MAP_CHLO_UNARY, COMMA),
ElementwiseOpConversion<mhlo::CompareOp>,
ElementwiseOpConversion<mhlo::SelectOp>>(context);
// clang-format on
#undef MAP_UNARY
#undef MAP_BINARY