From b9eecb7b8f199532a8b236a4d64e1fdba5efa42b Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer@nvidia.com>
Date: Thu, 6 Feb 2025 15:29:16 +0100
Subject: [PATCH] [mlir][IR] Clean up type constraints around
 `ValueSemanticsContainerOf`

---
 .../include/mlir/Dialect/Arith/IR/ArithOps.td | 17 +++-------
 .../Dialect/Polynomial/IR/PolynomialTypes.td  |  3 +-
 mlir/include/mlir/IR/CommonTypeConstraints.td | 33 ++++---------------
 3 files changed, 14 insertions(+), 39 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
index ea9b0f6509b80..d50b6aeca15c9 100644
--- a/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
+++ b/mlir/include/mlir/Dialect/Arith/IR/ArithOps.td
@@ -144,13 +144,6 @@ class Arith_CompareOp<string mnemonic, list<Trait> traits = []> :
   let assemblyFormat = "$predicate `,` $lhs `,` $rhs attr-dict `:` type($lhs)";
 }
 
-// Just like `Arith_CompareOp` but also admits 0-D vectors. Introduced
-// temporarily to allow gradual transition to 0-D vectors.
-class Arith_CompareOpOfAnyRank<string mnemonic, list<Trait> traits = []> :
-    Arith_CompareOp<mnemonic, traits> {
-  let results = (outs BoolLikeOfAnyRank:$result);
-}
-
 class Arith_IntBinaryOpWithOverflowFlags<string mnemonic, list<Trait> traits = []> :
     Arith_BinaryOp<mnemonic, traits #
       [Pure, DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>,
@@ -1426,9 +1419,9 @@ def Arith_BitcastOp : Arith_CastOp<"bitcast", BitcastTypeConstraint,
 // CmpIOp
 //===----------------------------------------------------------------------===//
 
-def Arith_CmpIOp
-  : Arith_CompareOpOfAnyRank<"cmpi",
-                             [DeclareOpInterfaceMethods<InferIntRangeInterface, ["inferResultRanges"]>]> {
+def Arith_CmpIOp : Arith_CompareOp<"cmpi",
+    [DeclareOpInterfaceMethods<InferIntRangeInterface,
+                               ["inferResultRanges"]>]> {
   let summary = "integer comparison operation";
   let description = [{
     The `cmpi` operation is a generic comparison for integer-like types. Its two
@@ -1495,8 +1488,8 @@ def Arith_CmpIOp
   }];
 
   let arguments = (ins Arith_CmpIPredicateAttr:$predicate,
-                       SignlessIntegerOrIndexLikeOfAnyRank:$lhs,
-                       SignlessIntegerOrIndexLikeOfAnyRank:$rhs);
+                       SignlessIntegerOrIndexLike:$lhs,
+                       SignlessIntegerOrIndexLike:$rhs);
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
diff --git a/mlir/include/mlir/Dialect/Polynomial/IR/PolynomialTypes.td b/mlir/include/mlir/Dialect/Polynomial/IR/PolynomialTypes.td
index 89e406183e0b0..cf33503764abb 100644
--- a/mlir/include/mlir/Dialect/Polynomial/IR/PolynomialTypes.td
+++ b/mlir/include/mlir/Dialect/Polynomial/IR/PolynomialTypes.td
@@ -26,7 +26,8 @@ def Polynomial_PolynomialType : Polynomial_Type<"Polynomial", "polynomial"> {
   let assemblyFormat = "`<` struct(params) `>`";
 }
 
-def PolynomialLike: TypeOrContainer<Polynomial_PolynomialType, "polynomial-like">;
+def PolynomialLike : TypeOrValueSemanticsContainer<
+    Polynomial_PolynomialType, "polynomial-like">;
 
 
 #endif // POLYNOMIAL_TYPES
diff --git a/mlir/include/mlir/IR/CommonTypeConstraints.td b/mlir/include/mlir/IR/CommonTypeConstraints.td
index 82e335e30b6fa..a18b32253d857 100644
--- a/mlir/include/mlir/IR/CommonTypeConstraints.td
+++ b/mlir/include/mlir/IR/CommonTypeConstraints.td
@@ -879,12 +879,6 @@ class NestedTupleOf<list<Type> allowedTypes> :
 //===----------------------------------------------------------------------===//
 // Common type constraints
 //===----------------------------------------------------------------------===//
-// Type constraint for types that are "like" some type or set of types T, that is
-// they're either a T, a vector of Ts, or a tensor of Ts.
-class TypeOrContainer<Type allowedType, string name> : TypeConstraint<Or<[
-  allowedType.predicate,
-  ValueSemanticsContainerOf<[allowedType]>.predicate]>,
-  name>;
 
 // Type constraint for types that are "like" some type or set of types T, that is
 // they're either a T or a mapable container of Ts.
@@ -894,36 +888,23 @@ class TypeOrValueSemanticsContainer<Type allowedType, string name>
   ValueSemanticsContainerOf<[allowedType]>.predicate]>,
   name>;
 
-// Temporary constraint to allow gradual transition to supporting 0-D vectors.
-// TODO: Remove this when all ops support 0-D vectors.
-class TypeOrContainerOfAnyRank<Type allowedType, string name> : TypeConstraint<Or<[
-  allowedType.predicate, VectorOfAnyRankOf<[allowedType]>.predicate,
-  TensorOf<[allowedType]>.predicate]>,
-  name>;
-
-
 // Type constraint for bool-like types: bools, vectors of bools, tensors of
 // bools.
-def BoolLike : TypeOrContainer<I1, "bool-like">;
+def BoolLike : TypeOrValueSemanticsContainer<I1, "bool-like">;
 
-def BoolLikeOfAnyRank : TypeOrContainerOfAnyRank<I1, "bool-like">;
-
-// Type constraint for signless-integer-like types: signless integers, 
-// vectors of signless integers or tensors of signless integers.
+// Type constraint for signless-integer-like types: signless integers or
+// value-semantics containers of signless integers.
 def SignlessIntegerLike : TypeOrValueSemanticsContainer<
     AnySignlessInteger, "signless-integer">;
 
 // Type constraint for signless-integer-like types: signless integers, indices,
-// vectors of signless integers or indices, tensors of signless integers.
+// or value-semantics containers of signless integers or indices.
 def SignlessIntegerOrIndexLike : TypeOrValueSemanticsContainer<
     AnySignlessIntegerOrIndex, "signless-integer-like">;
 
-def SignlessIntegerOrIndexLikeOfAnyRank : TypeOrContainerOfAnyRank<
-    AnySignlessIntegerOrIndex,
-    "signless-integer-like">;
-
-// Type constraint for float-like types: floats, vectors or tensors thereof.
-def FloatLike : TypeOrContainer<AnyFloat, "floating-point-like">;
+// Type constraint for float-like types: floats or value-semantics containers
+// of floats.
+def FloatLike : TypeOrValueSemanticsContainer<AnyFloat, "floating-point-like">;
 
 // Type constraint for signless-integer-or-index-like or float-like types.
 def SignlessIntegerOrFloatLike : TypeConstraint<Or<[