llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[torchscript] export extra library file name to user (#3203) 

* so that it could be specified by user.
pull/3222/head
Yuanqiang Liu 2024-04-24 11:34:02 +08:00 committed by GitHub
parent f77d88390a
commit 8a1dbbd597
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
projects/pt1/python/torch_mlir/torchscript.py
View File

@ -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):