feat(wafer): 接入晶圆级后端编译器，提供独立构建、打包与测试入口

.github/workflows/wafer-isolation.yml

@@ -0,0 +1,46 @@
+name: Wafer isolated build and offline checks
+
+# Separate, opt-in job: the Ascend workflow and its runner stay independent.
+# The configured runner needs the pinned LLVM, Wafer SDK, Python build tools
+# and pytest. This workflow submits no work to a Wafer device.
+on:
+  workflow_dispatch:
+
+jobs:
+  WaferOffline:
+    if: vars.WAFER_CI_RUNNER != ''
+    runs-on: ${{ vars.WAFER_CI_RUNNER }}
+    env:
+      LLVM_SYSPATH: ${{ vars.WAFER_CI_LLVM_SYSPATH }}
+      WAFER_DEPS_ROOT: ${{ vars.WAFER_CI_DEPS_ROOT }}
+      WAFER_BUILD_DIR: ${{ github.workspace }}/.wafer-build
+      CMAKE_BUILD_PARALLEL_LEVEL: '4'
+      DICP_BACKEND: wafer
+      USE_SIM_MODE: '1'
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          submodules: recursive
+      - name: Build both compiler trees
+        run: bash scripts/wafer/compile_wafer.sh
+      - name: Assemble and install an isolated wheel
+        run: |
+          python setup_on_wafer.py --build-dir "$WAFER_BUILD_DIR" --wheel-dir "$WAFER_BUILD_DIR/wheel"
+          python -m venv --system-site-packages "$WAFER_BUILD_DIR/venv"
+          "$WAFER_BUILD_DIR/venv/bin/python" -m pip install --no-deps --force-reinstall "$WAFER_BUILD_DIR"/wheel/triton-*.whl
+      - name: Check frontend, patch profiles, tools and loader contracts
+        run: |
+          cd "$RUNNER_TEMP"
+          "$WAFER_BUILD_DIR/venv/bin/python" -m pytest -q \
+            "$GITHUB_WORKSPACE/test/wafer/test_frontend_isolation.py" \
+            "$GITHUB_WORKSPACE/test/wafer/test_tle_frontend.py" \
+            "$GITHUB_WORKSPACE/test/wafer/test_external_interfaces.py" \
+            "$GITHUB_WORKSPACE/test/wafer/test_scalar_copy.py" \
+            "$GITHUB_WORKSPACE/test/wafer/test_loader_isolation.py" \
+            "$GITHUB_WORKSPACE/test/wafer/test_commonir_abi.py" \
+            "$GITHUB_WORKSPACE/test/wafer/test_patch_profiles.py"
+      - uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          name: wafer-build-identity
+          path: .wafer-build/wafer-build.json

backend/compiler.py

@@ -126,12 +126,17 @@ def __init__(self, target: str) -> None:
         elif self.driver.target == "maca":
             self.capability = 80
             self.binary_ext = "mcfatbin"
+        elif self.driver.target == "wafer":
+            from triton.backends.dicp_triton.wafer import WaferBackend
+
+            self._wafer_backend = WaferBackend(target)
+            self.binary_ext = self._wafer_backend.binary_ext
         else:
             raise RuntimeError(f"Target '{self.driver.target}' is not supported.")
 
     @staticmethod
     def supports_target(target: GPUTarget):
-        return target.backend in ["ascend", "mlu", "maca", "cpu"]
+        return target.backend in ["ascend", "mlu", "maca", "cpu", "wafer"]
 
     @staticmethod
     def make_ttir(mod, metadata, opt):
@@ -152,7 +157,9 @@ def make_ttir(mod, metadata, opt):
     def add_stages(self, stages, options, language=None):
         if self.driver.is_cpu_verify:
             return self._cpu_backend.add_stages(stages, options, language)
