[BACKEND] Sync triton-ascend release/3.2.1 updates

bin/RegisterTritonDialects.h

@@ -1,12 +1,11 @@
 #pragma once
-
-#ifdef __AMD__
-#include "amd/include/Dialect/TritonAMDGPU/IR/Dialect.h"
-#include "amd/include/TritonAMDGPUTransforms/Passes.h"
-#endif
-#ifdef __NVIDIA__
-#include "third_party/nvidia/include/Dialect/NVGPU/IR/Dialect.h"
-#endif
+#include "AutoBlockify/Passes.h"
+#include "TritonToHFusion/Passes.h"
+#include "TritonToHIVM/Passes.h"
+#include "TritonToLLVM/Passes.h"
+// #include "amd/include/Dialect/TritonAMDGPU/IR/Dialect.h"
+// #include "amd/include/TritonAMDGPUTransforms/Passes.h"
+// #include "third_party/nvidia/include/Dialect/NVGPU/IR/Dialect.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
 #ifdef __NVIDIA__
@@ -66,11 +65,13 @@ inline void registerTritonDialects(mlir::DialectRegistry &registry) {
 #endif
   mlir::triton::registerConvertTritonToTritonGPUPass();
   mlir::triton::registerAllocateSharedMemoryPass();
-#ifdef __NVIDIA__
-  mlir::triton::registerConvertTritonGPUToLLVMPass();
-  mlir::triton::registerConvertNVGPUToLLVMPass();
-  mlir::triton::registerDecomposeUnsupportedNVIDIAConversions();
-#endif
+  // mlir::triton::registerConvertTritonGPUToLLVMPass();
+  // mlir::triton::registerConvertNVGPUToLLVMPass();
+  // mlir::triton::registerDecomposeUnsupportedNVIDIAConversions();
+  mlir::triton::registerTritonToHIVMPasses();
+  mlir::triton::registerTritonToHFusionPasses();
+  mlir::triton::registerTritonToLLVMPasses();
+  mlir::triton::registerAutoBlockifyPasses();
   mlir::registerLLVMDIScope();
 
 #ifdef __AMD__

include/triton/Tools/Sys/GetEnv.hpp

@@ -20,6 +20,7 @@ inline const std::set<std::string> CACHE_INVALIDATING_ENV_VARS = {
     "DISABLE_PTXAS_OPT",
     "LLVM_IR_ENABLE_DUMP",
     "LLVM_ENABLE_TIMING",
+    "MLIR_DISABLE_MULTITHREADING",
     "LLVM_PASS_PLUGIN_PATH",
     "MLIR_ENABLE_DIAGNOSTICS",
     "MLIR_ENABLE_DUMP",

python/src/llvm.cc

@@ -25,6 +25,7 @@
 #include "llvm/Transforms/InstCombine/InstCombine.h"
 #include <csignal>
 #include <memory>
+#include <pybind11/gil.h>
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 #include <stdexcept>
@@ -172,7 +173,7 @@ void init_triton_llvm(py::module &&m) {
           [](llvm::Module::FunctionListType &s) {
             return py::make_iterator(s.begin(), s.end());
           },
-          py::keep_alive<0, 1>());
+          py::keep_alive<0, 1>(), py::call_guard<py::gil_scoped_release>());
 
   // Module Flag behavior. See
   // https://llvm.org/doxygen/classllvm_1_1Module.html#a0a5c55e12c97b80021330fe82b642293
@@ -388,7 +389,8 @@ void init_triton_llvm(py::module &&m) {
       // (optional) parameters
       py::arg("arch") = "", py::arg("features") = "",
       py::arg("flags") = std::vector<std::string>{},
-      py::arg("enable_fp_fusion") = false);
+      py::arg("enable_fp_fusion") = false,
+      py::call_guard<py::gil_scoped_release>());
 
   m.def(
       "translate_to_asm",

python/triton/__init__.py

@@ -23,6 +23,7 @@
     MockTensor,
 )
 from .runtime.jit import jit
