7610b13729
Retry landing https://github.com/llvm/llvm-project/pull/153373 ## Major changes from previous attempt - remove the test in CAPI because no existing tests in CAPI deal with sanitizer exemptions - update `mlir/docs/Dialects/GPU.md` to reflect the new behavior: load GPU binary in global ctors, instead of loading them at call site. - skip the test on Aarch64 since we have an issue with initialization there --------- Co-authored-by: Mehdi Amini <joker.eph@gmail.com>
73 lines
1.8 KiB
Python
73 lines
1.8 KiB
Python
# 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)
|