Added opencl pass options to approxhpvm.py

hpvm/tools/py-approxhpvm/CMakeLists.txt

@@ -22,16 +22,20 @@ set(DIRECT_LINK_LIBS ${OpenCL_LIBRARY} "$<TARGET_FILE:tensor_runtime>")
 # and does not export its file location.
 # Keep this in sync with hpvm/projects/hpvm-rt/CMakeLists.txt.
 set(HPVM_RT_PATH ${LLVM_BUILD_DIR}/tools/hpvm/projects/hpvm-rt/hpvm-rt.bc)
+# And this must be manually in sync with `hpvm/lib/Transforms/*`.
+# which is fine for because we don't have many passes for now.
 set(
     AVAILABLE_PASSES
     LLVMBuildDFG
-    LLVMInPlaceDFGAnalysis
+    LLVMClearDFG
     LLVMDFG2LLVM_CPU
     LLVMDFG2LLVM_CUDNN
+    LLVMDFG2LLVM_OpenCL
     LLVMDFG2LLVM_WrapperAPI
     LLVMFuseHPVMTensorNodes
-    LLVMClearDFG
     LLVMGenHPVM
+    LLVMInPlaceDFGAnalysis
+    LLVMLocalMem
 )
 
 # First resolve all `@symbol@` by configuring the file

hpvm/tools/py-approxhpvm/main.py.in

@@ -24,18 +24,30 @@ COMPILE_FLAGS = ["fno-exceptions", "std=c++11", "O3"]
 def compile_hpvm_c(
     input_file: PathLike,
     output_file: PathLike,
-    codegen_target: str = "tensor",
+    tensor_target: Optional[str],
+    opencl: bool,
     include: List[PathLike] = None,
     working_dir: PathLike = None,
     conf_file: PathLike = None,
 ):
     from subprocess import check_output
 
-    codegen_functions = {
-        "tensor": lambda i, o: opt_codegen_tensor(i, o, conf_file),
-        "cudnn": opt_codegen_cudnn
-    }
-    codegen_f = codegen_functions[codegen_target]
+    passes = ["LLVMBuildDFG", "LLVMInPlaceDFGAnalysis"]
+    flags = ["buildDFG", "inplace"]
+    if tensor_target == "tensor":
+        if conf_file is None:
+            raise ValueError("conf_file must be defined when tensor_target=='tensor'.")
+        passes += ["LLVMFuseHPVMTensorNodes", "LLVMDFG2LLVM_WrapperAPI"]
+        flags += ["hpvm-fuse", "dfg2llvm-wrapperapi", f"configuration-inputs-filename={conf_file}"]
+    elif tensor_target == "cudnn":
+        passes += ["LLVMDFG2LLVM_CUDNN"]
+        flags += ["dfg2llvm-cudnn"]
+    if opencl:
+        passes += ["LLVMDFG2LLVM_OpenCL"]
+        flags += ["dfg2llvm-opencl"]
+    passes += ["LLVMDFG2LLVM_CPU", "LLVMClearDFG", "LLVMGenHPVM"]
+    flags += ["dfg2llvm-cpu", "clearDFG"]
+
     working_dir = Path(working_dir or ".")
     if not working_dir.is_dir():
         os.makedirs(working_dir)
@@ -48,7 +60,7 @@ def compile_hpvm_c(
     commands = [
         hpvm_c_to_ll(input_file, ll_file, extra_includes=include),
         opt_codegen_hpvm(ll_file, hpvm_ll_file),
-        codegen_f(hpvm_ll_file, llvm_ll_file),
+        _run_opt(hpvm_ll_file, llvm_ll_file, passes, flags),
         link_hpvm_rt(llvm_ll_file, hpvm_rt_linked_file),
         link_binary(hpvm_rt_linked_file, output_file),
     ]
@@ -76,36 +88,6 @@ def opt_codegen_hpvm(src_file: PathLike, target_file: PathLike) -> List[str]:
     return _run_opt(src_file, target_file, ["LLVMGenHPVM"], ["genhpvm", "globaldce"])
 
 
-def opt_codegen_cudnn(src_file: PathLike, target_file: PathLike) -> List[str]:
-    passes = [
-        "LLVMBuildDFG", "LLVMInPlaceDFGAnalysis",
-        "LLVMDFG2LLVM_CUDNN", "LLVMDFG2LLVM_CPU",
-        "LLVMFuseHPVMTensorNodes", "LLVMClearDFG", "LLVMGenHPVM"
-    ]
-    flags = [
-        "buildDFG", "inplace", "hpvm-fuse",
-        "dfg2llvm-cudnn", "dfg2llvm-cpu", "clearDFG"
-    ]
-    return _run_opt(src_file, target_file, passes, flags)
-
-
-def opt_codegen_tensor(
-    src_file: PathLike, target_file: PathLike, conf_file: PathLike
-):
-    passes = [
-        "LLVMBuildDFG", "LLVMInPlaceDFGAnalysis",
-        "LLVMDFG2LLVM_WrapperAPI", "LLVMDFG2LLVM_CPU",
-        "LLVMFuseHPVMTensorNodes", "LLVMClearDFG", "LLVMGenHPVM"
-    ]
-    flags = [
-        "buildDFG", "inplace", "hpvm-fuse",
-        "dfg2llvm-wrapperapi",
-        f"configuration-inputs-filename={conf_file}",
-        "dfg2llvm-cpu", "clearDFG",
-    ]
-    return _run_opt(src_file, target_file, passes, flags)
-
-
 def link_hpvm_rt(src_file: PathLike, target_file: PathLike) -> List[str]:
     return [str(LLVM_BUILD_BIN / "llvm-link"), str(src_file), HPVM_RT_PATH, "-o", str(target_file)]
 
@@ -143,7 +125,7 @@ def _run_opt(
 ) -> List[str]:
     unavailable = set(pass_names) - set(AVAILABLE_PASSES)
     if unavailable:
-        raise ValueError(f"Passes {unavailable} are unavailable from CMake")
+        raise ValueError(f"Passes {unavailable} are unavailable for this compilation.")
     load_passes_strs = [s for pass_ in pass_names for s in ["-load", f"{pass_}.so"]]
     pass_flags_strs = [f"-{flag}" for flag in pass_flags]
     return [
@@ -158,18 +140,22 @@ def parse_args():
     parser.add_argument("output_file", type=Path, help="Path to generate binary to")
     parser.add_argument(
         "-t",
-        "--codegen-target",
+        "--tensor-target",
         type=str,
-        required=True,
         choices=["tensor", "cudnn"],
-        help="Backend to use",
+        help="Backend to use for tensor operators",
     )
     parser.add_argument(
-        "-d", "--working-dir", type=Path, help="Directory to generate temp files in"
+        "--conf-file", type=Path,
+        help="File to approximation configurations; required for tensor target 'tensor'"
     )
     parser.add_argument(
-        "--conf-file", type=Path,
-        help="File to approximation configurations; required for 'tensor' target"
+        "--opencl",
+        action="store_true",
+        help="Use OpenCL support. Requires HPVM built with OpenCL",
+    )
+    parser.add_argument(
+        "-d", "--working-dir", type=Path, help="Directory to generate temp files in"
     )
     parser.add_argument(
         "-I", "--include", type=Path, action="append",
@@ -177,9 +163,9 @@ def parse_args():
     )
 
     args = parser.parse_args()
-    if args.codegen_target == "tensor":
+    if args.tensor_target == "tensor":
         if args.conf_file is None:
-            parser.error('Codegen target "tensor" requires --conf-file argument')
+            parser.error('Tensor target "tensor" requires --conf-file argument')
     return args