-        if self.driver.target == "ascend":
+        if self.driver.target == "wafer":
+            return self._wafer_backend.add_stages(stages, options, language)
+        elif self.driver.target == "ascend":
             from triton.backends.dicp_triton.npu import (
                 make_ttir,
                 ttir_to_linalg_dicp,
@@ -245,7 +252,9 @@ def add_stages(self, stages, options, language=None):
     def load_dialects(self, ctx):
         if self.driver.is_cpu_verify:
             return self._cpu_backend.load_dialects(ctx)
-        if self.driver.target == "mlu":
+        if self.driver.target == "wafer":
+            return self._wafer_backend.load_dialects(ctx)
+        elif self.driver.target == "mlu":
             from triton._C.libtriton import mlu
 
             mlu.load_dialects(ctx)
@@ -263,7 +272,9 @@ def get_driver(self):
     def parse_options(self, options: dict) -> Any:
         if self.driver.is_cpu_verify:
             return self._cpu_backend.parse_options(options)
-        if self.target.backend == "ascend":
+        if self.target.backend == "wafer":
+            return self._wafer_backend.parse_options(options)
+        elif self.target.backend == "ascend":
             from triton.backends.dicp_triton.npu import NPUOptions
 
             args = {
@@ -340,6 +351,8 @@ def get_codegen_implementation(self, options=None):
         codegen_fns = dict()
         if self.driver.is_cpu_verify:
             return self._cpu_backend.get_codegen_implementation(options)
+        elif self.target.backend == "wafer":
+            return self._wafer_backend.get_codegen_implementation(options)
         elif self.target.backend == "ascend":
             from triton.backends.dicp_triton.npu import min_dot_size
 
@@ -366,7 +379,9 @@ def get_codegen_implementation(self, options=None):
     def pack_metadata(self, metadata):
         if self.driver.is_cpu_verify:
             return self._cpu_backend.pack_metadata(metadata)
-        if self.target.backend == "ascend":
+        if self.target.backend == "wafer":
+            return self._wafer_backend.pack_metadata(metadata)
+        elif self.target.backend == "ascend":
 
             KERNEL_NAME_MAX_LEN = 49
             kernel_name_orig = metadata.kernel_name
@@ -395,7 +410,9 @@ def pack_metadata(self, metadata):
     def hash(self):
         if self.driver.is_cpu_verify:
             return self._cpu_backend.hash()
-        if self.target.backend == "mlu":
+        if self.target.backend == "wafer":
+            return self._wafer_backend.hash()
+        elif self.target.backend == "mlu":
             from triton.backends.dicp_triton.mlu import get_cnas_version
 
             version = get_cnas_version()
@@ -405,7 +422,9 @@ def hash(self):
         return str(version_key)
 
     def get_module_map(self) -> Dict[str, ModuleType]:
-        if self.target.backend == "mlu":
+        if self.target.backend == "wafer":
+            return self._wafer_backend.get_module_map()
+        elif self.target.backend == "mlu":
             from triton.language.extra.mlu import libdevice
 
             return {"triton.language.extra.libdevice": libdevice}

backend/driver.py

@@ -170,6 +170,26 @@ def __init__(self, target=None):
             from .ascend_autotune_hooks import hook_autotune_for_ascend
 
             hook_autotune_for_ascend()
+        elif backend == "wafer":
+            from .wafer_runtime import (
+                SimulatorUtils,
+                WaferLauncher,
+                WaferUtils,
+                get_runtime,
+            )
+
+            self.target = "wafer"
+            if os.getenv("USE_SIM_MODE", "0").lower() in ("1", "true", "yes"):
+                self.utils = SimulatorUtils()
+                self.get_current_device = lambda: 0
+                self.set_current_device = lambda device: None
+            else:
+                self.utils = WaferUtils()
+                self.get_current_device = lambda: get_runtime().current_device()
+                self.set_current_device = lambda device: get_runtime().set_device(
+                    device
+                )
+            self.launcher_cls = WaferLauncher
         elif backend == "nvidia":
             from triton.backends.nvidia.driver import CudaLauncher, CudaUtils
 
@@ -197,6 +217,8 @@ def is_active():
 
     @classmethod
     def is_active(self):
+        if get_current_backend() == "wafer":
+            return True
         try:
             current_backend = get_current_backend()
             if current_backend == "ascend":
@@ -257,12 +279,21 @@ def get_device_capability(self):
             return ("maca", 0)
         elif self.target == "ascend":
             return ("ascend", 0)
+        elif self.target == "wafer":
+            return ("wafer", 0)
         elif self.target == "nvidia":
             capability = torch.cuda.get_device_capability(self.get_current_device())
             return ("cuda", capability)
         return ("dicp", 0)
 
     def get_current_stream(self, device):
+        if self.target == "wafer":
+            if os.getenv("USE_SIM_MODE", "0").lower() in ("1", "true", "yes"):
+                return None
+            from .wafer_runtime import get_runtime
+
+            stream = get_runtime().current_stream(device)
+            return None if stream is None else stream.txda_stream
         import torch
 
         if self.target == "mlu":
@@ -286,6 +317,8 @@ def get_current_stream(self, device):
         return None
 
     def get_current_device(self):
+        if self.target == "wafer":
+            return 0
         import torch
 
         # dicp doesn't have a device to return. Return something.
@@ -338,6 +371,8 @@ def get_current_target(self):
             arch = self.utils.get_arch()
             warp_size = 0
             return GPUTarget(backend, arch, warp_size)
+        elif self.target == "wafer":
+            return GPUTarget("wafer", "wafer", 32)
         elif self.target == "nvidia":
             device = self.get_current_device()
             capability = torch.cuda.get_device_capability(device)
@@ -350,7 +385,14 @@ def assemble_tensormap_to_arg(self, tensormaps_info, args):
         return args
 
     def get_device_interface(self):
-        if self.target == "ascend":
+        if self.target == "wafer":
+            from .wafer_runtime import get_runtime
+
+            runtime = get_runtime()
+            if not hasattr(runtime, "Event") or not hasattr(runtime, "synchronize"):
+                raise RuntimeError("Wafer benchmarking requires torch_txda Event and synchronize support")
+            return runtime
+        elif self.target == "ascend":
             import torch
 
             return torch.npu
@@ -362,7 +404,11 @@ def get_device_interface(self):
             assert False, f"Not implemented for {self.target}"
 
     def get_empty_cache_for_benchmark(self):
-        if self.target == "ascend":
+        if self.target == "wafer":
+            # no device cache flush.
+            # None distinguishes this policy from other backends' zeroable tensor.
+            return None
+        elif self.target == "ascend":
             import torch
 
             cache_size = 192 * 1024 * 1024
@@ -383,6 +429,12 @@ def get_active_torch_device(self):
 
         if self.is_cpu_verify:
             return self._cpu_driver.get_active_torch_device()
+        if self.target == "wafer" and os.getenv("USE_SIM_MODE", "0").lower() not in (
+            "1",
+            "true",
+            "yes",
+        ):
+            return torch.device("txda", self.get_current_device())
         return torch.device("cpu")
 
     def map_python_to_cpp_type(self, ty: str) -> str:
@@ -407,4 +459,5 @@ def map_python_to_cpp_type(self, ty: str) -> str:
 
     @classmethod
     def clear_cache(self, cache):
-        cache.zero_()
+        if cache is not None:
+            cache.zero_()

backend/utils.py

@@ -157,6 +157,11 @@ def get_current_backend():
     global backend
     if backend is not None:
         return backend
+    override = os.getenv("DICP_BACKEND", "").lower()
+    if override:
+        if override not in {"ascend", "mlu", "maca", "nvidia", "wafer"}:
+            raise RuntimeError(f"Unsupported DICP_BACKEND '{override}'.")
+        backend = override
     elif command_exists("npu-smi"):
         backend = "ascend"
     elif command_exists("cnmon"):
@@ -165,6 +170,8 @@ def get_current_backend():
         backend = "maca"
     elif command_exists("nvidia-smi"):
         backend = "nvidia"
+    elif command_exists("tsm_smi"):
+        backend = "wafer"
     else:
         backend = None
     return backend