Add ReferToOp attribute constraint for SymbolRefAttr

This constraint can be used to limit a SymbolRefAttr to point
to a specific kind of op in the closest parent with a symbol table.

PiperOrigin-RevId: 278001364

2 files changed, 12 insertions, 0 deletions

diff --git a/mlir/test/lib/TestDialect/TestDialect.cpp b/mlir/test/lib/TestDialect/TestDialect.cpp
index 496cfc6f627..5528e8f0daa 100644
--- a/mlir/test/lib/TestDialect/TestDialect.cpp
+++ b/mlir/test/lib/TestDialect/TestDialect.cpp
@@ -16,6 +16,8 @@
 // =============================================================================
 
 #include "TestDialect.h"
+#include "mlir/IR/Function.h"
+#include "mlir/IR/Module.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/FoldUtils.h"
diff --git a/mlir/test/lib/TestDialect/TestOps.td b/mlir/test/lib/TestDialect/TestOps.td
index 94b0250a436..cc1e22278ba 100644
--- a/mlir/test/lib/TestDialect/TestOps.td
+++ b/mlir/test/lib/TestDialect/TestOps.td
@@ -201,6 +201,16 @@ def UpdateFloatElementsAttr : Pat<
     $f64attr)>;
 
 //===----------------------------------------------------------------------===//
+// Test Attribute Constraints
+//===----------------------------------------------------------------------===//
+
+def SymbolRefOp : TEST_Op<"symbol_ref_attr"> {
+  let arguments = (ins
+    Confined<SymbolRefAttr, [ReferToOp<"FuncOp">]>:$symbol
+  );
+}
+
+//===----------------------------------------------------------------------===//
 // Test Regions
 //===----------------------------------------------------------------------===//