[TableGen] Add a `StaticShapeMemRefOf` trait.
The trait specifies that the `MemRefOf` has to have a static shape.
PiperOrigin-RevId: 264692758
diff --git a/include/mlir/IR/OpBase.td b/include/mlir/IR/OpBase.td
index fc909c5..37a0867 100644
--- a/include/mlir/IR/OpBase.td
+++ b/include/mlir/IR/OpBase.td
@@ -417,6 +417,13 @@
def F32MemRef : MemRefOf<[F32]>;
def F64MemRef : MemRefOf<[F64]>;
+// TODO(b/130064155) Have an easy way to add another constraint to a type.
+class StaticShapeMemRefOf<list<Type> allowedTypes>
+ : Type<And<[MemRefOf<allowedTypes>.predicate, HasStaticShapePred]>,
+ "statically shaped " # MemRefOf<allowedTypes>.description>;
+
+def AnyStaticShapeMemRef : StaticShapeMemRefOf<[AnyType]>;
+
// This represents a generic tuple without any constraints on element type.
def AnyTuple : Type<IsTupleTypePred, "tuple">;
diff --git a/test/lib/TestDialect/TestOps.td b/test/lib/TestDialect/TestOps.td
index 2efeb9f..976f6ac 100644
--- a/test/lib/TestDialect/TestOps.td
+++ b/test/lib/TestDialect/TestOps.td
@@ -46,6 +46,10 @@
let results = (outs NestedTupleOf<[I32, F32]>);
}
+def TakesStaticMemRefOp : TEST_Op<"takes_static_memref"> {
+ let arguments = (ins AnyStaticShapeMemRef:$x);
+}
+
//===----------------------------------------------------------------------===//
// Test Operands
//===----------------------------------------------------------------------===//
diff --git a/test/mlir-tblgen/types.mlir b/test/mlir-tblgen/types.mlir
index 4a94468..f51446a 100644
--- a/test/mlir-tblgen/types.mlir
+++ b/test/mlir-tblgen/types.mlir
@@ -182,3 +182,10 @@
"test.if_first_operand_is_none_then_so_is_second"(%arg1, %arg0) : (none, tensor<1x2xi32>) -> ()
return
}
+
+// -----
+
+func @does_not_have_static_memref(%arg0: memref<?xi32>) {
+ // expected-error@+1 {{'test.takes_static_memref' op operand #0 must be statically shaped memref of any type values}}
+ "test.takes_static_memref"(%arg0) : (memref<?xi32>) -> ()
+}
