From 5ff217d512e784c097dba011d2945cb69709952d Mon Sep 17 00:00:00 2001 From: Emilien Bauer Date: Mon, 7 Aug 2023 21:26:24 +0100 Subject: [PATCH] Add direct .so backdoor capability, and XDSL_SKIP_CLEAN env variable to skip temp file cleanup. --- devito/operator/xdsl_operator.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/devito/operator/xdsl_operator.py b/devito/operator/xdsl_operator.py index 31ab786a0c4..8188ad39b3e 100644 --- a/devito/operator/xdsl_operator.py +++ b/devito/operator/xdsl_operator.py @@ -12,7 +12,7 @@ from io import StringIO from devito.exceptions import InvalidOperator -from devito.logger import perf +from devito.logger import perf, info from devito.ir.iet import Callable, MetaCall from devito.ir.support import SymbolRegistry from devito.operator.operator import IRs @@ -62,8 +62,9 @@ class XDSLOperator(Operator): def __new__(cls, expressions, **kwargs): self = super(XDSLOperator, cls).__new__(cls, expressions, **kwargs) - self._tf = tempfile.NamedTemporaryFile(prefix="devito-jit-", suffix='.so') - self._interop_tf = tempfile.NamedTemporaryFile(prefix="devito-jit-interop-", suffix=".o") + delete=not os.getenv("XDSL_SKIP_CLEAN", False) + self._tf = tempfile.NamedTemporaryFile(prefix="devito-jit-", suffix='.so', delete=delete) + self._interop_tf = tempfile.NamedTemporaryFile(prefix="devito-jit-interop-", suffix=".o", delete=delete) self._make_interop_o() self.__class__ = cls return self @@ -140,6 +141,10 @@ def _jit_compile(self): # allow jit backdooring to provide your own xdsl code backdoor = os.getenv('XDSL_JIT_BACKDOOR') if backdoor is not None: + if os.path.splitext(backdoor)[1] == ".so": + info(f"JIT Backdoor: skipping compilation and using {backdoor}") + self._tf.name = backdoor + return print("JIT Backdoor: loading xdsl file from: " + backdoor) with open(backdoor, 'r') as f: module_str = f.read()