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

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,10 @@ if(TRITON_BUILD_PYTHON_MODULE)
     ${triton_plugins}
 
     # mlir
+    # Provides mlir::func::registerInlinerExtension, needed by the ascend
+    # backend's commonir-to-hivm lowering (the Inliner pass that follows it
+    # aborts if the func dialect's inliner interface is unregistered).
+    MLIRFuncInlinerExtension
     MLIRAMDGPUDialect
     MLIRNVVMDialect
     MLIRNVVMToLLVMIRTranslation

cmake/FlagTreeOptions.cmake

@@ -45,6 +45,10 @@ macro(flagtree_configure_options)
     set(FLAGCX_ENABLED ON)
   endif()
 
+  # TLE DSA (ascend-only) needs the third_party/tle/dsa subtree, which this
+  # branch does not carry; keep it an explicit opt-in.
+  set(FLAGTREE_TLE_DSA "$ENV{FLAGTREE_TLE_DSA}")
+
   set(FLAGTREE_TLE ON)
   if(FLAGTREE_BACKEND STREQUAL "xpu")
     set(FLAGTREE_TLE OFF)
@@ -82,6 +86,42 @@ 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++)
+    # The ascend backend consumes FLIR's incubated passes (TritonToLinalgIncubated).
+    set(FLIR_BUILD_INCUBATED ON CACHE BOOL "Build FLIR incubated passes" FORCE)
+    # 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: the ascend
+    # CMakeLists uses it to opt into linking MLIRMeshDialect, but LLVM 22
+    # renamed that 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")
+    # TLE DSA is an opt-in extension requiring the third_party/tle/dsa subtree,
+    # which this branch does not carry. The ascend C++ sources guard every use
+    # of it with #ifdef __TLE_DSA__, so leaving this off compiles those paths
+    # out and drops the corresponding CMake dependencies.
+    if(FLAGTREE_TLE_DSA)
+      if(NOT EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/third_party/tle/dsa/CMakeLists.txt")
+        message(FATAL_ERROR "FLAGTREE_TLE_DSA requires third_party/tle/dsa")
+      endif()
+      add_definitions(-D__TLE_DSA__)
+    endif()
+    # 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++)
@@ -328,12 +368,12 @@ macro(flagtree_configure_flir_dependency)
     add_subdirectory(third_party/flir/lib/Dialect/CommonIR)
   endif()
 
-  if(FLAGTREE_BACKEND STREQUAL "tsingmicro")
+  if(FLAGTREE_BACKEND MATCHES "^(tsingmicro|ascend)$")
     if(NOT EXISTS "${PROJECT_SOURCE_DIR}/third_party/flir/CMakeLists.txt")
       message(FATAL_ERROR "The ${FLAGTREE_BACKEND} backend requires third_party/flir")
     endif()
 
-    # TsingMicro only consumes FLIR's C++ targets; do not build its Python/CPU plugin.
+    # These backends only consume FLIR's C++ targets; do not build its Python/CPU plugin.
     set(TRITON_SHARED_BUILD_CPU_BACKEND OFF)
     list(REMOVE_ITEM TRITON_CODEGEN_BACKENDS "flir")
     if(NOT TARGET TritonSharedUtils)
@@ -346,6 +386,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 +630,12 @@ macro(flagtree_configure_tools_and_tests)
       flagtree_add_tle_generated_header_dependencies()
     endif()
     add_subdirectory(test)
+  elseif(FLAGTREE_BACKEND STREQUAL "ascend")
+    # The ascend backend adds its own bin/ (third_party/ascend/CMakeLists.txt),
+    # whose tools link TritonTestAnalysis/TritonTestDialect from test/lib.
+    # Add test/ so those libraries get defined; bin/ is not added here because
+    # the backend supplies its own.
+    add_subdirectory(test)
   elseif(FLAGTREE_BACKEND STREQUAL "iluvatar")
     option(FLAGTREE_ILUVATAR_BUILD_BIN
       "Build third_party/iluvatar/bin tools and lit tests" OFF)
@@ -683,6 +744,13 @@ macro(flagtree_added_python_src)
   endforeach()
   add_library(triton SHARED ${_flagtree_python_sources})
 
+  if(FLAGTREE_BACKEND STREQUAL "ascend")
+    # third_party/ascend/python/src/main.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(),
@@ -888,7 +889,12 @@ def uninstall_triton():
 
 cache = FlagTreeCache()
 
-download_flagtree_third_party("flir", condition=(flagtree_backend == "tsingmicro"), required=True)
+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": "bb13ef6120526aa3df6f950006af5ec0e585d2f6"
+    },
     {
         "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,16 @@
+import os
+import subprocess
+
+
+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/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,10 @@ 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):
+        # 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/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

python/triton/extension/__init__.py

@@ -0,0 +1,4 @@
+# flagtree backend path specialization
+from triton._flagtree_spec import spec_path
+
+spec_path(__path__)

python/triton/knobs.py

@@ -450,6 +450,9 @@ class autotuning_knobs(base_knobs):
     cache: env_bool = env_bool("TRITON_CACHE_AUTOTUNING")
     print: env_bool = env_bool("TRITON_PRINT_AUTOTUNING")
     adjust_block_size: env_bool = env_bool("FLAGTREE_AABS", True)  # flagtree aabs
+    # Opt-in to a backend-provided autotuner (consumed by the ascend runtime,
+    # which replaces triton.autotune when this is set).
+    use_backend_autotuning: env_bool = env_bool("FLAGTREE_USE_BACKEND_AUTOTUNING")
 
 
 class LaunchHook(Protocol):

python/triton/language/semantic.py

