From b32f885ad7e0b91331ea3f46f6bda84df9b8e159 Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Thu, 27 May 2021 02:43:43 -0700
Subject: [PATCH] [MLIR][KernelGen] Enable cluster rank specialization

Replace the previously used `TransformUnrankedHloPass` which rank-specializes
only one operation at a time. The new generalized rank specialization clusters
compatible operations and rank-specializes them collectively.

PiperOrigin-RevId: 376127752
---
 include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
index 33c8ecd..c4e9a84 100644
--- a/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
+++ b/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td
@@ -1029,7 +1029,7 @@ def HLO_TupleOp : HLO_Op<"tuple", [NoSideEffect]> {
 }
 
 def HLO_CompareOp: HLO_Op<"compare", [NoSideEffect, SameTypeOperands,
-    SameOperandsAndResultShape,
+    SameOperandsAndResultShape, Elementwise,
     DeclareOpInterfaceMethods<InferShapedTypeOpInterface,
         ["inferReturnTypeComponents", "reifyReturnTypeShapes"]>]> {
   let summary = "Comparison operator";