[CONTRIB] Enable create_staticlib to take in tar files

python/tvm/contrib/cc.py

@@ -17,9 +17,11 @@
 """Util to invoke C/C++ compilers in the system."""
 # pylint: disable=invalid-name
 import sys
+import shutil
 import os
 import subprocess
 
+from . import utils as _utils, tar as _tar
 from .._ffi.base import py_str
 
 
@@ -83,6 +85,54 @@ def create_shared(output, objects, options=None, cc=None):
         raise ValueError("Unsupported platform")
 
 
+def _linux_ar(output, inputs, ar):
+    ar = ar or "ar"
+
+    libname = os.path.basename(output)
+    if not libname.startswith("lib"):
+        libname = "lib" + libname
+    temp = _utils.tempdir()
+    temp_output = temp.relpath(libname)
+    cmd = [ar, "-crs", temp_output]
+
+    # handles the case where some input files are tar of objects
+    # unpack them and return the list of files inside
+    objects = _tar.normalize_file_list_by_unpacking_tars(temp, inputs)
+
+    cmd += objects
+    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
+    (out, _) = proc.communicate()
+    if proc.returncode != 0:
+        msg = "AR error:\n"
+        msg += py_str(out)
+        msg += "\nCommand line: " + " ".join(cmd)
+        raise RuntimeError(msg)
+
+    shutil.move(temp_output, output)
+
+
+def create_staticlib(output, inputs, ar=None):
+    """Create static library.
+
+    Parameters
+    ----------
+    output : str
+        The target shared library.
+
+    inputs : List[str]
+        List of inputs files. Each input file can be a tarball
+        of objects or an object file.
+
+    ar : Optional[str]
+        Path to the ar command to be used
+    """
+
+    if _is_linux_like():
+        return _linux_ar(output, inputs, ar)
+    else:
+        raise ValueError("Unsupported platform")
+
+
 def create_executable(output, objects, options=None, cc=None):
     """Create executable binary.
 

python/tvm/contrib/tar.py

@@ -84,3 +84,44 @@ def untar(tar_file, directory):
         msg = "Tar error:\n"
         msg += py_str(out)
         raise RuntimeError(msg)
+
+
+def normalize_file_list_by_unpacking_tars(temp, file_list):
+    """Normalize the file list by unpacking tars in list.
+
+    When a filename is a tar, it will untar it into an unique dir
+    in temp and return the list of files in the tar.
+    When a filename is a normal file, it will be simply added to the list.
+
+    This is useful to untar objects in tar and then turn
+    them into a library.
+
+    Parameters
+    ----------
+    temp: tvm.contrib.utils.TempDirectory
+        A temp dir to hold the untared files.
+
+    file_list: List[str]
+        List of path
+
+    Returns
+    -------
+    ret_list: List[str]
+        An updated list of files
+    """
+    temp_count = 0
+    ret_list = []
+    for file_path in file_list:
+        # enable extracting a tarball
+        if file_path.endswith(".tar"):
+            temp_dir = temp.relpath(f"temp{temp_count}")
+            temp_count += 1
+            os.mkdir(temp_dir)
+            untar(file_path, temp_dir)
+            # append all files inside
+            for root, _, files in os.walk(temp_dir):
+                for file in files:
+                    ret_list.append(os.path.join(root, file))
+        else:
+            ret_list.append(file_path)
+    return ret_list

tests/python/unittest/test_target_codegen_blob.py

@@ -19,7 +19,7 @@
 import numpy as np
 from tvm import relay
 import tvm.relay.testing
-from tvm.contrib import graph_executor, cc, utils, popen_pool
+from tvm.contrib import graph_executor, cc, utils, popen_pool, tar
 import tvm
 import tvm.testing
 from tvm.script import ir as I, tir as T
@@ -101,13 +101,16 @@ def my_inplace_update(x: T.Buffer((12), "float32")) -> None:
     libA = tvm.build(ModA, target=target)
     libB = tvm.build(ModB, target=target)
 
-    pathA = temp.relpath("libA.a")
-    pathB = temp.relpath("libB.a")
+    pathA = temp.relpath("libA.tar")
+    pathB = temp.relpath("libB.tar")
+    pathAll = temp.relpath("libAll.a")
+
     path_dso = temp.relpath("mylib.so")
-    libA.export_library(pathA, cc.create_staticlib)
-    libB.export_library(pathB, cc.create_staticlib)
+    libA.export_library(pathA, tar.tar)
+    libB.export_library(pathB, tar.tar)
+    cc.create_staticlib(pathAll, [pathA, pathB])
     # package two static libs together
-    cc.create_shared(path_dso, ["-Wl,--whole-archive", pathA, pathB, "-Wl,--no-whole-archive"])
+    cc.create_shared(path_dso, ["-Wl,--whole-archive", pathAll, "-Wl,--no-whole-archive"])
 
     def popen_check():
         # Load dll, will trigger system library registration