[Ascend] feat: add Ascend backend support for Triton 3.6.x

.pre-commit-config.yaml

@@ -11,6 +11,7 @@ repos:
       - id: check-toml
       - id: check-ast
       - id: check-added-large-files
+        exclude: ^third_party/ascend/backend/lib/libdevice\.10\.bc$
       - id: check-merge-conflict
       - id: check-executables-have-shebangs
       - id: check-shebang-scripts-are-executable

CMakeLists.txt

@@ -172,6 +172,12 @@ if(FLAGTREE_BACKEND AND EXISTS ${BACKEND_INCLUDE_DIR})
   # flagtree: include directories
   include_directories(${PROJECT_SOURCE_DIR}/third_party/${FLAGTREE_BACKEND}/include)
   include_directories(${PROJECT_BINARY_DIR}/third_party/${FLAGTREE_BACKEND}/include) # Tablegen'd files
+  if(FLAGTREE_BACKEND STREQUAL "ascend")
+    # The ascend backend builds the core include/lib from its own CMakeLists, so
+    # only the search paths are needed here.
+    include_directories(${PROJECT_SOURCE_DIR}/include)
+    include_directories(${PROJECT_BINARY_DIR}/include) # Tablegen'd files
+  endif()
 else()
   include_directories(${PROJECT_SOURCE_DIR}/include)
   include_directories(${PROJECT_BINARY_DIR}/include) # Tablegen'd files
@@ -212,6 +218,8 @@ if(TRITON_BUILD_PYTHON_MODULE)
   # FlagPrism: register component targets before backend codegen captures TRITON_LIBS.
   flagtree_add_flagprism_components()
   flagtree_configure_flir_dependency()
+  # Provide the BiShengIR targets before the backend subdirectory links them.
+  flagtree_configure_ascend_npu_ir_dependency()
 
   foreach(CODEGEN_BACKEND ${TRITON_CODEGEN_BACKENDS})
     add_subdirectory(third_party/${CODEGEN_BACKEND})
@@ -257,6 +265,7 @@ if(TRITON_BUILD_PYTHON_MODULE)
     ${triton_plugins}
 
     # mlir
+    MLIRFuncInlinerExtension
     MLIRAMDGPUDialect
     MLIRNVVMDialect
     MLIRNVVMToLLVMIRTranslation

cmake/FlagTreeOptions.cmake

@@ -37,6 +37,11 @@ macro(flagtree_configure_options)
     add_compile_definitions(__FLAGTREE_COMMON_IR__)
   endif()
 
+  # Ascend always requires CommonIR, without a C/C++ feature gate.
+  if(FLAGTREE_BACKEND STREQUAL "ascend")
+    set(FLAGTREE_COMMON_IR_ENABLED ON)
+  endif()
+
   set(FLAGCX_ENABLED OFF)
   set(FLAGCX_SUPPORT_BACKENDS nvidia iluvatar)
   if(NOT FLAGTREE_BACKEND OR
@@ -82,6 +87,29 @@ macro(flagtree_configure_options)
     set(CMAKE_CXX_COMPILER clang++-16)
     add_definitions(-D__NVIDIA__)
     add_definitions(-D__AMD__)
+  elseif(FLAGTREE_BACKEND STREQUAL "ascend")
+    set(CMAKE_C_COMPILER clang)
+    set(CMAKE_CXX_COMPILER clang++)
+    # AscendNPU-IR gates its sources on the LLVM major version. cmake/llvm-hash.txt
+    # pins LLVM 22 here.
+    #
+    # Deliberately leave LLVM_MAJOR_VERSION_21_COMPATIBLE unset: LLVM 22 renamed
+    # the Mesh dialect to Shard (only libMLIRShardDialect.a ships in the pinned
+    # package), so turning it on would reference a missing target.
+    set(LLVM_MAJOR_VERSION_22_OR_HIGHER ON)
+    set(LLVM_MAJOR_VERSION_22_COMPATIBLE ON CACHE BOOL
+        "AscendNPU-IR build with llvm major version 22 or later" FORCE)
+    add_definitions(-D__LLVM_MAJOR_VERSION_22_COMPATIBLE__)
+    add_compile_options("-Wno-deprecated-declarations")
+    add_compile_options("-Wno-error=deprecated-declarations")
+    # AscendNPU-IR: build only the standalone IR, and keep it out of the
+    # backend's own CMakeLists so the subdirectory is added exactly once.
+    set(BISHENGIR_ENABLE_A5_UNPUBLISHED_FEATURES ON)
+    set(BISHENGIR_BUILD_STANDALONE_IR_ONLY ON)
+    set(ASCENDNPU_IR_SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR}/third_party/ascend/AscendNPU-IR)
+    set(ASCENDNPU_IR_BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/third_party/ascend/AscendNPU-IR)
+    include_directories(${ASCENDNPU_IR_SRC_DIR}/bishengir/include)
+    include_directories(${ASCENDNPU_IR_BINARY_DIR}/bishengir/include)
   elseif(FLAGTREE_BACKEND STREQUAL "tsingmicro")
     set(CMAKE_C_COMPILER clang)
     set(CMAKE_CXX_COMPILER clang++)
