Merge pull request #133 from vatai/132-insert-polly-in-middle-of-o3-list

132 insert polly in middle of o3 list

Author: vatai

examples/inputs/fdepnodep.f90

@@ -1,6 +1,6 @@
 program main
   implicit none
-  integer, parameter :: N = 1000
+  integer, parameter :: N = 3000
   real(8) :: A(N, N)
   integer :: i, count_rate, count_start, count_end
   real(8) :: walltime

tadashi/isl.pxd

@@ -76,6 +76,7 @@ cdef extern from "isl/union_map.h":
     isl_union_map *isl_union_map_intersect_domain(isl_union_map *m1, isl_union_map *m2)
     isl_union_map *isl_union_map_intersect_domain_union_set(isl_union_map *map, isl_union_set *set)
     const char *isl_union_map_to_str(isl_union_map *map)
+ctypedef isl_union_map* union_map
 
 # ---
 # aff

tadashi/passesparser.py

@@ -0,0 +1,126 @@
+#!/bin/env python
+import subprocess
+from pprint import pprint
+from typing import Optional
+
+
+class PassParser:
+    _pass_tree: Optional[list[str | tuple]]
+
+    def __init__(self):
+        cmd = [
+            "opt",
+            "-S",
+            "-O3",
+            "--print-pipeline-passes",
+            "/dev/null",
+            "-o",
+            "/dev/null",
+        ]
+        proc = subprocess.run(cmd, capture_output=True, check=True)
+        self.passes_str = proc.stdout.decode().strip()
+        self._pass_tree = None
+
+    def pass_tree(self) -> list:
+        if self._pass_tree is None:
+            self._pass_tree = self.parse(0, len(self.passes_str))
+        return self._pass_tree
+
+    def parse(self, begin: int, end: int) -> list:
+        cur = begin
+        results = []
+        while cur < end:
+            if self.passes_str[cur] == ",":
+                results.append(self.passes_str[begin:cur])
+                begin = cur + 1
+            elif self.passes_str[cur] == "(":
+                key = self.passes_str[begin:cur]
+                begin = cur + 1
+                cur = self._closing_cur(cur)
+                assert self.passes_str[cur - 1] == ")"
+                subtree = self.parse(begin, cur - 1)
+                results.append((key, subtree))
+                begin = cur + 1
+            cur += 1
+        if begin < cur:
+            results.append(self.passes_str[begin:cur])
+        return results
+
+    def _closing_cur(self, cur: int):
+        cur += 1
+        num_open = 1
+        while num_open != 0:
+            if self.passes_str[cur] == "(":
+                num_open += 1
+            if self.passes_str[cur] == ")":
+                num_open -= 1
+            cur += 1
+        return cur
+
+    @staticmethod
+    def reassemble(passes: list | tuple):
+        if isinstance(passes, list):
+            flat = [
+                p if isinstance(p, str) else PassParser.reassemble(p) for p in passes
+            ]
+            return ",".join(flat)
+        elif isinstance(passes, tuple):
+            fn, subtree = passes
+            return f"{fn}({PassParser.reassemble(subtree)})"
+        else:
+            raise ValueError("This shouldn't happen")
+
+    def find(self, prefix: str):
+        return self._find(prefix, self.pass_tree())
+
+    @staticmethod
+    def _find(prefix: str, subtree: list[str | tuple]):
+        locs = []
+        for i, node in enumerate(subtree):
+            if isinstance(node, tuple):
+                rest = PassParser._find(prefix, node[1])
+                if rest:
+                    for r in rest:
+                        locs.append([i] + r)
+            else:
+                if node.startswith(prefix):
+                    locs.append([i])
+        return locs
+
+    def split(self, locs: list[int]):
+        return self._split(locs, self.pass_tree())
+
+    @staticmethod
+    def _split(locs: list[int], subtree: list[tuple]):
+        head, *tail = locs
+        l, r = [], []
+        if any(tail):
+            k, v = subtree[head]
+            l, r = PassParser._split(tail, v)
+            left = subtree[0:head] + [(k, l)]
+            right = [(k, r)] + subtree[head + 1 :]
+        else:
+            left = subtree[0:head]
+            right = subtree[head:]
+        return left, right
+
+
+def main():
+    pp = PassParser()
+    full = pp.pass_tree()
+    reassembled = pp.reassemble(full)
+    assert pp.passes_str == reassembled
+
+    locs = pp.find("loop-rotate")
+    print(reassembled)
+    print(locs)
+    l, r = pp.split(locs[1])
+    print(pp.reassemble(full))
+    print("----------")
+    print(pp.reassemble(l))
+    print("----------")
+    print(pp.reassemble(r))
+
+
+if __name__ == "__main__":
+    main()

