183 lines · python
1# RUN: %PYTHON %s | FileCheck %s2 3import gc4import io5import itertools6from mlir.ir import *7 8 9def run(f):10 print("\nTEST:", f.__name__)11 f()12 gc.collect()13 assert Context._get_live_count() == 014 return f15 16 17# CHECK-LABEL: TEST: testSymbolTableInsert18@run19def testSymbolTableInsert():20 with Context() as ctx:21 ctx.allow_unregistered_dialects = True22 m1 = Module.parse(23 """24 func.func private @foo()25 func.func private @bar()"""26 )27 m2 = Module.parse(28 """29 func.func private @qux()30 func.func private @foo()31 "foo.bar"() : () -> ()"""32 )33 34 symbol_table = SymbolTable(m1.operation)35 36 # CHECK: func private @foo37 # CHECK: func private @bar38 assert "foo" in symbol_table39 print(symbol_table["foo"])40 assert "bar" in symbol_table41 bar = symbol_table["bar"]42 print(symbol_table["bar"])43 44 assert "qux" not in symbol_table45 46 del symbol_table["bar"]47 try:48 symbol_table.erase(symbol_table["bar"])49 except KeyError:50 pass51 else:52 assert False, "expected KeyError"53 54 # CHECK: module55 # CHECK: func private @foo()56 print(m1)57 assert "bar" not in symbol_table58 59 bar._set_invalid()60 try:61 print(bar)62 except RuntimeError as e:63 if "the operation has been invalidated" not in str(e):64 raise65 else:66 assert False, "expected RuntimeError due to invalidated operation"67 68 qux = m2.body.operations[0]69 m1.body.append(qux)70 symbol_table.insert(qux)71 assert "qux" in symbol_table72 73 # Check that insertion actually renames this symbol in the symbol table.74 foo2 = m2.body.operations[0]75 m1.body.append(foo2)76 updated_name = symbol_table.insert(foo2)77 assert foo2.name.value != "foo"78 assert foo2.name == updated_name79 assert isinstance(updated_name, StringAttr)80 81 # CHECK: module82 # CHECK: func private @foo()83 # CHECK: func private @qux()84 # CHECK: func private @foo{{.*}}85 print(m1)86 87 try:88 symbol_table.insert(m2.body.operations[0])89 except ValueError as e:90 if "Expected operation to have a symbol name" not in str(e):91 raise92 else:93 assert False, "exepcted ValueError when adding a non-symbol"94 95 96# CHECK-LABEL: testSymbolTableRAUW97@run98def testSymbolTableRAUW():99 with Context() as ctx:100 m = Module.parse(101 """102 func.func private @foo() {103 call @bar() : () -> ()104 return105 }106 func.func private @bar()107 """108 )109 foo, bar = list(m.operation.regions[0].blocks[0].operations)[0:2]110 111 # Do renaming just within `foo`.112 SymbolTable.set_symbol_name(bar, "bam")113 SymbolTable.replace_all_symbol_uses("bar", "bam", foo)114 # CHECK: call @bam()115 # CHECK: func private @bam116 print(m)117 # CHECK: Foo symbol: StringAttr("foo")118 # CHECK: Bar symbol: StringAttr("bam")119 print(f"Foo symbol: {repr(SymbolTable.get_symbol_name(foo))}")120 print(f"Bar symbol: {repr(SymbolTable.get_symbol_name(bar))}")121 122 # Do renaming within the module.123 SymbolTable.set_symbol_name(bar, "baz")124 SymbolTable.replace_all_symbol_uses("bam", "baz", m.operation)125 # CHECK: call @baz()126 # CHECK: func private @baz127 print(m)128 # CHECK: Foo symbol: StringAttr("foo")129 # CHECK: Bar symbol: StringAttr("baz")130 print(f"Foo symbol: {repr(SymbolTable.get_symbol_name(foo))}")131 print(f"Bar symbol: {repr(SymbolTable.get_symbol_name(bar))}")132 133 134# CHECK-LABEL: testSymbolTableVisibility135@run136def testSymbolTableVisibility():137 with Context() as ctx:138 m = Module.parse(139 """140 func.func private @foo() {141 return142 }143 """144 )145 foo = m.operation.regions[0].blocks[0].operations[0]146 # CHECK: Existing visibility: StringAttr("private")147 print(f"Existing visibility: {repr(SymbolTable.get_visibility(foo))}")148 SymbolTable.set_visibility(foo, "public")149 # CHECK: func public @foo150 print(m)151 152 153# CHECK: testWalkSymbolTables154@run155def testWalkSymbolTables():156 with Context() as ctx:157 m = Module.parse(158 """159 module @outer {160 module @inner{161 }162 }163 """164 )165 166 def callback(symbol_table_op, uses_visible):167 print(f"SYMBOL TABLE: {uses_visible}: {symbol_table_op}")168 169 # CHECK: SYMBOL TABLE: True: module @inner170 # CHECK: SYMBOL TABLE: True: module @outer171 SymbolTable.walk_symbol_tables(m.operation, True, callback)172 173 # Make sure exceptions in the callback are handled.174 def error_callback(symbol_table_op, uses_visible):175 assert False, "Raised from python"176 177 try:178 SymbolTable.walk_symbol_tables(m.operation, True, error_callback)179 except RuntimeError as e:180 # CHECK: GOT EXCEPTION: Exception raised in callback:181 # CHECK: AssertionError: Raised from python182 print(f"GOT EXCEPTION: {e}")183