@@ -230,7 +258,7 @@ macro(flagtree_configure_core_source)
   endif()
 
   if(FLAGTREE_BACKEND MATCHES
-     "^(xpu|cambricon|aipu|tsingmicro|enflame|rpu|thrive|metax|tileir|ppu|spacemit|sunrise)$")
+     "^(xpu|cambricon|aipu|tsingmicro|enflame|rpu|thrive|metax|tileir|ppu|spacemit|sunrise|ascend)$")
     include_directories(${PROJECT_SOURCE_DIR}/include)
     include_directories(${PROJECT_BINARY_DIR}/include) # Tablegen'd files
     if(FLAGTREE_BACKEND STREQUAL "xpu")
@@ -320,7 +348,7 @@ endmacro()
 macro(flagtree_configure_flir_dependency)
   if(FLAGTREE_COMMON_IR_ENABLED)
     if(NOT EXISTS "${PROJECT_SOURCE_DIR}/third_party/flir/CMakeLists.txt")
-      message(FATAL_ERROR "FLAGTREE_COMMON_IR requires third_party/flir")
+      message(FATAL_ERROR "CommonIR support requires third_party/flir")
     endif()
     include_directories(${PROJECT_SOURCE_DIR}/third_party/flir/include)
     include_directories(${PROJECT_BINARY_DIR}/third_party/flir/include)
@@ -346,6 +374,21 @@ macro(flagtree_configure_flir_dependency)
 endmacro()
 
 
+# AscendNPU-IR provides the BiShengIR dialects the ascend backend links against.
+# Added from the top level (not from the backend's own CMakeLists) so the source
+# directory is only ever associated with one binary directory.
+macro(flagtree_configure_ascend_npu_ir_dependency)
+  if(FLAGTREE_BACKEND STREQUAL "ascend")
+    if(NOT EXISTS "${ASCENDNPU_IR_SRC_DIR}/CMakeLists.txt")
+      message(FATAL_ERROR
+        "The ascend backend requires AscendNPU-IR at ${ASCENDNPU_IR_SRC_DIR}. "
+        "Run the Python package build to download third-party dependencies.")
+    endif()
+    add_subdirectory(${ASCENDNPU_IR_SRC_DIR} ${ASCENDNPU_IR_BINARY_DIR})
+  endif()
+endmacro()
+
+
 macro(flagtree_configure_tle_plugin append_tle_plugin)
   if(FLAGTREE_TLE)
     if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/tle/CMakeLists.txt")
@@ -575,6 +618,10 @@ macro(flagtree_configure_tools_and_tests)
       flagtree_add_tle_generated_header_dependencies()
     endif()
     add_subdirectory(test)
+  elseif(FLAGTREE_BACKEND STREQUAL "ascend")
+    if(FLAGTREE_TLE)
+      flagtree_add_tle_generated_header_dependencies()
+    endif()
   elseif(FLAGTREE_BACKEND STREQUAL "iluvatar")
     option(FLAGTREE_ILUVATAR_BUILD_BIN
       "Build third_party/iluvatar/bin tools and lit tests" OFF)
@@ -683,6 +730,13 @@ macro(flagtree_added_python_src)
   endforeach()
   add_library(triton SHARED ${_flagtree_python_sources})
 
+  if(FLAGTREE_BACKEND STREQUAL "ascend")
+    # third_party/ascend/triton_ascend.cc registers the buffer_ir submodule, so
+    # its implementation must be part of the triton module.
+    target_sources(triton PRIVATE
+      ${CMAKE_CURRENT_SOURCE_DIR}/third_party/ascend/python/triton/extension/buffer/src/buffer_ir.cc)
+  endif()
+
   if(FLAGTREE_BACKEND STREQUAL "xpu")
     target_sources(triton PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/third_party/xpu/python/src/mlir_pass_abi_shim.cc)
     add_dependencies(triton TritonAMDGPUTableGen TritonAMDGPUAttrDefsIncGen)

