73 lines · python
1# UNSUPPORTED: target=aarch64{{.*}}, target=arm64{{.*}}2# RUN: %PYTHON %s 2>&1 | FileCheck %s3# REQUIRES: host-supports-jit4import gc, sys, os, tempfile5from mlir.ir import *6from mlir.passmanager import *7from mlir.execution_engine import *8from mlir.runtime import *9 10 11# Log everything to stderr and flush so that we have a unified stream to match12# errors/info emitted by MLIR to stderr.13def log(*args):14 print(*args, file=sys.stderr)15 sys.stderr.flush()16 17 18def run(f):19 log("\nTEST:", f.__name__)20 f()21 gc.collect()22 assert Context._get_live_count() == 023 24 25def lowerToLLVM(module):26 pm = PassManager.parse(27 "builtin.module(convert-func-to-llvm,reconcile-unrealized-casts)"28 )29 pm.run(module.operation)30 return module31 32 33# Test JIT callback in global constructor34# CHECK-LABEL: TEST: testJITCallbackInGlobalCtor35def testJITCallbackInGlobalCtor():36 init_cnt = 037 38 @ctypes.CFUNCTYPE(None)39 def initCallback():40 nonlocal init_cnt41 init_cnt += 142 43 with Context():44 module = Module.parse(45 r"""46llvm.mlir.global_ctors ctors = [@ctor], priorities = [0 : i32], data = [#llvm.zero]47llvm.func @ctor() {48 func.call @init_callback() : () -> ()49 llvm.return50}51func.func private @init_callback() attributes { llvm.emit_c_interface }52 """53 )54 55 # Setup execution engine56 execution_engine = ExecutionEngine(lowerToLLVM(module))57 58 # Validate initialization hasn't run yet59 assert init_cnt == 060 61 # # Register callback62 execution_engine.register_runtime("init_callback", initCallback)63 64 # # Initialize and verify65 execution_engine.initialize()66 assert init_cnt == 167 # # Second initialization should be no-op68 execution_engine.initialize()69 assert init_cnt == 170 71 72run(testJITCallbackInGlobalCtor)73