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

Add support for legalizing mhlo.transpose to lmhlo.transpose 

PiperOrigin-RevId: 330127758
This commit is contained in:
Eugene Burmako 2020-09-04 14:58:10 -07:00 committed by TensorFlow MLIR Team
parent b6b30698ad
commit f5d12604ed
3 changed files with 15 additions and 0 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

@ -72,6 +72,7 @@ MAP_HLO_TO_LHLO(SinOp);
MAP_HLO_TO_LHLO(SqrtOp);
MAP_HLO_TO_LHLO(SubOp);
MAP_HLO_TO_LHLO(TanhOp);
MAP_HLO_TO_LHLO(TransposeOp);
#undef MAP_HLO_TO_LHLO

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

@ -500,6 +500,7 @@ void populateHLOToLHLOConversionPattern(
HloToLhloOpConverter<mhlo::SqrtOp>,
HloToLhloOpConverter<mhlo::SubOp>,
HloToLhloOpConverter<mhlo::TanhOp>,
HloToLhloOpConverter<mhlo::TransposeOp>,
HloToLhloReduceOpConverter,
HloToLhloReturnOpConverter,
HloToLhloTensorLoadOpConverter,

13
tests/hlo-legalize-to-lhlo.mlir
View File

@ -525,3 +525,16 @@ func @reduce(%arg0: tensor<1x8xf32>, %arg1: tensor<f32>) -> tensor<1xf32> {
: (tensor<1x8xf32>, tensor<f32>) -> tensor<1xf32>
return %0 : tensor<1xf32>
}
// -----
// BOTH-LABEL: func @transpose
func @transpose(%operand: memref<2x2xf32>, %result: memref<2x2xf32>) {
%tensor_operand = tensor_load %operand : memref<2x2xf32>
%tensor_result = "mhlo.transpose"(%tensor_operand) {permutation = dense<[1, 0]> : tensor<2xi64>}
: (tensor<2x2xf32>) -> tensor<2x2xf32>
// BOTH: "lmhlo.transpose"(%{{.*}}, %{{.*}}) {permutation = dense<[1, 0]> : tensor<2xi64>}
// BOTH-NOT: tensor_store
tensor_store %tensor_result, %result : memref<2x2xf32>
return
}