python/requirements-ascend.txt

@@ -0,0 +1,3 @@
+--extra-index-url https://ascend.devcloud.huaweicloud.com/cann/pypi/simple/
+
+cann-shmem

python/setup_tools/setup_helper.py

@@ -68,7 +68,8 @@ def get_long_description():
 
 def init_backends(backend_installer):
     if flagtree_backend:
-        if flagtree_backend in ("aipu", "tsingmicro", "enflame", "rpu", "thrive", "sunrise", "tileir", "ppu", "amd"):
+        if flagtree_backend in ("aipu", "tsingmicro", "enflame", "rpu", "thrive", "sunrise", "tileir", "ppu", "amd",
+                                "ascend"):
             backends = [
                 *backend_installer.copy(configs.default_backends + tuple(configs.extend_backends)),
                 *backend_installer.copy_externals(),
@@ -233,6 +234,12 @@ def get_hook_instance(hook_name):
     return hook_instance if callable(hook_instance) else None
 
 
+def finalize_backend_build_tree(build_lib):
+    hook = get_hook_instance("finalize_build_tree")
+    if hook:
+        hook(build_lib)
+
+
 def enable_flagtree_third_party(name):
     return os.environ.get(f"USE_{name.upper()}", 'ON') == 'ON'
 
@@ -888,7 +895,14 @@ def uninstall_triton():
 
 cache = FlagTreeCache()
 
-download_flagtree_third_party("flir", condition=(flagtree_backend == "tsingmicro"), required=True)
+# FLIR supplies the CommonIR dialect, which the ascend backend requires for its
+# CommonIRToHIVM lowering and FLIR's triton-lane-vectorize pass.
+download_flagtree_third_party("flir", condition=(flagtree_backend in ("tsingmicro", "ascend")), required=True)
+
+# The ascend backend links the BiShengIR dialects built from AscendNPU-IR.
+if flagtree_backend == "ascend":
+    configs.activated_module.warn_stale_ascend_npu_ir_checkout(utils.flagtree_submodules["AscendNPU-IR"])
+    download_flagtree_third_party("AscendNPU-IR", condition=True, required=True)
 '''
    FlagCX is a third-party library adopted by the tle distributed system,
    refer to https://github.com/flagos-ai/FlagCX

python/setup_tools/utils/__init__.py

@@ -60,6 +60,10 @@ def _register_submodule(self, submodule):
     },
     {"name": "flir", "url": "https://github.com/flagos-ai/flir.git"},
     {"name": "FlagPrism", "url": "https://github.com/flagos-ai/FlagPrism.git"},
+    {
+        "name": "AscendNPU-IR", "url": "https://gitcode.com/Ascend/AscendNPU-IR.git", "relative_path":
+        "ascend/AscendNPU-IR", "commit_id": "efb7141581c009bbd1d676b18941ed15ede62a68"
+    },
     {
         "name": "flagcx", "url": "https://github.com/flagos-ai/FlagCX.git", "relative_path": "tle/third_party/flagcx",
         "commit_id": "c93438decaba8d935e7aa6ae8377e56fd3910b32"

python/setup_tools/utils/ascend.py

@@ -0,0 +1,48 @@
+import importlib.util
+import os
+import subprocess
+from pathlib import Path
+
+
+def install_extension(*args, **kargs):
+    """Copy producer tools next to libtriton, then strip only the copies."""
+    from python.build_helpers import get_cmake_dir
+
+    build_ext = kargs.get("build_ext") or (args[0] if args else None)
+    if build_ext is None:
+        raise RuntimeError("Ascend install_extension requires the build_ext command")
+
+    # Match setup.py's CMakeBuild.build_extension: ext.path, not ext.name.
+    ext = next(ext for ext in build_ext.extensions if ext.name == "triton")
+    extdir = Path(build_ext.get_ext_fullpath(ext.path)).resolve().parent
+    source = Path(__file__).resolve().parents[3] / "third_party/ascend/build/packaging.py"
+    _load_packaging(source)._copy_ascend_tools(extdir, get_cmake_dir())
+
+
+def _load_packaging(source):
+    spec = importlib.util.spec_from_file_location("_flagtree_ascend_packaging", source)
+    if spec is None or spec.loader is None:
+        raise RuntimeError(f"Cannot load Ascend packaging hook: {source}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def finalize_build_tree(build_lib):
+    """Delegate packaging to the backend without importing its runtime."""
+    source = Path(__file__).resolve().parents[3] / "third_party/ascend/build/packaging.py"
+    _load_packaging(source).finalize_build_tree(build_lib)
+
+
+def warn_stale_ascend_npu_ir_checkout(module):
+    # DownloadManager reuses dst_path as-is when it already exists, so a
+    # checkout left on a different commit would silently shadow the pin
+    # registered for this build.
+    if not module.commit_id or not os.path.isdir(os.path.join(module.dst_path, ".git")):
+        return
+    head = subprocess.run(["git", "rev-parse", "HEAD"], cwd=module.dst_path, capture_output=True, text=True)
+    if head.returncode != 0 or head.stdout.strip().startswith(module.commit_id):
+        return
+    print(f"[WARNING] {module.dst_path} is at {head.stdout.strip()[:8]} but this build pins "
+          f"AscendNPU-IR to {module.commit_id[:8]}; remove the directory (or run "
+          f"'git checkout {module.commit_id}' inside it) before rebuilding")

python/src/ir.cc

@@ -382,13 +382,19 @@ void init_triton_ir(py::module &&m) {
       .value("NONE", PropagateNan::NONE)
       .value("ALL", PropagateNan::ALL);
 
-  py::enum_<InputPrecision>(m, "INPUT_PRECISION", py::module_local())
+  auto inputPrecision =
+      py::enum_<InputPrecision>(m, "INPUT_PRECISION", py::module_local());
+  inputPrecision
       .value("TF32", InputPrecision::TF32)
       .value("TF32x3", InputPrecision::TF32x3)
       .value("IEEE", InputPrecision::IEEE)
       .value("BF16x3", InputPrecision::BF16x3)
-      .value("BF16x6", InputPrecision::BF16x6)
-      .export_values();
+      .value("BF16x6", InputPrecision::BF16x6);
+
+  // Backend-specialized enum definitions may additionally support HF32.
+  if (auto hf32 = symbolizeInputPrecision("hf32"))
+    inputPrecision.value("HF32", *hf32);
+  inputPrecision.export_values();
 
   py::enum_<ScaleDotElemType>(m, "ScaleDotElemTypeTY", py::module_local())
       .value("E4M3", ScaleDotElemType::E4M3)

python/triton/_utils.py

@@ -26,6 +26,11 @@
 from typing import Any, Callable, TYPE_CHECKING, Union, List, Dict
 
 from triton._flagtree_spec import spec_func
+from triton._flagtree_backend import FLAGTREE_BACKEND
+
+# The ascend backend lowers tensors through a SIMD pipeline that does not
+# require power-of-2 block shapes (see validate_block_shape below).
+_SKIP_POW2_BLOCK_SHAPE = FLAGTREE_BACKEND == "ascend"
 
 if TYPE_CHECKING:
     from .language import core
@@ -80,7 +85,13 @@ def validate_block_shape(shape: List[int]):
     for i, d in enumerate(shape):
         if not isinstance(d, int):
             raise TypeError(f"Shape element {i} must have type `constexpr[int]`, got `constexpr[{type(d)}]")
-        if not is_power_of_two(d):
+        # A negative dimension is invalid for every backend.
+        if d < 0:
+            raise ValueError(f"Shape element {i} must be non-negative")
+        # Backends that lower through a SIMD pipeline (e.g. ascend) support
+        # non-power-of-2 block shapes; this check has no builder in scope, so
+        # gate it on the active backend.
+        if not is_power_of_two(d) and not _SKIP_POW2_BLOCK_SHAPE:
             raise ValueError(f"Shape element {i} must be a power of 2")
         numel *= d
 

python/triton/backends/compiler.py

@@ -95,6 +95,17 @@ def get_tensor_descriptor_specialization(arg):
         """Return backend-specific TensorDescriptor cache-key attributes."""
         return None
 
+    @staticmethod
+    def use_alignment_specialization(options: dict) -> bool:
+        """Whether divisibility-by-16 specialization applies for these options.
+
+        Receives raw launch keyword options before parse_options, once per
+        binder call when alignment specialization can apply. The default
+        preserves existing cache keys. Overrides must agree with the options
+        used for compilation; parameter-level opt-outs still take precedence.
+        """
+        return True
+
     @staticmethod
     @abstractmethod
     def supports_target(target: GPUTarget):

python/triton/compiler/code_generator.py

@@ -390,6 +390,12 @@ def __init__(self, context, prototype, gscope, function_name, jit_fn: JITFunctio
         # convert_custom_types(input_tensor, dtype, fp_downcast_rounding=None, _builder=None)
         self.builder.codegen_fns = codegen_fns
         self.builder.module_map = {} if module_map is None else module_map
+        # Backend extensions are instance-local and must be ready before IR emission.
+        # Gluon keeps its own builder and semantics.
+        if not is_gluon:
+            initialize_builder = codegen_fns.get("initialize_builder")
+            if initialize_builder is not None:
+                initialize_builder(self.builder, context, options)
         self.module = self.builder.create_module() if module is None else module
         self.function_ret_types = {} if function_types is None else function_types
         self.prototype = prototype
@@ -1107,6 +1113,23 @@ def visit_IfExp(self, node):
                 return self.visit(node.orelse)
 
     def visit_With(self, node):
+        # Let backends lower specialized context managers before invoking
+        # the standard context manager protocol.
+        with_handlers = self.builder.codegen_fns.get("with_handlers", {})
+        if with_handlers and len(node.items) == 1:
+            item = node.items[0]
+            if isinstance(item.context_expr, ast.Call):
+                context_type = self.visit(item.context_expr.func)
+                handler = with_handlers.get(context_type)
+                if handler is not None:
+                    if item.optional_vars is not None:
+                        raise self._unsupported(
+                            node, "`as` bindings are not supported for backend-specific `with` statements")
+                    if ContainsReturnChecker(self.gscope).visit(node):
+                        raise self._unsupported(
+                            node, "Cannot have `return` statements inside `with` statements in triton ")
+                    return handler(self, node)
+
         # Lower `with` statements by constructing context managers and calling their enter/exit hooks
         # Instantiate each context manager with builder injection
         cm_list = []

python/triton/compiler/compiler.py

@@ -469,8 +469,12 @@ def __init__(self, src, metadata_group, hash):
         # stores the text of each level of IR that was generated during compilation
         asm_files = [Path(p) for c, p in metadata_group.items() if not c.endswith(".json")]
         binary_ext = backend.binary_ext
+        # A backend may emit more than one binary artifact (e.g. the ascend
+        # backend caches the MLIR bytecode stage as .mlirbc alongside .npubin);
+        # those must be read as bytes, not decoded as UTF-8 text.
+        binary_exts = set(getattr(backend, "binary_extensions", None) or {binary_ext})
         self.asm = AsmDict({
-            file.suffix[1:]: file.read_bytes() if file.suffix[1:] == binary_ext else file.read_text()
+            file.suffix[1:]: file.read_bytes() if file.suffix[1:] in binary_exts else file.read_text()
             for file in asm_files
         })
         self.metadata_group = metadata_group
@@ -516,8 +520,16 @@ def raise_(err):
         if knobs.runtime.kernel_load_start_hook is not None:
             knobs.runtime.kernel_load_start_hook(self.module, self.function, self.name, self.metadata_group, self.hash)
         # TODO: n_regs, n_spills should be metadata generated when calling `ptxas`
-        self.module, self.function, self.n_regs, self.n_spills, self.n_max_threads = kernel_driver.utils.load_binary(
-            self.name, self.kernel, self.metadata.shared, device)
+        # The ascend backend's load_binary needs the kernel's mix_mode (aiv/aic/mix)
+        # to pick the launch path, and keys the binary by the mangled kernel_name
+        # rather than the Python function name.
+        if hasattr(self.metadata, "mix_mode"):
+            self.module, self.function, self.n_regs, self.n_spills, self.n_max_threads = \
+                kernel_driver.utils.load_binary(self.metadata.kernel_name, self.kernel, self.metadata.shared, device,
+                                                self.metadata.mix_mode)
+        else:
+            self.module, self.function, self.n_regs, self.n_spills, self.n_max_threads = \
+                kernel_driver.utils.load_binary(self.name, self.kernel, self.metadata.shared, device)
         warp_size = self.metadata.target.warp_size
         if self.metadata.num_warps * warp_size > self.n_max_threads \
             and self.metadata.target.backend != "maca":  # flagtree metax