Add a test.symbol_scope operation that has the SymbolTable Traits to the Test dialect

PiperOrigin-RevId: 277741687

4 files changed, 23 insertions, 0 deletions

diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 6e7e0e8bef8..d8ad2e4dc21 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -1285,6 +1285,8 @@ def SameOperandsAndResultElementType :
   NativeOpTrait<"SameOperandsAndResultElementType">;
 // Op is a symbol.
 def Symbol : NativeOpTrait<"Symbol">;
+// Op defines a symbol table.
+def SymbolTable : NativeOpTrait<"SymbolTable">;
 // Op is a terminator.
 def Terminator       : NativeOpTrait<"IsTerminator">;
 
diff --git a/mlir/test/IR/traits.mlir b/mlir/test/IR/traits.mlir
index 60c2afac44b..69804f1f4f1 100644
--- a/mlir/test/IR/traits.mlir
+++ b/mlir/test/IR/traits.mlir
@@ -205,3 +205,13 @@ func @failedSingleBlockImplicitTerminator_missing_terminator() {
   }) : () -> ()
 }
 
+// -----
+
+// Test that operation with the SymbolTable Trait define a new symbol scope.
+"test.symbol_scope"() ({
+  func @foo() {
+  }
+  "test.finish" () : () -> ()
+}) : () -> ()
+func @foo() {
+}
diff --git a/mlir/test/lib/TestDialect/TestDialect.h b/mlir/test/lib/TestDialect/TestDialect.h
index ade0eb81c40..5be1c108f23 100644
--- a/mlir/test/lib/TestDialect/TestDialect.h
+++ b/mlir/test/lib/TestDialect/TestDialect.h
@@ -30,6 +30,7 @@
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/StandardTypes.h"
+#include "mlir/IR/SymbolTable.h"
 
 namespace mlir {
 
diff --git a/mlir/test/lib/TestDialect/TestOps.td b/mlir/test/lib/TestDialect/TestOps.td
index 1ce13c79468..09603e0fc18 100644
--- a/mlir/test/lib/TestDialect/TestOps.td
+++ b/mlir/test/lib/TestDialect/TestOps.td
@@ -86,6 +86,16 @@ def MultiTensorRankOf : TEST_Op<"multi_tensor_rank_of"> {
 // Test Operands
 //===----------------------------------------------------------------------===//
 
+def SymbolScopeOp : TEST_Op<"symbol_scope",
+    [SymbolTable, SingleBlockImplicitTerminator<"TerminatorOp">]> {
+  let summary =  "operation which defines a new symbol table";
+  let regions = (region SizedRegion<1>:$region);
+}
+
+//===----------------------------------------------------------------------===//
+// Test Operands
+//===----------------------------------------------------------------------===//
+
 def MixedNormalVariadicOperandOp : TEST_Op<
     "mixed_normal_variadic_operand", [SameVariadicOperandSize]> {
   let arguments = (ins