Make TupleOf take a list of allowed types.

    This better matches other container types. Seems better to do that, even though tuples are a little different, since they don't have a single element type.

    Also fixed its description to mention the element type.

--

PiperOrigin-RevId: 249730341

diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 6cffc49..cc6999f 100644 (file)
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -414,23 +414,34 @@ def F64MemRef  : MemRefOf<[F64]>;
 // This represents a generic tuple without any constraints on element type.
 def AnyTuple : Type<IsTupleTypePred, "tuple">;
 
-// TODO(b/132952417) Make this accept a list of types like the classes above.
-// A Tuple that only holds elements of a certain type. This cannot inherit from
-// ContainerType because tuples do not always have a single element type that
-// could be retrieved with elementTypeCall.
-class TupleOf<Type t> :
-    Type<And<[
-        IsTupleTypePred,
-        Concat<
-            [{
-              llvm::all_of(
-                $_self.cast<TupleType>().getTypes(),
-                [](Type t) {
-                  return
-            }],
-            SubstLeaves<"$_self", "t", t.predicate>,
-            "; })">
-    ]>, "tuple">;
+// TODO(b/130358239) Support typed tuples of arbitrary nesting.
+
+// A container type that has other types embedded in it, but (unlike
+// ContainerType) can hold elements with a mix of types. Requires a call that
+// produces a list of all elements' types.
+class MixedContainerType<Type etype, Pred containerPred, code elementTypesCall,
+                         string descr> :
+    Type<
+        And<[
+            containerPred,
+            Concat<
+                "llvm::all_of(" # elementTypesCall # ", [](Type t) { return ",
+                SubstLeaves<"$_self", "t", etype.predicate>,
+                "; })"
+            >
+        ]>,
+        descr # " with any combination of " # etype.description # " values"> {
+  // The type of elements in the container.
+  Type elementType = etype;
+
+  // Call to retrieve.
+  code getElementTypesCall = elementTypesCall;
+}
+
+// A Tuple that holds a mix of elements of the allowed types.
+class TupleOf<list<Type> allowedTypes>
+    : MixedContainerType<AnyTypeOf<allowedTypes>, IsTupleTypePred,
+                         "$_self.cast<TupleType>().getTypes()", "tuple">;
 
 //===----------------------------------------------------------------------===//
 // Common type constraints

diff --git a/mlir/test/TestDialect/TestDialect.h b/mlir/test/TestDialect/TestDialect.h
index 4746914..55151df 100644 (file)
--- a/mlir/test/TestDialect/TestDialect.h
+++ b/mlir/test/TestDialect/TestDialect.h
@@ -25,6 +25,7 @@
 
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/StandardTypes.h"
 
 namespace mlir {
 

diff --git a/mlir/test/TestDialect/TestOps.td b/mlir/test/TestDialect/TestOps.td
index c571f12..9bdfb63 100644 (file)
--- a/mlir/test/TestDialect/TestOps.td
+++ b/mlir/test/TestDialect/TestOps.td
@@ -46,4 +46,12 @@ def VUVFoldTwoResultOp : Pattern<(VUVTwoResultOp $input), [
         (replaceWithValue $input)
       ]>;
 
+//===----------------------------------------------------------------------===//
+// Test Types
+//===----------------------------------------------------------------------===//
+
+def TupleOp : TEST_Op<"tuple_32_bit"> {
+  let results = (outs TupleOf<[I32, F32]>);
+}
+
 #endif // TEST_OPS
\ No newline at end of file

diff --git a/mlir/test/TestDialect/tests/types.mlir b/mlir/test/TestDialect/tests/types.mlir
new file mode 100644 (file)
index 0000000..b402e46
--- /dev/null
+++ b/mlir/test/TestDialect/tests/types.mlir
@@ -0,0 +1,40 @@
+// RUN: mlir-test-opt %s -split-input-file -verify | FileCheck %s
+
+// -----
+
+// CHECK-LABEL: @tuple_success
+func @tuple_success() {
+  %0 = "test.tuple_32_bit"() : () -> (tuple<i32>)
+  return
+}
+
+// -----
+
+// CHECK-LABEL: @tuple_mixed_success
+func @tuple_mixed_success() {
+  %0 = "test.tuple_32_bit"() : () -> (tuple<i32, f32>)
+  return
+}
+
+// -----
+
+func @tuple_empty_success() {
+  %0 = "test.tuple_32_bit"() : () -> (tuple<>)
+  return
+}
+
+// -----
+
+func @tuple_wrong_type_scalar() {
+  // expected-error@+1 {{must be tuple with any combination of 32-bit integer or 32-bit float values}}
+  %0 = "test.tuple_32_bit"() : () -> (tuple<i64>)
+  return
+}
+
+// -----
+
+func @tuple_wrong_type_tensor() {
+  // expected-error@+1 {{must be tuple with any combination of 32-bit integer or 32-bit float values}}
+  %0 = "test.tuple_32_bit"() : () -> (tuple<tensor<i32>>)
+  return
+}
\ No newline at end of file