+from .runtime._async_compile import AsyncCompileMode
 from .compiler import compile, CompilationError
 from .errors import TritonError
 
@@ -31,6 +32,7 @@
 from . import tools
 
 __all__ = [
+    "AsyncCompileMode",
     "autotune",
     "cdiv",
     "CompilationError",

python/triton/backends/__init__.py

@@ -35,7 +35,12 @@ class Backend:
 def _discover_backends():
     backends = dict()
     root = os.path.dirname(__file__)
+    # The package does not ship the files required to load the
+    # upstream nvidia and amd backends, so skip discovering them here.
+    ignored_dirs = {"nvidia", "amd"}
     for name in os.listdir(root):
+        if name in ignored_dirs:
+            continue
         if not os.path.isdir(os.path.join(root, name)):
             continue
         if name.startswith('__'):

python/triton/extension/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2018-2020 Philippe Tillet
+# Copyright 2020-2022 OpenAI
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.

python/triton/extension/buffer/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2018-2020 Philippe Tillet
+# Copyright 2020-2022 OpenAI
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.

python/triton/extension/buffer/language/__init__.py

@@ -0,0 +1,44 @@
+# Copyright 2018-2020 Philippe Tillet
+# Copyright 2020-2022 OpenAI
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+__all__ = [
+    # core
+    "builtin",
+    "is_builtin",
+
+    # buffer
+    "buffer",
+
+    # base address space
+    "address_space",
+
+    # alloc
+    "alloc",
+
+    # to_buffer
+    "to_buffer",
+
+    # to_tensor
+    "to_tensor",
+    "subview",
+]
+
+from .core import builtin, is_builtin, address_space, buffer, alloc, to_buffer, to_tensor, subview

python/triton/extension/buffer/language/builder.py

@@ -0,0 +1,73 @@
+# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+"""
+buffer-specific builder utilities for code generation.
+"""
+
+__all__ = [
+    "create_builder_method_wrapper_with_buffer_builder",
+    "attach_builder_methods_with_buffer_builder",
+    "setup_unified_builder_with_buffer_builder",
+]
+
+
+def create_builder_method_wrapper_with_buffer_builder(main_builder, delegate_builder, method_name):
+    """
+    Create a wrapper that delegates a method call to another builder while
+    synchronizing insertion points and locations.
+    """
+    delegate_method = getattr(delegate_builder, method_name)
+
+    def wrapper(*args, **kwargs):
+        saved_ip = main_builder.get_insertion_point()
+        saved_loc = main_builder.get_loc()
+        delegate_builder.restore_insertion_point(saved_ip)
+        if saved_loc:
+            delegate_builder.set_loc(saved_loc)
+        result = delegate_method(*args, **kwargs)
+        main_builder.restore_insertion_point(saved_ip)
+        if saved_loc:
+            main_builder.set_loc(saved_loc)
+        return result
+
+    wrapper.__name__ = method_name
+    wrapper.__doc__ = getattr(delegate_method, '__doc__', None)
+    return wrapper
+
+
+def attach_builder_methods_with_buffer_builder(main_builder, delegate_builder, method_names):
+    """Attach multiple methods from a delegate builder to the main builder."""
+    for method_name in method_names:
+        wrapper = create_builder_method_wrapper_with_buffer_builder(main_builder, delegate_builder, method_name)
+        setattr(main_builder, method_name, wrapper)
+
+
+def setup_unified_builder_with_buffer_builder(main_builder, buffer_builder):
+    """Set up a unified builder interface by attaching methods from specialized builders."""
+    main_builder._buffer_builder = buffer_builder
+    buffer_methods = [
+        'get_null_attr',
+        'get_str_array_attr',
+        'alloc',
+        'to_buffer',
+        'to_tensor',
+        'subview',
+    ]
+    attach_builder_methods_with_buffer_builder(main_builder, buffer_builder, buffer_methods)