WIP: add npu support

xpu_timer/config_subparsers/hpu.py

@@ -0,0 +1,67 @@
+import textwrap
+
+from . import BaseBuildRender
+
+
+class BuildRender(BaseBuildRender):
+    @classmethod
+    def add_arguments(cls, parser):
+        BaseBuildRender.add_arguments(parser)
+
+    def __post_init__(self):
+        if self.args.sdk_path is None:
+            self.args.sdk_path = "/usr/local/Ascend/ascend-toolkit/latest/"
+
+    def rend_config_bzl(self):
+        hpu_config = textwrap.dedent(
+            """
+            XPU_TIMER_CONFIG = struct(
+                linkopt = [
+                    "-Wl,--version-script=$(location //xpu_timer/hpu:only_keep_hpu.lds)",
+                    "-L{hpu_path}/lib64",
+                ],
+                copt = [
+                    "-DXPU_HPU",
+                ],
+                deps = ["@hpu//:hpu_headers", "//xpu_timer/hpu:only_keep_hpu.lds"],
+                timer_deps = ["//xpu_timer/hpu:hpu_timer"],
+                hook_deps = ["//xpu_timer/hpu:hpu_hook"],
+            )"""
+        )
+
+        self.xpu_timer_config.append(hpu_config.format(hpu_path=self.sdk_path))
+        return "\n".join(self.xpu_timer_config)
+
+    def rend_bazelrc(self):
+        self.bazelrc_config.append(f"build --repo_env=HPU_HOME={self.sdk_path}")
+        return "\n".join(self.bazelrc_config)
+
+    def setup_files(self):
+        with open("WORKSPACE.template") as f:
+            workspace = f.read()
+
+        deps = textwrap.dedent(
+            """
+            load("//third_party/hpu:hpu_workspace.bzl", "hpu_workspace")
+            hpu_workspace()
+            """
+        )
+        return workspace + deps
+
+    def setup_platform_version(self):
+        # /usr/local/Ascend/ascend-toolkit/latest/version.cfg
+
+        version = None
+        path = f"{self.sdk_path}/version.cfg"
+        pattern = "toolkit_installed_version="
+        with open(path) as f:
+            for line in f:
+                if line.startswith(pattern):
+                    version = line.split(pattern)[-1]
+                    break
+        if version is None:
+            raise ValueError("Cannot found version")
+
+        major = version.split(":")[-1].replace("]", "").replace(".", "")
+        minor = ""
+        return f"hu{major}{minor}", "HPU"

xpu_timer/third_party/hpu/hpu.BUILD

@@ -0,0 +1,34 @@
+package(
+    default_visibility=["//visibility:public"]
+)
+
+
+cc_library(
+    name = "hpu_headers",
+    hdrs = glob(
+        include=["include/**/*.h", "include/**/*.hpp"],
+    ),
+    includes = ["include", "include/aclnn"],
+
+)
+
+cc_import(
+    name = "opapi",
+    shared_library = "lib64/libopapi.so",
+    includes = ["include"],
+    alwayslink = 1,
+)
+
+
+cc_import(
+    name = "ascendcl",
+    shared_library = "lib64/libascendcl.so",
+    includes = ["include"],
+    alwayslink = 1,
+)
+
+cc_import(
+    name = "hccl",
+    shared_library = "lib64/libhccl.so",
+    alwayslink = 1,
+)

xpu_timer/third_party/hpu/hpu_workspace.bzl

@@ -0,0 +1,10 @@
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
+load("//:workspace.bzl", "ali_code_repository", "dynamic_local_repository")
+
+def hpu_workspace():
+    dynamic_local_repository(
+        name = "hpu",
+        include_default_path = "/usr/local/Ascend/ascend-toolkit/latest",
+        build_file = "//third_party/hpu:hpu.BUILD",
+        include = "ASCEND_TOOLKIT_HOME",
+    )

xpu_timer/xpu_timer/hpu/BUILD.bazel

@@ -0,0 +1,42 @@
+load("@rules_python//python:defs.bzl", "py_binary")
+load("//:workspace.bzl", "xpu_cc_binary", "xpu_cc_library")
+
+package(default_visibility = ["//visibility:public"])
+
+exports_files(["only_keep_hpu.lds"])
+
+xpu_cc_library(
+    name = "hpu_timer",
+    srcs = [
+        "hpu_dtype_util.cc",
+        "hpu_timer.cc",
+    ],
+    hdrs = [
+        "hpu_dtype_util.h",
+        "hpu_timer.h",
+    ],
+    deps = [
+        "//xpu_timer/common:xpu_timer",
+        "//xpu_timer/protos:cc_hook_proto",
+    ],
+)
+
+xpu_cc_library(
+    name = "hpu_hook",
+    srcs = [
+        "hook.cc",
+    ],
+    hdrs = [
+        "hook.h",
+    ],
+    deps = [
+        "//xpu_timer/common:macro",
+        "//xpu_timer/common:manager",
+        "//xpu_timer/common:util",
+        "@hpu//:ascendcl",
+        "@hpu//:hccl",
+        "@hpu//:opapi",
+    ],
+    # force to keep all syms, if not, linker will remove useless symbols.
+    alwayslink = True,
+)