blob: 5020c00344a339d60e59f11bc900f85852702ce2 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
|
# UNSUPPORTED: target=aarch64{{.*}}, target=arm64{{.*}}
# RUN: %PYTHON %s 2>&1 | FileCheck %s
# REQUIRES: host-supports-jit
import gc, sys, os, tempfile
from mlir.ir import *
from mlir.passmanager import *
from mlir.execution_engine import *
from mlir.runtime import *
# Log everything to stderr and flush so that we have a unified stream to match
# errors/info emitted by MLIR to stderr.
def log(*args):
print(*args, file=sys.stderr)
sys.stderr.flush()
def run(f):
log("\nTEST:", f.__name__)
f()
gc.collect()
assert Context._get_live_count() == 0
def lowerToLLVM(module):
pm = PassManager.parse(
"builtin.module(convert-func-to-llvm,reconcile-unrealized-casts)"
)
pm.run(module.operation)
return module
# Test JIT callback in global constructor
# CHECK-LABEL: TEST: testJITCallbackInGlobalCtor
def testJITCallbackInGlobalCtor():
init_cnt = 0
@ctypes.CFUNCTYPE(None)
def initCallback():
nonlocal init_cnt
init_cnt += 1
with Context():
module = Module.parse(
r"""
llvm.mlir.global_ctors ctors = [@ctor], priorities = [0 : i32], data = [#llvm.zero]
llvm.func @ctor() {
func.call @init_callback() : () -> ()
llvm.return
}
func.func private @init_callback() attributes { llvm.emit_c_interface }
"""
)
# Setup execution engine
execution_engine = ExecutionEngine(lowerToLLVM(module))
# Validate initialization hasn't run yet
assert init_cnt == 0
# # Register callback
execution_engine.register_runtime("init_callback", initCallback)
# # Initialize and verify
execution_engine.initialize()
assert init_cnt == 1
# # Second initialization should be no-op
execution_engine.initialize()
assert init_cnt == 1
run(testJITCallbackInGlobalCtor)
|
