From f232da1f9d111cba246223353a08410701ea97a7 Mon Sep 17 00:00:00 2001
From: Rahul Joshi <jurahul@google.com>
Date: Thu, 10 Dec 2020 16:38:26 -0800
Subject: [PATCH] [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
---
 include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td   |  1 +
 .../mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td  | 21 +++++++++++++++++++
 .../mhlo/transforms/legalize_to_linalg.cc     |  3 ++-
 tests/lhlo_gpu_ops.mlir                       |  3 ++-
 4 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
index ba4749e..63fec24 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
@@ -902,6 +902,7 @@ def HLO_ConvOp : HLO_Op<"convolution", [NoSideEffect]>, BASE_HLO_ConvOp {
     ConvolutionAttributes.attributes);
 
   let results = (outs HLO_Tensor);
+  let hasCustomHLOConverter = 1;
 }
 
 def HLO_CopyOp: HLO_Op<"copy", [NoSideEffect, SameOperandsAndResultType]>, BASE_HLO_CopyOp {
diff --git a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td
index 9b1b126..57fdfb6 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops_base.td
@@ -958,6 +958,17 @@ def HLO_PrecisionConfigAttr:
     OptionalAttr<
           TypedArrayAttrBase<HLO_PrecisionAttr, "Precision Config attribute">>;
 
+def BoolElementsAttr :
+    ElementsAttrBase<
+      And<[CPred<"$_self.isa<::mlir::DenseIntOrFPElementsAttr>()">,
+           CPred<"$_self.cast<::mlir::DenseIntOrFPElementsAttr>().getType().getElementType().isInteger(1)">]>,
+      "constant boolean vector/tensor attribute"> {
+  let storageType = [{ ::mlir::DenseElementsAttr }];
+  let returnType = [{ ::mlir::DenseElementsAttr }];
+
+  let convertFromStorage = "$_self";
+}
+
 def ConvolutionAttributes {
   dag attributes = (ins
     // Default value: one for each of the spatial dimension.
@@ -968,6 +979,8 @@ def ConvolutionAttributes {
     OptionalAttr<I64ElementsAttr>:$lhs_dilation,
     // Default value: one for each of the spatial dimension.
     OptionalAttr<I64ElementsAttr>:$rhs_dilation,
+    // Default value: one for each of the spatial dimension.
+    OptionalAttr<BoolElementsAttr>:$window_reversal,
     ConvDimensionNumbers:$dimension_numbers,
     I64Attr:$feature_group_count,
     I64Attr:$batch_group_count,
@@ -983,6 +996,14 @@ class BASE_HLO_ConvOp {
 
     See https://www.tensorflow.org/xla/operation_semantics#conv_convolution.
   }];
+
+  code extraClassDeclaration = [{
+    bool hasWindowReversal() {
+      auto reversal = window_reversalAttr();
+      return reversal && llvm::any_of(reversal.getBoolValues(),
+                                      [](bool v) { return v; });
+    }
+  }];
 }
 
 class BASE_HLO_CopyOp {
diff --git a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
index 31b7a60..9beda13 100644
--- a/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
+++ b/lib/Dialect/mhlo/transforms/legalize_to_linalg.cc
@@ -243,7 +243,8 @@ struct ConvToLinalgConverter : public OpConversionPattern<lmhlo::ConvOp> {
     }
 
     // TODO: LHS dilation for deconvolution not supported yet.
-    if (op.lhs_dilation()) {
+    // TODO(jurahul): Window reversal is not supported yet.
+    if (op.lhs_dilation() || op.hasWindowReversal()) {
       return failure();
     }
 
diff --git a/tests/lhlo_gpu_ops.mlir b/tests/lhlo_gpu_ops.mlir
index 35bf59b..83327d4 100644
--- a/tests/lhlo_gpu_ops.mlir
+++ b/tests/lhlo_gpu_ops.mlir
@@ -103,7 +103,8 @@ func @conv_backinput(%input : memref<4x5x16x16xf64>, %filter : memref<5x3x7x7xf6
     precision_config = [],
     result_scale = 1.000000e+00 : f64,
     rhs_dilation = dense<1> : tensor<2xi64>,
-    window_strides = dense<1> : tensor<2xi64>}
+    window_strides = dense<1> : tensor<2xi64>,
+    window_reversal = dense<true>: tensor<2xi1>}
   : (memref<4x5x16x16xf64>, memref<5x3x7x7xf64>, memref<4x3x16x16xf64>, memref<32xui8>) -> ()
   return
 }