tadashi/translators.py

@@ -18,6 +18,8 @@
 from cython.cimports.tadashi.codegen import codegen
 from cython.cimports.tadashi.scop import Scop
 
+from .passesparser import PassParser
+
 ABC_ERROR_MSG = "Translator is an abstract base class, use a derived class."
 DOUBLE_SET_SOURCE = "Translator.set_source() should only be called once."
 
@@ -219,9 +221,16 @@ class Polly(Translator):
     compiler: str
     json_paths: list[Path]
     cwd: Path
+    before_polly_passes: str
+    after_polly_passes: str
 
     def __init__(self, compiler: str = "clang"):
         self.compiler = str(compiler)
+        pp = PassParser()
+        locs = pp.find("loop-rotate")
+        before, after = pp.split(locs[1])
+        self.before_polly_passes = pp.reassemble(before)
+        self.after_polly_passes = pp.reassemble(after)
 
     def _run(self, cmd: list[str], description: str, cwd: str = None):
         """cmd is command list, description is verb-ing, cwd defailts to self.cwd"""
@@ -259,43 +268,35 @@ def populate_ccscops(self, options: list[str]):
             raise MemoryError()
         timestamp = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
         self.cwd = Path(tempfile.mkdtemp(prefix=f"tadashi-{timestamp}-"))
-        self._get_preopt_bitcode(options)
+        self._get_pre_polly_bc(options)
         stderr = self._export_jscops(options)
         self.json_paths = self._fill_json_paths(stderr)
         for file in self.json_paths:
             with open(self.cwd / file) as fp:
                 jscop = json.load(fp)
             self._proc_jscop(jscop)
 
-    def _get_preopt_bitcode(self, options: list[str]) -> Path:
-        output = self.cwd / self.source.with_suffix(".O1.bc").name
-        if output.exists():
-            return output
+    def _get_pre_polly_bc(self, options: list[str]) -> Path:
+        compile_O0_bc = self.cwd / self.source.with_suffix(".pre_polly.bc").name
+        pre_polly_bc = self.cwd / self.source.with_suffix(".pre_polly.bc").name
+        if pre_polly_bc.exists():
+            return pre_polly_bc
 
         compiler_opts = {
-            "clang": [
-                "-O0",
-                "-Xclang",
-                "-disable-O0-optnone",
-            ],
-            "flang": ["-O1"],
+            "clang": ["-O0", "-Xclang", "-disable-O0-optnone"],
+            "flang": ["-O0"],
         }
-        cmd = [
-            self.compiler,
-            *options,
-            "-c",
-            "-emit-llvm",
-            str(self.source),
-            *compiler_opts[self.compiler[:5]],
-            "-o",
-            str(output),
-        ]
-        self._run(cmd, "parsing")
-        return output
+        compile_cmd = [self.compiler, *options, "-c", "-emit-llvm", str(self.source)]
+        compile_cmd += [*compiler_opts[self.compiler[:5]], "-o", str(compile_O0_bc)]
+        self._run(compile_cmd, "compiling with O0")
+        opt_cmd = ["opt", f"-passes={self.before_polly_passes}"]
+        opt_cmd += [str(compile_O0_bc), f"-o={str(pre_polly_bc)}"]
+        self._run(opt_cmd, "running pre polly opt passes")
+        return pre_polly_bc
 
     def _export_jscops(self, options: list[str]) -> str:
         cmd = self._polly() + [
-            str(self._get_preopt_bitcode(options)),
+            str(self._get_pre_polly_bc(options)),
             "-polly-export-jscop",
             "-o=/dev/null",
         ]
@@ -351,22 +352,28 @@ def _proc_jscop(self, jscop):
         self.ccscops.emplace_back(domain, sched, read, write)
 
     def legal(self) -> bool:
-        input_path = str(self._get_preopt_bitcode([]))
+        input_path = str(self._get_pre_polly_bc([]))
         cmd = self._polly() + [input_path, "-polly-import-jscop", "-o=/dev/null"]
         proc = self._run(cmd, "checking legality")
         return proc.returncode == 0
 
     def _import_jscops(self, options: list[str]) -> Path:
