From 8a1dbbd597fd52199802513eb1fe0abbb2cf36ef Mon Sep 17 00:00:00 2001 From: Yuanqiang Liu Date: Wed, 24 Apr 2024 11:34:02 +0800 Subject: [PATCH] [torchscript] export extra library file name to user (#3203) * so that it could be specified by user. --- projects/pt1/python/torch_mlir/torchscript.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/projects/pt1/python/torch_mlir/torchscript.py b/projects/pt1/python/torch_mlir/torchscript.py index 48b9066d2..acb487319 100644 --- a/projects/pt1/python/torch_mlir/torchscript.py +++ b/projects/pt1/python/torch_mlir/torchscript.py @@ -9,6 +9,7 @@ from enum import Enum import sys from io import StringIO import tempfile +import os from torch._functorch.compile_utils import strip_overloads import torch @@ -253,19 +254,20 @@ BACKEND_LEGAL_OPS = { } -def _canon_extra_library(extra_library): - extra_library_file_name = "" +def _canon_extra_library(extra_library, extra_library_file_name="custom_op_extra_library.mlir"): if len(extra_library) != 0: extra_library_dict = {} for library_func in extra_library: extra_library_dict[library_func.__name__] = library_func mlir_library = generate_library(extra_library_dict) - extra_library_file_name = \ - tempfile.gettempdir() + "/custom_op_extra_library.mlir" - with open(extra_library_file_name, "w") as f: + extra_library_file = \ + os.path.join(tempfile.gettempdir(), extra_library_file_name) + with open(extra_library_file, "w") as f: f.write(mlir_library) - return extra_library_file_name + return extra_library_file + else: + return "" def _lower_mlir_module(verbose, output_type, module):