@@ -601,7 +601,9 @@ def arange(self, start: int, end: int, *, ret_ty: tl.block_type = None) -> Tenso
         if end <= start:
             raise ValueError("arange's end argument must be greater than the start argument")
         range = end - start
-        if (range & (range - 1)) != 0:
+        # The power-of-2 restriction only applies in SIMT mode; backends whose
+        # builder reports SIMD (e.g. ascend) accept arbitrary ranges.
+        if (range & (range - 1)) != 0 and getattr(self.builder, "is_simt_mode", lambda: True)():
             raise ValueError("arange's range must be a power of 2")
         shape = [range]
         if ret_ty is None:

python/triton/spec/ascend

@@ -0,0 +1 @@
+../../../third_party/ascend/spec/

third_party/ascend/CMakeLists.txt

@@ -0,0 +1,97 @@
+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include)
+include_directories(${CMAKE_CURRENT_BINARY_DIR}/include)
+include_directories(${CMAKE_SOURCE_DIR}/third_party/flir/include)
+include_directories(${CMAKE_BINARY_DIR}/third_party/flir/include)
+
+# list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake")
+
+# Triton Ascend is dependent on AscendNPU IR
+# set(ASCENDNPU_IR_SRC_DIR "${PROJECT_SOURCE_DIR}/third_party/ascend/AscendNPU-IR")
+# set(ASCENDNPU_IR_BINARY_DIR "${PROJECT_BINARY_DIR}/third_party/ascend/AscendNPU-IR")
+
+# set(BISHENGIR_ENABLE_A5_UNPUBLISHED_FEATURES ON)
+# set(BISHENGIR_BUILD_STANDALONE_IR_ONLY ON)
+#
+# # Temporarily save and clear the RULE_LAUNCH_* attributes in the current directory to ensure AscendNPU-IR is not affected.
+# get_property(_saved_launch_compile DIRECTORY PROPERTY RULE_LAUNCH_COMPILE)
+# get_property(_saved_launch_link DIRECTORY PROPERTY RULE_LAUNCH_LINK)
+# set_property(DIRECTORY PROPERTY RULE_LAUNCH_COMPILE "")
+# set_property(DIRECTORY PROPERTY RULE_LAUNCH_LINK "")
+#
+# add_subdirectory(${ASCENDNPU_IR_SRC_DIR} ${ASCENDNPU_IR_BINARY_DIR})
+#
+# # restore properties
+# set_property(DIRECTORY PROPERTY RULE_LAUNCH_COMPILE ${_saved_launch_compile})
+# set_property(DIRECTORY PROPERTY RULE_LAUNCH_LINK ${_saved_launch_link})
+#
+# include_directories(${ASCENDNPU_IR_SRC_DIR}/bishengir/include)
+# include_directories(${ASCENDNPU_IR_BINARY_DIR}/bishengir/include) # Tablegen'd files
+
+add_subdirectory(backend/spec/lib)
+
+add_subdirectory(${PROJECT_SOURCE_DIR}/include
+                 ${PROJECT_BINARY_DIR}/include)
+add_subdirectory(${PROJECT_SOURCE_DIR}/lib
+                 ${PROJECT_BINARY_DIR}/lib)
+add_subdirectory(include)
+add_subdirectory(lib)
+add_subdirectory(bin)
+
+add_triton_library(Registrar Registrar.cc)
+
+if(LLVM_MAJOR_VERSION_21_COMPATIBLE)
+  set(_MLIRMeshDialect_LIB MLIRMeshDialect)
+else()
+  set(_MLIRMeshDialect_LIB "")
+endif()
+
+add_triton_plugin(TritonAscend
+  ${CMAKE_CURRENT_SOURCE_DIR}/triton_ascend.cc
+  ${CMAKE_CURRENT_SOURCE_DIR}/ascend_ir.cc
+  LINK_LIBS
+  TritonToLinalgIncubated
+  BiShengIRScopeDialect
+  BiShengIRHIVMDialect
+  ${_MLIRMeshDialect_LIB}
+)
+# AscendCustomOpsBitcode is built by the TLE DSA subtree; only require it when
+# TLE DSA is part of the build.
+if(FLAGTREE_TLE_DSA)
+  add_dependencies(TritonAscend AscendCustomOpsBitcode)
+endif()
+
+target_include_directories(TritonAscend BEFORE PRIVATE
+  ${CMAKE_CURRENT_SOURCE_DIR}/python/src)
+
+option(TRITON_ENABLE_COVERAGE_LLVM_COV "Enable code llvm-cov coverage tool for Ascend plugin " OFF)
+if(TRITON_ENABLE_COVERAGE_LLVM_COV)
+  message(STATUS "Enabling llvm-cov coverage tool flags for TritonAscend")
+  target_compile_options(TritonAscend PRIVATE
+    -fprofile-arcs
+    -ftest-coverage
+    -O0
+    -fprofile-update=atomic
+    --coverage
+  )
+  target_link_options(TritonAscend PRIVATE
+    --coverage
+    -lgcov
+  )
+  # branch coverage
+  target_compile_definitions(TritonAscend PRIVATE
+    COVERAGE_ENABLED=1
+  )
+endif()
+
+
+# To enable hitest coverage tool
+if(TRITON_ENABLE_COVERAGE_HITEST)
+  set_target_properties(TritonAscend PROPERTIES
+    RULE_LAUNCH_COMPILE "hitestwrapper"
+    RULE_LAUNCH_LINK "hitestwrapper"
+  )
+endif()
+
+if(TRITON_BUILD_UT)
+  add_subdirectory(unittest)
+endif()