-        input_path = str(self._get_preopt_bitcode(options))
-        output = self.cwd / self.source.with_suffix(".bc").name
-        cmd = self._polly() + [
-            input_path,
-            "-polly-import-jscop",
+        input_path = str(self._get_pre_polly_bc(options))
+        post_polly_bc = self.cwd / self.source.with_suffix(".post_opt.bc").name
+
+        polly_cmd = self._polly() + [input_path, "-polly-import-jscop"]
+        polly_cmd += [
             "-polly-codegen",
-            f"-o={str(output)}",
+            f"-o={str(post_polly_bc)}",
             "-disable-polly-legality",
         ]
-        self._run(cmd, "imnporting jscops")
+        self._run(polly_cmd, "importing jscops")
+
+        output = self.cwd / self.source.with_suffix(".bc").name
+
+        opt_cmd = ["opt", f"-passes={self.after_polly_passes}"]
+        opt_cmd += [str(post_polly_bc), f"-o={str(output)}"]
+        opt_proc = self._run(opt_cmd, "running opt after jscops import")
         return output
 
     @cython.ccall
@@ -375,39 +382,33 @@ def generate_code(
     ) -> Path:
         for scop_idx, jscop_path in enumerate(self.json_paths):
             jscop_path = self.cwd / jscop_path
-            ccscop = self.ccscops[scop_idx]
-            sched = isl.isl_schedule_node_get_schedule(ccscop.current_node)
-            umap = isl.isl_schedule_get_map(sched)
-            isl.isl_schedule_free(sched)
-            sched_str = isl.isl_union_map_to_str(umap).decode()
             backup_path = jscop_path.with_suffix(jscop_path.suffix + ".bak")
             shutil.copy2(jscop_path, backup_path)
-            with jscop_path.open("r", encoding="utf-8") as f:
-                jscop = json.load(f)
-            for stmt in jscop["statements"]:
-                name = stmt["name"]
-                uset = isl.isl_union_set_read_from_str(
-                    self.ctx, stmt["domain"].encode()
-                )
-                stmt_map = isl.isl_union_map_intersect_domain_union_set(
-                    isl.isl_union_map_copy(umap), uset
-                )
-                stmt["schedule"] = isl.isl_union_map_to_str(stmt_map).decode()
-                isl.isl_union_map_free(stmt_map)
-            isl.isl_union_map_free(umap)
-            with jscop_path.open("w", encoding="utf-8") as f:
-                json.dump(jscop, f, indent=2)
-        output_path = Path(output_path).with_suffix(".s")
-        self._generate_assembly(output_path, options)
-        return output_path
+            self._update_jscop(jscop_path, scop_idx)
+        return self._compile_to_obj(output_path, options)
 
-    def _generate_assembly(self, output: Path, options: list[str]):
+    @cython.cfunc
+    def _update_jscop(self, jscop_path: Path, scop_idx: int):
+        ccscop = self.ccscops[scop_idx]
+        sched = isl.isl_schedule_node_get_schedule(ccscop.current_node)
+        isl.isl_schedule_free(sched)
+        umap = isl.isl_schedule_get_map(sched)
+        with jscop_path.open("r", encoding="utf-8") as f:
+            jscop = json.load(f)
+        for stmt in jscop["statements"]:
+            uset = isl.isl_union_set_read_from_str(self.ctx, stmt["domain"].encode())
+            tmp = isl.isl_union_map_copy(umap)
+            stmt_map = isl.isl_union_map_intersect_domain_union_set(tmp, uset)
+            stmt["schedule"] = isl.isl_union_map_to_str(stmt_map).decode()
+            isl.isl_union_map_free(stmt_map)
+        with jscop_path.open("w", encoding="utf-8") as f:
+            json.dump(jscop, f, indent=2)
+        isl.isl_union_map_free(umap)
+
+    def _compile_to_obj(self, output: Path, options: list[str]):
+        output_path = Path(output).with_suffix(".o")
         input_path = str(self._import_jscops(options))
-        cmd = [
-            "llc",
-            "-relocation-model=pic",
-            # *options,
-            input_path,
-            f"-o={str(output)}",
-        ]
+        cmd = ["llc", "-relocation-model=pic", "--filetype=obj"]
+        cmd += [input_path, f"-o={str(output_path)}"]
         self._run(cmd, "generating assembly")
+        return output_path