From f5d12604ed11e1908ac7415479d53887aa798a78 Mon Sep 17 00:00:00 2001
From: Eugene Burmako <burmako@google.com>
Date: Fri, 4 Sep 2020 14:58:10 -0700
Subject: [PATCH] Add support for legalizing mhlo.transpose to lmhlo.transpose

PiperOrigin-RevId: 330127758
---
 .../Dialect/mhlo/transforms/map_hlo_to_lhlo_op.h    |  1 +
 lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc |  1 +
 tests/hlo-legalize-to-lhlo.mlir                     | 13 +++++++++++++
 3 files changed, 15 insertions(+)

diff --git a/include/mlir-hlo/Dialect/mhlo/transforms/map_hlo_to_lhlo_op.h b/include/mlir-hlo/Dialect/mhlo/transforms/map_hlo_to_lhlo_op.h
index bc08f8c..1f46ea1 100644
--- a/include/mlir-hlo/Dialect/mhlo/transforms/map_hlo_to_lhlo_op.h
+++ b/include/mlir-hlo/Dialect/mhlo/transforms/map_hlo_to_lhlo_op.h
@@ -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
 
diff --git a/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc b/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
index c7e2c81..efeead8 100644
--- a/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
+++ b/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc
@@ -500,6 +500,7 @@ void populateHLOToLHLOConversionPattern(
       HloToLhloOpConverter<mhlo::SqrtOp>,
       HloToLhloOpConverter<mhlo::SubOp>,
       HloToLhloOpConverter<mhlo::TanhOp>,
+      HloToLhloOpConverter<mhlo::TransposeOp>,
       HloToLhloReduceOpConverter,
       HloToLhloReturnOpConverter,
       HloToLhloTensorLoadOpConverter,
diff --git a/tests/hlo-legalize-to-lhlo.mlir b/tests/hlo-legalize-to-lhlo.mlir
index 7246492..cd26348 100644
--- a/tests/hlo-legalize-to-lhlo.mlir
+++ b/tests/hlo-legalize-to-lhlo.mlir
@@ -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
+}