[AOT] Save AOT module artifacts as zip archive (#5316)

* Save AOT module artifacts as zip archive

* Add version mark

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix linting

* Use tempfile for temp dir management

* Added test for AOT module archiving

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: PENGUINLIONG

python/taichi/aot/module.py

@@ -1,5 +1,9 @@
 from contextlib import contextmanager
+from glob import glob
 from pathlib import Path, PurePosixPath
+from shutil import rmtree
+from tempfile import mkdtemp
+from zipfile import ZipFile
 
 from taichi.aot.utils import (produce_injected_args,
                               produce_injected_args_from_template)
@@ -8,6 +12,8 @@
 from taichi.lang.matrix import MatrixField
 from taichi.types.annotations import template
 
+import taichi
+
 
 class KernelTemplate:
     def __init__(self, kernel_fn, aot_module):
@@ -88,6 +94,7 @@ def __init__(self, arch):
         rtm = impl.get_runtime()
         rtm._finalize_root_fb_for_aot()
         self._aot_builder = rtm.prog.make_aot_module_builder(arch)
+        self._content = []
 
     def add_field(self, name, field):
         """Add a taichi field to the AOT module.
@@ -147,8 +154,11 @@ def add_kernel(self, kernel_fn, template_args=None, name=None):
         # kernel AOT
         self._kernels.append(kernel)
 
+        self._content += ["kernel:" + kernel_name]
+
     def add_graph(self, name, graph):
         self._aot_builder.add_graph(name, graph._compiled_graph)
+        self._content += ["cgraph:" + name]
 
     @contextmanager
     def add_kernel_template(self, kernel_fn):
@@ -193,3 +203,29 @@ def save(self, filepath, filename):
         """
         filepath = str(PurePosixPath(Path(filepath)))
         self._aot_builder.dump(filepath, filename)
+
+    def archive(self, filepath: str):
+        """
+        Args:
+          filepath (str): path to the stored archive of aot artifacts, MUST
+            end with `.tcm`.
+        """
+        assert filepath.endswith(".tcm"), \
+            "AOT module artifact archive must ends with .tcm"
+        tcm_path = Path(filepath).absolute()
+        assert tcm_path.parent.exists(), "Output directory doesn't exist"
+
+        temp_dir = mkdtemp(prefix="tcm_")
+        # Save first as usual.
+        self.save(temp_dir, "")
+
+        # Package all artifacts into a zip archive and attach contend data.
+        with ZipFile(tcm_path, "w") as z:
+            z.writestr("__content__", '\n'.join(self._content))
+            z.writestr("__version__",
+                       '.'.join(str(x) for x in taichi.__version__))
+            for path in glob(f"{temp_dir}/*", recursive=True):
+                z.write(path, Path.relative_to(Path(path), temp_dir))
+
+        # Remove cached files
+        rmtree(temp_dir)

tests/python/test_aot.py

@@ -2,6 +2,7 @@
 import os
 import sys
 import tempfile
+import zipfile
 
 import numpy as np
 import pytest
@@ -555,3 +556,24 @@ def run(arr: ti.types.ndarray(), val1: ti.f32, val2: ti.template()):
             res = json.load(json_file)
             args_count = res['aot_data']['kernels']['run']['args_count']
             assert args_count == 2, res  # `arr` and `val1`
+
+
+@test_utils.test(arch=[ti.vulkan])
+def test_archive():
+    density = ti.field(float, shape=(4, 4))
+
+    @ti.kernel
+    def init():
+        for i, j in density:
+            density[i, j] = 1
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # note ti.aot.Module(ti.opengl) is no-op according to its docstring.
+        m = ti.aot.Module(ti.lang.impl.current_cfg().arch)
+        m.add_field('density', density)
+        m.add_kernel(init)
+        tcm_path = f"{tmpdir}/x.tcm"
+        m.archive(tcm_path)
+        with zipfile.ZipFile(tcm_path, 'r') as z:
+            assert z.read("__version__") == bytes(
+                '.'.join(str(x) for x in ti.__version__), 'utf-8')