[Mlir-commits] [mlir] 9788048 - [MLIR][Shape] Add `shape.shape_eq` operation

[Frederik Gossen]

Author: Frederik Gossen
Date: 2020-07-15T10:30:52Z
New Revision: 978804821e88a34d484a8ebab72d2888f869a086

URL: https://github.com/llvm/llvm-project/commit/978804821e88a34d484a8ebab72d2888f869a086
DIFF: https://github.com/llvm/llvm-project/commit/978804821e88a34d484a8ebab72d2888f869a086.diff

LOG: [MLIR][Shape] Add `shape.shape_eq` operation

Add `shape.shape_eq` operation to the shape dialect.
The operation allows to test shapes and extent tensors for equality.

Differential Revision: https://reviews.llvm.org/D82528

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td
    mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
    mlir/test/Dialect/Shape/ops.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td
index 8bf1e36c63e2..ef20b5a9813d 100644
--- a/mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td
+++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeBase.td
@@ -104,7 +104,13 @@ def Shape_ValueShapeType : DialectType<ShapeDialect,
   }];
 }
 
-def Shape_ShapeOrSizeType: AnyTypeOf<[Shape_SizeType, Shape_ShapeType],
+def Shape_ShapeOrSizeType : AnyTypeOf<[Shape_SizeType, Shape_ShapeType],
   "shape or size">;
 
+def Shape_ExtentTensorType : 1DTensorOf<[Index]>;
+
+def Shape_ShapeOrExtentTensorType : AnyTypeOf<[Shape_ShapeType,
+                                               Shape_ExtentTensorType],
+                                              "shape or extent tensor">;
+
 #endif // SHAPE_BASE_TD

diff  --git a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
index 6f263cd6db63..7b54616ad703 100644
--- a/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
+++ b/mlir/include/mlir/Dialect/Shape/IR/ShapeOps.td
@@ -138,6 +138,22 @@ def Shape_ConstSizeOp : Shape_Op<"const_size", [
   let hasFolder = 1;
 }
 
+def Shape_ShapeEqOp : Shape_Op<"shape_eq", [Commutative, NoSideEffect]> {
+  let summary = "Returns whether the input shapes or extent tensors are equal";
+  let description = [{
+    Takes two shape or extent tensor operands and determines whether they are
+    equal. When extent tensors are compared to shapes they are regarded as their
+    equivalent non-error shapes. Error shapes can be tested for equality like
+    any other shape value, meaning that the error value is equal to itself.
+  }];
+
+  let arguments = (ins Shape_ShapeOrExtentTensorType:$lhs,
+                       Shape_ShapeOrExtentTensorType:$rhs);
+  let results = (outs I1:$result);
+
+  let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs)";
+}
+
 def Shape_FromExtentsOp : Shape_Op<"from_extents", [NoSideEffect]> {
   let summary = "Creates a shape from extents";
   let description = [{

diff  --git a/mlir/test/Dialect/Shape/ops.mlir b/mlir/test/Dialect/Shape/ops.mlir
index 18953efa05f1..94323e856750 100644
--- a/mlir/test/Dialect/Shape/ops.mlir
+++ b/mlir/test/Dialect/Shape/ops.mlir
@@ -116,3 +116,18 @@ func @rank(%shape : !shape.shape) -> !shape.size {
   %rank = shape.rank %shape
   return %rank : !shape.size
 }
+
+func @shape_eq_on_shapes(%a : !shape.shape, %b : !shape.shape) -> i1 {
+  %result = shape.shape_eq %a, %b : !shape.shape, !shape.shape
+  return %result : i1
+}
+
+func @shape_eq_on_tensors(%a : tensor<?xindex>, %b : tensor<?xindex>) -> i1 {
+  %result = shape.shape_eq %a, %b : tensor<?xindex>, tensor<?xindex>
+  return %result : i1
+}
+
+func @shape_eq_on_mixed(%a : tensor<?xindex>, %b : !shape.shape) -> i1 {
+  %result = shape.shape_eq %a, %b : tensor<?xindex>, !shape.shape
+  return %result : i1
+}