diff --git a/.github/workflows/format.yaml b/.github/workflows/format.yaml index 1d35bcb4..a7475748 100644 --- a/.github/workflows/format.yaml +++ b/.github/workflows/format.yaml @@ -34,20 +34,33 @@ jobs: clang-format: needs: markdownlint runs-on: ubuntu-latest + steps: - uses: actions/checkout@v4 - - uses: cpp-linter/cpp-linter-action@v2 - id: cpp-lint with: - style: file - ignore: | - third_party/* - .git - tidy-checks: '-*' # disable clang tidy at this stage - version: 17 - - name: Fail test - if: steps.cpp-lint.outputs.checks-failed > 0 - run: echo "Some files failed the linting checks!" && exit 1 + fetch-depth: 0 + + - name: Install clang-format + run: | + sudo apt-get update + sudo apt-get install -y clang-format-17 + + - name: Check changed files + run: | + git diff --diff-filter=d --name-only \ + origin/${{ github.base_ref }}...HEAD \ + > all_changed_files.txt + + awk '/\.(cc|cpp|cxx|h|hpp)$/ && $0 !~ /^third_party\// { print }' \ + all_changed_files.txt \ + > changed_files.txt + + if [ ! -s changed_files.txt ]; then + echo "No C/C++ files changed" + exit 0 + fi + + xargs -a changed_files.txt clang-format-17 --dry-run -Werror python-black: needs: markdownlint @@ -58,4 +71,4 @@ jobs: fetch-depth: 0 # 确保能获取到 main 分支的历史 - name: Check formatting for PR files run: | - bash scripts/ci_python_black.sh \ No newline at end of file + bash scripts/ci_python_black.sh diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 8a2c501f..d993ca93 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -10,7 +10,7 @@ on: - main env: - CI_PATH: "${{ vars.CI_BASE_PATH }}/GitHub/${{ github.repository }}/${GITHUB_RUN_NUMBER}" + CI_PATH: "${{ vars.CI_BASE_PATH }}/GitHub/${{ github.repository }}/${{ github.run_number }}" THIRD_PARTY_PATH: "${{ vars.CI_BASE_PATH }}/data/DLCompiler/third_party" concurrency: @@ -31,23 +31,21 @@ jobs: - name: Create custom directory run: | set -ex - echo ${{ env.CI_PATH }} - mkdir -p ${{ env.CI_PATH }} + echo "$CI_PATH" + mkdir -p "$CI_PATH" - name: Clean custom directory run: | set -ex - if [ -d "${{ env.CI_PATH }}" ]; then - rm -rf ${{ env.CI_PATH }}/* - rm -rf ${{ env.CI_PATH }}/.github - fi + rm -rf "$CI_PATH" + mkdir -p "$CI_PATH" - name: Move code to custom directory run: | set -ex - mv $GITHUB_WORKSPACE/* ${{ env.CI_PATH }}/ - mv $GITHUB_WORKSPACE/.github ${{ env.CI_PATH }}/ - mv $GITHUB_WORKSPACE/.git ${{ env.CI_PATH }}/ + mv "$GITHUB_WORKSPACE"/* "$CI_PATH"/ + mv "$GITHUB_WORKSPACE/.github" "$CI_PATH"/ + mv "$GITHUB_WORKSPACE/.git" "$CI_PATH"/ - name: Setup Git Safe Directory run: | @@ -68,16 +66,16 @@ jobs: which conda echo "which conda? $(which conda)" conda activate dlcompiler - source /usr/local/Ascend/cann-8.5.0/set_env.sh + source /usr/local/Ascend/ascend-toolkit/set_env.sh + source /usr/local/Ascend/cann-9.1.0-beta.1/share/info/ascendnpu-ir/bin/set_env.sh cd ${{ env.CI_PATH }} - export JSON_PATH34=${{ vars.CI_BASE_PATH }}/data/v34/include.zip - export GOOGLETEST_DIR34=${{ vars.CI_BASE_PATH }}/data/v34/googletest - export LLVM_TGZ_PATH34=${{ vars.CI_BASE_PATH }}/data/v34/llvm-064f02da-ubuntu-arm64.tar.gz + export JSON_PATH35=${{ vars.CI_BASE_PATH }}/data/v35/include.zip + export GOOGLETEST_DIR35=${{ vars.CI_BASE_PATH }}/data/v35/googletest + export LLVM_TGZ_PATH35=${{ vars.CI_BASE_PATH }}/data/v35/llvm-7d5de303-ubuntu-arm64.tar.gz rm -rf ./third_party/* git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/ascendnpu-ir ./third_party/ascendnpu-ir git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/json ./third_party/json - git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/triton_shared ./third_party/triton_shared git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/triton ./third_party/triton echo "whoami before compile: $(whoami)" echo "which python? $(which python)" @@ -89,37 +87,39 @@ jobs: bash compile_shared.sh apply_patch=true ' - - name: Build and install tilelang-dlc - run: | - set -ex - # 切换到 root 用户执行 - sudo -E bash -c ' - source /home/dlc_ci/.bashrc - conda activate dlcompiler - source /usr/local/Ascend/cann-8.5.0/set_env.sh - cd ${{ env.CI_PATH }} - export TILELANG_DLC_PATH=${{ vars.CI_BASE_PATH }}/data/tilelang-dlc - export DLCOMPILER_SOURCE=${{ env.CI_PATH }} - export TILELANG_USE_DLCOMPILER=1 - echo "whoami? $(whoami)" - echo "which python? $(which python)" - bash scripts/install_tilelang-dlc.sh - ' - - - name: Run tilelang-dlc tests on ascend - run: | - set -ex - # 切换到 root 用户执行 - sudo -E bash -c ' - source /home/dlc_ci/.bashrc - conda activate dlcompiler - source /usr/local/Ascend/cann-8.5.0/set_env.sh - cd ${{env.CI_PATH }} - export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH - export ASCEND_RT_VISIBLE_DEVICES=7 - export TILELANG_USE_DLCOMPILER=1 - bash test/commonir/run_tests.sh - ' + # - name: Build and install tilelang-dlc + # run: | + # set -ex + # # 切换到 root 用户执行 + # sudo -E bash -c ' + # source /home/dlc_ci/.bashrc + # conda activate dlcompiler + # source /usr/local/Ascend/ascend-toolkit/set_env.sh + # source /usr/local/Ascend/cann-9.1.0-beta.1/share/info/ascendnpu-ir/bin/set_env.sh + # cd ${{ env.CI_PATH }} + # export TILELANG_DLC_PATH=${{ vars.CI_BASE_PATH }}/data/tilelang-dlc + # export DLCOMPILER_SOURCE=${{ env.CI_PATH }} + # export TILELANG_USE_DLCOMPILER=1 + # echo "whoami? $(whoami)" + # echo "which python? $(which python)" + # bash scripts/install_tilelang-dlc.sh + # ' + + # - name: Run tilelang-dlc tests on ascend + # run: | + # set -ex + # # 切换到 root 用户执行 + # sudo -E bash -c ' + # source /home/dlc_ci/.bashrc + # conda activate dlcompiler + # source /usr/local/Ascend/ascend-toolkit/set_env.sh + # source /usr/local/Ascend/cann-9.1.0-beta.1/share/info/ascendnpu-ir/bin/set_env.sh + # cd ${{env.CI_PATH }} + # export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH + # export ASCEND_RT_VISIBLE_DEVICES=7 + # export TILELANG_USE_DLCOMPILER=1 + # bash test/commonir/run_tests.sh + # ' - name: Run triton tests on ascend run: | @@ -128,11 +128,11 @@ jobs: sudo -E bash -c ' source /home/dlc_ci/.bashrc conda activate dlcompiler - source /usr/local/Ascend/cann-8.5.0/set_env.sh + source /usr/local/Ascend/ascend-toolkit/set_env.sh + source /usr/local/Ascend/cann-9.1.0-beta.1/share/info/ascendnpu-ir/bin/set_env.sh cd ${{env.CI_PATH }} echo "whoami? $(whoami)" echo "which python? $(which python)" - export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH export ASCEND_RT_VISIBLE_DEVICES=7 bash test/ascend/run_tests.sh ' @@ -144,11 +144,11 @@ jobs: sudo -E bash -c ' source /home/dlc_ci/.bashrc conda activate dlcompiler - source /usr/local/Ascend/cann-8.5.0/set_env.sh + source /usr/local/Ascend/ascend-toolkit/set_env.sh + source /usr/local/Ascend/cann-9.1.0-beta.1/share/info/ascendnpu-ir/bin/set_env.sh cd ${{env.CI_PATH }} echo "whoami? $(whoami)" echo "which python? $(which python)" - export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH export ASCEND_RT_VISIBLE_DEVICES=7 bash test/ascend/test_mlir.sh ' @@ -160,11 +160,11 @@ jobs: sudo -E bash -c ' source /home/dlc_ci/.bashrc conda activate dlcompiler - source /usr/local/Ascend/cann-8.5.0/set_env.sh + source /usr/local/Ascend/ascend-toolkit/set_env.sh + source /usr/local/Ascend/cann-9.1.0-beta.1/share/info/ascendnpu-ir/bin/set_env.sh cd ${{env.CI_PATH }} echo "whoami? $(whoami)" echo "which python? $(which python)" - export PATH=${{ vars.CI_BASE_PATH }}/data/bishengir_latest/:$PATH export ASCEND_RT_VISIBLE_DEVICES=7 bash test/dsl/run_tests.sh ' diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ae24cfc4..e9cbbf0b 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -6,7 +6,7 @@ on: - "v*.*.*" # 只对形如 v1.2.3 的 tag 触发 env: - RELEASE_PATH: "${{ vars.CI_BASE_PATH }}/GitHub/Release/${{ github.repository }}/${GITHUB_RUN_NUMBER}" + RELEASE_PATH: "${{ vars.CI_BASE_PATH }}/GitHub/Release/${{ github.repository }}/${{ github.run_number }}" THIRD_PARTY_PATH: "${{ vars.CI_BASE_PATH }}/data/DLCompiler/third_party" concurrency: @@ -27,23 +27,21 @@ jobs: - name: Create custom directory run: | set -ex - echo ${{ env.RELEASE_PATH }} - mkdir -p ${{ env.RELEASE_PATH }} + echo "$RELEASE_PATH" + mkdir -p "$RELEASE_PATH" - name: Clean custom directory run: | set -ex - if [ -d "${{ env.RELEASE_PATH }}" ]; then - rm -rf ${{ env.RELEASE_PATH }}/* - rm -rf ${{ env.RELEASE_PATH }}/.github - fi + rm -rf "$RELEASE_PATH" + mkdir -p "$RELEASE_PATH" - name: Move code to custom directory run: | set -ex - mv $GITHUB_WORKSPACE/* ${{ env.RELEASE_PATH }}/ - mv $GITHUB_WORKSPACE/.github ${{ env.RELEASE_PATH }}/ - mv $GITHUB_WORKSPACE/.git ${{ env.RELEASE_PATH }}/ + mv "$GITHUB_WORKSPACE"/* "$RELEASE_PATH"/ + mv "$GITHUB_WORKSPACE/.github" "$RELEASE_PATH"/ + mv "$GITHUB_WORKSPACE/.git" "$RELEASE_PATH"/ - name: Setup Git Safe Directory run: | @@ -60,15 +58,15 @@ jobs: source /home/dlc_ci/.bashrc source activate dlcompiler source /usr/local/Ascend/ascend-toolkit/set_env.sh + source /usr/local/Ascend/cann-9.1.0-beta.1/share/info/ascendnpu-ir/bin/set_env.sh cd ${{ env.RELEASE_PATH }} - export JSON_PATH34=${{ vars.CI_BASE_PATH }}/data/v34/include.zip - export GOOGLETEST_DIR34=${{ vars.CI_BASE_PATH }}/data/v34/googletest - export LLVM_TGZ_PATH34=${{ vars.CI_BASE_PATH }}/data/v34/llvm-064f02da-ubuntu-arm64.tar.gz + export JSON_PATH35=${{ vars.CI_BASE_PATH }}/data/v35/include.zip + export GOOGLETEST_DIR35=${{ vars.CI_BASE_PATH }}/data/v35/googletest + export LLVM_TGZ_PATH35=${{ vars.CI_BASE_PATH }}/data/v35/llvm-7d5de303-ubuntu-arm64.tar.gz rm -rf ./third_party/* git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/ascendnpu-ir ./third_party/ascendnpu-ir git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/json ./third_party/json - git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/triton_shared ./third_party/triton_shared git clone --no-hardlinks ${{ env.THIRD_PARTY_PATH }}/triton ./third_party/triton echo "whoami? $(whoami)" echo "which python? $(which python)" diff --git a/.gitignore b/.gitignore index eb685574..c86d2dd8 100644 --- a/.gitignore +++ b/.gitignore @@ -99,6 +99,4 @@ fusion_result.json launcher_cxx11abi* # package -backend/triton-shared-opt-v3* -backend/dicp_opt -third_party/triton-shared-opt \ No newline at end of file +backend/dicp_opt \ No newline at end of file diff --git a/.gitmodules b/.gitmodules index b46ce5b7..2aadc14b 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,14 +1,8 @@ -[submodule "third_party/triton_linalg"] - path = third_party/triton_linalg - url = git@github.com:Cambricon/triton-linalg.git [submodule "third_party/triton"] path = third_party/triton url = git@github.com:triton-lang/triton.git -[submodule "third_party/triton_shared"] - path = third_party/triton_shared - url = git@github.com:microsoft/triton-shared.git - depth = 1 + branch = v3.5.0 [submodule "third_party/ascendnpu-ir"] path = third_party/ascendnpu-ir url = https://gitcode.com/Ascend/AscendNPU-IR.git - depth = 1 + branch = ef9139b323e25e8dae0a812ac585f3b47ab5d955 \ No newline at end of file diff --git a/CMakeLists.txt b/CMakeLists.txt index 693c96bb..f14d6a15 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,15 +1,17 @@ set(DC_TRITON_SOURCE_DIR ${CMAKE_CURRENT_SOURCE_DIR}) set(DC_TRITON_BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}) -set(TRTION_SHARED_NPU_SPECIFIC_SOURCES "${DC_TRITON_SOURCE_DIR}/compiler/lib/Conversion/TritonToLinalgNPU") set(DC_TRITON_INCLUDE_DIR "") set(DC_TRITON_LINK_DIR "") include_directories(${CMAKE_CURRENT_BINARY_DIR}/compiler/include) include_directories(${CMAKE_CURRENT_SOURCE_DIR}/compiler/include) -include_directories(${CMAKE_CURRENT_BINARY_DIR}/third_party/triton_shared/include) include_directories(${CMAKE_CURRENT_SOURCE_DIR}/third_party/json/include) -include_directories(${CMAKE_CURRENT_SOURCE_DIR}/third_party/triton_shared/include) + +# ASCENDNPU_IR编译选项 +set(LLVM_MAJOR_VERSION_22_COMPATIBLE ON CACHE BOOL "NPUIR build with LLVM 22" FORCE) +add_definitions(-D__LLVM_MAJOR_VERSION_22_COMPATIBLE__) +add_compile_options(-Wno-switch) set(ASCENDNPU_IR_SRC_DIR ${CMAKE_CURRENT_SOURCE_DIR}/third_party/ascendnpu-ir) set(ASCENDNPU_IR_BINARY_DIR ${CMAKE_CURRENT_BINARY_DIR}/third_party/ascendnpu-ir) @@ -31,18 +33,17 @@ ${dialect_libs} ${conversion_libs} ) -list(APPEND CMAKE_MODULE_PATH +list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_LIST_DIR}/cmake/ ) include(utils) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/tools) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/compiler) -add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/third_party/triton_shared) if (TRITON_BUILD_PYTHON_MODULE) - add_triton_plugin(tritonDicpTriton ${CMAKE_CURRENT_SOURCE_DIR}/triton_dicp_triton.cc - LINK_LIBS + add_triton_plugin(tritonDicpTriton ${CMAKE_CURRENT_SOURCE_DIR}/triton_dicp_triton.cc + LINK_LIBS MLIRAffineToStandard MLIRIR @@ -50,16 +51,27 @@ if (TRITON_BUILD_PYTHON_MODULE) MLIRTransforms MLIRSupport MLIRBytecodeWriter - - TritonToLinalgNPUCoversion - - LinalgExtTransforms - TritonExtTransforms + MLIRFuncInlinerExtension + MLIRBufferizationDialect + MLIRMemRefDialect - LinalgToLinked - LinkedToHIVM - DiscreteMaskAccessConversion + # DICP NPU pass libraries + TritonToLinalg + TritonToStructured TritonToUnstructure + TritonToHIVM + TritonToAnnotation + TritonToHFusion + TritonToLLVM + TritonToGraph + AutoBlockify + AscendLegalize + CommonIRTransforms + DiscreteMaskAccessConversion + MLIRTritonNPUUtils + TritonAffinityOpt + TritonDicpIR + TritonStructuredIR ) target_link_libraries(tritonDicpTriton PRIVATE Python3::Module pybind11::headers) -endif() \ No newline at end of file +endif() diff --git a/README.md b/README.md index 5046c79b..8d541022 100755 --- a/README.md +++ b/README.md @@ -54,7 +54,7 @@ bash compile_shared.sh apply_patch=true # 如果不应用patch，可以直 ### 查看编译过程的mlir文件 ```bash -export DLC_DUMP_IR=1 # 默认在当前目录下 +export TRITON_DEBUG=1 # 默认在当前目录下 ``` ### 测试 diff --git a/backend/ascend_autotune_hooks.py b/backend/ascend_autotune_hooks.py new file mode 100644 index 00000000..f5f5eade --- /dev/null +++ b/backend/ascend_autotune_hooks.py @@ -0,0 +1,111 @@ +"""Ascend auto-tune runtime hooks: module-level proxy switch. + +Called from DICPDriver.__init__ when backend="ascend". + +Why module-level replacement instead of subclassing at import time? + - ``@triton.autotune()`` creates tuner instances at decoration time + - Driver init may happen AFTER decoration, so static subclassing won't + affect existing instances. + - Module-level replacement of ``triton.autotune`` / ``triton.max_autotune`` + affects ALL subsequently decorated kernels regardless of creation order. + +Design: flat module, no classes, no registry. + - ``_is_ascend_backend()`` dynamically checks whether the active driver + target is ``"ascend"``. + - Proxy functions check this at call time and dispatch to either + stock-triton or ascend-enhanced (lazy-import on first use). + - ``triton.autotune`` / ``triton.max_autotune`` are replaced **on import** + of this module, so the proxy is always in place before any user code + runs ``@triton.autotune``. +""" + +import triton + +# ------------------------------------------------------------------ +# State +# ------------------------------------------------------------------ + +# Lazy-loaded ascend implementations (filled on first ascend call). +_ASCEND_AUTOTUNE = None +_ASCEND_MAX_AUTOTUNE = None + + +def _is_ascend_backend(): + """Return True when the active driver target is 'ascend'. + + ``driver.active`` is a lazy property — the first access triggers + ``_create_driver()``, which discovers the NPU and constructs a + ``DICPDriver(target='ascend')``. We check the ``target`` attribute + rather than the driver type so the logic is duck-type-safe. + """ + try: + return getattr(triton.runtime.driver.active, "target", None) == "ascend" + except Exception: + return False + + +def _ascend_autotune_fn(): + """Lazy singleton accessor for ascend ``autotune``.""" + global _ASCEND_AUTOTUNE + if _ASCEND_AUTOTUNE is None: + from .ascend_autotune_runtime.autotuner import autotune as _fn + + _ASCEND_AUTOTUNE = _fn + return _ASCEND_AUTOTUNE + + +def _ascend_max_autotune_fn(): + """Lazy singleton accessor for ascend ``max_autotune``.""" + global _ASCEND_MAX_AUTOTUNE + if _ASCEND_MAX_AUTOTUNE is None: + from .ascend_autotune_runtime.autotuner import max_autotune as _fn + + _ASCEND_MAX_AUTOTUNE = _fn + return _ASCEND_MAX_AUTOTUNE + + +# ------------------------------------------------------------------ +# Proxies (installed once on import of this module) +# ------------------------------------------------------------------ +def _autotune_proxy(configs, key, **kwargs): + if _is_ascend_backend(): + return _ascend_autotune_fn()(configs=configs, key=key, **kwargs) + from triton.runtime.autotuner import autotune as _stock + + return _stock(configs=configs, key=key, **kwargs) + + +def _max_autotune_proxy(configs, key, kernel_type="mixcv", **kwargs): + if _is_ascend_backend(): + return _ascend_max_autotune_fn()( + configs=configs, key=key, kernel_type=kernel_type, **kwargs + ) + # Non-ascend: max_autotune is ascend-only; fall back to plain autotune + # (ascend-specific params such as kernel_type are silently ignored). + return triton.autotune(configs=configs, key=key, **kwargs) + + +# Replace immediately — proxy is always in place, check decides the path. +triton.autotune = _autotune_proxy +triton.max_autotune = _max_autotune_proxy + + +# ------------------------------------------------------------------ +# Public API — kept as no-ops for backward compatibility. +# ------------------------------------------------------------------ +def hook_autotune_for_ascend(): + """No-op: autotune auto-detects the ascend backend at call time. + + Retained for backward compatibility — callers such as + ``DICPDriver.__init__`` and test conftest files still invoke this, + but the proxy now checks ``driver.active.target`` dynamically. + """ + pass + + +def unhook_autotune_for_ascend(): + """No-op: autotune auto-detects the ascend backend at call time. + + Retained for backward compatibility with test teardown code. + """ + pass diff --git a/backend/ascend_autotune_runtime/__init__.py b/backend/ascend_autotune_runtime/__init__.py new file mode 100644 index 00000000..b5a99865 --- /dev/null +++ b/backend/ascend_autotune_runtime/__init__.py @@ -0,0 +1,62 @@ +"""Ascend autotune runtime: auto-tiling + compile-option search. + +Previously from triton-ascend/third_party/ascend/backend/runtime. +Adapted for DLCompiler (triton.backends.dicp_triton). +""" + +from .autoparser import ( + AutoParser, + AxesKeyParser, + SplitAxesParser, + TilingAxesParser, + ReductionAxesParser, + LowDimsAxesParser, + PtrNumsParser, +) +from .tile_generator import AxisInfo, BlockInfo, KernelMeta, TileGenerator +from .compile_options import ( + CompileOptionsSpec, + expand_compile_option_configs, + parse_compile_options_hint, +) +from .autotuner import ( + AutoTilingTuner, + autotune, + max_autotune, + get_max_configs, + BaseAutotuner, + CubeAutotuner, + MixcvAutotuner, + VectorAutotuner, + get_autotune_cube_config, + get_autotune_cv_config, + get_autotune_vector_config, +) + +__all__ = [ + "AutoParser", + "AxesKeyParser", + "SplitAxesParser", + "TilingAxesParser", + "ReductionAxesParser", + "LowDimsAxesParser", + "PtrNumsParser", + "AxisInfo", + "BlockInfo", + "KernelMeta", + "TileGenerator", + "CompileOptionsSpec", + "expand_compile_option_configs", + "parse_compile_options_hint", + "AutoTilingTuner", + "autotune", + "max_autotune", + "get_max_configs", + "BaseAutotuner", + "CubeAutotuner", + "MixcvAutotuner", + "VectorAutotuner", + "get_autotune_cube_config", + "get_autotune_cv_config", + "get_autotune_vector_config", +] diff --git a/backend/ascend_autotune_runtime/autoparser.py b/backend/ascend_autotune_runtime/autoparser.py new file mode 100644 index 00000000..6c129a2e --- /dev/null +++ b/backend/ascend_autotune_runtime/autoparser.py @@ -0,0 +1,1117 @@ +# 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. + +import ast +from typing import Dict, List, Union + + +class AutoParser(ast.NodeVisitor): + """ + Base class for parsing triton dsl kernel code using AST analysis. + + Provides common functionality for traversing the AST (abstract syntax tree) + of a triton dsl kernel function and identifying specific elements in the code. + + Subclassed should implement specific parsing logic by overriding the relevant + node visit methods. + """ + + def __init__(self, func_ast: ast.AST): + self.func_ast = func_ast + + def parse(self): + self.visit(self.func_ast) + + def contains_target_var(self, node, var): + """ + Recursively checks if a given AST node or its children contain a reference + to the specified variable. + + :param node: the AST node to check + :type node: ast.AST + :param var: the variable name to search for + :type var: str + :return: True if the variable is found, False otherwise + """ + if isinstance(node, ast.Name) and node.id == var: + return True + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + if self.contains_target_var(item, var): + return True + elif isinstance(value, ast.AST): + if self.contains_target_var(value, var): + return True + return False + + +class DotCallParser(AutoParser): + """ + Detects whether a Triton kernel source contains tensor-core dot calls. + """ + + def __init__(self, func_ast: ast.AST, scope=None, seen=None): + super().__init__(func_ast) + self.has_dot = False + self.scope = scope or {} + self.seen = seen or set() + + def parse(self): + super().parse() + return self.has_dot + + def visit_Call(self, node): + if self._is_tl_dot_call(node.func): + self.has_dot = True + return + if self._called_jit_has_dot(node.func): + self.has_dot = True + return + self.generic_visit(node) + + @staticmethod + def _is_tl_dot_call(func): + return ( + isinstance(func, ast.Attribute) + and func.attr in ("dot", "dot_scaled") + and isinstance(func.value, ast.Name) + and func.value.id == "tl" + ) + + def _called_jit_has_dot(self, func): + if not isinstance(func, ast.Name): + return False + callee = self.scope.get(func.id) + if callee is None or not callable(getattr(callee, "parse", None)): + return False + callee_id = id(callee) + if callee_id in self.seen: + return False + self.seen.add(callee_id) + callee_scope = ( + callee.get_capture_scope() + if callable(getattr(callee, "get_capture_scope", None)) + else self.scope + ) + return DotCallParser(callee.parse(), callee_scope, self.seen).parse() + + +class AxesKeyParser(AutoParser): + """ + A parser for extracting axis information from a given function's AST. + This class is designed to handle specific patterns in the function's code to + determine the axis associated with a given variable. It is particularly useful + for parsing triton DSL kernel code and identifying axis information. + It recursively processes assignment nodes and lessthan nodes to obtain the axes + corresponding to the specified var in the given function. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str]): + super().__init__(func_ast) + self.keys = keys + self.checked_vars = list() + + def get_axis(self, var: str, node=None): + """ + Traverse the AST using the provided variable name and mask-based less-than + operations to obtain the corresponding axis name. + + :param var: the variable name to get the corresponding axis. + :type var: str + """ + if var in self.checked_vars: + return None + axis = None + if not node: + node = self.func_ast + for child_node in ast.walk(node): + # handle compare node + if isinstance(child_node, ast.Compare): + axis = self.handle_lt_node(var, child_node) + elif isinstance(child_node, ast.Assign): + axis = self.handle_assign_node(var, child_node) + + elif isinstance(child_node, ast.BinOp) and isinstance( + child_node.op, ast.BitAnd + ): + + axis = self.handle_lt_node(var, child_node.left) + if axis is None: + axis = self.handle_lt_node(var, child_node.right) + + if axis is not None: + return axis + self.checked_vars.append(var) + return None + + def handle_assign_node(self, var, node): + if not isinstance(node, ast.Assign) or not isinstance(node.targets, list): + return None + if len(node.targets) != 1 or not isinstance(node.targets[0], ast.Name): + return None + + target = node.targets[0].id + if target in self.checked_vars: + return None + # Prevent cyclic assignment. + if var == target or not self.contains_target_var(node.value, var): + return None + + axis = self.get_axis(var, node.value) + if axis: + return axis + + axis = self.get_axis(target) + return axis + + def handle_lt_node(self, var, node): + if not isinstance(node, ast.Compare) or not isinstance(node.ops, list): + return None + if len(node.ops) != 1 or not isinstance(node.ops[0], ast.Lt): + return None + if not isinstance(node.comparators, list) or len(node.comparators) != 1: + return None + if not isinstance(node.left, ast.Name) or var != node.left.id: + return None + + comparator = node.comparators[0] + if not isinstance(comparator, ast.Name) and not ( + isinstance(comparator, ast.Call) + and isinstance(comparator.func, ast.Name) + and comparator.func.id == "min" + ): + return None + + for k, v in self.keys.items(): + if self.contains_target_var(comparator, v): + return k + return None + + +class SplitAxesParser(AxesKeyParser): + """ + Extracts the split axis parameters from triton kernel code. The parsing is based on the + `tl.program_id` statement. This class identifies potential split axes by analyzing the usage + of the `tl.program_id` variable in the program and its multiplication operations with other + variables(currently supporting scenarios where multiplication is either direct or indirect via + intermediate variables). It then filters these candidates based on a list of candidate parameters + (parameters not provided by the user). After that, it confirms the split axis corresponding to + the current parameter using mask comparison and the `keys` passed in `autotune`. + + Note: + 1. Split axis parameters must be multiplied with `tl.program_id`. + 2. Without mask comparision, it is impossible to confirm the exact split axis, which would lead + to parameter parsing failure. (eg. mask = offsets < n_elements) + 3. The identified split axes are limited to the list of candidated parameters, ensuring that + only those parameters that can be dynamically adjusted through the autotune process are considered. + """ + + def __init__( + self, func_ast: ast.AST, keys: Dict[str, str], candidates_params: List[str] + ): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the split axis corresponding to + the split axis parameters. + :type keys: Dict[str, str] + :param candidates_params: a list of parameters names that were not provided by the user when calling + triton kernel function. The parser will only consider these parameters as potential split axis + parameters. + :type candidates_params: List[str] + """ + super().__init__(func_ast, keys) + self.split_axes = dict() + self.program_id_vars = list() + self.program_id_var_dims = dict() + self.num_programs_var_dims = dict() + self.grid_stride_tiling_only = dict() + # axis_name -> program_id axis dim + self.split_axis_pid_dims = dict() + # axis_name -> program_id axis dim (includes axes inferred without split params) + self.axis_pid_dims = dict() + self.candidates_params = candidates_params + + def parse(self) -> Dict[str, str]: + super().parse() + return self.split_axes + + def visit_Assign(self, node): + pid_dim = self._get_program_id_dim(node.value) + if pid_dim is not None: + if ( + len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and node.targets[0].id not in self.program_id_vars + ): + self.program_id_vars.append(node.targets[0].id) + self.program_id_var_dims[node.targets[0].id] = pid_dim + num_programs_dim = self._get_num_programs_dim(node.value) + if num_programs_dim is not None: + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + self.num_programs_var_dims[node.targets[0].id] = num_programs_dim + self.generic_visit(node) + + def visit_BinOp(self, node): + if isinstance(node.op, ast.Mult): + split_axes_val = None + split_axis_pid_dim = None + if isinstance(node.left, ast.Name) and node.left.id in self.program_id_vars: + if isinstance(node.right, ast.Name): + split_axes_val = node.right.id + split_axis_pid_dim = self.program_id_var_dims.get(node.left.id) + elif isinstance(node.left, ast.Call) and isinstance( + node.left.func, ast.Attribute + ): + if ( + isinstance(node.left.func.value, ast.Name) + and node.left.func.value.id == "tl" + and node.left.func.attr == "program_id" + ): + if isinstance(node.right, ast.Name): + split_axes_val = node.right.id + split_axis_pid_dim = self._get_program_id_dim(node.left) + + if ( + isinstance(node.right, ast.Name) + and node.right.id in self.program_id_vars + ): + if isinstance(node.left, ast.Name): + split_axes_val = node.left.id + split_axis_pid_dim = self.program_id_var_dims.get(node.right.id) + elif isinstance(node.right, ast.Call) and isinstance( + node.right.func, ast.Attribute + ): + if ( + isinstance(node.right.func.value, ast.Name) + and node.right.func.value.id == "tl" + and node.right.func.attr == "program_id" + ): + if isinstance(node.left, ast.Name): + split_axes_val = node.left.id + split_axis_pid_dim = self._get_program_id_dim(node.right) + + if ( + split_axes_val in self.candidates_params + and split_axes_val not in self.split_axes.values() + ): + split_axes_key = self.get_axis(split_axes_val) + if split_axes_key and not self._is_tiling_only_split( + split_axes_key, split_axes_val + ): + self.split_axes[split_axes_key] = split_axes_val + if split_axis_pid_dim is not None: + self._record_axis_pid_dim(split_axes_key, split_axis_pid_dim) + self.generic_visit(node) + + def visit_For(self, node): + if not isinstance(node.iter, ast.Call): + self.generic_visit(node) + return + + iter_fn = node.iter.func + is_range = isinstance(iter_fn, ast.Name) and iter_fn.id == "range" + is_tl_range = ( + isinstance(iter_fn, ast.Attribute) + and isinstance(iter_fn.value, ast.Name) + and iter_fn.value.id == "tl" + and iter_fn.attr == "range" + ) + if not (is_range or is_tl_range): + self.generic_visit(node) + return + + if len(node.iter.args) == 0: + self.generic_visit(node) + return + + start = node.iter.args[0] if len(node.iter.args) >= 2 else None + stop = node.iter.args[1] if len(node.iter.args) >= 2 else node.iter.args[0] + pid_dim = self._extract_pid_dim_from_expr(start) + axis = self._axis_from_expr(stop) + if axis is not None and pid_dim is not None: + self._record_axis_pid_dim(axis, pid_dim) + if len(node.iter.args) >= 3: + step = node.iter.args[2] + loop_tiling_only_param = self._extract_grid_stride_split_param( + start, step, pid_dim + ) + if loop_tiling_only_param is not None: + self._mark_tiling_only_param(axis, loop_tiling_only_param) + + self.generic_visit(node) + + def _get_program_id_dim(self, node): + if not ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "tl" + and node.func.attr == "program_id" + ): + return None + + axis_dim = 0 + if len(node.args) > 0: + if isinstance(node.args[0], ast.Constant) and isinstance( + node.args[0].value, int + ): + axis_dim = node.args[0].value + else: + return None + + for kw in node.keywords: + if kw.arg == "axis": + if isinstance(kw.value, ast.Constant) and isinstance( + kw.value.value, int + ): + axis_dim = kw.value.value + else: + return None + break + return axis_dim + + def _get_num_programs_dim(self, node): + if not ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "tl" + and node.func.attr == "num_programs" + ): + return None + + axis_dim = 0 + if len(node.args) > 0: + if isinstance(node.args[0], ast.Constant) and isinstance( + node.args[0].value, int + ): + axis_dim = node.args[0].value + else: + return None + + for kw in node.keywords: + if kw.arg == "axis": + if isinstance(kw.value, ast.Constant) and isinstance( + kw.value.value, int + ): + axis_dim = kw.value.value + else: + return None + break + return axis_dim + + def _extract_pid_dim_from_expr(self, node): + if node is None: + return None + for child in ast.walk(node): + if isinstance(child, ast.Name) and child.id in self.program_id_var_dims: + return self.program_id_var_dims[child.id] + pid_dim = self._get_program_id_dim(child) + if pid_dim is not None: + return pid_dim + return None + + def _contains_pid_dim(self, node, pid_dim): + if node is None: + return False + for child in ast.walk(node): + if isinstance(child, ast.Name): + if self.program_id_var_dims.get(child.id, None) == pid_dim: + return True + if self._get_program_id_dim(child) == pid_dim: + return True + return False + + def _contains_num_programs_dim(self, node, pid_dim): + if node is None: + return False + for child in ast.walk(node): + if isinstance(child, ast.Name): + if self.num_programs_var_dims.get(child.id, None) == pid_dim: + return True + if self._get_num_programs_dim(child) == pid_dim: + return True + return False + + def _is_candidate_name(self, node, candidate_name): + return ( + isinstance(node, ast.Name) + and node.id == candidate_name + and candidate_name in self.candidates_params + ) + + def _extract_pid_multiplied_candidate(self, node, pid_dim): + if node is None: + return None + candidates = set() + for child in ast.walk(node): + if not isinstance(child, ast.BinOp) or not isinstance(child.op, ast.Mult): + continue + left = child.left + right = child.right + if ( + isinstance(left, ast.Name) + and left.id in self.candidates_params + and self._contains_pid_dim(right, pid_dim) + ): + candidates.add(left.id) + if ( + isinstance(right, ast.Name) + and right.id in self.candidates_params + and self._contains_pid_dim(left, pid_dim) + ): + candidates.add(right.id) + if len(candidates) == 1: + return next(iter(candidates)) + return None + + def _contains_num_programs_multiplied_candidate( + self, node, candidate_name, pid_dim + ): + if node is None: + return False + for child in ast.walk(node): + if not isinstance(child, ast.BinOp) or not isinstance(child.op, ast.Mult): + continue + if self._is_candidate_name(child.left, candidate_name): + if self._contains_num_programs_dim(child.right, pid_dim): + return True + if self._is_candidate_name(child.right, candidate_name): + if self._contains_num_programs_dim(child.left, pid_dim): + return True + return False + + def _extract_grid_stride_split_param(self, start, step, pid_dim): + if start is None or step is None: + return None + candidate_name = self._extract_pid_multiplied_candidate(start, pid_dim) + if candidate_name is None: + return None + if self._contains_num_programs_multiplied_candidate( + step, candidate_name, pid_dim + ): + return candidate_name + return None + + def _mark_tiling_only_param(self, axis, candidate_name): + self.grid_stride_tiling_only.setdefault(axis, set()).add(candidate_name) + if self.split_axes.get(axis, None) == candidate_name: + del self.split_axes[axis] + self.split_axis_pid_dims.pop(axis, None) + + def _is_tiling_only_split(self, axis, candidate_name): + return candidate_name in self.grid_stride_tiling_only.get(axis, set()) + + def _axis_from_expr(self, node): + if node is None: + return None + for k, v in self.keys.items(): + if self.contains_target_var(node, v): + return k + return None + + def _record_axis_pid_dim(self, axis, pid_dim): + self.axis_pid_dims[axis] = pid_dim + if axis in self.split_axes: + self.split_axis_pid_dims[axis] = pid_dim + + +class TilingAxesParser(AxesKeyParser): + """ + Extracts the tiling axis parameters from triton kernel code. The parsing is based on the + `tl.arange`, `tl.range` and `range()` statement. This class identifies potential tiling axes by analyzing + the usage of the `range` and `tl.range` within `for` loop in the program. Common parameters + between `range()` or `tl.range` and `tl.arange` are extracted. It then filters these candidates based on a + list of candidate parameters (parameters not provided by the user). After that, it confirms the + tiling axis corresponding to the current parameter using mask comparison and the `keys` passed + in `autotune`. + + Note: + 1. Tiling axis parameters must be calculated within the `tl.arange` function and the `for` loop + using `tl.range`. + 2. Without mask comparision, it is impossible to confirm the exact tiling axis, which would lead + to parameter parsing failure. (eg. mask = offsets < n_elements). + 3. The identified tiling axes are limited to the list of candidated parameters, ensuring that + only those parameters that can be dynamically adjusted through the autotune process are considered. + """ + + def __init__( + self, func_ast: ast.AST, keys: Dict[str, str], candidates_params: List[str] + ): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the tiling axis corresponding to + the tiling axis parameters. + :type keys: Dict[str, str] + :param candidates_params: a list of parameters names that were not provided by the user when calling + triton kernel function. The parser will only consider these parameters as potential tiling axis + parameters. + :type candidates_params: List[str] + """ + super().__init__(func_ast, keys) + self.tiling_axes = dict() + self.candidates_params = candidates_params + self.candidates_params_for_loop = list() + + def parse(self) -> Dict[str, str]: + super().parse() + return self.tiling_axes + + def visit_For(self, node): + if isinstance(node.iter, ast.Call) and len(node.iter.args) == 3: + step_expr = node.iter.args[2] + for_loop_param = self._extract_unique_candidate(step_expr) + if ( + for_loop_param is not None + and for_loop_param not in self.candidates_params_for_loop + ): + self.candidates_params_for_loop.append(for_loop_param) + self.generic_visit(node) + + def visit_Assign(self, node): + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + # handle FloorDiv + if isinstance(node.value, ast.BinOp) and isinstance( + node.value.op, ast.FloorDiv + ): + denominator = node.value.right + denominator_param = self._extract_unique_candidate(denominator) + if ( + denominator_param is not None + and denominator_param not in self.candidates_params_for_loop + ): + self.candidates_params_for_loop.append(denominator_param) + self.visit(self.func_ast) + + tiling_axes_val = self.get_tiling_axes_val(node.value) + if ( + tiling_axes_val is not None + and tiling_axes_val in self.candidates_params_for_loop + ): + tiling_axes_key = self.get_axis(tiling_axes_val) + if tiling_axes_key: + self.tiling_axes[tiling_axes_key] = tiling_axes_val + self.generic_visit(node) + + def get_tiling_axes_val(self, node): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + if ( + node.func.attr == "arange" + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "tl" + ): + if isinstance(node.args, list) and len(node.args) == 2: + for param in self.candidates_params_for_loop: + if self.contains_target_var(node.args[1], param): + return param + + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + val = self.get_tiling_axes_val(item) + if val: + return val + elif isinstance(value, ast.AST): + val = self.get_tiling_axes_val(value) + if val: + return val + return None + + def _extract_unique_candidate(self, expr): + """ + Extract a unique tiling candidate from an expression. + Return None when no candidate or ambiguous (more than one candidate) appears. + """ + if expr is None: + return None + candidates = [ + param + for param in self.candidates_params + if self.contains_target_var(expr, param) + ] + if len(candidates) == 1: + return candidates[0] + return None + + +class ReductionAxesParser(AxesKeyParser): + """ + Extracts the reduction axis from triton kernel code. The parsing is based on the + reduction function (eg. tl.max, tl.min, tl.sum, ...). This class identifies the + dimensions of reduction operations by analyzing the reduction function calls in + the program. After that, It confirms the reduction axis corresponding to the current + parameter using mask comparison and the keys passed in autotune. + + Note: + 1. The call to the reduction function must start with 'tl', meaning the function must + be a function from triton.language + 2. It's preferable to specify the reduction axis dimension in the reduction function + using keyword arguments(eg. axis=xxx). Otherwise, specifying it via positional + arguments may lead to errors. + 3. Mask comparison must be performed on the potential reduction axis length parameters, + and the comparison parameters or target parameters of the comparison expression + must be sliced. Otherwise, the correspondence between dimensions and axes cannot + be confirmed, which will lead to failure in parsing the reduction axis. + 4. The identified reduction axes are limited to the candidate list provided in the keys. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the reduction axis. + :type keys: Dict[str, str] + """ + super().__init__(func_ast, keys) + self.reduction_axes = list() + self.reduction_func = ( + "sum", + "xor_sum", + "max", + "min", + "argmax", + "argmin", + ) # tl.xxx + self.ndim = 1 + + def parse(self) -> List[str]: + super().parse() + return self.reduction_axes + + def visit_Assign(self, node): + self._scan_subscripts(node.value) + self.generic_visit(node) + + def _scan_subscripts(self, node): + if isinstance(node, ast.Subscript): + ndim = self._get_subscripts_ndim(node) + if ndim > self.ndim: + self.ndim = ndim + + for child in ast.iter_child_nodes(node): + self._scan_subscripts(child) + + def _get_subscripts_ndim(self, subscript_node): + slice_node = subscript_node.slice + + if isinstance(slice_node, ast.Tuple): + # e.g. [:, None] -> Tuple(elts=[Slice(), Constant(None)]) + return len(slice_node.elts) + elif isinstance( + slice_node, (ast.Slice, ast.Constant, ast.Name, ast.UnaryOp, ast.BinOp) + ): + # e.g. [0], [:], [i], [-1], [i+1] + return 1 + else: + # Fallback: treat as 1D + return 1 + + def visit_Call(self, node): + if not isinstance(node.func, ast.Attribute): + return + func = node.func + if not isinstance(func.value, ast.Name) or func.value.id != "tl": + self.generic_visit(node) + return + if func.attr not in self.reduction_func: + return + + axis_dim = None + args = node.args + if len(args) == 1: + # Axis passed as keyword argument + for keyword in node.keywords: + if keyword.arg == "axis": + axis_dim = self.get_axis_dim(keyword.value) + break + + elif len(args) == 2: + # Axis passed as positional argument. Check the second param + axis_dim = self.get_axis_dim(args[1]) + + else: + raise ValueError("Reduction funtions args error") + + if axis_dim is not None: + reduction_axis = self.get_axis(axis_dim) + if reduction_axis and reduction_axis not in self.reduction_axes: + self.reduction_axes.append(reduction_axis) + + def get_axis_dim(self, node): + if isinstance(node, ast.Constant): + axis_dim = node.value + elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub): + operand = node.operand + if isinstance(operand, ast.Constant): + axis_dim = self.ndim - operand.value + else: + raise ValueError(f"Reduction function axis error, got: {ast.dump(node)}") + + if not isinstance(axis_dim, int): + raise ValueError( + "Reduction function axis must be an integer, " + f"got {type(node.value).__name__}: {node.value}" + ) + return axis_dim + + def get_axis(self, axis_dim: int): + """ + Override the parent class method to accept an integer axis dimension + instead of a string. + + :param axis_dim: + :type axis_dim: int + """ + if axis_dim in self.checked_vars: + return None + self.checked_vars.append(axis_dim) + for node in ast.walk(self.func_ast): + if not isinstance(node, ast.Assign): + continue + reduction_axis = self.handle_assign_node(axis_dim, node) + if reduction_axis: + return reduction_axis + return None + + def handle_assign_node(self, axis_dim: int, node): + if not isinstance(node.value, ast.Compare): + return None + + # only support less than + if len(node.value.ops) != 1 or not isinstance(node.value.ops[0], ast.Lt): + return None + + target_axis_len = None + for axis_len in self.keys.values(): + if self.contains_target_var(node.value, axis_len): + target_axis_len = axis_len + break + if not target_axis_len: + return None + + # handel compare left var + if isinstance(node.value.left, ast.Name): + if self.check_compare_left(node.value.left.id, axis_dim): + reduction_axis = next( + (k for k, v in self.keys.items() if target_axis_len == v), None + ) + if reduction_axis and reduction_axis not in self.reduction_axes: + return reduction_axis + # handel compare target var + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + if self.check_compare_target(node.targets[0].id, axis_dim): + reduction_axis = next( + (k for k, v in self.keys.items() if target_axis_len == v), None + ) + if reduction_axis and reduction_axis not in self.reduction_axes: + return reduction_axis + return None + + def check_compare_left(self, var, axis_dim): + for node in ast.walk(self.func_ast): + if not isinstance(node, ast.Assign): + continue + if ( + len(node.targets) != 1 + or not isinstance(node.targets[0], ast.Name) + or node.targets[0].id != var + ): + continue + if self.is_current_dim_slice(node.value, axis_dim): + return True + return False + + def check_compare_target(self, var, axis_dim, node=None): + if not node: + node = self.func_ast + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + if self.check_compare_target(var, axis_dim, item): + return True + elif isinstance(value, ast.AST): + if isinstance(value, ast.Subscript): + if not isinstance(value.value, ast.Name) or value.value.id != var: + continue + if self.is_current_dim_slice(value, axis_dim): + return True + else: + if self.check_compare_target(var, axis_dim, value): + return True + return False + + def is_current_dim_slice(self, node, dim): + for node in ast.walk(node): + if not isinstance(node, ast.Subscript) or not isinstance( + node.slice, ast.Tuple + ): + continue + elts = node.slice.elts + if len(elts) != 0 and isinstance(elts[dim], ast.Slice): + return True + return False + + +class LowDimsAxesParser(AxesKeyParser): + """ + Extracts the low dimensions axis from triton kernel code. The parsing is based on the + `tl.arange` statement. This class identifies low dimensions axis by analyzing the usage + of the `tl.arange` in the program and extracts the variables computed by `tl.arange` and + their associated operations. Then it checks if these variables are involved in slicing + operations to determine dimension expansion and filters out variables that are expanded + in non-lowest dimensions. After that, it compares the extracted variables with the provided + `keys` to map them to specific low-dimensional axis. + + Note: + 1. low dimensions axis must be calculated within the `tl.arange` function and involved in + slicing operations to be identified. + 2. Without mask comparision, it is impossible to confirm the exact low dimensions axis, which + would lead to parameter parsing failure. (eg. mask = offsets < n_elements). + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to confirm the low-dimensional axis. + :type keys: Dict[str, str] + """ + super().__init__(func_ast, keys) + self.low_dims_axis = list() + self.keys = keys + self.checked_slice_vars = list() + + def parse(self) -> List[str]: + super().parse() + return self.low_dims_axis + + def visit_Assign(self, node): + if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): + tl_arange_node = self.get_tl_arange_node(node) + low_dims_axis = None + if isinstance(tl_arange_node, ast.Call): + partin_other_slice = [False] + if self.is_partin_low_dim_slice(node.targets[0].id, partin_other_slice): + low_dims_axis = self.get_axis(node.targets[0].id) + elif not partin_other_slice[0]: + low_dims_axis = self.get_axis(node.targets[0].id) + elif isinstance(tl_arange_node, ast.Subscript) and self.is_low_dim_slice( + tl_arange_node, [False] + ): + low_dims_axis = self.get_axis(node.targets[0].id) + + if low_dims_axis and low_dims_axis not in self.low_dims_axis: + self.low_dims_axis.append(low_dims_axis) + self.generic_visit(node) + + def get_tl_arange_node(self, node): + for _, value in ast.iter_fields(node): + if isinstance(value, list): + for item in value: + if self.is_tl_arange_call(item): + return item + node = self.get_tl_arange_node(item) + if node: + return node + elif isinstance(value, ast.AST): + if self.is_tl_arange_call(value): + return value + node = self.get_tl_arange_node(value) + if node: + return node + return None + + def is_tl_arange_call(self, node): + """ + Checks if the given AST node is a call to `tl.arange` or a subscript of `tl.arange`. + It supports direct calls to `tl.arange` and subscripts of `tl.arange`, such as + `tl.arange()[None, :]` + """ + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + if ( + node.func.attr == "arange" + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "tl" + ): + return True + elif isinstance(node, ast.Subscript): + return self.is_tl_arange_call(node.value) + return False + + def is_low_dim_slice(self, node: ast.Subscript, partin_other_slice): + if not isinstance(node.slice, ast.Tuple) or not isinstance( + node.slice.elts, list + ): + return False + elts = node.slice.elts + if len(elts) != 0 and not isinstance(elts[-1], ast.Slice): + partin_other_slice[0] = True + return False + return True + + def is_partin_low_dim_slice(self, var, partin_other_slice, node=None): + if not node: + node = self.func_ast + for child_node in ast.walk(node): + if isinstance(child_node, ast.Subscript) and isinstance( + child_node.value, ast.Name + ): + if var == child_node.value.id and self.is_low_dim_slice( + child_node, partin_other_slice + ): + return True + elif isinstance(child_node, ast.Assign): + if ( + len(child_node.targets) == 1 + and isinstance(child_node.targets[0], ast.Name) + and var != child_node.targets[0].id + ): # Prevent cyclic assignment. + if not self.contains_target_var(child_node.value, var): + continue + target_var = child_node.targets[0].id + if target_var in self.checked_slice_vars: + continue + + if self.is_partin_low_dim_slice( + var, partin_other_slice, child_node.value + ): + return True + if self.is_partin_low_dim_slice(target_var, partin_other_slice): + return True + + self.checked_slice_vars.append(var) + return False + + +class PtrNumsParser(AutoParser): + """ + Counts the number of pointer parameters from triton kernel code. The parsing of pointer-type + parameters is determined based on whether these parameters participate in memory access + statements such as `tl.load` and `tl.store`. + First, all input parameters in the kernel function are parsed, and then recursively, all variables + involved in the computation of each input parameter are identified. + If an input parameter directly participates in the computation of the first argument of `tl.load` + or `tl.store`, or if an intermediate variable computed from this input parameter indirectly + participates in the computation of the first argument of `tl.load` or `tl.store`, then this + parameter is considered a pointer-type parameter. + + Note: + 1. Variables modified with `tl.constexpr` are not pointer-type variables and will not be + further parsed. + 2. Only memory access statementes where the input parameter is directly involved or indirectly + involved through one level of computation are counted. Intermediate variables computed from + the input parameter through two or more levels of computation are not counted. + """ + + def __init__(self, func_ast: ast.AST, keys: Dict[str, str], miss_params: List[str]): + """ + :param func_ast: Abstract syntax tree of the triton kernel function + :type func_ast: ast.AST + :param keys: a dict of axis name: argument name, used to exclude potential ptr params. + :type keys: Dict[str, str] + :param miss_params: a list of parameters names that were not provided by the user when calling triton + kernel function. + :type miss_params: List[str] + """ + super().__init__(func_ast) + self.checked_vars = list() + self.ptr_nums = 0 + self.ptr_params = list() + self.keys = keys + self.miss_params = miss_params + self.constexpr_params = list() + + def parse(self): + super().parse() + return self.ptr_nums, self.ptr_params + + def visit_FunctionDef(self, node): + if isinstance(node.args, ast.arguments): + for arg in node.args.args: + if not isinstance(arg, ast.arg): + continue + + if isinstance(arg.annotation, ast.Attribute): + # var modified by tl.constexpr are not pointer type var, passed + is_tl = ( + isinstance(arg.annotation.value, ast.Name) + and arg.annotation.value.id == "tl" + ) + if is_tl and arg.annotation.attr == "constexpr": + if arg.arg not in self.constexpr_params: + self.constexpr_params.append(arg.arg) + continue + + if self.is_in_addr_calc(arg.arg) and arg.arg not in self.keys.values(): + self.ptr_params.append(arg.arg) + self.ptr_nums += 1 + + for miss_param in self.miss_params: + if miss_param not in self.constexpr_params: + print( + f"[WARNING] The parameter '{miss_param}' needs to be declared as tl.constexpr!" + ) + self.generic_visit(node) + + def is_in_addr_calc(self, var): + for node in ast.walk(self.func_ast): + if isinstance(node, ast.Call): + if isinstance(node.func, ast.Attribute) and isinstance( + node.func.value, ast.Name + ): + if node.func.value.id == "tl" and ( + node.func.attr == "load" or node.func.attr == "store" + ): + if [ + arg + for arg in node.args + if self.contains_target_var(arg, var) + ]: + return True + + elif isinstance(node, ast.Assign): + if ( + len(node.targets) == 1 + and isinstance(node.targets[0], ast.Name) + and var != node.targets[0].id + ): # Prevent cyclic assignment. + target_var = node.targets[0].id + if target_var in self.checked_vars: + continue + if isinstance(node.value, ast.BinOp) and isinstance( + node.value.op, ast.Add + ): + if ( + isinstance(node.value.left, ast.Name) + and node.value.left.id == var + ): + if self.is_in_addr_calc(node.targets[0].id): + return True + elif ( + isinstance(node.value.right, ast.Name) + and node.value.right.id == var + ): + if self.is_in_addr_calc(node.targets[0].id): + return True + self.checked_vars.append(var) + return False diff --git a/backend/ascend_autotune_runtime/autotuner.py b/backend/ascend_autotune_runtime/autotuner.py new file mode 100644 index 00000000..00a4f167 --- /dev/null +++ b/backend/ascend_autotune_runtime/autotuner.py @@ -0,0 +1,1569 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# 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. + +from __future__ import annotations + +import builtins +import copy +import functools +import ast +import inspect +import os +import time +from concurrent.futures import ThreadPoolExecutor +import itertools +from typing import Dict, List + +from torch import Tensor + +import triton +from triton.runtime.autotuner import Autotuner, Config +from triton.backends.dicp_triton.utils import is_compile_on_910_95 + +from .autoparser import ( + DotCallParser, + LowDimsAxesParser, + PtrNumsParser, + ReductionAxesParser, + SplitAxesParser, + TilingAxesParser, +) +from .compile_options import ( + expand_compile_option_configs, + format_compile_option_result, + get_compile_option_param_names, + parse_compile_options_hint, + summarize_compile_option_configs, +) +from .benchmark import select_benchmark_strategy +from .utils import get_byte_per_numel, is_valid_axis_name, valid_axis_names + + +def _make_config_compat(**kwargs): + supported_config_args = inspect.signature(Config).parameters + return Config( + **{key: value for key, value in kwargs.items() if key in supported_config_args} + ) + + +def _empty_npu_cache_after_failure(): + try: + import torch + + torch.npu.empty_cache() + except Exception: + pass + + +def _unwrap_parse_target(fn): + seen = set() + cur = fn + while cur is not None and id(cur) not in seen: + if callable(getattr(cur, "parse", None)): + return cur + + jit_function = getattr(cur, "jit_function", None) + if callable(getattr(jit_function, "parse", None)): + return jit_function + + seen.add(id(cur)) + cur = getattr(cur, "fn", None) + return fn + + +class AutoTilingTuner(Autotuner): + """ + Automatic generateing candidate tiling configs and evaluating their performance to get the best config. + """ + + def __init__( + self, + fn, + arg_names, + configs, + key, + reset_to_zero, + restore_value, + pre_hook=None, + post_hook=None, + prune_configs_by: Dict = None, + warmup=None, + rep=None, + use_cuda_graph=False, + do_bench=None, + cache_results=False, + auto_profile_dir=None, + hints=None, + ): + """ + :param key: a list of argument name, where the change of arguments in value will triger re-generating candidates configs and evaluating. + The parameters in the list will be assigned axis names in sequence, with the axis name being in + {'x','y','z','w','v','t','rx','ry','rz','rw','rv','rt}, where the prefix 'r' means a reduction axis. + Only the axis name in this param should add perfix 'r' if it's a reduction axis. + :type key: List[str] + """ + super().__init__( + fn, + arg_names, + configs, + key, + reset_to_zero, + restore_value, + pre_hook, + post_hook, + prune_configs_by, + warmup, + rep, + use_cuda_graph, + do_bench, + cache_results, + ) + self.user_defined_do_bench = do_bench is not None + if not hints: + self.hints = {} + else: + self.hints = dict(hints) + self.ast_fn = _unwrap_parse_target(self.fn) + split_params = self.hints.get("split_params", None) + tiling_params = self.hints.get("tiling_params", None) + low_dim_axes = self.hints.get("low_dim_axes", None) + reduction_axes = self.hints.get("reduction_axes", None) + self._init_axis_params( + key, + split_params, + tiling_params, + low_dim_axes, + reduction_axes, + ) + + self.auto_gen_config = not configs or self.hints.get("auto_gen_config", False) + self._infer_compile_options_hint_if_needed() + self.compile_options = parse_compile_options_hint( + self.hints.get("compile_options", None) + ) + self.gen_configs = [] # generated configs from TileGenerator + self.auto_profile_dir = auto_profile_dir + if not configs: + self.user_configs = [] + else: + self.user_configs = configs + self.is_simt_mode = False + self.simt_stack_limit = 8192 + self.user_specified_warps = None + self.user_specified_multibuffer = None + self.default_multibuffer = not is_compile_on_910_95 + self.print_autotuning = os.getenv("TRITON_PRINT_AUTOTUNING", None) == "1" + self.print_autotuning_timings = ( + os.getenv("TRITON_PRINT_AUTOTUNING_TIMINGS", None) == "1" + ) + # Compile kernels in parallel by default for triton.runtime.JITFunction, + # but not for others, e.g., LibEntry, since it's not compatible with AsyncCompileMode + self.compile_parallel = ( + isinstance(self.fn, triton.runtime.JITFunction) + and os.getenv("TRITON_AUTOTUNE_PARALLEL_COMPILE", "1") == "1" + ) + + def _infer_compile_options_hint_if_needed(self): + if "compile_options" in self.hints: + return + self.hints["compile_options"] = ( + "mixcv" if self._autoparse_has_dot() else "vector" + ) + + def _parse_ast(self): + parse = getattr(self.ast_fn, "parse", None) + if not callable(parse): + raise ValueError( + "Cannot parse Ascend autotune kernel AST. " + "Please pass explicit autotune hints for wrapped kernels." + ) + return parse() + + def _get_capture_scope(self): + get_capture_scope = getattr(self.ast_fn, "get_capture_scope", None) + if callable(get_capture_scope): + return get_capture_scope() + return {} + + def _autoparse_has_dot(self) -> bool: + try: + func_ast = self._parse_ast() + return DotCallParser( + func_ast, self._get_capture_scope(), {id(self.ast_fn)} + ).parse() + except Exception as e: + raise ValueError( + "Cannot infer Ascend compile_options from kernel AST. " + "Please pass hints={'compile_options': 'mixcv'} or " + "hints={'compile_options': 'vector'} explicitly." + ) from e + + def _expand_simt_num_warps_configs( + self, base_configs: List[Config] + ) -> List[Config]: + _default_cand_num_warps = [8, 16, 32, 64] + cand_num_warps = ( + _default_cand_num_warps + if self.user_specified_warps is None + else [self.user_specified_warps] + ) + + simt_configs = [] + for base_cfg in base_configs: + for num_warps in cand_num_warps: + new_cfg = copy.deepcopy(base_cfg) + new_cfg.num_warps = num_warps + simt_configs.append(new_cfg) + + if self.print_autotuning: + print( + f"Triton autotuning: Expanded to {len(simt_configs)} SIMT configs (with warps: {cand_num_warps})" + ) + return simt_configs + + def _expand_simd_multibuffer_configs( + self, base_configs: List[Config] + ) -> List[Config]: + if self.user_specified_multibuffer is not None: + if self.print_autotuning: + print( + "Triton autotuning: Skip SIMD multibuffer expansion because user " + f"specified multibuffer={self.user_specified_multibuffer}" + ) + return base_configs + + opposite_default_multibuffer = not self.default_multibuffer + simd_configs = [] + for base_cfg in base_configs: + simd_configs.append(base_cfg) + new_cfg = copy.deepcopy(base_cfg) + new_cfg.kwargs["multibuffer"] = opposite_default_multibuffer + simd_configs.append(new_cfg) + + if self.print_autotuning: + print( + "Triton autotuning: Expanded to " + f"{len(simd_configs)} SIMD configs (toggle multibuffer={opposite_default_multibuffer})" + ) + return simd_configs + + def _init_axis_params( + self, key, split_params, tiling_params, low_dim_axes, reduction_axes + ): + if isinstance(key, list): + if split_params or tiling_params or low_dim_axes or reduction_axes: + raise ValueError( + "If any axis-related parameters (split_params, tiling_params, low_dim_axes, reduction_axes)" + " are provided, 'key' must be a dict, not a list." + ) + if len(key) > len(valid_axis_names): + raise ValueError( + "Number of parameters exceeds the number of available axes." + ) + self.keys = {axis: param for axis, param in zip(valid_axis_names, key)} + elif isinstance(key, dict): + if not set(key.keys()).issubset(set(valid_axis_names)): + raise ValueError( + "All keys in 'key' must be valid axis names. Got unexpected keys." + ) + self.keys = key + if any([split_params, tiling_params, low_dim_axes, reduction_axes]) is None: + raise ValueError( + "If 'key' is a dict, all axis-related parameters (split_params, tiling_params, low_dim_axes," + " reduction_axes) must be provided." + ) + if not isinstance(split_params, dict): + raise ValueError( + "split_params must be a dict, got: {}".format(type(split_params)) + ) + if not isinstance(tiling_params, dict): + raise ValueError( + "tiling_params must be a dict, got: {}".format(type(tiling_params)) + ) + if not isinstance(low_dim_axes, list): + raise ValueError( + "low_dim_axes must be a list, got: {}".format(type(low_dim_axes)) + ) + if not isinstance(reduction_axes, list): + raise ValueError( + "reduction_axes must be a list, got: {}".format( + type(reduction_axes) + ) + ) + + used_axes = set(split_params.keys()).union( + tiling_params.keys(), + low_dim_axes, + reduction_axes, + ) + if not used_axes.issubset(self.keys.keys()): + raise ValueError( + "The following axes are used but not present in the 'key': {}".format( + used_axes - set(self.keys.keys()) + ) + ) + + self.split_params = split_params + self.all_split_params = {} + self.fixed_split_params = {} + self.tiling_params = tiling_params + self.low_dim_axes = low_dim_axes + self.reduction_axes = reduction_axes + self.fixed_grid_dims = set() + self.fixed_grid_dim_values = {} + self.split_axis_pid_dims = {} + self.axis_pid_dims = {} + self.dual_reduction = False + self.persistent_reduction = False + self.num_buffers = -1 + + def _autoparse_axis_params(self, all_args): + miss_params = [arg for arg in self.arg_names if arg not in all_args.keys()] + # parse pointer params nums + if self.num_buffers == -1: + self.num_buffers = self._autoparse_ptr_nums(all_args) + + # parse autotiling axes + # reduction axis must be parsed before other axes. it will alter the key + if not self.reduction_axes: + self.reduction_axes = self._autoparse_reduction_axes() + if len(self.reduction_axes) >= 2: + self.dual_reduction = True + + if not self.low_dim_axes: + self.low_dim_axes = self._autoparse_low_dim_axes() + + if len(self.reduction_axes) == 1: + reduction_axis = self.reduction_axes[0] + reduction_param = self.keys.get(reduction_axis, None) + reduction_numel = all_args.get(reduction_param, float("inf")) + persistent_threshold = self._get_persistent_reduction_threshold( + reduction_axis + ) + if reduction_numel <= persistent_threshold: + self.persistent_reduction = True + + if not self.split_params: + all_split_params = self._autoparse_split_params( + self._get_constexpr_candidates() + ) + self.all_split_params = dict(all_split_params) + self.fixed_split_params = {} + self.fixed_grid_dim_values = self._get_fixed_grid_dim_values( + all_args.get("grid", None), + all_args, + ) + self.fixed_grid_dims = set(self.fixed_grid_dim_values.keys()) + + fixed_grid_axes = { + axis + for axis, pid_dim in self.axis_pid_dims.items() + if pid_dim in self.fixed_grid_dims + } + + # Only missing constexpr params are tunable, and fixed-grid axes + # should not be tuned on split. + self.split_params = { + axis: param + for axis, param in all_split_params.items() + if param in miss_params and axis not in fixed_grid_axes + } + + # Fixed split is inferred only from fixed grid dims. + for axis, pid_dim in self.axis_pid_dims.items(): + if pid_dim not in self.fixed_grid_dims: + continue + core_num = self.fixed_grid_dim_values.get(pid_dim, 0) + axis_len_name = self.keys.get(axis, None) + axis_len = all_args.get(axis_len_name, None) + if not isinstance(core_num, int) or core_num <= 0: + continue + if not isinstance(axis_len, int) or axis_len <= 0: + continue + + self.fixed_split_params[axis] = (axis_len + core_num - 1) // core_num + elif not self.axis_pid_dims: + # When split axes are provided by hints, parse axis->program_id mapping + # independently for fixed-grid semantics and diagnostics. + self._autoparse_axis_pid_dims() + miss_params = [ + arg for arg in miss_params if arg not in self.split_params.values() + ] + if not self.tiling_params: + self.tiling_params = self._autoparse_tiling_params(miss_params) + miss_params = [ + arg for arg in miss_params if arg not in self.tiling_params.values() + ] + if miss_params: + raise ValueError( + f"Missing required arguments: {miss_params}. " + f"These arguments must be explicitly provided and cannot be automatically tuned. " + f"Please ensure that these arguments are passed when calling the function." + ) + + def _gen_tile_configs( + self, kv_dict: Dict[str, int], dtype: torch.dtype + ) -> List[Config]: + from .tile_generator import KernelMeta, TileGenerator + + axis_sizes = {} + for k, v in kv_dict.items(): + if not is_valid_axis_name(k): + continue + if not isinstance(v, int): + raise ValueError( + f"Not supported dim type: {type(v)}, `int` is the only supported type" + ) + axis_sizes[k] = v + + kernel_meta = KernelMeta( + axis_sizes, + self.split_params, + self.fixed_split_params, + self.tiling_params, + self.low_dim_axes, + dtype, + self.persistent_reduction, + self.dual_reduction, + self.num_buffers, + self.is_simt_mode, + ) + tile_gen = TileGenerator(kernel_meta=kernel_meta) + tile_gen.descend_split_tiling() + + self.gen_configs.clear() + self.gen_configs = tile_gen.configs + + if self.is_simt_mode: + self.gen_configs = self._expand_simt_num_warps_configs(self.gen_configs) + else: + self.gen_configs = self._expand_simd_multibuffer_configs(self.gen_configs) + + if len(self.gen_configs) == 0: + print( + "[WARNING] The generated candidate tiling configs are empty based on provided parameters!" + ) + + if self.print_autotuning: + print("Generated configs number: {}".format(len(self.gen_configs))) + + def generate_key_and_configs(self, *args, **kwargs): + self.nargs = dict(zip(self.arg_names, args)) + self.is_simt_mode = kwargs.get("force_simt_only", False) + if "num_warps" in kwargs and kwargs["num_warps"] is not None: + self.user_specified_warps = kwargs["num_warps"] + else: + self.user_specified_warps = None + if "multibuffer" in kwargs and kwargs["multibuffer"] is not None: + self.user_specified_multibuffer = kwargs["multibuffer"] + else: + self.user_specified_multibuffer = None + + # generate key + all_args = {**self.nargs, **kwargs} + _args = {k: v for (k, v) in all_args.items() if k in self.arg_names} + key = [_args[v] for _, v in self.keys.items() if v in _args] + + # Currently, we use the dtype with maximum byte length + dtype = None + for _, arg in _args.items(): + if hasattr(arg, "dtype"): + key.append(str(arg.dtype)) + dtype = ( + arg.dtype + if get_byte_per_numel(arg.dtype) >= get_byte_per_numel(dtype) + else dtype + ) + if dtype is None: + raise NotImplementedError("Not support for non-Tensor inputs") + + key = tuple(key) + if key not in self.cache: + if self.auto_gen_config: + self._autoparse_axis_params(all_args) + _kv_dict = {k: _args[v] for k, v in self.keys.items() if v in _args} + self._gen_tile_configs(_kv_dict, dtype) + fixed_compile_options = { + name: kwargs[name] + for name in get_compile_option_param_names(self.compile_options) + if name in kwargs and kwargs[name] is not None + } + self.fixed_compile_options = fixed_compile_options + gen_configs = expand_compile_option_configs( + self.gen_configs, + self.compile_options, + generated_tiling=True, + fixed_options=fixed_compile_options, + ) + user_configs = expand_compile_option_configs( + self.user_configs, + self.compile_options, + generated_tiling=False, + fixed_options=fixed_compile_options, + ) + if self.print_autotuning and self.compile_options.enabled: + print( + "Triton autotuning compile_options: " + f"kernel_type={self.compile_options.kernel_type}, " + f"manual_params={sorted(self.compile_options.params.keys())}, " + f"fixed_runtime_params={sorted(fixed_compile_options.keys())}, " + f"max_configs={self.compile_options.max_configs}, " + f"generated_tiling_configs={len(self.gen_configs)}->{len(gen_configs)}, " + f"user_configs={len(self.user_configs)}->{len(user_configs)}" + ) + for idx, sample in enumerate( + summarize_compile_option_configs(gen_configs + user_configs), 1 + ): + print(f"Triton autotuning compile_options sample[{idx}]: {sample}") + if len(gen_configs) == 0 and len(user_configs) == 0: + self.configs = [ + _make_config_compat( + kwargs={}, + num_warps=4, + num_stages=2, + num_ctas=1, + num_buffers_warp_spec=0, + num_consumer_groups=0, + reg_dec_producer=0, + reg_inc_consumer=0, + ) + ] + self.configs = expand_compile_option_configs( + self.configs, + self.compile_options, + generated_tiling=True, + fixed_options=fixed_compile_options, + ) + if self.print_autotuning and self.compile_options.enabled: + print( + "Triton autotuning compile_options fallback: " + f"configs=1->{len(self.configs)}" + ) + for idx, sample in enumerate( + summarize_compile_option_configs(self.configs), 1 + ): + print( + f"Triton autotuning compile_options sample[{idx}]: {sample}" + ) + else: + self.configs = gen_configs + user_configs + return key + + def run(self, *args, **kwargs): + key = self.generate_key_and_configs(*args, **kwargs) + if self.is_simt_mode and kwargs.get("simt_stack_limit", None) is None: + kwargs["simt_stack_limit"] = self.simt_stack_limit + used_cached_result = True + if key not in self.cache: + # prune configs + pruned_configs = self.prune_configs(kwargs) + if len(pruned_configs) > 1: + used_cached_result = False + + def benchmark(): + bench_start = time.time() + timings = self._batch_bench(*args, configs=pruned_configs, **kwargs) + bench_end = time.time() + self.bench_time = bench_end - bench_start + self.cache[key] = builtins.min(timings, key=timings.get) + full_nargs = { + **self.nargs, + **kwargs, + **self.cache[key].all_kwargs(), + } + self.pre_hook(full_nargs, reset_only=True) + self.configs_timings = timings + if self.print_autotuning_timings: + self._print_config_timings(timings) + + if self.cache_results: + used_cached_result = self.check_disk_cache( + key, pruned_configs, benchmark + ) + else: + benchmark() + config = self.cache[key] + else: + self.cache[key] = pruned_configs[0] + config = self.cache[key] + else: + config = self.cache[key] + + self.best_config = config + if self.print_autotuning and not used_cached_result: + print( + f"Triton autotuning for function {self.base_fn.__name__} finished after " + f"{self.bench_time:.2f}s; best config selected: " + f"{format_compile_option_result(self.best_config, self.compile_options, getattr(self, 'fixed_compile_options', {}))};" + ) + + if not used_cached_result and self.auto_profile_dir is not None: + self._profile(*args, config=self.best_config, **kwargs) + if config.pre_hook is not None: + full_nargs = {**self.nargs, **kwargs, **config.all_kwargs()} + config.pre_hook(full_nargs) + final_kwargs = dict(config.all_kwargs(), **kwargs) + ret = self.fn.run( + *args, + **final_kwargs, + ) + self.nargs = None + return ret + + @staticmethod + def _timing_sort_key(cost): + if isinstance(cost, (list, tuple)) and cost: + return cost[0] + return cost + + def _print_config_timings(self, timings): + sorted_timings = sorted( + timings.items(), key=lambda item: self._timing_sort_key(item[1]) + ) + for idx, (config, cost) in enumerate(sorted_timings, 1): + print( + "Triton autotuning timing" + f"[{idx}]: cost={cost}; " + f"{format_compile_option_result(config, self.compile_options, getattr(self, 'fixed_compile_options', {}))};", + flush=True, + ) + + def _batch_bench(self, *args, configs, **kwargs): + from triton.compiler.errors import CompilationError, CompileTimeAssertionFailure + from triton.runtime.errors import OutOfResources + + kernels_call = { + config: self._make_kernel_call(*args, config=config, **kwargs) + for config in configs + } + run_fns = {} + exc = None + exc_stack = "" + + if self.compile_parallel: + import psutil + + max_workers = min(psutil.cpu_count(logical=False) // 2, len(kernels_call)) + future_kernels = [] + try: + with ( + ThreadPoolExecutor(max_workers=max_workers) as executor, + triton.AsyncCompileMode(executor), + ): + for config, fn in kernels_call.items(): + future_kernels.append((config, fn(warmup=True))) + + for config, fut in future_kernels: + try: + if hasattr(fut, "result"): + fut = fut.result() + run_fns[config] = functools.partial( + kernels_call[config], warmup=False + ) + except ( + CompileTimeAssertionFailure, + CompilationError, + OutOfResources, + Exception, + ) as e: + _empty_npu_cache_after_failure() + import traceback + + exc_stack = traceback.format_exc() + exc = e + except Exception as e: + # ignore exception from __exit__() of AsyncCompileMode + triton.runtime._async_compile.active_mode.set(None) + if exc is None: + _empty_npu_cache_after_failure() + import traceback + + exc_stack = traceback.format_exc() + exc = e + else: + for config, fn in kernels_call.items(): + try: + fn(warmup=False) + run_fns[config] = functools.partial(fn, warmup=False) + except ( + CompileTimeAssertionFailure, + CompilationError, + OutOfResources, + Exception, + ) as e: + _empty_npu_cache_after_failure() + import traceback + + exc_stack = traceback.format_exc() + exc = e + + if len(run_fns) == 0: + raise RuntimeError( + f"No valid triton configs. {type(exc).__name__}: {exc} \nStack trace: {exc_stack}" + ) + + strategy = select_benchmark_strategy( + self.do_bench, + self.user_defined_do_bench, + len(run_fns), + ) + try: + return strategy.bench(run_fns) + except Exception: + _empty_npu_cache_after_failure() + return self._batch_bench_fallback(strategy, run_fns) + + def _batch_bench_fallback(self, strategy, run_fns): + timings = {} + for config, fn in run_fns.items(): + try: + cost = strategy.bench({config: fn}) + if isinstance(cost, dict): + timings[config] = cost.get(config, float("inf")) + elif isinstance(cost, (list, tuple)) and len(cost) == 1: + timings[config] = cost[0] + elif isinstance(cost, (int, float)): + timings[config] = cost + else: + timings[config] = float("inf") + except Exception: + _empty_npu_cache_after_failure() + timings[config] = float("inf") + return timings + + def _make_kernel_call(self, *args, config, **meta): + # check for conflicts, i.e. meta-parameters both provided + # as kwargs and by the autotuner + conflicts = meta.keys() & config.kwargs.keys() + if conflicts: + raise ValueError( + f"Conflicting meta-parameters: {', '.join(conflicts)}." + " Make sure that you don't re-define auto-tuned symbols." + ) + # augment meta-parameters with tunable ones + current = dict(meta, **config.all_kwargs()) + full_nargs = {**self.nargs, **current} + + def kernel_call(warmup): + if config.pre_hook: + config.pre_hook(full_nargs) + self.pre_hook(full_nargs) + try: + current.update({"warmup": warmup}) + res = self.fn.run( + *args, + **current, + ) + if warmup: + return res + except Exception as e: + try: + self.post_hook(full_nargs, exception=e) + finally: + # Throw exception raised by `self.fn.run` + raise + + self.post_hook(full_nargs, exception=None) + + return kernel_call + + def warmup(self, *args, **kwargs): + _ = self.generate_key_and_configs(*args, **kwargs) + pruned_configs = self.prune_configs(kwargs) + ret = [] + if self.compile_parallel: + import psutil + + max_workers = min(psutil.cpu_count(logical=False) // 2, len(pruned_configs)) + with ( + ThreadPoolExecutor(max_workers=max_workers) as executor, + triton.AsyncCompileMode(executor), + ): + for config in pruned_configs: + ret.append(self.fn.warmup(*args, **kwargs, **config.all_kwargs())) + else: + for config in pruned_configs: + ret.append(self.fn.warmup(*args, **kwargs, **config.all_kwargs())) + self.nargs = None + return ret + + def _profile(self, *args, config, **meta): + from ..testing import do_bench_npu + + kernel_call = self._make_kernel_call(*args, config=config, **meta) + fn = functools.partial(kernel_call, warmup=False) + do_bench_npu(fn, prof_dir=self.auto_profile_dir, keep_res=True) + + def _autoparse_split_params(self, candidates_params: List[str]) -> Dict[str, str]: + """ + Extracts the split axis parameters from triton kernel code. + """ + func_ast = self._parse_ast() + parser = SplitAxesParser(func_ast, self.keys, candidates_params) + split_axes = parser.parse() + self.split_axis_pid_dims = dict(getattr(parser, "split_axis_pid_dims", {})) + self.axis_pid_dims = dict(getattr(parser, "axis_pid_dims", {})) + if self.print_autotuning: + print( + f"Ascend autotuning parse split axes: {split_axes}, " + f"split axis pid dims: {self.split_axis_pid_dims}, " + f"axis pid dims: {self.axis_pid_dims}" + ) + return split_axes + + def _autoparse_axis_pid_dims(self) -> Dict[str, int]: + """ + Extract axis -> program_id dim mapping without relying on split-parameter + classification, so fixed-grid semantics can always consume it. + """ + func_ast = self._parse_ast() + parser = SplitAxesParser( + func_ast, + self.keys, + self._get_constexpr_candidates(), + ) + _ = parser.parse() + self.axis_pid_dims = dict(getattr(parser, "axis_pid_dims", {})) + self.split_axis_pid_dims = dict(getattr(parser, "split_axis_pid_dims", {})) + if self.print_autotuning: + print( + "Ascend autotuning parse axis pid dims (independent): " + f"{self.axis_pid_dims}" + ) + return self.axis_pid_dims + + def _get_constexpr_candidates(self) -> List[str]: + """ + Returns all constexpr parameter names from the kernel function definition. + """ + func_ast = self._parse_ast() + constexpr_names = [] + for node in ast.walk(func_ast): + if not isinstance(node, ast.FunctionDef): + continue + if not isinstance(node.args, ast.arguments): + continue + for arg in node.args.args: + if not isinstance(arg, ast.arg): + continue + ann = arg.annotation + if ( + isinstance(ann, ast.Attribute) + and isinstance(ann.value, ast.Name) + and ann.value.id == "tl" + and ann.attr == "constexpr" + ): + constexpr_names.append(arg.arg) + break + return constexpr_names + + def _get_fixed_grid_dim_values( + self, grid, all_args: Dict[str, object] = None + ) -> Dict[int, int]: + """ + Returns fixed grid dim -> value. + - Static tuple/list grid: direct extraction + - Callable grid: infer fixed dims by perturbing missing constexpr params + """ + if grid is None: + return {} + if callable(grid): + return self._infer_fixed_dims_from_callable_grid(grid, all_args or {}) + return self._extract_fixed_grid_dims(grid) + + def _extract_fixed_grid_dims(self, grid) -> Dict[int, int]: + if isinstance(grid, int): + grid = (grid,) + if not isinstance(grid, (tuple, list)): + return {} + fixed_dims = {} + for idx, dim in enumerate(grid): + if isinstance(dim, int) and dim > 0: + fixed_dims[idx] = dim + return fixed_dims + + def _normalize_grid_tuple(self, grid_out): + if isinstance(grid_out, int): + return (grid_out,) + if isinstance(grid_out, (tuple, list)): + return tuple(grid_out) + return None + + def _infer_fixed_dims_from_callable_grid( + self, grid_fn, all_args: Dict[str, object] + ) -> Dict[int, int]: + constexpr_candidates = self._get_constexpr_candidates() + base_meta = dict(all_args or {}) + + # Fill missing constexpr with stable probe defaults so grid(meta) can execute. + for name in constexpr_candidates: + if name not in base_meta: + base_meta[name] = 128 + + try: + base_grid_raw = grid_fn(dict(base_meta)) + except Exception: + return {} + + base_grid = self._normalize_grid_tuple(base_grid_raw) + if base_grid is None: + return {} + + dynamic_dims = set() + # Missing constexpr are tunable candidates. + tunable_probe_names = [ + name for name in constexpr_candidates if name not in (all_args or {}) + ] + probe_values = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512] + + for name in tunable_probe_names: + baseline = base_meta.get(name, 128) + for probe in probe_values: + if probe == baseline: + continue + probe_meta = dict(base_meta) + probe_meta[name] = probe + try: + probe_grid_raw = grid_fn(probe_meta) + except Exception: + continue + probe_grid = self._normalize_grid_tuple(probe_grid_raw) + if probe_grid is None: + continue + if len(probe_grid) != len(base_grid): + dynamic_dims.update(range(min(len(probe_grid), len(base_grid)))) + continue + for idx, (base_dim, probe_dim) in enumerate(zip(base_grid, probe_grid)): + if not (isinstance(base_dim, int) and isinstance(probe_dim, int)): + dynamic_dims.add(idx) + continue + if base_dim != probe_dim: + dynamic_dims.add(idx) + + fixed_dims = {} + for idx, dim in enumerate(base_grid): + if idx in dynamic_dims: + continue + if isinstance(dim, int) and dim > 0: + fixed_dims[idx] = dim + return fixed_dims + + def _autoparse_tiling_params(self, candidates_params: List[str]) -> Dict[str, str]: + """ + Extracts the tiling axis parameters from triton kernel code. + """ + func_ast = self._parse_ast() + parser = TilingAxesParser(func_ast, self.keys, candidates_params) + tiling_axes = parser.parse() + if self.print_autotuning: + print(f"Ascend autotuning parse tiling axes: {tiling_axes}") + return tiling_axes + + def _autoparse_reduction_axes(self) -> List[str]: + """ + Extracts the reduction axis parameters from triton kernel code. + """ + func_ast = self._parse_ast() + parser = ReductionAxesParser(func_ast, self.keys) + reduction_axes = parser.parse() + for axis in reduction_axes: + self.keys[f"r{axis}"] = self.keys.pop(axis) + reduction_axes = [f"r{axis}" for axis in reduction_axes] + + if self.print_autotuning: + print( + f"Ascend autotuning parse keys: {self.keys} \n" + f"Ascend autotuning parse reduction axes: {reduction_axes}" + ) + return reduction_axes + + def _autoparse_low_dim_axes(self) -> List[str]: + """ + Extracts the low dimension axis from triton kernel code. + """ + func_ast = self._parse_ast() + parser = LowDimsAxesParser(func_ast, self.keys) + low_dim_axes = parser.parse() + if len(low_dim_axes) < 1: + if self.print_autotuning: + print( + "[WARNING] Failed to parse low-dimensional axes, fallback to empty low_dim_axes." + ) + return [] + if self.print_autotuning: + print(f"Ascend autotuning parse low dimensional axes: {low_dim_axes}") + return low_dim_axes + + def _autoparse_ptr_nums(self, all_args: dict) -> int: + """ + Counts the number of pointer parameters from triton kernel code. + """ + ptr_nums = 0 + ptr_params = list() + for k, v in all_args.items(): + if isinstance(v, Tensor): + ptr_nums += 1 + ptr_params.append(k) + + if self.print_autotuning: + print( + f"Ascend autotuning parse pointer params: {ptr_params}, pointer nums: {ptr_nums}" + ) + return ptr_nums + + def _get_persistent_reduction_threshold(self, reduction_axis: str) -> int: + # Keep this heuristic aligned with inductor-style policy: + # inner reduction axis uses a larger threshold than other axes. + if self.low_dim_axes and reduction_axis == self.low_dim_axes[0]: + return 1024 + return 64 + + +def autotune( + configs, + key, + prune_configs_by=None, + reset_to_zero=None, + restore_value=None, + pre_hook=None, + post_hook=None, + warmup=None, + rep=None, + use_cuda_graph=False, + do_bench=None, + cache_results=False, + *, + auto_prof_dir=None, + hints=None, +): + """ + Decorator for auto-tuning a :code:`triton.jit`'d function. + + .. highlight:: python + .. code-block:: python + + @triton.autotune(configs=[ + triton.Config(kwargs={'BLOCK_SIZE': 128}, num_warps=4), + triton.Config(kwargs={'BLOCK_SIZE': 1024}, num_warps=8), + ], + key=['x_size'] # the two above configs will be evaluated anytime + # the value of x_size changes + ) + @triton.jit + def kernel(x_ptr, x_size, **META): + BLOCK_SIZE = META['BLOCK_SIZE'] + :note: When all the configurations are evaluated, the kernel will run multiple times. + This means that whatever value the kernel updates will be updated multiple times. + To avoid this undesired behavior, you can use the `reset_to_zero` argument, which + resets the value of the provided tensor to `zero` before running any configuration. + + If the environment variable :code:`TRITON_PRINT_AUTOTUNING` is set to + :code:`"1"`, Triton will print a message to stdout after autotuning each + kernel, including the time spent autotuning and the best configuration. + + :param configs: a list of :code:`triton.Config` objects + :type configs: list[triton.Config] + :param key: a list of argument names whose change in value will trigger the evaluation of all provided configs. + :type key: list[str] + :param prune_configs_by: a dict of functions that are used to prune configs, fields: + 'perf_model': performance model used to predicate running time with different configs, returns running time + 'top_k': number of configs to bench + 'early_config_prune'(optional): a function used to do early prune (eg, num_stages). It takes configs:List[Config] as its input, and returns pruned configs. + :param reset_to_zero: a list of argument names whose value will be reset to zero before evaluating any configs. + :type reset_to_zero: list[str] + :param restore_value: a list of argument names whose value will be restored after evaluating any configs. + :type restore_value: list[str] + :param pre_hook: a function that will be called before the kernel is called. + This overrides the default pre_hook used for 'reset_to_zero' and 'restore_value'. + 'kwargs': a dict of all arguments passed to the kernel. + 'reset_only': a boolean indicating whether the pre_hook is called to reset the values only, without a corresponding post_hook. + :type pre_hook: lambda args, reset_only + :param post_hook: a function that will be called after the kernel is called. + This overrides the default post_hook used for 'restore_value'. + 'kwargs': a dict of all arguments passed to the kernel. + 'exception': the exception raised by the kernel in case of a compilation or runtime error. + :type post_hook: lambda args, exception + :param warmup: warmup time (in ms) to pass to benchmarking (deprecated). + :type warmup: int + :param rep: repetition time (in ms) to pass to benchmarking (deprecated). + :type rep: int + :param do_bench: a benchmark function to measure the time of each run. + :type do_bench: lambda fn, quantiles + :param cache_results: whether to cache autotune timings to disk. + :type cache_results: bool + :param auto_prof_dir: the specified directory to store the profiling results of the best config. + If this parameter is None or the best config is retrieved from cache, the profiling process will be ignored. + :type auto_prof_dir: str + :param hints: a dict of autotune hint auguments passed to AutoTilingTuner. + """ + + def decorator(fn): + return AutoTilingTuner( + fn, + fn.arg_names, + configs, + key, + reset_to_zero, + restore_value, + pre_hook=pre_hook, + post_hook=post_hook, + prune_configs_by=prune_configs_by, + warmup=warmup, + rep=rep, + use_cuda_graph=use_cuda_graph, + do_bench=do_bench, + cache_results=cache_results, + auto_profile_dir=auto_prof_dir, + hints=hints, + ) + + return decorator + + +_ALL_PARAMS = { + "num_stages", + "unit_flag", + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "enable_hivm_auto_cv_balance", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + "enable_ubuf_saving", +} + +_DEFAULTS = { + "num_stages": [2], + "unit_flag": [False], + "limit_auto_multi_buffer_only_for_local_buffer": [False], + "limit_auto_multi_buffer_of_local_buffer": ["no-l0c"], + "set_workspace_multibuffer": [2, 4], + "enable_hivm_auto_cv_balance": [True], + "tile_mix_vector_loop": [2, 4], + "tile_mix_cube_loop": [2, 4], + "enable_ubuf_saving": [True], +} + +_VALID_VALUES = { + "num_stages": [1, 2], + "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], + "set_workspace_multibuffer": [2, 4], + "tile_mix_vector_loop": [2, 4, 8], + "tile_mix_cube_loop": [2, 4, 8], +} + +_CUBE_PARAMS = {"num_stages", "unit_flag", "limit_auto_multi_buffer_of_local_buffer"} + +_MIXCV_PARAMS = { + "num_stages", + "unit_flag", + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "enable_hivm_auto_cv_balance", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + "enable_ubuf_saving", +} + +_VECTOR_PARAMS = { + "num_stages", + "enable_ubuf_saving", +} + + +def _check_boolean_list(val, param_name): + return ( + isinstance(val, (list, tuple)) + and len(val) > 0 + and all(isinstance(x, bool) for x in val) + ) + + +def _check_string_in_set(val, valid_set, param_name): + return ( + isinstance(val, (list, tuple)) + and len(val) > 0 + and all(v in valid_set for v in val) + ) + + +def _check_int_in_set(val, valid_set, param_name): + return ( + isinstance(val, (list, tuple)) + and len(val) > 0 + and all(isinstance(v, int) and v in valid_set for v in val) + ) + + +_VALIDATION_RULES = { + "num_stages": { + "desc": f"must be one or more of: {_VALID_VALUES['num_stages']}", + "check": lambda val, p: _check_int_in_set(val, _VALID_VALUES["num_stages"], p), + }, + "unit_flag": { + "desc": "must be non-empty list/tuple of boolean values", + "check": _check_boolean_list, + }, + "limit_auto_multi_buffer_only_for_local_buffer": { + "desc": "must be non-empty list/tuple of boolean values", + "check": _check_boolean_list, + }, + "limit_auto_multi_buffer_of_local_buffer": { + "desc": f"must be one or more of: {_VALID_VALUES['limit_auto_multi_buffer_of_local_buffer']}", + "check": lambda val, p: _check_string_in_set( + val, _VALID_VALUES["limit_auto_multi_buffer_of_local_buffer"], p + ), + }, + "set_workspace_multibuffer": { + "desc": f"must be one or more of: {_VALID_VALUES['set_workspace_multibuffer']}", + "check": lambda val, p: _check_int_in_set( + val, _VALID_VALUES["set_workspace_multibuffer"], p + ), + }, + "enable_hivm_auto_cv_balance": { + "desc": "must be non-empty list/tuple of boolean values", + "check": _check_boolean_list, + }, + "tile_mix_vector_loop": { + "desc": f"must be one or more of: {_VALID_VALUES['tile_mix_vector_loop']}", + "check": lambda val, p: _check_int_in_set( + val, _VALID_VALUES["tile_mix_vector_loop"], p + ), + }, + "tile_mix_cube_loop": { + "desc": f"must be one or more of: {_VALID_VALUES['tile_mix_cube_loop']}", + "check": lambda val, p: _check_int_in_set( + val, _VALID_VALUES["tile_mix_cube_loop"], p + ), + }, + "enable_ubuf_saving": { + "desc": "must be non-empty list/tuple of boolean values", + "check": _check_boolean_list, + }, +} + + +class BaseAutotuner: + """ + Base class for generating auto-tuning configurations without block dimensions. + Users must provide fixed dimension parameters when calling the kernel. + """ + + def __init__( + self, operator_name, supported_params, default_params, validation_rules + ): + self.operator_name = operator_name + self.supported_params = supported_params + self.default_params = default_params + self.validation_rules = validation_rules + + def validate_parameters(self, **kwargs): + # Check for unsupported parameters + invalid_params = [k for k in kwargs.keys() if k not in _ALL_PARAMS] + if invalid_params: + print( + f"[ERROR] Invalid parameters for {self.operator_name}: {invalid_params}" + ) + return False + + for param, rule in self.validation_rules.items(): + if param in kwargs: + if not rule["check"](kwargs[param], param): + print( + f"[ERROR] Invalid value for '{param}' in {self.operator_name}: {kwargs[param]}" + ) + print(f" Expected: {rule['desc']}") + return False + return True + + def get_configs(self, **kwargs): + """ + Generate a list of Config objects. + Each parameter must be provided as a list (even for a single value). + The function produces the Cartesian product of all parameter lists. + - num_stages: each value will be set as Config.num_stages (not placed in kwargs) + - other parameters: each value will be placed in Config.kwargs + Returns a list of Config objects. + """ + if not self.validate_parameters(**kwargs): + return [] + + # Collect parameter values, using defaults for missing ones + param_values = {} + for p in sorted(self.supported_params): + if p in kwargs: + param_values[p] = kwargs[p] + else: + param_values[p] = self.default_params.get(p, [None]) + + keys = list(param_values.keys()) + values = list(param_values.values()) + combos = list(itertools.product(*values)) + + configs = [] + for combo in combos: + config_kwargs = {} + num_stages_val = None + for i, pname in enumerate(keys): + val = combo[i] + if pname == "num_stages": + num_stages_val = val + else: + config_kwargs[pname] = val + + configs.append( + Config( + kwargs=config_kwargs, + num_stages=num_stages_val if num_stages_val is not None else 2, + ) + ) + return configs + + +CubeAutotuner = BaseAutotuner( + operator_name="cube", + supported_params=_CUBE_PARAMS, + default_params=_DEFAULTS, + validation_rules=_VALIDATION_RULES, +) + +MixcvAutotuner = BaseAutotuner( + operator_name="mixcv", + supported_params=_MIXCV_PARAMS, + default_params=_DEFAULTS, + validation_rules=_VALIDATION_RULES, +) + +VectorAutotuner = BaseAutotuner( + operator_name="vector", + supported_params=_VECTOR_PARAMS, + default_params=_DEFAULTS, + validation_rules=_VALIDATION_RULES, +) + + +def get_autotune_cube_config(**kwargs: Any) -> List[triton.Config]: + """ + Generate autotune configuration for the cube operator. + Supported parameters: num_stages, unit_flag, limit_auto_multi_buffer_of_local_buffer. + """ + import triton + + return CubeAutotuner.get_configs(**kwargs) + + +def get_autotune_cv_config(**kwargs: Any) -> List[triton.Config]: + """ + Generate autotune configuration for the mixcv operator. + Supported parameters: num_stages, unit_flag, limit_auto_multi_buffer_only_for_local_buffer, + limit_auto_multi_buffer_of_local_buffer, set_workspace_multibuffer, + enable_hivm_auto_cv_balance, tile_mix_vector_loop, tile_mix_cube_loop, enable_ubuf_saving + """ + import triton + + return MixcvAutotuner.get_configs(**kwargs) + + +def get_autotune_vector_config(**kwargs: Any) -> List[triton.Config]: + """ + Generate autotune configuration for the vector operator. + Supported parameters: num_stages, enable_ubuf_saving + """ + import triton + + return VectorAutotuner.get_configs(**kwargs) + + +def get_max_configs(config, kernel_type="mixcv", **kwargs): + """ + Expand a single base Config by combining it with tuning parameters. + + :param config: A triton.Config object serving as the base. + :param kernel_type: Operator type, one of "cube", "mixcv", "vector". Default "mixcv". + :param kwargs: Tuning parameters, each provided as a list (e.g., enable_hivm_auto_cv_balance=[True, False]). + If a parameter is not provided, its value is taken from the base config (if present) + or from the defaults. + :return: List of expanded Config objects. + """ + # Determine the set of parameters supported by the current kernel_type + if kernel_type == "cube": + supported = _CUBE_PARAMS + elif kernel_type == "vector": + supported = _VECTOR_PARAMS + else: + supported = _MIXCV_PARAMS + + # Warn about unsupported parameters provided in kwargs + unsupported = [k for k in kwargs if k not in supported and k in _ALL_PARAMS] + if unsupported: + print( + f"[WARNING] The following parameters are not supported for kernel_type '{kernel_type}': {unsupported}. They will be ignored." + ) + + # Build value lists for each parameter (priority: kwargs > base config > defaults) + param_values = {} + base_kwargs = config.kwargs + base_num_stages = config.num_stages + + for param in sorted(supported): + if param in kwargs: + # User-provided list via tuning_params takes precedence + val_list = kwargs[param] + elif param == "num_stages": + # num_stages is an attribute of Config, not part of kwargs. + # Triton's default is 3, but Ascend only supports the local defaults. + # Treat an unsupported base value as "not fixed" so examples that use + # triton.Config(kwargs={...}) still follow the Ascend default table. + val_list = ( + [base_num_stages] + if base_num_stages in _VALID_VALUES["num_stages"] + else _DEFAULTS[param] + ) + elif param in base_kwargs: + # Parameter present in base config's kwargs -> fix to that single value + val_list = [base_kwargs[param]] + else: + # Otherwise fall back to defaults + val_list = _DEFAULTS.get(param, [None]) + + # Validate the value list + if param in _VALIDATION_RULES: + rule = _VALIDATION_RULES[param] + if not rule["check"](val_list, param): + raise ValueError( + f"Invalid value for '{param}': {val_list}. Expected: {rule['desc']}" + ) + param_values[param] = val_list + + # Cartesian product of all parameter lists + keys = list(param_values.keys()) + values = list(param_values.values()) + combos = list(itertools.product(*values)) + + new_configs = [] + for combo in combos: + # Start with a copy of the original config's kwargs + new_kwargs = config.kwargs.copy() + num_stages_val = None + + for i, pname in enumerate(keys): + val = combo[i] + if pname == "num_stages": + num_stages_val = val + else: + # Overwrite or add the parameter to kwargs + new_kwargs[pname] = val + + config_args = { + "kwargs": new_kwargs, + "num_warps": getattr(config, "num_warps", 4), + "num_stages": ( + num_stages_val + if num_stages_val is not None + else getattr(config, "num_stages", 2) + ), + "num_ctas": getattr(config, "num_ctas", 1), + "maxnreg": getattr(config, "maxnreg", None), + "pre_hook": getattr(config, "pre_hook", None), + "ir_override": getattr(config, "ir_override", None), + "num_buffers_warp_spec": getattr(config, "num_buffers_warp_spec", None), + "num_consumer_groups": getattr(config, "num_consumer_groups", None), + "reg_dec_producer": getattr(config, "reg_dec_producer", None), + "reg_inc_consumer": getattr(config, "reg_inc_consumer", None), + } + new_config = _make_config_compat(**config_args) + new_configs.append(new_config) + + return new_configs + + +def max_autotune( + configs, + key, + kernel_type="mixcv", + prune_configs_by=None, + reset_to_zero=None, + restore_value=None, + pre_hook=None, + post_hook=None, + warmup=None, + rep=None, + use_cuda_graph=False, + do_bench=None, + cache_results=False, + **tuning_params, +): + """ + Decorator that expands each base Config with tuning parameters before auto-tuning. + + Usage is similar to @triton.autotune, but allows automatic expansion of + additional tuning parameters (e.g., enable_hivm_auto_cv_balance, tile_mix_vector_loop, ...) + for each provided base configuration. + + :param configs: List of base triton.Config objects. + :param key: List of argument names whose change triggers re-tuning. + :param kernel_type: Operator type, one of "cube", "mixcv", "vector". Default "mixcv". + :param prune_configs_by: Same as in autotune. + :param reset_to_zero: Same as in autotune. + :param restore_value: Same as in autotune. + :param pre_hook: Same as in autotune. + :param post_hook: Same as in autotune. + :param warmup: Deprecated. + :param rep: Deprecated. + :param use_cuda_graph: Deprecated. + :param do_bench: Same as in autotune. + :param cache_results: Same as in autotune. + :param tuning_params: Additional tuning parameters as keyword arguments. + Each value must be a list; the Cartesian product of these lists + will be combined with each base config. + """ + + def decorator(fn): + if not configs or len(configs) == 0: + raise ValueError( + "[max_autotune] The argument 'configs' cannot be empty. " + "Please provide at least one base config. " + ) + # Expand each base config with the provided tuning parameters + expanded_configs = [] + for cfg in configs: + expanded = get_max_configs(cfg, kernel_type=kernel_type, **tuning_params) + expanded_configs.extend(expanded) + + # Call the original autotune decorator with the expanded configs + return autotune( + configs=expanded_configs, + key=key, + prune_configs_by=prune_configs_by, + reset_to_zero=reset_to_zero, + restore_value=restore_value, + pre_hook=pre_hook, + post_hook=post_hook, + warmup=warmup, + rep=rep, + use_cuda_graph=use_cuda_graph, + do_bench=do_bench, + cache_results=cache_results, + )(fn) + + return decorator diff --git a/backend/ascend_autotune_runtime/benchmark.py b/backend/ascend_autotune_runtime/benchmark.py new file mode 100644 index 00000000..25c28b37 --- /dev/null +++ b/backend/ascend_autotune_runtime/benchmark.py @@ -0,0 +1,27 @@ +"""Benchmark strategy for Ascend autotuning.""" + +from __future__ import annotations + +from typing import Mapping + + +class NpuProfilerBenchStrategy: + def bench(self, run_fns: Mapping): + from ..testing import do_bench_npu + + costs = do_bench_npu(list(run_fns.values()), clear_l2_cache=False) + if not isinstance(costs, (list, tuple)): + raise RuntimeError( + "do_bench_npu must return one timing per autotune config " + f"when benchmarking {len(run_fns)} configs, got {type(costs).__name__}." + ) + if len(costs) != len(run_fns): + raise RuntimeError( + "do_bench_npu returned mismatched timing count: " + f"expected {len(run_fns)}, got {len(costs)}." + ) + return {config: cost for config, cost in zip(run_fns.keys(), costs)} + + +def select_benchmark_strategy(*_unused): + return NpuProfilerBenchStrategy() diff --git a/backend/ascend_autotune_runtime/compile_options.py b/backend/ascend_autotune_runtime/compile_options.py new file mode 100644 index 00000000..1f4957a0 --- /dev/null +++ b/backend/ascend_autotune_runtime/compile_options.py @@ -0,0 +1,612 @@ +from __future__ import annotations + +import inspect +import itertools +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +from triton.runtime.autotuner import Config +from triton.backends.dicp_triton.utils import is_compile_on_910_95 + + +DEFAULT_MAX_CONFIGS = None + +_VALID_VALUES = { + "num_stages": [1, 2], + "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], + "set_workspace_multibuffer": [2, 4], + "tile_mix_vector_loop": [1, 2, 4, 8], + "tile_mix_cube_loop": [1, 2, 4, 8], +} + +_BOOLEAN_PARAMS = { + "enable_tuning_mode", + "multibuffer", + "unit_flag", + "limit_auto_multi_buffer_only_for_local_buffer", + "enable_hivm_auto_cv_balance", + "enable_ubuf_saving", + # "enable_preload", + "enable_auto_bind_sub_block", +} + +_SUPPORTED_PARAMS = { + "cube": { + "enable_tuning_mode", + "num_stages", + "unit_flag", + "limit_auto_multi_buffer_of_local_buffer", + }, + "mixcv": { + "enable_tuning_mode", + "num_stages", + "multibuffer", + "unit_flag", + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "enable_hivm_auto_cv_balance", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + "enable_ubuf_saving", + # "enable_preload", + "enable_auto_bind_sub_block", + }, + "vector": { + "enable_tuning_mode", + "num_stages", + "enable_ubuf_saving", + }, +} + +_ALL_PARAMS = set().union(*_SUPPORTED_PARAMS.values()) + +_MIXCV_910_95_UNSUPPORTED_PARAMS = { + "tile_mix_vector_loop", + "tile_mix_cube_loop", +} + +_AUTO_SEARCH_PRESETS = { + "cube": { + "enable_tuning_mode": [True], + "num_stages": [1, 2], + "unit_flag": [False, True], + "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], + }, + "mixcv": { + "enable_tuning_mode": [True], + "num_stages": [1, 2], + "unit_flag": [False, True], + "limit_auto_multi_buffer_only_for_local_buffer": [True, False], + "limit_auto_multi_buffer_of_local_buffer": ["no-limit", "no-l0c"], + "set_workspace_multibuffer": [2, 4], + "enable_hivm_auto_cv_balance": [True], + "tile_mix_vector_loop": [1, 2, 4], + "tile_mix_cube_loop": [1, 2, 4], + "enable_ubuf_saving": [False, True], + # "enable_preload": [False, True], + "enable_auto_bind_sub_block": [True], + }, + "vector": { + "enable_tuning_mode": [True], + "num_stages": [1, 2], + "enable_ubuf_saving": [True, False], + }, +} + + +def _is_mixcv_multi_buffer_auto_enabled( + num_stages: int, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], +) -> bool: + """Whether `enable_auto_multi_buffer` takes effect for this combination.""" + if num_stages == 1: + return False + multibuffer = _resolve_compile_option( + "multibuffer", combo, config, fixed_options, default=None + ) + return multibuffer is not False + + +def _is_mixcv_limit_to_local_only_active( + num_stages: int, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], +) -> bool: + return _is_mixcv_multi_buffer_auto_enabled(num_stages, combo, config, fixed_options) + + +def _is_mixcv_workspace_multibuffer_active( + num_stages: int, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], +) -> bool: + if not _is_mixcv_multi_buffer_auto_enabled( + num_stages, combo, config, fixed_options + ): + return False + limit_to_local_only = _resolve_compile_option( + "limit_auto_multi_buffer_only_for_local_buffer", + combo, + config, + fixed_options, + default=True, + ) + return limit_to_local_only is False + + +_MIXCV_OPTION_ACTIVITY_RULES = { + "limit_auto_multi_buffer_only_for_local_buffer": _is_mixcv_limit_to_local_only_active, + "limit_auto_multi_buffer_of_local_buffer": _is_mixcv_limit_to_local_only_active, + "set_workspace_multibuffer": _is_mixcv_workspace_multibuffer_active, + "tile_mix_vector_loop": _is_mixcv_workspace_multibuffer_active, + "tile_mix_cube_loop": _is_mixcv_workspace_multibuffer_active, +} + + +@dataclass +class CompileOptionsSpec: + enabled: bool = False + kernel_type: str = "mixcv" + params: Dict[str, List[Any]] = field(default_factory=dict) + max_configs: Optional[int] = DEFAULT_MAX_CONFIGS + + +def _normalize_kernel_type(kernel_type: str) -> str: + if kernel_type == "mix": + return "mixcv" + if kernel_type not in _SUPPORTED_PARAMS: + raise ValueError( + "compile_options kernel_type must be one of: cube, mix, mixcv, vector" + ) + return kernel_type + + +def _as_value_list(name: str, value: Any) -> List[Any]: + values = list(value) if isinstance(value, (list, tuple)) else [value] + if not values: + raise ValueError(f"compile_options parameter '{name}' must not be empty") + return values + + +def validate_compile_option_values(name: str, values: List[Any]) -> None: + if name in _BOOLEAN_PARAMS and not all(isinstance(v, bool) for v in values): + raise ValueError(f"compile_options parameter '{name}' expects boolean values") + + if name in _VALID_VALUES and not all(v in _VALID_VALUES[name] for v in values): + raise ValueError( + f"compile_options parameter '{name}' expects values in {_VALID_VALUES[name]}" + ) + + +def parse_compile_options_hint(hint: Any) -> CompileOptionsSpec: + if hint is None or hint is False: + return CompileOptionsSpec(enabled=False) + + if hint is True: + return CompileOptionsSpec(enabled=True) + + if isinstance(hint, str): + return CompileOptionsSpec( + enabled=True, + kernel_type=_normalize_kernel_type(hint), + ) + + if not isinstance(hint, dict): + raise TypeError("hints['compile_options'] must be bool, str, or dict") + + raw = dict(hint) + kernel_type = _normalize_kernel_type( + raw.pop("kernel_type", raw.pop("type", "mixcv")) + ) + max_configs = raw.pop("max_configs", DEFAULT_MAX_CONFIGS) + if max_configs is not None and ( + not isinstance(max_configs, int) or max_configs <= 0 + ): + raise ValueError( + "compile_options max_configs must be a positive integer or None" + ) + + nested_options = raw.pop("options", {}) + if nested_options: + if not isinstance(nested_options, dict): + raise TypeError("compile_options options must be a dict") + raw.update(nested_options) + + supported = _SUPPORTED_PARAMS[kernel_type] + params: Dict[str, List[Any]] = {} + for name, value in raw.items(): + if name not in _ALL_PARAMS: + raise ValueError(f"Unknown compile_options parameter: {name}") + if name not in supported: + print( + f"[WARNING] compile_options parameter '{name}' is not supported " + f"for kernel_type '{kernel_type}' and will be ignored." + ) + continue + values = _as_value_list(name, value) + validate_compile_option_values(name, values) + params[name] = values + + return CompileOptionsSpec( + enabled=True, + kernel_type=kernel_type, + params=params, + max_configs=max_configs, + ) + + +def _make_config_compat(**kwargs): + supported_config_args = inspect.signature(Config).parameters + return Config( + **{key: value for key, value in kwargs.items() if key in supported_config_args} + ) + + +def _value_space_for_config( + config: Config, spec: CompileOptionsSpec, *, generated_tiling: bool +) -> Dict[str, List[Any]]: + supported = _SUPPORTED_PARAMS[spec.kernel_type] + preset = _AUTO_SEARCH_PRESETS[spec.kernel_type] + + value_space = {} + for name in sorted(supported): + if name in spec.params: + values = spec.params[name] + elif name not in preset: + continue + else: + values = preset[name] + validate_compile_option_values(name, values) + value_space[name] = values + if spec.kernel_type == "mixcv" and is_compile_on_910_95: + for name in _MIXCV_910_95_UNSUPPORTED_PARAMS: + value_space.pop(name, None) + return value_space + + +def _is_inactive_reason( + name: str, + num_stages: int, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], +) -> Optional[str]: + if not _is_param_effective(name, num_stages, combo, config, fixed_options): + if name in { + "limit_auto_multi_buffer_only_for_local_buffer", + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + }: + return "depends on auto multi-buffer" + return None + + +def _is_param_effective( + name: str, + num_stages: int, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], +) -> bool: + if name not in _MIXCV_OPTION_ACTIVITY_RULES: + return True + return _MIXCV_OPTION_ACTIVITY_RULES[name](num_stages, combo, config, fixed_options) + + +def _resolve_compile_option( + name: str, + combo: Dict[str, Any], + config: Config, + fixed_options: Dict[str, Any], + default: Any, +) -> Any: + if name in combo: + return combo[name] + if name in fixed_options: + return fixed_options[name] + return default + + +def _effective_values( + name: str, + value_space: Dict[str, List[Any]], + fixed_options: Dict[str, Any], + default: Any, +) -> List[Any]: + if name in fixed_options: + return [fixed_options[name]] + if name in value_space: + return value_space[name] + return [default] + + +def _emit_values( + names: List[str], + value_space: Dict[str, List[Any]], + fixed_options: Dict[str, Any], +) -> List[tuple[str, List[Any]]]: + return [ + (name, value_space[name]) + for name in names + if name in value_space and name not in fixed_options + ] + + +def _product_dict(items: List[tuple[str, List[Any]]]): + if not items: + yield {} + return + names = [name for name, _ in items] + values = [values for _, values in items] + for combo in itertools.product(*values): + yield dict(zip(names, combo)) + + +def _mixcv_branch_items( + *, + num_stages: int, + value_space: Dict[str, List[Any]], + fixed_options: Dict[str, Any], +) -> List[List[tuple[str, List[Any]]]]: + independent_names = [ + "enable_tuning_mode", + "unit_flag", + "enable_ubuf_saving", + "enable_hivm_auto_cv_balance", + "enable_auto_bind_sub_block", + ] + independent_items = _emit_values(independent_names, value_space, fixed_options) + + multibuffer_values = _effective_values( + "multibuffer", value_space, fixed_options, None + ) + branches = [] + for multibuffer in multibuffer_values: + multibuffer_items = [] + if "multibuffer" in value_space and "multibuffer" not in fixed_options: + multibuffer_items = [("multibuffer", [multibuffer])] + branch_base = independent_items + multibuffer_items + + if num_stages == 1 or multibuffer is False: + branches.append(branch_base) + continue + + limit_only_values = _effective_values( + "limit_auto_multi_buffer_only_for_local_buffer", + value_space, + fixed_options, + True, + ) + for limit_only in limit_only_values: + limit_only_items = [] + if ( + "limit_auto_multi_buffer_only_for_local_buffer" in value_space + and "limit_auto_multi_buffer_only_for_local_buffer" not in fixed_options + ): + limit_only_items = [ + ("limit_auto_multi_buffer_only_for_local_buffer", [limit_only]) + ] + if limit_only is True: + branch_names = ["limit_auto_multi_buffer_of_local_buffer"] + else: + branch_names = [ + "limit_auto_multi_buffer_of_local_buffer", + "set_workspace_multibuffer", + "tile_mix_vector_loop", + "tile_mix_cube_loop", + ] + branches.append( + branch_base + + limit_only_items + + _emit_values(branch_names, value_space, fixed_options) + ) + + return branches + + +def _make_expanded_config( + config: Config, + spec: CompileOptionsSpec, + combo_value: Dict[str, Any], + num_stages: int, +) -> Config: + new_kwargs = dict(config.kwargs) + for name in _SUPPORTED_PARAMS[spec.kernel_type]: + new_kwargs.pop(name, None) + for name, value in combo_value.items(): + new_kwargs[name] = value + + return _make_config_compat( + kwargs=new_kwargs, + num_warps=getattr(config, "num_warps", 4), + num_stages=num_stages, + num_ctas=getattr(config, "num_ctas", 1), + maxnreg=getattr(config, "maxnreg", None), + pre_hook=getattr(config, "pre_hook", None), + ir_override=getattr(config, "ir_override", None), + num_buffers_warp_spec=getattr(config, "num_buffers_warp_spec", None), + num_consumer_groups=getattr(config, "num_consumer_groups", None), + reg_dec_producer=getattr(config, "reg_dec_producer", None), + reg_inc_consumer=getattr(config, "reg_inc_consumer", None), + ) + + +def _hashable_value(value: Any): + if isinstance(value, dict): + return tuple(sorted((key, _hashable_value(val)) for key, val in value.items())) + if isinstance(value, (list, tuple)): + return tuple(_hashable_value(item) for item in value) + if isinstance(value, set): + return tuple(sorted(_hashable_value(item) for item in value)) + try: + hash(value) + except TypeError: + return repr(value) + return value + + +def _config_key(config: Config) -> tuple: + return ( + tuple( + sorted( + (key, _hashable_value(value)) for key, value in config.kwargs.items() + ) + ), + getattr(config, "num_warps", 4), + getattr(config, "num_stages", None), + getattr(config, "num_ctas", 1), + getattr(config, "maxnreg", None), + id(getattr(config, "pre_hook", None)), + _hashable_value(getattr(config, "ir_override", None)), + getattr(config, "num_buffers_warp_spec", None), + getattr(config, "num_consumer_groups", None), + getattr(config, "reg_dec_producer", None), + getattr(config, "reg_inc_consumer", None), + ) + + +def expand_compile_option_configs( + configs: List[Config], + spec: CompileOptionsSpec, + *, + generated_tiling: bool, + fixed_options: Optional[Dict[str, Any]] = None, +) -> List[Config]: + if not spec.enabled or not configs: + return configs + + fixed_options = fixed_options or {} + expanded_configs = [] + emitted_config_keys = set() + for config in configs: + value_space = _value_space_for_config( + config, spec, generated_tiling=generated_tiling + ) + if "num_stages" in fixed_options: + num_stage_values = [fixed_options["num_stages"]] + value_space.pop("num_stages", None) + else: + num_stage_values = value_space.pop("num_stages") + + for name in fixed_options: + if name != "num_stages": + value_space.pop(name, None) + + for num_stages in num_stage_values: + if spec.kernel_type == "mixcv": + branch_items = _mixcv_branch_items( + num_stages=num_stages, + value_space=value_space, + fixed_options=fixed_options, + ) + combo_iter = itertools.chain.from_iterable( + _product_dict(items) for items in branch_items + ) + else: + combo_iter = _product_dict(list(value_space.items())) + + for combo_value in combo_iter: + new_config = _make_expanded_config( + config, spec, combo_value, num_stages + ) + config_key = _config_key(new_config) + if config_key in emitted_config_keys: + continue + emitted_config_keys.add(config_key) + + if ( + spec.max_configs is not None + and len(expanded_configs) >= spec.max_configs + ): + raise ValueError( + "compile_options generated more than " + f"{spec.max_configs} configs. Narrow the search space or raise max_configs." + ) + expanded_configs.append(new_config) + + return expanded_configs + + +def get_compile_option_param_names(spec: CompileOptionsSpec) -> set[str]: + if not spec.enabled: + return set() + return set(_SUPPORTED_PARAMS[spec.kernel_type]) + + +def format_compile_option_result( + config: Config, + spec: CompileOptionsSpec, + fixed_options: Optional[Dict[str, Any]] = None, +) -> str: + if not spec.enabled: + return str(config) + + fixed_options = fixed_options or {} + compile_param_names = _SUPPORTED_PARAMS[spec.kernel_type] - {"num_stages"} + selected_meta = { + key: value + for key, value in sorted(config.kwargs.items()) + if key not in compile_param_names + } + selected_meta["num_stages"] = getattr(config, "num_stages", None) + + effective = { + key: value + for key, value in sorted(config.kwargs.items()) + if key in compile_param_names + } + effective.update( + { + key: value + for key, value in sorted(fixed_options.items()) + if key in compile_param_names + } + ) + + if spec.kernel_type == "mixcv": + num_stages = selected_meta["num_stages"] + multibuffer = effective.get("multibuffer", None) + effective["enable_auto_multi_buffer"] = ( + False if multibuffer is False or num_stages == 1 else True + ) + for name in _SUPPORTED_PARAMS[spec.kernel_type]: + if name == "num_stages": + continue + reason = _is_inactive_reason( + name, num_stages, config.kwargs, config, fixed_options + ) + if reason is not None and name not in effective: + effective[name] = f"" + if reason is not None and name in effective: + effective[name] = f"" + if is_compile_on_910_95: + for name in _MIXCV_910_95_UNSUPPORTED_PARAMS: + effective[name] = "" + + selected_items = [f"{key}={value}" for key, value in selected_meta.items()] + effective_items = [f"{key}={value}" for key, value in sorted(effective.items())] + return ( + "selected_meta: " + + ", ".join(selected_items) + + "; effective_compile_options: " + + ", ".join(effective_items) + ) + + +def summarize_compile_option_configs( + configs: List[Config], limit: Optional[int] = None +) -> List[str]: + summary = [] + selected_configs = configs if limit is None else configs[:limit] + for config in selected_configs: + items = [f"{key}={value}" for key, value in sorted(config.kwargs.items())] + items.append(f"num_stages={getattr(config, 'num_stages', None)}") + summary.append(", ".join(items)) + return summary diff --git a/backend/ascend_autotune_runtime/tile_generator.py b/backend/ascend_autotune_runtime/tile_generator.py new file mode 100644 index 00000000..c21b5e05 --- /dev/null +++ b/backend/ascend_autotune_runtime/tile_generator.py @@ -0,0 +1,569 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# 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. + +from __future__ import annotations + +import functools +import sys +from dataclasses import dataclass +from typing import ( + Dict, + List, + Tuple, +) + +from triton.runtime.autotuner import Config + +from .utils import ( + get_byte_per_numel, + next_power_of_2, + num_vector_core, + ub_size_in_kbytes, + rf_size_in_kbytes, +) + + +@dataclass +class AxisInfo: + name: str + index: int + length: int + + prefix: str = "" + split_name: str = "" + tiling_name: str = "" + is_split_axis: bool = False + is_tunable_split_axis: bool = False + is_tiling_axis: bool = False + fixed_split_size: int = 0 + + @property + def is_reduction(self): + return self.prefix == "r" + + +class KernelMeta: + def __init__( + self, + axis_sizes: Dict[str, int], + split_params: Dict[str, str], + fixed_split_params: Dict[str, int], + tiling_params: Dict[str, str], + low_dims: List[str], + dtype: torch.dtype, + persistent_reduction: bool, + dual_reduction: bool, + num_buffers: int, + is_simt_mode: bool, + ): + """ + :param split_params: a dict of axis name: argument name, the argument is an adjustable parameter in a split axis, such as 'XBLOCK'. + The axis name must be in key's axis names. Do not add prefix 'r' before the axis name. + This param can be empty. Note that the auto tiling feature will be disabled when the split_params and tiling_params are both empty. + The split axis can usually be identified according to `tl.program_id()` expression. + :type split_params: Dict[str, str] + :param tiling_params: a dict of axis name: argument name, the argument is an adjustable parameter in a tiling axis, such as 'XBLOCK_SUB'. + The axis name must be in key's axis names. Do not add prefix 'r' before the axis name. + This param can be empty. Note that the auto tiling feature will be disabled when the split_params and tiling_params are both empty. + The tiling axis can usually be identified according to `tl.arange()` expression. + :type tiling_params: Dict[str, str] + :param low_dims: a list of axis name in which the corresponding axis is low dim aixs. + The axis name must be in key's axis names. Do not add prefix 'r' before the axis name. + :type low_dims: List[str] + :param dual_reduction: performing reduction on more than one axis. + :param persistent_reduction: there is no splitting in reduction axis. + """ + self._validate_axis( + axis_sizes, split_params, fixed_split_params, tiling_params, low_dims + ) + + axis_dict = {} + idx = 0 + for name, length in axis_sizes.items(): + prefix = "" + if name.startswith("r"): + prefix = "r" + + is_tunable_split_axis = name in split_params + fixed_split_size = fixed_split_params.get(name, 0) + is_split_axis = is_tunable_split_axis or fixed_split_size > 0 + is_tiling_axis = name in tiling_params + split_name = "" if not is_tunable_split_axis else split_params[name] + tiling_name = "" if name not in tiling_params else tiling_params[name] + + axis_dict[name] = AxisInfo( + name=name, + index=idx, + length=length, + prefix=prefix, + split_name=split_name, + tiling_name=tiling_name, + is_split_axis=is_split_axis, + is_tunable_split_axis=is_tunable_split_axis, + is_tiling_axis=is_tiling_axis, + fixed_split_size=fixed_split_size, + ) + idx += 1 + + self.axis_info = list(axis_dict.values()) + self.split_axis = [x for x in axis_dict.values() if x.is_split_axis] + self.tunable_split_axis = [ + x for x in axis_dict.values() if x.is_tunable_split_axis + ] + self.tiling_axis = [x for x in axis_dict.values() if x.is_tiling_axis] + self.low_dims_axis = [x for x in axis_dict.values() if x.name in low_dims] + self.dtype = dtype + self.persistent_reduction = persistent_reduction + self.dual_reduction = dual_reduction + self.num_buffers = num_buffers + self.is_simt_mode = is_simt_mode + + @classmethod + def _validate_axis( + cls, + axis_sizes: Dict[str, int], + split_params: Dict[str, str], + fixed_split_params: Dict[str, int], + tiling_params: Dict[str, str], + low_dims: List[str], + ) -> None: + for axis_name in axis_sizes.keys(): + if axis_name.startswith("r") and len(axis_name) == 1: + raise ValueError("The name of a reduction axis is empty!") + + def check_keys(params: List[str], context="parameter"): + for k in params: + if k not in axis_sizes and ("r" + k) not in axis_sizes: + raise KeyError( + f"{context} '{k}' not found in known axes: {axis_sizes.keys()}" + ) + + check_keys(split_params.keys(), "split axis") + check_keys(fixed_split_params.keys(), "fixed split axis") + check_keys(tiling_params.keys(), "tiling axis") + check_keys(low_dims, "low dim axis") + + +@dataclass +class BlockInfo: + block_name: str # e.g., XBLOCK + sub_block_name: str # e.g., XBLOCK_SUB + block_size: int + sub_block_size: int + + +""" +Generate possible candidate tiling configs for benchmarking +""" + + +class TileGenerator: + num_warps = 1 + num_stages = 1 + + def __init__(self, kernel_meta: KernelMeta): + self.kernel_meta = kernel_meta + self.persistent_reduction = self.kernel_meta.persistent_reduction + self.dual_reduction = self.kernel_meta.dual_reduction + + self.blocks = self.init_blocks_info(kernel_meta) + self.numels = [axis.length for axis in kernel_meta.axis_info] + self.candidate_blocks = [] + self.configs = [] + self.dtype_bytes = get_byte_per_numel(kernel_meta.dtype) + + self.num_buffers = ( + 3 if kernel_meta.num_buffers == 0 else min(kernel_meta.num_buffers, 3) + ) + self.is_simt_mode = kernel_meta.is_simt_mode + local_mem_size = rf_size_in_kbytes if self.is_simt_mode else ub_size_in_kbytes + self.max_numel_threshold = ( + local_mem_size * 1024 // self.dtype_bytes // self.num_buffers + ) + self.max_total_numel = ( + functools.reduce(lambda x, y: x * y, [x.block_size for x in self.blocks]) + if self.blocks + else 1 + ) + self.small_kernel = self.max_total_numel < 128 * 1024 + self.tiny_kernel = self.max_total_numel <= 32 * 1024 + self.stop_numel = ( + min(1024 // self.dtype_bytes, self.max_total_numel // (num_vector_core * 2)) + if self.small_kernel + else 1024 // self.dtype_bytes + ) + self.max_programs_num = 65535 + self.tiny_program_threshold = num_vector_core // 8 + self.tiny_per_program_cap = 1 + self.tiny_low_program_hist = { + p: 0 for p in range(1, self.tiny_program_threshold + 1) + } + self.tiny_low_program_active = False + self.tiny_low_program_tile_floor = 0 + + @classmethod + def init_blocks_info(cls, kernel_meta: KernelMeta) -> List[BlockInfo]: + blocks = [] + for axis in kernel_meta.axis_info: + block_name = axis.split_name + sub_block_name = axis.tiling_name + block_size = ( + axis.fixed_split_size if axis.fixed_split_size > 0 else axis.length + ) + sub_block_size = block_size + blocks.append( + BlockInfo(block_name, sub_block_name, block_size, sub_block_size) + ) + + return blocks + + @classmethod + def get_key_from_dict(cls, kwargs: Dict[str, int]): + return tuple(sorted(kwargs.items())) + + def calcu_last_split_blocks(self, axis_idx): + splits = 1 + for x in self.kernel_meta.split_axis: + if x.index != axis_idx: + splits = splits * ( + (self.numels[x.index] + self.blocks[x.index].block_size - 1) + // self.blocks[x.index].block_size + ) + else: + break + + last_splits = num_vector_core // splits + last_splits = max(1, last_splits) + last_blocks = (self.numels[axis_idx] + last_splits - 1) // last_splits + return last_blocks + + def aligned_numel(self, numel, align_bytes=32): + if self.is_simt_mode: + return next_power_of_2(numel) + + align_numel = align_bytes // self.dtype_bytes + if numel <= align_numel: + return numel + return ((numel + align_numel - 1) // align_numel) * align_numel + + def calculate_tile_numel(self): + tile_numel = 1 + for axis in self.kernel_meta.axis_info: + if axis.is_tiling_axis: + tile_numel *= self.blocks[axis.index].sub_block_size + else: + tile_numel *= self.blocks[axis.index].block_size + + return tile_numel + + def fill_config(self, cfg, candi_block): + for axis in self.kernel_meta.axis_info: + if not (axis.is_split_axis or axis.is_tiling_axis): + continue + block_info = self.blocks[axis.index] + if axis.is_split_axis: + curr_numel = candi_block[axis.index] + if not axis.is_tiling_axis: + curr_numel = self.aligned_numel(curr_numel) + if block_info.block_name: + cfg[block_info.block_name] = curr_numel + if axis.is_tiling_axis: + tiling_numel = self.aligned_numel(block_info.sub_block_size) + cfg[block_info.sub_block_name] = ( + tiling_numel + if self.is_simt_mode + else min(tiling_numel, candi_block[axis.index]) + ) + + def find_config(self, cfg): + for config_var in self.configs: + if config_var.kwargs == cfg: + return True + return False + + def _try_add_tiny_low_program_config(self, total_programs): + if ( + not self.tiny_kernel + or total_programs < 1 + or total_programs > self.tiny_program_threshold + ): + return + + if ( + self.tiny_low_program_hist.get(total_programs, 0) + >= self.tiny_per_program_cap + ): + return + + candi_block = tuple([x.block_size for x in self.blocks]) + if self.add_to_configs(list(candi_block)): + if candi_block not in self.candidate_blocks: + self.candidate_blocks.append(candi_block) + if not self.tiny_low_program_active: + self.tiny_low_program_active = True + self.tiny_low_program_tile_floor = self.calculate_tile_numel() + self.tiny_low_program_hist[total_programs] = ( + self.tiny_low_program_hist.get(total_programs, 0) + 1 + ) + + def _calc_total_programs(self, candi_block=None): + grids = [] + for axis in self.kernel_meta.split_axis: + numel = self.numels[axis.index] + block_size = ( + self.blocks[axis.index].block_size + if candi_block is None + else candi_block[axis.index] + ) + programs = (numel + block_size - 1) // block_size + grids.append(programs) + + total_programs = functools.reduce(lambda x, y: x * y, grids) if grids else 1 + return total_programs + + def add_to_configs(self, candi_block): + newcfg = {} + self.fill_config(newcfg, candi_block) + tile_numel = self.calculate_tile_numel() + stop_numel_threshold = ( + 0 if len(self.configs) < 10 or self.small_kernel else self.stop_numel + 100 + ) + if self.tiny_low_program_active and self.tiny_low_program_tile_floor > 0: + total_programs = self._calc_total_programs(candi_block) + program_threshold = ( + self.tiny_program_threshold + if self.small_kernel + else num_vector_core // 2 + ) + if total_programs <= program_threshold: + tiny_low_program_threshold = max( + self.stop_numel, self.tiny_low_program_tile_floor // 2 + ) + stop_numel_threshold = max( + stop_numel_threshold, tiny_low_program_threshold + ) + if ( + tile_numel <= self.max_numel_threshold + and tile_numel >= stop_numel_threshold + and not self.find_config(newcfg) + ): + self.configs.append(Config(newcfg, num_warps=1, num_stages=1)) + return True + return False + + def desecnd_all_low_dims_with_all_blocks(self): + restore_sub_blocks = {} + for axis in self.kernel_meta.low_dims_axis: + restore_sub_blocks[axis.index] = self.blocks[axis.index].sub_block_size + self.descend_all_low_dims() + for axis in self.kernel_meta.low_dims_axis: + self.blocks[axis.index].sub_block_size = restore_sub_blocks[axis.index] + + def descend_one_axis(self, axis_idx: int, is_split=False): + def calc_total_programs(): + grids = [] + for axis in self.kernel_meta.split_axis: + numel = self.numels[axis.index] + block_size = self.blocks[axis.index].block_size + programs = (numel + block_size - 1) // block_size + grids.append(programs) + + total_programs = functools.reduce(lambda x, y: x * y, grids) if grids else 1 + return total_programs + + reached_stop_numel = False + slow_decend_split = False + num_vector_core_tile = num_vector_core + max_programs_num = ( + num_vector_core_tile + if self.kernel_meta.tiling_axis + else self.max_programs_num + ) + if not is_split and len(self.candidate_blocks) == 0: + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + + axis = self.kernel_meta.axis_info[axis_idx] + while True: + for candi_block in self.candidate_blocks: + if self.add_to_configs(candi_block): + self.desecnd_all_low_dims_with_all_blocks() + + # tile numel reached threshold + tile_numel = self.calculate_tile_numel() + if tile_numel <= self.stop_numel: + if self.add_to_configs([x.block_size for x in self.blocks]): + self.desecnd_all_low_dims_with_all_blocks() + reached_stop_numel = True + break + + numel = ( + self.blocks[axis_idx].block_size + if is_split + else self.blocks[axis_idx].sub_block_size + ) + if numel == 1: + if self.add_to_configs([x.block_size for x in self.blocks]): + self.desecnd_all_low_dims_with_all_blocks() + break + + if is_split: + if self.persistent_reduction and axis.is_reduction: + reached_stop_numel = True + break + total_programs = calc_total_programs() + if total_programs > num_vector_core_tile: + if len(self.configs) == 0: + num_vector_core_tile = max_programs_num + slow_decend_split = total_programs > num_vector_core_tile // 2 + if total_programs > num_vector_core_tile: + last_blocks = self.calcu_last_split_blocks(axis_idx) + if last_blocks != self.blocks[axis_idx].block_size: + self.blocks[axis_idx].block_size = last_blocks + self.candidate_blocks.append( + tuple([x.block_size for x in self.blocks]) + ) + break + + program_threshold = ( + self.tiny_program_threshold + if self.small_kernel + else num_vector_core // 2 + ) + if self.tiny_kernel and total_programs <= program_threshold: + self._try_add_tiny_low_program_config(total_programs) + if total_programs > program_threshold or self.dual_reduction: + if len(self.candidate_blocks) > 2: + self.candidate_blocks.pop(0) + self.candidate_blocks.append( + tuple([x.block_size for x in self.blocks]) + ) + if self.small_kernel: + self.add_to_configs( + list(tuple([x.block_size for x in self.blocks])) + ) + slow_decend_split = total_programs > num_vector_core_tile // 2 + + if not slow_decend_split: + self.blocks[axis_idx].block_size = (numel + 1) // 2 + else: + step = (numel + 3) // 4 if (numel + 3) // 4 > 1 else 1 + self.blocks[axis_idx].block_size = numel - step + self.blocks[axis_idx].sub_block_size = self.blocks[axis_idx].block_size + total_programs = calc_total_programs() + if self.blocks[axis_idx].block_size == 1 and ( + total_programs > program_threshold or self.dual_reduction + ): + self.candidate_blocks.append( + tuple([x.block_size for x in self.blocks]) + ) + else: + if numel >= 32: + self.blocks[axis_idx].sub_block_size = next_power_of_2(numel // 2) + else: + self.blocks[axis_idx].sub_block_size = numel - 1 + return reached_stop_numel + + def descend_all_low_dims(self): + low_dim_numels = [ + self.blocks[x.index].sub_block_size for x in self.kernel_meta.low_dims_axis + ] + if not low_dim_numels: + return False + + def descend_all_axis(min_numel): + + for axis in self.kernel_meta.low_dims_axis: + if axis.is_reduction and self.persistent_reduction: + continue + + numel = self.blocks[axis.index].sub_block_size + if numel == 1: + continue + if min_numel > 1 and abs(numel - min_numel) / min_numel < 0.2: + continue + if numel >= 128: + self.blocks[axis.index].sub_block_size = next_power_of_2(numel // 2) + else: + numel = self.blocks[axis.index].sub_block_size + numel = numel // 2 + self.blocks[axis.index].sub_block_size = min( + self.aligned_numel(numel), next_power_of_2(numel) + ) + + if len(self.candidate_blocks) == 0: + # means there is no split axis and tiling_not_low_dim axis + # so we need to init the candidates_blk_sizes + self.candidate_blocks.append(tuple([x.block_size for x in self.blocks])) + + count = 0 + tile_numel = self.calculate_tile_numel() + while tile_numel > self.stop_numel and count < 100: + count += 1 + tile_numel = self.calculate_tile_numel() + for candi_block in self.candidate_blocks: + self.add_to_configs(candi_block) + min_numel = min(low_dim_numels) + descend_all_axis(min_numel) + new_tile_numel = self.calculate_tile_numel() + if tile_numel == new_tile_numel: + descend_all_axis(0) + + return tile_numel < self.stop_numel + + def descend_split_tiling(self): + + tiling_not_low_dims = [ + x + for x in self.kernel_meta.tiling_axis + if x not in self.kernel_meta.low_dims_axis + ] + + def descend_split_axis(): + for axis in self.kernel_meta.tunable_split_axis: + if self.descend_one_axis(axis.index, is_split=True): + return True + + return self.calculate_tile_numel() <= self.stop_numel + + def descend_tiling_not_low_dims(): + for axis in tiling_not_low_dims: + if axis.is_reduction and self.persistent_reduction: + continue + + if self.descend_one_axis(axis.index): + return True + return self.calculate_tile_numel() <= self.stop_numel + + while True: + # descend split axis + if descend_split_axis(): + break + if len(self.candidate_blocks) > 0: + candi_block = self.candidate_blocks[0] + for i, blk_size in enumerate(candi_block): + self.blocks[i].sub_block_size = blk_size + # descend tiling but not low dims + if descend_tiling_not_low_dims(): + break + # descend low dims, need to descend all axis at the same time + self.descend_all_low_dims() + break diff --git a/backend/ascend_autotune_runtime/utils.py b/backend/ascend_autotune_runtime/utils.py new file mode 100644 index 00000000..83f09bf7 --- /dev/null +++ b/backend/ascend_autotune_runtime/utils.py @@ -0,0 +1,138 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved. +# 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. + +import torch + +_cached_params = None + + +def _init_npu_params(): + global _cached_params + if _cached_params is not None: + return _cached_params + + from triton.runtime.driver import driver + + target = driver.active.get_current_target() + device = driver.active.get_current_device() + prop = driver.active.utils.get_device_properties(device) + + num_cube_core = prop["num_aicore"] + num_vector_core = prop["num_aicore"] + ub_size_in_kbytes = 192 + rf_size_in_kbytes = None + + ASCEND_VARIANTS = ["Ascend910B", "Ascend910_93", "Ascend910_95", "Ascend950"] + if any(variant in target.arch for variant in ASCEND_VARIANTS): + num_vector_core = num_cube_core * 2 + + if target.arch.startswith("Ascend910_95") or target.arch.startswith("Ascend950"): + ub_size_in_kbytes = 256 + rf_size_in_kbytes = 128 + + _cached_params = { + "target": target, + "device": device, + "prop": prop, + "num_cube_core": num_cube_core, + "num_vector_core": num_vector_core, + "ub_size_in_kbytes": ub_size_in_kbytes, + "rf_size_in_kbytes": rf_size_in_kbytes, + } + return _cached_params + + +def __getattr__(name): + if name in [ + "target", + "device", + "prop", + "num_cube_core", + "num_vector_core", + "ub_size_in_kbytes", + "rf_size_in_kbytes", + ]: + return _init_npu_params()[name] + raise AttributeError(f"module '{__name__}' has no attribute '{name}'") + + +# wrapper npu 32 bytes align, get and pass unalign info to triton meta +# then autotune choose tiling param and send them to bishengIR +byte_per_numel = { + torch.float32: 4, # torch.float32 or torch.float + torch.float64: 8, # torch.float64 or torch.double + torch.float16: 2, # torch.float16 or torch.half + torch.bfloat16: 2, # torch.bfloat16 + torch.int32: 4, # torch.int32 or torch.int + torch.int64: 8, # torch.int64 or torch.long + torch.int16: 2, # torch.int16 or torch.short + torch.int8: 1, # torch.int8 + torch.uint8: 1, # torch.uint8 + torch.bool: 1, # torch.bool + torch.complex32: 4, # torch.complex32 (not yet available in PyTorch as of the latest stable release) + torch.complex64: 8, # torch.complex64 + torch.complex128: 16, # torch.complex128 +} + +# Some PyTorch versions expose extra fp8 dtypes. Register them when available. +for fp8_dtype_name in ( + "float8_e4m3fn", + "float8_e4m3fnuz", + "float8_e5m2", + "float8_e5m2fnuz", +): + fp8_dtype = getattr(torch, fp8_dtype_name, None) + if fp8_dtype is not None: + byte_per_numel[fp8_dtype] = 1 + +valid_axis_names = [ + "x", + "y", + "z", + "w", + "v", + "t", +] + + +def get_byte_per_numel(dtype: torch.dtype) -> int: + return 1 if dtype is None else byte_per_numel[dtype] + + +def is_valid_axis_name(name: str) -> bool: + if name.startswith("r"): + return name[1:] in valid_axis_names + return name in valid_axis_names + + +# move to an appropriate place, currently duplicated with triton.__init__.py +def next_power_of_2(n: int): + """Return the smallest power of 2 greater than or equal to n""" + n -= 1 + n |= n >> 1 + n |= n >> 2 + n |= n >> 4 + n |= n >> 8 + n |= n >> 16 + n |= n >> 32 + n += 1 + return n diff --git a/backend/commonir/adapter.py b/backend/commonir/adapter.py index ad3642dd..bc46c360 100644 --- a/backend/commonir/adapter.py +++ b/backend/commonir/adapter.py @@ -61,7 +61,7 @@ def get_kernel_source(self) -> str: @classmethod def compile_and_create_adapter(cls, tilelang_module): - if os.environ.get("DLC_DUMP_IR", "0") == "1": + if os.environ.get("TRITON_DEBUG", "0") == "1": with tempfile.TemporaryDirectory() as tmpdir: dst_path = os.path.join(tmpdir, "kernel.tilelangir.mlir") cls._write_mlir_file(dst_path, str(tilelang_module)) diff --git a/backend/commonir/backend.py b/backend/commonir/backend.py index 7f49170c..4e9efa5f 100644 --- a/backend/commonir/backend.py +++ b/backend/commonir/backend.py @@ -1,10 +1,62 @@ import functools import os +import re +import tempfile +from pathlib import Path from typing import Any + +from triton._C.libtriton import ir, dicp_triton, passes from ..compiler import DICPOptions from ..driver import DICPDriver from ..utils import get_current_backend +replace_commonir_ir = os.environ.get("DLC_REPLACE_COMMON_IR_FILE", None) +replace_commonir_linked_ir = os.environ.get("DLC_REPLACE_COMMONIR_LINKED_IR_FILE", None) + + +def add_matmul_input_precision(commonir: str) -> str: + # TODO: Temporary string-level patch for TileLang-generated CommonIR. + # TileLang should eventually use a pybind-based codegen path, similar to + # Triton's codegen, to emit linalg.matmul attributes directly. + return re.sub( + r"\blinalg\.matmul\s+(?!\{)", + 'linalg.matmul {input_precision = "ieee"} ', + commonir, + ) + + +def commonir_to_linkedir(commonir, metadata, opt, named_ops=True): + if replace_commonir_ir is not None: + print(f"[DEBUG] Replace common ir with {replace_commonir_ir}") + commonir = Path(replace_commonir_ir).read_text() + + commonir = add_matmul_input_precision(commonir) + + compile_on_910_95 = metadata["compile_on_910_95"] + enable_nd2nz_on_vector = metadata["enable_nd2nz_on_vector"] + enable_select_analysis = metadata["enable_select_analysis"] + + with tempfile.TemporaryDirectory() as tmpdir: + src_path = os.path.join(tmpdir, "kernel.commonir.mlir") + Path(src_path).write_text(commonir) + context = ir.context() + commonir_backend.load_dialects(context) + mod = ir.parse_mlir_module(src_path, context) + pm = ir.pass_manager(context) + dicp_triton.passes.commonir.add_vectorize_parallel_loop(pm) + passes.common.add_cse(pm) + passes.common.add_canonicalizer(pm) + dicp_triton.passes.commonir.add_annotate_kernel_attrs(pm) + dicp_triton.passes.ttir.add_ascend_npu_ir_legalize(pm, False) + pm.run(mod) + content = str(mod) + print(content) + + if replace_commonir_linked_ir is not None: + print(f"[DEBUG] Replace Linkedir with {replace_commonir_linked_ir}") + return Path(replace_commonir_linked_ir).read_text() + return content + class CommonIRBackend: binary_ext = "ttlinalgdir" @@ -12,9 +64,8 @@ class CommonIRBackend: def __init__(self) -> None: target = get_current_backend() self.driver = DICPDriver(target) - if self.driver.target == "dicp": - self.binary_ext = "ttlinalgdir" - elif self.driver.target == "mlu": + self.target = target + if self.driver.target == "mlu": self.capability = target.arch assert isinstance(self.capability, int) self.binary_ext = "cnbin" @@ -24,32 +75,31 @@ def __init__(self) -> None: elif self.driver.target == "ascend": self.binary_ext = "npubin" else: - raise RuntimeError(f"Target '{self.target_type}' is not supported.") - - def get_attrs_descriptor(self, params, args): - if self.driver.target == "ascend": - from triton.backends.dicp_triton.npu import AscendAttrsDescriptor - - return AscendAttrsDescriptor(params, args) - else: - raise RuntimeError( - f"backend {self.driver.target} not supported for get_attrs_descriptor." - ) + raise RuntimeError(f"Target '{self.driver.target}' is not supported.") def add_stages(self, stages, options, language=None): - if self.driver.target == "ascend": from triton.backends.dicp_triton.npu import ( - commonir_to_linkedir, - linalg_to_bin_enable_npu_compile, + linalg_to_bin_enable_npu_compile_910_95, + linalg_to_bin_enable_npu_compile_A2_A3, ) stages["linkedir"] = lambda src, metadata: commonir_to_linkedir( src, metadata, options, named_ops=True ) - stages["npubin"] = lambda src, metadata: linalg_to_bin_enable_npu_compile( - src, metadata, options - ) + + if options.compile_on_910_95: + stages["npubin"] = ( + lambda src, metadata: linalg_to_bin_enable_npu_compile_910_95( + src, metadata, options + ) + ) + else: + stages["npubin"] = ( + lambda src, metadata: linalg_to_bin_enable_npu_compile_A2_A3( + src, metadata, options + ) + ) else: raise RuntimeError("backend not supported") @@ -58,6 +108,13 @@ def load_dialects(self, ctx): from triton._C.libtriton import mlu mlu.load_dialects(ctx) + elif self.driver.target == "ascend": + from triton._C.libtriton import dicp_triton + + dicp_triton.load_dialects(ctx) + from triton._C.libtriton import dicp_triton + + dicp_triton.ir.load_dialects(ctx) return def get_driver(self): @@ -164,30 +221,17 @@ def get_codegen_implementation(self, options=None): def pack_metadata(self, metadata): if self.driver.target == "ascend": - from triton.backends.dicp_triton.npu import TRITON_PROFILER_REGISTERED - - # collect necessary metadata to launch kernels - # TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1 could set unique name. - # Get this name as the kernel_name to CANN runtime. - # kernel_name is unique to Ascend backend and should not be public. - # CANN runtime limits the length of kernel name <= 50. - # Considering '\n' is appended, thus the real kernel name <= 49. KERNEL_NAME_MAX_LEN = 49 - kernel_name_orig, mix_mode = metadata.name.split() + kernel_name_orig = metadata.kernel_name if len(kernel_name_orig) > KERNEL_NAME_MAX_LEN: kernel_name = kernel_name_orig[-KERNEL_NAME_MAX_LEN:] - # import warnings - # # red = "\x1b[31;20m" - # # reset = "\x1b[0m" - # warnings.warn(kernel_name_orig + " is truncated to " + kernel_name) - # warnings.warn("because '" + kernel_name_orig + "' exceeds torchnpu profiler's length limit < 50") else: kernel_name = kernel_name_orig return { "kernel_name": kernel_name, "hash": metadata.hash, "debug": metadata.debug, - "profiler_registered": TRITON_PROFILER_REGISTERED, + "tensor_kinds": metadata.tensor_kinds, } elif self.driver.target == "mlu": return (metadata.num_warps,) diff --git a/backend/commonir/compiler.py b/backend/commonir/compiler.py index 10e22d2d..7516b319 100644 --- a/backend/commonir/compiler.py +++ b/backend/commonir/compiler.py @@ -1,6 +1,8 @@ import functools import hashlib import json +import re +import os from pathlib import Path from typing import Any, List from triton._C.libtriton import get_cache_invalidating_env_vars @@ -72,7 +74,11 @@ def _init_handles(self): self.n_regs, self.n_spills, ) = commonir_backend.get_driver().utils.load_binary( - self.name, self.kernel, self.metadata.shared, device + self.name, + self.kernel, + self.metadata.shared, + device, + mix_mode=self.metadata.mix_mode, ) @property @@ -129,17 +135,34 @@ def runner(*args, stream=None): return runner -class CommonIRCompiler(object): +def _inject_npu_attrs(module: str, metadata: dict) -> str: + target = metadata.get("target") + arch = target.arch if target and hasattr(target, "arch") else "" + if arch: + module = re.sub( + r'(module\s+attributes\s*\{dicp\.backend\s*=\s*"ascend")', + rf'\1, hacc.target = #hacc.target<"{arch}">', + module, + count=1, + ) + return module + +class CommonIRCompiler(object): def compile(self, commonir_src: CommonIRSource, options=None, _env_vars=None): target = commonir_backend.get_driver().get_current_target() assert isinstance(target, GPUTarget), "target must be of GPUTarget type" extra_options = {} - options = commonir_backend.parse_options( - dict(options or dict(), **extra_options) - ) + options = dict(options or dict(), **extra_options) + if "debug" not in options and ( + os.environ.get("TRITON_DEBUG", "0") == "1" + or os.environ.get("DEBUG", "0") == "1" + ): + options["debug"] = True + + options = commonir_backend.parse_options(options) # create cache manager env_vars = get_cache_invalidating_env_vars() if _env_vars is None else _env_vars @@ -162,6 +185,9 @@ def compile(self, commonir_src: CommonIRSource, options=None, _env_vars=None): stages = dict() commonir_backend.add_stages(stages, options) module = commonir_src.src + module = _inject_npu_attrs(module, metadata) + print(module) + ir_filename = f"{file_name}.source" metadata_group[ir_filename] = fn_cache_manager.put(module, ir_filename) diff --git a/backend/compiler.py b/backend/compiler.py index b5d0397f..0390763e 100644 --- a/backend/compiler.py +++ b/backend/compiler.py @@ -27,21 +27,13 @@ def _get_llvm_bin_path(bin_name: str) -> str: return os.path.join(path, bin_name) -def _get_triton_linalg_opt_path() -> str: - # path = os.getenv("TRITON_LINALG_OPT_PATH", "") - path = "triton-shared-opt" - if path == "": - raise Exception("TRITON_SHARED_OPT_PATH is not set.") - return path - - def _ttir_to_linalgdir(mod): ttir_code = str(mod) with tempfile.TemporaryDirectory() as tmpdir: src_path = os.path.join(tmpdir, "tt.mlir") dst_path = os.path.join(tmpdir, "triton_linalg.mlir") Path(src_path).write_text(ttir_code) - triton_linalg_opt_path = _get_triton_linalg_opt_path() + triton_linalg_opt_path = _get_dicp_triton_opt_path() subprocess.check_call( [triton_linalg_opt_path, src_path, "--triton-to-linalg", "-o", dst_path] ) @@ -125,8 +117,8 @@ def __init__(self, target: str) -> None: self._cpu_backend = CPUBackend(target) self.binary_ext = "obj" - elif self.driver.target == "dicp": - self.binary_ext = "ttlinalgdir" + elif self.driver.target == "ascend": + self.binary_ext = "npubin" elif self.driver.target == "mlu": self.capability = target.arch assert isinstance(self.capability, int) @@ -134,14 +126,12 @@ def __init__(self, target: str) -> None: elif self.driver.target == "maca": self.capability = 80 self.binary_ext = "mcfatbin" - elif self.driver.target == "ascend": - self.binary_ext = "npubin" else: - raise RuntimeError(f"Target '{self.target_type}' is not supported.") + raise RuntimeError(f"Target '{self.driver.target}' is not supported.") @staticmethod def supports_target(target: GPUTarget): - return target.backend in ["dicp", "mlu", "maca", "ascend", "cpu"] + return target.backend in ["ascend", "mlu", "maca", "cpu"] @staticmethod def make_ttir(mod, metadata, opt): @@ -159,26 +149,39 @@ def make_ttir(mod, metadata, opt): metadata["shared"] = 0 return mod - def get_attrs_descriptor(self, params, args): + 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": - from triton.backends.dicp_triton.npu import AscendAttrsDescriptor - - return AscendAttrsDescriptor(params, args) - else: - raise RuntimeError( - f"backend {self.driver.target} not supported for get_attrs_descriptor." + from triton.backends.dicp_triton.npu import ( + make_ttir, + ttir_to_linalg_dicp, + linalg_to_bin_enable_npu_compile_910_95, + linalg_to_bin_enable_npu_compile_A2_A3, + ttir_to_npubin, ) - def add_stages(self, stages, options, language=None): - if self.driver.target not in ["ascend", "mlu"]: - stages["ttir"] = lambda src, metadata: self.make_ttir( - src, metadata, options - ) - if self.driver.target == "dicp": - stages["ttlinalgdir"] = lambda src, metadata: _optimize_ttlinalgdir( - _ttir_to_linalgdir(src) + stages["ttir"] = lambda src, metadata: make_ttir(src, metadata, options) + if options.force_simt_only: + stages["npubin"] = lambda src, metadata: ttir_to_npubin( + src, metadata, options + ) + return + stages["dicp"] = lambda src, metadata: ttir_to_linalg_dicp( + src, metadata, options, named_ops=True ) - stages["fatbin"] = lambda src, metadata: _linalg_to_fatbin(src, metadata) + if options.compile_on_910_95: + stages["npubin"] = ( + lambda src, metadata: linalg_to_bin_enable_npu_compile_910_95( + src, metadata, options + ) + ) + else: + stages["npubin"] = ( + lambda src, metadata: linalg_to_bin_enable_npu_compile_A2_A3( + src, metadata, options + ) + ) elif self.driver.target == "mlu": from triton.backends.dicp_triton.mlu import ( onchip_mem_analysis, @@ -236,73 +239,22 @@ def add_stages(self, stages, options, language=None): stages["mcfatbin"] = lambda src, metadata: make_mcfatbin( src, metadata, options, self.capability ) - elif self.driver.target == "ascend": - from triton.backends.dicp_triton.npu import ( - make_ttir, - ttir_to_linalg, - ttir_to_ttsharedir_ascend, - ttsharedir_to_linkedir, - linalg_to_bin_enable_npu_compile, - ) - - stages["ttir"] = lambda src, metadata: make_ttir(src, metadata, options) - lower_by_ttshared = os.getenv("LOWER_BY_TTSHARED", "1") - if lower_by_ttshared == "0": - if options.enable_npu_compile: - stages["ttadapter"] = lambda src, metadata: ttir_to_linalg( - src, metadata, options, named_ops=True - ) - stages["npubin"] = ( - lambda src, metadata: linalg_to_bin_enable_npu_compile( - src, metadata, options - ) - ) - else: - if options.enable_npu_compile: - stages["ttshared"] = ( - lambda src, metadata: ttir_to_ttsharedir_ascend( - src, metadata, options, named_ops=True - ) - ) - stages["linkedir"] = lambda src, metadata: ttsharedir_to_linkedir( - src, - metadata, - options, - named_ops=True, - cpu_verify=self.driver.is_cpu_verify, - ) - if self.driver.is_cpu_verify: - from .cpu_backend import ( - _ttsharedir_to_llir, - _llir_to_bin, - _optimize_llir, - ) - - stages["llir"] = lambda src, metadata: _optimize_llir( - _ttsharedir_to_llir(src, metadata) - ) - stages["obj"] = lambda src, metadata: _llir_to_bin( - src, metadata - ) - else: - stages["npubin"] = ( - lambda src, metadata: linalg_to_bin_enable_npu_compile( - src, metadata, options - ) - ) else: - raise RuntimeError("backend not supported") + raise RuntimeError(f"backend {self.driver.target} not supported") def load_dialects(self, ctx): + if self.driver.is_cpu_verify: + return self._cpu_backend.load_dialects(ctx) if self.driver.target == "mlu": from triton._C.libtriton import mlu mlu.load_dialects(ctx) - return + elif self.driver.target == "ascend": + from triton._C.libtriton import dicp_triton - @functools.lru_cache() - def hash(self): - return self.target + dicp_triton.load_dialects(ctx) + dicp_triton.ir.load_dialects(ctx) + return def get_driver(self): return self.driver @@ -319,6 +271,7 @@ def parse_options(self, options: dict) -> Any: for k in NPUOptions.__dataclass_fields__.keys() if k in options } + args.setdefault("arch", self.target.arch) options = NPUOptions(**args) return options elif self.target.backend == "mlu": @@ -414,30 +367,18 @@ def pack_metadata(self, metadata): if self.driver.is_cpu_verify: return self._cpu_backend.pack_metadata(metadata) if self.target.backend == "ascend": - from triton.backends.dicp_triton.npu import TRITON_PROFILER_REGISTERED - - # collect necessary metadata to launch kernels - # TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1 could set unique name. - # Get this name as the kernel_name to CANN runtime. - # kernel_name is unique to Ascend backend and should not be public. - # CANN runtime limits the length of kernel name <= 50. - # Considering '\n' is appended, thus the real kernel name <= 49. + KERNEL_NAME_MAX_LEN = 49 - kernel_name_orig, mix_mode = metadata.name.split() + kernel_name_orig = metadata.kernel_name if len(kernel_name_orig) > KERNEL_NAME_MAX_LEN: kernel_name = kernel_name_orig[-KERNEL_NAME_MAX_LEN:] - # import warnings - # # red = "\x1b[31;20m" - # # reset = "\x1b[0m" - # warnings.warn(kernel_name_orig + " is truncated to " + kernel_name) - # warnings.warn("because '" + kernel_name_orig + "' exceeds torchnpu profiler's length limit < 50") else: kernel_name = kernel_name_orig return { "kernel_name": kernel_name, "hash": metadata.hash, "debug": metadata.debug, - "profiler_registered": TRITON_PROFILER_REGISTERED, + "tensor_kinds": metadata.tensor_kinds, } elif self.target.backend == "mlu": return (metadata.num_warps,) @@ -452,6 +393,8 @@ def pack_metadata(self, metadata): @functools.lru_cache() def hash(self): + if self.driver.is_cpu_verify: + return self._cpu_backend.hash() if self.target.backend == "mlu": from triton.backends.dicp_triton.mlu import get_cnas_version diff --git a/backend/cpu_backend.py b/backend/cpu_backend.py index c8a1a80b..7902b043 100644 --- a/backend/cpu_backend.py +++ b/backend/cpu_backend.py @@ -1,5 +1,4 @@ # CPU Backend for verification -# Merged from triton_shared backend/compiler.py and backend/driver.py from triton.backends.compiler import BaseBackend, GPUTarget from triton._C.libtriton import ir, passes @@ -23,13 +22,11 @@ import platform import triton.backends.dicp_triton.utils as dicp_utils -dump_ir = os.environ.get("DLC_DUMP_IR", "0") == "1" - -def _get_triton_shared_opt_path() -> str: - path = os.getenv("TRITON_SHARED_OPT_PATH", "") +def _get_dicp_triton_opt_path() -> str: + path = os.getenv("DICP_TRITON_OPT_PATH", "") if path == "": - raise Exception("TRITON_SHARED_OPT_PATH is not set.") + raise Exception("DICP_TRITON_OPT_PATH is not set.") return path @@ -41,7 +38,7 @@ def _get_llvm_bin_path(bin_name: str) -> str: def _dump_ir_if_needed(files): - path = os.getenv("TRITON_SHARED_DUMP_PATH", "") + path = os.getenv("DLC_DUMP_PATH", "") if not path: return for f in files: @@ -49,21 +46,21 @@ def _dump_ir_if_needed(files): def _get_sanitizer_type(): - sanitizer_type = os.getenv("TRITON_SHARED_SANITIZER_TYPE", "") + sanitizer_type = os.getenv("DLC_SANITIZER_TYPE", "") if sanitizer_type != "" and sanitizer_type != "asan" and sanitizer_type != "tsan": - raise Exception(f"TRITON_SHARED_SANITIZER_TYPE {sanitizer_type} is invalid.") + raise Exception(f"DLC_SANITIZER_TYPE {sanitizer_type} is invalid.") return sanitizer_type -def _ttir_to_ttsharedir(mod, metadata): +def _ttir_to_linalg(mod, metadata, opt): ttir_code = str(mod) with tempfile.TemporaryDirectory() as tmpdir: src_path = os.path.join(tmpdir, "tt.mlir") - dst_path = os.path.join(tmpdir, "ttshared.mlir") + dst_path = os.path.join(tmpdir, "linalg.mlir") Path(src_path).write_text(ttir_code) - triton_shared_opt_path = _get_triton_shared_opt_path() + triton_opt_path = _get_dicp_triton_opt_path() subprocess_args = [ - triton_shared_opt_path, + triton_opt_path, src_path, "--triton-to-linalg-experimental", "--mlir-print-debuginfo", @@ -75,28 +72,26 @@ def _ttir_to_ttsharedir(mod, metadata): subprocess_args.insert(2, "--add-llvm-debug-info") subprocess.check_call(subprocess_args) result = Path(dst_path).read_text() - if dump_ir: - dicp_utils._dump_stage_ir( - result, metadata["hash"], "kernel.ttsharedir.mlir" - ) + if opt.debug: + dicp_utils._dump_stage_ir(result, metadata["hash"], "kernel.linalg.mlir") return result -def _optimize_ttsharedir(ttsharedir: str): - return ttsharedir +def _optimize_linalg(linalg_ir: str): + return linalg_ir -def _ttsharedir_to_llir(ttsharedir: str, metadata): +def _linalg_to_llir(linalg_ir: str, metadata, opt): with tempfile.TemporaryDirectory() as tmpdir: - ttshared_path = os.path.join(tmpdir, "ttshared.mlir") + linalg_path = os.path.join(tmpdir, "linalg.mlir") llmlir_path = os.path.join(tmpdir, "ll.mlir") llir_path = os.path.join(tmpdir, "ll.ir") - Path(ttshared_path).write_text(ttsharedir) + Path(linalg_path).write_text(linalg_ir) mlir_opt_path = _get_llvm_bin_path("mlir-opt") subprocess.check_call( [ mlir_opt_path, - ttshared_path, + linalg_path, "--convert-elementwise-to-linalg", "--convert-linalg-to-affine-loops", "--empty-tensor-to-alloc-tensor", @@ -127,7 +122,7 @@ def _ttsharedir_to_llir(ttsharedir: str, metadata): [mlir_translate_path, llmlir_path, "--mlir-to-llvmir", "-o", llir_path] ) result = Path(llir_path).read_text() - if dump_ir: + if opt.debug: dicp_utils._dump_stage_ir(result, metadata["hash"], "kernel.llir.mlir") return result @@ -310,23 +305,23 @@ def make_ttir(mod, metadata, options): passes.ttir.add_loop_unroll(pm) passes.common.add_cse(pm) pm.run(mod) - if dump_ir: + if options.debug: dicp_utils._dump_stage_ir(str(mod), metadata["hash"], "kernel.ttir.mlir") return mod def add_stages(self, stages, options, language): stages["ttir"] = lambda src, metadata: self.make_ttir(src, metadata, options) - stages["ttsharedir"] = lambda src, metadata: _optimize_ttsharedir( - _ttir_to_ttsharedir(src, metadata) + stages["linalg"] = lambda src, metadata: _optimize_linalg( + _ttir_to_linalg(src, metadata, options) ) stages["llir"] = lambda src, metadata: _optimize_llir( - _ttsharedir_to_llir(src, metadata) + _linalg_to_llir(src, metadata, options) ) stages["obj"] = lambda src, metadata: _llir_to_bin(src, metadata) @functools.lru_cache() def hash(self): - return self.target + return str(self.target) def get_module_map(self) -> Dict[str, ModuleType]: return {} @@ -413,8 +408,8 @@ def _generate_launcher(constants, signature, kernel_name): #include #include #include -#include "ExecutionEngine/CRunnerUtils.h" -#include "ExecutionEngine/CRunnerUtils.cpp" +#include "CRunnerUtils.h" +#include "CRunnerUtils.cpp" extern "C" {{ // Pointer type (=Memref) becomes int64_t + MemRef struct @@ -539,13 +534,13 @@ def _generate_launcher(constants, signature, kernel_name): static struct PyModuleDef ModuleDef = {{ PyModuleDef_HEAD_INIT, - \"__triton_shared_ref_cpu_kernel_launcher\", + \"__dicp_cpu_kernel_launcher\", NULL, //documentation -1, //size ModuleMethods }}; -PyMODINIT_FUNC PyInit___triton_shared_ref_cpu_kernel_launcher(void) {{ +PyMODINIT_FUNC PyInit___dicp_cpu_kernel_launcher(void) {{ PyObject *m = PyModule_Create(&ModuleDef); if(m == NULL) {{ return NULL; @@ -575,7 +570,7 @@ def compile_module(launcher_src, kernel_placeholder_name): name="python", major=py_version.major, minor=py_version.minor ) cpu_backend_path = Path(__file__).resolve().parent - include_dir = os.path.join(cpu_backend_path, "include") + cpu_verify_dir = os.path.join(cpu_backend_path, "cpu_verify") def launch( gridX, @@ -594,7 +589,7 @@ def launch( src = launcher_src.replace(kernel_placeholder_name, kernel_name) key = hashlib.sha256(src.encode("utf-8") + kernel_obj).hexdigest() cache = get_cache_manager(key) - name = "__triton_shared_ref_cpu_kernel_launcher" + name = "__dicp_cpu_kernel_launcher" if platform.system() == "Windows": filename = f"{name}.pyd" else: @@ -606,7 +601,7 @@ def launch( if platform.system() == "Windows": if sanitizer_type != "": raise Exception( - "Sanitizers are not supported on Windows with triton-shared." + "Sanitizers are not supported on Windows with DLC." ) obj_path = os.path.join(tmpdir, "kernel.obj") launcher_src_path = os.path.join(tmpdir, "main.cxx") @@ -621,7 +616,7 @@ def launch( launcher_src_path, obj_path, f"-I{py_include_dir}", - f"-I{include_dir}", + f"-I{cpu_verify_dir}", "/link", f"/LIBPATH:{py_lib_dir}", "/link", @@ -643,7 +638,7 @@ def launch( launcher_src_path, obj_path, f"-I{py_include_dir}", - f"-I{include_dir}", + f"-I{cpu_verify_dir}", f"-L{py_lib_dir}", "-shared", f"-l{py_lib}", @@ -666,7 +661,7 @@ def launch( launcher_src_path, obj_path, f"-I{py_include_dir}", - f"-I{include_dir}", + f"-I{cpu_verify_dir}", f"-L{py_lib_dir}", "-shared", f"-l{py_lib}", diff --git a/backend/include/ExecutionEngine/CRunnerUtils.cpp b/backend/cpu_verify/CRunnerUtils.cpp similarity index 83% rename from backend/include/ExecutionEngine/CRunnerUtils.cpp rename to backend/cpu_verify/CRunnerUtils.cpp index 87e47027..7e780c9c 100644 --- a/backend/include/ExecutionEngine/CRunnerUtils.cpp +++ b/backend/cpu_verify/CRunnerUtils.cpp @@ -41,10 +41,6 @@ template void stdSort(uint64_t n, V *p) { std::sort(p, p + n); } } // namespace -// Small runtime support "lib" for vector.print lowering. -// By providing elementary printing methods only, this -// library can remain fully unaware of low-level implementation -// details of our vectors. Also useful for direct LLVM IR output. extern "C" void printI64(int64_t i) { fprintf(stdout, "%" PRId64, i); } extern "C" void printU64(uint64_t u) { fprintf(stdout, "%" PRIu64, u); } extern "C" void printF32(float f) { fprintf(stdout, "%g", f); } @@ -63,7 +59,6 @@ extern "C" void memrefCopy(int64_t elemSize, UnrankedMemRefType *srcArg, int64_t rank = src.rank; MLIR_MSAN_MEMORY_IS_INITIALIZED(src.sizes, rank * sizeof(int64_t)); - // Handle empty shapes -> nothing to copy. for (int rankp = 0; rankp < rank; ++rankp) if (src.sizes[rankp] == 0) return; @@ -80,7 +75,6 @@ extern "C" void memrefCopy(int64_t elemSize, UnrankedMemRefType *srcArg, int64_t *srcStrides = static_cast(alloca(sizeof(int64_t) * rank)); int64_t *dstStrides = static_cast(alloca(sizeof(int64_t) * rank)); - // Initialize index and scale strides. for (int rankp = 0; rankp < rank; ++rankp) { indices[rankp] = 0; srcStrides[rankp] = src.strides[rankp] * elemSize; @@ -89,22 +83,15 @@ extern "C" void memrefCopy(int64_t elemSize, UnrankedMemRefType *srcArg, int64_t readIndex = 0, writeIndex = 0; for (;;) { - // Copy over the element, byte by byte. memcpy(dstPtr + writeIndex, srcPtr + readIndex, elemSize); - // Advance index and read position. for (int64_t axis = rank - 1; axis >= 0; --axis) { - // Advance at current axis. auto newIndex = ++indices[axis]; readIndex += srcStrides[axis]; writeIndex += dstStrides[axis]; - // If this is a valid index, we have our next index, so continue copying. if (src.sizes[axis] != newIndex) break; - // We reached the end of this axis. If this is axis 0, we are done. if (axis == 0) return; - // Else, reset to 0 and undo the advancement of the linear index that - // this axis had. Then continue with the axis one outer. indices[axis] = 0; readIndex -= src.sizes[axis] * srcStrides[axis]; writeIndex -= dst.sizes[axis] * dstStrides[axis]; @@ -112,12 +99,10 @@ extern "C" void memrefCopy(int64_t elemSize, UnrankedMemRefType *srcArg, } } -/// Prints GFLOPS rating. extern "C" void printFlops(double flops) { fprintf(stderr, "%lf GFLOPS\n", flops / 1.0E9); } -/// Returns the number of seconds since Epoch 1970-01-01 00:00:00 +0000 (UTC). extern "C" double rtclock() { #ifndef _WIN32 struct timeval tp; @@ -137,8 +122,6 @@ extern "C" void *mlirAlignedAlloc(uint64_t alignment, uint64_t size) { #ifdef _WIN32 return _aligned_malloc(size, alignment); #elif defined(__APPLE__) - // aligned_alloc was added in MacOS 10.15. Fall back to posix_memalign to also - // support older versions. void *result = nullptr; (void)::posix_memalign(&result, alignment, size); return result; @@ -157,10 +140,7 @@ extern "C" void mlirAlignedFree(void *ptr) { #endif } -extern "C" void *rtsrand(uint64_t s) { - // Standard mersenne_twister_engine seeded with s. - return new std::mt19937(s); -} +extern "C" void *rtsrand(uint64_t s) { return new std::mt19937(s); } extern "C" uint64_t rtrand(void *g, uint64_t m) { std::mt19937 *generator = static_cast(g); diff --git a/backend/include/ExecutionEngine/CRunnerUtils.h b/backend/cpu_verify/CRunnerUtils.h similarity index 75% rename from backend/include/ExecutionEngine/CRunnerUtils.h rename to backend/cpu_verify/CRunnerUtils.h index 1e55ca92..bc114132 100644 --- a/backend/include/ExecutionEngine/CRunnerUtils.h +++ b/backend/cpu_verify/CRunnerUtils.h @@ -18,19 +18,16 @@ #ifdef _WIN32 #ifndef MLIR_CRUNNERUTILS_EXPORT #ifdef mlir_c_runner_utils_EXPORTS -// We are building this library #define MLIR_CRUNNERUTILS_EXPORT __declspec(dllexport) #define MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS #else -// We are using this library #define MLIR_CRUNNERUTILS_EXPORT __declspec(dllimport) -#endif // mlir_c_runner_utils_EXPORTS -#endif // MLIR_CRUNNERUTILS_EXPORT -#else // _WIN32 -// Non-windows: use visibility attributes. +#endif +#endif +#else #define MLIR_CRUNNERUTILS_EXPORT __attribute__((visibility("default"))) #define MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS -#endif // _WIN32 +#endif #include #include @@ -38,9 +35,6 @@ #include #include -//===----------------------------------------------------------------------===// -// Codegen-compatible structures for Vector type. -//===----------------------------------------------------------------------===// namespace mlir { namespace detail { @@ -64,8 +58,6 @@ template struct Vector1D { T vector[Dim]; }; -// 1-D vector, padded to the next power of 2 allocation. -// Specialization occurs to avoid zero size arrays (which fail in -Werror). template struct Vector1D { Vector1D() { static_assert(nextPowerOf2(sizeof(T[Dim])) > sizeof(T[Dim]), "size error"); @@ -82,7 +74,6 @@ template struct Vector1D { } // namespace detail } // namespace mlir -// N-D vectors recurse down to 1-D. template struct Vector { inline Vector &operator[](unsigned i) { return vector[i]; } inline const Vector &operator[](unsigned i) const { @@ -93,8 +84,6 @@ template struct Vector { Vector vector[Dim]; }; -// 1-D vectors in LLVM are automatically padded to the next power of 2. -// We insert explicit padding in to account for this. template struct Vector : public mlir::detail::Vector1D void dropFront(int64_t arr[N], int64_t *res) { *(res + i - 1) = arr[i]; } -//===----------------------------------------------------------------------===// -// Codegen-compatible structures for StridedMemRef type. -//===----------------------------------------------------------------------===// template class StridedMemrefIterator; -/// StridedMemRef descriptor type with static rank. template struct StridedMemRefType { T *basePtr; T *data; @@ -143,7 +128,6 @@ template struct StridedMemRefType { StridedMemrefIterator begin() { return {*this, offset}; } StridedMemrefIterator end() { return {*this, -1}; } - // This operator[] is extremely slow and only for sugaring purposes. StridedMemRefType operator[](int64_t idx) { StridedMemRefType res; res.basePtr = basePtr; @@ -155,7 +139,6 @@ template struct StridedMemRefType { } }; -/// StridedMemRef descriptor type specialized for rank 1. template struct StridedMemRefType { T *basePtr; T *data; @@ -177,7 +160,6 @@ template struct StridedMemRefType { T &operator[](int64_t idx) { return *(data + offset + idx * strides[0]); } }; -/// StridedMemRef descriptor type specialized for rank 0. template struct StridedMemRefType { T *basePtr; T *data; @@ -195,7 +177,6 @@ template struct StridedMemRefType { StridedMemrefIterator end() { return {*this, offset + 1}; } }; -/// Iterate over all elements in a strided memref. template class StridedMemrefIterator { public: using iterator_category = std::forward_iterator_tag; @@ -237,18 +218,11 @@ template class StridedMemrefIterator { } private: - /// Offset in the buffer. This can be derived from the indices and the - /// descriptor. int64_t offset = 0; - - /// Array of indices in the multi-dimensional memref. std::array indices = {}; - - /// Descriptor for the strided memref. StridedMemRefType *descriptor; }; -/// Iterate over all elements in a 0-ranked strided memref. template class StridedMemrefIterator { public: using iterator_category = std::forward_iterator_tag; @@ -268,11 +242,7 @@ template class StridedMemrefIterator { reference operator*() { return *elt; } pointer operator->() { return elt; } - // There are no indices for a 0-ranked memref, but this API is provided for - // consistency with the general case. const std::array &getIndices() { - // Since this is a 0-array of indices we can keep a single global const - // copy. static const std::array indices = {}; return indices; } @@ -286,25 +256,16 @@ template class StridedMemrefIterator { } private: - /// Pointer to the single element in the zero-ranked memref. T *elt; }; -//===----------------------------------------------------------------------===// -// Codegen-compatible structure for UnrankedMemRef type. -//===----------------------------------------------------------------------===// -// Unranked MemRef template struct UnrankedMemRefType { int64_t rank; void *descriptor; }; -//===----------------------------------------------------------------------===// -// DynamicMemRefType type. -//===----------------------------------------------------------------------===// template class DynamicMemRefIterator; -// A reference to one of the StridedMemRef types. template class DynamicMemRefType { public: int64_t rank; @@ -351,7 +312,6 @@ template class DynamicMemRefType { DynamicMemRefIterator begin() { return {*this, offset}; } DynamicMemRefIterator end() { return {*this, -1}; } - // This operator[] is extremely slow and only for sugaring purposes. DynamicMemRefType operator[](int64_t idx) { assert(rank > 0 && "can't make a subscript of a zero ranked array"); @@ -363,15 +323,12 @@ template class DynamicMemRefType { return res; } - // This operator* can be used in conjunction with the previous operator[] in - // order to access the underlying value in case of zero-ranked memref. T &operator*() { assert(rank == 0 && "not a zero-ranked memRef"); return data[offset]; } }; -/// Iterate over all elements in a dynamic memref. template class DynamicMemRefIterator { public: using iterator_category = std::forward_iterator_tag; @@ -423,27 +380,15 @@ template class DynamicMemRefIterator { } private: - /// Offset in the buffer. This can be derived from the indices and the - /// descriptor. int64_t offset = 0; - - /// Array of indices in the multi-dimensional memref. std::vector indices = {}; - - /// Descriptor for the dynamic memref. DynamicMemRefType *descriptor; }; -//===----------------------------------------------------------------------===// -// Small runtime support library for memref.copy lowering during codegen. -//===----------------------------------------------------------------------===// extern "C" MLIR_CRUNNERUTILS_EXPORT void memrefCopy(int64_t elemSize, ::UnrankedMemRefType *src, ::UnrankedMemRefType *dst); -//===----------------------------------------------------------------------===// -// Small runtime support library for vector.print lowering during codegen. -//===----------------------------------------------------------------------===// extern "C" MLIR_CRUNNERUTILS_EXPORT void printI64(int64_t i); extern "C" MLIR_CRUNNERUTILS_EXPORT void printU64(uint64_t u); extern "C" MLIR_CRUNNERUTILS_EXPORT void printF32(float f); @@ -454,25 +399,13 @@ extern "C" MLIR_CRUNNERUTILS_EXPORT void printClose(); extern "C" MLIR_CRUNNERUTILS_EXPORT void printComma(); extern "C" MLIR_CRUNNERUTILS_EXPORT void printNewline(); -//===----------------------------------------------------------------------===// -// Small runtime support library for timing execution and printing GFLOPS -//===----------------------------------------------------------------------===// extern "C" MLIR_CRUNNERUTILS_EXPORT void printFlops(double flops); extern "C" MLIR_CRUNNERUTILS_EXPORT double rtclock(); -//===----------------------------------------------------------------------===// -// Runtime support library for random number generation. -//===----------------------------------------------------------------------===// -// Uses a seed to initialize a random generator and returns the generator. extern "C" MLIR_CRUNNERUTILS_EXPORT void *rtsrand(uint64_t s); -// Returns a random number in the range of [0, m). extern "C" MLIR_CRUNNERUTILS_EXPORT uint64_t rtrand(void *, uint64_t m); -// Deletes the random number generator. extern "C" MLIR_CRUNNERUTILS_EXPORT void rtdrand(void *); -//===----------------------------------------------------------------------===// -// Runtime support library to allow the use of std::sort in MLIR program. -//===----------------------------------------------------------------------===// extern "C" MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_stdSortI64(uint64_t n, StridedMemRefType *vref); extern "C" MLIR_CRUNNERUTILS_EXPORT void diff --git a/backend/include/ExecutionEngine/Msan.h b/backend/cpu_verify/Msan.h similarity index 100% rename from backend/include/ExecutionEngine/Msan.h rename to backend/cpu_verify/Msan.h diff --git a/backend/driver.py b/backend/driver.py index 4bf69b20..098b75ce 100644 --- a/backend/driver.py +++ b/backend/driver.py @@ -131,7 +131,13 @@ def __init__(self, target=None): super().__init__() self.is_cpu_verify = os.environ.get("DLC_CPU_VERIFY", "0") == "1" - if target == "mlu": + # Extract backend name from GPUTarget or string + if hasattr(target, "backend"): + backend = target.backend + else: + backend = str(target) if target else get_current_backend() + + if backend == "mlu": from triton.backends.dicp_triton.mlu import BangLauncher, BangUtils self.target = "mlu" @@ -148,19 +154,23 @@ def __init__(self, target=None): self.is_linear_pointer = lambda ptr, device: self.utils.is_linear_pointer( ptr, device ) - elif target == "maca": + elif backend == "maca": from triton.backends.dicp_triton.maca import MacaLauncher, MacaUtils self.target = "maca" self.utils = MacaUtils() self.launcher_cls = MacaLauncher - elif target == "ascend": - from triton.backends.dicp_triton.npu import NPULauncher, NPUUtils + elif backend == "ascend": + from triton.backends.dicp_triton.npu_driver import NPULauncher, NPUUtils self.target = "ascend" self.utils = NPUUtils() self.launcher_cls = NPULauncher - elif target == "nvidia": + + from .ascend_autotune_hooks import hook_autotune_for_ascend + + hook_autotune_for_ascend() + elif backend == "nvidia": from triton.backends.nvidia.driver import CudaLauncher, CudaUtils self.target = "nvidia" @@ -192,7 +202,7 @@ def is_active(self): if current_backend == "ascend": def test_npucompiler(): - from triton.backends.dicp_triton.npu import _get_bisheng_path + from triton.backends.dicp_triton.utils import _get_bisheng_path npucompiler = _get_bisheng_path() targets = ( diff --git a/backend/include/ExecutionEngine/version.txt b/backend/include/ExecutionEngine/version.txt deleted file mode 100644 index c3f15e55..00000000 --- a/backend/include/ExecutionEngine/version.txt +++ /dev/null @@ -1 +0,0 @@ -https://github.com/llvm/llvm-project/commit/3be3883e6d67bf908fd12b51219075293ebb3dff diff --git a/backend/lib/libdevice.10.bc b/backend/lib/libdevice.10.bc new file mode 100644 index 00000000..3c01c89b Binary files /dev/null and b/backend/lib/libdevice.10.bc differ diff --git a/backend/npu.py b/backend/npu.py index 0c89c723..758bc936 100644 --- a/backend/npu.py +++ b/backend/npu.py @@ -5,52 +5,77 @@ import sysconfig import functools import hashlib +import logging from triton.runtime.cache import get_cache_manager, get_dump_manager from triton.backends.compiler import GPUTarget from triton._C.libtriton import ir, passes, dicp_triton -from triton.runtime.cache import get_dump_manager +from triton.compiler.errors import CompileTimeAssertionFailure import triton.backends.dicp_triton.utils as dicp_utils from dataclasses import dataclass from typing import Any, Union, Tuple, Dict import ctypes import re -import pybind11 -import shutil - -###################### utils.py start ###################### - -TRITON_PROFILER_REGISTERED = False -dump_ir = os.environ.get("DLC_DUMP_IR", "0") == "1" -replace_ttshared_ir = os.environ.get("DLC_REPLACE_TTSHARED_IR_FILE", None) -replace_linked_ir = os.environ.get("DLC_REPLACE_LINKED_IR_FILE", None) -replace_commonir_ir = os.environ.get("DLC_REPLACE_COMMON_IR_FILE", None) -replace_commonir_linked_ir = os.environ.get("DLC_REPLACE_COMMONIR_LINKED_IR_FILE", None) -if ( - dump_ir - or (replace_ttshared_ir is not None) - or (replace_linked_ir is not None) - or (replace_commonir_linked_ir is not None) -): - os.environ["TRITON_ALWAYS_COMPILE"] = "1" - dump_dir = "./tmp" - os.environ["TRITON_DUMP_DIR"] = os.environ.get("TRITON_DUMP_DIR", dump_dir) - if os.path.exists(dump_dir): - print(f"Directory **{dump_dir}** exists. Deleting the entire directory...") - shutil.rmtree(dump_dir) - -local_bishengir_path = os.path.join(os.path.dirname(__file__), "../../_C/bishengir") -bisheng_install_path = os.environ.get("BISHENG_INSTALL_PATH", None) -if ( - bisheng_install_path is None - and os.path.exists(local_bishengir_path) - and os.path.isdir(local_bishengir_path) - and os.path.exists(os.path.join(local_bishengir_path, "bishengir-compile")) - and os.path.exists(os.path.join(local_bishengir_path, "bishengir-hivm-compile")) - and os.path.exists(os.path.join(local_bishengir_path, "bishengir-opt")) - and os.path.exists(os.path.join(local_bishengir_path, "hivmc")) -): - os.environ["BISHENG_INSTALL_PATH"] = local_bishengir_path - os.environ["PATH"] = local_bishengir_path + os.pathsep + os.environ["PATH"] +import sys + +from .utils import ( + TRITON_PROFILER_REGISTERED, + replace_dicp_ir, + _get_npucompiler_path, + _get_bisheng_path, + _get_ascend_path, + _get_bishengir_opt_path, + _is_ascend_sanitizer_enabled, + _is_auto_map_parallel_blocks_enabled, + _check_bishengir_api_change, + _check_bishengir_able_save_ir, + _is_debug_line_info_disabled, + _enable_print_ub_bits, + _enable_dump_memory_info, + _enable_msdebug, + _enable_unpublished_feature, + _check_cxx11_abi, + _build_npu_ext, + convert_sigtype_to_int, + force_disable_ffts, + triton_enable_libdevice_simt, + get_cann_version, + is_compile_on_910_95, +) +from .npu_driver import NPUUtils + + +def _get_dicp_opt_path() -> str: + import triton + + triton_dir = os.path.dirname(os.path.realpath(triton.__file__)) + plat_name = sysconfig.get_platform() + python_version = sysconfig.get_python_version() + cmake_dir_name = f"cmake.{plat_name}-{sys.implementation.name}-{python_version}" + return os.path.join( + triton_dir, + os.pardir, + "build", + cmake_dir_name, + "third_party", + "dicp_triton", + "tools", + "dicp_triton_opt", + "dicp_opt", + ) + + +def _get_mlir_path(path: str, *paths) -> str: + root_path = os.getenv("MLIR_ROOT", "") + if root_path == "": + raise EnvironmentError("MLIR_ROOT is not set.") + return os.path.join(root_path, path, *paths) + + +def _get_llvm_path(path: str, *paths) -> str: + root_path = os.getenv("LLVM_ROOT", "") + if root_path == "": + raise EnvironmentError("LLVM_ROOT is not set.") + return os.path.join(root_path, path, *paths) def downgrade_llir(llir): @@ -59,19 +84,6 @@ def downgrade_llir(llir): return llir -def _replace_mod_ir_with_file(mod, filepath: str, stage_name: str): - p = Path(filepath) - if not p.exists(): - raise FileNotFoundError(f"Replacement MLIR file not found: {filepath}") - print(f"[DEBUG] replacing '{stage_name}' IR with file '{filepath}'") - try: - new_mod = ir.parse_mlir_module(str(p), mod.context) - new_mod.context = mod.context - return new_mod - except Exception as e: - raise RuntimeError(f"Failed to parse replacement MLIR file '{filepath}': {e}") - - def _downgrade_mem_attrs(llir: str): memory_pattern = r"memory\([^()]*\)" @@ -88,7 +100,7 @@ def replace_mem_attr(m): assert len(pair) <= 2 if len(pair) == 1: rw = rw_map[pair[0].strip()] - loc = loc_map["other"] # all location + loc = loc_map["other"] else: rw = rw_map[pair[1].strip()] loc_str = pair[0].strip() @@ -124,237 +136,103 @@ def _downgrade_stacksaverestore_intrinsics(llir: str): return llir -def _get_triton_adapter_opt_path() -> str: - path = os.path.dirname(__file__) - path = os.path.join(path, "triton-adapter-opt") - return path - - -def _get_dicp_opt_path() -> str: - base_path = os.path.dirname(__file__) - path = os.path.join(base_path, "../../_C", "dicp_opt") - return path - - -def _get_triton_shared_opt_path() -> str: - base_path = os.path.dirname(__file__) - path = os.path.join(base_path, "../../_C", "triton-shared-opt-v3_2") - path34 = os.path.join(base_path, "../../_C", "triton-shared-opt-v3_4") - if os.path.exists(path34): - path = path34 - path = os.getenv("TRITON_SHARED_OPT_PATH", path) # allow user override - if not os.path.exists(path): - raise EnvironmentError( - f"Couldn't find triton-shared-opt at {path}, set TRITON_SHARED_OPT_PATH to override" - ) - return path - - -def _get_mlir_path(path: str, *paths) -> str: - root_path = os.getenv("MLIR_ROOT", "") - if root_path == "": - raise EnvironmentError("MLIR_ROOT is not set.") - return os.path.join(root_path, path, *paths) - - -def _get_llvm_path(path: str, *paths) -> str: - root_path = os.getenv("LLVM_ROOT", "") - if root_path == "": - raise EnvironmentError("LLVM_ROOT is not set.") - return os.path.join(root_path, path, *paths) - - -def _get_npucompiler_path() -> str: - npu_compiler_path = shutil.which("bishengir-compile") - if npu_compiler_path is None: - npu_compiler_root = os.getenv("TRITON_NPU_COMPILER_PATH", "") - if npu_compiler_root is None: - raise EnvironmentError( - "Couldn't find executable bishengir-compile or TRITON_NPU_COMPILER_PATH." - ) - npu_compiler_path = os.path.join(npu_compiler_root, "npuc") - npu_compiler_path = os.path.abspath(npu_compiler_path) - return npu_compiler_path - - -def _get_bisheng_path() -> str: - bisheng_path = shutil.which("bisheng") - if bisheng_path is None: - npu_compiler_root = os.getenv("TRITON_NPU_COMPILER_PATH", "") - if npu_compiler_root is None: - raise EnvironmentError( - "Couldn't find executable bisheng or TRITON_NPU_COMPILER_PATH" - ) - bisheng_path = os.path.join(npu_compiler_root, "ccec") - return bisheng_path - - @functools.lru_cache(None) -def _get_ascend_path() -> str: - path = os.getenv("ASCEND_HOME_PATH", "") - if path == "": - raise EnvironmentError( - "ASCEND_HOME_PATH is not set, source /set_env.sh first" +def _get_bishengir_llvm_version() -> int: + try: + npu_compiler_path, _ = _get_npucompiler_path() + result = subprocess.run( + [npu_compiler_path, "--version"], + capture_output=True, + text=True, + timeout=10, ) - return Path(path) - - -def _is_ascend_sanitizer_enabled() -> bool: - return os.getenv("TRITON_ENABLE_SANITIZER", "false").lower() in ("true", "1") - - -def _is_auto_map_parallel_blocks_enabled() -> bool: - return os.getenv("TRITON_ALL_BLOCKS_PARALLEL", "false").lower() in ("true", "1") - - -def _build_npu_ext(obj_name: str, src_path, src_dir, *, kernel_launcher=None) -> str: - suffix = sysconfig.get_config_var("EXT_SUFFIX") - so_path = os.path.join(src_dir, f"{obj_name}{suffix}") - - cxx = os.environ.get("CC") - if cxx is None: - clangxx = shutil.which("clang++") - gxx = shutil.which("g++") - cxx = clangxx if clangxx is not None else gxx - if cxx is None: - raise RuntimeError("Failed to find C++ compiler") - cc_cmd = [cxx, src_path] - # disable all warnings - cc_cmd += [f"-w"] - # find the python library - if hasattr(sysconfig, "get_default_scheme"): - scheme = sysconfig.get_default_scheme() - else: - scheme = sysconfig._get_default_scheme() - # 'posix_local' is a custom scheme on Debian. However, starting Python 3.10, the default install - # path changes to include 'local'. This change is required to use triton with system-wide python. - if scheme == "posix_local": - scheme = "posix_prefix" - py_include_dir = sysconfig.get_paths(scheme=scheme)["include"] - cc_cmd += [f"-I{py_include_dir}"] - # device_print.h - cc_cmd += [f"-I{os.path.dirname(os.path.realpath(__file__))}"] - # find the ascend library - asc_path = _get_ascend_path() - cc_cmd += [ - f"-I{os.path.join(asc_path, 'include')}", - f"-I{os.path.join(asc_path, 'include/experiment')}", - f"-I{os.path.join(asc_path, 'include/experiment/msprof')}", - f"-I{pybind11.get_include()}", - f"-L{os.path.join(asc_path, 'lib64')}", - "-lruntime", - "-lascendcl", - ] - - if kernel_launcher == "torch": - import torch - import torch_npu - - torch_path = os.path.dirname(os.path.realpath(torch.__file__)) - torch_npu_path = os.path.dirname(os.path.realpath(torch_npu.__file__)) - use_cxx11_abi = _check_cxx11_abi() - cc_cmd += [ - f"-I{os.path.join(torch_path, 'include')}", - f"-I{os.path.join(torch_npu_path, 'include')}", - f"-L{os.path.join(torch_npu_path, 'lib')}", - "-ltorch_npu", - f"-D_GLIBCXX_USE_CXX11_ABI={use_cxx11_abi}", - ] - - cc_cmd += ["-std=c++17", "-shared", "-fPIC", "-o", so_path] - - ret = subprocess.check_call(cc_cmd) - - if ret == 0: - return so_path - else: - raise RuntimeError("Failed to compile " + src_path) - - -def _get_kernel_target(metadata: dict): - if "target" not in metadata: - raise Exception("No target provided!") - sub_target = metadata["target"].arch - assert isinstance(sub_target, str) - if sub_target.startswith("Ascend910B"): - mix_mode = metadata["mix_mode"] - if mix_mode.lower().strip("_").startswith("aiv"): - return "ascend_910b_vec", "c220-vec", "aiv" - elif mix_mode.lower().strip("_").startswith("aic"): - return "ascend_910b_cube", "c220-cube", "aic" - else: - return "ascend_910b", "c220", "mix" - elif sub_target.startswith("Ascend910"): - return "ascend_910", "c100", "mix" - else: - raise NotImplementedError(f"NPU subtarget {sub_target} not supported yet") - - -def _check_cxx11_abi(): - import torch - - return 1 if torch._C._GLIBCXX_USE_CXX11_ABI else 0 - - -def convert_sigtype_to_int(sigty: str): - MAP_SIGTYPE_TO_INT = { - # Boolean - "i1": 12, # BOOL - # Integer types - "i8": 2, # INT8 - "i16": 6, # INT16 - "i32": 3, # INT32 - "i64": 9, # INT64 - # Unsigned integer types - "u32": 8, # UINT32 - "u64": 10, # UINT64 - # Floating point types - "fp16": 1, # FLOAT16 - "bf16": 27, # DT_BF16 - "fp32": 0, # FLOAT - "fp64": 11, # DOUBLE - } - if sigty not in MAP_SIGTYPE_TO_INT: - raise ValueError(f"Unsupported data type: {sigty}") - - return MAP_SIGTYPE_TO_INT[sigty] + output = (result.stdout + result.stderr).lower() + m = re.search(r"llvm\s+(\d+)", output) + if m: + return int(m.group(1)) + except Exception: + pass + return 19 + + +# def _downgrade_mlir_for_legacy_llvm(content: str) -> str: +# content = content.replace("*xf", "?xf") +# content = content.replace("*xi", "?xi") +# content = content.replace("*xbf", "?xbf") +# # 匹配形如 "memref<...> to tensor<...>" 的模式 +# pattern = r"(memref\<.*?\>)\s+to\s+(tensor\<.*?\>)" +# # 使用正则替换，保留memref和tensor类型，中间插入注释 +# content = re.sub(pattern, r"\1 // to \2", content) +# if len(re.findall("hivm\.hir\.custom", content)) > 0: +# content = re.sub(r'"#hivm\.pipe<([A-Za-z0-9_]*)>"', r"#hivm.pipe<\1>", content) +# content = re.sub( +# r'"#hivm\.tcore_type<([A-Za-z0-9_]*)>"', r"#hivm.tcore_type<\1>", content +# ) +# content = re.sub( +# r'"#hivm\.vf_mode<([A-Za-z0-9_]*)>"', r"#hivm.vf_mode<\1>", content +# ) +# return content + + +def _downgrade_mlir_for_legacy_llvm(content: str) -> str: + _TO_BUFFER_RE = re.compile( + r"\bbufferization\.to_buffer\b[^\n:]*\s+(?P%[^\s:]+)[^\n:]*:\s*" + r"(?Ptensor<[^\n]+?>)\s+to\s+(?Pmemref<[^\n]+?>)", + ) + _TO_TENSOR_RE = re.compile( + r"\bbufferization\.to_tensor\b\s+(?P%[^\s:]+)" + r"(?P(?:\s+(?:restrict|writable))*)\s*:\s*" + r"(?Pmemref<[^\n]+?>)\s+to\s+tensor<[^\n]+?>", + ) + _HIVM_ATTR_RE = re.compile(r'"#hivm\.(pipe|tcore_type|vf_mode)<([A-Za-z0-9_]*)>"') + content = content.replace("*x", "?x") + content = _TO_BUFFER_RE.sub( + lambda m: f'bufferization.to_memref {m.group("value")} : {m.group("memref")}', + content, + ) + content = _TO_TENSOR_RE.sub( + lambda m: f'bufferization.to_tensor {m.group("value")}{m.group("attrs")} : {m.group("memref")}', + content, + ) + return _HIVM_ATTR_RE.sub(r"#hivm.\1<\2>", content) def _check_bishengir_is_regbased() -> bool: - bishengir_path = _get_npucompiler_path() + bishengir_path, _ = _get_npucompiler_path() try: result = subprocess.run( - f"{bishengir_path} --help | grep 'reg-based'", - shell=True, + [bishengir_path, "--help"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, ) - if result.returncode == 0: - # bishengir-compile is regbased version + if result.returncode == 0 and "reg-based" in result.stdout: return True - else: - # bishengir-compile is membased version - return False + return False except Exception as e: print(f"ERROR: {e}") return False -###################### utils.py end ###################### +# --------------------------------------------------------------------------- +# Compiler pipeline functions +# --------------------------------------------------------------------------- -# TODO: materialize the concrete min shape def min_dot_size(target: GPUTarget): - # return lambda lhsType, rhsType: (16, 16, 16) return lambda lhsType, rhsType: (1, 1, 1) def make_ttir(mod, metadata, opt): if "hash" not in metadata: - metadata["hash"] = hashlib.md5(f"{mod}-{metadata}".encode()).hexdigest() - mod.set_attr("dicp.backend", ir.builder(mod.context).get_string_attr("ascend")) - # the same optimize pass for triton-ir as all other backends + metadata["hash"] = hashlib.sha256(f"{mod}-{metadata}".encode()).hexdigest() + if opt.arch: + target_attr_str = f'#hacc.target<"{opt.arch}">' + try: + builder = dicp_triton.ir.dicp_npu_ir_builder(mod.context, opt.arch) + mod.set_attr("hacc.target", builder.parse_attr(target_attr_str)) + except Exception as e: + logging.warning(f"[DICP] Failed to set hacc.target: {e}") pm = ir.pass_manager(mod.context) pm.enable_debug() passes.common.add_inliner(pm) @@ -365,206 +243,98 @@ def make_ttir(mod, metadata, opt): passes.common.add_licm(pm) passes.common.add_symbol_dce(pm) pm.run(mod) - if opt.debug or dump_ir: + if opt.debug: dicp_utils._dump_stage_ir(str(mod), metadata["hash"], "kernel.ttir.mlir") return mod -def ttir_to_linalg(mod, metadata, opt, *, named_ops=True): - # use triton_adapter to lower Triton-MLIR to linalg - # Get Triton-MLIR as string - ttir_code = str(mod) - with tempfile.TemporaryDirectory() as tmpdir: - src_path = os.path.join(tmpdir, "kernel.ttir.mlir") - dst_path = os.path.join(tmpdir, "kernel.ttadapter.mlir") - Path(src_path).write_text(ttir_code) - triton_adapter_opt_path = _get_triton_adapter_opt_path() - - cmd_list = [ - triton_adapter_opt_path, - src_path, - f"--triton-to-linalg=global-kernel=false named-ops={named_ops}", - "-o", - dst_path, - ] - if _is_ascend_sanitizer_enabled(): - cmd_list += ["--mlir-print-debuginfo"] # pass debug info - - ret = subprocess.run(cmd_list, capture_output=True, check=True) - if opt.debug: - dump_manager = get_dump_manager(metadata["hash"]) - dump_manager.put( - Path(dst_path).read_text(), "kernel.ttadapter.mlir", binary=False - ) - - return Path(dst_path).read_text() - - -def ttir_to_ttsharedir_ascend(mod, metadata, opt, *, named_ops=False): +def ttir_to_linalg_dicp(mod, metadata, opt, *, named_ops=False): + """Lower TTIR to linalg using triton-dicp pipeline.""" pm = ir.pass_manager(mod.context) - dicp_triton.passes.triton_shared_ascend.add_discrete_mask_access_conversion( - pm, False, False - ) - dicp_triton.passes.triton_shared_ascend.add_triton_to_unstructure(pm) - dicp_triton.passes.triton_shared_ascend.add_bubble_up_operation(pm) - dicp_triton.passes.triton_shared_ascend.add_canonicalize_cmpi(pm) - dicp_triton.passes.triton_shared_ascend.add_canonicalize_triton_ir_ascend(pm) - dicp_triton.passes.triton_shared_ascend.add_triton_to_linalg_npu(pm) - pm.run(mod) - if opt.debug or dump_ir: - cmd_list = [ - _get_dicp_opt_path(), - "kernel.ttir.mlir", - "--discrete-mask-access-conversion", - "--triton-to-unstructure", - "--bubble-up-operation", - "--canonicalize-cmpi", - "--canonicalize-triton-ir-ascend", - "--triton-to-linalg-npu-conversion", - ] - dicp_utils._dump_stage_ir( - str(mod), metadata["hash"], "kernel.ttshared.mlir", cmd_list - ) - if replace_ttshared_ir is not None: - return _replace_mod_ir_with_file( - mod, replace_ttshared_ir, "ttir_to_ttsharedir_ascend" - ) - return mod - - -def commonir_to_linkedir(commonir, metadata, opt, *, named_ops=False): - if replace_commonir_ir is not None: - print(f"[DEBUG] Replace common ir with {replace_commonir_ir}") - commonir = Path(replace_commonir_ir).read_text() - - assert isinstance(commonir, str) - if opt.debug or dump_ir: - dicp_utils._dump_stage_ir(commonir, metadata["hash"], "kernel.commonir.mlir") - with tempfile.TemporaryDirectory() as tmpdir: - src_path = os.path.join(tmpdir, "kernel.commonir.mlir") - dst_path = os.path.join(tmpdir, "kernel.linked.mlir") - Path(src_path).write_text(commonir) - cmd_list = [ - _get_dicp_opt_path(), - src_path, - "--lower-affine", - "--normalize-slice-ops", - "--linalg-if-to-select", - "--linalg-generic-to-scf", - "--scalar-to-1d-tensor", - f"--linalg-to-linked=global-kernel=false named-ops=true", - "--linked-to-hivm", - "--vectorize-parallel-loop", - "-o", - dst_path, - ] - try: - ret = subprocess.run(cmd_list, capture_output=True, check=True) - except subprocess.CalledProcessError as e: - print(f"Error: code={e.returncode}, stdout:{e.stdout},stderr: {e.stderr}") - content = Path(dst_path).read_text() - - # TODO(zmz): 修改test_path 中内容，暂时在python中处理，bishengir-compile后续会支持，去掉这里逻辑。 - # 将"*xfxxx"替换成"?xfxxx" - content = content.replace("*xf", "?xf") - content = content.replace("*xi", "?xi") - content = content.replace("*xbf", "?xbf") - # 匹配形如 "memref<...> to tensor<...>" 的模式 - pattern = r"(memref\<.*?\>)\s+to\s+(tensor\<.*?\>)" - # 使用正则替换，保留memref和tensor类型，中间插入注释 - content = re.sub(pattern, r"\1 // to \2", content) - - if opt.debug or dump_ir: - cmd_list = [ - _get_dicp_opt_path(), - "kernel.ttshared.mlir", - "--lower-affine", - "--normalize-slice-ops", - "--linalg-if-to-select", - "--linalg-generic-to-scf", - "--scalar-to-1d-tensor", - f"--linalg-to-linked=global-kernel=false named-ops=true", - "--linked-to-hivm", - "--vectorize-parallel-loop", - ] - dicp_utils._dump_stage_ir( - content, metadata["hash"], "kernel.linkedir.mlir", cmd_list - ) - - if replace_commonir_linked_ir is not None: - print(f"[DEBUG] Replace Linkedir with {replace_commonir_linked_ir}") - return Path(replace_commonir_linked_ir).read_text() - return content + enable_mask_fallback = metadata["enable_mask_fallback_conversion"] + optimize_dynamic_offset = metadata["optimize_dynamic_offset"] + compile_on_910_95 = metadata["compile_on_910_95"] + force_simt_template = metadata["force_simt_template"] + enable_sync_block_lock = metadata["enable_sync_block_lock"] + enable_nd2nz_on_vector = metadata["enable_nd2nz_on_vector"] + enable_select_analysis = metadata["enable_select_analysis"] + auto_blockify_size = metadata["auto_blockify_size"] + if not _is_auto_map_parallel_blocks_enabled(): + auto_blockify_size = 1 + + dicp_triton.passes.ttir.add_ascend_legalize(pm) + dicp_triton.passes.ttir.add_auto_blockify(pm, auto_blockify_size) + + if metadata["add_auto_scheduling"]: + dicp_triton.passes.ttir.add_dag_sync(pm) + dicp_triton.passes.ttir.add_dag_scope(pm) + passes.common.add_cse(pm) + passes.common.add_canonicalizer(pm) + dicp_triton.passes.ttir.add_dag_ssbuffer(pm) + passes.common.add_cse(pm) + passes.common.add_canonicalizer(pm) + + dicp_triton.passes.ttir.add_triton_to_structure( + pm, enable_mask_fallback, optimize_dynamic_offset + ) + dicp_triton.passes.ttir.add_discrete_mask_access_conversion( + pm, compile_on_910_95, force_simt_template, enable_sync_block_lock + ) + dicp_triton.passes.ttir.add_triton_to_annotation(pm) + dicp_triton.passes.ttir.add_triton_to_unstructure( + pm, compile_on_910_95, force_simt_template + ) + dicp_triton.passes.ttir.add_triton_to_hivm(pm) + dicp_triton.passes.ttir.add_triton_to_hfusion(pm) + dicp_triton.passes.ttir.add_triton_to_llvm(pm) + dicp_triton.passes.ttir.add_bubble_up_operation(pm) + dicp_triton.passes.ttir.add_triton_to_structure( + pm, enable_mask_fallback, optimize_dynamic_offset + ) + dicp_triton.passes.ttir.add_triton_to_linalg( + pm, + False, + named_ops, + enable_nd2nz_on_vector, + enable_select_analysis, + compile_on_910_95, + ) + dicp_triton.passes.ttir.add_ascend_npu_ir_legalize(pm, False) -def ttsharedir_to_linkedir(mod, metadata, opt, *, named_ops=False, cpu_verify=False): - pm = ir.pass_manager(mod.context) - dicp_triton.passes.linked_npu.add_lower_affine(pm) - dicp_triton.passes.linked_npu.add_normalize_slice_ops(pm) - dicp_triton.passes.linked_npu.add_linalg_if_to_select(pm) - dicp_triton.passes.linked_npu.add_linalg_generic_to_scf(pm) - dicp_triton.passes.linked_npu.add_scalar_to_1d_tensor(pm) - dicp_triton.passes.linked_npu.add_linalg_to_linked(pm, named_ops, True, cpu_verify) - dicp_triton.passes.linked_npu.add_linked_to_hivm(pm) - if cpu_verify: - dicp_triton.passes.linked_npu.add_debug_cpu_verify(pm) - # TODO(zmz): 修改test_path 中内容，暂时在python中处理，bishengir-compile后续会支持，去掉这里逻辑。 pm.run(mod) + content = str(mod) - if cpu_verify: - return content - # 将"*xfxxx"替换成"?xfxxx" - content = content.replace("*xf", "?xf") - content = content.replace("*xi", "?xi") - content = content.replace("*xbf", "?xbf") - # 匹配形如 "memref<...> to tensor<...>" 的模式 - pattern = r"(memref\<.*?\>)\s+to\s+(tensor\<.*?\>)" - # 使用正则替换，保留memref和tensor类型，中间插入注释 - content = re.sub(pattern, r"\1 // to \2", content) - # 处理customop的attr - if len(re.findall("hivm\.hir\.custom", content)) > 0: - content = re.sub(r'"#hivm\.pipe<([A-Za-z0-9_]*)>"', r"#hivm.pipe<\1>", content) - content = re.sub( - r'"#hivm\.tcore_type<([A-Za-z0-9_]*)>"', r"#hivm.tcore_type<\1>", content - ) - content = re.sub( - r'"#hivm\.vf_mode<([A-Za-z0-9_]*)>"', r"#hivm.vf_mode<\1>", content - ) - if opt.debug or dump_ir: + if opt.debug: + pipeline_str = pm.get_pipeline_str() + dicp_opt_path = _get_dicp_opt_path() cmd_list = [ - _get_dicp_opt_path(), - "kernel.ttshared.mlir", - "--lower-affine", - "--normalize-slice-ops", - "--linalg-if-to-select", - "--linalg-generic-to-scf", - "--scalar-to-1d-tensor", - f"--linalg-to-linked=global-kernel=false named-ops=true", - "--linked-to-hivm", + dicp_opt_path, + "", + f"--pass-pipeline={pipeline_str}", + "--mlir-print-debuginfo", + "-o", + "/dev/null", ] dicp_utils._dump_stage_ir( - content, metadata["hash"], "kernel.linkedir.mlir", cmd_list + content, metadata["hash"], "kernel.dicp.mlir", cmd_list ) - if replace_linked_ir is not None: - print(f"[DEBUG] Replace Linkedir with {replace_linked_ir}") - return Path(replace_linked_ir).read_text() return content def linalg_to_llir(linalg: str, metadata, opt): with tempfile.TemporaryDirectory() as tmpdir: - ttadapter_path = os.path.join(tmpdir, "kernel.ttadapter.mlir") + dicp_path = os.path.join(tmpdir, "kernel.dicp.mlir") llmlir_path = os.path.join(tmpdir, "kernel.llir.mlir") llir_path = os.path.join(tmpdir, "kernel.ll") - Path(ttadapter_path).write_text(linalg) + Path(dicp_path).write_text(linalg) mlir_opt_path = _get_mlir_path("bin", "mlir-opt") - # TritonAdapter-MLIR to LLVM-MLIR subprocess.check_call( [ mlir_opt_path, - ttadapter_path, + dicp_path, "--convert-linalg-to-affine-loops", "--eliminate-empty-tensors", "--empty-tensor-to-alloc-tensor", @@ -582,12 +352,8 @@ def linalg_to_llir(linalg: str, metadata, opt): "--expand-strided-metadata", "--finalize-memref-to-llvm", "--convert-func-to-llvm", - # Lowering memrefs creates more affine.apply ops. - # Lowering these affine ops again creates further arith ops, - # so we have to run these two passes again here. "--lower-affine", "--convert-arith-to-llvm", - # Remove all unrealized casts created "--reconcile-unrealized-casts", "-o", llmlir_path, @@ -599,7 +365,6 @@ def linalg_to_llir(linalg: str, metadata, opt): Path(llmlir_path).read_text(), "kernel.llir.mlir", binary=False ) - # LLVM-MLIR to LLVM-IR mlir_translate_path = _get_mlir_path("bin", "mlir-translate") subprocess.check_call( [mlir_translate_path, llmlir_path, "--mlir-to-llvmir", "-o", llir_path] @@ -612,12 +377,7 @@ def linalg_to_llir(linalg: str, metadata, opt): def llir_to_cpuasm(llir: str, metadata, opt): - # add metadata at final stage - # Note: Compiled Kernel requires to estimate size of shared memory to occupy - # Currently, CPU backend requires no limit on shared memory size metadata["shared"] = 1 - # We can get a function name (C naming) from - # LLVM-IR by getting the first "define void @". fn_name = llir.split("define void @")[1].split("(")[0].strip() metadata["name"] = fn_name + " cpu" with tempfile.TemporaryDirectory() as tmpdir: @@ -657,7 +417,6 @@ def llir_to_cpuasm(llir: str, metadata, opt): dump_manager = get_dump_manager(metadata["hash"]) dump_manager.put(Path(dst_path).read_text(), "kernel.s", binary=False) - # Actually it's text-format assembly. Use read_text(). return Path(dst_path).read_text() @@ -671,22 +430,12 @@ def __get_metadata_attr_by_callback(lib, postfix: str, metadata, meta_key: str): def _parse_linalg_metadata(linalg: str, metadata: dict): - """ - Parse Linalg IR to extract metadata required for NPU compilation. - Extracts and updates the following fields in metadata: - - mix_mode - - kernel_name - - tensor_kinds - - shared (currently hardcoded) - - name (combined kernel_name and mix_mode) - Additionally, removes the mix_mode attribute from the IR. - """ - # --- Regular expressions and examples --- - # Example: mix_mode = "aiv" -> aiv + DISABLE_AUTO_TILE_AND_BIND_SUBBLOCK_REGEX = ( + r"hivm.disable_auto_tile_and_bind_subblock" + ) MIX_MODE_REGEX = r'mix_mode\s*=\s*"([^"]+)"' - # Example: func.func @gather_sorted_kernel(%arg0: ...) -> gather_sorted_kernel + PARALLEL_MODE_REGEX = r'parallel_mode\s*=\s*"([^"]+)"' KERNEL_NAME_REGEX = r"func\.func\s+@(\w+)" - # Example: %arg1: memref {tt.divisibility = 16 : i32, tt.tensor_kind = 0 : i32} -> ('1', '0') TENSOR_KIND_REGEX = ( r"%arg(\d+):[^,)]*?\{[^}]*?tt\.tensor_kind\s*=\s*([^:\s}]+)\s*:[^}]*?\}" ) @@ -699,211 +448,707 @@ def _parse_linalg_metadata(linalg: str, metadata: dict): # Note: Compiled Kernel requires to estimate size of shared memory to occupy # Currently, NPU backend does not limit on shared memory metadata["shared"] = 1 - # the mix mode is also encoded into metadata['name'] for runtime to distinguish + metadata["auto_tile_and_bind_subblock"] = not re.search( + DISABLE_AUTO_TILE_AND_BIND_SUBBLOCK_REGEX, linalg + ) metadata["mix_mode"] = re.search(MIX_MODE_REGEX, linalg).group(1) + metadata["parallel_mode"] = re.search(PARALLEL_MODE_REGEX, linalg).group(1) metadata["kernel_name"] = re.search(KERNEL_NAME_REGEX, linalg).group(1) - # Use while space to split kernel_name and mix_mode. - # Check the function load_binary in npu_driver.py. - metadata["name"] = metadata["kernel_name"] + " " + metadata["mix_mode"] - # Parse all tensor kinds from arguments + metadata["name"] = metadata["kernel_name"] metadata["tensor_kinds"] = [ int(kind) for _, kind in re.findall(TENSOR_KIND_REGEX, linalg) ] + metadata["required_ub_bits"] = 0 - # Parse all bitcode paths bitcodes = re.findall(BITCODES_REGEX, linalg) metadata["bitcodes"] = [val for group in bitcodes for val in group if val] - - # remove the mix_mode attribute - linalg = re.sub(REMOVE_MIX_MODE_REGEX, "", linalg) return linalg, metadata -def linalg_to_bin_enable_npu_compile(linalg: str, metadata, opt): +def _parse_ttir_metadata(ttir: str, metadata: dict): + KERNEL_NAME_REGEX = r"tt\.func\spublic\s+@(\w+)" + TENSOR_KIND_REGEX = ( + r"%arg(\d+):[^,)]*?\{[^}]*?tt\.tensor_kind\s*=\s*([^:\s}]+)\s*:[^}]*?\}" + ) + + metadata["shared"] = 1 + metadata["mix_mode"] = "aiv" + metadata["kernel_name"] = re.search(KERNEL_NAME_REGEX, ttir).group(1) + metadata["name"] = metadata["kernel_name"] + metadata["tensor_kinds"] = [ + int(kind) for _, kind in re.findall(TENSOR_KIND_REGEX, ttir) + ] + return metadata + + +def get_common_bishengir_compile_options(metadata): + bishengir_target = metadata["target"].arch + bishengir_target_opt = f"--target={bishengir_target}" + return [bishengir_target_opt] + + +def get_auto_bind_sub_block_option(metadata): + enable_auto_bind_sub_block = metadata["enable_auto_bind_sub_block"] + return True if enable_auto_bind_sub_block is None else enable_auto_bind_sub_block + + +def _save_npuir_debug_output( + stdout_bytes: bytes, stderr_bytes: bytes, tmpdir: str, metadata_hash: str +): + stdout = stdout_bytes.decode("utf-8") if stdout_bytes else "" + stderr = stderr_bytes.decode("utf-8") if stderr_bytes else "" + combined = stdout + stderr + if not combined.strip(): + combined = "No output captured." + output_path = os.path.join(tmpdir, "kernel.npuir.mlir") + with open(output_path, "w", encoding="utf-8") as f: + f.write(combined) + + dump_manager = get_dump_manager(metadata_hash) + dump_manager.put( + Path(output_path).read_text(encoding="utf-8"), "kernel.npuir.mlir", binary=False + ) + + +def get_libdevice(): + current = os.path.dirname(__file__) + return os.path.join(current, "lib/libdevice.10.bc") + + +# --------------------------------------------------------------------------- +# Shared NPU compilation orchestration +# --------------------------------------------------------------------------- + + +def _compile_linalg_to_npu_bin( + linalg, + metadata, + opt, + *, + build_options_fn, + bishengir_hivm_opt=None, + extra_cmd_args=None, + debug_stage_name="kernel.npuir_input.mlir", +): + """Shared orchestration for linalg → npubin compilation. + + Parameters + ---------- + build_options_fn : callable + ``(metadata, opt) -> list[str]`` that builds the initial + ``_compile_option_list`` for this platform. + bishengir_hivm_opt : str or None + If set, injected between ``--enable-hfusion-compile`` and + ``--enable-triton-kernel-compile`` in the bishengir-compile block + (A2/A3 only). + extra_cmd_args : callable or None + Optional ``(metadata, opt) -> list[str]`` for appending extra args + after ``-o bin_file`` (910_95: vf_merge_level + hfusion multi-consumer). + debug_stage_name : str + File name used when dumping the input IR in debug mode. + """ linalg, metadata = _parse_linalg_metadata(linalg, metadata) + if replace_dicp_ir is not None: + print(f"[DEBUG] Replace dicp ir with {replace_dicp_ir}") + linalg = Path(replace_dicp_ir).read_text() + if _get_bishengir_llvm_version() < 22: + linalg = _downgrade_mlir_for_legacy_llvm(linalg) + if opt.debug: + dicp_utils._dump_stage_ir(linalg, metadata["hash"], debug_stage_name) + with tempfile.TemporaryDirectory() as tmpdir: - ttadapter_path = os.path.join(tmpdir, "kernel.ttadapter.mlir") - lower_by_ttshared = os.getenv("LOWER_BY_TTSHARED", "1") - if lower_by_ttshared == "1": - ttadapter_path = os.path.join(tmpdir, "kernel.linkedir.mlir") - Path(ttadapter_path).write_text(linalg) + dicp_path = os.path.join(tmpdir, "kernel.dicp.mlir") + Path(dicp_path).write_text(linalg) bin_file = os.path.join(tmpdir, "kernel") - if _check_bishengir_is_regbased(): - bishengir_hivm_opt = "--reg-based=true" + if _check_bishengir_api_change(): + bin_file_with_ext = "kernel.o" else: - bishengir_hivm_opt = "--enable-hivm-compile=true" - bin_path = os.path.join(tmpdir, "kernel_reloc.o") + bin_file_with_ext = "kernel_reloc.o" + bin_path = os.path.join(tmpdir, bin_file_with_ext) callback_path = os.path.join(tmpdir, "libkernel.so") - multibuffer = metadata["multibuffer"] - _compile_option_list = [] - # if dump_ir: - # _compile_option_list += [f"--mlir-print-ir-before-all"] - # _compile_option_list += [f"--mlir-print-ir-after-all"] - _compile_option_list += [ - f"--enable-auto-multi-buffer={multibuffer}", - ] - - if _is_ascend_sanitizer_enabled(): - _compile_option_list += ["--enable-sanitizer=true"] - if _is_auto_map_parallel_blocks_enabled(): - _compile_option_list += ["--enable-auto-blockify-loop"] - - bitcodes = metadata["bitcodes"] - if bitcodes is not None: - for bitcodes in bitcodes: - _compile_option_list += [f"--link-aicore-bitcode={bitcodes}"] - - npu_compiler_path = _get_npucompiler_path() - # support bishengir-compile more version - if "8.2.RC1.alpha002" in npu_compiler_path: - bin_path = os.path.join(tmpdir, "kernel_reloc.o") - elif "8.2.RC1.alpha003" in npu_compiler_path: - bin_path = os.path.join(tmpdir, "kernel.o") - else: - bin_path = os.path.join(tmpdir, "kernel.o") + # --- platform-specific option list --- + _compile_option_list = build_options_fn(metadata, opt) + # --- bishengir-compile wrapper --- + npu_compiler_path, env = _get_npucompiler_path() if npu_compiler_path.endswith("bishengir-compile"): + _compile_option_list += ["--enable-hfusion-compile=true"] + if bishengir_hivm_opt: + _compile_option_list.append(bishengir_hivm_opt) + _compile_option_list += ["--enable-triton-kernel-compile=true"] + + # --- debug --- + if opt.debug: _compile_option_list += [ - "--enable-hfusion-compile=true", - bishengir_hivm_opt, - "--enable-triton-kernel-compile=true", + "--bishengir-print-ir-after=hivm-graph-sync-solver" ] - inject_barrier_all = metadata["inject_barrier_all"] + cmd_list = ( + [npu_compiler_path, dicp_path] + _compile_option_list + ["-o", bin_file] + ) + if extra_cmd_args: + cmd_list.extend(extra_cmd_args(metadata, opt)) + + if opt.debug: + print(f"[DEBUG] cmd_list: {' '.join(cmd_list)}") + + # --- execute --- + try: + ret = subprocess.run( + cmd_list, + env=env, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + check=True, + ) + except subprocess.CalledProcessError as e: + if opt.debug: + _save_npuir_debug_output(e.stdout, e.stderr, tmpdir, metadata["hash"]) + error_msg = e.stderr.decode("utf-8") if e.stderr else str(e) + raise CompileTimeAssertionFailure( + None, None, f"bishengir-compile failed: {error_msg}" + ) from e + + stdout_bytes = ret.stdout + stderr_bytes = ret.stderr + stdout_str = stdout_bytes.decode("utf-8") if stdout_bytes else "" + stderr_str = stderr_bytes.decode("utf-8") if stderr_bytes else "" + if opt.debug: + print( + f"[DEBUG] bishengir-compile stdout:\n{stdout_str if stdout_str else ''}" + ) + print( + f"[DEBUG] bishengir-compile stderr:\n{stderr_str if stderr_str else ''}" + ) + _save_npuir_debug_output( + stdout_bytes, stderr_bytes, tmpdir, metadata["hash"] + ) + + match = re.search(r"UB\s+size\s*=\s*(\d+)\s*bits", stdout_str) + if match: + metadata["required_ub_bits"] = int(match.group(1)) + + if not Path(bin_path).exists(): + error_msg = ret.stderr.decode("utf-8") if ret.stderr else "" + print(f"[DEBUG] {bin_path} is not found") + print(f"[DEBUG] Stderr:\n{error_msg}") + raise CompileTimeAssertionFailure( + None, None, f"bishengir-compile output not found: {error_msg}" + ) + + if Path(callback_path).is_file(): + lib = ctypes.CDLL(callback_path) + __get_metadata_attr_by_callback( + lib, "_infer_task_type_function", metadata, "bs_task_type" + ) + __get_metadata_attr_by_callback( + lib, "_infer_workspace_shape_function", metadata, "workspace_size" + ) + __get_metadata_attr_by_callback( + lib, "_infer_sync_block_lock_num_function", metadata, "lock_num" + ) + __get_metadata_attr_by_callback( + lib, "_infer_sync_block_lock_init_function", metadata, "lock_init_val" + ) + + return Path(bin_path).read_bytes() + + +# --------------------------------------------------------------------------- +# 910_95 compilation path +# --------------------------------------------------------------------------- + + +def linalg_to_bin_enable_npu_compile_910_95(linalg: str, metadata, opt): + def _build_options(m, o): + opts = get_common_bishengir_compile_options(m) + + multibuffer = m.get("multibuffer") + num_stages = m.get("num_stages") + multi_buffer_value = True + if multibuffer is not None and not multibuffer: + multi_buffer_value = False + elif num_stages is not None and num_stages == 1: + multi_buffer_value = False + opts.append(f"--enable-auto-multi-buffer={multi_buffer_value}") + + enable_tuning_mode = m["enable_tuning_mode"] + if enable_tuning_mode is not None: + opts.append(f"--enable-tuning-mode={enable_tuning_mode}") + + if m.get("disable_tightly_coupled_buffer_reuse"): + opts.append("--disable-tightly-coupled-buffer-reuse") + + opts.append(f"--enable-auto-bind-sub-block={get_auto_bind_sub_block_option(m)}") + + if force_disable_ffts(): + opts.append("--disable-ffts") + if _is_ascend_sanitizer_enabled(): + opts.append("--enable-sanitizer=true") + if not _is_debug_line_info_disabled(): + opts.append("--enable-debug-info=true") + if _enable_print_ub_bits(): + opts.append("--enable-print-memory-allocated-size") + + enable_hivm_auto_cv_balance = m["enable_hivm_auto_cv_balance"] + if enable_hivm_auto_cv_balance is not None: + opts.append(f"--enable-hivm-auto-cv-balance={enable_hivm_auto_cv_balance}") + + sync_solver = m["sync_solver"] + if sync_solver is not None: + opts.append(f"--enable-hivm-graph-sync-solver={sync_solver}") + + unit_flag = m["unit_flag"] + if unit_flag is not None: + opts.append(f"--enable-hivm-unit-flag-sync={unit_flag}") + + inject_barrier_all = m["inject_barrier_all"] if inject_barrier_all is not None: - _compile_option_list += [ - f"--enable-hivm-inject-barrier-all-sync={inject_barrier_all}" - ] + opts.append(f"--enable-hivm-inject-barrier-all-sync={inject_barrier_all}") + + inject_block_all = m["inject_block_all"] + if inject_block_all is not None: + opts.append(f"--enable-hivm-inject-block-all-sync={inject_block_all}") + + limit_auto_multi_buffer_only_for_local_buffer = m[ + "limit_auto_multi_buffer_only_for_local_buffer" + ] + if limit_auto_multi_buffer_only_for_local_buffer is not None: + opts.append( + f"--limit-auto-multi-buffer-only-for-local-buffer={limit_auto_multi_buffer_only_for_local_buffer}" + ) + + set_workspace_multibuffer = m["set_workspace_multibuffer"] + if set_workspace_multibuffer is not None: + opts.append(f"--set-workspace-multibuffer={set_workspace_multibuffer}") + + auto_multi_buffer = m["limit_auto_multi_buffer_of_local_buffer"] + if auto_multi_buffer is not None: + opts.append( + f"--limit-auto-multi-buffer-of-local-buffer={auto_multi_buffer}" + ) + + enable_mixed_cv = m["enable_mixed_cv"] + if enable_mixed_cv is not None: + opts.append(f"--enable-mixed-cv={enable_mixed_cv}") + + enable_cce_vf_auto_sync = m["enable_cce_vf_auto_sync"] + if enable_cce_vf_auto_sync is not None: + opts.append( + f"--append-bisheng-options=-mllvm --cce-vf-auto-sync={enable_cce_vf_auto_sync}" + ) + + enable_cce_vf_remove_membar = m["enable_cce_vf_remove_membar"] + if enable_cce_vf_remove_membar is not None: + opts.append( + f"--append-bisheng-options=-mllvm --cce-vf-remove-membar={enable_cce_vf_remove_membar}" + ) - disable_auto_inject_block_sync = metadata["disable_auto_inject_block_sync"] + env_vf = os.getenv("TRITON_ENABLE_VF_FUSION") + enable_vf_fusion = ( + env_vf.lower() in ("true", "1", "yes") + if env_vf is not None + else m.get("enable_vf_fusion", False) + ) + if enable_vf_fusion: + opts.append("--enable-vf-fusion") + + enable_drop_unit_dims = m["enable_drop_unit_dims"] + if enable_drop_unit_dims is not None: + opts.append(f"--enable-drop-unit-dims={enable_drop_unit_dims}") + + enable_flatten = m["enable_flatten"] + if enable_flatten is not None: + opts.append(f"--enable-flatten={enable_flatten}") + + enable_auto_vectorize_v2 = m["enable_auto_vectorize_v2"] + if enable_auto_vectorize_v2 is not None: + opts.append(f"--enable-auto-vectorize-v2={enable_auto_vectorize_v2}") + + auto_vectorize_v2_max_fused_ops_num = m["auto_vectorize_v2_max_fused_ops_num"] + if auto_vectorize_v2_max_fused_ops_num is not None: + opts.append( + f"--hfusion-max-fused-ops-in-auto-vectorize-v2={auto_vectorize_v2_max_fused_ops_num}" + ) + + prevec_max_fused_ops_num = m["prevec_max_fused_ops_num"] + if prevec_max_fused_ops_num is not None: + opts.append( + f"--hfusion-max-fused-elementwise-ops={prevec_max_fused_ops_num}" + ) + + disable_auto_inject_block_sync = m["disable_auto_inject_block_sync"] if disable_auto_inject_block_sync is not None: - _compile_option_list += [ + opts.append( f"--disable-auto-inject-block-sync={disable_auto_inject_block_sync}" - ] + ) + + bitcodes = m["bitcodes"] + if bitcodes is not None: + for bitcode in bitcodes: + opts.append(f"--link-aicore-bitcode={bitcode}") + + enable_auto_blockify = m["enable_auto_blockify"] + if _is_auto_map_parallel_blocks_enabled(): + if enable_auto_blockify is None or enable_auto_blockify: + opts.append("--enable-auto-blockify-loop") + elif enable_auto_blockify: + opts.append("--enable-auto-blockify-loop") + + bisheng_options = m["bisheng_options"] + if bisheng_options is not None: + opts.append(f"--append-bisheng-options={bisheng_options}") + + if o.mix_mode in ("aic",): + opts.append("--disable-hfusion-vectorize=true") + + return opts + + def _extra_cmd_args(m, o): + args = [] + vf_merge_level = m.get("vf_merge_level") + if vf_merge_level is not None and vf_merge_level != 1: + args.append(f"--enable-vf-merge-level={vf_merge_level}") + hfusion = m.get("hfusion_enable_multiple_consumer_fusion") + if hfusion: + args.append(f"--hfusion-enable-multiple-consumer-fusion={hfusion}") + return args + + return _compile_linalg_to_npu_bin( + linalg, + metadata, + opt, + build_options_fn=_build_options, + extra_cmd_args=_extra_cmd_args, + debug_stage_name="kernel.dicp.mlir", + ) - unit_flag = metadata["unit_flag"] + +# --------------------------------------------------------------------------- +# A2/A3 compilation path +# --------------------------------------------------------------------------- + + +def linalg_to_bin_enable_npu_compile_A2_A3(linalg: str, metadata, opt): + if _check_bishengir_is_regbased(): + bishengir_hivm_opt = "--reg-based=true" + else: + bishengir_hivm_opt = "--enable-hivm-compile=true" + + def _build_options(m, o): + opts = [f"--target={NPUUtils().get_arch()}"] + + multibuffer = m.get("multibuffer") + num_stages = m.get("num_stages") + multi_buffer_value = True + if multibuffer is not None and not multibuffer: + multi_buffer_value = False + elif num_stages is not None and num_stages == 1: + multi_buffer_value = False + opts.append(f"--enable-auto-multi-buffer={multi_buffer_value}") + + enable_tuning_mode = m["enable_tuning_mode"] + if enable_tuning_mode is not None: + opts.append(f"--enable-tuning-mode={enable_tuning_mode}") + + enable_ubuf_saving = m["enable_ubuf_saving"] + if enable_ubuf_saving is not None: + opts.append(f"--enable-ubuf-saving={enable_ubuf_saving}") + + enable_preload = m["enable_preload"] + if enable_preload is not None: + opts.append(f"--enable-preload={enable_preload}") + + opts.append(f"--enable-auto-bind-sub-block={get_auto_bind_sub_block_option(m)}") + + if _is_ascend_sanitizer_enabled(): + opts.append("--enable-sanitizer=true") + if not _is_debug_line_info_disabled(): + opts.append("--enable-debug-info=true") + if _enable_print_ub_bits(): + opts.append("--enable-print-memory-allocated-size") + if _enable_dump_memory_info(): + opts.append("--enable-memory-display=true") + if _enable_msdebug(): + opts.append("--enable-ms-debug=true") + + enable_hivm_auto_cv_balance = m["enable_hivm_auto_cv_balance"] + if enable_hivm_auto_cv_balance is not None: + opts.append(f"--enable-hivm-auto-cv-balance={enable_hivm_auto_cv_balance}") + + sync_solver = m["sync_solver"] + if sync_solver is not None: + opts.append(f"--enable-hivm-graph-sync-solver={sync_solver}") + opts.append(f"--enable-hivm-cross-core-gss={sync_solver}") + + unit_flag = m["unit_flag"] if unit_flag is not None: - _compile_option_list += [f"--enable-hivm-unit-flag-sync={unit_flag}"] + opts.append(f"--enable-hivm-unit-flag-sync={unit_flag}") - disable_auto_cv_work_space_manage = metadata[ - "disable_auto_cv_work_space_manage" - ] - if disable_auto_cv_work_space_manage is True: - _compile_option_list += [ - f"--disable-auto-cv-work-space-manage={disable_auto_cv_work_space_manage}" - ] + enable_drop_unit_dims = m["enable_drop_unit_dims"] + if enable_drop_unit_dims is not None: + opts.append(f"--enable-drop-unit-dims={enable_drop_unit_dims}") - enable_auto_bind_sub_block = metadata["enable_auto_bind_sub_block"] - if enable_auto_bind_sub_block is not None: - _compile_option_list += [ - f"--enable-auto-bind-sub-block={enable_auto_bind_sub_block}" - ] + enable_flatten = m["enable_flatten"] + if enable_flatten is not None: + opts.append(f"--enable-flatten={enable_flatten}") + + enable_auto_vectorize_v2 = m["enable_auto_vectorize_v2"] + if enable_auto_vectorize_v2 is not None: + opts.append(f"--enable-auto-vectorize-v2={enable_auto_vectorize_v2}") - limit_auto_multi_buffer_only_for_local_buffer = metadata[ + inject_barrier_all = m["inject_barrier_all"] + if inject_barrier_all is not None: + opts.append(f"--enable-hivm-inject-barrier-all-sync={inject_barrier_all}") + + inject_block_all = m["inject_block_all"] + if inject_block_all is not None: + opts.append(f"--enable-hivm-inject-block-all-sync={inject_block_all}") + + limit_auto_multi_buffer_only_for_local_buffer = m[ "limit_auto_multi_buffer_only_for_local_buffer" ] if limit_auto_multi_buffer_only_for_local_buffer is not None: - _compile_option_list += [ + opts.append( f"--limit-auto-multi-buffer-only-for-local-buffer={limit_auto_multi_buffer_only_for_local_buffer}" - ] + ) - set_workspace_multibuffer = metadata["set_workspace_multibuffer"] + set_workspace_multibuffer = m["set_workspace_multibuffer"] if set_workspace_multibuffer is not None: - _compile_option_list += [ - f"--set-workspace-multibuffer={set_workspace_multibuffer}" - ] + opts.append(f"--set-workspace-multibuffer={set_workspace_multibuffer}") - tile_mix_vector_loop = metadata["tile_mix_vector_loop"] + tile_mix_vector_loop = m["tile_mix_vector_loop"] if tile_mix_vector_loop is not None: - _compile_option_list += [f"--tile-mix-vector-loop={tile_mix_vector_loop}"] + opts.append(f"--tile-mix-vector-loop={tile_mix_vector_loop}") - tile_mix_cube_loop = metadata["tile_mix_cube_loop"] + tile_mix_cube_loop = m["tile_mix_cube_loop"] if tile_mix_cube_loop is not None: - _compile_option_list += [f"--tile-mix-cube-loop={tile_mix_cube_loop}"] + opts.append(f"--tile-mix-cube-loop={tile_mix_cube_loop}") + + auto_multi_buffer = m["limit_auto_multi_buffer_of_local_buffer"] + if auto_multi_buffer is not None: + opts.append( + f"--limit-auto-multi-buffer-of-local-buffer={auto_multi_buffer}" + ) + + disable_auto_inject_block_sync = m["disable_auto_inject_block_sync"] + if disable_auto_inject_block_sync is not None: + opts.append( + f"--disable-auto-inject-block-sync={disable_auto_inject_block_sync}" + ) + + bitcodes = m["bitcodes"] + if bitcodes is not None: + for bitcode in bitcodes: + opts.append(f"--link-aicore-bitcode={bitcode}") + + if m.get("disable_auto_cv_work_space_manage") is True: + opts.append("--disable-auto-cv-work-space-manage=True") + opts.append(f"--link-aicore-bitcode={get_libdevice()}") + + disable_size_align_for_cast = m["disable_size_align_for_cast"] + if disable_size_align_for_cast is not None: + opts.append(f"--disable-size-align-for-cast={disable_size_align_for_cast}") + + if _is_auto_map_parallel_blocks_enabled(): + opts.append("--enable-auto-blockify-loop") + + return opts + + return _compile_linalg_to_npu_bin( + linalg, + metadata, + opt, + build_options_fn=_build_options, + bishengir_hivm_opt=bishengir_hivm_opt, + ) + + +# --------------------------------------------------------------------------- +# SIMT-only TTIR -> npubin path +# --------------------------------------------------------------------------- + + +def ttir_to_npubin(mod, metadata, opt): + ttir_code = str(mod) + metadata = _parse_ttir_metadata(ttir_code, metadata) + with tempfile.TemporaryDirectory() as tmpdir: + src_path = os.path.join(tmpdir, "kernel.ttir.mlir") + Path(src_path).write_text(ttir_code) + bin_file = os.path.join(tmpdir, "kernel") + bin_path = os.path.join(tmpdir, "kernel.o") + _compile_option_list = get_common_bishengir_compile_options(metadata) + if opt.force_simt_only: + _compile_option_list += ["--enable-hivm-compile=false"] + _compile_option_list += ["--enable-triton-ir-compile"] + _compile_option_list += ["--pure-simt"] + _compile_option_list += [f"--num-warps={opt.num_warps}"] + _compile_option_list += [f"--threads-per-warp={opt.warp_size}"] + if opt.enable_bishengir_simt_optimization != 000: + _compile_option_list += [ + f"--enable-bishengir-simt-optimization={opt.enable_bishengir_simt_optimization}" + ] + if opt.simt_stack_limit: + _compile_option_list += [f"--simt-stack-limit={opt.simt_stack_limit}"] + if opt.shared_mem_dynamic_size is not None: + _compile_option_list += [ + f"--shared-mem-dynamic-size={opt.shared_mem_dynamic_size}" + ] + if opt.enable_simt_reorder_instruction: + _compile_option_list += ["--enable-simt-reorder-instruction=true"] + if opt.disable_fma: + _compile_option_list += [f"--disable-fma"] + enable_libdevice_simt = triton_enable_libdevice_simt() + if enable_libdevice_simt: + bisheng_options = metadata["bisheng_options"] + if bisheng_options is not None: + _compile_option_list += [ + f"--append-bisheng-options={bisheng_options}" + ] + + npu_compiler_path, env = _get_npucompiler_path() cmd_list = ( - [npu_compiler_path, ttadapter_path] - + _compile_option_list - + ["-o", bin_file] + [npu_compiler_path, src_path] + _compile_option_list + ["-o", bin_file] ) - if dump_ir: - print(f"DEBUG dump ir[bishengir-compile] command: {cmd_list}") try: - ret = subprocess.run(cmd_list, capture_output=True, check=True, text=True) - if dump_ir: - dicp_utils._dump_stage_ir(ret.stderr, metadata["hash"], "bisheng.mlir") + ret = subprocess.run(cmd_list, env=env, capture_output=True, check=True) except subprocess.CalledProcessError as e: - # Print compilation error details - print(f"bishengir-compile compilation failed with exit code {e.returncode}") - print(f"Stderr:\n{e.stderr}") - raise RuntimeError("bishengir-compile compilation failed") from e - - if not Path(bin_path).is_file(): - print(ret.stderr) - if Path(callback_path).is_file(): - lib = ctypes.CDLL(callback_path) - __get_metadata_attr_by_callback( - lib, "_infer_workspace_shape_function", metadata, "workspace_size" - ) - __get_metadata_attr_by_callback( - lib, "_infer_sync_block_lock_num_function", metadata, "lock_num" - ) - __get_metadata_attr_by_callback( - lib, "_infer_sync_block_lock_init_function", metadata, "lock_init_val" + error_msg = e.stderr.decode("utf-8") if e.stderr else str(e) + raise CompileTimeAssertionFailure( + None, None, f"bishengir-compile (SIMT) failed: {error_msg}" + ) from e + if not Path(bin_path).exists(): + error_msg = ret.stderr.decode("utf-8") + print(f"[DEBUG] {bin_path} is not found") + print(f"[DEBUG] Stderr:\n{error_msg}") + raise CompileTimeAssertionFailure( + None, None, f"bishengir-compile (SIMT) output not found: {error_msg}" ) return Path(bin_path).read_bytes() +# --------------------------------------------------------------------------- +# Options dataclasses +# --------------------------------------------------------------------------- + + @dataclass(frozen=True) class NPUOptions: debug: bool = False sanitize_overflow: bool = False - # sanitize_overflow: bool = True - llvm_version: int = 15 + llvm_version: int = 22 kernel_name: str = "triton_" + arch: str = "" cluster_dims: tuple = (1, 1, 1) - num_warps: int = -1 - num_ctas: int = -1 - num_stages: int = 2 + num_warps: int = 32 + num_ctas: int = 1 + num_stages: int = 1 if is_compile_on_910_95 else 2 + warp_size: int = 32 num_buffers_warp_spec: int = 0 num_consumer_groups: int = 0 reg_dec_producer: int = 0 reg_inc_consumer: int = 0 + auto_blockify_size: int = 1 + enable_auto_blockify: bool = None + compile_on_910_95: bool = is_compile_on_910_95 + optimize_dynamic_offset: bool = False + enable_mask_fallback_conversion: bool = False enable_warp_specialization: bool = False enable_nd2nz_on_vector: bool = False enable_persistent: bool = False optimize_epilogue: bool = False enable_fp_fusion: bool = True allow_fp8e4nv: bool = False + auto_tile_and_bind_subblock: bool = True + supported_fp8_dtypes: Tuple[str] = ( + "fp8e5", + "fp8e4b15", + "fp8e4nv", + "fp8e4b8", + "fp8e5b16", + ) + deprecated_fp8_dtypes: Tuple[str] = () + vf_merge_level: int = 1 + default_dot_input_precision: str = "ieee" allowed_dot_input_precisions: Tuple[str] = ("ieee", "hf32") - enable_npu_compile: bool = True - max_num_imprecise_acc_default: bool = None + max_num_imprecise_acc_default: int = 0 extern_libs: dict = None - multibuffer: bool = True - inject_barrier_all: bool = False - disable_auto_inject_block_sync: bool = False - unit_flag: bool = False - disable_auto_cv_work_space_manage: bool = False - enable_auto_bind_sub_block: bool = True - tile_mix_vector_loop: int = None - tile_mix_cube_loop: int = None + bisheng_options: str = "-cce-link-aicore-ll-module " + get_libdevice() + + enable_tuning_mode: bool = False + multibuffer: bool = not is_compile_on_910_95 + enable_ubuf_saving: bool = None + enable_preload: bool = None + enable_auto_bind_sub_block: bool = None + disable_tightly_coupled_buffer_reuse: bool = False + enable_select_analysis: bool = True + enable_hivm_auto_cv_balance: bool = None + sync_solver: bool = None + unit_flag: bool = None + enable_cce_vf_auto_sync: bool = None + enable_cce_vf_remove_membar: bool = None + enable_drop_unit_dims: bool = None + enable_flatten: bool = None + enable_auto_vectorize_v2: bool = None + auto_vectorize_v2_max_fused_ops_num: int = None + prevec_max_fused_ops_num: int = None + inject_barrier_all: bool = None + inject_block_all: bool = None + disable_size_align_for_cast: bool = None limit_auto_multi_buffer_only_for_local_buffer: bool = None + limit_auto_multi_buffer_of_local_buffer: str = None set_workspace_multibuffer: int = None + tile_mix_vector_loop: int = None + tile_mix_cube_loop: int = None + disable_auto_inject_block_sync: bool = None + disable_auto_cv_work_space_manage: bool = False + enable_mixed_cv: bool = None + enable_vf_fusion: bool = False + add_auto_scheduling: bool = False + hfusion_enable_multiple_consumer_fusion: bool = False stream: int = None + parallel_mode: str = "simd" + force_simt_only: bool = False + force_simt_template: bool = False + enable_sync_block_lock: bool = False + shared_mem_dynamic_size: int = None + enable_bishengir_simt_optimization: int = 000 + compile_mode: str = "simd" + mix_mode: str = "" + simt_stack_limit: int = None + enable_simt_reorder_instruction: bool = False + disable_fma: bool = False + + def __post_init__(self): + if self.compile_mode == "simd": + object.__setattr__(self, "parallel_mode", "simd") + elif self.compile_mode == "unstructured_in_simt": + object.__setattr__(self, "force_simt_template", True) + elif self.compile_mode == "simt_only": + object.__setattr__(self, "force_simt_only", True) + object.__setattr__(self, "parallel_mode", "simt") + + if self.force_simt_only: + if self.shared_mem_dynamic_size is None: + object.__setattr__(self, "shared_mem_dynamic_size", 122880) + else: + object.__setattr__(self, "shared_mem_dynamic_size", 221184) def hash(self): key = "_".join([f"{name}-{val}" for name, val in self.__dict__.items()]) - return hashlib.md5(key.encode("utf-8")).hexdigest() + key = "_".join([key, get_cann_version()]) + return hashlib.sha256(key.encode("utf-8")).hexdigest() @dataclass(frozen=True) class CPUOptions: debug: bool = False - llvm_version: int = 15 + llvm_version: int = 22 kernel_name: str = "triton_" cluster_dims: tuple = (1, 1, 1) @@ -922,687 +1167,3 @@ class CPUOptions: def hash(self): key = "_".join([f"{name}-{val}" for name, val in self.__dict__.items()]) return hashlib.md5(key.encode("utf-8")).hexdigest() - - -class NPUUtils(object): - def __new__(cls): - if not hasattr(cls, "instance"): - cls.instance = super(NPUUtils, cls).__new__(cls) - return cls.instance - - def __init__(self): - dirname = os.path.dirname(os.path.realpath(__file__)) - src = Path(os.path.join(dirname, "npu_utils.cpp")).read_text() - key = hashlib.md5(src.encode("utf-8")).hexdigest() - cache = get_cache_manager(key) - fname = "npu_utils.so" - cache_path = cache.get_file(fname) - if cache_path is None: - with tempfile.TemporaryDirectory() as tmpdir: - src_path = os.path.join(tmpdir, "npu_utils.cpp") - with open(src_path, "w") as f: - f.write(src) - so = _build_npu_ext("npu_utils", src_path, tmpdir) - with open(so, "rb") as f: - cache_path = cache.put(f.read(), fname, binary=True) - import importlib.util - - spec = importlib.util.spec_from_file_location("npu_utils", cache_path) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - self.npu_utils_mod = mod - - def load_binary(self, name, kernel, shared, device): - fnname, mix_mode = name.split() - return self.npu_utils_mod.load_kernel_binary( - fnname, kernel, shared, device, mix_mode - ) - - @functools.lru_cache() - def get_device_properties(self, device): - # temperoarily added "max_shared_mem" properties to avoid triton-compiler complain - # fetch available memory at runtime - num_aic = self.get_aicore_num() - num_aiv = num_aic * 2 - return {"max_shared_mem": 1, "num_aicore": num_aic, "num_vectorcore": num_aiv} - - @functools.lru_cache() - def get_arch(self): - # temporarily return empty arch descriptor - return self.npu_utils_mod.get_arch() - - @functools.lru_cache() - def get_aicore_num(self): - # temporarily return empty arch descriptor - return self.npu_utils_mod.get_aicore_num() - - @functools.lru_cache() - def get_aivector_core_num(self): - return self.get_device_properties("npu")["num_vectorcore"] - - -class NPULauncher(object): - def __init__(self, src, metadata): - debug_mode = metadata.debug - workspace_size = ( - int(metadata.workspace_size) if hasattr(metadata, "workspace_size") else -1 - ) - lock_init_value = ( - int(metadata.lock_init_value) if hasattr(metadata, "lock_init_value") else 0 - ) - lock_num = int(metadata.lock_num) if hasattr(metadata, "lock_num") else -1 - constants = src.constants if hasattr(src, "constants") else dict() - cst_key = lambda i: src.fn.arg_names.index(i) if isinstance(i, str) else i - constants = {cst_key(key): value for key, value in constants.items()} - signature = {cst_key(key): value for key, value in src.signature.items()} - mix_mode = metadata.mix_mode - wrapper_src = generate_npu_wrapper_src( - constants, signature, workspace_size, mix_mode, lock_num, lock_init_value - ) - self.so_launcher_path = make_npu_launcher_stub(wrapper_src, debug_mode) - # initialize launcher - import importlib.util - - spec = importlib.util.spec_from_file_location( - "__triton_launcher", self.so_launcher_path - ) - mod = importlib.util.module_from_spec(spec) - spec.loader.exec_module(mod) - self.launch = getattr(mod, "launch") - - def __call__(self, *args, **kwargs): - self.launch(*args, **kwargs) - - -def make_npu_launcher_stub(src, debug=False): - """ - Generate the launcher stub to launch the kernel - """ - # try to get cached file - so_cache_key = hashlib.sha256(src.encode("utf-8")).hexdigest() - so_cache_manager = get_cache_manager(so_cache_key) - # append the cxx11_abi value to the launcher name to avoid - # linking to a launcher with wrong cxx11_abi. - use_cxx11_abi = _check_cxx11_abi() - name = f"launcher_cxx11abi{use_cxx11_abi}" - suffix = sysconfig.get_config_var("EXT_SUFFIX") - so_name = f"{name}{suffix}" - - if debug: - dump_manager = get_dump_manager(so_cache_key) - print(f"Dumping {name}.cxx to {dump_manager.cache_dir}") - dump_manager.put(src, f"{name}.cxx", binary=False) - - cache_path = so_cache_manager.get_file(so_name) - if cache_path is not None: - return cache_path - - with tempfile.TemporaryDirectory() as tmpdir: - if debug: - so_cache_manager.put(src, f"{name}.cxx", binary=False) - src_path = os.path.join(tmpdir, f"{name}.cxx") - with open(src_path, "w") as f: - f.write(src) - enable_taskqueue = os.getenv("TRITON_ENABLE_TASKQUEUE", "true").lower() in ( - "true", - "1", - ) - if enable_taskqueue: - kernel_launcher_type = "torch" - else: - kernel_launcher_type = None - so = _build_npu_ext( - name, src_path, tmpdir, kernel_launcher=kernel_launcher_type - ) - if debug: - with open(so, "rb") as f: - return dump_manager.put(f.read(), so_name, binary=True) - with open(so, "rb") as f: - return so_cache_manager.put(f.read(), so_name, binary=True) - - -def extract_device_print_code_from_cann(): - from triton.backends.dicp_triton.npu import _get_bisheng_path - - ccec_compiler_bin_folder, _ = os.path.split(os.path.realpath(_get_bisheng_path())) - ccec_compiler_folder, _ = os.path.split(ccec_compiler_bin_folder) - clang_version = os.listdir(os.path.join(ccec_compiler_folder, "lib/clang/"))[0] - ccelib_path = os.path.join( - ccec_compiler_folder, f"lib/clang/{clang_version}/include/ccelib" - ) - - def read_header(header_path): - with open(os.path.join(ccelib_path, header_path), "r") as f: - code = f.read() - - # remove all #include "..." - lines = code.splitlines() - purged_lines = [] - for line in lines: - normalized_line = " ".join(line.split()) - if not normalized_line.startswith('#include "'): - purged_lines.append(line) - code = "\n".join(purged_lines) - - # remove [aicore] functions - aicore_positions = [] - for m in re.finditer("\[aicore\]", code): - aicore_positions.append(m.start()) - - def find_aicore_function_span(src, pos): - for i in range(pos - 1, -1, -1): - if ( - src[i] == "}" - ): # this relies on that all [aicore] functions come after normal functions - left = i + 1 - break - n = len(src) - brace_nest = 0 - for j in range(pos, n, 1): - if src[j] == "{": - brace_nest += 1 - elif src[j] == "}": - brace_nest -= 1 - if brace_nest == 0: - right = j - break - return left, right - - new_code = "" - segment_start = 0 - for pos in aicore_positions: - left, right = find_aicore_function_span(code, pos) - new_code += code[segment_start:left] - segment_start = right + 1 - new_code += code[segment_start:] - - # remove __gm__ and rename macros - new_code = new_code.replace("__gm__", " ") - new_code = new_code.replace("__CCELIB_RT_ERROR_NONE", "RT_ERROR_NONE") - new_code = new_code.replace("__CCELIB_RT_MEMORY_HBM", "RT_MEMORY_HBM") - new_code = new_code.replace( - "__CCELIB_RT_MEMCPY_HOST_TO_DEVICE", "RT_MEMCPY_HOST_TO_DEVICE" - ) - new_code = new_code.replace( - "__CCELIB_RT_MEMCPY_DEVICE_TO_HOST", "RT_MEMCPY_DEVICE_TO_HOST" - ) - return new_code - - # the following headers should be included in this order - headers_combined = "\n".join( - [ - read_header("common/common_impl.h"), - read_header("internal/debug_tunnel/payload.h"), - read_header("internal/debug_tunnel/payload_impl.h"), - read_header("internal/debug_tunnel/tunnel.h"), - read_header("internal/debug_tunnel/tunnel_impl.h"), - ] - ) - # Prepend the needed include so generated code has std::cout / std::endl available - return "#include \n" + headers_combined - - -# the template is from triton-adapter HEAD. Wrapping the generated kernel binary into a python module -def generate_npu_wrapper_src( - constants, signature, workspace_size, mix_mode, lock_num, lock_ini_val -): - import os - - # TODO(zmz)，临时方案signature 的value 中，如果有*u1，换成*i1 - signature = {k: v.replace("*u1", "*i1") for k, v in signature.items()} - - def _ty_to_cpp(ty): - if ty[0] == "*": - return "void*" - if ty == "constexpr": - return "PyObject*" - return { - "i1": "int32_t", - "i8": "int8_t", - "i16": "int16_t", - "i32": "int32_t", - "i64": "int64_t", - "u32": "uint32_t", - "u64": "uint64_t", - "fp16": "float", - "bf16": "float", - "fp32": "float", - "f32": "float", - "fp64": "double", - }[ty] - - def _extracted_ty(ty): - if ty[0] == "*": - return "PyObject*" - if ty == "constexpr": - return "PyObject*" - return { - "i1": "int32_t", - "i32": "int32_t", - "i64": "int64_t", - "u32": "uint32_t", - "u64": "uint64_t", - "fp16": "float", - "bf16": "float", - "fp32": "float", - "f32": "float", - "fp64": "double", - }[ty] - - def _format_of(ty): - return { - "PyObject*": "O", - "float": "f", - "double": "d", - "long": "l", - "uint32_t": "I", - "int32_t": "i", - "uint64_t": "K", - "int64_t": "L", - }[ty] - - arg_decls = ", ".join(f"{_ty_to_cpp(ty)} arg{i}" for i, ty in signature.items()) - """ - args: - int gridX, gridY, gridZ; - rtStream_t stream; - const void *functon; - PyObject* packed_metadata, *launch_metadata; - PyObject* launch_enter_hook, *launch_exit_hook; - *args_expand - """ - format = "iiiKKOOOO" + "".join( - [_format_of(_extracted_ty(ty)) for ty in signature.values()] - ) - - grid_info = {"X": "i32", "Y": "i32", "Z": "i32"} - - enable_device_print = os.getenv("TRITON_DEVICE_PRINT", "false").lower() in ( - "true", - "1", - ) - enable_taskqueue = os.getenv("TRITON_ENABLE_TASKQUEUE", "true").lower() in ( - "true", - "1", - ) - enable_auto_map_parallel_blocks = _is_auto_map_parallel_blocks_enabled() - npu_utils = NPUUtils() - num_physical_blocks = ( - npu_utils.get_aivector_core_num() - if mix_mode == "aiv" - else npu_utils.get_aicore_num() - ) - task_type = ( - "MSPROF_GE_TASK_TYPE_AIV" - if mix_mode == "aiv" - else "MSPROF_GE_TASK_TYPE_AI_CORE" - ) - LINE_CHANGE_CHAR = chr(10) # it is \n - - cpp_device_pointer = """ -typedef struct _DevicePtrInfo { - void *dev_ptr; - bool valid; -} DevicePtrInfo; - -static inline DevicePtrInfo getPointer(PyObject *obj, int idx) { - DevicePtrInfo ptr_info; - ptr_info.dev_ptr = 0; - ptr_info.valid = true; - if (PyLong_Check(obj)) { - ptr_info.dev_ptr = reinterpret_cast(PyLong_AsUnsignedLongLong(obj)); - return ptr_info; - } - if (obj == Py_None) { - // valid nullptr - return ptr_info; - } - PyObject *ptr = PyObject_GetAttrString(obj, "data_ptr"); - if(ptr){ - PyObject *empty_tuple = PyTuple_New(0); - PyObject *ret = PyObject_Call(ptr, empty_tuple, NULL); - Py_DECREF(empty_tuple); - Py_DECREF(ptr); - if (!PyLong_Check(ret)) { - PyErr_SetString(PyExc_TypeError, "data_ptr method of Pointer object must return 64-bit int"); - ptr_info.valid = false; - return ptr_info; - } - ptr_info.dev_ptr = reinterpret_cast(PyLong_AsUnsignedLongLong(ret)); - if(!ptr_info.dev_ptr) - return ptr_info; - Py_DECREF(ret); // Thanks ChatGPT! - return ptr_info; - } - PyErr_SetString(PyExc_TypeError, "Pointer argument must be either uint64 or have data_ptr method"); - return ptr_info; -} -""" - - cpp_msprof_extern = """ -extern "C" { - typedef int (* callback)(unsigned int type, void* data, unsigned int len); - extern int MsprofReportApi(unsigned int agingFlag, const MsprofApi *api); - extern unsigned long int MsprofSysCycleTime(); - extern int MsprofRegisterCallback(unsigned int moduleId, callback handle); - static unsigned int __MsprofFlagL0 = 0; - static unsigned int __MsprofFlagL1 = 0; - - int ProfCtrlHandle(unsigned int CtrlType, void* CtrlData, unsigned int DataLen) { - if ((CtrlData == nullptr) || (DataLen == 0U)) { - return 1; - } - - if (CtrlType == 1) { - MsprofCommandHandle* handle = (MsprofCommandHandle *)(CtrlData); - if (handle->type >= 6) // 6 is not used here - return 1; - if (handle->type == 1) { // init - 0 , start - 1 - __MsprofFlagL0 = ((0x00000800ULL & handle->profSwitch) == 0x00000800ULL) ? 1 : 0; - __MsprofFlagL1 = ((0x00000002ULL & handle->profSwitch) == 0x00000002ULL) ? 1 : 0; - } - } - return 0; - } -} -""" - - cpp_msprof_callback = """ - MsprofRegisterCallback(8, ProfCtrlHandle); // 8 - CCE defined in msprof headerfile slog.h -""" - - cpp_msprof_call_before_launch = """ - unsigned long int beginTime = 0; - unsigned long int endTime = 0; - unsigned long int opNameHashID = 0; - unsigned int threadId = 0; - char* _kernelName = const_cast(name.c_str()); - size_t length = name.length(); - if (__MsprofFlagL0 || __MsprofFlagL1) - { - beginTime = MsprofSysCycleTime(); - } -""" - - cpp_msprof_call_after_launch = f""" - if (__MsprofFlagL0 || __MsprofFlagL1) - {{ - endTime = MsprofSysCycleTime(); - opNameHashID = MsprofGetHashId(_kernelName, length); - threadId = (unsigned int)(syscall(SYS_gettid)); - MsprofApi info; - info.level = MSPROF_REPORT_NODE_LEVEL; - info.magicNumber = 0x5a5a; //MSPROF_REPORT_DATA_MAGIC_NUM - info.type = MSPROF_REPORT_NODE_LAUNCH_TYPE; - info.threadId = threadId; - info.reserve = 0; - info.beginTime = beginTime; - info.endTime = endTime; - info.itemId = opNameHashID; - MsprofReportApi(false, &info); - }} - if (__MsprofFlagL1) - {{ - MsprofCompactInfo nodeBasicInfo; - nodeBasicInfo.level = MSPROF_REPORT_NODE_LEVEL; - nodeBasicInfo.magicNumber = 0x5a5a; //MSPROF_REPORT_DATA_MAGIC_NUM - nodeBasicInfo.type = MSPROF_REPORT_NODE_BASIC_INFO_TYPE; - nodeBasicInfo.threadId = threadId; - nodeBasicInfo.timeStamp = endTime; - nodeBasicInfo.data.nodeBasicInfo.opName = opNameHashID; - nodeBasicInfo.data.nodeBasicInfo.opType = opNameHashID; - nodeBasicInfo.data.nodeBasicInfo.taskType = {task_type}; - nodeBasicInfo.data.nodeBasicInfo.blockDim = blockNum; - MsprofReportCompactInfo(0, static_cast(&nodeBasicInfo), sizeof(MsprofCompactInfo)); - - // Report tensor info - int max_tensors_num = tensorShapes.size() < MSPROF_GE_TENSOR_DATA_NUM ? tensorShapes.size() : MSPROF_GE_TENSOR_DATA_NUM; - MsprofAdditionalInfo tensorInfo; - tensorInfo.level = MSPROF_REPORT_NODE_LEVEL; - tensorInfo.type = MSPROF_REPORT_NODE_TENSOR_INFO_TYPE; - tensorInfo.threadId = threadId; - tensorInfo.timeStamp = endTime; - auto profTensorData = reinterpret_cast(tensorInfo.data); - profTensorData->opName = opNameHashID; - int tensorCount = 0; - int dataTypes[MSPROF_GE_TENSOR_DATA_NUM]; - if (tensorShapes.size() > 0) {{ - {LINE_CHANGE_CHAR.join( - f'dataTypes[{i}] = {convert_sigtype_to_int(ty[1:])};' - for i, ty in signature.items() - if ty.startswith("*") and i < 5 - )} - }} - for (int i = 0; i < tensorShapes.size() && tensorCount < MSPROF_GE_TENSOR_DATA_NUM; i++) {{ - auto fillTensorData = [&](int index, int tensorType) {{ - profTensorData->tensorData[index].tensorType = tensorType; - profTensorData->tensorData[index].format = 2; // GeDataFormat: ND = 2 - profTensorData->tensorData[index].dataType = dataTypes[i]; - int nDim = tensorShapes[i].size(); - nDim = nDim < MSPROF_GE_TENSOR_DATA_SHAPE_LEN ? nDim : MSPROF_GE_TENSOR_DATA_SHAPE_LEN; - for (int j = 0; j < nDim; j++) {{ - profTensorData->tensorData[index].shape[j] = tensorShapes[i][j]; - }} - for (int j = nDim; j < MSPROF_GE_TENSOR_DATA_SHAPE_LEN; j++) {{ - profTensorData->tensorData[index].shape[j] = 0; - }} - }}; - int tensorType = (i < tensorKinds.size()) ? tensorKinds[i] : 0; // DeFault tensor type is input - if (tensorType == TENSOR_KIND_INPUT || tensorType == TENSOR_KIND_INPUT_OUTPUT) {{ - fillTensorData(tensorCount, MSPROF_GE_TENSOR_TYPE_INPUT); - tensorCount++; - }} - if ((tensorType == TENSOR_KIND_OUTPUT || tensorType == TENSOR_KIND_INPUT_OUTPUT) && tensorCount < MSPROF_GE_TENSOR_DATA_NUM){{ - fillTensorData(tensorCount, MSPROF_GE_TENSOR_TYPE_OUTPUT); - tensorCount++; - }} - }} - profTensorData->tensorNum = tensorCount; - MsprofReportAdditionalInfo(false, static_cast(&tensorInfo), sizeof(MsprofAdditionalInfo)); - }} -""" - - return f""" -#include -#include -#include -#include -#include -#define PY_SSIZE_T_CLEAN -#include -{'#include ' if enable_taskqueue else ''} -#include "/usr/local/Ascend/cann/pkg_inc/runtime/runtime/rt.h" -{extract_device_print_code_from_cann() if enable_device_print else ''} - -#define TENSOR_KIND_INPUT 0 -#define TENSOR_KIND_OUTPUT 1 -#define TENSOR_KIND_INPUT_OUTPUT 2 - -{cpp_msprof_extern} - -{cpp_device_pointer} - -static void _launch(const char* kernelName, const void* func, rtStream_t stream, int gridX, int gridY, int gridZ, std::vector> &tensorShapes, std::vector &tensorKinds{', ' + arg_decls if len(signature) > 0 else ''}) {{ - // only 1D parallelization is supported for NPU - // Pointer type becomes flattend 1-D Memref tuple: base_ptr, data_ptr, offset, shape, stride - // base_ptr offset shape and stride are not used, arbitrarily set for now - std::string name = ""; - name.append(kernelName); - {'auto launch_call = [=]()' if enable_taskqueue else ''} {{ - uint32_t blockNum = gridX * gridY * gridZ; - {'if (blockNum > (uint32_t)' + str(num_physical_blocks) + ') { /* std::cout << "WARNING: Grid " << blockNum << " > physical limit ' + str(num_physical_blocks) + ', performance maybe reduced." << std::endl; */ if (blockNum > 65535 && !' + str(enable_auto_map_parallel_blocks).lower() + ') {std::cout << "Grid " << blockNum << " > 65535, Please set TRITON_ALL_BLOCKS_PARALLEL=1 to enable all blocks parallel execution." << std::endl; } }'} - - {'blockNum = std::min(blockNum, (uint32_t)' + str(num_physical_blocks) + ');' if enable_auto_map_parallel_blocks else ''} - {'cce::internal::DebugTunnelData *DTData = cce::internal::DebugTunnel::Open(blockNum);' if enable_device_print else ''} - rtError_t ret; - void *ffts_addr = NULL; - uint32_t ffts_len; ret = rtGetC2cCtrlAddr((uint64_t*)&ffts_addr, &ffts_len); - if (ret != RT_ERROR_NONE) {{ - return {'ret' if enable_taskqueue else ''}; - }} - // stub argument for workspace - void *syncBlockLock = NULL; - void *workspace_addr = NULL; - uint16_t ModuleId = 0; - {f''' - uint64_t syncBlockLockSize = {lock_num} * sizeof(int64_t); - ret = rtMalloc(reinterpret_cast(&syncBlockLock), - syncBlockLockSize, RT_MEMORY_HBM, 0); - if (ret != RT_ERROR_NONE) {{ - return {'ret' if enable_taskqueue else ''}; - }} - std::vector lockInitData({lock_num}, {lock_ini_val}); - ret = rtMemcpy(syncBlockLock, syncBlockLockSize, reinterpret_cast(lockInitData.data()), - syncBlockLockSize, RT_MEMCPY_HOST_TO_DEVICE); - if (ret != RT_ERROR_NONE) {{ - return {'ret' if enable_taskqueue else ''}; - }} - ''' if lock_num > 0 else ''} - {f''' - uint64_t totalWorkSpaceSize = {workspace_size} * blockNum; - ret = rtMalloc(reinterpret_cast(&workspace_addr), - totalWorkSpaceSize, RT_MEMORY_HBM, ModuleId); - if (ret != RT_ERROR_NONE) {{ - return {'ret' if enable_taskqueue else ''}; - }} - ''' if workspace_size > 0 else ''} - struct __attribute__((packed)) {{ - void* ffts_addr __attribute__((aligned(8))); - void* syncBlockLock __attribute__((aligned(8))); - void* workspace_addr __attribute__((aligned(8))); - {' '.join(f'{_ty_to_cpp(ty)} arg{i} __attribute__((aligned({4 if ty[0] != "*" and ty[-2:] != "64" else 8})));' for i, ty in signature.items() if ty != "constexpr")} - {' '.join(f'{_ty_to_cpp(ty)} grid{mark} __attribute__((aligned(4)));' for mark, ty in grid_info.items() if ty != "constexpr")} - {'void* DTData __attribute__((aligned(8)));' if enable_device_print else ''} - }} args = {{ - static_cast(ffts_addr), - static_cast(syncBlockLock), - static_cast(workspace_addr), - {(', '.join(f'static_cast<{_ty_to_cpp(ty)}>(arg{i})' for i, ty in signature.items() if ty != "constexpr") + ',') if len(signature) > 0 else ''} - {', '.join(f'static_cast<{_ty_to_cpp(ty)}>(grid{mark})' for mark, ty in grid_info.items() if ty != "constexpr")} - {', static_cast(DTData)' if enable_device_print else ''} - }}; - {cpp_msprof_call_before_launch} - ret = rtKernelLaunch(func, blockNum, static_cast(&args), sizeof(args), NULL, stream); - {'void *&stream_ref = const_cast(stream);' if enable_device_print else ''} - {'cce::internal::DebugTunnel::Close(DTData, stream_ref);' if enable_device_print else ''} - {cpp_msprof_call_after_launch} - {'return ret;' if enable_taskqueue else ''} - }}; - {'at_npu::native::OpCommand cmd; cmd.Name(name.c_str()).SetCustomHandler(launch_call).Run();' if enable_taskqueue else ''} - return; -}} - -// Extract tensor shape from PyObject -static std::vector _get_tensor_shape(PyObject *tensor) {{ - std::vector shape; - - // Early return if tensor is None or null - if (!tensor || tensor == Py_None) {{ - return shape; - }} - - // Calling tensor.size() - PyObject* size_result = PyObject_CallMethod(tensor, "size", NULL); - if (!size_result) {{ - return shape; - }} - // Using PySequence_Fast to improve access efficiency - PyObject* seq = PySequence_Fast(size_result, "Expected a sequence from tensor.size()"); - if (seq) {{ - Py_ssize_t len = PySequence_Fast_GET_SIZE(seq); - PyObject** items = PySequence_Fast_ITEMS(seq); - for (Py_ssize_t i = 0; i < len; ++i) {{ - PyObject* dim = items[i]; - if (PyLong_Check(dim)) {{ - shape.push_back(PyLong_AsLong(dim)); - }} - }} - }} - Py_DECREF(seq); - Py_DECREF(size_result); - return shape; -}} - -static PyObject* launch(PyObject* self, PyObject* args) {{ - int gridX, gridY, gridZ; - rtStream_t stream; - const void *function; - PyObject *packedMetadata = NULL; - PyObject *launch_metadata = NULL; - PyObject *launch_enter_hook = NULL; - PyObject *launch_exit_hook = NULL; - std::vector> tensorShapes; - {' '.join([f"{_extracted_ty(ty)} _arg{i}; " for i, ty in signature.items()])} - if(!PyArg_ParseTuple( - args, \"{format}\", - &gridX, &gridY, &gridZ, &stream, &function, - &packedMetadata, &launch_metadata, - &launch_enter_hook, &launch_exit_hook - {', ' + ', '.join(f"&_arg{i}" for i, ty in signature.items()) if len(signature) > 0 else ''} - ) - ) {{ - return NULL; - }} - if (__MsprofFlagL1) - {{ - { - LINE_CHANGE_CHAR.join( - f"{{ auto tmp = _get_tensor_shape(_arg{i}); if (!tmp.empty()) tensorShapes.push_back(tmp); }}" - for i, ty in signature.items() if ty[0] == "*" - ) - } - }} - - if (launch_enter_hook != Py_None && !PyObject_CallObject(launch_enter_hook, args)) {{ - return NULL; - }} - - // get kernel_name - PyObject *kernelNameObj = PyDict_GetItemString(packedMetadata, "kernel_name"); - const char *kernelName = PyUnicode_AsUTF8(kernelNameObj); - // get tensor_kinds - std::vector tensorKinds; - PyObject *tensorKindList = PyDict_GetItemString(packedMetadata, "tensor_kinds"); - if (tensorKindList) {{ - int size = PyObject_Size(tensorKindList); - for (int i = 0; i < size; i++) {{ - PyObject *kind = PySequence_GetItem(tensorKindList, i); - tensorKinds.push_back(PyLong_AsLong(kind)); - }} - }} - - // raise exception asap - {"; ".join([f"DevicePtrInfo ptr_info{i} = getPointer(_arg{i}, {i}); if (!ptr_info{i}.valid) return NULL;" if ty[0]=="*" else "" for i, ty in signature.items()])}; - _launch(kernelName, function, stream, gridX, gridY, gridZ, tensorShapes, tensorKinds{', ' + ', '.join(f"ptr_info{i}.dev_ptr" if ty[0]=="*" else f"_arg{i}" for i, ty in signature.items()) if len(signature) > 0 else ''}); - if (PyErr_Occurred()) {{ - return NULL; - }} - if (launch_exit_hook != Py_None && !PyObject_CallObject(launch_exit_hook, args)) {{ - return NULL; - }} - Py_RETURN_NONE; -}} - -static PyMethodDef ModuleMethods[] = {{ - {{"launch", launch, METH_VARARGS, "Entry point for all kernels with this signature"}}, - {{NULL, NULL, 0, NULL}} // sentinel -}}; - -static struct PyModuleDef ModuleDef = {{ - PyModuleDef_HEAD_INIT, - \"__triton_launcher\", - NULL, //documentation - -1, //size - ModuleMethods -}}; - -PyMODINIT_FUNC PyInit___triton_launcher(void) {{ - PyObject *m = PyModule_Create(&ModuleDef); - if(m == NULL) {{ - return NULL; - }} - PyModule_AddFunctions(m, ModuleMethods); - {cpp_msprof_callback} - return m; -}} -""" diff --git a/backend/npu_driver.py b/backend/npu_driver.py new file mode 100644 index 00000000..c53e4d10 --- /dev/null +++ b/backend/npu_driver.py @@ -0,0 +1,969 @@ +from pathlib import Path +import tempfile +import os +import re +import subprocess +import sysconfig +import functools +import hashlib +from triton.runtime.cache import get_cache_manager, get_dump_manager +from triton.backends.compiler import GPUTarget + +from .utils import ( + TRITON_PROFILER_REGISTERED, + _check_cxx11_abi, + _get_ascend_path, + _get_bisheng_path, + _build_npu_ext, + _precompile_npu_hash, + _precompile_npu_ext, + _precompile_npu_ext_with_lock, + _is_auto_map_parallel_blocks_enabled, + convert_sigtype_to_int, + get_ascend_arch_from_env, + is_ffts_supported, + force_disable_ffts, + get_backend_func, +) + +# --------------------------------------------------------------------------- +# NPUUtils — singleton that loads npu_utils.cpp +# --------------------------------------------------------------------------- + + +class NPUUtils(object): + def __new__(cls): + if not hasattr(cls, "instance"): + cls.instance = super(NPUUtils, cls).__new__(cls) + return cls.instance + + def __init__(self): + dirname = os.path.dirname(os.path.realpath(__file__)) + src = Path(os.path.join(dirname, "npu_utils.cpp")).read_text() + key = hashlib.md5(src.encode("utf-8")).hexdigest() + cache = get_cache_manager(key) + fname = "npu_utils.so" + cache_path = cache.get_file(fname) + if cache_path is None: + with tempfile.TemporaryDirectory() as tmpdir: + src_path = os.path.join(tmpdir, "npu_utils.cpp") + with open(src_path, "w") as f: + f.write(src) + so = _build_npu_ext("npu_utils", None, src_path, kernel_launcher=None) + with open(so, "rb") as f: + cache_path = cache.put(f.read(), fname, binary=True) + import importlib.util + + spec = importlib.util.spec_from_file_location("npu_utils", cache_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + self.npu_utils_mod = mod + + def load_binary(self, name, kernel, shared, device, mix_mode): + return self.npu_utils_mod.load_kernel_binary( + name, kernel, shared, device, mix_mode + ) + + @functools.lru_cache() + def get_device_properties(self, device): + num_aic = self.get_aicore_num() + num_aiv = num_aic * 2 + return {"max_shared_mem": 1, "num_aicore": num_aic, "num_vectorcore": num_aiv} + + @functools.lru_cache() + def get_arch(self): + return self.npu_utils_mod.get_arch() + + @functools.lru_cache() + def get_aicore_num(self): + return self.npu_utils_mod.get_aicore_num() + + @functools.lru_cache() + def get_aivector_core_num(self): + return self.get_device_properties("npu")["num_vectorcore"] + + +# --------------------------------------------------------------------------- +# Header source generation (precompiled separately from wrapper) +# --------------------------------------------------------------------------- + + +def generate_npu_header_src(): + enable_taskqueue = os.getenv("TRITON_ENABLE_TASKQUEUE", "true").lower() in ( + "true", + "1", + ) + return f""" +#ifndef TRITON_NPU_HEADERS +#define TRITON_NPU_HEADERS + +#include +#include +#include +#include +#include +#include +#include "runtime/runtime/rt.h" +#include +{get_backend_func("header_file", enable_taskqueue)} + +#endif +""" + + +# --------------------------------------------------------------------------- +# Device-print code extraction from CANN ccelib +# --------------------------------------------------------------------------- + + +def extract_device_print_code_from_cann(): + ccec_compiler_bin_folder, _ = os.path.split(os.path.realpath(_get_bisheng_path())) + ccec_compiler_folder, _ = os.path.split(ccec_compiler_bin_folder) + clang_version = os.listdir(os.path.join(ccec_compiler_folder, "lib/clang/"))[0] + ccelib_path = os.path.join( + ccec_compiler_folder, f"lib/clang/{clang_version}/include/ccelib" + ) + + def read_header(header_path): + with open(os.path.join(ccelib_path, header_path), "r") as f: + code = f.read() + + # remove all #include "..." + lines = code.splitlines() + purged_lines = [] + for line in lines: + normalized_line = " ".join(line.split()) + if not normalized_line.startswith('#include "'): + purged_lines.append(line) + code = "\n".join(purged_lines) + + # remove [aicore] functions + aicore_positions = [] + for m in re.finditer(r"\[aicore\]", code): + aicore_positions.append(m.start()) + + def find_aicore_function_span(src, pos): + for i in range(pos - 1, -1, -1): + if src[i] == "}": + left = i + 1 + break + n = len(src) + brace_nest = 0 + for j in range(pos, n, 1): + if src[j] == "{": + brace_nest += 1 + elif src[j] == "}": + brace_nest -= 1 + if brace_nest == 0: + right = j + break + return left, right + + new_code = "" + segment_start = 0 + for pos in aicore_positions: + left, right = find_aicore_function_span(code, pos) + new_code += code[segment_start:left] + segment_start = right + 1 + new_code += code[segment_start:] + + new_code = new_code.replace("__gm__", " ") + new_code = new_code.replace("__CCELIB_RT_ERROR_NONE", "RT_ERROR_NONE") + new_code = new_code.replace("__CCELIB_RT_MEMORY_HBM", "RT_MEMORY_HBM") + new_code = new_code.replace( + "__CCELIB_RT_MEMCPY_HOST_TO_DEVICE", "RT_MEMCPY_HOST_TO_DEVICE" + ) + new_code = new_code.replace( + "__CCELIB_RT_MEMCPY_DEVICE_TO_HOST", "RT_MEMCPY_DEVICE_TO_HOST" + ) + return new_code + + headers_combined = "\n".join( + [ + read_header("common/common_impl.h"), + read_header("internal/debug_tunnel/payload.h"), + read_header("internal/debug_tunnel/payload_impl.h"), + read_header("internal/debug_tunnel/tunnel.h"), + read_header("internal/debug_tunnel/tunnel_impl.h"), + ] + ) + return "#include \n" + headers_combined + + +# --------------------------------------------------------------------------- +# Type helpers (merged _ty_to_cpp + _extracted_ty) +# --------------------------------------------------------------------------- + + +def ty_to_cpp(ty): + if ty[0] == "*": + return "void*" + return { + "i1": "int32_t", + "i8": "int8_t", + "i16": "int16_t", + "i32": "int32_t", + "i64": "int64_t", + "u1": "uint32_t", + "u8": "uint8_t", + "u16": "uint16_t", + "u32": "uint32_t", + "u64": "uint64_t", + "fp16": "float", + "bf16": "float", + "fp32": "float", + "f32": "float", + "fp64": "double", + }[ty] + + +def extracted_ty(ty): + if ty[0] == "*": + return "PyObject*" + if ty == "constexpr": + return "PyObject*" + return ty_to_cpp(ty) + + +def format_of(ty): + if ty[0] == "*": + return "O" + if ty == "constexpr": + return "O" + return { + "float": "f", + "double": "d", + "long": "l", + "int8_t": "b", + "int16_t": "h", + "int32_t": "i", + "int64_t": "L", + "uint8_t": "B", + "uint16_t": "H", + "uint32_t": "I", + "uint64_t": "K", + }[ty_to_cpp(ty)] + + +def _format_of_msprof_task_type_ratio(bs_task_type, mix_mode): + default_task_type = ( + "MSPROF_GE_TASK_TYPE_AIV" + if mix_mode == "aiv" + else "MSPROF_GE_TASK_TYPE_AI_CORE" + ) + if not bs_task_type: + return default_task_type, 0 + task_type_num, mix_block_dim_ratio = divmod(int(bs_task_type), 10) + task_type_map = { + 1: "MSPROF_GE_TASK_TYPE_AIV", + 2: "MSPROF_GE_TASK_TYPE_AI_CORE", + 3: "MSPROF_GE_TASK_TYPE_MIX_AIC", + 4: "MSPROF_GE_TASK_TYPE_MIX_AIV", + } + return task_type_map.get(task_type_num, default_task_type), mix_block_dim_ratio + + +# --------------------------------------------------------------------------- +# Wrapper source generation +# --------------------------------------------------------------------------- + + +def generate_npu_wrapper_src( + constants, + signature, + workspace_size, + mix_mode, + lock_num, + lock_init_value, + *, + bs_task_type=0, + compile_on_910_95=False, + force_simt_only=False, + shared_mem_dynamic_size=0, + parallel_mode="simd", + enable_auto_blockify=None, +): + import os + + # TODO(zmz): temporary workaround — signature values with *u1 should become *i1 + signature = {k: v.replace("*u1", "*i1") for k, v in signature.items()} + + def _serialize_signature(sig): + if isinstance(sig, tuple): + return ",".join(map(_serialize_signature, sig)) + return sig + + def _extracted_type(ty): + if isinstance(ty, tuple): + val = ",".join(map(_extracted_type, ty)) + return f"[{val}]" + if ty[0] == "*": + return "PyObject*" + if ty == "constexpr": + return "PyObject*" + return ty_to_cpp(ty) + + def _format_of(ty): + if isinstance(ty, tuple): + val = "".join(map(_format_of, ty)) + return f"({val})" + if ty[0] == "*": + return "O" + if ty == "constexpr": + return "O" + if ty == "void*": + return "O" + return { + "float": "f", + "double": "d", + "long": "l", + "int8_t": "b", + "int16_t": "h", + "int32_t": "i", + "int64_t": "L", + "uint8_t": "B", + "uint16_t": "H", + "uint32_t": "I", + "uint64_t": "K", + }[ty_to_cpp(ty)] + + # Compute args_format BEFORE re-indexing (needs original signature types) + args_format = "".join([_format_of(ty) for ty in signature.values()]) + fmt = "iiiKKOOOO" + args_format + + # Serialize and re-index signature: flatten tuples, re-number from 0 + signature = ",".join(map(_serialize_signature, signature.values())) + signature = list(filter(bool, signature.split(","))) + signature = {i: s for i, s in enumerate(signature)} + + args_list = ( + ", " + ", ".join(f"&_arg{i}" for i, ty in signature.items()) + if len(signature) > 0 + else "" + ) + + arg_decls = ", ".join( + f"{ty_to_cpp(ty)} arg{i}" for i, ty in signature.items() if ty != "constexpr" + ) + internal_args_list = [] + for i, ty in signature.items(): + if ty[0] == "*": + internal_args_list.append(f"ptr_info{i}.dev_ptr") + elif ty != "constexpr": + internal_args_list.append(f"_arg{i}") + + grid_info = {"X": "i32", "Y": "i32", "Z": "i32"} + + newline = "\n " + ptr_decls = [ + f"DevicePtrInfo ptr_info{i} = getPointer(_arg{i}, {i}); if (!ptr_info{i}.valid) return NULL;" + for i, ty in signature.items() + if ty[0] == "*" + ] + + enable_device_print = os.getenv("TRITON_DEVICE_PRINT", "false").lower() in ( + "true", + "1", + ) + enable_taskqueue = os.getenv("TRITON_ENABLE_TASKQUEUE", "true").lower() in ( + "true", + "1", + ) + enable_grid_warn_print = os.getenv("TRITON_GRID_WARN_PRINT", "false").lower() in ( + "true", + "1", + ) + enable_auto_map_parallel_blocks = enable_auto_blockify + if enable_auto_map_parallel_blocks is None: + enable_auto_map_parallel_blocks = _is_auto_map_parallel_blocks_enabled() + npu_utils = NPUUtils() + num_physical_blocks = ( + npu_utils.get_aivector_core_num() + if mix_mode == "aiv" + else npu_utils.get_aicore_num() + ) + + task_type, mix_block_dim_ratio = _format_of_msprof_task_type_ratio( + bs_task_type, mix_mode + ) + is_mix_task_type = "true" if ("MIX" in task_type) else "false" + LINE_CHANGE_CHAR = chr(10) + alloc_success_code = "return 1;" + sync_lock_fail_code = ( + 'fprintf(stderr, "Error: syncBlockLock allocation failed\\n"); return;' + ) + workspace_fail_code = ( + 'fprintf(stderr, "Error: workspace allocation failed\\n"); return;' + ) + + enable_simt = ("simt" in parallel_mode) or force_simt_only + + arch = get_ascend_arch_from_env() + target_support_ffts = is_ffts_supported(arch) and (not force_disable_ffts()) + + cpp_device_pointer = """ +typedef struct _DevicePtrInfo { + void *dev_ptr; + bool valid; +} DevicePtrInfo; + +static inline DevicePtrInfo getPointer(PyObject *obj, int idx) { + DevicePtrInfo ptr_info; + ptr_info.dev_ptr = 0; + ptr_info.valid = true; + if (PyLong_Check(obj)) { + ptr_info.dev_ptr = reinterpret_cast(PyLong_AsUnsignedLongLong(obj)); + return ptr_info; + } + if (obj == Py_None) { + return ptr_info; + } + PyObject *ptr = PyObject_GetAttrString(obj, "data_ptr"); + if(ptr){ + PyObject *empty_tuple = PyTuple_New(0); + PyObject *ret = PyObject_Call(ptr, empty_tuple, NULL); + Py_DECREF(empty_tuple); + Py_DECREF(ptr); + if (!PyLong_Check(ret)) { + PyErr_SetString(PyExc_TypeError, "data_ptr method of Pointer object must return 64-bit int"); + ptr_info.valid = false; + return ptr_info; + } + ptr_info.dev_ptr = reinterpret_cast(PyLong_AsUnsignedLongLong(ret)); + if(!ptr_info.dev_ptr) + return ptr_info; + aclrtPtrAttributes attributes; + aclError status = aclrtPointerGetAttributes(ptr_info.dev_ptr, &attributes); + if (status == ACL_SUCCESS) { + if (attributes.location.type != ACL_MEM_LOCATION_TYPE_DEVICE && attributes.location.type != 4) { + Py_DECREF(ret); + PyErr_Format(PyExc_ValueError, + "Pointer argument (at %d) cannot be accessed from Triton (cpu tensor?)", idx); + ptr_info.valid = false; + return ptr_info; + } + } else { + Py_DECREF(ret); + PyErr_Format(PyExc_RuntimeError, + "Failed to query pointer attributes at argument %d. " + "Error code: %d. This may indicate invalid memory address " + "or NPU device error.", + idx, status); + ptr_info.valid = false; + return ptr_info; + } + Py_DECREF(ret); + return ptr_info; + } + PyErr_SetString(PyExc_TypeError, "Pointer argument must be either uint64 or have data_ptr method"); + ptr_info.valid = false; + return ptr_info; +} +""" + + cpp_msprof_extern = """ +extern "C" { + typedef int (* callback)(unsigned int type, void* data, unsigned int len); + extern int MsprofReportApi(unsigned int agingFlag, const MsprofApi *api); + extern unsigned long int MsprofSysCycleTime(); + extern int MsprofRegisterCallback(unsigned int moduleId, callback handle); + static unsigned int __MsprofFlagL0 = 0; + static unsigned int __MsprofFlagL1 = 0; + + int ProfCtrlHandle(unsigned int CtrlType, void* CtrlData, unsigned int DataLen) { + if ((CtrlData == nullptr) || (DataLen == 0U)) { + return 1; + } + + if (CtrlType == 1) { + MsprofCommandHandle* handle = (MsprofCommandHandle *)(CtrlData); + if (handle->type >= 6) + return 1; + if (handle->type == 1) { + __MsprofFlagL0 = ((0x00000800ULL & handle->profSwitch) == 0x00000800ULL) ? 1 : 0; + __MsprofFlagL1 = ((0x00000002ULL & handle->profSwitch) == 0x00000002ULL) ? 1 : 0; + } + } + return 0; + } +} +""" + + cpp_msprof_callback = """ + MsprofRegisterCallback(8, ProfCtrlHandle); +""" + + cpp_msprof_call_before_launch = """ + unsigned long int beginTime = 0; + unsigned long int endTime = 0; + unsigned long int opNameHashID = 0; + unsigned int threadId = 0; + char* _kernelName = const_cast(name.c_str()); + size_t length = name.length(); + if (__MsprofFlagL0 || __MsprofFlagL1) + { + beginTime = MsprofSysCycleTime(); + } +""" + + cpp_msprof_call_after_launch = f""" + if (__MsprofFlagL0 || __MsprofFlagL1) + {{ + endTime = MsprofSysCycleTime(); + opNameHashID = MsprofGetHashId(_kernelName, length); + threadId = (unsigned int)(syscall(SYS_gettid)); + MsprofApi info; + info.level = MSPROF_REPORT_NODE_LEVEL; + info.magicNumber = 0x5a5a; + info.type = MSPROF_REPORT_NODE_LAUNCH_TYPE; + info.threadId = threadId; + info.reserve = 0; + info.beginTime = beginTime; + info.endTime = endTime; + info.itemId = opNameHashID; + MsprofReportApi(false, &info); + }} + if (__MsprofFlagL1) + {{ + MsprofCompactInfo nodeBasicInfo; + nodeBasicInfo.level = MSPROF_REPORT_NODE_LEVEL; + nodeBasicInfo.magicNumber = 0x5a5a; + nodeBasicInfo.type = MSPROF_REPORT_NODE_BASIC_INFO_TYPE; + nodeBasicInfo.threadId = threadId; + nodeBasicInfo.timeStamp = endTime; + nodeBasicInfo.data.nodeBasicInfo.opName = opNameHashID; + nodeBasicInfo.data.nodeBasicInfo.opType = opNameHashID; + nodeBasicInfo.data.nodeBasicInfo.taskType = {task_type}; + nodeBasicInfo.data.nodeBasicInfo.blockDim = nodeBasicBlockDim; + MsprofReportCompactInfo(0, static_cast(&nodeBasicInfo), sizeof(MsprofCompactInfo)); + + // 'mix' kernel need to report the ctxID + if ({is_mix_task_type} > 0) {{ + MsprofAdditionalInfo info; + info.level = MSPROF_REPORT_NODE_LEVEL; + info.type = MSPROF_REPORT_NODE_CONTEXT_ID_INFO_TYPE; + info.threadId = threadId; + info.timeStamp = endTime; + MsprofContextIdInfo ctxId; + ctxId.opName = opNameHashID; + ctxId.ctxIdNum = 1; + for (uint32_t i = 0; i < ctxId.ctxIdNum; i++) {{ + ctxId.ctxIds[i] = i; + }} + size_t copyLen = sizeof(MsprofContextIdInfo); + if (copyLen > MSPROF_ADDTIONAL_INFO_DATA_LENGTH) {{ + copyLen = MSPROF_ADDTIONAL_INFO_DATA_LENGTH; + }} + memcpy(info.data, &ctxId, copyLen); + MsprofReportAdditionalInfo(false, static_cast(&info), sizeof(MsprofAdditionalInfo)); + }} + + // Report tensor info + int max_tensors_num = tensorShapes.size() < MSPROF_GE_TENSOR_DATA_NUM ? tensorShapes.size() : MSPROF_GE_TENSOR_DATA_NUM; + MsprofAdditionalInfo tensorInfo; + tensorInfo.level = MSPROF_REPORT_NODE_LEVEL; + tensorInfo.type = MSPROF_REPORT_NODE_TENSOR_INFO_TYPE; + tensorInfo.threadId = threadId; + tensorInfo.timeStamp = endTime; + auto profTensorData = reinterpret_cast(tensorInfo.data); + profTensorData->opName = opNameHashID; + int tensorCount = 0; + int dataTypes[MSPROF_GE_TENSOR_DATA_NUM]; + if (tensorShapes.size() > 0) {{ + {newline.join( + f'dataTypes[{i}] = {convert_sigtype_to_int(ty[1:])};' + for i, ty in signature.items() + if ty[0] == "*" and i < 5 + )} + }} + for (int i = 0; i < tensorShapes.size() && tensorCount < MSPROF_GE_TENSOR_DATA_NUM; i++) {{ + auto fillTensorData = [&](int index, int tensorType) {{ + profTensorData->tensorData[index].tensorType = tensorType; + profTensorData->tensorData[index].format = 2; + profTensorData->tensorData[index].dataType = dataTypes[i]; + int nDim = tensorShapes[i].size(); + nDim = nDim < MSPROF_GE_TENSOR_DATA_SHAPE_LEN ? nDim : MSPROF_GE_TENSOR_DATA_SHAPE_LEN; + for (int j = 0; j < nDim; j++) {{ + profTensorData->tensorData[index].shape[j] = tensorShapes[i][j]; + }} + for (int j = nDim; j < MSPROF_GE_TENSOR_DATA_SHAPE_LEN; j++) {{ + profTensorData->tensorData[index].shape[j] = 0; + }} + }}; + int tensorType = (i < tensorKinds.size()) ? tensorKinds[i] : 0; + if (tensorType == TENSOR_KIND_INPUT || tensorType == TENSOR_KIND_INPUT_OUTPUT) {{ + fillTensorData(tensorCount, MSPROF_GE_TENSOR_TYPE_INPUT); + tensorCount++; + }} + if ((tensorType == TENSOR_KIND_OUTPUT || tensorType == TENSOR_KIND_INPUT_OUTPUT) && tensorCount < MSPROF_GE_TENSOR_DATA_NUM){{ + fillTensorData(tensorCount, MSPROF_GE_TENSOR_TYPE_OUTPUT); + tensorCount++; + }} + }} + profTensorData->tensorNum = tensorCount; + MsprofReportAdditionalInfo(false, static_cast(&tensorInfo), sizeof(MsprofAdditionalInfo)); + }} +""" + + # Kernel launch: SIMT path for 910_95 + cpp_kernel_launch = f""" + ret = rtKernelLaunch(func, blockNum, static_cast(&args), sizeof(args), NULL, stream); +""" + if compile_on_910_95 and enable_simt: + cpp_kernel_launch = f""" + rtArgsEx_t argsInfo = {{}}; + argsInfo.args = static_cast(&args); + argsInfo.argsSize = sizeof(args); + rtTaskCfgInfo_t cfgInfo = {{}}; + cfgInfo.localMemorySize = {shared_mem_dynamic_size}; + ret = rtKernelLaunchWithFlagV2(func, blockNum, &argsInfo, NULL, stream, 0, &cfgInfo); +""" + + precompile_headers = """ +#include "precompiled.h" +""" + + return f""" +{precompile_headers} +{'#define __CCE_ENABLE_PRINT__' if enable_device_print else ''} +{extract_device_print_code_from_cann() if enable_device_print else ''} +#define PY_SSIZE_T_CLEAN +{'#define ENABLE_GRID_WARN_PRINT' if enable_grid_warn_print else ''} +#define TENSOR_KIND_INPUT 0 +#define TENSOR_KIND_OUTPUT 1 +#define TENSOR_KIND_INPUT_OUTPUT 2 + +{cpp_msprof_extern} + +{cpp_device_pointer} + +static void _launch(const char* kernelName, const void* func, rtStream_t stream, int gridX, int gridY, int gridZ, std::vector> &tensorShapes, std::vector &tensorKinds{', ' + arg_decls if len(signature) > 0 else ''}) {{ + std::string name = ""; + name.append(kernelName); + void *workspace_addr_ptr = NULL; + uint32_t blockNum4Workspace = gridX * gridY * gridZ; + {get_backend_func("pre_launch", True)} + {f''' + uint64_t totalWorkSpaceSize = {workspace_size} * blockNum4Workspace; + {get_backend_func("allocate_memory", "totalWorkSpaceSize", "stream")} + ''' if workspace_size > 0 else ''} + {'auto launch_call = [=]() -> rtError_t' if enable_taskqueue else ''} {{ + {get_backend_func("pre_launch", False)} + uint32_t blockNum = gridX * gridY * gridZ; + #ifdef ENABLE_GRID_WARN_PRINT + static bool warned = false; + if (!warned && blockNum > (uint32_t){num_physical_blocks}) {{ + printf("WARNING: Grid %u > physical limit {num_physical_blocks}, performance maybe reduced.\\n",blockNum); + warned = true; + }} + #endif + + {'blockNum = std::min(blockNum, (uint32_t)' + str(num_physical_blocks) + ');' if enable_auto_map_parallel_blocks else ''} + // set mixBlockDimRatio for nodeBasicBlockDim for msprof report + uint32_t mixBlockNumRation = {mix_block_dim_ratio}; + uint32_t nodeBasicBlockDim = (mixBlockNumRation << 16) + blockNum; + + {'cce::internal::DebugTunnelData *DTData = cce::internal::DebugTunnel::Open(blockNum);' if enable_device_print else ''} + rtError_t ret = RT_ERROR_NONE; + {'void *ffts_addr = NULL; uint32_t ffts_len; ret = rtGetC2cCtrlAddr((uint64_t*)&ffts_addr, &ffts_len);' if target_support_ffts else ''} + {'if (ret != RT_ERROR_NONE) return ret;' if (target_support_ffts and enable_taskqueue) else 'if (ret != RT_ERROR_NONE) return;' if (target_support_ffts and (not enable_taskqueue)) else ''} + // stub argument for syncBlockLock + void *syncBlockLock_ptr = NULL; + uint16_t ModuleId = 0; + {f''' + uint64_t syncBlockLockSize = {lock_num} * sizeof(int64_t); + {get_backend_func("allocate_sync_block_lock", "syncBlockLockSize", "stream")} + if (!syncBlockLock_ptr) {{ + {alloc_success_code if enable_taskqueue else sync_lock_fail_code} + }} + std::vector lockInitData({lock_num}, {lock_init_value}); + ret = rtMemcpy(syncBlockLock_ptr, syncBlockLockSize, reinterpret_cast(lockInitData.data()), + syncBlockLockSize, RT_MEMCPY_HOST_TO_DEVICE); + if (ret != RT_ERROR_NONE) {{ + return {'ret' if enable_taskqueue else ''}; + }} + ''' if lock_num > 0 else ''} + {'if (ret != RT_ERROR_NONE) return ret;' if (workspace_size > 0 and enable_taskqueue) else 'if (ret != RT_ERROR_NONE) return;' if (workspace_size > 0 and not enable_taskqueue) else ''} + struct __attribute__((packed)) {{ + {'void* ffts_addr __attribute__((aligned(8)));' if target_support_ffts else ''} + {'void* syncBlockLock __attribute__((aligned(8)));' if not force_simt_only else ''} + {'void* workspace_addr __attribute__((aligned(8)));' if not force_simt_only else ''} + {' '.join(f'{ty_to_cpp(ty)} arg{i} __attribute__((aligned({4 if ty[0] != "*" and ty[-2:] != "64" else 8})));' for i, ty in signature.items() if i not in constants and ty != "constexpr")} + {' '.join(f'{ty_to_cpp(ty)} grid{mark} __attribute__((aligned(4)));' for mark, ty in grid_info.items() if ty != "constexpr")} + {'void* DTData __attribute__((aligned(8)));' if enable_device_print else ''} + }} args = {{ + {'static_cast(ffts_addr),' if target_support_ffts else ''} + {('static_cast(syncBlockLock_ptr),' if lock_num > 0 else 'nullptr,') if not force_simt_only else ''} + {('static_cast(workspace_addr_ptr),' if workspace_size > 0 else 'nullptr,') if not force_simt_only else ''} + {(lambda _rt: (', '.join(_rt) + ',') if _rt else '')( + [f'static_cast<{ty_to_cpp(ty)}>(arg{i})' for i, ty in signature.items() if i not in constants and ty != "constexpr"] + )} + {', '.join(f'static_cast<{ty_to_cpp(ty)}>(grid{mark})' for mark, ty in grid_info.items() if ty != "constexpr")} + {', static_cast(DTData)' if enable_device_print else ''} + }}; + {cpp_msprof_call_before_launch} + {cpp_kernel_launch} + {'void *&stream_ref = const_cast(stream);' if enable_device_print else ''} + {'cce::internal::DebugTunnel::Close(DTData, stream_ref);' if enable_device_print else ''} + {cpp_msprof_call_after_launch} + {'return ret;' if enable_taskqueue else 'ret = rtStreamSynchronize(stream);'} + }}; + {f'''{get_backend_func("async_launch", "launch_call") if enable_taskqueue else ''}'''} + return; +}} + +// Extract tensor shape from PyObject +static std::vector _get_tensor_shape(PyObject *tensor) {{ + std::vector shape; + + if (!tensor || tensor == Py_None) {{ + return shape; + }} + + PyObject* size_result = PyObject_CallMethod(tensor, "size", NULL); + if (!size_result) {{ + return shape; + }} + PyObject* seq = PySequence_Fast(size_result, "Expected a sequence from tensor.size()"); + if (seq) {{ + Py_ssize_t len = PySequence_Fast_GET_SIZE(seq); + PyObject** items = PySequence_Fast_ITEMS(seq); + for (Py_ssize_t i = 0; i < len; ++i) {{ + PyObject* dim = items[i]; + if (PyLong_Check(dim)) {{ + shape.push_back(PyLong_AsLong(dim)); + }} + }} + }} + Py_DECREF(seq); + Py_DECREF(size_result); + return shape; +}} + +static PyObject* launch(PyObject* self, PyObject* args) {{ + int gridX, gridY, gridZ; + rtStream_t stream; + const void *function; + PyObject *packedMetadata = NULL; + PyObject *launch_metadata = NULL; + PyObject *launch_enter_hook = NULL; + PyObject *launch_exit_hook = NULL; + std::vector> tensorShapes; + {newline.join([f"{_extracted_type(ty)} _arg{i};" for i, ty in signature.items()])} + if(!PyArg_ParseTuple( + args, \"{fmt}\", + &gridX, &gridY, &gridZ, &stream, &function, + &packedMetadata, &launch_metadata, + &launch_enter_hook, &launch_exit_hook{args_list})) {{ + return NULL; + }} + if (__MsprofFlagL1) + {{ + { + newline.join( + f"{{ auto tmp = _get_tensor_shape(_arg{i}); if (!tmp.empty()) tensorShapes.push_back(tmp); }}" + for i, ty in signature.items() if ty[0] == "*" + ) + } + }} + + if (launch_enter_hook != Py_None){{ + PyObject* args = Py_BuildValue("(O)", launch_metadata); + PyObject* ret = PyObject_CallObject(launch_enter_hook, args); + Py_DECREF(args); + if (!ret) + return NULL; + }} + + // get kernel_name + PyObject *kernelNameObj = PyDict_GetItemString(packedMetadata, "kernel_name"); + const char *kernelName = PyUnicode_AsUTF8(kernelNameObj); + // get tensor_kinds + std::vector tensorKinds; + PyObject *tensorKindList = PyDict_GetItemString(packedMetadata, "tensor_kinds"); + if (tensorKindList) {{ + int size = PyObject_Size(tensorKindList); + for (int i = 0; i < size; i++) {{ + PyObject *kind = PySequence_GetItem(tensorKindList, i); + tensorKinds.push_back(PyLong_AsLong(kind)); + }} + }} + + // raise exception asap + {newline.join(ptr_decls)} + _launch(kernelName, function, stream, gridX, gridY, gridZ, tensorShapes, tensorKinds{', ' + ', '.join(internal_args_list) if len(internal_args_list) > 0 else ''}); + if (PyErr_Occurred()) {{ + return NULL; + }} + if(launch_exit_hook != Py_None){{ + PyObject* args = Py_BuildValue("(O)", launch_metadata); + PyObject* ret = PyObject_CallObject(launch_exit_hook, args); + Py_DECREF(args); + if (!ret) + return NULL; + }} + Py_RETURN_NONE; +}} + +static PyMethodDef ModuleMethods[] = {{ + {{"launch", launch, METH_VARARGS, "Entry point for all kernels with this signature"}}, + {{NULL, NULL, 0, NULL}} // sentinel +}}; + +static struct PyModuleDef ModuleDef = {{ + PyModuleDef_HEAD_INIT, + \"__triton_launcher\", + NULL, + -1, + ModuleMethods +}}; + +PyMODINIT_FUNC PyInit___triton_launcher(void) {{ + PyObject *m = PyModule_Create(&ModuleDef); + if(m == NULL) {{ + return NULL; + }} + PyModule_AddFunctions(m, ModuleMethods); + {cpp_msprof_callback} + return m; +}} +""" + + +# --------------------------------------------------------------------------- +# Launcher stub builder +# --------------------------------------------------------------------------- + + +def make_npu_launcher_stub(header_src, wrapper_src, debug=False): + enable_precompile = not os.getenv("TRITON_DISABLE_PRECOMPILE", "false").lower() in ( + "true", + "1", + ) + header_path = _precompile_npu_ext_with_lock(header_src, enable_precompile) + assert header_path is not None, "the precompiled.h path is empty." + + so_cache_key = hashlib.sha256(wrapper_src.encode("utf-8")).hexdigest() + so_cache_manager = get_cache_manager(so_cache_key) + use_cxx11_abi = _check_cxx11_abi() + name = f"launcher_cxx11abi{use_cxx11_abi}" + suffix = sysconfig.get_config_var("EXT_SUFFIX") + so_name = f"{name}{suffix}" + + if debug: + dump_manager = get_dump_manager(so_cache_key) + if header_path is not None: + print(f"Dumping precompiled.h to {dump_manager.cache_dir}") + dump_manager.put(header_src, "precompiled.h", binary=False) + print(f"Dumping {name}.cxx to {dump_manager.cache_dir}") + dump_manager.put(wrapper_src, f"{name}.cxx", binary=False) + + cache_path = so_cache_manager.get_file(so_name) + if cache_path is not None: + return cache_path + + kernel_launcher_type = "torch" + + with tempfile.TemporaryDirectory() as tmpdir: + src_path = os.path.join(tmpdir, f"{name}.cxx") + with open(src_path, "w") as f: + f.write(wrapper_src) + so_path = _build_npu_ext( + name, + header_path, + src_path, + kernel_launcher=kernel_launcher_type, + precompile=enable_precompile, + ) + if debug: + with open(so_path, "rb") as f: + return dump_manager.put(f.read(), so_name, binary=True) + with open(so_path, "rb") as f: + return so_cache_manager.put(f.read(), so_name, binary=True) + + +# --------------------------------------------------------------------------- +# NPULauncher +# --------------------------------------------------------------------------- + + +class NPULauncher(object): + def __init__(self, src, metadata): + self.compile_only = os.getenv("TRITON_COMPILE_ONLY", "false").lower() in ( + "true", + "1", + ) + self.enable_msprof_register_tensor = os.getenv( + "TRITON_REGISTER_TENSOR_MSPROF", "false" + ).lower() in ("true", "1") + debug_mode = metadata.debug + workspace_size = ( + int(metadata.workspace_size) if hasattr(metadata, "workspace_size") else -1 + ) + lock_init_value = ( + int(metadata.lock_init_value) if hasattr(metadata, "lock_init_value") else 0 + ) + lock_num = int(metadata.lock_num) if hasattr(metadata, "lock_num") else -1 + constants = src.constants if hasattr(src, "constants") else dict() + cst_key = lambda i: src.fn.arg_names.index(i) if isinstance(i, str) else i + constants = {cst_key(key): value for key, value in constants.items()} + signature = {cst_key(key): value for key, value in src.signature.items()} + mix_mode = metadata.mix_mode + + bs_task_type = getattr(metadata, "bs_task_type", 0) + compile_on_910_95 = getattr(metadata, "compile_on_910_95", False) + force_simt_only = getattr(metadata, "force_simt_only", False) + shared_mem_dynamic_size = getattr(metadata, "shared_mem_dynamic_size", 0) + parallel_mode = getattr(metadata, "parallel_mode", "simd") + enable_auto_blockify = getattr(metadata, "enable_auto_blockify", None) + + header_src = generate_npu_header_src() + wrapper_src = generate_npu_wrapper_src( + constants, + signature, + workspace_size, + mix_mode, + lock_num, + lock_init_value, + bs_task_type=bs_task_type, + compile_on_910_95=compile_on_910_95, + force_simt_only=force_simt_only, + shared_mem_dynamic_size=shared_mem_dynamic_size, + parallel_mode=parallel_mode, + enable_auto_blockify=enable_auto_blockify, + ) + self.so_launcher_path = make_npu_launcher_stub( + header_src, wrapper_src, debug_mode + ) + self.mix_mode = mix_mode + self.shared = metadata.shared if hasattr(metadata, "shared") else 1 + + import importlib.util + + spec = importlib.util.spec_from_file_location( + "__triton_launcher", self.so_launcher_path + ) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + self.launch = getattr(mod, "launch") + + def __call__(self, *args, **kwargs): + if self.compile_only: + cache_manager = get_cache_manager(args[5]["hash"]) + print("[INFO]: skip running kernel") + print(f"[INFO]: The compiled kernel cache is in {cache_manager.cache_dir}") + if self.enable_msprof_register_tensor: + # args[5] must be the packed metadata + args = list(args) + args[5]["tensor_params_shape"] = get_backend_func( + "get_tensor_params_shape", *args + ) + else: + if self.compile_only: + return + profiler_registered = self.launch(*args, **kwargs) + from . import utils as _backend_utils + + _backend_utils.TRITON_PROFILER_REGISTERED = ( + True if profiler_registered == 1 else False + ) diff --git a/backend/testing.py b/backend/testing.py new file mode 100644 index 00000000..6ba65d8b --- /dev/null +++ b/backend/testing.py @@ -0,0 +1,159 @@ +# 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. + +import builtins +import multiprocessing +import os +from datetime import datetime, timezone +from pathlib import Path + +import triton.runtime as runtime + + +def get_home_dir(): + return os.getenv("TRITON_HOME", Path.home()) + + +def do_bench_npu( + funcs, warmup=5, active=30, clear_l2_cache=False, prof_dir=None, keep_res=False +): + import torch + import torch_npu + + if not isinstance(funcs, list): + funcs = [funcs] + + # warmup kernel + for fn in funcs: + fn() + torch.npu.synchronize() + + experimental_config = torch_npu.profiler._ExperimentalConfig( + aic_metrics=torch_npu.profiler.AiCMetrics.PipeUtilization, + profiler_level=torch_npu.profiler.ProfilerLevel.Level1, + l2_cache=False, + data_simplification=False, + ) + + if prof_dir is not None: + torch_path = prof_dir + else: + process = multiprocessing.current_process() + pid = process.pid + process_name = process.name + timestamp = datetime.now(tz=timezone.utc).strftime("%Y%m%d_%H%M%S") + base_path = os.path.join(get_home_dir(), ".triton", "profile_results") + torch_path = os.path.join(base_path, f"prof_{timestamp}_{process_name}-{pid}") + + if clear_l2_cache: + buffer = runtime.driver.active.get_empty_cache_for_benchmark() + buffer = buffer.float() # to avoid type cast + buffer.sum() + torch.npu.synchronize() # shake out of any npu error + + total = warmup + active + with torch_npu.profiler.profile( + activities=[torch_npu.profiler.ProfilerActivity.NPU], + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(torch_path), + record_shapes=False, + profile_memory=False, + with_stack=False, + with_flops=False, + with_modules=False, + experimental_config=experimental_config, + ) as prof: + for fn in funcs: + for _ in builtins.range(total): + if clear_l2_cache: + buffer.sum() # use buffer read to clear l2 cache + torch.npu.synchronize() + fn() + torch.npu.synchronize() + if clear_l2_cache: + del buffer + + time_cost = _collect_prof_result(torch_path, funcs, warmup, active) + _rm_dic(keep_res, torch_path) + return time_cost + + +def _rm_dic(keep_res, torch_path): + if keep_res: + return + import shutil + + if os.path.exists(torch_path): + shutil.rmtree(torch_path) + + +def _collect_prof_result( + base_dir: str, funcs, num_warmup: int, num_active: int, key: str = None +): + """ + Collect kernel performance from kernel_details.csv, returned in millisecond. + The first `num_warmup` rows of each function are warmup data and will be ignored, the next `num_active` rows will be averaged. + + :param base_dir: the profiler path + :type base_dir: str + :param funcs: a list of Callable being profiled + :type funcs: List[Callable] + :param num_warmup: warmup count in kernel_details.csv of each fn + :type num_warmup: int + :param num_active: active count in kernel_details.csv of each fn + :type num_active: int + :param key: filter key for kernel name + :type key: str + """ + + import numpy as np + import pandas as pd + + kernel_details_file = None + for root, _, files in os.walk(base_dir): + for file in files: + if file == "kernel_details.csv": + kernel_details_file = os.path.join(root, file) + break + num_funcs = len(funcs) + if kernel_details_file is None: + if num_funcs == 1: + return float("inf") + else: + return [float("inf")] * num_funcs + + df = pd.read_csv(kernel_details_file) + # filter out l2 cache clearing operation + filter_cond = ~df["Type"].str.contains(r"^ReduceSum$", case=False, na=False) + filter_df = df[filter_cond] + if key is not None: + key_rows = filter_df[filter_df["Name"].str.contains(key, na=False)] + else: + key_rows = filter_df + time_cost = [0] * num_funcs + for func_idx in np.arange(0, num_funcs): + for active_index in np.arange(0, num_active): + row_index = func_idx * (num_warmup + num_active) + num_warmup + active_index + time_cost[func_idx] += key_rows.iloc[row_index]["Duration(us)"] + time_cost = [x / num_active / 1e3 for x in time_cost] + + if num_funcs == 1: + return time_cost[0] + else: + return time_cost diff --git a/backend/utils.py b/backend/utils.py index e1858193..4e999516 100644 --- a/backend/utils.py +++ b/backend/utils.py @@ -1,12 +1,122 @@ import contextlib import io import functools -from pathlib import Path +import hashlib import os +import platform +import re import shutil import subprocess import sys -from triton.runtime.cache import get_dump_manager +import sysconfig +from pathlib import Path +from triton.runtime.cache import get_cache_manager, get_dump_manager + +import pybind11 + + +# --------------------------------------------------------------------------- +# Backend function dispatch (torch_npu only, aligned with triton-ascend) +# --------------------------------------------------------------------------- + + +_BACKEND_POLICY = None + + +def get_backend_func(name, *args, **kwargs): + global _BACKEND_POLICY + if _BACKEND_POLICY is None: + import torch + import torch_npu + + _BACKEND_POLICY = "torch_npu" + return _TORCH_NPU_BACKEND_FUNCS[name](*args, **kwargs) + + +def _version_hash(): + import torch + import torch_npu + + return [torch.version.git_version, torch_npu.version.git_version] + + +def _cxx_abi(): + import torch + + return 1 if torch._C._GLIBCXX_USE_CXX11_ABI else 0 + + +def _header_file(enable_taskqueue): + taskqueue_include = ( + "#include " if enable_taskqueue else "" + ) + return f"""#include +#include +{taskqueue_include}""" + + +def _allocate_memory(size, stream): + return f"workspace_addr_ptr = const_cast(at::empty({size}, at::TensorOptions().device(at::kPrivateUse1).dtype(at::kByte)).storage().data());" + + +def _allocate_sync_block_lock(size, stream): + return f"syncBlockLock_ptr = const_cast(at_npu::native::allocate_workspace({size}, {stream}).storage().data());" + + +def _pre_launch(first_call): + return "" + + +def _async_launch(func): + return f"""at_npu::native::OpCommand cmd; + cmd.Name(name.c_str()).SetCustomHandler({func}).Run();""" + + +def _get_cc_cmd(build_pch): + import torch + import torch_npu + + torch_path = os.path.dirname(os.path.realpath(torch.__file__)) + torch_npu_path = os.path.dirname(os.path.realpath(torch_npu.__file__)) + cc_cmd = [ + f"-I{os.path.join(torch_path, 'include')}", + f"-I{os.path.join(torch_npu_path, 'include')}", + f"-D_GLIBCXX_USE_CXX11_ABI={_cxx_abi()}", + ] + if not build_pch: + cc_cmd += [ + f"-L{os.path.join(torch_npu_path, 'lib')}", + "-ltorch_npu", + ] + return cc_cmd + + +def _get_tensor_params_shape(*args): + import torch + + tensor_params = [arg for arg in args if isinstance(arg, torch.Tensor)] + tensor_params_shape = [] + for t in tensor_params: + tensor_params_shape.append([s for s in t.shape]) + return tensor_params_shape + + +_TORCH_NPU_BACKEND_FUNCS = { + "version_hash": _version_hash, + "cxx_abi": _cxx_abi, + "header_file": _header_file, + "allocate_memory": _allocate_memory, + "allocate_sync_block_lock": _allocate_sync_block_lock, + "pre_launch": _pre_launch, + "async_launch": _async_launch, + "get_cc_cmd": _get_cc_cmd, + "get_tensor_params_shape": _get_tensor_params_shape, +} + + +# --------------------------------------------------------------------------- +# Quiet context manager +# --------------------------------------------------------------------------- @contextlib.contextmanager @@ -36,6 +146,10 @@ def command_exists(cmd): return False +# --------------------------------------------------------------------------- +# Backend detection +# --------------------------------------------------------------------------- + backend = None @@ -67,6 +181,22 @@ def init_dicp_driver(): raise RuntimeError("No supported backend found.") +# --------------------------------------------------------------------------- +# Dump helpers +# --------------------------------------------------------------------------- + +TRITON_PROFILER_REGISTERED = False + +replace_dicp_ir = os.environ.get("DLC_REPLACE_DICP_IR_FILE", None) +if os.environ.get("TRITON_DEBUG", "0") == "1" or replace_dicp_ir is not None: + os.environ["TRITON_ALWAYS_COMPILE"] = "1" + dump_dir = "./tmp" + os.environ["TRITON_DUMP_DIR"] = os.environ.get("TRITON_DUMP_DIR", dump_dir) + if os.path.exists(dump_dir): + print(f"Directory **{dump_dir}** exists. Deleting the entire directory...") + shutil.rmtree(dump_dir) + + def _dump_stage_ir(ir_str, key, filename, cmd_list=None): dump_manager = get_dump_manager(key) print("Dumping intermediate results to " + dump_manager.cache_dir + "/" + filename) @@ -74,3 +204,512 @@ def _dump_stage_ir(ir_str, key, filename, cmd_list=None): if cmd_list: cmd_list[1] = dump_manager.cache_dir + "/" + filename print(f"DEBUG dump ir command: {cmd_list}") + + +# --------------------------------------------------------------------------- +# BishengIR path initialisation (local _C/bishengir override) +# --------------------------------------------------------------------------- + +local_bishengir_path = os.path.join(os.path.dirname(__file__), "../../_C/bishengir") +bisheng_install_path = os.environ.get("BISHENG_INSTALL_PATH", None) +if ( + bisheng_install_path is None + and os.path.exists(local_bishengir_path) + and os.path.isdir(local_bishengir_path) + and os.path.exists(os.path.join(local_bishengir_path, "bishengir-compile")) + and os.path.exists(os.path.join(local_bishengir_path, "bishengir-hivm-compile")) + and os.path.exists(os.path.join(local_bishengir_path, "bishengir-opt")) + and os.path.exists(os.path.join(local_bishengir_path, "hivmc")) +): + os.environ["BISHENG_INSTALL_PATH"] = local_bishengir_path + os.environ["PATH"] = local_bishengir_path + os.pathsep + os.environ["PATH"] + + +# --------------------------------------------------------------------------- +# Path resolution +# --------------------------------------------------------------------------- + + +def _get_npucompiler_path(): + ascend_dir = os.path.dirname(os.path.abspath(__file__)) + env = os.environ.copy() + npu_compiler_path = os.path.join(ascend_dir, "../../_C/bishengir/bishengir-compile") + if os.path.exists(npu_compiler_path) and os.access(npu_compiler_path, os.X_OK): + npuir_env_path = os.path.dirname(npu_compiler_path) + env["PATH"] = npuir_env_path + os.pathsep + env["PATH"] + else: + npu_compiler_path = shutil.which("bishengir-compile") + if npu_compiler_path is None: + npu_compiler_root = os.getenv("TRITON_NPU_COMPILER_PATH", None) + if npu_compiler_root is None: + raise EnvironmentError( + "Couldn't find executable bishengir-compile or TRITON_NPU_COMPILER_PATH." + ) + npu_compiler_path = os.path.join(npu_compiler_root, "npuc") + return os.path.abspath(npu_compiler_path), env + + +def _get_bishengir_opt_path(): + ascend_dir = os.path.dirname(os.path.abspath(__file__)) + env = os.environ.copy() + bishengir_opt_path = os.path.join(ascend_dir, "../../_C/bishengir/bishengir-opt") + if os.path.exists(bishengir_opt_path) and os.access(bishengir_opt_path, os.X_OK): + npuir_env_path = os.path.dirname(bishengir_opt_path) + env["PATH"] = npuir_env_path + os.pathsep + env["PATH"] + else: + bishengir_opt_path = shutil.which("bishengir-opt") + if bishengir_opt_path is None: + bishengir_opt_root = os.getenv("TRITON_NPU_COMPILER_PATH", None) + if bishengir_opt_root is None: + raise EnvironmentError( + "Couldn't find executable bishengir-opt or TRITON_NPU_COMPILER_PATH" + ) + bishengir_opt_path = os.path.join(bishengir_opt_root, "bishengir-opt") + return os.path.abspath(bishengir_opt_path), env + + +def _get_bisheng_path() -> str: + bisheng_path = shutil.which("bisheng") + if bisheng_path is None: + npu_compiler_root = os.getenv("TRITON_NPU_COMPILER_PATH", None) + if npu_compiler_root is None: + raise EnvironmentError( + "Couldn't find executable bisheng or TRITON_NPU_COMPILER_PATH" + ) + bisheng_path = os.path.join(npu_compiler_root, "ccec") + return bisheng_path + + +@functools.lru_cache(None) +def _get_ascend_path() -> Path: + path = os.getenv("ASCEND_HOME_PATH", "") + if path == "": + raise EnvironmentError( + "ASCEND_HOME_PATH is not set, source /set_env.sh first" + ) + return Path(path) + + +# --------------------------------------------------------------------------- +# Feature flags +# --------------------------------------------------------------------------- + + +def _is_ascend_sanitizer_enabled() -> bool: + return os.getenv("TRITON_ENABLE_SANITIZER", "false").lower() in ("true", "1") + + +def _is_auto_map_parallel_blocks_enabled() -> bool: + return os.getenv("TRITON_ALL_BLOCKS_PARALLEL", "false").lower() in ("true", "1") + + +def get_ascend_arch_from_env(): + return os.getenv("TRITON_ASCEND_ARCH", "") + + +def is_ffts_supported(arch: str) -> bool: + if is_compile_on_910_95: + return False + if arch in ("Ascend910A", "Ascend310B4"): + return False + return True + + +def force_disable_ffts() -> bool: + if is_compile_on_910_95: + return True + return os.getenv("TRITON_DISABLE_FFTS", "false").lower() in ("true", "1") + + +# --------------------------------------------------------------------------- +# C++ compilation helpers +# --------------------------------------------------------------------------- + + +def _check_cxx11_abi(): + return get_backend_func("cxx_abi") + + +def _get_cxx(): + cxx = os.environ.get("CC") + if cxx is not None: + return cxx + clangxx = shutil.which("clang++") + gxx = shutil.which("g++") + cxx = clangxx if clangxx is not None else gxx + if cxx is None: + raise RuntimeError("Failed to find C++ compiler") + return cxx + + +def _get_cxx_precompiled(header_path): + cxx = os.environ.get("CC") + if cxx is None: + clangxx = shutil.which("clang++") + gxx = shutil.which("g++") + if clangxx is not None: + return [clangxx, "-include", header_path] + elif gxx is not None: + return [gxx] + else: + raise RuntimeError("Failed to find C++ compiler") + return [cxx] + + +def _precompile_npu_hash(header_src): + cxx = _get_cxx() + py_version = sys.version + asc_path = str(_get_ascend_path()) + version_txt = [header_src, cxx, py_version, asc_path] + version_txt += get_backend_func("version_hash") + return hashlib.sha256("_".join(version_txt).encode("utf-8")).hexdigest() + + +def _precompile_npu_ext(header_path, gch_path): + cc_cmd = [_get_cxx(), "-x", "c++-header", header_path] + cc_cmd += ["-w"] + + if hasattr(sysconfig, "get_default_scheme"): + scheme = sysconfig.get_default_scheme() + else: + scheme = sysconfig._get_default_scheme() + if scheme == "posix_local": + scheme = "posix_prefix" + py_include_dir = sysconfig.get_paths(scheme=scheme)["include"] + cc_cmd += [f"-I{py_include_dir}"] + cc_cmd += [f"-I{os.path.dirname(os.path.realpath(__file__))}"] + + asc_path = _get_ascend_path() + rt_path = os.path.join(asc_path, "include/experiment/runtime/runtime/rt.h") + if not os.path.exists(rt_path): + cc_cmd += [ + f"-I{os.path.join(asc_path, 'pkg_inc')}", + f"-I{os.path.join(asc_path, 'pkg_inc/profiling')}", + ] + + cc_cmd += [ + f"-I{os.path.join(asc_path, 'include')}", + f"-I{os.path.join(asc_path, 'include/experiment')}", + f"-I{os.path.join(asc_path, 'include/experiment/msprof')}", + f"-I{pybind11.get_include()}", + ] + + cc_cmd += get_backend_func("get_cc_cmd", build_pch=True) + + cc_cmd += ["-std=c++17", "-shared", "-fPIC", "-o", gch_path] + + result = subprocess.run(cc_cmd, capture_output=True, text=True) + if result.returncode != 0: + raise RuntimeError( + f"Failed to precompile {gch_path}, error: {result.stderr}, cmd={cc_cmd}" + ) + return header_path + + +def _precompile_npu_ext_with_lock(header_src, enable_precompile): + import fcntl + + precompile_hash = _precompile_npu_hash(header_src) + cache = get_cache_manager(precompile_hash) + gch_path = cache.get_file("precompiled.h.gch") + header_path = cache.get_file("precompiled.h") + + if enable_precompile: + if header_path is not None and gch_path is not None: + return header_path + else: + if header_path is not None: + return header_path + + cache_dir = os.getenv("TRITON_CACHE_DIR", "").strip() + lock_path = os.path.join(cache_dir, f"{precompile_hash}.lock") + with open(lock_path, "a+") as f: + try: + fcntl.flock(f, fcntl.LOCK_EX) + header_path = cache.get_file("precompiled.h") + if enable_precompile: + gch_path = cache.get_file("precompiled.h.gch") + if header_path is not None and gch_path is not None: + return header_path + else: + if header_path is not None: + return header_path + header_path = cache.put(header_src, "precompiled.h", binary=False) + if not enable_precompile: + return header_path + src_dir = os.path.dirname(header_path) + gch_path = os.path.join(src_dir, "precompiled.h.gch") + _precompile_npu_ext(header_path, gch_path) + return header_path + finally: + fcntl.flock(f, fcntl.LOCK_UN) + + +def _build_npu_ext( + obj_name: str, header_path, src_path, *, kernel_launcher="torch", precompile=False +) -> str: + suffix = sysconfig.get_config_var("EXT_SUFFIX") + src_dir = os.path.dirname(src_path) + so_path = os.path.join(src_dir, f"{obj_name}{suffix}") + + if precompile: + cc_cmd = _get_cxx_precompiled(header_path) + cc_cmd += [src_path] + else: + cc_cmd = [_get_cxx(), src_path] + + cc_cmd += ["-w"] + + if hasattr(sysconfig, "get_default_scheme"): + scheme = sysconfig.get_default_scheme() + else: + scheme = sysconfig._get_default_scheme() + if scheme == "posix_local": + scheme = "posix_prefix" + py_include_dir = sysconfig.get_paths(scheme=scheme)["include"] + cc_cmd += [f"-I{py_include_dir}"] + cc_cmd += [f"-I{os.path.dirname(os.path.realpath(__file__))}"] + + asc_path = _get_ascend_path() + if header_path is not None: + cc_cmd += [f"-I{os.path.dirname(header_path)}"] + + rt_path = os.path.join(asc_path, "include/experiment/runtime/runtime/rt.h") + if not os.path.exists(rt_path): + cc_cmd += [ + f"-I{os.path.join(asc_path, 'pkg_inc')}", + f"-I{os.path.join(asc_path, 'pkg_inc/profiling')}", + ] + + cc_cmd += [ + f"-I{os.path.join(asc_path, 'include')}", + f"-I{os.path.join(asc_path, 'include/experiment')}", + f"-I{os.path.join(asc_path, 'include/experiment/msprof')}", + f"-I{pybind11.get_include()}", + f"-L{os.path.join(asc_path, 'lib64')}", + "-lruntime", + "-lascendcl", + ] + + if kernel_launcher: + cc_cmd += get_backend_func("get_cc_cmd", build_pch=False) + + cc_cmd += ["-std=c++17", "-shared", "-fPIC", "-Winvalid-pch", "-o", so_path] + + result = subprocess.run(cc_cmd, capture_output=True, text=True) + if result.returncode == 0: + return so_path + if "precompiled.h.gch" in result.stderr: + return _build_npu_ext( + obj_name, + header_path, + src_path, + kernel_launcher=kernel_launcher, + precompile=False, + ) + raise RuntimeError( + f"Failed to compile {src_path}, error: {result.stderr}, cmd={cc_cmd}" + ) + + +# --------------------------------------------------------------------------- +# Type mapping for msprof tensor reporting +# --------------------------------------------------------------------------- + + +def convert_sigtype_to_int(sigty: str): + MAP_SIGTYPE_TO_INT = { + "i1": 12, + "i4": 29, + "i8": 2, + "i16": 6, + "i32": 3, + "i64": 9, + "u1": 30, + "u8": 4, + "u16": 7, + "u32": 8, + "u64": 10, + "fp16": 1, + "bf16": 27, + "fp32": 0, + "fp64": 11, + "fp8e5": 35, + "fp8e4nv": 36, + } + if sigty not in MAP_SIGTYPE_TO_INT: + raise ValueError(f"Unsupported data type: {sigty}") + return MAP_SIGTYPE_TO_INT[sigty] + + +# --------------------------------------------------------------------------- +# BishengIR API detection +# --------------------------------------------------------------------------- + + +def _check_bishengir_api_change() -> bool: + bishengir_path, _ = _get_npucompiler_path() + try: + result = subprocess.run( + [bishengir_path, "--help"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + if result.returncode == 0 and "limit-auto-multi-buffer-buffer" in result.stdout: + return True + return False + except Exception as e: + print(f"ERROR: {e}") + return False + + +# --------------------------------------------------------------------------- +# Additional feature flags +# --------------------------------------------------------------------------- + + +def _is_debug_line_info_disabled() -> bool: + return os.getenv("TRITON_DISABLE_LINE_INFO", "true").lower() in ("true", "1") + + +def _enable_print_ub_bits() -> bool: + return os.getenv("ENABLE_PRINT_UB_BITS", "false").lower() in ("true", "1") + + +def _enable_dump_memory_info() -> bool: + return os.getenv("TRITON_MEMORY_DISPLAY", "false").lower() in ("true", "1") + + +def _enable_msdebug() -> bool: + return os.getenv("LLVM_EXTRACT_DI_LOCAL_VARIABLES", "false").lower() in ( + "true", + "1", + ) + + +def _enable_unpublished_feature() -> bool: + return os.getenv("ENABLE_UNPUBLISHED_FEATURE", "false").lower() in ("true", "1") + + +# --------------------------------------------------------------------------- +# 910_95 device detection +# --------------------------------------------------------------------------- + + +def _get_ascend_devices(): + import glob as _glob + + devices = [] + pci_path = "/sys/bus/pci/devices/*" + for dev in _glob.glob(pci_path): + try: + vendor_path = os.path.join(dev, "vendor") + device_path = os.path.join(dev, "device") + if os.path.exists(vendor_path): + with open(vendor_path, "r") as f: + vendor = f.read().strip() + if vendor == "0x19e5" and os.path.exists(device_path): + with open(device_path, "r") as f: + device = f.read().strip() + devices.append(device) + except (IOError, OSError): + continue + return devices + + +def _check_npu_smi_device(): + try: + result = subprocess.run( + ["npu-smi", "info"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + shell=False, + timeout=100, + ) + if result.returncode == 0: + output = result.stdout.lower() + return ( + "ascend910_95" in output + or "ascend950" in output + or "910_958b" in output + ) + return False + except Exception: + return False + + +_ascend_devices = _get_ascend_devices() +_pci_condition = any("0xd806" in dev for dev in _ascend_devices) +_npu_smi_condition = _check_npu_smi_device() +is_compile_on_910_95 = _pci_condition or _npu_smi_condition + + +# --------------------------------------------------------------------------- +# CANN version & libdevice helpers +# --------------------------------------------------------------------------- + + +def get_machine_arch(): + ARCHITECTURE_ALIASES = { + "x86_64": "x86_64", + "amd64": "x86_64", + "i386": "x86_64", + "i686": "x86_64", + "arm64": "aarch64", + "aarch64": "aarch64", + "armv7l": "aarch64", + "armv8l": "aarch64", + "arm": "aarch64", + } + system_arch = platform.machine() + return ARCHITECTURE_ALIASES.get(system_arch, system_arch) + + +def get_cann_version(): + ascend_path = _get_ascend_path() + arch = get_machine_arch() + cann_version_file_path = os.path.join( + ascend_path, arch + "-linux", "ascend_toolkit_install.info" + ) + if not os.path.exists(cann_version_file_path): + cann_version_file_path = os.path.join( + ascend_path, arch + "-linux", "ascend_all_cann_install.info" + ) + version = "" + innerversion = "" + with open(cann_version_file_path) as f: + for line in f: + line = line.strip() + if line.startswith("version="): + version = line.split("=")[1] + elif line.startswith("innerversion="): + innerversion = line.split("=")[1] + if version and innerversion: + return "CANN-" + version + "-" + innerversion + if version: + return "CANN-" + version + raise ValueError("get_cann_version is empty!") + + +def triton_enable_libdevice_simt(): + enable_libdevice_simt = os.getenv("TRITON_ENABLE_LIBDEVICE_SIMT", False) + return enable_libdevice_simt + + +def _check_bishengir_able_save_ir() -> bool: + bishengir_path, _ = _get_npucompiler_path() + try: + result = subprocess.run( + [bishengir_path, "--help"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + if result.returncode == 0 and "save-linked-ir" in result.stdout: + return True + return False + except Exception as e: + print(f"ERROR: {e}") + return False diff --git a/compile_shared.sh b/compile_shared.sh index 945b9d9e..6c811e3b 100644 --- a/compile_shared.sh +++ b/compile_shared.sh @@ -3,7 +3,7 @@ export LANG="zh_CN.UTF-8" export LC_ALL="zh_CN.UTF-8" home_path=$(pwd) -# compile triton shared library with patch +# compile DLCompiler apply_patch=false build_package=false @@ -23,12 +23,6 @@ echo "start compile ========================================" echo apply_patch: $apply_patch echo "======================================================" -echo "start apply ascendnpu-ir patch" -cd $home_path/third_party/ascendnpu-ir -git checkout . -git apply ../../patch/ascendnpu-ir.patch -echo "apply patch/ascendnpu-ir.patch success!" - # SET ENV # export JSON_PATH=/path/to/your/json/file # export GOOGLETEST_DIR=/path/to/your/googletest/directory @@ -47,19 +41,10 @@ check_npu() { check_npu + if [[ $apply_patch == true ]]; then # do dangerous stuff - echo "Apply triton and triton_shared patch" - echo "当前环境检测为：$([[ $is_npu == true ]] && echo 'ascend加速卡，使用适配patch' || echo '非ascend加速卡，不使用适配patch')" - if [[ $is_npu == true ]]; then - cd $TRITON_PLUGIN_DIRS/third_party/triton_shared/ - git checkout . - ls $TRITON_PLUGIN_DIRS/patch/ttshared/*.patch | xargs -n1 git apply - if [ $? -ne 0 ]; then - echo "Error: triton_shared git apply failed." >&2 - exit 1 - fi - fi + echo "Apply triton patch" cd $TRITON_PLUGIN_DIRS/third_party/triton/ git checkout . ls $TRITON_PLUGIN_DIRS/patch/triton/*.patch | xargs -n1 git apply diff --git a/compiler/CMakeLists.txt b/compiler/CMakeLists.txt index 5cd7900f..94560da4 100644 --- a/compiler/CMakeLists.txt +++ b/compiler/CMakeLists.txt @@ -1 +1,6 @@ -dicp_add_all_subdirs() \ No newline at end of file +# DICP NPU compiler passes +include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include) +include_directories(${CMAKE_CURRENT_BINARY_DIR}/include) + +add_subdirectory(include) +add_subdirectory(lib) diff --git a/compiler/include/CMakeLists.txt b/compiler/include/CMakeLists.txt index 5cd7900f..8ce02c99 100644 --- a/compiler/include/CMakeLists.txt +++ b/compiler/include/CMakeLists.txt @@ -1 +1 @@ -dicp_add_all_subdirs() \ No newline at end of file +add_subdirectory(dicp) diff --git a/compiler/include/dicp/AscendLegalize/AscendLegalizePass.h b/compiler/include/dicp/AscendLegalize/AscendLegalizePass.h new file mode 100644 index 00000000..f5721440 --- /dev/null +++ b/compiler/include/dicp/AscendLegalize/AscendLegalizePass.h @@ -0,0 +1,18 @@ +#ifndef TRITON_ADAPTER_ASCENDLEGALIZE_H +#define TRITON_ADAPTER_ASCENDLEGALIZE_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +#define GEN_PASS_DECL_ASCENDLEGALIZE +#include "dicp/AscendLegalize/Passes.h.inc" + +namespace mlir { +namespace triton { + +std::unique_ptr> createAscendLegalizePass(); + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_ASCENDLEGALIZE_H diff --git a/compiler/include/dicp/AscendLegalize/CMakeLists.txt b/compiler/include/dicp/AscendLegalize/CMakeLists.txt new file mode 100644 index 00000000..cf810ec3 --- /dev/null +++ b/compiler/include/dicp/AscendLegalize/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name AscendLegalize) +add_public_tablegen_target(AscendLegalizePassIncGen) diff --git a/compiler/include/dicp/AscendLegalize/Passes.h b/compiler/include/dicp/AscendLegalize/Passes.h new file mode 100644 index 00000000..ae8cf48d --- /dev/null +++ b/compiler/include/dicp/AscendLegalize/Passes.h @@ -0,0 +1,15 @@ +#ifndef TRITON_ADAPTER_ASCENDLEGALIZE_PASSES_H +#define TRITON_ADAPTER_ASCENDLEGALIZE_PASSES_H + +#include "dicp/AscendLegalize/AscendLegalizePass.h" + +namespace mlir { +namespace triton { + +#define GEN_PASS_REGISTRATION +#include "dicp/AscendLegalize/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_ASCENDLEGALIZE_PASSES_H diff --git a/compiler/include/dicp/AscendLegalize/Passes.td b/compiler/include/dicp/AscendLegalize/Passes.td new file mode 100644 index 00000000..e9a837ce --- /dev/null +++ b/compiler/include/dicp/AscendLegalize/Passes.td @@ -0,0 +1,15 @@ +#ifndef ASCENDLEGALIZE_PASSES +#define ASCENDLEGALIZE_PASSES + +include "mlir/Pass/PassBase.td" + +def AscendLegalize : Pass<"ascend-legalize", "mlir::ModuleOp"> { + let summary = "Legalize IR for Ascend backend (normalize cmpi predicates etc.)"; + let constructor = "triton::createAscendLegalizePass()"; + let dependentDialects = [ + "mlir::arith::ArithDialect", + "mlir::triton::TritonDialect" + ]; +} + +#endif // ASCENDLEGALIZE_PASSES diff --git a/compiler/include/dicp/AutoBlockify/AutoBlockify.h b/compiler/include/dicp/AutoBlockify/AutoBlockify.h new file mode 100644 index 00000000..3487c10d --- /dev/null +++ b/compiler/include/dicp/AutoBlockify/AutoBlockify.h @@ -0,0 +1,99 @@ + + +#pragma once + +#include "mlir/Pass/Pass.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/IR/PatternMatch.h" + +#define GEN_PASS_DECL_AUTOBLOCKIFY +#include "dicp/AutoBlockify/Passes.h.inc" + +#define GEN_PASS_DEF_AUTOBLOCKIFY +#include "dicp/AutoBlockify/Passes.h.inc" + +namespace mlir { +namespace triton { + +std::unique_ptr> +createAutoBlockifyPass(const AutoBlockifyOptions &options = {}); + +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace triton; + +class PropagateUnrealizedCastDown + : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + explicit PropagateUnrealizedCastDown(MLIRContext *context, + Value logicalBlockId, + Value logicalBlockNum, + int autoBlockifySize); + + LogicalResult matchAndRewrite(UnrealizedConversionCastOp op, + PatternRewriter &rewriter) const override; + +private: + void handleBlockifyLoop(scf::ForOp blockifyLoop, Operation *op, + PatternRewriter &rewriter) const; + void rewriteSplat(UnrealizedConversionCastOp op, triton::SplatOp splatOp, + PatternRewriter &rewriter) const; + void rewriteExpandDims(UnrealizedConversionCastOp op, + triton::ExpandDimsOp expandDimsOp, + PatternRewriter &rewriter) const; + void rewriteReduce(UnrealizedConversionCastOp op, triton::ReduceOp reduceOp, + PatternRewriter &rewriter) const; + void rewriteScan(UnrealizedConversionCastOp op, triton::ScanOp scanOp, + PatternRewriter &rewriter) const; + void rewriteLoad(UnrealizedConversionCastOp op, triton::LoadOp loadOp, + PatternRewriter &rewriter) const; + void rewriteStore(UnrealizedConversionCastOp op, triton::StoreOp storeOp, + PatternRewriter &rewriter) const; + void rewriteAtomicRMW(UnrealizedConversionCastOp op, + triton::AtomicRMWOp atomicRMWOp, + PatternRewriter &rewriter) const; + void rewriteAssert(UnrealizedConversionCastOp op, triton::AssertOp assertOp, + PatternRewriter &rewriter) const; + void rewriteExtractSlice(UnrealizedConversionCastOp op, + tensor::ExtractSliceOp extractSliceOp, + PatternRewriter &rewriter) const; + void rewriteInsertSlice(UnrealizedConversionCastOp op, + tensor::InsertSliceOp insertSliceOp, + PatternRewriter &rewriter) const; + void rewriteWhile(UnrealizedConversionCastOp op, scf::WhileOp whileOp, + PatternRewriter &rewriter) const; + void rewriteLoop(UnrealizedConversionCastOp op, LoopLikeOpInterface loopOp, + PatternRewriter &rewriter) const; + void rewriteIf(UnrealizedConversionCastOp &op, scf::IfOp ifOp, + ArrayRef indices, PatternRewriter &rewriter) const; + void rewriteYield(UnrealizedConversionCastOp &op, scf::YieldOp yieldOp, + PatternRewriter &rewriter) const; + void rewriteCondition(UnrealizedConversionCastOp op, + scf::ConditionOp conditionOp, + PatternRewriter &rewriter) const; + void rewriteGeneraleOp(UnrealizedConversionCastOp op, Operation *generalOp, + PatternRewriter &rewriter) const; + + Value logicalBlockId; + Value logicalBlockNum; + int autoBlockifySize; +}; + +class AutoBlockifyPass : public ::impl::AutoBlockifyBase { +public: + explicit AutoBlockifyPass(const AutoBlockifyOptions &options); + void runOnOperation() override; + +private: + bool checkBlockifiable(Value v); + void preProcess(triton::FuncOp func); + + DenseSet checkedValues; + Value logicalBlockId; + Value logicalBlockNum; +}; diff --git a/compiler/include/dicp/AutoBlockify/CMakeLists.txt b/compiler/include/dicp/AutoBlockify/CMakeLists.txt new file mode 100644 index 00000000..abae4f6e --- /dev/null +++ b/compiler/include/dicp/AutoBlockify/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name AutoBlockify) +add_public_tablegen_target(AutoBlockifyPassIncGen) \ No newline at end of file diff --git a/compiler/include/dicp/AutoBlockify/Passes.h b/compiler/include/dicp/AutoBlockify/Passes.h new file mode 100644 index 00000000..3e6cf3c7 --- /dev/null +++ b/compiler/include/dicp/AutoBlockify/Passes.h @@ -0,0 +1,17 @@ + + +#ifndef TRITON_ADAPTER_AUTO_BLOCKIFY_PASSES_H +#define TRITON_ADAPTER_AUTO_BLOCKIFY_PASSES_H + +#include "dicp/AutoBlockify/AutoBlockify.h" + +namespace mlir { +namespace triton { + +#define GEN_PASS_REGISTRATION +#include "dicp/AutoBlockify/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_AUTO_BLOCKIFY_PASSES_H diff --git a/compiler/include/dicp/AutoBlockify/Passes.td b/compiler/include/dicp/AutoBlockify/Passes.td new file mode 100644 index 00000000..56ab1587 --- /dev/null +++ b/compiler/include/dicp/AutoBlockify/Passes.td @@ -0,0 +1,21 @@ +#ifndef AUTO_BLOCKIFY_PASSES +#define AUTO_BLOCKIFY_PASSES + +include "mlir/Pass/PassBase.td" + +def AutoBlockify : Pass<"auto-blockify", "mlir::ModuleOp"> { + let summary = "Apply auto blockify v2"; + let constructor = "triton::createAutoBlockifyPass()"; + let dependentDialects = [ + "mlir::arith::ArithDialect", + "mlir::tensor::TensorDialect", + "mlir::triton::TritonDialect" + ]; + let options = [ + Option<"autoBlockifySize", "auto-blockify-size", "int", "1", + "Apply auto blockify v2 when TRITON_ALL_BLOCKS_PARALLEL is 1." + "Expand highest dimension with blockify size"> + ]; +} + +#endif // AUTO_BLOCKIFY_PASSES diff --git a/compiler/include/dicp/AutoBlockify/Utils.h b/compiler/include/dicp/AutoBlockify/Utils.h new file mode 100644 index 00000000..e9c7a365 --- /dev/null +++ b/compiler/include/dicp/AutoBlockify/Utils.h @@ -0,0 +1,45 @@ + + +#pragma once + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +using namespace mlir; +using namespace triton; + +constexpr llvm::StringLiteral autoBlockifySizeAttr = "auto_blockify_size"; +constexpr llvm::StringLiteral logicalBlockIdAttr = "logical_block_id"; +constexpr llvm::StringLiteral autoBlockifyLoopAttr = "auto_blockify_loop"; +constexpr llvm::StringLiteral autoBlockifyRegionOpAttr = + "auto_blockify_region_op"; + +RankedTensorType getExpandedType(Type type, UnrealizedConversionCastOp op); + +Value rewriteValue(Value value, UnrealizedConversionCastOp op, + OpBuilder &builder); + +void replaceValue(Operation *newOp, Operation *oldOp, Value newMask, + RewriterBase &rewriter, + ArrayRef replaceIndices = {}); + +Value createMask(Value mask, Value uccMask, ArrayRef targetShape, + RewriterBase &rewriter); + +void mapRegionIterArg(IRMapping &mapping, ValueRange oldArgs, + ValueRange newArgs, ArrayRef indices, Value mask, + OpBuilder &builder); + +void mapYieldedValue(IRMapping &mapping, scf::YieldOp yieldOp, + ArrayRef indices, UnrealizedConversionCastOp op, + OpBuilder &builder); + +Operation *createBlockifyLoop(Operation *targetOp, + UnrealizedConversionCastOp op, + Value logicalBlockId, Value logicalBlockNum, + int autoBlockifySize, RewriterBase &rewriter); + +std::optional getBlockifyLoop(Operation *op); \ No newline at end of file diff --git a/compiler/include/dicp/CMakeLists.txt b/compiler/include/dicp/CMakeLists.txt index 5cd7900f..d7619581 100644 --- a/compiler/include/dicp/CMakeLists.txt +++ b/compiler/include/dicp/CMakeLists.txt @@ -1 +1,14 @@ -dicp_add_all_subdirs() \ No newline at end of file +add_subdirectory(AutoBlockify) +add_subdirectory(AscendLegalize) +add_subdirectory(Dialect) +add_subdirectory(DiscreteMaskAccessConversion) +add_subdirectory(TritonAffinityOpt) +add_subdirectory(TritonToAnnotation) +add_subdirectory(TritonToGraph) +add_subdirectory(TritonToHFusion) +add_subdirectory(TritonToHIVM) +add_subdirectory(TritonToLLVM) +add_subdirectory(TritonToLinalg) +add_subdirectory(TritonToStructured) +add_subdirectory(TritonToUnstructure) +add_subdirectory(Utils) diff --git a/compiler/include/dicp/Conversion/CMakeLists.txt b/compiler/include/dicp/Conversion/CMakeLists.txt deleted file mode 100644 index 5cd7900f..00000000 --- a/compiler/include/dicp/Conversion/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -dicp_add_all_subdirs() \ No newline at end of file diff --git a/compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/Passes.h b/compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/Passes.h deleted file mode 100644 index d7c1cf1f..00000000 --- a/compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/Passes.h +++ /dev/null @@ -1,33 +0,0 @@ -#ifndef TRITON_DLC_DISCRETE_MASK_ACCESS_CONVERSION_PASSES_H -#define TRITON_DLC_DISCRETE_MASK_ACCESS_CONVERSION_PASSES_H - -#include "mlir/Pass/Pass.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include "mlir/IR/PatternMatch.h" - -#define GEN_PASS_DECL_DISCRETEMASKACCESSCONVERSION -#include "dicp/Conversion//DiscreteMaskAccessConversion/Passes.h.inc" - -#define GEN_PASS_DEF_DISCRETEMASKACCESSCONVERSION -#include "dicp/Conversion//DiscreteMaskAccessConversion/Passes.h.inc" - -namespace mlir { -namespace triton { - -std::unique_ptr> createDiscreteMaskAccessConversionPass( - const DiscreteMaskAccessConversionOptions &options = {}); - -} // namespace triton -} // namespace mlir - -namespace mlir { -namespace triton { - -#define GEN_PASS_REGISTRATION -#include "dicp/Conversion//DiscreteMaskAccessConversion/Passes.h.inc" - -} // namespace triton -} // namespace mlir - -#endif // TRITON_DLC_DISCRETE_MASK_ACCESS_CONVERSION_PASSES_H diff --git a/compiler/include/dicp/Conversion/LinalgToLinked/CMakeLists.txt b/compiler/include/dicp/Conversion/LinalgToLinked/CMakeLists.txt deleted file mode 100644 index d002213f..00000000 --- a/compiler/include/dicp/Conversion/LinalgToLinked/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls --name LinalgToLinked) -add_public_tablegen_target(LinalgToLinkedConversionPassIncGen) diff --git a/compiler/include/dicp/Conversion/LinalgToLinked/LinalgToLinked.h b/compiler/include/dicp/Conversion/LinalgToLinked/LinalgToLinked.h deleted file mode 100644 index 743cb567..00000000 --- a/compiler/include/dicp/Conversion/LinalgToLinked/LinalgToLinked.h +++ /dev/null @@ -1,12 +0,0 @@ -#pragma once - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Pass/Pass.h" - -namespace mlir::dicp::linked { - -std::unique_ptr> -createLinalgToLinkedPass(bool globalKernel = true, bool namedOps = true, - bool cpuVerify = false); - -} // namespace mlir::dicp::linked diff --git a/compiler/include/dicp/Conversion/LinalgToLinked/Passes.h b/compiler/include/dicp/Conversion/LinalgToLinked/Passes.h deleted file mode 100644 index 3f351e27..00000000 --- a/compiler/include/dicp/Conversion/LinalgToLinked/Passes.h +++ /dev/null @@ -1,12 +0,0 @@ - -#pragma once -#include "dicp/Conversion/LinalgToLinked/LinalgToLinked.h" - -namespace mlir::dicp::linked { - -std::unique_ptr> createDebugCPUVerifyPass(); - -#define GEN_PASS_REGISTRATION -#include "dicp/Conversion/LinalgToLinked/Passes.h.inc" - -} // namespace mlir::dicp::linked diff --git a/compiler/include/dicp/Conversion/LinalgToLinked/Passes.td b/compiler/include/dicp/Conversion/LinalgToLinked/Passes.td deleted file mode 100644 index a3d59ca7..00000000 --- a/compiler/include/dicp/Conversion/LinalgToLinked/Passes.td +++ /dev/null @@ -1,34 +0,0 @@ -#ifndef LINALG_TO_LINKED_CONVERSION_PASSES -#define LINALG_TO_LINKED_CONVERSION_PASSES - -include "mlir/Pass/PassBase.td" - -def LinalgToLinked : Pass<"linalg-to-linked", "mlir::ModuleOp"> { - let summary = "Convert Linalg to Linked dialect"; - let constructor = "linked::createLinalgToLinkedPass()"; - let options = [ - Option<"globalKernel", "global-kernel", - "bool", /*default*/"true", - "Generate a global kernel">, - Option<"namedOps", "named-ops", - "bool", /*default*/"true", - "Use linalg named ops instead of linalg.generic">, - Option<"cpuVerify", "cpu-verify", - "bool", /*default*/"false", - "Skip NPU workspace args for CPU verification"> - ]; -} - -def DebugCPUVerify : Pass<"debug-cpu-verify", "mlir::ModuleOp"> { - let summary = "Verify that only MLIR built-in dialects remain for CPU runner"; - let description = [{ - Verification pass that scans the module for operations belonging to - non-MLIR-upstream (external) dialects. Any such operation is reported - as an error, ensuring the IR has been fully lowered to standard MLIR - constructs before being handed off to a CPU runner for correctness - validation. - }]; - let constructor = "linked::createDebugCPUVerifyPass()"; -} - -#endif diff --git a/compiler/include/dicp/Conversion/LinalgToLinked/TritonOpConverter.h b/compiler/include/dicp/Conversion/LinalgToLinked/TritonOpConverter.h deleted file mode 100644 index ceb0e458..00000000 --- a/compiler/include/dicp/Conversion/LinalgToLinked/TritonOpConverter.h +++ /dev/null @@ -1,223 +0,0 @@ -#ifndef TRITON_DLC_TRITONOPCONVERTER_H -#define TRITON_DLC_TRITONOPCONVERTER_H - -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/Interfaces/FunctionInterfaces.h" -#include "mlir/Transforms/DialectConversion.h" - -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/TypeSwitch.h" -#include "llvm/Support/Debug.h" - -using namespace mlir; - -namespace mlir::dicp::linked { - -template -class ReductionOpBaseConverter : public OpConversionPattern { -public: - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(OpTy op, typename OpTy::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const final { - auto sourceType = - cast(adaptor.getOperands().front().getType()); - assert(sourceType.hasRank() && "Expected input is ranked"); - - int64_t axis = op.getAxis(); - assert(axis >= 0 && axis < sourceType.getRank() && - "Expected reduction axis is within operand's rank"); - - auto reductionOps = this->getRedOps(op); - if (reductionOps.size() == 1) { - return this->convertToTargetOp(op, adaptor, rewriter); - } - return this->convertToTargetOpExtended(op, adaptor, rewriter); - } - -protected: - llvm::SmallVector getRedOps(OpTy redOp) const { - auto redBody = redOp.getBody(); - return llvm::map_to_vector(redBody->without_terminator(), - [](Operation &op) { return &op; }); - } - - arith::ConstantOp getRedBaseConstOp(ConversionPatternRewriter &rewriter, - Operation *redOp, - Type constantType) const { - const int64_t bitWidth = constantType.getIntOrFloatBitWidth(); - - auto attr = - llvm::TypeSwitch(redOp) - .Case([&](arith::AddFOp) { - return rewriter.getFloatAttr(constantType, 0.f); - }) - .Case([&](arith::AddIOp) { - return rewriter.getIntegerAttr(constantType, 0); - }) - .Case([&](arith::MulFOp) { - return rewriter.getFloatAttr(constantType, 1.f); - }) - .template Case([&](auto) { - return rewriter.getFloatAttr( - constantType, -std::numeric_limits::infinity()); - }) - .template Case([&](auto) { - return rewriter.getFloatAttr( - constantType, std::numeric_limits::infinity()); - }) - .Case([&](arith::MinSIOp) { - return rewriter.getIntegerAttr(constantType, - llvm::maxIntN(bitWidth)); - }) - .Case([&](arith::MinUIOp) { - return rewriter.getIntegerAttr(constantType, - llvm::maxUIntN(bitWidth)); - }) - .Case([&](arith::MaxSIOp) { - return rewriter.getIntegerAttr(constantType, - llvm::minIntN(bitWidth)); - }) - .Case([&](arith::MaxUIOp) { - return rewriter.getIntegerAttr(constantType, 0); - }) - .Case([&](arith::OrIOp) { - return rewriter.getIntegerAttr(constantType, 0); - }) - .Case([&](arith::AndIOp) { - return rewriter.getIntegerAttr(constantType, 1); - }) - .Case([&](arith::XOrIOp) { - return rewriter.getIntegerAttr(constantType, 0); - }) - .Default([](Operation *op) { - op->dump(); - llvm_unreachable("Reduction op not supported yet"); - return nullptr; - }); - - return rewriter.create(redOp->getLoc(), constantType, - attr); - } - - bool requiresF32Conversion(const Type elemType, Operation *redOp) const { - return isa(elemType) && - elemType.getIntOrFloatBitWidth() < - Float32Type::get(elemType.getContext()) - .getIntOrFloatBitWidth() && - (isa(redOp) || isa(redOp)); - } - - Value getRedElement(Value lhs, Value rhs, const Location loc, - Operation *redOp, OpBuilder &b, - const bool convertLhsToF32Precision) const { - return llvm::TypeSwitch(redOp) - .template Case([&](auto redOp) { - if (convertLhsToF32Precision) { - lhs = b.create(loc, Float32Type::get(b.getContext()), - lhs); - } - return b.create(loc, lhs, rhs); - }) - .template Case( - [&](auto redOp) { - return b.create(loc, lhs, rhs); - }) - .Default([](Operation *op) { - op->dump(); - llvm_unreachable("Reduction op not yet supported"); - return nullptr; - }); - } - - virtual bool isReductionOpSupported(Operation *redOp) const = 0; - - virtual LogicalResult - convertToTargetOp(OpTy op, typename OpTy::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const = 0; - - virtual LogicalResult - convertToTargetOpExtended(OpTy op, typename OpTy::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const = 0; -}; - -class ReduceConverter : public ReductionOpBaseConverter { -public: - explicit ReduceConverter(MLIRContext *context) - : ReductionOpBaseConverter(context) {} - - using ReductionOpBaseConverter::ReductionOpBaseConverter; - -protected: - bool isReductionOpSupported(Operation *redOp) const override; - - LogicalResult - convertToTargetOp(triton::ReduceOp op, - typename triton::ReduceOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const override; - - LogicalResult - convertToTargetOpExtended(triton::ReduceOp op, - typename triton::ReduceOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const override; -}; - -class ScanConverter : public ReductionOpBaseConverter { -public: - explicit ScanConverter(MLIRContext *context) - : ReductionOpBaseConverter(context) {} - - using ReductionOpBaseConverter::ReductionOpBaseConverter; - -protected: - bool isReductionOpSupported(Operation *redOp) const override; - - LogicalResult - convertToTargetOp(triton::ScanOp op, typename triton::ScanOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const override; - - LogicalResult - convertToTargetOpExtended(triton::ScanOp op, - typename triton::ScanOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const override; -}; - -class DeviceAssertConverter : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - -private: - static constexpr llvm::StringRef printFuncNameBase = "triton_assert"; - static constexpr llvm::StringRef msgAttrName = "msg"; - -public: - LogicalResult - matchAndRewrite(triton::AssertOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override; -}; - -class DevicePrintConverter : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - -private: - static constexpr llvm::StringRef printFuncNameBase = "triton_print"; - static constexpr llvm::StringRef prefixAttrName = "prefix"; - static constexpr llvm::StringRef hexAttrName = "hex"; - -public: - LogicalResult - matchAndRewrite(triton::PrintOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override; -}; - -} // namespace mlir::dicp::linked - -#endif diff --git a/compiler/include/dicp/Conversion/LinalgToLinked/VerifyNoLinalgGenericPass.hpp b/compiler/include/dicp/Conversion/LinalgToLinked/VerifyNoLinalgGenericPass.hpp deleted file mode 100644 index d0603509..00000000 --- a/compiler/include/dicp/Conversion/LinalgToLinked/VerifyNoLinalgGenericPass.hpp +++ /dev/null @@ -1,59 +0,0 @@ -#pragma once - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/raw_ostream.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Location.h" -#include "mlir/IR/MLIRContext.h" -#include "mlir/IR/Operation.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Transforms/DialectConversion.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Pass/Pass.h" - -#include -#include -#include - -using namespace mlir; -using namespace mlir::func; - -namespace mlir { -namespace dicp { -namespace linked { - -struct VerifyNoLinalgGenericPass : public PassWrapper> { - StringRef getArgument() const final { return "verify-no-linalg-generic"; } - StringRef getDescription() const final { - return "Verify that no 'linalg.generic' operations exist"; - } - - void runOnOperation() override { - bool foundGeneric = false; - getOperation()->walk([&](linalg::GenericOp op) { - op.emitError() << "linalg.generic is not allowed in this pass pipeline."; - foundGeneric = true; - }); - - if (foundGeneric) { - signalPassFailure(); - } - } -}; - -inline std::unique_ptr createVerifyNoLinalgGenericPass() { - return std::make_unique(); -} - -} // namespace linked -} // namespace dicp -} // namespace mlir diff --git a/compiler/include/dicp/Conversion/LinalgToNPU/CMakeLists.txt b/compiler/include/dicp/Conversion/LinalgToNPU/CMakeLists.txt deleted file mode 100644 index d99c572d..00000000 --- a/compiler/include/dicp/Conversion/LinalgToNPU/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls --name LinalgToNPU) -add_public_tablegen_target(LinalgToNPUConversionPassIncGen) diff --git a/compiler/include/dicp/Conversion/LinalgToNPU/ConversionPatterns.hpp b/compiler/include/dicp/Conversion/LinalgToNPU/ConversionPatterns.hpp deleted file mode 100644 index ceb5217e..00000000 --- a/compiler/include/dicp/Conversion/LinalgToNPU/ConversionPatterns.hpp +++ /dev/null @@ -1,174 +0,0 @@ -#pragma once -#include "dicp/Dialect/NPU/IR/NPUDialect.h" -#include "dicp/Dialect/NPU/IR/NPUTypes.h" -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/Passes.h" -#include "llvm/ADT/SmallVectorExtras.h" -#include "llvm/ADT/TypeSwitch.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/FormatVariadic.h" -#include "llvm/Support/MathExtras.h" - -#include -#include -#include - -using namespace mlir; -using namespace dicp; - -namespace { - -template -static Value insertFront(T type, Location loc, - PatternRewriter &rewriter) { - auto tq = rewriter.create(loc, type); - auto *parentBlock = tq->getBlock(); - tq->moveBefore(&parentBlock->front()); - return tq; -} - -struct CopyConverter : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(memref::CopyOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto args = adaptor.getOperands(); - auto replacement = rewriter.create( - op.getLoc(), args[0], args[1]); - rewriter.replaceOp(op, replacement); - return success(); - } -}; - -struct SubViewConverter : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(memref::SubViewOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto args = adaptor.getOperands(); - auto replacement = rewriter.create( - op.getLoc(), op.getResult().getType(), args[0]); - rewriter.replaceOp(op, replacement); - return success(); - } -}; - -// struct AddFConverter : public OpConversionPattern { -// using OpConversionPattern::OpConversionPattern; - -// LogicalResult -// matchAndRewrite(arith::AddFOp op, OpAdaptor adaptor, -// ConversionPatternRewriter &rewriter) const override { -// auto loc = op.getLoc(); -// auto args = adaptor.getOperands(); -// Type resultType = op.getResult().getType(); - -// auto tPipType = npu::TPipType::get(resultType.getContext(), 1); -// auto tpip = rewriter.create(loc,tPipType);//, rewriter.getI32IntegerAttr(1)); - -// auto tQueueType = npu::TQueueType::get(resultType.getContext(), 1); -// rewriter.create(loc,tQueueType);//, rewriter.getI32IntegerAttr(1)); - -// Value replacement = rewriter.create( -// loc, resultType, args[0], args[1]); - -// rewriter.replaceOp(op, replacement); -// return success(); -// } -// }; - -struct LinalgGenericConverter : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(linalg::GenericOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto loc = op.getLoc(); - auto args = adaptor.getOperands(); - Value lhs = op->getOperand(0); - Value rhs = op->getOperand(1); - Operation* lDefOp = lhs.getDefiningOp(); - Operation* rDefOp = rhs.getDefiningOp(); - rewriter.setInsertionPointAfter(lDefOp); - // global tpip - auto tPipType = npu::TPipType::get(getContext()); - Value tpip = rewriter.create(loc,tPipType); - // global tqueue - auto vecInQueueType = npu::TQueueType::get(getContext(), 0, 2); - Value lQueue = rewriter.create(loc, vecInQueueType); - Value rQueue = rewriter.create(loc, vecInQueueType); - auto vecOutQueueType = npu::TQueueType::get(getContext(), 1, 2); - Value outQueue = rewriter.create(loc, vecOutQueueType); - // global tensor - auto GlobalTensorType = npu::GlobalTensorType::get(getContext(), 0); - Value lGlobalTensor = rewriter.create(loc, GlobalTensorType); - Value rGlobalTensor = rewriter.create(loc, GlobalTensorType); - Value outGlobalTensor = rewriter.create(loc, GlobalTensorType); - // init - lGlobalTensor = rewriter.create(loc, GlobalTensorType, lGlobalTensor); - rGlobalTensor = rewriter.create(loc, GlobalTensorType, rGlobalTensor); - outGlobalTensor = rewriter.create(loc, GlobalTensorType, outGlobalTensor); - lQueue = rewriter.create(loc, vecInQueueType, tpip, lQueue); - rQueue = rewriter.create(loc, vecInQueueType, tpip, rQueue); - outQueue = rewriter.create(loc, vecOutQueueType, tpip, outQueue); - - rewriter.setInsertionPointAfter(op); - // copy in - if (isa(rDefOp)) { - auto rLocal = rewriter.create(loc, rhs.getType(), rQueue); - auto *parentBlock = rLocal->getBlock(); - rLocal->moveBefore(&parentBlock->front()); - rewriter.replaceOp(rDefOp, rLocal); - for (auto user : rhs.getUsers()) { - if (isa(user)) { - rewriter.setInsertionPointAfter(user); - rQueue = rewriter.create(loc, vecInQueueType, rQueue, rLocal); - } - } - } - if (isa(lDefOp)) { - auto lLocal = rewriter.create(loc, lhs.getType(), lQueue); - auto *parentBlock = lLocal->getBlock(); - lLocal->moveBefore(&parentBlock->front()); - rewriter.replaceOp(lDefOp, lLocal); - for (auto user : lhs.getUsers()) { - if (isa(user)) { - rewriter.setInsertionPointAfter(user); - lQueue = rewriter.create(loc, vecInQueueType, lQueue, lLocal); - } - } - } - // compute - rewriter.setInsertionPointAfter(op); - Value lLocal = rewriter.create(loc, lhs.getType(), lQueue); - Value rLocal = rewriter.create(loc, rhs.getType(), rQueue); - ValueRange outputs = op.getOutputs(); - Value outLocal = rewriter.create(loc, outputs[0].getType(), outQueue); - Operation* outDefOp = outputs[0].getDefiningOp(); - if (isa(outDefOp)) { - rewriter.replaceOp(outDefOp, outLocal); - } - auto replacement = rewriter.create(loc, lLocal, rLocal, outLocal); - rewriter.replaceOp(op, replacement); - outQueue = rewriter.create(loc, vecOutQueueType, outQueue, outLocal); - lQueue = rewriter.create(loc, vecInQueueType, lQueue, lLocal); - rQueue = rewriter.create(loc, vecInQueueType, rQueue, rLocal); - // copy out - outLocal = rewriter.create(loc, outputs[0].getType(), outQueue); - for (auto user : outputs[0].getUsers()) { - if (isa(user)) { - user->setOperand(0, outLocal); - rewriter.setInsertionPointAfter(user); - outQueue = rewriter.create(loc, vecOutQueueType, outQueue, outLocal); - } - } - return success(); - } -}; - -} // namespace diff --git a/compiler/include/dicp/Conversion/LinalgToNPU/LinalgToNPU.h b/compiler/include/dicp/Conversion/LinalgToNPU/LinalgToNPU.h deleted file mode 100644 index 652a2cb7..00000000 --- a/compiler/include/dicp/Conversion/LinalgToNPU/LinalgToNPU.h +++ /dev/null @@ -1,11 +0,0 @@ -#pragma once - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Pass/Pass.h" - -namespace mlir::dicp::npu { -void populateLinalgToNPUConversionPatterns(RewritePatternSet &patterns); - -std::unique_ptr> createLinalgToNPUPass(); - -} // namespace mlir::dicp::npu diff --git a/compiler/include/dicp/Conversion/LinalgToNPU/Passes.h b/compiler/include/dicp/Conversion/LinalgToNPU/Passes.h deleted file mode 100644 index f13a824c..00000000 --- a/compiler/include/dicp/Conversion/LinalgToNPU/Passes.h +++ /dev/null @@ -1,10 +0,0 @@ - -#pragma once -#include "dicp/Conversion/LinalgToNPU/LinalgToNPU.h" - -namespace mlir::dicp::npu { - -#define GEN_PASS_REGISTRATION -#include "dicp/Conversion/LinalgToNPU/Passes.h.inc" - -} // namespace mlir::dicp::npu diff --git a/compiler/include/dicp/Conversion/LinalgToNPU/Passes.td b/compiler/include/dicp/Conversion/LinalgToNPU/Passes.td deleted file mode 100644 index 32196a64..00000000 --- a/compiler/include/dicp/Conversion/LinalgToNPU/Passes.td +++ /dev/null @@ -1,11 +0,0 @@ -#ifndef LINALG_TO_NPU_CONVERSION_PASSES -#define LINALG_TO_NPU_CONVERSION_PASSES - -include "mlir/Pass/PassBase.td" - -def LinalgToNPU : Pass<"linalg-to-npu", "mlir::ModuleOp"> { - let summary = "Convert Linalg to NPU dialect"; - let constructor = "npu::createLinalgToNPUPass()"; -} - -#endif diff --git a/compiler/include/dicp/Conversion/LinkedToHIVM/CMakeLists.txt b/compiler/include/dicp/Conversion/LinkedToHIVM/CMakeLists.txt deleted file mode 100644 index e2725728..00000000 --- a/compiler/include/dicp/Conversion/LinkedToHIVM/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls --name LinkedToHIVM) -add_public_tablegen_target(LinkedToHIVMConversionPassIncGen) diff --git a/compiler/include/dicp/Conversion/LinkedToHIVM/Passes.h b/compiler/include/dicp/Conversion/LinkedToHIVM/Passes.h deleted file mode 100644 index 7f8407ee..00000000 --- a/compiler/include/dicp/Conversion/LinkedToHIVM/Passes.h +++ /dev/null @@ -1,13 +0,0 @@ -#pragma once - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Pass/Pass.h" - -namespace mlir::dicp::linked { - -std::unique_ptr> createLinkedToHIVMPass(); - -#define GEN_PASS_REGISTRATION -#include "dicp/Conversion/LinkedToHIVM/Passes.h.inc" - -} // namespace mlir::dicp::linked diff --git a/compiler/include/dicp/Conversion/LinkedToHIVM/Passes.td b/compiler/include/dicp/Conversion/LinkedToHIVM/Passes.td deleted file mode 100644 index 9b9844af..00000000 --- a/compiler/include/dicp/Conversion/LinkedToHIVM/Passes.td +++ /dev/null @@ -1,12 +0,0 @@ -#ifndef LINKED_TO_HIVM_CONVERSION_PASSES -#define LINKED_TO_HIVM_CONVERSION_PASSES - -include "mlir/Pass/PassBase.td" - -def LinkedToHIVM : Pass<"linked-to-hivm", "mlir::ModuleOp"> { - let summary = "Convert Linked to HIVM dialect"; - let constructor = "linked::createLinkedToHIVMPass()"; - let dependentDialects = ["hivm::HIVMDialect"]; -} - -#endif // LINKED_TO_HIVM_CONVERSION_PASSES diff --git a/compiler/include/dicp/Conversion/TritonToLinalgNPU/CMakeLists.txt b/compiler/include/dicp/Conversion/TritonToLinalgNPU/CMakeLists.txt deleted file mode 100644 index 5cd7900f..00000000 --- a/compiler/include/dicp/Conversion/TritonToLinalgNPU/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -dicp_add_all_subdirs() \ No newline at end of file diff --git a/compiler/include/dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/CMakeLists.txt b/compiler/include/dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/CMakeLists.txt deleted file mode 100644 index 5d8615fe..00000000 --- a/compiler/include/dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls --name MemRefCopyGatherToTensorInsert) -add_public_tablegen_target(MemRefCopyGatherToTensorInsertPassIncGen) diff --git a/compiler/include/dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.h b/compiler/include/dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.h deleted file mode 100644 index d6664621..00000000 --- a/compiler/include/dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.h +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef MEMREF_COPY_GATHER_TO_TENSOR_INSERT_PASSES_H -#define MEMREF_COPY_GATHER_TO_TENSOR_INSERT_PASSES_H - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Pass/Pass.h" - -namespace mlir::dicp::linked { - -std::unique_ptr> -createMemRefCopyGatherToTensorInsertPass(); - -#define GEN_PASS_REGISTRATION -#include "dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.h.inc" - -} // namespace mlir::dicp::linked - -#endif diff --git a/compiler/include/dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.td b/compiler/include/dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.td deleted file mode 100644 index a0f9e4a4..00000000 --- a/compiler/include/dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.td +++ /dev/null @@ -1,24 +0,0 @@ -#ifndef MEMREF_COPY_GATHER_TO_TENSOR_INSERT_PASSES -#define MEMREF_COPY_GATHER_TO_TENSOR_INSERT_PASSES - -include "mlir/Pass/PassBase.td" - -def MemRefCopyGatherToTensorInsert : Pass<"discrete-gather-to-direct-insert", "mlir::ModuleOp"> { - let summary = "Converts alloc+copy based gather loops to direct tensor insertions"; - let description = [{ - Identifies patterns where a temporary memref is allocated, populated via - index-based gathers (subview + copy) in a loop, and then converted to a tensor. - Replaces this with a tensor.empty + scf.for (iter_args) + tensor.insert - sequence to avoid stack allocation and enable register-level operation. - }]; - let constructor = "::mlir::dicp::linked::createMemRefCopyGatherToTensorInsertPass()"; - let dependentDialects = [ - "mlir::scf::SCFDialect", - "mlir::memref::MemRefDialect", - "mlir::tensor::TensorDialect", - "mlir::arith::ArithDialect", - "mlir::bufferization::BufferizationDialect" - ]; -} - -#endif diff --git a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonArithToLinalg/ConversionPatterns.hpp b/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonArithToLinalg/ConversionPatterns.hpp deleted file mode 100644 index 3f64492c..00000000 --- a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonArithToLinalg/ConversionPatterns.hpp +++ /dev/null @@ -1,503 +0,0 @@ -#ifndef TRITON_NPU_CONVERSION_PATTERNS -#define TRITON_NPU_CONVERSION_PATTERNS - -#include "dicp/Utils/Utils.h" - -#include "triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp" -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include "bishengir/Dialect/Annotation/IR/Annotation.h" - -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/Passes.h" -#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" - -#include "llvm/ADT/SmallVectorExtras.h" -#include "llvm/ADT/TypeSwitch.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/FormatVariadic.h" -#include "llvm/Support/MathExtras.h" - -#include -#include -#include - -using namespace mlir; -using namespace triton; -using namespace mlir::dicp::linked; - -namespace { - -enum class InputPrecision : uint32_t { - TF32 = 0, - TF32x3 = 1, - IEEE = 2, - HF32 = 3, -}; -static ::llvm::StringRef stringifyInputPrecision(InputPrecision val) { - switch (val) { - case InputPrecision::TF32: - return "tf32"; - case InputPrecision::TF32x3: - return "tf32x3"; - case InputPrecision::IEEE: - return "ieee"; - case InputPrecision::HF32: - return "hf32"; - } - return ""; -} - -struct BroadcastNPUConverter : public OpConversionPattern { -private: - using OpConversionPattern::OpConversionPattern; - - SmallVector getBroadcastDims(RankedTensorType src, - RankedTensorType dst) const { - SmallVector broadcastDims; - auto srcShape = src.getShape(); - auto dstShape = dst.getShape(); - - for (size_t i = 0; i < srcShape.size(); i++) { - if (dstShape[i] != srcShape[i]) { - assert(srcShape[i] == 1); - broadcastDims.push_back(i); - } - } - assert(!broadcastDims.empty() && "cannot identify broadcast dimension"); - return broadcastDims; - } - - // Broadcasts input tensor based on TosaToLinalg's broadcastToShape - AffineMap getBroadcastAffineMap(MLIRContext *context, - ArrayRef inputShape, - ArrayRef broadcastToShape) const { - - assert(broadcastToShape.size() >= inputShape.size()); - - // Create affine map and shapes for tensor initialization. - SmallVector outExpr; - - size_t diff = broadcastToShape.size() - inputShape.size(); - for (size_t i = 0; i < broadcastToShape.size(); i++) { - if (i < diff) { - continue; - } - size_t j = i - diff; - if (inputShape[j] == 1) { - // Broadcast singleton dimension - outExpr.push_back(mlir::getAffineConstantExpr(0, context)); - continue; - } - // Non-broadcast case - outExpr.push_back(mlir::getAffineDimExpr(i, context)); - } - return AffineMap::get(broadcastToShape.size(), 0, outExpr, context); - } - -public: - // Dimensions of collapesd tensor is all unbroadcast dims - SmallVector getUnbroadcastDims(RankedTensorType src, - RankedTensorType dst) const { - SmallVector unbroadcastDims; - auto srcShape = src.getShape(); - auto dstShape = dst.getShape(); - - for (size_t i = 0; i < srcShape.size(); ++i) { - if (dstShape[i] == srcShape[i]) { - unbroadcastDims.emplace_back(srcShape[i]); - } - } - return unbroadcastDims; - } - // Here convert tt.broadcast to linalg.broadcast - // - // before - // %out = tt.broadcast %in : tensor<1x4x8xf32> -> tensor<128x4x8xf32> - // - // after - // %collpased = tensor.collapse_shape %in [[0, 1], [2]] : - // tensor<1x4x8xf32> into tensor<4x8xf32> - // %out = linalg.broadcast ins(%collpased : tensor<4x8xf32>) - // outs(%empty : tensor<128x4x8xf32>) dimensions = [0] - LogicalResult - matchAndRewrite(triton::BroadcastOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - assert(op->getNumResults() == 1 && "BroadcastOp assumes single result"); - - RankedTensorType sourceType = - cast(adaptor.getSrc().getType()); - RankedTensorType resultType = cast(op.getType()); - auto elementType = resultType.getElementType(); - size_t resultRank = resultType.getRank(); - auto loc = op.getLoc(); - - auto initEmpty = rewriter.create( - loc, resultType.getShape(), elementType); - - SmallVector broadcastDims = - getBroadcastDims(sourceType, resultType); - SmallVector unbroadcastDims = - getUnbroadcastDims(sourceType, resultType); - - SmallVector collapseReassociationIndices; - auto collapseReassociationIndicesOptional = - getReassociationIndicesForCollapse(sourceType.getShape(), - unbroadcastDims); - if (!collapseReassociationIndicesOptional.has_value()) { - return rewriter.notifyMatchFailure( - op, "Failure with getReassociationIndicesForCollapse call"); - } - collapseReassociationIndices = collapseReassociationIndicesOptional.value(); - - RankedTensorType collapseResultType = - RankedTensorType::get(unbroadcastDims, sourceType.getElementType()); - - auto collpasedOp = rewriter.create( - loc, collapseResultType, adaptor.getSrc(), - collapseReassociationIndices); - - auto broadcastOp = rewriter.create( - loc, collpasedOp, initEmpty, - rewriter.getDenseI64ArrayAttr(broadcastDims)); - - rewriter.replaceOp(op, broadcastOp.getResults()); - return success(); - } -}; - -struct MatmulNPUConverter : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - // true means tensor elements are zeros - // false means not zero or it cannot be determined - bool isZeroTensor(Value &v, bool integers) const { - if (auto splatOp = v.getDefiningOp()) { - if (auto constOp = splatOp.getSrc().getDefiningOp()) { - if (auto val = dyn_cast(constOp.getValue())) { - return val.getValueAsDouble() == 0.; - } - if (auto val = dyn_cast(constOp.getValue())) { - return val.getValue() == 0; - } - } - return false; - } - - if (auto constOp = v.getDefiningOp()) { - if (auto denseAttr = dyn_cast(constOp.getValue())) { - if (denseAttr.isSplat()) { - if (integers) - return denseAttr.getSplatValue().isZero(); - return denseAttr.getSplatValue().isZero(); - } - } - } - - return false; - } - - LogicalResult - matchAndRewrite(triton::DotOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto opa = adaptor.getA(); - auto opb = adaptor.getB(); - auto opc = adaptor.getC(); - auto dstType = cast(op.getType()); - auto inputPrec = op.getInputPrecision(); - - if (dstType.getRank() == 2) { - auto matmulOp = rewriter.replaceOpWithNewOp( - op, ValueRange{opa, opb}, ValueRange{opc}); - matmulOp->setAttr( - "input_precison", - rewriter.getStringAttr(stringifyInputPrecision(inputPrec))); - } else if (dstType.getRank() == 3) { - auto matmulOp = rewriter.replaceOpWithNewOp( - op, ValueRange{opa, opb}, ValueRange{opc}); - matmulOp->setAttr( - "input_precison", - rewriter.getStringAttr(stringifyInputPrecision(inputPrec))); - } else { - llvm_unreachable("Datatype of DotOp operands could only be 2D or 3D"); - } - return success(); - } -}; - -struct ReduceNPUConverter : public OpConversionPattern { - - ReduceNPUConverter(MLIRContext *context, bool transposeToRank0 = false, - PatternBenefit benefit = 1) - : OpConversionPattern(context, benefit), - transposeToRank0(transposeToRank0) {} - -private: - bool transposeToRank0; - - llvm::SmallVector getRedOps(triton::ReduceOp redOp) const { - auto reduceBlock = redOp.getBody(); - return llvm::map_to_vector(reduceBlock->without_terminator(), - [](Operation &op) { return &op; }); - } - - bool isReductionOpSupported(Operation *redOp) const { - return isa(redOp); - } - - arith::ConstantOp getRedBaseConstOp(ConversionPatternRewriter &rewriter, - Operation *redOp, - Type constantType) const { - const int64_t bitWidth = constantType.getIntOrFloatBitWidth(); - - auto attr = - llvm::TypeSwitch(redOp) - .Case([&](arith::AddFOp) { - return rewriter.getFloatAttr(constantType, 0.f); - }) - .Case([&](arith::AddIOp) { - return rewriter.getIntegerAttr(constantType, 0); - }) - .Case([&](auto) { - return rewriter.getFloatAttr( - constantType, -std::numeric_limits::infinity()); - }) - .Case([&](auto) { - return rewriter.getFloatAttr( - constantType, std::numeric_limits::infinity()); - }) - .Case([&](arith::MinSIOp) { - return rewriter.getIntegerAttr(constantType, - llvm::maxIntN(bitWidth)); - }) - .Case([&](arith::MinUIOp) { - return rewriter.getIntegerAttr(constantType, - llvm::maxUIntN(bitWidth)); - }) - .Case([&](arith::MaxSIOp) { - return rewriter.getIntegerAttr(constantType, - llvm::minIntN(bitWidth)); - }) - .Case( - [&](auto) { return rewriter.getIntegerAttr(constantType, 0); }) - .Case([&](arith::MulFOp) { - return rewriter.getFloatAttr(constantType, 1.f); - }) - .Case( - [&](auto) { return rewriter.getIntegerAttr(constantType, 1); }) - .Case([&](arith::OrIOp) { - return rewriter.getIntegerAttr(constantType, 0); - }) - .Default([](Operation *op) { - op->dump(); - llvm_unreachable("Reduction op not yet supported"); - return nullptr; - }); - - return rewriter.create(redOp->getLoc(), constantType, - attr); - } - - bool requiresF32Conversion(const Type elemType, Operation *redOp) const { - unsigned width = - cast(Float32Type::get(elemType.getContext())).getWidth(); - return isa(elemType) && - elemType.getIntOrFloatBitWidth() < width && - isa(redOp); - } - - Value getRedElement(Value lhs, Value rhs, const Location loc, - Operation *redOp, OpBuilder &b, - const bool convertLhsToF32Precision) const { - return llvm::TypeSwitch(redOp) - .Case([&](auto redOp) { - if (convertLhsToF32Precision) { - lhs = b.create(loc, Float32Type::get(b.getContext()), - lhs); - } - return b.create(loc, lhs, rhs); - }) - .Case([&](auto redOp) { - return b.create(loc, lhs, rhs); - }) - .Default([](Operation *op) { - op->dump(); - llvm_unreachable("Reduction op not yet supported"); - return nullptr; - }); - } - - LogicalResult - convertToLinalgReduce(triton::ReduceOp op, - typename triton::ReduceOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const { - auto source = adaptor.getOperands().front(); - auto sourceType = cast(source.getType()); - auto elemType = sourceType.getElementType(); - auto resType = op.getResult().front().getType(); - auto loc = op.getLoc(); - auto reductionOps = getRedOps(op); - - // Reduction of arbitrary operations isn't supported because using the first - // element across the reduction dimension requires us to iterate over a - // subview that skips over each first element. - if (reductionOps.size() != 1 || - !isReductionOpSupported(reductionOps.front())) { - return rewriter.notifyMatchFailure( - op, "Only support lowering reduction with body " - "containing 1 max(i/f), addf, ori, or mulf."); - } - - auto rop = reductionOps.front(); - auto axis = op.getAxis(); - auto rank = sourceType.getRank(); - auto isVectorReduce = (rank == 1); - - // For now we are transposing reductions from Triton Shared as an - // optimization. This should not be the job of Triton Shared so moving - // forward this will be removed. Doing the transpose here lacks a wider - // scope of analysis that might indicate that the transpose to a given axis - // is not optimal. - if (transposeToRank0) { - // if it is not a vector reduce, we can transpose the source - // so that the reduction axis is the first dimension. - if (!isVectorReduce && axis != 0) { - SmallVector order; - order.reserve(rank); - order.push_back(axis); - for (int i = 0; i < rank; ++i) { - if (i != axis) { - order.push_back(i); - } - } - source = getTransposedValue(source, op.getLoc(), rewriter, order); - axis = 0; - } - } - - bool convertToF32Precision = requiresF32Conversion(resType, rop); - - auto constantType = convertToF32Precision - ? Float32Type::get(rewriter.getContext()) - : elemType; - - auto accBaseConstOp = getRedBaseConstOp(rewriter, rop, constantType); - Value initTensor; - - if (isVectorReduce) { - // The affine vectorizer cannot vectorize affine loops generated from - // linalg.reduce for the vector reduce case, so we must rewrite the - // linalg.reduce to affine loops manually. Here we lower to AllocTensor - // directly instead of EmptyOp so that the subsequent pass can recognize - // the patterns (EmptyOp is susceptible to being CSE'd away, making it - // harder to match the patterns correctly). - initTensor = rewriter.create( - loc, RankedTensorType::get({}, constantType), ValueRange{}); - initTensor = rewriter.create(loc, accBaseConstOp, - initTensor, ValueRange{}); - } else { - Value init = rewriter.create( - loc, cast(resType).getShape(), constantType); - initTensor = rewriter - .create(loc, ValueRange{accBaseConstOp}, - ValueRange{init}) - .result(); - } - - Value finalResult = - rewriter - .create( - loc, ValueRange{source}, ValueRange{initTensor}, - SmallVector{axis}, - [&](OpBuilder &opBuilder, Location loc, ValueRange inputs) { - assert(inputs.size() == 2); - Value result = - getRedElement(inputs[0], inputs[1], loc, rop, opBuilder, - convertToF32Precision); - opBuilder.create(loc, result); - }) - .getResult(0); - - if (isVectorReduce) { - finalResult = - rewriter.create(loc, constantType, finalResult); - } - - if (convertToF32Precision) { - finalResult = rewriter.create(loc, resType, finalResult); - } - - rewriter.replaceOp(op, finalResult); - return success(); - } - -public: - LogicalResult - matchAndRewrite(triton::ReduceOp op, - typename triton::ReduceOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto sourceType = - cast(adaptor.getOperands().front().getType()); - assert(sourceType.hasRank() && "Expected input is " - "ranked"); - - int64_t axis = op.getAxis(); - assert(axis >= 0 && axis < sourceType.getRank() && - "Expected reduction " - "axis is within " - "operand's rank"); - - return convertToLinalgReduce(op, adaptor, rewriter); - } -}; - -class BitcastNPUConverter : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(triton::BitcastOp op, - PatternRewriter &rewriter) const { - if (op->hasAttr("Input_Arg_i1_Bitcast_To_i8")) { - return failure(); - } - - Value result; - if (auto resPointerType = dyn_cast(op.getType())) { - // TODO: use typeconverter - auto srcPointerType = cast(op.getSrc().getType()); - auto resType = MemRefType::get({ShapedType::kDynamic}, - resPointerType.getPointeeType()); - // Handling special case - // %0 = tt.bitcast %arg0 {MixUse} : !tt.ptr -> !tt.ptr - if (isa(op.getSrc()) && - srcPointerType.getPointeeType() == rewriter.getIntegerType(1) && - resPointerType.getPointeeType() == rewriter.getIntegerType(8)) { - rewriter.modifyOpInPlace(op, [&]() { - op->setAttr("Input_Arg_i1_Bitcast_To_i8", rewriter.getUnitAttr()); - }); - return success(); - } - result = - rewriter.create(op.getLoc(), resType, op.getSrc()); - } else { - result = rewriter.create(op.getLoc(), op.getType(), - op.getSrc()); - } - rewriter.replaceOp(op, result); - return success(); - } -}; - -} // namespace - -#endif diff --git a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonArithToLinalg/NPUSpecific.h b/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonArithToLinalg/NPUSpecific.h deleted file mode 100644 index 54e2917f..00000000 --- a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonArithToLinalg/NPUSpecific.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef TRITON_CONVERSION_TRITONARITHTOLINALGNPU_H -#define TRITON_CONVERSION_TRITONARITHTOLINALGNPU_H - -#include "mlir/Pass/Pass.h" -#include "mlir/Transforms/DialectConversion.h" - -#include "triton/Dialect/Triton/IR/Dialect.h" - -namespace mlir::dicp::linked { - -void populateTritonArithToLinalgNPUConversionPatterns( - bool pidsToFuncArgs, bool addptrToLinalg, bool assertToCf, - bool transposeReduceToRank0, RewritePatternSet &patterns); - -// NPU Specialization: Dynamically determines if an arith/math operation is -// legal for lowering. An operation is legal if it does NOT require -// conversion/lowering by the current pass. -bool isLegalConstantAndTensorArithmeticOpForNPU(Operation *op); - -} // namespace mlir::dicp::linked - -#endif // TRITON_CONVERSION_TritonArithToLinalgNPU_TritonArithToLinalgNPU_H diff --git a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/Passes.h b/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/Passes.h deleted file mode 100644 index 842ab27b..00000000 --- a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/Passes.h +++ /dev/null @@ -1,12 +0,0 @@ -#ifndef TRITON_TO_LINALG_NPU_CONVERSION_PASSES_H -#define TRITON_TO_LINALG_NPU_CONVERSION_PASSES_H - -#include "dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/TritonToLinalgNPUCoversion.h" - -namespace mlir::dicp::linked { -#define GEN_PASS_REGISTRATION -#include "dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/Passes.h.inc" - -} // namespace mlir::dicp::linked - -#endif diff --git a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/Passes.td b/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/Passes.td deleted file mode 100644 index e7bf0381..00000000 --- a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/Passes.td +++ /dev/null @@ -1,17 +0,0 @@ -#ifndef TRITON_TO_LINALG_NPU_CONVERSION_PASSES -#define TRITON_TO_LINALG_NPU_CONVERSION_PASSES - -include "mlir/Pass/PassBase.td" - -def TritonToLinalgNPUCoversion : Pass<"triton-to-linalg-npu-conversion", "mlir::ModuleOp"> { - let summary = "Convert Triton to Linalg dialect"; - let constructor = "::mlir::dicp::linked::createTritonToLinalgNPUCoversionPass()"; - let options = [ - Option<"enableMakeGatherScatterTensorPtr", "enable-make-gather-scatter", "bool", /*default*/"true", - "Enable make_gather_scatter_tptr support">, - Option<"enableCollapseShape", "enable-collapse-shape", "bool", /*default*/"false", - "Enable collapse shape pass">, - ]; -} - -#endif diff --git a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/TritonToLinalgNPUCoversion.h b/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/TritonToLinalgNPUCoversion.h deleted file mode 100644 index 7c5ebc80..00000000 --- a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/TritonToLinalgNPUCoversion.h +++ /dev/null @@ -1,13 +0,0 @@ -#ifndef TRITON_CONVERSION_TRITONTOLINALG_NPU_TRITONTOLINALGH -#define TRITON_CONVERSION_TRITONTOLINALG_NPU_TRITONTOLINALGH - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Pass/Pass.h" - -namespace mlir::dicp::linked { - -std::unique_ptr> createTritonToLinalgNPUCoversionPass(); - -} // namespace mlir::dicp::linked - -#endif // TRITON_CONVERSION_TRITONTOLINALG_NPU_TRITONTOLINALGH diff --git a/compiler/include/dicp/Conversion/TritonToUnstructure/BubbleUpOperation.h b/compiler/include/dicp/Conversion/TritonToUnstructure/BubbleUpOperation.h deleted file mode 100644 index 91990030..00000000 --- a/compiler/include/dicp/Conversion/TritonToUnstructure/BubbleUpOperation.h +++ /dev/null @@ -1,31 +0,0 @@ -#pragma once - -#include "mlir/Pass/Pass.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include "mlir/IR/PatternMatch.h" - -#define GEN_PASS_DECL_BUBBLEUPOPERATION -#include "dicp/Conversion/TritonToUnstructure/Passes.h.inc" - -#define GEN_PASS_DEF_BUBBLEUPOPERATION -#include "dicp/Conversion/TritonToUnstructure/Passes.h.inc" - -namespace mlir { -namespace triton { - -std::unique_ptr> -createBubbleUpOperationPass(const BubbleUpOperationOptions &options = {}); - -} // namespace triton -} // namespace mlir - -using namespace mlir; -using namespace triton; - -class BubbleUpOperationPass - : public ::impl::BubbleUpOperationBase { -public: - explicit BubbleUpOperationPass(const BubbleUpOperationOptions &options); - void runOnOperation() override; -}; diff --git a/compiler/include/dicp/Conversion/TritonToUnstructure/Passes.h b/compiler/include/dicp/Conversion/TritonToUnstructure/Passes.h deleted file mode 100644 index e30f6441..00000000 --- a/compiler/include/dicp/Conversion/TritonToUnstructure/Passes.h +++ /dev/null @@ -1,18 +0,0 @@ - - -#ifndef TRITON_DLC_TRITON_TO_UNSTRUCTURE_CONVERSION_PASSES_H -#define TRITON_DLC_TRITON_TO_UNSTRUCTURE_CONVERSION_PASSES_H - -#include "BubbleUpOperation.h" -#include "UnstructureConversionPass.h" - -namespace mlir { -namespace triton { - -#define GEN_PASS_REGISTRATION -#include "dicp/Conversion/TritonToUnstructure/Passes.h.inc" - -} // namespace triton -} // namespace mlir - -#endif // TRITON_DLC_TRITON_TO_UNSTRUCTURE_CONVERSION_PASSES_H diff --git a/compiler/include/dicp/Dialect/CommonIR/CMakeLists.txt b/compiler/include/dicp/Dialect/CommonIR/CMakeLists.txt new file mode 100644 index 00000000..43f764a7 --- /dev/null +++ b/compiler/include/dicp/Dialect/CommonIR/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name CommonIR) +add_public_tablegen_target(CommonIRPassIncGen) diff --git a/compiler/include/dicp/Dialect/CommonIR/Passes.h b/compiler/include/dicp/Dialect/CommonIR/Passes.h new file mode 100644 index 00000000..a4716394 --- /dev/null +++ b/compiler/include/dicp/Dialect/CommonIR/Passes.h @@ -0,0 +1,22 @@ +#ifndef TRITON_COMMONIR_PASSES_H_ +#define TRITON_COMMONIR_PASSES_H_ + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir::func { +class FuncOp; +} + +namespace mlir::dicp::CommonIR { + +std::unique_ptr> createVectorizeParallelLoopPass(); + +std::unique_ptr> createAnnotateKernelAttrsPass(); + +#define GEN_PASS_REGISTRATION +#include "dicp/Dialect/CommonIR/Passes.h.inc" + +} // namespace mlir::dicp::CommonIR + +#endif // TRITON_COMMONIR_PASSES_H_ diff --git a/compiler/include/dicp/Dialect/CommonIR/Passes.td b/compiler/include/dicp/Dialect/CommonIR/Passes.td new file mode 100644 index 00000000..0368a83e --- /dev/null +++ b/compiler/include/dicp/Dialect/CommonIR/Passes.td @@ -0,0 +1,33 @@ +#ifndef TRITON_COMMONIR_PASSES +#define TRITON_COMMONIR_PASSES + +include "mlir/Pass/PassBase.td" + +def VectorizeParallelLoop : Pass<"vectorize-parallel-loop", "func::FuncOp"> { + let summary = "Convert single-element parallel loops to vectorized batch processing."; + let constructor = "mlir::dicp::CommonIR::createVectorizeParallelLoopPass()"; + let dependentDialects = [ + "mlir::arith::ArithDialect", "mlir::memref::MemRefDialect", + "mlir::tensor::TensorDialect", "mlir::scf::SCFDialect", + "mlir::bufferization::BufferizationDialect", "mlir::func::FuncDialect" + ]; +} + +def AnnotateKernelAttrs : Pass<"annotate-kernel-attrs", "mlir::ModuleOp"> { + let summary = "Annotate kernel func.func with mix_mode/parallel_mode/global_kernel/SyncBlockLockArgIdx/WorkspaceArgIdx."; + let description = [{ + For CommonIR entry kernels that arrive as plain `func.func` (without a + `tt.func` upstream), the downstream TritonToLinalg lowering never gets a + chance to attach the kernel-level attributes that the Ascend host wrapper + parses (`mix_mode`, `parallel_mode`, `global_kernel`, + `SyncBlockLockArgIdx`, `WorkspaceArgIdx`). This pass attaches those + attributes and prepends the two stub arguments (syncBlockLock + workspace) + expected by the NPU runtime. + }]; + let constructor = "mlir::dicp::CommonIR::createAnnotateKernelAttrsPass()"; + let dependentDialects = [ + "mlir::func::FuncDialect", "mlir::memref::MemRefDialect" + ]; +} + +#endif // TRITON_COMMONIR_PASSES diff --git a/compiler/include/dicp/Dialect/LinalgExt/CMakeLists.txt b/compiler/include/dicp/Dialect/LinalgExt/CMakeLists.txt deleted file mode 100644 index 5cd7900f..00000000 --- a/compiler/include/dicp/Dialect/LinalgExt/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -dicp_add_all_subdirs() \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/LinalgExt/IR/CMakeLists.txt b/compiler/include/dicp/Dialect/LinalgExt/IR/CMakeLists.txt deleted file mode 100644 index 71c24091..00000000 --- a/compiler/include/dicp/Dialect/LinalgExt/IR/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -add_dicp_compiler_dialect(DICPLinalgExt LinalgExt linalgext) \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtOpBase.td b/compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtOpBase.td deleted file mode 100644 index f8ba4ba8..00000000 --- a/compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtOpBase.td +++ /dev/null @@ -1,36 +0,0 @@ -#ifndef LINALG_EXT_OPS_BASE -#define LINALG_EXT_OPS_BASE - -include "mlir/IR/OpBase.td" -include "mlir/IR/AttrTypeBase.td" -include "mlir/IR/EnumAttr.td" - -def LinalgExt_Dialect : Dialect { - let name = "linalgext"; - let cppNamespace = "::mlir::dicp::LinalgExt"; - let summary = [{Linalg Extensions}]; - let description = [{ - A dialect designed for experimenting with non-structured operations that - cannot be represented efficiently/directly by the Linalg dialect. - }]; - // let extraClassDeclaration = [{ void registerTypes(); }]; - // let useDefaultAttributePrinterParser = 0; -} - -class LinalgExt_Op traits = []> - : Op {} - -class RankedTensorOrMemRefOf allowedTypes> - : ShapedContainerType< - allowedTypes, - Or<[IsMemRefTypePred, And<[IsTensorTypePred, HasRankPred]>]>, - "ranked tensor or memref", "::mlir::ShapedType">; - -def AnyRankedTensorOrMemRefType : RankedTensorOrMemRefOf<[AnyType]>; - - -def LinalgExt_ElementwiseTrait : NativeOpTrait<"LinalgExt_ElementwiseTrait">; -def LinalgExt_ReduceTrait : NativeOpTrait<"LinalgExt_ReduceTrait">; -def LinalgExt_TensorCoreTrait : NativeOpTrait<"LinalgExt_TensorCoreTrait">; - -#endif // LINALG_EXT_OPS_BASE \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtOps.h b/compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtOps.h deleted file mode 100644 index d848c5e4..00000000 --- a/compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtOps.h +++ /dev/null @@ -1,34 +0,0 @@ -#ifndef DIALECT_LINALGEXT_IR_LINALGEXT_OPS_H_ -#define DIALECT_LINALGEXT_IR_LINALGEXT_OPS_H_ - -#include "dicp/Dialect/LinalgExt/IR/Traits.h" - -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/GPU/IR/GPUDialect.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/IR/LinalgInterfaces.h" -#include "mlir/IR/AffineMap.h" -#include "mlir/IR/Attributes.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Dialect.h" -#include "mlir/IR/MLIRContext.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/Operation.h" -#include "mlir/IR/Value.h" -#include "mlir/Interfaces/ControlFlowInterfaces.h" -#include "mlir/Interfaces/DestinationStyleOpInterface.h" -#include "mlir/Interfaces/InferTypeOpInterface.h" -#include "mlir/Interfaces/SideEffectInterfaces.h" -#include "mlir/Interfaces/TilingInterface.h" -#include "mlir/Support/LLVM.h" - -#include "dicp/Dialect/LinalgExt/IR/LinalgExtDialect.h.inc" - -#define GET_OP_CLASSES -#include "dicp/Dialect/LinalgExt/IR/LinalgExtOps.h.inc" -#define GET_TYPEDEF_CLASSES -#include "dicp/Dialect/LinalgExt/IR/LinalgExtTypes.h.inc" - -#endif // DIALECT_LINALGEXT_IR_LINALGEXT_OPS_H_ \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtOps.td b/compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtOps.td deleted file mode 100644 index 2131c073..00000000 --- a/compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtOps.td +++ /dev/null @@ -1,50 +0,0 @@ -#ifndef DIALECT_LINALGEXT_IR_LINALG_EXT_OPS -#define DIALECT_LINALGEXT_IR_LINALG_EXT_OPS - -include "LinalgExtOpBase.td" -include "mlir/IR/OpBase.td" -include "mlir/Interfaces/SideEffectInterfaces.td" -include "mlir/Dialect/Linalg/IR/LinalgInterfaces.td" -include "mlir/Interfaces/ControlFlowInterfaces.td" -include "mlir/Interfaces/DestinationStyleOpInterface.td" -include "mlir/Interfaces/TilingInterface.td" - -class VectorTwoOperands traits = []> - : LinalgExt_Op { - let arguments = (ins - Arg:$lhs, - Arg:$rhs - ); - let results = (outs AnyMemRef : $outs); - let assemblyFormat = - [{operands attr - dict `:` functional - type(operands, results)}]; -} - -class VectorOneOperands traits = []> - : LinalgExt_Op { - let arguments = (ins Arg : $lhs); - let results = (outs AnyMemRef : $outs); -} - -def LinalgExt_ReduceSumOp : VectorTwoOperands<"reduce_sum", [LinalgExt_ReduceTrait]> { - let summary = "Reduce sum operator"; - let description = [{ - Reduce a tensor along the given axis by computing the sum of the axis.}]; - - let arguments = (ins - AnyRankedTensorOrMemRefType:$input, - AnyRankedTensorOrMemRefType:$output, - DefaultValuedAttr:$axis - ); - - let results = (outs Variadic : $result); - let assemblyFormat = [{ - attr-dict - `reduce_axis` `(` $axis `)` - `ins` `(` $input `:` type($input) `)` - `outs` `(` $output `:` type($output) `)` - (`->` type($result)^)? - }]; -} - -#endif // DIALECT_LINALGEXT_IR_LINALG_EXT_OPS \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/LinalgExt/IR/Traits.h b/compiler/include/dicp/Dialect/LinalgExt/IR/Traits.h deleted file mode 100644 index c5e683a0..00000000 --- a/compiler/include/dicp/Dialect/LinalgExt/IR/Traits.h +++ /dev/null @@ -1,26 +0,0 @@ -#ifndef DIALECT_LINALGEXT_IR_TAITS_H_ -#define DIALECT_LINALGEXT_IR_TAITS_H_ - -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/Support/LogicalResult.h" - -namespace mlir { -namespace OpTrait { - -template -class LinalgExt_ElementwiseTrait - : public TraitBase {}; - -template -class LinalgExt_ReduceTrait - : public TraitBase {}; - -template -class LinalgExt_TensorCoreTrait - : public TraitBase {}; - -} // namespace OpTrait -} // namespace mlir - -#endif // DIALECT_LINALGEXT_IR_TAITS_H_ \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/LinalgExt/Transforms/CMakeLists.txt b/compiler/include/dicp/Dialect/LinalgExt/Transforms/CMakeLists.txt deleted file mode 100644 index 86132beb..00000000 --- a/compiler/include/dicp/Dialect/LinalgExt/Transforms/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls -name LinalgExt) -add_public_tablegen_target(LinalgExtTransformsIncGen) diff --git a/compiler/include/dicp/Dialect/LinalgExt/Transforms/Passes.h b/compiler/include/dicp/Dialect/LinalgExt/Transforms/Passes.h deleted file mode 100644 index ca63827d..00000000 --- a/compiler/include/dicp/Dialect/LinalgExt/Transforms/Passes.h +++ /dev/null @@ -1,34 +0,0 @@ -#ifndef LinalgEXT_DIALECT_Linalg_TRANSFORMS_PASSES_H_ -#define LinalgEXT_DIALECT_Linalg_TRANSFORMS_PASSES_H_ - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Pass/Pass.h" - -namespace mlir { -class TypeConverter; -class ConversionTarget; -namespace func { -class FuncOp; -} -} // namespace mlir - -namespace mlir::dicp::LinalgExt { - -std::unique_ptr> createLinalgIfToSelectPass(); - -std::unique_ptr> createLinalgGenericToSCFPass(); - -std::unique_ptr> createScalarTo1DTensorPass(); - -std::unique_ptr> -createNormalizeSliceOpsPass(); - -std::unique_ptr> -createVectorizeParallelLoopPass(); - -#define GEN_PASS_REGISTRATION -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h.inc" - -} // namespace mlir::dicp::LinalgExt - -#endif diff --git a/compiler/include/dicp/Dialect/LinalgExt/Transforms/Passes.td b/compiler/include/dicp/Dialect/LinalgExt/Transforms/Passes.td deleted file mode 100644 index 04041fea..00000000 --- a/compiler/include/dicp/Dialect/LinalgExt/Transforms/Passes.td +++ /dev/null @@ -1,84 +0,0 @@ -#ifndef LinalgEXT_PASSES -#define LinalgEXT_PASSES - -include "mlir/Pass/PassBase.td" - -def LinalgIfToSelect : Pass<"linalg-if-to-select", "mlir::ModuleOp"> { - let summary = "Convert scf.if inside parallel linalg.generic to arith.select " - "and hoist selects."; - let description = [{ - This pass converts conditional logic (`scf.if`) inside a parallel `linalg.generic` - operation into data-flow operations (`arith.select`). - - The conversion process involves two main stages: - 1. **LinalgIfToSelectPattern**: Converts `scf.if` statements (where the else - block yields a constant or is empty) into an equivalent sequence of - `arith.select` operations. For arguments used inside the `then` block, - it introduces an `arith.select` using a zero constant for the false case. - It also attempts to fix zero-value assignment for `memref.store` operations - that result from the conversion (FixLinalgSelectZeroForStorePattern). - 2. **LinalgLiftSelectPattern**: Lifts scalar `arith.select` operations that use - loop-invariant operands or block arguments outside the `linalg.generic` - operation, converting them into tensor-based `arith.select`s which are - then added as new inputs to the `linalg.generic`. This process helps to - simplify the body of the `linalg.generic` and enables further fusion. - }]; - let constructor = "mlir::dicp::LinalgExt::createLinalgIfToSelectPass()"; - let dependentDialects = [ - "mlir::linalg::LinalgDialect", "mlir::arith::ArithDialect", - "mlir::tensor::TensorDialect", "mlir::scf::SCFDialect", - "mlir::memref::MemRefDialect" - ]; -} - -def LinalgGenericToSCF : Pass<"linalg-generic-to-scf", "mlir::ModuleOp"> { - let summary = "Lower linalg.generic ops to SCF loops"; - let description = [{ - Converts linalg.generic operations into explicit scf.for loops, - generating structured loop nests that preserve the original - indexing maps and iterator types. - }]; - let constructor = "mlir::dicp::LinalgExt::createLinalgGenericToSCFPass()"; - let dependentDialects = ["mlir::linalg::LinalgDialect"]; -} - -def ScalarTo1DTensor : Pass<"scalar-to-1d-tensor", "mlir::func::FuncOp"> { - let summary = - "Convert scalar computations and memref load/store to tensor<1 x T> form"; - let description = - [{The ScalarTo1DTensor pass targets scalar computations and memory access - within a function and rewrites them into an explicit tensor<1 x T> - - based form - .This enables uniform handling of scalar values in subsequent - bufferization and lowering passes.}]; - - let constructor = "mlir::dicp::LinalgExt::createScalarTo1DTensorPass()"; - let dependentDialects = [ - "mlir::arith::ArithDialect", "mlir::memref::MemRefDialect", - "mlir::tensor::TensorDialect", "mlir::bufferization::BufferizationDialect", - "mlir::func::FuncDialect" - ]; -} - -def NormalizeSliceOps : Pass<"normalize-slice-ops", "func::FuncOp"> { - let summary = "Normalize Slice Ops."; - let constructor = "mlir::dicp::LinalgExt::createNormalizeSliceOpsPass()"; - let dependentDialects = ["mlir::tensor::TensorDialect"]; -} - -def VectorizeParallelLoop : Pass<"vectorize-parallel-loop", "func::FuncOp"> { - let summary = "Convert single-element parallel loops to vectorized batch " - "processing with step=size."; - let description = [{This pass transforms parallel loops that process single - elements into vectorized operations that can process - multiple elements simultaneously, - potentially increasing computational throughput.}]; - let constructor = "mlir::dicp::LinalgExt::createVectorizeParallelLoopPass()"; - let dependentDialects = [ - "mlir::arith::ArithDialect", "mlir::memref::MemRefDialect", - "mlir::tensor::TensorDialect", "mlir::scf::SCFDialect", - "mlir::bufferization::BufferizationDialect", "mlir::func::FuncDialect" - ]; -} - -#endif diff --git a/compiler/include/dicp/Dialect/LinalgExt/Transforms/Transforms.h b/compiler/include/dicp/Dialect/LinalgExt/Transforms/Transforms.h deleted file mode 100644 index 2eb34489..00000000 --- a/compiler/include/dicp/Dialect/LinalgExt/Transforms/Transforms.h +++ /dev/null @@ -1,46 +0,0 @@ -#ifndef DICP_LINALGEXT_TRANSFORMS_H_ -#define DICP_LINALGEXT_TRANSFORMS_H_ - -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h" -#include "mlir/Dialect/Vector/IR/VectorOps.h" -#include "mlir/IR/Operation.h" -#include "mlir/Pass/Pass.h" - -namespace mlir::dicp::LinalgExt { - -using ForControlFnRef = llvm::function_ref; - -/// Insert pattern to remove single iteration loop. The pattern will detect -/// single iteration loops based on the range returned ValueBoundsOpInterface. -void populateRemoveSingleIterationLoopPattern( - RewritePatternSet &patterns, ForControlFnRef controlFn = nullptr); - -/** - * Populates the given pattern set with rewrite patterns that lift - * scalar arithmetic select operations out of Linalg generic ops. - * - * This typically includes patterns like: - * - LiftScalarSelectToTensorPattern: Converts a scalar arith.select - * with loop-invariant operands inside a linalg.generic to a - * tensor arith.select placed outside the generic op. - * - LiftYieldSelectOutPattern: Converts an arith.select yielded by - * the linalg.generic body into an outside tensor arith.select. - * - */ -void populateLinalgLiftSelectPattern(RewritePatternSet &patterns); - -/// This pattern detects a chain of `tensor::InsertSliceOp` that together -/// implement an interleave write: multiple source tensors are inserted into the -/// same destination along the last dimension using different static offsets. -/// Once detected, the pattern normalizes this chain into a canonical form -/// where the last-dimension offsets are [0..channelNum-1] and the stride is -/// `channelNum`, making the interleave structure explicit and ready for further -/// fusion or replacement by a dedicated Interleave op. -void populateNormalizeInsertSliceOpInInterleavePattern( - RewritePatternSet &patterns); - -} // namespace mlir::dicp::LinalgExt - -#endif \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/NPU/IR/CMakeLists.txt b/compiler/include/dicp/Dialect/NPU/IR/CMakeLists.txt deleted file mode 100644 index ec3c889a..00000000 --- a/compiler/include/dicp/Dialect/NPU/IR/CMakeLists.txt +++ /dev/null @@ -1,14 +0,0 @@ -# add_dicp_compiler_dialect(DICPNPU NPU npu) -set(LLVM_TARGET_DEFINITIONS NPUOps.td) -mlir_tablegen(NPUDialect.h.inc -gen-dialect-decls -dialect=npu) -mlir_tablegen(NPUDialect.cpp.inc -gen-dialect-defs -dialect=npu) -mlir_tablegen(NPUOps.h.inc -gen-op-decls) -mlir_tablegen(NPUOps.cpp.inc -gen-op-defs) - - - -set(LLVM_TARGET_DEFINITIONS NPUTypes.td) -mlir_tablegen(NPUTypes.h.inc -gen-typedef-decls) -mlir_tablegen(NPUTypes.cpp.inc -gen-typedef-defs) - -add_public_tablegen_target(DICPNPUIncGen) \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/NPU/IR/NPUDialect.h b/compiler/include/dicp/Dialect/NPU/IR/NPUDialect.h deleted file mode 100644 index cd0b6fa9..00000000 --- a/compiler/include/dicp/Dialect/NPU/IR/NPUDialect.h +++ /dev/null @@ -1,21 +0,0 @@ -#pragma once - -#include "dicp/Dialect/NPU/IR/NPUTypes.h" -#include "mlir/Bytecode/BytecodeOpInterface.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Dialect.h" -#include "mlir/IR/MLIRContext.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/SymbolTable.h" -#include "mlir/IR/TypeSupport.h" -#include "mlir/IR/Types.h" -#include "mlir/Interfaces/CallInterfaces.h" -#include "mlir/Interfaces/FunctionInterfaces.h" -#include "mlir/Interfaces/SideEffectInterfaces.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -#define GET_OP_CLASSES -#include "dicp/Dialect/NPU/IR/NPUDialect.h.inc" -#define GET_OP_CLASSES -#include "dicp/Dialect/NPU/IR/NPUOps.h.inc" diff --git a/compiler/include/dicp/Dialect/NPU/IR/NPUDialect.td b/compiler/include/dicp/Dialect/NPU/IR/NPUDialect.td deleted file mode 100644 index 6bc7e16e..00000000 --- a/compiler/include/dicp/Dialect/NPU/IR/NPUDialect.td +++ /dev/null @@ -1,38 +0,0 @@ -#ifndef NPU_DIALECT -#define NPU_DIALECT - -include "mlir/IR/OpBase.td" -// include "mlir/Interfaces/FunctionInterfaces.td" -// include "mlir/IR/SymbolInterfaces.td" -include "mlir/Interfaces/SideEffectInterfaces.td" - -// Provide a definition of the 'NPU' dialect in the ODS framework so that we -// can define our operations. -def NPU_Dialect : Dialect { - let name = "npu"; - let cppNamespace = "::mlir::dicp::npu"; - let usePropertiesForAttributes = 1; - let useDefaultTypePrinterParser = 1; - - let extraClassDeclaration = [{ - void registerTypes(); - }]; -} - - -// def NPU_TPipType : I<4>; -// def NPU_TQueType : I<4>; - -// def NPU_TQueueType : DialectType< -// NPU_Dialect, -// CPred<"isa($_self)">, -// "tqueue"> { -// let description = [{ -// Allocates buffers for a particular device memory space. -// }]; -// let builderCall = "$_builder.getType()"; -// } - -include "dicp/Dialect/NPU/IR/NPUTypes.td" - -#endif // NPU_DIALECT \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/NPU/IR/NPUOps.td b/compiler/include/dicp/Dialect/NPU/IR/NPUOps.td deleted file mode 100644 index cce1e1ff..00000000 --- a/compiler/include/dicp/Dialect/NPU/IR/NPUOps.td +++ /dev/null @@ -1,264 +0,0 @@ -#ifndef NPU_DIALECT_OPS -#define NPU_DIALECT_OPS - -include "dicp/Dialect/NPU/IR/NPUDialect.td" -include "dicp/Dialect/NPU/IR/NPUTypes.td" -include "mlir/IR/OpBase.td" -include "mlir/IR/SymbolInterfaces.td" // SymbolUserOpInterface -include "mlir/IR/OpAsmInterface.td" // OpAsmOpInterface -include "mlir/Interfaces/CallInterfaces.td" // CallOpInterface -include "mlir/Interfaces/CastInterfaces.td" // CastOpInterface -include "mlir/Interfaces/FunctionInterfaces.td" // FunctionOpInterface -include "mlir/Interfaces/SideEffectInterfaces.td" // Pure -include "mlir/Interfaces/ControlFlowInterfaces.td" // BranchOpInterface -include "mlir/Interfaces/InferTypeOpInterface.td" // SameOperandsAndResultType -include "mlir/Interfaces/SideEffectInterfaces.td" // Pure -include "mlir/Interfaces/CastInterfaces.td" // CastOpInterface -include "mlir/Interfaces/CallInterfaces.td" // CallOpInterface -include "mlir/IR/BuiltinAttributeInterfaces.td" - -// Provide a definition of the 'NPU' dialect in the ODS framework so that we -// can define our operations. -// def NPU_Dialect : Dialect { -// let name = "npu"; -// let cppNamespace = "::mlir::dicp::npu"; -// let usePropertiesForAttributes = 1; -// } - -// Base class for NPU dialect operations. This operation inherits from the base -// `Op` class in OpBase.td, and provides: -// * The parent dialect of the operation. -// * The mnemonic for the operation, or the name without the dialect prefix. -// * A list of traits for the operation. -class NPU_Op traits = []> - : Op; - -//===----------------------------------------------------------------------===// -// NPU Operations -//===----------------------------------------------------------------------===// - -//===----------------------------------------------------------------------===// -// CreateTPipOp -//===----------------------------------------------------------------------===// -def CreateTPipOp : NPU_Op<"create_tpip"> { - let summary = "create npu tpip"; - let description = - [{Create npu tpip.}]; - let results = (outs NPU_TPipType:$result); - let assemblyFormat = [{ - attr-dict `:` type($result) - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} -def CreateTQueueOp : NPU_Op<"create_tqueue"> { - let summary = "create npu tqueue"; - let description = - [{Create npu tqueue.}]; - let results = (outs NPU_TQueueType : $result); - let assemblyFormat = [{ - attr-dict `:` type($result) - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} - -def CreateGlobalTensorOp : NPU_Op<"create_global_tensor"> { - let summary = "create npu global tensor"; - let description = - [{Create npu global tensor.}]; - let results = (outs NPU_GlobalTensorType : $result); - let assemblyFormat = [{ - attr-dict `:` type($result) - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} - -def SetGlobalBufferOp : NPU_Op<"set_global_buffer"> { - let summary = "set npu global buffer"; - let description = - [{set npu global buffer.}]; - let arguments = (ins Arg:$globalTensor - ); - let results = (outs NPU_GlobalTensorType : $result); - let assemblyFormat = [{ - attr-dict - `ins` `(` $globalTensor `:` type($globalTensor)`)` - `:` type($result) - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} - -def InitlBufferOp : NPU_Op<"init_buffer"> { - let summary = "init npu buffer"; - let description = - [{init npu buffer.}]; - let arguments = (ins Arg:$tpip, - Arg:$queue - ); - let results = (outs NPU_TQueueType : $result); - let assemblyFormat = [{ - attr-dict - `ins` `(` $tpip `:` type($tpip) `,` $queue `:` type($queue)`)` - `:` type($result) - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} - -def AllocLocalOp : NPU_Op<"alloc_local"> { - let summary = "alloc npu local tensor"; - let description = - [{alloc npu local tensor.}]; - let arguments = (ins Arg:$queue - ); - let results = (outs AnyRankedOrUnrankedMemRef : $result); - let assemblyFormat = [{ - attr-dict - `ins` `(` $queue `:` type($queue)`)` - `:` type($result) - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} - -def FreeLocalOp : NPU_Op<"free_local"> { - let summary = "free npu local tensor"; - let description = - [{free npu local tensor.}]; - let arguments = (ins Arg:$queue, - AnyRankedOrUnrankedMemRef : $local - ); - let results = (outs NPU_TQueueType : $result); - let assemblyFormat = [{ - attr-dict - `ins` `(` $queue `:` type($queue) `,` $local `:` type($local) `)` - `:` type($result) - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} - -def EnQueOp : NPU_Op<"en_queue"> { - let summary = "npu local tensor en queue"; - let description = - [{npu local tensor en queue.}]; - let arguments = (ins Arg:$queue, - AnyRankedOrUnrankedMemRef : $local - ); - let results = (outs NPU_TQueueType : $result); - let assemblyFormat = [{ - attr-dict - `ins` `(` $queue `:` type($queue) `,` $local `:` type($local)`)` - `:` type($result) - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} - -def DeQueOp : NPU_Op<"de_queue"> { - let summary = "npu local tensor de queue"; - let description = - [{npu local tensor de queue.}]; - let arguments = (ins Arg:$queue - ); - let results = (outs AnyRankedOrUnrankedMemRef : $result); - let assemblyFormat = [{ - attr-dict - `ins` `(` $queue `:` type($queue)`)` - `:` type($result) - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} - -def SliceOp : NPU_Op<"slice"> { - let description = [{npu slice.}]; - let arguments = (ins Arg:$source - ); - let results = (outs AnyRankedOrUnrankedMemRef : $result); - let assemblyFormat = [{ - attr-dict - `ins` `(` $source `:` type($source)`)` - `:` type($result) - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} -//===----------------------------------------------------------------------===// -// AddFOp -//===----------------------------------------------------------------------===// - -def AddFOp : NPU_Op<"addf"> { - let summary = "element-wise addition operation"; - let description = - [{The "add" operation performs element - - wise addition between two - tensors.The shapes of the tensor operands are expected to match.}]; - - let arguments = (ins Arg : $lhs, - Arg : $rhs, - AnyRankedOrUnrankedMemRef: $out); - // let results = (outs AnyRankedOrUnrankedMemRef); - - // Indicate that the operation has a custom parser and printer method. - // let hasCustomAssemblyFormat = 1; - - // Allow building an AddOp with from the two input operands. - // let builders = [OpBuilder<(ins "Value" : $lhs, "Value" : $rhs)>]; - let assemblyFormat = [{ - attr-dict - `ins` `(` $lhs `:` type($lhs) `,` $rhs `:` type($rhs) `)` - `outs` `(` $out `:` type($out) `)` - }]; - let hasCanonicalizer = 0; - let hasFolder = 0; -} - -//===----------------------------------------------------------------------===// -// CopyOp -//===----------------------------------------------------------------------===// - -def CopyOp : NPU_Op<"copy"> { - - let description = - [{Copies the data from the source to the destination. - Source and destination are expected to have the same element type and - shape.Otherwise, - the result is undefined.They may have different layouts.}]; - - let arguments = (ins Arg:$source, - Arg:$target - ); - - let assemblyFormat = [{ - attr-dict - `ins` `(` $source `:` type($source)`)` - `outs` `(` $target `:` type($target) `)` - }]; - // let builders = [ - // Build a SubViewOp with static entries and inferred result type. - // OpBuilder<(ins "MemRefType":$resultType, "Value":$source, - // "ArrayRef":$offsets, "ArrayRef":$sizes, - // "ArrayRef":$strides)> - // ]; - - let hasCanonicalizer = 0; - let hasFolder = 0; -} - -#endif // NPU_DIALECT_OPS \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/NPU/IR/NPUTypes.h b/compiler/include/dicp/Dialect/NPU/IR/NPUTypes.h deleted file mode 100644 index fea0b9ac..00000000 --- a/compiler/include/dicp/Dialect/NPU/IR/NPUTypes.h +++ /dev/null @@ -1,22 +0,0 @@ -#pragma once - -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/DialectImplementation.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/TypeSupport.h" -#include "mlir/IR/TypeUtilities.h" -#include "mlir/IR/Types.h" -#include "llvm/ADT/TypeSwitch.h" - -#define GET_TYPEDEF_CLASSES -#include "dicp/Dialect/NPU/IR/NPUTypes.h.inc" - -// namespace mlir::dicp::npu { - -// struct TQueueType : public Type::TypeBase { -// using Base::Base; - -// static constexpr StringLiteral name = "npu.tqueue"; -// }; -// } diff --git a/compiler/include/dicp/Dialect/NPU/IR/NPUTypes.td b/compiler/include/dicp/Dialect/NPU/IR/NPUTypes.td deleted file mode 100644 index b845d44c..00000000 --- a/compiler/include/dicp/Dialect/NPU/IR/NPUTypes.td +++ /dev/null @@ -1,52 +0,0 @@ -#ifndef NPU_DIALECT_TYPES -#define NPU_DIALECT_TYPES - -include "mlir/IR/AttrTypeBase.td" -include "mlir/IR/BuiltinTypeInterfaces.td" -include "dicp/Dialect/NPU/IR/NPUDialect.td" - -class NPUTypeDef traits = []> - : TypeDef { - // Used by printer/parser - let mnemonic = _mnemonic; -} - - -def NPU_TPipType : NPUTypeDef<"TPip", "tpip", []> { - let summary = "queue type (`::mlir::dicp::npu::TPipType`) in NPU IR type system"; - let description = [{tpip}]; - let hasCustomAssemblyFormat = 1; -} - -def NPU_GlobalTensorType : NPUTypeDef<"GlobalTensor", "globaltensor", []> { - let summary = "GlobalTensor type (`::mlir::dicp::npu::GlobalTensorType`) in NPU IR type system"; - let description = [{GlobalTensor}]; - let parameters = (ins - "uint32_t":$type - ); - let builders = [ - TypeBuilder<(ins "unsigned":$type), [{ - return $_get($_ctxt, type); - }]> - ]; - let hasCustomAssemblyFormat = 1; - let skipDefaultBuilders = 1; -} - -def NPU_TQueueType : NPUTypeDef<"TQueue", "tqueue", []> { - let summary = "queue type (`::mlir::dicp::npu::TQueueType`) in NPU IR type system"; - let description = [{queue}]; - let parameters = (ins - "uint32_t":$position, - "uint32_t":$bufferNumber - ); - let builders = [ - TypeBuilder<(ins "unsigned":$position, "unsigned":$bufferNumber), [{ - return $_get($_ctxt, position, bufferNumber); - }]> - ]; - let hasCustomAssemblyFormat = 1; - let skipDefaultBuilders = 1; -} - -#endif diff --git a/compiler/include/dicp/Dialect/NPU/CMakeLists.txt b/compiler/include/dicp/Dialect/TritonDicp/CMakeLists.txt similarity index 100% rename from compiler/include/dicp/Dialect/NPU/CMakeLists.txt rename to compiler/include/dicp/Dialect/TritonDicp/CMakeLists.txt diff --git a/compiler/include/dicp/Dialect/TritonDicp/IR/CMakeLists.txt b/compiler/include/dicp/Dialect/TritonDicp/IR/CMakeLists.txt new file mode 100644 index 00000000..9c8ad6d6 --- /dev/null +++ b/compiler/include/dicp/Dialect/TritonDicp/IR/CMakeLists.txt @@ -0,0 +1,15 @@ +set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR}) + +set(LLVM_TARGET_DEFINITIONS TritonDicpOps.td) +mlir_tablegen(TritonDicpDialect.h.inc -gen-dialect-decls -dialect=dicp) +mlir_tablegen(TritonDicpDialect.cpp.inc -gen-dialect-defs -dialect=dicp) +mlir_tablegen(TritonDicpOps.h.inc -gen-op-decls) +mlir_tablegen(TritonDicpOps.cpp.inc -gen-op-defs) +add_mlir_doc(TritonDicpDialect TritonDicpDialect dialects/ -gen-dialect-doc) +add_mlir_doc(TritonDicpOps TritonDicpOps dialects/ -gen-op-doc) +add_public_tablegen_target(TritonDicpTableGen) + +set(LLVM_TARGET_DEFINITIONS TritonDicpAttrDefs.td) +mlir_tablegen(TritonDicpOpsAttrDefs.h.inc -gen-attrdef-decls) +mlir_tablegen(TritonDicpOpsAttrDefs.cpp.inc -gen-attrdef-defs) +add_public_tablegen_target(TritonDicpAttrDefsIncGen) diff --git a/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpAttrDefs.td b/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpAttrDefs.td new file mode 100644 index 00000000..a0a71ef8 --- /dev/null +++ b/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpAttrDefs.td @@ -0,0 +1,25 @@ +//===-- TritonDicpAttrDefs.td - dialect attributes def. ----*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef TRITON_DICP_ATTRDEFS +#define TRITON_DICP_ATTRDEFS + +include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.td" + +include "mlir/IR/AttrTypeBase.td" +include "mlir/IR/EnumAttr.td" + +class TritonDicp_Attr traits = []> + : AttrDef { + let mnemonic = attrMnemonic; + let cppNamespace = "::mlir::triton::dicp"; +} + + + +#endif // TRITON_DICP_ATTRDEFS diff --git a/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h b/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h new file mode 100644 index 00000000..239e6326 --- /dev/null +++ b/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h @@ -0,0 +1,33 @@ +//===- TritonDicpDialect.h - MLIR TritonDicp dialect --------------*- C++ +//-*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This file defines the TritonDicp dialect in MLIR, containing DICP operations. +// +//===----------------------------------------------------------------------===// + +#ifndef TRITON_DIALECT_DICP_DIALECT_H +#define TRITON_DIALECT_DICP_DIALECT_H + +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/OpDefinition.h" + +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h.inc" + +#define GET_ATTRDEF_CLASSES +#include "dicp/Dialect/TritonDicp/IR/TritonDicpOpsAttrDefs.h.inc" + +#define GET_OP_CLASSES +#include "dicp/Dialect/TritonDicp/IR/TritonDicpOps.h.inc" + +namespace mlir::triton::dicp {} // namespace mlir::triton::dicp + +#endif // TRITON_DIALECT_DICP_DIALECT_H diff --git a/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpDialect.td b/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpDialect.td new file mode 100644 index 00000000..cff24a40 --- /dev/null +++ b/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpDialect.td @@ -0,0 +1,31 @@ +//===-- TritonDicpDialect.td - dialect op definitions -------*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#ifndef TRITON_DICP_DIALECT +#define TRITON_DICP_DIALECT + +include "mlir/IR/OpBase.td" + +def TritonDicp_Dialect : Dialect { + let name = "dicp"; + let cppNamespace = "::mlir::triton::dicp"; + let summary = "The TritonDicp dialect in Triton."; + + let description = [{ + TritonDicp is a dialect for representing operations on DICP NPUs. + }]; + + let dependentDialects = [ + "mlir::LLVM::LLVMDialect", + "triton::TritonDialect", + ]; + + let extraClassDeclaration = [{}]; +} + +#endif // TRITON_DICP_DIALECT diff --git a/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpOps.td b/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpOps.td new file mode 100644 index 00000000..f069849c --- /dev/null +++ b/compiler/include/dicp/Dialect/TritonDicp/IR/TritonDicpOps.td @@ -0,0 +1,473 @@ +//===-- TritonDicpOps.td - TritonDicp op definitions ---------*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This is the TritonDicp IR operation definition file. +// +//===----------------------------------------------------------------------===// + +#ifndef TRITON_DICP_OPS +#define TRITON_DICP_OPS + +include "mlir/IR/OpBase.td" +include "mlir/IR/EnumAttr.td" +include "mlir/Dialect/LLVMIR/LLVMTypes.td" +include "mlir/Interfaces/SideEffectInterfaces.td" +include "mlir/IR/OpAsmInterface.td" +include "mlir/Interfaces/InferTypeOpInterface.td" // SameOperandsAndResultType + +include "triton/Dialect/Triton/IR/TritonAttrDefs.td" +include "triton/Dialect/Triton/IR/TritonTypes.td" +include "triton/Dialect/Triton/IR/TritonInterfaces.td" + +include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.td" +include "dicp/Dialect/TritonDicp/IR/TritonDicpAttrDefs.td" + +//===----------------------------------------------------------------------===// +// TritonDicp op definitions +//===----------------------------------------------------------------------===// + +class TT_Dicp_Op traits = []> : + Op; + + +// +// Interfaces +// +def GlobalMemory : Resource<"::mlir::triton::GlobalMemory">; + + +// +// Annotation Op +// +def AnnotationOp : TT_Dicp_Op<"annotation", [Pure, MemoryEffects<[MemWrite]>]> { + let summary = "Annotate a tensor with key-value attribute pairs"; + let description = [{ + `dicp.annotation` operation can be used to annotate a tensor with + key-value attribute pairs. + + Example: + ```mlir + dicp.annotation %target {key : val} + ``` + }]; + let arguments = (ins TT_Tensor:$src); + let assemblyFormat = [{ + $src attr-dict `:` type($src) + }]; +} + + +// +// Mod Op +// +def ModOp : TT_Dicp_Op<"mod", [Pure]> { + let summary = "Mod operation (%) of input tensors."; + let description = [{ + Performs element-wise division with remainder of input tensors. + }]; + + let arguments = (ins TT_Tensor:$lhs, TT_Tensor:$rhs); + let results = (outs TT_Tensor:$result); + + let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) type($rhs) `->` type($result)"; +} + + +// +// IndexPut Op +// +def IndexPutOp : TT_Dicp_Op<"index_put", [ + MemoryEffects<[MemWrite]>, + SameVariadicOperandSize, +]> { + let summary = "Scatter store to a tensor pointer with embedding semantics"; + + let description = [{ + Index put values from a tensor into a destination tensor. + + The operation takes: + - ptr: pointer type, the destination tensor pointer (in GM) + - index: tensor, a index to scatter (in UB) + - value: tensor, a value to store (in UB) + - dim: int32, the dimension to scatter along + - index_boundary: int64, the upper boundary for index values + - end_offset: tuple of int, the offsets of each dimension for the end of the scatter region + - start_offset: tuple of int, the offsets of each dimension for the start of the scatter region + - dst_stride: tuple of int, the stride of each dimension of destination tensor + + + Constraints: + - `ptr` and `value` must have the same rank. + - `ptr.dtype` only supports `float16`, `bfloat16`, `float32` currently. + - `index` must be an integer tensor. If `index.rank` != 1, it will be reshaped to 1D. + - `index.numel` must equal `value.shape[dim]`. + - `value` support 2~5D tensors. + - `dim` must be valid (0 <= dim < rank(value) - 1). + }]; + + let arguments = ( + ins TT_Ptr:$ptr, + TT_Tensor:$index, + TT_Tensor:$value, + TT_Int:$dim, + TT_Int:$indexBoundary, + Variadic>:$endOffset, + Variadic>:$startOffset, + Variadic>:$dstStride + ); + + let assemblyFormat = [{ + $ptr `:` type($ptr) `,` $index `:` type($index) `,` + $value `:` type($value) `,` $dim `:` type($dim) `,` $indexBoundary `:` type($indexBoundary) `,` + `[` $endOffset `:` type($endOffset) `]` `,` `[` $startOffset `:` type($startOffset) `]` `,` + `[` $dstStride `:` type($dstStride) `]` + attr-dict + }]; +} + +// +// GatherOutToUb Op +// +def GatherOutToUbOp : TT_Dicp_Op<"gather_out_to_ub", [ + DeclareOpInterfaceMethods, + AttrSizedOperandSegments, +]> { + let summary = "Gather load from a tensor pointer with the embedding semantics"; + + let description = [{ + Gather from a source tensor in Global Memory (GM) to Unified Buffer (UB) + along a specified dimension with out-of-bound handling. + + The operation takes: + - src: pointer type, the source tensor pointer (in GM) + - index: tensor, a tensor to gather (in UB) + - index_boundary: int64, the upper boundary for index values + - dim: int32, the dimension to gather along + - src_stride: tuple of int64, the stride of each dimension of src tensor + - end_offset: tuple of int32, the end offsets of each dimension for index tensor + - start_offset: tuple of int32, the start offsets of each dimension for index tensor + - other(Optional): scalar value, the default value when index is out of boundary (in UB) + + Returns: + a tensor, with the same shape as `index.shape` (in UB) + + Constraints: + - `src` and `index` must have the same rank. + - `src.dtype` only supports `float16`, `bfloat16`, `float32` currently. + - `index` must be an integer tensor, with rank between 1 and 5. + - `dim` must be valid (0 <= dim < rank(index)). + - `other` must be a scalar value. + - For every dimension `i` not equal to `dim`, `index.size[i]` <= `src.size[i]`. + - The output shape is the same as `index.shape`. If `index` is None, \ + the output tensor will be an empty tensor with the same shape as `index`. + }]; + + let arguments = ( + ins TT_Ptr:$src, + TT_Tensor:$index, + TT_Int:$indexBoundary, + TT_Int:$dim, + Variadic>:$srcStride, + Variadic>:$endOffset, + Variadic>:$startOffset, + Optional:$other + ); + + let results = (outs TT_Tensor:$result); + + let assemblyFormat = [{ + $src `:` type($src) `,` $index `:` type($index) `,` + $indexBoundary `:` type($indexBoundary) `,` $dim `:` type($dim) `,` + `[` $srcStride `:` type($srcStride) `]` `,` `[` $endOffset `:` type($endOffset) `]` `,` + `[` $startOffset `:` type($startOffset) `]` (`,` $other^ `:` type($other))? + attr-dict `->` type($result) + }]; +} + +// +// ScatterUbToOut Op +// +def ScatterUbToOutOp : TT_Dicp_Op<"scatter_ub_to_out", [ + MemoryEffects<[MemWrite]>, + SameVariadicOperandSize, +]> { + let summary = "scatter store from a tensor pointer with the embedding semantics"; + + let description = [{ + Scatter a tile from Unified Buffer (UB) into a destination tensor in Global Memory (GM) + along a specified dimension, with index-boundary checking. + + The operation takes: + - ptr: pointer type, the destination tensor pointer (in GM) + - value: tensor, a tile value to store (in UB) + - index: tensor, a tile index to scatter (in UB) + - index_boundary: int, the upper boundary for index values + - dim: int, the dimension to scatter along + - dst_stride: tuple of int, the stride of each dimension of destination tensor + - end_offset: tuple of int32, the end offsets of each dimension for index tensor + - start_offset: tuple of int32, the start offsets of each dimension for index tensor + + Constraints: + - `ptr` and `index` must have the same rank. + - `ptr.dtype` only supports `float16`, `bfloat16`, `float32` currently. + - `index` must be an integer tensor, with rank between 1 and 5. + - `dim` must be valid (0 <= dim < rank(index)). + - For every dimension `i` not equal to `dim`, `index.size[i]` <= `ptr.size[i]`. + - The output shape is the same as `index.shape`. If `index` is None, \ + the output tensor will be an empty tensor with the same shape as `index`. + }]; + + let arguments = ( + ins TT_Ptr:$ptr, + TT_Tensor:$value, + TT_Tensor:$index, + TT_Int:$indexBoundary, + TT_Int:$dim, + Variadic>:$dstStride, + Variadic>:$endOffset, + Variadic>:$startOffset + ); + + let assemblyFormat = [{ + $ptr `:` type($ptr) `,` $value `:` type($value) `,` `,` $index `:` type($index) `,` + $indexBoundary `:` type($indexBoundary) `,` $dim `:` type($dim) `,` + `[` $dstStride `:` type($dstStride) `]` `,` `[` $endOffset `:` type($endOffset) `]` `,` + `[` $startOffset `:` type($startOffset) `]` + attr-dict + }]; +} + + +// +// IndexSelectSimd Op +// +def IndexSelectSimdOp : TT_Dicp_Op<"index_select_simd", [ + MemoryEffects<[MemRead]>, + DeclareOpInterfaceMethods, + AttrSizedOperandSegments +]> { + let summary = "Index select SIMD operation from global memory"; + + let description = [{ + Index select operation (SIMD version) that loads data from multiple indices along a + specified dimension. The operation selects data from GM and loads them + as tiles directly to UB with zero-copy semantics. + + The operation takes: + - src: Source pointer (in GM) + - index: 1D tensor of indices to select (already in UB) + - dim: The dimension along which to select + - src_shape: Complete shape of the source tensor + - src_offset: Starting offset for reading + - read_shape: Size to read (tile shape) + + Constraints: + - read_shape[dim] must be -1 + - src_offset[dim] can be -1 (will be ignored) + }]; + + let arguments = ( + ins + TT_PtrLike:$src, + TT_IntTensor:$index, + I32Attr:$dim, + Variadic:$src_shape, + Variadic:$src_offset, + DenseI32ArrayAttr:$read_shape + ); + + let results = (outs TT_Tensor:$result); + + let assemblyFormat = [{ + $src `,` $index `,` $dim `,` + `[` $src_shape `]` `,` + `[` $src_offset `]` `,` + $read_shape + attr-dict `:` type($src) `,` type($index) `->` type($result) + }]; +} + + +// +// Built-in: IndirectLoad Op +// +def IndirectLoadOp : TT_Dicp_Op<"indirect_load", [ + DeclareOpInterfaceMethods, + AttrSizedOperandSegments +]> { + let summary = "Built-in: indirect load from global memory using per-element offsets with optional mask/other"; + + let description = [{ + Built-in operation emitted by the compiler for unstructured (discrete) memory + accesses.These are not written directly in the user IR. + + Load values from global memory based on per-element offsets. If `mask` + is provided, false lanes return `other`. + + The operation takes: + - src: Source pointer + - offsets: Tensor of per-element offsets (relative to `src`) for accessing source memory + - mask (optional): if mask[idx] is false, do not load the data at address pointer[idx] + - other (optional): if mask[idx] is false, return other[idx] + }]; + + let arguments = ( + ins TT_Ptr:$src, + TT_IntTensor:$offsets, + Optional:$mask, + Optional:$other + ); + + let results = (outs TT_Tensor:$result); + + let assemblyFormat = [{ + $src `:` type($src) `,` + $offsets `:` type($offsets) + (`,` $mask^ `:` type($mask))? + (`,` $other^ `:` type($other))? + attr-dict `->` type($result) + }]; + +} + + +// +// Built-in: IndirectStore Op +// +def IndirectStoreOp : TT_Dicp_Op<"indirect_store", [ + MemoryEffects<[MemWrite]> +]> { + let summary = "Built-in: indirect store from UB using per-element offsets with optional mask/other"; + + let description = [{ + Built-in operation emitted by the compiler for unstructured (discrete) memory + accesses.These are not written directly in the user IR. + + Store values from UB based to GM on per-element offsets. + + The operation takes: + - src: Source pointer + - offsets: Tensor of per-element offsets (relative to `src`) for accessing source memory + - value: The tensor of elements to be stored + - mask (optional): If mask[idx] is false, do not store value[idx] at pointer[idx] + }]; + + let arguments = ( + ins TT_Ptr:$src, + TT_IntTensor:$offsets, + TT_Type:$value, + Optional:$mask + ); + + let assemblyFormat = [{ + $src `:` type($src) `,` + $offsets `:` type($offsets) `,` + $value `:` type($value) + (`,` $mask^ `:` type($mask))? + attr-dict + }]; + +} + +// +// Custom Op +// +def CustomOp : TT_Dicp_Op<"custom", [Pure, MemoryEffects<[MemWrite]>]> { + let summary = "self-defined custom operation"; + let description = [{ + `dicp.custom` triton custom op is designed to pass self-defined custom operation. + + Example: + ```dicp.custom {str_args = ["sync_block_wait", "cube"]} + ``` + }]; + let arguments = (ins StrAttr:$op_name, ArrayAttr:$str_args, Variadic:$args); + + let assemblyFormat = "$op_name attr-dict ($args^ `:` type($args))?"; +} + +def FlipOp : TT_Dicp_Op<"flip", [ + NoMemoryEffect, + DeclareOpInterfaceMethods +]> { + let summary = "Reverse a tensor along a given dimension"; + let description = [{ + Reverses the elements of the input tensor along the specified dimension. + The output tensor has the same shape and element type as the input. + }]; + + let arguments = (ins + TT_Tensor:$src, // Input tensor + I64Attr:$dim // Dimension to flip along + ); + + let results = (outs + TT_Tensor:$flipped // Flipped values + ); + + let assemblyFormat = + "$src `,` $dim attr-dict `:` type($src) `->` type($flipped)"; +} + +def SortOp : TT_Dicp_Op<"sort", [ + NoMemoryEffect, + DeclareOpInterfaceMethods +]> { + let summary = "Sorts a tensor along a given dimension and returns sorted values."; + let description = [{ + Sorts the elements of the input tensor along the specified dimension. + Returns one tensor: + The sorted tensor (same shape and element type as input). + }]; + + let arguments = (ins + TT_Tensor:$src, // Input tensor + I64Attr:$dim, // Dimension to sort along + BoolAttr:$descending // Sort order + ); + + let results = (outs + TT_Tensor:$sorted // Sorted values + ); + + let assemblyFormat = "$src `,` $dim `,` $descending attr-dict `:` type($src) `->` type($sorted)"; +} + +// +// Conv Op +// +def Conv1dOp : TT_Dicp_Op<"conv1d", [ + Pure, + DeclareOpInterfaceMethods +]> { + let summary = "1D convolution operation"; + let description = [{ + Performs a 1D convolution operation on the input tensor. + }]; + + let arguments = (ins + TT_Tensor:$input, + TT_Tensor:$weight, + Optional:$bias, + I64Attr:$stride, + I64Attr:$padding_size, + I64Attr:$dilation, + I64Attr:$groups + ); + + let results = (outs TT_Tensor:$result); + + let assemblyFormat = [{ + `(` operands `)` + attr-dict `:` functional-type(operands, $result) + }]; + + let hasVerifier = 1; +} + +#endif // TRITON_DICP_OPS diff --git a/compiler/include/dicp/Dialect/TritonExt/CMakeLists.txt b/compiler/include/dicp/Dialect/TritonExt/CMakeLists.txt deleted file mode 100644 index 5cd7900f..00000000 --- a/compiler/include/dicp/Dialect/TritonExt/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -dicp_add_all_subdirs() \ No newline at end of file diff --git a/compiler/include/dicp/Dialect/TritonExt/Transforms/CMakeLists.txt b/compiler/include/dicp/Dialect/TritonExt/Transforms/CMakeLists.txt deleted file mode 100644 index 86071529..00000000 --- a/compiler/include/dicp/Dialect/TritonExt/Transforms/CMakeLists.txt +++ /dev/null @@ -1,3 +0,0 @@ -set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls -name TritonExt) -add_public_tablegen_target(TritonExtTransformsIncGen) diff --git a/compiler/include/dicp/Dialect/TritonExt/Transforms/Passes.h b/compiler/include/dicp/Dialect/TritonExt/Transforms/Passes.h deleted file mode 100644 index 3203314b..00000000 --- a/compiler/include/dicp/Dialect/TritonExt/Transforms/Passes.h +++ /dev/null @@ -1,22 +0,0 @@ -#ifndef TRITONEXT_DIALECT_TRITON_TRANSFORMS_PASSES_H_ -#define TRITONEXT_DIALECT_TRITON_TRANSFORMS_PASSES_H_ - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Pass/Pass.h" - -namespace mlir::dicp::trtion_ext { -const unsigned INT_BIT_WIDTH = 32; -const unsigned SET_INIT_SIZE = 16; - -enum TensorKind { NONE = -1, INPUT = 0, OUTPUT = 1, INPUT_OUTPUT = 2 }; - -std::unique_ptr> createCanonicalizeTritonIRAscendPass(); - -std::unique_ptr> createCanonicalizeCmpiPass(); - -#define GEN_PASS_REGISTRATION -#include "dicp/Dialect/TritonExt/Transforms/Passes.h.inc" - -} // namespace mlir::dicp::trtion_ext - -#endif diff --git a/compiler/include/dicp/Dialect/TritonExt/Transforms/Passes.td b/compiler/include/dicp/Dialect/TritonExt/Transforms/Passes.td deleted file mode 100644 index 88a6f2ab..00000000 --- a/compiler/include/dicp/Dialect/TritonExt/Transforms/Passes.td +++ /dev/null @@ -1,43 +0,0 @@ -#ifndef TRITONEXT_PASSES -#define TRITONEXT_PASSES - -include "mlir/Pass/PassBase.td" - -def CanonicalizeTritonIRAscend : Pass<"canonicalize-triton-ir-ascend", "mlir::ModuleOp"> { - let summary = "Canonicalize Triton IR for NPU backend"; - - let description = [{ - Perform canonicalization and Triton-specific pattern rewrites to prepare - Triton IR for lowering to the NPU backend. This pass performs actions such - as: lowering atomic RMW/CAS to linalg.generic forms, handling masked - operations, moving/merging bitcasts, canonicalizing scalar/tensor store - patterns, rewrites arith.remf, and other device-specific cleanups required by the backend. - }]; - - let constructor = "mlir::dicp::trtion_ext::createCanonicalizeTritonIRAscendPass()"; - - let dependentDialects = [ - "mlir::triton::TritonDialect", - "mlir::arith::ArithDialect", - "mlir::scf::SCFDialect", - "mlir::memref::MemRefDialect", - "mlir::tensor::TensorDialect", - ]; -} - - -def CanonicalizeCmpi : Pass<"canonicalize-cmpi", "mlir::ModuleOp"> { - let summary = "Canonicalize selected arith.cmpi predicates to stricter forms (with RHS+1)."; - let description = [{ - Canonicalize comparisons to equivalent forms using RHS+1: - - signed less-or-equal (sle) -> signed less-than (slt) with RHS + 1 - - signed greater-than (sgt) -> signed greater-or-equal (sge) with RHS - 1 - - unsigned less-or-equal (ule)-> unsigned less-than (ult) with RHS + 1 - - The pass only transforms when the RHS is an integer scalar or a ranked tensor of integer elements. - }]; - let constructor = "mlir::dicp::trtion_ext::createCanonicalizeCmpiPass()"; - let dependentDialects = ["mlir::triton::TritonDialect","mlir::arith::ArithDialect"]; -} - -#endif diff --git a/compiler/include/dicp/Dialect/TritonStructured/CMakeLists.txt b/compiler/include/dicp/Dialect/TritonStructured/CMakeLists.txt new file mode 100644 index 00000000..f33061b2 --- /dev/null +++ b/compiler/include/dicp/Dialect/TritonStructured/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/compiler/include/dicp/Dialect/TritonStructured/IR/CMakeLists.txt b/compiler/include/dicp/Dialect/TritonStructured/IR/CMakeLists.txt new file mode 100644 index 00000000..9c32c97c --- /dev/null +++ b/compiler/include/dicp/Dialect/TritonStructured/IR/CMakeLists.txt @@ -0,0 +1,8 @@ +set(LLVM_TARGET_DEFINITIONS TritonStructuredDialect.td) +mlir_tablegen(TritonStructuredDialect.h.inc -gen-dialect-decls -dialect=tts) +mlir_tablegen(TritonStructuredDialect.cpp.inc -gen-dialect-defs -dialect=tts) +mlir_tablegen(TritonStructuredOps.h.inc -gen-op-decls) +mlir_tablegen(TritonStructuredOps.cpp.inc -gen-op-defs) + + +add_public_tablegen_target(TritonStructuredTableGen) diff --git a/compiler/include/dicp/Dialect/TritonStructured/IR/TritonStructuredDialect.h b/compiler/include/dicp/Dialect/TritonStructured/IR/TritonStructuredDialect.h new file mode 100644 index 00000000..08ccfd54 --- /dev/null +++ b/compiler/include/dicp/Dialect/TritonStructured/IR/TritonStructuredDialect.h @@ -0,0 +1,30 @@ +#ifndef MLIR_DIALECT_TRITON_STRUCTURED_IR_TRITON_STRUCTURED_DIALECT_H_ +#define MLIR_DIALECT_TRITON_STRUCTURED_IR_TRITON_STRUCTURED_DIALECT_H_ + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Dialect.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/SymbolTable.h" +#include "mlir/IR/TypeSupport.h" +#include "mlir/IR/Types.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/IR/Dialect.h" + +using namespace mlir; +using namespace mlir::triton; +//===----------------------------------------------------------------------===// +// TritonStructured Operations +//===----------------------------------------------------------------------===// +#include "dicp/Dialect/TritonStructured/IR/TritonStructuredDialect.h.inc" + +// Include the auto-generated header file containing the declarations of the +// TritonStructured operations. +#define GET_OP_CLASSES + +#include "dicp/Dialect/TritonStructured/IR/TritonStructuredOps.h.inc" + +#endif diff --git a/compiler/include/dicp/Dialect/TritonStructured/IR/TritonStructuredDialect.td b/compiler/include/dicp/Dialect/TritonStructured/IR/TritonStructuredDialect.td new file mode 100644 index 00000000..7d205da9 --- /dev/null +++ b/compiler/include/dicp/Dialect/TritonStructured/IR/TritonStructuredDialect.td @@ -0,0 +1,66 @@ +#ifndef TRITON_STRUCTURED_DIALECT +#define TRITON_STRUCTURED_DIALECT + +include "mlir/IR/OpBase.td" +include "triton/Dialect/Triton/IR/TritonDialect.td" +include "triton/Dialect/Triton/IR/TritonTypes.td" +include "triton/Dialect/Triton/IR/TritonAttrDefs.td" +//include "triton/Dialect/Triton/IR/TritonTypes.td" +include "mlir/Interfaces/SideEffectInterfaces.td" + + + +def Triton_Structured_Dialect : Dialect { + let name = "tts"; + + let cppNamespace = "::mlir::tts"; + + let summary = "Structured Triton operations"; + + let description = [{ + Triton Structured Dialect. + }]; + + let dependentDialects = [ + "triton::TritonDialect" + ]; + + let usePropertiesForAttributes = 1; +} + +// +// Op Base +// +class TTS_Op traits = []> : + Op { +} + + +// SameVariadicResultSize +// AttrSizedResultSegments +def TTS_GetStructuredStateOp : TTS_Op<"get_structured_state", [AttrSizedResultSegments, Pure]> { + let summary = "Placeholder for the structured pointer states computed during PtrAnalysis."; + let description = "Used to pass the offsets and strides to scf.for op to simplify IR rewrites."; + + let arguments = (ins AnyTypeOf<[TT_PtrLike, I32Tensor, I64Tensor,I16Tensor,I8Tensor,I1Tensor]>:$input); + let results = (outs AnyTypeOf<[TT_PtrLike, I32Tensor, I64Tensor,I16Tensor,I8Tensor,I1Tensor]>:$structured, Variadic:$offsets, Variadic:$strides); + + let builders = [ + OpBuilder<(ins "Value":$input)>, + ]; + + let extraClassDeclaration = [{ + static std::optional, SmallVector>> + getOffsetAndStrideTypes(MLIRContext *context, Type ptrLikeType); + + static std::optional> + getOffsetAndStrideSegmentSizes(Type ptrLikeType); + }]; + + let hasFolder = 0; + let hasVerifier = 1; +} + + + +#endif // TRITON_STRUCTURED_DIALECT diff --git a/compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/CMakeLists.txt b/compiler/include/dicp/DiscreteMaskAccessConversion/CMakeLists.txt similarity index 100% rename from compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/CMakeLists.txt rename to compiler/include/dicp/DiscreteMaskAccessConversion/CMakeLists.txt diff --git a/compiler/include/dicp/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.h b/compiler/include/dicp/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.h new file mode 100644 index 00000000..4b68c181 --- /dev/null +++ b/compiler/include/dicp/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.h @@ -0,0 +1,43 @@ + + +#ifndef TRITON_ADAPTER_DISCRETEMASKACCESSCONVERSION_H +#define TRITON_ADAPTER_DISCRETEMASKACCESSCONVERSION_H + +#include "mlir/Pass/Pass.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/IR/PatternMatch.h" + +#define GEN_PASS_DECL_DISCRETEMASKACCESSCONVERSION +#include "dicp/DiscreteMaskAccessConversion/Passes.h.inc" + +#define GEN_PASS_DEF_DISCRETEMASKACCESSCONVERSION +#include "dicp/DiscreteMaskAccessConversion/Passes.h.inc" + +namespace mlir { +namespace triton { + +std::unique_ptr> createDiscreteMaskAccessConversionPass( + const DiscreteMaskAccessConversionOptions &options = {}); + +} // namespace triton +} // namespace mlir + +namespace { + +using namespace mlir; +using namespace triton; + +class DiscreteMaskAccessConversionPass + : public ::impl::DiscreteMaskAccessConversionBase< + DiscreteMaskAccessConversionPass> { +public: + explicit DiscreteMaskAccessConversionPass( + const DiscreteMaskAccessConversionOptions &options); + void getDependentDialects(DialectRegistry ®istry) const override; + void runOnOperation() override; +}; + +} // namespace + +#endif // DISCRETE_MASK_ACCESS_CONVERSION_H diff --git a/compiler/include/dicp/DiscreteMaskAccessConversion/Passes.h b/compiler/include/dicp/DiscreteMaskAccessConversion/Passes.h new file mode 100644 index 00000000..09342c81 --- /dev/null +++ b/compiler/include/dicp/DiscreteMaskAccessConversion/Passes.h @@ -0,0 +1,17 @@ + + +#ifndef TRITON_ADAPTER_DISCRETE_MASK_ACCESS_CONVERSION_PASSES_H +#define TRITON_ADAPTER_DISCRETE_MASK_ACCESS_CONVERSION_PASSES_H + +#include "dicp/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.h" + +namespace mlir { +namespace triton { + +#define GEN_PASS_REGISTRATION +#include "dicp/DiscreteMaskAccessConversion/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_DISCRETE_MASK_ACCESS_CONVERSION_PASSES_H diff --git a/compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/Passes.td b/compiler/include/dicp/DiscreteMaskAccessConversion/Passes.td similarity index 65% rename from compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/Passes.td rename to compiler/include/dicp/DiscreteMaskAccessConversion/Passes.td index 343127b8..3adeb8bb 100644 --- a/compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/Passes.td +++ b/compiler/include/dicp/DiscreteMaskAccessConversion/Passes.td @@ -1,5 +1,3 @@ - - #ifndef DISCRETE_MASK_ACCESS_CONVERSION_PASSES #define DISCRETE_MASK_ACCESS_CONVERSION_PASSES @@ -9,12 +7,15 @@ def DiscreteMaskAccessConversion : Pass<"discrete-mask-access-conversion", "mlir let summary = "Recognize and convert discrete mask memory access"; let constructor = "triton::createDiscreteMaskAccessConversionPass()"; let options = [ - Option<"compileOn91095", "compile-on-910-95", + Option<"compileOn91095", "compile-on-910-95", "bool", /*default*/"false", "compile on 910_95">, - Option<"forceSimtTemplate", "force-simt-template", + Option<"forceSimtTemplate", "force-simt-template", "bool", /*default*/"false", - "force to use simt template"> + "force to use simt template">, + Option<"enableSyncBlockLock", "enable-sync-block-lock", + "bool", /*default*/"true", + "enable sync block lock/unlock in store"> ]; } diff --git a/compiler/include/dicp/TritonAffinityOpt/CMakeLists.txt b/compiler/include/dicp/TritonAffinityOpt/CMakeLists.txt new file mode 100644 index 00000000..4a804f07 --- /dev/null +++ b/compiler/include/dicp/TritonAffinityOpt/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonAffinityOpt) +add_public_tablegen_target(TritonAffinityOptConversionPassIncGen) \ No newline at end of file diff --git a/compiler/include/dicp/TritonAffinityOpt/DAG.h b/compiler/include/dicp/TritonAffinityOpt/DAG.h new file mode 100644 index 00000000..ee760d2d --- /dev/null +++ b/compiler/include/dicp/TritonAffinityOpt/DAG.h @@ -0,0 +1,304 @@ +#ifndef AffinityDAGDEF +#define AffinityDAGDEF +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "dicp/TritonAffinityOpt/Utils.hpp" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TinyPtrVector.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include +#include +#include +#include +#include +#include + +namespace mlir { +namespace AffinityDAG { + +enum class OpAbility { + PREFER_VECTOR = 1 << 0, + CUBE_ONLY = 1 << 1, + CUBE_AND_VECTOR = PREFER_VECTOR | CUBE_ONLY + +}; + +enum CoreType { + UNDETERMINED = 0, + VECTOR_ONLY = 1 << 0, + CUBE_ONLY = 1 << 1, + CUBE_AND_VECTOR = VECTOR_ONLY | CUBE_ONLY +}; + +inline constexpr CoreType toCoreType(OpAbility ct) { + using U = std::underlying_type_t; + return static_cast(static_cast(ct)); +} + +constexpr inline CoreType operator|(CoreType lhs, CoreType rhs) { + return enumOp(std::bit_or<>(), lhs, rhs); +} + +inline CoreType operator&(CoreType lhs, CoreType rhs) { + return enumOp(std::bit_and<>(), lhs, rhs); +} + +inline bool intersects(CoreType lhs, CoreType rhs) { + return (lhs & rhs) != CoreType::UNDETERMINED; +} + +inline CoreType operator&(OpAbility lhs, CoreType rhs) { + return toCoreType(lhs) & rhs; +} + +inline CoreType operator!(CoreType ct) { + CoreType newCt = UNDETERMINED; + if ((ct & CoreType::CUBE_ONLY) == UNDETERMINED) { + newCt = newCt | CoreType::CUBE_ONLY; + } + + if ((ct & CoreType::VECTOR_ONLY) == UNDETERMINED) { + newCt = newCt | CoreType::VECTOR_ONLY; + } + + return newCt; +} + +inline hivm::TCoreType toHivm(CoreType ct) { + switch (ct) { + case UNDETERMINED: + return hivm::TCoreType::CUBE_OR_VECTOR; + case CUBE_ONLY: + return hivm::TCoreType::CUBE; + case VECTOR_ONLY: + return hivm::TCoreType::VECTOR; + case CUBE_AND_VECTOR: + return hivm::TCoreType::CUBE_AND_VECTOR; + default: + llvm_unreachable("Invalid CoreType that cannot convert to hivm"); + } +} + +inline bool intersects(OpAbility lhs, CoreType rhs) { + return (lhs & rhs) != CoreType::UNDETERMINED; +} + +inline bool exactlyOneType(CoreType ct) { + return (ct == CUBE_ONLY) || (ct == VECTOR_ONLY); +} + +const char *literalCoreType(CoreType ct); + +class MoveOnly { +protected: + MoveOnly() = default; + ~MoveOnly() = default; + + MoveOnly(const MoveOnly &) = delete; + MoveOnly &operator=(const MoveOnly &) = delete; + + MoveOnly(MoveOnly &&) = default; + MoveOnly &operator=(MoveOnly &&) = default; +}; + +class Node; +class OpNode; +class ValueNode; + +ValueNode *getDataSource(OpNode *op); + +class Graph : MoveOnly { +public: + using OpMapRaw = llvm::DenseMap>; + using ValueMapRaw = llvm::DenseMap>; + using OpMap = std::shared_ptr; + using ValueMap = std::shared_ptr; + + Graph(Block *block, Graph *parent = nullptr, OpMap opMap = nullptr, + ValueMap valueMap = nullptr, bool inheritParent = true); + + static std::unique_ptr fromMultiBlockFunc(triton::FuncOp funcOp); + + OpMapRaw &getOpMap() const { return *opMap; } + + ValueMapRaw &getValueMap() const { return *valueMap; } + + // [DEBUG] start + std::unique_ptr> legacyOpMap = nullptr; + std::unique_ptr> legacyValueTypes = nullptr; + + inline llvm::DenseMap &getOpMapLegacy() { + if (!legacyOpMap) { + legacyOpMap = + std::move(std::make_unique>()); + for (auto &[key, val] : *opMap) { + (*legacyOpMap)[key] = val.get(); + } + } + + return *legacyOpMap; + } + + llvm::DenseMap &getValueTypes(); + + // [DEBUG] end + +private: + friend class Node; + friend class OpNode; + OpMap opMap; + ValueMap valueMap; + Block *block; + Graph *parent; + OpNode *terminator = nullptr; + size_t opCount = 0; + llvm::SmallVector blockArgs; +}; + +class Node : MoveOnly { +protected: + friend class Graph; + friend class ValueNode; + bool isUpstreamOfCubeMem = false; + virtual CoreType absorbImpl() = 0; + llvm::SmallVector outputs; + +public: + CoreType isOnPrivate = UNDETERMINED; + + enum NodeKind { NK_Op, NK_Value }; + + inline CoreType isOn() const { return isOnPrivate; } + + bool absorb() { + auto newCoreType = absorbImpl(); + auto changed = newCoreType != isOnPrivate; + isOnPrivate = newCoreType; + + return changed; + }; + + virtual llvm::SmallVector getAffected() const = 0; + virtual OpNode *getSourceOpNode() = 0; + + ArrayRef getOutputs() const { return outputs; } + + CoreType absorbCommon(); + +private: + const NodeKind kind; + +public: + NodeKind getKind() const { return kind; } + +protected: + Node(NodeKind kind) : kind(kind) {} +}; + +class OpNode : public Node { + friend class Graph; + friend class ValueNode; + llvm::SmallVector inputs; + llvm::SmallVector subgraphs; + virtual CoreType absorbImpl() override; + +public: + Operation *op; + + OpNode(Operation *op, Graph *graph); + OpAbility canRunOn() const; + inline ArrayRef getInputs() const { return inputs; } + + static bool classof(const Node *node) { return node->getKind() == NK_Op; } + + virtual llvm::SmallVector getAffected() const override { + llvm::SmallVector result(inputs.begin(), inputs.end()); + result.append(outputs.begin(), outputs.end()); + + return result; + } + + virtual OpNode *getSourceOpNode() override { return this; } +}; + +class ValueNode : public Node { + friend class Graph; + friend class OpNode; + virtual CoreType absorbImpl() override; + +public: + Node *source = nullptr; + Value value; + // ValueNode(OpResult value); + // ValueNode(BlockArgument value); + + ValueNode(Value value) : Node(NK_Value), value(value){}; + virtual OpNode *getSourceOpNode() override { + if (!source) { + return nullptr; + } + + return source->getSourceOpNode(); + } + static bool classof(const Node *node) { return node->getKind() == NK_Value; } + + virtual llvm::SmallVector getAffected() const override { + llvm::SmallVector result(outputs.begin(), outputs.end()); + if (source) + result.push_back(source); + + return result; + } +}; + +class GraphManager { +private: + llvm::DenseMap> graphs; + +public: + static GraphManager &getInstance() { + static GraphManager instance; + return instance; + } + + void registerGraph(llvm::StringRef funcName, + std::shared_ptr graph) { + graphs[funcName] = graph; + } + + AffinityDAG::Graph *getGraph(llvm::StringRef funcName) { + auto it = graphs.find(funcName); + return it != graphs.end() ? it->second.get() : nullptr; + } + + void removeGraph(llvm::StringRef funcName) { graphs.erase(funcName); } +}; + +inline llvm::DenseMap &Graph::getValueTypes() { + static std::mutex mtx; + std::lock_guard lock(mtx); + if (!legacyValueTypes) { + legacyValueTypes = + std::move(std::make_unique>()); + for (auto &[key, val] : *valueMap) { + llvm::dbgs() << key << "\n"; + llvm::dbgs().flush(); + (*legacyValueTypes)[key] = val.get()->isOn(); + } + } + + return *legacyValueTypes; +} + +} // namespace AffinityDAG +} // namespace mlir +#endif diff --git a/compiler/include/dicp/TritonAffinityOpt/Passes.h b/compiler/include/dicp/TritonAffinityOpt/Passes.h new file mode 100644 index 00000000..565f5f8f --- /dev/null +++ b/compiler/include/dicp/TritonAffinityOpt/Passes.h @@ -0,0 +1,27 @@ + + +#ifndef TRITON_ADAPTER_TRITON_AFFINITY_OPTIMIZATION_PASSES_H +#define TRITON_ADAPTER_TRITON_AFFINITY_OPTIMIZATION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert Triton dialect to Annotation dialect. +std::unique_ptr> createDAGSSBufferPass(); + +std::unique_ptr> createDAGSyncPass(); + +std::unique_ptr> createDAGScopePass(); + +#define GEN_PASS_REGISTRATION +#include "dicp/TritonAffinityOpt/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_AFFINITY_OPTIMIZATION_PASSES_H \ No newline at end of file diff --git a/compiler/include/dicp/TritonAffinityOpt/Passes.td b/compiler/include/dicp/TritonAffinityOpt/Passes.td new file mode 100644 index 00000000..1b2bdbf3 --- /dev/null +++ b/compiler/include/dicp/TritonAffinityOpt/Passes.td @@ -0,0 +1,24 @@ +#ifndef TRITON_AFFINITY_OPTIMIZATION_PASSES +#define TRITON_AFFINITY_OPTIMIZATION_PASSES + +include "mlir/Pass/PassBase.td" + +def DAGSSBuffer : Pass<"dag-ssbuf", "mlir::ModuleOp"> { + let summary = "Convert vector operations to shared storage buffer operations"; + let constructor = "triton::createDAGSSBufferPass()"; + let dependentDialects = ["hivm::HIVMDialect", "bufferization::BufferizationDialect", "scope::ScopeDialect", "annotation::AnnotationDialect"]; +} + +def DAGScope : Pass<"dag-scope", "mlir::ModuleOp"> { + let summary = "Convert native triton code to NPU-affine code"; + let constructor = "triton::createDAGScopePass()"; + let dependentDialects = ["hivm::HIVMDialect", "bufferization::BufferizationDialect", "scope::ScopeDialect", "annotation::AnnotationDialect"]; +} + +def DAGSync : Pass<"dag-sync", "mlir::ModuleOp"> { + let summary = "DAG sync"; + let constructor = "triton::createDAGSyncPass()"; + let dependentDialects = ["hivm::HIVMDialect", "bufferization::BufferizationDialect", "annotation::AnnotationDialect"]; +} + +#endif // TRITON_AFFINITY_OPTIMIZATION_PASSES \ No newline at end of file diff --git a/compiler/include/dicp/TritonAffinityOpt/Utils.hpp b/compiler/include/dicp/TritonAffinityOpt/Utils.hpp new file mode 100644 index 00000000..69571082 --- /dev/null +++ b/compiler/include/dicp/TritonAffinityOpt/Utils.hpp @@ -0,0 +1,64 @@ +#ifndef TRITON_AFFINITY_UTILS_HPP +#define TRITON_AFFINITY_UTILS_HPP + +#include +#include +#include +#include + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" + +namespace mlir::AffinityDAG { + +template +constexpr inline T enumOp(F &&func, T lhs, T rhs) { + static_assert(std::is_enum_v, "T must be an enum type"); + + using U = std::underlying_type_t; + + return static_cast(std::invoke(std::forward(func), static_cast(lhs), + static_cast(rhs))); +} + +// since we do not have llvm::set_intersects in this version... +template bool intersects(S1Ty &s1, S2Ty &s2) { + if (s1.size() > s2.size()) { + return intersects(s2, s1); + } + + return llvm::any_of(s1, [&](auto e) { return s2.count(e); }); +} + +/** + * @returns the pointer to the value wrapped by the pointer if the key is in the + * map, otherwise nullptr + * + * @safety The user is responsible for checking the nullity of the retured + * pointer; the lifespan of the pointer is valid as long as the value belongs to + * the map + */ +template +inline auto getFromSmartPtr(MapTy &map, const typename MapTy::key_type &key) + -> decltype(map.find(key)->second.get()) { + auto it = map.find(key); + return it != map.end() ? it->second.get() : nullptr; +} + +/** + * @returns the pointer to the value if the key is in the map, otherwise nullptr + * + * @safety The user is responsible for checking the nullity of the retured + * pointer; the lifespan of the pointer is valid as long as the the map is not + * modified + */ +template +inline auto getPtr(MapTy &map, const typename MapTy::key_type &key) + -> decltype(map.find(key)->second) * { + auto it = map.find(key); + return it != map.end() ? &it->second : nullptr; +} + +} // namespace mlir::AffinityDAG + +#endif diff --git a/compiler/include/dicp/TritonToAnnotation/CMakeLists.txt b/compiler/include/dicp/TritonToAnnotation/CMakeLists.txt new file mode 100644 index 00000000..69d72d21 --- /dev/null +++ b/compiler/include/dicp/TritonToAnnotation/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToAnnotation) +add_public_tablegen_target(TritonToAnnotationConversionPassIncGen) \ No newline at end of file diff --git a/compiler/include/dicp/TritonToAnnotation/Passes.h b/compiler/include/dicp/TritonToAnnotation/Passes.h new file mode 100644 index 00000000..c9bb391a --- /dev/null +++ b/compiler/include/dicp/TritonToAnnotation/Passes.h @@ -0,0 +1,23 @@ + + +#ifndef TRITON_ADAPTER_TRITON_TO_ANNOTATION_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_ANNOTATION_CONVERSION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert Triton dialect to Annotation dialect. +std::unique_ptr> createTritonToAnnotationPass(); + +#define GEN_PASS_REGISTRATION +#include "dicp/TritonToAnnotation/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_TO_ANNOTATION_CONVERSION_PASSES_H diff --git a/compiler/include/dicp/TritonToAnnotation/Passes.td b/compiler/include/dicp/TritonToAnnotation/Passes.td new file mode 100644 index 00000000..ac7ba68d --- /dev/null +++ b/compiler/include/dicp/TritonToAnnotation/Passes.td @@ -0,0 +1,12 @@ +#ifndef TRITON_TO_ANNOTATION_CONVERSION_PASSES +#define TRITON_TO_ANNOTATION_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def TritonToAnnotation : Pass<"triton-to-annotation", "mlir::ModuleOp"> { + let summary = "Convert Triton to Annotation dialect"; + let constructor = "triton::createTritonToAnnotationPass()"; + let dependentDialects = ["annotation::AnnotationDialect"]; +} + +#endif // TRITON_TO_ANNOTATION_CONVERSION_PASSES diff --git a/compiler/include/dicp/TritonToGraph/AliasAnalysis.h b/compiler/include/dicp/TritonToGraph/AliasAnalysis.h new file mode 100644 index 00000000..11982b2c --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/AliasAnalysis.h @@ -0,0 +1,140 @@ + + +#ifndef TRITON_TO_CFG_ALIAS_ANALYSIS_H +#define TRITON_TO_CFG_ALIAS_ANALYSIS_H + +#include "dicp/TritonToGraph/tensor.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/TypeUtilities.h" +#include "mlir/IR/Types.h" +#include "mlir/IR/Value.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "llvm/ADT/DenseMap.h" + +namespace mlir { +namespace triton { +namespace cfg { + +class ControlFlowGraph; + +// AliasAnalysis - Alias分析和指针跟踪 +// 用于分析pointer的alias关系，跟踪全局内存对象（如gm_obj）的别名链 +class AliasAnalysis { +public: + // 为整个CFG分析pointer别名 + void analyzePointerAliases(ControlFlowGraph &cfg); + + // 获取Value的base pointer（递归查找真实的基础指针） + Value getBasePointer(Value ptr) const; + + // 获取Value对应的TensorObject + TensorObject *getTensorObject(Value value) const { + auto it = baseTensorMap.find(value); + return (it != baseTensorMap.end()) ? it->second : nullptr; + } + + // 判断两个指针是否可能指向同一tensor + bool mayAlias(Value ptr1, Value ptr2) const { + return getBasePointer(ptr1) == getBasePointer(ptr2); + } + + // 添加alias关系 + void addAlias(Value ptr, Value base, TensorObject *tensor) { + aliasMap[ptr] = base; + baseTensorMap[ptr] = tensor; + } + + // 判断是否是指针类型 + static bool isPointerType(Type type) { + if (auto ptrType = + mlir::dyn_cast(getElementTypeOrSelf(type))) { + return true; + } + return false; + } + + // 判断是否是tensor pointer类型（指向tensor的指针） + static bool isTensorPointerType(Type type) { + if (auto ptrType = mlir::dyn_cast(type)) { + Type pointeeType = ptrType.getPointeeType(); + return mlir::isa(pointeeType); + } + return false; + } + + // 判断是否是标量指针类型（指向标量的指针） + static bool isScalarPointerType(Type type) { + if (auto ptrType = mlir::dyn_cast(type)) { + Type pointeeType = ptrType.getPointeeType(); + return !mlir::isa(pointeeType); + } + return false; + } + + // 判断是否是全局内存类型 + static bool isGlobalMemoryType(Type type) { + // 在Triton中，全局内存通常由!tt.ptr类型表示 + // 可以进一步根据地址空间等信息判断 + if (auto ptrType = mlir::dyn_cast(type)) { + // 地址空间1通常是全局内存 + // 需要根据具体硬件进行调整 + return ptrType.getAddressSpace() == 1 || ptrType.getAddressSpace() == 0; + } + return false; + } + + // 获取所有tracked的base pointer + const DenseMap &getAliasMap() const { return aliasMap; } + + // 获取所有tracked的tensor对象 + const DenseMap &getBaseTensorMap() const { + return baseTensorMap; + } + + // 打印所有alias信息 + void print(llvm::raw_ostream &os) const { + os << "=== Alias Analysis Result ===\n"; + os << "Total tracked aliases: " << aliasMap.size() << "\n"; + + for (const auto &entry : aliasMap) { + Value ptr = entry.first; + Value base = entry.second; + TensorObject *tensor = baseTensorMap.lookup(ptr); + + os << " " << ptr << " -> " << base; + if (tensor) { + os << " [" << tensor->getName() << "]"; + } + os << "\n"; + } + } + +private: + // 分析addptr操作 + void analyzeAddPtrOp(mlir::triton::AddPtrOp addptrOp); + + // 分析make_tensor_ptr操作 + void analyzeMakeTensorPtrOp(mlir::triton::MakeTensorPtrOp op); + + // 分析load操作 + void analyzeLoadOp(mlir::triton::LoadOp loadOp); + + // 分析store操作 + void analyzeStoreOp(mlir::triton::StoreOp storeOp); + + // 分析broadcast操作 + void analyzeBroadcastOp(mlir::triton::BroadcastOp broadcastOp); + + // 分析splat操作 + void analyzeSplatOp(mlir::triton::SplatOp splatOp); + + DenseMap aliasMap; // ptr -> base ptr + DenseMap baseTensorMap; // value -> tensor +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_ALIAS_ANALYSIS_H diff --git a/compiler/include/dicp/TritonToGraph/CMakeLists.txt b/compiler/include/dicp/TritonToGraph/CMakeLists.txt new file mode 100644 index 00000000..5ffe7ae0 --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToGraph) +add_public_tablegen_target(TritonToGraphPassIncGen) diff --git a/compiler/include/dicp/TritonToGraph/ControlFlowGraph.h b/compiler/include/dicp/TritonToGraph/ControlFlowGraph.h new file mode 100644 index 00000000..5ea3a32c --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/ControlFlowGraph.h @@ -0,0 +1,246 @@ + + +#ifndef TRITON_TO_CFG_CONTROL_FLOW_GRAPH_H +#define TRITON_TO_CFG_CONTROL_FLOW_GRAPH_H + +#include "dicp/TritonToGraph/MemorySSA.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Region.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/MapVector.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include +#include +#include + +namespace mlir { +namespace triton { +namespace cfg { + +// 前向声明 +class BasicBlock; +class ControlFlowGraph; + +// Instruction 结构体：表示一个指令 +class Instruction { +public: + Instruction(size_t id, Operation *op, BasicBlock *parentBlock) + : id(id), operation(op), parentBlock(parentBlock), memorySSAInfo() {} + + // 获取基本信息 + size_t getId() const { return id; } + Operation *getOperation() const { return operation; } + BasicBlock *getParentBlock() const { return parentBlock; } + + // 检查是否有子图（内部区域） + bool hasSubGraph() const { return subGraph != nullptr; } + ControlFlowGraph *getSubGraph() const { return subGraph.get(); } + void setSubGraph(std::unique_ptr graph) { + subGraph = std::move(graph); + } + + // 获取指令的字符串表示 + std::string getAsString() const; + + // 打印 + void print(raw_ostream &os, unsigned indent = 0) const; + void dump() const; + + // Memory SSA信息 + MemorySSAInfo &getMemorySSAInfo() { return memorySSAInfo; } + const MemorySSAInfo &getMemorySSAInfo() const { return memorySSAInfo; } + +private: + size_t id; // 唯一ID + Operation *operation; // 对应的 MLIR Operation + BasicBlock *parentBlock; // 所属的 BasicBlock + std::unique_ptr + subGraph; // 子图（用于 reduce 等有内部区域的操作） + MemorySSAInfo memorySSAInfo; // Memory SSA信息（用于tensor/pointer分析） +}; + +// 基本块类型 +enum class BlockType { + NORMAL, // 普通块（包含多个指令） + ENTRY, // 函数入口块 + EXIT, // 函数出口块 + IF_COND, // if 条件判断块（包含单个 scf.if 指令） + FOR_COND, // for 循环头块（包含单个 scf.for 指令） + WHILE_COND, // while 条件块（包含单个 scf.while 指令） + COND_BR, // cf.cond_br 条件分支块（包含单个 cf.cond_br 指令） + BR, // cf.br 无条件跳转块（包含单个 cf.br 指令） + LOOP_BODY, // 循环体块 + LOOP_EXIT, // 循环出口块 +}; + +// 基本块节点 +class BasicBlock { +public: + BasicBlock(size_t id, BlockType type, BasicBlock *parentStructure = nullptr) + : id(id), type(type), parentStructure(parentStructure) {} + + // 获取基本信息 + size_t getId() const { return id; } + BlockType getType() const { return type; } + void setType(BlockType t) { type = t; } + + // 外层结构（如果是嵌套在 loop/if 中） + BasicBlock *getParentStructure() const { return parentStructure; } + void setParentStructure(BasicBlock *parent) { parentStructure = parent; } + + // 对应控制流结构的出口块（用于 IF_COND/FOR_COND/WHILE_COND） + // 指向 if/for/while 结束后到达的基本块 + BasicBlock *getExitBlock() const { return exitBlock; } + void setExitBlock(BasicBlock *exit) { exitBlock = exit; } + + // Instruction 操作 + void addInstruction(std::unique_ptr inst); + Instruction *getInstruction(size_t idx) const; + size_t getNumInstructions() const { return instructions.size(); } + const SmallVector> &getInstructions() const { + return instructions; + } + + // 检查最后一条指令是否为 ReturnOp + bool endsWithReturnOp() const; + + // 获取后继和前驱 + ArrayRef getSuccessors() const { return successors; } + ArrayRef getPredecessors() const { return predecessors; } + + // 边的操作 + void addSuccessor(BasicBlock *succ); + void addPredecessor(BasicBlock *pred); + + size_t getNumSuccessors() const { return successors.size(); } + size_t getNumPredecessors() const { return predecessors.size(); } + + // 获取名称 + std::string getName() const; + + // 获取类型字符串 + StringRef getTypeString() const; + + // 打印 + void print(raw_ostream &os) const; + void dump() const; + + // 导出为 JSON（用于网页可视化） + void exportToJSON(raw_ostream &os, unsigned indent = 0) const; + +private: + size_t id; // 唯一ID + BlockType type; // 块类型 + BasicBlock *parentStructure; // 外层结构（loop/if）的 basic block + BasicBlock *exitBlock = + nullptr; // 对应控制流结构的出口块（用于 IF_COND/FOR_COND/WHILE_COND） + SmallVector> instructions; // 指令列表 + SmallVector successors; // 后继块指针列表 + SmallVector predecessors; // 前驱块指针列表 +}; + +// 控制流图 +class ControlFlowGraph { +public: + explicit ControlFlowGraph(triton::FuncOp func); + ~ControlFlowGraph(); + + // 获取函数 + triton::FuncOp getFunction() const { return function; } + + // 基本块操作 + BasicBlock *createBasicBlock(BlockType type, + BasicBlock *parentStructure = nullptr); + + BasicBlock *getBasicBlock(size_t id) { + if (id < basicBlocks.size()) + return basicBlocks[id].get(); + return nullptr; + } + const BasicBlock *getBasicBlock(size_t id) const { + if (id < basicBlocks.size()) + return basicBlocks[id].get(); + return nullptr; + } + + size_t getNumBlocks() const { return basicBlocks.size(); } + + // 获取入口和出口块 + BasicBlock *getEntryBlock() { return entryBlock; } + const BasicBlock *getEntryBlock() const { return entryBlock; } + BasicBlock *getExitBlock() { return exitBlock; } + const BasicBlock *getExitBlock() const { return exitBlock; } + + void setEntryBlock(BasicBlock *bb) { entryBlock = bb; } + void setExitBlock(BasicBlock *bb) { exitBlock = bb; } + + // 添加边 + void addEdge(BasicBlock *from, BasicBlock *to); + + // 遍历 + using BlockVisitor = llvm::function_ref; + void traverse(BlockVisitor visitor); + + // Operation 到 Instruction 的查询 + Instruction *getInstruction(Operation *op) const { + auto it = opToInstructionMap.find(op); + return (it != opToInstructionMap.end()) ? it->second : nullptr; + } + + // 添加 Operation 到 Instruction 的映射 + void addOpToInstruction(Operation *op, Instruction *inst) { + opToInstructionMap[op] = inst; + } + + bool isBackEdge(BasicBlock *from, BasicBlock *to) const; + + // 结构化搜索 API + // 从起始块开始，沿着 successor 顺序向下结构化搜索 + // NORMAL 块：遍历每个 Instruction，调用 callback 传入 Operation* + // IF_COND/FOR_COND/WHILE_COND 块：递归访问每个 successor + // 遇到 exitBlock 时停止 + using OperationVisitor = llvm::function_ref; + void searchNormalBlock(BasicBlock *block, OperationVisitor callback) const; + void searchCondBlock(BasicBlock *block, OperationVisitor callback) const; + void searchBlock(BasicBlock *block, OperationVisitor callback) const; + + // 打印 + void print(raw_ostream &os) const; + void dump() const; + + // 导出为 DOT 格式 + void exportDOT(raw_ostream &os) const; + + // 导出到文件 + llvm::Error exportToFile(StringRef filename) const; + + // 导出为 HTML（网页可视化） + llvm::Error exportToHTML(StringRef filename) const; + + // 导出为 JSON + void exportToJSON(raw_ostream &os) const; + +private: + triton::FuncOp function; // 所属函数 + SmallVector> basicBlocks; // 基本块节点 + BasicBlock *entryBlock = nullptr; // 入口块 + BasicBlock *exitBlock = nullptr; // 出口块 + size_t nextBlockId = 0; // 下一个块ID + + // Operation 到 Instruction 的映射（支持快速查询） + std::unordered_map opToInstructionMap; +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_CONTROL_FLOW_GRAPH_H diff --git a/compiler/include/dicp/TritonToGraph/ControlFlowGraphBuilder.h b/compiler/include/dicp/TritonToGraph/ControlFlowGraphBuilder.h new file mode 100644 index 00000000..7e98cac0 --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/ControlFlowGraphBuilder.h @@ -0,0 +1,222 @@ + + +#ifndef TRITON_TO_CFG_CONTROL_FLOW_GRAPH_BUILDER_H +#define TRITON_TO_CFG_CONTROL_FLOW_GRAPH_BUILDER_H + +#include "dicp/TritonToGraph/ControlFlowGraph.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Value.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassRegistry.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include +#include + +namespace mlir { +namespace triton { +namespace cfg { + +// 构建控制流图的 Pass +class BuildCFGPass + : public PassWrapper> { +public: + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(BuildCFGPass) + + BuildCFGPass() = default; + + StringRef getArgument() const override { return "build-cfg"; } + StringRef getDescription() const override { + return "Build Control Flow Graph from TTIR"; + } + + void runOnOperation() override; + + // Pass 选项 + std::string outputDir = "."; + + // 获取依赖的方言 + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + } + +protected: + // 构建单个函数的 CFG + std::unique_ptr buildForFunction(triton::FuncOp func); +}; + +// 添加数据结构定义 +// IF_COND 的 yield value 和 result value 的对应关系 +struct IfYieldResultMapping { + // then 分支的 yield values（来自 scf.yield 操作的操作数） + SmallVector thenYieldValues; + // else 分支的 yield values（如果有 else 分支） + SmallVector elseYieldValues; + // if 操作的 result values + SmallVector resultValues; + // 对应关系: resultValues[i] 对应 thenYieldValues[i] 或 elseYieldValues[i] +}; + +// FOR_COND 的 yield value 和 iter args value 的对应关系 +struct ForYieldIterArgMapping { + // yield 操作的 values（来自循环体末尾的 scf.yield） + SmallVector yieldValues; + // iter_args（循环初始参数，对应 for 操作的 iter_args） + SmallVector iterArgValues; + // for 操作的 result values + SmallVector resultValues; + // 对应关系: iterArgValues[i] 在循环体中使用时被更新，yieldValues[i] 是新的值 + // resultValues[i] 对应最后一次迭代的 yieldValues[i] +}; + +// COND_BR 的 true/false 分支信息 +struct CondBranchMapping { + // true 分支的目标块参数值（来自 cf.cond_br 的 trueOperands） + SmallVector trueOperands; + // false 分支的目标块参数值（来自 cf.cond_br 的 falseOperands） + SmallVector falseOperands; + // 条件值 + Value condition; + // true 分支目标块 + Block *trueDest; + // false 分支目标块 + Block *falseDest; +}; + +// BR 的分支信息 +struct BranchMapping { + // 目标块参数值（来自 cf.br 的 destOperands） + SmallVector destOperands; + // 目标块 + Block *dest; +}; + +// 独立的 CFG 构建器类（用于非 Pass 场景） +class ControlFlowGraphBuilder { +public: + // 为函数构建 CFG + std::unique_ptr build(triton::FuncOp func); + + // 为模块构建所有函数的 CFG + std::vector> + buildForModule(ModuleOp module); + + // 处理一个 region，返回该 region 的入口块和出口块 + struct RegionBlocks { + cfg::BasicBlock *entryBlock; + cfg::BasicBlock *exitBlock; + }; + + RegionBlocks buildForRegion(Region ®ion, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *entryBlock, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 block 中的操作，返回最后处理的基本块 + cfg::BasicBlock *processBlock(Block &block, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 scf.if 操作，返回 if 后面的基本块 + cfg::BasicBlock *handleIfOp(scf::IfOp ifOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 scf.for 操作，返回 for 后面的基本块 + cfg::BasicBlock *handleForOp(scf::ForOp forOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 scf.while 操作，返回 while 后面的基本块 + cfg::BasicBlock *handleWhileOp(scf::WhileOp whileOp, + cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 cf.cond_br 操作，返回条件分支后面的基本块 + cfg::BasicBlock * + handleCondBranchOp(cf::CondBranchOp condBrOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 处理 cf.br 操作，返回无条件跳转后面的基本块 + cfg::BasicBlock *handleBranchOp(cf::BranchOp brOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure = nullptr); + + // 创建一个新的指令并添加到 basic block + cfg::Instruction *createInstruction(Operation *op, + cfg::BasicBlock *parentBlock, + cfg::ControlFlowGraph &cfg); + + // 1. 快速收集所有的 IF_COND 基本块 + // 遍历 CFG 中所有基本块，返回类型为 IF_COND 的基本块列表 + SmallVector + collectIfCondBlocks(cfg::ControlFlowGraph &cfg); + + // 2. 快速收集所有的 FOR_COND 基本块 + // 遍历 CFG 中所有基本块，返回类型为 FOR_COND 的基本块列表 + SmallVector + collectForCondBlocks(cfg::ControlFlowGraph &cfg); + + // 3. 获取 IF_COND 对应的 yield value 和 result value 的对应关系 + // 参数: IF_COND 类型的基本块 + // 返回: IfYieldResultMapping 结构体，包含 then/else 的 yield values 和 result + // values + std::optional + getIfYieldResultMapping(cfg::BasicBlock *ifCondBB); + + // 4. 获取 FOR_COND 对应的 yield value 和 iter args value 的对应关系 + // 参数: FOR_COND 类型的基本块 + // 返回: ForYieldIterArgMapping 结构体，包含 yield values、iter_args 和 result + // values + std::optional + getForYieldIterArgMapping(cfg::BasicBlock *forCondBB); + + // 5. 快速收集所有的 COND_BR 基本块 + // 遍历 CFG 中所有基本块，返回类型为 COND_BR 的基本块列表 + SmallVector + collectCondBrBlocks(cfg::ControlFlowGraph &cfg); + + // 6. 获取 COND_BR 对应的条件分支信息 + // 参数: COND_BR 类型的基本块 + // 返回: CondBranchMapping 结构体，包含条件、目标块和参数信息 + std::optional + getCondBranchMapping(cfg::BasicBlock *condBrBB); + + // 7. 快速收集所有的 BR 基本块 + // 遍历 CFG 中所有基本块，返回类型为 BR 的基本块列表 + SmallVector collectBrBlocks(cfg::ControlFlowGraph &cfg); + + // 8. 获取 BR 对应的分支信息 + // 参数: BR 类型的基本块 + // 返回: BranchMapping 结构体，包含目标块和参数信息 + std::optional getBranchMapping(cfg::BasicBlock *brBB); + + // 获取下一个指令 ID + size_t getNextInstructionId() { return nextInstructionId++; } + + // MLIR Block 到 CFG BasicBlock 的映射（用于处理 cf.cond_br 等跳转指令） + DenseMap blockToBasicBlockMap; + + // 获取或创建 Block 对应的 BasicBlock + cfg::BasicBlock * + getOrCreateBasicBlockForBlock(Block *block, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *parentStructure = nullptr); + + // 注册 Block 到 BasicBlock 的映射 + void registerBlockMapping(Block *mlirBlock, cfg::BasicBlock *cfgBlock); + +private: + size_t nextInstructionId = 0; // 下一个指令 ID +}; + +// 创建 Pass 的工厂函数 +std::unique_ptr> createBuildCFGPass(); + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_CONTROL_FLOW_GRAPH_BUILDER_H diff --git a/compiler/include/dicp/TritonToGraph/DataflowGraph.h b/compiler/include/dicp/TritonToGraph/DataflowGraph.h new file mode 100644 index 00000000..cfe140ce --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/DataflowGraph.h @@ -0,0 +1,259 @@ + + +#ifndef TRITON_TO_CFG_DATAFLOW_GRAPH_H +#define TRITON_TO_CFG_DATAFLOW_GRAPH_H + +#include "dicp/TritonToGraph/AliasAnalysis.h" +#include "dicp/TritonToGraph/ControlFlowGraph.h" +#include "dicp/TritonToGraph/MemorySSA.h" +#include "dicp/TritonToGraph/MemorySsaBuilder.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Casting.h" +#include + +namespace mlir { +namespace triton { +namespace cfg { + +// 前向声明结果类 +class DataFlowResult; +class MemorySSAResult; +class SSAResult; + +// DataFlowInfo - 统一的数据流信息 +class DataFlowInfo { +public: + // 为函数入口创建参数定义 + void createParameterDefinitions(triton::FuncOp func); + + // Memory SSA接口 + MemorySSADef *getMemoryDefinition(Value value) const; + void addMemoryDefinition(Value value, MemorySSADef *def); + + SmallVector getMemoryUses(Value value) const; + void addMemoryUse(Value value, const MemorySSAUse &use); + + void removeMemoryDefinition(Value value); + void clearMemoryUses(Value value); + + // 传统SSA接口（复用MLIR原生功能） + Operation *getSSADefinition(Value value) const { + return value.getDefiningOp(); + } + + SmallVector getSSAUses(Value value) const { + SmallVector result; + for (OpOperand &use : value.getUses()) { + result.push_back(&use); + } + return result; + } + + // 循环Phi接口 + void addPhi(Value value, const PhiInfo &phiInfo) { Phis[value] = phiInfo; } + + PhiInfo &getPhi(Value value) { return Phis[value]; } + + bool hasPhi(Value value) const { return Phis.count(value) > 0; } + + // 统一查询接口 - 返回unique_ptr，使用LLVM RTTI进行类型判断 + std::unique_ptr queryDataFlow(Value value) const; + + // 查询某个定义的所有使用 + SmallVector getUses(MemorySSADef *def) const; + + // 查询某个操作的memory使用 + SmallVector getUsesByUserOp(Operation *userOp) const; + + // 遍历接口 + void + forEachDefinition(llvm::function_ref func) const; + void forEachUse(llvm::function_ref func) const; + + // 获取所有Memory SSA definitions + const DenseMap &getMemoryDefinitions() const { + return memoryDefinitions; + } + + // 获取所有循环phi信息 + const DenseMap &getPhis() const { return Phis; } + + // 构建def-use缓存 + void buildDefUseCache() const; + + // 打印信息（调试用） + void print(llvm::raw_ostream &os) const; + + // 导出到JSON + void exportToJSON(llvm::raw_ostream &os) const; + +private: + // Memory SSA映射 + DenseMap memoryDefinitions; + DenseMap> memoryUses; + + // Loop Phi映射 + DenseMap Phis; + + // Use-Def映射缓存（def -> uses） + mutable DenseMap> defUseCache; + mutable bool defUseCacheValid = false; + + void invalidateDefUseCache() { + defUseCacheValid = false; + defUseCache.clear(); + } +}; + +// DataFlowResult - 数据流查询结果的基类 +// 使用LLVM RTTI系统，支持isa<>和dyn_cast<>进行类型判断 +class DataFlowResult { +public: + enum class Kind { + MemorySSA, // Memory SSA结果（tensor/pointer） + SSA, // 传统SSA结果（标量） + NONE // 无数据流信息 + }; + + DataFlowResult(Kind kind, Operation *originOp) + : kind(kind), originOp(originOp) {} + virtual ~DataFlowResult() = default; + + Kind getKind() const { return kind; } + Operation *getOriginOp() const { return originOp; } + + SmallVector &getUses() { return uses; } + const SmallVector &getUses() const { return uses; } + + std::optional &getPhi() { return Phi; } + const std::optional &getPhi() const { return Phi; } + + // LLVM RTTI支持 + static bool classof(const DataFlowResult *) { return true; } + +protected: + Kind kind; + Operation *originOp; + SmallVector uses; // 所有uses + std::optional Phi; // Phi信息（如果有） +}; + +// MemorySSAResult - Memory SSA的结果 +class MemorySSAResult : public DataFlowResult { +public: + MemorySSAResult(Operation *originOp, MemorySSADef *definition) + : DataFlowResult(Kind::MemorySSA, originOp), definition(definition) {} + + MemorySSADef *getDefinition() const { return definition; } + + // LLVM RTTI支持 + static bool classof(const DataFlowResult *result) { + return result->getKind() == Kind::MemorySSA; + } + +private: + MemorySSADef *definition; // MEMORY_SSA时使用 +}; + +// SSAResult - 传统SSA的结果 +class SSAResult : public DataFlowResult { +public: + SSAResult(Operation *originOp, Operation *ssaDefinition) + : DataFlowResult(Kind::SSA, originOp), ssaDefinition(ssaDefinition) {} + + Operation *getSSADefinition() const { return ssaDefinition; } + + // LLVM RTTI支持 + static bool classof(const DataFlowResult *result) { + return result->getKind() == Kind::SSA; + } + +private: + Operation *ssaDefinition; // SSA时使用 +}; + +// NoneResult - 无数据流信息的结果 +class NoneResult : public DataFlowResult { +public: + NoneResult() : DataFlowResult(Kind::NONE, nullptr) {} + + // LLVM RTTI支持 + static bool classof(const DataFlowResult *result) { + return result->getKind() == Kind::NONE; + } +}; + +// DataFlowGraph - 数据流图 +class DataFlowGraph { +public: + explicit DataFlowGraph(ControlFlowGraph &cfg) : cfg(cfg) {} + + ~DataFlowGraph() = default; + + // 构建完整的数据流信息 + void build(); + + // 查询Value的数据流信息（使用LLVM RTTI判断具体类型） + std::unique_ptr queryDataFlow(Value value) const { + return dataFlowInfo.queryDataFlow(value); + } + + // 获取所有Memory SSA definitions + SmallVector getAllDefinitions() const { + SmallVector result; + for (const auto &entry : dataFlowInfo.getMemoryDefinitions()) { + result.push_back(entry.second); + } + return result; + } + + // 获取definition的所有uses + SmallVector getUses(MemorySSADef *def) const { + return dataFlowInfo.getUses(def); + } + + // 获取操作的所有uses + SmallVector getUsesByUserOp(Operation *userOp) const { + return dataFlowInfo.getUsesByUserOp(userOp); + } + + // 获取CFG + ControlFlowGraph &getCFG() { return cfg; } + const ControlFlowGraph &getCFG() const { return cfg; } + + // 获取DataFlowInfo + DataFlowInfo &getDataFlowInfo() { return dataFlowInfo; } + const DataFlowInfo &getDataFlowInfo() const { return dataFlowInfo; } + + // 导出数据流信息到JSON + void exportToJSON(llvm::raw_ostream &os) const; + + // 导出def-use链到DOT格式 + void exportDefUseToDOT(llvm::raw_ostream &os) const; + + // 打印所有数据流信息（调试用） + void print(llvm::raw_ostream &os) const; + void dump() const; + +private: + ControlFlowGraph &cfg; + + // 组件 + std::unique_ptr aliasAnalysis; + std::unique_ptr memorySSABuilder; + + // 数据流信息 + DataFlowInfo dataFlowInfo; + + // 构建def-use图 + void buildDefUseGraph(); +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_DATAFLOW_GRAPH_H diff --git a/compiler/include/dicp/TritonToGraph/GraphAnalysis.h b/compiler/include/dicp/TritonToGraph/GraphAnalysis.h new file mode 100644 index 00000000..8208dc49 --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/GraphAnalysis.h @@ -0,0 +1,400 @@ + + +#ifndef TRITON_TO_GRAPH_GRAPH_ANALYSIS_H +#define TRITON_TO_GRAPH_GRAPH_ANALYSIS_H + +#include "dicp/TritonToGraph/ControlFlowGraph.h" +#include "dicp/TritonToGraph/DataflowGraph.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SmallVector.h" +#include + +namespace mlir { +namespace triton { +namespace cfg { + +//===----------------------------------------------------------------------===// +// Traversal Context - 遍历时维护上下文 +//===----------------------------------------------------------------------===// + +struct TraversalContext { + // 当前所处的嵌套结构栈（从外到内） + SmallVector structureStack; + + // 当前深度 + int depth = 0; + + void push(BasicBlock *structure) { structureStack.push_back(structure); } + void pop() { + if (!structureStack.empty()) + structureStack.pop_back(); + } + BasicBlock *currentStructure() const { + return structureStack.empty() ? nullptr : structureStack.back(); + } + bool isEmpty() const { return structureStack.empty(); } +}; + +/// CFGTraversalBase - CFG 遍历基类 +/// 用户继承此类，实现 preVisit/postVisit 回调 +/// 可在子类中添加自定义状态字段 +class CFGTraversalBase { +public: + virtual ~CFGTraversalBase() = default; + + // 访问基本块前调用，返回 false 则跳过此块 + virtual bool preVisitBlock(BasicBlock *block, TraversalContext &ctx); + + // 访问基本块后调用 + virtual void postVisitBlock(BasicBlock *block, TraversalContext &ctx) {} + + // 访问指令后调用 + virtual void VisitInstruction(Instruction *inst, TraversalContext &ctx) {} + + // 遇到控制流结构（for/if/while）时调用 + virtual void onEnterStructure(BasicBlock *structure, TraversalContext &ctx) {} + virtual void onExitStructure(BasicBlock *structure, TraversalContext &ctx) {} + + // 遇到回边时调用 + virtual void onBackEdge(BasicBlock *from, BasicBlock *to, + TraversalContext &ctx) {} +}; + +/// CFGTraverser - CFG 遍历器（双向支持） +class CFGTraverser { +public: + explicit CFGTraverser(ControlFlowGraph &cfg) : cfg(cfg) {} + + //===----------------------------------------------------------------------=== + // Forward Traversal (沿后继节点向下遍历) + //===----------------------------------------------------------------------=== + + /// 从入口开始的 DFS + void dfsForward(CFGTraversalBase &visitor); + + /// 从指定块开始的 DFS + void dfsForward(BasicBlock *start, CFGTraversalBase &visitor); + + /// BFS 遍历 + void bfsForward(CFGTraversalBase &visitor); + void bfsForward(BasicBlock *start, CFGTraversalBase &visitor); + + //===----------------------------------------------------------------------=== + // Backward Traversal (沿前驱节点向上遍历) + //===----------------------------------------------------------------------=== + + /// 从指定块反向 DFS（沿前驱遍历） + void dfsBackward(BasicBlock *start, CFGTraversalBase &visitor); + + /// 从指定块反向 BFS + void bfsBackward(BasicBlock *start, CFGTraversalBase &visitor); + +private: + ControlFlowGraph &cfg; + + void dfsForwardImpl(BasicBlock *block, DenseSet &visited, + TraversalContext &ctx, CFGTraversalBase &visitor); + void dfsBackwardImpl(BasicBlock *block, DenseSet &visited, + TraversalContext &ctx, CFGTraversalBase &visitor); +}; + +//===----------------------------------------------------------------------===// +// Curly Recursive Template Pattern for DFG Traversal +//===----------------------------------------------------------------------===// + +/// DFGTraversalBase - DFG 遍历基类 +class DFGTraversalBase { +public: + virtual ~DFGTTraversalBase() = default; + + // 访问定义前调用（反向遍历 value -> def） + virtual bool VisitDef(Value value, Operation *defOp, int depth); + + // 访问使用前调用（正向遍历 value -> use） + virtual bool VisitUse(Value value, OpOperand *use, int depth); + + // 遇到 phi/iter_arg 时调用 + virtual void onPhi(Value phiValue, const PhiInfo &phiInfo, int depth) {} +}; + +/// DFGTraverser - DFG 遍历器（双向 SSA/MemorySSA） +class DFGTraverser { +public: + struct Options { + bool useMemorySSA = false; // false=传统 SSA, true=Memory SSA + bool followPhi = true; // 是否跨越 phi/iter_arg + int maxDepth = -1; // -1 = 无限制 + DenseSet stopOps; // 遇到停止的操作 + }; + + explicit DFGTraverser(DataFlowGraph &dfg) : dfg(dfg) {} + + //===----------------------------------------------------------------------=== + // Backward Traversal (value -> definitions) + //===----------------------------------------------------------------------=== + + /// 从 value 开始反向追踪所有定义 + void dfsBackward(Value seed, DFGTraversalBase &visitor, + const Options &opts = {}); + + /// 多起点反向 DFS + void dfsBackward(ArrayRef seeds, DFGTraversalBase &visitor, + const Options &opts = {}); + + /// BFS 反向追踪 + void bfsBackward(Value seed, DFGTraversalBase &visitor, + const Options &opts = {}); + + //===----------------------------------------------------------------------=== + // Forward Traversal (value -> uses) + //===----------------------------------------------------------------------=== + + /// 从 value 开始正向追踪所有使用 + void dfsForward(Value seed, DFGTraversalBase &visitor, + const Options &opts = {}); + + /// 多起点正向 DFS + void dfsForward(ArrayRef seeds, DFGTraversalBase &visitor, + const Options &opts = {}); + + /// BFS 正向追踪 + void bfsForward(Value seed, DFGTraversalBase &visitor, + const Options &opts = {}); + + //===----------------------------------------------------------------------=== + // Bidirectional + //===----------------------------------------------------------------------=== + + /// 从 value 双向遍历（先 backward 到根，再 forward 到所有 uses） + void traverseBidirectional(Value seed, DFGTraversalBase &visitor, + const Options &opts = {}); + +private: + DataFlowGraph &dfg; + + void dfsBackwardImpl(Value value, DFGTraversalBase &visitor, + DenseSet &visited, const Options &opts, + int depth); + void dfsForwardImpl(Value value, DFGTraversalBase &visitor, + DenseSet &visited, const Options &opts, + int depth); +}; + +//===----------------------------------------------------------------------===// +// Region Abstraction and Analysis +//===----------------------------------------------------------------------===// + +/// Region - 指令集合（替代原始代码中的 SmallVector） +class Region { +public: + explicit Region(StringRef name = "") : name_(name.str()) {} + + void add(Instruction *inst); + void add(Operation *op, ControlFlowGraph &cfg); + void addAll(ArrayRef insts); + + bool contains(Instruction *inst) const; + bool contains(Operation *op) const; + + void remove(Instruction *inst); + void clear(); + + size_t size() const { return instSet_.size(); } + bool empty() const { return instSet_.empty(); } + + // 获取按块内顺序排序的指令列表 + SmallVector orderedInstructions() const; + + // 获取所有操作 + SmallVector operations() const; + + StringRef name() const { return name_; } + void setName(StringRef name) { name_ = name.str(); } + + // 迭代器支持 + auto begin() const { return instSet_.begin(); } + auto end() const { return instSet_.end(); } + +private: + std::string name_; + DenseSet instSet_; +}; + +/// RegionAnalyzer - Region 分析器 +class RegionAnalyzer { +public: + struct Dependency { + enum Type { DATA, CONTROL }; + Type type; + Value value; + Instruction *from; + Instruction *to; + }; + + struct ExternalDeps { + // 外部定义 -> 内部使用 + struct Input { + Value value; + Instruction *externalDef; // null 表示函数参数 + SmallVector internalUses; + }; + // 内部定义 -> 外部使用 + struct Output { + Value value; + Instruction *internalDef; + SmallVector externalUses; + }; + SmallVector inputs; + SmallVector outputs; + }; + + explicit RegionAnalyzer(DataFlowGraph &dfg, ControlFlowGraph &cfg) + : dfg(dfg), cfg(cfg) {} + + // 检查 region 之间是否有依赖 + bool hasDependency(const Region &from, const Region &to) const; + + // 获取两个 region 间的所有依赖 + SmallVector getDependencies(const Region &from, + const Region &to) const; + + // 分析 region 的外部依赖 + ExternalDeps analyzeExternalDeps(const Region ®ion) const; + + // 检查依赖是否为循环依赖（region A 依赖 B，B 又依赖 A） + bool isCyclicDependency(const Region &a, const Region &b) const; + +private: + DataFlowGraph &dfg; + ControlFlowGraph &cfg; +}; + +//===----------------------------------------------------------------------===// +// Program Slicing +//===----------------------------------------------------------------------===// + +/// SliceCriterion - 切片准则 +struct SliceCriterion { + enum Direction { BACKWARD, FORWARD, BIDIRECTIONAL }; + SmallVector seeds; + Direction dir = BACKWARD; + DFGTraverser::Options dfgOpts; +}; + +/// ProgramSlice - 程序切片 +class ProgramSlice { +public: + void add(Instruction *inst) { instructions_.insert(inst); } + void addAll(const Region ®ion); + + bool contains(Instruction *inst) const { + return instructions_.contains(inst); + } + + size_t size() const { return instructions_.size(); } + bool empty() const { return instructions_.empty(); } + + // 获取入口点（没有前驱在切片中） + SmallVector entryPoints(DataFlowGraph &dfg) const; + + // 获取出口点（没有后继在切片中） + SmallVector exitPoints(DataFlowGraph &dfg) const; + + // 集合操作 + void merge(const ProgramSlice &other); + void intersect(const ProgramSlice &other); + void subtract(const ProgramSlice &other); + + // 转换为 Region + Region toRegion(StringRef name = "") const; + + // 迭代器 + auto begin() const { return instructions_.begin(); } + auto end() const { return instructions_.end(); } + +private: + DenseSet instructions_; +}; + +/// ProgramSlicer - 程序切片器 +class ProgramSlicer { +public: + ProgramSlicer(DataFlowGraph &dfg, ControlFlowGraph &cfg) + : dfg(dfg), cfg(cfg) {} + + // 计算切片 + ProgramSlice compute(const SliceCriterion &criterion); + + // 从 yield values 计算切片（常用场景） + ProgramSlice sliceFromYields(ArrayRef yields, + SliceCriterion::Direction dir); + + // 多切片操作 + static ProgramSlice merge(ArrayRef slices); + static ProgramSlice intersect(ArrayRef slices); + + // 检查切片间依赖 + struct SliceDependency { + const ProgramSlice *from; + const ProgramSlice *to; + SmallVector values; + }; + SmallVector + computeDependencies(ArrayRef slices); + +private: + DataFlowGraph &dfg; + ControlFlowGraph &cfg; +}; + +//===----------------------------------------------------------------------===// +// Region Absorption +//===----------------------------------------------------------------------===// + +/// AbsorptionPolicy - 吸收策略 +struct AbsorptionPolicy { + enum Direction { UPSTREAM, DOWNSTREAM, BOTH }; + Direction dir = BOTH; + int maxDepth = -1; + bool crossRegionBoundary = false; + DenseSet stopOps; + std::function shouldStop = nullptr; +}; + +/// RegionAbsorber - Region 吸收器 +class RegionAbsorber { +public: + RegionAbsorber(DataFlowGraph &dfg, ControlFlowGraph &cfg) + : dfg(dfg), cfg(cfg) {} + + // 从种子指令开始吸收 + void absorb(Region ®ion, ArrayRef seeds, + const AbsorptionPolicy &policy); + + // 从 value 的 def/use 链吸收 + void absorbFromValue(Region ®ion, Value value, + const AbsorptionPolicy &policy); + + // 吸收直到遇到边界 + void absorbUntilBoundary(Region ®ion, ArrayRef seeds, + std::function isBoundary); + +private: + DataFlowGraph &dfg; + ControlFlowGraph &cfg; + + void absorbUpstream(Region ®ion, Instruction *inst, + const AbsorptionPolicy &policy, + DenseSet &visited, int depth); + void absorbDownstream(Region ®ion, Instruction *inst, + const AbsorptionPolicy &policy, + DenseSet &visited, int depth); +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_GRAPH_GRAPH_ANALYSIS_H diff --git a/compiler/include/dicp/TritonToGraph/InterProceduralCFG.h b/compiler/include/dicp/TritonToGraph/InterProceduralCFG.h new file mode 100644 index 00000000..c7694684 --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/InterProceduralCFG.h @@ -0,0 +1,88 @@ + + +#ifndef TRITON_TO_CFG_INTER_PROCEDURAL_CFG_H +#define TRITON_TO_CFG_INTER_PROCEDURAL_CFG_H + +#include "dicp/TritonToGraph/ControlFlowGraph.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Operation.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include + +namespace mlir { +namespace triton { +namespace cfg { + +// 过程间控制流图 (ICFG) +class InterProceduralCFG { +public: + explicit InterProceduralCFG(ModuleOp module); + ~InterProceduralCFG(); + + // 为每个函数构建CFG + void build(); + + // 获取函数的CFG + ControlFlowGraph *getFunctionCFG(triton::FuncOp func); + const ControlFlowGraph *getFunctionCFG(triton::FuncOp func) const; + ControlFlowGraph *getFunctionCFG(StringRef funcName); + + // 在调用点连接调用者和被调用者的CFG + struct CallSite { + Operation *callOp; // 调用操作 + triton::FuncOp caller; // 调用者函数 + triton::FuncOp callee; // 被调用者函数 + BasicBlock *callBlock; // 调用点的基本块 + Instruction *callInst; // 调用指令 + }; + + // 获取所有调用点 + const SmallVector &getCallSites() const { return callSites; } + + // 连接ICFG中的调用边 + void connectCallGraph(); + + // 查询函数调用关系 + SmallVector getCallees(triton::FuncOp caller) const; + SmallVector getCallers(triton::FuncOp callee) const; + + // 全局可达性分析 + void computeReachability(); + bool isReachable(triton::FuncOp from, triton::FuncOp to) const; + + // 可视化 + void dumpToDot(const std::string &filename) const; + void print(raw_ostream &os) const; + + // 导出到 HTML（包含所有函数） + llvm::Error exportToHTML(const std::string &filename) const; + +private: + ModuleOp module; + + // 函数到CFG的映射 + DenseMap> functionCFGs; + + // 调用点列表 + SmallVector callSites; + + // 调用图 (函数级别) + DenseMap> callGraph; + DenseMap> reverseCallGraph; + + // 可达性矩阵 + DenseMap> reachability; +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_INTER_PROCEDURAL_CFG_H diff --git a/compiler/include/dicp/TritonToGraph/MemorySSA.h b/compiler/include/dicp/TritonToGraph/MemorySSA.h new file mode 100644 index 00000000..dab1a4ca --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/MemorySSA.h @@ -0,0 +1,274 @@ + + +#ifndef TRITON_TO_CFG_MEMORY_SSA_H +#define TRITON_TO_CFG_MEMORY_SSA_H + +#include "dicp/TritonToGraph/tensor.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace triton { +namespace cfg { + +// Forward declarations +class MemorySSADef; +class MemorySSAUse; + +class BasicBlock; +class ControlFlowGraph; + +// Memory SSA Definition - 表示tensor/pointer的定义 +class MemorySSADef { +public: + // 构造函数 + MemorySSADef(TensorObject *tensor, Operation *defOp, unsigned version = 0) + : tensor(tensor), defOp(defOp), version(version) {} + + // 获取tensor对象 + TensorObject *getTensor() const { return tensor; } + + // 获取创建该definition的操作 + // - 对于入参，返回nullptr + // - 对于其他操作，返回对应的Operation指针 + Operation *getDefOp() const { return defOp; } + + // 获取版本号（函数内全局唯一序数） + unsigned getVersion() const { return version; } + + // 获取唯一标识: tensor_name,version 或 tensor_name,param + std::string getId() const { + return tensor->getName() + "," + + (defOp ? std::to_string(version) : std::to_string(0)); + } + + // 判断是否是入参（函数参数） + bool isParameter() const { return defOp == nullptr; } + + // 判断是否是phi节点（控制流合并） + bool isPhi() const { + return defOp && (isa(defOp) || isa(defOp) || + isa(defOp)); + } + + // 打印信息 + void print(llvm::raw_ostream &os) const { + os << "Definition[" << getId() << ", tensor=" << tensor->getName(); + if (defOp) { + os << ", op=" << defOp->getName(); + } else { + os << ", param"; + } + os << "]"; + } + +private: + TensorObject *tensor; // 对应的tensor对象 + Operation *defOp; // 创建该definition的操作 + unsigned version; // 版本号（函数内全局递增） +}; + +// Memory SSA Use - 表示tensor/pointer的使用 +class MemorySSAUse { +public: + // 构造函数 + MemorySSAUse(MemorySSADef *definition, Operation *userOp, unsigned operandIdx) + : definition(definition), userOp(userOp), operandIdx(operandIdx) { + // 缓存value以提高查询性能 + operandValue = userOp->getOperand(operandIdx); + } + + // 获取使用的definition + MemorySSADef *getDefinition() const { return definition; } + + // 获取使用该definition的操作 + Operation *getUserOp() const { return userOp; } + + // 获取operand序号 + unsigned getOperandIdx() const { return operandIdx; } + + // 获取Value（从userOp的operand） + Value getValue() const { return operandValue; } + + // 获取用户操作的名称 + std::string getUserOpName() const { + return userOp->getName().getStringRef().str(); + } + + // 打印信息 + void print(llvm::raw_ostream &os) const { + os << "Use["; + if (definition) { + os << definition->getId(); + } else { + os << "null"; + } + os << " in " << userOp->getName() << ", operand #" << operandIdx << "]"; + } + +private: + MemorySSADef *definition; // 使用的definition + Operation *userOp; // 使用该definition的操作 + unsigned operandIdx; // operand序号 + Value operandValue; // 缓存的operand value +}; + +// PhiInfo - 循环Phi信息 +struct PhiInfo { + // Phi类型 + enum Type { + ITER_ARG, // scf.for的iter_arg + IF_RESULT, // scf.if的result + WHILE_ARG // scf.while的arg + }; + + Type type; + BasicBlock *loopHeader; // 循环头基本块 + + // Phi值的来源 + struct { + MemorySSADef *initialValue; // 初始值（初始iteration） + MemorySSADef *yieldValue; // yield的值（后续iteration） + } comingFrom; + + // 是否是第一次迭代 + bool isInitial() const { return comingFrom.yieldValue == nullptr; } + + // 获取当前definition（根据上下文决定） + MemorySSADef *getCurrentDefinition(int iteration) const { + return (iteration == 0) ? comingFrom.initialValue : comingFrom.yieldValue; + } +}; + +// MemorySSAInfo - 指令的Memory SSA信息 +struct MemorySSAInfo { + // 指令的operands使用的definitions + SmallVector uses; + + // 指令的results创建的definitions + SmallVector definitions; + + // Alias信息（仅对pointer相关操作） + struct AliasInfo { + Value aliasee; // 别名的源value + TensorObject *baseTensor; // 对应的tensor对象 + }; + std::optional aliasInfo; + + // 快速查询接口 + bool hasDefinition(Value value) const { + // 通过 getDefiningOp() 获取定义该 value 的操作 + Operation *defOp = value.getDefiningOp(); + if (!defOp) + return false; + + // 获取该操作的所有 results + auto results = defOp->getResults(); + for (size_t i = 0; i < definitions.size() && i < results.size(); ++i) { + if (results[i] == value) { + return definitions[i] != nullptr; + } + } + return false; + } + + MemorySSADef *getDefinition(Value value) const { + // 查找该value在results中的索引 + for (auto result : llvm::enumerate(value.getDefiningOp()->getResults())) { + if (result.value() == value) { + size_t idx = result.index(); + if (idx < definitions.size()) { + return definitions[idx]; + } + break; + } + } + return nullptr; + } + + bool hasUse(Value value) const { + for (const MemorySSAUse &use : uses) { + if (use.getValue() == value) { + return true; + } + } + return false; + } + + SmallVector getUses(Value value) const { + SmallVector result; + for (const MemorySSAUse &use : uses) { + if (use.getValue() == value) { + result.push_back(use); + } + } + return result; + } + + // 判断是否创建了新的definition + bool hasNewDefinitions() const { return !definitions.empty(); } + + // 判断是否是写入操作 + bool isMemoryWriter() const { + for (MemorySSADef *def : definitions) { + if (def && !def->isParameter()) { + return true; + } + } + return false; + } + + // 遍历定义 + void forEachDefinition(llvm::function_ref func) const { + for (MemorySSADef *def : definitions) { + if (def) + func(def); + } + } + + void forEachUse(llvm::function_ref func) const { + for (const MemorySSAUse &use : uses) { + if (use.getDefinition()) + func(use); + } + } + + // 清空所有信息 + void clear() { + uses.clear(); + definitions.clear(); + aliasInfo.reset(); + } + + // 打印信息 + void print(llvm::raw_ostream &os) const { + os << "MemorySSAInfo[\n"; + os << " Uses: " << uses.size() << "\n"; + for (const MemorySSAUse &use : uses) { + os << " "; + use.print(os); + os << "\n"; + } + os << " Definitions: " << definitions.size() << "\n"; + for (MemorySSADef *def : definitions) { + if (def) { + os << " "; + def->print(os); + os << "\n"; + } + } + if (aliasInfo) { + os << " Alias: " << aliasInfo->aliasee << " -> " + << aliasInfo->baseTensor->getName() << "\n"; + } + os << "]"; + } +}; + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_MEMORY_SSA_H diff --git a/compiler/include/dicp/TritonToGraph/MemorySsaBuilder.h b/compiler/include/dicp/TritonToGraph/MemorySsaBuilder.h new file mode 100644 index 00000000..4513f1b8 --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/MemorySsaBuilder.h @@ -0,0 +1,155 @@ + + +#ifndef TRITON_TO_CFG_MEMORY_SSA_BUILDER_H +#define TRITON_TO_CFG_MEMORY_SSA_BUILDER_H + +#include "dicp/TritonToGraph/AliasAnalysis.h" +#include "dicp/TritonToGraph/MemorySSA.h" +#include "mlir/IR/Types.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" + +namespace mlir { +namespace triton { +namespace cfg { + +// 前向声明 +class DataFlowInfo; +class ControlFlowGraph; +class Instruction; + +// MemorySSABuilder - Memory SSA构建器 +// 构建整个CFG的Memory SSA信息，包括创建definitions、uses、处理控制流phi节点 +class MemorySSABuilder { +public: + MemorySSABuilder(ControlFlowGraph &cfg, AliasAnalysis &aliasAnalysis, + DataFlowInfo &dataFlowInfo) + : cfg(cfg), aliasAnalysis(aliasAnalysis), dataFlowInfo(dataFlowInfo) {} + + ~MemorySSABuilder(); + + // 构建整个CFG的Memory SSA + void build(); + +private: + // 处理单个BasicBlock + void processBasicBlock(BasicBlock *bb); + + // 处理单个指令 + void processInstruction(Instruction *inst); + + // 处理scf.if的phi节点 + void processIfOp(scf::IfOp ifOp, Instruction *inst, BasicBlock *thenEntryBB, + BasicBlock *elseEntryBB); + + // 处理scf.for的iter_args + void processForOp(scf::ForOp forOp, Instruction *inst, + BasicBlock *loopBodyEntryBB); + + // 处理scf.while的args + void processWhileOp(scf::WhileOp whileOp, Instruction *inst, + BasicBlock *beforeEntryBB, BasicBlock *afterEntryBB); + + // 判断是否是tensor类型 + bool isTensorType(Type type) const { + return mlir::isa(type) || + mlir::isa(type); + } + + // 根据操作创建tensor对象 + TensorObject *createTensorObject(Operation *op); + + // 创建tensor definition + MemorySSADef *createDefinition(TensorObject *tensor, Operation *op); + + // 创建use + MemorySSAUse createUse(MemorySSADef *def, Operation *userOp, + unsigned operandIdx); + + // 判断是否是入参 + bool isParameter(Operation *op) const { + return op == nullptr; // 入参的defOp为nullptr + } + + // 判断是否是返回新Tensor的操作（根据返回值类型判断，排除load） + bool isTensorWriter(Operation *op) const; + + // 判断是否是修改内存的操作（有副作用） + bool isMemoryWriter(Operation *op) const { + // 只有：tt.store（写入内存） + return isa(op); + } + + // 判断是否是修改读取的操作（有副作用） + bool isMemoryReader(Operation *op) const { + // 只有：tt.store（写入内存） + return isa(op); + } + + // 判断是否是创建指针的操作 + bool isPointerOp(Operation *op) const { + // 返回指针类型：addptr（偏移指针）、make_tensor_ptr（创建张量指针） + if (isa(op)) + return true; + + if (isPointerBroadcastOrSplat(op)) + return true; + + return false; + } + + bool isPointerBroadcastOrSplat(mlir::Operation *op) const; + + // 创建函数的参数定义 + void createParameterDefinitions(); + + // 获取或创建tensor对象 + TensorObject *getOrCreateTensorObject(Value value); + + // 获取操作的字符串名称（用于生成tensor名称） + std::string getOpName(Operation *op); + + // 成员变量 + ControlFlowGraph &cfg; + AliasAnalysis &aliasAnalysis; + DataFlowInfo &dataFlowInfo; + // size_t nextVersionId; // 下一个可用的版本号（用于MemorySSADef） + // size_t nextTensorId; // 下一个可用的tensor ID（用于TensorObject命名） + + std::map nextVersion; + std::map nextTensor; + + // 所有创建的definitions（用于内存管理） + SmallVector allDefinitions; + + // tensor对象缓存 + DenseMap tensorObjectCache; +}; + +// MemorySSABuilderHelper - 辅助函数 +namespace MemorySSABuilderHelper { +// 获取操作的结果类型 +Type getResultType(Operation *op, unsigned resultIdx); + +// 获取Value的形状信息 +SmallVector getShapeFromValue(Value value); + +// 判断两个shape是否相同 +bool shapesEqual(ArrayRef shape1, ArrayRef shape2); + +// 获取scf.if的region yield操作 +Operation *getYieldOp(Region ®ion); + +// 为tensor创建唯一名称 +std::string createUniqueTensorName(StringRef prefix, size_t id); + +// 判断是否需要为操作创建新版本 +bool shouldCreateNewVersion(Operation *op, MemorySSADef *currentDef); +} // namespace MemorySSABuilderHelper + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_MEMORY_SSA_BUILDER_H diff --git a/compiler/include/dicp/TritonToGraph/Passes.h b/compiler/include/dicp/TritonToGraph/Passes.h new file mode 100644 index 00000000..f8f17b7a --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/Passes.h @@ -0,0 +1,24 @@ + + +#ifndef TRITON_TO_CFG_PASSES_H +#define TRITON_TO_CFG_PASSES_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +namespace mlir { +namespace triton { +namespace cfg { + +// 创建 BuildCFG pass 的工厂函数 +std::unique_ptr> createBuildCFGPass(); + +// 注册所有 CFG 相关的 passes +#define GEN_PASS_REGISTRATION +#include "dicp/TritonToGraph/Passes.h.inc" + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_PASSES_H diff --git a/compiler/include/dicp/TritonToGraph/Passes.td b/compiler/include/dicp/TritonToGraph/Passes.td new file mode 100644 index 00000000..78230c76 --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/Passes.td @@ -0,0 +1,25 @@ + + +#ifndef TRITON_TO_CFG_PASSES +#define TRITON_TO_CFG_PASSES + +include "mlir/Pass/PassBase.td" + +def BuildCFG : Pass<"build-cfg", "mlir::ModuleOp"> { + let summary = "Build Control Flow Graph from TTIR"; + let description = [{ + This pass analyzes the TTIR and builds a Control Flow Graph (CFG) + representation. It exports the CFG to multiple formats (text, DOT, JSON, HTML) + for visualization and analysis. + + The output files will be named as _cfg. and placed in + the directory specified by the output-dir option. + }]; + let constructor = "mlir::triton::cfg::createBuildCFGPass()"; + let options = [ + Option<"output-dir", "outputDir", "std::string", "\".\"", + "Directory where CFG files will be generated">, + ]; +} + +#endif // TRITON_TO_CFG_PASSES diff --git a/compiler/include/dicp/TritonToGraph/tensor.h b/compiler/include/dicp/TritonToGraph/tensor.h new file mode 100644 index 00000000..76e18cc9 --- /dev/null +++ b/compiler/include/dicp/TritonToGraph/tensor.h @@ -0,0 +1,146 @@ + + +#ifndef TRITON_TO_CFG_TENSOR_H +#define TRITON_TO_CFG_TENSOR_H + +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Types.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/StringRef.h" + +namespace mlir { +namespace triton { +namespace cfg { + +// Compute类型（用于指令分类） +enum class ComputeType { + CUBE, // shape为2维 + VECTOR, // shape为1维 + SCALAR // 其他标量 +}; + +// TensorObject - Tensor对象定义 +class TensorObject { +public: + // Tensor类型分类 + enum class TensorKind { + GLOBAL_MEMORY, // 全局内存tensor（如gm_obj） + L2, // L2 Cache + L1, // L1 Cache + UB // Unified Buffer + }; + + // 构造函数 + TensorObject(StringRef name, ArrayRef shape, Type type, + Type elementType, TensorKind kind = TensorKind::GLOBAL_MEMORY) + : name(name.str()), shape(shape.begin(), shape.end()), type(type), + elementType(elementType), kind(kind) {} + + // 获取tensor名称 + const std::string &getName() const { return name; } + + // 设置tensor名称 + void setName(StringRef newName) { name = newName.str(); } + + // 获取shape + ArrayRef getShape() const { return shape; } + + // 获取MLIR类型 + Type getType() const { return type; } + + // 获取tensor种类 + TensorKind getKind() const { return kind; } + + // 设置tensor种类 + void setKind(TensorKind newKind) { kind = newKind; } + + // 获取元素数据类型 + Type getElementType() const { return elementType; } + + // 设置元素数据类型 + void setElementType(Type newElementType) { elementType = newElementType; } + + // 获取维度数 + size_t getRank() const { return shape.size(); } + + // 获取tensor大小（元素总数） + int64_t getSize() const { + int64_t size = 1; + for (int64_t dim : shape) { + size *= dim; + } + return size; + } + + // 判断是否是一维tensor + bool isVector() const { return shape.size() == 1; } + + // 判断是否是二维tensor + bool isMatrix() const { return shape.size() == 2; } + + // 判断是否是标量 + bool isScalar() const { return shape.empty(); } + + // 打印信息 + void print(llvm::raw_ostream &os) const { + os << "Tensor[" << name << ", shape=["; + for (size_t i = 0; i < shape.size(); ++i) { + if (i > 0) + os << ", "; + os << shape[i]; + } + os << "], element=" << elementType << ", kind=" << getKindString() << "]"; + } + + // 获取kind的字符串表示 + std::string getKindString() const { + switch (kind) { + case TensorKind::GLOBAL_MEMORY: + return "GLOBAL_MEMORY"; + case TensorKind::L2: + return "L2"; + case TensorKind::L1: + return "L1"; + case TensorKind::UB: + return "UB"; + } + return "UNKNOWN"; + } + +private: + std::string name; // Tensor名称，如"gm_obj_0" + SmallVector shape; + Type type; // 完整类型（如tensor<64x64xf32>） + Type elementType; // 元素数据类型（如f32, i8, f16等） + TensorKind kind; +}; + +// 从类型中提取shape和element type +inline void extractShapeAndElementType(Type type, + SmallVectorImpl &shape, + Type &elementType) { + if (auto rankedType = mlir::dyn_cast(type)) { + shape.append(rankedType.getShape().begin(), rankedType.getShape().end()); + elementType = rankedType.getElementType(); + } else if (auto ptrType = mlir::dyn_cast(type)) { + Type pointeeType = ptrType.getPointeeType(); + if (auto rankedType = mlir::dyn_cast(pointeeType)) { + shape.append(rankedType.getShape().begin(), rankedType.getShape().end()); + elementType = rankedType.getElementType(); + } else { + // 标量指针 + elementType = pointeeType; + } + } else { + // 默认值 + elementType = type; + } +} + +} // namespace cfg +} // namespace triton +} // namespace mlir + +#endif // TRITON_TO_CFG_TENSOR_H diff --git a/compiler/include/dicp/TritonToHFusion/CMakeLists.txt b/compiler/include/dicp/TritonToHFusion/CMakeLists.txt new file mode 100644 index 00000000..e4b37b61 --- /dev/null +++ b/compiler/include/dicp/TritonToHFusion/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToHFusion) +add_public_tablegen_target(TritonToHFusionConversionPassIncGen) \ No newline at end of file diff --git a/compiler/include/dicp/TritonToHFusion/Passes.h b/compiler/include/dicp/TritonToHFusion/Passes.h new file mode 100644 index 00000000..b53a35fd --- /dev/null +++ b/compiler/include/dicp/TritonToHFusion/Passes.h @@ -0,0 +1,23 @@ + + +#ifndef TRITON_ADAPTER_TRITON_TO_HFUSION_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_HFUSION_CONVERSION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert Triton dialect to HFusion dialect. +std::unique_ptr> createTritonToHFusionPass(); + +#define GEN_PASS_REGISTRATION +#include "dicp/TritonToHFusion/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_TO_HFUSION_CONVERSION_PASSES_H diff --git a/compiler/include/dicp/TritonToHFusion/Passes.td b/compiler/include/dicp/TritonToHFusion/Passes.td new file mode 100644 index 00000000..d4bd6d31 --- /dev/null +++ b/compiler/include/dicp/TritonToHFusion/Passes.td @@ -0,0 +1,12 @@ +#ifndef TRITON_TO_HFUSION_CONVERSION_PASSES +#define TRITON_TO_HFUSION_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def TritonToHFusion : Pass<"triton-to-hfusion", "mlir::ModuleOp"> { + let summary = "Convert Triton to HFusion dialect"; + let constructor = "triton::createTritonToHFusionPass()"; + let dependentDialects = ["hfusion::HFusionDialect"]; +} + +#endif // TRITON_TO_HFUSION_CONVERSION_PASSES diff --git a/compiler/include/dicp/TritonToHIVM/CMakeLists.txt b/compiler/include/dicp/TritonToHIVM/CMakeLists.txt new file mode 100644 index 00000000..4db98b26 --- /dev/null +++ b/compiler/include/dicp/TritonToHIVM/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToHIVM) +add_public_tablegen_target(TritonToHIVMConversionPassIncGen) \ No newline at end of file diff --git a/compiler/include/dicp/TritonToHIVM/Passes.h b/compiler/include/dicp/TritonToHIVM/Passes.h new file mode 100644 index 00000000..2c7f95cd --- /dev/null +++ b/compiler/include/dicp/TritonToHIVM/Passes.h @@ -0,0 +1,23 @@ + + +#ifndef TRITON_ADAPTER_TRITON_TO_HIVM_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_HIVM_CONVERSION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert Triton dialect to HIVM dialect. +std::unique_ptr> createTritonToHIVMPass(); + +#define GEN_PASS_REGISTRATION +#include "dicp/TritonToHIVM/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_TO_HIVM_CONVERSION_PASSES_H diff --git a/compiler/include/dicp/TritonToHIVM/Passes.td b/compiler/include/dicp/TritonToHIVM/Passes.td new file mode 100644 index 00000000..db1650c6 --- /dev/null +++ b/compiler/include/dicp/TritonToHIVM/Passes.td @@ -0,0 +1,12 @@ +#ifndef TRITON_TO_HIVM_CONVERSION_PASSES +#define TRITON_TO_HIVM_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def TritonToHIVM : Pass<"triton-to-hivm", "mlir::ModuleOp"> { + let summary = "Convert Triton to HIVM dialect"; + let constructor = "triton::createTritonToHIVMPass()"; + let dependentDialects = ["hivm::HIVMDialect"]; +} + +#endif // TRITON_TO_HIVM_CONVERSION_PASSES diff --git a/compiler/include/dicp/TritonToLLVM/CMakeLists.txt b/compiler/include/dicp/TritonToLLVM/CMakeLists.txt new file mode 100644 index 00000000..fb5a9b29 --- /dev/null +++ b/compiler/include/dicp/TritonToLLVM/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToLLVM) +add_public_tablegen_target(TritonToLLVMConversionPassIncGen) diff --git a/compiler/include/dicp/TritonToLLVM/Passes.h b/compiler/include/dicp/TritonToLLVM/Passes.h new file mode 100644 index 00000000..c36abd20 --- /dev/null +++ b/compiler/include/dicp/TritonToLLVM/Passes.h @@ -0,0 +1,21 @@ +#ifndef TRITON_ADAPTER_TRITON_TO_LLVM_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_LLVM_CONVERSION_PASSES_H + +#include "mlir/Pass/Pass.h" + +namespace mlir { +// Forward declarations. +class ModuleOp; + +namespace triton { + +/// Creates a pass to convert Triton dialect to LLVM dialect. +std::unique_ptr> createTritonToLLVMPass(); + +#define GEN_PASS_REGISTRATION +#include "dicp/TritonToLLVM/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_TO_LLVM_CONVERSION_PASSES_H diff --git a/compiler/include/dicp/TritonToLLVM/Passes.td b/compiler/include/dicp/TritonToLLVM/Passes.td new file mode 100644 index 00000000..56e5dc2a --- /dev/null +++ b/compiler/include/dicp/TritonToLLVM/Passes.td @@ -0,0 +1,12 @@ +#ifndef TRITON_TO_LLVM_CONVERSION_PASSES +#define TRITON_TO_LLVM_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def TritonToLLVM : Pass<"triton-to-llvm", "mlir::ModuleOp"> { + let summary = "Convert Triton to LLVM dialect"; + let constructor = "triton::createTritonToLLVMPass()"; + let dependentDialects = ["LLVM::LLVMDialect", "tensor::TensorDialect", "arith::ArithDialect"]; +} + +#endif // TRITON_TO_LLVM_CONVERSION_PASSES diff --git a/compiler/include/dicp/TritonToLinalg/ArgMinMaxConverter.h b/compiler/include/dicp/TritonToLinalg/ArgMinMaxConverter.h new file mode 100644 index 00000000..a4c6ed26 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/ArgMinMaxConverter.h @@ -0,0 +1,333 @@ + + +#ifndef TRITON_ADAPTER_ARGMINMAXCONVERTER_H +#define TRITON_ADAPTER_ARGMINMAXCONVERTER_H + +#include "dicp/TritonToLinalg/ConversionPatterns.h" +#include "dicp/Utils/Utils.h" + +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/Interfaces/FunctionInterfaces.h" +#include "mlir/Transforms/DialectConversion.h" + +#define DEBUG_TYPE "triton-to-linalg" + +#include "llvm/Support/Debug.h" +#include "llvm/Support/LogicalResult.h" + +#include + +namespace TTOpConverters { +using namespace mlir; +using namespace triton; + +template +class ArgMinMaxBaseConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchTieBreakResult(Value currValue, Value currIndex, + Value reduceValue, Value reduceIndex, + mlir::Block::iterator &it, + Value &tileBreakValue) const { + LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n"); + auto eqCmpOp = dyn_cast(*it); + if (eqCmpOp) { + if (eqCmpOp.getPredicate() != arith::CmpFPredicate::OEQ || + currValue != eqCmpOp.getLhs() || reduceValue != eqCmpOp.getRhs()) { + return failure(); + } + } + + auto eqCmpIOp = dyn_cast(*it++); + if (eqCmpIOp) { + if (eqCmpIOp.getPredicate() != arith::CmpIPredicate::eq || + currValue != eqCmpIOp.getLhs() || reduceValue != eqCmpIOp.getRhs()) { + return failure(); + } + } + + LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n"); + auto sltCmpOp = dyn_cast(*it++); + if (!sltCmpOp || sltCmpOp.getPredicate() != arith::CmpIPredicate::slt || + currIndex != sltCmpOp.getLhs() || reduceIndex != sltCmpOp.getRhs()) { + return failure(); + } + + // matching: %13 = arith.andi %11, %12 : i1 + LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n"); + auto andOp = dyn_cast(*it++); + + Value cmpOp; + if (eqCmpOp) + cmpOp = eqCmpOp; + else + cmpOp = eqCmpIOp; + + if (!andOp || andOp.getLhs() != cmpOp || andOp.getRhs() != sltCmpOp) { + return failure(); + } + + tileBreakValue = andOp; + return success(); + } + + LogicalResult matchShouldUpdateValue(Value currValue, Value currIndex, + Value reduceValue, Value reduceIndex, + mlir::Block::iterator &it, + Value &shouldUpdate) const { + Value tieResult; + if (failed(matchTieBreakResult(currValue, currIndex, reduceValue, + reduceIndex, it, tieResult))) { + LLVM_DEBUG(llvm::dbgs() << "Tie break result match failed\n"); + return failure(); + } + + Value comparisonResult; + if (failed(T::matchComparisonResult(currValue, currIndex, reduceValue, + reduceIndex, it, comparisonResult))) { + LLVM_DEBUG(llvm::dbgs() << "Comparison result match failed\n"); + return failure(); + } + + LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n"); + auto orOp = dyn_cast(*it++); + if (!orOp || orOp.getLhs() != comparisonResult || + orOp.getRhs() != tieResult) { + return failure(); + } + + shouldUpdate = orOp; + return success(); + } + + Value getInitTensor(ConversionPatternRewriter &rewriter, + ArrayRef shape, Value fillValue, + Location loc) const { + Value initTensor = + rewriter.create(loc, shape, fillValue.getType()); + return rewriter + .create(loc, ValueRange{fillValue}, + ValueRange{initTensor}) + .result(); + } + +public: + ArgMinMaxBaseConverter(MLIRContext *context) : OpConversionPattern(context) {} + + LogicalResult + matchAndRewrite(triton::ReduceOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + if (op.getBody()->getNumArguments() != 4) { + return failure(); + } + + auto block = op.getBody(); + auto ops = block->without_terminator(); + + Value currValue = block->getArgument(0); + Value currIndex = block->getArgument(1); + Value reduceValue = block->getArgument(2); + Value reduceIndex = block->getArgument(3); + + auto opsIt = ops.begin(); + Value shouldUpdate; + if (failed(matchShouldUpdateValue(currValue, currIndex, reduceValue, + reduceIndex, opsIt, shouldUpdate))) { + return failure(); + } + + LLVM_DEBUG(llvm::dbgs() << "Matching: " << *opsIt << "\n"); + auto valueSelectOp = dyn_cast(*opsIt++); + if (!valueSelectOp || valueSelectOp.getCondition() != shouldUpdate || + currValue != valueSelectOp.getTrueValue() || + reduceValue != valueSelectOp.getFalseValue()) { + return failure(); + } + + LLVM_DEBUG(llvm::dbgs() << "Matching: " << *opsIt << "\n"); + auto indexSelectOp = dyn_cast(*opsIt++); + if (!indexSelectOp || indexSelectOp.getCondition() != shouldUpdate || + currIndex != indexSelectOp.getTrueValue() || + reduceIndex != indexSelectOp.getFalseValue()) { + return failure(); + } + + LLVM_DEBUG(llvm::dbgs() << "Matching: " << *opsIt << "\n"); + auto termOp = dyn_cast(*opsIt++); + if (!(termOp && termOp == block->getTerminator() && + termOp.getOperands() == + ArrayRef{valueSelectOp, indexSelectOp})) { + return failure(); + } + + // Rewrite phase: perform the actual conversion + auto loc = op.getLoc(); + auto elemTypes = op.getElementTypes(); + + auto valueType = elemTypes[0]; + // tl.argmin reorder + bool isUnsigned = false; + if (isa(valueType)) { + arith::CmpFOp cmpFOp; + block->walk([&](arith::CmpFOp cmpOp) { + auto pred = cmpOp.getPredicate(); + if (pred == arith::CmpFPredicate::OEQ || + pred == arith::CmpFPredicate::ONE || + pred == arith::CmpFPredicate::UEQ || + pred == arith::CmpFPredicate::UNE) { + return WalkResult::advance(); + } else if (pred == arith::CmpFPredicate::OGT || + pred == arith::CmpFPredicate::OLT || + pred == arith::CmpFPredicate::UGT || + pred == arith::CmpFPredicate::ULT) { + cmpFOp = cmpOp; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + cmpFOp->moveBefore(block, block->getOperations().begin()); + } else if (isa(valueType)) { + arith::CmpIOp cmpIOp; + block->walk([&](arith::CmpIOp cmpOp) { + auto pred = cmpOp.getPredicate(); + if (pred == arith::CmpIPredicate::ugt || + pred == arith::CmpIPredicate::ult) { + isUnsigned = true; + } + if (pred == arith::CmpIPredicate::eq || + pred == arith::CmpIPredicate::ne) { + return WalkResult::advance(); + } else if (pred == arith::CmpIPredicate::sgt || + pred == arith::CmpIPredicate::slt || + pred == arith::CmpIPredicate::ugt || + pred == arith::CmpIPredicate::ult) { + if (cmpOp.getLhs() == block->getArgument(0) && + cmpOp.getRhs() == block->getArgument(2)) { + cmpIOp = cmpOp; + return WalkResult::interrupt(); + } + } + return WalkResult::advance(); + }); + cmpIOp->moveBefore(block, block->getOperations().begin()); + } + + TypedAttr valueAttr; + if (isa(valueType)) { + valueAttr = rewriter.getFloatAttr(valueType, T::getBaseReductionValue()); + } else if (isa(valueType)) { + if (isUnsigned) { + valueAttr = + rewriter.getIntegerAttr(valueType, T::getBaseReductionUIntValue()); + } else { + valueAttr = + rewriter.getIntegerAttr(valueType, T::getBaseReductionIntValue()); + } + } + + auto reduceWithIndexParams = getReduceWithIndexParams(op); + auto valuesAccBaseVal = + rewriter.create(loc, valueType, valueAttr); + int indicesInitValue = + (llvm::succeeded(reduceWithIndexParams) && + reduceWithIndexParams->tieBreakType == TieBreakType::RIGHT) + ? -1 + : std::numeric_limits::max(); + + auto indexType = elemTypes[1]; + auto indicesAccBaseVal = rewriter.create( + loc, indexType, rewriter.getIntegerAttr(indexType, indicesInitValue)); + + auto valueResultType = dyn_cast(op.getType(0)); + const auto isScalarReduce = valueResultType == nullptr; + SmallVector reductionResultShape{ + isScalarReduce ? SmallVector{} + : SmallVector(valueResultType.getShape())}; + + SmallVector outputs{ + getInitTensor(rewriter, reductionResultShape, valuesAccBaseVal, loc), + getInitTensor(rewriter, reductionResultShape, indicesAccBaseVal, loc)}; + + auto linalgOp = rewriter.create( + loc, adaptor.getOperands(), outputs, + SmallVector{adaptor.getAxis()}, + [&](OpBuilder &b, Location loc, ValueRange inputs) { + assert(inputs.size() == 4); + + auto tritonReduceBlock = op.getBody(); + IRMapping mapping; + mapping.map(tritonReduceBlock->getArguments(), inputs); + + for (auto &op : tritonReduceBlock->without_terminator()) { + b.clone(op, mapping); + } + + auto tritonYield = tritonReduceBlock->getTerminator(); + auto results = + llvm::map_to_vector(tritonYield->getOperands(), [&](Value val) { + return mapping.lookup(val); + }); + b.create(loc, results); + }); + + // before we rewrite the argmax reduce op, we know it has return value + // so addReduceWithIndexAttrIfNeeded won't fail + // but ignoring it will lead to compiling failure + if (llvm::succeeded(reduceWithIndexParams) && + reduceWithIndexParams->tieBreakType != TieBreakType::None) { + addReduceWithIndexAttr(*reduceWithIndexParams, rewriter, linalgOp); + } + + if (isScalarReduce) { + SmallVector reduceResults{ + rewriter.create( + loc, valueType, linalgOp.getResults()[0], ValueRange{}), + rewriter.create( + loc, indexType, linalgOp.getResults()[1], ValueRange{})}; + rewriter.replaceOp(op, reduceResults); + } else { + rewriter.replaceOp(op, linalgOp); + } + + return success(); + } +}; + +class ArgMinConverter : public ArgMinMaxBaseConverter { +public: + static LogicalResult matchComparisonResult(Value currValue, Value currIndex, + Value reduceValue, + Value reduceIndex, + mlir::Block::iterator &it, + Value &comparisonResult); + + static float getBaseReductionValue(); + + static int8_t getBaseReductionIntValue(); + static uint8_t getBaseReductionUIntValue(); + + ArgMinConverter(MLIRContext *context) : ArgMinMaxBaseConverter(context) {} +}; + +class ArgMaxConverter : public ArgMinMaxBaseConverter { +public: + static LogicalResult matchComparisonResult(Value currValue, Value currIndex, + Value reduceValue, + Value reduceIndex, + mlir::Block::iterator &it, + Value &comparisonResult); + + static float getBaseReductionValue(); + + static int8_t getBaseReductionIntValue(); + static uint8_t getBaseReductionUIntValue(); + + ArgMaxConverter(MLIRContext *context) : ArgMinMaxBaseConverter(context) {} +}; + +} // namespace TTOpConverters + +#endif diff --git a/compiler/include/dicp/TritonToLinalg/AscendNPUIRLegalizePass.h b/compiler/include/dicp/TritonToLinalg/AscendNPUIRLegalizePass.h new file mode 100644 index 00000000..98eb9cf0 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/AscendNPUIRLegalizePass.h @@ -0,0 +1,35 @@ + + +#ifndef TRITON_ADAPTER_ASCEND_NPU_IR_LEGALIZE_PASS_H +#define TRITON_ADAPTER_ASCEND_NPU_IR_LEGALIZE_PASS_H + +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Pass/Pass.h" + +#define GEN_PASS_DECL_ASCENDNPUIRLEGALIZE +#define GEN_PASS_DEF_ASCENDNPUIRLEGALIZE +#include "dicp/TritonToLinalg/Passes.h.inc" + +namespace mlir { +namespace triton { + +std::unique_ptr> createAscendNPUIRLegalizePass(); + +std::unique_ptr> +createAscendNPUIRLegalizePass(const AscendNPUIRLegalizeOptions &options); + +} // namespace triton +} // namespace mlir + +class AscendNPUIRLegalizePass + : public ::impl::AscendNPUIRLegalizeBase { +public: + AscendNPUIRLegalizePass() = default; + + explicit AscendNPUIRLegalizePass(const AscendNPUIRLegalizeOptions &options) + : AscendNPUIRLegalizeBase(options) {} + + void runOnOperation() override; +}; + +#endif // TRITON_ADAPTER_ASCEND_NPU_IR_LEGALIZE_PASS_H diff --git a/compiler/include/dicp/TritonToLinalg/BlockPtrAnalysis.h b/compiler/include/dicp/TritonToLinalg/BlockPtrAnalysis.h new file mode 100644 index 00000000..2e457fe9 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/BlockPtrAnalysis.h @@ -0,0 +1,298 @@ + + +#ifndef TRITON_ANALYSIS_BLOCKPTRANALYSIS_H +#define TRITON_ANALYSIS_BLOCKPTRANALYSIS_H +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" + +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Value.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" + +#include +namespace mlir { + +class ConversionPatternRewriter; + +namespace triton { + +enum class MemAccVal { Undefined = 0, StrucMemAcc = 1, UnstrucMemAcc = 2 }; + +struct MemAccType { + + MemAccVal value; + + explicit constexpr MemAccType(MemAccVal v = MemAccVal::Undefined) + : value(v) {} + + constexpr operator MemAccVal() const { return value; } + explicit operator bool() = delete; + + constexpr bool isUndefined() const { return value == MemAccVal::Undefined; } + constexpr bool isStructured() const { + return value == MemAccVal::StrucMemAcc; + } + constexpr bool isUnstructured() const { + return value == MemAccVal::UnstrucMemAcc; + } + + void merge(MemAccType &other) { + this->value = (this->value > other.value) ? this->value : other.value; + } + + std::string_view toString() const { + static constexpr std::string_view names[] = {"Undefined", "StrucMemAcc", + "UnstrucMemAcc"}; + return names[static_cast(value)]; + } +}; + +class BlockData { +public: + SmallVector &getOffsetsRef(); + SmallVector &getSizesRef(); + SmallVector &getStridesRef(); + Value &getSourceRef(); + OpFoldResult &getScalarRef(); + Type &getResElemTyRef(); + MemAccType &getMemAccTypeRef(); + + SmallVector getOffsets() const; + SmallVector getSizes() const; + SmallVector getStrides() const; + Type getResElemTy() const; + OpFoldResult getOffset(int) const; + OpFoldResult getSize(int) const; + OpFoldResult getStride(int) const; + OpFoldResult getScalar() const; + Value getSource() const; + MemAccType getMemAccType() const; + + bool isScalar() const; + bool isEmpty() const; + bool hasSource() const; + bool hasResElemTy() const; + void removeSource(); + + int64_t getRank() const; + MemRefType getResultMemrefType(int64_t offset, + ArrayRef resultShape) const; + + void addBlock(BlockData &lBlock, BlockData &rBlock, Location loc, + ConversionPatternRewriter &rewriter); + void subBlock(BlockData &lBlock, BlockData &rBlock, Location loc, + ConversionPatternRewriter &rewriter); + void mulBlock(BlockData &lBlock, BlockData &rBlock, Location loc, + ConversionPatternRewriter &rewriter); + void divBlock(BlockData &lBlock, BlockData &rBlock, Location loc, + ConversionPatternRewriter &rewriter); + + memref::ReinterpretCastOp createCastOp(ArrayRef resultShape, + const Location &loc, + OpBuilder &builder) const; + + void setResElemTy(const Type &); + void setSource(const Value &); + void setScalar(const OpFoldResult &); + void setOffsets(const SmallVector &); + void setStrides(const SmallVector &); + void setSizes(const SmallVector &); + void setMemAccTy(const MemAccType &); + void setMemAccVal(const MemAccVal); + + void dump() const; + +private: + SmallVector offsets; + SmallVector sizes; + SmallVector strides; + Value source; + // `Scalar` is a shortcut used when the entire blockdata describes a single + // scalar value + OpFoldResult scalar; + Type resElemTy; + MemAccType memAccTy; + + // Accumulate offsets of each dimension in BlockData to get a total offset + // from source ptr, which is used in memref::ReinterpretCastOp + OpFoldResult inferBlockOffset(const Location &loc, OpBuilder &builder) const; +}; + +class BlockDataParser { +public: + static Value getScalarMemRef(Value ptr, Value memref, const Location &loc, + ConversionPatternRewriter &rewriter); + + static void parse(Value operand, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseAdd(arith::AddIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseSub(arith::SubIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseMul(arith::MulIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseDiv(arith::DivSIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseRem(arith::RemSIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void + parseUnrealizedCast(UnrealizedConversionCastOp op, BlockData &data, + const Location &loc, ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void + parseMakeRange(triton::MakeRangeOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseLinalgGenericFromMakeRange( + linalg::GenericOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void + parseExpandDims(triton::ExpandDimsOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseBitcast(triton::BitcastOp op, BlockData &data, + const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseExtSI(arith::ExtSIOp op, BlockData &data, + const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void + parseBroadcast(triton::BroadcastOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseSplat(triton::SplatOp op, BlockData &data, + const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void + parseConstSplat(arith::ConstantOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + template + static std::enable_if_t || + std::is_same_v> + parseTensorPtr(T op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseAddPtr(triton::AddPtrOp op, BlockData &data, + const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void + parseExtractSlice(tensor::ExtractSliceOp op, BlockData &data, + const Location &loc, ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void + parseReinterpretCast(memref::ReinterpretCastOp op, BlockData &data, + const Location &loc, ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseReduce(triton::ReduceOp op, BlockData &data, + const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void + parseAtomicRmw(triton::AtomicRMWOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseFill(linalg::FillOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseSelect(arith::SelectOp op, BlockData &data, + const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + static void parseCustomOp(hivm::CustomOp op, BlockData &data, + const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known, + unsigned resultIdx); + + static void rewriteAddPtr(triton::AddPtrOp op, + triton::AddPtrOp::Adaptor &adaptor, + ConversionPatternRewriter &rewriter, + llvm::SmallDenseMap &known); + + static void + rewriteMakeTensorPtrOp(triton::MakeTensorPtrOp op, Value base, + ConversionPatternRewriter &rewriter, + llvm::SmallDenseMap &known); + + static void rewriteAdvanceOp(triton::AdvanceOp op, + ConversionPatternRewriter &rewriter, + llvm::SmallDenseMap &known); + + static void + rewriteCustomOp(hivm::CustomOp op, hivm::CustomOp::Adaptor &adaptor, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known); + + template + static std::enable_if_t || + std::is_same_v> + rewriteTerminator(T op, ConversionPatternRewriter &rewriter, + const llvm::SmallDenseSet &blockArgIdxSet, + ArrayRef iterArgIdxMap, + const llvm::SmallDenseMap &known); + + /// @param known is mainly designed for `rewriteLoop`, and is just non-const + /// in `rewriteLoop`, `rewriteAddPtr` and `rewriteAdvance` + static void rewriteLoopOp(LoopLikeOpInterface op, + ConversionPatternRewriter &rewriter, + llvm::SmallDenseMap &known); + + static void rewriteAddPtrToUnstrucMemAcc(triton::AddPtrOp op, + triton::AddPtrOp::Adaptor &adaptor, + ConversionPatternRewriter &rewriter, + BlockData &data); +}; + +template +void parseIndirectLoad(OpTy op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known, + unsigned resultIdx = 0); + +} // namespace triton + +} // namespace mlir + +#endif diff --git a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/CMakeLists.txt b/compiler/include/dicp/TritonToLinalg/CMakeLists.txt similarity index 56% rename from compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/CMakeLists.txt rename to compiler/include/dicp/TritonToLinalg/CMakeLists.txt index e25e3fbf..57914a55 100644 --- a/compiler/include/dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/CMakeLists.txt +++ b/compiler/include/dicp/TritonToLinalg/CMakeLists.txt @@ -1,3 +1,3 @@ set(LLVM_TARGET_DEFINITIONS Passes.td) -mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToLinalgNPUCoversion) -add_public_tablegen_target(TritonToLinalgNPUCoversionPassIncGen) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToLinalg) +add_public_tablegen_target(TritonToLinalgConversionPassIncGen) \ No newline at end of file diff --git a/compiler/include/dicp/TritonToLinalg/ConversionPatterns.h b/compiler/include/dicp/TritonToLinalg/ConversionPatterns.h new file mode 100644 index 00000000..b11c08ba --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/ConversionPatterns.h @@ -0,0 +1,113 @@ + + +#ifndef CONVERSIONPATTERNS_H +#define CONVERSIONPATTERNS_H + +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Passes.h" + +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/MathExtras.h" + +#include +#include +#include + +using namespace mlir; +using namespace triton; + +//===----------------------------------------------------------------------===// +// Utilities +//===----------------------------------------------------------------------===// + +static Value getScalarValue(Value operand, Location loc, + ConversionPatternRewriter &rewriter) { + SmallVector ops; + + auto reconstructScalarValue = [&](Value src) { + for (auto op = ops.rbegin(); op != ops.rend(); ++op) { + src = TypeSwitch(*op) + .Case([&](Operation *op) { + auto resType = op->getResults()[0].getType(); + if (auto shapedType = dyn_cast(resType)) { + resType = shapedType.getElementType(); + } + return rewriter.create(loc, resType, src); + }) + .Case([&](Operation *op) { + auto resType = op->getResults()[0].getType(); + if (auto shapedType = dyn_cast(resType)) { + resType = shapedType.getElementType(); + } + return rewriter.create(loc, resType, src); + }) + .Default([](Operation *op) { + llvm_unreachable("unsupported op in generating "); + return nullptr; + }); + } + return src; + }; + + while (true) { + if (!dyn_cast(operand.getType())) { + return reconstructScalarValue(operand); + } else if (auto op = operand.getDefiningOp()) { + if (auto attr = dyn_cast(op.getValue())) { + if (!attr.isSplat()) { + InFlightDiagnostic diag = emitError(loc) + << "other value used in masked load " + "produced by unsupported instruction"; + return nullptr; + } + auto elemValue = attr.getSplatValue(); + auto constOp = arith::ConstantOp::materialize( + rewriter, elemValue, attr.getElementType(), op.getLoc()); + return reconstructScalarValue(constOp.getResult()); + } + } else if (auto op = operand.getDefiningOp()) { + operand = op.getSrc(); + } else if (auto op = operand.getDefiningOp()) { + ops.push_back(op.getOperation()); + operand = op.getIn(); + } else if (auto op = operand.getDefiningOp()) { + ops.push_back(op.getOperation()); + operand = op.getIn(); + } else { + InFlightDiagnostic diag = emitError(loc) + << "other value used in masked load produced " + "by unsupported instruction"; + return nullptr; + } + } + return nullptr; +} + +static SmallVector getNParallelLoopsAttrs(unsigned n) { + return SmallVector(n, utils::IteratorType::parallel); +} + +// for IntLike and FloatLike types +static std::optional getBitWidth(Type a) { + if (auto type = dyn_cast(a)) { + auto elementType = type.getElementType(); + if (elementType.isIntOrFloat()) { + return type.getElementType().getIntOrFloatBitWidth(); + } + return std::nullopt; + } + + if (a.isIntOrFloat()) { + return a.getIntOrFloatBitWidth(); + } + return std::nullopt; +} +#endif // CONVERSIONPATTERNS_H diff --git a/compiler/include/dicp/TritonToLinalg/DescriptorConverter.h b/compiler/include/dicp/TritonToLinalg/DescriptorConverter.h new file mode 100644 index 00000000..18203ae5 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/DescriptorConverter.h @@ -0,0 +1,85 @@ + + +#ifndef TRITON_ADAPTER_DESCRIPTORCONVERTER_H +#define TRITON_ADAPTER_DESCRIPTORCONVERTER_H + +#include "dicp/TritonToLinalg/BlockPtrAnalysis.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/Interfaces/FunctionInterfaces.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Debug.h" + +namespace DescriptorConverter { +using namespace mlir; +using namespace triton; + +struct Descriptor { + Value base; + SmallVector shape; + SmallVector strides; + triton::PaddingOptionAttr padding; +}; + +bool hasATensorDescriptorType(mlir::TypeRange types); + +class DescriptorLoadConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::DescriptorLoadOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class DescriptorStoreConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::DescriptorStoreOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class DescriptorGatherConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::DescriptorGatherOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class DescriptorScatterConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::DescriptorScatterOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class DescriptorReduceConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::DescriptorReduceOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +} // end of namespace DescriptorConverter + +#endif // TRITON_ADAPTER_DESCRIPTORCONVERTER_H diff --git a/compiler/include/dicp/TritonToLinalg/FunctionConverter.h b/compiler/include/dicp/TritonToLinalg/FunctionConverter.h new file mode 100644 index 00000000..ee28aa01 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/FunctionConverter.h @@ -0,0 +1,40 @@ + + +#ifndef TRITON_ADAPTER_FUNCTIONCONVERTER_H +#define TRITON_ADAPTER_FUNCTIONCONVERTER_H + +#include "mlir/Interfaces/FunctionInterfaces.h" +#include "mlir/Transforms/DialectConversion.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +namespace FunctionConverter { +using namespace mlir; +using namespace triton; + +class GetProgramIDConverter + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + static uint32_t constexpr LAUNCH_GRID_RANK = + getMaxEnumValForProgramIDDim() + 1; + +public: + LogicalResult + matchAndRewrite(triton::GetProgramIdOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class GetNumProgramsConverter + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + static uint32_t constexpr LAUNCH_GRID_RANK = + getMaxEnumValForProgramIDDim() + 1; + +public: + LogicalResult + matchAndRewrite(triton::GetNumProgramsOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; +} // namespace FunctionConverter +#endif diff --git a/compiler/include/dicp/TritonToLinalg/HoistBroadcast.h b/compiler/include/dicp/TritonToLinalg/HoistBroadcast.h new file mode 100644 index 00000000..4bbd887e --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/HoistBroadcast.h @@ -0,0 +1,60 @@ + + +#ifndef TRITON_ADAPTER_TRITONTOLINALG_HOISTBROADCAST_H +#define TRITON_ADAPTER_TRITONTOLINALG_HOISTBROADCAST_H + +#include "dicp/TritonToLinalg/BlockPtrAnalysis.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/Interfaces/FunctionInterfaces.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "triton-to-linalg" + +namespace HoistBroadcast { +using namespace mlir; +using namespace triton; + +class BroadcastConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::BroadcastOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class BroadcastHoister { +public: + BroadcastHoister(triton::BroadcastOp op); + LogicalResult parse(Value operand, const Location &loc, + ConversionPatternRewriter &rewriter); + LogicalResult parseAddptr(triton::AddPtrOp op, const Location &loc, + ConversionPatternRewriter &rewriter); + LogicalResult parseBroadcast(triton::BroadcastOp op, const Location &loc, + ConversionPatternRewriter &rewriter); + LogicalResult parseSplat(triton::SplatOp op, const Location &loc, + ConversionPatternRewriter &rewriter); + LogicalResult findSrc(Value operand); + LogicalResult replaceBroadcastOp(triton::BroadcastOp op, + ConversionPatternRewriter &rewriter); + bool canBroadcast(); + +private: + Value source; + triton::BroadcastOp opToHoist; + SmallVector tensorSizes; + llvm::SmallDenseMap broadcastMap; +}; +} // namespace HoistBroadcast + +#endif // TRITON_ADAPTER_TRITONTOLINALG_HOISTBROADCAST_H diff --git a/compiler/include/dicp/TritonToLinalg/ImplicitPermute.h b/compiler/include/dicp/TritonToLinalg/ImplicitPermute.h new file mode 100644 index 00000000..b17f7201 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/ImplicitPermute.h @@ -0,0 +1,102 @@ + + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "dicp/TritonToStructured/MaskAnalysis.h" +#include "dicp/TritonToStructured/PtrAnalysis.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/AffineMap.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LogicalResult.h" +#include "mlir/Transforms/DialectConversion.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +namespace ImplicitPermute { + +using namespace mlir; +using namespace triton; + +class LoadConverter : public OpRewritePattern { +public: + explicit LoadConverter(MLIRContext *context) + : OpRewritePattern(context){}; + + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::LoadOp op, + PatternRewriter &rewriter) const override; +}; + +class StoreConverter : public OpRewritePattern { +public: + explicit StoreConverter(MLIRContext *context) + : OpRewritePattern(context){}; + + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::StoreOp op, + PatternRewriter &rewriter) const override; +}; + +class AtomicRMWConverter : public OpRewritePattern { +public: + explicit AtomicRMWConverter(MLIRContext *context) + : OpRewritePattern(context){}; + + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::AtomicRMWOp op, + PatternRewriter &rewriter) const override; +}; + +class AtomicCASConverter : public OpRewritePattern { +public: + explicit AtomicCASConverter(MLIRContext *context) + : OpRewritePattern(context){}; + + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::AtomicCASOp op, + PatternRewriter &rewriter) const override; +}; + +class MemOpTransformer { +public: + TritonToStructured::PtrState ptrState; + TritonToStructured::MaskState maskState; + + enum class MemType { load, store, deafaultType }; + + MemType currentType = MemType::deafaultType; + + MemOpTransformer(MemType memType) : currentType(memType) {} + + Value materializeImplicitPermute(Value srcTensor, const Location loc, + PatternRewriter &rewriter); + + Value createNewAddPtr(Value oldPtr, const Location loc, + PatternRewriter &rewriter); + + Value createNewAdvancePtr(Value oldPtr, const Location loc, + PatternRewriter &rewriter); + + Value createNewTensorPtr(Value oldPtr, const Location loc, + PatternRewriter &rewriter); + + Value createNewMask(Value oldPtr, const Location loc, + PatternRewriter &rewriter); + + Value createNewOther(Value oldOther, const Location loc, + PatternRewriter &rewriter); + + SmallVector + getBoundaryCheck(ArrayRef oldBoundaryCheck) const; + + bool applyPermuteOnMask(); +}; + +} // namespace ImplicitPermute diff --git a/compiler/include/dicp/Dialect/TritonExt/Transforms/CanonicalizerPattern.h b/compiler/include/dicp/TritonToLinalg/LoadStoreConverter.h similarity index 64% rename from compiler/include/dicp/Dialect/TritonExt/Transforms/CanonicalizerPattern.h rename to compiler/include/dicp/TritonToLinalg/LoadStoreConverter.h index b6967c2a..aebb5e3a 100644 --- a/compiler/include/dicp/Dialect/TritonExt/Transforms/CanonicalizerPattern.h +++ b/compiler/include/dicp/TritonToLinalg/LoadStoreConverter.h @@ -1,5 +1,7 @@ -#ifndef TRITON_DLC_LOADSTORECONVERTER_H -#define TRITON_DLC_LOADSTORECONVERTER_H + + +#ifndef TRITON_ADAPTER_LOADSTORECONVERTER_H +#define TRITON_ADAPTER_LOADSTORECONVERTER_H #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/IR/AffineMap.h" @@ -17,10 +19,48 @@ #include "triton/Dialect/Triton/IR/Dialect.h" +namespace LoadStoreConverter { + using namespace mlir; using namespace triton; -namespace mlir::dicp::trtion_ext { +class AddPtrConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::AddPtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class LoadConverter : public OpConversionPattern { +private: + void propagateWasBoolToInt8Attr(Operation *srcLoadOp, Operation *dstOp, + PatternRewriter &rewriter) const; + + LogicalResult toTensorAndReplace(triton::LoadOp &op, + RankedTensorType &tensorType, Value localMem, + bool mayImplicitTransposeWithLastAxis, + const Location &loc, + ConversionPatternRewriter &rewriter) const; + + LogicalResult checkModifiedByAddPtrConverter(triton::LoadOp &op) const; + + LogicalResult + continueModifyFromAddPtrConverter(triton::LoadOp &op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const; + + void + fillTensorWithOtherForMaskScenario(Value other, Value localMem, + ArrayRef maskDim, + ConversionPatternRewriter &rewriter) const; + +public: + explicit LoadConverter(MLIRContext *context); + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; // tempate class's impl must in header file template @@ -74,6 +114,17 @@ class ScalarStoreCanonicalizer : public OpRewritePattern { PatternRewriter &rewriter) const override; }; +class StoreConverter : public OpConversionPattern { +public: + explicit StoreConverter(MLIRContext *context); + + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::StoreOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + class ScalarAtomicRMWCanonicalizer : public OpRewritePattern { using OpRewritePattern::OpRewritePattern; @@ -135,6 +186,10 @@ class AtomicRMWConverter : public OpConversionPattern { } } else if (rmwOp == triton::RMWOp::XCHG) { binaryOp = rhs; + } else if (rmwOp == triton::RMWOp::UMAX) { + binaryOp = builder.create(loc, lhs, rhs); + } else if (rmwOp == triton::RMWOp::UMIN) { + binaryOp = builder.create(loc, lhs, rhs); } else { op.emitOpError("unsupported atomic RMW operation: "); llvm_unreachable( @@ -176,104 +231,15 @@ class AtomicMaxMinCanonicalizer : public OpRewritePattern { PatternRewriter &rewriter) const override; }; -class SelectCanonicalizer : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(arith::SelectOp op, - PatternRewriter &rewriter) const override; -}; - -/* - * Move tt.bitcast to a previous location if tt.bitcast is not directly applied - * on function arguments - */ -class BitcastCanonicalizer : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(triton::BitcastOp bitcastOp, - PatternRewriter &rewriter) const override; -}; - -template -class ScalarMathCanonicalizer : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(MathOp op, - PatternRewriter &rewriter) const override { - if (op->getNumResults() != 1) { - return rewriter.notifyMatchFailure( - op, "ScalarMathCanonicalizer expects single scalar output."); - } - if (!op->getResult(0).getType().isIntOrIndexOrFloat()) { - return rewriter.notifyMatchFailure( - op, "ScalarMathCanonicalizer handles scalar load scene."); - } - if (auto linalgOp = op->template getParentOfType()) { - return rewriter.notifyMatchFailure( - op, "ScalarMathCanonicalizer handles op not within tt.reduce."); - } - if (auto linalgOp = op->template getParentOfType()) { - return rewriter.notifyMatchFailure( - op, "ScalarMathCanonicalizer handles op not within tt.scan."); - } - auto loc = op.getLoc(); - llvm::SmallVector inputs; - for (auto input : op->getOperands()) { - auto blkTy = RankedTensorType::get({(int64_t)1}, input.getType()); - auto inputSplat = rewriter.create(loc, blkTy, input); - inputs.push_back(inputSplat.getResult()); - } - auto blkOp = rewriter.create(loc, inputs); - Value offset = - rewriter.create(loc, rewriter.getIndexAttr(0)); - auto extractOp = - rewriter.create(loc, blkOp.getResult(), offset); - rewriter.replaceOp(op, extractOp); - return success(); - } -}; - -/* - * Rewrite tt.make_tensor_ptr with non-contiguous order to - * tt.make_tensor_ptr + tt.load + tt.trans. - */ -class MakeTensorPtrCanonicalizer - : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(triton::MakeTensorPtrOp op, - PatternRewriter &rewriter) const override; -}; - -class ReduceSingleCanonicalizer : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(triton::ReduceOp reduceOp, - PatternRewriter &rewriter) const override; -}; - -/** - * @brief Rewrites arith.remf: remf(a, b) = a - b * floor(a / b) - */ -class RemfToBasicArithmetic final : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(arith::RemFOp op, - PatternRewriter &rewriter) const override; -}; - -/** - * @brief Rewrites arith.remsi: remsi(a, b) = a - b * (a / b) - */ -class RemSIToBasicArithmetic final : public OpRewritePattern { +class ReinterpretCastStrideCanonicalizer + : public OpRewritePattern { public: - using OpRewritePattern::OpRewritePattern; + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(arith::RemSIOp op, + LogicalResult matchAndRewrite(memref::ReinterpretCastOp op, PatternRewriter &rewriter) const override; + static bool hasFixableZeroStride(memref::ReinterpretCastOp op); }; -} // namespace mlir::dicp::trtion_ext +} // namespace LoadStoreConverter #endif diff --git a/compiler/include/dicp/TritonToLinalg/MarkTensorKindPass.h b/compiler/include/dicp/TritonToLinalg/MarkTensorKindPass.h new file mode 100644 index 00000000..43881b91 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/MarkTensorKindPass.h @@ -0,0 +1,36 @@ + + +#ifndef TRITON_ADAPTER_CONVERSION_MARKTENSORKINDPASS_H +#define TRITON_ADAPTER_CONVERSION_MARKTENSORKINDPASS_H + +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#define GEN_PASS_DEF_MARKTENSORKIND +#include "dicp/TritonToLinalg/Passes.h.inc" + +namespace mlir { +namespace triton { + +std::unique_ptr> createMarkTensorKindPass(); + +enum TensorKind { NONE = -1, INPUT = 0, OUTPUT = 1, INPUT_OUTPUT = 2 }; + +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace triton; + +class MarkTensorKindPass + : public ::impl::MarkTensorKindBase { +public: + MarkTensorKindPass() = default; + + void runOnOperation() override; +}; + +#endif // TRITON_ADAPTER_CONVERSION_MARKTENSORKINDPASS_H \ No newline at end of file diff --git a/compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/MaskAnalysis.h b/compiler/include/dicp/TritonToLinalg/MaskAnalysis.h similarity index 66% rename from compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/MaskAnalysis.h rename to compiler/include/dicp/TritonToLinalg/MaskAnalysis.h index 10ba119e..80b35a9d 100644 --- a/compiler/include/dicp/Conversion/DiscreteMaskAccessConversion/MaskAnalysis.h +++ b/compiler/include/dicp/TritonToLinalg/MaskAnalysis.h @@ -1,4 +1,5 @@ + #ifndef TRITON_ANALYSIS_MASKANALYSIS_H #define TRITON_ANALYSIS_MASKANALYSIS_H @@ -7,6 +8,7 @@ #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "triton/Dialect/Triton/IR/Dialect.h" +#include #include namespace mlir { @@ -14,9 +16,11 @@ namespace mlir { // this class helps build Operations class OpBuilder; -namespace dicp { +namespace triton { // use to decode the pattern in a mask used for load and store +enum class MaskPosition { Head, Tail, Middle, Unknown }; + class MaskState { public: OpFoldResult start; @@ -30,22 +34,70 @@ class MaskState { return dims.size(); } + MaskPosition getMaskPosition(llvm::ArrayRef &tensorShape) { + if (getRank() != tensorShape.size()) { + return MaskPosition::Unknown; + } + + bool isHead = true; + int dynIndex = -1; + + for (int i = 0; i < getRank(); ++i) { + auto offsetVal = mlir::getConstantIntValue(offsets[i]); + if (!offsetVal.has_value() || offsetVal.value() != 0) { + isHead = false; + if (dynIndex == -1) { + dynIndex = i; + } else { // temporarily support only one dyn dim + return MaskPosition::Unknown; + } + } + } + + if (isHead) { + return MaskPosition::Head; + } + + for (int i = 0; i < getRank(); ++i) { + auto dimVal = mlir::getConstantIntValue(dims[i]); + if (i == dynIndex) { + continue; + } + if (!dimVal.has_value() || dimVal.value() != tensorShape[i]) { + return MaskPosition::Unknown; + } + } + return MaskPosition::Middle; + } + bool isEmpty() const { return getRank() == 0 && !scalar && !start && !end; } bool isMask() const { return !start && !end && !scalar && dims.size() != 0 && offsets.size() != 0; } + bool isMemrefSubviewValid(Value source, OpBuilder &builder) const; + // parse value recursively LogicalResult parse(Value operand, const Location &loc, OpBuilder &builder); tensor::ExtractSliceOp getExtractSlice(Value source, const Location &loc, OpBuilder &builder) const; + tensor::ExtractSliceOp getExtractSlice(Value source, const Location &loc, + OpBuilder &builder, + SmallVector offsets, + SmallVector dims) const; + tensor::InsertSliceOp getInsertSlice(Value source, Value dest, const Location &loc, OpBuilder &builder) const; + tensor::InsertSliceOp getInsertSlice(Value source, Value dest, + const Location &loc, OpBuilder &builder, + SmallVector offsets, + SmallVector dims) const; + memref::SubViewOp getSubview(Value source, const Location &loc, OpBuilder &builder) const; @@ -70,6 +122,10 @@ class MaskState { LogicalResult minStates(const MaskState &lhsState, const MaskState &rhsState, const Location &loc, OpBuilder &builder); + OpFoldResult clampToNonNegativeIndex(const OpFoldResult value, + const Location &loc, + OpBuilder &builder) const; + // Helper functions to parse values to populate MaskState LogicalResult parseConstant(arith::ConstantOp constOp, const Location &loc, @@ -112,12 +168,18 @@ class MaskState { LogicalResult parseSplat(triton::SplatOp splatOp, const Location &loc, OpBuilder &builder); + // Operand is the result of tensor.insert + LogicalResult parseInsert(tensor::InsertOp insertOp, const Location &loc, + OpBuilder &builder); + // Operand is the result of expand_dims LogicalResult parseExpandDims(triton::ExpandDimsOp expandDimsOp, const Location &loc, OpBuilder &builder); }; -} // namespace dicp +std::optional runMaskAnalysis(Operation *op, OpBuilder &builder); + +} // namespace triton } // namespace mlir diff --git a/compiler/include/dicp/TritonToLinalg/Passes.h b/compiler/include/dicp/TritonToLinalg/Passes.h new file mode 100644 index 00000000..a1dc6e1f --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/Passes.h @@ -0,0 +1,17 @@ + + +#ifndef TRITON_ADAPTER_TRITON_TO_LINALG_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_LINALG_CONVERSION_PASSES_H + +#include "dicp/TritonToLinalg/AscendNPUIRLegalizePass.h" +#include "dicp/TritonToLinalg/MarkTensorKindPass.h" +#include "dicp/TritonToLinalg/TritonToLinalgPass.h" + +namespace mlir::triton { + +#define GEN_PASS_REGISTRATION +#include "dicp/TritonToLinalg/Passes.h.inc" + +} // namespace mlir::triton + +#endif // TRITON_ADAPTER_TRITON_TO_LINALG_CONVERSION_PASSES_H diff --git a/compiler/include/dicp/TritonToLinalg/Passes.td b/compiler/include/dicp/TritonToLinalg/Passes.td new file mode 100644 index 00000000..9384ff40 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/Passes.td @@ -0,0 +1,44 @@ +#ifndef TRITON_TO_LINALG_CONVERSION_PASSES +#define TRITON_TO_LINALG_CONVERSION_PASSES + +include "mlir/Pass/PassBase.td" + +def TritonToLinalg : Pass<"triton-to-linalg", "mlir::ModuleOp"> { + let summary = "Convert Triton to Linalg dialect"; + let constructor = "triton::createTritonToLinalgPass()"; + let options = [ + Option<"globalKernel", "global-kernel", + "bool", /*default*/"true", + "generate a global kernel">, + Option<"namedOps", "named-ops", + "bool", /*default*/"false", + "use linalg named ops instead of linalg.generic">, + Option<"enableNd2nzOnVector", "enable-nd2nz-on-vector", + "bool", /*default*/"false", + "enable nd2nz on vector">, + Option<"enableSelectAnalysis", "enable-select-analysis", + "bool", /*default*/"true", + "enable select analysis">, + Option<"compileOn91095", "compile-on-910-95", + "bool", /*default*/"false", + "compile on 910_95"> + ]; +} + +def MarkTensorKind : Pass<"mark-tensor-kind", "mlir::ModuleOp"> { + let summary = "Mark tensor kind (INPUT/OUTPUT/INPUT_OUTPUT) on Triton function arguments"; + let constructor = "triton::createMarkTensorKindPass()"; +} + +def AscendNPUIRLegalize : Pass<"ascend-npu-ir-legalize", "mlir::ModuleOp"> { + let summary = "Legalize Ascend NPU IR: reify arith.mului_extended (i32) " + "as arith.mulsi_extended with sign-bit correction"; + let constructor = "triton::createAscendNPUIRLegalizePass()"; + let options = [ + Option<"unsafeMode", "unsafe-mode", + "bool", /*default*/"false", + "Skip sign-bit correction (unsafe but faster)"> + ]; +} + +#endif // TRITON_TO_LINALG_CONVERSION_PASSES diff --git a/compiler/include/dicp/TritonToLinalg/TritonOpConverter.h b/compiler/include/dicp/TritonToLinalg/TritonOpConverter.h new file mode 100644 index 00000000..c8c073b3 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/TritonOpConverter.h @@ -0,0 +1,723 @@ + + +#ifndef TRITON_ADAPTER_TRITONOPCONVERTER_H +#define TRITON_ADAPTER_TRITONOPCONVERTER_H + +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" +#include "dicp/TritonToLinalg/BlockPtrAnalysis.h" +#include "dicp/TritonToLinalg/TritonToLinalgPass.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/Interfaces/FunctionInterfaces.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "triton-to-linalg" + +namespace TTOpConverters { +using namespace mlir; +using namespace triton; + +static constexpr unsigned kFuncNameCap = 128; + +/* +Convert `tt.precise_div` operation to `arith.divf` operation. +tensor_x / tensor_y + +```ttir + %11 = tt.precise_divf %7, %10 : tensor<100xf32> +``` + +converts to: + +```mlir + %11 = arith.divf %7, %10 : tensor<100xf32> +``` +*/ +struct PreciseDivConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::PreciseDivFOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +/* +Convert `tt.fp_to_fp` operation with RTNE (default) rounding mode to +`arith.truncf` or `arith.extf` operation. + +For fp8 conversions with default RTNE rounding: +- downcast: tt.fp_to_fp -> arith.truncf +- upcast: tt.fp_to_fp -> arith.extf + +Note: Non-RTNE rounding modes (e.g., RTZ) are handled by TritonToHFusion pass. +*/ +struct FpToFpCanonicalizer : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(triton::FpToFpOp op, + PatternRewriter &rewriter) const override; +}; + +class SelectCanonicalizer : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(arith::SelectOp op, + PatternRewriter &rewriter) const override; +}; + +/* + * Move tt.bitcast to a previous location if tt.bitcast is not directly applied + * on function arguments + */ +class BitcastCanonicalizer : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(triton::BitcastOp bitcastOp, + PatternRewriter &rewriter) const override; +}; + +template +class ScalarMathCanonicalizer : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(MathOp op, + PatternRewriter &rewriter) const override { + if (op->getNumResults() != 1) { + return rewriter.notifyMatchFailure( + op, "ScalarMathCanonicalizer expects single scalar output."); + } + if (!op->getResult(0).getType().isIntOrIndexOrFloat()) { + return rewriter.notifyMatchFailure( + op, "ScalarMathCanonicalizer handles scalar load scene."); + } + if (auto linalgOp = op->template getParentOfType()) { + return rewriter.notifyMatchFailure( + op, "ScalarMathCanonicalizer handles op not within tt.reduce."); + } + if (auto linalgOp = op->template getParentOfType()) { + return rewriter.notifyMatchFailure( + op, "ScalarMathCanonicalizer handles op not within tt.scan."); + } + auto loc = op.getLoc(); + llvm::SmallVector inputs; + for (auto input : op->getOperands()) { + auto blkTy = RankedTensorType::get({(int64_t)1}, input.getType()); + auto inputSplat = rewriter.create(loc, blkTy, input); + inputs.push_back(inputSplat.getResult()); + } + auto blkOp = rewriter.create(loc, inputs); + Value offset = + rewriter.create(loc, rewriter.getIndexAttr(0)); + auto extractOp = + rewriter.create(loc, blkOp.getResult(), offset); + rewriter.replaceOp(op, extractOp); + return success(); + } +}; + +/* + * Rewrite tt.make_tensor_ptr with non-contiguous order to + * tt.make_tensor_ptr + tt.load + tt.trans. + */ +class MakeTensorPtrCanonicalizer + : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(triton::MakeTensorPtrOp op, + PatternRewriter &rewriter) const override; +}; + +class ReduceSingleCanonicalizer : public OpRewritePattern { +public: + using OpRewritePattern::OpRewritePattern; + LogicalResult matchAndRewrite(triton::ReduceOp reduceOp, + PatternRewriter &rewriter) const override; +}; + +class DenseConstantConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(arith::ConstantOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class MakeRangeConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::MakeRangeOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class SplatConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::SplatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class UnsplatConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::UnsplatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class ReshapeConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::ReshapeOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class ExpandDimsConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::ExpandDimsOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class ClampFConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::ClampFOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class BroadcastConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::BroadcastOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +template +class ReductionOpBaseConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(OpTy op, typename OpTy::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const final { + auto sourceType = + cast(adaptor.getOperands().front().getType()); + assert(sourceType.hasRank() && "Expected input is ranked"); + + int64_t axis = op.getAxis(); + assert(axis >= 0 && axis < sourceType.getRank() && + "Expected reduction axis is within operand's rank"); + + auto realReductionOps = this->getRealReductionOps(op); + if (realReductionOps.size() == 1) { + return this->convertToTargetOp(op, adaptor, rewriter); + } + return this->convertToTargetOpExtended(op, adaptor, rewriter); + } + +protected: + llvm::SmallVector getReductionOps(OpTy reductionOp) const { + auto reductionBody = reductionOp.getBody(); + return llvm::map_to_vector(reductionBody->without_terminator(), + [](Operation &op) { return &op; }); + } + + llvm::SmallVector getRealReductionOps(OpTy reductionOp) const { + llvm::SmallVector realOps; + for (Operation &bodyOp : reductionOp.getBody()->without_terminator()) { + // Skips non-reduce operations, including type conversion operations (this + // can be extended as needed). + if (isa(&bodyOp)) + continue; + realOps.push_back(&bodyOp); + } + return realOps; + } + + arith::ConstantOp + getMultiOpReductionBaseConstOp(ConversionPatternRewriter &rewriter, OpTy op, + Location loc, Type constantType) const { + // for multiop reduce of 1 element result is defined as exactly this element + // for multiop reduce of tensor with N elements default value + // is not involved in result computation + auto reductionOps = this->getReductionOps(op); + assert(reductionOps.size() == 1); + auto reductionOp = reductionOps.front(); + + assert(constantType.isIntOrFloat()); + + if (constantType.isInteger()) { + return rewriter.create( + loc, constantType, rewriter.getIntegerAttr(constantType, 0)); + } + return rewriter.create( + loc, constantType, rewriter.getFloatAttr(constantType, 0.f)); + } + + arith::ConstantOp getReductionBaseConstOp(ConversionPatternRewriter &rewriter, + Operation *reductionOp, + Type constantType) const { + const int64_t bitWidth = constantType.getIntOrFloatBitWidth(); + + auto attr = + llvm::TypeSwitch(reductionOp) + .Case([&](arith::AddFOp) { + return rewriter.getFloatAttr(constantType, 0.f); + }) + .Case([&](arith::AddIOp) { + return rewriter.getIntegerAttr(constantType, 0); + }) + .Case([&](arith::MulFOp) { + return rewriter.getFloatAttr(constantType, 1.f); + }) + .Case([&](arith::MulIOp) { + return rewriter.getIntegerAttr(constantType, 1); + }) + .template Case([&](auto) { + return rewriter.getFloatAttr( + constantType, -std::numeric_limits::infinity()); + }) + .template Case([&](auto) { + return rewriter.getFloatAttr( + constantType, std::numeric_limits::infinity()); + }) + .Case([&](arith::MinSIOp) { + return rewriter.getIntegerAttr(constantType, + llvm::maxIntN(bitWidth)); + }) + .Case([&](arith::MinUIOp) { + return rewriter.getIntegerAttr(constantType, + llvm::maxUIntN(bitWidth)); + }) + .Case([&](arith::MaxSIOp) { + return rewriter.getIntegerAttr(constantType, + llvm::minIntN(bitWidth)); + }) + .Case([&](arith::MaxUIOp) { + return rewriter.getIntegerAttr(constantType, 0); + }) + .Case([&](arith::OrIOp) { + return rewriter.getIntegerAttr(constantType, 0); + }) + .Case([&](arith::AndIOp) { + return rewriter.getIntegerAttr(constantType, 1); + }) + .Case([&](arith::XOrIOp) { + return rewriter.getIntegerAttr(constantType, 0); + }) + .Default([](Operation *op) { + op->dump(); + llvm_unreachable("Reduction op not supported yet"); + return nullptr; + }); + + return rewriter.create(reductionOp->getLoc(), + constantType, attr); + } + + bool requiresF32Conversion(const Type elemType, + Operation *reductionOp) const { + return isa(elemType) && + elemType.getIntOrFloatBitWidth() < + Float32Type::get(elemType.getContext()) + .getIntOrFloatBitWidth() && + (isa(reductionOp) || isa(reductionOp)); + } + + Value getReductionElement(Value lhs, Value rhs, const Location loc, + Operation *reductionOp, OpBuilder &b, + const bool convertLhsToF32Precision) const { + return llvm::TypeSwitch(reductionOp) + .template Case([&](auto reductionOp) { + if (convertLhsToF32Precision) { + lhs = b.create(loc, Float32Type::get(b.getContext()), + lhs); + } + return b.create(loc, lhs, rhs); + }) + .template Case([&](auto reductionOp) { + return b.create(loc, lhs, rhs); + }) + .Default([](Operation *op) { + op->dump(); + llvm_unreachable("Reduction op not yet supported"); + return nullptr; + }); + } + + virtual bool isReductionOpSupported(Operation *reductionOp) const = 0; + + virtual LogicalResult + convertToTargetOp(OpTy op, typename OpTy::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const = 0; + + virtual LogicalResult + convertToTargetOpExtended(OpTy op, typename OpTy::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const = 0; +}; + +class ReduceConverter : public ReductionOpBaseConverter { +public: + explicit ReduceConverter(MLIRContext *context) + : ReductionOpBaseConverter(context) {} + + using ReductionOpBaseConverter::ReductionOpBaseConverter; + +protected: + bool isReductionOpSupported(Operation *reductionOp) const override; + + static bool isMultiReductionOpSupported(Operation *reductionOp); + + Value cloneReduceOps(OpBuilder &builder, Value in, Value out, Value opIns, + Value opOuts, triton::ReduceOp op) const; + + void + checkIsNotCallOp(const llvm::SmallVector &reductionOps) const; + + bool isSCFOpReduce(const llvm::SmallVector &reductionOps) const; + + bool + isMultiOpReduce(const llvm::SmallVector &reductionOps) const; + + Value computeReduceResultWithCompileFlag( + OpBuilder &opBuilder, Location loc, Value lhs, Value rhs, Value source, + Value initTensor, triton::ReduceOp reductionOp, + bool compileOn91095Flag = false) const; + + LogicalResult + convertToTargetOp(triton::ReduceOp op, + typename triton::ReduceOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override; + + LogicalResult + convertToTargetOpExtended(triton::ReduceOp op, + typename triton::ReduceOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class ScanConverter : public ReductionOpBaseConverter { +public: + explicit ScanConverter(MLIRContext *context) + : ReductionOpBaseConverter(context) {} + + using ReductionOpBaseConverter::ReductionOpBaseConverter; + +protected: + bool isReductionOpSupported(Operation *reductionOp) const override; + + LogicalResult + convertToTargetOp(triton::ScanOp op, typename triton::ScanOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override; + + LogicalResult + convertToTargetOpExtended(triton::ScanOp op, + typename triton::ScanOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class ExternElementwiseClOpConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::ExternElementwiseOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class UnrealizedCastConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(UnrealizedConversionCastOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class JoinConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::JoinOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class SplitConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::SplitOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class CatConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::CatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class GatherConverter : public OpConversionPattern { +private: + static constexpr llvm::StringRef gatherFuncNameBase = "triton_gather"; + +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::GatherOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class YieldConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(scf::YieldOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +template || + std::is_same_v>> +class LoopConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(LoopOpTy op, + typename OpConversionPattern::OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + llvm::SmallDenseMap known; + + op->removeAttr("UnhandledLoopOp"); + BlockDataParser::rewriteLoopOp(op, rewriter, known); + return success(); + } +}; + +class AdvanceConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::AdvanceOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class MakeTensorPtrConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + explicit MakeTensorPtrConverter(MLIRContext *context) + : OpConversionPattern(context) {} + + LogicalResult + matchAndRewrite(triton::MakeTensorPtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class TransposeConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::TransOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class BitcastConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::BitcastOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class TritonMulhiuiConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::MulhiUIOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class TritonPreciseSqrtConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::PreciseSqrtOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class DeviceAssertConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + +private: + static constexpr llvm::StringRef printFuncNameBase = "triton_assert"; + static constexpr llvm::StringRef msgAttrName = "msg"; + +public: + LogicalResult + matchAndRewrite(triton::AssertOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class DevicePrintConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + +private: + static constexpr llvm::StringRef printFuncNameBase = "triton_print"; + static constexpr llvm::StringRef prefixAttrName = "prefix"; + static constexpr llvm::StringRef hexAttrName = "hex"; + +public: + LogicalResult + matchAndRewrite(triton::PrintOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +struct MatmulConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::DotOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +struct FlipOpConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::dicp::FlipOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; + + static constexpr StringRef baseFuncName = "triton_flip"; +}; + +struct SortOpConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::dicp::SortOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +struct DotScaledConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::DotScaledOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class PtrToIntConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::PtrToIntOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +class IndexPutConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::dicp::IndexPutOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; + +private: + static constexpr llvm::StringRef funcNameBase = "triton_index_put"; +}; + +class GatherOutToUbConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::dicp::GatherOutToUbOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; + +private: + static constexpr llvm::StringRef funcNameBase = "triton__gather_out_to_ub"; +}; + +class ScatterUbToOutConverter + : public OpConversionPattern { +public: + using OpConversionPattern< + triton::dicp::ScatterUbToOutOp>::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::dicp::ScatterUbToOutOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; + +private: + static constexpr llvm::StringRef funcNameBase = "triton_scatter_ub_to_out"; +}; + +class IndirectLoadConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::dicp::IndirectLoadOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; + +private: + static constexpr llvm::StringRef funcNameBase = "triton_indirect_load"; +}; + +class IndirectStoreConverter + : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(triton::dicp::IndirectStoreOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; + +private: + static constexpr llvm::StringRef funcNameBase = "triton_indirect_store"; +}; + +class IndexSelectSimdConverter + : public OpConversionPattern { +public: + explicit IndexSelectSimdConverter(MLIRContext *context); + using OpConversionPattern< + triton::dicp::IndexSelectSimdOp>::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::dicp::IndexSelectSimdOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override; +}; + +} // end of namespace TTOpConverters + +#endif diff --git a/compiler/include/dicp/TritonToLinalg/TritonToLinalgPass.h b/compiler/include/dicp/TritonToLinalg/TritonToLinalgPass.h new file mode 100644 index 00000000..dd1f9f05 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/TritonToLinalgPass.h @@ -0,0 +1,93 @@ + + +#ifndef TRITON_ADAPTER_CONVERSION_TRITONTOLINALG_H +#define TRITON_ADAPTER_CONVERSION_TRITONTOLINALG_H + +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#define GEN_PASS_CLASSES +#include "dicp/TritonToLinalg/Passes.h.inc" + +extern int nd2nzFlag; +extern bool compileOn91095Flag; +extern bool existDotFlag; + +namespace mlir { +namespace triton { + +std::unique_ptr> createTritonToLinalgPass(); + +std::unique_ptr> +createTritonToLinalgPass(bool, bool, bool, bool, bool); + +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace triton; +const std::string globalKernelAttr = "global_kernel"; +const std::string kernelMixModeName = "mix_mode"; +const std::string kernelParallelModeName = "parallel_mode"; + +class TritonTypeConverter : public mlir::TypeConverter { +public: + explicit TritonTypeConverter(); +}; + +class TritonToLinalgPass : public TritonToLinalgBase { + + static auto constexpr LAUNCH_GRID_RANK = getMaxEnumValForProgramIDDim() + 1; + static unsigned int constexpr TRITON_PROGRAM_INFO_ARG_COUNT = + LAUNCH_GRID_RANK * 2; + +private: + // grid构造 num_programs 3维, program_id 3维 + // remember 'xxxOp' is usually a Pointer, so that we can change target memory + // without giving a reference argument + void addProgramInfo(triton::FuncOp func, bool globalKernel); + + void convertTTFunc(triton::FuncOp func, const bool existDot, + const bool existSIMTOp); + + LogicalResult convertMultipleBlockControlFlow(Operation *funcOp, + OpBuilder &builder); + // 处理嵌套的if/else + scf::IfOp transformNestedIfElse(Operation &nestedBranch, OpBuilder &builder); + + void addDynamicLegal(ConversionTarget &target, + TritonTypeConverter &tritonTypeConverter); + + void + populateTritonToLinalgCanonicalizationPatterns(RewritePatternSet &patterns); + + void populateTritonToLinalgConversionPatterns(TypeConverter &typeConverter, + RewritePatternSet &patterns, + unsigned int launchGridRank); + + LogicalResult processDescriptorOperations(ModuleOp moduleOp); + LogicalResult processPtrBroadcastOperations(ModuleOp moduleOp); + LogicalResult processImplicitPermuteOperations(ModuleOp moduleOp); + LogicalResult processLegalStrideOperations(ModuleOp moduleOp); + +public: + TritonToLinalgPass() = default; + + TritonToLinalgPass(bool globalKernel, bool namedOps, bool enableNd2nzOnVector, + bool enableSelectAnalysis, bool compileOn91095) { + this->globalKernel = globalKernel; + this->namedOps = namedOps; + this->enableNd2nzOnVector = enableNd2nzOnVector; + this->enableSelectAnalysis = enableSelectAnalysis; + this->compileOn91095 = compileOn91095; + }; + + void getDependentDialects(DialectRegistry ®istry) const override; + + void runOnOperation() override; +}; + +#endif // TRITON_ADAPTER_CONVERSION_TRITONTOLINALG_H diff --git a/compiler/include/dicp/TritonToLinalg/UseAnalysis.h b/compiler/include/dicp/TritonToLinalg/UseAnalysis.h new file mode 100644 index 00000000..22f9ed40 --- /dev/null +++ b/compiler/include/dicp/TritonToLinalg/UseAnalysis.h @@ -0,0 +1,124 @@ + + +#ifndef TRITON_ANALYSIS_USEANALYSIS_H +#define TRITON_ANALYSIS_USEANALYSIS_H + +#include "mlir/Analysis/DataFlow/SparseAnalysis.h" + +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" +#include "mlir/Transforms/DialectConversion.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +namespace mlir { +namespace triton { + +enum class UseType { + Undefined, // Initial state + DataUse, // value used for tensor computation only + MetaUse, // value used for metadata only + MixUse // value used for both tensor computation and metadata +}; + +struct UseInfo : public dataflow::AbstractSparseLattice { + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UseInfo) + using AbstractSparseLattice::AbstractSparseLattice; + + // Lattice state transfer function + ChangeResult meetUseType(const UseType &other) { + if (other == UseType::Undefined) { + return ChangeResult::NoChange; + } + + switch (type) { + case UseType::Undefined: + type = other; + return ChangeResult::Change; + case UseType::DataUse: + case UseType::MetaUse: + if (type == other) { + return ChangeResult::NoChange; + } else { + type = UseType::MixUse; + return ChangeResult::Change; + } + case UseType::MixUse: + return ChangeResult::NoChange; + default: + llvm_unreachable("bad type"); + } + } + + ChangeResult meet(const AbstractSparseLattice &other) override { + auto rhs = reinterpret_cast(&other); + return meetUseType(rhs->type); + } + + void print(raw_ostream &os) const override { + switch (type) { + case UseType::DataUse: + os << "DataUse"; + break; + case UseType::MetaUse: + os << "MetaUse"; + break; + case UseType::MixUse: + os << "MixUse"; + break; + default: + os << "Undefined"; + } + } + + UseType type = UseType::Undefined; +}; + +class UseAnalysis : public dataflow::SparseBackwardDataFlowAnalysis { +public: + using SparseBackwardDataFlowAnalysis::SparseBackwardDataFlowAnalysis; + +#if LLVM_VERSION_MAJOR >= 20 + LogicalResult visitOperation(Operation *op, ArrayRef operands, + ArrayRef results) override; +#else + void visitOperation(Operation *op, ArrayRef operands, + ArrayRef results) override; +#endif + + void visitBranchOperand(OpOperand &operand) override { return; } + + void visitCallOperand(OpOperand &operand) override { return; } + + void setToExitState(UseInfo *lattice) override { + lattice->type = UseType::Undefined; + } + +private: + void propagateUse(UseInfo *lattice, const UseType &type) { + auto changed = lattice->meetUseType(type); + propagateIfChanged(lattice, changed); + } + + void propagateResults(UseInfo *lattice, ArrayRef results) { + auto changed = ChangeResult::NoChange; + for (auto result : results) { + changed |= lattice->meet(*result); + } + propagateIfChanged(lattice, changed); + } +}; + +class MetaUseEraser : public RewritePattern { +public: + MetaUseEraser(MLIRContext *context); + + LogicalResult matchAndRewrite(Operation *op, + PatternRewriter &rewriter) const final; +}; + +LogicalResult runUseAnalysis(triton::FuncOp &funcOp); + +} // namespace triton + +} // namespace mlir + +#endif // TRITON_CONVERSION_TRITONTOAFFINE_TRITONUSEANALYSIS_H diff --git a/compiler/include/dicp/TritonToStructured/CMakeLists.txt b/compiler/include/dicp/TritonToStructured/CMakeLists.txt new file mode 100644 index 00000000..b685f4e7 --- /dev/null +++ b/compiler/include/dicp/TritonToStructured/CMakeLists.txt @@ -0,0 +1,3 @@ +set(LLVM_TARGET_DEFINITIONS Passes.td) +mlir_tablegen(Passes.h.inc -gen-pass-decls --name TritonToStructured) +add_public_tablegen_target(TritonToStructuredConversionPassIncGen) \ No newline at end of file diff --git a/compiler/include/dicp/TritonToStructured/CannonicalizerConverter.h b/compiler/include/dicp/TritonToStructured/CannonicalizerConverter.h new file mode 100644 index 00000000..5b22559b --- /dev/null +++ b/compiler/include/dicp/TritonToStructured/CannonicalizerConverter.h @@ -0,0 +1,338 @@ + +#ifndef TRITON_ADAPTER_CANNONICALIZERCONVERTER_H +#define TRITON_ADAPTER_CANNONICALIZERCONVERTER_H + +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/AffineMap.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LogicalResult.h" +#include "mlir/Transforms/DialectConversion.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +namespace CannonicalizerConverter { + +using namespace mlir; +using namespace triton; + +class CmpConverter : public OpRewritePattern { +public: + explicit CmpConverter(MLIRContext *context) + : OpRewritePattern(context) {} + + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(arith::CmpIOp cmpOp, + PatternRewriter &rewriter) const override; +}; + +class SplatCmpConverter : public OpRewritePattern { +public: + explicit SplatCmpConverter(MLIRContext *context) + : OpRewritePattern(context) {} + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(arith::CmpIOp op, + PatternRewriter &rewriter) const override; +}; + +class AddPtrSplatConverter : public OpRewritePattern { +public: + explicit AddPtrSplatConverter(MLIRContext *context) + : OpRewritePattern(context) {} + + LogicalResult matchAndRewrite(triton::AddPtrOp op, + PatternRewriter &rewriter) const override; +}; + +class LoadBroadcastConverter : public OpRewritePattern { +public: + explicit LoadBroadcastConverter(MLIRContext *context) + : OpRewritePattern(context) {} + + LogicalResult matchAndRewrite(triton::LoadOp op, + PatternRewriter &rewriter) const override; +}; + +class IfYieldAddHoistConverter : public OpRewritePattern { +public: + explicit IfYieldAddHoistConverter(MLIRContext *context) + : OpRewritePattern(context) {} + + LogicalResult matchAndRewrite(scf::IfOp ifOp, + PatternRewriter &rewriter) const override; + +private: + bool isSupportedTensorResultType(Type type) const; + + bool isDefinedOutsideIf(Value value, scf::IfOp ifOp) const; + + bool extractAddendFromAddExpr(Value maybeAddExpr, Value baseValue, + Value &addendOut) const; + + Value buildZeroTensorLikeType(Type laneType, Location loc, + PatternRewriter &rewriter) const; + + bool tryRewriteSingleLane(unsigned laneIdx, Value baseBranchYield, + Value addExprBranchYield, bool baseInThenBranch, + Type laneType, scf::IfOp ifOp, + PatternRewriter &rewriter, + SmallVectorImpl &updatedThenYieldOperands, + SmallVectorImpl &updatedElseYieldOperands, + SmallVectorImpl &hoistedBasePerLane, + SmallVectorImpl &laneRewrittenFlags) const; +}; + +class PromotePointerIterArgsPattern : public OpRewritePattern { +public: + explicit PromotePointerIterArgsPattern(MLIRContext *context) + : OpRewritePattern(context) {} + + LogicalResult matchAndRewrite(scf::ForOp forOp, + PatternRewriter &rewriter) const override; + +private: + // Information about a pointer iteration argument to be promoted + struct PointerArgInfo { + unsigned oldIndex; // Original index in the iteration arguments + Value basePointer; // Base pointer value passed as init arg + Value offsetValue; // Offset value used in addptr operation + Value newIterArg; // New integer iteration argument + Value addPtrValue; // The addptr operation result that updates the pointer + // Offset value used in advancePtr operation (with explicit inlined + // capacity) + SmallVector offsetValues; + SmallVector newInitArgs; + SmallVector newIterArgTypes; + }; + + LogicalResult matchAndRewriteForAddPtr(scf::ForOp forOp, + PatternRewriter &rewriter) const; + + // Check if the loop meets basic transformation conditions + LogicalResult matchLoop(scf::ForOp forOp) const; + + // Collect all pointer iteration arguments that match the promotion pattern + SmallVector collectPointerIterArgs(scf::ForOp forOp) const; + + // Check if a value has pointer tensor type + bool isPointerIterArg(Value iterArg) const; + + // Analyze a pointer iteration argument to determine if it matches the + // promotion pattern + std::optional analyzePointerIterArg(Value iterArg, + Block &loopBody) const; + + // Check if an index corresponds to a pointer argument being promoted + bool isPointerArgIndex(ArrayRef pointerArgs, + unsigned idx) const; + + // Get pointer argument information for a specific index + const PointerArgInfo *getPointerArgInfo(ArrayRef pointerArgs, + unsigned idx) const; + + // Create a new for loop with updated iteration argument types + scf::ForOp createNewForLoop(scf::ForOp forOp, ArrayRef newInitArgs, + ArrayRef newIterArgTypes, + PatternRewriter &rewriter) const; + + // Rewrite the loop body to use integer iteration arguments instead of + // pointers + LogicalResult rewriteLoopBody(scf::ForOp oldForOp, scf::ForOp newForOp, + SmallVector &pointerArgs, + DenseMap &indexMap, + PatternRewriter &rewriter) const; + + // Create new iteration arguments by replacing pointers with integer offsets + std::tuple, SmallVector, + DenseMap> + createNewIterArgs(scf::ForOp forOp, ArrayRef pointerArgs, + PatternRewriter &rewriter) const; + + // Create IR mapping for cloning operations, rebuilding pointers from integer + // offsets + IRMapping createIRMapping(scf::ForOp oldForOp, scf::ForOp newForOp, + SmallVector &pointerArgs, + DenseMap &indexMap, + PatternRewriter &rewriter) const; + + // Reconstruct a pointer value from base pointer and integer offset + Value rebuildPointer(scf::ForOp forOp, ArrayRef pointerArgs, + unsigned idx, PatternRewriter &rewriter) const; + + // Clone instructions from old loop body to new loop body, skipping + // transformed addptr ops + LogicalResult cloneInstructions(Block &oldBody, Block &newBody, + ArrayRef pointerArgs, + DenseMap &indexMap, + IRMapping &mapping, + PatternRewriter &rewriter) const; + + // Clone and transform the yield operation, converting pointer updates to + // integer additions + LogicalResult cloneYieldOp(scf::YieldOp yieldOp, + ArrayRef pointerArgs, + DenseMap &indexMap, + IRMapping &mapping, + PatternRewriter &rewriter) const; + + // Create integer addition for pointer offset updates in the yield operation + Value createIntegerAdd(unsigned idx, ArrayRef pointerArgs, + DenseMap &indexMap, + PatternRewriter &rewriter) const; + + // Extract constant integer value from offset (handles both scalar and tensor + // constants) + std::optional extractConstantOffset(Value offsetValue) const; + + // Replace the original loop results with reconstructed pointers from integer + // results + LogicalResult replaceResults(scf::ForOp oldForOp, scf::ForOp newForOp, + ArrayRef pointerArgs, + DenseMap &indexMap, + PatternRewriter &rewriter) const; + + // Reconstruct final pointer from integer result after the loop + Value reconstructPointer(scf::ForOp forOp, unsigned idx, Value intResult, + ArrayRef pointerArgs, + PatternRewriter &rewriter) const; + + LogicalResult matchAndRewriteAdvancePtr(scf::ForOp forOp, + PatternRewriter &rewriter) const; + + SmallVector + collectPointerIterArgsForAdvancePtr(scf::ForOp forOp) const; + + std::optional + analyzePointerIterArgForAdvancePtr(Value iterArg, Block &loopBody) const; + + std::tuple, SmallVector, + DenseMap> + createNewIterArgsForAdvancePtr(scf::ForOp forOp, + SmallVector &pointerArgs, + PatternRewriter &rewriter) const; + + IRMapping + createIRMappingForAdvancePtr(scf::ForOp oldForOp, scf::ForOp newForOp, + SmallVector &pointerArgs, + DenseMap &indexMap, + PatternRewriter &rewriter) const; + + Value rebuildPointerForAdvancePtr(scf::ForOp forOp, + ArrayRef pointerArgs, + unsigned idx, + PatternRewriter &rewriter) const; + + SmallVector + createOffsetsForAdvancePtr(unsigned idx, ArrayRef pointerArgs, + DenseMap &indexMap, + PatternRewriter &rewriter) const; + + LogicalResult cloneInstructionsForAdvancePtr( + Block &oldBody, Block &newBody, ArrayRef pointerArgs, + DenseMap &indexMap, IRMapping &mapping, + PatternRewriter &rewriter) const; + + LogicalResult + rewriteLoopBodyForAdvancePtr(scf::ForOp oldForOp, scf::ForOp newForOp, + SmallVector &pointerArgs, + DenseMap &indexMap, + PatternRewriter &rewriter) const; + + LogicalResult cloneYieldOpForAdvancePtr( + scf::YieldOp yieldOp, ArrayRef pointerArgs, + DenseMap &indexMap, IRMapping &mapping, + PatternRewriter &rewriter) const; + + SmallVector + reconstructPointerForAdvance(scf::ForOp forOp, unsigned idx, Value intResult, + ArrayRef pointerArgs, + PatternRewriter &rewriter) const; +}; + +class SimplifyTensorIterArgsPattern : public OpRewritePattern { +public: + explicit SimplifyTensorIterArgsPattern(MLIRContext *context) + : OpRewritePattern(context) {} + + LogicalResult matchAndRewrite(scf::ForOp forOp, + PatternRewriter &rewriter) const override; + +private: + static constexpr llvm::StringLiteral kSimplifiedAttr = + "tts.simplify_tensor_iter_args.done"; + static constexpr llvm::StringLiteral kFailedAttr = + "tts.simplify_tensor_iter_args.failed"; + static constexpr llvm::StringLiteral kIncompleteAttr = + "tts.simplify_tensor_iter_args.incomplete"; + struct ShapeChainInfo { + Value base; + SmallVector chain; // from base -> iter shape + void dump() const; + }; + + struct RelayMapM1 { + unsigned innerIdx; + unsigned outerInitIdx; // from inner.initArg block-arg mapping + unsigned outerYieldIdx; // from outer.yield operand position of inner result + }; + + struct CandidateInfo { + unsigned idx; + ShapeChainInfo shapeInfo; + SmallVector arithOps; // in execution order + std::optional relayMap; + void dump() const; + }; + + bool isBlockArgumentFromAnotherForLoop(Value v) const; + std::optional getRelayMapM1(scf::ForOp innerFor, + scf::ForOp outerFor, + unsigned innerIdx) const; + void splitCandidatesByRelay(SmallVector all, + SmallVector &locals, + SmallVector &relays) const; + + Value cloneShapeChain(Location loc, Value base, ArrayRef chain, + PatternRewriter &rewriter) const; + Value normalizeInitArgForShapePeel(Value v) const; + std::optional peelShapeChain(Value v) const; + bool isArithWithConst(Operation *op, Value curVal, Value &nextVal, + Value &constVal) const; + Value getNewConstLikeOperand(Value cst, Type targetTy, + PatternRewriter &rewriter) const; + bool canBuildConstLikeOperand(Value cst, Type targetTy) const; + LogicalResult collectReverseLinearYieldPath( + Value yielded, Value iterArg, + SmallVectorImpl &opsInExecOrder) const; + + bool extractBinaryArithOperands(Operation *op, Value &lhs, Value &rhs) const; + Value createSameBinaryArithOp(Operation *oldOp, Location loc, Value lhs, + Value rhs, PatternRewriter &rewriter) const; + bool + isSafeToRewriteLanesByResultUses(scf::ForOp forOp, + ArrayRef candidates) const; + FailureOr rewriteForWithLocalCandidates( + scf::ForOp forOp, ArrayRef candidates, + const IRMapping *outerCaptureMap, PatternRewriter &rewriter) const; + LogicalResult precheckRelayCandidates(scf::ForOp innerFor, + ArrayRef relayCandidates, + scf::ForOp &outerForOut) const; + FailureOr rewriteInnerForWithRelayCandidates( + scf::ForOp innerFor, ArrayRef relayCandidates, + const IRMapping *outerCaptureMap, PatternRewriter &rewriter) const; + FailureOr rewriteOuterForWithRelayCandidates( + scf::ForOp innerFor, scf::ForOp oldInnerFor, scf::ForOp outerFor, + ArrayRef relayCandidates, PatternRewriter &rewriter) const; + LogicalResult + rewriteForWithRelayCandidates(scf::ForOp newfor, scf::ForOp oldFor, + ArrayRef relayCandidates, + PatternRewriter &rewriter) const; +}; +} // namespace CannonicalizerConverter + +#endif diff --git a/compiler/include/dicp/TritonToStructured/MaskAnalysis.h b/compiler/include/dicp/TritonToStructured/MaskAnalysis.h new file mode 100644 index 00000000..8a37f32e --- /dev/null +++ b/compiler/include/dicp/TritonToStructured/MaskAnalysis.h @@ -0,0 +1,110 @@ + +#ifndef TRITON_TO_STRUCTURED_MASKANALYSIS_H +#define TRITON_TO_STRUCTURED_MASKANALYSIS_H + +#include +#include + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Support/LogicalResult.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +namespace TritonToStructured { +using namespace mlir; +using namespace triton; + +struct dimInfo { + OpFoldResult offset; + OpFoldResult shape; + OpFoldResult rhs; + size_t dimIndex; + bool hasBroadCast = false; + + enum class CompareType { slt, sge, ult, uge, deafaultType }; + + CompareType currentType = CompareType::deafaultType; + + dimInfo(size_t dimIndex = 0, bool hasBroadCast = false) + : dimIndex(dimIndex), hasBroadCast(hasBroadCast) {} + + dimInfo(OpFoldResult offset, OpFoldResult shape, size_t dimIndex = 0, + bool hasBroadCast = false, + CompareType Type = CompareType::deafaultType, + OpFoldResult rhs = nullptr) + : offset(offset), shape(shape), dimIndex(dimIndex), + hasBroadCast(hasBroadCast), currentType(Type), rhs(rhs) {} + + bool setType(arith::CmpIPredicate Type); + bool compareTypeIsLess() const; + void dump() const; +}; + +struct MaskState { + SmallVector stateInfo; + OpFoldResult scalar; + Value newMask; + + // Recursively parse a Value; call the corresponding function based on the + // defining operation and Value type + LogicalResult parse(Value operand, const Location loc, OpBuilder &builder); + + bool isEmpty() const { return stateInfo.empty() && !scalar; } + void dump() const; + + // Operand is the result of a constant + // Get the value of the constant and assign it to scalar. + LogicalResult parseConstant(arith::ConstantOp constOp, const Location loc, + OpBuilder &builder); + + LogicalResult parseIntScalar(Value scalar, const Location loc, + OpBuilder &builder); + + LogicalResult parseMakeRange(triton::MakeRangeOp rangeOp, const Location loc, + OpBuilder &builder); + + LogicalResult parseExtSI(arith::ExtSIOp op, const Location loc, + OpBuilder &builder); + + LogicalResult parseSplat(triton::SplatOp splatOp, const Location loc, + OpBuilder &builder); + + LogicalResult parseExpandDims(triton::ExpandDimsOp expandDimsOp, + const Location loc, OpBuilder &builder); + + LogicalResult parseAdd(arith::AddIOp addOp, const Location loc, + OpBuilder &builder); + + LogicalResult parseBroadcast(triton::BroadcastOp broadcastOp, + const Location loc, OpBuilder &builder); + + LogicalResult addStates(const MaskState &lhsState, const MaskState &rhsState, + Location loc, OpBuilder &builder); + + LogicalResult addStateScalar(const MaskState &state, + const OpFoldResult scalar, Location loc, + OpBuilder &builder); + + LogicalResult parseCmp(arith::CmpIOp cmpOp, const Location loc, + OpBuilder &builder); + + LogicalResult parseRem(arith::RemSIOp remOp, const Location loc, + OpBuilder &builder); + + LogicalResult parseDiv(arith::DivSIOp divOp, const Location loc, + OpBuilder &builder); + + LogicalResult parseAnd(arith::AndIOp andOp, const Location loc, + OpBuilder &builder); + + LogicalResult analysisMask(Value operand); + + Value createNewMask(const Location loc, OpBuilder &builder); +}; + +} // namespace TritonToStructured + +#endif \ No newline at end of file diff --git a/compiler/include/dicp/TritonToStructured/MemOpConverter.h b/compiler/include/dicp/TritonToStructured/MemOpConverter.h new file mode 100644 index 00000000..eaac50a8 --- /dev/null +++ b/compiler/include/dicp/TritonToStructured/MemOpConverter.h @@ -0,0 +1,108 @@ + +#ifndef TRITON_ADAPTER_MEMOPCONVERTER_H +#define TRITON_ADAPTER_MEMOPCONVERTER_H + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "dicp/TritonToStructured/MaskAnalysis.h" +#include "dicp/TritonToStructured/PtrAnalysis.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/AffineMap.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LogicalResult.h" +#include "mlir/Transforms/DialectConversion.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +namespace MemOpConverter { + +using namespace mlir; +using namespace triton; + +class LoadConverter : public OpRewritePattern { +public: + explicit LoadConverter(MLIRContext *context, + bool optimizeDynamicOffset = false, + bool enableMaskFallbackConversion = false) + : OpRewritePattern(context), + optimizeDynamicOffset(optimizeDynamicOffset), + enableMaskFallbackConversion(enableMaskFallbackConversion){}; + + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::LoadOp op, + PatternRewriter &rewriter) const override; + +private: + bool optimizeDynamicOffset; + bool enableMaskFallbackConversion; +}; + +class StoreConverter : public OpRewritePattern { +public: + explicit StoreConverter(MLIRContext *context, + bool optimizeDynamicOffset = false, + bool enableMaskFallbackConversion = false) + : OpRewritePattern(context), + optimizeDynamicOffset(optimizeDynamicOffset), + enableMaskFallbackConversion(enableMaskFallbackConversion){}; + + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::StoreOp op, + PatternRewriter &rewriter) const override; + +private: + bool optimizeDynamicOffset; + bool enableMaskFallbackConversion; +}; + +class MemOpTransformer { +public: + TritonToStructured::PtrState ptrState; + TritonToStructured::MaskState maskState; + + enum class MemType { load, store, deafaultType }; + + bool optimizeDynamicOffset; + + MemType currentType = MemType::deafaultType; + + MemOpTransformer(MemType memType, bool optimizeDynamicOffset = false) + : currentType(memType), optimizeDynamicOffset(optimizeDynamicOffset) {} + + Value materializeImplicitBroadcast(Value srcTensor, const Location loc, + PatternRewriter &rewriter); + + Value materializeImplicitReshape(Value srcTensor, const Location loc, + PatternRewriter &rewriter); + + Value materializeImplicitSelect(Value srcTensor, Value mask, Value other, + const Location loc, + PatternRewriter &rewriter); + + Value materializeImplicitPermute(Value srcTensor, const Location loc, + PatternRewriter &rewriter); + + Value createNewPtr(Value oldPtr, const Location loc, + PatternRewriter &rewriter); + + Value createNewMask(Value oldPtr, const Location loc, + PatternRewriter &rewriter); + + Value createNewOther(Value oldOther, const Location loc, + PatternRewriter &rewriter); + + bool applyPermuteOnMask(); +}; + +// Create local lock var +hivm::CreateSyncBlockLockOp createSyncBlockLockVar(OpBuilder &builder, + Location loc); + +} // namespace MemOpConverter + +#endif \ No newline at end of file diff --git a/compiler/include/dicp/TritonToStructured/Passes.h b/compiler/include/dicp/TritonToStructured/Passes.h new file mode 100644 index 00000000..948b6dcb --- /dev/null +++ b/compiler/include/dicp/TritonToStructured/Passes.h @@ -0,0 +1,17 @@ + + +#ifndef TRITON_ADAPTER_TRITON_TO_STRUCTURE_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_STRUCTURE_CONVERSION_PASSES_H + +#include "dicp/TritonToStructured/TritonToStructuredPass.h" + +namespace mlir { +namespace triton { + +#define GEN_PASS_REGISTRATION +#include "dicp/TritonToStructured/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_TO_UNSTRUCTURE_CONVERSION_PASSES_H diff --git a/compiler/include/dicp/TritonToStructured/Passes.td b/compiler/include/dicp/TritonToStructured/Passes.td new file mode 100644 index 00000000..5c1397a9 --- /dev/null +++ b/compiler/include/dicp/TritonToStructured/Passes.td @@ -0,0 +1,19 @@ +#ifndef TRITON_TO_STRUCTURED_PASSES +#define TRITON_TO_STRUCTURED_PASSES + +include "mlir/Pass/PassBase.td" + +def TritonToStructured : Pass<"triton-to-structured", "mlir::ModuleOp"> { + let summary = "remove reminder/divider and reproduce addptr/mask expression "; + let constructor = "triton::createTritonToStructuredPass()"; + let options = [ + Option<"enableMaskFallbackConversion", "enable-mask-fallback-conversion", + "bool", /*default*/"false", + "If enabled, select will perform a fallback conversion when mask matching fails.">, + Option<"optimizeDynamicOffset", "optimize-dynamic-offset", + "bool", /*default*/"false", + "Enable dynamic offset feature"> + ]; +} + +#endif // TRITON_TO_STRUCTURED_PASSES \ No newline at end of file diff --git a/compiler/include/dicp/TritonToStructured/PtrAnalysis.h b/compiler/include/dicp/TritonToStructured/PtrAnalysis.h new file mode 100644 index 00000000..a73f1723 --- /dev/null +++ b/compiler/include/dicp/TritonToStructured/PtrAnalysis.h @@ -0,0 +1,178 @@ + +#ifndef TRITON_TO_STRUCTURED_PTRANALYSIS_H +#define TRITON_TO_STRUCTURED_PTRANALYSIS_H + +#include +#include + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Support/LogicalResult.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +namespace TritonToStructured { +using namespace mlir; +using namespace triton; + +struct StateInfo { + OpFoldResult stride; + OpFoldResult shape; // rem value + size_t dimIndex; + + StateInfo() : dimIndex(0) {} + StateInfo(OpFoldResult stride, OpFoldResult shape, size_t dimIndex = 0) + : stride(stride), shape(shape), dimIndex(dimIndex) {} + void dump() const; +}; + +struct PtrState { + SmallVector + stateInfo; // shape info when load, maintained with visitOps + SmallVector sizes; // original shape, maintained with visitOps + SmallVector permuteIds; + SmallVector + order; // the order of the original data format, only used for block_ptr + SmallVector + dimOffsets; // the offsets per dimension, only used for block_ptr + + Value source; // base address (ptr), maintained with visitOps + OpFoldResult offset; // scalar offset (int), maintained with visitOps + + // whether the record needs to be processed in the current pass, when ignore + // is true, it indicates that this scenario should not be processed within the + // current pass + bool shouldLinearize = false; + bool isPermuted = false; + + void dump() const; + bool isEmpty() const; + bool isScalar() const; + bool hasSource() const; + bool isBlockPtr() const; + bool isSameSizeAs(const PtrState &x) const; + + void generateOriginPermuteIds(); + + // Formula of "contiguous axes" + // - axis i is contiguous if stride[i] == product(shape[0..i-1]). + size_t countContiguousAxes(SmallVector stateInfo) const; + void analyzePermute(); + + void updatePtrState(SmallVector stateInfo, + SmallVector sizes, Value source, + OpFoldResult offset, const Location loc, + OpBuilder &builder, bool shouldLinearize = false); + + void normalizeState(const Location loc, OpBuilder &builder); + + LogicalResult mulState(const PtrState &lhsState, const PtrState &rhsState, + Operation *op, OpBuilder &builder); + + LogicalResult subState(const PtrState &lhsState, const PtrState &rhsState, + Operation *op, OpBuilder &builder); + + LogicalResult addState(PtrState &lhsState, PtrState &rhsState, Operation *op, + OpBuilder &builder); + + triton::AddPtrOp createAddPtrOp(OpBuilder &builder, Location loc); + + triton::MakeTensorPtrOp createMakeTensorPtrOp(OpBuilder &builder, + Location loc); +}; + +class PtrAnalysis { +public: + // AddptrOp result -> PtrState + llvm::SmallDenseMap knownPtrs; + IRMapping ptrMap; + + bool operandIsScalar(Value operand); + + bool optimizeDynamicOffset; + + PtrAnalysis(bool optimizeDynamicOffset = false) + : optimizeDynamicOffset(optimizeDynamicOffset) {} + + LogicalResult initStateByScalar(Value operand, PtrState &state, + const Location loc, OpBuilder &builder); + + LogicalResult initStateByPointer(Value operand, PtrState &state, + const Location loc, OpBuilder &builder); + + LogicalResult visitOperandMul(arith::MulIOp mulOp, PtrState &state, + const Location loc, OpBuilder &builder); + + LogicalResult visitOperandSub(arith::SubIOp subOp, PtrState &state, + const Location loc, OpBuilder &builder); + + LogicalResult visitOperandMakeRange(triton::MakeRangeOp rangeOp, + PtrState &state, Location loc, + OpBuilder &builder); + + LogicalResult visitOperandBroadcast(triton::BroadcastOp broadcastOp, + PtrState &state, const Location loc, + OpBuilder &builder); + + LogicalResult visitOperandSplat(triton::SplatOp splatOp, PtrState &state, + const Location loc, OpBuilder &builder); + + LogicalResult visitOperandExpandDims(triton::ExpandDimsOp expandDimsOp, + PtrState &state, const Location loc, + OpBuilder &builder); + + LogicalResult visitOperandConstSplat(arith::ConstantOp op, PtrState &state, + const Location loc, OpBuilder &builder); + + LogicalResult visitOperandExtSI(arith::ExtSIOp extOp, PtrState &state, + const Location loc, OpBuilder &builder); + + LogicalResult visitOperandRem(arith::RemSIOp remOp, PtrState &state, + const Location loc, OpBuilder &builder); + + LogicalResult visitOperandDiv(arith::DivSIOp divOp, PtrState &state, + const Location loc, OpBuilder &builder); + + LogicalResult visitOperandAdd(arith::AddIOp addOp, PtrState &state, + const Location loc, OpBuilder &builder); + + // Recursively parse a Value; call the corresponding + // function based on the defining operation and argument type. + LogicalResult visitOperand(Value operand, PtrState &state, const Location loc, + OpBuilder &builder); + + // Operand is the result of addptr. + // Main assumptions: + // - The ptr field should populate the source field + // - ptr and offset fields should result in same rank + // Expected result: + // - The resulting state for ptr and offset wil be added + LogicalResult visitOperandAddptr(triton::AddPtrOp addptrOp, PtrState &state, + const Location loc, OpBuilder &builder); + + // Operand is the result of tt.make_tensor_ptr. + // Expected result: + // Parse source pointer and grab results + LogicalResult + visitOperandMakeTensorPtr(triton::MakeTensorPtrOp makeTensorPtrOp, + PtrState &state, const Location loc, + OpBuilder &builder); + + // Parse the state of AddPtrOp, insert any instruction needed to + // calculate strides and offsets, build PtrState for this operand, and record + // PtrState for knownPtrs. + LogicalResult rewriteAddptrOp(triton::AddPtrOp op); +}; + +bool isMultiple(const OpFoldResult ÷nd, const OpFoldResult &divisor); +bool isEqual(const OpFoldResult &ofr1, const OpFoldResult &ofr2); +bool isLess(const OpFoldResult &ofs1, const OpFoldResult &ofs2); +bool isGreater(const OpFoldResult &ofs1, const OpFoldResult &ofs2); +std::optional +extractDivisibilityFromOpFoldResult(mlir::OpFoldResult ofr); + +} // namespace TritonToStructured + +#endif diff --git a/compiler/include/dicp/TritonToStructured/TritonToStructuredPass.h b/compiler/include/dicp/TritonToStructured/TritonToStructuredPass.h new file mode 100644 index 00000000..8ffc8050 --- /dev/null +++ b/compiler/include/dicp/TritonToStructured/TritonToStructuredPass.h @@ -0,0 +1,52 @@ + +#ifndef TRITON_ADAPTER_CONVERSION_TRITONTOSTRUCTURED_H +#define TRITON_ADAPTER_CONVERSION_TRITONTOSTRUCTURED_H + +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#define GEN_PASS_CLASSES +#include "dicp/TritonToStructured/Passes.h.inc" + +namespace mlir { +namespace triton { + +std::unique_ptr> createTritonToStructuredPass(); + +std::unique_ptr> createTritonToStructuredPass(bool, + bool); + +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace triton; + +class TritonToStructuredPass + : public TritonToStructuredBase { +public: + TritonToStructuredPass() = default; + + TritonToStructuredPass(bool enableMaskFallbackConversion, + bool optimizeDynamicOffset) { + this->enableMaskFallbackConversion = enableMaskFallbackConversion; + this->optimizeDynamicOffset = optimizeDynamicOffset; + }; + void getDependentDialects(DialectRegistry ®istry) const override; + void runOnOperation() override; + +private: + void populateTritonToStructuredCanonicalizationPatterns( + RewritePatternSet &patterns); + + void populateTritonToStructuredPatterns(RewritePatternSet &patterns, + bool optimizeDynamicOffset, + bool enableMaskFallbackConversion); + + LogicalResult processSplatBinaryOperations(ModuleOp moduleOp); +}; + +#endif // TRITON_ADAPTER_CONVERSION_TRITONTOSTRUCTURED_H diff --git a/compiler/include/dicp/TritonToUnstructure/BubbleUpOperation.h b/compiler/include/dicp/TritonToUnstructure/BubbleUpOperation.h new file mode 100644 index 00000000..e0ed2f91 --- /dev/null +++ b/compiler/include/dicp/TritonToUnstructure/BubbleUpOperation.h @@ -0,0 +1,92 @@ + + +#pragma once + +#include "mlir/Pass/Pass.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/IR/PatternMatch.h" + +#define GEN_PASS_DECL_BUBBLEUPOPERATION +#include "dicp/TritonToUnstructure/Passes.h.inc" + +#define GEN_PASS_DEF_BUBBLEUPOPERATION +#include "dicp/TritonToUnstructure/Passes.h.inc" + +namespace mlir { +namespace triton { + +std::unique_ptr> +createBubbleUpOperationPass(const BubbleUpOperationOptions &options = {}); + +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace triton; + +template +class BubbleUpExtract : public OpRewritePattern { + static_assert(std::is_same_v || + std::is_same_v); + +public: + using OpRewritePattern::OpRewritePattern; + + explicit BubbleUpExtract(MLIRContext *context, bool enableAggressiveMode); + + LogicalResult matchAndRewrite(ExtractOpTy op, + PatternRewriter &rewriter) const override; + +private: + Value createExtractOp(ExtractOpTy op, Value value, Location loc, + PatternRewriter &rewriter) const; + template + void bubbleUpIntBinaryOp(ExtractOpTy op, BinOpTy binOp, Location loc, + PatternRewriter &rewriter) const; + template + void bubbleUpFloatBinaryOp(ExtractOpTy op, BinOpTy binOp, Location loc, + PatternRewriter &rewriter) const; + + void bubbleUpOperation(ExtractOpTy op, arith::ExtSIOp parentOp, Location loc, + PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, arith::CmpIOp parentOp, Location loc, + PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, arith::TruncFOp parentOp, Location loc, + PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, arith::ExtFOp parentOp, Location loc, + PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, arith::FPToSIOp parentOp, Location loc, + PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, arith::SIToFPOp parentOp, Location loc, + PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, triton::ClampFOp parentOp, + Location loc, PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, arith::CmpFOp parentOp, Location loc, + PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, triton::BroadcastOp parentOp, + Location loc, PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, triton::ExpandDimsOp parentOp, + Location loc, PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, triton::SplatOp parentOp, Location loc, + PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, triton::MakeRangeOp parentOp, + Location loc, PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, triton::AddPtrOp parentOp, + Location loc, PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, math::FloorOp parentOp, Location loc, + PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, math::CeilOp parentOp, Location loc, + PatternRewriter &rewriter) const; + void bubbleUpOperation(ExtractOpTy op, tensor::ExtractSliceOp parentOp, + Location loc, PatternRewriter &rewriter) const; + + bool enableAggressiveMode; +}; + +class BubbleUpOperationPass + : public ::impl::BubbleUpOperationBase { +public: + explicit BubbleUpOperationPass(const BubbleUpOperationOptions &options); + void runOnOperation() override; +}; diff --git a/compiler/include/dicp/Conversion/TritonToUnstructure/CMakeLists.txt b/compiler/include/dicp/TritonToUnstructure/CMakeLists.txt similarity index 100% rename from compiler/include/dicp/Conversion/TritonToUnstructure/CMakeLists.txt rename to compiler/include/dicp/TritonToUnstructure/CMakeLists.txt diff --git a/compiler/include/dicp/Conversion/TritonToUnstructure/OffsetAnalysis.h b/compiler/include/dicp/TritonToUnstructure/OffsetAnalysis.h similarity index 84% rename from compiler/include/dicp/Conversion/TritonToUnstructure/OffsetAnalysis.h rename to compiler/include/dicp/TritonToUnstructure/OffsetAnalysis.h index 2c13880f..d724a75b 100644 --- a/compiler/include/dicp/Conversion/TritonToUnstructure/OffsetAnalysis.h +++ b/compiler/include/dicp/TritonToUnstructure/OffsetAnalysis.h @@ -1,12 +1,15 @@ + + #ifndef TRITON_ANALYSIS_OFFSETANALYSIS_H #define TRITON_ANALYSIS_OFFSETANALYSIS_H - +#include "bishengir/Dialect/HIVM/IR/HIVM.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/IR/BuiltinOps.h" #include "mlir/IR/OpDefinition.h" #include "mlir/IR/PatternMatch.h" #include "mlir/IR/Value.h" +#include "mlir/Transforms/DialectConversion.h" #include "triton/Dialect/Triton/IR/Dialect.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/STLExtras.h" @@ -51,17 +54,19 @@ struct PtrOffsetInfo { */ public: + enum class AxisInfo { unstructured, structured, scalarlike, scalar }; + explicit PtrOffsetInfo(); PtrOffsetInfo(const PtrOffsetInfo &other); explicit PtrOffsetInfo(const Value &ptr); - explicit PtrOffsetInfo(ArrayRef structured); - explicit PtrOffsetInfo(const Value &ptr, bool structured); - explicit PtrOffsetInfo(const Value &ptr, ArrayRef structured); + explicit PtrOffsetInfo(ArrayRef structured); + explicit PtrOffsetInfo(const Value &ptr, AxisInfo structured); + explicit PtrOffsetInfo(const Value &ptr, ArrayRef structured); explicit PtrOffsetInfo(const Value &ptr, const Value &offset, - bool structured); + AxisInfo structured); explicit PtrOffsetInfo(const Value &ptr, const Value &offset, - ArrayRef structured); + ArrayRef structured); PtrOffsetInfo &operator=(const PtrOffsetInfo &other); @@ -70,8 +75,8 @@ struct PtrOffsetInfo { SmallVector getOffsets() const; SmallVector &getOffsetsRef(); bool isScalarLike() const; - SmallVector &getStructuredRef(); - const SmallVector &getStructured() const; + SmallVector &getStructuredRef(); + const SmallVector &getStructured() const; int getRank() const; void setPtr(const Value &ptr); @@ -79,15 +84,17 @@ struct PtrOffsetInfo { void setOffsets(ValueRange offsets); void setStructured(); void setStructured(int rank); + void setStructured(int rank, AxisInfo info); void setUnstructured(); void setUnstructured(int rank); - void setStructured(ArrayRef structured); + void setStructured(ArrayRef structured); void setStructured(const PtrOffsetInfo &other); void setScalarLike(bool scalarLike); bool isStructured(int dim) const; bool isStructured() const; bool isUnstructured() const; + bool isUnstructuredOrScalarlike() const; void setZeroOffset(); @@ -98,7 +105,7 @@ struct PtrOffsetInfo { bool scalarLike = false; - SmallVector structured; + SmallVector structured; }; PtrOffsetInfo combineInfo(const PtrOffsetInfo &lhs, const PtrOffsetInfo &rhs); @@ -189,9 +196,10 @@ void parseSIToFP(arith::SIToFPOp op, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap); -void parseMakeTensorDesc(triton::MakeTensorDescOp op, const Location &loc, - RewriterBase &rewriter, - llvm::DenseMap &offsetMap); +// FIXME:Z|wait triton version upgrade to 3.5 +// void parseMakeTensorDesc(triton::MakeTensorDescOp op, const Location &loc, +// RewriterBase &rewriter, +// llvm::DenseMap &offsetMap); void parseMakeTensorPtr(triton::MakeTensorPtrOp op, const Location &loc, RewriterBase &rewriter, @@ -223,13 +231,26 @@ void parseExtractSlice(tensor::ExtractSliceOp op, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap); +void parseInsertSlice(tensor::InsertSliceOp op, const Location &loc, + RewriterBase &rewriter, + llvm::DenseMap &offsetMap); + void parseExtract(tensor::ExtractOp op, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap); +void parseInsert(tensor::InsertOp op, const Location &loc, + RewriterBase &rewriter, + llvm::DenseMap &offsetMap); + void parseIntToPtr(triton::IntToPtrOp op, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap); + +void parseCustomOp(hivm::CustomOp op, const Location &loc, + RewriterBase &rewriter, + llvm::DenseMap &offsetMap, + unsigned resultIdx); } // namespace triton } // namespace mlir diff --git a/compiler/include/dicp/TritonToUnstructure/Passes.h b/compiler/include/dicp/TritonToUnstructure/Passes.h new file mode 100644 index 00000000..63bbecd1 --- /dev/null +++ b/compiler/include/dicp/TritonToUnstructure/Passes.h @@ -0,0 +1,18 @@ + + +#ifndef TRITON_ADAPTER_TRITON_TO_UNSTRUCTURE_CONVERSION_PASSES_H +#define TRITON_ADAPTER_TRITON_TO_UNSTRUCTURE_CONVERSION_PASSES_H + +#include "dicp/TritonToUnstructure/BubbleUpOperation.h" +#include "dicp/TritonToUnstructure/UnstructureConversionPass.h" + +namespace mlir { +namespace triton { + +#define GEN_PASS_REGISTRATION +#include "dicp/TritonToUnstructure/Passes.h.inc" + +} // namespace triton +} // namespace mlir + +#endif // TRITON_ADAPTER_TRITON_TO_UNSTRUCTURE_CONVERSION_PASSES_H diff --git a/compiler/include/dicp/Conversion/TritonToUnstructure/Passes.td b/compiler/include/dicp/TritonToUnstructure/Passes.td similarity index 75% rename from compiler/include/dicp/Conversion/TritonToUnstructure/Passes.td rename to compiler/include/dicp/TritonToUnstructure/Passes.td index a9abee65..4833a248 100644 --- a/compiler/include/dicp/Conversion/TritonToUnstructure/Passes.td +++ b/compiler/include/dicp/TritonToUnstructure/Passes.td @@ -8,7 +8,13 @@ def TritonToUnstructure : Pass<"triton-to-unstructure", "mlir::ModuleOp"> { let constructor = "triton::createTritonToUnstructurePass()"; let options = [ Option<"forceScalarizeMode", "force-scalarize-mode", "bool", "false", - "Scalarize unstructured memory access even if structured dimensions are mixed."> + "Scalarize unstructured memory access even if structured dimensions are mixed.">, + Option<"compileOn91095", "compile-on-910-95", + "bool", /*default*/"false", + "compile on 910_95">, + Option<"forceSimtTemplate", "force-simt-template", + "bool", /*default*/"false", + "force to use simt template"> ]; } diff --git a/compiler/include/dicp/Conversion/TritonToUnstructure/UnstructureConversionPass.h b/compiler/include/dicp/TritonToUnstructure/UnstructureConversionPass.h similarity index 84% rename from compiler/include/dicp/Conversion/TritonToUnstructure/UnstructureConversionPass.h rename to compiler/include/dicp/TritonToUnstructure/UnstructureConversionPass.h index e7e07f4a..9f480265 100644 --- a/compiler/include/dicp/Conversion/TritonToUnstructure/UnstructureConversionPass.h +++ b/compiler/include/dicp/TritonToUnstructure/UnstructureConversionPass.h @@ -1,22 +1,29 @@ -#ifndef TRITON_DLC_UNSTRUCTURECONVERSION_H -#define TRITON_DLC_UNSTRUCTURECONVERSION_H -#include "dicp/Conversion/TritonToUnstructure/OffsetAnalysis.h" + +#ifndef TRITON_ADAPTER_UNSTRUCTURECONVERSION_H +#define TRITON_ADAPTER_UNSTRUCTURECONVERSION_H + +#include "dicp/TritonToUnstructure/OffsetAnalysis.h" #include "mlir/Pass/Pass.h" #include "triton/Dialect/Triton/IR/Dialect.h" +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" #include "mlir/IR/PatternMatch.h" #define GEN_PASS_DECL_TRITONTOUNSTRUCTURE -#include "dicp/Conversion/TritonToUnstructure/Passes.h.inc" +#include "dicp/TritonToUnstructure/Passes.h.inc" #define GEN_PASS_DEF_TRITONTOUNSTRUCTURE -#include "dicp/Conversion/TritonToUnstructure/Passes.h.inc" +#include "dicp/TritonToUnstructure/Passes.h.inc" + +extern bool compileOn91095Flag; +extern bool forceSimtTemplateFlag; namespace mlir { namespace triton { -std::unique_ptr> createTritonToUnstructurePass(); +std::unique_ptr> +createTritonToUnstructurePass(const TritonToUnstructureOptions &options = {}); } // namespace triton } // namespace mlir @@ -97,6 +104,7 @@ class UnstructuredMemAccessConverter : public OpRewritePattern { class TritonToUnstructurePass : public ::impl::TritonToUnstructureBase { public: + explicit TritonToUnstructurePass(const TritonToUnstructureOptions &options); void getDependentDialects(DialectRegistry ®istry) const override; void runOnOperation() override; @@ -113,8 +121,13 @@ class TritonToUnstructurePass llvm::DenseMap offsetMap; llvm::DenseMap offsetMapForLoopArgs; llvm::SmallDenseMap fromTensorArg; + + LogicalResult processIfYieldAddHoistOperations(ModuleOp moduleOp); }; } // namespace -#endif // TRITON_DLC_UNSTRUCTURECONVERSION_H +void replacePtrArguments(triton::FuncOp funcOp, + llvm::DenseMap &offsetMap); + +#endif // TRITON_ADAPTER_UNSTRUCTURECONVERSION_H diff --git a/compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtInterface.td b/compiler/include/dicp/Utils/CMakeLists.txt similarity index 100% rename from compiler/include/dicp/Dialect/LinalgExt/IR/LinalgExtInterface.td rename to compiler/include/dicp/Utils/CMakeLists.txt diff --git a/compiler/include/dicp/Utils/InterleaveOptimization.h b/compiler/include/dicp/Utils/InterleaveOptimization.h new file mode 100644 index 00000000..deb66df5 --- /dev/null +++ b/compiler/include/dicp/Utils/InterleaveOptimization.h @@ -0,0 +1,74 @@ + + +#pragma once + +#include "dicp/TritonToLinalg/BlockPtrAnalysis.h" +#include "dicp/TritonToLinalg/MaskAnalysis.h" +#include "dicp/TritonToLinalg/UseAnalysis.h" +#include "dicp/Utils/Utils.h" + +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/GPU/IR/GPUDialect.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Passes.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Matchers.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LogicalResult.h" + +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/MathExtras.h" + +#include +#include +#include +#include +#include + +namespace mlir { +namespace triton { + +enum class IndexMode : int { EVEN_MODE = 0, ODD_MODE = 1 }; + +MemRefType expandInterleaveMemRefType(MemRefType originType); + +bool checkIsCaseOffsetValid(OpFoldResult originOffset); + +std::pair +recountReinterpretCastOffset(OpFoldResult originOffset, Builder &builder); + +LogicalResult +DeinterleaveStatusOptimization(triton::LoadOp op, + triton::LoadOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter); + +LogicalResult DeinterleaveStatusWithMaskOptimization( + triton::LoadOp op, triton::LoadOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter, MaskState &mstate, Value localMem); + +LogicalResult +InterleaveStatusOptimization(SmallVector materializeVec); + +LogicalResult +InterleaveStatusWithMaskOptimization(SmallVector materializeVec); + +} // namespace triton +} // namespace mlir diff --git a/compiler/include/dicp/Utils/Utils.h b/compiler/include/dicp/Utils/Utils.h index e50253dd..8674d499 100644 --- a/compiler/include/dicp/Utils/Utils.h +++ b/compiler/include/dicp/Utils/Utils.h @@ -1,5 +1,7 @@ -#ifndef TRITON_UTILS_H -#define TRITON_UTILS_H + + +#ifndef DICP_UTILS_UTILS_H +#define DICP_UTILS_UTILS_H #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" @@ -8,83 +10,112 @@ #include "mlir/Dialect/Utils/StructuredOpsUtils.h" #include "mlir/IR/OpDefinition.h" #include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" #include "mlir/Transforms/DialectConversion.h" - +#include "triton/Dialect/Triton/IR/Dialect.h" #include "llvm/ADT/ArrayRef.h" -#include "llvm/ADT/StringSwitch.h" +#include "llvm/Support/LogicalResult.h" #include #include +#include -// Dispatch conversion pattern handlers based on backend string. Executes -// ASCEND_HANDLER when backend == "ascend", otherwise DEFAULT_HANDLER. -#define DISPATCH_BACKEND_CONVERSION_PATTERNS(BACKEND_STR, ASCEND_HANDLER, \ - DEFAULT_HANDLER) \ - do { \ - auto populatePatterns = \ - llvm::StringSwitch>(BACKEND_STR) \ - .Case("ascend", [&] { ASCEND_HANDLER; }) \ - .Default([&] { DEFAULT_HANDLER; }); \ - populatePatterns(); \ - } while (0) +namespace mlir { -namespace mlir::dicp { +namespace ConverterUtils { -// Tags used for marking specific operations for later processing or -// identification. const std::string GeneratedByMakeTensorPtrTAG = "GeneratedByMakeTensorPtr"; -const std::string MayImplicitTransposeWithLastAxisTAG = - "MayImplicitTransposeWithLastAxis"; const std::string discreteMaskAttrName = "DiscreteMask"; const std::string discreteAttrName = "DiscreteMemAccess"; - -// Gets the string attribute "dicp.backend" from the module if it exists. -llvm::StringRef getBackend(ModuleOp module); - -bool isAscendBackend(ModuleOp module); +const std::string continuousAttrName = "ContinuousMemAccess"; +const std::string customSrcPtrIndexAttrName = "SrcPtrIndex"; bool isaPermutedMemRefType(MemRefType); -// Retrieves the last (innermost) stride of a memref::ReinterpretCastOp if it is -// a constant. std::optional getLastStrideOfReinterpretCastOp(memref::ReinterpretCastOp op); -// Creates a new tensor by transposing the 'source' value according to the -// 'order'. Value getTransposedValue(Value source, const Location loc, ConversionPatternRewriter &rewriter, llvm::ArrayRef order); -// Returns a vector of `n` `utils::IteratorType::parallel` attributes. SmallVector getNParallelLoopsAttrs(unsigned n); -// Reconstructs the scalar value from an operand that might be a tensor/vector -// containing a single splat value, handling implicit casts like `sitofp` or -// `truncf`. Value getScalarValue(Value operand, Location loc, ConversionPatternRewriter &rewriter); -// Identifies the dimensions in the source tensor that are broadcast to match -// the destination tensor's shape (where source dim size is 1). +memref::SubViewOp makeSubViewOp(Value src, + const llvm::SmallVector &offsets, + const llvm::SmallVector &sizes, + const Location &loc, + ConversionPatternRewriter &rewriter); + +tensor::ExtractSliceOp +makeExtractSliceOp(Value src, const llvm::SmallVector &offsets, + const llvm::SmallVector &sizes, + const Location &loc, ConversionPatternRewriter &rewriter); + +std::optional getFullShapeOp(Value val, + ConversionPatternRewriter &rewriter); + +SmallVector +getBoundarySizes(llvm::ArrayRef boundaryCheck, Value ptr, + const Location &loc, ConversionPatternRewriter &rewriter); + SmallVector getBroadcastDims(RankedTensorType src, RankedTensorType dst); -// Identifies the dimensions that are NOT broadcast (i.e., source shape matches -// destination shape). SmallVector getUnbroadcastDims(RankedTensorType src, RankedTensorType dst); -// Enumeration for types of operations that interact with memory indirectly -// (e.g., loads/computations on pointers). +} // namespace ConverterUtils + +class ConversionPatternRewriter; + +namespace triton { + enum class IndirectLoadInterfaceOpType { Undefined = 0, Load = 1, Calc = 2 }; -// Traces back from a 'rootOp' through its operands' definitions to find the -// first operation that satisfies the specified 'condFn'. +// Traceback from rootOp to find the targetOp with the specified condition mlir::Operation * findFirstMatchingOperandDef(mlir::Operation *rootOp, const std::function &condFn); +void traverseBackwardUpdateOperandChainIf( + Operation *op, std::function conditionFn, + std::function stopFn, + std::function actionFn, OpBuilder &builder, + DenseSet &handledOperation); + +void traverseBackwardUpdateOperandChainIf( + Operation *rootOp, std::function conditionFn, + std::function stopFn, + std::function actionFn); + +void traverseForwardUpdateUserChainIf( + Operation *op, std::function conditionFn, + std::function stopFn, + std::function actionFn, OpBuilder &builder, + llvm::SmallPtrSet &stopOps); + +void traverseForwardUpdateUserChainIf( + Operation *rootOp, std::function conditionFn, + std::function stopFn, + std::function actionFn, + llvm::SmallPtrSet &stopOps); + +// UseAnalysis will tag operations whose results are used only as meta-data +// with "MetaUse" tag. +bool isMetaUse(Operation *op); + +bool isMixUse(Operation *op); + +IndirectLoadInterfaceOpType getIndirectLoadInterfaceOpType(Operation *op); + +bool opIsIndirectLoad(Operation *op); + +bool opIsIndirectCalc(Operation *op); + /// Maximum expected rank for loop tiling in tensor operations. static constexpr int kMaxTiledRank = 4; @@ -105,14 +136,36 @@ static constexpr int kMaxTiledRank = 4; template void createSimpleNestedLoops(OpBuilder &rewriter, Location loc, Value target, ArrayRef loopDims, Func bodyFunc) { - // Implementation details omitted in header but provided in the question's - // context. - // ... + MemRefType type = cast(target.getType()); + int rank = type.getRank(); + + Value zero = rewriter.create(loc, 0); + Value one = rewriter.create(loc, 1); + + llvm::SmallVector loops; + llvm::SmallVector ivs; + + for (int dim : loopDims) { + Value ub; + if (type.isDynamicDim(dim)) { + ub = rewriter.create(loc, target, dim).getResult(); + } else { + ub = rewriter.create(loc, type.getDimSize(dim)); + } + + auto forOp = rewriter.create(loc, zero, ub, one); + rewriter.setInsertionPointToStart(forOp.getBody()); + loops.push_back(forOp); + ivs.push_back(forOp.getInductionVar()); + } + + bodyFunc(ivs); + + if (!loops.empty()) { + rewriter.setInsertionPointAfter(loops.front()); + } } -// Recursively creates a potentially nested structure of `scf.for` loops. -// This allows for defining complex loop nests where the body is generated by -// 'bodyBuilder'. scf::ForOp createNestedLoops( OpBuilder &builder, Location loc, unsigned currentDim, unsigned totalDims, ValueRange LBs, ValueRange UBs, ValueRange steps, SmallVector &ivs, @@ -120,11 +173,97 @@ scf::ForOp createNestedLoops( function_ref &, ValueRange)> bodyBuilder); +ModuleOp getModuleOpFromOperation(Operation *op); + +bool isTensorPtrType(Type type); + +} // namespace triton + +class OpBuilder; + +OpFoldResult addOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b); + +OpFoldResult subOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b); + +OpFoldResult mulOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b); + +OpFoldResult divOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b); + +OpFoldResult remOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b); + +OpFoldResult minOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b); + +OpFoldResult maxOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b); + +enum class ReduceWithIndexType { MAX, MIN, None }; +enum class TieBreakType { LEFT, RIGHT, None }; + +struct ReduceWithIndexParams { + ReduceWithIndexType withIndexType = ReduceWithIndexType::None; + TieBreakType tieBreakType = TieBreakType::None; + bool isUnsignedSrc; +}; + +llvm::FailureOr +getReduceWithIndexParams(triton::ReduceOp op); + +void addReduceWithIndexAttr(ReduceWithIndexParams params, + ConversionPatternRewriter &rewriter, + linalg::ReduceOp reduceOp); + +OpFoldResult getOpFoldResultOfLayoutInfo(Value value, OpBuilder &builder); + enum class TypelessValue { Undefined = 0, Zero = 1, Min = 2, Max = 3 }; +FailureOr specializeTypelessValueToAttr(TypelessValue, Type, + OpBuilder &); + FailureOr specializeTypelessValueToConstant(TypelessValue, Type, Location, OpBuilder &); +std::optional getIntAttr(const OpFoldResult ofr); + +Value materializeValue(OpBuilder &builder, Location loc, OpFoldResult ofr); + +bool isZero(const OpFoldResult ofr); + +bool isOne(const OpFoldResult ofr); + +Value convertToIndexIfNeeded(Value intValue, const Location &loc, OpBuilder &b); + +RankedTensorType getExtractSlicedType(ArrayRef shape, + const llvm::SmallBitVector &droppedDims, + Type elemType); + +bool checkStructureAnnotated(Operation *op, RewriterBase &rewriter); + +// Dispatch conversion pattern handlers based on backend string. Executes +// DICP_HANDLER when backend == "dicp", otherwise DEFAULT_HANDLER. +#define DISPATCH_BACKEND_CONVERSION_PATTERNS(BACKEND_STR, DICP_HANDLER, \ + DEFAULT_HANDLER) \ + do { \ + auto populatePatterns = \ + llvm::StringSwitch>(BACKEND_STR) \ + .Case("dicp", [&] { DICP_HANDLER; }) \ + .Default([&] { DEFAULT_HANDLER; }); \ + populatePatterns(); \ + } while (0) + +} // namespace mlir + +namespace mlir::dicp { + +llvm::StringRef getBackend(ModuleOp module); + +bool isDicpBackend(ModuleOp module); + } // namespace mlir::dicp -#endif // TRITONNPU_UTILS_UTILS_H \ No newline at end of file +#endif // DICP_UTILS_UTILS_H diff --git a/compiler/lib/AscendLegalize/AscendLegalizePass.cpp b/compiler/lib/AscendLegalize/AscendLegalizePass.cpp new file mode 100644 index 00000000..7ed17b60 --- /dev/null +++ b/compiler/lib/AscendLegalize/AscendLegalizePass.cpp @@ -0,0 +1,277 @@ +#include "dicp/AscendLegalize/AscendLegalizePass.h" + +#include "bishengir/Dialect/HACC/IR/HACC.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/Passes.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/STLExtras.h" + +#include + +#define DEBUG_TYPE "ascend-legalize" + +using namespace mlir; +using namespace triton; + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_ASCENDLEGALIZE +#include "dicp/AscendLegalize/Passes.h.inc" +} // namespace triton +} // namespace mlir + +namespace { + +static bool isNonNegativeIntegerConstant(Value value, unsigned depth = 0) { + if (depth > 4) + return false; + + if (auto splatOp = value.getDefiningOp()) + return isNonNegativeIntegerConstant(splatOp.getSrc(), depth + 1); + + if (auto bitcastOp = value.getDefiningOp()) + return isNonNegativeIntegerConstant(bitcastOp.getIn(), depth + 1); + + auto constantOp = value.getDefiningOp(); + if (!constantOp) + return false; + + Attribute attr = constantOp.getValue(); + if (auto intAttr = dyn_cast(attr)) + return !intAttr.getValue().isNegative(); + + if (auto denseAttr = dyn_cast(attr)) { + return llvm::all_of(denseAttr.getValues(), + [](const APInt &value) { return !value.isNegative(); }); + } + + return false; +} + +static bool isBoolExtSumReduce(triton::ReduceOp reduceOp) { + if (reduceOp.getSrcs().size() != 1) + return false; + + auto extUIOp = reduceOp.getSrcs()[0].getDefiningOp(); + if (!extUIOp) + return false; + + auto boolSrcType = dyn_cast(extUIOp.getOperand().getType()); + if (!boolSrcType || !boolSrcType.getElementType().isInteger(1)) + return false; + + Block &body = reduceOp.getCombineOp().front(); + if (body.getNumArguments() != 2) + return false; + + auto termOp = dyn_cast(body.getTerminator()); + if (!termOp || termOp.getOperands().size() != 1) + return false; + + auto addOp = termOp.getOperands()[0].getDefiningOp(); + if (!addOp) + return false; + + Value lhs = addOp.getLhs(); + Value rhs = addOp.getRhs(); + return (lhs == body.getArgument(0) && rhs == body.getArgument(1)) || + (lhs == body.getArgument(1) && rhs == body.getArgument(0)); +} + +static std::optional +getSignedPredicate(arith::CmpIPredicate predicate) { + switch (predicate) { + case arith::CmpIPredicate::ugt: + return arith::CmpIPredicate::sgt; + case arith::CmpIPredicate::uge: + return arith::CmpIPredicate::sge; + case arith::CmpIPredicate::ult: + return arith::CmpIPredicate::slt; + case arith::CmpIPredicate::ule: + return arith::CmpIPredicate::sle; + default: + return std::nullopt; + } +} + +/// Pattern: Flip cmpi predicates that are incompatible with +/// MaskState::parseCmp. +/// +/// MaskState::parseCmp requires lhs = tensor (has range/offset) and rhs = +/// scalar (splat constant). When the IR produces sge(scalar, tensor) instead, +/// this pattern rewrites it to the mathematically equivalent +/// sle(tensor, scalar) by swapping operands and flipping the predicate. +/// +/// Supported flips: +/// sge(A, B) -> sle(B, A) +/// sgt(A, B) -> slt(B, A) +struct FlipCmpiPredicatePattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(arith::CmpIOp cmpOp, + PatternRewriter &rewriter) const override { + auto predicate = cmpOp.getPredicate(); + + arith::CmpIPredicate flippedPredicate; + switch (predicate) { + case arith::CmpIPredicate::sge: + flippedPredicate = arith::CmpIPredicate::sle; + break; + case arith::CmpIPredicate::sgt: + flippedPredicate = arith::CmpIPredicate::slt; + break; + default: + return failure(); + } + + auto lhs = cmpOp.getLhs(); + auto rhs = cmpOp.getRhs(); + + auto lhsSplat = lhs.getDefiningOp(); + if (!lhsSplat) + return failure(); + + auto rhsSplat = rhs.getDefiningOp(); + if (rhsSplat) + return failure(); + + auto newCmp = rewriter.create(cmpOp.getLoc(), + flippedPredicate, rhs, lhs); + rewriter.replaceOp(cmpOp, newCmp.getResult()); + return success(); + } +}; + +/// Pattern: Replace arith::MaxNumFOp (NaN-quiet) with arith::MaximumFOp +/// (NaN-propagating). +/// +/// On Ascend NPU, online-softmax reductions need NaN propagation so that +/// m_ij = max(m_i, max(qk)) correctly propagates NaN through the reduce +/// region. MaxNumFOp silently swallows NaN, leading to wrong exp() results. +struct MaxNumFToMaximumFPattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(arith::MaxNumFOp op, + PatternRewriter &rewriter) const override { + rewriter.replaceOpWithNewOp(op, op.getType(), + op.getLhs(), op.getRhs()); + return success(); + } +}; + +/// Pattern: Give masked tt.load an explicit zero `other` operand when it is +/// omitted. +/// +/// The TA Triton frontend defaults masked tl.load(..., mask=...) to +/// care_padding=True, which lowers to tt.load(ptr, mask, zero). The upstream +/// Triton 3.5 frontend leaves `other` absent instead. Ascend's downstream +/// unstructure/linalg lowering uses the explicit zero-fill operand to preserve +/// masked lanes, especially for dot operands where it becomes zero-padding +/// slices. Normalize here so both frontends feed the same IR shape to DICP. +struct AddZeroOtherToMaskedLoadPattern + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::LoadOp loadOp, + PatternRewriter &rewriter) const override { + if (!loadOp.getMask() || loadOp.getOther()) + return failure(); + + auto zeroAttr = rewriter.getZeroAttr(loadOp.getType()); + if (!zeroAttr) + return failure(); + + Location loc = loadOp.getLoc(); + Value zeroOther = + rewriter.create(loc, loadOp.getType(), zeroAttr); + auto newLoad = rewriter.create( + loc, loadOp.getPtr(), loadOp.getMask(), zeroOther, + loadOp.getBoundaryCheck(), loadOp.getPadding(), loadOp.getCache(), + loadOp.getEvict(), loadOp.getIsVolatile()); + + for (auto attr : loadOp->getAttrs()) { + if (!newLoad->hasAttr(attr.getName())) + newLoad->setAttr(attr.getName(), attr.getValue()); + } + + rewriter.replaceOp(loadOp, newLoad.getResult()); + return success(); + } +}; + +/// Pattern: Rewrite unsigned comparisons on bool-sum counts to signed +/// comparisons. +/// +/// Triton 3.5's tl.sum defaults integer inputs narrower than i32 to an i32 +/// accumulator with the same signedness. For bool input, that produces the IR +/// shape: +/// +/// arith.extui i1 -> i32 +/// tt.reduce(add i32) +/// arith.cmpi ugt/uge/ult/ule reduce, non_negative_constant +/// +/// TA keeps the same kernel on a bool/signed-friendly path, while Ascend +/// BiSheng HIR can fail on the unsigned route with an unsupported +/// uint32_t_to_uint64_t_rintmode vcast. The reduced value here is a count of +/// bool lanes, so it is non-negative and bounded by the reduction axis. For +/// non-negative signed constants, signed and unsigned comparisons are +/// equivalent. Canonicalize only this proven bool-count pattern and leave +/// ordinary unsigned integer comparisons untouched. +struct BoolSumUnsignedCmpToSignedPattern + : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(arith::CmpIOp cmpOp, + PatternRewriter &rewriter) const override { + auto signedPredicate = getSignedPredicate(cmpOp.getPredicate()); + if (!signedPredicate) + return failure(); + + if (!isNonNegativeIntegerConstant(cmpOp.getRhs())) + return failure(); + + auto reduceOp = cmpOp.getLhs().getDefiningOp(); + if (!reduceOp || !isBoolExtSumReduce(reduceOp)) + return failure(); + + auto newCmp = rewriter.create( + cmpOp.getLoc(), *signedPredicate, cmpOp.getLhs(), cmpOp.getRhs()); + rewriter.replaceOp(cmpOp, newCmp.getResult()); + return success(); + } +}; + +/// The pass: AscendLegalizePass +struct AscendLegalizePass + : public ::mlir::triton::impl::AscendLegalizeBase { + using AscendLegalizeBase::AscendLegalizeBase; + + void runOnOperation() override { + ModuleOp moduleOp = getOperation(); + RewritePatternSet patterns(&getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + + // On Ascend NPU, replace arith::MaxNumFOp (NaN-quiet) with + // arith::MaximumFOp (NaN-propagating) so that online-softmax + // reductions propagate NaN correctly through the reduce region. + if (auto targetAttr = + moduleOp->getAttrOfType(hacc::TargetAttr::name)) { + patterns.add(patterns.getContext()); + } + + if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { + moduleOp->emitError("failed to apply ascend-legalize patterns"); + signalPassFailure(); + } + } +}; + +} // namespace + +std::unique_ptr> triton::createAscendLegalizePass() { + return std::make_unique(); +} diff --git a/compiler/lib/AscendLegalize/CMakeLists.txt b/compiler/lib/AscendLegalize/CMakeLists.txt new file mode 100644 index 00000000..d79100c5 --- /dev/null +++ b/compiler/lib/AscendLegalize/CMakeLists.txt @@ -0,0 +1,15 @@ +add_triton_library( + AscendLegalize + AscendLegalizePass.cpp + + DEPENDS + AscendLegalizePassIncGen + + LINK_LIBS PUBLIC + BiShengIRHACCDialect + MLIRArithDialect + MLIRIR + MLIRPass + TritonIR + MLIRTransforms +) diff --git a/compiler/lib/AutoBlockify/AutoBlockify.cpp b/compiler/lib/AutoBlockify/AutoBlockify.cpp new file mode 100644 index 00000000..865c0ca4 --- /dev/null +++ b/compiler/lib/AutoBlockify/AutoBlockify.cpp @@ -0,0 +1,343 @@ + + +#include "dicp/AutoBlockify/AutoBlockify.h" +#include "dicp/AutoBlockify/Utils.h" +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" +#include "dicp/Utils/Utils.h" + +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/Passes.h" + +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "auto-blockify" + +using namespace mlir; +using namespace triton; + +PropagateUnrealizedCastDown::PropagateUnrealizedCastDown(MLIRContext *context, + Value logicalBlockId, + Value logicalBlockNum, + int autoBlockifySize) + : OpRewritePattern(context), + logicalBlockId(logicalBlockId), logicalBlockNum(logicalBlockNum), + autoBlockifySize(autoBlockifySize) {} + +LogicalResult +PropagateUnrealizedCastDown::matchAndRewrite(UnrealizedConversionCastOp op, + PatternRewriter &rewriter) const { + if (op.getInputs().size() != 2) + return failure(); + auto funcOp = op->getParentOfType(); + auto input = op.getInputs()[0]; + auto res = op->getResult(0); + SmallPtrSet users(op->user_begin(), op->user_end()); + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "Handling UnrealizedConversionCastOp:\n" << op << "\n"; + os << "Users:\n"; + for (auto *user : users) + os << *user << "\n"; + }); + for (auto *user : users) { + PatternRewriter::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(user); + if (auto uccOp = dyn_cast(user)) { + if (uccOp->getResultTypes()[0] != input.getType()) { + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << *user << "\n"; + }); + return op.emitError("UnrealizedConversionCastOp cannot be resolved\n"); + } + rewriter.replaceOp(user, input); + } else if (auto blockifyLoop = getBlockifyLoop(user)) { + handleBlockifyLoop(blockifyLoop.value(), user, rewriter); + } else if (auto splatOp = dyn_cast(user)) { + rewriteSplat(op, splatOp, rewriter); + } else if (auto expandDimsOp = dyn_cast(user)) { + rewriteExpandDims(op, expandDimsOp, rewriter); + } else if (auto reduceOp = dyn_cast(user)) { + rewriteReduce(op, reduceOp, rewriter); + } else if (auto scanOp = dyn_cast(user)) { + rewriteScan(op, scanOp, rewriter); + } else if (auto loadOp = dyn_cast(user)) { + rewriteLoad(op, loadOp, rewriter); + } else if (auto storeOp = dyn_cast(user)) { + rewriteStore(op, storeOp, rewriter); + } else if (auto atomicRMWOp = dyn_cast(user)) { + rewriteAtomicRMW(op, atomicRMWOp, rewriter); + } else if (auto assertOp = dyn_cast(user)) { + rewriteAssert(op, assertOp, rewriter); + } else if (auto extractSliceOp = dyn_cast(user)) { + rewriteExtractSlice(op, extractSliceOp, rewriter); + } else if (auto insertSliceOp = dyn_cast(user)) { + rewriteInsertSlice(op, insertSliceOp, rewriter); + } else if (auto whileOp = dyn_cast(user)) { + rewriteWhile(op, whileOp, rewriter); + } else if (auto loopOp = dyn_cast(user)) { + rewriteLoop(op, loopOp, rewriter); + } else if (auto yieldOp = dyn_cast(user)) { + rewriteYield(op, yieldOp, rewriter); + } else if (auto conditionOp = dyn_cast(user)) { + rewriteCondition(op, conditionOp, rewriter); + } else if (user->hasTrait() || + isa(user)) { + rewriteGeneraleOp(op, user, rewriter); + } else if (isa(user)) { + auto *newOp = + createBlockifyLoop(user, op, logicalBlockId, logicalBlockNum, + autoBlockifySize, rewriter); + rewriter.setInsertionPoint(newOp); + handleBlockifyLoop(*getBlockifyLoop(newOp), newOp, rewriter); + } else { + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "Unhandled Op\n" << *user << "\n"; + }); + llvm_unreachable("Unhandled operation"); + } + } + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "After successful conversion\n"; + os << funcOp << "\n"; + }); + rewriter.eraseOp(op); + return success(); +} + +AutoBlockifyPass::AutoBlockifyPass(const AutoBlockifyOptions &options) + : AutoBlockifyBase(options) {} + +bool AutoBlockifyPass::checkBlockifiable(Value v) { + if (!checkedValues.insert(v).second) + return true; + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "Checking blockifiable:\n" << v << "\n"; + }); + for (auto &use : v.getUses()) { + auto *user = use.getOwner(); + auto opNum = use.getOperandNumber(); + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "User:\n" << *user << "\n"; + }); + if (isa( + user) || + llvm::any_of(user->getOperandTypes(), isTensorPtrType)) + return false; + if (auto ifOp = dyn_cast(user)) { + user->setAttr(autoBlockifyRegionOpAttr, UnitAttr::get(v.getContext())); + return true; + } else if (auto whileOp = dyn_cast(user)) { + if (!checkBlockifiable(whileOp.getBeforeArguments()[opNum])) + return false; + } else if (auto loopOp = dyn_cast(user)) { + auto regionIterArg = loopOp.getTiedLoopRegionIterArg(&use); + auto loopResult = loopOp.getTiedLoopResult(&use); + if (!regionIterArg || !loopResult) { + user->setAttr(autoBlockifyRegionOpAttr, UnitAttr::get(v.getContext())); + return true; + } + if (!checkBlockifiable(regionIterArg) || !checkBlockifiable(loopResult)) + return false; + } else if (auto conditionOp = dyn_cast(user)) { + auto whileOp = cast(user->getParentOp()); + if (opNum == 0) { + whileOp->setAttr(autoBlockifyRegionOpAttr, + UnitAttr::get(v.getContext())); + return true; + } + if (!checkBlockifiable(whileOp.getAfterArguments()[opNum - 1]) || + !checkBlockifiable(whileOp->getResult(opNum - 1))) + return false; + } else if (auto conditionOp = dyn_cast(user)) { + if (auto loopOp = dyn_cast(user->getParentOp()); + loopOp && !checkBlockifiable(loopOp.getInits()[opNum])) + return false; + } else { + for (auto res : user->getResults()) { + if (!checkBlockifiable(res)) + return false; + } + } + } + return true; +} + +void AutoBlockifyPass::preProcess(triton::FuncOp func) { + IRRewriter rewriter(func.getContext()); + rewriter.setInsertionPointToStart(&func.getBody().front()); + auto loc = rewriter.getUnknownLoc(); + // Get logical block num + auto xNum = + rewriter.create(loc, triton::ProgramIDDim::X); + auto yNum = + rewriter.create(loc, triton::ProgramIDDim::Y); + auto zNum = + rewriter.create(loc, triton::ProgramIDDim::Z); + auto yzNum = rewriter.create(loc, yNum, zNum); + logicalBlockNum = rewriter.create(loc, yzNum, xNum); + + // Get logical block id + auto xDim = + rewriter.create(loc, triton::ProgramIDDim::X); + auto yDim = + rewriter.create(loc, triton::ProgramIDDim::Y); + auto zDim = + rewriter.create(loc, triton::ProgramIDDim::Z); + xDim->setAttr(logicalBlockIdAttr, rewriter.getUnitAttr()); + yDim->setAttr(logicalBlockIdAttr, rewriter.getUnitAttr()); + zDim->setAttr(logicalBlockIdAttr, rewriter.getUnitAttr()); + auto xFlatten = rewriter.create(loc, xDim, yzNum); + auto yFlatten = rewriter.create(loc, yDim, zNum); + logicalBlockId = rewriter.create(loc, xFlatten, yFlatten); + logicalBlockId = rewriter.create(loc, logicalBlockId, zDim); + + // get blockified block id + auto blockifyTensorType = + RankedTensorType::get({autoBlockifySize}, rewriter.getI32Type()); + auto blockfyRange = rewriter.create( + loc, blockifyTensorType, 0, autoBlockifySize); + auto splatedLogicalBlockId = rewriter.create( + loc, blockfyRange.getType(), logicalBlockId); + Value blockifiedId = + rewriter.create(loc, splatedLogicalBlockId, blockfyRange); + + // get mask + auto splatedBlockNum = rewriter.create( + loc, blockfyRange.getType(), logicalBlockNum); + auto upperboundMask = rewriter.create( + loc, arith::CmpIPredicate::slt, blockifiedId, splatedBlockNum); + auto splatedZero = rewriter.create( + loc, DenseElementsAttr::get(blockifyTensorType, + rewriter.getI32IntegerAttr(0))); + auto lowerboundMask = rewriter.create( + loc, arith::CmpIPredicate::sge, blockifiedId, splatedZero); + Value blockifiedIdMask = + rewriter.create(loc, upperboundMask, lowerboundMask); + + blockifiedId = rewriter + .create( + loc, logicalBlockId.getType(), + ValueRange({blockifiedId, blockifiedIdMask})) + ->getResult(0); + + // replace program id to be computed from blockified id + SmallVector toReplace; + func.walk([&](triton::GetProgramIdOp id) { + if (id->hasAttr(logicalBlockIdAttr)) + return; + toReplace.push_back(id); + }); + for (auto id : toReplace) { + rewriter.setInsertionPoint(id); + Value newId; + if (id.getAxis() == triton::ProgramIDDim::X) { + newId = rewriter.create(id.getLoc(), blockifiedId, yzNum); + newId = rewriter.create(id.getLoc(), newId, xNum); + } else if (id.getAxis() == triton::ProgramIDDim::Y) { + newId = rewriter.create(id.getLoc(), blockifiedId, zNum); + newId = rewriter.create(id.getLoc(), newId, yNum); + } else { + newId = rewriter.create(id.getLoc(), blockifiedId, zNum); + } + rewriter.replaceOp(id, newId); + } + + // Create for loop for region ops + func.walk([&](Operation *op) { + if (op->hasAttr(autoBlockifyRegionOpAttr)) { + auto *newOp = createBlockifyLoop( + op, blockifiedId.getDefiningOp(), + logicalBlockId, logicalBlockNum, autoBlockifySize, rewriter); + newOp->removeAttr(autoBlockifyRegionOpAttr); + return WalkResult::skip(); + } + return WalkResult::advance(); + }); +} + +void AutoBlockifyPass::runOnOperation() { + if (autoBlockifySize == 1) + return; + ModuleOp moduleOp = getOperation(); + if (autoBlockifySize <= 0) { + moduleOp->emitWarning("[AutoBlockify V2] AutoBlockifySize cannot be " + "negative integer, skipping."); + return signalPassFailure(); + } + + MLIRContext *ctx = &getContext(); + + moduleOp.walk([&](triton::FuncOp func) { + LogicalResult result = success(); + func.walk([&](triton::GetProgramIdOp id) { + if (!checkBlockifiable(id.getResult())) { + result = failure(); + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + if (failed(result)) { + func->emitWarning("Cannot apply auto blockify"); + return WalkResult::skip(); + } + preProcess(func); + + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "After preprocess:\n" << func << "\n"; + }); + + RewritePatternSet patterns(ctx); + patterns.add( + ctx, logicalBlockId, logicalBlockNum, autoBlockifySize); + + if (failed(applyPatternsGreedily(func, std::move(patterns)))) { + moduleOp->emitError("failed to apply Patterns"); + signalPassFailure(); + return WalkResult::interrupt(); + } + + IRRewriter rewriter(ctx); + func->walk([&](UnrealizedConversionCastOp op) { + rewriter.setInsertionPoint(op); + auto input = op.getInputs()[0]; + auto resType = cast(op->getResultTypes()[0]); + if (auto constantOp = input.getDefiningOp()) { + Attribute val = constantOp.getValue(); + if (auto denseAttr = dyn_cast(val)) + val = denseAttr.getSplatValue(); + rewriter.replaceOpWithNewOp( + op, DenseElementsAttr::get(resType, val)); + } else if (auto tensorType = + dyn_cast(input.getType())) { + input = rewriter.create(input.getLoc(), input, 0); + rewriter.replaceOpWithNewOp(op, resType, input); + } else { + rewriter.replaceOpWithNewOp(op, resType, input); + } + }); + func->setAttr(autoBlockifySizeAttr, + rewriter.getI32IntegerAttr(autoBlockifySize)); + return WalkResult::skip(); + }); + + PassManager pm(&getContext(), moduleOp.getOperationName()); + pm.addPass(createCSEPass()); + pm.addPass(createCanonicalizerPass()); + if (failed(runPipeline(pm, moduleOp))) { + signalPassFailure(); + } +} + +std::unique_ptr> +triton::createAutoBlockifyPass(const AutoBlockifyOptions &options) { + return std::make_unique(options); +} \ No newline at end of file diff --git a/compiler/lib/AutoBlockify/CMakeLists.txt b/compiler/lib/AutoBlockify/CMakeLists.txt new file mode 100644 index 00000000..a0ccd59b --- /dev/null +++ b/compiler/lib/AutoBlockify/CMakeLists.txt @@ -0,0 +1,22 @@ +add_triton_library(AutoBlockify + AutoBlockify.cpp + RewriteOperation.cpp + Utils.cpp + + DEPENDS + AutoBlockifyPassIncGen + + LINK_LIBS PUBLIC + MLIRArithDialect + MLIRDialectUtils + MLIRIR + MLIRMathDialect + MLIRPass + MLIRTensorDialect + TritonIR + TritonTransforms + TritonAnalysis + MLIRTransforms + MLIRSupport + MLIRSCFTransforms +) \ No newline at end of file diff --git a/compiler/lib/AutoBlockify/RewriteOperation.cpp b/compiler/lib/AutoBlockify/RewriteOperation.cpp new file mode 100644 index 00000000..48246c61 --- /dev/null +++ b/compiler/lib/AutoBlockify/RewriteOperation.cpp @@ -0,0 +1,492 @@ + + +#include "dicp/AutoBlockify/AutoBlockify.h" +#include "dicp/AutoBlockify/Utils.h" +#include "dicp/Utils/Utils.h" + +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "auto-blockify-rewrite-operation" + +using namespace mlir; +using namespace triton; + +void PropagateUnrealizedCastDown::handleBlockifyLoop( + scf::ForOp blockifyLoop, Operation *op, PatternRewriter &rewriter) const { + SmallVector newOperands; + for (auto opr : op->getOperands()) { + auto uccOp = opr.getDefiningOp(); + if (!uccOp) { + newOperands.push_back(opr); + continue; + } + auto input = uccOp.getInputs()[0]; + auto tensorType = cast(input.getType()); + Value newOperand; + if (tensorType.getRank() > 1) { + SmallVector offsets(tensorType.getRank(), + rewriter.getIndexAttr(0)); + SmallVector sizes(1, rewriter.getIndexAttr(1)); + SmallVector strides(tensorType.getRank(), + rewriter.getIndexAttr(1)); + offsets[0] = blockifyLoop.getInductionVar(); + for (auto dim : llvm::drop_begin(tensorType.getShape())) + sizes.push_back(rewriter.getIndexAttr(dim)); + newOperand = rewriter.create( + input.getLoc(), cast(opr.getType()), input, offsets, + sizes, strides); + } else { + newOperand = rewriter.create( + input.getLoc(), input, ValueRange{blockifyLoop.getInductionVar()}); + if (isa(opr.getType())) { + newOperand = rewriter.create( + input.getLoc(), rewriter.getIndexType(), newOperand); + } + } + newOperands.push_back(newOperand); + } + rewriter.modifyOpInPlace(op, [&]() { op->setOperands(newOperands); }); +} + +void PropagateUnrealizedCastDown::rewriteGeneraleOp( + UnrealizedConversionCastOp op, Operation *generalOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto res = op->getResult(0); + auto inputType = cast(input.getType()); + SmallVector newOperands; + SmallVector newResults; + SmallVector newResultTypes; + + for (auto operand : generalOp->getOperands()) + newOperands.push_back(rewriteValue(operand, op, rewriter)); + for (auto resType : generalOp->getResultTypes()) { + newResultTypes.push_back(getExpandedType(resType, op)); + } + auto *newOp = + rewriter.create(generalOp->getLoc(), generalOp->getName().getIdentifier(), + newOperands, newResultTypes, generalOp->getAttrs()); + replaceValue(newOp, generalOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteSplat( + UnrealizedConversionCastOp op, triton::SplatOp splatOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto resType = cast(splatOp.getResult().getType()); + auto curShape = + llvm::to_vector(cast(input.getType()).getShape()); + auto splatedShape = resType.getShape(); + for (auto dim : splatedShape) { + input = rewriter.create(input.getLoc(), input, + curShape.size()); + curShape.push_back(dim); + input = rewriter.create( + input.getLoc(), + RankedTensorType::get(curShape, getElementTypeOrSelf(input)), input); + } + replaceValue(input.getDefiningOp(), splatOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteExpandDims( + UnrealizedConversionCastOp op, triton::ExpandDimsOp expandDimsOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto newOp = rewriter.create( + expandDimsOp.getLoc(), input, expandDimsOp.getAxis() + 1); + for (auto attr : expandDimsOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, expandDimsOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteReduce( + UnrealizedConversionCastOp op, triton::ReduceOp reduceOp, + PatternRewriter &rewriter) const { + auto mask = op.getInputs()[1]; + auto srcs = llvm::map_to_vector(reduceOp.getSrcs(), [&](Value src) { + return rewriteValue(src, op, rewriter); + }); + auto newOp = rewriter.create(reduceOp.getLoc(), srcs, + reduceOp.getAxis() + 1); + auto &newCombineOp = newOp.getCombineOp(); + rewriter.cloneRegionBefore(reduceOp.getCombineOp(), newCombineOp, + newCombineOp.end()); + for (auto attr : reduceOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, reduceOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteScan(UnrealizedConversionCastOp op, + triton::ScanOp scanOp, + PatternRewriter &rewriter) const { + auto mask = op.getInputs()[1]; + auto srcs = llvm::map_to_vector(scanOp.getSrcs(), [&](Value src) { + return rewriteValue(src, op, rewriter); + }); + auto newOp = rewriter.create( + scanOp.getLoc(), srcs, scanOp.getAxis() + 1, scanOp.getReverse()); + auto &newCombineOp = newOp.getCombineOp(); + rewriter.cloneRegionBefore(scanOp.getCombineOp(), newCombineOp, + newCombineOp.end()); + for (auto attr : scanOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, scanOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteLoad(UnrealizedConversionCastOp op, + triton::LoadOp loadOp, + PatternRewriter &rewriter) const { + auto uccMask = op.getInputs()[1]; + auto ptr = rewriteValue(loadOp.getPtr(), op, rewriter); + auto other = rewriteValue(loadOp.getOther(), op, rewriter); + auto mask = rewriteValue(loadOp.getMask(), op, rewriter); + auto res = loadOp.getResult(); + auto resType = getExpandedType(res.getType(), op); + if (!other) { + other = rewriter.create( + rewriter.getUnknownLoc(), + DenseElementsAttr::get( + resType, rewriter.getZeroAttr(getElementTypeOrSelf(res)))); + } + mask = createMask(mask, uccMask, resType.getShape(), rewriter); + auto boundaryCheck = llvm::map_to_vector(loadOp.getBoundaryCheck(), + [](int32_t idx) { return idx + 1; }); + auto newOp = rewriter.create( + loadOp.getLoc(), ptr, mask, other, boundaryCheck, loadOp.getPadding(), + loadOp.getCache(), loadOp.getEvict(), loadOp.getIsVolatile()); + for (auto attr : loadOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, loadOp, uccMask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteStore( + UnrealizedConversionCastOp op, triton::StoreOp storeOp, + PatternRewriter &rewriter) const { + auto uccMask = op.getInputs()[1]; + auto ptr = rewriteValue(storeOp.getPtr(), op, rewriter); + auto value = rewriteValue(storeOp.getValue(), op, rewriter); + auto mask = rewriteValue(storeOp.getMask(), op, rewriter); + auto ptrShape = cast(ptr.getType()).getShape(); + mask = createMask(mask, uccMask, ptrShape, rewriter); + auto boundaryCheck = llvm::map_to_vector(storeOp.getBoundaryCheck(), + [](int32_t idx) { return idx + 1; }); + auto newOp = rewriter.create( + storeOp.getLoc(), ptr, value, mask, boundaryCheck, storeOp.getCache(), + storeOp.getEvict()); + for (auto attr : storeOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(storeOp, newOp); +} + +void PropagateUnrealizedCastDown::rewriteAtomicRMW( + UnrealizedConversionCastOp op, triton::AtomicRMWOp atomicRMWOp, + PatternRewriter &rewriter) const { + auto uccMask = op.getInputs()[1]; + auto ptr = rewriteValue(atomicRMWOp.getPtr(), op, rewriter); + auto val = rewriteValue(atomicRMWOp.getVal(), op, rewriter); + auto mask = rewriteValue(atomicRMWOp.getMask(), op, rewriter); + auto resType = getExpandedType(atomicRMWOp.getResult().getType(), op); + mask = createMask(mask, uccMask, resType.getShape(), rewriter); + auto newOp = rewriter.create( + atomicRMWOp.getLoc(), resType, atomicRMWOp.getAtomicRmwOp(), ptr, val, + mask, atomicRMWOp.getSem(), atomicRMWOp.getScope()); + for (auto attr : atomicRMWOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, atomicRMWOp, uccMask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteAssert( + UnrealizedConversionCastOp op, triton::AssertOp assertOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto inputShape = cast(input.getType()).getShape(); + auto conditionType = cast(mask.getType()); + auto oneAttr = rewriter.getIntegerAttr(getElementTypeOrSelf(mask), 1); + auto one = rewriter.create( + mask.getLoc(), DenseElementsAttr::get(conditionType, oneAttr)); + Value condition = rewriter.create(input.getLoc(), mask, one); + condition = createMask(nullptr, condition, inputShape, rewriter); + condition = + rewriter.create(condition.getLoc(), condition, input); + auto newOp = rewriter.create(assertOp.getLoc(), condition, + assertOp.getMessage()); + for (auto attr : assertOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(assertOp, newOp); +} + +void PropagateUnrealizedCastDown::rewriteExtractSlice( + UnrealizedConversionCastOp op, tensor::ExtractSliceOp extractSliceOp, + PatternRewriter &rewriter) const { + auto mask = op.getInputs()[1]; + auto src = rewriteValue(extractSliceOp.getSource(), op, rewriter); + auto offsets = llvm::to_vector(extractSliceOp.getMixedOffsets()); + auto sizes = llvm::to_vector(extractSliceOp.getMixedSizes()); + auto strides = llvm::to_vector(extractSliceOp.getMixedStrides()); + auto srcType = cast(src.getType()); + offsets.insert(offsets.begin(), rewriter.getIndexAttr(0)); + sizes.insert(sizes.begin(), rewriter.getIndexAttr(srcType.getShape()[0])); + strides.insert(strides.begin(), rewriter.getIndexAttr(1)); + auto newOp = rewriter.create( + extractSliceOp.getLoc(), src, offsets, sizes, strides); + auto newMask = rewriter.create( + mask.getLoc(), mask, offsets, sizes, strides); + for (auto attr : extractSliceOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, extractSliceOp, newMask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteInsertSlice( + UnrealizedConversionCastOp op, tensor::InsertSliceOp insertSliceOp, + PatternRewriter &rewriter) const { + auto mask = op.getInputs()[1]; + auto src = rewriteValue(insertSliceOp.getSource(), op, rewriter); + auto dst = rewriteValue(insertSliceOp.getDest(), op, rewriter); + auto offsets = llvm::to_vector(insertSliceOp.getMixedOffsets()); + auto sizes = llvm::to_vector(insertSliceOp.getMixedSizes()); + auto strides = llvm::to_vector(insertSliceOp.getMixedStrides()); + auto srcType = cast(src.getType()); + offsets.insert(offsets.begin(), rewriter.getIndexAttr(0)); + sizes.insert(sizes.begin(), rewriter.getIndexAttr(srcType.getShape()[0])); + strides.insert(strides.begin(), rewriter.getIndexAttr(1)); + auto newOp = rewriter.create( + insertSliceOp.getLoc(), src, dst, offsets, sizes, strides); + for (auto attr : insertSliceOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + replaceValue(newOp, insertSliceOp, mask, rewriter); +} + +void PropagateUnrealizedCastDown::rewriteWhile( + UnrealizedConversionCastOp op, scf::WhileOp whileOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto res = op->getResult(0); + SmallVector indices; + SmallVector newInits; + IRMapping mapping; + for (auto [idx, init] : llvm::enumerate(whileOp.getInits())) { + if (init == res) { + indices.push_back(idx); + newInits.push_back(input); + } else { + newInits.push_back(init); + } + } + auto newOp = rewriter.create( + whileOp.getLoc(), whileOp->getResultTypes(), newInits, + [&](OpBuilder &b, Location loc, ValueRange args) { + mapRegionIterArg(mapping, whileOp.getBeforeArguments(), args, indices, + mask, b); + for (auto &bodyOp : *whileOp.getBeforeBody()) + b.clone(bodyOp, mapping); + }, + [&](OpBuilder &b, Location loc, ValueRange args) { + mapRegionIterArg(mapping, whileOp.getAfterArguments(), args, {}, mask, + b); + for (auto &bodyOp : whileOp.getAfterBody()->without_terminator()) + b.clone(bodyOp, mapping); + auto yieldOp = + cast(whileOp.getAfterBody()->getTerminator()); + mapYieldedValue(mapping, yieldOp, indices, op, b); + }); + for (auto attr : whileOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(whileOp, newOp); +} + +void PropagateUnrealizedCastDown::rewriteLoop(UnrealizedConversionCastOp op, + LoopLikeOpInterface loopOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto res = op->getResult(0); + SmallVector indices; + SmallVector newInits; + IRMapping mapping; + for (auto [idx, init] : llvm::enumerate(loopOp.getInits())) { + if (init == res) { + indices.push_back(idx); + newInits.push_back(input); + } else { + newInits.push_back(init); + } + } + LoopLikeOpInterface newOp; + if (auto forOp = dyn_cast(loopOp.getOperation())) { + newOp = rewriter.create( + forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), + forOp.getStep(), newInits, + [&](OpBuilder &b, Location loc, Value iv, ValueRange args) { + mapping.map(forOp.getInductionVar(), iv); + mapRegionIterArg(mapping, forOp.getRegionIterArgs(), args, indices, + mask, b); + for (auto &bodyOp : forOp.getBody()->without_terminator()) + b.clone(bodyOp, mapping); + auto yieldOp = cast(forOp.getBody()->getTerminator()); + mapYieldedValue(mapping, yieldOp, indices, op, b); + }); + for (auto attr : forOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + } else { + llvm_unreachable("Unhandled loopOp"); + } + replaceValue(newOp, loopOp, mask, rewriter, indices); +} + +void PropagateUnrealizedCastDown::rewriteIf(UnrealizedConversionCastOp &op, + scf::IfOp ifOp, + ArrayRef indices, + PatternRewriter &rewriter) const { + IRMapping mapping; + auto mask = op.getInputs()[1]; + auto thenBlockBuilder = [&](OpBuilder &b, Location loc) { + for (auto &bodyOp : *ifOp.thenBlock()) + b.clone(bodyOp, mapping); + }; + auto elseBlockBuilder = [&](OpBuilder &b, Location loc) { + for (auto &bodyOp : *ifOp.elseBlock()) + b.clone(bodyOp, mapping); + }; + scf::IfOp newOp; + if (ifOp.elseBlock()) { + newOp = rewriter.create(ifOp.getLoc(), ifOp.getCondition(), + thenBlockBuilder, elseBlockBuilder); + } else { + newOp = rewriter.create(ifOp.getLoc(), ifOp.getCondition(), + thenBlockBuilder, nullptr); + } + for (auto attr : ifOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + if (mapping.contains(op)) + op = cast(mapping.lookup(op)); + replaceValue(newOp, ifOp, mask, rewriter, indices); +} + +void PropagateUnrealizedCastDown::rewriteYield( + UnrealizedConversionCastOp &op, scf::YieldOp yieldOp, + PatternRewriter &rewriter) const { + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto res = op->getResult(0); + SmallVector indices; + auto newOperands = llvm::to_vector(yieldOp.getOperands()); + for (auto [idx, opr] : llvm::enumerate(newOperands)) { + if (opr == res) + indices.push_back(idx); + } + if (auto loopOp = dyn_cast(yieldOp->getParentOp())) { + auto uccOp = rewriter.create( + op.getLoc(), res.getType(), ValueRange({input})); + for (auto curIdx : indices) + newOperands[curIdx] = uccOp->getResult(0); + auto newOp = rewriter.create(yieldOp.getLoc(), newOperands); + for (auto attr : yieldOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(yieldOp, newOp); + rewriter.setInsertionPoint(loopOp); + for (auto curIdx : indices) { + auto &initArg = loopOp.getInitsMutable()[curIdx]; + auto initVal = initArg.get(); + uccOp = rewriter.create( + initVal.getLoc(), input.getType(), ValueRange({initVal})); + uccOp = rewriter.create( + initVal.getLoc(), initVal.getType(), + ValueRange({uccOp->getResult(0), mask})); + rewriter.modifyOpInPlace(loopOp, + [&]() { initArg.set(uccOp->getResult(0)); }); + } + } else if (auto ifOp = dyn_cast(yieldOp->getParentOp())) { + for (auto curIdx : indices) + newOperands[curIdx] = input; + auto newOp = rewriter.create(yieldOp.getLoc(), newOperands); + for (auto attr : yieldOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(yieldOp, newOp); + yieldOp = ifOp.thenYield() == yieldOp ? ifOp.elseYield() : ifOp.thenYield(); + if (yieldOp) { + rewriter.setInsertionPoint(yieldOp); + newOperands = llvm::to_vector(yieldOp.getOperands()); + for (auto curIdx : indices) { + auto uccOp = rewriter.create( + op.getLoc(), input.getType(), ValueRange({newOperands[curIdx]})); + newOperands[curIdx] = uccOp->getResult(0); + } + rewriter.replaceOpWithNewOp(yieldOp, newOperands); + } + rewriter.setInsertionPoint(ifOp); + rewriteIf(op, ifOp, indices, rewriter); + } +} + +void PropagateUnrealizedCastDown::rewriteCondition( + UnrealizedConversionCastOp op, scf::ConditionOp conditionOp, + PatternRewriter &rewriter) const { + auto whileOp = cast(conditionOp->getParentOp()); + auto input = op.getInputs()[0]; + auto mask = op.getInputs()[1]; + auto res = op->getResult(0); + int64_t curIdx = -1; + auto args = llvm::to_vector(conditionOp.getArgs()); + for (auto [idx, opr] : llvm::enumerate(args)) { + if (opr == res) + curIdx = idx; + } + args[curIdx] = input; + auto newOp = rewriter.create( + conditionOp.getLoc(), conditionOp.getCondition(), args); + for (auto attr : conditionOp->getAttrs()) { + if (!newOp->hasAttr(attr.getName())) + newOp->setAttr(attr.getName(), attr.getValue()); + } + rewriter.replaceOp(conditionOp, newOp); + + res = whileOp->getResult(curIdx); + auto oldResType = res.getType(); + auto newResType = getExpandedType(oldResType, op); + rewriter.modifyOpInPlace(whileOp, [&]() { res.setType(newResType); }); + rewriter.setInsertionPointAfter(whileOp); + auto newUccOp = rewriter.create( + res.getLoc(), oldResType, ValueRange({res, mask})); + rewriter.replaceAllUsesExcept(res, newUccOp->getResult(0), newUccOp); + auto arg = whileOp.getAfterArguments()[curIdx]; + auto oldArgType = arg.getType(); + auto newArgType = getExpandedType(oldArgType, op); + rewriter.modifyOpInPlace(whileOp, [&]() { arg.setType(newArgType); }); + rewriter.setInsertionPointToStart(whileOp.getAfterBody()); + newUccOp = rewriter.create( + arg.getLoc(), oldArgType, ValueRange({arg, mask})); + rewriter.replaceAllUsesExcept(arg, newUccOp->getResult(0), newUccOp); +} \ No newline at end of file diff --git a/compiler/lib/AutoBlockify/Utils.cpp b/compiler/lib/AutoBlockify/Utils.cpp new file mode 100644 index 00000000..5efe21c7 --- /dev/null +++ b/compiler/lib/AutoBlockify/Utils.cpp @@ -0,0 +1,191 @@ + + +#include "dicp/AutoBlockify/Utils.h" +#include "dicp/Utils/Utils.h" + +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "auto-blockify-utils" + +using namespace mlir; +using namespace triton; + +RankedTensorType getExpandedType(Type type, UnrealizedConversionCastOp op) { + auto target = op.getInputs()[0]; + auto targetType = cast(target.getType()); + SmallVector targetShape{targetType.getShape()[0]}; + if (auto valueType = dyn_cast(type)) { + targetShape.append(valueType.getShape().begin(), + valueType.getShape().end()); + } + return RankedTensorType::get(targetShape, getElementTypeOrSelf(type)); +} + +Value rewriteValue(Value value, UnrealizedConversionCastOp op, + OpBuilder &builder) { + if (value == nullptr) + return nullptr; + if (value == op->getResult(0)) + return op.getInputs()[0]; + return builder + .create( + value.getLoc(), getExpandedType(value.getType(), op), value) + ->getResult(0); +} + +void replaceValue(Operation *newOp, Operation *oldOp, Value newMask, + RewriterBase &rewriter, ArrayRef replaceIndices) { + int64_t idx = 0; + for (auto [res, oldRes] : + llvm::zip_equal(newOp->getResults(), oldOp->getResults())) { + if (replaceIndices.empty() || + llvm::find(replaceIndices, idx) != replaceIndices.end()) { + auto resType = res.getType(); + auto newUccOp = rewriter.create( + newOp->getLoc(), oldRes.getType(), ValueRange({res, newMask})); + rewriter.replaceAllUsesExcept(oldRes, newUccOp->getResult(0), newUccOp); + } else { + rewriter.replaceAllUsesWith(oldRes, res); + } + idx++; + } + rewriter.eraseOp(oldOp); +} + +Value createMask(Value mask, Value uccMask, ArrayRef targetShape, + RewriterBase &rewriter) { + SmallVector curShape{targetShape[0]}; + for (auto [idx, dim] : llvm::drop_begin(llvm::enumerate(targetShape))) { + curShape.push_back(dim); + uccMask = + rewriter.create(uccMask.getLoc(), uccMask, idx); + uccMask = rewriter.create( + uccMask.getLoc(), + RankedTensorType::get(curShape, getElementTypeOrSelf(uccMask)), + uccMask); + } + if (mask) { + mask = rewriter.create(mask.getLoc(), mask, uccMask); + } else { + mask = uccMask; + } + return mask; +} + +void mapRegionIterArg(IRMapping &mapping, ValueRange oldArgs, + ValueRange newArgs, ArrayRef indices, Value mask, + OpBuilder &builder) { + auto newArgIter = newArgs.begin(); + for (auto [idx, oldArg] : llvm::enumerate(oldArgs)) { + if (llvm::find(indices, idx) != indices.end()) { + auto newUccOp = builder.create( + oldArg.getLoc(), oldArg.getType(), ValueRange({*newArgIter, mask})); + mapping.map(oldArg, newUccOp->getResult(0)); + } else { + mapping.map(oldArg, *newArgIter); + } + ++newArgIter; + } +} + +void mapYieldedValue(IRMapping &mapping, scf::YieldOp yieldOp, + ArrayRef indices, UnrealizedConversionCastOp op, + OpBuilder &builder) { + SmallVector newOperands; + for (auto [idx, operand] : llvm::enumerate(yieldOp.getOperands())) { + operand = mapping.lookup(operand); + if (llvm::find(indices, idx) != indices.end()) + newOperands.push_back(rewriteValue(operand, op, builder)); + else + newOperands.push_back(operand); + } + builder.create(yieldOp.getLoc(), newOperands); +} + +Operation *createBlockifyLoop(Operation *targetOp, + UnrealizedConversionCastOp op, + Value logicalBlockId, Value logicalBlockNum, + int autoBlockifySize, RewriterBase &rewriter) { + auto loc = targetOp->getLoc(); + rewriter.setInsertionPoint(targetOp); + auto initVal = + rewriter.create(loc, rewriter.getIndexAttr(0)); + auto stepVal = + rewriter.create(loc, rewriter.getIndexAttr(1)); + auto blockifySizeVal = rewriter.create( + loc, rewriter.getIndexAttr(autoBlockifySize)); + Value upperBound = + rewriter.create(loc, logicalBlockNum, logicalBlockId); + auto i32Zero = + rewriter.create(loc, rewriter.getI32IntegerAttr(0)); + upperBound = rewriter.create(loc, upperBound, i32Zero); + upperBound = rewriter.create(loc, rewriter.getIndexType(), + upperBound); + upperBound = + rewriter.create(loc, upperBound, blockifySizeVal); + SmallVector inits; + if (auto loopOp = dyn_cast(targetOp)) { + inits = llvm::map_to_vector(loopOp.getInits(), + [&rewriter, &op](Value v) -> Value { + return rewriteValue(v, op, rewriter); + }); + } else { + auto resultTypes = + llvm::map_to_vector(targetOp->getResultTypes(), [&op](Type type) { + return getExpandedType(type, op); + }); + inits = + llvm::map_to_vector(resultTypes, [&rewriter, &loc](Type type) -> Value { + auto tensorType = cast(type); + return rewriter.create(loc, tensorType.getShape(), + tensorType.getElementType()); + }); + } + auto mask = op.getInputs()[1]; + Operation *newOp; + auto blockifyLoop = rewriter.create( + loc, initVal, upperBound, stepVal, inits, + [&](OpBuilder &b, Location loc, Value iv, ValueRange args) { + newOp = b.clone(*targetOp); + + SmallVector newResults; + for (auto [arg, res] : llvm::zip_equal(args, newOp->getResults())) { + auto tensorType = cast(arg.getType()); + auto rank = tensorType.getRank(); + Value newRes; + if (rank > 1) { + SmallVector offsets(tensorType.getRank(), + b.getIndexAttr(0)); + SmallVector sizes(1, b.getIndexAttr(1)); + SmallVector strides(tensorType.getRank(), + b.getIndexAttr(1)); + offsets[0] = iv; + for (auto dim : llvm::drop_begin(tensorType.getShape())) + sizes.push_back(b.getIndexAttr(dim)); + newRes = b.create(loc, res, arg, offsets, + sizes, strides); + } else { + newRes = b.create(loc, res, arg, ValueRange{iv}); + } + newResults.push_back(newRes); + } + b.create(loc, newResults); + }); + + replaceValue(blockifyLoop, targetOp, mask, rewriter); + blockifyLoop->setAttr(autoBlockifyLoopAttr, rewriter.getUnitAttr()); + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "After creating blockify loop:\n" << blockifyLoop << "\n"; + }); + return newOp; +} + +std::optional getBlockifyLoop(Operation *op) { + while (auto forOp = op->getParentOfType()) { + if (forOp->hasAttr(autoBlockifyLoopAttr)) + return forOp; + op = forOp; + } + return std::nullopt; +} \ No newline at end of file diff --git a/compiler/lib/CMakeLists.txt b/compiler/lib/CMakeLists.txt index 5cd7900f..1d2bc5f1 100644 --- a/compiler/lib/CMakeLists.txt +++ b/compiler/lib/CMakeLists.txt @@ -1 +1,44 @@ -dicp_add_all_subdirs() \ No newline at end of file +add_subdirectory(AutoBlockify) +add_subdirectory(AscendLegalize) +add_subdirectory(Dialect) +add_subdirectory(DiscreteMaskAccessConversion) +add_subdirectory(TritonAffinityOpt) +add_subdirectory(TritonToAnnotation) +add_subdirectory(TritonToGraph) +add_subdirectory(TritonToHFusion) +add_subdirectory(TritonToHIVM) +add_subdirectory(TritonToLLVM) +add_subdirectory(TritonToLinalg) +add_subdirectory(TritonToStructured) +add_subdirectory(TritonToUnstructure) +add_subdirectory(Utils) + +if(TRITON_ENABLE_COVERAGE_HITEST) + set(_instrument_targets + DiscreteMaskAccessConversion + TritonToAnnotation + TritonToHFusion + TritonToHIVM + TritonToLinalg + TritonToLLVM + TritonToStructured + TritonToUnstructure + MLIRTritonNPUUtils + TritonDicpIR + TritonStructuredIR + AutoBlockify + TritonAffinityOpt + ) + + foreach(_target ${_instrument_targets}) + if(TARGET ${_target}) + set_target_properties(${_target} PROPERTIES + RULE_LAUNCH_COMPILE "hitestwrapper" + RULE_LAUNCH_LINK "hitestwrapper" + ) + message(STATUS "Enabled hitestwrapper for target: ${_target}") + else() + message(WARNING "Target ${_target} not found, please check the actual target name") + endif() + endforeach() +endif() diff --git a/compiler/lib/Conversion/CMakeLists.txt b/compiler/lib/Conversion/CMakeLists.txt deleted file mode 100644 index 5cd7900f..00000000 --- a/compiler/lib/Conversion/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -dicp_add_all_subdirs() \ No newline at end of file diff --git a/compiler/lib/Conversion/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.cpp b/compiler/lib/Conversion/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.cpp deleted file mode 100644 index c82a378a..00000000 --- a/compiler/lib/Conversion/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.cpp +++ /dev/null @@ -1,181 +0,0 @@ -#include "dicp/Conversion/DiscreteMaskAccessConversion/MaskAnalysis.h" -#include "dicp/Conversion/DiscreteMaskAccessConversion/Passes.h" - -#include "dicp/Utils/Utils.h" -#include "mlir/IR/Attributes.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Transforms/DialectConversion.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/LogicalResult.h" - -bool compileOn91095Flag = false; -bool forceSimtTemplateFlag = false; - -namespace mlir { -namespace triton { -#define GEN_PASS_DEF_DISCRETEMASKACCESSCONVERSION -#include "dicp/Conversion/DiscreteMaskAccessConversion/Passes.h.inc" -} // namespace triton -} // namespace mlir - -using namespace mlir; -using namespace mlir::triton; -using namespace mlir::dicp; - -LogicalResult isDiscreteMask(Operation *op, Value mask, - PatternRewriter &rewriter) { - if (!mask) - return failure(); - - mlir::dicp::MaskState mstate; - auto isContMask = mstate.parse(mask, op->getLoc(), rewriter); - if (!isContMask.failed()) { - mstate.eraseInsertedOps(op, rewriter); - return failure(); - } - return success(); -} - -struct DiscreteMaskStoreConversion : OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(triton::StoreOp op, - PatternRewriter &rewriter) const final { - auto mask = op.getMask(); - auto loc = op.getLoc(); - auto dst = op.getPtr(); - auto src = op.getValue(); - - if (failed(isDiscreteMask(op, mask, rewriter))) - return failure(); - - auto loadFromDstOp = rewriter.create( - loc, dst, op.getCache(), op.getEvict(), false); - - auto selOp = rewriter.create(loc, mask, src, - loadFromDstOp.getResult()); - auto newStore = rewriter.create( - loc, dst, selOp, op.getCache(), op.getEvict()); - newStore->setAttr(mlir::dicp::discreteMaskAttrName, - UnitAttr::get(rewriter.getContext())); - rewriter.replaceOp(op, newStore); - return success(); - } -}; - -struct DiscreteMaskLoadConversion : OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(triton::LoadOp op, - PatternRewriter &rewriter) const final { - auto loc = op.getLoc(); - auto other = op.getOther(); - auto mask = op.getMask(); - auto ptr = op.getPtr(); - - if (failed(isDiscreteMask(op, mask, rewriter))) - return failure(); - if (compileOn91095Flag && forceSimtTemplateFlag) - return failure(); - - if (!other) { - FailureOr constant = specializeTypelessValueToConstant( - TypelessValue::Zero, ptr.getType(), loc, rewriter); - // TODO: fix me - if (failed(constant)) { - ptr.getType().dump(); - op->emitRemark() << " Unsupported type for constant creation"; - return failure(); - } - other = *constant; - } - - auto newLoadOp = rewriter.create( - loc, ptr, op.getCache(), op.getEvict(), op.getIsVolatile()); - auto discreteMaskOp = - rewriter.create(loc, mask, newLoadOp, other); - rewriter.replaceOp(op, discreteMaskOp); - return success(); - } -}; - -struct DiscreteMaskAtomicConversion : OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(triton::AtomicRMWOp op, - PatternRewriter &rewriter) const final { - auto loc = op.getLoc(); - auto ptr = op.getPtr(); - auto src = op.getVal(); - auto mask = op.getMask(); - auto rmwOp = op.getAtomicRmwOp(); - - if (failed(isDiscreteMask(op, mask, rewriter))) - return failure(); - - const std::map initMap = { - {RMWOp::FADD, TypelessValue::Zero}, - {RMWOp::ADD, TypelessValue::Zero}, - {RMWOp::UMAX, TypelessValue::Zero}, - {RMWOp::OR, TypelessValue::Zero}, - {RMWOp::MIN, TypelessValue::Max}, - {RMWOp::UMIN, TypelessValue::Max}, - {RMWOp::AND, TypelessValue::Max}, - {RMWOp::MAX, TypelessValue::Min}, - {RMWOp::XOR, TypelessValue::Zero}, - {RMWOp::XCHG, TypelessValue::Undefined}, - }; - assert(initMap.find(rmwOp) != initMap.end()); - auto typelessVal = initMap.at(rmwOp); - if (typelessVal == TypelessValue::Undefined) { - // Undefined default value atomic op will be decomposed in AscendNPU-IR - op->setAttr(mlir::dicp::discreteMaskAttrName, - UnitAttr::get(rewriter.getContext())); - return failure(); - } - - FailureOr fill = specializeTypelessValueToConstant( - typelessVal, src.getType(), loc, rewriter); - if (failed(fill)) - op->emitError("Unsupported atomic operation."); - - auto maskedValue = rewriter.create(loc, mask, src, *fill); - auto newAtomicOp = rewriter.create( - loc, src.getType(), rmwOp, ptr, maskedValue, mlir::Value(), op.getSem(), - op.getScope()); - rewriter.replaceOp(op, newAtomicOp); - return success(); - } -}; - -struct DiscreteMaskAccessConversionPass - : mlir::triton::impl::DiscreteMaskAccessConversionBase< - DiscreteMaskAccessConversionPass> { - - DiscreteMaskAccessConversionPass( - const DiscreteMaskAccessConversionOptions &options) - : DiscreteMaskAccessConversionBase(options) {} - - void runOnOperation() override { - compileOn91095Flag = this->compileOn91095; - forceSimtTemplateFlag = this->forceSimtTemplate; - - auto moduleOp = getOperation(); - - RewritePatternSet patterns(&getContext()); - patterns.add(patterns.getContext()); - if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { - moduleOp->emitError("failed to apply discrete mask access patterns"); - signalPassFailure(); - } - } -}; - -std::unique_ptr> -mlir::triton::createDiscreteMaskAccessConversionPass( - const DiscreteMaskAccessConversionOptions &options) { - return std::make_unique(options); -} diff --git a/compiler/lib/Conversion/LinalgToLinked/CMakeLists.txt b/compiler/lib/Conversion/LinalgToLinked/CMakeLists.txt deleted file mode 100644 index 7bc99c81..00000000 --- a/compiler/lib/Conversion/LinalgToLinked/CMakeLists.txt +++ /dev/null @@ -1,32 +0,0 @@ -add_triton_library(LinalgToLinked - LinalgToLinkedPass.cpp - DebugCPUVerifyPass.cpp - VerifyNoLinalgGenericPass.cpp - TritonOpConverter.cpp - - DEPENDS - DICPNPUIncGen - LinalgToLinkedConversionPassIncGen - - LINK_LIBS PUBLIC - BiShengIRAnnotationDialect - BiShengIRHIVMDialect - TritonTilingExtIR - MLIRArithDialect - MLIRDialectUtils - MLIRIR - MLIRMathDialect - MLIRPass - MLIRTensorDialect - MLIRTransforms - MLIRSupport - TritonAnalysis - TritonIR - TritonTransforms - TritonSharedAnalysis - DICPNPU - - TritonArithToLinalg - StructuredToMemref - TritonToStructured -) diff --git a/compiler/lib/Conversion/LinalgToLinked/DebugCPUVerifyPass.cpp b/compiler/lib/Conversion/LinalgToLinked/DebugCPUVerifyPass.cpp deleted file mode 100644 index 8b45cf09..00000000 --- a/compiler/lib/Conversion/LinalgToLinked/DebugCPUVerifyPass.cpp +++ /dev/null @@ -1,73 +0,0 @@ -#include "dicp/Conversion/LinalgToLinked/Passes.h" - -#include "bishengir/Dialect/Annotation/IR/Annotation.h" -#include "bishengir/Dialect/HIVM/IR/HIVM.h" -#include "triton-shared/Dialect/TritonTilingExt/IR/TritonTilingExtDialect.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include "mlir/IR/BuiltinOps.h" -#include "mlir/Pass/Pass.h" -#include "llvm/Support/Debug.h" - -#define DEBUG_TYPE "debug-cpu-verify" - -using namespace mlir; - -#define GEN_PASS_CLASSES -#include "dicp/Conversion/LinalgToLinked/Passes.h.inc" - -namespace { - -/// Returns true if the operation belongs to an external (non-MLIR-upstream) -/// dialect that should have been lowered away before CPU verification. -static bool isExternalDialectOp(Operation *op) { - Dialect *dialect = op->getDialect(); - if (!dialect) - return false; - return isa( - dialect); -} - -/// Returns true if the operation is a hivm.hir.sync_block_* op that should be -/// removed before CPU verification. -static bool isSyncBlockOp(Operation *op) { - return isa(op); -} - -class DebugCPUVerifyPass : public DebugCPUVerifyBase { -public: - void runOnOperation() override { - // First pass: remove all hivm.hir.sync_block_* operations - SmallVector syncBlockOpsToErase; - getOperation()->walk([&](Operation *op) { - if (isSyncBlockOp(op)) - syncBlockOpsToErase.push_back(op); - }); - for (Operation *op : syncBlockOpsToErase) - op->erase(); - - // Second pass: verify no external dialect operations remain - bool failed = false; - getOperation()->walk([&](Operation *op) { - if (!isExternalDialectOp(op)) - return; - op->emitError() << "external dialect op '" << op->getName() - << "' must be lowered before CPU verification"; - failed = true; - }); - - if (failed) - return signalPassFailure(); - - LLVM_DEBUG(llvm::dbgs() << "[debug-cpu-verify] PASSED — no external " - "dialect operations found\n"); - } -}; - -} // namespace - -std::unique_ptr> -mlir::dicp::linked::createDebugCPUVerifyPass() { - return std::make_unique(); -} diff --git a/compiler/lib/Conversion/LinalgToLinked/LinalgToLinkedPass.cpp b/compiler/lib/Conversion/LinalgToLinked/LinalgToLinkedPass.cpp deleted file mode 100644 index bf459391..00000000 --- a/compiler/lib/Conversion/LinalgToLinked/LinalgToLinkedPass.cpp +++ /dev/null @@ -1,567 +0,0 @@ -#include "bishengir/Dialect/Annotation/IR/Annotation.h" - -#include "dicp/Conversion/LinalgToLinked/LinalgToLinked.h" -#include "dicp/Conversion/LinalgToLinked/TritonOpConverter.h" - -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/Dialect/Utils/StaticValueUtils.h" -#include "mlir/IR/Location.h" -#include "mlir/IR/OperationSupport.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/IR/TypeUtilities.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" -#include "mlir/Transforms/Passes.h" - -#include "triton/Dialect/Triton/IR/Dialect.h" - -#define DEBUG_TYPE "linalg-to-linked" -#include "llvm/Support/Debug.h" -#include - -using namespace mlir; -using namespace dicp; -using namespace linked; - -#define GEN_PASS_CLASSES -#include "dicp/Conversion/LinalgToLinked/Passes.h.inc" - -namespace { - -const std::string globalKernelAttr = "global_kernel"; -const std::string kernelMixModeName = "mix_mode"; -inline constexpr unsigned getMaxEnumValForProgramIDDim() { return 2; } -static auto constexpr LAUNCH_GRID_RANK = getMaxEnumValForProgramIDDim() + 1; -static unsigned int constexpr TRITON_PROGRAM_INFO_ARG_COUNT = - LAUNCH_GRID_RANK * 2; - -class TritonTypeConverter : public mlir::TypeConverter { -public: - explicit TritonTypeConverter() { - addConversion([](Type type) { return type; }); - - addConversion([](triton::PointerType ptrType) { - return MemRefType::get({ShapedType::kDynamic}, ptrType.getPointeeType()); - }); - - addConversion([](TensorType tensorType) -> Type { - auto elemType = tensorType.getElementType(); - if (auto ptrType = dyn_cast(elemType)) { - elemType = ptrType.getPointeeType(); - } - return MemRefType::get(tensorType.getShape(), elemType); - }); - } -}; - -// TritonTypeConverter::TritonTypeConverter() { -// addConversion([](Type type) { return type; }); - -// addConversion([](triton::PointerType ptrType) { -// return MemRefType::get({ShapedType::kDynamic}, ptrType.getPointeeType()); -// }); - -// addConversion([](TensorType tensorType) -> Type { -// auto elemType = tensorType.getElementType(); -// if (auto ptrType = dyn_cast(elemType)) { -// elemType = ptrType.getPointeeType(); -// } -// return MemRefType::get(tensorType.getShape(), elemType); -// }); -// } - -static LogicalResult convertMultipleBlockControlFlow(Operation *funcOp, - OpBuilder &builder) { - if (!isa(funcOp)) { - funcOp->emitError( - "convertMultipleBlockControlFlow can only process func::FuncOp!"); - return failure(); - } - - SmallVector candidate; - SmallVector eraseBlocks; - for (Block &block : dyn_cast(funcOp).getBody()) { - auto curTerminator = block.getTerminator(); - if (isa(curTerminator)) - candidate.push_back(curTerminator); - else if (isa(curTerminator)) { - if (candidate.empty()) { - curTerminator->emitError( - "funcOp has more than one Block but got a early 'tt.return' Op."); - return failure(); - } - } else - return failure(); - - if (!block.isEntryBlock()) - eraseBlocks.push_back(&block); - } - - if (candidate.empty()) { - funcOp->emitError("funcOp has more than one Block but no candidate " - "Terminator was found!"); - return failure(); - } - - llvm::BitVector visitFlag(candidate.size(), false); - - // Recursive function to convert all cf::CondBranchOp to scf::IfOp - std::function convertToSCF = - [&](Operation *op, Operation *insertPosOp) -> void { - auto condBranchOp = dyn_cast_if_present(op); - auto iter = llvm::find(candidate, condBranchOp); - if (!(condBranchOp && iter != candidate.end())) { - op->emitError( - "convertToSCF must process with condBranchOp in candidates!"); - return; - } - visitFlag.set(iter - candidate.begin()); - - OpBuilder::InsertionGuard guard(builder); - builder.setInsertionPointAfter(insertPosOp); - - // Well, here force to destory original control flow - builder.create( - condBranchOp->getLoc(), condBranchOp.getCondition(), - /*thenBuilder=*/ - [&](OpBuilder &builder, Location loc) { - SmallVector movedOps = llvm::map_to_vector( - condBranchOp.getTrueDest()->without_terminator(), - [](Operation &op) { return &op; }); - for (auto *innerOp : movedOps) { - innerOp->moveBefore(builder.getInsertionBlock(), - builder.getInsertionPoint()); - } - - auto blockTerm = condBranchOp.getTrueDest()->getTerminator(); - if (isa(blockTerm)) { - if (movedOps.empty()) { - blockTerm->emitError( - "movedOps can not be empty before entering convertToSCF!"); - return; - } - convertToSCF(blockTerm, movedOps.back()); - } - - builder.create(loc); - }, - /*elseBuilder=*/ - [&](OpBuilder &builder, Location loc) { - SmallVector movedOps = llvm::map_to_vector( - condBranchOp.getFalseDest()->without_terminator(), - [](Operation &op) { return &op; }); - for (auto *innerOp : movedOps) { - innerOp->moveBefore(builder.getInsertionBlock(), - builder.getInsertionPoint()); - } - - auto blockTerm = condBranchOp.getFalseDest()->getTerminator(); - if (isa(blockTerm)) { - if (movedOps.empty()) { - blockTerm->emitError( - "movedOps can not be empty before entering convertToSCF!"); - return; - } - convertToSCF(blockTerm, movedOps.back()); - } - - builder.create(loc); - }); - }; - - Block::iterator insertOp(candidate.front()); - --insertOp; - convertToSCF(candidate.front(), &(*insertOp)); - - if (!visitFlag.all()) - return failure(); - - OpBuilder::InsertionGuard guard(builder); - builder.setInsertionPoint(candidate.front()); - builder.create(candidate.front()->getLoc()); - - for (Operation *eachTerm : candidate) - eachTerm->erase(); - for (Block *block : llvm::reverse(eraseBlocks)) - block->erase(); - - return success(); -} - -void convertFuncFunc(func::FuncOp func, const bool existDot) { - OpBuilder builder(func); - - auto name = func.getName(); - auto type = func.getFunctionType(); - - SmallVector argAttrs, resAttrs; - func.getAllArgAttrs(argAttrs); - func.getAllResultAttrs(resAttrs); - - // bit-casted tt.ptr的特殊处理 - SmallVector inputTypes{type.getInputs()}; - SmallVector retTypes{type.getResults()}; - if (func.getSymVisibility() == "public" && !func.isDeclaration()) { - for (size_t i = 0; i < func.getNumArguments(); ++i) { - auto arg = func.getArgument(i); - // Special method for i1 arg - if (!isa(arg.getType()) || - dyn_cast(arg.getType()).getElementTypeBitWidth() != - 1) { - continue; - } - - SmallVector argVaildUser{arg.getUsers()}; - llvm::erase_if(argVaildUser, [](Operation *op) -> bool { - return isOpTriviallyDead(op); - }); - - if (!argVaildUser.empty()) { - LLVM_DEBUG({ - auto &os = llvm::dbgs(); - os << arg << " has users:\n"; - int cnt = 0; - for (auto it : argVaildUser) { - os << "users[" << cnt++ << "] = " << *it; - } - }); - if (llvm::all_of(argVaildUser, [](Operation *userOp) { - return isa(userOp); - })) { - auto castOp = cast(*argVaildUser.begin()); - if (castOp.getInputs().size() == 1 && - castOp.getOutputs().size() == 1) { - arg.setType(castOp.getOutputs()[0].getType()); - inputTypes[i] = arg.getType(); - } - } else { - func->emitError(Twine("Unsupported use of func arg at index ") + - Twine(i)); - } - } else { - // Process unused bool ptr type specially, which guarantees bool pointer - // argument's type is realistic and don't mislead backend compiler. - // realistic memory layout of bool pointer is 8 bit width - auto memType = dyn_cast(arg.getType()) - .cloneWith(std::nullopt, builder.getI8Type()); - arg.setType(memType); - inputTypes[i] = arg.getType(); - } - } - } - auto castType = FunctionType::get(func.getContext(), inputTypes, retTypes); - - auto funcFunc = builder.create(func.getLoc(), name, castType); - funcFunc.setAllArgAttrs(argAttrs); - funcFunc.setAllResultAttrs(resAttrs); - auto kernelAttr = func->getAttr(globalKernelAttr); - if (kernelAttr) { - funcFunc->setAttr(globalKernelAttr, kernelAttr); - } - std::string kernelMixMode = "aiv"; - if (existDot) { - // mix also works for pure cube kernel by using the same MAGIC_ELF keyword - kernelMixMode = "mix"; - } - // Set mix_mode in the func attrs so that the backend could know - // the mix_mode by parse the func attrs. - // The backend needs to know the mix_mode because the host wrapper - // needs to set the devbin.magic. Check npu_utils.cpp. - funcFunc->setAttr(kernelMixModeName, builder.getStringAttr(kernelMixMode)); - - auto &funcFuncBody = funcFunc.getBody(); - auto &funcBody = func.getBody(); - - IRMapping map; - funcBody.cloneInto(&funcFuncBody, map); - - // 消除多个return，否则会在one-shot buffer报错 - if (!funcFuncBody.hasOneBlock()) { - if (failed(convertMultipleBlockControlFlow(funcFunc, builder))) { - llvm_unreachable("Encounter unsupported control flow"); - } - } - - for (Block &block : funcFuncBody.getBlocks()) { - auto term = block.getTerminator(); - builder.setInsertionPoint(term); - builder.create(func.getLoc(), term->getOperands()); - term->erase(); - } - func.erase(); -} - -void addProgramInfo(func::FuncOp func, bool globalKernel) { - OpBuilder b(func); - - auto origFuncType = func.getFunctionType(); - auto origInputTypes = origFuncType.getInputs(); - SmallVector newInputTypes(origInputTypes); - - if (globalKernel) { - func->setAttr(globalKernelAttr, b.getStringAttr("")); - } else { - func->setAttr(globalKernelAttr, b.getStringAttr("local")); - } -} - -class TritonAnnotationConverter - : public OpConversionPattern { -public: - using OpConversionPattern::OpConversionPattern; - - LogicalResult matchAndRewrite(triton::AnnotationOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const { - auto markOp = rewriter.create(op.getLoc(), op.getSrc()); - // Forward all annotations. - markOp->setAttrs(op->getAttrs()); - rewriter.eraseOp(op); - return success(); - } -}; - -class ExternElementwiseClOpConverter - : public OpConversionPattern { -public: - using OpConversionPattern::OpConversionPattern; - - LogicalResult matchAndRewrite(triton::ExternElementwiseOp op, - OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const { - auto loc = op.getLoc(); - if (!op.getPure()) { - op->emitWarning() << "impure elementwise op!"; - return failure(); - } - if (op.getSymbol().contains("__hmf_")) { - // 1. get or create the declaration of external elementwise function - Type dstTy = op.getResult().getType(); - bool isDstScalar = !isa(dstTy); - Type dstElemTy = - isDstScalar ? dstTy : cast(dstTy).getElementType(); - SmallVector srcElemTys; - SmallVector srcs; - for (auto src : op.getSrcs()) { - if (!isa(src.getType())) { - src = rewriter.create( - op.getLoc(), RankedTensorType::get({(int64_t)1}, src.getType()), - src); - } - srcs.push_back(src); - srcElemTys.push_back( - cast(src.getType()).getElementType()); - } - - FunctionType elemFuncType = - FunctionType::get(rewriter.getContext(), srcElemTys, {dstElemTy}); - auto mod = SymbolTable::getNearestSymbolTable(op); - auto extFunc = dyn_cast_or_null( - SymbolTable::lookupSymbolIn(mod, op.getSymbol())); - if (!extFunc) { - OpBuilder::InsertionGuard guard(rewriter); - rewriter.setInsertionPointToStart(&mod->getRegion(0).front()); - extFunc = rewriter.create(rewriter.getUnknownLoc(), - op.getSymbol(), elemFuncType); - extFunc.setPrivate(); - extFunc->setAttr(LLVM::LLVMDialect::getReadnoneAttrName(), - UnitAttr::get(rewriter.getContext())); - } - assert(isa( - SymbolTable::lookupSymbolIn(mod, op.getSymbol()))); - // 2. prepare the output tensor - Value output; - if (isDstScalar) { - dstTy = RankedTensorType::get({(int64_t)1}, dstElemTy); - } - bool found = false; - for (Value v : srcs) { - if (v.getType() == dstTy) { - found = true; - output = v; - break; - } - } - if (!found) { - output = rewriter.create( - op.getLoc(), cast(dstTy).getShape(), dstElemTy); - } - // 3. create the linalg.map op - auto mapOp = rewriter.create( - loc, - /*inputs=*/srcs, - /*init=*/output, - /*bodyBuilder=*/ - [&](OpBuilder &builder, Location loc, ValueRange regionArgs) { - auto elemOp = builder.create(loc, - /*name=*/op.getSymbol(), - /*resultType=*/dstElemTy, - /*operands=*/regionArgs); - builder.create(loc, elemOp->getResults()); - }); - if (isDstScalar) { - // need to convert tensor back to scalar - auto indexType = rewriter.getIndexType(); - Value zeroConstant = rewriter.create( - loc, indexType, rewriter.getIntegerAttr(indexType, 0)); - auto extractOp = rewriter.create( - loc, mapOp.getResults()[0], zeroConstant); - rewriter.replaceOp(op, extractOp); - } else { - rewriter.replaceOp(op, mapOp); - } - return success(); - } - return failure(); - } -}; - -class LinalgToLinkedPass : public LinalgToLinkedBase { -public: - explicit LinalgToLinkedPass(bool globalKernel, bool namedOps, - bool cpuVerify) { - this->globalKernel = globalKernel; - this->namedOps = namedOps; - this->cpuVerify = cpuVerify; - } - - void getDependentDialects(DialectRegistry ®istry) const override { - registry - .insert(); - } - - void populateLinalgToLinkedConversionPatterns(TypeConverter &typeConverter, - RewritePatternSet &patterns) { - populateFunctionOpInterfaceTypeConversionPattern( - patterns, typeConverter); - - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - if (!this->namedOps) { - linalg::populateElementwiseToLinalgConversionPatterns(patterns); - } - } - - void runOnOperation() override { - auto moduleOp = getOperation(); - MLIRContext *context = &getContext(); - RewritePatternSet patterns(&getContext()); - ConversionTarget target(getContext()); - TritonTypeConverter tritonTypeConverter{}; - - target.addLegalDialect< - func::FuncDialect, arith::ArithDialect, math::MathDialect, - linalg::LinalgDialect, affine::AffineDialect, scf::SCFDialect, - cf::ControlFlowDialect, tensor::TensorDialect, LLVM::LLVMDialect, - bufferization::BufferizationDialect, memref::MemRefDialect, - annotation::AnnotationDialect>(); - target.addLegalOp(); - // 根据条件判断需要转换的OP - target.addDynamicallyLegalOp( - [](mlir::Operation *op) { - if (op->use_empty()) { - return false; - } else { - return true; - } - }); - - target.addDynamicallyLegalOp([&](triton::FuncOp op) { - return tritonTypeConverter.isSignatureLegal(op.getFunctionType()); - }); - // Check if the kernel contains tl.dot. Without tl.dot, - // the kernel would be pure AIV kernel. - bool existDot = false; - moduleOp.walk([&](linalg::MatmulOp dotOp) { - existDot = true; - return WalkResult::interrupt(); - }); - this->populateLinalgToLinkedConversionPatterns(tritonTypeConverter, - patterns); - size_t tritonFuncCount = 0; - for (auto func : getOperation().getOps()) { - ++tritonFuncCount; - } - - size_t funcOpCount = 0; - // 遍历 func::FuncOp 操作 - for (auto func : getOperation().getOps()) { - ++funcOpCount; - } - - // 遍历kernel中的function，修改program id、number of programs参数 - for (auto func : getOperation().getOps()) { - addProgramInfo(func, globalKernel); - } - // 函数头尾转换 - moduleOp.walk([&](func::FuncOp func) { convertFuncFunc(func, existDot); }); - - PassManager pm(context); - if (failed(pm.run(moduleOp))) { - signalPassFailure(); - } - - // Insert NPU workspace args unless in CPU verify mode - if (failed(insertWorkspaceArgs(moduleOp))) { - signalPassFailure(); - return; - } - - target.addIllegalOp(); - target.addIllegalOp(); - if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) { - moduleOp->emitError("failed to apply Convertion Patterns"); - signalPassFailure(); - } - } - -private: - /// Inserts syncBlockLock and workspace args when not in CPU verify mode. - LogicalResult insertWorkspaceArgs(ModuleOp module) { - if (cpuVerify) - return success(); - - MLIRContext *ctx = module.getContext(); - auto memrefType = - MemRefType::get({ShapedType::kDynamic}, IntegerType::get(ctx, 8)); - - for (auto func : module.getOps()) { - if (!func->hasAttr(globalKernelAttr)) - continue; - - if (failed(func.insertArgument(0, memrefType, nullptr, func.getLoc()))) - return func.emitError("failed to insert syncBlockLock"); - func->setAttr("SyncBlockLockArgIdx", - IntegerAttr::get(IntegerType::get(ctx, 64), 0)); - - if (failed(func.insertArgument(1, memrefType, nullptr, func.getLoc()))) - return func.emitError("failed to insert workspace"); - func->setAttr("WorkspaceArgIdx", - IntegerAttr::get(IntegerType::get(ctx, 64), 1)); - } - return success(); - } -}; - -} // namespace - -std::unique_ptr> -linked::createLinalgToLinkedPass(bool globalKernel, bool namedOps, - bool cpuVerify) { - return std::make_unique(globalKernel, namedOps, - cpuVerify); -} diff --git a/compiler/lib/Conversion/LinalgToLinked/TritonOpConverter.cpp b/compiler/lib/Conversion/LinalgToLinked/TritonOpConverter.cpp deleted file mode 100644 index f69ea288..00000000 --- a/compiler/lib/Conversion/LinalgToLinked/TritonOpConverter.cpp +++ /dev/null @@ -1,791 +0,0 @@ -#include "dicp/Conversion/LinalgToLinked/TritonOpConverter.h" -#include - -#include "llvm/ADT/SmallVectorExtras.h" -#include "llvm/ADT/TypeSwitch.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/ErrorHandling.h" -#include "llvm/Support/LogicalResult.h" -#include "llvm/Support/raw_ostream.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/MemRef/Transforms/Passes.h" -#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/ValueRange.h" - -#include "triton-shared/Analysis/MaskAnalysis.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include -#define DEBUG_TYPE "triton-to-linked-npu" - -using namespace mlir; -using namespace mlir::dicp::linked; - -/// This function generates a series of `scf.for` loops for the given dimensions -/// in `loopDims`. Although the loops are created sequentially, nesting is -/// simulated by adjusting the insertion point to the body of the last created -/// loop. This allows the `bodyFunc` to be inserted into the innermost scope. -/// -/// \param rewriter The MLIR OpBuilder used to create operations. -/// \param loc The source location information for debuggability. -/// \param target The memref value whose dimensions are being looped over. -/// \param loopDims An array of dimension indices to create loops for. -/// \param bodyFunc A callable that defines the operations to insert in the -/// innermost loop. It takes a SmallVector of induction variables (one per -/// loop). -/// -template -static void createSimpleNestedLoops(OpBuilder &rewriter, Location loc, - Value target, ArrayRef loopDims, - Func bodyFunc) { - MemRefType type = cast(target.getType()); - int rank = type.getRank(); - - Value zero = rewriter.create(loc, 0); - Value one = rewriter.create(loc, 1); - - llvm::SmallVector loops; - llvm::SmallVector ivs; - - for (int dim : loopDims) { - Value ub; - if (type.isDynamicDim(dim)) { - ub = rewriter.create(loc, target, dim).getResult(); - } else { - ub = rewriter.create(loc, type.getDimSize(dim)); - } - - auto forOp = rewriter.create(loc, zero, ub, one); - rewriter.setInsertionPointToStart(forOp.getBody()); - loops.push_back(forOp); - ivs.push_back(forOp.getInductionVar()); - } - - bodyFunc(ivs); - - if (!loops.empty()) { - rewriter.setInsertionPointAfter(loops.front()); - } -} - -bool ReduceConverter::isReductionOpSupported(Operation *redOp) const { - return isa(redOp); -} - -LogicalResult -ReduceConverter::convertToTargetOp(triton::ReduceOp op, - typename triton::ReduceOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const { - auto source = adaptor.getOperands().front(); - auto sourceType = cast(source.getType()); - auto elemType = sourceType.getElementType(); - auto resType = op.getResult().front().getType(); - auto loc = op.getLoc(); - auto reductionOps = this->getRedOps(op); - - // Reduction of arbitrary operations isn't supported because using the first - // element across the reduction dimension requires us to iterate over a - // subview that skips over each first element. - if (!this->isReductionOpSupported(reductionOps.front())) { - return rewriter.notifyMatchFailure( - op, "Only support lowering reduction with single op and limited types " - "of reducetion"); - } - - auto rop = reductionOps.front(); - auto axis = op.getAxis(); - auto isVectorReduce = sourceType.getRank() == 1; - - auto constantType = elemType; - - auto accBaseConstOp = this->getRedBaseConstOp(rewriter, rop, constantType); - Value initTensor; - - if (isVectorReduce) { - auto holder = rewriter.create( - loc, RankedTensorType::get({}, constantType), ValueRange{}); - initTensor = rewriter - .create(loc, accBaseConstOp.getResult(), - holder.getResult()) - .getResult(0); - } else { - Value init = rewriter.create( - loc, cast(resType).getShape(), constantType); - initTensor = - rewriter.create(loc, accBaseConstOp.getResult(), init) - .getResult(0); - } - - Value finalResult = - rewriter - .create( - loc, ValueRange{source}, ValueRange{initTensor}, - SmallVector{axis}, - [&](OpBuilder &opBuilder, Location loc, ValueRange inputs) { - assert(inputs.size() == 2); - Value result = this->getRedElement(inputs[0], inputs[1], loc, - rop, opBuilder, false); - opBuilder.create(loc, result); - }) - .getResult(0); - - if (sourceType.getRank() == 1) { - finalResult = - rewriter.create(loc, constantType, finalResult); - } - - rewriter.replaceOp(op, finalResult); - return success(); -} - -static LogicalResult -addReduceWithIndexAttrIfNeeded(ConversionPatternRewriter &rewriter, - linalg::ReduceOp reduceOp) { - // To verify whether the operation of the reduceOp is ReduceWithIndex - // TODO: maybe a better way of judging? - Block &body = reduceOp.getCombiner().front(); - auto yieldOp = dyn_cast(body.getTerminator()); - - auto yieldValue = yieldOp.getValues(); - if (yieldValue.size() == 0) { - return failure(); - } - - const StringRef reduceRef = "reduce_mode"; - const StringRef tieBreakLeftRef = "tie_break_left"; - // INT - // Composite predicate to pick index of min (or max) element have to be - // written in following form: value1 < value2 or (value1 == value2 and index1 - // < index2) - for leftmost element (value1 == value2 and index1 < index2) or - // value1 < value2 - for leftmost element value1 < value2 or (value1 == value2 - // and index1 > index2) - for rightmost element (value1 == value2 and index1 > - // index2) or value1 < value2 - for rightmost element table below encodes all - // possible cases of sequences of predicates for min/max and - // leftmost/rightmost elements - std::map, - std::pair> - m{ - {{arith::CmpIPredicate::eq, arith::CmpIPredicate::sgt, - arith::CmpIPredicate::sgt}, - {"max_with_index", "false"}}, - {{arith::CmpIPredicate::eq, arith::CmpIPredicate::sgt, - arith::CmpIPredicate::slt}, - {"min_with_index", "false"}}, - {{arith::CmpIPredicate::eq, arith::CmpIPredicate::slt, - arith::CmpIPredicate::sgt}, - {"max_with_index", "true"}}, - {{arith::CmpIPredicate::eq, arith::CmpIPredicate::slt, - arith::CmpIPredicate::slt}, - {"min_with_index", "true"}}, - {{arith::CmpIPredicate::sgt, arith::CmpIPredicate::eq, - arith::CmpIPredicate::sgt}, - {"max_with_index", "false"}}, - {{arith::CmpIPredicate::slt, arith::CmpIPredicate::eq, - arith::CmpIPredicate::sgt}, - {"min_with_index", "false"}}, - {{arith::CmpIPredicate::sgt, arith::CmpIPredicate::eq, - arith::CmpIPredicate::slt}, - {"max_with_index", "true"}}, - {{arith::CmpIPredicate::slt, arith::CmpIPredicate::eq, - arith::CmpIPredicate::slt}, - {"min_with_index", "true"}}, - - {{arith::CmpIPredicate::eq, arith::CmpIPredicate::ugt, - arith::CmpIPredicate::ugt}, - {"max_with_index", "false"}}, - {{arith::CmpIPredicate::eq, arith::CmpIPredicate::ugt, - arith::CmpIPredicate::ult}, - {"min_with_index", "false"}}, - {{arith::CmpIPredicate::eq, arith::CmpIPredicate::ult, - arith::CmpIPredicate::ugt}, - {"max_with_index", "true"}}, - {{arith::CmpIPredicate::eq, arith::CmpIPredicate::ult, - arith::CmpIPredicate::ult}, - {"min_with_index", "true"}}, - {{arith::CmpIPredicate::ugt, arith::CmpIPredicate::eq, - arith::CmpIPredicate::ugt}, - {"max_with_index", "false"}}, - {{arith::CmpIPredicate::ult, arith::CmpIPredicate::eq, - arith::CmpIPredicate::ugt}, - {"min_with_index", "false"}}, - {{arith::CmpIPredicate::ugt, arith::CmpIPredicate::eq, - arith::CmpIPredicate::ult}, - {"max_with_index", "true"}}, - {{arith::CmpIPredicate::ult, arith::CmpIPredicate::eq, - arith::CmpIPredicate::ult}, - {"min_with_index", "true"}}, - }; - - std::vector preds; - using arith::CmpIPredicate; - std::unordered_set allowed{ - arith::CmpIPredicate::slt, arith::CmpIPredicate::sgt, - arith::CmpIPredicate::eq, arith::CmpIPredicate::ult, - arith::CmpIPredicate::ugt}; - // collect predicates under consideration - for (auto it = body.begin(); it != body.end(); ++it) { - if (auto op = dyn_cast(*it)) { - auto pred = op.getPredicate(); - if (allowed.find(pred) == allowed.end()) { - continue; - } - preds.push_back(pred); - } - } - // check if sequence of predicates matches any sequence for min/max - // leftmost/rightmost - if (m.find(preds) != m.end()) { - auto [type, tie_break] = m[preds]; - reduceOp->setAttr(reduceRef, rewriter.getStringAttr(type)); - reduceOp->setAttr(tieBreakLeftRef, rewriter.getStringAttr(tie_break)); - } - - // FLOAT - // For float case it's enough to check for OGT/OLT (comparison of elements) - // and for sgt/slt (comparison for indices) - std::string floatType; - std::string tieBreakLeftFloat; - for (auto it = body.begin(); it != body.end(); ++it) { - if (auto op = dyn_cast(*it)) { - if (op.getPredicate() != arith::CmpFPredicate::OGT && - op.getPredicate() != arith::CmpFPredicate::OLT) { - continue; - } - floatType = op.getPredicate() == arith::CmpFPredicate::OGT - ? "max_with_index" - : "min_with_index"; - } - if (auto op = dyn_cast(*it)) { - if (op.getPredicate() != arith::CmpIPredicate::sgt && - op.getPredicate() != arith::CmpIPredicate::slt) { - continue; - } - tieBreakLeftFloat = - op.getPredicate() == arith::CmpIPredicate::sgt ? "false" : "true"; - } - } - if (!floatType.empty() && !tieBreakLeftFloat.empty()) { - reduceOp->setAttr(reduceRef, rewriter.getStringAttr(floatType)); - reduceOp->setAttr(tieBreakLeftRef, - rewriter.getStringAttr(tieBreakLeftFloat)); - } - - return success(); -} - -LogicalResult ReduceConverter::convertToTargetOpExtended( - triton::ReduceOp op, typename triton::ReduceOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const { - auto loc = op.getLoc(); - auto elemTypes = op.getElementTypes(); - - auto valueResultType = dyn_cast(op.getType(0)); - const auto isScalarReduce = valueResultType == nullptr; - - SmallVector outputs; - for (auto i = 0; i < op.getResult().size() && i < elemTypes.size(); i++) { - auto result = dyn_cast(op.getType(i)); - SmallVector resultShape{ - isScalarReduce ? SmallVector{} - : SmallVector(result.getShape())}; - outputs.push_back( - rewriter.create(loc, resultShape, elemTypes[i])); - } - - auto linalgOp = rewriter.create( - loc, adaptor.getOperands(), outputs, - SmallVector{adaptor.getAxis()}, - [&](OpBuilder &b, Location loc, ValueRange inputs) { - auto tritonReduceBlock = op.getBody(); - IRMapping mapping; - mapping.map(tritonReduceBlock->getArguments(), inputs); - - for (auto &op : tritonReduceBlock->without_terminator()) { - b.clone(op, mapping); - } - - auto tritonYield = tritonReduceBlock->getTerminator(); - auto results = - llvm::map_to_vector(tritonYield->getOperands(), - [&](Value val) { return mapping.lookup(val); }); - b.create(loc, results); - }); - - if (failed(addReduceWithIndexAttrIfNeeded(rewriter, linalgOp))) { - return rewriter.notifyMatchFailure(op, "meaningless reduce operation"); - } - - if (isScalarReduce) { - SmallVector reduceResults; - for (auto i = 0; i < linalgOp.getResults().size() && i < elemTypes.size(); - i++) { - reduceResults.push_back(rewriter.create( - loc, elemTypes[i], linalgOp.getResults()[i], ValueRange{})); - } - rewriter.replaceOp(op, reduceResults); - } else { - rewriter.replaceOp(op, linalgOp); - } - return success(); -} - -bool ScanConverter::isReductionOpSupported(Operation *redOp) const { - return isa(redOp); -} - -LogicalResult -ScanConverter::convertToTargetOp(triton::ScanOp op, - typename triton::ScanOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const { - auto reductionOps = this->getRedOps(op); - if (reductionOps.empty()) { - return rewriter.notifyMatchFailure(op, - "No reduction op found in scan body"); - } - - bool reverse = op.getReverse(); - if (reverse) { - op.emitError("reverse=True is not yet supported for scan op"); - return failure(); - } - - llvm::SmallString<64> funcName; - auto rop = reductionOps.front(); - if (this->isReductionOpSupported(reductionOps.front())) { - if (isa(rop)) { - funcName = "triton_cumsum"; - } else if (isa(rop)) { - funcName = "triton_cumprod"; - } - - auto moduleOp = op->getParentOfType(); - rewriter.setInsertionPoint(moduleOp.getBody(), - std::prev(moduleOp.getBody()->end())); - - auto loc = op.getLoc(); - auto src = adaptor.getOperands().front(); - auto resTy = op.getResult().front().getType(); - auto libFnType = rewriter.getFunctionType( - {src.getType(), rewriter.getI32Type(), rewriter.getI1Type()}, {resTy}); - auto funcOp = rewriter.create(loc, funcName.str(), libFnType); - - SymbolTable symTab(moduleOp); - auto maybePrintFuncNameAttr = symTab.renameToUnique(funcOp, {&symTab}); - if (failed(maybePrintFuncNameAttr)) { - return op->emitError( - "failed to create a unique func name for device_print"); - } - SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); - - rewriter.setInsertionPoint(op); - auto scanAxis = op.getAxis(); - auto scanReverse = op.getReverse(); - Value axis = rewriter.create(loc, scanAxis, 32); - Value reverseVal = - rewriter.create(loc, scanReverse, 1); - auto callOp = rewriter.create( - loc, funcOp.getSymNameAttr(), TypeRange({resTy}), - ValueRange({src, axis, reverseVal})); - - rewriter.replaceOp(op, callOp); - - return success(); - } else { - // This branch is the associative_scan op. - auto loc = op.getLoc(); - - Value scanInput = op.getOperand(0); - - auto srcType = mlir::dyn_cast(scanInput.getType()); - if (!srcType) { - return rewriter.notifyMatchFailure( - op, "Expected RankedTensorType input for associative_scan"); - } - - auto elementType = srcType.getElementType(); - auto shape = srcType.getShape(); - int rank = shape.size(); - int axis = op.getAxis(); - - if (axis < 0 || axis >= rank) { - return rewriter.notifyMatchFailure(op, "Invalid scan axis: " + - std::to_string(axis)); - } - - if (op->getNumRegions() < 1 || op->getRegion(0).empty()) { - return rewriter.notifyMatchFailure(op, "Missing combine region"); - } - - OpBuilder::InsertionGuard guard(rewriter); - - auto memrefType = MemRefType::get(shape, elementType); - Value inputMemRef = - rewriter.create(loc, memrefType, scanInput); - Value outputMemRef = rewriter.create(loc, memrefType); - - auto processDimension = [&](ArrayRef baseIdxsArray) { - llvm::SmallVector baseIdxs(baseIdxsArray.begin(), - baseIdxsArray.end()); - llvm::SmallVector firstIdx = baseIdxs; - if (axis <= firstIdx.size()) { - firstIdx.insert(firstIdx.begin() + axis, - rewriter.create(loc, 0)); - } else { - firstIdx.push_back(rewriter.create(loc, 0)); - } - - Value firstVal = - rewriter.create(loc, inputMemRef, firstIdx); - rewriter.create(loc, firstVal, outputMemRef, firstIdx); - - Value axisSize = - rewriter.create(loc, inputMemRef, axis).getResult(); - Value one = rewriter.create(loc, 1); - - Value cmp = rewriter.create(loc, arith::CmpIPredicate::sgt, - axisSize, one); - auto ifOp = rewriter.create(loc, cmp, false); - - // Create a loop only when the axis size is greater than 1. - rewriter.setInsertionPointToStart(ifOp.thenBlock()); - - auto forOp = rewriter.create(loc, one, axisSize, one); - rewriter.setInsertionPointToStart(forOp.getBody()); - - Value k = forOp.getInductionVar(); - llvm::SmallVector currIdx = baseIdxs; - if (axis <= currIdx.size()) { - currIdx.insert(currIdx.begin() + axis, k); - } else { - currIdx.push_back(k); - } - - Value km1 = rewriter.create(loc, k, one); - llvm::SmallVector prevIdx = baseIdxs; - if (axis <= prevIdx.size()) { - prevIdx.insert(prevIdx.begin() + axis, km1); - } else { - prevIdx.push_back(km1); - } - - Value currentVal = - rewriter.create(loc, inputMemRef, currIdx); - Value prevResult = - rewriter.create(loc, outputMemRef, prevIdx); - - Region &combineRegion = op->getRegion(0); - Block &combineBlock = combineRegion.front(); - IRMapping mapping; - mapping.map(combineBlock.getArgument(0), prevResult); - mapping.map(combineBlock.getArgument(1), currentVal); - - for (Operation &innerOp : combineBlock.without_terminator()) { - rewriter.clone(innerOp, mapping); - } - - Operation *yieldOp = combineBlock.getTerminator(); - Value resultVal = mapping.lookup(yieldOp->getOperand(0)); - - rewriter.create(loc, resultVal, outputMemRef, currIdx); - - rewriter.setInsertionPointAfter(ifOp); - }; - - // Constructing loops for non-scanning dimensions - llvm::SmallVector nonScanDims; - for (int i = 0; i < rank; ++i) { - if (i != axis) - nonScanDims.push_back(i); - } - - createSimpleNestedLoops(rewriter, loc, outputMemRef, nonScanDims, - processDimension); - - rewriter.setInsertionPointAfter(op); - - MemRefType memrefTy = cast(outputMemRef.getType()); - Type tensorTy = - RankedTensorType::get(memrefTy.getShape(), memrefTy.getElementType()); - Value outputTensor = rewriter.create( - loc, tensorTy, outputMemRef, true, true); - rewriter.replaceOp(op, outputTensor); - return success(); - } -} - -LogicalResult ScanConverter::convertToTargetOpExtended( - triton::ScanOp op, typename triton::ScanOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const { - auto loc = op.getLoc(); - bool reverse = op.getReverse(); - if (reverse) { - return op.emitError( - "reverse=True is not yet supported for extended scan op"); - } - - // 1. Extract all input tensors (supports multiple inputs) - auto operands = op->getOperands(); - if (operands.empty()) { - return rewriter.notifyMatchFailure(op, - "No input operands for extended scan"); - } - - // 2. Validate all inputs are of RankedTensorType - llvm::SmallVector inputTensTypes; - for (auto operand : operands) { - auto tensorTy = dyn_cast(operand.getType()); - if (!tensorTy) { - return rewriter.notifyMatchFailure(op, - "All inputs must be RankedTensorType"); - } - inputTensTypes.push_back(tensorTy); - } - - // 3. Validate all input tensors have the same shape (scan operation requires - // matching input dimensions) - auto baseShape = inputTensTypes[0].getShape(); - int rank = baseShape.size(); - int axis = op.getAxis(); - if (axis < 0 || axis >= rank) { - return rewriter.notifyMatchFailure(op, "Invalid scan axis: " + - std::to_string(axis)); - } - for (size_t i = 1; i < inputTensTypes.size(); ++i) { - if (inputTensTypes[i].getShape() != baseShape) { - return rewriter.notifyMatchFailure(op, - "All inputs must have the same shape"); - } - } - - // 4. Prepare MemRefs for multiple inputs/outputs - llvm::SmallVector inputMemRefs; - llvm::SmallVector outputMemRefs; - llvm::SmallVector memRefTypes; - for (size_t i = 0; i < inputTensTypes.size(); ++i) { - auto &tensorTy = inputTensTypes[i]; - auto memRefTy = - MemRefType::get(tensorTy.getShape(), tensorTy.getElementType()); - memRefTypes.push_back(memRefTy); - // Convert input tensors to MemRefs - inputMemRefs.push_back( - rewriter.create(loc, memRefTy, operands[i])); - // Allocate MemRefs for outputs - outputMemRefs.push_back(rewriter.create(loc, memRefTy)); - } - - // 5. Define scanning logic for multiple inputs/outputs - LogicalResult loopResult = success(); - auto processDimension = [&](ArrayRef baseIdxsArray) { - llvm::SmallVector baseIdxs(baseIdxsArray.begin(), - baseIdxsArray.end()); - llvm::SmallVector firstIdx = baseIdxs; - // Insert start index (0) for the scan axis - if (axis <= firstIdx.size()) { - firstIdx.insert(firstIdx.begin() + axis, - rewriter.create(loc, 0)); - } else { - firstIdx.push_back(rewriter.create(loc, 0)); - } - - // 5.1 Process the first element: directly copy multiple inputs to multiple - // outputs (initialize cumulative results) - for (size_t i = 0; i < inputMemRefs.size(); ++i) { - Value firstVal = - rewriter.create(loc, inputMemRefs[i], firstIdx); - rewriter.create(loc, firstVal, outputMemRefs[i], - firstIdx); - } - - // 5.2 Calculate the size of the scan axis; create a loop only if the axis - // size > 1 - Value axisSize = - rewriter.create(loc, inputMemRefs[0], axis).getResult(); - Value one = rewriter.create(loc, 1); - Value cmp = rewriter.create(loc, arith::CmpIPredicate::sgt, - axisSize, one); - auto ifOp = rewriter.create(loc, cmp, false); - - rewriter.setInsertionPointToStart(ifOp.thenBlock()); - // Loop variable k: ranges from 1 to axisSize-1 - auto forOp = rewriter.create(loc, one, axisSize, one); - rewriter.setInsertionPointToStart(forOp.getBody()); - Value k = forOp.getInductionVar(); - - // 5.3 Calculate current index (k) and previous index (k-1) - llvm::SmallVector currIdx = baseIdxs; - if (axis <= currIdx.size()) { - currIdx.insert(currIdx.begin() + axis, k); - } else { - currIdx.push_back(k); - } - Value km1 = rewriter.create(loc, k, one); - llvm::SmallVector prevIdx = baseIdxs; - if (axis <= prevIdx.size()) { - prevIdx.insert(prevIdx.begin() + axis, km1); - } else { - prevIdx.push_back(km1); - } - - // 5.4 Load current elements and previous cumulative results - llvm::SmallVector currentVals; - llvm::SmallVector prevResults; - for (size_t i = 0; i < inputMemRefs.size(); ++i) { - currentVals.push_back( - rewriter.create(loc, inputMemRefs[i], currIdx)); - prevResults.push_back( - rewriter.create(loc, outputMemRefs[i], prevIdx)); - } - - // 5.5 Bind parameters for custom reduction logic - Region &combineRegion = op->getRegion(0); - if (combineRegion.empty()) { - op->emitError("Missing combine region in extended scan"); - loopResult = failure(); - return; - } - Block &combineBlock = combineRegion.front(); - // Validate that the number of reduction region arguments matches (number of - // previous results + number of current elements) - if (combineBlock.getNumArguments() != 2 * inputMemRefs.size()) { - op->emitError("Combine region arguments mismatch with input count"); - loopResult = failure(); - return; - } - IRMapping mapping; - for (size_t i = 0; i < inputMemRefs.size(); ++i) { - // Bind previous results (previous value of the i-th output) to the i-th - // argument of the reduction region - mapping.map(combineBlock.getArgument(i), prevResults[i]); - // Bind current elements (current value of the i-th input) to the i+N-th - // argument of the reduction region (N is the number of inputs) - mapping.map(combineBlock.getArgument(i + inputMemRefs.size()), - currentVals[i]); - } - - // 5.6 Clone all operations within the reduction region - for (Operation &innerOp : combineBlock.without_terminator()) { - rewriter.clone(innerOp, mapping); - } - - // 5.7 Extract reduction results and store them in outputMemRef - Operation *yieldOp = combineBlock.getTerminator(); - if (yieldOp->getNumOperands() != outputMemRefs.size()) { - op->emitError("Combine region returns mismatch with output count"); - loopResult = failure(); - return; - } - for (size_t i = 0; i < outputMemRefs.size(); ++i) { - Value resultVal = mapping.lookup(yieldOp->getOperand(i)); - rewriter.create(loc, resultVal, outputMemRefs[i], - currIdx); - } - - rewriter.setInsertionPointAfter(ifOp); - }; - - // 6. Generate nested loops for non-scan dimensions - llvm::SmallVector nonScanDims; - for (int i = 0; i < rank; ++i) { - if (i != axis) - nonScanDims.push_back(i); - } - createSimpleNestedLoops(rewriter, loc, outputMemRefs[0], nonScanDims, - processDimension); - - if (failed(loopResult)) { - return failure(); - } - - // 7. Convert multiple output MemRefs back to tensors and replace the original - // tt.scan operation - llvm::SmallVector outputTensors; - for (auto outputMemRef : outputMemRefs) { - MemRefType memrefTy = cast(outputMemRef.getType()); - Type tensorTy = - RankedTensorType::get(memrefTy.getShape(), memrefTy.getElementType()); - outputTensors.push_back(rewriter.create( - loc, tensorTy, outputMemRef, true, true)); - } - rewriter.replaceOp(op, outputTensors); - - return success(); -} - -LogicalResult DevicePrintConverter::matchAndRewrite( - triton::PrintOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const { - auto moduleOp = op->getParentOfType(); - rewriter.setInsertionPoint(moduleOp.getBody(), - std::prev(moduleOp.getBody()->end())); - SmallVector inputTypes; - for (auto arg : op.getArgs()) { - inputTypes.push_back(arg.getType()); - } - auto libFnType = rewriter.getFunctionType(inputTypes, {}); - auto funcOp = - rewriter.create(op.getLoc(), printFuncNameBase, libFnType); - SymbolTable symTab(moduleOp); - auto maybePrintFuncNameAttr = symTab.renameToUnique(funcOp, {&symTab}); - if (failed(maybePrintFuncNameAttr)) { - return op->emitError( - "failed to create a unique func name for device_print"); - } - SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); - auto prefixAttr = op.getPrefixAttr(); - funcOp->setAttr(prefixAttrName, prefixAttr); - auto hexAttr = op.getHexAttr(); - funcOp->setAttr(hexAttrName, hexAttr); - - rewriter.setInsertionPoint(op); - rewriter.create(op.getLoc(), funcOp, op.getArgs()); - - rewriter.eraseOp(op); - return success(); -} - -LogicalResult DeviceAssertConverter::matchAndRewrite( - triton::AssertOp op, OpAdaptor adaptor, - mlir::ConversionPatternRewriter &rewriter) const { - auto msgAttr = op.getMessageAttr(); - // Filter out automatically inserted assert ops - if (auto strAttr = mlir::dyn_cast(msgAttr)) { - llvm::StringRef msg = strAttr.getValue(); - if (msg.contains("overflow detected for operation")) { - rewriter.eraseOp(op); - return success(); - } - } - - auto moduleOp = op->getParentOfType(); - rewriter.setInsertionPoint(moduleOp.getBody(), - std::prev(moduleOp.getBody()->end())); - auto conditionType = op.getCondition().getType(); - - auto libFnType = rewriter.getFunctionType({conditionType}, {}); - auto funcOp = - rewriter.create(op.getLoc(), printFuncNameBase, libFnType); - mlir::SymbolTable symTab(moduleOp); - auto maybePrintFuncNameAttr = symTab.renameToUnique(funcOp, {&symTab}); - if (failed(maybePrintFuncNameAttr)) { - return op->emitError( - "failed to create a unique func name for device_assert"); - } - SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); - funcOp->setAttr(msgAttrName, msgAttr); - - rewriter.setInsertionPoint(op); - rewriter.create(op.getLoc(), funcOp, - ValueRange{op.getCondition()}); - - rewriter.eraseOp(op); - return success(); -} \ No newline at end of file diff --git a/compiler/lib/Conversion/LinalgToLinked/VerifyNoLinalgGenericPass.cpp b/compiler/lib/Conversion/LinalgToLinked/VerifyNoLinalgGenericPass.cpp deleted file mode 100644 index 429aefb8..00000000 --- a/compiler/lib/Conversion/LinalgToLinked/VerifyNoLinalgGenericPass.cpp +++ /dev/null @@ -1,6 +0,0 @@ -#include "dicp/Conversion/LinalgToLinked/VerifyNoLinalgGenericPass.hpp" - -using namespace mlir; -using namespace mlir::func; - -namespace mlir::dicp::linked {} // namespace mlir::dicp::linked diff --git a/compiler/lib/Conversion/LinalgToNPU/CMakeLists.txt b/compiler/lib/Conversion/LinalgToNPU/CMakeLists.txt deleted file mode 100644 index e48b1cba..00000000 --- a/compiler/lib/Conversion/LinalgToNPU/CMakeLists.txt +++ /dev/null @@ -1,28 +0,0 @@ -add_triton_library(LinalgToNPU - LinalgToNPUPass.cpp - LinalgToNPU.cpp - - DEPENDS - DICPNPUIncGen - LinalgToNPUConversionPassIncGen - - LINK_LIBS PUBLIC - TritonTilingExtIR - MLIRArithDialect - MLIRDialectUtils - MLIRIR - MLIRMathDialect - MLIRPass - MLIRTensorDialect - MLIRTransforms - MLIRSupport - TritonAnalysis - TritonIR - TritonTransforms - TritonSharedAnalysis - DICPNPU - - TritonArithToLinalg - StructuredToMemref - TritonToStructured -) diff --git a/compiler/lib/Conversion/LinalgToNPU/LinalgToNPU.cpp b/compiler/lib/Conversion/LinalgToNPU/LinalgToNPU.cpp deleted file mode 100644 index de5fdb7e..00000000 --- a/compiler/lib/Conversion/LinalgToNPU/LinalgToNPU.cpp +++ /dev/null @@ -1,33 +0,0 @@ -#include "dicp/Conversion/LinalgToNPU/LinalgToNPU.h" -#include "dicp/Dialect/NPU/IR/NPUDialect.h" -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/Passes.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include "llvm/ADT/SmallVectorExtras.h" -#include "llvm/ADT/TypeSwitch.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/FormatVariadic.h" -#include "llvm/Support/MathExtras.h" - -#include -#include - -#define DEBUG_TYPE "linalg-to-npu" -#include "dicp/Conversion/LinalgToNPU/ConversionPatterns.hpp" - -using namespace mlir; -using namespace dicp; - -#define GEN_PASS_CLASSES -#include "dicp/Conversion/LinalgToNPU/Passes.h.inc" - -void npu::populateLinalgToNPUConversionPatterns(RewritePatternSet &patterns) { - // patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); -} diff --git a/compiler/lib/Conversion/LinalgToNPU/LinalgToNPUPass.cpp b/compiler/lib/Conversion/LinalgToNPU/LinalgToNPUPass.cpp deleted file mode 100644 index de959807..00000000 --- a/compiler/lib/Conversion/LinalgToNPU/LinalgToNPUPass.cpp +++ /dev/null @@ -1,59 +0,0 @@ -#include "dicp/Conversion/LinalgToNPU/LinalgToNPU.h" -#include "dicp/Dialect/NPU/IR/NPUDialect.h" - -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Transforms/Passes.h" - -#define DEBUG_TYPE "linalg-to-npu" -#include "dicp/Conversion/LinalgToNPU/ConversionPatterns.hpp" - -using namespace mlir; -using namespace dicp; - -#define GEN_PASS_CLASSES -#include "dicp/Conversion/LinalgToNPU/Passes.h.inc" - -namespace { - -class LinalgToNPUPass : public LinalgToNPUBase { -public: - void getDependentDialects(DialectRegistry ®istry) const override { - registry - .insert(); - } - - void runOnOperation() override { - auto moduleOp = getOperation(); - // OpBuilder builder(moduleOp.getBodyRegion()); - // builder.setInsertionPointToStart(builder.getBlock()); - // builder.create(builder.getUnknownLoc()); - - RewritePatternSet patterns(&getContext()); - ConversionTarget target(getContext()); - target.addLegalDialect(); - // target.addDynamicallyLegalOp([](arith::AddFOp op) { - // return !isa(op.getResult().getType()); - // }); - // target.addIllegalOp(); - - npu::populateLinalgToNPUConversionPatterns(patterns); - // patterns.add(patterns.getContext()); - - if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) { - signalPassFailure(); - } - } -}; - -} // namespace - -std::unique_ptr> npu::createLinalgToNPUPass() { - return std::make_unique(); -} diff --git a/compiler/lib/Conversion/TritonToLinalgNPU/CMakeLists.txt b/compiler/lib/Conversion/TritonToLinalgNPU/CMakeLists.txt deleted file mode 100644 index 5cd7900f..00000000 --- a/compiler/lib/Conversion/TritonToLinalgNPU/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -dicp_add_all_subdirs() \ No newline at end of file diff --git a/compiler/lib/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/CMakeLists.txt b/compiler/lib/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/CMakeLists.txt deleted file mode 100644 index 9f9fd662..00000000 --- a/compiler/lib/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/CMakeLists.txt +++ /dev/null @@ -1,13 +0,0 @@ -add_triton_library(MemRefCopyGatherToTensorInsert - MemRefCopyGatherToTensorInsert.cpp - - DEPENDS - MemRefCopyGatherToTensorInsertPassIncGen - - LINK_LIBS - MLIRIR - MLIRPass - MLIRTransforms - MLIRSupport - TritonIR -) diff --git a/compiler/lib/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/MemRefCopyGatherToTensorInsert.cpp b/compiler/lib/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/MemRefCopyGatherToTensorInsert.cpp deleted file mode 100644 index 467adbeb..00000000 --- a/compiler/lib/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/MemRefCopyGatherToTensorInsert.cpp +++ /dev/null @@ -1,283 +0,0 @@ -#include "dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Rewrite/FrozenRewritePatternSet.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" - -using namespace mlir; - -namespace mlir::dicp::linked { -#define GEN_PASS_DEF_MEMREFCOPYGATHERTOTENSORINSERT -#include "dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.h.inc" -} // namespace mlir::dicp::linked - -namespace { - -/// Helper to convert OpFoldResult to Value. -/// If it is an attribute, creates an arith.constant. -static Value getValueOrCreateConstantIndexOp(PatternRewriter &rewriter, - Location loc, OpFoldResult ofr) { - if (auto val = dyn_cast(ofr)) - return val; - return rewriter.create( - loc, cast(cast(ofr)).getInt()); -} - -/// Helper function to check if a Value is derived from a Tensor ExtractOp. -/// It supports two chains: -/// 1. ExtractOp -> Value (Target) -/// 2. ExtractOp -> IndexCastOp -> Value (Target) -/// Returns the defining ExtractOp if a match is found involving the loop IV. -static tensor::ExtractOp findSourceExtractOp(Value val, Value loopIV) { - Operation *defOp = val.getDefiningOp(); - if (!defOp) - return nullptr; - - // Case 1: Direct ExtractOp - if (auto extractOp = dyn_cast(defOp)) { - for (Value idx : extractOp.getIndices()) { - if (idx == loopIV) - return extractOp; - } - return nullptr; - } - - // Case 2: ExtractOp -> IndexCastOp - if (auto indexCastOp = dyn_cast(defOp)) { - if (auto extractOp = - indexCastOp.getIn().getDefiningOp()) { - for (Value idx : extractOp.getIndices()) { - if (idx == loopIV) - return extractOp; - } - } - } - - return nullptr; -} - -/// Pattern to convert a specific memory gather loop into a tensor insertion -/// loop. -/// -/// Source Pattern: -/// scf.for %iv ... { -/// %idx = tensor.extract[%iv] <-- (Optional IndexCast) -/// %view = memref.reinterpret_cast %src to offset: [%idx] ... -/// %subview = memref.subview %alloc[%iv] ... -/// memref.copy %view, %subview -/// } -/// -/// Target Pattern: -/// %res = scf.for %iv ... iter_args(%acc = %empty) { -/// %idx = tensor.extract[%iv] -/// %cast_idx = arith.index_cast %idx <-- (If needed) -/// %view = memref.reinterpret_cast %src to offset: [%cast_idx] ... -/// %val = memref.load %view[0, 0, ...] <-- Matches rank -/// %next = tensor.insert %val into %acc[%iv, ...] <-- Matches subview -/// offsets scf.yield %next -/// } -struct MemRefCopyGatherToTensorInsertPattern - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(scf::ForOp forOp, - PatternRewriter &rewriter) const override { - // 1. Analyze Loop Body: Must contain exactly one memref.copy - memref::CopyOp copyOp; - int copyCount = 0; - forOp.getBody()->walk([&](memref::CopyOp op) { - copyOp = op; - copyCount++; - }); - - if (copyCount != 1 || !copyOp) - return failure(); - - // 2. Validate Target Semantics (Write Side) - // Expectation: Copy -> SubView -> Alloc -> ToTensor - auto subViewOp = copyOp.getTarget().getDefiningOp(); - if (!subViewOp) - return failure(); - - auto allocOp = subViewOp.getSource().getDefiningOp(); - if (!allocOp) - return failure(); - - bufferization::ToTensorOp toTensorOp; - for (Operation *user : allocOp->getUsers()) { - if (auto op = dyn_cast(user)) { - toTensorOp = op; - break; - } - } - - if (!toTensorOp || toTensorOp->getBlock() != forOp->getBlock() || - !forOp->isBeforeInBlock(toTensorOp)) { - return failure(); - } - - // 3. Validate Source Semantics (Read Side) - // Expectation: (Extract -> Optional Cast) -> ReinterpretCast -> Copy Source - auto reinterpretOp = - copyOp.getSource().getDefiningOp(); - if (!reinterpretOp) - return failure(); - - tensor::ExtractOp extractOp; - bool patternFound = false; - - // Check dynamic offsets to find the one driven by the loop induction - // variable. - for (OpFoldResult ofr : reinterpretOp.getOffsets()) { - if (auto val = dyn_cast(ofr)) { - extractOp = findSourceExtractOp(val, forOp.getInductionVar()); - if (extractOp) { - patternFound = true; - break; - } - } - } - - if (!patternFound) - return failure(); - - // ==================================================== - // Rewrite Phase - // ==================================================== - Location loc = forOp.getLoc(); - - // A. Prepare Accumulator (tensor.empty) - auto resultType = cast(toTensorOp.getResult().getType()); - Value initTensor = rewriter.create( - loc, resultType.getShape(), resultType.getElementType()); - - // B. Create New Loop - auto newForOp = rewriter.create( - loc, forOp.getLowerBound(), forOp.getUpperBound(), forOp.getStep(), - ValueRange{initTensor}); - - newForOp->setAttr("ExtractedLoadOrStore", rewriter.getUnitAttr()); - - // C. Populate New Loop Body - rewriter.setInsertionPointToStart(newForOp.getBody()); - - Value iv = newForOp.getInductionVar(); - Value acc = newForOp.getRegionIterArgs()[0]; - - // C.1. Recreate Index Calculation - // We clone the indices from the original extractOp. - // If an index was the old loop's IV, replace it with the new loop's IV. - SmallVector extractIndices; - for (Value idx : extractOp.getIndices()) { - if (idx == forOp.getInductionVar()) - extractIndices.push_back(iv); - else - extractIndices.push_back(idx); - } - - Value newExtract = rewriter.create( - extractOp.getLoc(), extractOp.getTensor(), extractIndices); - - Value newOffsetIdx = newExtract; - if (!newOffsetIdx.getType().isIndex()) { - newOffsetIdx = rewriter.create( - loc, rewriter.getIndexType(), newExtract); - } - - // C.2. Recreate ReinterpretCast - // Map the matched dynamic offset to our new calculated index. - OpFoldResult newOffsetOfr = rewriter.getIndexAttr(0); - - if (!reinterpretOp.getMixedOffsets().empty()) { - OpFoldResult oldOfr = reinterpretOp.getMixedOffsets()[0]; - - // If the old offset matches our pattern, replace it. - bool isTargetOffset = false; - if (auto val = dyn_cast(oldOfr)) { - if (findSourceExtractOp(val, forOp.getInductionVar())) { - isTargetOffset = true; - } - } - newOffsetOfr = isTargetOffset ? newOffsetIdx : oldOfr; - } - - Value newSrcMemref = rewriter.create( - reinterpretOp.getLoc(), reinterpretOp.getType(), - reinterpretOp.getSource(), newOffsetOfr, reinterpretOp.getMixedSizes(), - reinterpretOp.getMixedStrides()); - - // C.3. Load from Source - // FIX: Ensure we provide indices for ALL dimensions of the reinterpreted - // memref. We assume we are loading from the base (0, 0, ...). - auto memRefType = cast(newSrcMemref.getType()); - Value c0 = rewriter.create(loc, 0); - SmallVector loadIndices(memRefType.getRank(), c0); - - Value loadedVal = - rewriter.create(loc, newSrcMemref, loadIndices); - - // C.4. Insert into Accumulator - // FIX: Derive insertion indices from the original subview offsets. - // This handles multi-dimensional tensors correctly by mapping the subview - // logic. - SmallVector insertIndices; - for (OpFoldResult ofr : subViewOp.getMixedOffsets()) { - if (auto val = dyn_cast(ofr)) { - // If the offset is the old IV, use the new IV. - if (val == forOp.getInductionVar()) - insertIndices.push_back(iv); - else - insertIndices.push_back(val); - } else { - // Materialize static offset as a constant index. - insertIndices.push_back( - getValueOrCreateConstantIndexOp(rewriter, loc, ofr)); - } - } - - Value nextTensor = - rewriter.create(loc, loadedVal, acc, insertIndices); - - // C.5. Yield - auto yieldOp = rewriter.create(loc, nextTensor); - yieldOp->setAttr("DiscreteMemAccess", rewriter.getUnitAttr()); - - // D. Finalize - rewriter.replaceOp(toTensorOp, newForOp.getResult(0)); - - // Cleanup - rewriter.eraseOp(forOp); - rewriter.eraseOp(allocOp); - - return success(); - } -}; - -struct MemRefCopyGatherToTensorInsertPass - : mlir::dicp::linked::impl::MemRefCopyGatherToTensorInsertBase< - MemRefCopyGatherToTensorInsertPass> { - - void runOnOperation() override { - auto moduleOp = getOperation(); - RewritePatternSet patterns(&getContext()); - patterns.add(&getContext()); - - if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { - signalPassFailure(); - } - } -}; - -} // namespace - -std::unique_ptr> -mlir::dicp::linked::createMemRefCopyGatherToTensorInsertPass() { - return std::make_unique(); -} \ No newline at end of file diff --git a/compiler/lib/Conversion/TritonToLinalgNPU/TritonArithToLinalg/NPUSpecific.cpp b/compiler/lib/Conversion/TritonToLinalgNPU/TritonArithToLinalg/NPUSpecific.cpp deleted file mode 100644 index b04f83aa..00000000 --- a/compiler/lib/Conversion/TritonToLinalgNPU/TritonArithToLinalg/NPUSpecific.cpp +++ /dev/null @@ -1,106 +0,0 @@ -#include "dicp/Conversion/TritonToLinalgNPU/TritonArithToLinalg/NPUSpecific.h" - -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/Passes.h" - -#include "triton-shared/Conversion/TritonArithToLinalg/TritonArithToLinalg.h" -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include "llvm/ADT/SmallVectorExtras.h" -#include "llvm/ADT/TypeSwitch.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/FormatVariadic.h" -#include "llvm/Support/MathExtras.h" - -#include -#include - -#define DEBUG_TYPE "triton-arith-to-linalg-npu" -#include "dicp/Conversion/TritonToLinalgNPU/TritonArithToLinalg/ConversionPatterns.hpp" - -using namespace mlir; -using namespace triton; -using namespace mlir::dicp::linked; - -void mlir::dicp::linked::populateTritonArithToLinalgNPUConversionPatterns( - bool pidsToFuncArgs, bool addptrToLinalg, bool assertToCf, - bool transposeReduceToRank0, RewritePatternSet &patterns) { - - if (pidsToFuncArgs) { - patterns.add( - patterns.getContext()); - } - if (addptrToLinalg) { - patterns.add(patterns.getContext()); - } - if (assertToCf) { - patterns.add(patterns.getContext()); - } - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - - populateExternElementwiseOpToMLIROps(patterns); - - // Reduce converters - // Triton's reduce op is idential to linalg.reduce op, so we can clone - // `tt.reduce` body to `linalg.reduce`. Unfortunately, we still need to - // perform pattern matching to know what reduce ops we are dealing with - // so that we know how to initialize the initial reduce values correctly. - // - // We can do this in a generic way without pattern matching by always using - // the first elements along the reduction axis and perform the reduction on - // the remaining elements. However, this results in creatings sub-tensors that - // aren't always multiple of 2s, which are sub-optimal for certain hardwares. - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext(), - transposeReduceToRank0); - - // linalg::populateElementwiseToLinalgConversionPatterns(patterns); -} - -bool mlir::dicp::linked::isLegalConstantAndTensorArithmeticOpForNPU( - Operation *op) { - // Check for arith::ConstantOp - if (auto constOp = dyn_cast(op)) { - // 1. Scalar constants are always legal (handled elsewhere). - if (!isa(constOp.getResult().getType())) { - return true; - } - // 2. RankedTensor constant check: - if (auto denseAttr = dyn_cast(constOp.getValue())) { - // Dense splat constants of float/integer type are ILLEGAL (must be - // lowered, e.g., to linalg.fill). - if (denseAttr.isSplat() && - isa(denseAttr.getElementType())) { - return false; // **ILLEGAL**: needs lowering (e.g., to linalg.fill). - } - } - // All other RankedTensor constants (non-splat, non-float/int elements) are - // legal. - return true; - } - // 3. All non-constant arith/math operations are currently considered legal - // (i.e., they should not be lowered by the constant-related pass). - return true; -} \ No newline at end of file diff --git a/compiler/lib/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/CMakeLists.txt b/compiler/lib/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/CMakeLists.txt deleted file mode 100644 index 52072492..00000000 --- a/compiler/lib/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/CMakeLists.txt +++ /dev/null @@ -1,33 +0,0 @@ -add_triton_library(TritonToLinalgNPUCoversion - TritonToLinalgNPUConversionPass.cpp - - - DEPENDS - TritonToLinalgNPUCoversionPassIncGen - - LINK_LIBS PUBLIC - TritonTilingExtIR - MLIRArithDialect - MLIRDialectUtils - MLIRIR - MLIRMathDialect - MLIRPass - MLIRTensorDialect - MLIRTransforms - MLIRSupport - TPtrIR - TritonIR - TritonTransforms - TritonSharedAnalysis - TritonSharedUtils - TritonToLinalgExperimental - - TritonArithToLinalg - StructuredToMemref - TritonToStructured - TritonToUnstructured - TritonPtrToMemref - UnstructuredToMemref - MemRefCopyGatherToTensorInsert - TritonToUnstructure -) diff --git a/compiler/lib/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/TritonToLinalgNPUConversionPass.cpp b/compiler/lib/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/TritonToLinalgNPUConversionPass.cpp deleted file mode 100644 index 237d36eb..00000000 --- a/compiler/lib/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/TritonToLinalgNPUConversionPass.cpp +++ /dev/null @@ -1,93 +0,0 @@ -#include "dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.h" -#include "dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/TritonToLinalgNPUCoversion.h" -#include "dicp/Conversion/TritonToUnstructure/BubbleUpOperation.h" - -#include "triton-shared/Conversion/StructuredToMemref/StructuredToMemref.h" -#include "triton-shared/Conversion/TritonArithToLinalg/TritonArithToLinalg.h" -#include "triton-shared/Conversion/TritonPtrToMemref/TritonPtrToMemref.h" -#include "triton-shared/Conversion/TritonToLinalgExperimental/CollapseShape.h" -#include "triton-shared/Conversion/TritonToLinalgExperimental/ReconcilePtrCasts.h" -#include "triton-shared/Conversion/TritonToLinalgExperimental/TritonToPtr.h" -#include "triton-shared/Conversion/TritonToStructured/TritonToStructured.h" -#include "triton-shared/Conversion/TritonToUnstructured/TritonToUnstructured.h" -#include "triton-shared/Conversion/UnstructuredToMemref/UnstructuredToMemref.h" -#include "triton-shared/Dialect/TPtr/IR/TPtrDialect.h" -#include "triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.h" -#include "triton-shared/Dialect/TritonTilingExt/IR/TritonTilingExtDialect.h" - -#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h" -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/Ptr/IR/PtrDialect.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Transforms/Passes.h" - -using namespace mlir; -using namespace triton; -using namespace mlir::dicp::linked; - -#define GEN_PASS_CLASSES -#include "dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/Passes.h.inc" - -namespace { - -class TritonToLinalgNPUCoversionPass - : public TritonToLinalgNPUCoversionBase { - -public: - void getDependentDialects(DialectRegistry ®istry) const override { - registry.insert(); - } - - void runOnOperation() override { - auto moduleOp = getOperation(); - PassManager pm(&getContext(), moduleOp.getOperationName()); - - pm.addPass(createTritonToStructuredPass(enableMakeGatherScatterTensorPtr)); - - // Erase dead code and fold constants created during lowering - pm.addPass(createCSEPass()); - pm.addPass(createCanonicalizerPass()); - - pm.addPass(createTritonToUnstructuredPass()); - pm.addPass(createTritonArithToLinalgPass(true, false)); - - pm.addPass(createStructuredToMemrefPass()); - pm.addPass(createMemRefCopyGatherToTensorInsertPass()); - - pm.addPass(createUnstructuredToMemrefPass()); - pm.addPass(createTritonPtrToMemrefPass()); - pm.addPass(createTritonToPtrPass()); - pm.addPass(createReconcileUnrealizedCastsPass()); - pm.addPass(createReconcilePtrCastsPass()); - - // Now that remove-dead-values fully works with linalg ops, clean up the IR - // again, particularly unused loop iter-args that were created - // during triton-to-structured. - pm.addPass(createRemoveDeadValuesPass()); - pm.addPass(createCSEPass()); - pm.addPass(createCanonicalizerPass()); - if (enableCollapseShape) { - // Canonicalizer pass will rewrite tensor.expand_shape(linalg.fill) to - // linalg.fill(tensor.expand_shape) so we need to run it before - // collapseShape pass - pm.addPass(createCollapseShapePass()); - } - if (failed(runPipeline(pm, getOperation()))) { - signalPassFailure(); - } - } -}; -} // namespace - -std::unique_ptr> -mlir::dicp::linked::createTritonToLinalgNPUCoversionPass() { - return std::make_unique(); -} diff --git a/compiler/lib/Dialect/CommonIR/AnnotateKernelAttrsPass.cpp b/compiler/lib/Dialect/CommonIR/AnnotateKernelAttrsPass.cpp new file mode 100644 index 00000000..86a519fa --- /dev/null +++ b/compiler/lib/Dialect/CommonIR/AnnotateKernelAttrsPass.cpp @@ -0,0 +1,63 @@ +#include "dicp/Dialect/CommonIR/Passes.h" + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/BuiltinDialect.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassRegistry.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "llvm/Support/Debug.h" + +using namespace mlir; + +namespace mlir::dicp::CommonIR { +#define GEN_PASS_DEF_ANNOTATEKERNELATTRSPASS +#include "dicp/Dialect/CommonIR/Passes.h.inc" +} // namespace mlir::dicp::CommonIR + +#define DEBUG_TYPE "annotate-kernel-attrs-pass" + +namespace { + +static constexpr llvm::StringRef kGlobalKernel = "global_kernel"; +static constexpr llvm::StringRef kMixMode = "mix_mode"; +static constexpr llvm::StringRef kParallelMode = "parallel_mode"; + +struct AnnotateKernelAttrsPass + : public PassWrapper> { + StringRef getArgument() const final { return "annotate-kernel-attrs"; } + StringRef getDescription() const final { + return "Annotate kernel func.func with mix_mode / parallel_mode / " + "global_kernel so downstream lowering inserts the workspace and " + "syncBlockLock stub arguments."; + } + + void getDependentDialects(mlir::DialectRegistry ®istry) const override { + registry.insert(); + } + + void runOnOperation() override { + auto module = getOperation(); + + for (auto func : module.getOps()) { + if (func->hasAttr(kGlobalKernel)) + continue; + if (func.isDeclaration()) + continue; + + auto *ctx = func.getContext(); + func->setAttr(kGlobalKernel, StringAttr::get(ctx, "local")); + func->setAttr(kMixMode, StringAttr::get(ctx, "aiv")); + func->setAttr(kParallelMode, StringAttr::get(ctx, "simd")); + } + } + + MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AnnotateKernelAttrsPass) +}; + +} // namespace + +namespace mlir::dicp::CommonIR { +std::unique_ptr> createAnnotateKernelAttrsPass() { + return std::make_unique(); +} +} // namespace mlir::dicp::CommonIR diff --git a/compiler/lib/Dialect/CommonIR/CMakeLists.txt b/compiler/lib/Dialect/CommonIR/CMakeLists.txt new file mode 100644 index 00000000..099ad482 --- /dev/null +++ b/compiler/lib/Dialect/CommonIR/CMakeLists.txt @@ -0,0 +1,20 @@ +add_triton_library(CommonIRTransforms + VectorizeParallelLoopPass.cpp + AnnotateKernelAttrsPass.cpp + + DEPENDS + CommonIRPassIncGen + + LINK_LIBS PUBLIC + MLIRArithDialect + MLIRBufferizationDialect + MLIRFuncDialect + MLIRMathDialect + MLIRMemRefDialect + MLIRSCFDialect + MLIRTensorDialect + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport +) diff --git a/compiler/lib/Dialect/LinalgExt/Transforms/VectorizeParallelLoopPass.cpp b/compiler/lib/Dialect/CommonIR/VectorizeParallelLoopPass.cpp similarity index 92% rename from compiler/lib/Dialect/LinalgExt/Transforms/VectorizeParallelLoopPass.cpp rename to compiler/lib/Dialect/CommonIR/VectorizeParallelLoopPass.cpp index 8fdffb90..076a9b70 100644 --- a/compiler/lib/Dialect/LinalgExt/Transforms/VectorizeParallelLoopPass.cpp +++ b/compiler/lib/Dialect/CommonIR/VectorizeParallelLoopPass.cpp @@ -1,3 +1,5 @@ +#include "dicp/Dialect/CommonIR/Passes.h" + #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/Func/IR/FuncOps.h" @@ -10,19 +12,16 @@ #include "mlir/IR/PatternMatch.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h" + #include "llvm/Support/Debug.h" #include "llvm/Support/raw_ostream.h" using namespace mlir; -namespace mlir { -namespace dicp { -namespace LinalgExt { +namespace mlir::dicp::CommonIR { #define GEN_PASS_DEF_VECTORIZEPARALLELLOOPPASS -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h.inc" -} // namespace LinalgExt -} // namespace dicp -} // namespace mlir +#include "dicp/Dialect/CommonIR/Passes.h.inc" +} // namespace mlir::dicp::CommonIR #define DEBUG_TYPE "vectorize-parallel-loop-pass" @@ -110,6 +109,10 @@ struct VectorizeParallelLoopPattern : public OpRewritePattern { // scalarToTensorMap: 用于数据流向量化 (标量 Value -> 向量 Tensor Value) DenseMap scalarToTensorMap; + // storeTargetToTensor: 记录 local alloc store 的目标 memref → tensor + // 用于在循环处理完毕后，将 post-loop memref.copy 替换为 materialize + DenseMap storeTargetToTensor; + LLVM_DEBUG( llvm::dbgs() << "[VectorizeParallelLoop] Starting to process body operations...\n"); @@ -564,13 +567,19 @@ struct VectorizeParallelLoopPattern : public OpRewritePattern { if (isLocalAlloc) { LLVM_DEBUG(llvm::dbgs() << " Storing to local alloc memref.\n"); - // 直接 materialize 到目标 memref - auto matOp = - rewriter.create( - op.getLoc(), vectorResult, destMemref); - matOp.setWritable(true); - LLVM_DEBUG(llvm::dbgs() << " Created Materialize to local " - "memref (writable=true).\n"); + // 显式搬运: to_buffer + copy 将 tensor 数据写入 destMemref + // 避免 materialize_in_destination(writable=true) 与 post-loop + // memref.copy 之间因隐式 aliasing 产生缓冲区冲突 + auto toBufferOp = rewriter.create( + op.getLoc(), destMemref.getType(), vectorResult, nullptr); + rewriter.create(op.getLoc(), toBufferOp.getResult(), + destMemref); + LLVM_DEBUG(llvm::dbgs() + << " Created ToBuffer + Copy for " + "explicit tensor->memref data movement.\n"); + + // 记录 store 目标 memref → tensor，后续替换 post-loop memref.copy + storeTargetToTensor[destMemref] = vectorResult; // 重要：记录这个 memref 现在包含向量化的数据 // 后续的 load 可以直接使用这个 tensor @@ -646,6 +655,32 @@ struct VectorizeParallelLoopPattern : public OpRewritePattern { << inst.getName() << "\n"); } + // 3.5 扫描 post-loop 操作：将 memref.copy(storeTarget -> output) 替换为 + // materialize_in_destination(tensor -> output)，消除隐式 aliasing + if (!storeTargetToTensor.empty()) { + Block *parentBlock = op->getBlock(); + for (auto it = std::next(op->getIterator()); it != parentBlock->end(); + ++it) { + auto copyOp = dyn_cast(&*it); + if (!copyOp) + continue; + Value copySource = copyOp.getSource(); + auto tensorIt = storeTargetToTensor.find(copySource); + if (tensorIt == storeTargetToTensor.end()) + continue; + // 替换 post-loop memref.copy 为 materialize_in_destination + rewriter.setInsertionPoint(copyOp); + auto matOp = rewriter.create( + op.getLoc(), tensorIt->second, copyOp.getTarget()); + matOp.setWritable(true); + LLVM_DEBUG(llvm::dbgs() + << "[VectorizeParallelLoop] Replaced post-loop memref.copy " + "with materialize_in_destination (writable=true).\n"); + rewriter.eraseOp(copyOp); + break; + } + } + // 打印当前op LLVM_DEBUG({ llvm::dbgs() << "[VectorizeParallelLoop] Current Op: "; @@ -708,8 +743,8 @@ struct VectorizeParallelLoopPass } // namespace -namespace mlir::dicp::LinalgExt { +namespace mlir::dicp::CommonIR { std::unique_ptr> createVectorizeParallelLoopPass() { return std::make_unique(); } -} // namespace mlir::dicp::LinalgExt \ No newline at end of file +} // namespace mlir::dicp::CommonIR \ No newline at end of file diff --git a/compiler/lib/Dialect/LinalgExt/CMakeLists.txt b/compiler/lib/Dialect/LinalgExt/CMakeLists.txt deleted file mode 100644 index 5cd7900f..00000000 --- a/compiler/lib/Dialect/LinalgExt/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -dicp_add_all_subdirs() \ No newline at end of file diff --git a/compiler/lib/Dialect/LinalgExt/IR/CMakeLists.txt b/compiler/lib/Dialect/LinalgExt/IR/CMakeLists.txt deleted file mode 100644 index 07168992..00000000 --- a/compiler/lib/Dialect/LinalgExt/IR/CMakeLists.txt +++ /dev/null @@ -1,10 +0,0 @@ -add_dicp_triton_library(DICPLinalgExt - Dialect.cpp - Ops.cpp - - DEPENDS - DICPLinalgExtIncGen - - LINK_LIBS PUBLIC - MLIRLinalgDialect -) \ No newline at end of file diff --git a/compiler/lib/Dialect/LinalgExt/IR/Dialect.cpp b/compiler/lib/Dialect/LinalgExt/IR/Dialect.cpp deleted file mode 100644 index 9383cf0e..00000000 --- a/compiler/lib/Dialect/LinalgExt/IR/Dialect.cpp +++ /dev/null @@ -1,27 +0,0 @@ -#include "dicp/Dialect/LinalgExt/IR/LinalgExtOps.h" - -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/DialectImplementation.h" -#include "mlir/IR/IRMapping.h" -#include "mlir/IR/MLIRContext.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/OpImplementation.h" -#include "mlir/IR/Value.h" -#include "mlir/Support/LogicalResult.h" -#include "mlir/Transforms/InliningUtils.h" - -#include "llvm/ADT/SmallVector.h" - -#include "dicp/Dialect/LinalgExt/IR/LinalgExtDialect.cpp.inc" - -using namespace mlir; -namespace mlir::dicp::LinalgExt { - -void LinalgExtDialect::initialize() { - addOperations< -#define GET_OP_LIST -#include "dicp/Dialect/LinalgExt/IR/LinalgExtOps.cpp.inc" - >(); -} -} // namespace mlir::dicp::LinalgExt diff --git a/compiler/lib/Dialect/LinalgExt/IR/Ops.cpp b/compiler/lib/Dialect/LinalgExt/IR/Ops.cpp deleted file mode 100644 index 5396d700..00000000 --- a/compiler/lib/Dialect/LinalgExt/IR/Ops.cpp +++ /dev/null @@ -1,27 +0,0 @@ -#include "dicp/Dialect/LinalgExt/IR/LinalgExtOps.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/IR/AffineMap.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/IRMapping.h" -#include "mlir/IR/Location.h" -#include "mlir/IR/Value.h" -#include "mlir/IR/ValueRange.h" -#include "mlir/Interfaces/SideEffectInterfaces.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Support/LogicalResult.h" -#include "mlir/Transforms/InliningUtils.h" - -#include "llvm/ADT/SmallVector.h" - -#include - -#define GET_OP_CLASSES -#include "dicp/Dialect/LinalgExt/IR/LinalgExtOps.cpp.inc" - -using namespace mlir; -namespace mlir::dicp::LinalgExt {} \ No newline at end of file diff --git a/compiler/lib/Dialect/LinalgExt/Transforms/CMakeLists.txt b/compiler/lib/Dialect/LinalgExt/Transforms/CMakeLists.txt deleted file mode 100644 index f5d52f63..00000000 --- a/compiler/lib/Dialect/LinalgExt/Transforms/CMakeLists.txt +++ /dev/null @@ -1,30 +0,0 @@ -add_triton_library(LinalgExtTransforms - LinalgIfToSelect.cpp - LinalgGenericToSCF.cpp - ScalarTo1DTensorPass.cpp - RemoveSingleIterationLoop.cpp - TensorTransform.cpp - VectorizeParallelLoopPass.cpp - - DEPENDS - LinalgExtTransformsIncGen - - LINK_LIBS PUBLIC - TritonTilingExtIR - MLIRArithDialect - MLIRDialectUtils - MLIRIR - MLIRMathDialect - MLIRPass - MLIRTensorDialect - MLIRTransforms - MLIRSupport - TritonAnalysis - TritonIR - TritonTransforms - TritonSharedAnalysis - - TritonArithToLinalg - StructuredToMemref - TritonToStructured -) diff --git a/compiler/lib/Dialect/LinalgExt/Transforms/LinalgGenericToSCF.cpp b/compiler/lib/Dialect/LinalgExt/Transforms/LinalgGenericToSCF.cpp deleted file mode 100644 index 083ea25f..00000000 --- a/compiler/lib/Dialect/LinalgExt/Transforms/LinalgGenericToSCF.cpp +++ /dev/null @@ -1,336 +0,0 @@ -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h" -#include "dicp/Dialect/LinalgExt/Transforms/Transforms.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/Linalg/Utils/Utils.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Operation.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" -#include "mlir/Transforms/Passes.h" - -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/Support/Casting.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/LogicalResult.h" -#include "llvm/Support/raw_ostream.h" -#include - -#define DEBUG_TYPE "linalg-generic-to-scf" - -using namespace mlir; -using namespace dicp; -using namespace LinalgExt; - -namespace mlir::dicp::LinalgExt { -#define GEN_PASS_DEF_LINALGGENERICTOSCF -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h.inc" -} // namespace mlir::dicp::LinalgExt - -namespace { - -/** - * @brief Converts a multi-dimensional parallel linalg.generic (on tensors) into - * nested scf.for loops. - * - * Handles cases with and without outputs. - * - * Case 1: With Outputs (Accumulator/Map style) - * %res = linalg.generic ... ins(...) outs(%out) - * -> - * %res = scf.for ... iter_args(%curr = %out) { - * ... - * %new = tensor.insert ... - * scf.yield %new - * } - * - * Case 2: Without Outputs (Side-effecting/Void style) - * linalg.generic ... ins(...) outs() - * -> - * scf.for ... { - * ... - * scf.yield - * } - */ -struct LinalgGenericToScfForPattern - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(linalg::GenericOp linalgOp, - PatternRewriter &rewriter) const override { - // --- 0. Pre-check for specific IndexCast case --- - auto linalgYield = mlir::dyn_cast_or_null( - linalgOp.getBody()->getTerminator()); - if (linalgYield && linalgYield.getNumOperands() == 1) { - Value yieldValue = linalgYield.getOperand(0); - if (llvm::isa_and_nonnull( - yieldValue.getDefiningOp())) { - return rewriter.notifyMatchFailure( - linalgOp, "Conversion skipped: index increment detected."); - } - } - - // --- 1. Check LinalgOp Properties --- - // Check: all iterator types must be "parallel" - if (!llvm::all_of(linalgOp.getIteratorTypesArray(), - linalg::isParallelIterator)) { - return rewriter.notifyMatchFailure( - linalgOp, "Failed match: expected all parallel iterators"); - } - - // Check: all indexing maps must be identity - for (AffineMap map : linalgOp.getIndexingMapsArray()) { - if (!map.isIdentity()) { - return rewriter.notifyMatchFailure( - linalgOp, "Conversion failed: expected all maps to be identity."); - } - } - - // Check Output presence - // We allow 0 outputs or exactly 1 output. - if (linalgOp.getNumDpsInits() > 1) { - return rewriter.notifyMatchFailure( - linalgOp, "Conversion failed: >1 outputs not supported yet."); - } - - bool hasOutput = (linalgOp.getNumDpsInits() == 1); - Value outputTensor = nullptr; - - if (hasOutput) { - outputTensor = linalgOp.getDpsInitOperand(0)->get(); - // Verify output comes from tensor.empty if present - auto emptyOp = outputTensor.getDefiningOp(); - if (!emptyOp) { - return rewriter.notifyMatchFailure( - linalgOp, - "Conversion failed: output must originate from tensor.empty."); - } - } - - // --- 2. Determine Loop Bounds --- - // Use output shape if available, otherwise use the first input shape. - ShapedType boundsType; - if (hasOutput) { - boundsType = mlir::cast(outputTensor.getType()); - } else { - if (linalgOp.getNumDpsInputs() == 0) { - return rewriter.notifyMatchFailure( - linalgOp, - "Conversion failed: no inputs or outputs to determine bounds."); - } - Value firstInput = linalgOp.getDpsInputOperand(0)->get(); - boundsType = mlir::cast(firstInput.getType()); - } - - if (!boundsType.hasStaticShape()) { - return rewriter.notifyMatchFailure(linalgOp, - "Dynamic shapes are not supported."); - } - ArrayRef shape = boundsType.getShape(); - int64_t rank = boundsType.getRank(); - - if (rank == 0) { - return rewriter.notifyMatchFailure(linalgOp, "Rank 0 not supported."); - } - - Location loc = linalgOp.getLoc(); - Value c0 = rewriter.create(loc, 0); - Value c1 = rewriter.create(loc, 1); - - // --- 3. Build Nested Loops --- - SmallVector loops; - SmallVector ivs; - - // The 'current' tensor being threaded through iter_args. - Value currentThreadedTensor = hasOutput ? outputTensor : nullptr; - - // Create loops from outer (dim 0) to inner (dim rank-1). - for (int64_t i = 0; i < rank; ++i) { - Value ub = rewriter.create(loc, shape[i]); - - // Prepare iter_args for this specific loop level - SmallVector iterArgs; - if (hasOutput) { - iterArgs.push_back(currentThreadedTensor); - } - - auto forOp = rewriter.create(loc, c0, ub, c1, iterArgs); - - loops.push_back(forOp); - ivs.push_back(forOp.getInductionVar()); - - rewriter.setInsertionPointToStart(forOp.getBody()); - - // Update threaded tensor for the next inner loop - if (hasOutput) { - currentThreadedTensor = forOp.getRegionIterArg(0); - } - - if (linalgOp->hasAttr("ExtractedLoadOrStore")) { - forOp->setAttr("ExtractedLoadOrStore", rewriter.getUnitAttr()); - } - } - - // --- 4. Build Innermost Body --- - IRMapping map; - Block *linalgBody = linalgOp.getBody(); - - unsigned inputIdx = 0; - for (BlockArgument &arg : linalgBody->getArguments()) { - Value newSource; - - if (arg.getArgNumber() < linalgOp.getNumDpsInputs()) { - // --- Map Input --- - Value inputTensor = linalgOp.getDpsInputOperand(inputIdx++)->get(); - auto extractOp = - rewriter.create(loc, inputTensor, ivs); - newSource = extractOp; - } else { - // --- Map Output (Accumulator) --- - auto extractOp = - rewriter.create(loc, currentThreadedTensor, ivs); - newSource = extractOp; - } - map.map(arg, newSource); - } - - // Clone linalg.generic operations (excluding linalg.yield) - for (Operation &op : linalgBody->without_terminator()) { - rewriter.clone(op, map); - } - - // --- 5. Handle Yield (Innermost) --- - SmallVector innerYieldOperands; - - if (hasOutput) { - // If there is an output, insert the calculated scalar back into the - // tensor. - Value computedScalar = map.lookup(linalgYield.getOperand(0)); - Value insertedTensor = rewriter.create( - loc, computedScalar, currentThreadedTensor, ivs); - innerYieldOperands.push_back(insertedTensor); - } - if (!innerYieldOperands.empty()) - auto innerYield = rewriter.create(loc, innerYieldOperands); - - // --- 6. Handle Yields for Outer Loops --- - // Walk back up the loop nest. - for (int64_t i = rank - 2; i >= 0; --i) { - scf::ForOp innerLoop = loops[i + 1]; - - // Set insertion point to after the inner loop - rewriter.setInsertionPointAfter(innerLoop); - - // Yield the result of the inner loop (if any) - rewriter.create(loc, innerLoop.getResults()); - } - - // --- 7. Finalize Replacement --- - if (hasOutput) { - rewriter.replaceOp(linalgOp, loops[0].getResults()); - } else { - rewriter.eraseOp(linalgOp); - } - - return success(); - } -}; - -/// Pattern that matches: tensor.empty -> linalg.fill -> tensor.extract -/// and optimizes it to directly use the scalar value. -class SimplifySingleElementFillExtractPattern - : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(tensor::ExtractOp extractOp, - PatternRewriter &rewriter) const override { - // 1. Verify all indices are constant zero - for (Value index : extractOp.getIndices()) { - auto constIndex = index.getDefiningOp(); - if (!constIndex || constIndex.value() != 0) { - return failure(); - } - } - - // 2. Get the source tensor and check if it comes from linalg.fill - auto fillOp = extractOp.getTensor().getDefiningOp(); - if (!fillOp) { - return failure(); - } - - // 3. Verify the filled tensor has exactly one element - auto filledTensorType = - dyn_cast(fillOp.getOutputs()[0].getType()); - if (!filledTensorType || filledTensorType.getNumElements() != 1) { - return failure(); - } - - // 4. Verify fill value is a scalar (not a tensor) - Value fillValue = fillOp.getInputs()[0]; - if (isa(fillValue.getType())) { - return failure(); - } - - // 5. For safety, ensure the filled tensor is only used by this extract - // This prevents breaking other uses of the same tensor - if (!fillOp.getResult(0).hasOneUse()) { - return failure(); - } - - // 6. Try to find the tensor.empty operation - // We'll check its uses after we replace the extract operation - auto emptyOp = - fillOp.getDpsInitOperand(0)->get().getDefiningOp(); - - // 7. Replace extract with the scalar fill value - rewriter.replaceOp(extractOp, fillValue); - - // 8. Now that extract is replaced, fillOp's result should have no uses - // We can safely erase it - rewriter.eraseOp(fillOp); - - // 9. After fillOp is erased, check if emptyOp still exists and has no uses - if (emptyOp && emptyOp->use_empty()) { - rewriter.eraseOp(emptyOp); - } - - return success(); - } -}; - -struct LinalgGenericToSCFPass - : mlir::dicp::LinalgExt::impl::LinalgGenericToSCFBase< - LinalgGenericToSCFPass> { - - void runOnOperation() override { - auto moduleOp = getOperation(); - MLIRContext *context = &getContext(); - { - RewritePatternSet patterns(context); - patterns.add(context); - patterns.add(context); - populateRemoveSingleIterationLoopPattern(patterns); - if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { - signalPassFailure(); - } - } - } -}; - -} // namespace - -std::unique_ptr> -mlir::dicp::LinalgExt::createLinalgGenericToSCFPass() { - return std::make_unique(); -} \ No newline at end of file diff --git a/compiler/lib/Dialect/LinalgExt/Transforms/LinalgIfToSelect.cpp b/compiler/lib/Dialect/LinalgExt/Transforms/LinalgIfToSelect.cpp deleted file mode 100644 index 564ce2a8..00000000 --- a/compiler/lib/Dialect/LinalgExt/Transforms/LinalgIfToSelect.cpp +++ /dev/null @@ -1,885 +0,0 @@ -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h" -#include "dicp/Dialect/LinalgExt/Transforms/Transforms.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/Linalg/Utils/Utils.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Operation.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" -#include "mlir/Transforms/Passes.h" - -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/Support/Casting.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/LogicalResult.h" -#include "llvm/Support/raw_ostream.h" -#include - -#define DEBUG_TYPE "linalg-if-to-select" - -using namespace mlir; -using namespace dicp; -using namespace LinalgExt; - -namespace mlir::dicp::LinalgExt { -#define GEN_PASS_DEF_LINALGIFTOSELECT -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h.inc" -} // namespace mlir::dicp::LinalgExt - -namespace { - -// --- Lift scalar arith.select (inside linalg.generic) to -// tensor arith.select before the linalg.generic and add it as an extra -// `ins` operand of generic. --- -struct LiftScalarSelectToTensorPattern - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(linalg::GenericOp linalgOp, - PatternRewriter &rewriter) const override { - if (!linalgOp->hasAttr("ExtractedLoadOrStore")) - return rewriter.notifyMatchFailure( - linalgOp, "missing 'ExtractedLoadOrStore' attribute"); - - // Find a scalar arith.select in the body whose operands are either - // block arguments or constants. - Block *body = linalgOp.getBody(); - arith::SelectOp scalarSelect; - for (Operation &op : body->getOperations()) { - if (auto s = dyn_cast(op)) { - // Ensure it is scalar (not tensor) - if (isa(s.getType())) - continue; - // Only consider select whose operands (true/false/cond) are - // either block args or loop-invariant - Value t = s.getTrueValue(); - Value f = s.getFalseValue(); - Value c = s.getCondition(); - - // 'v' is allowed if it's a BlockArgument OR defined outside the - // linalg op's body (making it a loop-invariant). - auto isAllowed = [&](Value v) -> bool { - // Case 1: Block argument (e.g., %arg0) - if (isa(v)) - return true; - - // Case 2: Any SSA value defined *outside* the linalg body. - // This includes arith.constant and other loop-invariant values. - Operation *defOp = v.getDefiningOp(); - if (defOp && defOp->getBlock() != body) { - return true; - } - - return false; - }; - - if (isAllowed(t) && isAllowed(f) && isAllowed(c)) { - scalarSelect = s; - break; - } - } - } - - if (!scalarSelect) - return rewriter.notifyMatchFailure( - linalgOp, - "could not find a liftable scalar arith.select in the body"); - - Location loc = linalgOp.getLoc(); - - // Determine the tensor shape to materialize the lifted operands. - // Prefer first input's shape. - if (linalgOp.getNumDpsInputs() == 0) - return rewriter.notifyMatchFailure( - linalgOp, "linalg.generic has no input operands to determine shape"); - - Value firstInput = linalgOp.getDpsInputOperand(0)->get(); - auto firstTy = dyn_cast(firstInput.getType()); - if (!firstTy || !firstTy.hasStaticShape()) - return rewriter.notifyMatchFailure( - linalgOp, "first input is not a RankedTensorType or does not have " - "static shape"); - ArrayRef shape = firstTy.getShape(); - - // Utility to materialize a tensor from either a constant scalar or a - // block-argument-mapped input tensor. - auto materializeTensor = [&](Value v, Type desiredElemType) -> Value { - // Case 1: BlockArgument (maps to an existing input tensor) - if (auto ba = dyn_cast(v)) { - unsigned argNo = ba.getArgNumber(); - // block args are ordered as inputs then outputs. - if (argNo >= linalgOp.getNumDpsInputs()) { - // Block argument refers to an output tensor, which is not an allowed - // input for select in this pattern. - return nullptr; - } - Value tensorOperand = linalgOp.getDpsInputOperand(argNo)->get(); - // TODO: Ensure it has same element type or cast. For now, just return. - return tensorOperand; - } - - // Case 2: Loop-invariant SSA scalar value (from outside) - // This now correctly includes arith.constant and any other op. - Operation *defOp = v.getDefiningOp(); - if (defOp && defOp->getBlock() != body) { - // 'v' is the loop-invariant scalar value. Broadcast it. - Value empty = - rewriter.create(loc, shape, desiredElemType); - - // 'v' is the scalar to fill with. - auto fill = rewriter.create(loc, v, empty); - return fill.getResult(0); - } - - // 'v' is not a BlockArgument and not defined outside. - return nullptr; - }; - - // Materialize cond tensor, true tensor, false tensor - Value condTensor = materializeTensor(scalarSelect.getCondition(), - scalarSelect.getCondition().getType()); - if (!condTensor) - return rewriter.notifyMatchFailure( - linalgOp, "failed to materialize condition tensor"); - - Type trueElemTy = scalarSelect.getTrueValue().getType(); - Value trueTensor = - materializeTensor(scalarSelect.getTrueValue(), trueElemTy); - if (!trueTensor) - return rewriter.notifyMatchFailure( - linalgOp, "failed to materialize true-value tensor"); - - Type falseElemTy = scalarSelect.getFalseValue().getType(); - Value falseTensor = - materializeTensor(scalarSelect.getFalseValue(), falseElemTy); - if (!falseTensor) - return rewriter.notifyMatchFailure( - linalgOp, "failed to materialize false-value tensor"); - - // Create the lifted tensor select before the linalg.generic - rewriter.setInsertionPoint(linalgOp); - RankedTensorType condType = - dyn_cast(condTensor.getType()); - RankedTensorType trueType = - dyn_cast(trueTensor.getType()); - if (!condType || !trueType) - return rewriter.notifyMatchFailure( - linalgOp, "materialized tensor type is not a RankedTensorType"); - - // arith.select signature: (cond, true, false) -> result type (trueType) - Value tensorSelect = rewriter.create( - loc, /*result type*/ trueTensor.getType(), condTensor, trueTensor, - falseTensor); - - // --- Fix: Rebuild LinalgOp --- - unsigned numInputs = linalgOp.getNumDpsInputs(); - unsigned numOutputs = linalgOp.getNumDpsInits(); - - // 1. Prepare new operand list - auto newIns = llvm::to_vector<4>( - llvm::map_range(linalgOp.getDpsInputOperands(), - [](OpOperand *opnd) { return opnd->get(); })); - - newIns.push_back( - tensorSelect); // Add the new tensor select as the last input - ValueRange newOuts = linalgOp.getDpsInits(); - - // 2. Prepare new indexing_maps attribute - SmallVector newMapAttrs; - for (Attribute a : linalgOp.getIndexingMaps()) - newMapAttrs.push_back(a); - - // Add an identity map for the new input (reusing the first map or creating - // new) - if (!linalgOp.getIndexingMaps().empty()) { - newMapAttrs.push_back(linalgOp.getIndexingMaps()[0]); - } else { - // Fallback: create an identity map if original maps were empty - // This is necessary because we checked for firstInput above. - if (!firstTy) - return rewriter.notifyMatchFailure( - linalgOp, "cannot determine rank for identity map fallback"); - newMapAttrs.push_back( - AffineMapAttr::get(AffineMap::getMultiDimIdentityMap( - firstTy.getRank(), rewriter.getContext()))); - } - - ArrayAttr newMapsAttr = rewriter.getArrayAttr(newMapAttrs); - - // 3. Prepare new operandSegmentSizes attribute - auto newSegmentSizes = rewriter.getDenseI32ArrayAttr( - {static_cast(numInputs + 1), // New number of Inputs - static_cast(numOutputs)}); - - // 4. Create the new linalg.generic op and rebuild the region using - // bodyBuilder - auto newLinalgOp = rewriter.create( - loc, - /*resultTensorTypes=*/linalgOp.getResultTypes(), - /*inputs=*/newIns, - /*outputs=*/newOuts, - /*indexing_maps=*/newMapsAttr, - /*iterator_types=*/linalgOp.getIteratorTypes(), - /*doc=*/nullptr, - /*library_call=*/nullptr, - /*bodyBuilder=*/ - [&](OpBuilder &b, Location bodyLoc, ValueRange newBlockArgs) { - // newBlockArgs now contains (old_ins, new_tensor_select_arg, - // old_outs) - IRMapping bvm; - unsigned oldIdx = 0; - unsigned newIdx = 0; - - // Map old input block args - for (oldIdx = 0; oldIdx < numInputs; ++oldIdx, ++newIdx) { - bvm.map(linalgOp.getBody()->getArgument(oldIdx), - newBlockArgs[newIdx]); - } - - // This is the new block arg corresponding to new_tensor_select - Value newArgForTensorSelect = newBlockArgs[newIdx++]; - - // Map old output block args - for (; oldIdx < linalgOp.getBody()->getNumArguments(); - ++oldIdx, ++newIdx) { - bvm.map(linalgOp.getBody()->getArgument(oldIdx), - newBlockArgs[newIdx]); - } - - // Clone the linalgOp body, skipping the old scalarSelect - for (Operation &op : linalgOp.getBody()->without_terminator()) { - // Check if this op is the scalarSelect we are replacing - if (auto oldSel = dyn_cast(op)) { - if (oldSel == scalarSelect) { - // Yes! Map the result of the old select to the new block arg - bvm.map(oldSel.getResult(), newArgForTensorSelect); - // Do not clone this op - continue; - } - } - // Clone all other ops - b.clone(op, bvm); - } - - // Clone the terminator (linalg.yield) - b.clone(*linalgOp.getBody()->getTerminator(), bvm); - }); - - // 5. Copy other attributes from linalgOp (e.g., "ExtractedLoadOrStore") - for (const auto &attr : linalgOp->getAttrs()) { - if (attr.getName() == linalgOp.getOperandSegmentSizesAttrName() || - attr.getName() == linalgOp.getIndexingMapsAttrName()) { - continue; - } - newLinalgOp->setAttr(attr.getName(), attr.getValue()); - } - - // 6. Replace the old op results with the new op results - rewriter.replaceOp(linalgOp, newLinalgOp.getResults()); - - return success(); - } -}; - -// --- If the linalg.yield returns a scalar select that chooses -// between a constant and a scalar op inside the region, lift the select to -// after the linalg.generic as a tensor select. --- -struct LiftYieldSelectOutPattern : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(linalg::GenericOp linalgOp, - PatternRewriter &rewriter) const override { - if (!linalgOp->hasAttr("ExtractedLoadOrStore")) - return rewriter.notifyMatchFailure( - linalgOp, "missing 'ExtractedLoadOrStore' attribute"); - - Block *body = linalgOp.getBody(); - auto linalgYield = - mlir::dyn_cast_or_null(body->getTerminator()); - if (!linalgYield || linalgYield.getNumOperands() == 0) - return rewriter.notifyMatchFailure( - linalgOp, - "body terminator is not a linalg.yield or yields no values"); - - Value yielded = linalgYield.getOperand(0); - auto sel = yielded.getDefiningOp(); - if (!sel) - return rewriter.notifyMatchFailure( - linalgOp, "yielded value is not an arith.select op"); - - // Conditions: select is single-use (used only by linalg.yield), condition - // comes from a block arg, and one operand is a constant while the other is - // a scalar op inside the region with one use. - if (!llvm::hasSingleElement(sel->getUsers())) - return rewriter.notifyMatchFailure( - linalgOp, - "arith.select has more than one user (not only linalg.yield)"); - - Value cond = sel.getCondition(); - if (!isa(cond)) - return rewriter.notifyMatchFailure( - linalgOp, "select condition is not a block argument"); - - rewriter.setInsertionPointAfter(linalgOp); - - Value trueV = sel.getTrueValue(); - Value falseV = sel.getFalseValue(); - - // Identify which is constant and which is an internal op - Value constScalar; - Operation *internalOp = nullptr; - if (trueV.getDefiningOp()) { - constScalar = trueV; - internalOp = falseV.getDefiningOp(); - } else if (falseV.getDefiningOp()) { - constScalar = falseV; - internalOp = trueV.getDefiningOp(); - } else { - return rewriter.notifyMatchFailure( - linalgOp, "neither select operand is an arith.constant op"); - } - - if (!internalOp) - return rewriter.notifyMatchFailure( - linalgOp, "internal operand of select is not a defined operation"); - - // internalOp must be single-use and produce a scalar compatible with yield - if (!internalOp->hasOneUse()) - return rewriter.notifyMatchFailure(linalgOp, - "internal op has more than one use"); - if (internalOp->getNumResults() != 1) - return rewriter.notifyMatchFailure( - linalgOp, "internal op does not have exactly one result"); - - Type scalarTy = internalOp->getResult(0).getType(); - if (isa(scalarTy)) - return rewriter.notifyMatchFailure( - linalgOp, "internal op result is a TensorType (expected scalar)"); - - // Replace yield's operand with the internalOp's result so that the generic - // returns the "internal" values instead of the select results. - linalgYield->setOperand(0, internalOp->getResult(0)); - - // After the linalg.generic, create a tensor constant from constScalar and - // then a tensor select between the generic's result tensor and the - // constant tensor using the cond input tensor. - Location loc = linalgOp.getLoc(); - - // Determine result type (assume first result) and element type - if (linalgOp.getOperation()->getNumResults() == 0) - return rewriter.notifyMatchFailure(linalgOp, - "linalg.generic op has no results"); - - Value genericResult = linalgOp.getResult(0); - auto resultTy = dyn_cast(genericResult.getType()); - if (!resultTy) - return rewriter.notifyMatchFailure( - linalgOp, "first result of linalg.generic is not a RankedTensorType"); - - ArrayRef shape = resultTy.getShape(); - Type elemTy = resultTy.getElementType(); - - // materialize tensor from scalar constant - Value empty = rewriter.create(loc, shape, elemTy); - // scalar constant - auto constOp = constScalar.getDefiningOp(); - if (!constOp) - return rewriter.notifyMatchFailure( - linalgOp, "constant scalar is not a valid arith.constant op"); - - Value scalarConst = - rewriter.create(loc, elemTy, constOp.getValueAttr()); - Value filled = - rewriter.create(loc, scalarConst, empty).getResult(0); - - // cond tensor: map block argument to corresponding input tensor - auto condBA = cast(cond); - unsigned condArgNo = condBA.getArgNumber(); - if (condArgNo >= linalgOp.getNumDpsInputs()) - return rewriter.notifyMatchFailure( - linalgOp, - "condition block argument index is out of bounds for dps_inputs"); - - Value condTensor = linalgOp.getDpsInputOperand(condArgNo)->get(); - - // Insert the select after the generic op - Value tensorSelect = rewriter.create( - loc, /*result type*/ genericResult.getType(), condTensor, genericResult, - filled); - - // Replace uses of the generic result (outside the generic) with the - // tensorSelect result. Do not replace uses inside the generic (there - // shouldn't be any, but guard anyway). - genericResult.replaceAllUsesExcept(tensorSelect, - tensorSelect.getDefiningOp()); - - return success(); - } -}; - -/** - * @brief Converts scf.if within linalg.generic to arith.select. - * - * This pattern matches an scf.if inside a linalg.generic op. - * - * Requirements (Request 1 & 2): - * 1. The if op must be inside a linalg.generic op. - * 2. The linalg.generic op must have all "parallel" iterators and all - * "identity" maps. - * 3. All operations within the if regions (then/else) must produce scalar -values. * 4. If the if op has results and an else block: the else block must -contain * only one scf.yield whose yielded value is a constant. * 5. If the if -op has no results: the else block must be empty (or contain * only scf.yield). - * - * Conversion Logic: - * 1. Iterate over all operations (op) in the `then` block. - * 2. For each operand (val) of the op, check if val is a block argument (arg) - * of the linalg.generic. - * 3. If it is and the arg hasn't been mapped yet: - * a. Create a zero constant %zero (of the same type as arg). - * b. Create %sel_arg = arith.select %cond, %arg, %zero {recheck}. - * c. Store the mapping { %arg -> %sel_arg } in IRMapping (bvm). - * 4. Clone the body of the `then` block (excluding the terminator) before - * the ifOp using bvm for value mapping. - * 5. (Case A: ifOp has results, see Example 1) - * a. Get the `then` block's yield value (%then_val) and the `else` block's - * constant yield value (%else_const). - * b. Find the mapped value for %then_val in bvm (i.e., %cloned_then_val). - * c. Create the final select: - * `%final_sel = arith.select %cond, %cloned_then_val, %else_const {recheck}` - * d. Replace all uses of the ifOp with %final_sel and erase the ifOp. - * 6. (Case B: ifOp has no results, see Example 2) - * a. Erase the ifOp directly (since the `then` block's body has been cloned). - */ -struct LinalgIfToSelectPattern : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(scf::IfOp ifOp, - PatternRewriter &rewriter) const final { - - // --- 0. Find the parent linalg.generic operation --- - auto linalgOp = ifOp->getParentOfType(); - if (!linalgOp) { - // This scf.if is not inside a linalg.generic, match fails - return failure(); - } - Block *linalgBody = linalgOp.getBody(); - - // --- 1. Check linalg.generic constraints (Request 2) --- - // Linalg op must have identity maps - if (!llvm::all_of(linalgOp.getIteratorTypesArray(), - linalg::isParallelIterator)) { - return rewriter.notifyMatchFailure( - linalgOp, "Failed match: expected all parallel iterators."); - } - for (AffineMap map : linalgOp.getIndexingMapsArray()) { - if (!map.isIdentity()) { - return rewriter.notifyMatchFailure( - linalgOp, "Failed match: expected identity maps."); - } - } - - // --- 2. Check scf.if constraints --- - // Check: all results of ops in the if are scalar - bool allScalars = true; - ifOp.walk([&](Operation *innerOp) { - // Skip terminators (scf.yield) and the IfOp itself - if (isa(innerOp) || innerOp == ifOp.getOperation()) - return; - - for (Value res : innerOp->getResults()) { - if (isa(res.getType())) { - allScalars = false; - } - } - }); - if (!allScalars) { - return rewriter.notifyMatchFailure( - linalgOp, "Failed match: if region contains non-scalar ops."); - } - - // --- 3. Get condition, location, and set insertion point --- - Location loc = ifOp.getLoc(); - Value cond = ifOp.getCondition(); - // Set insertion point *before* the ifOp - rewriter.setInsertionPoint(ifOp); - - // --- 4. Create arith.select only for linalg block args *used* in 'then' - // block --- - IRMapping bvm; - Block *thenBlock = ifOp.thenBlock(); - - // Iterate over ops in 'then' block - for (Operation &op : thenBlock->without_terminator()) { - // Iterate over all operands of the op - for (Value val : op.getOperands()) { - auto ba = dyn_cast(val); - // Check: 1. Is a block argument - // 2. Belongs to linalgOp (not a nested op) - // 3. Select has not been created for this block argument yet - if (ba && ba.getOwner() == linalgBody && !bvm.contains(ba)) { - Type argType = ba.getType(); - // Create %false_value (zero constant) - Value falseValue; - if (isa(argType)) { - falseValue = rewriter.create(loc, 0); - - } else if (auto intTy = dyn_cast(argType)) { - falseValue = rewriter.create( - loc, argType, rewriter.getIntegerAttr(argType, 0)); - - } else if (auto floatTy = dyn_cast(argType)) { - falseValue = rewriter.create( - loc, argType, rewriter.getFloatAttr(argType, 0.0)); - - } else { - return rewriter.notifyMatchFailure( - linalgOp, - "Failed match: unsupported block arg type for zero init."); - return failure(); - } - - // Create arith.select cond, v, false_value - Value newSel = rewriter.create(loc, /*type*/ argType, - cond, ba, falseValue); - newSel.getDefiningOp()->setAttr("recheck", rewriter.getUnitAttr()); - // Map: old_arg -> new_select - bvm.map(ba, newSel); - } - } - } - - // --- 5. Perform transformation (based on whether ifOp has results) --- - if (ifOp.getNumResults() > 0) { - // --- Case 1: Has results (e.g., Example 1) --- - - // Check else block: must exist, contain only one scf.yield, and yield - // a constant (Request 4) - Block *elseBlock = ifOp.elseBlock(); - if (!elseBlock || elseBlock->empty()) { - return rewriter.notifyMatchFailure( - linalgOp, "Failed match: if with results has no else block."); - } - if (!llvm::hasSingleElement(*elseBlock)) { - return rewriter.notifyMatchFailure( - linalgOp, "Failed match: else block has more than one op."); - } - auto elseYield = dyn_cast(elseBlock->getTerminator()); - if (!elseYield || elseYield.getNumOperands() == 0) { - return rewriter.notifyMatchFailure( - linalgOp, "Failed match: else block has no valid yield."); - } - - Value elseValue = elseYield.getOperand(0); - if (!elseValue.getDefiningOp()) { - return rewriter.notifyMatchFailure( - linalgOp, "Failed match: else block does not yield a constant."); - return failure(); - } - - // Get the `then` block's yield value - auto thenYield = dyn_cast(thenBlock->getTerminator()); - if (!thenYield || thenYield->getNumOperands() == 0) { - return rewriter.notifyMatchFailure( - linalgOp, "Failed match: then block has no valid yield."); - return failure(); - } - Value thenValue = thenYield->getOperand(0); - - // Clone the operations of the 'then' block (excluding terminator) before - // ifOp - // BVM will automatically replace block-args - // BVM will also be updated during cloning: old_op_result -> new_op_result - for (Operation &op : thenBlock->without_terminator()) { - rewriter.clone(op, bvm); - } - - // Find the (potentially remapped) 'then' yield value - Value newThenValue = bvm.lookupOrDefault(thenValue); - - // Create the final arith.select (to replace ifOp) - Value finalSel = rewriter.create( - loc, newThenValue.getType(), cond, newThenValue, elseValue); - - // Replace ifOp results and erase ifOp - rewriter.replaceOp(ifOp, finalSel); - - } else { - // --- Case 2: No results (e.g., Example 2) --- - - // Check else block: must be empty (or only yield) - Block *elseBlock = ifOp.elseBlock(); - if (elseBlock && !elseBlock->empty()) { - // If the else block is not empty, it *must* have only one scf.yield - if (!isa(elseBlock->front()) || - !llvm::hasSingleElement(*elseBlock)) { - return rewriter.notifyMatchFailure( - linalgOp, - "Failed match: if with no results has non-empty else block."); - } - } - - // Clone the operations of the 'then' block (excluding terminator) before - // ifOp - // BVM will automatically replace block-args - for (Operation &op : thenBlock->without_terminator()) { - rewriter.clone(op, bvm); - } - - // Erase the if op - rewriter.eraseOp(ifOp); - } - - return success(); - } -}; - -/** - * @brief Fixes Linalg select's false value being zero for a store operation. - * - * Motivation: - * The `LinalgIfToSelectPattern` transforms `scf.if` into `arith.select`s. - * In this conversion, it defaults the 'false' value (when the condition is -false) * for a select operation to a zero constant (0). * * This zero value -is incorrect when the result of this `select` is used as the * *value* for a -`memref.store`. * * The correct 'false' value should be the *original* value -in the memref, as * a false condition means we should not modify the memory -location (i.e., we * store the old value). * * This pattern aims to correct -this error. It finds the "value select" (the select * providing the value to -store), finds the corresponding "index select" (providing * the index), and -replaces the incorrect zero with: * `memref.load(%memref, %false_index)` * -(i.e., the original value at the "false" index). * * Before Transformation -(IR): * * // "index select" (provides index for store) * %sel_idx = -arith.select %cond, %idx_t, %idx_f {recheck} * %idx_cast = arith.index_cast -%sel_idx : i32 to index * * // "value select" (provides value for store) * -%cst_0 = arith.constant 0.0 : f32 // <-- Incorrect 0 * %sel_val = arith.select -%cond, %val_t, %cst_0 {recheck} * * memref.store %sel_val, %memref[%idx_cast] - * - * After Transformation (Corrected IR): - * - * %sel_idx = arith.select %cond, %idx_t, %idx_f {recheck} - * %idx_cast = arith.index_cast %sel_idx : i32 to index - * - * // Correction: - * %idx_f_cast = arith.index_cast %idx_f : i32 to index - * // Load original value at the "false" index, hoisted outside linalg if -invariant * %new_val_f = memref.load %memref[%idx_f_cast] * // Replace 0 with -the original value * %sel_val = arith.select %cond, %val_t, %new_val_f * * -memref.store %sel_val, %memref[%idx_cast] */ -struct FixLinalgSelectZeroForStorePattern - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(arith::SelectOp valueSel, - PatternRewriter &rewriter) const final { - // 1. Must have the "recheck" attribute - if (!valueSel->hasAttr("recheck")) - return failure(); - - // 2. Check for "skip": if a user is arith.index_cast, this is an index - // select, not a value select - if (llvm::any_of(valueSel->getUsers(), - - [](Operation *u) { return isa(u); })) { - return rewriter.notifyMatchFailure(valueSel, - "Has index_cast user, skipping."); - } - - // 3. Find the "memref.store" use - memref::StoreOp storeOp; - for (Operation *user : valueSel->getUsers()) { - if (auto s = dyn_cast(user)) { - if (s.getValue() == valueSel.getResult()) { - storeOp = s; - break; - } - } - } - if (!storeOp) { - return rewriter.notifyMatchFailure( - storeOp, "Not used as value in a memref.store."); - return failure(); - } - - // 4. Analyze store indices (assume 1D) - if (storeOp.getIndices().size() != 1) { - return rewriter.notifyMatchFailure(storeOp, "Store is not 1D, skipping."); - } - Value storeIndex = storeOp.getIndices().front(); - - // 5. Find the "index select" - arith::SelectOp indexSel; - if (auto castOp = storeIndex.getDefiningOp()) { - indexSel = castOp.getIn().getDefiningOp(); - - } else { - indexSel = storeIndex.getDefiningOp(); - } - if (!indexSel) { - return rewriter.notifyMatchFailure( - storeOp, "Store index does not come from a select or cast(select)."); - } - - // --- 6. Check invariants and set insertion point --- - Value falseIndex = indexSel.getFalseValue(); - Value memref = storeOp.getMemRef(); - Location loc = valueSel.getLoc(); - - auto linalgOp = valueSel->getParentOfType(); - - // Helper function: Check if a value v is defined *outside* the linalgOp's - // body - auto isInvariant = [&](Value v) { - if (!linalgOp) - return false; // If not inside linalg, cannot be invariant w.r.t linalg - // body - Operation *defOp = v.getDefiningOp(); - if (!defOp) { - // Is a BlockArgument - auto ba = cast(v); - // It's invariant if its owner is the linalg.generic's body - return ba.getOwner() == linalgOp.getBody(); - } - // If defined by an op, the op must be outside the linalgOp's region - return !linalgOp.getRegion().isProperAncestor(defOp->getParentRegion()); - }; - - if (linalgOp && isInvariant(memref) && isInvariant(falseIndex)) { - // **Hoisting Path**: Insertion point is before linalg.generic - rewriter.setInsertionPoint(linalgOp); - - } else { - // **Internal Path**: Insertion point is before valueSel - rewriter.setInsertionPoint(valueSel); - } - - // 7. Create load index (ensure it is 'index' type) - Type indexType = rewriter.getIndexType(); - Value loadIndex; - if (falseIndex.getType() == indexType) { - loadIndex = falseIndex; - - } else if (isa(falseIndex.getType())) { - loadIndex = - rewriter.create(loc, indexType, falseIndex); - - } else { - return valueSel.emitError("False index type is not integer or index."); - } - - // 8. Create memref.load (position determined by setInsertionPoint) - // %new_val_f = memref.load %memref[%idx_f_cast] - Value newLoad = - rewriter.create(loc, memref, ValueRange{loadIndex}); - rewriter.setInsertionPoint(valueSel); - // 9. Replace the old select - // replaceOpWithNewOp automatically inserts the new op at the *original* - // op's (valueSel) location even if the rewriter's "default" insertion point - // is elsewhere - rewriter.replaceOpWithNewOp( - valueSel, valueSel.getCondition(), valueSel.getTrueValue(), newLoad); - - return success(); - } -}; - -struct LinalgIfToSelectPass - : mlir::dicp::LinalgExt::impl::LinalgIfToSelectBase { - void runOnOperation() override { - auto moduleOp = getOperation(); - MLIRContext *context = &getContext(); - UnitAttr unitAttr = UnitAttr::get(context); - moduleOp.walk([&](linalg::GenericOp linalgOp) { - // Skip if already tagged - if (linalgOp->hasAttr("ExtractedLoadOrStore")) - return; - - // Condition 1: Must be all "parallel" iterators - if (!llvm::all_of(linalgOp.getIteratorTypesArray(), - linalg::isParallelIterator)) { - return; - } - - // Condition 2: All indexing maps must be identity - bool allIdentity = true; - for (AffineMap map : linalgOp.getIndexingMapsArray()) { - if (!map.isIdentity()) { - return; - } - } - - // Condition 3 & 4: Must contain scf.if and (load/store) inside the body - bool hasScfIf = false; - bool hasLoadStore = false; - - linalgOp.getBody()->walk([&](Operation *innerOp) { - // Check for scf.if - if (isa(innerOp)) { - hasScfIf = true; - } - - // Check for various load/store operations - if (isa(innerOp)) { - hasLoadStore = true; - } - }); - - // Tag the op if all conditions are met - if (hasScfIf && hasLoadStore) { - linalgOp->setAttr("ExtractedLoadOrStore", unitAttr); - } - }); - RewritePatternSet patterns(context); - patterns.add(context); - patterns.add(context); - linalg::populateEraseUnusedOperandsAndResultsPatterns(patterns); - // Apply Patterns - if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { - signalPassFailure(); - } - - { - RewritePatternSet patterns(context); - populateLinalgLiftSelectPattern(patterns); - // Apply Patterns - if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { - signalPassFailure(); - } - - PassManager pm(&getContext(), moduleOp.getOperationName()); - // Erase dead code and fold constants created during lowering - pm.addPass(createCSEPass()); - pm.addPass(createCanonicalizerPass()); - pm.addPass(createSymbolDCEPass()); - if (failed(runPipeline(pm, getOperation()))) { - signalPassFailure(); - } - } - - } // namespace -}; - -} // namespace - -void mlir::dicp::LinalgExt::populateLinalgLiftSelectPattern( - RewritePatternSet &patterns) { - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - linalg::populateEraseUnusedOperandsAndResultsPatterns(patterns); -} - -std::unique_ptr> -mlir::dicp::LinalgExt::createLinalgIfToSelectPass() { - return std::make_unique(); -} \ No newline at end of file diff --git a/compiler/lib/Dialect/LinalgExt/Transforms/RemoveSingleIterationLoop.cpp b/compiler/lib/Dialect/LinalgExt/Transforms/RemoveSingleIterationLoop.cpp deleted file mode 100644 index 95bec078..00000000 --- a/compiler/lib/Dialect/LinalgExt/Transforms/RemoveSingleIterationLoop.cpp +++ /dev/null @@ -1,141 +0,0 @@ -#include "dicp/Dialect/LinalgExt/Transforms/Transforms.h" - -#include "mlir/Dialect/Affine/Utils.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Interfaces/ValueBoundsOpInterface.h" -#define DEBUG_TYPE "remove-single-iteration" - -using namespace mlir; -using namespace dicp; -using namespace LinalgExt; - -namespace mlir::dicp::LinalgExt { - -/** - * Checks if the scf::ForOp loop body will never execute a second iteration. - * - * This holds true if: - * 1. The upper bound (ub) is less than or equal to the step size. (ub <= step) - * 2. The lower bound (lb) is non-negative. (lb >= 0) - * - * If lb >= 0 and ub <= step, then lb + step >= ub, guaranteeing termination - * after the first iteration, assuming the loop runs at least once. - * - * @param op The scf::ForOp to analyze. - * @return True if the loop runs at most one time. - */ -static bool neverRunsSecondIteration(scf::ForOp op) { - // Can't perform the analysis if the loops's bounds aren't index-typed. - if (!op.getInductionVar().getType().isIndex()) - return false; - // If the upper bound (ub) is less than or equal to the loop step, then - // lower bound + step must be greater than the upper bound, assuming the - // lower bound is non-negative. - FailureOr isUbUnderStep = ValueBoundsConstraintSet::compare( - getAsOpFoldResult(op.getUpperBound()), ValueBoundsConstraintSet::LE, - getAsOpFoldResult(op.getStep())); - FailureOr isLbNonNegative = ValueBoundsConstraintSet::compare( - getAsOpFoldResult(op.getLowerBound()), ValueBoundsConstraintSet::GE, - getAsIndexOpFoldResult(op.getContext(), 0)); - return isUbUnderStep.value_or(false) && isLbNonNegative.value_or(false); -} - -// --- - -/** - * Checks if the scf::ForOp loop body is guaranteed to execute at least once. - * - * This holds true if the lower bound (lb) is strictly less than the - * upper bound (ub). (lb < ub) - * - * @param op The scf::ForOp to analyze. - * @return True if the loop is guaranteed to run the first iteration. - */ -static bool alwaysRunsFirstIteration(scf::ForOp op) { - // Can't perform the analysis if the loops's bounds aren't index-typed. - if (!op.getInductionVar().getType().isIndex()) - return false; - FailureOr isLb = ValueBoundsConstraintSet::compare( - getAsOpFoldResult(op.getLowerBound()), ValueBoundsConstraintSet::LT, - getAsOpFoldResult(op.getUpperBound())); - return isLb.value_or(false); -} - -/// Replaces the given op with the contents of the given single-block region, -/// using the operands of the block terminator to replace operation results. -static void replaceOpWithRegion(PatternRewriter &rewriter, scf::ForOp op, - ValueRange blockArgs = {}) { - Block *block = op.getBody(); - Operation *terminator = block->getTerminator(); - ValueRange results = terminator->getOperands(); - rewriter.inlineBlockBefore(block, op, blockArgs); - rewriter.replaceOp(op, results); - rewriter.eraseOp(terminator); -} - -/// Same as `replaceOpWithRegion` function but within an scf.if region. -static void replaceForWithIf(PatternRewriter &rewriter, scf::ForOp op, - ValueRange blockArgs = {}) { - Block *block = op.getBody(); - ValueRange initArgs = op.getInitArgs(); - Value count = - rewriter.create(op->getLoc(), arith::CmpIPredicate::sgt, - op.getUpperBound(), op.getLowerBound()); - auto ifOp = - rewriter.create(op->getLoc(), op.getResultTypes(), count, - /*withElseRegion=*/initArgs.size() != 0); - Operation *terminator = block->getTerminator(); - rewriter.inlineBlockBefore(block, &ifOp.getThenRegion().front(), - ifOp.getThenRegion().front().begin(), blockArgs); - if (initArgs.size() == 0) { - rewriter.eraseOp(terminator); - } else { - rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front()); - rewriter.create(ifOp.getLoc(), initArgs); - } - rewriter.replaceOp(op, ifOp); -} - -namespace { -/// Rewriting pattern that replaces single-iteration loops with their bodies. -struct SimplifyTrivialLoops : public OpRewritePattern { - - SimplifyTrivialLoops(MLIRContext *context, ForControlFnRef controlFn) - : OpRewritePattern(context), controlFn(controlFn) {} - - LogicalResult matchAndRewrite(scf::ForOp op, - PatternRewriter &rewriter) const override { - if (controlFn && !controlFn(op)) { - return rewriter.notifyMatchFailure( - op, "doesn't match according to the the control function"); - } - if (!neverRunsSecondIteration(op)) { - return rewriter.notifyMatchFailure(op, - "is not a single-iteration for loop"); - } - // The second iteration is never run so the loop atmost can have 1 - // iteration. Inline its body and remove the loop. - SmallVector blockArgs; - blockArgs.reserve(op.getInitArgs().size() + 1); - blockArgs.push_back(op.getLowerBound()); - llvm::append_range(blockArgs, op.getInitArgs()); - if (alwaysRunsFirstIteration(op)) { - replaceOpWithRegion(rewriter, op, blockArgs); - } else { - replaceForWithIf(rewriter, op, blockArgs); - } - return success(); - } - -private: - ForControlFnRef controlFn; -}; - -} // namespace - -void populateRemoveSingleIterationLoopPattern(RewritePatternSet &patterns, - ForControlFnRef controlFn) { - patterns.add(patterns.getContext(), controlFn); -} - -} // namespace mlir::dicp::LinalgExt diff --git a/compiler/lib/Dialect/LinalgExt/Transforms/ScalarTo1DTensorPass.cpp b/compiler/lib/Dialect/LinalgExt/Transforms/ScalarTo1DTensorPass.cpp deleted file mode 100644 index 1efbbbec..00000000 --- a/compiler/lib/Dialect/LinalgExt/Transforms/ScalarTo1DTensorPass.cpp +++ /dev/null @@ -1,203 +0,0 @@ -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/Dialect/SCF/IR/SCF.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" -#include "mlir/Transforms/Passes.h" - -#include "llvm/ADT/SmallVector.h" -#include "llvm/Support/Debug.h" - -#define DEBUG_TYPE "scalar-to-1d-tensor" - -using namespace mlir; -using namespace mlir::bufferization; -using namespace mlir::tensor; -using namespace mlir::memref; -using namespace mlir::arith; -using namespace mlir::func; - -namespace mlir { -namespace dicp { -namespace LinalgExt { -#define GEN_PASS_DEF_SCALARTO1DTENSOR -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h.inc" -} // namespace LinalgExt -} // namespace dicp -} // namespace mlir - -namespace { - -/// Pattern: Rewrite `memref.store` of a scalar into a sequence using 1D tensors -/// and bufferization. -/// -/// The scalar store is replaced with: -///  - tensor.empty<1 x elemTy> (at function entry) -///  - tensor.insert scalar into that empty tensor -///  - memref.reinterpret_cast (to create a 1D view at the store index) -///  - bufferization.materialize_in_destination (to write the tensor view) -/// -/// Match conditions: -///  - The store uses a single index (1D logical access). -///  - The target memref is either: -///     (A) a function entry block argument. -///     (B) a memref.cast or memref.reinterpret_cast whose source is (A). -struct MemrefStoreToMaterializePattern - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(memref::StoreOp storeOp, - PatternRewriter &rewriter) const override { - Location loc = storeOp.getLoc(); - - // Find the enclosing func op. - auto func = storeOp->getParentOfType(); - if (!func) { - return rewriter.notifyMatchFailure(storeOp, - "store not inside func, skipping"); - } - - // Only handle 1D index stores. - if (storeOp.getIndices().size() != 1) { - return rewriter.notifyMatchFailure(storeOp, "store not 1D, skipping"); - } - Value storeIndex = storeOp.getIndices().front(); - Value storedValue = storeOp.getValue(); - Value memrefVal = storeOp.getMemRef(); - - // Determine qualification (Case A or B) - bool qualifies = false; - Value castSrcOrArg = nullptr; - Block *entry = &func.getBody().front(); - - // Case A: direct block argument from function entry - if (auto ba = dyn_cast_or_null(memrefVal)) { - if (ba.getOwner() == entry) { - qualifies = true; - castSrcOrArg = memrefVal; - } - } - - // Case B: memref.cast or memref.reinterpret_cast whose source is func entry - // block argument. - if (!qualifies) { - if (auto castOp = memrefVal.getDefiningOp()) { - Value src = castOp.getSource(); - if (auto srBa = dyn_cast_or_null(src)) { - if (srBa.getOwner() == entry) { - qualifies = true; - castSrcOrArg = castOp.getSource(); // underlying source arg - } - } - } - } - - if (!qualifies) { - if (auto castOp = memrefVal.getDefiningOp()) { - Value src = castOp.getSource(); - if (auto srBa = dyn_cast_or_null(src)) { - if (srBa.getOwner() == entry) { - qualifies = true; - castSrcOrArg = castOp.getSource(); // underlying source arg - } - } - } - } - - if (!qualifies) { - return rewriter.notifyMatchFailure( - storeOp, "memref.store's memref is neither func arg " - "nor a cast from one; skipping"); - } - - // Determine element type and construct tensor<1 x elemTy> - Type elemTy = storedValue.getType(); - RankedTensorType oneTensorTy = RankedTensorType::get({1}, elemTy); - - // Create tensor.empty at function entry insertion point. - // We always create one, trading duplication for simplicity. - OpBuilder::InsertionGuard guard(rewriter); - rewriter.setInsertionPointToStart(&func.getBody().front()); - Value emptyTensor = - rewriter.create(loc, oneTensorTy, ValueRange{}); - - // Reset insertion point to before the storeOp - rewriter.setInsertionPoint(storeOp); - Value c0 = rewriter.create(loc, 0); - - // Insert the scalar into the 1-element tensor - Value inserted = rewriter.create( - loc, storedValue, emptyTensor, ValueRange{c0}); - - // Prepare result memref type: memref<1 x elemTy> - // MemRefType resultMemRefTy = MemRefType::get({1}, elemTy); // Unused - - // Use the non-casted source if available. - Value reinterpretSource = memrefVal; - if (auto castOp = memrefVal.getDefiningOp()) - reinterpretSource = castOp.getSource(); - - // Convert sizes/strides to OpFoldResult (using Attribute for static 1) - SmallVector sizesOf; - sizesOf.push_back(rewriter.getIndexAttr(1)); // Static size 1 - - SmallVector stridesOf; - stridesOf.push_back(rewriter.getIndexAttr(1)); // Static stride 1 - - // Use the dynamic store index (storeIndex) as the offset OpFoldResult - OpFoldResult offsetOf = storeIndex; - - // Call the ReinterpretCastOp creation method that accepts OpFoldResults. - // This overload correctly sets static_sizes/static_strides attributes. - Value reinterpretCast = rewriter.create( - loc, /*source*/ reinterpretSource, /*offset*/ offsetOf, - /*sizes*/ sizesOf, /*strides*/ stridesOf, - /*attrs*/ ArrayRef{}); - - // Materialize the inserted tensor into the 1D view of the memref. - rewriter.setInsertionPoint(storeOp); - rewriter.create( - loc, TypeRange{}, inserted, reinterpretCast, false, true); - - // Erase the original store. - rewriter.eraseOp(storeOp); - - return success(); - } -}; - -/// The pass: ScalarTo1DTensorPass -struct ScalarTo1DTensorPass - : public mlir::dicp::LinalgExt::impl::ScalarTo1DTensorBase< - ScalarTo1DTensorPass> { - ScalarTo1DTensorPass() = default; - void runOnOperation() override { - FuncOp func = getOperation(); - MLIRContext *context = &getContext(); - - // Build patterns - RewritePatternSet patterns(context); - patterns.add(context); - - // Apply patterns greedily on this function - if (failed(applyPatternsGreedily(func, std::move(patterns)))) { - signalPassFailure(); - } - } -}; - -} // namespace - -std::unique_ptr> -mlir::dicp::LinalgExt::createScalarTo1DTensorPass() { - return std::make_unique(); -} \ No newline at end of file diff --git a/compiler/lib/Dialect/LinalgExt/Transforms/TensorTransform.cpp b/compiler/lib/Dialect/LinalgExt/Transforms/TensorTransform.cpp deleted file mode 100644 index 5844c32e..00000000 --- a/compiler/lib/Dialect/LinalgExt/Transforms/TensorTransform.cpp +++ /dev/null @@ -1,340 +0,0 @@ -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h" -#include "dicp/Dialect/LinalgExt/Transforms/Transforms.h" - -#include "mlir/Dialect/Affine/Utils.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/Tensor/IR/Tensor.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Interfaces/ValueBoundsOpInterface.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" - -#define DEBUG_TYPE "tensor-transform" - -using namespace mlir; -using namespace dicp; -using namespace LinalgExt; - -namespace mlir { -namespace dicp { -namespace LinalgExt { -#define GEN_PASS_DEF_NORMALIZESLICEOPS -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h.inc" -} // namespace LinalgExt -} // namespace dicp -} // namespace mlir - -namespace { - -/// Compute the dropped dimensions of a rank-reducing tensor.extract_slice op or -/// rank-extending tensor.insert_slice op. -static llvm::SmallBitVector -getDroppedDimsForInterleave(ArrayRef reducedShape, - ArrayRef mixedSizes) { - // TODO this is old community function, delete it afterwards - llvm::SmallBitVector droppedDims(mixedSizes.size()); - int64_t shapePos = 0; - - for (const auto &size : enumerate(mixedSizes)) { - // Rank-reduced dims must have a static unit dimension. - bool isStaticUnitSize = - isa(size.value()) && - llvm::cast(cast(size.value())).getInt() == 1; - - if (shapePos == static_cast(reducedShape.size())) { - // There are no more dims in the reduced shape. All remaining sizes must - // be rank-reduced dims. - assert(isStaticUnitSize && "expected unit dim"); - droppedDims.set(size.index()); - continue; - } - - // Dim is preserved if the size is not a static 1. - if (!isStaticUnitSize) { - ++shapePos; - continue; - } - - // Dim is preserved if the reduced shape dim is also 1. - if (reducedShape[shapePos] == 1) { - ++shapePos; - continue; - } - - // Otherwise: Dim is dropped. - droppedDims.set(size.index()); - } - - return droppedDims; -} - -/// This pattern detects a chain of `tensor::InsertSliceOp` that together -/// implement an interleave write: multiple source tensors are inserted into the -/// same destination along the last dimension using different static offsets. -/// Once detected, the pattern normalizes this chain into a canonical form -/// where the last-dimension offsets are [0..channelNum-1] and the stride is -/// `channelNum`, making the interleave structure explicit and ready for further -/// fusion or replacement by a dedicated Interleave op. -struct NormalizeInsertSliceOpToInterleaveOp - : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - std::optional getInterLeaveChannelIdx(Operation *op) const { - if (auto insertSliceOp = llvm::dyn_cast(op)) { - if (insertSliceOp.getStaticOffsets().empty()) - return std::nullopt; - // dynamic offset return INT64_MIN - if (insertSliceOp.getStaticOffsets().back() == INT64_MIN) - return std::nullopt; - return insertSliceOp.getStaticOffsets().back(); - } - return std::nullopt; - } - - bool isInterLeavePartialPattern(tensor::InsertSliceOp insertSliceOp, - Value curSource, - int64_t interLeaveChannelNums) const { - SmallVector srcShape( - dyn_cast(curSource.getType()).getShape()); - SmallVector dstShape( - dyn_cast(insertSliceOp.getDest().getType()).getShape()); - - if (dstShape.size() != srcShape.size()) - return false; - - // Correspond to rule 1&2 - if (!std::equal(srcShape.begin(), srcShape.end() - 1, dstShape.begin()) || - ShapedType::isDynamic(dstShape.back()) || - ShapedType::isDynamic(srcShape.back()) || - dstShape.back() != srcShape.back() * interLeaveChannelNums) { - return false; - } - - // Collect layout offset - std::optional interLeaveChannelIdx = - getInterLeaveChannelIdx(insertSliceOp); - if (!interLeaveChannelIdx.has_value()) { - return false; - } - if (interLeaveChannelIdx > interLeaveChannelNums) { - return false; - } - - // Correspond to rule 3 - if (insertSliceOp.getStaticStrides().back() != 1) { - return false; - } - return true; - } - - // Actually, tensor::InsertSliceOp could express any dimension insertion - // flexibly with layout info. - // - // While current hfusion::InterleaveOp just wanna match pattern where - // 1. non-last diemsnion shape of input and output should be equal - // 2. tail shape scale between dst and src equals channelNum, which also - // means - // last dimension of all types could only be static - // 3. In layout, last dimension stride equals channelNum - // 4. In layout, last dimension offsets of all candidate InsertSliceOp - // should - // make a range of [0, channelNum) - // - // And for rank-extended state of tensor::InsertSliceOp, where source rank - // may less than destination, InterleaveOp just support extended rank is - // last dimension and size equals channelNum, then explicitly expand shape - // on src before create InterleaveOp - LogicalResult traceInterLeavePattern(tensor::InsertSliceOp insertSliceOp, - llvm::BitVector &findChannels, - SmallVector &inputs, - int64_t interLeaveChannelNums, - PatternRewriter &rewriter) const { - Value curSrc = insertSliceOp.getSource(); - - llvm::SmallBitVector extendedRankRecord = - getDroppedDimsForInterleave(insertSliceOp.getSourceType().getShape(), - insertSliceOp.getMixedSizes()); - if (extendedRankRecord.any()) { - if (extendedRankRecord.find_first() != - static_cast(extendedRankRecord.size()) - 1) - return rewriter.notifyMatchFailure( - insertSliceOp, "extended rank could only be last dimension"); - - // Extented rank size of destination must be interLeaveChannelNums - if (dyn_cast(insertSliceOp.getDest().getType()) - .getShape() - .back() != interLeaveChannelNums) - return rewriter.notifyMatchFailure( - insertSliceOp, - "size of extended rank axis should equal channel num"); - - auto originType = llvm::dyn_cast(curSrc.getType()); - SmallVector shape(originType.getShape()); - - // Here represents last dimension which is only extended rank - shape.push_back(1); - RankedTensorType newType = - RankedTensorType::get(shape, originType.getElementType()); - - std::optional> reassociation = - getReassociationIndicesForReshape(originType, newType); - assert(reassociation.has_value()); - - auto expandOp = rewriter.create( - insertSliceOp.getLoc(), newType, curSrc, reassociation.value()); - curSrc = expandOp.getResult(); - } - - if (!isInterLeavePartialPattern(insertSliceOp, curSrc, - interLeaveChannelNums)) - return rewriter.notifyMatchFailure( - insertSliceOp, - "current tensor::InsertSliceOp layout doen't satisfy condition " - "to be converted to hfusion::InterleaveOp"); - - auto channelIdxMaybe = getInterLeaveChannelIdx(insertSliceOp); - if (!channelIdxMaybe.has_value()) { - return failure(); - } - int channelIdx = channelIdxMaybe.value(); - // set channelIdx-bit to findChannels and push corresponding input - findChannels[channelIdx] = true; - inputs[channelIdx] = curSrc; - - // findChannels all true - // Correspond to rule 4 - if (findChannels == llvm::BitVector(findChannels.size(), true)) { - return success(); - } - - // trace further - auto dstDefiningOp = - insertSliceOp->getOperand(1).getDefiningOp(); - if (!dstDefiningOp) { - return rewriter.notifyMatchFailure( - insertSliceOp, - "tensor::InsertSliceOp chain from current op can't reach " - "interLeave channel num"); - } - return traceInterLeavePattern(dstDefiningOp, findChannels, inputs, - interLeaveChannelNums, rewriter); - } - - LogicalResult matchAndRewrite(tensor::InsertSliceOp insertSliceOp, - PatternRewriter &rewriter) const override { - if (!insertSliceOp.hasPureTensorSemantics()) { - return failure(); - } - - // TODO: find interLeaveChannelNums greedily. - const int64_t interLeaveChannelNums = 2; - llvm::BitVector findChannels(interLeaveChannelNums, false); - SmallVector inputs(interLeaveChannelNums); - - // 1. Trace the pattern and collect inputs - if (traceInterLeavePattern(insertSliceOp, findChannels, inputs, - interLeaveChannelNums, rewriter) - .failed()) { - return failure(); - } - - // 2. Find the root destination tensor. - // traceInterLeavePattern has confirmed the chain exists. - // We traverse back (interLeaveChannelNums - 1) times to find the - // initial buffer that the first slice was inserted into. - Value accumulatedDest = insertSliceOp.getDest(); - for (int i = 0; i < interLeaveChannelNums - 1; ++i) { - auto parentOp = accumulatedDest.getDefiningOp(); - // Use cast because traceInterLeavePattern ensured the chain exists - if (!parentOp) - break; - accumulatedDest = parentOp.getDest(); - } - - Location loc = insertSliceOp.getLoc(); - - // 3. Rebuild the InsertSliceOp chain with normalized strides. - for (int i = 0; i < interLeaveChannelNums; ++i) { - Value src = inputs[i]; - if (!src) { - // Should not happen if traceInterLeavePattern returns success - // and findChannels is fully set. - return failure(); - } - - auto srcType = llvm::cast(src.getType()); - // Not support dynamic shape - if (!srcType.hasStaticShape()) { - return failure(); - } - int64_t rank = srcType.getRank(); - - // Prepare Offsets, Sizes, Strides - SmallVector offsets, sizes, strides; - ArrayRef srcShape = srcType.getShape(); - - // Populate dimensions 0 to Rank-2 (standard dims) - // and Rank-1 (last dim/channel dim) - for (int64_t d = 0; d < rank; ++d) { - // Offset: 0 for all dims, except last dim is the channel index 'i' - if (d == rank - 1) { - offsets.push_back(rewriter.getIndexAttr(i)); - } else { - offsets.push_back(rewriter.getIndexAttr(0)); - } - - // Size: Get from src shape (Static only) - sizes.push_back(rewriter.getIndexAttr(srcShape[d])); - - // Stride: 1 for all dims, except last dim is 'interLeaveChannelNums' - if (d == rank - 1) { - strides.push_back(rewriter.getIndexAttr(interLeaveChannelNums)); - } else { - strides.push_back(rewriter.getIndexAttr(1)); - } - } - - // Create the new standardized InsertSliceOp - auto newInsertOp = rewriter.create( - loc, src, accumulatedDest, offsets, sizes, strides); - - // Update accumulatedDest for the next iteration - accumulatedDest = newInsertOp.getResult(); - } - - // 4. Replace the original op with the result of the new chain - rewriter.replaceOp(insertSliceOp, accumulatedDest); - return success(); - } -}; // struct NormalizeInsertSliceOpInInterleavePattern - -struct NormalizeSliceOpsPass - : public mlir::dicp::LinalgExt::impl::NormalizeSliceOpsBase< - NormalizeSliceOpsPass> { - NormalizeSliceOpsPass() = default; - void runOnOperation() override { - mlir::func::FuncOp func = getOperation(); - MLIRContext *context = &getContext(); - - // Build patterns - RewritePatternSet patterns(context); - populateNormalizeInsertSliceOpInInterleavePattern(patterns); - - // Apply patterns greedily on this function - if (failed(applyPatternsGreedily(func, std::move(patterns)))) { - signalPassFailure(); - } - } -}; - -} // namespace - -void mlir::dicp::LinalgExt::populateNormalizeInsertSliceOpInInterleavePattern( - RewritePatternSet &patterns) { - patterns.add(patterns.getContext()); -} - -std::unique_ptr> -mlir::dicp::LinalgExt::createNormalizeSliceOpsPass() { - return std::make_unique(); -} \ No newline at end of file diff --git a/compiler/lib/Dialect/NPU/CMakeLists.txt b/compiler/lib/Dialect/NPU/CMakeLists.txt deleted file mode 100644 index 5cd7900f..00000000 --- a/compiler/lib/Dialect/NPU/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -dicp_add_all_subdirs() \ No newline at end of file diff --git a/compiler/lib/Dialect/NPU/IR/CMakeLists.txt b/compiler/lib/Dialect/NPU/IR/CMakeLists.txt deleted file mode 100644 index 6faab514..00000000 --- a/compiler/lib/Dialect/NPU/IR/CMakeLists.txt +++ /dev/null @@ -1,12 +0,0 @@ -add_dicp_triton_library(DICPNPU - NPUOps.cpp - NPUDialect.cpp - NPUTypes.cpp - - DEPENDS - DICPNPUIncGen - - LINK_LIBS PUBLIC - TritonIR - MLIRIR - ) diff --git a/compiler/lib/Dialect/NPU/IR/NPUDialect.cpp b/compiler/lib/Dialect/NPU/IR/NPUDialect.cpp deleted file mode 100644 index 277851dd..00000000 --- a/compiler/lib/Dialect/NPU/IR/NPUDialect.cpp +++ /dev/null @@ -1,52 +0,0 @@ -#include "dicp/Dialect/NPU/IR/NPUDialect.h" -#include "dicp/Dialect/NPU/IR/NPUDialect.cpp.inc" -#include "dicp/Dialect/NPU/IR/NPUTypes.h" -#include "mlir/IR/DialectImplementation.h" // required by `Types.cpp.inc` -#include "llvm/ADT/StringSwitch.h" -#include "llvm/ADT/TypeSwitch.h" // required by `Types.cpp.inc` - -#include "llvm/Support/raw_ostream.h" - -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" - -#include "mlir/Transforms/InliningUtils.h" - -using namespace mlir; -using namespace mlir::dicp::npu; - -// Type parseType(DialectAsmParser &parser) const { -// StringRef typeKind; -// if (parser.parseKeyword(&typeKind)) -// return {}; -// auto type = llvm::StringSwitch(typeKind) -// .Case("tqueue", TQueueType::get(getContext())) -// .Default(nullptr); -// if (!type) { -// parser.emitError(parser.getCurrentLocation()) -// << "unknown NPU type: " << typeKind; -// } -// return type; -// } - -// void printType(Type type, DialectAsmPrinter &p) const { -// if (llvm::isa(type)) { -// p << "tqueue"; -// } else { -// assert(false && "unknown tqueue type"); -// } -// } - -/// Dialect creation, the instance will be owned by the context. This is the -/// point of registration of custom types and operations for the dialect. -void NPUDialect::initialize() { - registerTypes(); - - addOperations< -#define GET_OP_LIST -#include "dicp/Dialect/NPU/IR/NPUOps.cpp.inc" - >(); -} - -//===----------------------------------------------------------------------===// -// TableGen'd op method definitions -//===----------------------------------------------------------------------===// diff --git a/compiler/lib/Dialect/NPU/IR/NPUOps.cpp b/compiler/lib/Dialect/NPU/IR/NPUOps.cpp deleted file mode 100644 index 829f7301..00000000 --- a/compiler/lib/Dialect/NPU/IR/NPUOps.cpp +++ /dev/null @@ -1,135 +0,0 @@ -#include "dicp/Dialect/NPU/IR/NPUDialect.h" -#include "dicp/Dialect/NPU/IR/NPUTypes.h" -#include "mlir/Bytecode/BytecodeOpInterface.h" -#include "mlir/IR/Attributes.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/OpImplementation.h" -#include "mlir/IR/Operation.h" -#include "mlir/IR/OperationSupport.h" -#include "mlir/IR/Value.h" -#include "mlir/Interfaces/FunctionImplementation.h" -#include "mlir/Support/LLVM.h" -#include "mlir/Support/LogicalResult.h" -#include "llvm/ADT/ArrayRef.h" -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/Casting.h" -#include -#include - -// #define GET_OP_CLASSES -// #include "dicp/Dialect/NPU/IR/NPUOps.h.inc" - -// #define GET_OP_CLASSES -// #include "dicp/Dialect/NPU/IR/NPUDialect.cpp.inc" -#define GET_OP_CLASSES -#include "dicp/Dialect/NPU/IR/NPUOps.cpp.inc" - -using namespace mlir; -using namespace mlir::dicp; - -namespace mlir::dicp::npu { - -/// A generalized parser for binary operations. This parses the different forms -/// of 'printBinaryOp' below. -static mlir::ParseResult parseBinaryOp(mlir::OpAsmParser &parser, - mlir::OperationState &result) { - SmallVector operands; - SMLoc operandsLoc = parser.getCurrentLocation(); - Type type; - if (parser.parseOperandList(operands, /*requiredOperandCount=*/2) || - parser.parseOptionalAttrDict(result.attributes) || - parser.parseColonType(type)) - return mlir::failure(); - - // If the type is a function type, it contains the input and result types of - // this operation. - if (FunctionType funcType = llvm::dyn_cast(type)) { - if (parser.resolveOperands(operands, funcType.getInputs(), operandsLoc, - result.operands)) - return mlir::failure(); - result.addTypes(funcType.getResults()); - return mlir::success(); - } - - // Otherwise, the parsed type is the type of both operands and results. - if (parser.resolveOperands(operands, type, result.operands)) - return mlir::failure(); - result.addTypes(type); - return mlir::success(); -} - -/// A generalized printer for binary operations. It prints in two different -/// forms depending on if all of the types match. -static void printBinaryOp(mlir::OpAsmPrinter &printer, mlir::Operation *op) { - printer << " " << op->getOperands(); - printer.printOptionalAttrDict(op->getAttrs()); - printer << " : "; - - // If all of the types are the same, print the type directly. - Type resultType = *op->result_type_begin(); - if (llvm::all_of(op->getOperandTypes(), - [=](Type type) { return type == resultType; })) { - printer << resultType; - return; - } - - // Otherwise, print a functional type. - printer.printFunctionalType(op->getOperandTypes(), op->getResultTypes()); -} - -//===----------------------------------------------------------------------===// -// AddOp -//===----------------------------------------------------------------------===// - -// void AddFOp::build(mlir::OpBuilder &builder, mlir::OperationState &state, -// mlir::Value lhs, mlir::Value rhs) { -// state.addTypes(UnrankedTensorType::get(builder.getF16Type())); -// state.addOperands({lhs, rhs}); -// } - -// mlir::ParseResult AddFOp::parse(mlir::OpAsmParser &parser, -// mlir::OperationState &result) { -// return parseBinaryOp(parser, result); -// } - -// void AddFOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); } - -// void CreateTQueueOp::print(mlir::OpAsmPrinter &p) { -// // printBinaryOp(p, *this); -// p.printFunctionalType((*this)->getResultTypes()); -// } - -// void Copy::build(OpBuilder &b, OperationState &state, -// ArrayRef offsets, -// ArrayRef sizes, -// ArrayRef strides -// ) { -// SmallVector staticStrides, staticOffsets, staticShape; -// SmallVector dynamicStrides, dynamicOffsets, dynamicShape; - -// dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets); -// dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); -// dispatchIndexOpFoldResults(shape, dynamicShape, staticShape); - -// Type resType; -// auto basePtr = cast(base.getType()); -// auto elemType = basePtr.getPointeeType(); -// // non-block pointer -// if (order.empty()) { -// resType = RankedTensorType::get(sizes, basePtr); -// } -// // block pointer -// else { -// resType = triton::PointerType::get(RankedTensorType::get(sizes, -// elemType), -// basePtr.getAddressSpace()); -// } - -// build(b, state, resType, base, sizes, dynamicStrides, dynamicOffsets, -// dynamicShape, b.getDenseI64ArrayAttr(staticStrides), -// b.getDenseI64ArrayAttr(staticOffsets), -// b.getDenseI64ArrayAttr(staticShape), order); -// } -} // namespace mlir::dicp::npu diff --git a/compiler/lib/Dialect/NPU/IR/NPUTypes.cpp b/compiler/lib/Dialect/NPU/IR/NPUTypes.cpp deleted file mode 100644 index 25154271..00000000 --- a/compiler/lib/Dialect/NPU/IR/NPUTypes.cpp +++ /dev/null @@ -1,81 +0,0 @@ -#include "dicp/Dialect/NPU/IR/NPUTypes.h" -#include "dicp/Dialect/NPU/IR/NPUDialect.h" -#include "mlir/IR/DialectImplementation.h" // required by `Types.cpp.inc` -#include "mlir/Support/LLVM.h" -#include "llvm/ADT/TypeSwitch.h" // required by `Types.cpp.inc` - -using namespace mlir; -using namespace mlir::dicp::npu; - -#define GET_TYPEDEF_CLASSES -#include "dicp/Dialect/NPU/IR/NPUTypes.cpp.inc" - -Type TQueueType::parse(AsmParser &parser) { - if (parser.parseLess()) - return Type(); - - int position = 1; - if (parser.parseInteger(position)) { - return Type(); - } - int bufferNumber = 1; - if (parser.parseInteger(bufferNumber)) { - return Type(); - } - - if (parser.parseGreater()) { - return Type(); - } - - return TQueueType::get(parser.getContext(), position, bufferNumber); -} - -void TQueueType::print(AsmPrinter &printer) const { - printer << "npu.tqueue<" << getPosition() << ", " << getBufferNumber() << ">"; -} - -Type TPipType::parse(AsmParser &parser) { - if (parser.parseLess()) - return Type(); - - // int type = 1; - // if (parser.parseInteger(type)) { - // return Type(); - // } - - if (parser.parseGreater()) { - return Type(); - } - - return TPipType::get(parser.getContext()); -} - -void TPipType::print(AsmPrinter &printer) const { printer << "npu.tpip"; } - -Type GlobalTensorType::parse(AsmParser &parser) { - if (parser.parseLess()) - return Type(); - - int type = 1; - if (parser.parseInteger(type)) { - return Type(); - } - - if (parser.parseGreater()) { - return Type(); - } - - return GlobalTensorType::get(parser.getContext(), type); -} - -void GlobalTensorType::print(AsmPrinter &printer) const { - printer << "npu.global_tensor<" << getType() << ">"; -} - -//===----------------------------------------------------------------------===// -void NPUDialect::registerTypes() { - addTypes< -#define GET_TYPEDEF_LIST -#include "dicp/Dialect/NPU/IR/NPUTypes.cpp.inc" - >(); -} diff --git a/compiler/lib/Dialect/TritonDicp/CMakeLists.txt b/compiler/lib/Dialect/TritonDicp/CMakeLists.txt new file mode 100644 index 00000000..f33061b2 --- /dev/null +++ b/compiler/lib/Dialect/TritonDicp/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/compiler/lib/Dialect/TritonDicp/IR/CMakeLists.txt b/compiler/lib/Dialect/TritonDicp/IR/CMakeLists.txt new file mode 100644 index 00000000..9317cbdf --- /dev/null +++ b/compiler/lib/Dialect/TritonDicp/IR/CMakeLists.txt @@ -0,0 +1,14 @@ +add_triton_library(TritonDicpIR + TritonDicpAttrs.cpp + TritonDicpDialect.cpp + TritonDicpOps.cpp + + DEPENDS + TritonDicpTableGen + TritonDicpAttrDefsIncGen + + LINK_LIBS PUBLIC + MLIRIR + MLIRLLVMDialect + TritonIR +) diff --git a/compiler/lib/Dialect/TritonDicp/IR/TritonDicpAttrs.cpp b/compiler/lib/Dialect/TritonDicp/IR/TritonDicpAttrs.cpp new file mode 100644 index 00000000..e7557ed5 --- /dev/null +++ b/compiler/lib/Dialect/TritonDicp/IR/TritonDicpAttrs.cpp @@ -0,0 +1,15 @@ +//===- TritonDicpAttrs.cpp - TritonDicp Attributes Definition +//--------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" + +#include "llvm/ADT/SmallSet.h" +#include "llvm/ADT/TypeSwitch.h" + +namespace mlir::triton::dicp {} // namespace mlir::triton::dicp diff --git a/compiler/lib/Dialect/TritonDicp/IR/TritonDicpDialect.cpp b/compiler/lib/Dialect/TritonDicp/IR/TritonDicpDialect.cpp new file mode 100644 index 00000000..e1b7e039 --- /dev/null +++ b/compiler/lib/Dialect/TritonDicp/IR/TritonDicpDialect.cpp @@ -0,0 +1,44 @@ +//===- TritonDicpDialect.cpp - TritonDicp Dialect registration +//--------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" + +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/DialectImplementation.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Operation.h" + +#include "llvm/ADT/StringExtras.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/AsmParser/Parser.h" +#include "llvm/IR/Function.h" +#include "llvm/Support/SourceMgr.h" + +using namespace mlir; +using namespace mlir::triton::dicp; + +void TritonDicpDialect::initialize() { + addOperations< +#define GET_OP_LIST +#include "dicp/Dialect/TritonDicp/IR/TritonDicpOps.cpp.inc" + >(); + addAttributes< +#define GET_ATTRDEF_LIST +#include "dicp/Dialect/TritonDicp/IR/TritonDicpOpsAttrDefs.cpp.inc" + >(); +} + +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.cpp.inc" +#define GET_ATTRDEF_CLASSES +#include "dicp/Dialect/TritonDicp/IR/TritonDicpOpsAttrDefs.cpp.inc" +#define GET_OP_CLASSES +#include "dicp/Dialect/TritonDicp/IR/TritonDicpOps.cpp.inc" diff --git a/compiler/lib/Dialect/TritonDicp/IR/TritonDicpOps.cpp b/compiler/lib/Dialect/TritonDicp/IR/TritonDicpOps.cpp new file mode 100644 index 00000000..8b62a433 --- /dev/null +++ b/compiler/lib/Dialect/TritonDicp/IR/TritonDicpOps.cpp @@ -0,0 +1,230 @@ +//===- TritonDicpOps.cpp - TritonDicp dialect operations +//--------------------===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" +#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "triton/Tools/Sys/GetEnv.hpp" +#include "llvm/ADT/STLExtras.h" +#include + +using namespace mlir; +using namespace mlir::triton; + +namespace mlir::triton::dicp { + +void GatherOutToUbOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSrcMutable(), + triton::GlobalMemory::get()); +} + +void IndirectLoadOp::getEffects( + SmallVectorImpl> + &effects) { + effects.emplace_back(MemoryEffects::Read::get(), &getSrcMutable(), + triton::GlobalMemory::get()); +} + +//-- IndexSelectSimdOp -- +LogicalResult IndexSelectSimdOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + + // Get operands using adaptor + IndexSelectSimdOpAdaptor adaptor(operands, attributes, properties, regions); + + // Get element type from src pointer + Type elemType; + if (auto ptrType = + dyn_cast(adaptor.getSrc().getType())) { + elemType = ptrType.getPointeeType(); + } else { + return failure(); + } + + // Get index shape to determine the size of dim + auto indicesType = dyn_cast(adaptor.getIndex().getType()); + if (!indicesType) + return failure(); + int64_t numIndices = indicesType.getShape()[0]; + + // Use adaptor to get attributes - this is the compatible way + int32_t dim = adaptor.getDim(); + auto readShapeAttr = adaptor.getReadShape(); + + // Build result shape: read_shape but with dim replaced by numIndices + SmallVector resultShape; + for (size_t i = 0; i < readShapeAttr.size(); ++i) { + if (i == static_cast(dim)) { + resultShape.push_back(numIndices); + } else { + resultShape.push_back(readShapeAttr[i]); + } + } + + // Create result tensor type + inferredReturnTypes.push_back(RankedTensorType::get(resultShape, elemType)); + + return success(); +} + +// FlipOp +LogicalResult +FlipOp::inferReturnTypes(MLIRContext *context, std::optional location, + ValueRange operands, DictionaryAttr attributes, + OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + auto inputTy = dyn_cast(operands[0].getType()); + if (!inputTy) { + if (location) + return emitOptionalError(location, + "expected ranked tensor for flip input"); + return failure(); + } + inferredReturnTypes.push_back(inputTy); + return success(); +} + +//-- SortOp -- +LogicalResult +SortOp::inferReturnTypes(MLIRContext *context, std::optional location, + ValueRange operands, DictionaryAttr attributes, + OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + if (operands.size() != 1) { + return emitOptionalError(location, + "expected exactly one operand for SortOp"); + } + + if (!isa(operands[0].getType())) { + return emitOptionalError(location, + "operand must be a ranked tensor type for SortOp"); + } + + Value src = operands[0]; + auto srcTy = cast(src.getType()); + auto srcShape = srcTy.getShape(); + auto srcEnc = srcTy.getEncoding(); + + if (srcShape.empty()) { + return emitOptionalError(location, "input tensor must have rank >= 1"); + } + + Type sortedTy = + RankedTensorType::get(srcShape, srcTy.getElementType(), srcEnc); + + inferredReturnTypes.push_back(sortedTy); + + return success(); +} + +//-- Conv1dOp -- +LogicalResult Conv1dOp::verify() { + auto inputType = dyn_cast(getInput().getType()); + auto weightType = dyn_cast(getWeight().getType()); + + constexpr int64_t dim2 = 2; + constexpr int64_t dim3 = 3; + auto inputRank = inputType.getShape().size(); + if (inputRank != dim2 && inputRank != dim3) { + return emitOpError("input tensor must be 2D or 3D, but got rank ") + << inputRank; + } + if (weightType.getShape().size() != dim3) { + return emitOpError("weight tensor must be 3D, but got rank ") + << weightType.getShape().size(); + } + + Value biasValue = getBias(); + if (biasValue) { + auto biasType = dyn_cast(biasValue.getType()); + if (!biasType || biasType.getRank() != 1) { + return emitOpError("bias must be a 1D ranked tensor"); + } + if (biasType.getElementType() != inputType.getElementType()) { + return emitOpError( + "bias must have the same element type as input and weight"); + } + if (biasType.getShape()[0] != weightType.getShape()[0]) { + return emitOpError( + "bias size must match weight's output channel dimension"); + } + } + + int64_t C_in = inputType.getShape()[inputRank - 2]; + int64_t C_out = weightType.getShape()[0]; + int64_t weight_C_in = weightType.getShape()[1]; + + if (C_in % getGroups() != 0) { + return emitOpError("input channels must be divisible by groups"); + } + if (C_out % getGroups() != 0) { + return emitOpError("output channels must be divisible by groups"); + } + if (weight_C_in != C_in / getGroups()) { + return emitOpError( + "weight's input channel dimension must equal input channels / groups"); + } + + return success(); +} + +LogicalResult Conv1dOp::inferReturnTypes( + MLIRContext *context, std::optional location, ValueRange operands, + DictionaryAttr attributes, OpaqueProperties properties, RegionRange regions, + SmallVectorImpl &inferredReturnTypes) { + Conv1dOpAdaptor adaptor(operands, attributes, properties, regions); + + auto inputType = dyn_cast(adaptor.getInput().getType()); + auto weightType = dyn_cast(adaptor.getWeight().getType()); + if (!inputType || !weightType) { + return failure(); + } + + ArrayRef inputShape = inputType.getShape(); + bool isBatched = inputShape.size() == 3; + + int64_t C_in = inputShape[isBatched ? 1 : 0]; + int64_t L_in = inputShape[isBatched ? 2 : 1]; + + ArrayRef weightShape = weightType.getShape(); + int64_t C_out = weightShape[0]; + int64_t K = weightShape[2]; + + int64_t stride = adaptor.getStride(); + int64_t padding_size = adaptor.getPaddingSize(); + int64_t dilation = adaptor.getDilation(); + + if (stride == 0) { + return failure(); + } + double l_out_double = + static_cast(L_in + 2 * padding_size - dilation * (K - 1) - 1) / + stride + + 1; + int64_t L_out = static_cast(std::floor(l_out_double)); + + constexpr int64_t dim3 = 3; + SmallVector outputShape; + if (isBatched) { + outputShape.push_back(inputShape[0]); + } + outputShape.push_back(weightShape[0]); + outputShape.push_back(L_out); + + Type elementType = inputType.getElementType(); + auto returnType = RankedTensorType::get(outputShape, elementType); + inferredReturnTypes.push_back(returnType); + return success(); +} + +} // namespace mlir::triton::dicp diff --git a/compiler/lib/Dialect/TritonExt/CMakeLists.txt b/compiler/lib/Dialect/TritonExt/CMakeLists.txt deleted file mode 100644 index 5cd7900f..00000000 --- a/compiler/lib/Dialect/TritonExt/CMakeLists.txt +++ /dev/null @@ -1 +0,0 @@ -dicp_add_all_subdirs() \ No newline at end of file diff --git a/compiler/lib/Dialect/TritonExt/Transforms/BoolTritonPtrPromotionPass.cpp b/compiler/lib/Dialect/TritonExt/Transforms/BoolTritonPtrPromotionPass.cpp deleted file mode 100644 index f115dced..00000000 --- a/compiler/lib/Dialect/TritonExt/Transforms/BoolTritonPtrPromotionPass.cpp +++ /dev/null @@ -1,297 +0,0 @@ -#include "dicp/Dialect/TritonExt/Transforms/Passes.h" - -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/Triton/IR/Types.h" - -#include "mlir/IR/Block.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/MLIRContext.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/Operation.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/IR/Types.h" -#include "mlir/IR/Value.h" -#include "mlir/Interfaces/FunctionInterfaces.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Support/LogicalResult.h" -#include "mlir/Transforms/DialectConversion.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" - -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallVectorExtras.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/ADT/TypeSwitch.h" -#include "llvm/Support/Casting.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/ErrorHandling.h" -#include "llvm/Support/LogicalResult.h" -#include "llvm/Support/raw_ostream.h" - -#define DEBUG_TYPE "bool-triton-ptr-promotion" - -using namespace mlir; -using namespace dicp; -using namespace trtion_ext; -using namespace triton; - -namespace mlir::dicp::trtion_ext { -#define GEN_PASS_DEF_BOOLTRITONPTRPROMOTION -#include "dicp/Dialect/TritonExt/Transforms/Passes.h.inc" -} // namespace mlir::dicp::trtion_ext - -namespace { - -/* - * Move tt.bitcast to a previous location if tt.bitcast is not directly applied - * on function arguments - */ -class BitcastCanonicalizer : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - /* - * Move tt.bitcast to a previous location if tt.bitcast is not directly - * applied on function arguments - */ - LogicalResult matchAndRewrite(triton::BitcastOp bitcastOp, - PatternRewriter &rewriter) const { - Value castSrc = bitcastOp.getSrc(); - Value castRes = bitcastOp.getResult(); - Type castSrcTy = castSrc.getType(); - Type castSrcPtrTy = isa(castSrcTy) - ? cast(castSrcTy).getElementType() - : castSrcTy; - if (!isa(castSrcPtrTy)) - return failure(); - - auto origBitwidth = getPointeeBitWidth(castSrc.getType()); - auto castBitwidth = getPointeeBitWidth(castRes.getType()); - - if (origBitwidth == 1) - origBitwidth = 8; - if (castBitwidth == 1) - castBitwidth = 8; - if (origBitwidth != castBitwidth) { - bitcastOp.emitError() << "Casting pointers with unmatched bitwidth!\n"; - return failure(); - } - - Operation *beforeCastOp = castSrc.getDefiningOp(); - if (beforeCastOp == nullptr) { - return failure(); - } - - auto newRes = - TypeSwitch>(beforeCastOp) - // before: addptr - bitcast - load/store - // after: bitcast - addptr - load/store - .Case([&](triton::AddPtrOp addptrOp) { - auto newCastOp = rewriter.create( - bitcastOp.getLoc(), castRes.getType(), addptrOp.getPtr()); - return rewriter.create( - bitcastOp.getLoc(), castRes.getType(), newCastOp.getResult(), - addptrOp.getOffset()); - }) - .Case([&](triton::SplatOp splatOp) { - Type newCastSrcTy = - cast(castRes.getType()).getElementType(); - - Value splatSrc = splatOp.getSrc(); - Type splatSrcTy = splatSrc.getType(); - if (auto splatSrcTensorTy = - dyn_cast(splatSrcTy)) - newCastSrcTy = - splatSrcTensorTy.cloneWith(std::nullopt, newCastSrcTy); - auto newCastOp = rewriter.create( - bitcastOp.getLoc(), newCastSrcTy, splatSrc); - return rewriter.create( - bitcastOp.getLoc(), castRes.getType(), newCastOp); - }) - // before: bitcast - bitcast - // after(fusion optimization): bitcast - .Case([&](triton::BitcastOp prevCastOp) { - return rewriter.create( - bitcastOp.getLoc(), castRes.getType(), prevCastOp.getSrc()); - }) - .Default([&](Operation *op) { - return rewriter.notifyMatchFailure(bitcastOp, - "Unknown bitcast pattern"); - }); - if (succeeded(newRes)) { - rewriter.replaceOp(bitcastOp, newRes.value()); - if (beforeCastOp->use_empty()) { - rewriter.eraseOp(beforeCastOp); - } - LLVM_DEBUG({ - auto &os = llvm::dbgs(); - os << "BitcastCanonicalizer s has users:\n"; - }); - return success(); - } - LLVM_DEBUG({ - auto &os = llvm::dbgs(); - os << "BitcastCanonicalizer f has users:\n"; - }); - return failure(); - } -}; - -class BitcastConverter : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(triton::BitcastOp op, - PatternRewriter &rewriter) const { - if (op->hasAttr("Input_Arg_i1_Bitcast_To_i8")) { - return failure(); - } - - Value result; - if (auto resPointerType = dyn_cast(op.getType())) { - // TODO: use typeconverter - auto srcPointerType = cast(op.getSrc().getType()); - auto resType = MemRefType::get({ShapedType::kDynamic}, - resPointerType.getPointeeType()); - // Handling special case - // %0 = tt.bitcast %arg0 {MixUse} : !tt.ptr -> !tt.ptr - if (isa(op.getSrc()) && - srcPointerType.getPointeeType() == rewriter.getIntegerType(1) && - resPointerType.getPointeeType() == rewriter.getIntegerType(8)) { - rewriter.modifyOpInPlace(op, [&]() { - op->setAttr("Input_Arg_i1_Bitcast_To_i8", rewriter.getUnitAttr()); - }); - return success(); - } - result = - rewriter.create(op.getLoc(), resType, op.getSrc()); - } else { - result = rewriter.create(op.getLoc(), op.getType(), - op.getSrc()); - } - rewriter.replaceOp(op, result); - return success(); - } -}; - -static void converti1ArgInTTFunc(triton::FuncOp func) { - OpBuilder builder(func); - - auto name = func.getName(); - auto type = func.getFunctionType(); - - SmallVector argAttrs, resAttrs; - func.getAllArgAttrs(argAttrs); - func.getAllResultAttrs(resAttrs); - - // bit-casted tt.ptr的特殊处理 - SmallVector inputTypes{type.getInputs()}; - SmallVector retTypes{type.getResults()}; - if (func.getSymVisibility() == "public" && !func.isDeclaration()) { - for (size_t i = 0; i < func.getNumArguments(); ++i) { - auto arg = func.getArgument(i); - // Special method for i1 arg - if (!isa(arg.getType()) || - dyn_cast(arg.getType()).getPointeeType() != - builder.getIntegerType(1)) { - continue; - } - - // 收集当前参数的 users（拷贝到 vector，方便之后 erase） - SmallVector argVaildUser{arg.getUsers()}; - bool isAllBitcastI1ToI8 = - llvm::all_of(argVaildUser, [](Operation *userOp) { - return userOp->hasAttr("Input_Arg_i1_Bitcast_To_i8"); - }); - - if (!isAllBitcastI1ToI8) { - LLVM_DEBUG({ - auto &os = llvm::dbgs(); - os << arg << " has users:\n"; - int cnt = 0; - for (auto it : argVaildUser) { - os << "users[" << cnt++ << "] = " << *it; - } - }); - func->emitError("The parameters of type i1 input to the function " - "cannot be processed."); - return; - } - - // 创建 new pointer type (i8 pointer, 复用原 address space) - auto ttPtrType = triton::PointerType::get( - builder.getI8Type(), - mlir::dyn_cast(arg.getType()).getAddressSpace()); - - // 先把参数类型改成 i8 pointer（这样替换时类型一致） - arg.setType(ttPtrType); - inputTypes[i] = arg.getType(); - - // 把 argVaildUser 的结果的 use 全部替换为 arg，然后把这些 op erase 掉 - for (Operation *userOp : argVaildUser) { - // userOp 可能已经被移除或不在模块中，检查父操作存在性 - if (!userOp || !userOp->getParentOp()) - continue; - - // 对 userOp 的每个 result 做替换（如果类型一致），若不一致则报错并返回 - for (Value res : userOp->getResults()) { - // 如果类型不一致，直接报错并返回（不做自动 cast） - if (res.getType() != arg.getType()) { - userOp->emitError("Result type of this operation is incompatible " - "with the new parameter type; cannot replace " - "uses safely."); - return; - } - - // 类型一致则替换所有 uses - if (!res.use_empty()) - res.replaceAllUsesWith(arg); - } - - // 删除原来的中间操作（例如原来的 tt.bitcast） - userOp->erase(); - } - } - } - - auto castType = FunctionType::get(func.getContext(), inputTypes, retTypes); - - auto funcFunc = builder.create(func.getLoc(), name, castType); - funcFunc.setAllArgAttrs(argAttrs); - funcFunc.setAllResultAttrs(resAttrs); - - auto &funcFuncBody = funcFunc.getBody(); - auto &funcBody = func.getBody(); - - IRMapping map; - funcBody.cloneInto(&funcFuncBody, map); - - func.erase(); -} - -struct BoolTritonPtrPromotionPass - : mlir::dicp::trtion_ext::impl::BoolTritonPtrPromotionBase< - BoolTritonPtrPromotionPass> { - void runOnOperation() override; -}; -} // namespace - -void BoolTritonPtrPromotionPass::runOnOperation() { - auto moduleOp = getOperation(); - auto ctx = &getContext(); - { - RewritePatternSet patterns(ctx); - patterns.add(ctx); - patterns.add(ctx); - if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { - moduleOp->emitError("failed to apply Canonicalizer Patterns"); - signalPassFailure(); - } - } - moduleOp.walk([&](triton::FuncOp func) { converti1ArgInTTFunc(func); }); -} - -std::unique_ptr> -trtion_ext::createBoolTritonPtrPromotionPass() { - return std::make_unique(); -} diff --git a/compiler/lib/Dialect/TritonExt/Transforms/CMakeLists.txt b/compiler/lib/Dialect/TritonExt/Transforms/CMakeLists.txt deleted file mode 100644 index 5a1967b3..00000000 --- a/compiler/lib/Dialect/TritonExt/Transforms/CMakeLists.txt +++ /dev/null @@ -1,23 +0,0 @@ -add_triton_library(TritonExtTransforms - CanonicalizeTritonIRAscend.cpp - CanonicalizeCmpiPass.cpp - CanonicalizerPattern.cpp - - DEPENDS - TritonExtTransformsIncGen - - LINK_LIBS PUBLIC - MLIRArithDialect - MLIRDialectUtils - MLIRIR - MLIRMathDialect - MLIRPass - MLIRTensorDialect - MLIRTransforms - MLIRSupport - - TritonIR - TritonSharedAnalysis - TritonArithToLinalg - TritonToStructured -) diff --git a/compiler/lib/Dialect/TritonExt/Transforms/CanonicalizeCmpiPass.cpp b/compiler/lib/Dialect/TritonExt/Transforms/CanonicalizeCmpiPass.cpp deleted file mode 100644 index e3282733..00000000 --- a/compiler/lib/Dialect/TritonExt/Transforms/CanonicalizeCmpiPass.cpp +++ /dev/null @@ -1,237 +0,0 @@ -#include "dicp/Dialect/TritonExt/Transforms/Passes.h" - -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/Triton/IR/Types.h" - -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Operation.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" -#include "mlir/Transforms/Passes.h" - -#include "llvm/ADT/STLExtras.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/raw_ostream.h" - -#define DEBUG_TYPE "canonical-cmp" - -using namespace mlir; -using namespace dicp; -using namespace trtion_ext; -using namespace triton; - -namespace mlir::dicp::trtion_ext { -#define GEN_PASS_DEF_CANONICALIZECMPI -#include "dicp/Dialect/TritonExt/Transforms/Passes.h.inc" -} // namespace mlir::dicp::trtion_ext - -namespace { - -/// Create a constant of value `1` of the given type -/// `tyNewElem`: return an IntegerAttr constant (arith.constant) -static Value buildConstOne(OpBuilder &b, Location loc, IntegerType intTy) { - MLIRContext *ctx = b.getContext(); - APInt one(intTy.getWidth(), 1); - auto intAttr = IntegerAttr::get(intTy, one); - return b.create(loc, intTy, intAttr); -} - -/// The rewrite pattern for arith::CmpIOp -struct CmpISemanticRewritePattern : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(arith::CmpIOp cmp, - PatternRewriter &rewriter) const final { - auto pred = cmp.getPredicate(); - - // ttshared只支持slt,ult,sge，因此将sle,sgt,ule转化为slt,ult,sge - bool isSLE = (pred == arith::CmpIPredicate::sle); - bool isSGT = (pred == arith::CmpIPredicate::sgt); - bool isULE = (pred == arith::CmpIPredicate::ule); - if (!isSLE && !isULE) - return failure(); - // 过滤：保留那些是 triton.load/triton.store 且 cmp.result() 被用作它们的 - // mask 的 users - SmallVector cmpUsers{cmp->getUsers()}; - llvm::erase_if(cmpUsers, [&](Operation *userOp) { - if (auto loadOp = dyn_cast(userOp)) { - // 请根据你工程中生成的 accessor 名称调整下面这一行 - // mask 可能为 nullptr / null Value，先检查再比较 - if (loadOp.getMask() && loadOp.getMask() == cmp.getResult()) - return false; // keep - } else if (auto storeOp = dyn_cast(userOp)) { - if (storeOp.getMask() && storeOp.getMask() == cmp.getResult()) - return false; // keep - } - return true; // erase everything else - }); - // 不存在则，不处理 - if (cmpUsers.empty()) { - LLVM_DEBUG(llvm::dbgs() - << "[CmpISemanticRewritePattern] Cmp op: " << cmp << "\n"; - llvm::dbgs() << " Predicate: " - << static_cast(cmp.getPredicate()) << "\n"; - llvm::dbgs() << " Users of cmp result:\n"; - - // 遍历 cmp.getResult() 的所有 users（直接用 getResult() 而非 - // getUsers()） - for (Operation *userOp - : cmp.getResult().getUsers()) { - llvm::dbgs() << " -> " << userOp->getName() << " ("; - userOp->print(llvm::dbgs()); - llvm::dbgs() << ")\n"; - }); - return failure(); - } - Value rhs = cmp.getRhs(); - auto splatOp = rhs.getDefiningOp(); - if (splatOp == nullptr) - return failure(); - - Value scrVal = splatOp.getSrc(); - IntegerType scrValType = mlir::dyn_cast(scrVal.getType()); - if (scrValType == nullptr) - return failure(); - - Location loc = cmp.getLoc(); - // Build constant 1 of extType - Value cstOne = buildConstOne(rewriter, loc, scrValType); - - Value newSplatSrcVal = nullptr; - // Map predicate - arith::CmpIPredicate newPred; - rewriter.setInsertionPoint(splatOp); - if (isSLE) { - newSplatSrcVal = - rewriter.create(loc, scrValType, scrVal, cstOne); - newPred = arith::CmpIPredicate::slt; - - } else if (isSGT) { - newSplatSrcVal = - rewriter.create(loc, scrValType, scrVal, cstOne); - newPred = arith::CmpIPredicate::sge; - - } else /* isULE */ { - newSplatSrcVal = - rewriter.create(loc, scrValType, scrVal, cstOne); - newPred = arith::CmpIPredicate::ult; - } - - triton::SplatOp newSplatOp = rewriter.create( - loc, splatOp.getType(), newSplatSrcVal); - rewriter.replaceOp(splatOp, newSplatOp); - - rewriter.setInsertionPoint(cmp); - Value lhs = cmp.getLhs(); - // Create new cmp on expanded types - Value newCmp = rewriter.create(loc, newPred, lhs, - newSplatOp.getResult()); - - // If the original result type is the same, just replace. - // NOTE: arith.cmpi returns an integer or vector of i1; our newCmp has same - // result shape. - rewriter.replaceOp(cmp, newCmp); - return success(); - } -}; - -/// Reverse a CmpIPredicate for operand swapping: -/// a < b <=> b > a, etc. -/// eq/ne are symmetric. -static arith::CmpIPredicate reversePredicate(arith::CmpIPredicate p) { - switch (p) { - case arith::CmpIPredicate::eq: - return arith::CmpIPredicate::eq; - case arith::CmpIPredicate::ne: - return arith::CmpIPredicate::ne; - case arith::CmpIPredicate::slt: - return arith::CmpIPredicate::sgt; - case arith::CmpIPredicate::sle: - return arith::CmpIPredicate::sge; - case arith::CmpIPredicate::sgt: - return arith::CmpIPredicate::slt; - case arith::CmpIPredicate::sge: - return arith::CmpIPredicate::sle; - case arith::CmpIPredicate::ult: - return arith::CmpIPredicate::ugt; - case arith::CmpIPredicate::ule: - return arith::CmpIPredicate::uge; - case arith::CmpIPredicate::ugt: - return arith::CmpIPredicate::ult; - case arith::CmpIPredicate::uge: - return arith::CmpIPredicate::ule; - default: - // For safety, return the original if unknown - return p; - } -} - -/// If cmp.lhs is defined by triton::SplatOp and cmp.rhs is NOT a -/// triton::SplatOp, swap lhs/rhs and reverse the predicate. -struct SwapSplatLhsCmpPattern : public OpRewritePattern { - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(arith::CmpIOp cmp, - PatternRewriter &rewriter) const final { - // Check lhs is splat - Value lhs = cmp.getLhs(); - Value rhs = cmp.getRhs(); - - auto lhsSplat = lhs.getDefiningOp(); - auto rhsSplat = rhs.getDefiningOp(); - if (!lhsSplat || rhsSplat) - return failure(); - - Value lhsSplatScrVal = lhsSplat.getSrc(); - if (!mlir::isa(lhsSplatScrVal.getType())) - return failure(); - - // Swap operands and reverse predicate - Location loc = cmp.getLoc(); - arith::CmpIPredicate oldPred = cmp.getPredicate(); - arith::CmpIPredicate newPred = reversePredicate(oldPred); - - // Create the new cmp with swapped operands - rewriter.setInsertionPoint(cmp); - Value newCmp = rewriter.create(loc, newPred, rhs, lhs); - - rewriter.replaceOp(cmp, newCmp); - return success(); - } -}; - -/// The pass that applies the pattern -struct CanonicalizeCmpiPass - : mlir::dicp::trtion_ext::impl::CanonicalizeCmpiBase { - - void runOnOperation() override { - auto moduleOp = getOperation(); - auto ctx = &getContext(); - { - RewritePatternSet patterns(ctx); - patterns.add(ctx); - patterns.add(ctx); - if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { - moduleOp->emitError("failed to apply Canonicalizer Patterns"); - signalPassFailure(); - } - } - PassManager pm(&getContext(), moduleOp.getOperationName()); - // Erase dead code and fold constants created during lowering - pm.addPass(createCSEPass()); - pm.addPass(createCanonicalizerPass()); - if (failed(runPipeline(pm, getOperation()))) { - signalPassFailure(); - } - } -}; -} // namespace - -std::unique_ptr> -mlir::dicp::trtion_ext::createCanonicalizeCmpiPass() { - return std::make_unique(); -} diff --git a/compiler/lib/Dialect/TritonExt/Transforms/CanonicalizeTritonIRAscend.cpp b/compiler/lib/Dialect/TritonExt/Transforms/CanonicalizeTritonIRAscend.cpp deleted file mode 100644 index 936c09d2..00000000 --- a/compiler/lib/Dialect/TritonExt/Transforms/CanonicalizeTritonIRAscend.cpp +++ /dev/null @@ -1,340 +0,0 @@ -#include "dicp/Dialect/TritonExt/Transforms/CanonicalizerPattern.h" -#include "dicp/Dialect/TritonExt/Transforms/Passes.h" - -#include "triton/Dialect/Triton/IR/Dialect.h" -#include "triton/Dialect/Triton/IR/Types.h" - -#include "mlir/IR/Block.h" -#include "mlir/IR/Builders.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/MLIRContext.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/Operation.h" -#include "mlir/IR/PatternMatch.h" -#include "mlir/IR/Types.h" -#include "mlir/IR/Value.h" -#include "mlir/Interfaces/FunctionInterfaces.h" -#include "mlir/Pass/Pass.h" -#include "mlir/Pass/PassManager.h" -#include "mlir/Support/LogicalResult.h" -#include "mlir/Transforms/DialectConversion.h" -#include "mlir/Transforms/GreedyPatternRewriteDriver.h" - -#include "llvm/ADT/STLExtras.h" -#include "llvm/ADT/SmallPtrSet.h" -#include "llvm/ADT/SmallVector.h" -#include "llvm/ADT/StringRef.h" -#include "llvm/Support/Casting.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/LogicalResult.h" -#include "llvm/Support/raw_ostream.h" - -#define DEBUG_TYPE "canonicalize-triton-ir-ascend" - -using namespace mlir; -using namespace dicp; -using namespace trtion_ext; -using namespace triton; - -namespace mlir::dicp::trtion_ext { -#define GEN_PASS_DEF_CANONICALIZETRITONIRASCEND -#include "dicp/Dialect/TritonExt/Transforms/Passes.h.inc" -} // namespace mlir::dicp::trtion_ext - -namespace { - -class BitcastConverter : public OpRewritePattern { -public: - using OpRewritePattern::OpRewritePattern; - - LogicalResult matchAndRewrite(triton::BitcastOp op, - PatternRewriter &rewriter) const { - if (op->hasAttr("Input_Arg_i1_Bitcast_To_i8")) { - return failure(); - } - - Value result; - if (auto resPointerType = dyn_cast(op.getType())) { - // TODO: use typeconverter - auto srcPointerType = cast(op.getSrc().getType()); - auto resType = MemRefType::get({ShapedType::kDynamic}, - resPointerType.getPointeeType()); - // Handling special case - // %0 = tt.bitcast %arg0 {MixUse} : !tt.ptr -> !tt.ptr - if (isa(op.getSrc()) && - srcPointerType.getPointeeType() == rewriter.getIntegerType(1) && - resPointerType.getPointeeType() == rewriter.getIntegerType(8)) { - rewriter.modifyOpInPlace(op, [&]() { - op->setAttr("Input_Arg_i1_Bitcast_To_i8", rewriter.getUnitAttr()); - }); - return success(); - } - result = - rewriter.create(op.getLoc(), resType, op.getSrc()); - } else { - result = rewriter.create(op.getLoc(), op.getType(), - op.getSrc()); - } - rewriter.replaceOp(op, result); - return success(); - } -}; - -static void converti1ArgInTTFunc(triton::FuncOp func) { - OpBuilder builder(func); - - auto name = func.getName(); - auto type = func.getFunctionType(); - - SmallVector argAttrs, resAttrs; - func.getAllArgAttrs(argAttrs); - func.getAllResultAttrs(resAttrs); - - // bit-casted tt.ptr的特殊处理 - SmallVector inputTypes{type.getInputs()}; - SmallVector retTypes{type.getResults()}; - if (func.getSymVisibility() == "public" && !func.isDeclaration()) { - for (size_t i = 0; i < func.getNumArguments(); ++i) { - auto arg = func.getArgument(i); - // Special method for i1 arg - if (!isa(arg.getType()) || - dyn_cast(arg.getType()).getPointeeType() != - builder.getIntegerType(1)) { - continue; - } - - // 收集当前参数的 users（拷贝到 vector，方便之后 erase） - SmallVector argVaildUser{arg.getUsers()}; - bool isAllBitcastI1ToI8 = - llvm::all_of(argVaildUser, [](Operation *userOp) { - return userOp->hasAttr("Input_Arg_i1_Bitcast_To_i8"); - }); - - if (!isAllBitcastI1ToI8) { - LLVM_DEBUG({ - auto &os = llvm::dbgs(); - os << arg << " has users:\n"; - int cnt = 0; - for (auto it : argVaildUser) { - os << "users[" << cnt++ << "] = " << *it; - } - }); - func->emitError("The parameters of type i1 input to the function " - "cannot be processed."); - return; - } - - // 创建 new pointer type (i8 pointer, 复用原 address space) - auto ttPtrType = triton::PointerType::get( - builder.getI8Type(), - mlir::dyn_cast(arg.getType()).getAddressSpace()); - - // 先把参数类型改成 i8 pointer（这样替换时类型一致） - arg.setType(ttPtrType); - inputTypes[i] = arg.getType(); - - // 把 argVaildUser 的结果的 use 全部替换为 arg，然后把这些 op erase 掉 - for (Operation *userOp : argVaildUser) { - // userOp 可能已经被移除或不在模块中，检查父操作存在性 - if (!userOp || !userOp->getParentOp()) - continue; - - // 对 userOp 的每个 result 做替换（如果类型一致），若不一致则报错并返回 - for (Value res : userOp->getResults()) { - // 如果类型不一致，直接报错并返回（不做自动 cast） - if (res.getType() != arg.getType()) { - userOp->emitError("Result type of this operation is incompatible " - "with the new parameter type; cannot replace " - "uses safely."); - return; - } - - // 类型一致则替换所有 uses - if (!res.use_empty()) - res.replaceAllUsesWith(arg); - } - - // 删除原来的中间操作（例如原来的 tt.bitcast） - userOp->erase(); - } - } - } - - auto castType = FunctionType::get(func.getContext(), inputTypes, retTypes); - - auto funcFunc = builder.create(func.getLoc(), name, castType); - funcFunc.setAllArgAttrs(argAttrs); - funcFunc.setAllResultAttrs(resAttrs); - - auto &funcFuncBody = funcFunc.getBody(); - auto &funcBody = func.getBody(); - - IRMapping map; - funcBody.cloneInto(&funcFuncBody, map); - - func.erase(); -} - -struct CanonicalizeTritonIRAscendPass - : mlir::dicp::trtion_ext::impl::CanonicalizeTritonIRAscendBase< - CanonicalizeTritonIRAscendPass> { - void runOnOperation() override; - - void - populateTritonToLinalgCanonicalizationPatterns(RewritePatternSet &patterns); - template - void addTensorKindToArguments(OpTy op, triton::FuncOp func, - TensorKind tensorKind); -}; - -} // namespace - -static void setBlockArgumentAttr(BlockArgument blockArg, triton::FuncOp func, - TensorKind tensorKind) { - unsigned argIdx = blockArg.getArgNumber(); - auto existingAttr = - func.getArgAttrOfType(argIdx, "tt.tensor_kind"); - TensorKind oldVal = existingAttr - ? static_cast(existingAttr.getInt()) - : TensorKind::NONE; - - TensorKind finalVal = tensorKind; - if ((oldVal == TensorKind::INPUT && tensorKind == TensorKind::OUTPUT) || - (oldVal == TensorKind::OUTPUT && tensorKind == TensorKind::INPUT)) { - finalVal = TensorKind::INPUT_OUTPUT; - } else if (oldVal == TensorKind::INPUT_OUTPUT) { - finalVal = oldVal; - } - - func.setArgAttr( - argIdx, "tt.tensor_kind", - IntegerAttr::get(IntegerType::get(func.getContext(), INT_BIT_WIDTH), - static_cast(finalVal))); -} - -template -void CanonicalizeTritonIRAscendPass::addTensorKindToArguments( - OpTy op, triton::FuncOp func, TensorKind tensorKind) { - Value ptr = op.getPtr(); - if (!ptr) - return; - - Value cur = ptr; - llvm::SmallPtrSet visited; - // 回溯 def-use 链，找到起源 BlockArgument - while (visited.insert(cur).second) { - // 如果是 BlockArgument，则尝试设置属性 - if (auto blockArg = dyn_cast(cur)) { - if (blockArg.getOwner() == &func.getBody().front()) { - auto type = blockArg.getType(); - // 检查是否是 triton::PointerType - if (!isa(type)) - break; - setBlockArgumentAttr(blockArg, func, tensorKind); - break; - } - } - - Operation *defOp = cur.getDefiningOp(); - if (!defOp) - break; - cur = defOp->getOperand(0); - } -} - -void CanonicalizeTritonIRAscendPass:: - populateTritonToLinalgCanonicalizationPatterns( - RewritePatternSet &patterns) { - - patterns.add( - patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer, - ScalarMathCanonicalizer>( - patterns.getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); -} - -void CanonicalizeTritonIRAscendPass::runOnOperation() { - auto moduleOp = getOperation(); - auto ctx = &getContext(); - - // Check if the kernel contains tl.dot. Without tl.dot, - // the kernel would be pure AIV kernel. - bool existDot = false; - moduleOp.walk([&](triton::DotOp dotOp) { - existDot = true; - return WalkResult::interrupt(); - }); - moduleOp.walk([&](triton::DotScaledOp dotScaledOp) { - existDot = true; - return WalkResult::interrupt(); - }); - - // Traverse all the triton::FuncOp to add tensor_kind attribute - moduleOp.walk([&](triton::FuncOp func) { - func.walk([&](triton::LoadOp loadOp) { - addTensorKindToArguments(loadOp, func, TensorKind::INPUT); - }); - func.walk([&](triton::StoreOp storeOp) { - addTensorKindToArguments(storeOp, func, TensorKind::OUTPUT); - }); - func.walk([&](triton::AtomicRMWOp atomicOp) { - addTensorKindToArguments(atomicOp, func, TensorKind::INPUT_OUTPUT); - }); - func.walk([&](triton::AtomicCASOp atomicOp) { - addTensorKindToArguments(atomicOp, func, TensorKind::INPUT_OUTPUT); - }); - }); - - { - RewritePatternSet canonicalizerPatterns(&getContext()); - populateTritonToLinalgCanonicalizationPatterns(canonicalizerPatterns); - if (failed(applyPatternsGreedily(moduleOp, - std::move(canonicalizerPatterns)))) { - moduleOp->emitError("failed to apply Canonicalizer Patterns"); - signalPassFailure(); - } - } - - { - RewritePatternSet patterns(ctx); - patterns.add(ctx); - if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { - moduleOp->emitError("failed to apply Canonicalizer Patterns"); - signalPassFailure(); - } - } - moduleOp.walk([&](triton::FuncOp func) { converti1ArgInTTFunc(func); }); -} - -std::unique_ptr> -trtion_ext::createCanonicalizeTritonIRAscendPass() { - return std::make_unique(); -} diff --git a/compiler/lib/Dialect/TritonExt/Transforms/CanonicalizerPattern.cpp b/compiler/lib/Dialect/TritonExt/Transforms/CanonicalizerPattern.cpp deleted file mode 100644 index ec09b292..00000000 --- a/compiler/lib/Dialect/TritonExt/Transforms/CanonicalizerPattern.cpp +++ /dev/null @@ -1,1012 +0,0 @@ - -#include "dicp/Dialect/TritonExt/Transforms/CanonicalizerPattern.h" -#include "dicp/Dialect/TritonExt/Transforms/Passes.h" - -#include "mlir/Dialect/Affine/IR/AffineOps.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Arith/Utils/Utils.h" -#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" -#include "mlir/Dialect/LLVMIR/LLVMDialect.h" -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/Passes.h" -#include "mlir/Dialect/Math/IR/Math.h" -#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" -#include "mlir/Dialect/Utils/StaticValueUtils.h" -#include "mlir/IR/Attributes.h" -#include "mlir/IR/BuiltinAttributes.h" -#include "mlir/IR/BuiltinOps.h" -#include "mlir/IR/BuiltinTypeInterfaces.h" -#include "mlir/IR/BuiltinTypes.h" -#include "mlir/IR/Location.h" -#include "mlir/IR/OpDefinition.h" -#include "mlir/IR/Value.h" - -#include "triton-shared/Analysis/MaskAnalysis.h" - -#include "triton/Dialect/Triton/IR/Dialect.h" - -#include "llvm/ADT/DenseMap.h" -#include "llvm/ADT/SmallVectorExtras.h" -#include "llvm/ADT/TypeSwitch.h" -#include "llvm/Support/Casting.h" -#include "llvm/Support/Debug.h" -#include "llvm/Support/ErrorHandling.h" -#include "llvm/Support/FormatVariadic.h" -#include "llvm/Support/MathExtras.h" - -#include -#include -#include - -#define DEBUG_TYPE "triton-load-store-converter" - -using namespace mlir; -using namespace triton; - -namespace mlir::dicp::trtion_ext { -const std::string GeneratedByMakeTensorPtrTAG = "GeneratedByMakeTensorPtr"; -static SmallVector getNParallelLoopsAttrs(unsigned n) { - return SmallVector(n, utils::IteratorType::parallel); -} - -AtomicRMWConverter::AtomicRMWConverter(MLIRContext *context) - : OpConversionPattern(context) {} - -// lowering tt.atomicRMW to linalg.generic -// If atomic op's return value is used by other op as it's the old value stored -// at the ptrwe will use tt.load to get it -// -// example: -// input: -// %return_value = tt.atomic_rmw fadd, acq_rel, gpu, -// %output_memref, %input_tensor, %mask : -// (tensor<256x!tt.ptr>, tensor<256xf32>, tensor<256xi1>) -// -> tensor<256xf32> -// -// output: -// memref.copy %output_memref, %ub_buf : memref to memref -// %17 = bufferization.to_tensor %alloc_3 restrict writable : memref<256xf32> -// linalg.generic -// {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} -// ins(%output_memref, %masked_input_memref : memref, memref) -// outs(%subview_2 : memref) -// attrs = {GenericAtomicRMW = "fadd", MemSemantic = "acq_rel", -// MemSyncScope = "gpu"} { -// ^bb0(%in: f32, %in_9: f32, %out: f32): -// %25 = arith.addf %in, %in_9 : f32 -// linalg.yield %25 : f32 -// } -LogicalResult -AtomicRMWConverter::matchAndRewrite(triton::AtomicRMWOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const { - // If the result of AtomicRMWOp is not used, we don't need to load the old - // data stored at the ptr - auto ptr = adaptor.getPtr(); - auto val = op.getVal(); - auto loc = op.getLoc(); - - auto resType = dyn_cast(op.getResult().getType()); - if (!resType) { - return rewriter.notifyMatchFailure( - op, "atomicRMWConverter: scalar will be handled by " - "ScalarAtomicRMWCanonicalizer"); - } - - auto rmwOp = op.getAtomicRmwOp(); - if (rmwOp == triton::RMWOp::UMAX || rmwOp == triton::RMWOp::UMIN) { - return rewriter.notifyMatchFailure( - op, "AtomicRMWConverter: unsupported atomic kind for now"); - } - - // 1. Simple case where no mask is used. - auto type = dyn_cast(ptr.getType()); - if (!type) { - // Seen when implicit broadcasting is done late in a chain of - // operations. The workaround is to broadcast the pointers early in the - // address calculation. A proper fix is complicated, but at least we can - // provide a better error message. - return rewriter.notifyMatchFailure( - op, "AtomicRMWOp expects a memref, not a memref of pointers"); - } - - auto dstMemref = ptr; - // Well, linalg structure op wouldn't support mixed tensor/buffer semantics - // any more in latest LLVM(triton LLVM dependency has involed this), so we - // need to convert tensor to buffer early. - auto dstOriType = cast(dstMemref.getType()); - MemRefType dstType = - MemRefType::get(dstOriType.getShape(), dstOriType.getElementType()); - Value inputMemref = - rewriter.create(loc, dstType, val); - - // 2. handle the mask for the atomic op - // When the dsl do not pass the mask to this op like - // `tl.atomic_add(out_ptr0 + xindex, tmp2)`, it will create a constant mask - // for this op by default, which is not supported by maskAnalysis, so we - // need to handle this situation - // - // This logic come from semantic.py: - // - // if not mask: - // mask_ir = builder.get_int1(True) - // mask_ty = tl.int1 - // if ptr.type.is_block(): - // mask_ir = \ - // builder.create_splat(mask_ir, ptr.type.get_block_shapes()) - // mask_ty = tl.block_type(tl.int1, ptr.type.get_block_shapes()) - // mask = tl.tensor(mask_ir, mask_ty) - // - // ... - // - // return ptr, val, mask - // - if (auto mask = op.getMask()) { - triton::MaskState mstate; - auto constantMask = mask.getDefiningOp(); - if (!constantMask) { - auto isContMask = mstate.parse(mask, loc, rewriter); - - if (isContMask.failed()) { - return rewriter.notifyMatchFailure( - op, "Cannot lower continuous masked loads"); - } - dstMemref = mstate.getSubview(ptr, loc, rewriter); - inputMemref = mstate.getSubview(inputMemref, loc, rewriter); - } else { - if (!isConstantMaskTrue(mask)) { - rewriter.eraseOp(op); - return success(); - } - } - } - - // create element-wise map - int64_t rank = type.getRank(); - SmallVector inputDims; - auto context = rewriter.getContext(); - - for (int i = 0; i < rank; i++) { - inputDims.push_back(getAffineDimExpr(i, context)); - } - - SmallVector indexingMaps; - // As mask has been erased for now - // the number of input must be 2 - // the input memref is also the output memref - // Thus, there are a total of three inputs and outputs. - // so here we have 3 map to create - for (int i = 0; i < 3; i++) { - indexingMaps.push_back(AffineMap::get(rank, 0, inputDims, context)); - } - - Value tensorToReplace; - if (!op.getResult().use_empty()) { - auto tensorType = - RankedTensorType::get(type.getShape(), type.getElementType()); - auto alloc = rewriter.create( - loc, MemRefType::get(type.getShape(), type.getElementType())); - // For the return value, don't need to care about mask for now - // this op don't support other, so we best not fill it - rewriter.create(loc, ptr, alloc); - tensorToReplace = rewriter.create( - loc, tensorType, alloc, true /* restrict */, true /* writable */); - } - - auto linalgOp = rewriter.create( - loc, /* operands */ ValueRange{dstMemref, inputMemref}, - ValueRange{dstMemref}, indexingMaps, getNParallelLoopsAttrs(rank), - [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) { - Value opResult = createAtomicBinaryOps(nestedBuilder, nestedLoc, op, - type.getElementType(), - blockArgs[0], blockArgs[1]); - nestedBuilder.create(nestedLoc, opResult); - }); - - // "library_call" - // indicating the actual semantic of this op - // TODO: If the hardware support the MemSemantic/MemSyncScope - // We pass them down - // otherwise they need to be deleted - const StringRef genericAtomicRMW = "GenericAtomicRMW"; - const StringRef memSemantic = "MemSemantic"; - const StringRef memSyncScope = "MemSyncScope"; - linalgOp->setAttr(genericAtomicRMW, - rewriter.getStringAttr(stringifyEnum(op.getAtomicRmwOp()))); - linalgOp->setAttr(memSemantic, - rewriter.getStringAttr(stringifyEnum(op.getSem()))); - linalgOp->setAttr(memSyncScope, - rewriter.getStringAttr(stringifyEnum(op.getScope()))); - - // Mark atomic_and/or/xor specially which need software simulation in terms - // of backend restriction - if (softwareAtomicKinds.contains(op.getAtomicRmwOp())) - linalgOp->setAttr("Software", rewriter.getUnitAttr()); - - // tt.atomicRMW op has two part of feature - // 1. load the old data at the ptr - // 2. atomically store the data on ub to the ptr - // at the same time it perform the action it has been assigned - // So we lower this op to load + atomically store - // - // The first part is not necessary when the returned value of atomic op - // is not used, it will be deleted cause it's meaningless - // Here, we preemptively determine whether it will be used - // and decide whether it is necessary to create the load process based on - // this assessment. - // - // logic of handling is copied - // TODO: decoupling the logic of load, put it in the Utils - if (!op.getResult().use_empty()) { - rewriter.replaceOp(op, tensorToReplace); - } else { - rewriter.eraseOp(op); - } - return success(); -} - -LogicalResult -AtomicCASConverter::matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const { - // If the result of AtomicCASOp is not used, we don't need to load the old - // data stored at the ptr - auto ptr = adaptor.getPtr(); - auto cmp = op.getCmp(); - auto val = op.getVal(); - auto loc = op.getLoc(); - - auto resType = dyn_cast(op.getResult().getType()); - if (!resType) { - return rewriter.notifyMatchFailure( - op, "atomicCASConverter: scalar will be handled by " - "ScalarAtomicCASCanonicalizer"); - } - - // 1. Simple case where no mask is used. - auto type = dyn_cast(ptr.getType()); - if (!type) { - // Seen when implicit broadcasting is done late in a chain of - // operations. The workaround is to broadcast the pointers early in the - // address calculation. A proper fix is complicated, but at least we can - // provide a better error message. - return rewriter.notifyMatchFailure( - op, "AtomicCASOp expects a memref, not a memref of pointers"); - } - - auto dstMemref = ptr; - // Well, linalg structure op wouldn't support mixed tensor/buffer semantics - // any more in latest LLVM(triton LLVM dependency has involed this), so we - // need to convert tensor to buffer early. - auto dstOriType = cast(dstMemref.getType()); - MemRefType dstType = - MemRefType::get(dstOriType.getShape(), dstOriType.getElementType()); - Value inputMemref = - rewriter.create(loc, dstType, val); - - Value cmpMemref = - rewriter.create(loc, dstType, cmp); - - // create element-wise map - int64_t rank = type.getRank(); - SmallVector inputDims; - auto context = rewriter.getContext(); - - for (int i = 0; i < rank; i++) { - inputDims.push_back(getAffineDimExpr(i, context)); - } - - SmallVector indexingMaps; - // As mask has been erased for now - // the number of input must be 2 - // the input memref is also the output memref - // Thus, there are a total of four inputs and outputs. - // so here we have 4 map to create - for (int i = 0; i < 4; i++) { // 4: 3 input and 1 output - indexingMaps.push_back(AffineMap::get(rank, 0, inputDims, context)); - } - - auto linalgOp = rewriter.create( - loc, ValueRange{dstMemref, cmpMemref, inputMemref}, - mlir::ValueRange{dstMemref}, indexingMaps, getNParallelLoopsAttrs(rank), - [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) { - Value lhs = blockArgs[0]; - Value rhs = blockArgs[1]; - Value setValue = blockArgs[2]; - Value cond; - if (mlir::isa(lhs.getType())) { - cond = nestedBuilder.create( - nestedLoc, arith::CmpFPredicate::UEQ, lhs, rhs); - } else { - cond = nestedBuilder.create( - nestedLoc, arith::CmpIPredicate::eq, lhs, rhs); - } - auto ifOp = nestedBuilder.create( - nestedLoc, TypeRange{setValue.getType()}, cond, true); - { - OpBuilder::InsertionGuard guard(nestedBuilder); - nestedBuilder.setInsertionPointToEnd(&ifOp.getThenRegion().front()); - nestedBuilder.create(nestedLoc, setValue); - } - { - OpBuilder::InsertionGuard guard(nestedBuilder); - nestedBuilder.setInsertionPointToEnd(&ifOp.getElseRegion().front()); - nestedBuilder.create(nestedLoc, lhs); - } - nestedBuilder.setInsertionPointToEnd(nestedBuilder.getBlock()); - nestedBuilder.create(nestedLoc, - ifOp.getResult(0)); - }); - - const StringRef genericAtomicRMW = "GenericAtomicRMW"; - const StringRef memSemantic = "MemSemantic"; - const StringRef memSyncScope = "MemSyncScope"; - auto attr = mlir::StringAttr::get(context, "cas"); - - linalgOp->setAttr(genericAtomicRMW, attr); - linalgOp->setAttr(memSemantic, - rewriter.getStringAttr(stringifyEnum(op.getSem()))); - linalgOp->setAttr(memSyncScope, - rewriter.getStringAttr(stringifyEnum(op.getScope()))); - - linalgOp->setAttr("Software", rewriter.getUnitAttr()); - - // tt.atomicRMW op has two part of feature - // 1. load the old data at the ptr - // 2. atomically store the data on ub to the ptr - // at the same time it perform the action it has been assigned - // So we lower this op to load + atomically store - // - // The first part is not necessary when the returned value of atomic op - // is not used, it will be deleted cause it's meaningless - // Here, we preemptively determine whether it will be used - // and decide whether it is necessary to create the load process based on - // this assessment. - // - // logic of handling is copied - if (!op.getResult().use_empty()) { - auto tensorType = - RankedTensorType::get(type.getShape(), type.getElementType()); - auto alloc = rewriter.create( - loc, MemRefType::get(type.getShape(), type.getElementType())); - - // For the return value, don't need to care about mask for now - // this op don't support other, so we best not fill it - rewriter.create(loc, ptr, alloc); - Value tensor = rewriter.create( - loc, tensorType, alloc, true /* restrict */, true /* writable */); - rewriter.replaceOp(op, tensor); - } else { - rewriter.eraseOp(op); - } - return success(); -} - -LogicalResult -ScalarStoreCanonicalizer::matchAndRewrite(triton::StoreOp op, - PatternRewriter &rewriter) const { - if (!op.getValue().getType().isIntOrIndexOrFloat()) { - return rewriter.notifyMatchFailure( - op, "ScalarStoreCanonicalizer handles scalar store scene!"); - } - auto ptr = op.getPtr(); - auto mask = op.getMask(); - auto value = op.getValue(); - if (mask) { - rewriter.replaceOpWithNewOp( - op, mask, [&](OpBuilder &b, Location loc) { - b.create(loc, ptr, value, op.getCache(), - op.getEvict()); - b.create(loc); - }); - return success(); - } - - auto ptrTy = RankedTensorType::get({(int64_t)1}, ptr.getType()); - auto ptrSplat = rewriter.create(op.getLoc(), ptrTy, ptr); - auto valTy = RankedTensorType::get({(int64_t)1}, value.getType()); - auto valSplat = rewriter.create(op.getLoc(), valTy, value); - auto newStoreOp = rewriter.create( - op.getLoc(), ptrSplat, valSplat, op.getCache(), op.getEvict()); - rewriter.replaceOp(op, newStoreOp); - return success(); -} - -LogicalResult -ScalarAtomicRMWCanonicalizer::matchAndRewrite(triton::AtomicRMWOp op, - PatternRewriter &rewriter) const { - if (!op.getVal().getType().isIntOrIndexOrFloat()) { - return rewriter.notifyMatchFailure( - op, "ScalarAtomicRMWCanonicalizer handles scalar atomic rmw op scene!"); - } - - auto ptr = op.getPtr(); - auto ptrTy = RankedTensorType::get({(int64_t)1}, ptr.getType()); - auto ptrSplat = rewriter.create(op.getLoc(), ptrTy, ptr); - auto valTy = RankedTensorType::get({(int64_t)1}, op.getVal().getType()); - auto valSplat = - rewriter.create(op.getLoc(), valTy, op.getVal()); - auto maskTy = RankedTensorType::get({(int64_t)1}, op.getMask().getType()); - auto maskSplat = - rewriter.create(op.getLoc(), maskTy, op.getMask()); - - auto newAtomicOp = rewriter.create( - op.getLoc(), valTy, op.getAtomicRmwOp(), ptrSplat, valSplat, maskSplat, - op.getSem(), op.getScope()); - auto idxZero = - rewriter.create(op.getLoc(), rewriter.getIndexAttr(0)); - rewriter.replaceOpWithNewOp(op, newAtomicOp, - ValueRange({idxZero})); - return success(); -} - -LogicalResult -ScalarAtomicCASCanonicalizer::matchAndRewrite(triton::AtomicCASOp op, - PatternRewriter &rewriter) const { - if (!op.getVal().getType().isIntOrIndexOrFloat() && - !op.getCmp().getType().isIntOrIndexOrFloat()) { - return rewriter.notifyMatchFailure( - op, "ScalarAtomicCASCanonicalizer handles scalar atomic cas op scene!"); - } - - auto ptr = op.getPtr(); - auto ptrTy = RankedTensorType::get({(int64_t)1}, ptr.getType()); - auto ptrSplat = rewriter.create(op.getLoc(), ptrTy, ptr); - auto cmpTy = RankedTensorType::get({(int64_t)1}, op.getCmp().getType()); - auto cmpSplat = - rewriter.create(op.getLoc(), cmpTy, op.getCmp()); - auto valTy = RankedTensorType::get({(int64_t)1}, op.getVal().getType()); - auto valSplat = - rewriter.create(op.getLoc(), valTy, op.getVal()); - - auto newAtomicOp = rewriter.create( - op.getLoc(), valTy, ptrSplat, cmpSplat, valSplat, op.getSem(), - op.getScope()); - auto idxZero = - rewriter.create(op.getLoc(), rewriter.getIndexAttr(0)); - rewriter.replaceOpWithNewOp(op, newAtomicOp, - ValueRange({idxZero})); - return success(); -} - -// The atomic max op with float input will be devided into -// two atomic max ops with integer input -// One handles the part of the tensor greater than zero -// the other deals with the part less than zero -// It will lead to maskAnalysis failure -// So here we need to revert the procedures in semantics.py -// The triton IR is like -// -// %cst_0 = arith.constant dense<0.000000e+00> : tensor<1x256xf32> -// %1 = tt.bitcast %value : tensor<1x256xf32> -> tensor<1x256xi32> -// %2 = tt.bitcast %ptr : tensor<1x256x!tt.ptr> -> -// tensor<1x256x!tt.ptr> %3 = arith.cmpf oge, %1, %cst_0 %4 = arith.cmpf -// olt, %1, %cst_0 %5 = arith.andi %8, %3 %6 = tt.atomic_rmw max, acq_rel, gpu, -// %2, %1, %5 : -// (tensor<1x256x!tt.ptr>, tensor<1x256xi32>, tensor<1x256xi1>) -> -// tensor<1x256xi32> -// %7 = arith.andi %8, %4 -// %8 = tt.atomic_rmw umin, acq_rel, gpu, %2, %1, %7 : -// (tensor<1x256x!tt.ptr>, tensor<1x256xi32>, tensor<1x256xi1>) -> -// tensor<1x256xi32> -// -// it's hard to handle and meaningless complicated for our device -// so we revert it to -// %0 = tt.atomic_rmw max, acq_rel, gpu, %23, %21, %8 : -// (tensor<1x256x!tt.ptr>, tensor<1x256xf32>, tensor<1x256xi1>) -> -// tensor<1x256xf32> -LogicalResult -AtomicMaxMinCanonicalizer::matchAndRewrite(triton::AtomicRMWOp op, - PatternRewriter &rewriter) const { - // Revert the op to its original form - auto ptrBitcastOp = op.getPtr().getDefiningOp(); - auto valueBitcastOp = op.getVal().getDefiningOp(); - if (!ptrBitcastOp || !valueBitcastOp) { - return failure(); - } - - // We only need to handle the op when the element type is float - auto elementType = - dyn_cast(valueBitcastOp.getSrc().getType()).getElementType(); - if (!isa(elementType)) { - return failure(); - } - - auto rmwOp = op.getAtomicRmwOp(); - // here we know that atomic UMAX/UMIN - // is created by special logic of triton right now - // so we can simply delete it - if (rmwOp == triton::RMWOp::UMAX || rmwOp == triton::RMWOp::UMIN) { - // if the return value of op is used, we can't simply erase it - if (op.getResult().use_empty()) { - rewriter.eraseOp(op); - return success(); - } - return failure(); - } - - if (rmwOp != triton::RMWOp::MAX && rmwOp != triton::RMWOp::MIN) { - return failure(); - } - - // 1. Though semantic interpreter will generate full true tensor as original - // mask if atomicrmwOp don't have it, above float devision process will also - // generate positive and negative comparison mask, which will cause to fold - // true mask. - // 2. While if atomicrmwOp has original mask, there exists andiop between - // original mask and positive/negative comparison mask - // - // Here wanna extract original mask - Value originalMask = op.getMask(); - if (auto andOp = originalMask.getDefiningOp()) - // LHS is convention in semantic interpreter - originalMask = andOp.getLhs(); - else if (auto cmpOp = originalMask.getDefiningOp()) { - if (cmpOp.getPredicate() != mlir::arith::CmpFPredicate::OGE || - !matchPattern(cmpOp.getRhs(), - /*positive float zero matcher*/ m_PosZeroFloat())) - // Here recheck frontend interpreter generation in no manual mask state - return op->emitError("Illegal mask for atomicrmwOp of float type"); - // Restore original true mask - originalMask = rewriter.create( - op->getLoc(), - /*typed attr*/ DenseElementsAttr::get( - cast(originalMask.getType()), true)); - } else - return op->emitError("Illegal mask for atomicrmwOp of float type"); - - auto originAtomicOp = rewriter.create( - op.getLoc(), valueBitcastOp.getSrc().getType(), op.getAtomicRmwOp(), - ptrBitcastOp.getSrc(), valueBitcastOp.getSrc(), originalMask, op.getSem(), - op.getScope()); - - // if the return value of op is used - // we need to handle its usage - // In semantic.py, if the atomic Max/Min with float input is used - // It will use select + bitcast to get float value - // so here we need to revert it too - // - // For example: - // %0 = tt.atomic_rmw max, acq_rel, gpu, %gm, %input, %mask1 : - // (tensor<32x!tt.ptr>... %1 = tt.atomic_rmw umin, acq_rel, gpu, %gm, - // %input, %mask2 : (tensor<32x!tt.ptr>... %2 = arith.select - // %devidedMask, %0, %1 : tensor<32xi1>, tensor<32xi32> %3 = tt.bitcast %2 : - // tensor<32xi32> -> tensor<32xf32> tt.store %outputMemref, %3 : - // tensor<32x!tt.ptr> - // - // will be revert to: - // %0 = tt.atomic_rmw max, acq_rel, gpu, %gm, %input, %mask : - // (tensor<32x!tt.ptr>... tt.store %outputMemref, %0 : - // tensor<32x!tt.ptr> - // - if (!op.getResult().use_empty()) { - for (OpOperand &use : op->getUses()) { - auto selectOp = dyn_cast(use.getOwner()); - if (!selectOp) - continue; - - for (OpOperand &selectUse : selectOp->getUses()) { - if (auto bitcastOp = - dyn_cast(selectUse.getOwner())) { - bitcastOp.getResult().replaceAllUsesWith(originAtomicOp); - } - } - } - rewriter.replaceOp(op, originAtomicOp); - } else { - rewriter.eraseOp(op); - } - - return success(); -} - -/* - * Move tt.bitcast to a previous location if tt.bitcast is not directly applied - * on function arguments - */ -LogicalResult -BitcastCanonicalizer::matchAndRewrite(triton::BitcastOp bitcastOp, - PatternRewriter &rewriter) const { - Value castSrc = bitcastOp.getSrc(); - Value castRes = bitcastOp.getResult(); - Type castSrcTy = castSrc.getType(); - Type castSrcPtrTy = isa(castSrcTy) - ? cast(castSrcTy).getElementType() - : castSrcTy; - if (!isa(castSrcPtrTy)) - return failure(); - - auto origBitwidth = getPointeeBitWidth(castSrc.getType()); - auto castBitwidth = getPointeeBitWidth(castRes.getType()); - - if (origBitwidth == 1) - origBitwidth = 8; - if (castBitwidth == 1) - castBitwidth = 8; - if (origBitwidth != castBitwidth) { - bitcastOp.emitError() << "Casting pointers with unmatched bitwidth!\n"; - return failure(); - } - - Operation *beforeCastOp = castSrc.getDefiningOp(); - if (beforeCastOp == nullptr) { - return failure(); - } - - auto newRes = - TypeSwitch>(beforeCastOp) - // before: addptr - bitcast - load/store - // after: bitcast - addptr - load/store - .Case([&](triton::AddPtrOp addptrOp) { - auto newCastOp = rewriter.create( - bitcastOp.getLoc(), castRes.getType(), addptrOp.getPtr()); - return rewriter.create( - bitcastOp.getLoc(), castRes.getType(), newCastOp.getResult(), - addptrOp.getOffset()); - }) - .Case([&](triton::SplatOp splatOp) { - Type newCastSrcTy = - cast(castRes.getType()).getElementType(); - - Value splatSrc = splatOp.getSrc(); - Type splatSrcTy = splatSrc.getType(); - if (auto splatSrcTensorTy = dyn_cast(splatSrcTy)) - newCastSrcTy = - splatSrcTensorTy.cloneWith(std::nullopt, newCastSrcTy); - auto newCastOp = rewriter.create( - bitcastOp.getLoc(), newCastSrcTy, splatSrc); - return rewriter.create( - bitcastOp.getLoc(), castRes.getType(), newCastOp); - }) - // before: bitcast - bitcast - // after(fusion optimization): bitcast - .Case([&](triton::BitcastOp prevCastOp) { - return rewriter.create( - bitcastOp.getLoc(), castRes.getType(), prevCastOp.getSrc()); - }) - .Default([&](Operation *op) { - return rewriter.notifyMatchFailure(bitcastOp, - "Unknown bitcast pattern"); - }); - if (succeeded(newRes)) { - rewriter.replaceOp(bitcastOp, newRes.value()); - if (beforeCastOp->use_empty()) { - rewriter.eraseOp(beforeCastOp); - } - return success(); - } - return failure(); -} - -void rewriteUserWithNewOrder( - mlir::OpOperand *use, PatternRewriter &rewriter, - llvm::SmallVector &blkShapeI64, // 8: container size - mlir::Location &loc, llvm::ArrayRef &order, size_t &orderSize) { - Operation *user = use->getOwner(); - rewriter.setInsertionPointAfter(user); - if (auto loadOp = dyn_cast(user)) { - auto loadResTy = loadOp.getResult().getType(); - auto loadResShapedTy = cast(loadResTy); - auto newLoadTy = loadResShapedTy.cloneWith( - blkShapeI64, loadResShapedTy.getElementType()); - auto newLoadOp = rewriter.create( - loc, newLoadTy, loadOp->getOperands(), loadOp->getAttrs()); - newLoadOp->setAttr(GeneratedByMakeTensorPtrTAG, - UnitAttr::get(rewriter.getContext())); - rewriter.replaceOp(loadOp, newLoadOp); - // load contiguous data then permute. thus the permute order is as - // follows. - SmallVector permuteOrder; // 8: container size - for (auto [i, v] : llvm::enumerate(order)) { - permuteOrder.push_back(orderSize - 1 - order[i]); - } - auto permuteOp = rewriter.create( - loc, newLoadOp.getResult(), - DenseI32ArrayAttr::get(loadOp.getContext(), permuteOrder)); - newLoadOp.getResult().replaceAllUsesExcept(permuteOp.getResult(), - permuteOp); - } else if (auto storeOp = dyn_cast(user)) { - // permute to contiguous then store. thus the permute order is as follows. - SmallVector permuteOrder; // 8: container size - for (auto [i, v] : llvm::enumerate(order)) { - permuteOrder.push_back(order[orderSize - 1 - i]); - } - auto permuteOp = rewriter.create( - loc, storeOp.getValue(), - DenseI32ArrayAttr::get(storeOp.getContext(), permuteOrder)); - storeOp.getValue().replaceAllUsesExcept(permuteOp.getResult(), permuteOp); - auto newStoreOp = rewriter.create( - loc, storeOp.getPtr(), storeOp.getValue(), storeOp.getMask(), - storeOp.getBoundaryCheck(), storeOp.getCache(), storeOp.getEvict()); - rewriter.replaceOp(storeOp, newStoreOp); - } else if (auto advanceOp = dyn_cast(user)) { - auto advanceResPtrTy = - cast(advanceOp.getResult().getType()); - auto advanceResShapedTy = - cast(advanceResPtrTy.getPointeeType()); - auto newAdvanceResShapedTy = advanceResShapedTy.cloneWith( - blkShapeI64, advanceResShapedTy.getElementType()); - auto newAdvanceResPtrTy = triton::PointerType::get( - newAdvanceResShapedTy, advanceResPtrTy.getAddressSpace()); - auto advanceOffsets = advanceOp.getOffsets(); - llvm::SmallVector newAdvanceOffsets; // 8: container size - for (int i = orderSize - 1; i >= 0; i--) { - newAdvanceOffsets.push_back(advanceOffsets[order[i]]); - } - SmallVector resUses; - for (auto &use : advanceOp->getUses()) - resUses.push_back(&use); - auto newAdvanceOp = rewriter.create( - loc, newAdvanceResPtrTy, advanceOp.getPtr(), newAdvanceOffsets); - rewriter.replaceOp(advanceOp, newAdvanceOp); - for (auto resUse : resUses) - rewriteUserWithNewOrder(resUse, rewriter, blkShapeI64, loc, order, - orderSize); - } else if (auto loopOp = dyn_cast(user)) { - auto initArg = use->get(); - auto iterArg = loopOp.getTiedLoopRegionIterArg(use); - auto resultValue = loopOp.getTiedLoopResult(use); - iterArg.setType(initArg.getType()); - resultValue.setType(initArg.getType()); - for (auto &argUse : iterArg.getUses()) - rewriteUserWithNewOrder(&argUse, rewriter, blkShapeI64, loc, order, - orderSize); - for (auto &resUse : resultValue.getUses()) - rewriteUserWithNewOrder(&resUse, rewriter, blkShapeI64, loc, order, - orderSize); - } else if (isa(user)) { - return; - } else { - llvm_unreachable( - "[MakeTensorPtrCanonicalizer] tt.make_tensor_ptr's result is " - "not used by load/store/advance op"); - } -} - -void markLoadUsers(mlir::OpOperand *use, PatternRewriter &rewriter) { - Operation *user = use->getOwner(); - if (auto loadOp = dyn_cast(user)) { - loadOp->setAttr(GeneratedByMakeTensorPtrTAG, - UnitAttr::get(rewriter.getContext())); - } else if (auto storeOp = dyn_cast(user)) { - return; - } else if (auto advanceOp = dyn_cast(user)) { - SmallVector resUses; - for (auto &use : advanceOp->getUses()) - resUses.push_back(&use); - for (auto resUse : resUses) - markLoadUsers(resUse, rewriter); - } else if (auto loopOp = dyn_cast(user)) { - auto initArg = use->get(); - auto iterArg = loopOp.getTiedLoopRegionIterArg(use); - auto resultValue = loopOp.getTiedLoopResult(use); - iterArg.setType(initArg.getType()); - resultValue.setType(initArg.getType()); - for (auto &argUse : iterArg.getUses()) - markLoadUsers(&argUse, rewriter); - for (auto &resUse : resultValue.getUses()) - markLoadUsers(&resUse, rewriter); - } else if (isa(user)) { - return; - } else { - llvm_unreachable( - "[MakeTensorPtrCanonicalizer] tt.make_tensor_ptr's result is " - "not used by load/store/advance op"); - } -} - -LogicalResult -MakeTensorPtrCanonicalizer::matchAndRewrite(triton::MakeTensorPtrOp op, - PatternRewriter &rewriter) const { - auto order = op.getOrder(); - auto orderSize = order.size(); - if (orderSize == 1) { - return rewriter.notifyMatchFailure( - op, "make_tensor_ptr's order has single value."); - } - - bool isPermuted = false; - for (auto [first, second] : llvm::zip(order.slice(0, orderSize - 1), - order.slice(1, orderSize - 1))) { - if (first != second + 1) { - isPermuted = true; - break; - } - } - - auto loc = op.getLoc(); - auto base = op.getBase(); - auto shape = op.getShape(); - auto strides = op.getStrides(); - auto offsets = op.getOffsets(); - auto result = op.getResult(); - SmallVector opUses; - - for (auto &use : result.getUses()) - opUses.push_back(&use); - for (auto use : opUses) - markLoadUsers(use, rewriter); - - if (!isPermuted) { - return rewriter.notifyMatchFailure( - op, "make_tensor_ptr's order is contiguous."); - } - - llvm::SmallVector blkShapeI32; - llvm::SmallVector blkShapeI64; - auto resPtrType = cast(result.getType()); - if (auto resShapedTy = dyn_cast(resPtrType.getPointeeType())) { - auto resBlkShape = resShapedTy.getShape(); - for (auto [i, v] : llvm::enumerate(resBlkShape)) { - auto reverseI = orderSize - 1 - i; - blkShapeI32.push_back(resBlkShape[order[reverseI]]); - blkShapeI64.push_back(resBlkShape[order[reverseI]]); - } - } - - llvm::SmallVector newShape; - llvm::SmallVector newStrides; - llvm::SmallVector newOffsets; - for (int i = orderSize - 1; i >= 0; i--) { - newShape.push_back(shape[order[i]]); - newStrides.push_back(strides[order[i]]); - newOffsets.push_back(offsets[order[i]]); - } - - llvm::SmallVector contiguousOrder; - for (int i = orderSize - 1; i >= 0; i--) - contiguousOrder.push_back(i); - - rewriter.setInsertionPoint(op); - auto newMakeTensorPtrOp = rewriter.create( - loc, base, ValueRange(newShape), ValueRange(newStrides), - ValueRange(newOffsets), blkShapeI32, contiguousOrder); - rewriter.replaceOp(op, newMakeTensorPtrOp); - for (auto use : opUses) - rewriteUserWithNewOrder(use, rewriter, blkShapeI64, loc, order, orderSize); - return success(); -} - -LogicalResult -ReduceSingleCanonicalizer::matchAndRewrite(triton::ReduceOp reduceOp, - PatternRewriter &rewriter) const { - auto srcs = reduceOp.getSrcs(); - bool allSrcSingleElem = true; - for (auto src : srcs) { - auto srcType = cast(src.getType()); - auto srcShape = srcType.getShape(); - int64_t numel = 1; - for (auto s : srcShape) { - numel *= s; - } - if (numel != 1) { - allSrcSingleElem = false; - break; - } - } - - if (!allSrcSingleElem) { - return rewriter.notifyMatchFailure( - reduceOp, "reduce's srcs are not all with single element"); - } - - auto results = reduceOp.getResult(); - auto loc = reduceOp->getLoc(); - auto zero = rewriter - .create( - loc, rewriter.getIndexType(), - rewriter.getIntegerAttr(rewriter.getIndexType(), 0)) - .getResult(); - for (int i = 0; i < srcs.size(); i++) { - auto src = srcs[i]; - auto srcType = cast(src.getType()); - auto srcRank = srcType.getRank(); - auto res = results[i]; - Value extracted; - if (srcRank == 1) { - // vector reduce generates a scalar result - extracted = - rewriter.create(loc, src, zero).getResult(); - } else { - auto srcShape = srcType.getShape(); - auto resType = cast(res.getType()); - auto resShape = resType.getShape(); - auto collapseReassociationIndicesOptional = - getReassociationIndicesForCollapse(srcShape, resShape); - if (!collapseReassociationIndicesOptional.has_value()) { - return rewriter.notifyMatchFailure( - reduceOp, "Failure with getReassociationIndicesForCollapse call"); - } - auto collapseReassociationIndices = - collapseReassociationIndicesOptional.value(); - extracted = rewriter - .create( - loc, src, collapseReassociationIndices) - .getResult(); - } - res.replaceAllUsesWith(extracted); - } - - return success(); -} - -/** - * @brief Wrapper to check if any operand of the given operation traces back to - * triton::LoadOp. - */ -static bool anyOperandFromTritonLoad(Operation *op) { - const int kMaxTraceDepth = 3; - - std::function trace = [&](Value value, int depth) -> bool { - // Base case 1: Reached maximum depth. - if (depth > kMaxTraceDepth) - return false; - - Operation *defOp = value.getDefiningOp(); - // Base case 2: Value is a block argument (no defining op), stop tracing. - if (!defOp) - return false; - - // Base case 3: Found the target operation! - if (dyn_cast(defOp)) - return true; - - // Recursive step: Check all operands of the current defining operation. - // We increment the depth for the next level. - for (Value operand : defOp->getOperands()) { - if (trace(operand, depth + 1)) - return true; - } - return false; - }; - - // Iterate over the operands of the top-level operation (op) - for (Value operand : op->getOperands()) { - // Start tracing from depth 1 (first hop from the current op) - if (trace(operand, 1)) - return true; - } - return false; -} - -LogicalResult -RemfToBasicArithmetic::matchAndRewrite(arith::RemFOp op, - PatternRewriter &rewriter) const { - if (!anyOperandFromTritonLoad(op)) { - return rewriter.notifyMatchFailure( - op, "None of the operands are defined by triton::LoadOp."); - } - Value lhs = op.getLhs(); // %a - Value rhs = op.getRhs(); // %b - Location loc = op.getLoc(); - // Implementation: a - b * floor(a / b) - // %div = arith.divf %a, %b - Value divOp = rewriter.create(loc, lhs, rhs); - // %flr = math.floor %div - Value floorOp = rewriter.create(loc, divOp); - // %mul = arith.mulf %b, %flr - Value mulOp = rewriter.create(loc, rhs, floorOp); - // %result = arith.subf %a, %mul - rewriter.replaceOpWithNewOp(op, lhs, mulOp); - return success(); -} - -LogicalResult -RemSIToBasicArithmetic::matchAndRewrite(arith::RemSIOp op, - PatternRewriter &rewriter) const { - if (!anyOperandFromTritonLoad(op)) { - return rewriter.notifyMatchFailure( - op, "None of the operands are defined by triton::LoadOp."); - } - Value lhs = op.getLhs(); // %a - Value rhs = op.getRhs(); // %b - Location loc = op.getLoc(); - // Implementation: a - b * (a / b) - // %div = arith.divsi %a, %b (Signed truncating integer division) - Value divOp = rewriter.create(loc, lhs, rhs); - // %mul = arith.muli %b, %div - Value mulOp = rewriter.create(loc, rhs, divOp); - // %result = arith.subi %a, %mul - rewriter.replaceOpWithNewOp(op, lhs, mulOp); - return success(); -} - -} // namespace mlir::dicp::trtion_ext diff --git a/compiler/lib/Dialect/TritonStructured/CMakeLists.txt b/compiler/lib/Dialect/TritonStructured/CMakeLists.txt new file mode 100644 index 00000000..f33061b2 --- /dev/null +++ b/compiler/lib/Dialect/TritonStructured/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(IR) diff --git a/compiler/lib/Dialect/TritonStructured/IR/CMakeLists.txt b/compiler/lib/Dialect/TritonStructured/IR/CMakeLists.txt new file mode 100644 index 00000000..27aac38f --- /dev/null +++ b/compiler/lib/Dialect/TritonStructured/IR/CMakeLists.txt @@ -0,0 +1,11 @@ +add_triton_library(TritonStructuredIR + TritonStructuredOps.cpp + TritonStructuredDialect.cpp + + DEPENDS + TritonStructuredTableGen + + LINK_LIBS PUBLIC + TritonIR + MLIRIR + ) diff --git a/compiler/lib/Dialect/TritonStructured/IR/TritonStructuredDialect.cpp b/compiler/lib/Dialect/TritonStructured/IR/TritonStructuredDialect.cpp new file mode 100644 index 00000000..2edf4281 --- /dev/null +++ b/compiler/lib/Dialect/TritonStructured/IR/TritonStructuredDialect.cpp @@ -0,0 +1,22 @@ +#include "dicp/Dialect/TritonStructured/IR/TritonStructuredDialect.h" + +using namespace mlir; +using namespace mlir::tts; + +/// Dialect creation, the instance will be owned by the context. This is the +/// point of registration of custom types and operations for the dialect. +void TritonStructuredDialect::initialize() { + addOperations< +#define GET_OP_LIST +#include "dicp/Dialect/TritonStructured/IR/TritonStructuredOps.cpp.inc" + >(); +} + +//===----------------------------------------------------------------------===// +// TableGen'd op method definitions +//===----------------------------------------------------------------------===// + +#define GET_OP_CLASSES +#include "dicp/Dialect/TritonStructured/IR/TritonStructuredOps.cpp.inc" + +#include "dicp/Dialect/TritonStructured/IR/TritonStructuredDialect.cpp.inc" diff --git a/compiler/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp b/compiler/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp new file mode 100644 index 00000000..da8cf059 --- /dev/null +++ b/compiler/lib/Dialect/TritonStructured/IR/TritonStructuredOps.cpp @@ -0,0 +1,109 @@ +#include "mlir/Bytecode/BytecodeOpInterface.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Support/LogicalResult.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/LogicalResult.h" +#include +#include +#include + +#define GET_OP_CLASSES +#include "dicp/Dialect/TritonStructured/IR/TritonStructuredDialect.h" +using namespace mlir; +using namespace mlir::tts; + +namespace mlir { +namespace tts { + +LogicalResult GetStructuredStateOp::verify() { + auto expectedOffsetAndStrideTypes = + getOffsetAndStrideTypes(getContext(), getInput().getType()); + + if (!expectedOffsetAndStrideTypes.has_value()) { + return failure(); + } + + auto [expectedOffsetTypes, expectedStrideTypes] = + *expectedOffsetAndStrideTypes; + + return success(expectedOffsetTypes.size() == getOffsets().size() && + llvm::equal(expectedOffsetTypes, getOffsets().getTypes()) && + expectedStrideTypes.size() == getStrides().size() && + llvm::equal(expectedStrideTypes, getStrides().getTypes())); +} + +void GetStructuredStateOp::build(OpBuilder &b, OperationState &state, + Value val) { + auto type = val.getType(); + + // Builder cannot fail, so we default to empty offset and stride types. + // The invalid op will be rejected by the verifier later. + auto [offsetTypes, strideTypes] = + getOffsetAndStrideTypes(b.getContext(), type) + .value_or(std::make_pair(SmallVector{}, SmallVector{})); + + build(b, state, val.getType(), offsetTypes, strideTypes, val); +} + +std::optional, SmallVector>> +GetStructuredStateOp::getOffsetAndStrideTypes(MLIRContext *context, Type type) { + auto sizes = getOffsetAndStrideSegmentSizes(type); + if (!sizes.has_value()) { + return std::nullopt; + } + return std::make_pair( + SmallVector(sizes->first, IndexType::get(context)), + SmallVector(sizes->second, IndexType::get(context))); +} + +std::optional> +GetStructuredStateOp::getOffsetAndStrideSegmentSizes(Type type) { + int32_t offsetSegmentSize = 0; + int32_t strideSegmentSize = 0; + + if (auto tensorType = llvm::dyn_cast(type)) { + if (tensorType.getElementType().isIntOrIndex()) { + // Tensors of offsets + // Important note: + // We only care about tensor of index / int (in addition to pointer type) + // because only values of int and index type can potentially be part of a + // pointer arithmetic sequence. + offsetSegmentSize = strideSegmentSize = tensorType.getRank(); + } else if (auto ptrType = + dyn_cast(tensorType.getElementType())) { + // Unstructured pointers (tensor>) + // Each tensor of rank k gets k values for its offsets and k values for + // its strides, all of which has Index type. + offsetSegmentSize = strideSegmentSize = tensorType.getRank(); + } + } + // Block pointers (!tt.ptr> or !tt.ptr) + else if (auto ptrType = llvm::dyn_cast(type)) { + if (auto tensorType = + llvm::dyn_cast(ptrType.getPointeeType())) { + // Each tensor of rank k gets k values for its offsets and k values for + // its strides, all of which has Index type. + offsetSegmentSize = strideSegmentSize = tensorType.getRank(); + } else { + // The only relevant state that can be updated in loops for scalar + // pointers are offset. No need to include stride here. + offsetSegmentSize = 1; + } + } else { + return std::nullopt; + } + + return std::make_pair(offsetSegmentSize, strideSegmentSize); +} + +} // namespace tts +} // namespace mlir diff --git a/compiler/lib/Conversion/DiscreteMaskAccessConversion/CMakeLists.txt b/compiler/lib/DiscreteMaskAccessConversion/CMakeLists.txt similarity index 81% rename from compiler/lib/Conversion/DiscreteMaskAccessConversion/CMakeLists.txt rename to compiler/lib/DiscreteMaskAccessConversion/CMakeLists.txt index a5b764ae..4acb2e7f 100644 --- a/compiler/lib/Conversion/DiscreteMaskAccessConversion/CMakeLists.txt +++ b/compiler/lib/DiscreteMaskAccessConversion/CMakeLists.txt @@ -1,15 +1,14 @@ add_triton_library(DiscreteMaskAccessConversion DiscreteMaskAccessConversionPass.cpp - MaskAnalysis.cpp DEPENDS DiscreteMaskAccessConversionPassIncGen - LINK_LIBS + LINK_LIBS + BiShengIRHIVMDialect MLIRIR MLIRPass MLIRTransforms MLIRSupport TritonIR - MLIRAnalysis ) diff --git a/compiler/lib/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.cpp b/compiler/lib/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.cpp new file mode 100644 index 00000000..127d04b6 --- /dev/null +++ b/compiler/lib/DiscreteMaskAccessConversion/DiscreteMaskAccessConversionPass.cpp @@ -0,0 +1,366 @@ + + +#include "dicp/DiscreteMaskAccessConversion/Passes.h" +#include "dicp/Utils/Utils.h" + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" +#include "dicp/TritonToLinalg/MaskAnalysis.h" +#include "dicp/TritonToStructured/MemOpConverter.h" +#include "dicp/TritonToUnstructure/OffsetAnalysis.h" +#include "mlir/IR/Attributes.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/Passes.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/LogicalResult.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_DISCRETEMASKACCESSCONVERSION +#include "dicp/DiscreteMaskAccessConversion/Passes.h.inc" +} // namespace triton +} // namespace mlir + +#define DEBUG_TYPE "discrete-mask-access-conversion" + +using namespace mlir; +using namespace hivm; + +// File-scope flags set by DiscreteMaskAccessConversionPass::runOnOperation() +// before pattern application, so that OpRewritePattern subclasses can read +// them. +static bool compileOn91095Flag = false; +static bool forceSimtTemplateFlag = false; +static bool enableSyncBlockLockFlag = true; + +LogicalResult isDiscreteMask(Operation *op, Value mask, + PatternRewriter &rewriter) { + if (!mask) + return failure(); + + MaskState mstate; + auto isContMask = mstate.parse(mask, op->getLoc(), rewriter); + if (!isContMask.failed()) { + mstate.eraseInsertedOps(op, rewriter); + return failure(); + } + return success(); +} + +// Recursively collect all leaf operands of a nested arith::AndIOp tree. +// This function also normalizes masks by distributing broadcast over andi +// broadcast(andi(a, b)) = andi(broadcast(a), broadcast(b)) +// so that inner AND operands nested inside a broadcast are still reachable. +static void collectAndLeaves(Value mask, SmallVectorImpl &leaves, + Location loc, PatternRewriter &rewriter) { + if (auto andOp = mask.getDefiningOp()) { + collectAndLeaves(andOp.getLhs(), leaves, loc, rewriter); + collectAndLeaves(andOp.getRhs(), leaves, loc, rewriter); + } else if (auto broadcastOp = mask.getDefiningOp()) { + // Distribute broadcast over andi so we can inspect each factor separately. + if (auto innerAnd = broadcastOp.getSrc().getDefiningOp()) { + Type dstType = mask.getType(); + Value broadcastA = + rewriter.create(loc, dstType, innerAnd.getLhs()) + .getResult(); + Value broadcastB = + rewriter.create(loc, dstType, innerAnd.getRhs()) + .getResult(); + collectAndLeaves(broadcastA, leaves, loc, rewriter); + collectAndLeaves(broadcastB, leaves, loc, rewriter); + } else { + leaves.push_back(mask); + } + } else { + leaves.push_back(mask); + } +} + +struct MaskDecomposition { + // AND of all leaves that MaskState::parse() can analyze as a rectangle mask. + // nullptr when no such leaves exist. + Value contMask; + // AND of all leaves that MaskState::parse() cannot analyze + // (discrete/runtime). nullptr when no such leaves exist. + Value discMask; +}; + +// Decompose an AND-tree mask into its continuous and discrete leaf components +// so that we can use contMask to bound GM accesses while discMask still drives +// the per-element selection. +static MaskDecomposition decomposeAndMask(Operation *op, Value mask, + const Location &loc, + PatternRewriter &rewriter) { + SmallVector leaves; + collectAndLeaves(mask, leaves, loc, rewriter); + + SmallVector contLeaves; + SmallVector discLeaves; + + for (Value leaf : leaves) { + MaskState st; + if (st.parse(leaf, loc, rewriter).succeeded()) { + if (st.isMask()) + contLeaves.push_back(leaf); + else + discLeaves.push_back(leaf); + } else { + discLeaves.push_back(leaf); + } + } + + Value contMask = nullptr; + for (Value v : contLeaves) + contMask = + contMask ? rewriter.create(loc, contMask, v).getResult() + : v; + + Value discMask = nullptr; + for (Value v : discLeaves) + discMask = + discMask ? rewriter.create(loc, discMask, v).getResult() + : v; + + return {contMask, discMask}; +} + +struct DiscreteMaskStoreConversion : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::StoreOp op, + PatternRewriter &rewriter) const final { + auto mask = op.getMask(); + auto loc = op.getLoc(); + auto dst = op.getPtr(); + auto src = op.getValue(); + + if (failed(isDiscreteMask(op, mask, rewriter))) + return failure(); + + if (compileOn91095Flag && forceSimtTemplateFlag) { + llvm::DenseMap offsetMap; + mlir::triton::parse(dst, loc, rewriter, offsetMap); + + if (!offsetMap.contains(dst)) { + return failure(); + } + + auto &info = offsetMap[dst]; + Value basePtr = info.getPtr(); + Value totalOffset = info.getOffset(); + + rewriter.create(loc, basePtr, totalOffset, + src, mask); + rewriter.eraseOp(op); + return success(); + } + + // When mask = contMask & discMask, use contMask to bound GM accesses and + // discMask to select the final per-element value. This prevents the + // unguarded full-load from reading past the tail-block boundary. + auto [contMask, discMask] = decomposeAndMask(op, mask, loc, rewriter); + if (contMask && discMask) { + // insert sync_block_lock + auto lockVar = MemOpConverter::createSyncBlockLockVar(rewriter, loc); + if (enableSyncBlockLockFlag) { + rewriter.create(loc, lockVar); + } + auto safeLoad = rewriter.create( + loc, dst, contMask, op.getCache(), op.getEvict(), false); + auto selOp = rewriter.create(loc, discMask, src, + safeLoad.getResult()); + auto newStore = rewriter.create( + loc, dst, selOp, contMask, op.getCache(), op.getEvict()); + newStore->setAttr(ConverterUtils::discreteMaskAttrName, + UnitAttr::get(rewriter.getContext())); + if (enableSyncBlockLockFlag) { + rewriter.create(loc, lockVar); + } + rewriter.replaceOp(op, newStore); + return success(); + } + + // Fallback: original full load + select (contMask absent, pure discrete). + // insert sync_block_lock + auto lockVar = MemOpConverter::createSyncBlockLockVar(rewriter, loc); + if (enableSyncBlockLockFlag) { + rewriter.create(loc, lockVar); + } + auto loadFromDstOp = rewriter.create( + loc, dst, op.getCache(), op.getEvict(), false); + auto selOp = rewriter.create(loc, mask, src, + loadFromDstOp.getResult()); + auto newStore = rewriter.create( + loc, dst, selOp, op.getCache(), op.getEvict()); + newStore->setAttr(ConverterUtils::discreteMaskAttrName, + UnitAttr::get(rewriter.getContext())); + if (enableSyncBlockLockFlag) { + rewriter.create(loc, lockVar); + } + rewriter.replaceOp(op, newStore); + return success(); + } +}; + +struct DiscreteMaskLoadConversion : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::LoadOp op, + PatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + auto other = op.getOther(); + auto mask = op.getMask(); + auto ptr = op.getPtr(); + + if (failed(isDiscreteMask(op, mask, rewriter))) + return failure(); + + const std::string isDiscreteMaskTag = "is_discrete_mask"; + op->setAttr(isDiscreteMaskTag, rewriter.getUnitAttr()); + + if (compileOn91095Flag && forceSimtTemplateFlag) + return failure(); + + // When mask = contMask & discMask, load only the safe range defined by + // contMask and use discMask for the per-element select, avoiding OOB reads. + auto [contMask, discMask] = decomposeAndMask(op, mask, loc, rewriter); + if (contMask && discMask) { + if (!other) { + FailureOr constant = specializeTypelessValueToConstant( + TypelessValue::Zero, ptr.getType(), loc, rewriter); + if (failed(constant)) { + llvm_unreachable("Unsupported type for constant creation"); + } + other = *constant; + } + auto safeLoad = rewriter.create( + loc, ptr, contMask, op.getCache(), op.getEvict(), op.getIsVolatile()); + // Use combined mask to select the result, avoid the uninitialized memory + // access. + auto combinedMask = + rewriter.create(loc, contMask, discMask); + auto discreteMaskOp = rewriter.create( + loc, combinedMask, safeLoad.getResult(), other); + rewriter.replaceOp(op, discreteMaskOp); + return success(); + } + + // Fallback: original full load + select (contMask absent, pure discrete). + if (!other) { + FailureOr constant = specializeTypelessValueToConstant( + TypelessValue::Zero, ptr.getType(), loc, rewriter); + if (failed(constant)) + llvm_unreachable("Unsupported type for constant creation"); + other = *constant; + } + + auto newLoadOp = rewriter.create( + loc, ptr, op.getCache(), op.getEvict(), op.getIsVolatile()); + auto discreteMaskOp = + rewriter.create(loc, mask, newLoadOp, other); + rewriter.replaceOp(op, discreteMaskOp); + return success(); + } +}; + +struct DiscreteMaskAtomicConversion + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(mlir::triton::AtomicRMWOp op, + PatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + auto ptr = op.getPtr(); + auto src = op.getVal(); + auto mask = op.getMask(); + RMWOp rmwOp = op.getAtomicRmwOp(); + + if (failed(isDiscreteMask(op, mask, rewriter))) + return failure(); + + const std::map initMap = { + {RMWOp::FADD, TypelessValue::Zero}, + {RMWOp::ADD, TypelessValue::Zero}, + {RMWOp::UMAX, TypelessValue::Zero}, + {RMWOp::OR, TypelessValue::Zero}, + {RMWOp::MIN, TypelessValue::Max}, + {RMWOp::UMIN, TypelessValue::Max}, + {RMWOp::AND, TypelessValue::Max}, + {RMWOp::MAX, TypelessValue::Min}, + {RMWOp::XOR, TypelessValue::Zero}, + {RMWOp::XCHG, TypelessValue::Undefined}, + }; + assert(initMap.find(rmwOp) != initMap.end()); + auto typelessVal = initMap.at(rmwOp); + if (typelessVal == TypelessValue::Undefined) { + // Undefined default value atomic op will be decomposed in BiShengIR + op->setAttr(ConverterUtils::discreteMaskAttrName, + UnitAttr::get(rewriter.getContext())); + return failure(); + } + + FailureOr fill = specializeTypelessValueToConstant( + typelessVal, src.getType(), loc, rewriter); + if (failed(fill)) + op->emitError("Unsupported atomic operation."); + + auto maskedValue = rewriter.create(loc, mask, src, *fill); + auto newAtomicOp = rewriter.create( + loc, src.getType(), rmwOp, ptr, maskedValue, mlir::Value(), op.getSem(), + op.getScope()); + rewriter.replaceOp(op, newAtomicOp); + return success(); + } +}; + +DiscreteMaskAccessConversionPass::DiscreteMaskAccessConversionPass( + const DiscreteMaskAccessConversionOptions &options) + : DiscreteMaskAccessConversionBase(options) {} + +void DiscreteMaskAccessConversionPass::runOnOperation() { + compileOn91095Flag = this->compileOn91095; + forceSimtTemplateFlag = this->forceSimtTemplate; + enableSyncBlockLockFlag = this->enableSyncBlockLock; + auto moduleOp = getOperation(); + + RewritePatternSet patterns(&getContext()); + patterns.add(patterns.getContext()); + if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { + moduleOp->emitError("failed to apply discrete mask access patterns"); + signalPassFailure(); + } + + // Clean up dead analysis ops left behind by MaskState::parse(). + // These are trivially-dead auxiliary ops (constants, arithmetic) with no + // users that parse() creates as side effects of mask analysis. + PassManager pm(&getContext(), moduleOp.getOperationName()); + pm.addPass(createCSEPass()); + pm.addPass(createCanonicalizerPass()); + if (failed(runPipeline(pm, getOperation()))) { + moduleOp->emitWarning( + "DiscreteMaskAccessConversion: dead-code cleanup failed"); + } + + LLVM_DEBUG({ + llvm::dbgs() << "==============================================\n"; + llvm::dbgs() << "After DiscreteMaskAccessConversionPass:\n" << moduleOp; + llvm::dbgs() << "\n==============================================\n"; + }); +} + +void DiscreteMaskAccessConversionPass::getDependentDialects( + DialectRegistry ®istry) const { + registry + .insert(); +} + +std::unique_ptr> +mlir::triton::createDiscreteMaskAccessConversionPass( + const DiscreteMaskAccessConversionOptions &options) { + return std::make_unique(options); +} diff --git a/compiler/lib/ExecutionEngine/compile_lib.cpp b/compiler/lib/ExecutionEngine/compile_lib.cpp deleted file mode 100644 index e69de29b..00000000 diff --git a/compiler/lib/Pipeline/pipelines.cpp b/compiler/lib/Pipeline/pipelines.cpp deleted file mode 100644 index e69de29b..00000000 diff --git a/compiler/lib/TritonAffinityOpt/CMakeLists.txt b/compiler/lib/TritonAffinityOpt/CMakeLists.txt new file mode 100644 index 00000000..2f49f3f0 --- /dev/null +++ b/compiler/lib/TritonAffinityOpt/CMakeLists.txt @@ -0,0 +1,19 @@ +add_triton_library(TritonAffinityOpt + DAGSSBuffer.cpp + DAG.cpp + DAGSync.cpp + DAGScope.cpp + + DEPENDS + TritonAffinityOptConversionPassIncGen + + LINK_LIBS + BiShengIRHIVMDialect + BiShengIRScopeDialect + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport + TritonIR + MLIRSCFDialect +) \ No newline at end of file diff --git a/compiler/lib/TritonAffinityOpt/DAG.cpp b/compiler/lib/TritonAffinityOpt/DAG.cpp new file mode 100644 index 00000000..c736523f --- /dev/null +++ b/compiler/lib/TritonAffinityOpt/DAG.cpp @@ -0,0 +1,518 @@ +#include "dicp/TritonAffinityOpt/DAG.h" +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/IR/TypeUtilities.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Interfaces/ControlFlowInterfaces.h" +#include "mlir/Interfaces/CopyOpInterface.h" +#include "mlir/Interfaces/LoopLikeInterface.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/raw_ostream.h" +#include +#include +#include +#include +#include +#include + +namespace mlir { +namespace AffinityDAG { + +const auto printFlags = + OpPrintingFlags().enableDebugInfo(true, true).skipRegions(); + +const char *literalCoreType(CoreType ct) { + switch (ct) { + case VECTOR_ONLY: + return "VECTOR_ONLY"; + case CUBE_ONLY: + return "CUBE_ONLY"; + case CUBE_AND_VECTOR: + return "CUBE_AND_VECTOR"; + case UNDETERMINED: + return "UNDETERMINED"; + } + return "Unknown"; +} + +bool opIsScf(Operation *op) { + if (!llvm::isa(op->getDialect())) + return false; + return true; +} + +Graph::Graph(Block *block, Graph *parent, OpMap opMap, ValueMap valueMap, + bool inheritParent) + : block(block), parent(parent), opMap(opMap), valueMap(valueMap) { + + if (parent && inheritParent) { + if (!this->opMap) { + this->opMap = parent->opMap; + } + + if (!this->valueMap) { + this->valueMap = parent->valueMap; + } + } + + if (!this->opMap) { + this->opMap = std::make_shared(); + } + + if (!this->valueMap) { + this->valueMap = std::make_shared(); + } + + for (auto blockArg : block->getArguments()) { + (*this->valueMap)[blockArg] = std::make_unique(blockArg); + blockArgs.push_back((*this->valueMap)[blockArg].get()); + } + + for (auto &opRef : block->getOperations()) { + opCount += 1; + auto op = &opRef; + auto opNodeUnique = std::make_unique(op, this); + auto opNode = opNodeUnique.get(); + (*this->opMap)[op] = std::move(opNodeUnique); + + if (block->mightHaveTerminator() && op == block->getTerminator()) { + terminator = opNode; + } + + for (auto &subgraph : opNode->subgraphs) { + opCount += subgraph.opCount; + } + } +}; + +bool valueIsScalar(Value value) { + auto type = value.getType(); + + if (type.isIntOrIndexOrFloat()) { + return true; + } + + if (auto tensorType = llvm::dyn_cast(type)) { + return tensorType.getRank() == 0; + } + + if (auto _ = llvm::dyn_cast(type)) { + return true; + } + + return false; +} + +bool valueIsTensorOfPtr(Value value) { + auto type = value.getType(); + if (auto tensorType = llvm::dyn_cast(type)) { + auto elementType = tensorType.getElementType(); + if (llvm::isa(elementType)) { + return true; + } + } + + return false; +} + +OpAbility OpNode::canRunOn() const { + if (opIsScf(op)) { + return OpAbility::CUBE_AND_VECTOR; + } + return llvm::TypeSwitch(op) + .Case([](auto) { return OpAbility::CUBE_ONLY; }) + .Case([](auto) { return OpAbility::CUBE_AND_VECTOR; }) + .Case([](arith::SelectOp op) { + // when cond is vector, selectOp should be vector, otherwise scalar + return (valueIsScalar(op.getCondition()) ? OpAbility::CUBE_AND_VECTOR + : OpAbility::PREFER_VECTOR); + }) + .Default([](Operation *op) { + auto isVector = false; + for (auto operand : op->getOperands()) { + if (!valueIsScalar(operand)) { + // if (valueIsTensorOfPtr(operand)) { + // return SCALAR; + // } + isVector = true; + } + } + + for (auto result : op->getResults()) { + if (!valueIsScalar(result)) { + // if (valueIsTensorOfPtr(result)) { + // return SCALAR; + // } + isVector = true; + } + } + + if (isVector) { + return OpAbility::PREFER_VECTOR; + } + + return OpAbility::CUBE_AND_VECTOR; + }); +} + +OpNode::OpNode(Operation *op, Graph *graph) : Node(Node::NK_Op), op(op) { + if (op == nullptr) { + return; + } + + llvm::outs() << op << "\n"; + + auto &valueMap = *graph->valueMap.get(); + auto &opMap = *graph->opMap.get(); + for (const auto operand : op->getOperands()) { + auto valueNode = valueMap.at(operand).get(); + valueNode->outputs.push_back(this); + inputs.push_back(valueNode); + } + + for (const auto &result : op->getResults()) { + auto valueNodeUnique = std::make_unique(result); + auto valueNode = valueNodeUnique.get(); + valueMap[result] = std::move(valueNodeUnique); + valueNode->source = this; + outputs.push_back(valueNode); + } + + // if (!op->hasTrait()) { + // llvm::dbgs() << "Not building subgraph because op is not SingleBlock: " + // << op << '\n'; return; + // } + + if (auto branchOp = llvm::dyn_cast(op)) { + + OpNode *terminator = nullptr; + llvm::SmallVector, 2> validRegions; + + for (auto ®ion : branchOp->getRegions()) { + if (region.getBlocks().empty()) + continue; + subgraphs.emplace_back(®ion.getBlocks().front(), graph); + validRegions.emplace_back(region, subgraphs.back()); + } + + for (auto [region, subgraph] : validRegions) { + SmallVector succRegions; + + branchOp.getSuccessorRegions(region, succRegions); + if (auto currTerminator = dyn_cast( + subgraph.terminator->op)) { + for (auto &succ : succRegions) { + auto forwardedVal = currTerminator.getSuccessorOperands(succ); + if (succ.isParent()) { + // Step1: first yield to parent -> results: double direction + if (!terminator && subgraph.terminator) { + terminator = subgraph.terminator; + for (auto [forwardedVal, resultNode] : + llvm::zip_equal(forwardedVal, outputs)) { + auto resultValueNode = llvm::dyn_cast(resultNode); + assert(resultValueNode && + "Output of a OpNode should be ValueNode!"); + auto forwardedNode = valueMap[forwardedVal].get(); + resultValueNode->source = forwardedNode; + forwardedNode->outputs.push_back(resultNode); + } + } + + } else { + // Step2: Region terminator -> Succ Operands + auto succRegion = succ.getSuccessor(); + + for (auto [operand, succInput] : + llvm::zip_equal(forwardedVal, succ.getSuccessorInputs())) { + auto forwardedNode = valueMap[operand].get(); + auto succNode = valueMap[succInput].get(); + forwardedNode->outputs.push_back(succNode); + succNode->source = forwardedNode; + } + } + } + } + } + + if (auto loopOp = llvm::dyn_cast(op)) { + // Step3: inits->iter_args (single directional) (should be handled in step + // 2: ) last terminator -> iter_args (bidirectional) + for (auto [init, iterArgVal] : + llvm::zip_equal(loopOp.getInits(), loopOp.getRegionIterArgs())) { + auto &initNode = valueMap[init]; + auto &iterArgNode = valueMap[iterArgVal]; + initNode->outputs.push_back(iterArgNode.get()); + } + // for(auto [init, iterArgVal, yieldNode] : + // llvm::zip_equal(loopOp.getInits(), loopOp.getRegionIterArgs(), + // terminator->outputs)) { + // auto& initNode = valueMap[init]; + // auto& iterArgNode = valueMap[iterArgVal]; + // initNode->outputs.push_back(iterArgNode.get()); + // yieldNode->outputs.push_back(iterArgNode.get()); + // iterArgNode->source = yieldNode; + // } + } + } +} + +// llvm::SmallVector getWriteOperandPriority(OpNode* op) { + +// llvm::SmallVector result(op->getInputs()); + +// auto getPriority = [](ValueNode* node) { +// auto typ = getElementTypeOrSelf(node->value); +// if (typ.isInteger(1)) { +// return 2; +// } +// if (llvm::isa(typ)) { +// return 1; +// } +// return 0; +// }; + +// std::stable_sort(result.begin(), result.end(), [&](ValueNode* a, ValueNode* +// b) { +// return getPriority(a) < getPriority(b); +// }); + +// return result; +// } + +ValueNode *getWriteDataSource(OpNode *op) { + auto inputRange = op->getInputs(); + for (auto node : inputRange.drop_front()) { + auto typ = getElementTypeOrSelf(node->value); + if (!typ.isInteger(1)) { + return node; + } + }; + + return nullptr; +} + +enum class MemPolicy { NONE, READ, WRITE }; + +CoreType Node::absorbCommon() { + + auto sourceNode = getSourceOpNode(); + auto op = sourceNode ? sourceNode->op : nullptr; + + if (!sourceNode || !op) { + CoreType newCoreType = isOnPrivate; + for (auto output : outputs) { + newCoreType = newCoreType | output->isOn(); + isUpstreamOfCubeMem = isUpstreamOfCubeMem || output->isUpstreamOfCubeMem; + } + return newCoreType; + } + + CoreType newCoreType = sourceNode->isOn(); + + OpAbility ability = sourceNode->canRunOn(); + + if (ability == OpAbility::CUBE_ONLY) { + return CUBE_ONLY; + } + + auto memIface = llvm::dyn_cast(op); + auto memPolicy = MemPolicy::NONE; + + if (memIface) { + // Possible improvements: Determine the policy to use based on shapes, + // inputs and outputs, etc + if (memIface.hasEffect()) { + memPolicy = MemPolicy::WRITE; + } else if (memIface.hasEffect()) { + memPolicy = MemPolicy::READ; + } + } + + if (memPolicy == MemPolicy::WRITE) { + if (auto data = getWriteDataSource(sourceNode)) { + auto currCt = data->isOn(); + if (exactlyOneType(currCt)) { + if (currCt == CUBE_ONLY) { + isUpstreamOfCubeMem = true; + } + return currCt; + } + } + + // data is not cube_only + return VECTOR_ONLY; + } + + for (auto output : outputs) { + switch (output->isOn()) { + case CUBE_AND_VECTOR: + newCoreType = newCoreType | VECTOR_ONLY; + // not breaking the switch because we need to handle cube + case CUBE_ONLY: + if (ability != OpAbility::PREFER_VECTOR || output->isUpstreamOfCubeMem || + memPolicy == MemPolicy::READ) { + isUpstreamOfCubeMem = + (isUpstreamOfCubeMem || output->isUpstreamOfCubeMem || + memPolicy == MemPolicy::READ); + newCoreType = newCoreType | CUBE_ONLY; + } + break; + case VECTOR_ONLY: + newCoreType = newCoreType | VECTOR_ONLY; + default: // UNDETERMINED, skip + break; + }; + } + + return newCoreType; +} + +CoreType OpNode::absorbImpl() { + if (opIsScf(op)) { + return CUBE_AND_VECTOR; + } + + auto newCoreType = absorbCommon(); + + // if (canRunOn() == OpAbility::CUBE_AND_VECTOR) { + // for (auto input : inputs) { + // newCoreType = newCoreType | input->isOn(); + // } + // } + + return newCoreType; +} + +CoreType ValueNode::absorbImpl() { return absorbCommon(); } + +std::unique_ptr Graph::fromMultiBlockFunc(triton::FuncOp funcOp) { + + auto dummyBlock = new Block(); + auto dummyGraph = std::make_unique(dummyBlock); + auto dummyNode = std::make_unique(nullptr, dummyGraph.get()); + size_t opCount = 0; + + for (auto &block : funcOp.getBody()) { + auto &subgraph = + dummyNode->subgraphs.emplace_back(&block, dummyGraph.get()); + opCount += subgraph.opCount; + } + + auto &opMap = *dummyGraph->opMap.get(); + auto &valueMap = *dummyGraph->valueMap.get(); + + llvm::SmallVector nodes; + nodes.reserve(opMap.size() + valueMap.size()); + + for (auto &[_, node] : opMap) { + if (node.get()) + nodes.push_back(node.get()); + } + + for (auto &[_, node] : valueMap) { + if (node.get()) + nodes.push_back(node.get()); + } + + auto diffuse = [&]() { + // Not sure if determinism is required + llvm::SmallSetVector worklist(nodes.begin(), nodes.end()); + + size_t threshold = worklist.size() * 5; + + for (size_t i = 0; i < threshold; i++) { + if (worklist.empty()) { + break; + } + + auto node = worklist.pop_back_val(); + + if (node->absorb()) { + auto affected = node->getAffected(); + worklist.insert(affected.begin(), affected.end()); + } + } + }; + + diffuse(); + + for (auto node : nodes) { + if (node->isOn() == UNDETERMINED) { + node->isOnPrivate = VECTOR_ONLY; + } + } + + diffuse(); + + OpPrintingFlags flags; + flags.skipRegions(); + + for (auto [idx, node] : llvm::enumerate(nodes)) { + llvm::TypeSwitch(node) + .Case([&, idx = idx](OpNode *node) { + if (node->op) { + llvm::dbgs() << llvm::formatv( + "\n\n====== OpNode on: {1} @ {0} ======\n", node->op, + literalCoreType(node->isOn())); + node->op->print(llvm::dbgs(), flags); + llvm::dbgs() << "\nAbility: " + << literalCoreType(toCoreType(node->canRunOn())); + llvm::dbgs() << llvm::formatv("\n====== {0} ======\n", node->op); + } + }) + .Case([&, idx = idx](ValueNode *node) { + if (node->value) { + llvm::dbgs() << llvm::formatv( + "\n\n====== ValueNode on {1} @ {0} ======\n", node->value, + literalCoreType(node->isOn())); + node->value.print(llvm::dbgs(), flags); + llvm::dbgs() << llvm::formatv("\n====== {0} ======\n", node->value); + } + }); + // if (auto opNode = llvm::dyn_cast(node)) { + // if (auto forOp = llvm::dyn_cast_if_present(opNode->op)) { + // llvm::dbgs() << "\n==== ForOp ====\n"; + // llvm::dbgs() << forOp << "\n"; + // llvm::dbgs() << "\n---- IterArgs ----\n"; + // for(auto iterArg : forOp.getRegionIterArgs()) { + // auto& valueNode = valueMap[iterArg]; + // llvm::dbgs() << llvm::formatv( + // "{0}: {1} upstream: {2} definingOp: {3} \n", + // iterArg.getArgNumber(), + // literalCoreType(valueNode->isOn()), + // literalCoreType(valueNode->source->isOn()), + // valueNode->getSourceOp()->op + // ); + // } + // llvm::dbgs() << "\n---- Results ----\n"; + // for(auto result : forOp.getResults()) { + // llvm::dbgs() << result.getResultNumber() << ' ' << + // literalCoreType(valueMap[result]->isOn()) << '\n'; + // } + // } + // } + } + + return dummyGraph; +}; + +} // namespace AffinityDAG +} // namespace mlir diff --git a/compiler/lib/TritonAffinityOpt/DAGSSBuffer.cpp b/compiler/lib/TritonAffinityOpt/DAGSSBuffer.cpp new file mode 100644 index 00000000..66fd75d3 --- /dev/null +++ b/compiler/lib/TritonAffinityOpt/DAGSSBuffer.cpp @@ -0,0 +1,5584 @@ + + +#include "dicp/TritonAffinityOpt/Passes.h" + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/HIVM/IR/HIVMImpl.h" +#include "bishengir/Dialect/HIVM/IR/HIVMInterfaces.h" +#include "bishengir/Dialect/HIVM/Transforms/Passes.h" +#include "bishengir/Dialect/HIVM/Utils/Utils.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h" +#include "mlir/Dialect/Linalg/Utils/Utils.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "dicp/Utils/Utils.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include + +// #include "mlir/Pass/Pass.h" +// #include "mlir/Pass/PassManager.h" + +// #include "mlir/Transforms/Canonicalizer.h" +// #include "mlir/Support/LogicalResult.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_DAGSSBUFFER +#include "dicp/TritonAffinityOpt/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace hivm; + +namespace { +struct DAGSSBufferPass + : public mlir::triton::impl::DAGSSBufferBase { + void runOnOperation() override; + void getDependentDialects(DialectRegistry ®istry) const override { + registry.insert(); + registry.insert(); + } +}; +} // namespace + +void ControlSsbufV2(ModuleOp module) { + mlir::OpBuilder builder(module.getContext()); + // 用于记录已经处理过的scope.scope操作 + llvm::DenseSet processedScopes; + + auto aiCAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + int cubeControlIndex = 15; + int vectorControlIndex = 14; + + llvm::DenseSet processedScopes2; + module->walk([&](SyncBlockWaitOp op) { + auto pipeS = hivm::PipeAttr::get(op->getContext(), hivm::PIPE::PIPE_S); + + // 向上查找父scope.scope操作 + mlir::Operation *parentOp = op->getParentOp(); + mlir::Operation *scopeOp = nullptr; + mlir::Operation *forOp = nullptr; + + // 向上遍历查找scope.scope操作 + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + parentOp = op->getParentOp(); + while (parentOp) { + if (dyn_cast(parentOp)) { + forOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + // 如果没有找到scope.scope操作，则跳过 + if (!scopeOp) { + return; + } + if (!forOp) { + return; + } + + // 如果该scope已经处理过，则跳过 + if (processedScopes2.count(forOp) > 0) + return; + + // 标记该scope为已处理 + processedScopes2.insert(forOp); + }); + bool firstSet = true; + bool firstWait = true; + for (auto forOp : processedScopes2) { + mlir::Operation *parentOp = forOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + + // 向上遍历查找scope.scope操作 + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + bool isAIC = false; + // 1. 先检查操作是否有这个属性 + + if (scopeOp->hasAttr("hivm.tcore_type")) { + auto attr = scopeOp->getAttr("hivm.tcore_type"); + if (attr == aiCAttr) { + isAIC = true; + } + } + + if (isAIC) { + // 在for循环的开头插入代码 + builder.setInsertionPoint(scopeOp); + // %ssb_ready_addr = llvm.mlir.constant(0 : i64) : i64 + auto i64Type = builder.getIntegerType(64); + auto i32Type = builder.getIntegerType(32); + + builder.setInsertionPointToStart(&forOp->getRegion(0).front()); + // %ssb_ready_addr = llvm.mlir.constant(0 : i64) : i64 + // add sync_block_wait + auto coreAttr = + hivm::TCoreTypeAttr::get(module.getContext(), hivm::TCoreType::CUBE); + auto setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto flagId = + builder.getIntegerAttr(builder.getI64Type(), vectorControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + + // 在循环末尾（yield之前）插入代码 + auto &loopBody = forOp->getRegion(0).front(); + // 找到循环体的terminator（应该是yield操作） + auto *terminator = loopBody.getTerminator(); + builder.setInsertionPoint(terminator); + + // add sync_block_set + coreAttr = + hivm::TCoreTypeAttr::get(module.getContext(), hivm::TCoreType::CUBE); + setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + flagId = builder.getIntegerAttr(builder.getI64Type(), cubeControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + + if (firstWait) { + auto &scopeBlock = scopeOp->getRegion(0).front(); + auto *scope_terminator = scopeBlock.getTerminator(); + builder.setInsertionPoint(scope_terminator); + // add sync_block_wait + coreAttr = hivm::TCoreTypeAttr::get(module.getContext(), + hivm::TCoreType::CUBE); + setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + flagId = + builder.getIntegerAttr(builder.getI64Type(), vectorControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + firstWait = false; + } + } else { + // 1. 在scopeop的开头插入代码 + // 假设scopeOp是一个具有区域的操作，我们获取其第一个块 + if (firstSet) { + auto &scopeBlock = scopeOp->getRegion(0).front(); + builder.setInsertionPointToStart(&scopeBlock); + + // add sync_block_wait + auto coreAttr = hivm::TCoreTypeAttr::get(module.getContext(), + hivm::TCoreType::VECTOR); + auto setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto flagId = + builder.getIntegerAttr(builder.getI64Type(), vectorControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + firstSet = false; + } + + auto i64Type = builder.getIntegerType(64); + auto i32Type = builder.getIntegerType(32); + + // 创建需要的常量 + auto c32ConstAttr = mlir::IntegerAttr::get(i64Type, 32); + auto c32ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c32ConstAttr); + + auto c0i64ConstAttr = mlir::IntegerAttr::get(i64Type, 0); + auto c0i64ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c0i64ConstAttr); + + auto c0i32ConstAttr = mlir::IntegerAttr::get(i32Type, 0); + auto c0i32ConstOp = builder.create( + scopeOp->getLoc(), i32Type, c0i32ConstAttr); + + auto c1i32ConstAttr = mlir::IntegerAttr::get(i32Type, 1); + auto c1i32ConstOp = builder.create( + scopeOp->getLoc(), i32Type, c1i32ConstAttr); + + // %sub_id = hivm.hir.get_sub_block_idx -> i64 + // 这里假设有一个getSubBlockIdxOp操作 + auto subIdOp = + builder.create(scopeOp->getLoc(), i64Type); + + // %ssb_addr_offset = arith.muli %sub_id, %c32_i64 : i64 + auto ssbAddrOffsetOp = builder.create( + scopeOp->getLoc(), subIdOp.getResult(), c32ConstOp.getResult()); + + // %ssb_addr = arith.addi %ssb_addr_offset, %c32_i64 : i64 + auto ssbAddrOp = builder.create( + scopeOp->getLoc(), ssbAddrOffsetOp.getResult(), + c32ConstOp.getResult()); + + // %vec_id = arith.cmpi eq, %sub_id, %c0_i64 : i64 + auto vecIdOp = builder.create( + scopeOp->getLoc(), mlir::arith::CmpIPredicate::eq, + subIdOp.getResult(), c0i64ConstOp.getResult()); + + // 2. 在parentop的开头插入代码 + builder.setInsertionPointToStart(&forOp->getRegion(0).front()); + + // add sync_block_wait + auto coreAttr = hivm::TCoreTypeAttr::get(module.getContext(), + hivm::TCoreType::VECTOR); + auto setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto flagId = + builder.getIntegerAttr(builder.getI64Type(), cubeControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + + // 在循环末尾（yield之前）插入代码 + auto &loopBody = forOp->getRegion(0).front(); + // 找到循环体的terminator（应该是yield操作） + auto *terminator = loopBody.getTerminator(); + builder.setInsertionPoint(terminator); + + // add sync_block_wait + coreAttr = hivm::TCoreTypeAttr::get(module.getContext(), + hivm::TCoreType::VECTOR); + setPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + waitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + flagId = builder.getIntegerAttr(builder.getI64Type(), vectorControlIndex); + builder.create(forOp->getLoc(), coreAttr, setPipe, + waitPipe, flagId); + } + } + + auto i64Type = builder.getIntegerType(64); + auto i32Type = builder.getIntegerType(32); + auto initPtrType = mlir::LLVM::LLVMPointerType::get(builder.getContext(), 11); + SmallVector scopeOps; + module->walk([&](mlir::Operation *op) { + // 检查是否为目标操作 + if (auto scopeOp = dyn_cast(op)) { + scopeOps.push_back(scopeOp); + } + }); + if (!scopeOps.empty()) { + auto scopeOp = scopeOps[0]; + builder.setInsertionPoint(scopeOp); + auto c0i64ConstAttr = mlir::IntegerAttr::get(i64Type, 0); + auto c0i64ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c0i64ConstAttr); + auto c32i64ConstAttr = mlir::IntegerAttr::get(i64Type, 32); + auto c32i64ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c32i64ConstAttr); + auto c64i64ConstAttr = mlir::IntegerAttr::get(i64Type, 64); + auto c64i64ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c64i64ConstAttr); + auto c96i64ConstAttr = mlir::IntegerAttr::get(i64Type, 96); + auto c96i64ConstOp = builder.create( + scopeOp->getLoc(), i64Type, c96i64ConstAttr); + auto c0i32ConstAttr = mlir::IntegerAttr::get(i32Type, 0); + auto c0i32ConstOp = builder.create( + scopeOp->getLoc(), i32Type, c0i32ConstAttr); + + auto c0initInttoptrOp = builder.create( + scopeOp->getLoc(), initPtrType, c0i64ConstOp.getResult()); + auto c32initInttoptrOp = builder.create( + scopeOp->getLoc(), initPtrType, c32i64ConstOp.getResult()); + auto c64initInttoptrOp = builder.create( + scopeOp->getLoc(), initPtrType, c64i64ConstOp.getResult()); + auto c96initInttoptrOp = builder.create( + scopeOp->getLoc(), initPtrType, c96i64ConstOp.getResult()); + + builder.create(scopeOp->getLoc(), c0i32ConstOp, + c0initInttoptrOp); + builder.create(scopeOp->getLoc(), c0i32ConstOp, + c32initInttoptrOp); + builder.create(scopeOp->getLoc(), c0i32ConstOp, + c64initInttoptrOp); + builder.create(scopeOp->getLoc(), c0i32ConstOp, + c96initInttoptrOp); + } +} + +scf::ForOp transformLoop(scf::ForOp forOp, OpBuilder &builder) { + + // 1. 获取原始循环的信息 + Value originalLowerBound = forOp.getLowerBound(); + Value originalUpperBound = forOp.getUpperBound(); + Value originalStep = forOp.getStep(); + SmallVector iterArgs; + for (auto arg : forOp.getInitArgs()) { + iterArgs.push_back(arg); + } + auto yields = forOp.getBody()->getTerminator(); + + // 2. 检查循环体中是否有特定操作 + int hasTargetOps = 0; + forOp.walk([&](Operation *op) { + if (auto ifOp = dyn_cast(op)) { + if (ifOp->hasAttr("ssbuffer")) { + hasTargetOps++; + } + } + }); + // 3. 如果存在目标操作，在迭代参数中添加计数器 + Value counterInit = nullptr; + mlir::Operation *parentOp = forOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + // 向上遍历查找scope.scope操作 + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + + builder.setInsertionPoint(scopeOp); + for (int i = 0; i < hasTargetOps; i++) { + Location loc = forOp.getLoc(); + auto argType = originalLowerBound.getType(); + + // 添加到迭代参数列表 + iterArgs.push_back(originalLowerBound); + } + // 2. 创建新的上界：originalUpperBound * 2 + Location loc = forOp.getLoc(); + Type ubType = originalStep.getType(); + builder.setInsertionPoint(forOp); + + int count = 0; + for (auto &op : forOp.getBody()->getOperations()) { + if (auto ifOp = dyn_cast(op)) { + auto parentOp = ifOp->getParentOp(); + if (parentOp == forOp && ifOp->hasAttr("ssbuffer")) { + count++; + } + } + } + + Value two; + if (ubType.isIndex()) { + two = builder.create(loc, count - 1); + } else if (auto intType = dyn_cast(ubType)) { + // 对于整数类型，创建相应类型的常数2 + two = builder.create(loc, intType, count - 1); + } else { + // 其他类型可能需要特殊处理 + llvm::errs() << "Warning: Unexpected type for upper bound: " << ubType + << "\n"; + // 尝试创建索引类型的2然后转换 + auto indexTwo = builder.create(loc, count - 1); + two = builder.create(loc, ubType, indexTwo); + } + + auto steps = builder.create(forOp.getLoc(), originalStep, two); + + auto nowUpperBound = + builder.create(forOp.getLoc(), originalUpperBound, steps); + + // 3. Create a new for loop + auto newForOp = + builder.create(forOp.getLoc(), originalLowerBound, + nowUpperBound, originalStep, iterArgs); + + // 4. 设置IR映射表，将旧循环的变量映射到新循环 + IRMapping mapper; + + // 映射迭代变量 + mapper.map(forOp.getInductionVar(), newForOp.getInductionVar()); + + // 映射迭代参数 + for (auto [oldArg, newArg] : + llvm::zip(forOp.getRegionIterArgs(), newForOp.getRegionIterArgs())) { + mapper.map(oldArg, newArg); + } + + SmallVector newCounterArgs; + for (int i = forOp.getRegionIterArgs().size(); + i < newForOp.getRegionIterArgs().size(); i++) { + newCounterArgs.push_back(newForOp.getRegionIterArgs()[i]); + } + // 5. 克隆循环体内容到新循环 + auto &newLoopBody = *newForOp.getBody(); + builder.setInsertionPointToStart(&newLoopBody); + + for (auto &op : forOp.getBody()->without_terminator()) { + builder.clone(op, mapper); + } + + // 6. 克隆yield操作 + if (auto yieldOp = dyn_cast(yields)) { + SmallVector newYieldOperands; + for (auto operand : yieldOp.getOperands()) { + newYieldOperands.push_back(mapper.lookupOrDefault(operand)); + } + if (hasTargetOps != 0) { + for (auto currentCounter : newCounterArgs) { + // 将更新后的计数器添加到yield操作数中 + newYieldOperands.push_back(currentCounter); + } + } + builder.create(yieldOp.getLoc(), newYieldOperands); + } + + // 7. 替换原循环的结果 + if (hasTargetOps != 0) { + // 新循环有额外的计数器结果，但原循环没有对应结果 + // 我们可以选择只替换原循环对应的结果，或者忽略计数器结果 + unsigned numOriginalResults = forOp.getNumResults(); + SmallVector originalResults; + for (unsigned i = 0; i < numOriginalResults; i++) { + originalResults.push_back(newForOp.getResult(i)); + } + forOp.replaceAllUsesWith(originalResults); + } else { + forOp.replaceAllUsesWith(newForOp.getResults()); + } + + // 8. 删除原循环 + forOp.erase(); + return newForOp; +} + +// Find the first occurrence of convert_layout or fixpipe operation after the +// specified operation +Operation * +findFirstTargetOpAfterWait(SyncBlockWaitOp waitOp, + SmallVector &excludedValues) { + bool startSearching = false; + + for (Operation &op : waitOp->getBlock()->getOperations()) { + // meet waitop, start searching + if (&op == waitOp) { + startSearching = true; + continue; + } + + // have not meet waitOp, skip + if (!startSearching) { + continue; + } + + Operation *res = nullptr; + + if (isa(op)) { + res = op.getOperands()[0].getDefiningOp(); + } else if (isa(op)) { + res = op.getOperands()[1].getDefiningOp(); + } else if (isa(op)) { + res = op.getOperands()[1].getDefiningOp(); + } else if (isa(op)) { + res = op.getOperands()[0].getDefiningOp(); + } + + // get res + if (res) { + // in excludedValues → skip，search next one + if (llvm::is_contained(excludedValues, res)) { + continue; + } + // not in excludedValues → note and return + excludedValues.push_back(res); + return res; + } + } + + return nullptr; +} + +void getWaitType(std::string CoreType, scf::ForOp forOp, + SmallVector &waitTypes, + SmallVector &allocTypes) { + auto scalarWaitPipe = PipeAttr::get(forOp.getContext(), hivm::PIPE::PIPE_S); + auto cubeWaitPipe = PipeAttr::get(forOp.getContext(), hivm::PIPE::PIPE_FIX); + auto vectorWaitPipe = + PipeAttr::get(forOp.getContext(), hivm::PIPE::PIPE_MTE3); + SmallVector excludedValues; + forOp.walk([&](Operation *op) { + if (auto waitOp = dyn_cast(op)) { + auto parentOp = op->getParentOp(); + if (isa(parentOp) && parentOp->hasAttr("ssbuffer")) { + auto ifOp = dyn_cast(parentOp); + if (forOp == ifOp->getParentOp()) { + auto waitPipe = waitOp.getPipe(); + auto setPipe = waitOp.getTpipe(); + if (waitPipe == scalarWaitPipe || setPipe == scalarWaitPipe) { + bool hasBackward = waitOp->hasAttr("ssbuf.backward"); + if (hasBackward) { + waitTypes.push_back(0); + allocTypes.push_back(waitOp); + } else { + waitTypes.push_back(1); + allocTypes.push_back(waitOp); + } + } else { + if ((waitPipe == cubeWaitPipe && CoreType == "cube") || + (waitPipe == vectorWaitPipe && CoreType == "vector")) { + auto allocOp = findFirstTargetOpAfterWait(waitOp, excludedValues); + waitTypes.push_back(0); + allocTypes.push_back(allocOp); + } else { + auto allocOp = findFirstTargetOpAfterWait(waitOp, excludedValues); + waitTypes.push_back(1); + allocTypes.push_back(allocOp); + } + } + } + } + } + }); +} + +DenseMap getCounterOffset(scf::ForOp forOp) { + int i = 0; + DenseMap bufferMap; + auto scalarWaitPipe = PipeAttr::get(forOp.getContext(), hivm::PIPE::PIPE_S); + forOp.walk([&](Operation *op) { + bufferMap[i] = 0; + auto ifOp = dyn_cast(op); + if (ifOp && ifOp->hasAttr("ssbuffer") && ifOp->getParentOp() == forOp) { + ifOp.walk([&](Operation *op) { + if (auto waitOp = dyn_cast(op)) { + if (auto waitIfOp = dyn_cast(op->getParentOp())) { + if (waitIfOp == ifOp) { + auto waitPipe = waitOp.getPipe(); + auto setPipe = waitOp.getTpipe(); + if ((waitPipe != scalarWaitPipe || setPipe != scalarWaitPipe)) { + bufferMap[i]++; + } + } + } + } + }); + i++; + } + }); + return bufferMap; +} + +SmallVector addBufValLoop(scf::ForOp forOp, + DenseMap VecBitMap, + DenseMap CubeBitMap, + OpBuilder &builder) { + auto aiCAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + bool isAIC = false; + // 向上查找父scope.scope操作 + mlir::Operation *parentOp = forOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + // 向上遍历查找scope.scope操作 + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + if (scopeOp->hasAttr("hivm.tcore_type")) { + auto attr = scopeOp->getAttr("hivm.tcore_type"); + if (attr == aiCAttr) { + isAIC = true; + } + } + auto bufferMap = getCounterOffset(forOp); + SmallVector buf_vals; + SmallVector if_conditions; + builder.setInsertionPointToStart(&scopeOp->getRegion(0).front()); + + // 1. 提取并处理end值 + Value startValue = forOp.getLowerBound(); + Value endValue = forOp.getUpperBound(); + // 2. 提取并处理step值 + Value stepValue = forOp.getStep(); + builder.setInsertionPoint(forOp); + Location loc = forOp.getLoc(); + int count = 0; + for (auto &op : forOp.getBody()->getOperations()) { + if (auto ifOp = dyn_cast(op)) { + auto parentOp = ifOp->getParentOp(); + if (parentOp == forOp && ifOp->hasAttr("ssbuffer")) { + count++; + } + } + } + + Value two; + Type ubType = stepValue.getType(); + if (ubType.isIndex()) { + two = builder.create(loc, count - 1); + } else if (auto intType = dyn_cast(ubType)) { + // 对于整数类型，创建相应类型的常数2 + two = builder.create(loc, intType, count - 1); + } else { + // 其他类型可能需要特殊处理 + llvm::errs() << "Warning: Unexpected type for upper bound: " << ubType + << "\n"; + // 尝试创建索引类型的2然后转换 + auto indexTwo = builder.create(loc, count - 1); + two = builder.create(loc, ubType, indexTwo); + } + + auto steps = builder.create(forOp.getLoc(), endValue.getType(), + stepValue, two); + + auto subLoopValue = builder.create( + forOp.getLoc(), endValue.getType(), endValue, steps); + + SmallVector WaitType; + SmallVector AllocType; + SmallVector bufferPtrs; + if (isAIC) { + builder.setInsertionPointToStart(&forOp->getRegion(0).front()); + // 创建常量32和64 + Value c0 = + builder.create(forOp.getLoc(), 0, 32 // 值32，64位 + ); + Value c32 = builder.create(forOp.getLoc(), 32, + 64 // 值32，64位 + ); + Value c64 = builder.create(forOp.getLoc(), 64, + 64 // 值64，64位 + ); + // 创建inttoptr操作 + Value ssb_vec0_ptr = builder.create( + forOp.getLoc(), + LLVM::LLVMPointerType::get(builder.getContext(), 11), // 地址空间11 + c32); + Value ssb_vec1_ptr = builder.create( + forOp.getLoc(), + LLVM::LLVMPointerType::get(builder.getContext(), 11), // 地址空间11 + c64); + bufferPtrs.push_back(ssb_vec0_ptr); + bufferPtrs.push_back(ssb_vec1_ptr); + // 创建load操作 + Value status_vec0 = builder.create( + forOp.getLoc(), builder.getI32Type(), ssb_vec0_ptr); + + Value status_vec1 = builder.create( + forOp.getLoc(), builder.getI32Type(), ssb_vec1_ptr); + + getWaitType("cube", forOp, WaitType, AllocType); + + for (auto i = 0; i < WaitType.size(); i++) { + auto correnspondAlloc = CubeBitMap[AllocType[i]]; + auto i32ConstAttr = + mlir::IntegerAttr::get(builder.getI32Type(), 1 << correnspondAlloc); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + Value bufi_vec0_val = builder.create( + forOp.getLoc(), status_vec0, buf_constant_set); + Value bufi_vec1_val = builder.create( + forOp.getLoc(), status_vec1, buf_constant_set); + Value flag_bufi_vec0; + Value flag_bufi_vec1; + // 创建比较操作 + if (WaitType[i] == 0) { + flag_bufi_vec0 = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_vec0_val, c0); + flag_bufi_vec1 = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_vec1_val, c0); + } else { + flag_bufi_vec0 = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_vec0_val, + buf_constant_set); + flag_bufi_vec1 = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_vec1_val, + buf_constant_set); + } + // 创建最终的and操作 + Value bufi_val = builder.create( + forOp.getLoc(), flag_bufi_vec0, flag_bufi_vec1); + buf_vals.push_back(bufi_val); + } + + } else { + builder.setInsertionPointToStart(&scopeOp->getRegion(0).front()); + Value c0 = + builder.create(forOp.getLoc(), 0, 32 // 值32，64位 + ); + auto i64Type = builder.getIntegerType(64); + // %sub_id = hivm.hir.get_sub_block_idx -> i64 + // 这里假设有一个getSubBlockIdxOp操作 + Value subId = builder.create(scopeOp->getLoc(), i64Type); + auto i64ConstAttr = mlir::IntegerAttr::get(i64Type, 32); + Value cstOffset = builder.create( + scopeOp->getLoc(), i64Type, i64ConstAttr); + Value ssbAddrOffset = + builder.create(scopeOp->getLoc(), subId, cstOffset); + Value ssbAddr = builder.create(scopeOp->getLoc(), + ssbAddrOffset, cstOffset); + builder.setInsertionPointToStart(&forOp->getRegion(0).front()); + // 创建inttoptr操作 + Value ssb_cube_ptr = builder.create( + forOp.getLoc(), + LLVM::LLVMPointerType::get(builder.getContext(), 11), // 地址空间11 + ssbAddr); + bufferPtrs.push_back(ssb_cube_ptr); + // 创建load操作 + Value status_cube = builder.create( + forOp.getLoc(), builder.getI32Type(), ssb_cube_ptr); + + getWaitType("vector", forOp, WaitType, AllocType); + for (auto i = 0; i < WaitType.size(); i++) { + auto correnspondAlloc = VecBitMap[AllocType[i]]; + auto i32ConstAttr = + mlir::IntegerAttr::get(builder.getI32Type(), 1 << correnspondAlloc); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + Value bufi_cube_val = builder.create( + forOp.getLoc(), status_cube, buf_constant_set); + + Value flag_bufi_cube; + // 创建比较操作 + if (WaitType[i] == 0) { + flag_bufi_cube = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_cube_val, c0); + } else { + flag_bufi_cube = builder.create( + forOp.getLoc(), arith::CmpIPredicate::eq, bufi_cube_val, + buf_constant_set); + } + buf_vals.push_back(flag_bufi_cube); + } + } + int bufIdx = 0; + int groupIdx = 0; + + for (const auto &pair : bufferMap) { + if (bufferMap[groupIdx] == 0) { + continue; + } + + // 获取对应的region迭代参数 + Value cnti = builder.create( + forOp.getLoc(), arith::CmpIPredicate::slt, + forOp.getRegionIterArgs()[forOp.getRegionIterArgs().size() - + (bufferMap.size() - 1 - groupIdx)], + subLoopValue); + + // 计算该组中所有buffer值的AND + Value finalBufVal = buf_vals[bufIdx]; + for (int count = 1; count < bufferMap[groupIdx]; count++) { + finalBufVal = builder.create(forOp.getLoc(), finalBufVal, + buf_vals[bufIdx + count]); + } + + auto cond = + builder.create(forOp.getLoc(), finalBufVal, cnti); + if_conditions.push_back(cond); + + // 更新索引 + bufIdx += bufferMap[groupIdx]; + groupIdx++; + } + int ifIndex = 0; + int acc = 0; + int bufferBit = 0; + for (int i = 0; i < CubeBitMap.size(); i++) { + bufferBit += (1 << i); + } + forOp.getBody()->walk([&](Operation *op) { + auto ifOp = dyn_cast(op); + if (ifOp && ifOp->hasAttr("ssbuffer")) { + // 获取then区域 + Block *thenBlock = &ifOp.getThenRegion().front(); + + // 找到then区域中的yield操作 + Operation *yieldOp = nullptr; + for (auto &op : *thenBlock) { + if (isa(op)) { + yieldOp = &op; + break; + } + } + if (yieldOp) { + builder.setInsertionPoint(yieldOp); + + if (isAIC) { + // 创建插入的语句 + // %status_v2 = llvm.load %ssb_ptr : !llvm.ptr<11> -> i32 + Value status_v2_0 = builder.create( + yieldOp->getLoc(), + builder.getIntegerType(32), // i32类型 + bufferPtrs[0] // 假设ssb_ptr已在作用域中定义 + ); + Value status_v2_1 = builder.create( + yieldOp->getLoc(), + builder.getIntegerType(32), // i32类型 + bufferPtrs[1] // 假设ssb_ptr已在作用域中定义 + ); + Value buf_val_new_0 = status_v2_0; + Value buf_val_new_1 = status_v2_1; + auto bufferNum = bufferMap[ifIndex]; + for (int i = 0; i < bufferNum; i++) { + if (WaitType[acc + i] == 0) { + auto correnspondAlloc = CubeBitMap[AllocType[acc + i]]; + auto i32ConstAttr = mlir::IntegerAttr::get(builder.getI32Type(), + 1 << correnspondAlloc); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + buf_val_new_0 = builder.create( + yieldOp->getLoc(), buf_val_new_0, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + buf_val_new_1 = builder.create( + yieldOp->getLoc(), buf_val_new_1, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + } else { + auto correnspondAlloc = CubeBitMap[AllocType[acc + i]]; + int bitPos = correnspondAlloc; + int basePattern = bufferBit; + int finalValue = basePattern ^ (1 << bitPos); + auto i32ConstAttr = + mlir::IntegerAttr::get(builder.getI32Type(), finalValue); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + buf_val_new_0 = builder.create( + yieldOp->getLoc(), buf_val_new_0, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + buf_val_new_1 = builder.create( + yieldOp->getLoc(), buf_val_new_1, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + } + } + acc += bufferNum; + builder.create(yieldOp->getLoc(), buf_val_new_0, + bufferPtrs[0]); + builder.create(yieldOp->getLoc(), buf_val_new_1, + bufferPtrs[1]); + + } else { + // 创建插入的语句 + // %status_v2 = llvm.load %ssb_ptr : !llvm.ptr<11> -> i32 + Value status_v2 = builder.create( + yieldOp->getLoc(), + builder.getIntegerType(32), // i32类型 + bufferPtrs[0] // 假设ssb_ptr已在作用域中定义 + ); + Value buf_val_new = status_v2; + auto bufferNum = bufferMap[ifIndex]; + for (int i = 0; i < bufferNum; i++) { + if (WaitType[acc + i] == 0) { + auto correnspondAlloc = VecBitMap[AllocType[acc + i]]; + auto i32ConstAttr = mlir::IntegerAttr::get(builder.getI32Type(), + 1 << correnspondAlloc); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + buf_val_new = builder.create( + yieldOp->getLoc(), buf_val_new, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + } else { + auto correnspondAlloc = VecBitMap[AllocType[acc + i]]; + int bitPos = correnspondAlloc; + int basePattern = bufferBit; + int finalValue = basePattern ^ (1 << bitPos); + auto i32ConstAttr = + mlir::IntegerAttr::get(builder.getI32Type(), finalValue); + auto buf_constant_set = builder.create( + scopeOp->getLoc(), builder.getI32Type(), i32ConstAttr); + buf_val_new = builder.create( + yieldOp->getLoc(), buf_val_new, + buf_constant_set // 假设buf3_clear已在作用域中定义 + ); + } + } + acc += bufferNum; + builder.create(yieldOp->getLoc(), buf_val_new, + bufferPtrs[0]); + } + ifIndex++; + } + } + }); + + return if_conditions; +} + +void ReplaceIf(scf::ForOp forOp, SmallVector conditions, + SmallVector &opsToErase, + DenseMap &ifArgMap, OpBuilder &builder, + ModuleOp moduleOp) { + SmallVector ifToProcess; + llvm::outs() << "enter replaceif\n"; + Value step = forOp.getStep(); + auto aiCAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + forOp.getBody()->walk([&](Operation *op) { + auto ifOp = dyn_cast(op); + if (ifOp && ifOp->hasAttr("ssbuffer") && forOp == ifOp->getParentOp()) { + ifToProcess.push_back(ifOp); + } + }); + + IRMapping IRMap; + for (int i = 0; i < ifToProcess.size(); i++) { + auto ifOp = ifToProcess[i]; + auto parentOp = ifOp->getParentOp(); + auto loc = ifOp.getLoc(); + // 获取for循环的iterargs（迭代参数） + auto iterArgs = forOp.getRegionIterArgs(); + if (iterArgs.size() < conditions.size()) { + return; + } + auto thenYieldOp = + dyn_cast(ifOp.getThenRegion().front().getTerminator()); + SmallVector thenResults; + if (thenYieldOp) { + // 如果已有返回值，保留它们 + for (auto result : thenYieldOp.getResults()) { + thenResults.push_back(result); + } + } + // 创建新的else区域，返回两个迭代参数 + SmallVector elseResults; + scf::YieldOp elseYieldOp = nullptr; + bool hasElse = false; + if (!ifOp.getElseRegion().empty()) { + elseYieldOp = + dyn_cast(ifOp.getElseRegion().front().getTerminator()); + hasElse = true; + } + if (elseYieldOp) { + for (auto result : elseYieldOp.getResults()) { + elseResults.push_back(result); + } + } + // 获取最后两个迭代参数 + Value iterArgMinus = iterArgs[iterArgs.size() - (conditions.size() - i)]; + // 创建新的then区域，返回两个迭代参数 + thenResults.push_back(iterArgMinus); + elseResults.push_back(iterArgMinus); + + // 保存原有的操作，以便后续克隆 + SmallVector thenOps; + for (auto &op : ifOp.getThenRegion().front()) { + thenOps.push_back(&op); + } + + SmallVector elseOps; + if (!ifOp.getElseRegion().empty()) { + for (auto &op : ifOp.getElseRegion().front()) { + elseOps.push_back(&op); + } + } + SmallVector resultTypes; + for (auto val : thenResults) { + resultTypes.push_back(val.getType()); + } + // 创建新的scf.if操作 + builder.setInsertionPoint(ifOp); + auto newIfOp = builder.create(loc, resultTypes, conditions[i], + /*withElseRegion=*/true); + newIfOp->setAttr("ssbuffer", builder.getUnitAttr()); + // 处理then区域 + auto &newThenBlock = newIfOp.getThenRegion().front(); + builder.setInsertionPointToStart(&newThenBlock); + + // 克隆then区域的操作 + for (auto op : thenOps) { + if (auto yieldOp = dyn_cast(op)) { + // 处理yield的操作数映射 + SmallVector mappedOperands; + for (auto operand : yieldOp->getOperands()) { + mappedOperands.push_back(IRMap.lookupOrDefault(operand)); + } + // 获取最后两个迭代参数 + Value iterArgMinus = + iterArgs[iterArgs.size() - (conditions.size() - i)]; + + // %ssb_addr = arith.addi %ssb_addr_offset, %c32_i64 : i64 + auto AddIOp = builder.create(forOp->getLoc(), + iterArgMinus, step); + // 这里加个add1 + mappedOperands.push_back(AddIOp); + builder.create(loc, mappedOperands); + } else { + auto newOp = builder.clone(*op, IRMap); + IRMap.map(op->getResults(), newOp->getResults()); + } + } + + // 处理else区域 + auto &newElseBlock = newIfOp.getElseRegion().front(); + builder.setInsertionPointToStart(&newElseBlock); + // 克隆else区域的操作 + if (hasElse) { + for (auto op : elseOps) { + if (auto yieldOp = dyn_cast(op)) { + // 处理yield的操作数映射 + SmallVector mappedOperands; + for (auto operand : yieldOp->getOperands()) { + mappedOperands.push_back(IRMap.lookupOrDefault(operand)); + } + Value iterArgMinus = + iterArgs[iterArgs.size() - (conditions.size() - i)]; + mappedOperands.push_back(iterArgMinus); + builder.create(loc, mappedOperands); + } else { + auto newOp = builder.clone(*op, IRMap); + IRMap.map(op->getResults(), newOp->getResults()); + } + } + } else { + SmallVector cntOperands; + cntOperands.push_back(iterArgMinus); + builder.create(loc, cntOperands); + } + + // 替换原有if操作的使用 + // 首先，将原if操作的结果替换为新if操作的对应结果 + for (unsigned j = 0; j < ifOp.getNumResults(); ++j) { + ifOp.getResult(j).replaceAllUsesWith(newIfOp.getResult(j)); + } + // 获取新if操作所在的块 + Block *newIfBlock = ifOp->getBlock(); + // 在for循环体内替换迭代参数的使用 + forOp.getBody()->walk([&](Operation *op) { + // 检查操作是否与新ifOp在同一个块中 + Block *opBlock = op->getBlock(); + if (opBlock != newIfBlock) { + // 不在同一个块中，跳过 + return; + } + if (op->isBeforeInBlock(newIfOp)) { + return; // 只处理if操作之后的use + } + for (unsigned j = 0; j < op->getNumOperands(); ++j) { + for (auto argIndex = 0; argIndex < conditions.size(); argIndex++) { + // 获取最后两个迭代参数 + Value iterArgMinus = + iterArgs[iterArgs.size() - (conditions.size() - i)]; + if (op->getOperand(j) == iterArgMinus) { + op->setOperand(j, + newIfOp.getResults()[newIfOp.getNumResults() - 1]); + } + } + } + }); + + // // 删除原有的if操作 + opsToErase.push_back(ifOp); + if (ifArgMap.find(newIfOp) == ifArgMap.end()) { + ifArgMap[newIfOp] = iterArgMinus; + } + } +} + +int getNestingDepth(scf::ForOp forOp) { + int depth = 0; + Operation *op = forOp.getOperation(); + while (op) { + if (op->getDialect() && op->getDialect()->getNamespace() == "scf") { + ++depth; + } + op = op->getParentOp(); + } + return depth; +} + +void printDenseMap(const mlir::DenseMap &Map) { + for (const auto &pair : Map) { + auto op = pair.first; + auto bitValue = pair.second; + + // ✅ 正确打印 Operation 的方法：print(llvm::outs()) + llvm::outs() << "Operation: \n"; + llvm::outs().flush(); + llvm::outs() << "Operation Name: " << *op << "\n"; + + // 打印对应的值 + llvm::outs() << " | Bit Value: " << bitValue << " | allocmap\n"; + llvm::outs().flush(); + llvm::outs() << "------------------------------------------------\n\n"; + llvm::outs().flush(); + } + llvm::outs().flush(); +} + +void getAllocBit(ModuleOp module, DenseMap &VecBitMap, + DenseMap &CubeBitMap, OpBuilder builder) { + auto aiCAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + auto scalarWaitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_S); + auto cubeWaitPipe = PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_FIX); + auto vectorWaitPipe = + PipeAttr::get(module.getContext(), hivm::PIPE::PIPE_MTE3); + + SmallVector scopeOpToEdit; + module.walk( + [&](scope::ScopeOp scopeOp) { scopeOpToEdit.push_back(scopeOp); }); + + int cubeAcc = 0; + int vecAcc = 0; + for (auto scopeOp : scopeOpToEdit) { + SmallVector excludedValues; + if (!scopeOp->hasAttr("hivm.tcore_type")) + continue; + + auto attr = scopeOp->getAttr("hivm.tcore_type"); + if (attr == aiCAttr) { + // CUBE 类型：分配 cubeAcc 自增 ID + scopeOp.walk([&](SyncBlockWaitOp waitOp) { + auto parentOp = waitOp->getParentOp(); + if (isa(parentOp) && parentOp->hasAttr("ssbuffer")) { + auto waitPipe = waitOp.getPipe(); + if (waitPipe != scalarWaitPipe) { + auto allocOp = findFirstTargetOpAfterWait(waitOp, excludedValues); + if (VecBitMap.count(allocOp)) { + CubeBitMap[allocOp] = VecBitMap[allocOp]; + cubeAcc++; + } else { + CubeBitMap[allocOp] = cubeAcc++; + } + } + } + }); + } else { + // VEC 类型：分配 vecAcc 自增 ID + scopeOp.walk([&](SyncBlockWaitOp waitOp) { + auto parentOp = waitOp->getParentOp(); + if (isa(parentOp) && parentOp->hasAttr("ssbuffer")) { + auto waitPipe = waitOp.getPipe(); + if (waitPipe != scalarWaitPipe) { + auto allocOp = findFirstTargetOpAfterWait(waitOp, excludedValues); + if (!VecBitMap.count(allocOp)) { + VecBitMap[allocOp] = vecAcc++; + } + } + } + }); + } + } + + // 记录【自增ID的最大值】，用于后面偏移计算 + const int maxVecId = vecAcc - 2; + const int maxCubeId = cubeAcc - 2; + // ======================== 阶段 2：分配【偏移ID】给 + // waitOp（绝对不重复）======================== + for (auto scopeOp : scopeOpToEdit) { + if (!scopeOp->hasAttr("hivm.tcore_type")) + continue; + + auto attr = scopeOp->getAttr("hivm.tcore_type"); + if (attr == aiCAttr) { + // CUBE：waitOp = maxCubeId + 1 + flagid + scopeOp.walk([&](SyncBlockWaitOp waitOp) { + auto parentOp = waitOp->getParentOp(); + if (isa(parentOp) && parentOp->hasAttr("ssbuffer")) { + auto setPipe = waitOp.getTpipe(); + auto waitPipe = waitOp.getPipe(); + if (setPipe != scalarWaitPipe && waitPipe == scalarWaitPipe) { + uint32_t flagid = 0; + if (auto attr = + waitOp->getAttrOfType("ssbuf.flagid")) { + flagid = attr.getUInt(); + } + // 核心修改：偏移分配，不再写死 16 + CubeBitMap[waitOp] = + (maxCubeId >= 0 ? (maxCubeId + 1) : 0) + flagid; + } + } + }); + } else { + // VEC：waitOp = maxVecId + 1 + flagid + scopeOp.walk([&](SyncBlockWaitOp waitOp) { + auto parentOp = waitOp->getParentOp(); + if (isa(parentOp) && parentOp->hasAttr("ssbuffer")) { + auto setPipe = waitOp.getTpipe(); + auto waitPipe = waitOp.getPipe(); + if (setPipe != scalarWaitPipe && waitPipe == scalarWaitPipe) { + uint32_t flagid = 0; + if (auto attr = + waitOp->getAttrOfType("ssbuf.flagid")) { + flagid = attr.getUInt(); + } + // 核心修改：偏移分配 + VecBitMap[waitOp] = (maxVecId >= 0 ? (maxVecId + 1) : 0) + flagid; + } + } + }); + } + } +} + +void modifyForIterargDeps(scf::ForOp forOp, + DenseMap ifCounters) { + Value iterArg = forOp.getInductionVar(); + + for (Operation &op : forOp.getBody()->without_terminator()) { + if (auto ifOp = dyn_cast(op)) { + if (ifCounters.find(ifOp) != ifCounters.end()) { + Value counter = ifCounters[ifOp]; + + ifOp.walk([&](Operation *opInIf) { + for (auto [i, operand] : llvm::enumerate(opInIf->getOperands())) { + if (operand == iterArg) { + opInIf->setOperand(i, counter); + } + } + }); + } + } + } +} + +void FlowSssbuf(ModuleOp module) { + mlir::OpBuilder builder(module.getContext()); + // 收集所有需要转换的循环 + SmallVector targetLoops; + llvm::outs() << "enter flowsssbuf\n\n"; + module.walk([&](Operation *op) { + if (auto forOp = dyn_cast(op)) { + // 检查循环是否包含特定的 sync_block_set 操作 + bool hasSyncBlockSet = false; + forOp.walk([&](Operation *op) { + if (isa(op)) { + if (auto ifOp = dyn_cast(op->getParentOp())) { + if (forOp == ifOp->getParentOp() && ifOp->hasAttr("ssbuffer")) { + hasSyncBlockSet = true; + } + } + } + }); + + if (hasSyncBlockSet) { + if (llvm::find(targetLoops, forOp) == targetLoops.end()) { + targetLoops.push_back(forOp); + } + } + } + }); + llvm::outs() << "enter flowsssbuf\n\n"; + + SmallVector transformLoops; + // 转换每个目标循环 + for (scf::ForOp forOp : targetLoops) { + auto newforOp = transformLoop(forOp, builder); + } + + module.walk([&](Operation *op) { + if (auto forOp = dyn_cast(op)) { + // 检查循环是否包含特定的 sync_block_set 操作 + bool hasSyncBlockSet = false; + forOp.walk([&](Operation *op) { + if (isa(op)) { + if (auto ifOp = dyn_cast(op->getParentOp())) { + if (forOp == ifOp->getParentOp() && ifOp->hasAttr("ssbuffer")) { + hasSyncBlockSet = true; + } + } + } + }); + + if (hasSyncBlockSet) { + if (llvm::find(transformLoops, forOp) == transformLoops.end()) { + transformLoops.push_back(forOp); + } + } + } + }); + + llvm::sort(transformLoops, [](scf::ForOp a, scf::ForOp b) { + return getNestingDepth(a) > getNestingDepth(b); + }); + DenseMap VecBitMap; + DenseMap CubeBitMap; + getAllocBit(module, VecBitMap, CubeBitMap, builder); + printDenseMap(CubeBitMap); + printDenseMap(VecBitMap); + SmallVector opsToErase; + for (scf::ForOp forOp : transformLoops) { + DenseMap ifArgMap; + llvm::outs() << "before replaceif\n"; + auto bufvals = addBufValLoop(forOp, VecBitMap, CubeBitMap, builder); + ReplaceIf(forOp, bufvals, opsToErase, ifArgMap, builder, module); + llvm::outs() << "after replaceif\n"; + for (const auto &pair : ifArgMap) { + auto val = pair.first; + auto bitValue = pair.second; + llvm::outs() << val << " " << bitValue << " ifargmrp\n\n\n"; + llvm::outs().flush(); + } + + modifyForIterargDeps(forOp, ifArgMap); + } + for (auto op : opsToErase) { + op->erase(); + } +} + +bool isTransOp(mlir::Operation *op) { + auto fixpipeOp = dyn_cast(op); + if (fixpipeOp) + return true; + + auto copyOp = dyn_cast(op); + if (!copyOp) + return false; + else { + + Value copySrc = copyOp.getODSOperands(0).front(); + MemRefType copySrcTy = dyn_cast(copySrc.getType()); + auto SrcAddrSpace = + dyn_cast_or_null(copySrcTy.getMemorySpace()); + bool isSrcUbSpace = + SrcAddrSpace.getAddressSpace() == hivm::AddressSpace::UB; + + Value copyDst = copyOp.getODSOperands(1).front(); + MemRefType copyDstTy = dyn_cast(copyDst.getType()); + auto DstAddrSpace = + dyn_cast_or_null(copyDstTy.getMemorySpace()); + bool isDstCbufSpace = + DstAddrSpace.getAddressSpace() == hivm::AddressSpace::L1; + + return isSrcUbSpace && isDstCbufSpace; + } +} + +void FindAndMarkBuffer(ModuleOp module) { + OpBuilder builder(module.getContext()); + unsigned int BufferIdx = 0; + Type idxType = builder.getI32Type(); + StringAttr setFlagAttr = builder.getStringAttr("Set flag"); + StringAttr waitFlagAttr = builder.getStringAttr("Wait flag"); + IntegerAttr idxAttr = builder.getI32IntegerAttr(BufferIdx); + + module.walk([&](mlir::Operation *op) { + if (isTransOp(op)) { + llvm::outs() << "Buffer idx" << BufferIdx << "\n"; + llvm::outs() << "Trans Op" << *op << "\n"; + Value SharedBuffer; + if (auto fixpipeOp = dyn_cast(op)) { + SharedBuffer = fixpipeOp.getODSOperands(1).front(); + } else { + auto copyOp = dyn_cast(op); + SharedBuffer = copyOp.getODSOperands(1).front(); + } + llvm::outs() << "SharedBuffer" << SharedBuffer << "\n"; + + if (!SharedBuffer) { + op->emitWarning("fixpipe op has empty output operand!"); + return; + } + + // 在Buffer的生产op后set flag标记，在Buffer消费op前增加wait flag标记 + op->setAttr("Buffer idx", builder.getI32IntegerAttr(BufferIdx)); + op->setAttr("Wait Flag", builder.getI32IntegerAttr(0)); + op->setAttr("Set Flag", builder.getI32IntegerAttr(1)); + + for (Operation *consumerOp : SharedBuffer.getUsers()) { + if (consumerOp == op) + continue; + if (!consumerOp) + continue; + + llvm::outs() << "consumerOp: " << *consumerOp << "\n"; + + consumerOp->setAttr("Buffer idx", builder.getI32IntegerAttr(BufferIdx)); + consumerOp->setAttr("Wait Flag", builder.getI32IntegerAttr(0)); + } + BufferIdx++; + } + }); +} + +// 结构体存 wait-set 区块信息 +struct WaitSetRegion { + Operation *waitOp; + Operation *lastSetOp; + SmallVector opsToMove; + bool hasCopyOrFixpipe = false; +}; + +struct MergedRegion { + SmallVector regions; + SmallVector opsToMove; + SmallVector yieldValues; + SmallVector resultTypes; +}; + +void MoveIterArgUsersIntoIf(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // iter_arg -> mergedRegion index + DenseMap iterArgToRegion; + + for (int r = 0; r < mergedRegions.size(); ++r) { + MergedRegion &mr = mergedRegions[r]; + + for (Operation *op : mr.opsToMove) { + for (Value v : op->getOperands()) { + if (auto barg = mlir::dyn_cast(v)) { + if (barg.getOwner() == &body) { + iterArgToRegion.try_emplace(barg, r); + } + } + } + } + } + + if (iterArgToRegion.empty()) + return; + + // 找最后一个 mergedRegion 的最后一个 op + Operation *lastOp = nullptr; + for (MergedRegion &mr : mergedRegions) + lastOp = mr.opsToMove.back(); + + if (!lastOp) + return; + + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) + opIndex[&op] = idx++; + + int startIdx = opIndex[lastOp] + 1; + + // 扫描 for body 尾部 op + for (Operation &op : body) { + if (opIndex[&op] < startIdx) + continue; + + llvm::SmallDenseSet usedRegions; + for (Value v : op.getOperands()) { + if (auto barg = mlir::dyn_cast(v)) { + auto it = iterArgToRegion.find(barg); + if (it != iterArgToRegion.end()) + usedRegions.insert(it->second); + } + } + + // 必须且只能依赖一个 mergedRegion + if (usedRegions.size() != 1) + continue; + + int target = *usedRegions.begin(); + + mergedRegions[target].opsToMove.push_back(&op); + } +} + +void ComputeYieldForMergedRegion(MergedRegion &mr, Block &body) { + + mr.yieldValues.clear(); + mr.resultTypes.clear(); + + SmallPtrSet inRegion(mr.opsToMove.begin(), + mr.opsToMove.end()); + + for (Operation *op : mr.opsToMove) { + for (Value res : op->getResults()) { + bool usedOutside = false; + + for (OpOperand &use : res.getUses()) { + Operation *user = use.getOwner(); + + // 不在同一个 for body，交给外层处理（通常不会出现） + if (user->getBlock() != &body) + continue; + + // 只要有一个 use 在 region 外，就必须 yield + if (!inRegion.contains(user)) { + usedOutside = true; + break; + } + } + + if (usedOutside) { + mr.yieldValues.push_back(res); + mr.resultTypes.push_back(res.getType()); + } + } + } +} + +static void ComputeYieldForMergedRegionV2(MergedRegion &mr, Block &body) { + + mr.yieldValues.clear(); + mr.resultTypes.clear(); + + // 当前 region 内的 ops + SmallPtrSet inRegion(mr.opsToMove.begin(), + mr.opsToMove.end()); + + for (Operation *op : mr.opsToMove) { + for (Value res : op->getResults()) { + + bool usedOutside = false; + + for (OpOperand &use : res.getUses()) { + Operation *user = use.getOwner(); + + // 如果使用在 region 内部 op，跳过 + if (inRegion.contains(user)) + continue; + + // 使用在 region 外部，包括嵌套 region 内部的 block + usedOutside = true; + break; + } + + if (usedOutside) { + mr.yieldValues.push_back(res); + mr.resultTypes.push_back(res.getType()); + } + } + } +} + +static void ComputeYieldForMergedRegionV3(MergedRegion &mr) { + mr.yieldValues.clear(); + mr.resultTypes.clear(); + + // 用 DenseSet 暂存当前 region 的所有 ops + DenseSet regionOps(mr.opsToMove.begin(), mr.opsToMove.end()); + + for (Operation *op : mr.opsToMove) { + for (Value res : op->getResults()) { + + bool needsYield = false; + + for (OpOperand &use : res.getUses()) { + Operation *user = use.getOwner(); + + // 如果 user 不在当前 region，则需要 yield + if (!regionOps.contains(user)) { + needsYield = true; + break; + } + } + + if (needsYield) { + mr.yieldValues.push_back(res); + mr.resultTypes.push_back(res.getType()); + } + } + } +} + +// 递归收集 op 和它所有 region 内的 ops +static void CollectAllNestedOps(Operation *op, + DenseSet ®ionOps) { + if (!op) + return; + + if (regionOps.contains(op)) + return; // 已经收集过 + + regionOps.insert(op); + + // 遍历所有 region，递归收集 + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Operation &nestedOp : block) { + CollectAllNestedOps(&nestedOp, regionOps); + } + } + } +} + +static void ComputeYieldForMergedRegionV4(MergedRegion &mr) { + mr.yieldValues.clear(); + mr.resultTypes.clear(); + + // 用 DenseSet 暂存当前 region 的所有 ops + // 初始 DenseSet: 顶层 opsToMove + DenseSet regionOps; + for (Operation *op : mr.opsToMove) { + CollectAllNestedOps(op, regionOps); // 完整展开嵌套 + } + + for (Operation *op : mr.opsToMove) { + for (Value res : op->getResults()) { + + bool needsYield = false; + + for (OpOperand &use : res.getUses()) { + Operation *user = use.getOwner(); + + // 如果 user 不在当前 region，则需要 yield + if (!regionOps.contains(user)) { + needsYield = true; + break; + } + } + + if (needsYield) { + mr.yieldValues.push_back(res); + mr.resultTypes.push_back(res.getType()); + } + } + } +} + +int findTargetRegion(Operation *startOp, Block &body, + DenseMap &opToRegion) { + + SmallVector worklist{startOp}; + SmallPtrSet visited; + + while (!worklist.empty()) { + Operation *op = worklist.pop_back_val(); + if (!visited.insert(op).second) + continue; + + auto it = opToRegion.find(op); + if (it != opToRegion.end()) + return it->second; + + for (Value operand : op->getOperands()) { + if (isa(operand)) + continue; + + Operation *defOp = operand.getDefiningOp(); + if (defOp && defOp->getBlock() == &body) + worklist.push_back(defOp); + } + } + + return -1; +} + +void greedyAbsorbToRegion(Operation *startOp, int regionIdx, int lowerBound, + Block &body, DenseMap &opIndex, + DenseMap &opToRegion, + SmallVector &mergedRegions) { + + auto &mr = mergedRegions[regionIdx]; + + SmallVector worklist; + SmallPtrSet visited(mr.opsToMove.begin(), + mr.opsToMove.end()); + + // 先把 startOp 本身吸收（如果还没被吸收） + if (!opToRegion.count(startOp)) { + mr.opsToMove.push_back(startOp); + opToRegion[startOp] = regionIdx; + visited.insert(startOp); + } + + worklist.push_back(startOp); + + while (!worklist.empty()) { + Operation *op = worklist.pop_back_val(); + + for (Value operand : op->getOperands()) { + if (isa(operand)) + continue; + + Operation *defOp = operand.getDefiningOp(); + if (!defOp || defOp->getBlock() != &body) + continue; + + int defIdx = opIndex[defOp]; + + // 超过前一个 region 的末尾 + if (defIdx < lowerBound) + continue; + + auto it = opToRegion.find(defOp); + + // 不能跨到其他 region + if (it != opToRegion.end() && it->second != regionIdx) + continue; + + // 去重 + if (!visited.insert(defOp).second) + continue; + + // 吸收 defOp + mr.opsToMove.push_back(defOp); + opToRegion[defOp] = regionIdx; + worklist.push_back(defOp); + } + } +} + +SmallVector +getOperationInput(Operation *op, SmallVector dependValues, + DenseMap>> + &collectDepValueMap) { + // Analyse each Op's input + DenseSet opInput; + if (isa(op) || isa(op)) { + SmallVector regionBlocks; + if (auto ifOp = dyn_cast(op)) { + regionBlocks.push_back(&(ifOp.getThenRegion().front())); + regionBlocks.push_back(&(ifOp.getElseRegion().front())); + } else { + auto forOp = dyn_cast(op); + regionBlocks.push_back(forOp.getBody()); + } + + // recursively walk scf op + for (Block *curBlock : regionBlocks) { + for (auto &curOp : *curBlock) { + for (auto operand : + getOperationInput(&curOp, dependValues, collectDepValueMap)) { + Operation *defOp; + if (auto blockArg = dyn_cast(operand)) { + Block *ownerBlock = blockArg.getOwner(); + defOp = ownerBlock->getParentOp(); + } else { + defOp = operand.getDefiningOp(); + } + Block *defBlock = defOp->getBlock(); + + if (!(defOp == op || llvm::is_contained(regionBlocks, defBlock))) { + opInput.insert(operand); + } + } + } + } + SmallVector retVector(opInput.begin(), opInput.end()); + return retVector; + } else { + SmallVector operands = op->getOperands(); + // store ifresult value that will be replaced + for (auto operand : operands) { + if (llvm::is_contained(dependValues, operand)) { + if (collectDepValueMap.find(operand) != collectDepValueMap.end()) { + collectDepValueMap[operand].second.push_back(op); + } else { + SmallVector userOps; + userOps.push_back(op); + collectDepValueMap[operand] = {operand, userOps}; + } + } + } + return operands; + } +} + +SmallVector collectDepValuesCalculation( + DenseSet forRegionOps, DenseSet regionOps, + Operation *op, SmallVector dependValues, + DenseMap>> + &collectDepValueMap) { + DenseSet collectOps; + std::deque opStack; + bool flag = false; + + opStack.push_back(op); + while (opStack.size()) { + Operation *curOp = opStack.front(); + opStack.pop_front(); + + for (auto operand : + getOperationInput(curOp, dependValues, collectDepValueMap)) { + if (llvm::is_contained(dependValues, operand)) { + flag = true; + } + + Operation *parentOp = operand.getDefiningOp(); + if (llvm::is_contained(regionOps, parentOp)) { + opStack.push_back(parentOp); + continue; + } else if (llvm::is_contained(forRegionOps, parentOp)) { + opStack.push_back(parentOp); + collectOps.insert(parentOp); + } + } + } + + if (flag) { + SmallVector retVector(collectOps.begin(), collectOps.end()); + return retVector; + } else { + collectDepValueMap.clear(); + SmallVector emptyVector; + emptyVector.clear(); + return emptyVector; + } +} + +void copyOpsToMergedRegion( + scf::ForOp forOp, SmallVector collectOps, + MergedRegion &mergedRegion, + DenseMap>> + &collectDepValueMap) { + Block *forBodyBlock = forOp.getBody(); + OpBuilder builder(forOp); + SmallVector clonedOps; + IRMapping mapper; + + // copy calculation of ifreult value related to load/store op + int cnt = 0; + for (Operation &origOp : forBodyBlock->without_terminator()) { + if (cnt >= collectOps.size()) + break; + + if (llvm::is_contained(collectOps, &origOp)) { + builder.setInsertionPointAfter(&origOp); + + Operation *clonedOp = (&origOp)->clone(mapper); + builder.insert(clonedOp); + mapper.map(&origOp, clonedOp); + + clonedOps.push_back(clonedOp); + cnt++; + + // replace the ifresult value by new cloned op's result + SmallVector results = origOp.getResults(); + for (auto [idx, result] : llvm::enumerate(origOp.getResults())) { + if (collectDepValueMap.find(result) != collectDepValueMap.end()) { + collectDepValueMap[result].first = clonedOp->getResult(idx); + } + } + } + } + + DenseSet mergedRegionOps; + for (Operation *op : mergedRegion.opsToMove) { + CollectAllNestedOps(op, mergedRegionOps); + } + + // replace the ifresult value by new cloned op's result + for (Operation *op : mergedRegionOps) { + for (auto [idx, operand] : llvm::enumerate(op->getOperands())) { + if (collectDepValueMap.find(operand) != collectDepValueMap.end()) { + op->setOperand(idx, collectDepValueMap[operand].first); + } + } + } + + // update MergedRegion + clonedOps.append(mergedRegion.opsToMove); + mergedRegion.opsToMove = clonedOps; +} + +void copyLoadCalculation(scf::ForOp forOp, SmallVector dependValues, + SmallVector &mergedRegions) { + mlir::Operation *parentOp = forOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + auto coreTypeAttr = + scopeOp->getAttrOfType(hivm::TCoreTypeAttr::name); + // only process the vector core + if (coreTypeAttr.getTcoretype() == hivm::TCoreType::CUBE) { + return; + } + + // recursively collect all op in forOp + DenseSet forRegionOps; + for (Operation &op : forOp.getBody()->without_terminator()) { + CollectAllNestedOps(&op, forRegionOps); + } + + for (MergedRegion &mr : mergedRegions) { + DenseSet regionOps; + for (Operation *op : mr.opsToMove) { + CollectAllNestedOps(op, regionOps); + } + + for (Operation *op : regionOps) { + if (isa(op) || isa(op)) { + // recusively check that whether load/store op's operands originated + // from if results + DenseMap>> + collectDepValueMap; + SmallVector collectOps = collectDepValuesCalculation( + forRegionOps, regionOps, op, dependValues, collectDepValueMap); + copyOpsToMergedRegion(forOp, collectOps, mr, collectDepValueMap); + } + } + } +} + +// 以 forOp 的 yield value 为中心 +// 决定它应该归属哪个 mergedRegion, 然后再向前吸 operand +void ExpandMergedRegionOpsForAIV(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // 记录 block 中 op 顺序 + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) + opIndex[&op] = idx++; + + // 建立 op -> region 映射 + DenseMap opToRegion; + for (int r = 0; r < mergedRegions.size(); ++r) + for (Operation *op : mergedRegions[r].opsToMove) + opToRegion[op] = r; + + // 取 scf.yield + auto yieldOp = cast(body.getTerminator()); + + // 依次处理每个 yield value（按编号顺序） + for (Value yv : yieldOp.getOperands()) { + + Operation *defOp = yv.getDefiningOp(); + if (!defOp || defOp->getBlock() != &body) + continue; + + int targetRegion = -1; + + // 如果已经在 region 中 + auto it = opToRegion.find(defOp); + if (it != opToRegion.end()) { + targetRegion = it->second; + } else { + // 否则向前搜索确定归属 + targetRegion = findTargetRegion(defOp, body, opToRegion); + } + + if (targetRegion == -1) + continue; + + // 计算边界 lowerBound + int lowerBound = 0; + + if (targetRegion > 0) { + Operation *prevLast = mergedRegions[targetRegion - 1].opsToMove.back(); + lowerBound = opIndex[prevLast] + 1; + } + + // 真正贪心吸收 + greedyAbsorbToRegion(defOp, targetRegion, lowerBound, body, opIndex, + opToRegion, mergedRegions); + } + + // 每个 region 内按 block 顺序排序 + for (auto &mr : mergedRegions) { + llvm::sort(mr.opsToMove, [&](Operation *a, Operation *b) { + return opIndex[a] < opIndex[b]; + }); + } +} + +// 以 mergedRegion 为中心, 向前吸 operand +void ExpandMergedRegionOpsForAIC(scf::ForOp forOp, + SmallVector &mergedRegions) { + Block &body = forOp.getRegion().front(); + + // 记录每个 mergedRegion 的起始 op index + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) { + opIndex[&op] = idx++; + } + + for (int r = 0; r < mergedRegions.size(); ++r) { + MergedRegion &mr = const_cast(mergedRegions[r]); + + // 本 mergedRegion 的最早 op + Operation *firstOp = mr.opsToMove.front(); + int lowerBound = 0; + + // 边界: 前一个 mergedRegion 的最后一个 op + if (r > 0) { + Operation *prevLast = mergedRegions[r - 1].opsToMove.back(); + lowerBound = opIndex[prevLast] + 1; + } + + SmallVector worklist(mr.opsToMove.begin(), mr.opsToMove.end()); + SmallPtrSet visited(mr.opsToMove.begin(), + mr.opsToMove.end()); + + while (!worklist.empty()) { + Operation *op = worklist.pop_back_val(); + + // 往前吸收operand + for (Value operand : op->getOperands()) { + // BlockArgument + if (mlir::isa(operand)) + continue; + + Operation *defOp = operand.getDefiningOp(); + if (!defOp) + continue; + + // 不在 for body + if (defOp->getBlock() != &body) + continue; + + int defIdx = opIndex[defOp]; + + // 超出允许向前吸收的边界 + if (defIdx < lowerBound) + continue; + + // 已经在 opsToMove + if (!visited.insert(defOp).second) + continue; + + // 吸收这个 defOp + mr.opsToMove.push_back(defOp); + worklist.push_back(defOp); + } + } + + // 最后按原 block 顺序排序 + llvm::sort(mr.opsToMove, [&](Operation *a, Operation *b) { + return opIndex[a] < opIndex[b]; + }); + } +} + +static void pullInRegionDependencies(Operation *regionOp, int regionId, + DenseMap &opToRegion, + Block &body) { + + SmallVector worklist; + + // 先把 region 内的 op 放进去 + for (Region ®ion : regionOp->getRegions()) + for (Block &block : region) + for (Operation &inner : block) + worklist.push_back(&inner); + + SmallPtrSet visited; + + while (!worklist.empty()) { + Operation *innerOp = worklist.pop_back_val(); + + if (!visited.insert(innerOp).second) + continue; + + // operand 的 defining op + for (Value operand : innerOp->getOperands()) { + + Operation *def = operand.getDefiningOp(); + if (!def) + continue; + + if (def->getBlock() != &body) + continue; + + if (!opToRegion.count(def)) { + + opToRegion[def] = regionId; + + // 如果 def 也是 region-op，继续扩展 + if (def->getNumRegions() > 0) + worklist.push_back(def); + } + } + + // 继续遍历 region + for (Region &r : innerOp->getRegions()) + for (Block &b : r) + for (Operation &child : b) + worklist.push_back(&child); + } +} + +// BFS 查找某个 op 最早被哪个 region 使用 +static int findEarliestRegion(Operation *startOp, + const DenseMap &seedRegionMap, + Block &body) { + + SmallVector worklist{startOp}; + SmallPtrSet visited; + int earliestRegion = -1; + + while (!worklist.empty()) { + Operation *op = worklist.pop_back_val(); + + if (!visited.insert(op).second) + continue; + + for (Value result : op->getResults()) { + for (OpOperand &use : result.getUses()) { + Operation *user = use.getOwner(); + + if (user->getBlock() != &body) + continue; + + auto it = seedRegionMap.find(user); + if (it != seedRegionMap.end()) { + int region = it->second; + if (earliestRegion == -1 || region < earliestRegion) + earliestRegion = region; + } else { + worklist.push_back(user); + } + } + } + } + + return earliestRegion; +} + +void ExpandMergedRegionOpsForAll(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // block 内 op 顺序 + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) + opIndex[&op] = idx++; + + // seed region map + DenseMap seedRegionMap; + for (int r = 0; r < mergedRegions.size(); r++) { + for (Operation *op : mergedRegions[r].opsToMove) { + seedRegionMap[op] = r; + } + } + + // 最终 op -> region + DenseMap opToRegion = seedRegionMap; + + // ---------- Step1 顺序扫描 ---------- + for (Operation &op : body) { + + if (isa(&op)) + continue; + + if (opToRegion.count(&op)) + continue; + + int region = findEarliestRegion(&op, seedRegionMap, body); + + if (region != -1) + opToRegion[&op] = region; + } + + // ---------- Step2 region-op 依赖补全 ---------- + for (Operation &op : body) { + + auto it = opToRegion.find(&op); + if (it == opToRegion.end()) + continue; + + if (op.getNumRegions() == 0) + continue; + + pullInRegionDependencies(&op, it->second, opToRegion, body); + } + + // ---------- Step3 append op ---------- + SmallPtrSet seen; + + for (Operation &op : body) { + + auto it = opToRegion.find(&op); + if (it == opToRegion.end()) + continue; + + if (!seen.insert(&op).second) + continue; + + int region = it->second; + mergedRegions[region].opsToMove.push_back(&op); + } + + // ---------- Step4 排序 ---------- + for (auto &mr : mergedRegions) { + + llvm::sort(mr.opsToMove, [&](Operation *a, Operation *b) { + return opIndex[a] < opIndex[b]; + }); + } +} + +void ExpandMergedRegionOpsByInput(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // block 内 op 顺序 + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) + opIndex[&op] = idx++; + + // seed region map + DenseMap seedRegionMap; + for (int r = 0; r < mergedRegions.size(); r++) { + for (Operation *op : mergedRegions[r].opsToMove) { + seedRegionMap[op] = r; + } + } + + // 最终 op -> region + DenseMap opToRegion = seedRegionMap; + + // ---------- Step1 顺序扫描 ---------- + for (Operation &op : body) { + + if (isa(&op)) + continue; + + if (opToRegion.count(&op)) + continue; + + int region = findEarliestRegion(&op, seedRegionMap, body); + + if (region != -1) + opToRegion[&op] = region; + } + + // ---------- Step2 region-op 依赖补全 ---------- + for (Operation &op : body) { + + auto it = opToRegion.find(&op); + if (it == opToRegion.end()) + continue; + + if (op.getNumRegions() == 0) + continue; + + pullInRegionDependencies(&op, it->second, opToRegion, body); + } + + // ---------- Step3 append op ---------- + SmallPtrSet seen; + + for (Operation &op : body) { + + auto it = opToRegion.find(&op); + if (it == opToRegion.end()) + continue; + + if (!seen.insert(&op).second) + continue; + + int region = it->second; + mergedRegions[region].opsToMove.push_back(&op); + } + + // ---------- Step4 排序 ---------- + for (auto &mr : mergedRegions) { + + llvm::sort(mr.opsToMove, [&](Operation *a, Operation *b) { + return opIndex[a] < opIndex[b]; + }); + } +} + +static void +ExpandMergedRegionOpsByOutput(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // block 顺序（保持 IR 顺序） + DenseMap opOrder; + int idx = 0; + for (Operation &op : body) + opOrder[&op] = idx++; + + for (auto &merged : mergedRegions) { + + // 收集 region 当前产生的 value + SmallPtrSet regionValues; + + for (Operation *op : merged.opsToMove) + for (Value res : op->getResults()) + regionValues.insert(res); + + bool changed = true; + + while (changed) { + changed = false; + + for (Operation &op : body) { + + if (isa(op) || isa(op)) + continue; + + if (llvm::is_contained(merged.opsToMove, &op)) + continue; + + bool depends = false; + + for (Value operand : op.getOperands()) { + if (regionValues.contains(operand)) { + depends = true; + break; + } + } + + if (!depends) + continue; + + // 加入 region + merged.opsToMove.push_back(&op); + + // 更新 regionValues + for (Value res : op.getResults()) + regionValues.insert(res); + + changed = true; + } + } + + // 排序保持原 block 顺序 + llvm::sort(merged.opsToMove, [&](Operation *a, Operation *b) { + return opOrder[a] < opOrder[b]; + }); + } +} + +static void MoveIndependentOpsIntoIf(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // 记录哪些 op 已经在 region 里 + SmallPtrSet alreadyAssigned; + + for (auto &mr : mergedRegions) + for (Operation *op : mr.opsToMove) + alreadyAssigned.insert(op); + + // 记录 iter_arg -> region + DenseMap iterArgToRegion; + + for (int r = 0; r < mergedRegions.size(); r++) { + for (Operation *op : mergedRegions[r].opsToMove) { + + for (Value operand : op->getOperands()) { + + if (auto barg = mlir::dyn_cast(operand)) { + + if (barg.getOwner() == &body) + iterArgToRegion[barg] = r; + } + } + } + } + + // block 顺序 + DenseMap opIndex; + int idx = 0; + for (Operation &op : body) + opIndex[&op] = idx++; + + // 扫描所有 op + for (Operation &op : body) { + + if (isa(op) || isa(op)) + continue; + + if (alreadyAssigned.contains(&op)) + continue; + + int targetRegion = -1; + + // 看 operand 是否来自 iter_arg + for (Value operand : op.getOperands()) { + + if (auto barg = mlir::dyn_cast(operand)) { + + if (barg.getOwner() != &body) + continue; + + auto it = iterArgToRegion.find(barg); + if (it != iterArgToRegion.end()) { + + targetRegion = it->second; + break; + } + } + } + + if (targetRegion == -1) + continue; + + mergedRegions[targetRegion].opsToMove.push_back(&op); + alreadyAssigned.insert(&op); + } + + // 排序保持 block 顺序 + for (auto &mr : mergedRegions) { + + llvm::sort(mr.opsToMove, [&](Operation *a, Operation *b) { + return opIndex[a] < opIndex[b]; + }); + } +} + +// 暴力包裹 +static void +ExpandMergedRegionOpsGreedyMaximum(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // 记录哪些 op 已经属于 region + DenseSet regionOps; + + for (auto ®ion : mergedRegions) + for (Operation *op : region.opsToMove) + regionOps.insert(op); + + // block op 列表 + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + DenseMap opIndex; + for (int i = 0; i < ops.size(); i++) + opIndex[ops[i]] = i; + + for (auto ®ion : mergedRegions) { + + if (region.opsToMove.empty()) + continue; + + // 找到 region 在 block 中的范围 + int start = ops.size(); + int end = -1; + + for (Operation *op : region.opsToMove) { + int idx = opIndex[op]; + start = std::min(start, idx); + end = std::max(end, idx); + } + + SmallVector newOps; + + // ---------- backward 扩展 ---------- + for (int i = start - 1; i >= 0; i--) { + Operation *op = ops[i]; + + if (isa(op)) + break; + + if (regionOps.contains(op)) + break; + + newOps.push_back(op); + } + + // ---------- forward 扩展 ---------- + for (int i = end + 1; i < ops.size(); i++) { + Operation *op = ops[i]; + + if (isa(op)) + break; + + if (regionOps.contains(op)) + break; + + newOps.push_back(op); + } + + // 加入 region + for (Operation *op : newOps) { + region.opsToMove.push_back(op); + regionOps.insert(op); + } + } + + // 最后保持 block 顺序 + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +static void CollectForYieldRelatedOps(scf::ForOp forOp, + SmallVector &mergedRegions, + DenseSet &yieldRelatedOps) { + + Block &body = forOp.getRegion().front(); + + // 已经属于 region 的 op + DenseSet regionOps; + for (auto ®ion : mergedRegions) + for (Operation *op : region.opsToMove) + regionOps.insert(op); + + auto yield = cast(body.getTerminator()); + + SmallVector worklist; + DenseSet visited; + + // 初始化 worklist + for (Value v : yield.getOperands()) + worklist.push_back(v); + + while (!worklist.empty()) { + Value v = worklist.pop_back_val(); + + if (!visited.insert(v).second) + continue; + + Operation *def = v.getDefiningOp(); + if (!def) + continue; + + // 只处理 for body 内的 op + if (def->getBlock() != &body) + continue; + + // 已经在 region 内 + if (regionOps.contains(def)) + continue; + + // 记录 + if (yieldRelatedOps.insert(def).second) { + + // 继续向上找依赖 + for (Value operand : def->getOperands()) + worklist.push_back(operand); + } + } +} + +// 贪心吸收region前后的op +static void +ExpandMergedRegionOpsGreedy(scf::ForOp forOp, + SmallVector &mergedRegions, + DenseSet &skipOps) { + + Block &body = forOp.getRegion().front(); + + // 记录哪些 op 已经属于 region + DenseSet regionOps; + for (auto ®ion : mergedRegions) + for (Operation *op : region.opsToMove) + regionOps.insert(op); + + // block op 列表 + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + // op -> index + DenseMap opIndex; + for (int i = 0; i < ops.size(); i++) + opIndex[ops[i]] = i; + + for (auto ®ion : mergedRegions) { + + if (region.opsToMove.empty()) + continue; + + // 找到 region 在 block 中的范围 + int start = ops.size(); + int end = -1; + + for (Operation *op : region.opsToMove) { + int idx = opIndex[op]; + start = std::min(start, idx); + end = std::max(end, idx); + } + + SmallVector newOps; + + // ---------- backward 扩展 ---------- + for (int i = start - 1; i >= 0; i--) { + Operation *op = ops[i]; + + // block terminator + if (isa(op)) + break; + + // 遇到其他 region 的 op + if (regionOps.contains(op)) + break; + + // yield 关联 op，跳过但继续扫描 + if (skipOps.contains(op)) + continue; + + newOps.push_back(op); + } + + // ---------- forward 扩展 ---------- + for (int i = end + 1; i < ops.size(); i++) { + Operation *op = ops[i]; + + // block terminator + if (isa(op)) + break; + + // 遇到其他 region 的 op + if (regionOps.contains(op)) + break; + + // yield 关联 op，跳过 + if (skipOps.contains(op)) + continue; + + newOps.push_back(op); + } + + // 加入 region + for (Operation *op : newOps) { + region.opsToMove.push_back(op); + regionOps.insert(op); + } + } + + // 最后保持 block 顺序 + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +// 贪心吸收region前面的op +static void +ExpandMergedRegionOpsGreedyV2(scf::ForOp forOp, + SmallVector &mergedRegions, + DenseSet &skipOps) { + + Block &body = forOp.getRegion().front(); + + // block op 列表 + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + // op -> index + DenseMap opIndex; + for (int i = 0; i < ops.size(); i++) + opIndex[ops[i]] = i; + + // 记录哪些 op 已经属于 region + DenseSet regionOps; + for (auto ®ion : mergedRegions) + for (Operation *op : region.opsToMove) + regionOps.insert(op); + + for (int r = 0; r < mergedRegions.size(); r++) { + + auto ®ion = mergedRegions[r]; + if (region.opsToMove.empty()) + continue; + + // ---------- 当前 region block 范围 ---------- + int start = ops.size(); + int end = -1; + + for (Operation *op : region.opsToMove) { + int idx = opIndex[op]; + start = std::min(start, idx); + end = std::max(end, idx); + } + + // ---------- 前一个 region 的末尾 ---------- + int prevEnd = -1; + + if (r > 0 && !mergedRegions[r - 1].opsToMove.empty()) { + for (Operation *op : mergedRegions[r - 1].opsToMove) { + prevEnd = std::max(prevEnd, opIndex[op]); + } + } + + SmallVector newOps; + + // ---------- backward expand ---------- + for (int i = start - 1; i > prevEnd; i--) { + + Operation *op = ops[i]; + + // terminator + if (isa(op)) + break; + + // 已属于 region + if (regionOps.contains(op)) + break; + + // yield chain op + if (skipOps.contains(op)) + continue; + + newOps.push_back(op); + } + + // ---------- 加入 region ---------- + for (Operation *op : newOps) { + region.opsToMove.push_back(op); + regionOps.insert(op); + } + } + + // ---------- 保持 block 顺序 ---------- + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +// 贪心吸收region前面的op +static void +ExpandMergedRegionOpsGreedyV2ForAIC(scf::ForOp forOp, + SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + + // block op 列表 + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + // op -> index + DenseMap opIndex; + for (int i = 0; i < ops.size(); i++) + opIndex[ops[i]] = i; + + // 记录哪些 op 已经属于 region + DenseSet regionOps; + for (auto ®ion : mergedRegions) + for (Operation *op : region.opsToMove) + regionOps.insert(op); + + for (int r = 0; r < mergedRegions.size(); r++) { + + auto ®ion = mergedRegions[r]; + if (region.opsToMove.empty()) + continue; + + // ---------- 当前 region block 范围 ---------- + int start = ops.size(); + int end = -1; + + for (Operation *op : region.opsToMove) { + int idx = opIndex[op]; + start = std::min(start, idx); + end = std::max(end, idx); + } + + // ---------- 前一个 region 的末尾 ---------- + int prevEnd = -1; + + if (r > 0 && !mergedRegions[r - 1].opsToMove.empty()) { + for (Operation *op : mergedRegions[r - 1].opsToMove) { + prevEnd = std::max(prevEnd, opIndex[op]); + } + } + + SmallVector newOps; + + // ---------- backward expand ---------- + for (int i = start - 1; i > prevEnd; i--) { + + Operation *op = ops[i]; + + // terminator + if (isa(op)) + break; + + // 已属于 region + if (regionOps.contains(op)) + break; + + newOps.push_back(op); + } + + // ---------- 加入 region ---------- + for (Operation *op : newOps) { + region.opsToMove.push_back(op); + regionOps.insert(op); + } + } + + // ---------- 保持 block 顺序 ---------- + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +static void MoveForYieldOpIntoRegion(scf::ForOp forOp, + DenseSet &yieldRelatedOps, + SmallVector &mergedRegions) { + + DenseMap opToRegion; + + for (int i = 0; i < mergedRegions.size(); i++) + for (Operation *op : mergedRegions[i].opsToMove) + opToRegion[op] = i; + + auto yield = cast(forOp.getBody()->getTerminator()); + + for (int i = 0; i < yield.getNumOperands(); i++) { + + Value iterArg = forOp.getRegionIterArgs()[i]; + Value yieldVal = yield.getOperand(i); + + Operation *def = yieldVal.getDefiningOp(); + if (!def) + continue; + + if (!yieldRelatedOps.contains(def)) + continue; + + int targetRegion = -1; + + for (Operation *user : iterArg.getUsers()) { + + if (opToRegion.count(user)) { + targetRegion = opToRegion[user]; + break; + } + } + + if (targetRegion == -1) + continue; + + SmallVector stack; + stack.push_back(def); + + while (!stack.empty()) { + Operation *op = stack.pop_back_val(); + + if (!yieldRelatedOps.contains(op)) + continue; + + mergedRegions[targetRegion].opsToMove.push_back(op); + + yieldRelatedOps.erase(op); + + for (Value operand : op->getOperands()) { + if (Operation *dep = operand.getDefiningOp()) + stack.push_back(dep); + } + } + } + + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +static void +MoveRemainingYieldOpsToPrevRegion(scf::ForOp forOp, + DenseSet &yieldRelatedOps, + SmallVector &mergedRegions) { + + if (yieldRelatedOps.empty()) + return; + + Block &body = forOp.getRegion().front(); + + // op -> region index + DenseMap opToRegion; + for (int i = 0; i < mergedRegions.size(); i++) + for (Operation *op : mergedRegions[i].opsToMove) + opToRegion[op] = i; + + // block 顺序 + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + DenseMap opIndex; + for (int i = 0; i < ops.size(); i++) + opIndex[ops[i]] = i; + + for (Operation *op : yieldRelatedOps) { + + if (op->getBlock() != &body) + continue; + + int idx = opIndex[op]; + + int targetRegion = -1; + + // 向前找最近的 region + for (int i = idx - 1; i >= 0; i--) { + Operation *prev = ops[i]; + + if (opToRegion.count(prev)) { + targetRegion = opToRegion[prev]; + break; + } + } + + if (targetRegion == -1) + continue; + + mergedRegions[targetRegion].opsToMove.push_back(op); + } + + // 排序 + 去重 + for (auto ®ion : mergedRegions) { + + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +static void MoveIndependentOpsIntoRegionBackwardV2( + scf::ForOp forOp, SmallVector &mergedRegions) { + + Block &body = forOp.getRegion().front(); + SmallVector ops; + for (Operation &op : body) + ops.push_back(&op); + + DenseMap opToRegion; + for (int i = 0; i < mergedRegions.size(); i++) + for (Operation *op : mergedRegions[i].opsToMove) + opToRegion[op] = i; + + // ----------- 收集移动计划 ----------- + DenseMap movePlan; + + for (int i = 0; i < mergedRegions.size(); i++) { + MergedRegion ®ion = mergedRegions[i]; + if (region.opsToMove.empty()) + continue; + + Operation *firstOp = region.opsToMove.front(); + Operation *lastOp = region.opsToMove.back(); + auto itFirst = std::find(ops.begin(), ops.end(), firstOp); + auto itLast = std::find(ops.begin(), ops.end(), lastOp); + if (itFirst == ops.end() || itLast == ops.end()) + continue; + + int startIdx = std::distance(ops.begin(), itFirst); + int endIdx = std::distance(ops.begin(), itLast); + + // ----------- 收集 wait-set 区间 ----------- + SmallVector> waitIntervals; + bool inWait = false; + int begin = -1; + for (int j = startIdx; j <= endIdx; j++) { + Operation *op = ops[j]; + if (op->getName().getStringRef().contains("sync_block_wait")) { + inWait = true; + begin = j + 1; + continue; + } + if (op->getName().getStringRef().contains("sync_block_set") && inWait) { + inWait = false; + waitIntervals.push_back({begin, j - 1}); + } + } + auto isInWaitSet = [&](int idx) { + for (auto &p : waitIntervals) + if (idx >= p.first && idx <= p.second) + return true; + return false; + }; + + // ----------- 从后往前扫描 region 内的 op ----------- + for (int j = endIdx; j >= startIdx; j--) { + Operation *op = ops[j]; + if (isa(op) || isInWaitSet(j)) + continue; + + // ---------- operand 是否依赖本 region ---------- + bool dependCurrentRegion = false; + for (Value operand : op->getOperands()) { + Operation *def = operand.getDefiningOp(); + if (!def) + continue; + if (std::find(region.opsToMove.begin(), region.opsToMove.end(), def) != + region.opsToMove.end()) { + dependCurrentRegion = true; + break; + } + } + if (dependCurrentRegion) + continue; + + // ---------- 当前 region 后续是否使用 ---------- + bool usedLaterInSameRegion = false; + for (Value result : op->getResults()) + for (Operation *user : result.getUsers()) + if (std::find(region.opsToMove.begin(), region.opsToMove.end(), + user) != region.opsToMove.end() && + std::find(region.opsToMove.begin(), region.opsToMove.end(), op) < + std::find(region.opsToMove.begin(), region.opsToMove.end(), + user)) { + usedLaterInSameRegion = true; + break; + } + if (usedLaterInSameRegion) + continue; + + // ---------- 找使用该 op 的后续 region ---------- + int targetRegion = -1; + for (int k = i + 1; k < mergedRegions.size(); ++k) { + for (Operation *candidate : mergedRegions[k].opsToMove) + for (Value operand : candidate->getOperands()) + if (operand.getDefiningOp() == op) { + targetRegion = k; + break; + } + if (targetRegion != -1) + break; + if (targetRegion != -1) + break; + } + if (targetRegion == -1) + continue; + + movePlan[op] = targetRegion; + // llvm::outs() << "MJ: plan move " << *op + // << " -> region " << targetRegion << "\n"; + } + } + + // ----------- 统一应用移动 ----------- + for (auto &it : movePlan) { + Operation *op = it.first; + int targetRegionIdx = it.second; + MergedRegion &targetRegion = mergedRegions[targetRegionIdx]; + // 更新数据结构 + targetRegion.opsToMove.push_back(op); + + llvm::outs() << "MJ: move " << *op << " -> region " << targetRegionIdx + << "\n"; + } + + // ----------- 更新原 region 的 opsToMove ----------- + for (int i = 0; i < mergedRegions.size(); ++i) { + MergedRegion ®ion = mergedRegions[i]; + SmallVector newOps; + for (Operation *op : region.opsToMove) { + auto it = movePlan.find(op); + if (it == movePlan.end() || it->second == i) { + // 没有移动计划，或者移动的目标就是自己，保留 + newOps.push_back(op); + } + } + region.opsToMove.swap(newOps); + } + + // ----------- 排序 + 去重 ----------- + for (auto ®ion : mergedRegions) { + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +// // debug: 如果一个forop的第一个region的最后3条op是%27 = tt.expand_dims %25#1 +// {axis = 1 : i32} : tensor<64xf32> -> tensor<64x1xf32> +// %28 = tt.broadcast %27 : tensor<64x1xf32> -> tensor<64x128xf32> +// %29 = arith.mulf %arg10, %28 : tensor<64x128xf32> +// 直接放到第2个region里 +static void TempChange(scf::ForOp forOp, + SmallVector &mergedRegions) { + + if (mergedRegions.size() < 2) + return; + + auto &srcRegion = mergedRegions[0]; + auto &dstRegion = mergedRegions[1]; + + if (srcRegion.opsToMove.size() < 3) + return; + + Operation *op1 = srcRegion.opsToMove[srcRegion.opsToMove.size() - 3]; + Operation *op2 = srcRegion.opsToMove[srcRegion.opsToMove.size() - 2]; + Operation *op3 = srcRegion.opsToMove[srcRegion.opsToMove.size() - 1]; + + // ---------- pattern 匹配 ---------- + if (!op1->getName().getStringRef().contains("tt.expand_dims")) + return; + + if (!op2->getName().getStringRef().contains("tt.broadcast")) + return; + + if (!op3->getName().getStringRef().contains("arith.mulf")) + return; + + llvm::outs() << "TempChange triggered\n"; + + SmallVector opsToMove = {op1, op2, op3}; + + // ---------- 移动到 region2 末尾 ---------- + for (Operation *op : opsToMove) { + dstRegion.opsToMove.push_back(op); + llvm::outs() << "TempChange move: " << *op << "\n"; + } + + // ---------- 从 region1 删除 ---------- + srcRegion.opsToMove.resize(srcRegion.opsToMove.size() - 3); + + // ---------- 排序 ---------- + for (auto ®ion : mergedRegions) { + llvm::sort(region.opsToMove, [](Operation *a, Operation *b) { + return a->isBeforeInBlock(b); + }); + + region.opsToMove.erase( + std::unique(region.opsToMove.begin(), region.opsToMove.end()), + region.opsToMove.end()); + } +} + +static void sortOperationsByDataFlow(llvm::SmallVector &ops) { + llvm::DenseSet visited; + llvm::SmallVector result; + + std::function dfs = [&](Operation *op) { + if (!visited.insert(op).second) + return; + + for (Value operand : op->getOperands()) { + if (Operation *def = operand.getDefiningOp()) { + if (llvm::is_contained(ops, def)) + dfs(def); + } + } + + result.push_back(op); + }; + + for (Operation *op : ops) + dfs(op); + + ops.assign(result.begin(), result.end()); +} + +static void rewriteOperandsRecursively(Operation *op, + DenseMap &valueMap) { + + // 1 rewrite 当前 op 的 operands + for (OpOperand &operand : op->getOpOperands()) { + Value v = operand.get(); + auto it = valueMap.find(v); + if (it != valueMap.end()) + operand.set(it->second); + } + + // 2 递归进入 region + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Operation &nestedOp : block) { + rewriteOperandsRecursively(&nestedOp, valueMap); + } + } + } +} + +static void CopyOpsToAfterwardRegions( + SmallVector &mergedRegions, + DenseMap &yieldMap, + DenseMap &cloneAndOriYieldMap, + SmallVector &copiedForOps) { + + if (mergedRegions.size() <= 1) + return; + + // 先整理一个 set，方便判断哪些 op 是 yield defining op + DenseSet yieldDefOps; + for (auto &it : yieldMap) + yieldDefOps.insert(it.second); + + // 倒序遍历 region + for (int i = mergedRegions.size() - 1; i >= 0; --i) { + MergedRegion &curRegion = mergedRegions[i]; + + DenseMap valueMap; + SmallVector clonedOps; + + // 遍历前面的 region + for (int k = 0; k < i; ++k) { + MergedRegion &prevRegion = mergedRegions[k]; + + int waitSetLevel = 0; + + for (Operation *op : prevRegion.opsToMove) { + + if (isa(op)) { + waitSetLevel++; + continue; + } + + if (isa(op)) { + waitSetLevel = std::max(waitSetLevel - 1, 0); + continue; + } + + if (waitSetLevel > 0) + continue; + + if (isa(op)) + continue; + + IRMapping mapper; + + for (auto result : op->getResults()) + if (valueMap.count(result)) + mapper.map(result, valueMap[result]); + + Operation *insertPoint = + curRegion.opsToMove.empty() ? nullptr : curRegion.opsToMove.front(); + + OpBuilder builder(insertPoint ? insertPoint : op); + + Operation *cloned = builder.clone(*op, mapper); + + // 记录 result mapping + for (auto it : llvm::zip(op->getResults(), cloned->getResults())) + valueMap[std::get<0>(it)] = std::get<1>(it); + + // 如果这个 op 是 yield defining op，记录 clone -> original + if (yieldDefOps.contains(op)) { + cloneAndOriYieldMap[cloned] = op; + } + + // 记录copy的for op + if (auto forOp = dyn_cast(cloned)) { + copiedForOps.push_back(forOp); + } + + clonedOps.push_back(cloned); + } + } + + // 插入到当前 region 开头 + curRegion.opsToMove.insert(curRegion.opsToMove.begin(), clonedOps.begin(), + clonedOps.end()); + + // rebuild SSA + for (Operation *op : curRegion.opsToMove) { + rewriteOperandsRecursively(op, valueMap); + } + + // 排序保证拓扑顺序 + sortOperationsByDataFlow(curRegion.opsToMove); + } +} + +/// 记录 forOp 的 yield value 与其原始生成的 op 的映射 +static void GetYieldMap(scf::ForOp forOp, + DenseMap &yieldMap) { + yieldMap.clear(); + + // 取 forOp body 的 scf.yield + auto yieldOp = dyn_cast(forOp.getBody()->getTerminator()); + if (!yieldOp) + return; + + for (Value yieldVal : yieldOp.getOperands()) { + // 获取生成 yieldVal 的原始 op + Operation *defOp = yieldVal.getDefiningOp(); + + // 对 block arg（可能是 iter_arg）没有 definingOp 的情况，可以跳过或直接记录 + // nullptr + if (!defOp) + continue; + + yieldMap[yieldVal] = defOp; + } +} + +static Value findIterArgForAIC(Value v, scf::ForOp forOp) { + while (true) { + if (auto arg = dyn_cast(v)) { + if (arg.getOwner() == forOp.getBody()) + return v; + return Value(); + } + + Operation *def = v.getDefiningOp(); + if (!def) + return Value(); + + if (def->getNumOperands() == 0) + return Value(); + + v = def->getOperand(0); + } +} + +static Operation * +findCloneOfYieldOp(Operation *oriYieldOp, + DenseMap &cloneAndOriYieldMap, + MergedRegion ®ion) { + + for (Operation *op : region.opsToMove) { + auto it = cloneAndOriYieldMap.find(op); + if (it != cloneAndOriYieldMap.end() && it->second == oriYieldOp) + return op; + } + return nullptr; +} + +static void RebuildForYielValuesForAIC( + scf::ForOp forOp, SmallVector &mergedRegions, + DenseMap &yieldMap, + DenseMap &cloneAndOriYieldMap) { + + auto yieldOp = cast(forOp.getBody()->getTerminator()); + + for (MergedRegion ®ion : mergedRegions) { + + triton::DotOp dotOp = nullptr; + + for (Operation *op : region.opsToMove) { + if (auto d = dyn_cast(op)) { + dotOp = d; + break; + } + } + + if (!dotOp) + continue; + + // 处理 dot operand + for (Value operand : dotOp->getOperands()) { + + Value iterArg = findIterArgForAIC(operand, forOp); + if (!iterArg) + continue; + + auto arg = cast(iterArg); + int idx = arg.getArgNumber(); + + if (idx >= yieldOp.getNumOperands()) + continue; + + Value oriYieldValue = yieldOp.getOperand(idx); + + auto it = yieldMap.find(oriYieldValue); + if (it == yieldMap.end()) + continue; + + Operation *oriYieldOp = it->second; + + Operation *cloneOp = + findCloneOfYieldOp(oriYieldOp, cloneAndOriYieldMap, region); + + if (!cloneOp) + continue; + + yieldOp.setOperand(idx, cloneOp->getResult(0)); + } + } +} + +void ExpandMergedRegionOps(scf::ForOp forOp, + SmallVector &mergedRegions, + SmallVector &copiedForOps) { + bool isInAIV = false; + auto scopeOp = forOp->getParentOfType(); + if (!scopeOp) + return; + + auto coreTypeAttr = + scopeOp->getAttrOfType(hivm::TCoreTypeAttr::name); + + if (coreTypeAttr.getTcoretype() == hivm::TCoreType::VECTOR) { + isInAIV = true; + } + + if (isInAIV) { + DenseSet yieldRelatedOps; + + // 1 收集 yield 相关 op + CollectForYieldRelatedOps(forOp, mergedRegions, yieldRelatedOps); + + // 2 greedy 扩展 + // ExpandMergedRegionOpsGreedy(forOp, mergedRegions, yieldRelatedOps); + ExpandMergedRegionOpsGreedyV2(forOp, mergedRegions, yieldRelatedOps); + + // 3 与前面wait-set region独立的op应该被放入后面的关联的region + MoveIndependentOpsIntoRegionBackwardV2(forOp, mergedRegions); + + // 4 根据 iter_arg 使用位置放入 region + MoveForYieldOpIntoRegion(forOp, yieldRelatedOps, mergedRegions); + + // 5 剩余 yield chain 放入前一个 region + MoveRemainingYieldOpsToPrevRegion(forOp, yieldRelatedOps, mergedRegions); + } else { // AIC单独处理, 避免出现CUBE内的tensor变量依赖 + // 用Map记录原始的for yield op的的映射 + DenseMap yieldMap; + GetYieldMap(forOp, yieldMap); + + llvm::outs() << "YieldMap:\n"; + for (auto it : yieldMap) { + llvm::outs() << *(it.second) << "\n"; + } + + // 2 greedy 扩展, yield value后续处理 + ExpandMergedRegionOpsGreedyV2ForAIC(forOp, mergedRegions); + + // 复制当前region的除tt.dot、以及[wait - + // set]之间的op到后续的所有MergedRegion 倒序实现 + // 记录clone和original的yield对应op的map + DenseMap cloneAndOriYieldMap; + CopyOpsToAfterwardRegions(mergedRegions, yieldMap, cloneAndOriYieldMap, + copiedForOps); + + // 4 + // 先确定每个MergedRegion的tt.dot的operand的来源是for的哪个iter_arg(递归查找), + // 假设为%arg0, 依据yieldMap可以得到oriYield 遍历当前MergedRegion的所有op, + // 确定哪条op对应的cloneAndOriYieldMap的second是oriYield, 假设为%45 + // 最后替换for yield op对应位置的operand为%45 + RebuildForYielValuesForAIC(forOp, mergedRegions, yieldMap, + cloneAndOriYieldMap); + } +} + +void MergeWaitSetRegions(SmallVector ®ions, + SmallVector &merged) { + for (int i = 0; i < regions.size();) { + MergedRegion mr; + mr.regions.push_back(®ions[i]); + mr.opsToMove.append(regions[i].opsToMove); + + int j = i; + while (!regions[j].hasCopyOrFixpipe && j + 1 < regions.size()) { + j++; + mr.regions.push_back(®ions[j]); + mr.opsToMove.append(regions[j].opsToMove); + } + + merged.push_back(std::move(mr)); + i = j + 1; + } + + for (MergedRegion &mr : merged) { + SmallPtrSet regionValues; + SmallPtrSet opSet; + + for (Operation *op : mr.opsToMove) + opSet.insert(op); + + for (Operation *op : mr.opsToMove) { + for (Value v : op->getResults()) { + bool usedOutside = false; + for (OpOperand &use : v.getUses()) { + Operation *user = use.getOwner(); + if (!opSet.contains(user) && user->getBlock() == op->getBlock()) { + usedOutside = true; + break; + } + } + if (usedOutside) { + mr.yieldValues.push_back(v); + mr.resultTypes.push_back(v.getType()); + } + } + } + } +} + +void GetBlockInfos(SmallVector ®ions, Block &body) { + for (auto it = body.begin(); it != body.end();) { + Operation *op = &*it; + + auto waitOp = dyn_cast(op); + if (!waitOp) { + it++; + continue; + } + + auto pipeS = hivm::PipeAttr::get(op->getContext(), hivm::PIPE::PIPE_S); + if (auto syncWait = dyn_cast(op)) { + if ((syncWait.getTpipe() == pipeS || syncWait.getPipe() == pipeS) && + (std::next(it) != body.end())) { + if (isa(&*std::next(it)) || + isa(&*std::next(it))) + return; + } + } + + Operation *lastSetOp = nullptr; + + // 扫描到下一个 wait, 收集所有 set + auto curIt = std::next(it); + auto endIt = curIt; + int setOpCount = 0; + SmallVector opsInRegion; + for (; curIt != body.end(); ++curIt) { + Operation *curOp = &*curIt; + if (isa(curOp) && setOpCount >= 1) + break; + if (isa(curOp)) { + setOpCount++; + endIt = curIt; // setop的位置 + lastSetOp = curOp; // 最后一个 set + } + } + + if (!lastSetOp) { + it = curIt; + continue; + } // 没有 set, 不包 + + // 收集 [wait, ..., lastSet] 之间的 ops + bool hasCopyOrFixpipe = false; + for (auto it2 = it; it2 != std::next(endIt); ++it2) { + Operation *curOp = &*it2; + opsInRegion.push_back(curOp); + if (isa(curOp) || isa(curOp)) { + hasCopyOrFixpipe = true; + } + } + + it = endIt++; + regions.push_back({waitOp, lastSetOp, opsInRegion, hasCopyOrFixpipe}); + } +} + +Value findIterArg(Value v, Type t) { + SmallVector worklist = {v}; + SmallPtrSet visited; + + while (!worklist.empty()) { + Value cur = worklist.front(); + worklist.erase(worklist.begin()); + if (!visited.insert(cur).second) + continue; + + // 匹配scf.for原始迭代参数, 直接返回 + if (auto b = mlir::dyn_cast(cur)) { + auto forOp = mlir::dyn_cast(b.getOwner()->getParentOp()); + if (forOp && b.getType() == t) { + for (Value iterArg : forOp.getRegionIterArgs()) { + if (iterArg.getAsOpaquePointer() == b.getAsOpaquePointer()) { + return b; + } + } + } + } + + Operation *defOp = cur.getDefiningOp(); + if (!defOp) + continue; + + // 核心逻辑：如果当前值是scf.if的结果 + // 进入then块找源头 + if (auto ifOp = mlir::dyn_cast(defOp)) { + Block &thenBlock = ifOp.getThenRegion().front(); + // 找到then块最后一个op（scf.yield） + // 取其operands（即ifOp结果的源头值） + for (auto &innerOp : llvm::reverse(thenBlock)) { + if (auto yieldOp = mlir::dyn_cast(&innerOp)) { + // 按索引匹配: cur是ifOp的第n个结果, 取yieldOp的第n个operand + for (auto [idx, res] : llvm::enumerate(ifOp.getResults())) { + if (res.getAsOpaquePointer() == cur.getAsOpaquePointer()) { + Value srcVal = yieldOp.getOperand(idx); + if (!visited.count(srcVal)) + worklist.push_back(srcVal); + break; + } + } + break; // 找到yield即退出, 无需遍历其他op + } + } + } else { + // 非if结果值 + // 正常往前追溯operands + for (Value operand : defOp->getOperands()) { + if (!visited.count(operand)) + worklist.push_back(operand); + } + } + } + + llvm::outs() << "未找到迭代参数, 返回原值: "; + v.print(llvm::outs()); + llvm::outs() << "\n"; + return v; +} + +// 如果 v 最终被 scf.for 的 yield 使用 +// → 返回对应的 forOp 的 iter_arg +// 如果 v 只是流向后面的 wait-set region / 其他 op +// → 直接返回原值 v +Value findIterArgForAll(Value v, Type t) { + for (Operation *user : v.getUsers()) { + + if (auto yieldOp = dyn_cast(user)) { + + if (auto forOp = dyn_cast(yieldOp->getParentOp())) { + + for (auto [idx, operand] : llvm::enumerate(yieldOp.getOperands())) { + + if (operand.getAsOpaquePointer() == v.getAsOpaquePointer()) { + + Value iterArg = forOp.getRegionIterArgs()[idx]; + + if (iterArg.getType() == t) + return iterArg; + } + } + } + } + } + + return v; +} + +void FindDependValues(SmallVector &dependValues, + SmallVector mergedRegions) { + dependValues.clear(); + for (auto &curMR : mergedRegions) { + for (Value yieldValue : curMR.yieldValues) { + // llvm::outs() << "yieldValue: "<< yieldValue << "\n"; + // 遍历当前区域的yieldValue的所有user OP，判断是否存在依赖关系 + for (OpOperand &use : yieldValue.getUses()) { + Operation *userOp = use.getOwner(); + + // llvm::outs() << "userOp: "<< *userOp << "\n"; + bool isUserInOtherRegion = false; + for (auto &otherMR : mergedRegions) { + // 跳过当前区域，只检查yieldValue是否被其他区域使用 + if (&otherMR == &curMR) + continue; + + // 只要有一个 userOp在 otherMR 的 opsToMove + // 列表中，就认为是dependValue llvm::outs() << "judge comtain\n"; for + // (size_t k = 0; k < otherMR.opsToMove.size(); k++) { + // llvm::outs() << "otherMR op: " << *(otherMR.opsToMove[k]) << + // "\n"; + // } + // llvm::outs() << "otherMR end\n"; + + // if (llvm::is_contained(otherMR.opsToMove, userOp)) { + // isUserInOtherRegion = true; + // llvm::outs() << "is_contained\n"; + // break; + // } + + // 用 DenseSet 暂存当前 region 的所有 ops + // 初始 DenseSet: 顶层 opsToMove + DenseSet otherOps; + for (Operation *op : otherMR.opsToMove) { + CollectAllNestedOps(op, otherOps); // 完整展开嵌套 + } + if (otherOps.contains(userOp)) { + isUserInOtherRegion = true; + break; + } + } + + // 无重复的添加依赖变量 + if (isUserInOtherRegion) { + if (!llvm::is_contained(dependValues, yieldValue)) { + dependValues.push_back(yieldValue); + } + break; + } + } + } + } +} + +void UpdateMergedRegionsWithNewForOp(SmallVector &mergedRegions, + IRMapping &mapper) { + for (auto &mr : mergedRegions) { + // WaitSetRegion 后续已经不使用了，直接释放，否则会出现野指针 + SmallVector newRegions; + newRegions.clear(); + mr.regions = newRegions; + // // 更新 opsToMove 列表 + // llvm::outs() << "before \n"; + // for (auto &op : mr.opsToMove) { + // llvm::outs() << "opsToMove: " << op << ", " << *op << '\n'; + // } + SmallVector newOpsToMove; + newOpsToMove.clear(); + for (Operation *op : mr.opsToMove) { + if (op) { + Operation *newOp = mapper.lookupOrNull(op); + newOpsToMove.push_back(newOp); + } + } + mr.opsToMove = newOpsToMove; + // llvm::outs() << "after \n"; + // for (auto &op : mr.opsToMove) { + // llvm::outs() << "opsToMove: " << op << ", " << *op << '\n'; + // } + // 更新 yieldValues 列表 + SmallVector newYieldValues; + newYieldValues.clear(); + for (Value v : mr.yieldValues) { + if (v) { + newYieldValues.push_back(mapper.lookupOrNull(v)); + } + } + mr.yieldValues = newYieldValues; + // resultTypes 是type 类型，无需更新 + } +} + +void AddArgsForDependValues(scf::ForOp forOp, SmallVector &dependValues, + SmallVector &mergedRegions, + ModuleOp module) { + OpBuilder moduleBuilder(module.getContext()); + SmallVector valueTypes; + valueTypes.clear(); + + if (dependValues.empty()) { + return; + } else { + for (Value v : dependValues) { + Type valueType = v.getType(); + valueTypes.push_back(valueType); + } + } + + // 为每个 dependValue 创建一个初始值（可能不存在相同shape和type的常量tensor） + SmallVector initTensors; + initTensors.clear(); + module.walk([&](Operation *op) { + if (auto constOp = dyn_cast(op)) { + moduleBuilder.setInsertionPoint(constOp); + for (Type valueType : valueTypes) { + auto tensorType = dyn_cast(valueType); + triton::PointerType ptrType; + ptrType = + (tensorType) + ? dyn_cast(tensorType.getElementType()) + : dyn_cast(valueType); + if (ptrType) { + // 如果依赖变量是一个ptr类型 + // 1. 创建 i64 0 + // 2. cast 成 !tt.ptr<...> + Value zero = moduleBuilder.create( + constOp.getLoc(), 0, 64); + Value ptrValue = moduleBuilder.create( + constOp.getLoc(), ptrType, zero); + if (tensorType) { + // 3. splat 成 tensor<...x!tt.ptr<...>> + Value ptrTensor = moduleBuilder.create( + constOp.getLoc(), tensorType, ptrValue); + initTensors.push_back(ptrTensor); + } else { + initTensors.push_back(ptrValue); + } + } else if (auto memrefType = dyn_cast(valueType)) { + // 如果中间变量是一个memref类型，为iterarg创建一个 alloc = memref + // 仅支持#hivm.address_space，对于#hivm.address_space，不存在 + // copy cbuf to cbuf 行为 + auto spaceAttr = + cast(memrefType.getMemorySpace()); + if (spaceAttr && + spaceAttr.getAddressSpace() == hivm::AddressSpace::L1) { + llvm::dbgs() << "AddArgsForDependValues: dependValue type is a " + "memref hivm::AddressSpace::L1 type!!!\n"; + return mlir::WalkResult::interrupt(); + } else { + mlir::Value alloc = moduleBuilder.create( + constOp.getLoc(), memrefType); + initTensors.push_back(alloc); + } + } else { + // 非 ptr 类型创建零值常量 + auto zeroAttr = moduleBuilder.getZeroAttr(valueType); + Value zeroTensor = moduleBuilder.create( + constOp.getLoc(), zeroAttr); + initTensors.push_back(zeroTensor); + } + } + return mlir::WalkResult::interrupt(); + } + return mlir::WalkResult::advance(); + }); + + auto initArgs = forOp.getInitArgs(); + + // 构建新的初始化参数列表 + SmallVector newInitArgs(initArgs.begin(), initArgs.end()); + // 添加 dependValue 的初始化参数 + for (Value initTensor : initTensors) { + newInitArgs.push_back(initTensor); + } + + // 获取原循环的边界和步长 + Value lb = forOp.getLowerBound(); + Value ub = forOp.getUpperBound(); + Value step = forOp.getStep(); + + // 创建新的 ForOp，插入点位于原操作之前 + OpBuilder builder(forOp); + auto newForOp = + builder.create(forOp.getLoc(), lb, ub, step, newInitArgs); + + // 获取新循环的 region 块（已自动包含循环索引和迭代参数） + Block &newBlock = newForOp.getRegion().front(); + Block &oldBlock = forOp.getRegion().front(); + + // 建立块参数的映射：原块参数 -> 新块参数 + IRMapping mapper; + for (unsigned i = 0; i < oldBlock.getNumArguments(); ++i) { + mapper.map(oldBlock.getArgument(i), newBlock.getArgument(i)); + } + // 将原循环体中的操作（不包括终结符）克隆到新块中 + // 同时按照顺序克隆新的 dependValues + SmallVector newDependValues = dependValues; + int cnt = 0; + builder.setInsertionPointToStart(&newBlock); + for (auto &op : oldBlock) { + auto newOp = builder.clone(op, mapper); + // dependValue 的定义OP 可能有多个 result + for (size_t i = 0; i < dependValues.size(); i++) { + Operation *defineOp = dependValues[i].getDefiningOp(); + if (defineOp == &op) { + unsigned int index = cast(dependValues[i]).getResultNumber(); + newDependValues[i] = newOp->getResult(index); + cnt++; + break; + } + } + } + // 判断是否找到了所有的 dependValue + if (newDependValues.size() != cnt) { + llvm::outs() << "can not find the depend value! \n"; + return; + } + dependValues = newDependValues; + + // 更新 mergedRegions 中的 op 为新的for循环的 op + UpdateMergedRegionsWithNewForOp(mergedRegions, mapper); + + // 创建新的循环 yield 操作：原操作数 + dependValues + auto oldYield = cast(newBlock.getTerminator()); + SmallVector newYieldOps(oldYield.getOperands()); + // 按顺序增加找到的 dependvalue + for (Value v : newDependValues) { + newYieldOps.push_back(v); + } + builder.setInsertionPointToEnd(&newBlock); + builder.create(oldYield.getLoc(), newYieldOps); + oldYield.erase(); + + // 将原 forOp 的所有使用替换为新 forOp + int oldResultNum = forOp->getResults().size(); + for (auto it : llvm::zip(forOp->getResults(), + newForOp->getResults().take_front(oldResultNum))) { + std::get<0>(it).replaceAllUsesWith(std::get<1>(it)); + } + forOp.erase(); +} + +void ComputeElseYieldValues(MergedRegion mergedRegion, + SmallVector &elseYieldValues, + SmallVector dependValues) { + int idx = 0; + for (Value v : mergedRegion.yieldValues) { + Type yieldType = mergedRegion.resultTypes[idx]; + elseYieldValues.push_back(findIterArg(v, yieldType)); + idx++; + } +} + +void ComputeElseYieldValuesV2(MergedRegion mergedRegion, + SmallVector &elseYieldValues, + SmallVector dependValues) { + // 对于yieldValues，其中的 yield value 一定是被 for op yield + // 所引用，或者被其他 region 所使用 + auto forOp = dyn_cast( + mergedRegion.yieldValues[0].getDefiningOp()->getBlock()->getParentOp()); + if (!forOp) { + llvm::outs() << "define op's parent is not ForOp \n"; + return; + } + auto iterArgs = forOp.getRegionIterArgs(); + auto forYieldValues = forOp.getYieldedValues(); + + // 新增的与 dependvalue 相关的 initarg + // 是接在原本for循环args后面，数量与dependvalue数量相等 + int baseDependIdx = iterArgs.size() - dependValues.size(); + + int idx = 0; + for (Value v : mergedRegion.yieldValues) { + Type yieldType = mergedRegion.resultTypes[idx]; + // yieldValue 中是dependvalue 的情况下 + // else yield value 使用对应的新增 iterargs + if (llvm::is_contained(dependValues, v)) { + int dependIdx = 0; + for (; dependIdx < dependValues.size(); dependIdx++) { + if (v == dependValues[dependIdx]) { + break; + } + } + // llvm::outs()<<"v2for:"< newYieldValues; + SmallVector newResultTypes; + + SmallPtrSet seen; + + for (auto [idx, v] : llvm::enumerate(region.yieldValues)) { + if (seen.insert(v).second) { + newYieldValues.push_back(v); + newResultTypes.push_back(region.resultTypes[idx]); + } + } + + region.yieldValues.swap(newYieldValues); + region.resultTypes.swap(newResultTypes); +} + +static void replaceExternalIfOpUses(scf::IfOp ifOp, + ArrayRef oldYieldValues) { + + for (size_t i = 0; i < oldYieldValues.size(); ++i) { + Value oldVal = oldYieldValues[i]; + Value newVal = ifOp.getResult(i); + + SmallVector usesToReplace; + + for (OpOperand &use : llvm::make_early_inc_range(oldVal.getUses())) { + + Operation *user = use.getOwner(); + + // 跳过 ifOp 内部的使用（then / else region） + if (ifOp->isAncestor(user)) + continue; + + // 只替换 ifOp 之后的使用 + if (user->getBlock() == ifOp->getBlock()) { + if (!ifOp->isBeforeInBlock(user)) + continue; + } + + usesToReplace.push_back(&use); + } + + for (OpOperand *use : usesToReplace) + use->set(newVal); + } +} + +void CreateIfOps(SmallVector &mergedRegions, + SmallVector dependValues) { + for (auto ®ion : mergedRegions) { + + // 去重yieldvalues + RemoveRedundantYieldValues(region); + + Operation *insertPt = region.opsToMove.front(); + OpBuilder builder(insertPt); + Location loc = insertPt->getLoc(); + Value cond = builder.create(loc, builder.getI1Type(), + builder.getBoolAttr(true)); + + bool needsYield = !region.yieldValues.empty(); + scf::IfOp ifOp; + if (needsYield) + ifOp = builder.create(loc, region.resultTypes, cond, true); + else + ifOp = builder.create(loc, TypeRange{}, cond, false); + + // 加标记 + ifOp->setAttr("ssbuffer", builder.getUnitAttr()); + + // 获取if yield value 在 else块 返回值 + SmallVector elseYieldValues; + + llvm::outs() << "before ComputeElseYieldValuesV2" + << "\n"; + if (needsYield) { + // ComputeElseYieldValues(region, elseYieldValues, dependValues); + ComputeElseYieldValuesV2(region, elseYieldValues, dependValues); + } + + llvm::outs() << "after ComputeElseYieldValuesV2" + << "\n"; + // 将op移进then块 + Block &thenBlock = ifOp.getThenRegion().front(); + for (Operation *m : llvm::reverse(region.opsToMove)) { + m->moveBefore(&thenBlock, thenBlock.begin()); + } + + // 创建 then/else yield + if (needsYield) { + OpBuilder thenBuilder(builder.getContext()); + thenBuilder.setInsertionPointToEnd(&thenBlock); + thenBuilder.create(loc, region.yieldValues); + + // else block + Block &elseBlock = ifOp.getElseRegion().front(); + OpBuilder elseBuilder(&elseBlock, elseBlock.end()); + elseBuilder.create(loc, elseYieldValues); + + // 替换外部使用 + + replaceExternalIfOpUses(ifOp, region.yieldValues); + + // 旧的逻辑 + // Block *block = ifOp->getBlock(); + // auto ifIt = Block::iterator(ifOp); + + // for (size_t i = 0; i < region.yieldValues.size(); ++i) { + // Value oldVal = region.yieldValues[i]; + // Value newVal = ifOp.getResult(i); + + // SmallVector usesToReplace; + + // for (OpOperand &use : llvm::make_early_inc_range(oldVal.getUses())) { + // Operation *user = use.getOwner(); + // // 同一个 block, user 必须在 ifOp 之后, 不能在 ifOp 内部（then / + // else） if (user->getBlock() != ifOp->getBlock() || + // !ifOp->isBeforeInBlock(user) || user->getParentOp() == ifOp) + // continue; + // usesToReplace.push_back(&use); + // } + + // for (OpOperand *use : usesToReplace) + // use->set(newVal); + // } + } + + llvm::outs() << "Create ifOp: " << *ifOp << "\n"; + } +} + +void CreateIfOpsOrigin(SmallVector &mergedRegions) { + for (auto ®ion : mergedRegions) { + + // 去重yieldvalues + RemoveRedundantYieldValues(region); + + Operation *insertPt = region.opsToMove.front(); + OpBuilder builder(insertPt); + Location loc = insertPt->getLoc(); + Value cond = builder.create(loc, builder.getI1Type(), + builder.getBoolAttr(true)); + + bool needsYield = !region.yieldValues.empty(); + scf::IfOp ifOp; + if (needsYield) + ifOp = builder.create(loc, region.resultTypes, cond, true); + else + ifOp = builder.create(loc, TypeRange{}, cond, false); + + // 加标记 + ifOp->setAttr("ssbuffer", builder.getUnitAttr()); + + // 将op移进then块 + Block &thenBlock = ifOp.getThenRegion().front(); + for (Operation *m : llvm::reverse(region.opsToMove)) { + m->moveBefore(&thenBlock, thenBlock.begin()); + } + + // 创建 then/else yield + if (needsYield) { + OpBuilder thenBuilder(builder.getContext()); + thenBuilder.setInsertionPointToEnd(&thenBlock); + thenBuilder.create(loc, region.yieldValues); + + // else block + SmallVector elseYieldValues; + int idx = 0; + for (Value v : region.yieldValues) { + Type yieldType = region.resultTypes[idx]; + elseYieldValues.push_back(findIterArgForAll(v, yieldType)); + idx++; + } + Block &elseBlock = ifOp.getElseRegion().front(); + OpBuilder elseBuilder(&elseBlock, elseBlock.end()); + elseBuilder.create(loc, elseYieldValues); + + // 替换外部使用 + Block *block = ifOp->getBlock(); + auto ifIt = Block::iterator(ifOp); + + for (size_t i = 0; i < region.yieldValues.size(); ++i) { + Value oldVal = region.yieldValues[i]; + Value newVal = ifOp.getResult(i); + + SmallVector usesToReplace; + + for (OpOperand &use : llvm::make_early_inc_range(oldVal.getUses())) { + Operation *user = use.getOwner(); + // 同一个 block, user 必须在 ifOp 之后, 不能在 ifOp 内部（then / + // else） + if (user->getBlock() != ifOp->getBlock() || + !ifOp->isBeforeInBlock(user) || user->getParentOp() == ifOp) + continue; + usesToReplace.push_back(&use); + } + + for (OpOperand *use : usesToReplace) + use->set(newVal); + } + } + + llvm::outs() << "Create ifOp: " << *ifOp << "\n"; + } +} + +void AddIfCondition(ModuleOp module) { + SmallVector copiedForOps; + SmallVector forOpList; + SmallVector, 1> regionList; + + module.walk([&](scf::ForOp forOp) { + Block &body = forOp.getRegion().front(); + SmallVector regions; + + // 获取基本的wait-set分块信息 + GetBlockInfos(regions, body); + + SmallVector mergedRegions; + // 合并wait-set块, 依据copyop / fixpipeop合并 + MergeWaitSetRegions(regions, mergedRegions); + + // 扩展if包裹的op范围 + // AIV、AIC处理有区别 + ExpandMergedRegionOps(forOp, mergedRegions, copiedForOps); + + // 处理forop的末尾对于iter_arg的自增操作, 如tt.advance, 移进对应的if op + MoveIterArgUsersIntoIf(forOp, mergedRegions); + + // 获取if yield的value, 并更新if内op的user为yield value + for (MergedRegion &mr : mergedRegions) { + // ComputeYieldForMergedRegion(mr, body); + ComputeYieldForMergedRegionV4(mr); + } + + // // 创建最终的if op + // CreateIfOpsOrigin(mergedRegions); + // }); + + forOpList.push_back(forOp); + regionList.push_back(mergedRegions); + }); + + llvm::outs() << "CopyForOp:\n"; + for (auto op : copiedForOps) { + llvm::outs() << *op << "\n"; + } + + SmallVector tmpOps; + for (auto copiedOp : copiedForOps) { + Block &body = copiedOp.getRegion().front(); + SmallVector regions; + + // 获取基本的wait-set分块信息 + GetBlockInfos(regions, body); + + SmallVector mergedRegions; + // 合并wait-set块, 依据copyop / fixpipeop合并 + MergeWaitSetRegions(regions, mergedRegions); + + // 扩展if包裹的op范围 + // AIV、AIC处理有区别 + ExpandMergedRegionOps(copiedOp, mergedRegions, tmpOps); + + // 处理forop的末尾对于iter_arg的自增操作, 如tt.advance, 移进对应的if op + MoveIterArgUsersIntoIf(copiedOp, mergedRegions); + + // 获取if yield的value, 并更新if内op的user为yield value + for (MergedRegion &mr : mergedRegions) { + // ComputeYieldForMergedRegion(mr, body); + ComputeYieldForMergedRegionV4(mr); + } + + // // 创建最终的if op + // CreateIfOpsOrigin(mergedRegions); + // }); + + forOpList.push_back(copiedOp); + regionList.push_back(mergedRegions); + } + + for (size_t i = 0; i < forOpList.size(); ++i) { + scf::ForOp oldForOp = forOpList[i]; + SmallVector newMergedRegions = regionList[i]; + + // 找到所有的VV或CC依赖 + SmallVector dependValues; + llvm::outs() << "FindDependValues! \n "; + FindDependValues(dependValues, newMergedRegions); + + if (dependValues.size() != 0) { + copyLoadCalculation(oldForOp, dependValues, newMergedRegions); + + // repeat previous operations + for (MergedRegion &mr : newMergedRegions) { + mr.yieldValues.clear(); + mr.resultTypes.clear(); + ComputeYieldForMergedRegionV4(mr); + } + FindDependValues(dependValues, newMergedRegions); + } + + // 如果存在VV或CC依赖，更新ForOp添加新的对应args + if (dependValues.size() != 0) { + AddArgsForDependValues(oldForOp, dependValues, newMergedRegions, module); + } + + // 创建最终的if op + llvm::outs() << "before create if ops" << '\n'; + CreateIfOps(newMergedRegions, dependValues); + } +} + +void ChangeAdvanceOpForm(ModuleOp module) { + module.walk([&](scf::ForOp forOp) { + Block &body = forOp.getRegion().front(); + constexpr int num = 8; + SmallVector ifOps; + for (Operation &op : body) + if (auto ifOp = dyn_cast(&op)) + ifOps.push_back(ifOp); + + for (scf::IfOp ifOp : ifOps) { + // 找 then region 中的 advance + triton::AdvanceOp advanceOp; + for (Operation &thenOp : ifOp.getThenRegion().front()) { + if (auto adv = dyn_cast(thenOp)) { + advanceOp = adv; + break; + } + } + if (!advanceOp) + continue; + + // base 必须是 for的iter_arg + Value base = advanceOp.getPtr(); + auto barg = dyn_cast(base); + if (!barg || barg.getOwner() != &body) + continue; + + // yield 去掉 advance 的返回值 + auto thenYield = + cast(ifOp.getThenRegion().front().getTerminator()); + auto elseYield = + cast(ifOp.getElseRegion().front().getTerminator()); + + int advanceIdx = -1; + for (auto it : llvm::enumerate(thenYield.getOperands())) { + if (it.value() == advanceOp.getResult()) { + advanceIdx = it.index(); + break; + } + } + + if (advanceIdx == -1) + continue; + + // 删除 advance + SmallVector thenOps(thenYield.getOperands().begin(), + thenYield.getOperands().end()); + SmallVector elseOps(elseYield.getOperands().begin(), + elseYield.getOperands().end()); + + thenOps.erase(thenOps.begin() + advanceIdx); + elseOps.erase(elseOps.begin() + advanceIdx); + + thenYield->setOperands(thenOps); + elseYield->setOperands(elseOps); + + // 重建 ifOp（去掉 advance 对应的 result） + OpBuilder ifBuilder(ifOp); + ifBuilder.setInsertionPoint(ifOp); + + // 构造新的 result types + SmallVector newResultTypes; + for (int i = 0; i < ifOp.getNumResults(); ++i) { + if (i != advanceIdx) + newResultTypes.push_back(ifOp.getResult(i).getType()); + } + + // 创建新的 if + auto newIf = ifBuilder.create(ifOp.getLoc(), newResultTypes, + ifOp.getCondition(), + /*withElseRegion=*/true); + newIf->setAttr("ssbuffer", ifBuilder.getUnitAttr()); + // 把已经修改过 yield 的 region 搬过去 + newIf.getThenRegion().takeBody(ifOp.getThenRegion()); + newIf.getElseRegion().takeBody(ifOp.getElseRegion()); + + // 替换if result的user + int newIdx = 0; + for (int oldIdx = 0; oldIdx < ifOp.getNumResults(); ++oldIdx) { + if (oldIdx == advanceIdx) + continue; + ifOp.getResult(oldIdx).replaceAllUsesWith(newIf.getResult(newIdx++)); + } + + OpBuilder builder(newIf); + builder.setInsertionPointAfter(newIf); + + Value flag = newIf.getCondition(); + + SmallVector newOffsets; + for (Value off : advanceOp.getOffsets()) { + auto intTy = cast(off.getType()); + auto zero = builder.create(newIf.getLoc(), 0, + intTy.getWidth()); + auto sel = + builder.create(newIf.getLoc(), flag, off, zero); + newOffsets.push_back(sel); + } + + auto newAdvance = builder.create( + newIf.getLoc(), base.getType(), base, newOffsets); + + // 原 if 的 advance result 的 users，接到 newAdvance + ifOp.getResult(advanceIdx).replaceAllUsesWith(newAdvance.getResult()); + + // 删除旧的ifOp和advance + advanceOp.erase(); + ifOp.erase(); + } + }); +} + +void processRedudantIf(ModuleOp module) { + SmallVector forOps; + llvm::outs() << module << " wwwww\n\n\n"; + module.walk([&](scf::ForOp forOp) { + auto initArgs = forOp.getInitArgs(); + if (initArgs.size() == 5) { + forOps.push_back(forOp); + } + }); + + for (auto forOp : forOps) { + auto initArgs = forOp.getInitArgs(); + Value newInit = initArgs[2]; + + // 构建新的初始化参数列表 + SmallVector newInitArgs(initArgs.begin(), initArgs.end()); + newInitArgs.push_back(newInit); + + // 获取原循环的边界和步长 + Value lb = forOp.getLowerBound(); + Value ub = forOp.getUpperBound(); + Value step = forOp.getStep(); + + // 创建新的 ForOp，插入点位于原操作之前 + OpBuilder builder(forOp); + auto newForOp = + builder.create(forOp.getLoc(), lb, ub, step, newInitArgs); + + // 获取新循环的 region 块（已自动包含循环索引和迭代参数） + Block &newBlock = newForOp.getRegion().front(); + Block &oldBlock = forOp.getRegion().front(); + + // 建立块参数的映射：原块参数 -> 新块参数（前6个对应） + IRMapping mapper; + for (unsigned i = 0; i < oldBlock.getNumArguments(); ++i) { + mapper.map(oldBlock.getArgument(i), newBlock.getArgument(i)); + } + // 将原循环体中的操作（不包括终结符）克隆到新块中 + builder.setInsertionPointToStart(&newBlock); + for (auto &op : oldBlock) { + auto newOp = builder.clone(op, mapper); + } + + // 在新块中查找第一个 scf::IfOp（即原代码中的第一个 if） + scf::IfOp firstIfOp = nullptr; + for (auto &op : newBlock.getOperations()) { + if (auto ifOp = dyn_cast(&op)) { + firstIfOp = ifOp; + break; + } + } + assert(firstIfOp && "Expected at least one if op in the loop body"); + + // 修改第一个 if 的 else 分支的 yield 操作： + // 将其第二个操作数（索引1）从原来的 %arg9 改为新迭代参数（新块参数索引6） + Block &elseBlock = firstIfOp.getElseRegion().front(); + auto elseYield = cast(elseBlock.getTerminator()); + SmallVector newElseYieldOps(elseYield.getOperands()); + newElseYieldOps[1] = newBlock.getArgument(6); // 新迭代参数 + builder.setInsertionPoint(elseYield); + builder.create(elseYield.getLoc(), newElseYieldOps); + elseYield->erase(); + + // 创建新的循环 yield 操作：原5个操作数 + 第一个 if 的第二个结果 + auto oldYield = cast(newBlock.getTerminator()); + SmallVector newYieldOps(oldYield.getOperands()); + newYieldOps.push_back(firstIfOp.getResult(1)); // 第一个 if 的第二个结果 + builder.setInsertionPointToEnd(&newBlock); + builder.create(oldYield.getLoc(), newYieldOps); + oldYield.erase(); + + // 将原 forOp 的所有使用替换为新 forOp 的前5个结果 + for (auto it : + llvm::zip(forOp->getResults(), newForOp->getResults().take_front(5))) { + std::get<0>(it).replaceAllUsesWith(std::get<1>(it)); + } + } + for (auto forOp : forOps) { + forOp.erase(); + } +} +// 针对依赖变量，对原本的for op增加double buffer相关的迭代参数 +scf::ForOp addDoubleBuffForArgs(ModuleOp module, SmallVector uniqueDeps, + int bufferNum) { + mlir::OpBuilder builder(module.getContext()); + SmallVector depValueForIdxs; + + // ========== 找到scf.if所在的scf::ForOp ========== + if (!isa(uniqueDeps[0].getDefiningOp()->getParentOp())) { + llvm::errs() << "Error: parent op of scf.if is not scf.for"; + } + scf::ForOp forOp = + dyn_cast(uniqueDeps[0].getDefiningOp()->getParentOp()); + + for (Value dependencyValue : uniqueDeps) { + // ========== 步骤1：验证目标Value是scf.if的返回值，并找到对应的scf::IfOp + // ========== + Operation *ifOp = dependencyValue.getDefiningOp(); + if (!ifOp || !isa(ifOp)) { + llvm::errs() << "Error: 目标Value不是scf.if的返回值\n"; + return nullptr; + } + scf::IfOp targetIfOp = dyn_cast(ifOp); + + // 确认当前Value是scf.if的第几个返回值 + int64_t depValueIdx = -1; + for (auto [idx, result] : llvm::enumerate(targetIfOp.getResults())) { + if (result == dependencyValue) { + depValueIdx = idx; + break; + } + } + + // ========== 步骤2：找到%38#2关联的scf.for迭代参数以及索引 ========== + // %38#2对应scf.if else分支yield的第2个操作数 → 即%arg10 + Operation *elseYield = targetIfOp.elseYield(); + Value dependencyArg = elseYield->getOperand( + depValueIdx); // depValueIdx=2，对应else yield的第2个参数 + + int64_t depValueForIdx = -1; + for (auto [idx, result] : llvm::enumerate(forOp.getRegionIterArgs())) { + if (result == dependencyArg) { + depValueForIdx = idx; + break; + } + } + depValueForIdxs.push_back(depValueForIdx); + llvm::outs() << "depValueForIdx: " << depValueForIdx << '\n'; + } + + llvm::outs() << "oldFor: " << forOp << '\n'; + + // 获取原始循环的信息 + Value originalLowerBound = forOp.getLowerBound(); + Value originalUpperBound = forOp.getUpperBound(); + Value originalStep = forOp.getStep(); + SmallVector originalInitArgs = forOp.getInitArgs(); + SmallVector iterArgs; + for (auto arg : originalInitArgs) { + iterArgs.push_back(arg); + } + auto yields = forOp.getBody()->getTerminator(); + + // 创建计数器初始零值 + Value counterInit = nullptr; + mlir::Operation *parentOp = forOp->getParentOp(); + mlir::Operation *scopeOp = nullptr; + // 向上遍历查找scope.scope操作 + while (parentOp) { + if (dyn_cast(parentOp)) { + scopeOp = parentOp; + break; + } + parentOp = parentOp->getParentOp(); + } + + builder.setInsertionPoint(scopeOp); + Location loc = forOp.getLoc(); + auto boundType = originalLowerBound.getType(); + counterInit = builder.create(loc, boundType, 0); + + // 添加和depValueForIdxs相同的迭代参数和计数器 + for (int64_t idx : depValueForIdxs) { + for (int i = 0; i < bufferNum - 1; i++) { + iterArgs.push_back(originalInitArgs[idx]); + } + + // 在迭代参数中添加计数器 + for (int i = 0; i < 2; i++) { + iterArgs.push_back(counterInit); + } + } + + builder.setInsertionPoint(forOp); + // 创建新的for循环 + auto newForOp = + builder.create(forOp.getLoc(), originalLowerBound, + originalUpperBound, originalStep, iterArgs); + + // 设置IR映射表，将旧循环的变量映射到新循环 + IRMapping mapper; + + // 映射迭代变量 + mapper.map(forOp.getInductionVar(), newForOp.getInductionVar()); + + // 映射迭代参数 + for (auto [oldArg, newArg] : + llvm::zip(forOp.getRegionIterArgs(), newForOp.getRegionIterArgs())) { + mapper.map(oldArg, newArg); + } + + SmallVector newArgs; + for (int i = forOp.getRegionIterArgs().size(); + i < newForOp.getRegionIterArgs().size(); i++) { + newArgs.push_back(newForOp.getRegionIterArgs()[i]); + } + // 克隆循环体内容到新循环 + auto &newLoopBody = *newForOp.getBody(); + builder.setInsertionPointToStart(&newLoopBody); + + for (auto &op : forOp.getBody()->without_terminator()) { + builder.clone(op, mapper); + } + + // 克隆yield操作 + if (auto yieldOp = dyn_cast(yields)) { + SmallVector newYieldOperands; + for (auto operand : yieldOp.getOperands()) { + newYieldOperands.push_back(mapper.lookupOrDefault(operand)); + } + // 将新增的迭代参数添加到yield操作数中 + for (auto currentCounter : newArgs) { + newYieldOperands.push_back(currentCounter); + } + builder.create(yieldOp.getLoc(), newYieldOperands); + } + + // 替换原循环的结果 + unsigned numOriginalResults = forOp.getNumResults(); + SmallVector originalResults; + for (unsigned i = 0; i < numOriginalResults; i++) { + originalResults.push_back(newForOp.getResult(i)); + } + forOp.replaceAllUsesWith(originalResults); + + // 8. 删除原循环 + forOp.erase(); + + llvm::outs() << "for op erased!\n"; + return newForOp; +} + +static bool isScalarValue(Value value) { + if (!value) + return false; + Type type = value.getType(); + + return type.isIntOrIndex() || mlir::isa(type) || + mlir::isa(type) || + mlir::isa(type); +} + +SmallVector buildNBufferProducer(OpBuilder &builder, Location loc, + Value frontCnt, Value newDepVal, + ArrayRef buffs, + ArrayRef constants) { + // N-buffer producer: determines which buffer is written to newDepVal based on + // frontCnt % N + const int N = buffs.size(); + SmallVector results; + + // idx = frontCnt % N + Value bufferIndex = + builder.create(loc, frontCnt, constants[N]); + + // 1. buffer0: handle the first buffer separately + Value isBuffer0 = builder.create(loc, arith::CmpIPredicate::eq, + bufferIndex, constants[0]); + + auto dstShapedType = mlir::dyn_cast(newDepVal.getType()); + Value mask; + if (isScalarValue(isBuffer0) == 0) { + auto maskType = + RankedTensorType::get(dstShapedType.getShape(), isBuffer0.getType()); + mask = builder.create(loc, maskType, isBuffer0); + llvm::outs() << "build NBufferProducer debug15\n"; + llvm::outs().flush(); + } else { + mask = isBuffer0; + } + Value newBuff0 = + builder.create(loc, mask, newDepVal, buffs[0]); + + results.push_back(newBuff0); + + // 2. Double-buffer specialization (when N == 2, a direct select is + // sufficient) + if (N == 2) { + + Value newBuff1 = + builder.create(loc, mask, buffs[1], newDepVal); + + auto nextCnt = builder.create(loc, frontCnt, constants[1]); + + results.push_back(newBuff1); + results.push_back(nextCnt.getResult()); + + return results; + } + + // 3. Build the root IF: when idx == 0, + // use the first buffer; otherwise enter the nestedIf chain to use other + // buffers + SmallVector resultTypes; + for (int i = 1; i < N; ++i) + resultTypes.push_back(buffs[i].getType()); + + auto rootIf = builder.create(loc, resultTypes, isBuffer0, true); + + // ---- THEN: buffers are directly forwarded ---- + { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(&rootIf.getThenRegion().front()); + + SmallVector unchangedBuffers(buffs.begin() + 1, buffs.end()); + + builder.create(loc, unchangedBuffers); + } + + // 4. Construct the nested-if chain, updating one buffer at each level + Block *currentElseBlock = &rootIf.getElseRegion().front(); + + scf::IfOp parentIf = rootIf; + + for (int i = 1; i < N - 1; ++i) { + + builder.setInsertionPointToStart(currentElseBlock); + + // Check whether the current buffer is selected + Value isCurrent = builder.create( + loc, arith::CmpIPredicate::eq, bufferIndex, constants[i]); + + // Update buffer[i] + dstShapedType = mlir::dyn_cast(newDepVal.getType()); + auto maskType = + RankedTensorType::get(dstShapedType.getShape(), isCurrent.getType()); + mask = builder.create(loc, maskType, isCurrent); + Value updatedBuffer = + builder.create(loc, mask, newDepVal, buffs[i]); + + // If this is the last level: directly yield both buffers + if (i == N - 2) { + + dstShapedType = mlir::dyn_cast(newDepVal.getType()); + maskType = + RankedTensorType::get(dstShapedType.getShape(), isCurrent.getType()); + mask = builder.create(loc, maskType, isCurrent); + Value lastBuffer = + builder.create(loc, mask, buffs[N - 1], newDepVal); + + builder.create(loc, ValueRange{updatedBuffer, lastBuffer}); + + break; + } + + // Create the next nested if + SmallVector subResultTypes; + for (int j = i + 1; j < N; ++j) + subResultTypes.push_back(buffs[j].getType()); + + auto nextIf = + builder.create(loc, subResultTypes, isCurrent, true); + + // THEN: forward the remaining buffers + { + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointToStart(&nextIf.getThenRegion().front()); + + SmallVector remainingBuffers(buffs.begin() + i + 1, buffs.end()); + + builder.create(loc, remainingBuffers); + } + + // Update the else yield + builder.setInsertionPointToEnd(&parentIf.getElseRegion().front()); + + SmallVector yields; + yields.push_back(updatedBuffer); + yields.append(nextIf.getResults().begin(), nextIf.getResults().end()); + + builder.create(loc, yields); + + parentIf = nextIf; + currentElseBlock = &nextIf.getElseRegion().front(); + } + + // 5. Update the frontCnt counter + builder.setInsertionPointAfter(rootIf); + + auto nextCnt = builder.create(loc, frontCnt, constants[1]); + + // Collect results + results.append(rootIf.getResults().begin(), rootIf.getResults().end()); + + results.push_back(nextCnt.getResult()); + + return results; +} + +SmallVector buildNBufferConsumer(OpBuilder &builder, Location loc, + Value postCnt, ArrayRef oldBuffs, + ArrayRef constants) { + // Consumer: selects which buffer to read based on postCnt % N + const int bufferNum = oldBuffs.size(); + SmallVector results; + + // idx = postCnt % N + Value bufferIndex = + builder.create(loc, postCnt, constants[bufferNum]); + + Value isBuffer0 = builder.create(loc, arith::CmpIPredicate::eq, + bufferIndex, constants[0]); + auto dstShapedType = mlir::dyn_cast(oldBuffs[0].getType()); + Value mask; + if (isScalarValue(isBuffer0) == 0) { + auto maskType = + RankedTensorType::get(dstShapedType.getShape(), isBuffer0.getType()); + mask = builder.create(loc, maskType, isBuffer0); + } else { + mask = isBuffer0; + } + + // 1. Double-buffer specialization (avoid generating scf.if) + if (bufferNum == 2) { + Value selected = + builder.create(loc, mask, oldBuffs[0], oldBuffs[1]); + auto nextCnt = builder.create(loc, postCnt, constants[1]); + + results.push_back(selected); + results.push_back(nextCnt); + + return results; + } + + // 2. Build the root IF: + // when idx == 0, use the first buffer; otherwise enter the nestedIf chain to + // use other buffers + SmallVector resultTypes{oldBuffs[0].getType()}; + + auto rootIf = builder.create(loc, resultTypes, isBuffer0, true); + + // ---- THEN: directly return buffer0 ---- + { + builder.setInsertionPointToStart(&rootIf.getThenRegion().front()); + + builder.create(loc, oldBuffs[0]); + } + + // 3. Construct the nested-if chain + Block *currentElse = &rootIf.getElseRegion().front(); + + for (int i = 1; i < bufferNum - 2; ++i) { + + builder.setInsertionPointToStart(currentElse); + + Value isCurrent = builder.create( + loc, arith::CmpIPredicate::eq, bufferIndex, constants[i]); + + auto nestedIf = builder.create( + loc, TypeRange{oldBuffs[0].getType()}, isCurrent, true); + + // THEN → return the current buffer + { + builder.setInsertionPointToStart(&nestedIf.getThenRegion().front()); + + builder.create(loc, oldBuffs[i]); + } + + // ELSE → yield nested result + builder.setInsertionPointToEnd(currentElse); + builder.create(loc, nestedIf.getResult(0)); + + // Enter the next else branch + currentElse = &nestedIf.getElseRegion().front(); + } + + // 4. Final level (use select to finish) + builder.setInsertionPointToStart(currentElse); + + int last = bufferNum - 2; + + Value isLast = builder.create(loc, arith::CmpIPredicate::eq, + bufferIndex, constants[last]); + + auto maskType = RankedTensorType::get({}, isLast.getType()); + dstShapedType = mlir::dyn_cast(oldBuffs[last].getType()); + maskType = RankedTensorType::get(dstShapedType.getShape(), isLast.getType()); + mask = builder.create(loc, maskType, isLast); + + Value finalSelect = builder.create(loc, mask, oldBuffs[last], + oldBuffs[last + 1]); + + builder.create(loc, finalSelect); + + // rootIf result = selected buffer + results.push_back(rootIf.getResult(0)); + + // 5. Update the postCnt counter + builder.setInsertionPointAfter(rootIf); + + auto nextCnt = builder.create(loc, postCnt, constants[1]); + + results.push_back(nextCnt); + + return results; +} + +void replaceDepsMap(scf::IfOp oldIfOp, scf::IfOp newIfOp, + SmallVector &newDeps, bool isFront, + DenseMap> &newIfResultDeps) { + mlir::IRMapping valueMap; + + // old result -> new result + for (unsigned i = 0; i < oldIfOp.getNumResults(); ++i) { + valueMap.map(oldIfOp.getResult(i), newIfOp.getResult(i)); + } + + if (isFront) { + for (int i = 0; i < newDeps.size(); i++) { + Value v = newDeps[i]; + if (valueMap.contains(v)) + newDeps[i] = valueMap.lookup(v); + } + } + + // rewrite deps in-place + for (auto &it : newIfResultDeps) { + auto &deps = it.second; + + for (auto &value : deps) { + if (auto mapped = valueMap.lookupOrNull(value)) + value = mapped; + } + } +} + +scf::IfOp addResultsForFrontIfOp( + scf::IfOp frontIfOp, OpBuilder builder, int bufferNum, Value depValue, + SmallVector constants, SmallVector buffs, Value frontCnt, + Value postCnt, SmallVector &extraResultIndices, + SmallVector &newDeps, + DenseMap> &newIfResultDeps) { + OpBuilder::InsertionGuard guard(builder); + + Location loc = frontIfOp.getLoc(); + Value cond = frontIfOp.getCondition(); + + auto &oldThenBlock = frontIfOp.getThenRegion().front(); + auto &oldElseBlock = frontIfOp.getElseRegion().front(); + + // New result types = old results + extra buffers + counter + SmallVector newResultTypes(frontIfOp.getResultTypes().begin(), + frontIfOp.getResultTypes().end()); + + for (int i = 1; i < bufferNum; ++i) + newResultTypes.push_back(buffs[i].getType()); + + newResultTypes.push_back(frontCnt.getType()); + + unsigned oldNumResults = frontIfOp.getNumResults(); + + // Create new IfOp + builder.setInsertionPoint(frontIfOp); + auto newIfOp = + builder.create(loc, newResultTypes, cond, /*hasElse=*/true); + + SmallVector bufferIndices(bufferNum); + SmallVector newBuffs; + int frontCntIndex = -1; + + // THEN region + { + mlir::IRMapping mapping; + Block &newThenBlock = newIfOp.getThenRegion().front(); + + builder.setInsertionPointToStart(&newThenBlock); + + // Clone original then body + for (auto &op : oldThenBlock.without_terminator()) + builder.clone(op, mapping); + + // Update dependency value position inf ifOp results + auto result = dyn_cast(depValue); + if (!result) { + llvm::outs() << "depValue is not a result Value!\n"; + return nullptr; + } + + int depIdx = result.getResultNumber(); + Value depYieldValue = frontIfOp.thenYield()->getOperand(depIdx); + + Value newDepVal = mapping.contains(depYieldValue) + ? mapping.lookup(depYieldValue) + : depYieldValue; + + builder.setInsertionPointAfter(newDepVal.getDefiningOp()); + + // Create N buffer + SmallVector produced = buildNBufferProducer( + builder, loc, frontCnt, newDepVal, buffs, constants); + + // Last value in newBuffs is the counter + newBuffs.append(produced.begin(), produced.end() - 1); + + // Rebuild new yield + SmallVector thenOperands; + + for (Value v : oldThenBlock.getTerminator()->getOperands()) { + Value mapped = mapping.lookupOrDefault(v); + + // Replace first buffer + if (mapped == newDepVal) { + thenOperands.push_back(newBuffs[0]); + bufferIndices[0] = thenOperands.size() - 1; + } else { + thenOperands.push_back(mapped); + } + } + + // Replace other buffer + for (int i = 1; i < bufferNum; ++i) { + thenOperands.push_back(newBuffs[i]); + bufferIndices[i] = thenOperands.size() - 1; + } + + // Add counter + thenOperands.push_back(produced.back()); + frontCntIndex = thenOperands.size() - 1; + + builder.setInsertionPointToEnd(&newThenBlock); + builder.create(loc, thenOperands); + + // record new result indices + for (int idx : bufferIndices) + extraResultIndices.push_back(idx); + + extraResultIndices.push_back(frontCntIndex); + } + + // ELSE region + { + mlir::IRMapping mapping; + Block &newElseBlock = newIfOp.getElseRegion().front(); + + builder.setInsertionPointToStart(&newElseBlock); + + // Clone original else body + for (auto &op : oldElseBlock.without_terminator()) + builder.clone(op, mapping); + + builder.setInsertionPointToEnd(&newElseBlock); + + SmallVector elseOperands; + + for (Value v : oldElseBlock.getTerminator()->getOperands()) + elseOperands.push_back(mapping.lookupOrDefault(v)); + + // Add buffer + for (int i = 1; i < bufferNum; ++i) + elseOperands.push_back(buffs[i]); + + // Add counter + elseOperands.push_back(frontCnt); + + builder.create(loc, elseOperands); + } + + // Update dependency value + replaceDepsMap(frontIfOp, newIfOp, newDeps, true, newIfResultDeps); + + // Replace old ifOp + frontIfOp.replaceAllUsesWith(newIfOp.getResults().take_front(oldNumResults)); + + frontIfOp.erase(); + + return newIfOp; +} + +scf::IfOp addResultsForPostIfOp( + scf::IfOp postIfOp, scf::IfOp newfrontIfOp, OpBuilder builder, + int bufferNum, Value newDepValue, SmallVector constants, + SmallVector buffs, Value frontCnt, Value postCnt, + SmallVector &extraResultIndices, SmallVector &newDeps, + DenseMap> &newIfResultDeps) { + // 1. Parse the extra result indices produced by frontIf (added buffers and + // counters) + SmallVector bufferIndices(extraResultIndices.begin(), + extraResultIndices.end() - 1); + int frontCntIndex = extraResultIndices[bufferNum]; + + Location ifLoc = postIfOp.getLoc(); + Value cond = postIfOp.getCondition(); + + auto &oldThenBlock = postIfOp.getThenRegion().front(); + auto &oldElseBlock = postIfOp.getElseRegion().front(); + + // 2. Create a new IfOp (add a new postCnt result) + SmallVector newResultTypes(postIfOp.getResultTypes().begin(), + postIfOp.getResultTypes().end()); + newResultTypes.push_back(postCnt.getType()); + + builder.setInsertionPoint(postIfOp); + auto newIfOp = builder.create(ifLoc, newResultTypes, cond, + /*hasElse=*/true); + + mlir::IRMapping mapping; + + // 3. THEN region: clone the original logic, insert the multibuffer consumer + // and update dependency buffers + auto &newThenBlock = newIfOp.getThenRegion().front(); + builder.setInsertionPointToStart(&newThenBlock); + + // clone then body + for (auto &op : oldThenBlock.without_terminator()) + builder.clone(op, mapping); + builder.setInsertionPointToStart(&newThenBlock); + + // Find dependency uses that need to be replaced (located inside the current + // IfOp) + SmallVector replaceUses; + for (auto &use : newDepValue.getUses()) { + if (newIfOp == dyn_cast(use.getOwner()->getParentOp())) { + replaceUses.push_back(&use); + } + } + + // Collect buffers produced by frontIf + SmallVector oldBuffers; + for (int i = 0; i < bufferIndices.size(); ++i) + oldBuffers.push_back(newfrontIfOp.getResult(bufferIndices[i])); + + // Multibuffer consumer caculation + SmallVector consumerResults = + buildNBufferConsumer(builder, ifLoc, postCnt, oldBuffers, constants); + + Value selectedBuffer = consumerResults[0]; + Value nextPostCnt = consumerResults[1]; + + // Replace dependent buffer + for (auto *usePtr : replaceUses) { + usePtr->set(selectedBuffer); + } + + // Create then yield + SmallVector thenOperands; + for (auto v : oldThenBlock.getTerminator()->getOperands()) + thenOperands.push_back(mapping.lookupOrDefault(v)); + + int postCntIndex = thenOperands.size(); + thenOperands.push_back(nextPostCnt); + + builder.setInsertionPointToEnd(&newThenBlock); + builder.create(ifLoc, thenOperands); + extraResultIndices.push_back(postCntIndex); + + // 4. ELSE region：forward counter directly + auto &newElseBlock = newIfOp.getElseRegion().front(); + + for (auto &op : oldElseBlock.without_terminator()) + builder.clone(op, mapping); + + builder.setInsertionPointToEnd(&newElseBlock); + + SmallVector elseOperands; + for (auto v : oldElseBlock.getTerminator()->getOperands()) + elseOperands.push_back(mapping.lookupOrDefault(v)); + + elseOperands.push_back(postCnt); + + builder.create(ifLoc, elseOperands); + + // 5. Replace old ifOp with new one + auto oldNumResults = postIfOp.getNumResults(); + + // Update depency value + replaceDepsMap(postIfOp, newIfOp, newDeps, false, newIfResultDeps); + + postIfOp.replaceAllUsesWith(newIfOp.getResults().take_front(oldNumResults)); + + postIfOp.erase(); + + return newIfOp; +} + +void addMultiBuffCaculate(ModuleOp module, SmallVector newUniqueDeps, + DenseMap> &ifResultDeps, + scf::ForOp &newForOp, int bufferNum) { + + // ============================================================ + // Overall Idea + // + // For each dependency Value: + // 1. Find the front IfOp that produces it + // 2. Add multi-buffer results to the front IfOp + // 3. Find the post IfOp that consumes the result and extend it accordingly + // 4. Update the for-loop yield so that buffer states are correctly propagated + // ============================================================ + + OpBuilder builder(module.getContext()); + int processedDepCount = 0; + + SmallVector postIfOps; + newForOp.walk([&](scf::IfOp postIfOp) { postIfOps.push_back(postIfOp); }); + for (auto postIfOp : postIfOps) { + if (!ifResultDeps.count(postIfOp)) { + continue; + } + auto newDeps = ifResultDeps[postIfOp]; + for (int depValueIdx = 0; depValueIdx < newDeps.size(); depValueIdx++) { + Value depValue = newDeps[depValueIdx]; + + // Step 1. Locate the front IfOp that produces depValue + Operation *defOp = depValue.getDefiningOp(); + if (!defOp || !isa(defOp)) { + llvm::outs() << "Error: depValue is not produced by scf.if\n"; + break; + } + + scf::IfOp frontIfOp = cast(defOp); + + // Position of depValue in the IfOp results + auto result = dyn_cast(depValue); + if (!result) { + llvm::outs() << "depValue is not an OpResult!\n"; + return; + } + + int64_t depResultIndex = result.getResultNumber(); + + // Position of depValue in the IfOp results + Value depYieldValue = frontIfOp.thenYield()->getOperand(depResultIndex); + + // Step 2. Find the multi-buffer position in the ForOp + int64_t extraArgBaseIdx = + newForOp.getRegionIterArgs().size() - + (2 + bufferNum - 1) * (newUniqueDeps.size() - processedDepCount++); + + // Collect all buffers + SmallVector buffers; + + // buffer0 来自 else yield + buffers.push_back(frontIfOp.elseYield()->getOperand(depResultIndex)); + + // Other buffers come from for iter args + for (int i = 1; i < bufferNum; ++i) { + buffers.push_back( + newForOp.getRegionIterArgs()[extraArgBaseIdx + i - 1]); + } + + // Two counters + Value frontCnt = + newForOp.getRegionIterArgs()[extraArgBaseIdx + bufferNum - 1]; + Value postCnt = newForOp.getRegionIterArgs()[extraArgBaseIdx + bufferNum]; + + // Step 3. Create constants (0 ~ bufferNum) for rem / cmp buffer selection + // logic + SmallVector constants; + builder.setInsertionPoint(frontIfOp); + + auto dataType = frontCnt.getType(); + for (int i = 0; i <= bufferNum; ++i) { + constants.push_back(builder.create( + frontIfOp.getLoc(), dataType, builder.getIntegerAttr(dataType, i))); + } + + // Record the positions of newly added results in the IfOp + SmallVector extraResultIndices(bufferNum + 1); + extraResultIndices.clear(); + + // Step 4. Extend the front IfOp + scf::IfOp newFrontIfOp = addResultsForFrontIfOp( + frontIfOp, builder, bufferNum, depValue, constants, buffers, frontCnt, + postCnt, extraResultIndices, newDeps, ifResultDeps); + + // buffer result indices + SmallVector bufferResultIndices(extraResultIndices.begin(), + extraResultIndices.end() - 1); + + int frontCntResultIndex = extraResultIndices[bufferNum]; + + Value newDepValue = newFrontIfOp.getResult(depResultIndex); + + // Step 5. Find the post IfOp that consumes the dependency value + scf::IfOp postIfOp = nullptr; + + for (auto &use : newDepValue.getUses()) { + if (auto candidate = + dyn_cast(use.getOwner()->getParentOp())) { + postIfOp = candidate; + break; + } + } + + if (!postIfOp) { + llvm::outs() << "Error: no consuming IfOp found.\n"; + return; + } + + // Step 6. Extend the post IfOp + + scf::IfOp newPostIfOp = addResultsForPostIfOp( + postIfOp, newFrontIfOp, builder, bufferNum, newDepValue, constants, + buffers, frontCnt, postCnt, extraResultIndices, newDeps, + ifResultDeps); + + llvm::outs() << "after addResultsForPostIfOp.\n"; + + int postCntResultIndex = extraResultIndices.back(); + + // Step 7. Update the ForOp yield (buffer propagation) + auto forYield = cast(newForOp.getBody()->getTerminator()); + + // Update buffer1 ~ bufferN + for (int i = 1; i < bufferNum; ++i) { + + int yieldIdx = extraArgBaseIdx + (i - 1); + + if (yieldIdx < forYield->getNumOperands() && + bufferResultIndices[i] < newFrontIfOp.getNumResults()) { + + forYield->setOperand(yieldIdx, + newFrontIfOp.getResult(bufferResultIndices[i])); + + llvm::outs() << "Replaced yield operand " << yieldIdx << "\n"; + } else { + llvm::errs() << "Warning: index out of range\n"; + } + } + + // Step 8. Update frontCnt + OpOperand *frontCntYieldUse = nullptr; + + for (auto &use : frontCnt.getUses()) { + if (isa(use.getOwner()) && + newForOp == use.getOwner()->getParentOp()) { + frontCntYieldUse = &use; + break; + } + } + + frontCntYieldUse->set(newFrontIfOp.getResult(frontCntResultIndex)); + + // Step 9. Update postCnt + OpOperand *postCntYieldUse = nullptr; + + for (auto &use : postCnt.getUses()) { + if (isa(use.getOwner()) && + newForOp == use.getOwner()->getParentOp()) { + postCntYieldUse = &use; + break; + } + } + + postCntYieldUse->set(newPostIfOp.getResult(postCntResultIndex)); + } + } + + llvm::outs() << "multibuffer end!\n"; +} + +// Compute the nesting level of an ifOp within the specified forOp +static int computeIfLevel(scf::IfOp ifOp, scf::ForOp rootForOp) { + int level = 1; + + Operation *parent = ifOp->getParentOp(); + + while (parent && parent != rootForOp.getOperation()) { + if (isa(parent)) + level++; + + parent = parent->getParentOp(); + } + + return level; +} + +int assignIfOpLevels(scf::ForOp forOp) { + SmallVector targetIfOps; + int maxLevel = 0; + // Collect all ifOp assigned with ssbuffer tag + forOp.walk([&](scf::IfOp ifOp) { + if (ifOp->hasAttr("ssbuffer")) { + targetIfOps.push_back(ifOp); + } + }); + + // Caculate buffer levels + for (auto ifOp : targetIfOps) { + int level = computeIfLevel(ifOp, forOp); + maxLevel = std::max(level, maxLevel); + Builder builder(ifOp.getContext()); + ifOp->setAttr("ssbuffer.level", builder.getI32IntegerAttr(level)); + } + return maxLevel; +} + +static bool hasSSBufferIf(scf::ForOp forOp) { + bool found = false; + + forOp.walk([&](scf::IfOp ifOp) { + if (ifOp->hasAttr("ssbuffer")) { + found = true; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + + return found; +} + +static bool hasAncestorSSBufferFor(scf::ForOp forOp) { + Operation *parent = forOp->getParentOp(); + + while (parent) { + if (auto parentFor = dyn_cast(parent)) { + if (hasSSBufferIf(parentFor)) + return true; + } + parent = parent->getParentOp(); + } + + return false; +} + +static bool hasAncestorRootFor(scf::ForOp forOp) { + Operation *parent = forOp->getParentOp(); + + while (parent) { + if (auto parentFor = dyn_cast(parent)) { + if (hasSSBufferIf(parentFor)) + return true; + } + parent = parent->getParentOp(); + } + return false; +} + +SmallVector +collectIfInfo(scf::ForOp &curForOp, + DenseMap> &ifDeps, int level) { + // Find all dependency variables based on the inputs and outputs of ifOp + SmallVector allDeps; + DenseSet producedValues; + scf::ForOp newForOp = nullptr; + curForOp.walk([&](scf::IfOp ifOp) { + auto attr = ifOp->getAttrOfType("ssbuffer.level"); + // No level or level mismatch → continue searching + if (!attr || attr.getInt() != level) + return WalkResult::advance(); + + // Levels match → check the direct parent + if (auto parentFor = dyn_cast(ifOp->getParentOp())) { + newForOp = parentFor; // 更新 + } + + // Stop walking regardless of whether the parent is a for-loop + return WalkResult::interrupt(); + }); + + if (newForOp) + curForOp = newForOp; + + // Step 1: Collect first to preserve order + SmallVector ifOps; + curForOp.walk([&](scf::IfOp ifOp) { + auto curLevel = ifOp->getAttrOfType("ssbuffer.level"); + if (!curLevel || curLevel.getInt() != level) { + return WalkResult::advance(); + } + ifOps.push_back(ifOp); + return WalkResult::advance(); + }); + llvm::outs() << "ifOps:" << ifOps.size() << "\n"; + + int miniDepNum = 2; + if (ifOps.size() < miniDepNum) { + return allDeps; + } + // Step 2: Process in order + for (auto ifOp : ifOps) { + llvm::outs() << "ifOp->getOperands():" << ifOp->getOperands().size() + << "\n"; + SmallVector deps; + if (producedValues.empty()) { + llvm::outs() << "producedValues为空!" + << "\n"; + } + + // inputs + Region &thenRegion = ifOp.getThenRegion(); + for (Operation &op : thenRegion.front()) { + for (Value operand : op.getOperands()) { + for (Value v : producedValues) { + if (operand == v && !llvm::is_contained(deps, operand)) { + deps.push_back(operand); + } + } + } + } + + // outputs + for (Value result : ifOp.getResults()) { + producedValues.insert(result); + } + + if (!deps.empty()) { + ifDeps[ifOp] = deps; + allDeps.append(deps.begin(), deps.end()); + } + } + llvm::outs().flush(); + return allDeps; +} + +bool isCube(scope::ScopeOp scope) { + bool ret = false; + scope.walk([&](Operation *op) { + if (isa(op)) { + ret = true; + } + }); + return ret; +} + +// Traverse each Vector scope, find the outer ForOp, and process internal IfOps +void WalkAIVNestedForAndProcess( + ModuleOp module, DenseMap> &ifResultDeps, + int bufferNum) { + if (bufferNum < 2) { + return; + } + + module.walk([&](scope::ScopeOp scope) { + if (isCube(scope)) { + return; + } + + // Traverse ForOps inside the Cube scope (outer loops) + SmallVector targetFors; + + scope.walk([&](scf::ForOp forOp) { + // Must contain an ssbuffer if + if (!hasSSBufferIf(forOp)) + return WalkResult::advance(); + + // Skip if an ancestor is already the root + if (hasAncestorRootFor(forOp)) + return WalkResult::advance(); + + // Find rootForOp + targetFors.push_back(forOp); + + return WalkResult::advance(); + }); + llvm::outs() << "targetFors: " << targetFors.size(); + int maxLevels; + for (auto outerFor : targetFors) { + ifResultDeps.clear(); + scf::ForOp currentFor = outerFor; + maxLevels = assignIfOpLevels(currentFor); + for (int level = 1; level <= maxLevels; level++) { + auto uniqueDeps = collectIfInfo(currentFor, ifResultDeps, level); + llvm::outs() << "maxLevels:" << maxLevels << "\n"; + if (uniqueDeps.empty()) { + continue; + } + llvm::outs() << "uniqueDeps:" << uniqueDeps.size() << "\n"; + auto newForOp = addDoubleBuffForArgs(module, uniqueDeps, bufferNum); + DenseMap> newIfResultDeps; + auto uniqueList = collectIfInfo(newForOp, newIfResultDeps, level); + addMultiBuffCaculate(module, uniqueList, newIfResultDeps, newForOp, + bufferNum); + } + } + }); +} + +void DAGSSBufferPass::runOnOperation() { + auto module = getOperation(); + + AddIfCondition(module); + + FlowSssbuf(module); + ControlSsbufV2(module); + + // advance不能出现在if里, 规避处理 + ChangeAdvanceOpForm(module); + + DenseMap> ifResultDeps; + WalkAIVNestedForAndProcess(module, ifResultDeps, 2); + + return; +} + +std::unique_ptr> mlir::triton::createDAGSSBufferPass() { + return std::make_unique(); +} diff --git a/compiler/lib/TritonAffinityOpt/DAGScope.cpp b/compiler/lib/TritonAffinityOpt/DAGScope.cpp new file mode 100644 index 00000000..838a8d87 --- /dev/null +++ b/compiler/lib/TritonAffinityOpt/DAGScope.cpp @@ -0,0 +1,1084 @@ + + +#include "dicp/TritonAffinityOpt/Passes.h" + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/HIVM/IR/HIVMImpl.h" +#include "bishengir/Dialect/HIVM/IR/HIVMInterfaces.h" +#include "bishengir/Dialect/HIVM/Transforms/Passes.h" +#include "bishengir/Dialect/HIVM/Utils/Utils.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Operation.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "dicp/Utils/Utils.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include + +#include "dicp/TritonAffinityOpt/DAG.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_DAGSCOPE +#include "dicp/TritonAffinityOpt/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace hivm; + +namespace { +struct DAGScopePass : public mlir::triton::impl::DAGScopeBase { + void runOnOperation() override; +}; +} // namespace + +static std::pair +encapsulateWithScope(triton::FuncOp funcOp) { + Block &entryBlock = funcOp.getBody().front(); + Block &lastBlock = funcOp.getBody().back(); + Operation *terminator = lastBlock.getTerminator(); + + // 辅助函数：判断操作是否应该被跳过 + auto shouldSkipOp = [](Operation *op) -> bool { + return isa(op) || isa(op) || + isa(op); + }; + + // 第三步：准备要移动的操作列表（按顺序） + SmallVector opsToMove; + DenseMap opOrder; + int order = 0; + + // 记录原始顺序并收集需要移动的操作 + for (Operation &op : lastBlock.without_terminator()) { + opOrder[&op] = order++; + if (!shouldSkipOp(&op)) { + opsToMove.push_back(&op); + } + } + + // 按原始顺序排序 + std::sort( + opsToMove.begin(), opsToMove.end(), + [&](Operation *a, Operation *b) { return opOrder[a] < opOrder[b]; }); + + if (opsToMove.empty()) { + return std::make_pair(nullptr, nullptr); + } + + // 第四步：创建scope操作并移动操作 + Operation *lastOpToMove = opsToMove.back(); + OpBuilder builder(&lastBlock, ++lastOpToMove->getIterator()); + + // 创建第一个scope + auto scopeOp = builder.create(builder.getUnknownLoc(), + llvm::ArrayRef{}); + scopeOp.getBodyRegion().emplaceBlock(); + Block *scopeBody = &scopeOp.getBodyRegion().front(); + + // 移动操作到scope中 + OpBuilder scopeBuilder(scopeBody, scopeBody->end()); + DenseMap valueMapping; + + for (Operation *op : opsToMove) { + SmallVector originalResults = op->getResults(); + op->remove(); + scopeBuilder.insert(op); + + // 更新值的映射 + for (size_t i = 0; i < originalResults.size(); ++i) { + valueMapping[originalResults[i]] = op->getResult(i); + } + } + + // 添加return操作 + scopeBuilder.create(builder.getUnknownLoc()); + + // 创建第二个scope（如果需要） + scopeBuilder.setInsertionPointAfter(scopeOp); + auto newScopeOp = scopeBuilder.create( + builder.getUnknownLoc(), llvm::ArrayRef{}); + newScopeOp.getRegion().emplaceBlock(); + + OpBuilder newScopeBuilder(&newScopeOp.getRegion().front(), + newScopeOp.getRegion().front().begin()); + newScopeBuilder.create(scopeOp->getLoc()); + + // 设置属性 + auto vecAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + auto aicAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + + scopeOp->setAttr(hivm::TCoreTypeAttr::name, vecAttr); + newScopeOp->setAttr(hivm::TCoreTypeAttr::name, aicAttr); + + return std::make_pair(scopeOp, newScopeOp); +} + +struct OpMoveInfo { + Operation *op; + Operation *targetParent; // 目标父操作（nullptr表示aicScope本身） +}; + +// 递归遍历函数 - 优化版本 +void collectOpsToMove(Operation *op, AffinityDAG::Graph &graph, + Operation *parentFor, + llvm::SmallVector &aivToMove, + llvm::SmallVector &cubeToMove) { + // 检查当前操作是否需要移动 + bool needsMoveAiv = false; + bool needsMoveCube = false; + auto &valueTypes = graph.getValueTypes(); + // 检查结果类型 + int i = 0; + for (auto res : op->getResults()) { + i++; + if (AffinityDAG::intersects(valueTypes[res], + AffinityDAG::CoreType::VECTOR_ONLY)) { + needsMoveAiv = true; + } + if (AffinityDAG::intersects(valueTypes[res], + AffinityDAG::CoreType::CUBE_ONLY)) { + needsMoveCube = true; + } + } + + if (isa(op)) { + auto res = op->getOperand(0); + if (AffinityDAG::intersects(valueTypes[res], + AffinityDAG::CoreType::VECTOR_ONLY)) { + needsMoveAiv = true; + } + if (AffinityDAG::intersects(valueTypes[res], + AffinityDAG::CoreType::CUBE_ONLY)) { + needsMoveCube = true; + } + } + // 检查特定操作类型 + if (isa(op)) { + needsMoveAiv = true; + } + + // 检查特定操作类型 + if (isa(op)) { + needsMoveCube = true; + } + + // 检查特定操作类型 + if (isa(op) || isa(op) || isa(op)) { + needsMoveAiv = true; + needsMoveCube = true; + } + + if (isa(op)) { + if (auto storeOp = dyn_cast(op)) { + // 获取所有操作数列表 + auto operands = storeOp.getOperands(); + bool typeMatched = false; + + // 按顺序检查第1个、第0个、第2个操作数 + std::vector checkOrder = {1, 0, 2}; + for (size_t idx : checkOrder) { + // 先判断操作数索引是否有效，避免越界访问 + if (idx >= operands.size()) { + continue; + } + auto operand = operands[idx]; + auto coreType = valueTypes[operand]; + + if (coreType == AffinityDAG::CoreType::VECTOR_ONLY) { + needsMoveAiv = true; + typeMatched = true; + } else if (coreType == AffinityDAG::CoreType::CUBE_ONLY) { + needsMoveCube = true; + typeMatched = true; + } + } + // 所有指定操作数都不匹配时，执行原else逻辑 + if (!typeMatched) { + needsMoveAiv = true; + needsMoveCube = true; + } + } + } + + if (isa(op)) { + if (auto assertOp = dyn_cast(op)) { + // 获取所有操作数列表 + auto operand = assertOp.getCondition(); + + auto coreType = valueTypes[operand]; + if (coreType == AffinityDAG::CoreType::VECTOR_ONLY) { + needsMoveAiv = true; + } else if (coreType == AffinityDAG::CoreType::CUBE_ONLY) { + needsMoveCube = true; + } else { + needsMoveAiv = true; + needsMoveCube = true; + } + } + } + + // 检查 Sync 操作的 tcore_type 属性 + if ((isa(op) || isa(op))) { + mlir::OpBuilder builder(op); + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + if (op->getAttr("tcore_type") == coreAttr) { + needsMoveCube = true; + } else { + needsMoveAiv = true; + } + } + + // 如果不需要移动，直接返回 + if (!needsMoveAiv && !needsMoveCube) { + llvm::outs() << "Unsupport Op: " << *op << " \n"; + } + + // 处理 for 循环 + if (auto forOp = dyn_cast(op)) { + // 确定父级 for 循环 + Operation *targetParent = parentFor != nullptr ? parentFor : nullptr; + aivToMove.push_back({op, targetParent}); + cubeToMove.push_back({op, targetParent}); + + // 递归处理循环体 + for (auto &block : forOp.getRegion()) { + for (auto &innerOp : block) { + collectOpsToMove(&innerOp, graph, forOp, aivToMove, cubeToMove); + } + } + } else if (auto ifOp = dyn_cast(op)) { + // 确定父级 for 循环 + Operation *targetParent = parentFor != nullptr ? parentFor : nullptr; + aivToMove.push_back({op, targetParent}); + cubeToMove.push_back({op, targetParent}); + + // 递归处理循环体 + for (auto &block : ifOp.getThenRegion()) { + for (auto &innerOp : block) { + collectOpsToMove(&innerOp, graph, ifOp, aivToMove, cubeToMove); + } + } + + // 检查并遍历IfOp的else分支（如果存在） + for (auto &block : ifOp.getElseRegion()) { + for (auto &innerOp : block) { + collectOpsToMove(&innerOp, graph, ifOp, aivToMove, cubeToMove); + } + } + } else { + if (needsMoveAiv) { + // 处理其他操作 + aivToMove.push_back({op, parentFor}); + } + if (needsMoveCube) { + cubeToMove.push_back({op, parentFor}); + } + } +} + +mlir::Block *getBlockByIndex(mlir::Region ®ion, int blockIndex) { + // 边界校验：索引非法时返回nullptr + if (blockIndex < 0) + return nullptr; + + int currentIdx = 0; + for (auto &block : region) { + if (currentIdx == blockIndex) { + return █ // 找到对应索引的Block，直接返回 + } + currentIdx++; + } + // 索引越界时返回nullptr + return nullptr; +} + +void processOperationToMove( + const OpMoveInfo &info, + llvm::DenseMap &parentMap, + mlir::OpBuilder &builder, mlir::IRMapping &mapper, mlir::Block *aivBlock, + mlir::Operation *terminator, AffinityDAG::Graph &graph, int MoveType) { + // llvm::outs()<<*info.op<<" ssss\n\n\n"; + // llvm::outs().flush(); + // 获取原始Block信息并计算索引 + mlir::Block *originalBlock = info.op->getBlock(); + int originalRegionIndex = -1; + int originalBlockIndex = -1; + int blockCounter = 0; + auto &valueTypes = graph.getValueTypes(); + if (originalBlock) { + mlir::Operation *parentOp = info.op->getParentOp(); // 原始父操作 + if (parentOp) { // 确保父操作存在 + // 老版本MLIR用 getParent() 替代 getParentRegion()，返回值就是Region* + mlir::Region *blockBelongsToRegion = originalBlock->getParent(); + int regionCounter = 0; + for (auto ®ion : parentOp->getRegions()) { // 遍历父操作的所有region + // 直接对比指针，判断当前region是否是block所属的region + if (®ion == blockBelongsToRegion) { + originalRegionIndex = regionCounter; + break; + } + regionCounter++; + } + } + } + + if (originalBlock) { + for (auto &block : originalBlock->getParent()->getBlocks()) { + if (&block == originalBlock) { + originalBlockIndex = blockCounter; + break; + } + blockCounter++; + } + } + + if (originalBlockIndex == -1) { + originalBlockIndex = 0; + } + if (originalRegionIndex == -1) { + originalRegionIndex = 0; + } + + // 处理 scf::ForOp 类型操作 + if (mlir::isa(info.op)) { + auto forOp = mlir::cast(info.op); + + auto getMapped = [&](mlir::Value v) { return mapper.lookupOrDefault(v); }; + auto inputs = forOp.getInitArgs(); + auto outputs = forOp.getResults(); + + // 分离需要移动到aivScope的参数 + llvm::SmallVector aivInputs; + llvm::DenseMap aivInputsMap; + int aivIndex = 1; + + for (int i = 0; i < inputs.size(); ++i) { + if (valueTypes[outputs[i]] != MoveType) { + aivInputs.push_back(inputs[i]); + aivInputsMap[i + 1] = aivIndex; + aivIndex++; + } + } + + // 创建新的for循环 + auto aivForOp = builder.create( + forOp.getLoc(), getMapped(forOp.getLowerBound()), + getMapped(forOp.getUpperBound()), getMapped(forOp.getStep()), + llvm::to_vector(llvm::map_range(aivInputs, getMapped))); + + // 清空循环体 + if (!aivForOp.getBody()->empty()) { + aivForOp.getBody()->getTerminator()->erase(); + } + + // 处理原始循环的yield操作 + auto oldBody = forOp.getBody(); + auto oldYield = + mlir::dyn_cast(oldBody->getTerminator()); + assert(oldYield && "scf::ForOp must have a yield terminator"); + + llvm::SmallVector aivYieldOperands; + for (int i = 0; i < inputs.size(); ++i) { + if (valueTypes[outputs[i]] != MoveType) { + aivYieldOperands.push_back(oldYield.getOperand(i)); + } + } + + // 映射循环参数 + auto oldBodyArgs = forOp.getBody()->getArguments(); + auto aivBodyArgs = aivForOp.getBody()->getArguments(); + + for (auto it = aivInputsMap.begin(); it != aivInputsMap.end(); ++it) { + int oldInputIndex = it->first; + int mappedNewIndex = it->second; + mapper.map(oldBodyArgs[oldInputIndex], aivBodyArgs[mappedNewIndex]); + mapper.map((*info.op).getResults()[oldInputIndex - 1], + aivForOp->getResults()[mappedNewIndex - 1]); + } + mapper.map(oldBodyArgs[0], aivBodyArgs[0]); + + // 将新循环移动到目标位置 + if (info.targetParent == nullptr) { + mlir::Block *targetBlock = aivBlock; + if (terminator) { + aivForOp->moveBefore(terminator); + } else { + aivForOp->moveBefore(targetBlock, targetBlock->end()); + } + parentMap[forOp] = aivForOp; + } else { + auto targetParent = parentMap[info.targetParent]; + auto ®ion = targetParent->getRegion(originalRegionIndex); + + if (region.empty()) { + region.push_back(new mlir::Block()); + } + + mlir::Block *targetBlock = getBlockByIndex(region, originalBlockIndex); + if (targetBlock) { + aivForOp->moveBefore(targetBlock, targetBlock->end()); + parentMap[forOp] = aivForOp; + } else { + llvm::outs() << "Can't find block by index\n"; + } + } + } + + // 处理 scf::YieldOp 类型操作 + else if (mlir::isa(info.op)) { + auto yieldOp = mlir::cast(info.op); + + // 处理父节点为 scf::ForOp 的情况 + if (auto parentForOp = + mlir::dyn_cast(info.targetParent)) { + auto it = parentMap.find(parentForOp); + if (it == parentMap.end()) { + return; + } + auto targetOp = it->second; + auto newForOp = mlir::cast(targetOp); + + auto oldInputs = parentForOp.getInitArgs(); + auto oldOutputs = parentForOp.getResults(); + auto oldYieldOperands = yieldOp.getOperands(); + + llvm::SmallVector newYieldOperands; + for (int i = 0; i < oldInputs.size(); ++i) { + if (valueTypes[oldOutputs[i]] != MoveType) { + mlir::Value oldOperand = oldYieldOperands[i]; + mlir::Value newOperand = mapper.lookupOrDefault(oldOperand); + newYieldOperands.push_back(newOperand); + } + } + + auto newYieldOp = builder.create(yieldOp.getLoc(), + newYieldOperands); + auto ®ion = newForOp->getRegion(0); + mlir::Block *targetBlock = ®ion.front(); + newYieldOp->moveBefore(targetBlock, targetBlock->end()); + } + // 处理父节点为 scf::IfOp 的情况 + else if (auto parentIfOp = + mlir::dyn_cast(info.targetParent)) { + auto it = parentMap.find(parentIfOp); + if (it == parentMap.end()) { + return; + } + auto targetOp = it->second; + auto newIfOp = mlir::cast(targetOp); + + auto oldInputs = parentIfOp.getResults(); + auto oldOutputs = parentIfOp.getResults(); + auto oldYieldOperands = yieldOp.getOperands(); + + llvm::SmallVector newYieldOperands; + for (int i = 0; i < oldInputs.size(); ++i) { + if (valueTypes[oldOutputs[i]] != MoveType) { + mlir::Value oldOperand = oldYieldOperands[i]; + mlir::Value newOperand = mapper.lookupOrDefault(oldOperand); + newYieldOperands.push_back(newOperand); + } + } + + auto ®ion = newIfOp->getRegion(originalRegionIndex); + auto newYieldOp = builder.create(yieldOp.getLoc(), + newYieldOperands); + mlir::Block *targetBlock = getBlockByIndex(region, originalBlockIndex); + if (targetBlock) { + newYieldOp->moveBefore(targetBlock, targetBlock->end()); + } else { + llvm::outs() << "Can't find block by index\n"; + } + } + } + + // 处理 scf::IfOp 类型操作 + else if (mlir::isa(info.op)) { + auto ifOp = mlir::cast(info.op); + + auto getMapped = [&](mlir::Value v) { return mapper.lookupOrDefault(v); }; + mlir::Value condition = ifOp.getCondition(); + + // 分离需要移动到aivScope的结果 + llvm::SmallVector aivResults; + llvm::SmallVector aivResultTypes; + llvm::DenseMap aivResultMap; + int aivResultIndex = 0; + + for (int i = 0; i < ifOp.getNumResults(); ++i) { + mlir::Value result = ifOp.getResult(i); + if (valueTypes[result] != MoveType) { + aivResults.push_back(result); + aivResultTypes.push_back(result.getType()); + aivResultMap[i] = aivResultIndex; + aivResultIndex++; + } + } + + // 创建新的if操作 + auto aivIfOp = builder.create( + ifOp.getLoc(), aivResultTypes, getMapped(condition)); + + // 映射if操作结果 + for (auto &[oldIdx, newIdx] : aivResultMap) { + mapper.map(ifOp.getResult(oldIdx), aivIfOp.getResult(newIdx)); + } + + // 初始化then和else区域 + mlir::Region &thenRegion = aivIfOp.getThenRegion(); + mlir::Block *thenBlock = new mlir::Block(); + thenRegion.push_back(thenBlock); + + mlir::Region &elseRegion = ifOp.getElseRegion(); + if (!elseRegion.empty()) { + mlir::Region &elseRegion = aivIfOp.getElseRegion(); + mlir::Block *elseBlock = new mlir::Block(); + elseRegion.push_back(elseBlock); + } + + // 将新if操作移动到目标位置 + if (info.targetParent == nullptr) { + mlir::Block *targetBlock = aivBlock; + if (terminator) { + aivIfOp->moveBefore(terminator); + } else { + aivIfOp->moveBefore(targetBlock, targetBlock->end()); + } + parentMap[ifOp] = aivIfOp; + } else { + auto ®ion = + parentMap[info.targetParent]->getRegion(originalRegionIndex); + if (region.empty()) { + region.push_back(new mlir::Block()); + } + + mlir::Block *targetBlock = getBlockByIndex(region, originalBlockIndex); + if (targetBlock) { + aivIfOp->moveBefore(targetBlock, targetBlock->end()); + parentMap[ifOp] = aivIfOp; + } else { + llvm::outs() << "Can't find block by index\n"; + } + } + } + + // 处理其他类型操作（克隆） + else { + auto clonedOp = builder.clone(*info.op, mapper); + auto numberRes = clonedOp->getNumResults(); + for (auto i = 0; i < numberRes; i++) { + mapper.map((*info.op).getResults()[i], clonedOp->getResults()[i]); + } + + if (info.targetParent == nullptr) { + mlir::Block *targetBlock = aivBlock; + clonedOp->moveBefore(terminator); + parentMap[info.op] = clonedOp; + } else { + auto parentIt = parentMap.find(info.targetParent); + auto mappedParentOp = parentIt->second; + auto ®ion = mappedParentOp->getRegion(originalRegionIndex); + + if (region.empty()) { + region.push_back(new mlir::Block()); + } + + mlir::Block *targetBlock = getBlockByIndex(region, originalBlockIndex); + if (targetBlock) { + clonedOp->moveBefore(targetBlock, targetBlock->end()); + } else { + llvm::outs() << "Can't find block by index\n"; + } + } + } +} + +static void SplitScope(triton::FuncOp funcOp, AffinityDAG::Graph &graph, + Operation *aivScope, Operation *aicScope, + ModuleOp module) { + llvm::SmallVector aivToMove; + llvm::SmallVector cubeToMove; + for (auto &block : aivScope->getRegion(0)) { + for (auto &op : block) { + collectOpsToMove(&op, graph, nullptr, aivToMove, cubeToMove); + } + } + mlir::IRMapping aivmapper; + mlir::OpBuilder builder(aivScope); + llvm::DenseMap aivparentMap; + + // 第二遍：实际移动操作 + // 先移动for循环 + mlir::Block *aivBlock = + &aivScope->getRegion(0).front(); // 或者使用合适的block + SmallVector deleteOp; + auto *terminator = aivBlock->getTerminator(); + // 如果操作已被使用，直接跳过 + llvm::SmallVector + aivUsedOp; // 改为函数内静态，保持原有逻辑 + for (const auto &info : aivToMove) { + if (std::find(aivUsedOp.begin(), aivUsedOp.end(), info.op) != + aivUsedOp.end()) { + return; + } + aivUsedOp.push_back(info.op); + processOperationToMove(info, aivparentMap, builder, aivmapper, aivBlock, + terminator, graph, AffinityDAG::CoreType::CUBE_ONLY); + } + + llvm::DenseMap aicparentMap; + mlir::IRMapping aicmapper; + mlir::Block *aicBlock = + &aicScope->getRegion(0).front(); // 或者使用合适的block + terminator = aicBlock->getTerminator(); + llvm::SmallVector + aicUsedOp; // 改为函数内静态，保持原有逻辑 + for (const auto &info : cubeToMove) { + if (std::find(aicUsedOp.begin(), aicUsedOp.end(), info.op) != + aicUsedOp.end()) { + return; + } + aicUsedOp.push_back(info.op); + processOperationToMove(info, aicparentMap, builder, aicmapper, aicBlock, + terminator, graph, + AffinityDAG::CoreType::VECTOR_ONLY); + } + + for (const auto &info : aivToMove) { + if (std::find(deleteOp.begin(), deleteOp.end(), info.op) == + deleteOp.end()) { + deleteOp.push_back(info.op); + } + } + for (const auto &info : cubeToMove) { + if (std::find(deleteOp.begin(), deleteOp.end(), info.op) == + deleteOp.end()) { + deleteOp.push_back(info.op); + } + } + + // llvm::outs() << "\n" << module<<" ====== ddd ====== \n\n\n"; + // llvm::outs().flush(); + for (auto it = deleteOp.rbegin(); it != deleteOp.rend(); ++it) { + (*it)->erase(); // 解引用反向迭代器，调用 erase 方法 + } + return; +} + +/// 创建setop +static hivm::SyncBlockSetOp +createSyncBlockSetOp(OpBuilder &builder, Location loc, hivm::TCoreType coreType, + hivm::PIPE setPipeEnum, hivm::PIPE waitPipeEnum, + int64_t flag) { + MLIRContext *ctx = builder.getContext(); + auto coreAttr = hivm::TCoreTypeAttr::get(ctx, coreType); + auto setPipe = hivm::PipeAttr::get(ctx, setPipeEnum); + auto waitPipe = hivm::PipeAttr::get(ctx, waitPipeEnum); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + return builder.create(loc, coreAttr, setPipe, waitPipe, + flagId); +} + +/// 创建waitop +static hivm::SyncBlockWaitOp +createSyncBlockWaitOp(OpBuilder &builder, Location loc, + hivm::TCoreType coreType, hivm::PIPE setPipeEnum, + hivm::PIPE waitPipeEnum, int64_t flag) { + MLIRContext *ctx = builder.getContext(); + auto coreAttr = hivm::TCoreTypeAttr::get(ctx, coreType); + auto setPipe = hivm::PipeAttr::get(ctx, setPipeEnum); + auto waitPipe = hivm::PipeAttr::get(ctx, waitPipeEnum); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + return builder.create(loc, coreAttr, setPipe, waitPipe, + flagId); +} + +// 在scope return前插入wait +static void insertWaitBeforeFinalReturn(Region *region, OpBuilder &builder, + int64_t flag, bool coretypebool) { + for (Block &block : *region) { + if (auto returnOp = + dyn_cast_or_null(block.getTerminator())) { + builder.setInsertionPoint(returnOp); + if (coretypebool) { + createSyncBlockWaitOp(builder, returnOp->getLoc(), + hivm::TCoreType::CUBE, hivm::PIPE::PIPE_V, + hivm::PIPE::PIPE_FIX, flag); + return; + } else { + createSyncBlockWaitOp(builder, returnOp->getLoc(), + hivm::TCoreType::VECTOR, hivm::PIPE::PIPE_M, + hivm::PIPE::PIPE_MTE3, flag); + return; + } + } + } +} + +/// 在scope内起始位置加上set +static void insertSetAtRegionStart(Region *region, OpBuilder &builder, + int64_t flag, bool coretypebool) { + if (!region->empty()) { + Block &entry = region->front(); + Location loc = entry.empty() ? region->getParentOp()->getLoc() + : entry.front().getLoc(); + builder.setInsertionPointToStart(&entry); + if (coretypebool) { + createSyncBlockSetOp(builder, loc, hivm::TCoreType::VECTOR, + hivm::PIPE::PIPE_V, hivm::PIPE::PIPE_FIX, flag); + } else { + createSyncBlockSetOp(builder, loc, hivm::TCoreType::CUBE, + hivm::PIPE::PIPE_M, hivm::PIPE::PIPE_MTE3, flag); + } + } +} + +static Operation *findNextSyncBlockSetAfter(Operation *startOp) { + Block *block = startOp->getBlock(); + auto it = ++startOp->getIterator(); + for (; it != block->end(); ++it) { + if (isa(*it)) + return &*it; + } + return nullptr; +} + +static hivm::SyncBlockWaitOp findWaitOpInRegionWithFlag(Region *region, + int64_t flag) { + hivm::SyncBlockWaitOp result; + region->walk([&](hivm::SyncBlockWaitOp op) { + auto flagAttr = op->getAttrOfType("static_flag_id"); + if (flagAttr && flagAttr.getInt() == flag) { + result = op; + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + return result; +} + +static Operation *findInsertionPointAfterWaitForAIV(Operation *waitOp) { + Block *block = waitOp->getBlock(); + auto it = ++waitOp->getIterator(); + + for (; it != block->end(); ++it) { + if (isa(*it) || isa(*it)) { + break; + } + } + + while (it != block->begin()) { + auto prevIt = std::prev(it); + if (isa(*prevIt)) { + it = prevIt; + } else { + break; + } + } + + return &*it; +} + +static Operation *findInsertionPointAfterWaitForAIC(Operation *waitOp) { + Block *block = waitOp->getBlock(); + auto it = ++waitOp->getIterator(); + for (; it != block->end(); ++it) { + if (auto fixpipe = dyn_cast(*it)) { + if (it != block->begin()) { + auto prev = std::prev(it); + if (isa(*prev)) + return &*prev; + } + return &*it; + } + if (isa(*it)) + return &*it; + } + return nullptr; +} + +// 查找 FixpipeOp 下一行的 sync_block_set 操作的 flag 值 +static int findFixPipeFlagSafe(hivm::FixpipeOp fixpipeOp) { + mlir::Operation *fixpipeOperation = fixpipeOp.getOperation(); + if (!fixpipeOperation || !fixpipeOperation->getBlock()) { + return -1; + } + + // 获取 FixpipeOp 的迭代器 + auto it = ++fixpipeOperation->getIterator(); + + // 遍历后续操作直到找到 sync_block_set + while (it != fixpipeOperation->getBlock()->end()) { + mlir::Operation &op = *it++; + + if (op.getName().getStringRef() == "hivm.hir.sync_block_set") { + auto staticFlagAttr = + op.getAttrOfType("static_flag_id"); + return staticFlagAttr.getInt(); + break; + } + } + + return -1; +} + +/// cube处理逻辑 +static void processFixpipeOpsInAIC(Region *aicRegion, Region *aivRegion) { + + MLIRContext *ctx = aicRegion->getContext(); + OpBuilder builder(ctx); + SmallVector fixpipes; + aicRegion->walk([&](hivm::FixpipeOp op) { fixpipes.push_back(op); }); + + for (auto fixpipeOp : fixpipes) { + + auto newflag = findFixPipeFlagSafe(fixpipeOp); + // 1. 在 FixpipeOp 前插 Wait + builder.setInsertionPoint(fixpipeOp); + createSyncBlockWaitOp(builder, fixpipeOp->getLoc(), hivm::TCoreType::CUBE, + hivm::PIPE::PIPE_V, hivm::PIPE::PIPE_FIX, newflag); + bool coretypebool = true; + + // 2. 在 aicRegion 末尾 Return 前插 Wait + insertWaitBeforeFinalReturn(aicRegion, builder, newflag, coretypebool); + + // 3. 在 aivRegion 开头插 Set + insertSetAtRegionStart(aivRegion, builder, newflag, coretypebool); + + // 4. 在 aicRegion 向后找 SyncBlockSetOp + if (auto *nextSetOp = findNextSyncBlockSetAfter(fixpipeOp)) { + auto setFlagAttr = + nextSetOp->getAttrOfType("static_flag_id"); + // 调试：打印set + // llvm::dbgs() << "aicnextSetOp:"; + // nextSetOp->dump(); + if (!setFlagAttr) { + llvm::dbgs() << "AIC can not find setop in aic\n"; + continue; + } + int64_t setflag = setFlagAttr.getInt(); + + // 5. 在 aivRegion 中找 flag=setflag 的 WaitOp + auto targetWait = findWaitOpInRegionWithFlag(aivRegion, setflag); + if (!targetWait) { + llvm::dbgs() << "AIC can not find waitop in aiv\n"; + continue; + } + + // 调试：打印wait + // llvm::dbgs() << "aictargetWait:"; + // llvm::dbgs() << targetWait << "\n"; + + // 6. 从该 Wait 向下找 ToMemrefOp 或 Yield，插 Set(newflag) + if (auto *insertPt = findInsertionPointAfterWaitForAIV(targetWait)) { + builder.setInsertionPoint(insertPt); + createSyncBlockSetOp(builder, fixpipeOp->getLoc(), + hivm::TCoreType::VECTOR, hivm::PIPE::PIPE_V, + hivm::PIPE::PIPE_FIX, newflag); + } + } + } +} + +// 查找 copyOp 下一行的 sync_block_set 操作的 flag 值 +static int findCopyFlagSafe(bufferization::ToBufferOp toMemrefOp) { + mlir::Operation *toMemrefOperation = toMemrefOp.getOperation(); + if (!toMemrefOperation || !toMemrefOperation->getBlock()) { + return -1; + } + + // 获取 copyOp 的迭代器 + auto it = ++toMemrefOperation->getIterator(); + + // 遍历后续操作直到找到 sync_block_set + while (it != toMemrefOperation->getBlock()->end()) { + mlir::Operation &op = *it++; + + if (op.getName().getStringRef() == "hivm.hir.sync_block_set") { + auto staticFlagAttr = + op.getAttrOfType("static_flag_id"); + return staticFlagAttr.getInt(); + break; + } + } + + return -1; +} +/// vector处理逻辑 +static void processToMemrefOpsInAIV(Region *aivRegion, Region *aicRegion) { + + MLIRContext *ctx = aivRegion->getContext(); + OpBuilder builder(ctx); + SmallVector toMemrefs; + aivRegion->walk( + [&](bufferization::ToBufferOp op) { toMemrefs.push_back(op); }); + + for (auto toMemrefOp : toMemrefs) { + auto newflag = findCopyFlagSafe(toMemrefOp); + + // 1. 在 ToMemrefOp 前插 Wait + builder.setInsertionPoint(toMemrefOp); + createSyncBlockWaitOp(builder, toMemrefOp->getLoc(), + hivm::TCoreType::VECTOR, hivm::PIPE::PIPE_M, + hivm::PIPE::PIPE_MTE3, newflag); + bool coretypebool = false; + + // 2. 在 aivRegion 末尾 Return 前插 Wait + insertWaitBeforeFinalReturn(aivRegion, builder, newflag, coretypebool); + + // 3. 在 aicRegion 开头插 Set + insertSetAtRegionStart(aicRegion, builder, newflag, coretypebool); + + // 4. 在 aivRegion 向后找 SyncBlockSetOp + if (auto *nextSetOp = findNextSyncBlockSetAfter(toMemrefOp)) { + auto setFlagAttr = + nextSetOp->getAttrOfType("static_flag_id"); + // 调试：打印set及其所有attribute + // llvm::dbgs() << "aivnextSetOp:"; + // nextSetOp->dump(); + // llvm::dbgs() << "Attributes:\n"; + // for (auto namedAttr : nextSetOp->getAttrs()) { + // llvm::dbgs() << " " << namedAttr.getName() << " = "; + // namedAttr.getValue().print(llvm::dbgs()); + // llvm::dbgs() << "\n"; + // } + if (!setFlagAttr) { + llvm::dbgs() << "AIV can not find setop in aiv\n"; + continue; + } + int64_t setflag = setFlagAttr.getInt(); + + // 5. 在 aicRegion 中找 flag=setflag 的 WaitOp + auto targetWait = findWaitOpInRegionWithFlag(aicRegion, setflag); + + if (!targetWait) { + llvm::dbgs() << "AIV can not find waitop in aic\n"; + continue; + } + + // 调试：打印wait + // llvm::dbgs() << "aivtargetWait:"; + // llvm::dbgs() << targetWait << "\n"; + + // 6. 从该 Wait 向下找 Fixpipe 前 Wait 或 Yield，插 Set(newflag) + if (auto *insertPt = findInsertionPointAfterWaitForAIC(targetWait)) { + builder.setInsertionPoint(insertPt); + createSyncBlockSetOp(builder, toMemrefOp->getLoc(), + hivm::TCoreType::CUBE, hivm::PIPE::PIPE_M, + hivm::PIPE::PIPE_MTE3, newflag); + } + } + } +} + +/// 同步点增强 +void addSyncOpsForBufferWait(ModuleOp module) { + for (auto funcOp : + llvm::make_early_inc_range(module.getOps())) { + if (funcOp.getBody().empty()) { + continue; + } + + Region *aicRegion = nullptr; + Region *aivRegion = nullptr; + + funcOp.walk([&](scope::ScopeOp scopeOp) { + auto coreTypeAttr = scopeOp->getAttrOfType( + hivm::TCoreTypeAttr::name); + if (!coreTypeAttr) + return; + + if (coreTypeAttr.getTcoretype() == hivm::TCoreType::CUBE) { + aicRegion = &scopeOp.getRegion(); + } + if (coreTypeAttr.getTcoretype() == hivm::TCoreType::VECTOR) { + aivRegion = &scopeOp.getRegion(); + } + }); + + if (!aicRegion || !aivRegion) { + continue; + } + + processFixpipeOpsInAIC(aicRegion, aivRegion); + processToMemrefOpsInAIV(aivRegion, aicRegion); + } +} + +void DAGScopePass::runOnOperation() { + auto module = getOperation(); + // llvm::outs()<())) { + // skip invalid function + if (funcOp.getBody().empty()) { + continue; + } + + // 收集所有 memref.alloc 操作 + llvm::SmallVector allocOps; + + // 遍历函数中的所有操作（包括嵌套区域中的操作） + funcOp.walk([&](mlir::Operation *op) { + if (mlir::isa(op)) { + allocOps.push_back(op); + } + }); + + mlir::Block &entryBlock = funcOp.getBody().front(); + mlir::Block::iterator insertPos = entryBlock.begin(); + + // 将 alloc 操作移动到函数的最前面 + for (mlir::Operation *allocOp : allocOps) { + // 如果 alloc 操作已经是最前面的操作，跳过 + if (allocOp->getBlock() == &entryBlock && + allocOp->isBeforeInBlock(&*insertPos)) { + continue; + } + + // 将 alloc 操作移动到指定位置 + allocOp->moveBefore(&entryBlock, insertPos); + } + + auto funcName = funcOp.getName(); + auto *graph_ptr = + AffinityDAG::GraphManager::getInstance().getGraph(funcName); + if (!graph_ptr) { + continue; + } + auto &main_graph = *graph_ptr; + + auto ScopeList = encapsulateWithScope(funcOp); + auto aivScope = ScopeList.first; // 第一个元素 + auto aicScope = ScopeList.second; // 第二个元素 + + SplitScope(funcOp, main_graph, aivScope, aicScope, module); + } + + addSyncOpsForBufferWait(module); + // llvm::outs()<> mlir::triton::createDAGScopePass() { + return std::make_unique(); +} diff --git a/compiler/lib/TritonAffinityOpt/DAGSync.cpp b/compiler/lib/TritonAffinityOpt/DAGSync.cpp new file mode 100644 index 00000000..9bf0a70f --- /dev/null +++ b/compiler/lib/TritonAffinityOpt/DAGSync.cpp @@ -0,0 +1,1618 @@ +#include "dicp/TritonAffinityOpt/Passes.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/HIVM/IR/HIVMImpl.h" +#include "bishengir/Dialect/HIVM/IR/HIVMInterfaces.h" +#include "bishengir/Dialect/HIVM/Transforms/Passes.h" +#include "bishengir/Dialect/HIVM/Utils/Utils.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" +#include "dicp/TritonAffinityOpt/Utils.hpp" +#include "mlir/Analysis/AliasAnalysis.h" +#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h" +#include "mlir/Analysis/DataFlowFramework.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Block.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Dominance.h" +#include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" +#include "mlir/Interfaces/LoopLikeInterface.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallBitVector.h" +#include "llvm/Support/Casting.h" + +#include "dicp/Utils/Utils.h" +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/FormatVariadic.h" +#include +#include +#include + +#include "dicp/TritonAffinityOpt/DAG.h" +#include "triton/Dialect/Triton/IR/Types.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_DAGSYNC +#include "dicp/TritonAffinityOpt/Passes.h.inc" +} // namespace triton +} // namespace mlir + +#define DEBUG_TYPE "triton-affinity-dag-sync" + +// 使用 DAG 命名空间 +using namespace mlir; +using namespace hivm; +using namespace AffinityDAG; + +constexpr size_t MAX_FLAG_ID = 14; +llvm::DenseMap *valueTypes; +// 修改类声明，将数据搬运逻辑集成到同步插入中 +namespace { +struct DAGSyncPass : public mlir::triton::impl::DAGSyncBase { + void runOnOperation() override; + +private: + // 原有的辅助函数 + CoreType getNodeDeviceType(OpNode *node, + llvm::DenseMap *valueTypes); + bool needVectorCubeSync(CoreType src, CoreType dst); + + // 修改后的同步插入函数，包含数据搬运 + void insertSyncAndMovement(mlir::Operation *srcOp, mlir::Operation *dstOp, + CoreType srcType, CoreType dstType, + mlir::OpBuilder &builder, int flag, + llvm::DenseMap *valueMap, + Graph &mainGraph); + + // 新增：处理跨 block 的同步和数据搬运 + void insertSyncAndMovementForCrossBlock( + mlir::Operation *srcOp, mlir::Operation *dstOp, CoreType srcType, + CoreType dstType, mlir::OpBuilder &builder, int flag, + bool dstIsInnerBlock, llvm::DenseMap *valueMap, + Graph &mainGraph); + + // 新增：处理 scf.for 循环迭代参数的同步 + void processScfForSync(mlir::scf::ForOp forOp, Node *forNode, + llvm::DenseMap *valueTypes, + mlir::OpBuilder &builder, int &flag); + + // 数据搬运相关的辅助函数 + void insertCubeToVectorDataMovement(mlir::Operation *srcOp, + mlir::Operation *dstOp, + mlir::Value srcResult, + mlir::OpBuilder &builder, + mlir::Location loc, mlir::Value iterArgs); + + void + insertVectorToCubeDataMovement(mlir::Operation *srcOp, mlir::Operation *dstOp, + Operation *posOp, mlir::Value srcResult, + mlir::OpBuilder &builder, mlir::Location loc, + llvm::DenseMap *valueMap); + + // 获取或创建合适的 memref.alloc + mlir::Value getOrCreateAllocation(mlir::Operation *op, mlir::Type tensorType, + hivm::AddressSpace addressSpace, + mlir::OpBuilder &builder, + mlir::Location loc); + + // 获取 tensor 的形状和元素类型 + mlir::RankedTensorType getTensorType(mlir::Value tensorValue); + + // 替换 dstOp 中使用 srcResult 的操作数 + void replaceOperandWithNewValue(mlir::Operation *dstOp, mlir::Value oldValue, + mlir::Value newValue); + + // Find sync position + Operation *FindLastestPosition(Operation *srcOp, Graph &mainGraph, + OpBuilder &builder); + Operation *FindEarliestPosition(Operation *dstOp, Graph &mainGraph, + OpBuilder &builder); +}; +} // namespace + +void DAGSyncPass::processScfForSync( + mlir::scf::ForOp forOp, Node *forNode, + llvm::DenseMap *valueTypes, mlir::OpBuilder &builder, + int &flag) { + + mlir::Block *loopBody = forOp.getBody(); + mlir::scf::YieldOp yieldOp = nullptr; + for (mlir::Operation &op : *loopBody) { + if (auto yield = mlir::dyn_cast(&op)) { + yieldOp = yield; + break; + } + } + Location loc = forOp.getLoc(); + + for (int i = 0; i < forOp.getInitArgs().size(); i++) { + mlir::BlockArgument iterArg = loopBody->getArgument(i + 1); + // 找到首次使用 + mlir::Operation *firstUser = nullptr; + + for (mlir::Operation &op : *loopBody) { + // 跳过 yield 操作 + if (mlir::isa(&op)) { + continue; + } + + // 检查是否使用该迭代参数 + bool usesIterArg = false; + for (mlir::Value operand : op.getOperands()) { + if (operand == iterArg) { + usesIterArg = true; + break; + } + } + + if (usesIterArg) { + firstUser = &op; + break; + } + } + // map 内找到对应的iterType，iterType由首次在loop内使用到的op定义 + if (!firstUser) { + continue; + } + CoreType iterType = CoreType::CUBE_AND_VECTOR; + if (valueTypes->find(firstUser->getResult(0)) != valueTypes->end()) { + iterType = valueTypes->find(firstUser->getResult(0))->second; + } + + // 获取对应yield + mlir::Value yieldOperand = yieldOp->getOperand(i); + CoreType yieldType = CoreType::CUBE_AND_VECTOR; + if (valueTypes->find(yieldOperand) != valueTypes->end()) { + yieldType = valueTypes->find(yieldOperand)->second; + } + mlir::Operation *yieldDefiningOp = yieldOperand.getDefiningOp(); + + if (yieldType == CoreType::CUBE_ONLY && iterType == CoreType::VECTOR_ONLY) { + + // 2. 插入同步指令 + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_FIX); + auto waitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_V); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + + // set 在 yieldDefiningOp 后 + builder.setInsertionPointAfter(yieldDefiningOp); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + mlir::Value srcResult = yieldDefiningOp->getResult(0); + + // // 1. 插入数据搬运 + insertCubeToVectorDataMovement(yieldDefiningOp, firstUser, srcResult, + builder, loc, iterArg); + + // wait 在 firstUser 前 + builder.setInsertionPoint(firstUser); + coreAttr = hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + // llvm::outs() << "yieldOp" << yieldDefiningOp << "iterargs" << firstUser + // << "\n"; llvm::outs() << "Inserted CUBE->VECTOR sync and data movement + // (flag=" << flag << ")\n"; + } + // VECTOR -> CUBE + else if (yieldType == CoreType::VECTOR_ONLY && + iterType == CoreType::CUBE_ONLY) { + + // 2. 插入同步指令 + auto coreAttr = hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE3); + auto waitPipe = + PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE1); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + + // set 在 yieldDefiningOp 后 + builder.setInsertionPointAfter(yieldDefiningOp); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // 1. 插入数据搬运 + // insertVectorToCubeDataMovement(yieldDefiningOp, firstUser, srcResult, + // builder, loc, iterArg); + + // wait 在 firstUser 前 + builder.setInsertionPoint(firstUser); + coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + // llvm::outs() << "yieldOp" << yieldDefiningOp << "iterargs" << firstUser + // << "\n"; llvm::outs() << "Inserted VECTOR->CUBE sync and data movement + // (flag=" << flag << ")\n"; + } + } +} + +// 获取节点的设备类型 +CoreType DAGSyncPass::getNodeDeviceType( + OpNode *node, llvm::DenseMap *valueTypes) { + if (!node || !node->op) { + return CoreType::CUBE_AND_VECTOR; + } + + // 尝试从节点的结果中获取设备类型 + // 通常使用第一个结果来代表节点的设备类型 + if (node->op->getNumResults() > 0) { + mlir::Value result = node->op->getResult(0); + auto it = valueTypes->find(result); + if (it != valueTypes->end()) { + return it->second; + } + } + + // 如果没有找到，检查操作数 + // for (mlir::Value operand : node->op->getOperands()) { + // auto it = valueTypes->find(operand); + // if (it != valueTypes->end()) { + // return it->second; + // } + // } + + return CoreType::CUBE_AND_VECTOR; // 默认 +} + +// 判断是否需要vector<->cube同步 +bool DAGSyncPass::needVectorCubeSync(CoreType src, CoreType dst) { + return (src == CoreType::VECTOR_ONLY && dst == CoreType::CUBE_ONLY) || + (src == CoreType::CUBE_ONLY && dst == CoreType::VECTOR_ONLY); +} + +// 获取 tensor 类型 +mlir::RankedTensorType DAGSyncPass::getTensorType(mlir::Value tensorValue) { + if (auto tensorType = + dyn_cast(tensorValue.getType())) { + return tensorType; + } + return nullptr; +} + +// 替换操作数 +void DAGSyncPass::replaceOperandWithNewValue(mlir::Operation *dstOp, + mlir::Value oldValue, + mlir::Value newValue) { + for (unsigned i = 0; i < dstOp->getNumOperands(); ++i) { + if (dstOp->getOperand(i) == oldValue) { + dstOp->setOperand(i, newValue); + // llvm::outs() << "Replaced operand " << i << " of " << + // dstOp->getName().getStringRef() + // << " with new value\n"; + } + } +} + +// 修改 getOrCreateAllocation 函数，将 alloc 提到函数最外层 +mlir::Value DAGSyncPass::getOrCreateAllocation(mlir::Operation *op, + mlir::Type tensorType, + hivm::AddressSpace addressSpace, + mlir::OpBuilder &builder, + mlir::Location loc) { + auto rankedTensorType = cast(tensorType); + auto elementType = rankedTensorType.getElementType(); + auto shape = rankedTensorType.getShape(); + + auto addressSpaceAttr = + hivm::AddressSpaceAttr::get(builder.getContext(), addressSpace); + auto memrefType = mlir::MemRefType::get(shape, elementType, + /*layout=*/nullptr, addressSpaceAttr); + + // 查找是否已经存在相同类型的 allocation（在函数的 entry block 中） + mlir::Operation *funcOp = op; + while (funcOp && !mlir::isa(funcOp)) { + funcOp = funcOp->getParentOp(); + } + + if (auto func = mlir::dyn_cast(funcOp)) { + // 在函数的 entry block 中查找现有的 allocation + mlir::Block &entryBlock = func.getBody().front(); + // for (auto& blockOp : entryBlock) { + // if (auto allocOp = mlir::dyn_cast(&blockOp)) { + // if (allocOp.getType() == memrefType) { + // // 找到匹配的 allocation，直接复用 + // llvm::outs() << "Reusing existing allocation: " << allocOp << + // "\n"; return allocOp.getResult(); + // } + // } + // } + + // 没有找到现有的 allocation，在函数开头创建新的 + builder.setInsertionPointToStart(&entryBlock); + return builder.create(loc, memrefType); + } + + // 如果没有找到函数，回退到原逻辑 + builder.setInsertionPoint(op); + return builder.create(loc, memrefType); +} + +// 插入 CUBE -> VECTOR 数据搬运 +void DAGSyncPass::insertCubeToVectorDataMovement( + mlir::Operation *srcOp, mlir::Operation *dstOp, mlir::Value srcResult, + mlir::OpBuilder &builder, mlir::Location loc, mlir::Value iterArgs) { + auto srcTensorType = getTensorType(srcResult); + if (!srcTensorType) { + return; + } + + // 1. 在 srcOp 之后创建 UB 空间的 memref.alloc + builder.setInsertionPointAfter(srcOp); + mlir::Value ubAlloc = getOrCreateAllocation( + srcOp, srcTensorType, hivm::AddressSpace::UB, builder, loc); + + // 2. 创建 fixpipe 指令 + builder.setInsertionPointAfter(srcOp); + FixpipeDMAModeAttr dmaModeAttr = + FixpipeDMAModeAttr::get(builder.getContext(), FixpipeDMAMode::NZ2ND); + + auto fixpipeOp = + builder.create(loc, mlir::TypeRange{}, // 没有返回值 + srcResult, // src + ubAlloc, // dst + /*unit_flag_cond=*/mlir::ValueRange{}, + /*dma_mode=*/dmaModeAttr, + /*dual_dst_mode=*/nullptr, + /*pre_quant=*/nullptr, + /*pre_relu=*/nullptr, + /*channel_split=*/nullptr, + /*unit_flag_mode=*/mlir::ArrayAttr{}, + /*quant_scale=*/nullptr); + + llvm::outs() << "Inserted fixpipe after " << srcOp->getName().getStringRef() + << " for CUBE->VECTOR data movement\n"; + + // 3. 在 dstOp 前创建 memory_space_cast 和 to_tensor + builder.setInsertionPoint(dstOp); + + // memory_space_cast（如果需要） + mlir::Value plainMemref = ubAlloc; + auto memrefType = cast(ubAlloc.getType()); + if (memrefType.getMemorySpace()) { + auto plainMemrefType = mlir::MemRefType::get(memrefType.getShape(), + memrefType.getElementType()); + plainMemref = builder.create( + loc, plainMemrefType, ubAlloc); + (*valueTypes)[plainMemref] = CoreType::VECTOR_ONLY; + } + + // 4. 创建 to_tensor + auto toTensorOp = builder.create( + loc, + srcTensorType, // 原始的 tensor 类型 + plainMemref, + /*restrict=*/true, + /*writable=*/true); + (*valueTypes)[toTensorOp.getResult()] = CoreType::VECTOR_ONLY; + + // 5. 替换 dstOp 的操作数 + if (!iterArgs) { + replaceOperandWithNewValue(dstOp, srcResult, toTensorOp.getResult()); + } else { + replaceOperandWithNewValue(dstOp, iterArgs, toTensorOp.getResult()); + } +} + +static uint64_t getElemBytesForAlign(Type t) { + if (auto ft = dyn_cast(t)) + return (uint64_t)((ft.getWidth() + 7) / 8); + if (auto it = dyn_cast(t)) + return (uint64_t)((it.getWidth() + 7) / 8); + if (isa(t)) + return 8ULL; + if (auto ct = dyn_cast(t)) + return 2ULL * getElemBytesForAlign(ct.getElementType()); + return 0ULL; +} + +static FailureOr getBlockElemsFor32BAlign(Type elemType) { + constexpr uint64_t kAlignBytes = 32; + uint64_t elemBytes = getElemBytesForAlign(elemType); + if (elemBytes <= 0) + return failure(); + if (elemBytes >= kAlignBytes) + return 1; + if (kAlignBytes % elemBytes != 0) + return failure(); + return kAlignBytes / elemBytes; +} + +static std::optional> +newCbubAllocShape(memref::AllocOp allocOp) { + auto type = dyn_cast(allocOp.getType()); + // 仅支持静态 2D MemRef + if (!type || type.getRank() != 2) + return std::nullopt; + + auto shape = type.getShape(); + int64_t M = shape[0]; + int64_t N = shape[1]; + auto elemType = type.getElementType(); + auto blkOr = getBlockElemsFor32BAlign(elemType); + int64_t blk = (int64_t)*blkOr; + // 必须是静态且 16 对齐 + if (ShapedType::isDynamic(M) || ShapedType::isDynamic(N)) + return std::nullopt; + if (M % 16 != 0) + return std::nullopt; + + // 新 shape: (N/16, M/16, 16, 16) + SmallVector newShape = {N / blk, M / 16, 16, blk}; + + return newShape; +} + +// 修改 VECTOR->CUBE 数据搬运函数 +void DAGSyncPass::insertVectorToCubeDataMovement( + mlir::Operation *srcOp, mlir::Operation *dstOp, Operation *posOp, + mlir::Value srcResult, mlir::OpBuilder &builder, mlir::Location loc, + llvm::DenseMap *valueMap) { + auto srcTensorType = getTensorType(srcResult); + if (!srcTensorType) { + return; + } + if (isa(srcOp) && isa(dstOp)) { + return; + } + + // 1. 在 srcOp 之后创建 UB 空间的 memref.alloc（用于 to_memref） + builder.setInsertionPointAfter(srcOp); + + // 首先创建 UB 空间的 memref type + auto ubSpaceAttr = + hivm::AddressSpaceAttr::get(builder.getContext(), hivm::AddressSpace::UB); + auto ubMemrefType = mlir::MemRefType::get(srcTensorType.getShape(), + srcTensorType.getElementType(), + /*layout=*/nullptr, ubSpaceAttr); + + // 创建 bufferization.to_memref + if (srcOp->getBlock() == dstOp->getBlock()) { + builder.setInsertionPoint(posOp); + } + auto toBufferOp = + builder.create(loc, ubMemrefType, srcResult); + + // 2. 创建 CBUF 空间的 memref.alloc（用于 copy 的目标） + mlir::Value cbufAllocOld = getOrCreateAllocation( + srcOp, srcTensorType, hivm::AddressSpace::L1, builder, loc); + auto cbufShape = *newCbubAllocShape( + dyn_cast(cbufAllocOld.getDefiningOp())); + // 获取旧的memref类型并创建新的类型 + auto oldType = dyn_cast(cbufAllocOld.getType()); + + // 获取新的维度数量 + unsigned newRank = cbufShape.size(); + + // 方法1：创建新的恒等布局映射 + AffineMap identityMap = builder.getMultiDimIdentityMap(newRank); + MemRefLayoutAttrInterface layout = AffineMapAttr::get(identityMap); + + // 方法2：如果旧类型有布局，尝试调整它（更安全的选择） + // 先检查旧类型是否有布局 + if (auto oldLayout = oldType.getLayout()) { + if (auto affineMapAttr = dyn_cast(oldLayout)) { + // 如果旧布局是AffineMap，尝试创建新的恒等映射 + // 因为维度改变，旧的affine map可能不再有效 + layout = AffineMapAttr::get(identityMap); + } else { + // 对于其他类型的布局，可能需要特殊处理 + layout = oldLayout; + } + } + + // 创建新的alloc类型 + auto newAllocType = MemRefType::get(cbufShape, oldType.getElementType(), + layout, // 使用新创建的布局 + oldType.getMemorySpace()); + + builder.setInsertionPoint(cbufAllocOld.getDefiningOp()); + // 创建新的alloc操作 + auto cbufAlloc = builder.create( + cbufAllocOld.getDefiningOp()->getLoc(), newAllocType); + + builder.setInsertionPointAfter(toBufferOp); + // 3. 创建 copy 指令（src 是 ub memref，dst 是 cbuf memref） + auto copyOp = + builder.create(loc, mlir::TypeRange{}, // 没有返回值 + toBufferOp.getResult(), // src (memref in UB) + cbufAlloc // dst (memref in CBUF) + ); + + // llvm::outs() << "Inserted copy after " << srcOp->getName().getStringRef() + // << " for VECTOR->CUBE data movement\n"; + + // 4. 在 dstOp 前创建 convert_layout + builder.setInsertionPoint(dstOp); + auto ndLayout = + hivm::DataLayoutAttr::get(builder.getContext(), hivm::DataLayout::ND); + // 创建 convert_layout + auto convertLayoutOp = builder.create( + loc, + cbufAllocOld.getType(), // 输出类型与输入相同 + cbufAlloc, + ndLayout, // srcLayout + ndLayout // dstLayout + ); + (*valueTypes)[convertLayoutOp.getResult()] = CoreType::CUBE_ONLY; + + // 5. 创建 memory_space_cast + auto cbufMemrefType = cast(convertLayoutOp.getType()); + auto plainMemrefType = mlir::MemRefType::get(cbufMemrefType.getShape(), + cbufMemrefType.getElementType()); + + auto memspaceCastOp = builder.create( + loc, plainMemrefType, convertLayoutOp.getResult()); + (*valueTypes)[memspaceCastOp.getResult()] = CoreType::CUBE_ONLY; + + // 6. 创建 to_tensor + auto toTensorOp = builder.create( + loc, + srcTensorType, // 原始的 tensor 类型 + memspaceCastOp.getResult(), + /*restrict=*/true, + /*writable=*/true); + (*valueTypes)[toTensorOp.getResult()] = CoreType::CUBE_ONLY; + + // 7. 替换 dstOp 的操作数 + replaceOperandWithNewValue(dstOp, srcResult, toTensorOp.getResult()); +} + +Operation *DAGSyncPass::FindLastestPosition(Operation *srcOp, Graph &mainGraph, + OpBuilder &builder) { + Operation *insertPos = nullptr; + auto opMap = mainGraph.getOpMapLegacy(); + auto valueTypes = &mainGraph.getValueTypes(); + // Find the first cube-dependent vector core operation. + for (auto nextOp = srcOp->getNextNode(); nextOp != nullptr; + nextOp = nextOp->getNextNode()) { + auto nextType = getNodeDeviceType(opMap[nextOp], valueTypes); + if (nextType == CoreType::CUBE_ONLY) + continue; + // No memref ops in IR yet; directly tracing operands + for (auto operand : nextOp->getOperands()) { + auto defOp = operand.getDefiningOp(); + auto defType = getNodeDeviceType(opMap[defOp], valueTypes); + if (defType == CoreType::CUBE_ONLY) { + // To prevent UB overflow, we need to break the dependency at the point + // where the result shape is minimized + // — i.e., trace upward to find the first broadcast. + for (auto prevOp = nextOp->getPrevNode(); + prevOp != nullptr && prevOp != srcOp; + prevOp = prevOp->getPrevNode()) { + if (isa(prevOp)) { + if (prevOp->getPrevNode() && + isa(prevOp->getPrevNode())) { + return prevOp->getPrevNode(); + } + return prevOp; + } + } + // Can't find the result shape is minimized + return nextOp; + } + } + + // Once meet SyncBlockWaitOp, return now! + if (auto waitOp = dyn_cast(nextOp)) { + if (waitOp.getTcoreType() == + hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR)) { + return nextOp; + } + } + insertPos = nextOp; + } + return insertPos; +} + +Operation *DAGSyncPass::FindEarliestPosition(Operation *dstOp, Graph &mainGraph, + OpBuilder &builder) { + auto insertPos = dstOp; + auto opMap = mainGraph.getOpMapLegacy(); + auto valueTypes = &mainGraph.getValueTypes(); + for (auto prevOp = dstOp->getPrevNode(); prevOp != nullptr; + prevOp = prevOp->getPrevNode()) { + if (dstOp->getBlock() != prevOp->getBlock()) + continue; + // Once meet SyncBlockSetOp, return now! + if (auto waitOp = dyn_cast(prevOp)) { + if (waitOp.getTcoreType() == + hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR)) { + return insertPos; + } + } + insertPos = prevOp; + } + return insertPos; +} + +// 主要的同步和数据搬运插入函数 +void DAGSyncPass::insertSyncAndMovement( + mlir::Operation *srcOp, mlir::Operation *dstOp, CoreType srcType, + CoreType dstType, mlir::OpBuilder &builder, int flag, + llvm::DenseMap *valueMap, Graph &mainGraph) { + mlir::Location loc = srcOp->getLoc(); + // 保存当前的插入点 + mlir::OpBuilder::InsertionGuard guard(builder); + + // 检查是否是跨 block + mlir::Block *srcBlock = srcOp->getBlock(); + mlir::Block *dstBlock = dstOp->getBlock(); + bool sameBlock = (srcBlock == dstBlock); + + if (!sameBlock) { + // 检查是否是外层到内层的依赖 + bool dstIsInnerBlock = false; + mlir::Operation *dstParentOp = dstBlock->getParentOp(); + while (dstParentOp) { + if (dstParentOp->getBlock() == srcBlock) { + dstIsInnerBlock = true; + break; + } + if (dstParentOp->getBlock()) { + dstParentOp = dstParentOp->getBlock()->getParentOp(); + } else { + break; + } + } + + if (dstIsInnerBlock) { + insertSyncAndMovementForCrossBlock(srcOp, dstOp, srcType, dstType, + builder, flag, true, valueMap, + mainGraph); + return; + } + } + + // 同一 block 内的处理 + // 获取 srcOp 的输出（假设第一个结果） + if (srcOp->getNumResults() == 0) { + return; + } + mlir::Value srcResult = srcOp->getResult(0); + + // CUBE -> VECTOR + if (srcType == CoreType::CUBE_ONLY && dstType == CoreType::VECTOR_ONLY) { + + // 2. 插入同步指令 + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_FIX); + auto waitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_V); + auto lastSetPipe = + PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE3); + auto lastWaitPipe = + PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE1); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + auto flagAddId = builder.getIntegerAttr(builder.getI64Type(), flag * 2); + auto lastFlagAddId = + builder.getIntegerAttr(builder.getI64Type(), (flag - 1) * 2); + + // set 在 srcOp 后 + builder.setInsertionPointAfter(srcOp); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // wait 在 dstOp 前 + + auto posOp = FindEarliestPosition(dstOp, mainGraph, builder); + builder.setInsertionPoint(posOp); + coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // 1. 插入数据搬运 + insertCubeToVectorDataMovement(srcOp, dstOp, srcResult, builder, loc, + nullptr); + + // llvm::outs() << "Inserted CUBE->VECTOR sync and data movement (flag=" << + // flag << ")\n"; + } + // VECTOR -> CUBE + else if (srcType == CoreType::VECTOR_ONLY && dstType == CoreType::CUBE_ONLY) { + + // 2. 插入同步指令 + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE3); + auto waitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE1); + auto lastSetPipe = + PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_FIX); + auto lastWaitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_V); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + auto flagAddId = builder.getIntegerAttr(builder.getI64Type(), flag * 2); + auto lastFlagAddId = + builder.getIntegerAttr(builder.getI64Type(), (flag - 1) * 2); + + // set 在 srcOp 后 + // builder.setInsertionPointAfter(srcOp); + auto posOp = FindLastestPosition(srcOp, mainGraph, builder); + if (posOp) { + builder.setInsertionPoint(posOp); + } else { + builder.setInsertionPointAfter(srcOp); + } + auto setOp = builder.create(loc, coreAttr, setPipe, + waitPipe, flagId); + + // wait 在 dstOp 前 + builder.setInsertionPoint(dstOp); + coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // 1. 插入数据搬运 + insertVectorToCubeDataMovement(srcOp, dstOp, setOp, srcResult, builder, loc, + valueMap); + + // llvm::outs() << "Inserted VECTOR->CUBE sync and data movement (flag=" << + // flag << ")\n"; + } +} + +// 跨 block 的同步和数据搬运 +void DAGSyncPass::insertSyncAndMovementForCrossBlock( + mlir::Operation *srcOp, mlir::Operation *dstOp, CoreType srcType, + CoreType dstType, mlir::OpBuilder &builder, int flag, bool dstIsInnerBlock, + llvm::DenseMap *valueMap, Graph &mainGraph) { + if (!dstIsInnerBlock) { + insertSyncAndMovement(srcOp, dstOp, srcType, dstType, builder, flag, + valueMap, mainGraph); + return; + } + + mlir::Location loc = srcOp->getLoc(); + mlir::Block *dstBlock = dstOp->getBlock(); + + // 获取 srcOp 的输出 + if (srcOp->getNumResults() == 0) { + return; + } + mlir::Value srcResult = srcOp->getResult(0); + + // CUBE -> VECTOR + if (srcType == CoreType::CUBE_ONLY && dstType == CoreType::VECTOR_ONLY) { + + // 2. 插入同步指令（跨 block 特殊处理） + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_FIX); + auto waitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_V); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + + // set 在 srcOp 后（外层） + builder.setInsertionPointAfter(srcOp); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // 1. 插入数据搬运（同 block 内逻辑） + insertCubeToVectorDataMovement(srcOp, dstOp, srcResult, builder, loc, + nullptr); + + // wait 在内层 block 入口前 + mlir::Operation *parentOp = dstBlock->getParentOp(); + while (srcOp->getBlock() != parentOp->getBlock()) { + parentOp = parentOp->getBlock()->getParentOp(); + } + if (parentOp) { + builder.setInsertionPoint(parentOp); + coreAttr = hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + } else { + builder.setInsertionPoint(dstOp); + coreAttr = hivm::TCoreTypeAttr::get(builder.getContext(), + hivm::TCoreType::VECTOR); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + } + + } + // VECTOR -> CUBE + else if (srcType == CoreType::VECTOR_ONLY && dstType == CoreType::CUBE_ONLY) { + + // 2. 插入同步指令（跨 block 特殊处理） + auto coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::VECTOR); + auto setPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE3); + auto waitPipe = PipeAttr::get(builder.getContext(), hivm::PIPE::PIPE_MTE1); + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + + // set 在 srcOp 后（外层） + builder.setInsertionPointAfter(srcOp); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + + // 1. 插入数据搬运（同 block 内逻辑） + insertVectorToCubeDataMovement(srcOp, dstOp, srcOp, srcResult, builder, loc, + valueMap); + + // wait 在内层 block 入口前 + mlir::Operation *parentOp = dstBlock->getParentOp(); + if (parentOp) { + builder.setInsertionPoint(parentOp); + coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + } else { + builder.setInsertionPoint(dstOp); + coreAttr = + hivm::TCoreTypeAttr::get(builder.getContext(), hivm::TCoreType::CUBE); + builder.create(loc, coreAttr, setPipe, waitPipe, flagId); + } + + // llvm::outs() << "Inserted cross-block VECTOR->CUBE sync and data movement + // (flag=" << flag << ")\n"; + } +} + +void LegalizeDot(triton::FuncOp funcOp) { + mlir::OpBuilder builder(funcOp); + funcOp.walk([&](triton::DotOp dotOp) { + // 获取dot操作的输入 + Value a = dotOp.getOperands()[0]; + Value b = dotOp.getOperands()[1]; + Value c = dotOp.getOperands()[2]; // 累加器参数 + + // 检查累加器是否为全零常量 + bool isZeroAccumulator = false; + + // 检查是否直接是arith.constant 0 + if (auto constantOp = c.getDefiningOp()) { + if (auto denseAttr = dyn_cast(constantOp.getValue())) { + if (denseAttr.isSplat() && + denseAttr.getSplatValue().getValueAsDouble() == 0.0) { + isZeroAccumulator = true; + } + } + } + + if (!isZeroAccumulator) { + // 创建新的零累加器 + Location loc = dotOp.getLoc(); + auto resultType = dotOp.getResult().getType(); + + Value originalResult = dotOp.getResult(); + builder.setInsertionPoint(dotOp); + // 创建全零张量 + auto zeroAttr = DenseElementsAttr::get( + dyn_cast(resultType), APFloat(0.0f)); + auto zeroConstant = builder.create(loc, zeroAttr); + + // 创建新的dot操作，使用零作为累加器 + auto newDot = + builder.create(loc, resultType, a, b, zeroConstant); + + // 创建加法操作，将新的dot结果与原来的累加器c相加 + auto addOp = builder.create(loc, newDot, c); + + // 用addOp替换原来的dotOp + originalResult.replaceAllUsesWith(addOp.getResult()); + + // 删除原dotOp（如果它没有其他用途） + if (dotOp.use_empty()) { + dotOp.erase(); + } + } + }); +} + +static void rewriteCopyChainForCbub(hivm::CopyOp copyOp, + ArrayRef newShape, + OpBuilder &builder) { + + // 获取 copy 的输入（ins），应为 to_memref 的结果 + Value insVal = copyOp.getOperands()[0]; + auto toBufferOp = insVal.getDefiningOp(); + if (!toBufferOp) + return; + + Value inputTensor = toBufferOp.getTensor(); + auto inputTensorType = dyn_cast(inputTensor.getType()); + if (!inputTensorType || inputTensorType.getRank() != 2) + return; + + // blk = 32/位宽 + // 中间 reshape 形状：[M/16, 16, N/ blk, blk] + int64_t M = inputTensorType.getShape()[0]; + int64_t N = inputTensorType.getShape()[1]; + auto elemType = inputTensorType.getElementType(); + auto blkOr = getBlockElemsFor32BAlign(elemType); + int64_t blk = (int64_t)*blkOr; + SmallVector intermediateShape3D = {M, N / blk, blk}; + SmallVector intermediateShapetrans = {N / blk, M, blk}; + auto elementType = inputTensorType.getElementType(); + auto interTensor3DType = + RankedTensorType::get(intermediateShape3D, elementType); + auto interTensortransType = + RankedTensorType::get(intermediateShapetrans, elementType); + + auto finalTensorType = RankedTensorType::get(newShape, elementType); + + auto loc = inputTensor.getLoc(); + + // Set insertion point before copyOp (or toBufferOp) + auto tensorOp = inputTensor.getDefiningOp(); + builder.setInsertionPointAfter(tensorOp); + + // 插入 triton.reshape 将 2D tensor 展开为 3D + auto reshape3DOp = + builder.create(loc, interTensor3DType, inputTensor); + (*valueTypes)[reshape3DOp.getResult()] = CoreType::VECTOR_ONLY; + + // nark tiling dim for reshapeop + auto markOp3d = builder.create(loc, reshape3DOp); + auto tilingDimAttr3d = builder.getDictionaryAttr(SmallVector{ + NamedAttribute(builder.getStringAttr("1"), builder.getIndexAttr(1))}); + markOp3d->setAttr("tiling_dim_mapping", tilingDimAttr3d); + + // 插入 triton.trans 调整维度顺序 Insert tt.trans {order = [1, 0, 2]} + SmallVector order = {1, 0, 2}; + auto orderAttr = + builder.getDenseI32ArrayAttr(order); // OpBuilder supports this + auto transOp = builder.create( + loc, interTensortransType, reshape3DOp.getResult(), orderAttr); + (*valueTypes)[transOp.getResult()] = CoreType::VECTOR_ONLY; + + // 插入 triton.reshape 将 3D tensor 展开为 4D + auto reshape4DOp = builder.create(loc, finalTensorType, + transOp.getResult()); + (*valueTypes)[reshape4DOp.getResult()] = CoreType::VECTOR_ONLY; + + // nark tiling dim for reshapeop + auto markOp4d = builder.create(loc, reshape4DOp); + auto tilingDimAttr4d = builder.getDictionaryAttr(SmallVector{ + NamedAttribute(builder.getStringAttr("1"), builder.getIndexAttr(1))}); + markOp4d->setAttr("tiling_dim_mapping", tilingDimAttr4d); + + // Create new to_memref + builder.setInsertionPoint(toBufferOp); + auto oldBufferType = cast(toBufferOp.getResult().getType()); + auto newMemRefType = MemRefType::get(newShape, elementType, mlir::AffineMap{}, + oldBufferType.getMemorySpace()); + auto newtoBufferOp = builder.create( + toBufferOp.getLoc(), newMemRefType, reshape4DOp.getResult()); + (*valueTypes)[newtoBufferOp.getResult()] = CoreType::VECTOR_ONLY; + + // Create NEW copyOp (replacing the old one) + builder.setInsertionPoint(copyOp); + auto resultTypes = copyOp->getResultTypes(); + auto newCopyOp = + builder.create(copyOp.getLoc(), + resultTypes, // TypeRange + reshape4DOp.getResult(), // src (ins) + copyOp.getOperands()[1] // dst (outs) + ); + + // 替换 uses 并清理旧 op + copyOp.replaceAllUsesWith(newCopyOp); + copyOp.erase(); + toBufferOp.erase(); + + return; +} + +bool valueIsPtrOrShapedPtr(Value value) { + auto type = value.getType(); + if (llvm::isa(type)) { + return true; + } + + if (auto tensorType = llvm::dyn_cast(type)) { + return llvm::isa(tensorType.getElementType()); + } + + return false; +} + +using AliasSet = llvm::SmallBitVector; +using AliasInfo = llvm::DenseMap; +constexpr size_t EXPECTED_MAX_ROOT_PTR_COUNT = 16; + +/** + * Reworked alias analysis inplace of triton::SharedMemoryAliasAnalysis or + * mlir::AliasAnalysis where both failed to establish aliases from function + * arguments (i.e. global memory): + * 1. triton::SharedMemoryAliasAnalysis - does not reach global memory + * 2. mlir::AliasAnalysis - marks all function arguments as MayAlias or + * PartialAlias; we may be able to get around by marking all arguments as + * noalias, but it may affect other analyses... + */ +AliasInfo getAlias(triton::FuncOp funcOp, Graph::ValueMapRaw &valueMap) { + AliasInfo aliasInfo; + + struct WorkNode { + Value value; + Value cause; + }; + + llvm::SmallVector worklist; + + auto addDownstreamToWorklist = [&worklist](ValueNode *node) { + if (node == nullptr) { + return; + } + for (auto *output : node->getOutputs()) { + if (auto *opNode = dyn_cast(output)) { + for (auto result : opNode->op->getResults()) { + if (valueIsPtrOrShapedPtr(result)) { + worklist.push_back({result, node->value}); + } + } + } else if (auto *valueNode = dyn_cast(output)) { + worklist.push_back({valueNode->value, node->value}); + } + } + }; + + SmallVector rootPtrs; + + for (auto arg : funcOp.getBody().getArguments()) { + if (valueIsPtrOrShapedPtr(arg)) { + rootPtrs.push_back(arg); + } + } + funcOp.walk([&](Operation *op) { + // If an op generates ptr from nowhere (e.g. tt.empty()), then we consider + // it as a root + bool generatesRootPtr = + !llvm::any_of(op->getOperands(), valueIsPtrOrShapedPtr); + if (generatesRootPtr) { + for (auto result : op->getResults()) { + if (valueIsPtrOrShapedPtr(result)) { + rootPtrs.push_back(result); + } + } + } + }); + + size_t numRootPtrs = rootPtrs.size(); + + for (auto [index, value] : llvm::enumerate(rootPtrs)) { + auto &bitset = aliasInfo[value]; + bitset.resize(numRootPtrs); + bitset[index] = true; + + addDownstreamToWorklist(getFromSmartPtr(valueMap, value)); + } + + while (!worklist.empty()) { + auto [value, upstream] = worklist.pop_back_val(); + if (!valueIsPtrOrShapedPtr(value)) { + continue; + } + auto node = getFromSmartPtr(valueMap, value); + if (!node) { + continue; + } + + /** + * Safety: + * 1. aliasInfo[value] : may insert into densemap and cause resize + * 2. aliasInfo[upstream] : steady reference without any insertion into the + * densemap, will not cause resize, since upstream is in the map + * 3. after aliasInfo[upstream] : currAlias is still valid + */ + auto &currAlias = aliasInfo[value]; + auto &causeAlias = aliasInfo[upstream]; + + if (!causeAlias.test(currAlias)) { + continue; + } + + currAlias |= causeAlias; + addDownstreamToWorklist(node); + } + return aliasInfo; +} + +/** + * @returns whether two values share any root ptr + */ +bool mayAlias(AliasInfo &aliasInfo, Value valueA, Value valueB) { + auto aliasesA = getPtr(aliasInfo, valueA); + auto aliasesB = getPtr(aliasInfo, valueB); + if (!valueIsPtrOrShapedPtr(valueA) || !valueIsPtrOrShapedPtr(valueB) || + aliasesA == nullptr || aliasesB == nullptr) { + return false; + } + + return aliasesA->anyCommon(*aliasesB); +} + +template +OpTy createBlockSync(OpBuilder builder, hivm::TCoreType coreType, + hivm::PIPE srcPipe, hivm::PIPE dstPipe, int flag, + Operation *cause) { + auto flagId = builder.getIntegerAttr(builder.getI64Type(), flag); + auto coreAttr = hivm::TCoreTypeAttr::get(builder.getContext(), coreType); + auto setPipe = PipeAttr::get(builder.getContext(), srcPipe); + auto waitPipe = PipeAttr::get(builder.getContext(), dstPipe); + return builder.create(cause->getLoc(), coreAttr, setPipe, waitPipe, + flagId); +} + +const size_t MAX_EXPECTED_PARENTS_COUNT = 8; +std::optional> +findAncestorCommonBlock(mlir::Operation *opA, mlir::Operation *opB) { + if (opA->getBlock() == opB->getBlock()) { + return std::make_pair(opA, opB); + } + + // record all ancestors of opA + llvm::SmallPtrSet ancestorsA; + mlir::Operation *curr = opA; + while (curr) { + ancestorsA.insert(curr); + curr = curr->getParentOp(); + } + + // find the last ancestor of opB which is also the ancestor of opA + mlir::Operation *commonAncOp = nullptr; + curr = opB; + while (curr) { + if (ancestorsA.count(curr)) { + commonAncOp = curr; + break; + } + curr = curr->getParentOp(); + } + + if (!commonAncOp) { + return std::nullopt; + } + + // find the ancestors in the given region + for (mlir::Region ®ion : commonAncOp->getRegions()) { + for (mlir::Block &block : region) { + auto *ancA = block.findAncestorOpInBlock(*opA); + auto *ancB = block.findAncestorOpInBlock(*opB); + if (ancA && ancB) { + return std::make_pair(ancA, ancB); + } + } + } + return std::nullopt; +} + +struct SyncCandidate { + CoreType srcCoreType; + Operation *setCause; + Operation *setAfter; + Operation *waitCause; + Operation *waitBefore; +}; + +// setOp, waitOp +std::pair +createBlockSyncBetween(OpBuilder builder, hivm::PIPE srcPipe, + hivm::PIPE dstPipe, SyncCandidate candidate, int flag) { + auto srcCoreType = toHivm(candidate.srcCoreType); + auto dstCoreType = toHivm(!candidate.srcCoreType); + + builder.setInsertionPointAfter(candidate.setAfter); + auto setOp = createBlockSync( + builder, srcCoreType, srcPipe, dstPipe, flag, candidate.setCause); + builder.setInsertionPoint(candidate.waitBefore); + auto waitOp = createBlockSync( + builder, dstCoreType, srcPipe, dstPipe, flag, candidate.waitCause); + return {setOp, waitOp}; +}; + +/** + * @returns the source/destination of a memory op + * + * Strangely, triton memory ops do not register their pointers to memory + * effects, so we need some special cases as fallback here + */ +Value getMemoryAddress(Operation *op, MemoryEffects::EffectInstance effect) { + if (mlir::Value v = effect.getValue()) + return v; + + return TypeSwitch(op) + .Case([](auto op) { return op.getPtr(); }) + .Default([](auto) { return nullptr; }); +} + +void addMemEffectsSync(triton::FuncOp funcOp, Graph *graph, OpBuilder &builder, + int &syncFlag) { + DominanceInfo domInfo(funcOp); + PostDominanceInfo postDomInfo(funcOp); + + // [(node, EffectInstance, LinearisationPt)] + llvm::SmallVector> memOps; + + // [(setAfter, waitBefore, srcOP, dstOp)][CoreType] + llvm::SmallVector candidates; + llvm::SmallVector backwardCandidates; + + auto aliasInfo = getAlias(funcOp, graph->getValueMap()); + + funcOp.walk([&](MemoryEffectOpInterface memIface) { + auto *op = memIface.getOperation(); + if (llvm::isa(op)) { + return; + } + + auto *currNode = graph->getOpMap()[op].get(); + SmallVector effects; + + memIface.getEffects(effects); + + for (auto &effect : effects) { + if (!isa(effect.getEffect())) { + continue; + } + auto currVal = getMemoryAddress(op, effect); + if (!currVal) { + op->emitWarning("Mem effect src/dst is unknown!"); + continue; + } + memOps.emplace_back(currNode, effect); + bool isWrite = isa(effect.getEffect()); + for (auto &[prevNode, prevEffect] : memOps) { + auto prevVal = getMemoryAddress(prevNode->op, prevEffect); + if ((isa(prevEffect.getEffect()) || isWrite) && + mayAlias(aliasInfo, prevVal, currVal) && + prevNode->isOn() != + currNode->isOn() // write is forced on single core type, so we + // are safe to judge based on whether the core + // types are different + ) { + CoreType srcCoreType = isWrite ? !currNode->isOn() : prevNode->isOn(); + CoreType dstCoreType = !srcCoreType; + auto opPair = findAncestorCommonBlock(prevNode->op, currNode->op); + if (!opPair.has_value()) { + op->emitWarning(llvm::formatv( + "Unable to find ancestors in common block with {0}\n", + *prevNode->op)); + continue; + } + auto [setAfter, waitBefore] = opPair.value(); + if (setAfter == waitBefore) { + continue; + } + candidates.push_back(SyncCandidate{srcCoreType, prevNode->op, + setAfter, op, waitBefore}); + + if (op->getParentOfType()) { // need to insert + // backward sync + backwardCandidates.push_back(SyncCandidate{ + dstCoreType, op, waitBefore, prevNode->op, setAfter}); + } + } + } + } + }); + + /** + * Adds block sync operations for the candidate + * Specially, for backward syncs, we also perform a kick-start for first wait + * and clean-up to avoid affecting possible future syncs + */ + auto addBlockSyncCommon = [&builder, &syncFlag, + &aliasInfo](SyncCandidate cand) { + LLVM_DEBUG(llvm::dbgs() << "\n\n=== Insert sync between ===\n" + << *cand.setAfter << "\n" + << "Cause: " << *cand.setCause << "\n"); + + LLVM_DEBUG(llvm::dbgs() << "----------" + << "\n" + << *cand.waitBefore << "\n" + << "Cause: " << *cand.waitCause << "\n"); + + LLVM_DEBUG(llvm::dbgs() << "=== Insert Sync End ===\n\n"); + + auto srcPipe = cand.srcCoreType == CoreType::CUBE_ONLY + ? hivm::PIPE::PIPE_FIX + : hivm::PIPE::PIPE_MTE2; + auto dstPipe = hivm::PIPE::PIPE_S; + auto [forwardSet, forwardWait] = createBlockSyncBetween( + builder, srcPipe, dstPipe, cand, syncFlag % MAX_FLAG_ID); + forwardSet->setAttr("ssbuf.flagid", builder.getUI32IntegerAttr(syncFlag)); + forwardWait->setAttr("ssbuf.flagid", builder.getUI32IntegerAttr(syncFlag)); + + auto loopAncestor = cand.waitBefore->getParentOfType(); + if (loopAncestor) { // backward sync, need to kick-start and clean-up; only + // need to ensure the first and last one, because we + // ensure once a flag is consumed, it will be set by + // construction (wait-set in the same block) + while (auto newAncestor = + loopAncestor->getParentOfType()) { + loopAncestor = newAncestor; + } + + auto srcPipe = cand.srcCoreType != CoreType::CUBE_ONLY + ? hivm::PIPE::PIPE_FIX + : hivm::PIPE::PIPE_MTE2; + auto backendFlag = syncFlag; + syncFlag++; + auto backwardCandidate = + SyncCandidate{!cand.srcCoreType, cand.waitCause, cand.waitBefore, + cand.setCause, cand.setAfter}; + auto [backwardSet, backwardWait] = createBlockSyncBetween( + builder, srcPipe, dstPipe, backwardCandidate, syncFlag % MAX_FLAG_ID); + + backwardSet->setAttr("ssbuf.backward", builder.getUnitAttr()); + backwardWait->setAttr("ssbuf.backward", builder.getUnitAttr()); + backwardSet->setAttr("ssbuf.flagid", + builder.getUI32IntegerAttr(backendFlag)); + backwardWait->setAttr("ssbuf.flagid", + builder.getUI32IntegerAttr(backendFlag)); + + builder.setInsertionPoint(loopAncestor); + auto kickstart = createBlockSync( + builder, toHivm(backwardCandidate.srcCoreType), srcPipe, dstPipe, + syncFlag, backwardCandidate.setCause); + + builder.setInsertionPointAfter(loopAncestor); + auto cleanup = createBlockSync( + builder, toHivm(!backwardCandidate.srcCoreType), srcPipe, dstPipe, + syncFlag, backwardCandidate.waitCause); + } + + syncFlag++; + }; + + if (candidates.empty()) { + return; + } + + auto setAfterDominate = [&domInfo](Operation *a, Operation *b) { + if (domInfo.dominates(a, b)) { + return true; + } + if (domInfo.dominates(b, a)) { + return false; + } + if (a->isAncestor(b)) { + return false; + } + if (b->isAncestor(a)) { + return true; + } + return false; + }; + + auto waitBeforePostDominate = [&postDomInfo](Operation *a, Operation *b) { + if (postDomInfo.postDominates(a, b)) { + return true; + } + if (postDomInfo.postDominates(b, a)) { + return false; + } + if (a->isAncestor(b)) { + return true; + } + if (b->isAncestor(a)) { + return false; + } + return false; + }; + + /** + * Sorts the candidates by (setAfter, waitBefore), and then removes redundant + * sync candidates that fully contains another + */ + auto selectCandidatesToInsertSync = + [&setAfterDominate, &waitBeforePostDominate, + &addBlockSyncCommon](llvm::SmallVectorImpl &candidates) { + llvm::sort( + candidates, [&](const SyncCandidate &a, const SyncCandidate &b) { + if (a.setAfter != b.setAfter) { + return setAfterDominate(a.setAfter, b.setAfter); + } + + if (a.waitBefore != b.waitBefore) { + return waitBeforePostDominate(a.waitBefore, b.waitBefore); + } + + return false; + }); + + for (auto [i, cand] : llvm::enumerate(candidates)) { + bool shouldInsert = true; + for (auto otherCand : ArrayRef(candidates).drop_front(i + 1)) { + bool duplicated = (cand.waitBefore == otherCand.waitBefore && + cand.setAfter == otherCand.setAfter && + cand.srcCoreType == otherCand.srcCoreType); + bool containsOther = + (cand.srcCoreType == otherCand.srcCoreType && + setAfterDominate(cand.setAfter, otherCand.setAfter) && + waitBeforePostDominate(cand.waitBefore, otherCand.waitBefore)); + if (duplicated || containsOther) { + shouldInsert = false; + break; + } + } + + if (shouldInsert) { + addBlockSyncCommon(cand); + } + } + }; + + selectCandidatesToInsertSync(candidates); +} + +void DAGSyncPass::runOnOperation() { + auto module = getOperation(); + mlir::OpBuilder builder(&getContext()); + + // 遍历所有函数 + for (auto funcOp : + llvm::make_early_inc_range(module.getOps())) { + // 跳过无效函数 + LegalizeDot(funcOp); + if (funcOp.getBody().empty()) { + continue; + } + + // llvm::outs() << "\n====================================\n"; + // llvm::outs() << "处理函数: " << funcOp.getName() << "\n"; + // llvm::outs() << "====================================\n"; + + auto unique_graph = Graph::fromMultiBlockFunc(funcOp); + std::shared_ptr shared_graph = std::move(unique_graph); + auto &main_graph = *shared_graph; + + auto funcName = funcOp.getName(); + + // 获取 DAG 图的映射 + auto opMapRaw = main_graph.getOpMapLegacy(); + valueTypes = &main_graph.getValueTypes(); + auto *opMap = &opMapRaw; + for (const auto &pair : *opMap) { + Operation *op = pair.first; // 键：Operation 指针 + Node *node = pair.second; // 值：Node 指针 + + // 打印指针地址（最直接的方式） + // llvm::outs() << "Operation*: " << *op + // << "\n"; + for (auto res : op->getResults()) { + llvm::outs() << "Value: " << (*valueTypes)[res] << "\n"; + } + } + + if (!opMap) { + llvm::errs() << "Warning: Failed to create DAG graph for function " + << funcOp.getName() << "\n"; + continue; + } + + // 用于避免重复插入同步 + llvm::DenseSet> + processedPairs; + int syncFlag = 1; + addMemEffectsSync(funcOp, shared_graph.get(), builder, syncFlag); + + // 3. 使用 walk 遍历函数中的所有操作 + funcOp.walk([&](mlir::Operation *op) { + // 查找当前操作对应的 Node + auto nodeIt = opMap->find(op); + if (nodeIt == opMap->end()) { + // 这个操作不在 entry block 的 DAG 图中 + // 可能是嵌套在控制流内部的操作 + return; + } + + OpNode *currentNode = nodeIt->second; + + // 检查是否是 scf.for 操作 + if (auto forOp = mlir::dyn_cast(op)) { + // 处理 scf.for 循环的特殊同步逻辑 + int temp = syncFlag % MAX_FLAG_ID; + processScfForSync(forOp, currentNode, valueTypes, builder, temp); + } + + // 获取当前节点的设备类型 + CoreType currentType = getNodeDeviceType(currentNode, valueTypes); + + // 打印操作信息（可选） + // if (!llvm::isa(op->getDialect())) { + // llvm::outs() << "操作: " << *op + // << " 设备类型: " + // << (currentType == CoreType::VECTOR_ONLY ? "VECTOR" : + // currentType == CoreType::CUBE_ONLY ? "CUBE" : + // "SCALAR") + // << "\n"; + // } + + // 4. 遍历当前节点的所有输入节点 + for (ValueNode *inputValNode : currentNode->getInputs()) { + auto inputOp = inputValNode->value.getDefiningOp(); + if (!inputOp && opMap->contains(inputOp)) { + continue; + } + + auto inputNode = (*opMap)[inputOp]; + + // 获取输入节点的设备类型 + CoreType inputType = getNodeDeviceType(inputNode, valueTypes); + + // 5. 判断是否需要插入同步和数据搬运 + if (needVectorCubeSync(inputType, currentType)) { + // 检查是否已经处理过这对操作 + auto opPair = std::make_pair(inputOp, op); + if (processedPairs.insert(opPair).second) { + // 插入同步和数据搬运指令 + // 检查是否是跨 block 的依赖 + mlir::Block *srcBlock = inputOp->getBlock(); + mlir::Block *dstBlock = op->getBlock(); + + if (srcBlock == dstBlock) { + // 同一 block 内 + insertSyncAndMovement(inputOp, op, inputType, currentType, + builder, syncFlag % 14, valueTypes, + main_graph); + syncFlag++; + } else { + // 跨 block，判断是否是外层到内层 + llvm::outs() << "#########\n"; + bool dstIsInnerBlock = false; + mlir::Operation *dstParentOp = dstBlock->getParentOp(); + + // 向上查找，看 dstBlock 是否在 srcBlock 的区域内 + while (dstParentOp) { + if (dstParentOp->getBlock() == srcBlock) { + dstIsInnerBlock = true; + break; + } + if (dstParentOp->getBlock()) { + dstParentOp = dstParentOp->getBlock()->getParentOp(); + } else { + break; + } + } + if (dstIsInnerBlock) { + + insertSyncAndMovementForCrossBlock( + inputOp, op, inputType, currentType, builder, syncFlag % 14, + dstIsInnerBlock, valueTypes, main_graph); + syncFlag++; + } + } + } + } + } + }); + + // llvm::outs() << "\n函数 " << funcOp.getName() << " 统计:\n"; + // llvm::outs() << " - 插入的总同步操作数: " << syncFlag << "\n"; + funcOp.walk([&](hivm::CopyOp copyOp) { + llvm::outs() << copyOp << " sss\n\n\n\n"; + rewriteCopyChainForCbub( + copyOp, + dyn_cast(copyOp.getOperands()[1].getType()).getShape(), + builder); + }); + GraphManager::getInstance().registerGraph(funcName, shared_graph); + } + + // llvm::outs()<> mlir::triton::createDAGSyncPass() { + return std::make_unique(); +} diff --git a/compiler/lib/TritonToAnnotation/CMakeLists.txt b/compiler/lib/TritonToAnnotation/CMakeLists.txt new file mode 100644 index 00000000..b8f47ced --- /dev/null +++ b/compiler/lib/TritonToAnnotation/CMakeLists.txt @@ -0,0 +1,15 @@ +add_triton_library(TritonToAnnotation + TritonToAnnotation.cpp + + DEPENDS + TritonToAnnotationConversionPassIncGen + + LINK_LIBS + BiShengIRAnnotationDialect + BiShengIRDialectUtils + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport + TritonIR +) diff --git a/compiler/lib/TritonToAnnotation/TritonToAnnotation.cpp b/compiler/lib/TritonToAnnotation/TritonToAnnotation.cpp new file mode 100644 index 00000000..b149d475 --- /dev/null +++ b/compiler/lib/TritonToAnnotation/TritonToAnnotation.cpp @@ -0,0 +1,56 @@ + + +#include "dicp/TritonToAnnotation/Passes.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_TRITONTOANNOTATION +#include "dicp/TritonToAnnotation/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; + +namespace { +struct TritonToAnnotationPass + : public mlir::triton::impl::TritonToAnnotationBase< + TritonToAnnotationPass> { + void runOnOperation() override; +}; +} // namespace + +struct TritonAnnotationConversionPattern + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(mlir::triton::dicp::AnnotationOp op, + PatternRewriter &rewriter) const final { + auto markOp = rewriter.create(op.getLoc(), op.getSrc()); + // Forward all annotations. + markOp->setAttrs(op->getAttrs()); + rewriter.eraseOp(op); + return success(); + } +}; + +void TritonToAnnotationPass::runOnOperation() { + auto module = getOperation(); + ConversionTarget target(getContext()); + target.addLegalDialect(); + + RewritePatternSet patterns(&getContext()); + patterns.add(patterns.getContext()); + if (failed(applyPartialConversion(module, target, std::move(patterns)))) { + signalPassFailure(); + } +} + +std::unique_ptr> +mlir::triton::createTritonToAnnotationPass() { + return std::make_unique(); +} diff --git a/compiler/lib/TritonToGraph/AliasAnalysis.cpp b/compiler/lib/TritonToGraph/AliasAnalysis.cpp new file mode 100644 index 00000000..1bf3efc7 --- /dev/null +++ b/compiler/lib/TritonToGraph/AliasAnalysis.cpp @@ -0,0 +1,228 @@ + + +#include "dicp/TritonToGraph/AliasAnalysis.h" +#include "dicp/TritonToGraph/ControlFlowGraph.h" +#include "dicp/TritonToGraph/MemorySSA.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "triton/Dialect/Triton/IR/OpInterfaces.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "alias-analysis" + +using namespace mlir; +using namespace triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// AliasAnalysis Implementation +//===----------------------------------------------------------------------===// + +void AliasAnalysis::analyzePointerAliases(ControlFlowGraph &cfg) { + LLVM_DEBUG(llvm::dbgs() << "=== Starting Pointer Alias Analysis ===\n"); + + // 步骤1: 识别全局内存参数 + triton::FuncOp func = cfg.getFunction(); + + LLVM_DEBUG(llvm::dbgs() << "Analyzing function: " << func.getName() << "\n"); + + for (BlockArgument arg : func.getArguments()) { + Type argType = arg.getType(); + + if (isPointerType(argType) && isGlobalMemoryType(argType)) { + // 为入参创建tensor对象 + std::string paramName = "param_" + std::to_string(arg.getArgNumber()); + + SmallVector shape; + Type elementType; + + // 使用辅助函数提取shape和element type + extractShapeAndElementType(argType, shape, elementType); + + TensorObject *tensor = + new TensorObject(paramName, shape, argType, elementType, + TensorObject::TensorKind::GLOBAL_MEMORY); + + // 记录alias关系 [param, param, tensor] + addAlias(arg, arg, tensor); + + LLVM_DEBUG(llvm::dbgs() + << " Found global memory parameter " << arg.getArgNumber() + << ": " << tensor->getName() << "\n"); + } + } + + LLVM_DEBUG(llvm::dbgs() << "Found " << aliasMap.size() + << " global memory parameters\n"); + + // 步骤2: 分析alias操作 + size_t aliasOpsFound = 0; + cfg.traverse([&](BasicBlock &bb) { + for (const auto &instPtr : bb.getInstructions()) { + const Instruction *inst = instPtr.get(); + Operation *op = inst->getOperation(); + + if (!op) + continue; + + if (auto addptrOp = dyn_cast(op)) { + analyzeAddPtrOp(addptrOp); + aliasOpsFound++; + } else if (auto makeTensorPtrOp = dyn_cast(op)) { + analyzeMakeTensorPtrOp(makeTensorPtrOp); + aliasOpsFound++; + } else if (auto loadOp = dyn_cast(op)) { + analyzeLoadOp(loadOp); + } else if (auto storeOp = dyn_cast(op)) { + analyzeStoreOp(storeOp); + } else if (auto broadcastOp = dyn_cast(op)) { + analyzeBroadcastOp(broadcastOp); + aliasOpsFound++; + } else if (auto splatOp = dyn_cast(op)) { + analyzeSplatOp(splatOp); + aliasOpsFound++; + } + } + }); + + LLVM_DEBUG( + llvm::dbgs() << "=== Alias Analysis Complete ===\n" + << "Total tracked aliases: " << aliasMap.size() << "\n" + << "Alias operations found: " << aliasOpsFound << "\n"); +} + +Value AliasAnalysis::getBasePointer(Value ptr) const { + // 递归查找base pointer + auto it = aliasMap.find(ptr); + if (it == aliasMap.end()) { + // 没有找到，返回自身（可能是原始base pointer） + return ptr; + } + + Value base = it->second; + + // 如果base不是自己，继续递归查找 + if (base != ptr) { + return getBasePointer(base); + } + + return base; +} + +void AliasAnalysis::analyzeAddPtrOp(mlir::triton::AddPtrOp addptrOp) { + // %new = tt.addptr %ptr, %offset + Value ptr = addptrOp.getPtr(); + Value result = addptrOp.getResult(); + + // result是ptr的alias，指向同一个tensor + Value basePtr = getBasePointer(ptr); + TensorObject *tensor = getTensorObject(basePtr); + + if (tensor) { + addAlias(result, basePtr, tensor); + + LLVM_DEBUG(llvm::dbgs() << " AddPtr: " << result << " -> " << basePtr + << " [" << tensor->getName() << "]\n"); + } +} + +void AliasAnalysis::analyzeMakeTensorPtrOp(mlir::triton::MakeTensorPtrOp op) { + // %tensor_ptr = tt.make_tensor_ptr %base, ... + Value basePtr = op.getBase(); + Value result = op.getResult(); + + TensorObject *tensor = getTensorObject(basePtr); + + if (tensor) { + addAlias(result, basePtr, tensor); + + LLVM_DEBUG(llvm::dbgs() << " MakeTensorPtr: " << result << " -> " + << basePtr << " [" << tensor->getName() << "]\n"); + } +} + +void AliasAnalysis::analyzeLoadOp(mlir::triton::LoadOp loadOp) { + // tt.load操作通常不改变alias关系， + // 但可以验证load的ptr是否被正确跟踪 + + Value ptr = loadOp.getPtr(); + + // 验证ptr是否被alias analysis跟踪 + TensorObject *tensor = getTensorObject(ptr); + + if (tensor) { + LLVM_DEBUG(llvm::dbgs() + << " Load from tracked tensor: " << tensor->getName() << "\n"); + } else { + LLVM_DEBUG(llvm::dbgs() << " Load from untracked pointer\n"); + } +} + +void AliasAnalysis::analyzeStoreOp(mlir::triton::StoreOp storeOp) { + // tt.store操作不改变alias关系， + // 但可以验证store的目标是否被正确跟踪 + + Value ptr = storeOp.getPtr(); + + // 验证ptr是否被alias analysis跟踪 + TensorObject *tensor = getTensorObject(ptr); + + if (tensor) { + LLVM_DEBUG(llvm::dbgs() + << " Store to tracked tensor: " << tensor->getName() << "\n"); + } else { + LLVM_DEBUG(llvm::dbgs() << " Store to untracked pointer\n"); + } +} + +void AliasAnalysis::analyzeBroadcastOp(mlir::triton::BroadcastOp broadcastOp) { + // %out = tt.broadcast %src : tensor> -> tensor> + // 输入和输出都是指针tensor，输出tensor的每个元素与输入tensor对应行/列元素是别名关系 + + Value src = broadcastOp.getSrc(); + Value result = broadcastOp.getResult(); + + // 只处理元素类型为指针的情况 + Type srcElemType = getElementTypeOrSelf(src.getType()); + if (!isPointerType(srcElemType)) + return; + + // 查找输入tensor关联的TensorObject + // 注意：src本身是指针tensor，需要获取其base pointer对应的tensor + Value srcBasePtr = getBasePointer(src); + TensorObject *tensor = getTensorObject(srcBasePtr); + + if (tensor) { + // broadcast操作保持指针值不变（只是复制到更多位置）， + // 因此result与src指向相同的底层tensor对象 + addAlias(result, srcBasePtr, tensor); + + LLVM_DEBUG(llvm::dbgs() << " Broadcast: " << result << " -> " << srcBasePtr + << " [" << tensor->getName() << "]\n"); + } +} + +void AliasAnalysis::analyzeSplatOp(mlir::triton::SplatOp splatOp) { + // %out = tt.splat %src : !tt.ptr -> tensor> + // 将标量指针广播到tensor的每个位置，所有输出元素都是输入指针的别名 + + Value src = splatOp.getSrc(); // 标量指针 + Value result = splatOp.getResult(); // 指针tensor + + // 确认输入是指针类型 + if (!isPointerType(src.getType())) + return; + + // 获取标量指针的base pointer和对应的TensorObject + Value basePtr = getBasePointer(src); + TensorObject *tensor = getTensorObject(basePtr); + + if (tensor) { + // splat操作将同一个指针值复制到tensor的每个元素， + // 因此result tensor中所有元素都与src是别名关系（指向同一内存对象） + addAlias(result, basePtr, tensor); + + LLVM_DEBUG(llvm::dbgs() << " Splat: " << result << " -> " << basePtr + << " [" << tensor->getName() << "]\n"); + } +} diff --git a/compiler/lib/TritonToGraph/CMakeLists.txt b/compiler/lib/TritonToGraph/CMakeLists.txt new file mode 100644 index 00000000..56e9387c --- /dev/null +++ b/compiler/lib/TritonToGraph/CMakeLists.txt @@ -0,0 +1,24 @@ +add_triton_library(TritonToGraph + AliasAnalysis.cpp + ControlFlowGraph.cpp + ControlFlowGraphBuilder.cpp + DataflowGraph.cpp + MemorySsaBuilder.cpp + Passes.cpp + + DEPENDS + TritonToGraphPassIncGen + + LINK_LIBS PUBLIC + MLIRArithDialect + MLIRControlFlowDialect + MLIRDialectUtils + MLIRIR + MLIRFuncDialect + MLIRPass + MLIRSCFDialect + MLIRSupport + MLIRTransforms + TritonIR + TritonTransforms +) diff --git a/compiler/lib/TritonToGraph/ControlFlowGraph.cpp b/compiler/lib/TritonToGraph/ControlFlowGraph.cpp new file mode 100644 index 00000000..7b54037f --- /dev/null +++ b/compiler/lib/TritonToGraph/ControlFlowGraph.cpp @@ -0,0 +1,1120 @@ + + +#include "dicp/TritonToGraph/ControlFlowGraph.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" + +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/Path.h" +#include "llvm/Support/raw_ostream.h" + +#include +#include +#include +#include + +using namespace mlir; +using namespace triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// Instruction +//===----------------------------------------------------------------------===// + +std::string Instruction::getAsString() const { + if (!operation) + return ""; + std::string str; + llvm::raw_string_ostream os(str); + operation->print(os); + return str; +} + +void Instruction::print(raw_ostream &os, unsigned indent) const { + std::string instStr = getAsString(); + os << "Inst[" << id << "]: "; + + if (operation) { + BlockType parentType = parentBlock->getType(); + if (parentType == BlockType::IF_COND || parentType == BlockType::FOR_COND || + parentType == BlockType::WHILE_COND || + parentType == BlockType::COND_BR || parentType == BlockType::BR) { + size_t newlinePos = instStr.find('\n'); + if (newlinePos != std::string::npos) { + instStr = instStr.substr(0, newlinePos) + " ..."; + } + } + os.indent(indent) << instStr; + } else { + os << "\n"; + } +} + +void Instruction::dump() const { print(llvm::errs()); } + +//===----------------------------------------------------------------------===// +// BasicBlock +//===----------------------------------------------------------------------===// + +void BasicBlock::addInstruction(std::unique_ptr inst) { + instructions.push_back(std::move(inst)); +} + +Instruction *BasicBlock::getInstruction(size_t idx) const { + if (idx < instructions.size()) { + return instructions[idx].get(); + } + return nullptr; +} + +void BasicBlock::addSuccessor(BasicBlock *succ) { + if (!succ) + return; + // 避免重复添加 + for (auto *s : successors) { + if (s == succ) + return; + } + successors.push_back(succ); + succ->addPredecessor(this); +} + +void BasicBlock::addPredecessor(BasicBlock *pred) { + if (!pred) + return; + // 避免重复添加 + for (auto *p : predecessors) { + if (p == pred) + return; + } + predecessors.push_back(pred); +} + +std::string BasicBlock::getName() const { + std::string name = "BB"; + name += std::to_string(id); + return name; +} + +bool BasicBlock::endsWithReturnOp() const { + // 空块检查 + if (instructions.empty()) { + return false; + } + + // 获取最后一条指令 + const Instruction *lastInst = instructions.back().get(); + if (!lastInst) { + return false; + } + + // 获取对应的 Operation + Operation *op = lastInst->getOperation(); + if (!op) { + return false; + } + + // 检查是否为 triton::ReturnOp + return isa(op); +} + +StringRef BasicBlock::getTypeString() const { + switch (type) { + case BlockType::NORMAL: + return "NORMAL"; + case BlockType::ENTRY: + return "ENTRY"; + case BlockType::EXIT: + return "EXIT"; + case BlockType::IF_COND: + return "IF_COND"; + case BlockType::FOR_COND: + return "FOR_COND"; + case BlockType::WHILE_COND: + return "WHILE_COND"; + case BlockType::COND_BR: + return "COND_BR"; + case BlockType::BR: + return "BR"; + case BlockType::LOOP_BODY: + return "LOOP_BODY"; + case BlockType::LOOP_EXIT: + return "LOOP_EXIT"; + } + return "UNKNOWN"; +} + +void BasicBlock::print(raw_ostream &os) const { + os << "============================================================\n"; + os << "BasicBlock " << getName() << " [" << getTypeString() << "]\n"; + if (parentStructure) { + os << " Parent Structure: " << parentStructure->getName() << "\n"; + } + + // 打印前驱 + if (!predecessors.empty()) { + os << " Predecessors: ["; + for (size_t i = 0; i < predecessors.size(); ++i) { + if (i > 0) + os << ", "; + os << predecessors[i]->getName(); + } + os << "]\n"; + } + + // 打印后继 + if (!successors.empty()) { + os << " Successors: ["; + for (size_t i = 0; i < successors.size(); ++i) { + if (i > 0) + os << ", "; + os << successors[i]->getName(); + } + os << "]\n"; + } + + // 打印指令 - COND节点只打印第一行 + os << " Instructions (" << instructions.size() << "):\n"; + for (const auto &inst : instructions) { + std::string instStr = inst->getAsString(); + + // 对COND节点截断：只保留第一行 + if (type == BlockType::IF_COND || type == BlockType::FOR_COND || + type == BlockType::WHILE_COND) { + size_t newlinePos = instStr.find('\n'); + if (newlinePos != std::string::npos) { + instStr = instStr.substr(0, newlinePos) + " ..."; + } + } + + os.indent(4) << "Inst[" << inst->getId() << "]: " << instStr << "\n"; + } + os << "\n"; +} + +void BasicBlock::dump() const { print(llvm::errs()); } + +void BasicBlock::exportToJSON(raw_ostream &os, unsigned indent) const { + std::string ind(indent, ' '); + os << ind << "{\n"; + os << ind << " \"id\": " << id << ",\n"; + os << ind << " \"name\": \"" << getName() << "\",\n"; + os << ind << " \"type\": \"" << getTypeString() << "\",\n"; + os << ind << " \"parentStructure\": " + << (parentStructure ? std::to_string(parentStructure->getId()) : "null") + << ",\n"; + + // 前驱 + os << ind << " \"predecessors\": ["; + for (size_t i = 0; i < predecessors.size(); ++i) { + if (i > 0) + os << ", "; + os << predecessors[i]->getId(); + } + os << "],\n"; + + // 后继 + os << ind << " \"successors\": ["; + for (size_t i = 0; i < successors.size(); ++i) { + if (i > 0) + os << ", "; + os << successors[i]->getId(); + } + os << "],\n"; + + // 指令 + os << ind << " \"instructions\": [\n"; + for (size_t i = 0; i < instructions.size(); ++i) { + os << ind << " {\n"; + os << ind << " \"id\": " << instructions[i]->getId() << ",\n"; + // 转义字符串用于 JSON + std::string instStr = instructions[i]->getAsString(); + + // COND节点只取第一行 + if (type == BlockType::IF_COND || type == BlockType::FOR_COND || + type == BlockType::WHILE_COND) { + size_t newlinePos = instStr.find('\n'); + if (newlinePos != std::string::npos) { + instStr = instStr.substr(0, newlinePos) + " ..."; + } + } + + // 简单的 JSON 字符串转义 + std::string escaped; + for (char c : instStr) { + if (c == '"') + escaped += "\\\""; + else if (c == '\\') + escaped += "\\\\"; + else if (c == '\n') + escaped += "\\n"; + else if (c == '\r') + escaped += "\\r"; + else if (c == '\t') + escaped += "\\t"; + else if ((unsigned char)c < 0x20) { + char buf[8]; + snprintf(buf, sizeof(buf), "\\u%04x", c); + escaped += buf; + } else { + escaped += c; + } + } + os << ind << " \"operation\": \"" << escaped << "\"\n"; + os << ind << " }"; + if (i < instructions.size() - 1) + os << ","; + os << "\n"; + } + os << ind << " ]\n"; + os << ind << "}"; +} + +//===----------------------------------------------------------------------===// +// ControlFlowGraph +//===----------------------------------------------------------------------===// + +ControlFlowGraph::ControlFlowGraph(triton::FuncOp func) : function(func) {} + +ControlFlowGraph::~ControlFlowGraph() = default; + +BasicBlock *ControlFlowGraph::createBasicBlock(BlockType type, + BasicBlock *parentStructure) { + auto bb = std::make_unique(nextBlockId++, type, parentStructure); + BasicBlock *bbPtr = bb.get(); + basicBlocks.push_back(std::move(bb)); + + // 自动设置入口/出口块 + if (type == BlockType::ENTRY) + entryBlock = bbPtr; + if (type == BlockType::EXIT) + exitBlock = bbPtr; + + return bbPtr; +} + +void ControlFlowGraph::addEdge(BasicBlock *from, BasicBlock *to) { + if (!from || !to) + return; + from->addSuccessor(to); +} + +// 判断是否为回边（从循环体回到循环头） +bool ControlFlowGraph::isBackEdge(BasicBlock *from, BasicBlock *to) const { + // 回边定义：指向 FOR_COND/WHILE_COND 且 from 是该循环的后代 + if (to->getType() != BlockType::FOR_COND && + to->getType() != BlockType::WHILE_COND) + return false; + + // 检查 from 是否属于以 to 为头的循环 + // 方法：检查 from 的 parentStructure 链是否包含 to + BasicBlock *current = from; + while (current) { + if (current == to) + return true; + current = current->getParentStructure(); + } + return false; +} + +void ControlFlowGraph::searchNormalBlock(BasicBlock *block, + OperationVisitor callback) const { + for (const auto &inst : block->getInstructions()) { + if (Operation *op = inst->getOperation()) { + callback(op); + } + } +} + +void ControlFlowGraph::searchCondBlock(BasicBlock *block, + OperationVisitor callback) const { + if (!block) + return; + + // 获取该 Cond block 对应的 exit block（停止条件） + BasicBlock *exitBlock = block->getExitBlock(); + + // 遍历 Cond block 的每一个 successor + for (BasicBlock *succ : block->getSuccessors()) { + if (!succ) + continue; + + // 从每个 successor 开始，沿着 block 链条向下遍历 + // 使用队列进行 BFS 遍历该分支 + std::deque workList; + workList.push_back(succ); + + while (!workList.empty()) { + BasicBlock *current = workList.front(); + workList.pop_front(); + + // 如果碰到 exit block，停止该分支的遍历 + if (current == exitBlock) + break; + + // 对当前 block 调用 searchBlock + searchBlock(current, callback); + + if (current->getType() == BlockType::NORMAL) { + // 将后继加入队列（排除回边） + for (BasicBlock *next : current->getSuccessors()) { + // 跳过回边避免无限循环 + if (isBackEdge(current, next)) + continue; + workList.push_back(next); + } + } else if (block->getType() == BlockType::IF_COND || + block->getType() == BlockType::FOR_COND || + block->getType() == BlockType::WHILE_COND) { + BasicBlock *next = block->getExitBlock(); + if (isBackEdge(current, next)) + continue; + workList.push_back(next); + } + } + } +} + +void ControlFlowGraph::searchBlock(BasicBlock *block, + OperationVisitor callback) const { + if (block->getType() == BlockType::NORMAL) { + searchNormalBlock(block, callback); + } else if (block->getType() == BlockType::IF_COND || + block->getType() == BlockType::FOR_COND || + block->getType() == BlockType::WHILE_COND) { + searchCondBlock(block, callback); + } +} + +void ControlFlowGraph::traverse(BlockVisitor visitor) { + // 从入口块开始进行拓扑排序 + std::vector topoOrder; + DenseSet visited; + + // 使用DFS进行拓扑排序 + std::function dfs = [&](BasicBlock *bb) { + if (!bb || visited.contains(bb)) + return; + + visited.insert(bb); + + // 递归访问所有后继块 + for (BasicBlock *succ : bb->getSuccessors()) { + dfs(succ); + } + + // 在后序位置添加，最后需要反转 + topoOrder.push_back(bb); + }; + + // 从入口块开始DFS + dfs(entryBlock); + + // 反转得到拓扑序 + std::reverse(topoOrder.begin(), topoOrder.end()); + + // 按照拓扑序遍历 + for (BasicBlock *bb : topoOrder) { + visitor(*bb); + } +} + +void ControlFlowGraph::print(raw_ostream &os) const { + auto funcName = const_cast(function).getName(); + os << "=================================================================\n"; + os << "Control Flow Graph for function '" + << (funcName.empty() ? "unnamed" : funcName) << "'\n"; + os << "Number of blocks: " << basicBlocks.size() << "\n"; + os << "=================================================================\n\n"; + + for (const auto &bb : basicBlocks) { + bb->print(os); + } +} + +void ControlFlowGraph::dump() const { print(llvm::errs()); } + +void ControlFlowGraph::exportDOT(raw_ostream &os) const { + auto funcName = const_cast(function).getName(); + std::string funcNameStr = funcName.empty() ? "unnamed" : funcName.str(); + + // 清理函数名用于 DOT 标识符 + std::string cleanFuncName; + for (char c : funcNameStr) { + if (isalnum(c) || c == '_') + cleanFuncName += c; + else + cleanFuncName += '_'; + } + + os << "digraph CFG_" << cleanFuncName << " {\n"; + os << " label=\"CFG for " << funcNameStr << "\";\n"; + os << " labelloc=t;\n"; + os << " rankdir=TB;\n"; + os << " splines=true;\n"; // 使用曲线边 + os << " overlap=false;\n"; // 防止节点重叠 + os << " nodesep=0.6;\n"; // 节点水平间距 + os << " ranksep=1.2;\n"; // 层间距 + os << " fontsize=12;\n\n"; + + // 设置节点样式 + for (const auto &bb : basicBlocks) { + os << " \"" << bb->getName() << "\" ["; + os << "label=\"" << bb->getName() << "\\n(" << bb->getTypeString() << ")"; + + // 指令摘要：COND节点只取第一行，其他节点最多3条 + size_t numInsts = bb->getNumInstructions(); + if (numInsts > 0) { + os << "\\n"; + + // COND节点只显示第一条指令的第一行 + if (bb->getType() == BlockType::IF_COND || + bb->getType() == BlockType::FOR_COND || + bb->getType() == BlockType::WHILE_COND) { + std::string instStr = bb->getInstruction(0)->getAsString(); + size_t newlinePos = instStr.find('\n'); + if (newlinePos != std::string::npos) { + instStr = instStr.substr(0, newlinePos); + } + if (instStr.length() > 40) + instStr = instStr.substr(0, 37) + "..."; + + // 转义 + std::string escaped; + for (char c : instStr) { + if (c == '"') + escaped += "\\\""; + else if (c == '\\') + escaped += "\\\\"; + else if (c == '\n') + escaped += "\\n"; + else + escaped += c; + } + os << escaped; + if (numInsts > 1) + os << "\\n... (" << (numInsts - 1) << " more)"; + } else { + // 其他节点显示最多3条 + for (size_t i = 0; i < std::min(numInsts, (size_t)3); ++i) { + std::string instStr = bb->getInstruction(i)->getAsString(); + // 只取第一行 + size_t newlinePos = instStr.find('\n'); + if (newlinePos != std::string::npos) { + instStr = instStr.substr(0, newlinePos) + "..."; + } + if (instStr.length() > 40) + instStr = instStr.substr(0, 37) + "..."; + + std::string escaped; + for (char c : instStr) { + if (c == '"') + escaped += "\\\""; + else if (c == '\\') + escaped += "\\\\"; + else if (c == '\n') + escaped += "\\n"; + else + escaped += c; + } + os << escaped << "\\n"; + } + if (numInsts > 3) + os << "... (" << (numInsts - 3) << " more)\\n"; + } + } + + os << "\", "; + + // 形状和颜色 + switch (bb->getType()) { + case BlockType::ENTRY: + os << "style=filled, fillcolor=lightgreen, shape=ellipse"; + break; + case BlockType::EXIT: + os << "style=filled, fillcolor=lightcoral, shape=ellipse"; + break; + case BlockType::IF_COND: + os << "style=filled, fillcolor=lightyellow, shape=diamond"; + break; + case BlockType::FOR_COND: + case BlockType::WHILE_COND: + os << "style=filled, fillcolor=lightblue, shape=diamond"; + break; + case BlockType::LOOP_BODY: + os << "style=filled, fillcolor=lightcyan, shape=box"; + break; + case BlockType::LOOP_EXIT: + os << "style=filled, fillcolor=lightpink, shape=box"; + break; + default: + os << "shape=box"; + break; + } + os << ", fontsize=10];\n"; + } + + os << "\n"; + + // 输出边：回边红色，其他绿色 + for (const auto &bb : basicBlocks) { + for (auto *succ : bb->getSuccessors()) { + bool isBack = isBackEdge(bb.get(), succ); + os << " \"" << bb->getName() << "\" -> \"" << succ->getName() << "\""; + os << " [color=" << (isBack ? "red" : "green"); + os << ", penwidth=" << (isBack ? "2.5" : "1.5") << ""; + if (isBack) { + os << ", style=dashed"; // 回边用虚线 + } + os << "];\n"; + } + } + + os << "}\n"; +} + +llvm::Error ControlFlowGraph::exportToFile(StringRef filename) const { + std::error_code ec; + llvm::raw_fd_ostream os(filename, ec, llvm::sys::fs::OF_Text); + + if (ec) { + return llvm::createStringError(llvm::inconvertibleErrorCode(), + "Failed to open file: " + filename); + } + + // 判断文件扩展名 + if (filename.ends_with(".dot")) { + exportDOT(os); + } else if (filename.ends_with(".json")) { + exportToJSON(os); + } else { + print(os); + } + + os.close(); + return llvm::Error::success(); +} + +void ControlFlowGraph::exportToJSON(raw_ostream &os) const { + auto funcName = const_cast(function).getName(); + std::string funcNameStr = funcName.empty() ? "unnamed" : funcName.str(); + + os << "{\n"; + os << " \"functionName\": \"" << funcNameStr << "\",\n"; + os << " \"numBlocks\": " << basicBlocks.size() << ",\n"; + os << " \"blocks\": [\n"; + + for (size_t i = 0; i < basicBlocks.size(); ++i) { + basicBlocks[i]->exportToJSON(os, 4); + if (i < basicBlocks.size() - 1) + os << ","; + os << "\n"; + } + + // 添加边信息（带类型标记） + os << " ],\n"; + os << " \"edges\": [\n"; + + bool first = true; + for (const auto &bb : basicBlocks) { + for (auto *succ : bb->getSuccessors()) { + if (!first) + os << ",\n"; + first = false; + + bool isBack = isBackEdge(bb.get(), succ); + os << " {\"from\": " << bb->getId() << ", \"to\": " << succ->getId() + << ", \"fromName\": \"" << bb->getName() << "\"" + << ", \"toName\": \"" << succ->getName() << "\"" + << ", \"type\": \"" << (isBack ? "back" : "normal") << "\"}"; + } + } + os << "\n ]\n"; + os << "}\n"; +} + +llvm::Error ControlFlowGraph::exportToHTML(StringRef filename) const { + std::error_code ec; + llvm::raw_fd_ostream os(filename, ec, llvm::sys::fs::OF_Text); + + if (ec) { + return llvm::createStringError(llvm::inconvertibleErrorCode(), + "Failed to open file: " + filename); + } + + auto funcName = const_cast(function).getName(); + std::string funcNameStr = funcName.empty() ? "unnamed" : funcName.str(); + + os << R"html( + + + + CFG for )html" + << funcNameStr << R"html( + + + + + + +
+ +
+
+
正常控制流 (向下)
+
循环回边 (向上)
+
+ + + + +)html"; + + os.close(); + return llvm::Error::success(); +} \ No newline at end of file diff --git a/compiler/lib/TritonToGraph/ControlFlowGraphBuilder.cpp b/compiler/lib/TritonToGraph/ControlFlowGraphBuilder.cpp new file mode 100644 index 00000000..09de959e --- /dev/null +++ b/compiler/lib/TritonToGraph/ControlFlowGraphBuilder.cpp @@ -0,0 +1,883 @@ + + +#include "dicp/TritonToGraph/ControlFlowGraphBuilder.h" +#include "dicp/TritonToGraph/DataflowGraph.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/Support/Debug.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/Path.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "build-cfg" + +using namespace mlir; +using namespace mlir::triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// BuildCFGPass Implementation +//===----------------------------------------------------------------------===// + +void BuildCFGPass::runOnOperation() { + auto module = getOperation(); + llvm::errs() << "Building CFG for module\n"; + + // 获取输出目录（从命令行选项，默认为当前目录） + // TableGen 生成的选项成员变量名为 outputDir + std::string outputDir = this->outputDir; + + // 创建输出目录 + llvm::SmallString<128> outputPath(outputDir); + llvm::sys::fs::create_directories(outputPath); + + llvm::errs() << "CFG output directory: " << outputPath << "\n"; + + // 遍历模块中的所有 Triton 函数 (tt.func) + for (triton::FuncOp func : module.getOps()) { + llvm::errs() << "Processing function: " << func.getName() << "\n"; + + auto cfg = buildForFunction(func); + if (!cfg) { + func.emitError() << "Failed to build CFG for function"; + signalPassFailure(); + return; + } + + // 打印 CFG 到标准输出 + // cfg->print(llvm::outs()); + + // 导出到文件 + std::string baseName = func.getName().str(); + + // 导出文本格式 + llvm::SmallString<128> textPath(outputPath); + llvm::sys::path::append(textPath, baseName + "_cfg.txt"); + if (auto err = cfg->exportToFile(textPath)) { + llvm::errs() << "Failed to export CFG to " << textPath << "\n"; + } else { + llvm::errs() << "Exported CFG to " << textPath << "\n"; + } + + // 导出 DOT 格式 + llvm::SmallString<128> dotPath(outputPath); + llvm::sys::path::append(dotPath, baseName + "_cfg.dot"); + if (auto err = cfg->exportToFile(dotPath)) { + llvm::errs() << "Failed to export CFG to " << dotPath << "\n"; + } else { + llvm::errs() << "Exported CFG to " << dotPath << "\n"; + } + + // 导出 JSON 格式 + llvm::SmallString<128> jsonPath(outputPath); + llvm::sys::path::append(jsonPath, baseName + "_cfg.json"); + if (auto err = cfg->exportToFile(jsonPath)) { + llvm::errs() << "Failed to export CFG to " << jsonPath << "\n"; + } else { + llvm::errs() << "Exported CFG to " << jsonPath << "\n"; + } + + // 导出 HTML 格式（网页可视化） + llvm::SmallString<128> htmlPath(outputPath); + llvm::sys::path::append(htmlPath, baseName + "_cfg.html"); + if (auto err = cfg->exportToHTML(htmlPath)) { + llvm::errs() << "Failed to export CFG to " << htmlPath << "\n"; + } else { + llvm::errs() << "Exported CFG to " << htmlPath << "\n"; + } + + // 构建 DataFlowGraph（包含 Memory SSA 分析） + llvm::errs() << " Building DataFlowGraph with Memory SSA...\n"; + DataFlowGraph dataFlowGraph(*cfg); + dataFlowGraph.build(); + + // 导出 DataFlowGraph + std::error_code ec; + llvm::SmallString<128> dataflowPath(outputPath); + llvm::sys::path::append(dataflowPath, baseName + "_dataflow.json"); + llvm::raw_fd_ostream dfOs(dataflowPath, ec); + if (!ec) { + dataFlowGraph.exportToJSON(dfOs); + llvm::errs() << " Exported DataFlowGraph to " << dataflowPath << "\n"; + } + } +} + +std::unique_ptr +BuildCFGPass::buildForFunction(triton::FuncOp func) { + ControlFlowGraphBuilder cfgBuilder; + return cfgBuilder.build(func); +} + +ControlFlowGraphBuilder::RegionBlocks ControlFlowGraphBuilder::buildForRegion( + Region ®ion, cfg::ControlFlowGraph &cfg, cfg::BasicBlock *entryBlock, + cfg::BasicBlock *parentStructure) { + cfg::BasicBlock *currentBB = entryBlock; + cfg::BasicBlock *lastBlock = entryBlock; + + // 首先为 region 中的所有 block 创建对应的 BasicBlock 映射 + // 这样可以确保在处理 cf.cond_br 等跳转指令时目标块已存在 + for (Block &block : region) { + if (!blockToBasicBlockMap.count(&block)) { + auto *bb = cfg.createBasicBlock(cfg::BlockType::NORMAL, parentStructure); + registerBlockMapping(&block, bb); + } + } + + // 遍历 region 中的所有 block + for (Block &block : region) { + // 获取该 block 对应的 BasicBlock + cfg::BasicBlock *blockBB = blockToBasicBlockMap[&block]; + + // 如果是第一个 block，合并到 entryBlock + if (blockBB == blockToBasicBlockMap.lookup(®ion.front())) { + // 将 entryBlock 的指令移动到 blockBB（或者反过来） + // 这里简化处理：使用 entryBlock 继续处理 + currentBB = processBlock(block, cfg, currentBB, parentStructure); + } else { + // 确保从上一个 block 的结尾连接到这个 block + if (lastBlock && lastBlock != blockBB) { + // 检查是否已经有边连接 + bool hasEdge = false; + for (auto *succ : lastBlock->getSuccessors()) { + if (succ == blockBB) { + hasEdge = true; + break; + } + } + if (!hasEdge && !lastBlock->endsWithReturnOp()) { + cfg.addEdge(lastBlock, blockBB); + } + } + currentBB = processBlock(block, cfg, blockBB, parentStructure); + } + + if (currentBB) { + lastBlock = currentBB; + if (lastBlock->endsWithReturnOp()) { + cfg.addEdge(lastBlock, cfg.getExitBlock()); + } + } + } + + return {entryBlock, lastBlock}; +} + +cfg::BasicBlock * +ControlFlowGraphBuilder::processBlock(Block &block, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *currentBB, + cfg::BasicBlock *parentStructure) { + if (!currentBB) + return nullptr; + + // 遍历 block 中的所有操作 + for (Operation &op : block) { + // 检查是否是控制流操作 + if (isa(op)) { + // 为 if 条件创建单独的 basic block + auto *ifCondBB = + cfg.createBasicBlock(BlockType::IF_COND, parentStructure); + + // 将当前 if 指令添加到 ifCondBB + createInstruction(&op, ifCondBB, cfg); + + // 连接当前块到 if 条件块 + cfg.addEdge(currentBB, ifCondBB); + + // 处理 if 操作，返回 if 后面的块 + currentBB = + handleIfOp(cast(op), cfg, ifCondBB, parentStructure); + } else if (isa(op)) { + // 为 for 条件创建单独的 basic block + auto *forCondBB = + cfg.createBasicBlock(BlockType::FOR_COND, parentStructure); + + // 将当前 for 指令添加到 forCondBB + createInstruction(&op, forCondBB, cfg); + + // 连接当前块到 for 条件块 + cfg.addEdge(currentBB, forCondBB); + + // 处理 for 操作，返回 for 后面的块 + currentBB = + handleForOp(cast(op), cfg, forCondBB, parentStructure); + } else if (isa(op)) { + // 为 while 条件创建单独的 basic block + auto *whileCondBB = + cfg.createBasicBlock(BlockType::WHILE_COND, parentStructure); + + // 将当前 while 指令添加到 whileCondBB + createInstruction(&op, whileCondBB, cfg); + + // 连接当前块到 while 条件块 + cfg.addEdge(currentBB, whileCondBB); + + // 处理 while 操作，返回 while 后面的块 + currentBB = handleWhileOp(cast(op), cfg, whileCondBB, + parentStructure); + } else if (isa(op)) { + // yield 操作：创建指令并继续（后续由循环处理逻辑连接） + createInstruction(&op, currentBB, cfg); + } else if (isa(op)) { + // condition 操作（while 循环条件）：创建指令并继续 + createInstruction(&op, currentBB, cfg); + } else if (isa(op)) { + // cf.cond_br 条件分支 - 创建专门的 COND_BR 块并处理 + auto *condBrBB = + cfg.createBasicBlock(BlockType::COND_BR, parentStructure); + + // 将 cond_br 指令添加到 condBrBB + createInstruction(&op, condBrBB, cfg); + + // 连接当前块到 COND_BR 块 + cfg.addEdge(currentBB, condBrBB); + + // 处理 cond_br 操作，返回后续的基本块 + currentBB = handleCondBranchOp(cast(op), cfg, condBrBB, + parentStructure); + } else if (isa(op)) { + // cf.br 无条件跳转 - 创建专门的 BR 块并处理 + auto *brBB = cfg.createBasicBlock(BlockType::BR, parentStructure); + + // 将 br 指令添加到 brBB + createInstruction(&op, brBB, cfg); + + // 连接当前块到 BR 块 + cfg.addEdge(currentBB, brBB); + + // 处理 br 操作 + currentBB = + handleBranchOp(cast(op), cfg, brBB, parentStructure); + } else if (isa(op)) { + // return 操作 + createInstruction(&op, currentBB, cfg); + } else if (op.getNumRegions() > 0) { + // 有内部区域的 Triton 操作 (如 tt.reduce, tt.scan 等) + // 先创建指令 + auto *inst = createInstruction(&op, currentBB, cfg); + + // 为该操作创建子图 + auto subGraph = + std::make_unique(cfg.getFunction()); + auto *subEntry = subGraph->createBasicBlock(cfg::BlockType::ENTRY); + subGraph->setEntryBlock(subEntry); + auto *subExit = subGraph->createBasicBlock(cfg::BlockType::EXIT); + subGraph->setExitBlock(subExit); + + // 用于子图中生成唯一指令 ID + size_t subInstId = 0; + + // 遍历所有区域 + for (size_t regionIdx = 0; regionIdx < op.getNumRegions(); ++regionIdx) { + Region ®ion = op.getRegion(regionIdx); + if (!region.empty()) { + // 为每个区域构建 CFG + cfg::BasicBlock *regionEntryBB = nullptr; + cfg::BasicBlock *regionLastBB = nullptr; + + for (Block ®ionBlock : region) { + auto *bb = subGraph->createBasicBlock(cfg::BlockType::NORMAL); + if (!regionEntryBB) + regionEntryBB = bb; + + // 将区域中的操作添加到子图的基本块 + for (Operation ®ionOp : regionBlock) { + auto regionInst = std::make_unique( + subInstId++, ®ionOp, bb); + cfg::Instruction *instPtr = regionInst.get(); + bb->addInstruction(std::move(regionInst)); + + // 添加到op到instruction的映射 + subGraph->addOpToInstruction(®ionOp, instPtr); + } + + // 连接基本块 + if (regionLastBB) { + subGraph->addEdge(regionLastBB, bb); + } + regionLastBB = bb; + } + + // 连接区域入口到子图入口 + if (regionEntryBB) { + subGraph->addEdge(subEntry, regionEntryBB); + } + // 连接区域出口到子图出口 + if (regionLastBB) { + subGraph->addEdge(regionLastBB, subExit); + } + } + } + + // 设置子图 + inst->setSubGraph(std::move(subGraph)); + } else { + // 普通操作，直接添加到当前 basic block + createInstruction(&op, currentBB, cfg); + } + } + + return currentBB; +} + +cfg::BasicBlock * +ControlFlowGraphBuilder::handleIfOp(scf::IfOp ifOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *ifCondBB, + cfg::BasicBlock *parentStructure) { + // 创建 if 后面的汇合块 + auto *mergeBB = cfg.createBasicBlock(BlockType::NORMAL, parentStructure); + + // 设置 ifCondBB 的出口块为 mergeBB + ifCondBB->setExitBlock(mergeBB); + + // 处理 then 分支 + cfg::BasicBlock *thenExitBB = nullptr; + if (!ifOp.getThenRegion().empty()) { + // 创建 then 区域的入口块 + auto *thenEntryBB = cfg.createBasicBlock(BlockType::NORMAL, ifCondBB); + cfg.addEdge(ifCondBB, thenEntryBB); + + // 构建 then 区域的 CFG + auto result = + buildForRegion(ifOp.getThenRegion(), cfg, thenEntryBB, ifCondBB); + thenExitBB = result.exitBlock; + } + + // 处理 else 分支 + cfg::BasicBlock *elseExitBB = nullptr; + bool hasElse = !ifOp.getElseRegion().empty(); + + if (hasElse) { + // 创建 else 区域的入口块 + auto *elseEntryBB = cfg.createBasicBlock(BlockType::NORMAL, ifCondBB); + cfg.addEdge(ifCondBB, elseEntryBB); + + // 构建 else 区域的 CFG + auto result = + buildForRegion(ifOp.getElseRegion(), cfg, elseEntryBB, ifCondBB); + elseExitBB = result.exitBlock; + } + + // 连接 then 分支到汇合块 + if (thenExitBB) { + cfg.addEdge(thenExitBB, mergeBB); + } else { + // 空的 then 分支，直接从 ifCondBB 连接到 mergeBB + cfg.addEdge(ifCondBB, mergeBB); + } + + // 连接 else 分支到汇合块 + if (hasElse) { + if (elseExitBB) { + cfg.addEdge(elseExitBB, mergeBB); + } else { + // 空的 else 分支，直接从 ifCondBB 连接到 mergeBB + cfg.addEdge(ifCondBB, mergeBB); + } + } else { + // 没有 else 分支，ifCondBB 直接连接到 mergeBB（else 路径） + cfg.addEdge(ifCondBB, mergeBB); + } + + return mergeBB; +} + +cfg::BasicBlock *ControlFlowGraphBuilder::handleForOp( + scf::ForOp forOp, cfg::ControlFlowGraph &cfg, cfg::BasicBlock *forCondBB, + cfg::BasicBlock *parentStructure) { + // 创建循环体入口块 + auto *loopBodyEntryBB = cfg.createBasicBlock(BlockType::LOOP_BODY, forCondBB); + cfg.addEdge(forCondBB, loopBodyEntryBB); + + // 创建循环出口块 + auto *loopExitBB = + cfg.createBasicBlock(BlockType::LOOP_EXIT, parentStructure); + + // 设置 forCondBB 的出口块为 loopExitBB + forCondBB->setExitBlock(loopExitBB); + + // 构建循环体的 CFG + auto result = + buildForRegion(forOp.getRegion(), cfg, loopBodyEntryBB, forCondBB); + + // 循环体结束需要回到循环头（通过 yield 操作） + if (result.exitBlock) { + cfg.addEdge(result.exitBlock, forCondBB); + } + + // 循环出口 + cfg.addEdge(forCondBB, loopExitBB); + + return loopExitBB; +} + +cfg::BasicBlock *ControlFlowGraphBuilder::handleWhileOp( + scf::WhileOp whileOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *whileCondBB, cfg::BasicBlock *parentStructure) { + // while 操作有两个区域：before（条件）和 after（循环体） + // 控制流： + // whileCondBB (包含 scf.while 指令) + // ↓ + // beforeEntryBB (条件计算区域) + // ↓ + // scf.condition 分支：真 → afterEntryBB, 假 → loopExitBB + // ↓ + // afterEntryBB (循环体) + // ↓ + // scf.yield + // ↓ + // whileCondBB (回到 while 头，重新进入 before) + + // 创建 before 区域的入口块（条件计算） + auto *beforeEntryBB = cfg.createBasicBlock(BlockType::LOOP_BODY, whileCondBB); + cfg.addEdge(whileCondBB, beforeEntryBB); + + // 创建 after 区域的入口块（循环体） + auto *afterEntryBB = cfg.createBasicBlock(BlockType::LOOP_BODY, whileCondBB); + + // 创建循环出口块 + auto *loopExitBB = + cfg.createBasicBlock(BlockType::LOOP_EXIT, parentStructure); + + // 设置 whileCondBB 的出口块为 loopExitBB + whileCondBB->setExitBlock(loopExitBB); + + // 构建 before 区域的 CFG（条件计算区域） + // before 区域以一个 scf.condition 操作结束 + auto beforeResult = + buildForRegion(whileOp.getBefore(), cfg, beforeEntryBB, whileCondBB); + + // 构建 after 区域的 CFG（循环体区域） + // after 区域以 scf.yield 结束，yield 的参数会传递给 before 区域的参数 + auto afterResult = + buildForRegion(whileOp.getAfter(), cfg, afterEntryBB, whileCondBB); + + // 处理 before 区域结束后的分支 + // before 区域应该以一个 scf.condition 操作结束 + // 该操作决定是进入 after 区域还是退出循环 + if (beforeResult.exitBlock) { + // 从 before 出口连接到 after 入口（条件为真时） + cfg.addEdge(beforeResult.exitBlock, afterEntryBB); + + // 从 before 出口连接到循环出口（条件为假时） + // 注意：在实际的 scf.condition 中，条件为假会直接退出循环 + cfg.addEdge(beforeResult.exitBlock, loopExitBB); + } + + // after 区域结束后回到 whileCondBB（重新进入 before 区域进行条件检查） + if (afterResult.exitBlock) { + cfg.addEdge(afterResult.exitBlock, whileCondBB); + } + + return loopExitBB; +} + +cfg::Instruction *ControlFlowGraphBuilder::createInstruction( + Operation *op, cfg::BasicBlock *parentBlock, cfg::ControlFlowGraph &cfg) { + if (!op || !parentBlock) + return nullptr; + + auto inst = std::make_unique(getNextInstructionId(), op, + parentBlock); + cfg::Instruction *instPtr = inst.get(); + parentBlock->addInstruction(std::move(inst)); + + // 将 Operation 到 Instruction 的映射添加到 CFG + cfg.addOpToInstruction(op, instPtr); + + return instPtr; +} + +//===----------------------------------------------------------------------===// +// ControlFlowGraphBuilder Implementation +//===----------------------------------------------------------------------===// + +std::unique_ptr +ControlFlowGraphBuilder::build(triton::FuncOp func) { + auto cfg = std::make_unique(func); + + if (func.getBody().empty()) { + // 空函数，只创建入口和出口 + auto *entry = cfg->createBasicBlock(BlockType::ENTRY); + auto *exit = cfg->createBasicBlock(BlockType::EXIT); + cfg->addEdge(entry, exit); + return cfg; + } + + // 创建入口块 + auto *entryBlock = cfg->createBasicBlock(BlockType::ENTRY); + cfg->setEntryBlock(entryBlock); + + // 创建出口块 + auto *exitBlock = cfg->createBasicBlock(BlockType::EXIT); + cfg->setExitBlock(exitBlock); + + // 为函数体构建 CFG + auto result = buildForRegion(func.getBody(), *cfg, entryBlock, nullptr); + + // 连接函数体出口到函数出口 + if (result.exitBlock) { + cfg->addEdge(result.exitBlock, exitBlock); + } + + return cfg; +} + +std::vector> +ControlFlowGraphBuilder::buildForModule(ModuleOp module) { + std::vector> cfgs; + + for (triton::FuncOp func : module.getOps()) { + auto cfg = build(func); + if (cfg) { + cfgs.push_back(std::move(cfg)); + } + } + + return cfgs; +} + +cfg::BasicBlock *ControlFlowGraphBuilder::handleCondBranchOp( + cf::CondBranchOp condBrOp, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *condBrBB, cfg::BasicBlock *parentStructure) { + // 创建汇合块（用于 cond_br 之后的代码） + auto *mergeBB = cfg.createBasicBlock(BlockType::NORMAL, parentStructure); + + // 设置 condBrBB 的出口块为 mergeBB + condBrBB->setExitBlock(mergeBB); + + // 获取条件值 + Value condition = condBrOp.getCondition(); + + // 获取 true 分支的目标块和参数 + Block *trueDest = condBrOp.getTrueDest(); + SmallVector trueOperands(condBrOp.getTrueDestOperands()); + + // 获取 false 分支的目标块和参数 + Block *falseDest = condBrOp.getFalseDest(); + SmallVector falseOperands(condBrOp.getFalseDestOperands()); + + LLVM_DEBUG(llvm::dbgs() << " CondBr: condition=" << condition << "\n"); + LLVM_DEBUG(llvm::dbgs() << " True dest: " << trueDest << "\n"); + LLVM_DEBUG(llvm::dbgs() << " False dest: " << falseDest << "\n"); + + // 为 true 分支创建入口块（如果目标块还没有对应的 BasicBlock） + cfg::BasicBlock *trueEntryBB = + getOrCreateBasicBlockForBlock(trueDest, cfg, parentStructure); + + // 为 false 分支创建入口块 + cfg::BasicBlock *falseEntryBB = + getOrCreateBasicBlockForBlock(falseDest, cfg, parentStructure); + + // 连接 COND_BR 块到两个分支 + cfg.addEdge(condBrBB, trueEntryBB); + cfg.addEdge(condBrBB, falseEntryBB); + + // 存储分支信息到指令的 MemorySSAInfo 中（用于后续查询） + if (condBrBB->getNumInstructions() > 0) { + cfg::Instruction *inst = condBrBB->getInstruction(0); + // 可以通过 inst->getMemorySSAInfo() 存储额外信息 + } + + // 返回汇合块，后续代码将在此块中继续 + return mergeBB; +} + +cfg::BasicBlock *ControlFlowGraphBuilder::handleBranchOp( + cf::BranchOp brOp, cfg::ControlFlowGraph &cfg, cfg::BasicBlock *brBB, + cfg::BasicBlock *parentStructure) { + // 无条件跳转没有汇合块，直接连接到目标块 + + // 获取目标块和参数 + Block *dest = brOp.getDest(); + SmallVector destOperands(brOp.getDestOperands()); + + LLVM_DEBUG(llvm::dbgs() << " Br: unconditional branch\n"); + LLVM_DEBUG(llvm::dbgs() << " Dest: " << dest << "\n"); + + // 获取或创建目标块对应的 BasicBlock + cfg::BasicBlock *destBB = + getOrCreateBasicBlockForBlock(dest, cfg, parentStructure); + + // 连接 BR 块到目标块 + cfg.addEdge(brBB, destBB); + + // 无条件跳转没有后续代码，返回 nullptr 表示当前路径结束 + return nullptr; +} + +cfg::BasicBlock *ControlFlowGraphBuilder::getOrCreateBasicBlockForBlock( + Block *block, cfg::ControlFlowGraph &cfg, + cfg::BasicBlock *parentStructure) { + // 检查是否已经有对应的 BasicBlock + auto it = blockToBasicBlockMap.find(block); + if (it != blockToBasicBlockMap.end()) { + return it->second; + } + + // 创建新的 BasicBlock + auto *bb = cfg.createBasicBlock(BlockType::NORMAL, parentStructure); + + // 注册映射关系 + registerBlockMapping(block, bb); + + return bb; +} + +void ControlFlowGraphBuilder::registerBlockMapping(Block *mlirBlock, + cfg::BasicBlock *cfgBlock) { + blockToBasicBlockMap[mlirBlock] = cfgBlock; +} + +SmallVector +ControlFlowGraphBuilder::collectCondBrBlocks(cfg::ControlFlowGraph &cfg) { + SmallVector condBrBlocks; + + // 遍历 CFG 中的所有基本块 + for (size_t i = 0; i < cfg.getNumBlocks(); ++i) { + cfg::BasicBlock *bb = cfg.getBasicBlock(i); + if (bb && bb->getType() == BlockType::COND_BR) { + condBrBlocks.push_back(bb); + } + } + + return condBrBlocks; +} + +std::optional +ControlFlowGraphBuilder::getCondBranchMapping(cfg::BasicBlock *condBrBB) { + // 验证输入基本块类型 + if (!condBrBB || condBrBB->getType() != BlockType::COND_BR) { + return std::nullopt; + } + + // 获取 COND_BR 块中的指令（应该包含 cf.cond_br 操作） + if (condBrBB->getNumInstructions() == 0) { + return std::nullopt; + } + + cfg::Instruction *inst = condBrBB->getInstruction(0); + Operation *op = inst->getOperation(); + + // 确保是 cf.cond_br 操作 + auto condBrOp = dyn_cast(op); + if (!condBrOp) { + return std::nullopt; + } + + CondBranchMapping mapping; + + // 收集条件值 + mapping.condition = condBrOp.getCondition(); + + // 收集 true 分支信息 + mapping.trueDest = condBrOp.getTrueDest(); + for (Value operand : condBrOp.getTrueDestOperands()) { + mapping.trueOperands.push_back(operand); + } + + // 收集 false 分支信息 + mapping.falseDest = condBrOp.getFalseDest(); + for (Value operand : condBrOp.getFalseDestOperands()) { + mapping.falseOperands.push_back(operand); + } + + return mapping; +} + +SmallVector +ControlFlowGraphBuilder::collectBrBlocks(cfg::ControlFlowGraph &cfg) { + SmallVector brBlocks; + + // 遍历 CFG 中的所有基本块 + for (size_t i = 0; i < cfg.getNumBlocks(); ++i) { + cfg::BasicBlock *bb = cfg.getBasicBlock(i); + if (bb && bb->getType() == BlockType::BR) { + brBlocks.push_back(bb); + } + } + + return brBlocks; +} + +std::optional +ControlFlowGraphBuilder::getBranchMapping(cfg::BasicBlock *brBB) { + // 验证输入基本块类型 + if (!brBB || brBB->getType() != BlockType::BR) { + return std::nullopt; + } + + // 获取 BR 块中的指令（应该包含 cf.br 操作） + if (brBB->getNumInstructions() == 0) { + return std::nullopt; + } + + cfg::Instruction *inst = brBB->getInstruction(0); + Operation *op = inst->getOperation(); + + // 确保是 cf.br 操作 + auto brOp = dyn_cast(op); + if (!brOp) { + return std::nullopt; + } + + BranchMapping mapping; + + // 收集目标块信息 + mapping.dest = brOp.getDest(); + for (Value operand : brOp.getDestOperands()) { + mapping.destOperands.push_back(operand); + } + + return mapping; +} + +std::unique_ptr> +mlir::triton::cfg::createBuildCFGPass() { + return std::unique_ptr>(new BuildCFGPass()); +} + +//===----------------------------------------------------------------------===// +// ControlFlowGraphBuilderWithQueries Implementation +//===----------------------------------------------------------------------===// + +SmallVector +ControlFlowGraphBuilder::collectIfCondBlocks(cfg::ControlFlowGraph &cfg) { + SmallVector ifCondBlocks; + + // 遍历 CFG 中的所有基本块 + for (size_t i = 0; i < cfg.getNumBlocks(); ++i) { + cfg::BasicBlock *bb = cfg.getBasicBlock(i); + if (bb && bb->getType() == BlockType::IF_COND) { + ifCondBlocks.push_back(bb); + } + } + + return ifCondBlocks; +} + +SmallVector +ControlFlowGraphBuilder::collectForCondBlocks(cfg::ControlFlowGraph &cfg) { + SmallVector forCondBlocks; + + // 遍历 CFG 中的所有基本块 + for (size_t i = 0; i < cfg.getNumBlocks(); ++i) { + cfg::BasicBlock *bb = cfg.getBasicBlock(i); + if (bb && bb->getType() == BlockType::FOR_COND) { + forCondBlocks.push_back(bb); + } + } + + return forCondBlocks; +} + +std::optional +ControlFlowGraphBuilder::getIfYieldResultMapping(cfg::BasicBlock *ifCondBB) { + // 验证输入基本块类型 + if (!ifCondBB || ifCondBB->getType() != BlockType::IF_COND) { + return std::nullopt; + } + + // 获取 IF_COND 块中的指令（应该包含 scf.if 操作） + if (ifCondBB->getNumInstructions() == 0) { + return std::nullopt; + } + + cfg::Instruction *inst = ifCondBB->getInstruction(0); + Operation *op = inst->getOperation(); + + // 确保是 scf.if 操作 + auto ifOp = dyn_cast(op); + if (!ifOp) { + return std::nullopt; + } + + IfYieldResultMapping mapping; + + // 收集 result values + for (Value result : ifOp.getResults()) { + mapping.resultValues.push_back(result); + } + + // 处理 then 分支的 yield values + if (!ifOp.getThenRegion().empty()) { + Block &thenBlock = ifOp.getThenRegion().back(); + // 查找 then 区域末尾的 scf.yield 操作 + for (Operation &thenOp : thenBlock) { + if (auto yieldOp = dyn_cast(thenOp)) { + for (Value operand : yieldOp.getOperands()) { + mapping.thenYieldValues.push_back(operand); + } + break; + } + } + } + + // 处理 else 分支的 yield values（如果有） + if (!ifOp.getElseRegion().empty()) { + Block &elseBlock = ifOp.getElseRegion().back(); + // 查找 else 区域末尾的 scf.yield 操作 + for (Operation &elseOp : elseBlock) { + if (auto yieldOp = dyn_cast(elseOp)) { + for (Value operand : yieldOp.getOperands()) { + mapping.elseYieldValues.push_back(operand); + } + break; + } + } + } + + return mapping; +} + +std::optional +ControlFlowGraphBuilder::getForYieldIterArgMapping(cfg::BasicBlock *forCondBB) { + // 验证输入基本块类型 + if (!forCondBB || forCondBB->getType() != BlockType::FOR_COND) { + return std::nullopt; + } + + // 获取 FOR_COND 块中的指令（应该包含 scf.for 操作） + if (forCondBB->getNumInstructions() == 0) { + return std::nullopt; + } + + cfg::Instruction *inst = forCondBB->getInstruction(0); + Operation *op = inst->getOperation(); + + // 确保是 scf.for 操作 + auto forOp = dyn_cast(op); + if (!forOp) { + return std::nullopt; + } + + ForYieldIterArgMapping mapping; + + // 收集 iter_args（循环初始参数） + for (Value iterArg : forOp.getRegionIterArgs()) { + mapping.iterArgValues.push_back(iterArg); + } + + // 收集 result values + for (Value result : forOp.getResults()) { + mapping.resultValues.push_back(result); + } + + // 处理循环体的 yield values + if (!forOp.getRegion().empty()) { + Block &loopBlock = forOp.getRegion().back(); + // 查找循环体末尾的 scf.yield 操作 + for (Operation &loopOp : loopBlock) { + if (auto yieldOp = dyn_cast(loopOp)) { + for (Value operand : yieldOp.getOperands()) { + mapping.yieldValues.push_back(operand); + } + break; + } + } + } + + return mapping; +} diff --git a/compiler/lib/TritonToGraph/DataflowGraph.cpp b/compiler/lib/TritonToGraph/DataflowGraph.cpp new file mode 100644 index 00000000..76f88c5a --- /dev/null +++ b/compiler/lib/TritonToGraph/DataflowGraph.cpp @@ -0,0 +1,294 @@ + + +#include "dicp/TritonToGraph/DataflowGraph.h" +#include "dicp/TritonToGraph/AliasAnalysis.h" +#include "dicp/TritonToGraph/ControlFlowGraph.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "dataflow-graph" + +using namespace mlir; +using namespace triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// DataFlowInfo +//===----------------------------------------------------------------------===// + +// Memory SSA相关接口 +MemorySSADef *DataFlowInfo::getMemoryDefinition(Value value) const { + auto it = memoryDefinitions.find(value); + return (it != memoryDefinitions.end()) ? it->second : nullptr; +} + +void DataFlowInfo::addMemoryDefinition(Value value, MemorySSADef *def) { + memoryDefinitions[value] = def; + // 确保没有旧的uses冲突 + memoryUses.erase(value); +} + +SmallVector DataFlowInfo::getMemoryUses(Value value) const { + auto it = memoryUses.find(value); + if (it != memoryUses.end()) { + return it->second; + } + return SmallVector(); +} + +void DataFlowInfo::addMemoryUse(Value value, const MemorySSAUse &use) { + memoryUses[value].push_back(use); +} + +void DataFlowInfo::removeMemoryDefinition(Value value) { + memoryDefinitions.erase(value); + memoryUses.erase(value); + invalidateDefUseCache(); +} + +void DataFlowInfo::clearMemoryUses(Value value) { + memoryUses[value].clear(); + invalidateDefUseCache(); +} + +std::unique_ptr DataFlowInfo::queryDataFlow(Value value) const { + // 1. 优先查询Memory SSA + if (MemorySSADef *def = getMemoryDefinition(value)) { + auto result = std::make_unique(def->getDefOp(), def); + result->getUses() = getSSAUses(value); + return result; + } + + // 2. 查询传统SSA + if (Operation *defOp = value.getDefiningOp()) { + auto result = std::make_unique(defOp, defOp); + result->getUses() = getSSAUses(value); + return result; + } + + // 3. 入参 + auto result = std::make_unique(nullptr, nullptr); + result->getUses() = getSSAUses(value); + return result; +} + +SmallVector DataFlowInfo::getUses(MemorySSADef *def) const { + SmallVector result; + + // 遍历所有uses，查找使用该definition的 + for (const auto &entry : memoryUses) { + for (const MemorySSAUse &use : entry.second) { + if (use.getDefinition() == def) { + result.push_back(use); + } + } + } + + return result; +} + +SmallVector +DataFlowInfo::getUsesByUserOp(Operation *userOp) const { + SmallVector result; + + for (const auto &entry : memoryUses) { + for (const MemorySSAUse &use : entry.second) { + if (use.getUserOp() == userOp) { + result.push_back(use); + } + } + } + + return result; +} + +void DataFlowInfo::buildDefUseCache() const { + if (defUseCacheValid) + return; + + defUseCache.clear(); + + for (const auto &entry : memoryUses) { + for (const MemorySSAUse &use : entry.second) { + MemorySSADef *def = use.getDefinition(); + if (def) { + defUseCache[def].push_back(use); + } + } + } + + defUseCacheValid = true; +} + +void DataFlowInfo::forEachDefinition( + llvm::function_ref func) const { + for (const auto &entry : memoryDefinitions) { + func(entry.first, entry.second); + } +} + +void DataFlowInfo::forEachUse( + llvm::function_ref func) const { + for (const auto &entry : memoryUses) { + for (const MemorySSAUse &use : entry.second) { + func(use); + } + } +} + +void DataFlowInfo::print(llvm::raw_ostream &os) const { + os << "=== Data Flow Information ===" + << "\n"; + + os << "Memory Definitions: " << memoryDefinitions.size() << "\n"; + for (const auto &entry : memoryDefinitions) { + os << " " << entry.first << " -> "; + entry.second->print(os); + os << "\n"; + } + + os << "Memory Uses: " + << "\n"; + for (const auto &entry : memoryUses) { + os << " " << entry.first << ": "; + for (const MemorySSAUse &use : entry.second) { + os << "[" << use.getDefinition()->getId() << "] "; + } + os << "\n"; + } + + os << "Phis: " << Phis.size() << "\n"; + for (const auto &entry : Phis) { + os << " " << entry.first << ": "; + switch (entry.second.type) { + case PhiInfo::ITER_ARG: + os << "ITER_ARG"; + break; + case PhiInfo::IF_RESULT: + os << "IF_RESULT"; + break; + case PhiInfo::WHILE_ARG: + os << "WHILE_ARG"; + break; + } + os << "\n"; + } +} + +void DataFlowInfo::exportToJSON(llvm::raw_ostream &os) const { + os << "{\n"; + os << " \"memoryDefinitions\": {\n"; + bool first = true; + for (const auto &entry : memoryDefinitions) { + if (!first) + os << ",\n"; + first = false; + os << " \"" << entry.first << "\": {\n"; + os << " \"id\": \"" << entry.second->getId() << "\",\n"; + os << " \"tensor\": \"" << entry.second->getTensor()->getName() + << "\",\n"; + os << " \"version\": " << entry.second->getVersion(); + os << "\n }"; + } + os << "\n },\n"; + + os << " \"Phis\": {\n"; + first = true; + for (const auto &entry : Phis) { + if (!first) + os << ",\n"; + first = false; + os << " \"" << entry.first << "\": {\n"; + os << " \"type\": " << entry.second.type << "\n"; + os << " }"; + } + os << "\n }\n"; + os << "}\n"; +} + +//===----------------------------------------------------------------------===// +// DataFlowGraph +//===----------------------------------------------------------------------===// + +void DataFlowGraph::build() { + LLVM_DEBUG(llvm::dbgs() << "=== Starting Data Flow Graph Build ===" + << "\n"); + + // 步骤1: 构建Alias分析 + aliasAnalysis = std::make_unique(); + aliasAnalysis->analyzePointerAliases(cfg); + // aliasAnalysis->print(llvm::outs()); + + LLVM_DEBUG(llvm::dbgs() << "Alias analysis complete" + << "\n"); + + // 步骤2: 构建Memory SSA + memorySSABuilder = + std::make_unique(cfg, *aliasAnalysis, dataFlowInfo); + memorySSABuilder->build(); + + LLVM_DEBUG(llvm::dbgs() << "Memory SSA build complete" + << "\n"); + + // 步骤3: 构建def-use图 + buildDefUseGraph(); + + LLVM_DEBUG(llvm::dbgs() << "=== Data Flow Graph Build Complete ===" + << "\n"); +} + +void DataFlowGraph::buildDefUseGraph() { + // 构建def-use图（在DataFlowInfo中实现） + dataFlowInfo.buildDefUseCache(); + + LLVM_DEBUG(llvm::dbgs() << "Def-use graph built" + << "\n"); +} + +void DataFlowGraph::print(llvm::raw_ostream &os) const { + os << "=== Data Flow Graph ===" + << "\n"; + dataFlowInfo.print(os); +} + +void DataFlowGraph::dump() const { print(llvm::errs()); } + +void DataFlowGraph::exportToJSON(llvm::raw_ostream &os) const { + auto funcOp = cfg.getFunction(); + auto funcName = funcOp.getName(); + std::string funcNameStr = funcName.empty() ? "unnamed" : funcName.str(); + + os << "{\n"; + os << " \"function\": \"" << funcNameStr << "\"," + << "\n"; + os << " \"dataFlow\": "; + dataFlowInfo.exportToJSON(os); + os << "}\n"; +} + +void DataFlowGraph::exportDefUseToDOT(llvm::raw_ostream &os) const { + os << "digraph DefUseGraph {\n"; + os << " rankdir=TB;\n"; + os << " node [shape=box];\n\n"; + + // 遍历所有definitions并导出 + size_t nodeId = 0; + DenseMap defToNode; + + cfg.traverse([&](const BasicBlock &bb) { + for (const auto &instPtr : bb.getInstructions()) { + const Instruction *inst = instPtr.get(); + const MemorySSAInfo &ssaInfo = inst->getMemorySSAInfo(); + + for (MemorySSADef *def : ssaInfo.definitions) { + if (def && !defToNode.count(def)) { + defToNode[def] = nodeId++; + os << " node_" << defToNode[def] << " [label=\"" << def->getId() + << "\"];\n"; + } + } + } + }); + + os << "}\n"; +} diff --git a/compiler/lib/TritonToGraph/GraphAnalysis.cpp b/compiler/lib/TritonToGraph/GraphAnalysis.cpp new file mode 100644 index 00000000..78c77868 --- /dev/null +++ b/compiler/lib/TritonToGraph/GraphAnalysis.cpp @@ -0,0 +1,626 @@ + + +#include "dicp/TritonToGraph/GraphAnalysis.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "graph-analysis" + +using namespace mlir; +using namespace triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// CFGTraverser Implementation +//===----------------------------------------------------------------------===// + +void CFGTraverser::dfsForward(CFGTraversalBase &visitor) { + DenseSet visited; + TraversalContext ctx; + dfsForwardImpl(cfg.getEntryBlock(), visited, ctx, visitor); +} + +void CFGTraverser::dfsForward(BasicBlock *start, CFGTraversalBase &visitor) { + DenseSet visited; + TraversalContext ctx; + dfsForwardImpl(start, visited, ctx, visitor); +} + +void CFGTraverser::dfsForwardImpl(BasicBlock *block, + DenseSet &visited, + TraversalContext &ctx, + CFGTraversalBase &visitor) { + if (!block || visited.contains(block)) + return; + + visited.insert(block); + + // pre-visit block + visitor.preVisitBlock(block, ctx); + + // visit instructions in block + for (auto &instPtr : block->getInstructions()) { + Instruction *inst = instPtr.get(); + + // check for structure entry + if (inst->hasSubGraph()) { + visitor.onEnterStructure(block, ctx); + ctx.push(block); + } + + visitor.VisitInstruction(inst, ctx); + + // check for structure exit (last instruction in structure block) + if (inst->hasSubGraph()) { + ctx.pop(); + visitor.onExitStructure(block, ctx); + } + } + + // visit successors + for (BasicBlock *succ : block->getSuccessors()) { + if (visited.contains(succ)) + continue; + + // check for back edge + if (cfg.isBackEdge(block, succ)) { + visitor.onBackEdge(block, succ, ctx); + } + + dfsForwardImpl(succ, visited, ctx, visitor); + } + + visitor.postVisitBlock(block, ctx); +} + +void CFGTraverser::dfsBackward(BasicBlock *start, CFGTraversalBase &visitor) { + DenseSet visited; + TraversalContext ctx; + dfsBackwardImpl(start, visited, ctx, visitor); +} + +void CFGTraverser::dfsBackwardImpl(BasicBlock *block, + DenseSet &visited, + TraversalContext &ctx, + CFGTraversalBase &visitor) { + if (!block || visited.contains(block)) + return; + + visited.insert(block); + + visitor.preVisitBlock(block, ctx); + + // visit instructions in reverse order + auto &insts = block->getInstructions(); + for (auto it = insts.rbegin(); it != insts.rend(); ++it) { + Instruction *inst = it->get(); + visitor.VisitInstruction(inst, ctx); + } + + // visit predecessors + for (BasicBlock *pred : block->getPredecessors()) { + if (cfg.isBackEdge(pred, block)) { + visitor.onBackEdge(pred, block, ctx); + } + dfsBackwardImpl(pred, visited, ctx, visitor); + } + + visitor.postVisitBlock(block, ctx); +} + +void CFGTraverser::bfsForward(CFGTraversalBase &visitor) { + bfsForward(cfg.getEntryBlock(), visitor); +} + +void CFGTraverser::bfsForward(BasicBlock *start, CFGTraversalBase &visitor) { + DenseSet visited; + SmallVector> worklist; + + worklist.push_back({start, TraversalContext()}); + visited.insert(start); + + while (!worklist.empty()) { + auto [block, ctx] = worklist.pop_back_val(); + + preVisitBlock(block, const_cast(ctx)); + + for (auto &instPtr : block->getInstructions()) { + Instruction *inst = instPtr.get(); + visitor.VisitInstruction(inst, const_cast(ctx)); + } + + visitor.postVisitBlock(block, const_cast(ctx)); + + for (BasicBlock *succ : block->getSuccessors()) { + if (!visited.contains(succ)) { + visited.insert(succ); + worklist.push_back({succ, ctx}); + } + } + } +} + +void CFGTraverser::bfsBackward(BasicBlock *start, CFGTraversalBase &visitor) { + DenseSet visited; + SmallVector> worklist; + + worklist.push_back({start, TraversalContext()}); + visited.insert(start); + + while (!worklist.empty()) { + auto [block, ctx] = worklist.pop_back_val(); + + visitor.preVisitBlock(block, const_cast(ctx)); + + auto &insts = block->getInstructions(); + for (auto it = insts.rbegin(); it != insts.rend(); ++it) { + Instruction *inst = it->get(); + visitor.VisitInstruction(inst, ctx); + } + + visitor.postVisitBlock(block, const_cast(ctx)); + + for (BasicBlock *pred : block->getPredecessors()) { + if (!visited.contains(pred)) { + visited.insert(pred); + worklist.push_back({pred, ctx}); + } + } + } +} + +//===----------------------------------------------------------------------===// +// DFGTraverser Implementation +//===----------------------------------------------------------------------===// + +void DFGTraverser::dfsBackward(Value seed, DFGTraversalBase &visitor, + const Options &opts) { + DenseSet visited; + dfsBackwardImpl(seed, visitor, visited, pts, 0); +} + +void DFGTraverser::dfsBackward(ArrayRef seeds, DFGTraversalBase &visitor, + const Options &opts) { + DenseSet visited; + for (Value seed : seeds) { + dfsBackwardImpl(seed, visitor, visited, opts, 0); + } +} + +void DFGTraverser::dfsBackwardImpl(Value value, DFGTraversalBase &visitor, + DenseSet &visited, + const Options &opts, int depth) { + if (opts.maxDepth >= 0 && depth > opts.maxDepth) + return; + + // get definition + Operation *defOp = nullptr; + if (opts.useMemorySSA) { + auto result = dfg.queryDataFlow(value); + if (auto *memResult = dyn_cast(result.get())) { + defOp = memResult->getDefinition()->getDefOp(); + } + } else { + defOp = value.getDefiningOp(); + } + + if (!defOp) + return; + + if (visited.contains(defOp)) + return; + + if (opts.stopOps.contains(defOp)) + return; + + visited.insert(defOp); + + visitor.VisitDef(value, defOp, depth); + + // recursively visit operands + for (Value operand : defOp->getOperands()) { + dfsBackwardImpl(operand, visitor, visited, opts, depth + 1); + } + + // handle phi/iter_arg + if (opts.followPhi) { + auto &dataFlowInfo = dfg.getDataFlowInfo(); + if (dataFlowInfo.hasPhi(value)) { + auto &phiInfo = dataFlowInfo.getPhi(value); + visitor.onPhi(value, phiInfo, depth); + } + } +} + +void DFGTraverser::dfsForward(Value seed, DFGTraversalBase &visitor, + const Options &opts) { + DenseSet visited; + dfsForwardImpl(seed, visitor, visited, ctx, opts, 0); +} + +void DFGTraverser::dfsForwardImpl(Value value, DFGTraversalBase &visitor, + DenseSet &visited, + const Options &opts, int depth) { + if (opts.maxDepth >= 0 && depth > opts.maxDepth) + return; + + SmallVector uses; + if (opts.useMemorySSA) { + uses = dfg.getDataFlowInfo().getSSAUses(value); + } else { + for (OpOperand &use : value.getUses()) { + uses.push_back(&use); + } + } + + for (OpOperand *use : uses) { + Operation *userOp = use->getOwner(); + + if (visited.contains(userOp)) + continue; + + if (opts.stopOps.contains(userOp)) + continue; + + visited.insert(userOp); + + visitor.VisitUse(value, use, depth); + + // visit results of this operation + for (Value result : userOp->getResults()) { + dfsForwardImpl(result, visitor, visited, opts, depth + 1); + } + } +} + +//===----------------------------------------------------------------------===// +// Region Implementation +//===----------------------------------------------------------------------===// + +void Region::add(Instruction *inst) { instSet_.insert(inst); } + +void Region::add(Operation *op, ControlFlowGraph &cfg) { + if (Instruction *inst = cfg.getInstruction(op)) { + add(inst); + } +} + +void Region::addAll(ArrayRef insts) { + for (Instruction *inst : insts) { + add(inst); + } +} + +bool Region::contains(Instruction *inst) const { + return instSet_.contains(inst); +} + +bool Region::contains(Operation *op) const { + // note: requires cfg to be available - caller must ensure this + // use the version that takes cfg as parameter for proper lookup + return false; +} + +void Region::remove(Instruction *inst) { instSet_.erase(inst); } + +void Region::clear() { instSet_.clear(); } + +SmallVector Region::orderedInstructions() const { + SmallVector result(instSet_.begin(), instSet_.end()); + // sort by block id then instruction index + llvm::sort(result, [](Instruction *a, Instruction *b) { + auto *bbA = a->getParentBlock(); + auto *bbB = b->getParentBlock(); + if (bbA != bbB) + return bbA->getId() < bbB->getId(); + // within same block, find position + // this is approximate - exact order requires full block scan + return a->getId() < b->getId(); + }); + return result; +} + +SmallVector Region::operations() const { + SmallVector result; + for (Instruction *inst : instSet_) { + if (Operation *op = inst->getOperation()) { + result.push_back(op); + } + } + return result; +} + +//===----------------------------------------------------------------------===// +// RegionAnalyzer Implementation +//===----------------------------------------------------------------------===// + +bool RegionAnalyzer::hasDependency(const Region &from, const Region &to) const { + auto deps = getDependencies(from, to); + return !deps.empty(); +} + +SmallVector +RegionAnalyzer::getDependencies(const Region &from, const Region &to) const { + SmallVector deps; + + // check for data dependencies: from defines, to uses + for (Instruction *fromInst : from) { + for (Value result : fromInst->getOperation()->getResults()) { + for (OpOperand &use : result.getUses()) { + Operation *userOp = use.getOwner(); + if (Instruction *toInst = cfg.getInstruction(userOp)) { + if (to.contains(toInst)) { + deps.push_back({Dependency::DATA, result, fromInst, toInst}); + } + } + } + } + } + + return deps; +} + +RegionAnalyzer::ExternalDeps +RegionAnalyzer::analyzeExternalDeps(const Region ®ion) const { + ExternalDeps result; + + // find inputs: external definitions used inside region + for (Instruction *inst : region) { + for (Value operand : inst->getOperation()->getOperands()) { + Operation *defOp = operand.getDefiningOp(); + if (!defOp) { + // block argument - treat as external input + bool alreadyTracked = false; + for (auto &input : result.inputs) { + if (input.value == operand) { + input.internalUses.push_back(inst); + alreadyTracked = true; + break; + } + } + if (!alreadyTracked) { + result.inputs.push_back({operand, nullptr, {inst}}); + } + } else if (Instruction *defInst = cfg.getInstruction(defOp)) { + if (!region.contains(defInst)) { + // external definition + bool alreadyTracked = false; + for (auto &input : result.inputs) { + if (input.value == operand) { + input.internalUses.push_back(inst); + alreadyTracked = true; + break; + } + } + if (!alreadyTracked) { + result.inputs.push_back({operand, defInst, {inst}}); + } + } + } + } + } + + // find outputs: internal definitions used outside region + for (Instruction *inst : region) { + for (Value result : inst->getOperation()->getResults()) { + SmallVector externalUses; + for (OpOperand &use : result.getUses()) { + Operation *userOp = use.getOwner(); + if (Instruction *userInst = cfg.getInstruction(userOp)) { + if (!region.contains(userInst)) { + externalUses.push_back(userInst); + } + } + } + if (!externalUses.empty()) { + result.outputs.push_back({result, inst, externalUses}); + } + } + } + + return result; +} + +//===----------------------------------------------------------------------===// +// ProgramSlicer Implementation +//===----------------------------------------------------------------------===// + +ProgramSlice ProgramSlicer::compute(const SliceCriterion &criterion) { + ProgramSlice slice; + + DFGTraverser dfgTraverser(dfg); + + for (Value seed : criterion.seeds) { + class SliceBuilder : public DFGTraversalBase { + public: + SliceBuilder(ProgramSlice &slice, ControlFlowGraph &cfg) + : slice(slice), cfg(cfg) {} + + bool VisitDef(Value value, Operation *defOp, int depth) override { + if (Instruction *inst = cfg.getInstruction(defOp)) { + slice.add(inst); + } + return true; + } + + bool VisitUse(Value value, OpOperand *use, int depth) override { + Operation *userOp = use->getOwner(); + if (Instruction *inst = cfg.getInstruction(userOp)) { + slice.add(inst); + } + return true; + } + + private: + ProgramSlice &slice; + ControlFlowGraph &cfg; + }; + + SliceBuilder builder(slice, cfg); + + switch (criterion.dir) { + case SliceCriterion::BACKWARD: + dfgTraverser.dfsBackward(seed, builder, criterion.dfgOpts); + break; + case SliceCriterion::FORWARD: + dfgTraverser.dfsForward(seed, builder, criterion.dfgOpts); + break; + case SliceCriterion::BIDIRECTIONAL: + dfgTraverser.dfsBackward(seed, builder, criterion.dfgOpts); + dfgTraverser.dfsForward(seed, builder, criterion.dfgOpts); + break; + } + } + + return slice; +} + +ProgramSlice ProgramSlicer::sliceFromYields(ArrayRef yields, + SliceCriterion::Direction dir) { + SliceCriterion criterion; + criterion.seeds = SmallVector(yields); + criterion.dir = dir; + return compute(criterion); +} + +void ProgramSlice::merge(const ProgramSlice &other) { + for (Instruction *inst : other) { + instructions_.insert(inst); + } +} + +void ProgramSlice::intersect(const ProgramSlice &other) { + DenseSet toRemove; + for (Instruction *inst : instructions_) { + if (!other.contains(inst)) { + toRemove.insert(inst); + } + } + for (Instruction *inst : toRemove) { + instructions_.erase(inst); + } +} + +void ProgramSlice::subtract(const ProgramSlice &other) { + for (Instruction *inst : other) { + instructions_.erase(inst); + } +} + +Region ProgramSlice::toRegion(StringRef name) const { + Region region(name); + for (Instruction *inst : instructions_) { + region.add(inst); + } + return region; +} + +//===----------------------------------------------------------------------===// +// RegionAbsorber Implementation +//===----------------------------------------------------------------------===// + +void RegionAbsorber::absorb(Region ®ion, ArrayRef seeds, + const AbsorptionPolicy &policy) { + DenseSet visited; + + for (Instruction *seed : seeds) { + region.add(seed); + + if (policy.dir == AbsorptionPolicy::UPSTREAM || + policy.dir == AbsorptionPolicy::BOTH) { + absorbUpstream(region, seed, policy, visited, 0); + } + + if (policy.dir == AbsorptionPolicy::DOWNSTREAM || + policy.dir == AbsorptionPolicy::BOTH) { + absorbDownstream(region, seed, policy, visited, 0); + } + } +} + +void RegionAbsorber::absorbUpstream(Region ®ion, Instruction *inst, + const AbsorptionPolicy &policy, + DenseSet &visited, + int depth) { + if (policy.maxDepth >= 0 && depth > policy.maxDepth) + return; + + if (visited.contains(inst)) + return; + visited.insert(inst); + + if (policy.shouldStop && policy.shouldStop(inst)) + return; + + if (policy.stopOps.contains(inst->getOperation())) + return; + + // add to region + region.add(inst); + + // visit operands + for (Value operand : inst->getOperation()->getOperands()) { + if (Operation *defOp = operand.getDefiningOp()) { + if (Instruction *defInst = cfg.getInstruction(defOp)) { + if (!policy.crossRegionBoundary && region.contains(defInst)) + continue; + + absorbUpstream(region, defInst, policy, visited, depth + 1); + } + } + } +} + +void RegionAbsorber::absorbDownstream(Region ®ion, Instruction *inst, + const AbsorptionPolicy &policy, + DenseSet &visited, + int depth) { + if (policy.maxDepth >= 0 && depth > policy.maxDepth) + return; + + if (visited.contains(inst)) + return; + visited.insert(inst); + + if (policy.shouldStop && policy.shouldStop(inst)) + return; + + if (policy.stopOps.contains(inst->getOperation())) + return; + + // add to region + region.add(inst); + + // visit uses + for (Value result : inst->getOperation()->getResults()) { + for (OpOperand &use : result.getUses()) { + Operation *userOp = use.getOwner(); + if (Instruction *userInst = cfg.getInstruction(userOp)) { + if (!policy.crossRegionBoundary && region.contains(userInst)) + continue; + + absorbDownstream(region, userInst, policy, visited, depth + 1); + } + } + } +} + +void RegionAbsorber::absorbFromValue(Region ®ion, Value value, + const AbsorptionPolicy &policy) { + Operation *defOp = value.getDefiningOp(); + if (!defOp) + return; + + Instruction *inst = cfg.getInstruction(defOp); + if (!inst) + return; + + absorb(region, {inst}, policy); +} + +void RegionAbsorber::absorbUntilBoundary( + Region ®ion, ArrayRef seeds, + std::function isBoundary) { + AbsorptionPolicy policy; + policy.shouldStop = isBoundary; + absorb(region, seeds, policy); +} diff --git a/compiler/lib/TritonToGraph/InterProceduralCFG.cpp b/compiler/lib/TritonToGraph/InterProceduralCFG.cpp new file mode 100644 index 00000000..3e22737e --- /dev/null +++ b/compiler/lib/TritonToGraph/InterProceduralCFG.cpp @@ -0,0 +1,331 @@ + + +#include "dicp/TritonToGraph/InterProceduralCFG.h" +#include "dicp/TritonToGraph/ControlFlowGraphBuilder.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "icfg" + +using namespace mlir; +using namespace mlir::triton; + +namespace mlir { +namespace triton { +namespace cfg { + +InterProceduralCFG::InterProceduralCFG(ModuleOp module) : module(module) {} + +InterProceduralCFG::~InterProceduralCFG() = default; + +void InterProceduralCFG::build() { + LLVM_DEBUG(llvm::dbgs() << "Building ICFG for module\n"); + + ControlFlowGraphBuilder builder; + + // 遍历模块中的所有函数 + for (triton::FuncOp func : module.getOps()) { + LLVM_DEBUG(llvm::dbgs() + << "Building CFG for function: " << func.getName() << "\n"); + + // 为每个函数构建CFG + auto cfg = builder.build(func); + functionCFGs[func] = std::move(cfg); + } +} + +ControlFlowGraph *InterProceduralCFG::getFunctionCFG(triton::FuncOp func) { + auto it = functionCFGs.find(func); + if (it != functionCFGs.end()) { + return it->second.get(); + } + return nullptr; +} + +const ControlFlowGraph * +InterProceduralCFG::getFunctionCFG(triton::FuncOp func) const { + auto it = functionCFGs.find(func); + if (it != functionCFGs.end()) { + return it->second.get(); + } + return nullptr; +} + +ControlFlowGraph *InterProceduralCFG::getFunctionCFG(StringRef funcName) { + for (auto &[func, cfg] : functionCFGs) { + if (func.getName() == funcName) { + return cfg.get(); + } + } + return nullptr; +} + +void InterProceduralCFG::connectCallGraph() { + LLVM_DEBUG(llvm::dbgs() << "Connecting call graph\n"); + + // 遍历所有操作，查找调用点 + module.walk([&](Operation *op) { + // 检查是否是调用操作 + if (auto callInterface = dyn_cast(op)) { + // 尝试获取被调用函数 + auto callable = callInterface.getCallableForCallee(); + + // 从符号引用获取函数名 + StringRef calleeName; + if (auto symbolRef = callable.dyn_cast()) { + calleeName = symbolRef.getRootReference(); + } else if (auto flatSymbolRef = + op->getAttrOfType("callee")) { + calleeName = flatSymbolRef.getValue(); + } + + if (!calleeName.empty()) { + // 在模块中查找被调用函数 + if (auto calleeFunc = module.lookupSymbol(calleeName)) { + // 获取调用者的函数 + Operation *parentOp = op->getParentOp(); + while (parentOp && !isa(parentOp)) { + parentOp = parentOp->getParentOp(); + } + + if (auto callerFunc = dyn_cast(parentOp)) { + auto callerCFG = getFunctionCFG(callerFunc); + + if (callerCFG) { + // 找到调用点的基本块 + BasicBlock *callBlock = nullptr; + Instruction *callInst = nullptr; + + // 遍历查找包含该操作的 basic block 和 instruction + for (size_t i = 0; i < callerCFG->getNumBlocks(); ++i) { + auto *bb = callerCFG->getBasicBlock(i); + for (size_t j = 0; j < bb->getNumInstructions(); ++j) { + auto *inst = bb->getInstruction(j); + if (inst->getOperation() == op) { + callBlock = bb; + callInst = inst; + break; + } + } + if (callBlock) + break; + } + + if (callBlock) { + // 创建调用点记录 + CallSite callSite; + callSite.callOp = op; + callSite.caller = callerFunc; + callSite.callee = calleeFunc; + callSite.callBlock = callBlock; + callSite.callInst = callInst; + + callSites.push_back(callSite); + + // 添加到调用图 + callGraph[callerFunc].push_back(calleeFunc); + reverseCallGraph[calleeFunc].push_back(callerFunc); + + LLVM_DEBUG(llvm::dbgs() + << "Found call: " << callerFunc.getName() << " -> " + << calleeFunc.getName() << "\n"); + } + } + } + } + } + } + }); +} + +SmallVector +InterProceduralCFG::getCallees(triton::FuncOp caller) const { + auto it = callGraph.find(caller); + if (it != callGraph.end()) { + return it->second; + } + return SmallVector(); +} + +SmallVector +InterProceduralCFG::getCallers(triton::FuncOp callee) const { + auto it = reverseCallGraph.find(callee); + if (it != reverseCallGraph.end()) { + return it->second; + } + return SmallVector(); +} + +void InterProceduralCFG::computeReachability() { + LLVM_DEBUG(llvm::dbgs() << "Computing reachability\n"); + + // 使用简单的递归DFS计算可达性 + std::function &)> dfs = + [&](triton::FuncOp func, DenseSet &visited) { + if (visited.contains(func)) { + return; + } + visited.insert(func); + + for (auto callee : getCallees(func)) { + reachability[func].insert(callee); + dfs(callee, reachability[func]); + } + }; + + // 对每个函数计算可达性 + for (auto &[func, cfg] : functionCFGs) { + reachability[func].clear(); + dfs(func, reachability[func]); + } +} + +bool InterProceduralCFG::isReachable(triton::FuncOp from, + triton::FuncOp to) const { + auto it = reachability.find(from); + if (it != reachability.end()) { + return it->second.contains(to); + } + return false; +} + +void InterProceduralCFG::dumpToDot(const std::string &filename) const { + std::error_code ec; + llvm::raw_fd_ostream file(filename, ec); + + if (ec) { + llvm::errs() << "Failed to open file: " << filename << "\n"; + return; + } + + file << "digraph ICFG {\n"; + file << " node [shape=box];\n\n"; + + // 输出所有函数节点 + for (const auto &[func, cfg] : functionCFGs) { + StringRef name = const_cast(func).getName(); + file << " \"" << name << "\" [label=\"" << name << "\"];\n"; + } + + file << "\n"; + + // 输出调用边 + for (const auto &[caller, callees] : callGraph) { + StringRef callerName = const_cast(caller).getName(); + for (auto callee : callees) { + StringRef calleeName = const_cast(callee).getName(); + file << " \"" << callerName << "\" -> \"" << calleeName + << "\" [label=\"call\"];\n"; + } + } + + file << "}\n"; +} + +void InterProceduralCFG::print(raw_ostream &os) const { + os << "InterProcedural Control Flow Graph:\n"; + os << "Functions: " << functionCFGs.size() << "\n"; + os << "CallSites: " << callSites.size() << "\n"; + + os << "\nCall Graph:\n"; + for (const auto &[caller, callees] : callGraph) { + StringRef callerName = const_cast(caller).getName(); + os << " " << callerName << " -> "; + for (auto callee : callees) { + StringRef calleeName = const_cast(callee).getName(); + os << calleeName << " "; + } + os << "\n"; + } +} + +llvm::Error +InterProceduralCFG::exportToHTML(const std::string &filename) const { + std::error_code ec; + llvm::raw_fd_ostream os(filename, ec, llvm::sys::fs::OF_Text); + + if (ec) { + return llvm::createStringError(llvm::inconvertibleErrorCode(), + "Failed to open file: " + filename); + } + + os << R"html( + + + + Inter-Procedural CFG + + + + +

Inter-Procedural CFG

+
+ + + +)html"; + + os.close(); + return llvm::Error::success(); +} + +} // namespace cfg +} // namespace triton +} // namespace mlir diff --git a/compiler/lib/TritonToGraph/MemorySsaBuilder.cpp b/compiler/lib/TritonToGraph/MemorySsaBuilder.cpp new file mode 100644 index 00000000..908686ed --- /dev/null +++ b/compiler/lib/TritonToGraph/MemorySsaBuilder.cpp @@ -0,0 +1,605 @@ + + +#include "dicp/TritonToGraph/MemorySsaBuilder.h" +#include "dicp/TritonToGraph/ControlFlowGraph.h" +#include "dicp/TritonToGraph/DataflowGraph.h" +#include "dicp/TritonToGraph/tensor.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "triton/Dialect/Triton/IR/OpInterfaces.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/raw_ostream.h" + +#define DEBUG_TYPE "memory-ssa-builder" + +using namespace mlir; +using namespace triton; +using namespace cfg; + +//===----------------------------------------------------------------------===// +// MemorySSABuilder +//===----------------------------------------------------------------------===// + +MemorySSABuilder::~MemorySSABuilder() { + // 清理所有创建的tensor objects（如果有动态分配的） + // 这里假设TensorObject由外部管理生命周期 +} + +void MemorySSABuilder::build() { + LLVM_DEBUG(llvm::dbgs() << "=== Starting Memory SSA Build ===\n"); + + // 步骤1: 拓扑排序，确定处理顺序 + std::vector topoOrder; + { + // 简单的拓扑排序实现 + DenseSet visited; + std::function dfs = [&](BasicBlock *bb) { + if (visited.contains(bb)) + return; + visited.insert(bb); + for (BasicBlock *succ : bb->getSuccessors()) { + dfs(succ); + } + topoOrder.push_back(bb); + }; + + // 从入口块开始DFS + dfs(cfg.getEntryBlock()); + std::reverse(topoOrder.begin(), topoOrder.end()); + } + + LLVM_DEBUG(llvm::dbgs() << "Topological order: " << topoOrder.size() + << " blocks\n"); + + // 步骤2: 初始化函数的参数 + createParameterDefinitions(); + + LLVM_DEBUG(llvm::dbgs() << "Created parameter definitions\n"); + + // 步骤3: 按拓扑序遍历每个BasicBlock + for (BasicBlock *bb : topoOrder) { + LLVM_DEBUG(llvm::dbgs() << "Processing BB" << bb->getId() << "\n"); + + // 处理block + processBasicBlock(bb); + } + + LLVM_DEBUG(llvm::dbgs() << "=== Memory SSA Build Complete ===\n" + << "Processed blocks: " << topoOrder.size() << "\n" + << "Total definitions: " << allDefinitions.size() + << "\n"); +} + +void MemorySSABuilder::createParameterDefinitions() { + triton::FuncOp func = cfg.getFunction(); + + LLVM_DEBUG(llvm::dbgs() << "Creating parameter definitions for " + << func.getName() << "\n"); + + // 遍历函数参数 + for (BlockArgument arg : func.getArguments()) { + Type argType = arg.getType(); + + // 检查是否是我们关心的类型 + if (isTensorType(argType)) { + // 创建参数名称 + std::string paramName = "param_" + std::to_string(arg.getArgNumber()); + + // 如果是指针类型，从aliasAnalysis获取tensor对象 + TensorObject *tensor = nullptr; + if (aliasAnalysis.isPointerType(argType)) { + tensor = aliasAnalysis.getTensorObject(arg); + } + + // 如果没有找到tensor对象，创建一个新的 + if (!tensor) { + // 从类型推断shape和element type + SmallVector shape; + Type elementType; + + extractShapeAndElementType(argType, shape, elementType); + + tensor = new TensorObject(paramName, shape, argType, elementType, + TensorObject::TensorKind::GLOBAL_MEMORY); + } + + tensor->print(llvm::outs()); + llvm::outs() << "\n"; + + // 为入参创建definition + MemorySSADef *def = createDefinition(tensor, nullptr); + + // 记录到dataFlowInfo + dataFlowInfo.addMemoryDefinition(arg, def); + + LLVM_DEBUG(llvm::dbgs() + << " Created parameter definition: " << def->getId() << "\n"); + } + } +} + +void MemorySSABuilder::processBasicBlock(BasicBlock *bb) { + if (!bb) + return; + + // 处理block内的所有指令 + for (auto &instPtr : bb->getInstructions()) { + Instruction *inst = instPtr.get(); + inst->print(llvm::outs()); + llvm::outs() << "\n"; + processInstruction(inst); + MemorySSAInfo &ssaInfo = inst->getMemorySSAInfo(); + ssaInfo.print(llvm::outs()); + llvm::outs() << "\n"; + } +} + +void MemorySSABuilder::processInstruction(Instruction *inst) { + if (!inst) + return; + + Operation *op = inst->getOperation(); + if (!op) + return; + + LLVM_DEBUG(llvm::dbgs() << "Processing: " << op->getName() << "\n"); + + MemorySSAInfo &ssaInfo = inst->getMemorySSAInfo(); + + // 1. 处理operands：创建uses + LLVM_DEBUG(llvm::dbgs() << " Processing operands...\n"); + + for (OpOperand &operand : op->getOpOperands()) { + Value operandValue = operand.get(); + unsigned operandIdx = operand.getOperandNumber(); + + // 检查是否是tensor或pointer类型 + if (isTensorType(operandValue.getType()) || + aliasAnalysis.isPointerType(operandValue.getType())) { + + // 查找operand的definition + MemorySSADef *def = dataFlowInfo.getMemoryDefinition(operandValue); + + if (def) { + // 创建use + MemorySSAUse use = createUse(def, op, operandIdx); + ssaInfo.uses.push_back(use); + + // 记录到全局map + dataFlowInfo.addMemoryUse(operandValue, use); + + LLVM_DEBUG(llvm::dbgs() + << " MemorySSAUse: " << def->getId() << " in " + << op->getName() << " [operand #" << operandIdx << "]\n"); + } + } + } + + // 2. 按操作类型处理：store、load、tensor writer、pointer op + LLVM_DEBUG(llvm::dbgs() << " Processing by operation type...\n"); + + // 对store操作（内存写入）- 为第一个operand指向的tensor创建新definition + if (isMemoryWriter(op)) { + LLVM_DEBUG(llvm::dbgs() << " Memory writer: " << op->getName() << "\n"); + + // store的第一个operand是ptr（第二个是value） + // 针对Triton算子场景做了简化，此处不是流敏感分析 + if (op->getNumOperands() >= 2) { + Value ptr = op->getOperand(0); + if (isTensorType(ptr.getType()) || + aliasAnalysis.isPointerType(ptr.getType())) { + // 获取ptr当前的definition（修改前的状态） + MemorySSADef *oldDef = dataFlowInfo.getMemoryDefinition(ptr); + if (oldDef) { + // store会修改内存，为同一个tensor创建新的definition + TensorObject *tensor = oldDef->getTensor(); + MemorySSADef *newDef = createDefinition(tensor, op); + + // store之后，ptr指向的内存状态改变，更新definition + dataFlowInfo.addMemoryDefinition(ptr, newDef); + + LLVM_DEBUG(llvm::dbgs() << " Store from: " << oldDef->getId() + << " to: " << newDef->getId() << "\n"); + } + } + } + } + // 对load操作（内存读取）- 第一个operand的definition作为result的definition + else if (isMemoryReader(op)) { + LLVM_DEBUG(llvm::dbgs() << " Memory reader: " << op->getName() << "\n"); + + if (op->getNumOperands() > 0 && op->getNumResults() > 0) { + Value ptr = op->getOperand(0); + Value result = op->getResult(0); + + // 获取ptr的definition + MemorySSADef *ptrDef = dataFlowInfo.getMemoryDefinition(ptr); + if (ptrDef) { + ssaInfo.definitions.push_back(ptrDef); + dataFlowInfo.addMemoryDefinition(result, ptrDef); + + LLVM_DEBUG(llvm::dbgs() + << " Load creates new def: " << ptrDef->getId() + << " from " << ptrDef->getId() << "\n"); + } + } + } + // 对返回新Tensor的操作（tensor writer） + else if (isTensorWriter(op)) { + LLVM_DEBUG(llvm::dbgs() << " Tensor writer: " << op->getName() << "\n"); + + for (Value result : op->getResults()) { + Type resultType = result.getType(); + + // 检查是否是tensor类型 + if (isTensorType(resultType)) { + // 创建tensor对象 + TensorObject *tensor = createTensorObject(op); + + // Tensor writer：创建新definition + MemorySSADef *newDef = createDefinition(tensor, op); + ssaInfo.definitions.push_back(newDef); + dataFlowInfo.addMemoryDefinition(result, newDef); + + LLVM_DEBUG(llvm::dbgs() << " Tensor definition: " << newDef->getId() + << " for " << result << "\n"); + } + } + } + // 对pointer操作（addptr, make_tensor_ptr, [broadcast, splat]） + else if (isPointerOp(op)) { + LLVM_DEBUG(llvm::dbgs() + << " Pointer operation: " << op->getName() << "\n"); + + for (Value result : op->getResults()) { + Type resultType = result.getType(); + + // 检查是否是pointer类型 + if (aliasAnalysis.isPointerType(resultType)) { + // Pointer op：复用base pointer的definition（alias） + Value basePtr = aliasAnalysis.getBasePointer(result); + MemorySSADef *baseDef = dataFlowInfo.getMemoryDefinition(basePtr); + + if (baseDef) { + ssaInfo.definitions.push_back(baseDef); + dataFlowInfo.addMemoryDefinition(result, baseDef); + + LLVM_DEBUG(llvm::dbgs() << " Pointer alias: " << baseDef->getId() + << " for " << result << "\n"); + } + } + } + } + + // 3. 特殊处理控制流操作 + if (auto ifOp = dyn_cast(op)) { + processIfOp(ifOp, inst, nullptr, nullptr); + } else if (auto forOp = dyn_cast(op)) { + processForOp(forOp, inst, nullptr); + } else if (auto whileOp = dyn_cast(op)) { + // processWhileOp(whileOp, inst, nullptr, nullptr); + } + + LLVM_DEBUG(llvm::dbgs() << " Done\n"); +} + +void MemorySSABuilder::processIfOp(scf::IfOp ifOp, Instruction *inst, + BasicBlock *thenEntryBB, + BasicBlock *elseEntryBB) { + // 实现scf.if的phi节点处理 + // 从then和else区域收集yield的values，创建phi definitions + + LLVM_DEBUG(llvm::dbgs() << "Processing IfOp: " << ifOp << "\n"); + + // 获取if指令 + // scf::IfOp ifOp = + // cast(ifCondBB->getInstruction(0)->getOperation()); + + // 为每个result创建phi节点 + for (size_t i = 0; i < ifOp.getNumResults(); ++i) { + Value ifResult = ifOp.getResult(i); + Type resultType = ifResult.getType(); + + // 只处理tensor/pointer类型 + if (!isTensorType(resultType) && !aliasAnalysis.isPointerType(resultType)) { + continue; + } + + // 从then区域获取yield的value + Operation *thenYield = + MemorySSABuilderHelper::getYieldOp(ifOp.getThenRegion()); + Value thenValue = thenYield ? thenYield->getOperand(i) : Value(); + MemorySSADef *thenDef = + thenValue ? dataFlowInfo.getMemoryDefinition(thenValue) : nullptr; + + // 从else区域获取yield的value + Operation *elseYield = + MemorySSABuilderHelper::getYieldOp(ifOp.getElseRegion()); + Value elseValue = elseYield ? elseYield->getOperand(i) : Value(); + MemorySSADef *elseDef = + elseValue ? dataFlowInfo.getMemoryDefinition(elseValue) : nullptr; + + // 如果then和else都返回相同的definition，可以直接使用 + if (thenDef && elseDef && thenDef == elseDef) { + dataFlowInfo.addMemoryDefinition(ifResult, thenDef); + + LLVM_DEBUG(llvm::dbgs() + << " If result #" << i + << " uses same definition: " << thenDef->getId() << "\n"); + continue; + } + + // 创建phi definition + if (thenDef || elseDef) { + TensorObject *tensor = thenDef ? thenDef->getTensor() + : elseDef ? elseDef->getTensor() + : nullptr; + + if (tensor) { + std::string phiName = "phi_" + std::to_string(ifOp.getNumResults()) + + "_" + std::to_string(i); + TensorObject *phiTensor = + new TensorObject(phiName, tensor->getShape(), resultType, + tensor->getElementType(), tensor->getKind()); + + MemorySSADef *phiDef = createDefinition(phiTensor, ifOp.getOperation()); + + // 记录if result的definition + dataFlowInfo.addMemoryDefinition(ifResult, phiDef); + + // 为phi节点的operands创建uses + if (thenDef && thenValue) { + MemorySSAUse thenUse(thenDef, ifOp.getOperation(), /*operandIdx=*/i); + dataFlowInfo.addMemoryUse(thenValue, thenUse); + } + if (elseDef && elseValue) { + MemorySSAUse elseUse(elseDef, ifOp.getOperation(), /*operandIdx=*/i); + dataFlowInfo.addMemoryUse(elseValue, elseUse); + } + + // 创建phi信息 + PhiInfo phiInfo; + phiInfo.type = PhiInfo::IF_RESULT; + phiInfo.loopHeader = nullptr; + phiInfo.comingFrom.initialValue = thenDef; + phiInfo.comingFrom.yieldValue = elseDef; + + dataFlowInfo.addPhi(ifResult, phiInfo); + + LLVM_DEBUG(llvm::dbgs() << " Phi: " << phiDef->getId() + << " for if result #" << i << "\n"); + } + } + } +} + +void MemorySSABuilder::processForOp(scf::ForOp forOp, Instruction *inst, + BasicBlock *loopBodyEntryBB) { + // 实现scf.for的iter_args处理 + // iter_args在所有迭代中共享同一个definition + + LLVM_DEBUG(llvm::dbgs() << "Processing ForOp: " << forOp << "\n"); + + unsigned numIterArgs = forOp.getInitArgs().size(); + + for (unsigned i = 0; i < numIterArgs; i++) { + Value iterArg = forOp.getRegionIterArg(i); + Value initValue = forOp.getInitArgs()[i]; + + // 查找initValue的definition + MemorySSADef *initDef = dataFlowInfo.getMemoryDefinition(initValue); + + if (initDef) { + // iter_arg使用同一个definition，不创建新版本 + dataFlowInfo.addMemoryDefinition(iterArg, initDef); + + // 创建PhiInfo（用于跟踪循环依赖） + PhiInfo phiInfo; + phiInfo.type = PhiInfo::ITER_ARG; + phiInfo.loopHeader = nullptr; // 需要根据CFG确定 + phiInfo.comingFrom.initialValue = initDef; + phiInfo.comingFrom.yieldValue = nullptr; // 将在处理yield时更新 + + dataFlowInfo.addPhi(iterArg, phiInfo); + + // 为initValue创建use + MemorySSAUse initUse(initDef, forOp.getOperation(), /*operandIdx=*/i); + dataFlowInfo.addMemoryUse(initValue, initUse); + + LLVM_DEBUG(llvm::dbgs() << " IterArg #" << i << ": " << iterArg << " -> " + << initDef->getId() << "\n"); + } + } + + // 处理yield操作（在循环体中） + Operation *yieldOp = MemorySSABuilderHelper::getYieldOp(forOp.getRegion()); + if (yieldOp) { + for (unsigned i = 0; i < yieldOp->getNumOperands(); i++) { + Value yieldedValue = yieldOp->getOperand(i); + Value iterArg = forOp.getRegionIterArg(i); + + MemorySSADef *yieldedDef = dataFlowInfo.getMemoryDefinition(yieldedValue); + + if (yieldedDef) { + // 更新PhiInfo + if (PhiInfo *phiInfo = &dataFlowInfo.getPhi(iterArg)) { + phiInfo->comingFrom.yieldValue = yieldedDef; + + // 为yieldedValue创建use + MemorySSAUse yieldUse(yieldedDef, yieldOp, /*operandIdx=*/i); + dataFlowInfo.addMemoryUse(yieldedValue, yieldUse); + + LLVM_DEBUG(llvm::dbgs() << " Yield #" << i << ": " << yieldedValue + << " updates iter_arg\n"); + } + } + } + } +} + +MemorySSADef *MemorySSABuilder::createDefinition(TensorObject *tensor, + Operation *op) { + unsigned version = isParameter(op) ? 0 : ++nextVersion[tensor]; + auto *def = new MemorySSADef(tensor, op, version); + allDefinitions.push_back(def); + return def; +} + +MemorySSAUse MemorySSABuilder::createUse(MemorySSADef *def, Operation *userOp, + unsigned operandIdx) { + return MemorySSAUse(def, userOp, operandIdx); +} + +TensorObject *MemorySSABuilder::createTensorObject(Operation *op) { + if (!op) + return nullptr; + + // 根据操作创建tensor对象，使用独立的Tensor ID确保唯一性 + std::string name = getOpName(op); + Type resultType = op->getResultTypes().front(); + + SmallVector shape; + Type elementType; + + // 使用Tensor.h中的辅助函数提取shape和element type + extractShapeAndElementType(resultType, shape, elementType); + + // 设置默认的kind（可以根据操作类型推断） + TensorObject::TensorKind kind = TensorObject::TensorKind::GLOBAL_MEMORY; + + auto *tensor = new TensorObject(name, shape, resultType, elementType, kind); + + // 缓存tensor对象 + if (!op->getResults().empty()) { + tensorObjectCache[op->getResult(0)] = tensor; + } + + return tensor; +} + +std::string MemorySSABuilder::getOpName(Operation *op) { + if (!op) + return "unknown"; + + // 基于操作类型和独立的Tensor ID生成名称 + std::string opName = op->getName().getStringRef().str(); + std::replace(opName.begin(), opName.end(), '.', '_'); + + // 使用独立的Tensor ID生成器，确保唯一性 + return opName + "_tensor_" + std::to_string(++nextTensor[opName]); +} + +//===----------------------------------------------------------------------===// +// MemorySSABuilderHelper +//===----------------------------------------------------------------------===// + +namespace mlir { +namespace triton { +namespace cfg { +namespace MemorySSABuilderHelper { + +Type getResultType(Operation *op, unsigned resultIdx) { + if (!op || resultIdx >= op->getNumResults()) + return Type(); + return op->getResultTypes()[resultIdx]; +} + +SmallVector getShapeFromValue(Value value) { + Type type = value.getType(); + SmallVector shape; + + if (auto rankedType = mlir::dyn_cast(type)) { + shape.append(rankedType.getShape().begin(), rankedType.getShape().end()); + } + + return shape; +} + +bool shapesEqual(ArrayRef shape1, ArrayRef shape2) { + if (shape1.size() != shape2.size()) + return false; + return std::equal(shape1.begin(), shape1.end(), shape2.begin()); +} + +Operation *getYieldOp(Region ®ion) { + if (region.empty()) + return nullptr; + + Block &block = region.back(); + if (block.empty()) + return nullptr; + + Operation &lastOp = block.back(); + if (isa(&lastOp)) { + return &lastOp; + } + + return nullptr; +} + +std::string createUniqueTensorName(StringRef prefix, size_t id) { + return prefix.str() + "_" + std::to_string(id); +} + +bool shouldCreateNewVersion(Operation *op, MemorySSADef *currentDef) { + if (!op || !currentDef) + return true; + + // 如果操作会修改tensor内容，则创建新版本 + // 例如：tt.store、tt.trans等 + if (isa(op)) + return true; + if (isa(op)) + return true; + + // 其他操作可能复用当前definition + return false; +} + +} // namespace MemorySSABuilderHelper +} // namespace cfg +} // namespace triton +} // namespace mlir + +bool MemorySSABuilder::isPointerBroadcastOrSplat(mlir::Operation *op) const { + if (auto broadcastOp = mlir::dyn_cast(op)) { + Type elemType = getElementTypeOrSelf(broadcastOp.getResult().getType()); + return mlir::isa(elemType); + } + + // 处理 SplatOp: 检查第一个 operand (src) 是否是指针 + if (auto splatOp = mlir::dyn_cast(op)) { + return mlir::isa(splatOp.getSrc().getType()); + } + + return false; +} + +// 判断是否是返回新Tensor的操作（根据返回值类型判断，排除load） +bool MemorySSABuilder::isTensorWriter(Operation *op) const { + if (!op || op->getNumResults() == 0) + return false; + + // 检查返回值类型 + for (Value result : op->getResults()) { + Type resultType = result.getType(); + // 如果是RankedTensorType（不是指针），则是TensorWriter + if (mlir::isa(resultType)) { + // 但排除load（虽然load返回tensor，但它是从内存读取，不是"写入"或"创建"） + if (mlir::isa(op) || mlir::isa(op)) + return false; + + if (mlir::isa(op)) + return false; + + if (isPointerBroadcastOrSplat(op)) + return false; + + if (auto ptrType = mlir::dyn_cast( + getElementTypeOrSelf(resultType))) { + return false; + } + return true; + } + } + return false; +} diff --git a/compiler/lib/TritonToGraph/Passes.cpp b/compiler/lib/TritonToGraph/Passes.cpp new file mode 100644 index 00000000..44534d19 --- /dev/null +++ b/compiler/lib/TritonToGraph/Passes.cpp @@ -0,0 +1,18 @@ + + +#include "dicp/TritonToGraph/Passes.h" +#include "dicp/TritonToGraph/ControlFlowGraphBuilder.h" + +#include "mlir/Pass/PassRegistry.h" + +#define GEN_PASS_REGISTRATION +#include "dicp/TritonToGraph/Passes.h.inc" + +namespace mlir { +namespace triton { + +// registerTritonToCFGPasses() 由 Passes.h.inc 生成 +// createBuildCFGPass() 的实现需要在 ControlFlowGraphBuilder.cpp 中 + +} // namespace triton +} // namespace mlir diff --git a/compiler/lib/TritonToHFusion/CMakeLists.txt b/compiler/lib/TritonToHFusion/CMakeLists.txt new file mode 100644 index 00000000..f69c1bfd --- /dev/null +++ b/compiler/lib/TritonToHFusion/CMakeLists.txt @@ -0,0 +1,14 @@ +add_triton_library(TritonToHFusion + TritonToHFusion.cpp + + DEPENDS + TritonToHFusionConversionPassIncGen + + LINK_LIBS + BiShengIRHFusionDialect + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport + TritonIR +) diff --git a/compiler/lib/TritonToHFusion/TritonToHFusion.cpp b/compiler/lib/TritonToHFusion/TritonToHFusion.cpp new file mode 100644 index 00000000..fed57c9e --- /dev/null +++ b/compiler/lib/TritonToHFusion/TritonToHFusion.cpp @@ -0,0 +1,233 @@ +#include "dicp/TritonToHFusion/Passes.h" + +#include "bishengir/Dialect/HFusion/IR/HFusion.h" +#include "bishengir/Dialect/HFusion/IR/HFusionImpl.h" +#include "bishengir/Dialect/Tensor/IR/TensorImpl.h" +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Attributes.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/LogicalResult.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_TRITONTOHFUSION +#include "dicp/TritonToHFusion/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; +using namespace hfusion; + +namespace { +struct TritonModToHFusionConversion : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::dicp::ModOp op, + PatternRewriter &rewriter) const final { + auto lhsType = dyn_cast(op.getLhs().getType()); + auto rhsType = dyn_cast(op.getRhs().getType()); + if (!lhsType || !rhsType) { + return failure(); + } + + auto emptyTensor = rewriter.create( + op.getLoc(), lhsType.getShape(), lhsType.getElementType()); + auto newOp = + hfusion::createBinaryOp( + rewriter, op.getLoc(), hfusion::BinaryFn::mod, + ValueRange({op.getLhs(), op.getRhs()}), + ValueRange({emptyTensor.getResult()})); + + rewriter.replaceOp(op, newOp->getResult(0)); + return success(); + } +}; + +struct TritonHistogramToHFusionConversion + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::HistogramOp op, + PatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + Value input = op.getSrc(); + auto resultType = op.getResult().getType(); + + int64_t numBins = 256; // 256 is default fallback. + if (auto rankedTy = dyn_cast(resultType)) + if (rankedTy.hasStaticShape() && rankedTy.getNumElements() > 0) + numBins = rankedTy.getNumElements(); + + auto numBinsAttr = rewriter.getI64IntegerAttr(numBins); + + auto newOp = rewriter.create(loc, resultType, input, + numBinsAttr, Value()); + + rewriter.replaceOp(op, newOp.getResult()); + return success(); + } +}; + +struct TritonFpToFpToHFusionConversion : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::FpToFpOp op, + PatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + Value input = op.getSrc(); + auto resultType = op.getResult().getType(); + + // Only handle float-to-float conversions with non-RTNE rounding modes + // RTNE (default) rounding is handled by TritonToLinalg pass using + // arith.truncf/extf + auto srcType = cast(input.getType()); + auto dstType = cast(resultType); + if (!srcType.getElementType().isIntOrFloat() || + !dstType.getElementType().isIntOrFloat()) { + return failure(); + } + + // Check if this has a non-RTNE rounding mode + auto roundingMode = op.getRounding(); + if (!roundingMode.has_value() || + roundingMode.value() == triton::RoundingMode::RTNE) { + // RTNE or no rounding mode specified: let TritonToLinalg handle it + return failure(); + } + + // Map non-RTNE rounding modes to HFusion rounding mode + hfusion::RoundMode hfusionRoundMode; + switch (roundingMode.value()) { + case triton::RoundingMode::RTZ: + hfusionRoundMode = hfusion::RoundMode::TRUNC; + break; + default: + return op.emitError("Unsupported rounding mode for HFusion conversion"); + } + // Note: Only RTZ (and potential future non-RTNE modes) reach here + + // Get or create destination tensor (destination-style) + SmallVector dsts; + if (failed(tensor::getOrCreateDestinations(rewriter, loc, op, dsts))) + return failure(); + + // Create the HFusion cast operation with round_mode attribute + auto roundModeAttr = + hfusion::RoundModeAttr::get(rewriter.getContext(), hfusionRoundMode); + auto modeAttr = rewriter.getNamedAttr("mode", roundModeAttr); + + rewriter.replaceOpWithNewOp( + op, ValueRange{input}, ValueRange{dsts}, ArrayRef{modeAttr}); + + return success(); + } +}; + +struct TritonConv1dToHFusionConversion + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::dicp::Conv1dOp op, + PatternRewriter &rewriter) const final { + auto loc = op.getLoc(); + + Value input = op.getInput(); + Value weight = op.getWeight(); + Value biasValue = op.getBias(); + int64_t stride = op.getStride(); + int64_t padding_size = op.getPaddingSize(); + int64_t dilation = op.getDilation(); + int64_t groups = op.getGroups(); + + auto inputType = mlir::cast(input.getType()); + auto weightType = mlir::cast(weight.getType()); + if (!inputType.hasStaticShape() || !weightType.hasStaticShape()) { + return failure(); + } + + ArrayRef inputShape = inputType.getShape(); + ArrayRef weightShape = weightType.getShape(); + + bool isBatched = inputShape.size() == 3; + int64_t N; + if (isBatched) + N = inputShape[0]; + int64_t L_in = inputShape[isBatched ? 2 : 1]; + int64_t C_out = weightShape[0]; + int64_t kernel_size = weightShape[2]; + + if (stride == 0) { + return failure(); + } + int64_t L_out = + (L_in + 2 * padding_size - dilation * (kernel_size - 1) - 1) / stride + + 1; + + auto resultType = mlir::cast(op.getResult().getType()); + Type resultElementType = resultType.getElementType(); + + constexpr int64_t dim2 = 2; + constexpr int64_t dim3 = 3; + Value initTensor; + if (isBatched) { + SmallVector outputShape{N, C_out, L_out}; + initTensor = + rewriter.create(loc, outputShape, resultElementType); + } else { + SmallVector outputShape{C_out, L_out}; + initTensor = + rewriter.create(loc, outputShape, resultElementType); + } + + SmallVector ins; + ins.push_back(input); + ins.push_back(weight); + if (biasValue) { + ins.push_back(biasValue); + } + + auto newOp = rewriter.create( + loc, ins, initTensor, stride, padding_size, dilation, groups); + + rewriter.replaceOp(op, newOp.getResult()); + + return success(); + } +}; +} // namespace + +namespace { +struct TritonToHFusionPass + : public mlir::triton::impl::TritonToHFusionBase { + void runOnOperation() override; +}; +} // namespace + +void TritonToHFusionPass::runOnOperation() { + auto module = getOperation(); + + // Use greedy pattern rewriter for simpler pattern matching + // Patterns decide themselves whether to convert (via returning + // success/failure) + RewritePatternSet patterns(&getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + + // Apply patterns with greedy rewriting + // This allows patterns to return failure() without causing pass failure + if (failed(applyPatternsGreedily(module, std::move(patterns)))) { + signalPassFailure(); + } +} + +std::unique_ptr> +mlir::triton::createTritonToHFusionPass() { + return std::make_unique(); +} diff --git a/compiler/lib/Conversion/LinkedToHIVM/CMakeLists.txt b/compiler/lib/TritonToHIVM/CMakeLists.txt similarity index 56% rename from compiler/lib/Conversion/LinkedToHIVM/CMakeLists.txt rename to compiler/lib/TritonToHIVM/CMakeLists.txt index 9f411810..52a3d532 100644 --- a/compiler/lib/Conversion/LinkedToHIVM/CMakeLists.txt +++ b/compiler/lib/TritonToHIVM/CMakeLists.txt @@ -1,8 +1,8 @@ -add_triton_library(LinkedToHIVM - LinkedToHIVM.cpp +add_triton_library(TritonToHIVM + TritonToHIVM.cpp DEPENDS - LinkedToHIVMConversionPassIncGen + TritonToHIVMConversionPassIncGen LINK_LIBS BiShengIRHIVMDialect diff --git a/compiler/lib/Conversion/LinkedToHIVM/LinkedToHIVM.cpp b/compiler/lib/TritonToHIVM/TritonToHIVM.cpp similarity index 72% rename from compiler/lib/Conversion/LinkedToHIVM/LinkedToHIVM.cpp rename to compiler/lib/TritonToHIVM/TritonToHIVM.cpp index dbbb7da4..d06c7e7c 100644 --- a/compiler/lib/Conversion/LinkedToHIVM/LinkedToHIVM.cpp +++ b/compiler/lib/TritonToHIVM/TritonToHIVM.cpp @@ -1,5 +1,8 @@ + + +#include "dicp/TritonToHIVM/Passes.h" + #include "bishengir/Dialect/HIVM/IR/HIVM.h" -#include "dicp/Conversion/LinkedToHIVM/Passes.h" #include "mlir/IR/Attributes.h" #include "mlir/Pass/Pass.h" #include "mlir/Transforms/DialectConversion.h" @@ -8,14 +11,18 @@ #include "llvm/ADT/StringRef.h" #include "llvm/Support/LogicalResult.h" +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_TRITONTOHIVM +#include "dicp/TritonToHIVM/Passes.h.inc" +} // namespace triton +} // namespace mlir + using namespace mlir; -using namespace dicp; -using namespace linked; using namespace hivm; -#define GEN_PASS_CLASSES -#include "dicp/Conversion/LinkedToHIVM/Passes.h.inc" - namespace { struct CoreAndPipes { @@ -68,21 +75,27 @@ static CoreAndPipes GetCoreAndPipes(MLIRContext *ctx, llvm::StringRef opName, return {core, producer, consumer}; } -struct LinkedToHIVMPass : public LinkedToHIVMBase { +} // end anonymous namespace +namespace { +struct TritonToHIVMPass + : public mlir::triton::impl::TritonToHIVMBase { void runOnOperation() override; }; +} // namespace -struct TritonCustomSyncOpToHIVMSyncOpConversion - : OpRewritePattern { - using OpRewritePattern::OpRewritePattern; +struct TritonCustomOpToHIVMSyncOpConversion + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; - LogicalResult matchAndRewrite(triton::CustomSyncOp op, + LogicalResult matchAndRewrite(triton::dicp::CustomOp op, PatternRewriter &rewriter) const final { auto *ctx = op->getContext(); auto loc = op->getLoc(); + auto args = op.getStrArgs(); + auto argAttr = dyn_cast(args[0]); + auto id = dyn_cast(args[1]).getInt(); llvm::StringRef opName = op.getOpName(); - llvm::StringRef arg = op.getModeOrSender(); - auto id = op.getId(); + llvm::StringRef arg = argAttr.getValue(); if (opName == "sync_block_all") { if (arg == "all_cube") { @@ -123,38 +136,19 @@ struct TritonCustomSyncOpToHIVMSyncOpConversion } }; -// Convert CustomOp after operand type changed, -// for example tt.ptr changed to memref. -class TritonCustomOpToHIVMCustomOpConversion - : public OpConversionPattern { -public: - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(triton::CustomOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const override { - auto res_types = adaptor.getOutputs().getTypes(); - auto new_op = rewriter.create( - op->getLoc(), res_types, adaptor.getOperands(), op->getAttrs()); - rewriter.replaceOp(op, new_op); - return success(); - } -}; - -void LinkedToHIVMPass::runOnOperation() { +void TritonToHIVMPass::runOnOperation() { auto module = getOperation(); ConversionTarget target(getContext()); target.addLegalDialect(); RewritePatternSet patterns(&getContext()); - patterns.add(patterns.getContext()); - patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); if (failed(applyPartialConversion(module, target, std::move(patterns)))) { signalPassFailure(); } } -} // namespace -std::unique_ptr> linked::createLinkedToHIVMPass() { - return std::make_unique(); -} \ No newline at end of file +std::unique_ptr> +mlir::triton::createTritonToHIVMPass() { + return std::make_unique(); +} diff --git a/compiler/lib/TritonToLLVM/CMakeLists.txt b/compiler/lib/TritonToLLVM/CMakeLists.txt new file mode 100644 index 00000000..4ec73291 --- /dev/null +++ b/compiler/lib/TritonToLLVM/CMakeLists.txt @@ -0,0 +1,13 @@ +add_triton_library(TritonToLLVM + TritonToLLVM.cpp + + DEPENDS + TritonToLLVMConversionPassIncGen + + LINK_LIBS + MLIRIR + MLIRPass + MLIRTransforms + MLIRSupport + TritonIR +) diff --git a/compiler/lib/TritonToLLVM/TritonToLLVM.cpp b/compiler/lib/TritonToLLVM/TritonToLLVM.cpp new file mode 100644 index 00000000..33d5374b --- /dev/null +++ b/compiler/lib/TritonToLLVM/TritonToLLVM.cpp @@ -0,0 +1,269 @@ +#include "dicp/TritonToLLVM/Passes.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/IR/Attributes.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/LogicalResult.h" + +namespace mlir { +namespace triton { +#define GEN_PASS_DEF_TRITONTOLLVM +#include "dicp/TritonToLLVM/Passes.h.inc" +} // namespace triton +} // namespace mlir + +using namespace mlir; + +namespace { +struct TritonToLLVMPass + : public mlir::triton::impl::TritonToLLVMBase { + void runOnOperation() override; +}; +} // namespace + +namespace { + +static Type getElementType(Value value) { + auto type = value.getType(); + if (auto tensorType = dyn_cast(type)) + return tensorType.getElementType(); + return type; +} + +static int64_t getTensorNumElements(Value tensor) { + auto type = mlir::cast(tensor.getType()); + return type.getNumElements(); +} + +static Value getInt32Value(RewriterBase &rewriter, Location loc, int val) { + Type ty = rewriter.getI32Type(); + return rewriter.create(loc, ty, + rewriter.getIntegerAttr(ty, val)); +} + +// If operand size is smaller than 32 bits, pack in groups of 32 bits. +SmallVector packOperands(mlir::triton::ElementwiseInlineAsmOp op, + const SmallVector> &operands, + RewriterBase &rewriter, Location loc) { + SmallVector packedOperands; + unsigned numPackedElements = op.getPackedElement(); + for (int i = 0, e = op.getNumOperands(); i < e; i++) { + Type elemTy = getElementType(op.getOperand(i)); + unsigned bitWidth = + elemTy.isIntOrFloat() ? elemTy.getIntOrFloatBitWidth() : 64; + unsigned numElementPerReg = std::max(32 / bitWidth, 1u); + numElementPerReg = std::min(numElementPerReg, numPackedElements); + for (int j = 0; j < numPackedElements; j += numElementPerReg) { + if (numElementPerReg == 1) { + packedOperands.push_back(operands[j][i]); + continue; + } + Type t = VectorType::get(numElementPerReg, elemTy); + Value packed = rewriter.create(loc, t); + for (int k = 0; k < numElementPerReg; k++) { + packed = rewriter.create( + loc, packed, operands[j + k][i], getInt32Value(rewriter, loc, k)); + } + packedOperands.push_back(packed); + } + } + return packedOperands; +} + +static SmallVector unpackElements(Location loc, Value packedValues, + RewriterBase &rewriter) { + auto type = mlir::cast(packedValues.getType()); + auto elementType = type.getElementType(); + auto shape = type.getShape(); + + int64_t numElements = type.getNumElements(); + + SmallVector result; + for (int64_t linearIdx = 0; linearIdx < numElements; linearIdx++) { + SmallVector indexes(shape.size()); + int64_t remaining = linearIdx; + for (int64_t dim = shape.size() - 1; dim >= 0; dim--) { + indexes[dim] = + rewriter.create(loc, remaining % shape[dim]); + remaining /= shape[dim]; + } + Value extracted = rewriter.create(loc, elementType, + packedValues, indexes); + result.push_back(extracted); + } + + return result; +} + +static SmallVector> +createDestOps(triton::ElementwiseInlineAsmOp op, RewriterBase &rewriter, + const SmallVector> operands, Location loc) { + auto ctx = op->getContext(); + if (operands.size() % op.getPackedElement() != 0) + llvm::report_fatal_error("Inline asm op has more packed elements than " + "number of elements per thread."); + + // Pack elems smaller than 32 bits into 32-bit registers. + SmallVector packedOperands = packOperands(op, operands, rewriter, loc); + + // Types returned by the LLVM asm op. If there's more than one, they'll be + // wrapped in a type tuple. + SmallVector asmRetTypes; + for (auto result : op.getResult()) { + auto ty = getElementType(result); + + // Pack return elements into 32-bits. + unsigned bitWidth = ty.isIntOrFloat() ? ty.getIntOrFloatBitWidth() : 64; + unsigned numElemsPerReg = + std::min(std::max(32 / bitWidth, 1u), op.getPackedElement()); + assert(op.getPackedElement() % numElemsPerReg == 0); + if (numElemsPerReg > 1) { + ty = VectorType::get(numElemsPerReg, ty); + } + for (unsigned i = 0; i < op.getPackedElement() / numElemsPerReg; i++) { + asmRetTypes.push_back(ty); + } + } + Type asmRetType = asmRetTypes.size() > 1 + ? LLVM::LLVMStructType::getLiteral(ctx, asmRetTypes) + : asmRetTypes[0]; + + Value asmResults = + rewriter + .create( + loc, asmRetType, + /*operands=*/packedOperands, + /*asm_string=*/op.getAsmString(), + /*constraints=*/op.getConstraints(), + /*has_side_effects=*/!op.getPure(), + /*is_align_stack=*/false, + /*tail_call_kind=*/LLVM::tailcallkind::TailCallKind::None, + /*asm_dialect=*/ + LLVM::AsmDialectAttr::get(rewriter.getContext(), + LLVM::AsmDialect::AD_ATT), + /*operand_attrs=*/ArrayAttr()) + ->getResult(0); + + // asmResults is a flat struct; pack its values into + // [return_value][op.getPackedElement()]. + SmallVector> ret(op->getNumResults()); + int structIdx = 0; + for (int i = 0; i < op->getNumResults(); i++) { + for (int j = 0; j < op.getPackedElement(); j++) { + Value val; + if (asmRetTypes.size() > 1) { + val = + rewriter.create(loc, asmResults, structIdx++); + } else { + val = asmResults; + } + if (auto vectorTy = dyn_cast(val.getType())) { + for (int k = 0; k < vectorTy.getNumElements(); k++) { + ret[i].push_back(rewriter.create( + loc, val, getInt32Value(rewriter, loc, k))); + } + j += vectorTy.getNumElements() - 1; + } else { + ret[i].push_back(val); + } + } + } + return ret; +} + +static LogicalResult processScalarInlineAsm(triton::ElementwiseInlineAsmOp op, + PatternRewriter &rewriter) { + Location loc = op.getLoc(); + + auto outsWrapped = createDestOps(op, rewriter, {}, loc); + + SmallVector outs; + for (const auto &resWrapped : outsWrapped) { + outs.push_back(resWrapped[0]); + } + rewriter.replaceOp(op, outs); + + return success(); +} + +static LogicalResult processVectorInlineAsm(triton::ElementwiseInlineAsmOp op, + PatternRewriter &rewriter) { + Location loc = op.getLoc(); + + SmallVector> unpackedOperands; + for (auto operand : op.getOperands()) { + auto unpackedOperand = unpackElements(loc, operand, rewriter); + unpackedOperands.push_back(unpackedOperand); + } + + int64_t resultLength = getTensorNumElements(op->getResult(0)); + if (resultLength % op.getPackedElement()) { + op.emitError("Result tensor should be diveded to pack"); + return failure(); + } + + SmallVector> unpackedResults(op->getNumResults()); + for (int64_t i = 0; i < resultLength; i += op.getPackedElement()) { + // Block of elements to process with one call to the inline asm. This is + // ordered opposite `unpackedResults`: The outer dim is + // op.getPackedElement(), and the inner dim is the operand. + SmallVector> block(op.getPackedElement()); + for (auto &os : unpackedOperands) { + for (int j = 0; j < op.getPackedElement(); j++) { + block[j].push_back(os[i + j]); + } + } + auto cur = createDestOps(op, rewriter, block, loc); + assert(cur.size() == unpackedResults.size()); + for (unsigned j = 0; j < cur.size(); j++) { + unpackedResults[j].insert(unpackedResults[j].end(), cur[j].begin(), + cur[j].end()); + } + } + // Reorder and pack the results. + SmallVector outs; + for (int i = 0; i < unpackedResults.size(); i++) { + outs.push_back(rewriter.create( + loc, op->getResult(i).getType(), unpackedResults[i])); + } + rewriter.replaceOp(op, outs); + + return success(); +} + +} // namespace + +struct ElementwiseInlineAsmOpConversion + : OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(triton::ElementwiseInlineAsmOp op, + PatternRewriter &rewriter) const final { + return op.getOperands().empty() ? processScalarInlineAsm(op, rewriter) + : processVectorInlineAsm(op, rewriter); + } +}; + +void TritonToLLVMPass::runOnOperation() { + auto module = getOperation(); + ConversionTarget target(getContext()); + target.addLegalDialect(); + + RewritePatternSet patterns(&getContext()); + patterns.add(patterns.getContext()); + if (failed(applyPartialConversion(module, target, std::move(patterns)))) { + signalPassFailure(); + } +} + +std::unique_ptr> +mlir::triton::createTritonToLLVMPass() { + return std::make_unique(); +} diff --git a/compiler/lib/TritonToLinalg/ArgMinMaxConverter.cpp b/compiler/lib/TritonToLinalg/ArgMinMaxConverter.cpp new file mode 100644 index 00000000..05c45535 --- /dev/null +++ b/compiler/lib/TritonToLinalg/ArgMinMaxConverter.cpp @@ -0,0 +1,93 @@ + + +#include "dicp/TritonToLinalg/ArgMinMaxConverter.h" +#include +#include + +namespace TTOpConverters { +using namespace mlir; +using namespace triton; + +// ArgMinConverter functions +LogicalResult ArgMinConverter::matchComparisonResult( + Value currValue, Value currIndex, Value reduceValue, Value reduceIndex, + mlir::Block::iterator &it, Value &comparisonResult) { + LLVM_DEBUG(llvm::dbgs() << "Matching: " << *it << "\n"); + + auto cmpOp = dyn_cast(*it); + auto cmpIOp = dyn_cast(*it++); + if (!cmpOp && !cmpIOp) + return failure(); + + if (cmpOp) { + if (cmpOp.getPredicate() != arith::CmpFPredicate::OLT || + currValue != cmpOp.getLhs() || reduceValue != cmpOp.getRhs()) { + return failure(); + } + comparisonResult = cmpOp; + } + + if (cmpIOp) { + if ((cmpIOp.getPredicate() != arith::CmpIPredicate::slt && + cmpIOp.getPredicate() != arith::CmpIPredicate::ult) || + currValue != cmpIOp.getLhs() || reduceValue != cmpIOp.getRhs()) { + return failure(); + } + comparisonResult = cmpIOp; + } + + return success(); +} + +float ArgMinConverter::getBaseReductionValue() { + return std::numeric_limits::infinity(); +} + +int8_t ArgMinConverter::getBaseReductionIntValue() { + return std::numeric_limits::max(); +} +uint8_t ArgMinConverter::getBaseReductionUIntValue() { + return std::numeric_limits::max(); +} + +// ArgMaxConverter functions +LogicalResult ArgMaxConverter::matchComparisonResult( + Value currValue, Value currIndex, Value reduceValue, Value reduceIndex, + mlir::Block::iterator &it, Value &comparisonResult) { + auto cmpOp = dyn_cast(*it); + auto cmpIOp = dyn_cast(*it++); + if (!cmpOp && !cmpIOp) + return failure(); + + if (cmpOp) { + if (cmpOp.getPredicate() != arith::CmpFPredicate::OGT || + currValue != cmpOp.getLhs() || reduceValue != cmpOp.getRhs()) { + return failure(); + } + comparisonResult = cmpOp; + } + + if (cmpIOp) { + if ((cmpIOp.getPredicate() != arith::CmpIPredicate::sgt && + cmpIOp.getPredicate() != arith::CmpIPredicate::ugt) || + currValue != cmpIOp.getLhs() || reduceValue != cmpIOp.getRhs()) { + return failure(); + } + comparisonResult = cmpIOp; + } + + return success(); +} + +float ArgMaxConverter::getBaseReductionValue() { + return -std::numeric_limits::infinity(); +} + +int8_t ArgMaxConverter::getBaseReductionIntValue() { + return std::numeric_limits::min(); +} +uint8_t ArgMaxConverter::getBaseReductionUIntValue() { + return std::numeric_limits::min(); +} + +} // namespace TTOpConverters diff --git a/compiler/lib/TritonToLinalg/AscendNPUIRLegalizePass.cpp b/compiler/lib/TritonToLinalg/AscendNPUIRLegalizePass.cpp new file mode 100644 index 00000000..4ea6e9d0 --- /dev/null +++ b/compiler/lib/TritonToLinalg/AscendNPUIRLegalizePass.cpp @@ -0,0 +1,87 @@ + + +#include "dicp/TritonToLinalg/AscendNPUIRLegalizePass.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +using namespace mlir; +using namespace triton; + +struct ReifyUnsignedMulViaSignedPattern + : public OpRewritePattern { + bool unsafeMode; + + ReifyUnsignedMulViaSignedPattern(MLIRContext *ctx, bool unsafeMode) + : OpRewritePattern(ctx, /*benefit=*/1), + unsafeMode(unsafeMode) {} + + LogicalResult matchAndRewrite(arith::MulUIExtendedOp op, + PatternRewriter &rewriter) const override { + Value lhs = op.getLhs(); + Value rhs = op.getRhs(); + + // Only handle i32 element type + Type elemType; + if (auto shapedType = dyn_cast(lhs.getType())) { + if (!shapedType.getElementType().isInteger(32)) + return failure(); + elemType = shapedType.getElementType(); + } else { + if (!lhs.getType().isInteger(32)) + return failure(); + elemType = lhs.getType(); + } + + auto mulsiOp = rewriter.create( + op.getLoc(), op.getResultTypes(), lhs, rhs); + + if (unsafeMode) { + rewriter.replaceOp(op, {mulsiOp.getLow(), mulsiOp.getHigh()}); + return success(); + } + + // high_unsigned = high_signed + (a >> 31) * b + (b >> 31) * a + // Build a constant matching the operand type (scalar or tensor). + auto makeConst = [&](int64_t val) -> Value { + if (auto shaped = dyn_cast(lhs.getType())) + return rewriter.create( + op.getLoc(), shaped, + DenseElementsAttr::get( + shaped, APInt(shaped.getElementTypeBitWidth(), val))); + return rewriter.create( + op.getLoc(), rewriter.getIntegerAttr(elemType, val)); + }; + Value c31 = makeConst(31); + Value sA = rewriter.create(op.getLoc(), lhs, c31); + Value sB = rewriter.create(op.getLoc(), rhs, c31); + Value corrA = rewriter.create(op.getLoc(), sA, rhs); + Value corrB = rewriter.create(op.getLoc(), sB, lhs); + Value tmp = + rewriter.create(op.getLoc(), mulsiOp.getHigh(), corrA); + Value highUnsigned = + rewriter.create(op.getLoc(), tmp, corrB); + + rewriter.replaceOp(op, {mulsiOp.getLow(), highUnsigned}); + return success(); + } +}; + +void AscendNPUIRLegalizePass::runOnOperation() { + RewritePatternSet patterns(&getContext()); + patterns.add(&getContext(), unsafeMode); + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) + signalPassFailure(); +} + +std::unique_ptr> +triton::createAscendNPUIRLegalizePass() { + return std::make_unique(); +} + +std::unique_ptr> triton::createAscendNPUIRLegalizePass( + const AscendNPUIRLegalizeOptions &options) { + return std::make_unique(options); +} diff --git a/compiler/lib/TritonToLinalg/BlockPtrAnalysis.cpp b/compiler/lib/TritonToLinalg/BlockPtrAnalysis.cpp new file mode 100644 index 00000000..7e731b30 --- /dev/null +++ b/compiler/lib/TritonToLinalg/BlockPtrAnalysis.cpp @@ -0,0 +1,2416 @@ + + +#include "dicp/TritonToLinalg/BlockPtrAnalysis.h" +#include "dicp/TritonToLinalg/TritonToLinalgPass.h" +#include "dicp/Utils/Utils.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/IR/Types.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Transforms/DialectConversion.h" +#include "triton/Dialect/Triton/IR/Types.h" + +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/FormatVariadic.h" +#include +#include + +#define DEBUG_TYPE "triton-block-ptr-analysis" +namespace mlir { +namespace triton { + +// MemAccType selectMaxMemAccTy(const MemAccType &v1, const MemAccType &v2) { +// return (v1 > v2) ? v1 : v2; +// } + +SmallVector &BlockData::getOffsetsRef() { return this->offsets; } + +SmallVector &BlockData::getSizesRef() { return this->sizes; } + +SmallVector &BlockData::getStridesRef() { return this->strides; } + +Value &BlockData::getSourceRef() { return this->source; } + +OpFoldResult &BlockData::getScalarRef() { return this->scalar; } + +SmallVector BlockData::getOffsets() const { + return this->offsets; +} + +SmallVector BlockData::getSizes() const { return this->sizes; } + +SmallVector BlockData::getStrides() const { + return this->strides; +} + +OpFoldResult BlockData::getOffset(int index) const { + return this->offsets[index]; +} + +OpFoldResult BlockData::getSize(int index) const { return this->sizes[index]; } + +OpFoldResult BlockData::getStride(int index) const { + return this->strides[index]; +} + +OpFoldResult BlockData::getScalar() const { return this->scalar; } + +Value BlockData::getSource() const { return this->source; } + +MemAccType BlockData::getMemAccType() const { return this->memAccTy; }; + +MemAccType &BlockData::getMemAccTypeRef() { return this->memAccTy; }; + +bool BlockData::isScalar() const { return !(this->scalar).isNull(); } + +bool BlockData::isEmpty() const { + return !(this->getRank() || this->source || !(this->scalar).isNull()); +} + +bool BlockData::hasSource() const { return this->source != nullptr; } + +void BlockData::removeSource() { this->source = nullptr; }; + +bool BlockData::hasResElemTy() const { return this->resElemTy != nullptr; } + +Type &BlockData::getResElemTyRef() { return this->resElemTy; } + +Type BlockData::getResElemTy() const { return this->resElemTy; } + +int64_t BlockData::getRank() const { + assert(offsets.size() == sizes.size() && offsets.size() == strides.size()); + return this->offsets.size(); +} + +void BlockData::setResElemTy(const Type &Ty) { this->resElemTy = Ty; } + +void BlockData::setScalar(const OpFoldResult &scalar) { this->scalar = scalar; } + +void BlockData::setSource(const Value &src) { this->source = src; } + +void BlockData::setOffsets(const SmallVector &offsets) { + this->offsets = offsets; +} + +void BlockData::setStrides(const SmallVector &strides) { + this->strides = strides; +} + +void BlockData::setSizes(const SmallVector &szs) { + this->sizes = szs; +} + +void BlockData::setMemAccTy(const MemAccType &v) { this->memAccTy = v; } + +void BlockData::setMemAccVal(const MemAccVal v) { this->memAccTy.value = v; } + +OpFoldResult BlockData::inferBlockOffset(const Location &loc, + OpBuilder &builder) const { + OpFoldResult retOffset = builder.getIndexAttr(0); + for (auto ofr : offsets) { + retOffset = addOpFoldResult(retOffset, ofr, loc, builder); + } + return retOffset; +} + +MemRefType BlockData::getResultMemrefType(int64_t offset, + ArrayRef resultShape) const { + SmallVector staticStrides; + SmallVector dynamicStrides; + dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides); + + auto baseMemrefType = dyn_cast(this->source.getType()); + assert(baseMemrefType && + "Invalid element type. It should be a base memref type."); + auto elementType = baseMemrefType.getElementType(); + auto layout = + StridedLayoutAttr::get(this->source.getContext(), offset, staticStrides); + return MemRefType::get(resultShape, elementType, layout); +} + +void BlockData::addBlock(BlockData &lBlock, BlockData &rBlock, Location loc, + ConversionPatternRewriter &rewriter) { + assert(this->isEmpty() && lBlock.getRank() == rBlock.getRank()); + // When both left block and right block have source, it is indirect load. + assert(!(lBlock.hasSource() && rBlock.hasSource()) && + "Don't support each BlockData has own base source pointer"); + this->source = + lBlock.hasSource() ? lBlock.getSourceRef() : rBlock.getSourceRef(); + + assert(!(lBlock.hasResElemTy() && rBlock.hasResElemTy())); + if (lBlock.hasResElemTy()) { + assert(lBlock.hasSource()); + this->resElemTy = lBlock.getResElemTyRef(); + } else if (rBlock.hasResElemTy()) { + assert(rBlock.hasSource()); + this->resElemTy = rBlock.getResElemTyRef(); + } + + // Acctually `scalar` should be accumulated into `offset` and `stride` finally + // In addBlock, just pass `scalar` when: + // 1. both lhs and rhs have `scalar` + // 2. otherwise, both lhs and rhs are scalar type with rank 0 + // Except above, original `scalar` has been fused into `offset` under add. + if (lBlock.isScalar() && rBlock.isScalar()) { + auto addScalar = addOpFoldResult(lBlock.getScalarRef(), + rBlock.getScalarRef(), loc, rewriter); + this->scalar = addScalar; + } else if (lBlock.getRank() == 0) { + // When both lhs and rhs are scalar type with rank 0, just try passing + // potential `scalar` + this->scalar = + lBlock.isScalar() ? lBlock.getScalarRef() : rBlock.getScalarRef(); + } + + for (const auto &[lOffset, rOffset] : + llvm::zip(lBlock.getOffsetsRef(), rBlock.getOffsetsRef())) { + this->offsets.push_back(addOpFoldResult(lOffset, rOffset, loc, rewriter)); + } + + for (const auto &[lStride, rStride] : + llvm::zip(lBlock.getStridesRef(), rBlock.getStridesRef())) { + this->strides.push_back(addOpFoldResult(lStride, rStride, loc, rewriter)); + } + + // Both sizes are same implicitly under `add` + this->sizes = lBlock.getSizesRef(); + + this->getMemAccTypeRef().merge(lBlock.getMemAccTypeRef()); + this->getMemAccTypeRef().merge(rBlock.getMemAccTypeRef()); + // this->setMemAccTy(selectMaxMemAccTy(lBlock.getMemAccType(), + // rBlock.getMemAccType())); +} + +void BlockData::subBlock(BlockData &lBlock, BlockData &rBlock, Location loc, + ConversionPatternRewriter &rewriter) { + assert(this->isEmpty() && lBlock.getRank() == rBlock.getRank()); + + if (lBlock.isScalar() && rBlock.isScalar()) { + auto subScalar = subOpFoldResult(lBlock.getScalarRef(), + rBlock.getScalarRef(), loc, rewriter); + this->scalar = subScalar; + } else if (lBlock.getRank() == 0) { + // When both lhs and rhs are scalar type with rank 0, just try passing + // potential `scalar` + this->scalar = + lBlock.isScalar() ? lBlock.getScalarRef() : rBlock.getScalarRef(); + } + + for (const auto &[lOffset, rOffset] : + llvm::zip(lBlock.getOffsetsRef(), rBlock.getOffsetsRef())) { + this->offsets.push_back(subOpFoldResult(lOffset, rOffset, loc, rewriter)); + } + + for (const auto &[lStride, rStride] : + llvm::zip(lBlock.getStridesRef(), rBlock.getStridesRef())) { + this->strides.push_back(subOpFoldResult(lStride, rStride, loc, rewriter)); + } + + // Both sizes are same implicitly under `sub` + this->sizes = lBlock.getSizesRef(); + + this->getMemAccTypeRef().merge(lBlock.getMemAccTypeRef()); + this->getMemAccTypeRef().merge(rBlock.getMemAccTypeRef()); + // this->setMemAccTy(selectMaxMemAccTy(lBlock.getMemAccType(), + // rBlock.getMemAccType())); +} + +void BlockData::mulBlock(BlockData &lBlock, BlockData &rBlock, Location loc, + ConversionPatternRewriter &rewriter) { + assert(this->isEmpty() && lBlock.getRank() == rBlock.getRank()); + + assert(!(lBlock.hasSource() && rBlock.hasSource())); + + if (lBlock.isScalar() && rBlock.isScalar()) { + LLVM_DEBUG({ + llvm::dbgs() << "lBlock.scalar:" << lBlock.getScalar() + << " rBlbock.scalar:" << rBlock.getScalar() << "\n"; + }); + + auto scalar = + mulOpFoldResult(lBlock.getScalar(), rBlock.getScalar(), loc, rewriter); + this->scalar = scalar; + } + + // assert( + // (lBlock.isScalar() ^ rBlock.isScalar()) && + // "Currently only support one and only one scalar in function + // mulBlock()"); + + BlockData *lb = &lBlock; + BlockData *rb = &rBlock; + if (lb->isScalar()) { + std::swap(lb, rb); + } + + // In mulBlock, `scalar` will be accumulated into `offset` and `stride` + OpFoldResult rScalar = rb->getScalarRef(); + for (const auto &lOffset : lb->getOffsetsRef()) { + this->offsets.push_back(mulOpFoldResult(lOffset, rScalar, loc, rewriter)); + } + + for (const auto &lStride : lb->getStridesRef()) { + this->strides.push_back(mulOpFoldResult(lStride, rScalar, loc, rewriter)); + } + + this->sizes = lb->getSizesRef(); + + this->getMemAccTypeRef().merge(lBlock.getMemAccTypeRef()); + this->getMemAccTypeRef().merge(rBlock.getMemAccTypeRef()); + // this->setMemAccTy(selectMaxMemAccTy(lBlock.getMemAccType(), + // rBlock.getMemAccType())); +} + +void BlockData::divBlock(BlockData &lBlock, BlockData &rBlock, Location loc, + ConversionPatternRewriter &rewriter) { + assert(this->isEmpty() && lBlock.getRank() == rBlock.getRank()); + + assert(!(lBlock.hasSource() && rBlock.hasSource())); + assert(lBlock.isScalar() && rBlock.isScalar()); + + auto rScalar = rBlock.getScalar(); + this->scalar = divOpFoldResult(lBlock.getScalar(), rScalar, loc, rewriter); + + for (auto lOffset : lBlock.getOffsetsRef()) { + this->offsets.push_back(divOpFoldResult(lOffset, rScalar, loc, rewriter)); + } + + for (auto lStride : lBlock.getStridesRef()) { + this->strides.push_back(divOpFoldResult(lStride, rScalar, loc, rewriter)); + } + + this->sizes = lBlock.getSizesRef(); + + this->getMemAccTypeRef().merge(lBlock.getMemAccTypeRef()); + this->getMemAccTypeRef().merge(rBlock.getMemAccTypeRef()); + // this->setMemAccTy(selectMaxMemAccTy(lBlock.getMemAccType(), + // rBlock.getMemAccType())); +} + +memref::ReinterpretCastOp BlockData::createCastOp(ArrayRef resultShape, + const Location &loc, + OpBuilder &builder) const { + OpFoldResult resOffset = this->inferBlockOffset(loc, builder); + auto resultType = this->getResultMemrefType( + isa(resOffset) ? getConstantIntValue(resOffset).value() + : ShapedType::kDynamic, + resultShape); + + SmallVector strides(this->strides); + for (size_t i = 0; i < strides.size(); i++) { + if (resultShape[i] == 1) { + if (auto strideValue = dyn_cast(strides[i])) { + auto oneIdx = + builder.create(loc, builder.getIndexAttr(1)); + strides[i] = builder.create(loc, strideValue, oneIdx) + .getResult(); + } + } + } + + return builder.create( + loc, resultType, this->source, resOffset, this->sizes, strides); +} + +void BlockData::dump() const { + llvm::outs() << "[INFO][BEG] BlockData info\n"; + llvm::outs() << "offsets has " << offsets.size() << " items\n"; + int cnt = 0; + for (auto it = offsets.begin(); it != offsets.end(); ++it) { + llvm::outs() << "offsets[" << cnt++ << "] = " << *it << "\n"; + } + llvm::outs() << "sizes has " << sizes.size() << " items\n"; + cnt = 0; + for (auto it = sizes.begin(); it != sizes.end(); ++it) { + llvm::outs() << "sizes[" << cnt++ << "] = " << *it << "\n"; + } + llvm::outs() << "strides has " << strides.size() << " items\n"; + cnt = 0; + for (auto it = strides.begin(); it != strides.end(); ++it) { + llvm::outs() << "strides[" << cnt++ << "] = " << *it << "\n"; + } + llvm::outs() << "source = " << source << "\n"; + llvm::outs() << "scalar = " << scalar << "\n"; + llvm::outs() << "resElemTy = " << resElemTy << "\n"; + llvm::outs() << "memAccTy = " << memAccTy.toString() << "\n"; + llvm::outs() << "[INFO][END] BlockData info\n"; +} + +Value BlockDataParser::getScalarMemRef(Value ptr, Value memref, + const Location &loc, + ConversionPatternRewriter &rewriter) { + assert(isa(ptr.getType()) && "expect a scalar pointer"); + if (ptr.getDefiningOp()) { + if (auto castOp = memref.getDefiningOp()) { + return castOp.getResult(); + } else { + llvm_unreachable("pointer value is defined by an unexpected op"); + } + } + + assert(isa(ptr) && + "pointer should be produced by addptr or block argument"); + BlockData data; + data.setSource(memref); + data.getOffsetsRef().push_back(rewriter.getIndexAttr(0)); + data.getSizesRef().push_back(rewriter.getIndexAttr(1)); + data.getStridesRef().push_back(rewriter.getIndexAttr(1)); + auto castOp = data.createCastOp(SmallVector(1, 1), loc, rewriter); + return castOp.getResult(); +} + +void BlockDataParser::parse( + Value operand, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + if (known.find(operand) != known.end()) { + return data = known.lookup(operand), void(); + } + + if (isa(operand.getType())) { + data.setScalar(getOpFoldResultOfLayoutInfo(operand, rewriter)); + return; + } + + // + if (isa(operand.getType())) { + // Just consider two state: ptr and ptr> + auto remappedPtr = rewriter.getRemappedValue(operand); + assert(remappedPtr); + if (auto op = operand.getDefiningOp()) { + if (auto addPtrOp = dyn_cast(op)) { + parseAddPtr(addPtrOp, data, loc, rewriter, known); + } else if (auto bitcastOp = dyn_cast(op)) { + parseBitcast(bitcastOp, data, loc, rewriter, known); + } else if (auto makeTensorPtrOp = dyn_cast(op)) { + parseTensorPtr(makeTensorPtrOp, data, loc, rewriter, known); + } else if (auto advanceOp = dyn_cast(op)) { + // To support + // ptr_0 = tl.advance(ptr) + // ptr_1 = tl.advance(ptr_0) + parseTensorPtr(advanceOp, data, loc, rewriter, known); + } else if (auto intToPtrOp = dyn_cast(op)) { + data.setSource(remappedPtr); + } else if (auto customOp = dyn_cast(op)) { + data.setSource(remappedPtr); + } else { + LLVM_DEBUG({ llvm::dbgs() << operand << "\n"; }); + llvm_unreachable("Unexpected operand defining operation, a scalar " + "pointer can only be produced by AddPtrOp or direct " + "block ptr or hivm CustomOp"); + } + } else { + data.setSource(remappedPtr); + } + return; + } + + // not a scalar pointer + if (auto addOp = operand.getDefiningOp()) { + parseAdd(addOp, data, loc, rewriter, known); + } else if (auto subOp = operand.getDefiningOp()) { + parseSub(subOp, data, loc, rewriter, known); + } else if (auto mulOp = operand.getDefiningOp()) { + parseMul(mulOp, data, loc, rewriter, known); + } else if (auto addPtrOp = operand.getDefiningOp()) { + parseAddPtr(addPtrOp, data, loc, rewriter, known); + } else if (auto constOp = operand.getDefiningOp()) { + parseConstSplat(constOp, data, loc, rewriter, known); + } else if (auto broadcastOp = operand.getDefiningOp()) { + parseBroadcast(broadcastOp, data, loc, rewriter, known); + } else if (auto splatOp = operand.getDefiningOp()) { + parseSplat(splatOp, data, loc, rewriter, known); + } else if (auto expandDimsOp = + operand.getDefiningOp()) { + parseExpandDims(expandDimsOp, data, loc, rewriter, known); + } else if (auto remOp = operand.getDefiningOp()) { + parseRem(remOp, data, loc, rewriter, known); + } else if (auto bitcastOp = operand.getDefiningOp()) { + parseBitcast(bitcastOp, data, loc, rewriter, known); + } else if (auto extsiOp = operand.getDefiningOp()) { + parseExtSI(extsiOp, data, loc, rewriter, known); + } else if (auto divOp = operand.getDefiningOp()) { + parseDiv(divOp, data, loc, rewriter, known); + } else if (auto makeRangeOp = operand.getDefiningOp()) { + parseMakeRange(makeRangeOp, data, loc, rewriter, known); + } else if (auto reduceOp = operand.getDefiningOp()) { + parseReduce(reduceOp, data, loc, rewriter, known); + } else if (auto loadOp = operand.getDefiningOp()) { + parseIndirectLoad(loadOp, data, loc, rewriter, known); + } else if (auto castOp = operand.getDefiningOp()) { + parseIndirectLoad(castOp, data, loc, rewriter, known); + } else if (auto extractSliceOp = + operand.getDefiningOp()) { + parseExtractSlice(extractSliceOp, data, loc, rewriter, known); + } else if (auto forOp = operand.getDefiningOp()) { + auto opResult = dyn_cast(operand); + assert(opResult && "expected OpResult for scf.for result"); + unsigned resultIdx = opResult.getResultNumber(); + parseIndirectLoad(forOp, data, loc, rewriter, known, resultIdx); + } else if (auto tensorCastOp = operand.getDefiningOp()) { + // Used for identity operation. + parse(tensorCastOp.getSource(), data, loc, rewriter, known); + } else if (auto fillOp = operand.getDefiningOp()) { + parseFill(fillOp, data, loc, rewriter, known); + } else if (auto selectOp = operand.getDefiningOp()) { + parseSelect(selectOp, data, loc, rewriter, known); + } else if (auto customOp = operand.getDefiningOp()) { + auto opResult = dyn_cast(operand); + assert(opResult && "Expected operand to be an OpResult"); + unsigned resultIdx = opResult.getResultNumber(); + parseCustomOp(customOp, data, loc, rewriter, known, resultIdx); + } else if (auto genericOp = operand.getDefiningOp()) { + if (genericOp->hasAttr("tt.from_make_range")) { + parseLinalgGenericFromMakeRange(genericOp, data, loc, rewriter, known); + } else { + operand.dump(); + llvm_unreachable( + "encountered AddPtrOp produced by unsupported operation"); + } + } else if (auto atomicRMWOp = operand.getDefiningOp()) { + parseAtomicRmw(atomicRMWOp, data, loc, rewriter, known); + } else { + operand.dump(); + llvm_unreachable("encountered AddPtrOp produced by unsupported operation"); + } +} + +void BlockDataParser::parseAtomicRmw( + triton::AtomicRMWOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + auto opRes = op->getResult(0); + auto opResTy = opRes.getType(); + std::vector resShape; + if (auto shapedResTy = dyn_cast(opResTy)) { + resShape = shapedResTy.getShape().vec(); + if (resShape.size() == 1 && resShape[0] == 1) { + Value zeroIdx = rewriter.create(loc, 0); + Value extracted = + rewriter.create(loc, opRes, ValueRange{zeroIdx}); + Value scalarIdx = rewriter.create( + loc, rewriter.getIndexType(), extracted); + data.setMemAccVal(MemAccVal::StrucMemAcc); + data.setScalar(scalarIdx); + data.getSizesRef().push_back(rewriter.getIndexAttr(1)); + data.getStridesRef().push_back(rewriter.getIndexAttr(0)); + data.getOffsetsRef().push_back(scalarIdx); + return; + } + // For now, we consider this is UnstrucMemAcc because we have no other info. + // Visiting other ops may change the type due to more info. + data.setMemAccVal(MemAccVal::UnstrucMemAcc); + } else { + data.setMemAccVal(MemAccVal::StrucMemAcc); + resShape.push_back(1); + } + for (auto &s : resShape) { + data.getOffsetsRef().push_back(rewriter.getIndexAttr(0)); + data.getSizesRef().push_back(rewriter.getIndexAttr(s)); + data.getStridesRef().push_back(rewriter.getIndexAttr(1)); + } + data.setSource(opRes); +} + +void BlockDataParser::parseAdd( + arith::AddIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + BlockData lBlock, rBlock; + parse(op.getLhs(), lBlock, loc, rewriter, known); + parse(op.getRhs(), rBlock, loc, rewriter, known); + data.addBlock(lBlock, rBlock, loc, rewriter); +} + +void BlockDataParser::parseSub( + arith::SubIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + BlockData lBlock, rBlock; + parse(op.getLhs(), lBlock, loc, rewriter, known); + parse(op.getRhs(), rBlock, loc, rewriter, known); + data.subBlock(lBlock, rBlock, loc, rewriter); +} + +void BlockDataParser::parseMul( + arith::MulIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + BlockData lBlock, rBlock; + parse(op.getLhs(), lBlock, loc, rewriter, known); + parse(op.getRhs(), rBlock, loc, rewriter, known); + + data.mulBlock(lBlock, rBlock, loc, rewriter); +} + +void BlockDataParser::parseDiv( + arith::DivSIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + BlockData lBlock, rBlock; + parse(op.getLhs(), lBlock, loc, rewriter, known); + parse(op.getRhs(), rBlock, loc, rewriter, known); + data.divBlock(lBlock, rBlock, loc, rewriter); +} + +// TODO : support modulos +void BlockDataParser::parseRem( + arith::RemSIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(false && "Address expression with modulo is not supported yet, it " + "shall be analysis at linearize."); +} + +void BlockDataParser::parseMakeRange( + triton::MakeRangeOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + auto shape = dyn_cast(op.getType()).getShape(); + + auto start = op.getStart(); + auto end = op.getEnd(); + auto stride = (end >= start) && (end - start <= shape[0]); + assert(stride == 1 && + "make_range op should always return a tensor of stride 1"); + + data.getOffsetsRef().push_back(rewriter.getIndexAttr(start)); + data.getSizesRef().push_back(rewriter.getIndexAttr(shape[0])); + data.getStridesRef().push_back(rewriter.getIndexAttr(stride)); +} + +void BlockDataParser::parseLinalgGenericFromMakeRange( + linalg::GenericOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + assert(op->hasAttr("tt.from_make_range") && + "expected tt.from_make_range attribute"); + + auto offsetAttr = op->getAttr("tt.make_range_offset"); + auto sizeAttr = op->getAttr("tt.make_range_size"); + assert(offsetAttr && sizeAttr && + "tt.make_range_offset and tt.make_range_size required"); + + int64_t offset = cast(offsetAttr).getInt(); + int64_t size = cast(sizeAttr).getInt(); + + data.getOffsetsRef().push_back(rewriter.getIndexAttr(offset)); + data.getSizesRef().push_back(rewriter.getIndexAttr(size)); + data.getStridesRef().push_back(rewriter.getIndexAttr(1)); +} + +void BlockDataParser::parseExpandDims( + triton::ExpandDimsOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + + parse(op.getSrcMutable().get(), data, loc, rewriter, known); + auto resShape = dyn_cast(op.getResult().getType()).getShape(); + auto axis = op.getAxis(); + + assert(resShape[axis] == 1 && + "The destiny shape of changed dimension should be 1"); + + data.getOffsetsRef().insert(data.getOffsetsRef().begin() + axis, + rewriter.getIndexAttr(0)); + data.getSizesRef().insert(data.getSizesRef().begin() + axis, + rewriter.getIndexAttr(1)); + data.getStridesRef().insert(data.getStridesRef().begin() + axis, + rewriter.getIndexAttr(0)); +} + +void BlockDataParser::parseExtractSlice( + tensor::ExtractSliceOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + const std::string scenarioMessages = + "PtsAnalysis supports indirectly block load in the " + "following scenario\n" + "B = tl.load(Aptr + Aoffset) # B is 1D tensor\n" + "s = tl.extract_slice(indices, offsets= (i,), sizes= " + "(1,), strides= (1,)) # s is a tensor<1x$dtype>\n" + "D = tl.load(Cptr + s + Coffset) # s is used as the " + "scalar offset\n"; // tensor<2x$dtype> will be support + // soon + + auto extract_src = op->getOperand(0); + BlockData srcBlock; + parse(extract_src, srcBlock, loc, rewriter, known); + if (!srcBlock.hasSource()) { + llvm_unreachable(scenarioMessages.c_str()); + } + // Use isa_and_nonnull for LLVM 21 compatibility + if (!isa_and_nonnull(srcBlock.getSource().getDefiningOp())) { + llvm_unreachable(scenarioMessages.c_str()); + } + + auto extract_result = op->getResult(0); + auto shaped_ty = dyn_cast(extract_result.getType()); + auto shape = shaped_ty.getShape(); + if (shape.size() > 1 || shape[0] > 1) { + llvm_unreachable(scenarioMessages.c_str()); + } + auto castOp = rewriter.create( + loc, RankedTensorType::get(shape, rewriter.getIndexType()), + extract_result); + auto offset = castOp.getResult(); + if (data.isEmpty()) { + data.getOffsetsRef().push_back(offset); + data.getSizesRef().push_back(rewriter.getIndexAttr(shape[0])); + data.getStridesRef().push_back(rewriter.getIndexAttr(1)); + } else { + llvm_unreachable( + "parseExtractSlice with offset already setup not yet supported"); + } +} + +void BlockDataParser::parseBitcast( + triton::BitcastOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + parse(op.getSrc(), data, loc, rewriter, known); + + auto resType = op.getResult().getType(); + Type resElemPointeeTy = nullptr; + if (auto resShapedTy = dyn_cast(resType)) { + auto resElemTy = resShapedTy.getElementType(); + resElemPointeeTy = + dyn_cast(resElemTy).getPointeeType(); + } else { + auto srcPointeeType = + cast(op.getSrc().getType()).getPointeeType(); + auto resPointeeType = cast(resType).getPointeeType(); + + // Handling special case + // If Op is MetaUse or src is i1 block argument and dst is i8, + // it should be converted to UnrealizedConversionCast + if (op->hasAttr("MetaUse") || + (isa(op.getSrc()) && + srcPointeeType == rewriter.getIntegerType(1) && + resPointeeType == rewriter.getIntegerType(8))) { + resElemPointeeTy = resPointeeType; + } else { + auto remappedValue = rewriter.getRemappedValue(op); + data.setSource(remappedValue); + LLVM_DEBUG({ + llvm::dbgs() << "Remapping bitcastOp:\n"; + llvm::dbgs() << op << "\nto \n"; + llvm::dbgs() << remappedValue << "\n"; + }); + } + } + data.setResElemTy(resElemPointeeTy); +} + +void BlockDataParser::parseExtSI( + arith::ExtSIOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + parse(op.getIn(), data, loc, rewriter, known); +} + +void BlockDataParser::parseBroadcast( + triton::BroadcastOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + + auto src = op.getSrcMutable().get(); + auto dst = op.getResult(); + assert(isa(src.getType()) && + "tt.broadcast's input should be a tensor"); + + auto srcShape = dyn_cast(src.getType()).getShape(); + auto dstShape = dyn_cast(dst.getType()).getShape(); + assert(srcShape.size() == dstShape.size() && + "rank of source shoule be equal to destnation"); + + parse(src, data, loc, rewriter, known); + + for (const auto &[idx, src_dst] : + llvm::enumerate(llvm::zip(srcShape, dstShape))) { + const auto &[srcAxis, dstAxis] = src_dst; + if (srcAxis == dstAxis) { + continue; + } + assert(srcAxis < dstAxis && + "srcShape of broadcastOp must be less than dstShape."); + data.getSizesRef()[idx] = rewriter.getIndexAttr(dstAxis); + } +} + +void BlockDataParser::parseSplat( + triton::SplatOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + auto src = op.getSrc(); + auto dst = op.getResult(); + auto dstShape = dyn_cast(dst.getType()).getShape(); + + parse(src, data, loc, rewriter, known); + + if (isa(src.getType()) || + isa(src.getType())) { + if (!data.isEmpty()) { + data.getOffsetsRef().clear(); + data.getSizesRef().clear(); + data.getStridesRef().clear(); + } + for (auto dstAxis : dstShape) { + data.getOffsetsRef().push_back(rewriter.getIndexAttr(0)); + data.getSizesRef().push_back(rewriter.getIndexAttr(dstAxis)); + data.getStridesRef().push_back(rewriter.getIndexAttr(0)); + } + } else { + op->emitError("Block data Analysis: unsupported splat pattern"); + return; + } + if (data.isScalar()) { + data.getOffsetsRef()[0] = data.getScalarRef(); + } +} + +void BlockDataParser::parseConstSplat( + arith::ConstantOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + + DenseElementsAttr denseAttr = dyn_cast(op.getValue()); + assert(denseAttr && denseAttr.isSplat() && + isa(denseAttr.getElementType())); + + auto innerVal = denseAttr.getValues()[0].getValue(); + auto innerValIndexAttr = rewriter.getIndexAttr(innerVal.getSExtValue()); + + // for mul state + data.setScalar(innerValIndexAttr); + + auto resType = dyn_cast(op.getResult().getType()); + size_t loopLimit = resType.getShape().size(); + for (auto i = 0; i < loopLimit; i++) { + // Add original dense val to first dim offset for add state + if (i == 0) { + data.getOffsetsRef().push_back(innerValIndexAttr); + } else { + data.getOffsetsRef().push_back(rewriter.getIndexAttr(0)); + } + data.getSizesRef().push_back(rewriter.getIndexAttr(resType.getShape()[i])); + data.getStridesRef().push_back(rewriter.getIndexAttr(0)); + } +} + +template +std::enable_if_t || + std::is_same_v> +BlockDataParser::parseTensorPtr( + T op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + + Value remappedValue = rewriter.getRemappedValue(op); + if (auto castOp = remappedValue.getDefiningOp()) { + parseReinterpretCast(castOp, data, loc, rewriter, known); + } else { + llvm_unreachable("the value should be mapped to memref.reinterpret_cast"); + } +} + +void BlockDataParser::parseAddPtr( + triton::AddPtrOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + + BlockData ptrBlock, offsetBlock; + parse(op.getPtr(), ptrBlock, op.getLoc(), rewriter, known); + parse(op.getOffset(), offsetBlock, op.getLoc(), rewriter, known); + + assert(ptrBlock.hasSource() && + "Ptr field should provide source/base pointer"); + // offset has source means offset is from tl.load and other ops(TODO) + if (offsetBlock.hasSource()) { + ptrBlock.setMemAccTy(offsetBlock.getMemAccType()); + offsetBlock.removeSource(); + } + + // handle for loop & scalar + if (ptrBlock.getRank() == 1 && offsetBlock.getRank() == 0) { + offsetBlock.getSizesRef().push_back(rewriter.getIndexAttr(1)); + offsetBlock.getOffsetsRef().push_back(offsetBlock.getScalarRef()); + offsetBlock.getStridesRef().push_back(rewriter.getIndexAttr(0)); + } + + assert(ptrBlock.getRank() == offsetBlock.getRank() && + "ptr and offset should have same rank"); + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "[parseAddPtr][BEG] =========================\n"; + os << "[parseAddPtr] op is " << op << "\n"; + for (int i = 0; i < ptrBlock.getRank(); i++) { + os << "ptrBlock.getOffsetsRef()[" << i + << "] = " << ptrBlock.getOffsetsRef()[i] << "\n"; + os << "ptrBlock.getSizesRef()[" << i + << "] = " << ptrBlock.getSizesRef()[i] << "\n"; + os << "ptrBlock.getStridesRef()[" << i + << "] = " << ptrBlock.getStridesRef()[i] << "\n"; + os << "offsetBlock.getOffsetsRef()[" << i + << "] = " << offsetBlock.getOffsetsRef()[i] << "\n"; + os << "offsetBlock.getSizesRef()[" << i + << "] = " << offsetBlock.getSizesRef()[i] << "\n"; + os << "offsetBlock.getStridesRef()[" << i + << "] = " << offsetBlock.getStridesRef()[i] << "\n"; + } + os << "[parseAddPtr][END] -------------------------\n"; + }); + data.addBlock(ptrBlock, offsetBlock, op.getLoc(), rewriter); +} + +void BlockDataParser::parseReinterpretCast( + memref::ReinterpretCastOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + + data.setOffsets(op.getMixedOffsets()); + data.setSizes(op.getMixedSizes()); + data.setStrides(op.getMixedStrides()); + data.setSource(op.getSource()); + + // In memref::ReinterpretCastOp, offset means the total of collapsing multiple + // dimensions, which corresponds to first dim offset in block data. + // Here populate the rest of the dimensions with zeroes. + assert(data.getOffsetsRef().size() == 1); + size_t loopLimit = data.getSizesRef().size(); + for (size_t i = 1; i < loopLimit; i++) { + data.getOffsetsRef().push_back(rewriter.getIndexAttr(0)); + } +} + +void BlockDataParser::parseReduce( + triton::ReduceOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + + const std::string scenarioMessages = + "PtsAnalysis supports indirectly block load in the following scenario\n" + "B = tl.load(Aptr + Aoffset) # B is 1D tensor\n" + "s = tl.min(B) # s is a scalar\n" + "D = tl.load(Cptr + s + Coffset) # s is used as the scalar offset\n"; + + auto reduce_src = op->getOperand(0); + BlockData srcBlock; + parse(reduce_src, srcBlock, loc, rewriter, known); + if (!srcBlock.hasSource()) { + llvm_unreachable(scenarioMessages.c_str()); + } + // Use isa_and_nonnull for LLVM 21 compatibility + if (!isa_and_nonnull(srcBlock.getSource().getDefiningOp())) { + llvm_unreachable(scenarioMessages.c_str()); + } + + auto reduce_result = op->getResult(0); + auto shaped_ty = dyn_cast(reduce_result.getType()); + auto shape = shaped_ty.getShape(); + auto ops = llvm::map_to_vector(op.getBody()->without_terminator(), + [](Operation &op) { return &op; }); + // Support only the case: scalar = tl.load(1D tensor) + if (shape.size() != 1 || op.getAxis() != 0 || ops.size() != 1 || + !isa(ops.front())) { + llvm_unreachable(scenarioMessages.c_str()); + } + + auto castOp = rewriter.create( + loc, RankedTensorType::get(shape, rewriter.getIndexType()), + reduce_result); + auto offset = castOp.getResult(); + if (data.isEmpty()) { + data.getOffsetsRef().push_back(offset); + data.getSizesRef().push_back(rewriter.getIndexAttr(shape[0])); + data.getStridesRef().push_back(rewriter.getIndexAttr(1)); + } else { + llvm_unreachable("parseReduce with offset already setup not yet supported"); + } +} + +template +void parseIndirectLoad(OpTy op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known, + unsigned resultIdx) { + assert(resultIdx < op->getNumResults() && + "resultIdx out of range for parseIndirectLoad"); + auto opRes = op->getResult(resultIdx); + auto opResTy = opRes.getType(); + std::vector resShape; + if (auto shapedResTy = dyn_cast(opResTy)) { + // For now, we consider this is UnstrucMemAcc because we have no other info. + // Visiting other ops may change the type due to more info. + resShape = shapedResTy.getShape().vec(); + auto numOperands = 3; + if (resShape.size() == 1 && resShape[0] == 1 && + op->getNumOperands() == numOperands) { + Value zeroIdx = rewriter.create(loc, 0); + Value extracted = + rewriter.create(loc, opRes, ValueRange{zeroIdx}); + Value scalarIdx = rewriter.create( + loc, rewriter.getIndexType(), extracted); + data.setMemAccVal(MemAccVal::StrucMemAcc); + data.setScalar(scalarIdx); + data.getSizesRef().push_back(rewriter.getIndexAttr(1)); + data.getStridesRef().push_back(rewriter.getIndexAttr(0)); + data.getOffsetsRef().push_back(scalarIdx); + return; + } + data.setMemAccVal(MemAccVal::UnstrucMemAcc); + } else { + // scalar load means this is used as offset. It is StrucMemAcc. + data.setMemAccVal(MemAccVal::StrucMemAcc); + resShape.push_back(1); + } + for (auto &s : resShape) { + data.getOffsetsRef().push_back(rewriter.getIndexAttr(0)); + data.getSizesRef().push_back(rewriter.getIndexAttr(s)); + data.getStridesRef().push_back(rewriter.getIndexAttr(1)); + } + // set the source in BlockData so that we know an indirect-load op exists in + // the chain. + data.setSource(opRes); +} + +void BlockDataParser::parseCustomOp( + hivm::CustomOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known, unsigned resultIdx) { + auto srcValArrayAttr = op->getAttrOfType( + ConverterUtils::customSrcPtrIndexAttrName); + assert(srcValArrayAttr && + "structure hivm.custom op should present src tensor"); + auto srcValArray = srcValArrayAttr.asArrayRef(); + assert(srcValArray[resultIdx] != -1 && + "tensor result should map to src tensor"); + parse(op->getOperand(srcValArray[resultIdx]), data, loc, rewriter, known); + data.setSource(rewriter.getRemappedValue(op->getResult(resultIdx))); +} + +void BlockDataParser::parseFill( + linalg::FillOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + auto src = op.getInputs()[0]; + auto dst = op.getResult(0); + auto dstShape = dyn_cast(dst.getType()).getShape(); + + parse(src, data, loc, rewriter, known); + + if (isa(src.getType())) { + if (!data.isEmpty()) { + data.getOffsetsRef().clear(); + data.getSizesRef().clear(); + data.getStridesRef().clear(); + } + for (auto dstAxis : dstShape) { + data.getOffsetsRef().push_back(rewriter.getIndexAttr(0)); + data.getSizesRef().push_back(rewriter.getIndexAttr(dstAxis)); + data.getStridesRef().push_back(rewriter.getIndexAttr(0)); + } + } else { + op->emitError("Block data Analysis: unsupported fillOp pattern"); + return; + } + if (data.isScalar()) { + data.getOffsetsRef()[0] = data.getScalarRef(); + } +} + +void BlockDataParser::parseSelect( + arith::SelectOp op, BlockData &data, const Location &loc, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + assert(data.isEmpty()); + + auto res = op.getResult(); + auto resType = dyn_cast(res.getType()); + assert(resType && "arith.select result should be a ShapedType"); + assert(isa(resType.getElementType()) || + isa(resType.getElementType())); + + OpFoldResult indexOfr; + size_t loopLimit = resType.getShape().size(); + + Value cond = op.getCondition(); + bool condIsScalarI1 = isa(cond.getType()) && + cast(cond.getType()).getWidth() == 1 && + !isa(cond.getType()); + + auto trueConst = + dyn_cast(op.getTrueValue().getDefiningOp()); + auto falseConst = + dyn_cast(op.getFalseValue().getDefiningOp()); + auto trueDense = trueConst ? dyn_cast(trueConst.getValue()) + : DenseElementsAttr(); + auto falseDense = falseConst + ? dyn_cast(falseConst.getValue()) + : DenseElementsAttr(); + + bool denseConstCase = condIsScalarI1 && trueDense && falseDense; + + if (denseConstCase) { + // if cond is scalar i1 and both true and false value are splat dense const, + // we can directly use the value of the dense const to create scalar select + // op. + Attribute trueFirst = *trueDense.value_begin(); + Attribute falseFirst = *falseDense.value_begin(); + + Value trueScalar = nullptr; + Value falseScalar = nullptr; + if (auto tInt = dyn_cast(trueFirst)) { + trueScalar = rewriter.create(loc, tInt).getResult(); + } else { + llvm_unreachable("unsupported true dense element attr in parseSelect"); + } + + if (auto fInt = dyn_cast(falseFirst)) { + falseScalar = rewriter.create(loc, fInt).getResult(); + } else { + llvm_unreachable("unsupported false dense element attr in parseSelect"); + } + + assert(trueScalar.getType() == falseScalar.getType() && + "scalarized true/false type mismatch"); + + auto scalarSelect = rewriter.create( + loc, trueScalar.getType(), cond, trueScalar, falseScalar); + + indexOfr = getOpFoldResultOfLayoutInfo(scalarSelect.getResult(), rewriter); + } else { + assert(llvm::all_of(resType.getShape(), + [](int64_t dim) { return dim == 1; }) && + "parseSelect currently supports all-ones shape unless cond=i1 with " + "dense constants"); + + SmallVector indices; + indices.reserve(loopLimit); + for (size_t i = 0; i < loopLimit; ++i) { + indices.push_back(rewriter.create(loc, 0)); + } + + auto extractOp = rewriter.create(loc, res, indices); + indexOfr = extractOp.getResult(); + if (isa(extractOp.getType())) { + indexOfr = getOpFoldResultOfLayoutInfo(extractOp.getResult(), rewriter); + } + } + + // Set scalar for mul state + data.setScalar(indexOfr); + + for (size_t i = 0; i < loopLimit; ++i) { + // Add scalar to first dim offset for add state + if (i == 0) { + data.getOffsetsRef().push_back(indexOfr); + } else { + data.getOffsetsRef().push_back(rewriter.getIndexAttr(0)); + } + data.getSizesRef().push_back(rewriter.getIndexAttr(resType.getShape()[i])); + data.getStridesRef().push_back(rewriter.getIndexAttr(0)); + } +} + +void BlockDataParser::rewriteAddPtr( + triton::AddPtrOp op, triton::AddPtrOp::Adaptor &adaptor, + ConversionPatternRewriter &rewriter, + llvm::SmallDenseMap &known) { + auto insertPoint = rewriter.saveInsertionPoint(); + rewriter.setInsertionPoint(op); + + BlockData data; + parseAddPtr(op, data, op.getLoc(), rewriter, known); + + if (auto src = data.getSource(); + data.getMemAccTypeRef().isUnstructured() && + !(src && isa_and_nonnull(src.getDefiningOp()))) { + // TODO: Based on more info, try to create a performant IR + rewriteAddPtrToUnstrucMemAcc(op, adaptor, rewriter, data); + LLVM_DEBUG({ llvm::dbgs() << *getModuleOpFromOperation(op) << "\n"; }); + return; + } + + if (data.getSizesRef().size() == 0) { + data.getSizesRef().push_back(rewriter.getIndexAttr(1)); + data.getStridesRef().push_back(rewriter.getIndexAttr(0)); + data.getOffsetsRef().push_back(data.getScalarRef()); + } + + ArrayRef resultShape; + // shape {1,} is stub for single ptr + SmallVector stubScalarTypeShape(1, 1); + if (auto shapedType = dyn_cast(op.getResult().getType())) { + resultShape = shapedType.getShape(); + } else { + assert(data.getRank() == 1); + resultShape = stubScalarTypeShape; + } + + known[op.getResult()] = data; + + // If there are dimensions with size 1 and stride 0, replace 0 stride with the + // product of sizes of all lower dimensions. This avoids creating memref with + // zero stride. + // And here store the unmodified state into known ptrs, since any following + // pointer arithmetic operations should still use the original 0 stride. + auto inferedSize = 1; + auto hoistDim = op->getAttrOfType("hoist_dim"); + for (int i = data.getSizesRef().size() - 1; i >= 0; i--) { + auto strideConst = getConstantIntValue(data.getStridesRef()[i]); + auto sizeConst = getConstantIntValue(data.getSizesRef()[i]); + assert(sizeConst.has_value()); + bool shouldReplaceStride = + (sizeConst.value() == 1) || (hoistDim && hoistDim.getValue() == i); + if (shouldReplaceStride && strideConst && strideConst.value() == 0) { + data.getStridesRef()[i] = rewriter.getIndexAttr(inferedSize); + } + inferedSize *= sizeConst.value(); + } + + // Use dyn_cast_or_null to safely handle nullptr from getDefiningOp() + // This is necessary for LLVM 21 compatibility where dyn_cast asserts on + // nullptr + if (auto intToPtrOp = dyn_cast_or_null( + data.getSourceRef().getDefiningOp())) { + auto rtype = cast(intToPtrOp.getResult().getType()); + auto memrefType = + MemRefType::get({ShapedType::kDynamic}, rtype.getPointeeType()); + auto hivmPointCastOp = rewriter.create( + intToPtrOp.getLoc(), memrefType, ValueRange{intToPtrOp.getSrc()}); + data.setSource(hivmPointCastOp.getResult()); + } + + if (data.hasResElemTy()) { + // Handle bitcast scenario + auto memrefType = dyn_cast(data.getSourceRef().getType()) + .cloneWith(std::nullopt, data.getResElemTyRef()); + UnrealizedConversionCastOp castOp = + rewriter.create( + op.getLoc(), memrefType, data.getSourceRef()); + data.setSource(castOp.getOutputs()[0]); + } + + // ToDo: need to handle module scenario + + memref::ReinterpretCastOp castOp = + data.createCastOp(resultShape, op.getLoc(), rewriter); + Value src = castOp.getResult(); + LLVM_DEBUG({ + llvm::dbgs() << "cast MemRefType:\n"; + castOp.getOperation()->print(llvm::dbgs(), + OpPrintingFlags().printGenericOpForm()); + llvm::dbgs() << "\n"; + }); + + rewriter.replaceOp(op, src); + rewriter.restoreInsertionPoint(insertPoint); +} + +OpFoldResult +accumulatePotentialOffsetOnBase(triton::MakeTensorPtrOp op, Value base, + OpFoldResult offset, + ConversionPatternRewriter &rewriter) { + if (auto baseRecast = base.getDefiningOp()) { + assert(isa(op.getBase().getDefiningOp()) && + "base of MakeTensorPtrOp only comes from native ptr or AddPtrOp"); + + return addOpFoldResult(offset, baseRecast.getConstifiedMixedOffset(), + op.getLoc(), rewriter); + } + + return offset; +} + +void BlockDataParser::rewriteCustomOp( + hivm::CustomOp op, hivm::CustomOp::Adaptor &adaptor, + ConversionPatternRewriter &rewriter, + const llvm::SmallDenseMap &known) { + auto ip = rewriter.saveInsertionPoint(); + rewriter.setInsertionPoint(op); + auto loc = op.getLoc(); + llvm::SmallVector newInputs; + llvm::SmallVector newOutputs; + auto convertIntToPtr = [&rewriter](BlockData &data) { + if (auto intToPtrOp = dyn_cast_or_null( + data.getSourceRef().getDefiningOp())) { + auto rtype = cast(intToPtrOp.getResult().getType()); + auto memrefType = + MemRefType::get({ShapedType::kDynamic}, rtype.getPointeeType()); + auto hivmPointCastOp = rewriter.create( + intToPtrOp.getLoc(), memrefType, ValueRange{intToPtrOp.getSrc()}); + if (data.getSizesRef().size() == 0) { + data.getSizesRef().push_back(rewriter.getIndexAttr(1)); + if (data.getScalarRef().isNull()) { + data.getOffsetsRef().push_back(rewriter.getIndexAttr(0)); + } else { + data.getOffsetsRef().push_back(data.getScalarRef()); + } + data.getStridesRef().push_back(rewriter.getIndexAttr(1)); + } + data.setSource(hivmPointCastOp.getResult()); + } + }; + for (auto in : op.getInputs()) { + in = rewriter.getRemappedValue(in); + BlockData blockData; + auto curInput = in; + if (llvm::isa(in.getType())) { + parse(in, blockData, loc, rewriter, known); + convertIntToPtr(blockData); + curInput = blockData.createCastOp({ShapedType::kDynamic}, loc, rewriter); + } else if (auto tensor = llvm::dyn_cast(in.getType())) { + if (llvm::isa(tensor.getElementType())) { + parse(in, blockData, loc, rewriter, known); + convertIntToPtr(blockData); + curInput = blockData.createCastOp(tensor.getShape(), loc, rewriter); + } + } + newInputs.emplace_back(curInput); + } + for (auto out : op.getOutputs()) { + auto tensorTy = llvm::cast(out.getType()); + if (llvm::isa(tensorTy.getElementType())) { + // simd library shouldn't output tensor + // after rewrite, delete the tensor output value + continue; + } + newOutputs.emplace_back(rewriter.getRemappedValue(out)); + } + llvm::SmallVector resultTypes; + for (auto ty : op->getResultTypes()) { + if (auto ptrTy = llvm::dyn_cast(ty)) { + resultTypes.emplace_back( + MemRefType::get({ShapedType::kDynamic}, ptrTy.getPointeeType())); + continue; + } + if (auto tensorTy = llvm::dyn_cast(ty)) { + if (auto ptrTy = + llvm::dyn_cast(tensorTy.getElementType())) { + resultTypes.emplace_back( + MemRefType::get(tensorTy.getShape(), ptrTy.getPointeeType())); + continue; + } + } + resultTypes.emplace_back(ty); + } + auto newCustomOp = + rewriter.create(loc, resultTypes, op.getName(), newInputs, + newOutputs, adaptor.getTempBuffers()); + auto operandSegmentSizesAttr = newCustomOp->getAttr("operandSegmentSizes"); + newCustomOp->setAttrs(op->getAttrs()); + newCustomOp->setAttr("operandSegmentSizes", operandSegmentSizesAttr); + rewriter.replaceOp(op, newCustomOp.getResults()); + rewriter.restoreInsertionPoint(ip); +} + +// Design for load/store boundary_check. +memref::ReinterpretCastOp createRedundantOp(triton::MakeTensorPtrOp op, + ConversionPatternRewriter &rewriter, + BlockData &data) { + auto loc = op.getLoc(); + // to do boundary_check in tt.load, we need to keep the parent tensor's + // shape info in the IR. + // use the parent tensor's shape to create a cast + auto resultSizes = data.getSizes(); + auto resultOffsets = data.getOffsets(); + data.getSizesRef().clear(); + data.getOffsetsRef().clear(); + data.getSizesRef() = + std::move(llvm::map_to_vector(op.getShape(), [&](Value v) { + return getOpFoldResultOfLayoutInfo(v, rewriter); + })); + + // This redundant ReinterpretCastOp is to describe full tensor_ptr, so each + // dim offset from base is initialized as zero. + SmallVector curOffsets(op.getOffsets().size(), + rewriter.getIndexAttr(0)); + // Just accumulate base potential offset + curOffsets.front() = accumulatePotentialOffsetOnBase( + op, rewriter.getRemappedValue(op.getBase()), curOffsets.front(), + rewriter); + + for (auto offset : curOffsets) { + data.getOffsetsRef().push_back(offset); + } + + SmallVector staticShapes; + SmallVector dynamicShapes; + dispatchIndexOpFoldResults(data.getSizesRef(), dynamicShapes, staticShapes); + auto castOp = data.createCastOp(staticShapes, loc, rewriter); + // restore sizes and offsets + data.getSizesRef().clear(); + for (auto &s : resultSizes) { + data.getSizesRef().push_back(s); + } + data.getOffsetsRef().clear(); + for (auto &offset : resultOffsets) { + data.getOffsetsRef().push_back(offset); + } + return castOp; +} + +void BlockDataParser::rewriteMakeTensorPtrOp( + triton::MakeTensorPtrOp op, Value base, ConversionPatternRewriter &rewriter, + llvm::SmallDenseMap &known) { + Location loc = op.getLoc(); + BlockData data; + + auto orderSize = op.getOrder().size(); + + // Handle base is defined by tt.bitcast + BlockDataParser::parse(op.getBase(), data, loc, rewriter, known); + if (data.hasResElemTy()) { + auto memrefType = dyn_cast(data.getSourceRef().getType()) + .cloneWith(std::nullopt, data.getResElemTyRef()); + UnrealizedConversionCastOp castOp = + rewriter.create(loc, memrefType, + data.getSourceRef()); + data.setSource(castOp.getOutputs()[0]); + } else { + data.setSource(rewriter.getRemappedValue(op.getBase())); + } + + data.getOffsetsRef() = + std::move(llvm::map_to_vector(op.getOffsets(), [&](Value v) { + auto zeroVal = rewriter.create( + loc, rewriter.getI32IntegerAttr(0)); + v = rewriter.create(loc, v, zeroVal); + return getOpFoldResultOfLayoutInfo(v, rewriter); + })); + data.getStridesRef() = + std::move(llvm::map_to_vector(op.getStrides(), [&](Value v) { + return getOpFoldResultOfLayoutInfo(v, rewriter); + })); + + SmallVector newOffsets; + for (auto [offset, stride] : + llvm::zip(data.getOffsetsRef(), data.getStridesRef())) + newOffsets.push_back(mulOpFoldResult(offset, stride, loc, rewriter)); + + // 1. Consider that current base ptr may comes from `triton::AddPtrOp`, + // which have been converted to `memref::ReinterpretCastOp` with 1D + // shape([1,]) by `AddPtrConverter`. + // 2. While here would also convert `triton::MakeTensorPtrOp` to + // `memref::ReinterpretCastOp`, it will create use-def on double recast + // which means offset&size&stride info of first one will be dropped in terms + // of memref recast op fold specification. + // + // Conclusion with above two: + // Base of MakeTensorPtrOp has been seen as origin base, so it should + // reserve offset of first recast if it exists. + // Here extract the offset of first recast and add it to highest dimension + newOffsets.front() = + accumulatePotentialOffsetOnBase(op, base, newOffsets.front(), rewriter); + + data.getOffsetsRef().clear(); + + for (auto offset : newOffsets) { + data.getOffsetsRef().push_back(offset); + } + + ArrayRef resultShape; + auto pointerType = cast(op.getResult().getType()); + if (auto shapedType = dyn_cast(pointerType.getPointeeType())) { + resultShape = shapedType.getShape(); + data.getSizesRef().clear(); + for (auto dim_size : resultShape) { + data.getSizesRef().push_back( + IntegerAttr::get(IntegerType::get(op.getContext(), 64), dim_size)); + } + } else { + // scalar pointer, should produce a one dimensional memref + SmallVector scalarShape(1, 1); + resultShape = scalarShape; + assert(data.getRank() == 1); + } + + // special handling for davinci + // create redundant reinterpret_cast op for record shape info + auto redundantOp = createRedundantOp(op, rewriter, data); + redundantOp->setAttr("tensor_ptr_full_shape", rewriter.getUnitAttr()); + + // create reinterpret_cast op for the target block + data.setSource(redundantOp.getResult()); + known[op.getResult()] = data; + auto castOp = data.createCastOp(resultShape, loc, rewriter); + rewriter.replaceOp(op, castOp.getResult()); + + if (nd2nzFlag) { + auto basePtr = castOp.getResult(); + int original_rank = op.getShape().size() + 1; + std::string shapeStr; + + auto baseMemrefType = mlir::dyn_cast(basePtr.getType()); + assert(baseMemrefType && "basePtr is not a memref type"); + auto shape = baseMemrefType.getShape(); + + if (auto memrefType = mlir::dyn_cast(basePtr.getType())) { + for (auto dim : memrefType.getShape()) { + shapeStr += llvm::formatv("_{0}", dim); + } + } + std::string elemTypeName; + Type elemType = baseMemrefType.getElementType(); + if (auto intType = mlir::dyn_cast(elemType)) { + elemTypeName = llvm::formatv("i{0}", intType.getWidth()); + } else if (auto floatType = mlir::dyn_cast(elemType)) { + std::string floatTypeName; + llvm::raw_string_ostream os(floatTypeName); + floatType.print(os); + os.flush(); + elemTypeName = floatTypeName; + } else { + std::string typeName; + llvm::raw_string_ostream os(typeName); + elemType.print(os); + os.flush(); + elemTypeName = typeName; + } + + std::string memrefTypeStr; + llvm::raw_string_ostream os(memrefTypeStr); + baseMemrefType.print(os); + os.flush(); + + std::string laydbgsuffix; + for (char c : memrefTypeStr) { + if ((c >= '0' && c <= '9') || (c >= 'a' && c <= 'z') || + (c >= 'A' && c <= 'Z') || c == '_' || c == ',' || c == '[' || + c == ']') { + laydbgsuffix += c; + } + } + auto funcName = rewriter.getStringAttr( + llvm::formatv("__hmf_original_shape{0}d{1}_{2}_{3}", original_rank, + shapeStr, elemTypeName, laydbgsuffix)); + MemRefType targetMemrefType = MemRefType::get( + baseMemrefType.getShape(), baseMemrefType.getElementType(), + baseMemrefType.getLayout()); + const int vectorSize = 4; + SmallVector srcElemTys; + for (auto sz : op.getShape()) { + srcElemTys.push_back(sz.getType()); + } + srcElemTys.push_back(targetMemrefType); + Type dstElemTy = rewriter.getNoneType(); + FunctionType hintFuncType = + FunctionType::get(rewriter.getContext(), srcElemTys, {dstElemTy}); + + auto mod = SymbolTable::getNearestSymbolTable(op); + auto extFunc = dyn_cast_or_null( + SymbolTable::lookupSymbolIn(mod, funcName)); + SmallVector args; + for (auto sz : op.getShape()) { + args.push_back(sz); + } + args.push_back(basePtr); + if (!extFunc) { + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(&mod->getRegion(0).front()); + extFunc = rewriter.create(rewriter.getUnknownLoc(), + funcName, hintFuncType); + extFunc.setPrivate(); + extFunc->setAttr(LLVM::LLVMDialect::getReadnoneAttrName(), + UnitAttr::get(rewriter.getContext())); + rewriter.setInsertionPoint(op); + } + rewriter.create(loc, funcName, dstElemTy, args); + } +} + +void BlockDataParser::rewriteAdvanceOp( + triton::AdvanceOp op, ConversionPatternRewriter &rewriter, + llvm::SmallDenseMap &known) { + OpBuilder::InsertionGuard insertionGuard(rewriter); + rewriter.setInsertionPoint(op); + auto loc = op.getLoc(); + + BlockData blockData; + parse(op.getOperand(0), blockData, loc, rewriter, known); + + // region [BUGFIX] Add the code block below following the same logic as + // 'BlockDataParser::rewriteAddPtr' function. + known[op.getResult()] = blockData; + auto inferedSize = 1; + for (int i = blockData.getSizesRef().size() - 1; i >= 0; i--) { + auto strideConst = getConstantIntValue(blockData.getStridesRef()[i]); + auto sizeConst = getConstantIntValue(blockData.getSizesRef()[i]); + assert(sizeConst.has_value()); + if (sizeConst.value() == 1 && strideConst && strideConst.value() == 0) { + blockData.getStridesRef()[i] = rewriter.getIndexAttr(inferedSize); + } + inferedSize *= sizeConst.value(); + } + // endregion + + SmallVector incrementOffsets = + llvm::map_to_vector(op.getOffsets(), [&](Value offset) { + return getOpFoldResultOfLayoutInfo(offset, rewriter); + }); + + SmallVector newOffsets; + for (const auto [increment, originalOffset, stride] : + llvm::zip(incrementOffsets, blockData.getOffsetsRef(), + blockData.getStridesRef())) { + auto curDimOffset = + addOpFoldResult(mulOpFoldResult(increment, stride, loc, rewriter), + originalOffset, loc, rewriter); + + newOffsets.push_back(curDimOffset); + } + + blockData.getOffsetsRef().clear(); + + for (auto offset : newOffsets) + blockData.getOffsetsRef().push_back(offset); + + SmallVector scalarShape(1, 1); // Stub shape + ArrayRef resultShape; + auto pointerType = cast(op.getResult().getType()); + + if (auto shapedType = dyn_cast(pointerType.getPointeeType())) { + resultShape = shapedType.getShape(); + } else { + // scalar pointer, should produce a one dimensional memref + resultShape = scalarShape; + assert(blockData.getRank() == 1); + } + + auto newOp = blockData.createCastOp(resultShape, loc, rewriter); + rewriter.replaceOp(op, newOp.getResult()); + + known[newOp.getResult()] = blockData; +} + +template +std::enable_if_t || + std::is_same_v> +BlockDataParser::rewriteTerminator( + T op, ConversionPatternRewriter &rewriter, + const llvm::SmallDenseSet &blockArgIdxSet, + ArrayRef iterArgIdxMap, + const llvm::SmallDenseMap &known) { + // Any inserted instruction should be before this yield + OpBuilder::InsertionGuard insertionGuard{rewriter}; + rewriter.setInsertionPoint(op); + + auto adaptor = typename T::Adaptor(op); + ValueRange args; + if constexpr (std::is_same_v) { + args = adaptor.getOperands(); + } else { + args = adaptor.getArgs(); + } + + SmallVector initArgState; + SmallVector operands; + + operands.reserve(op->getNumOperands()); + for (const auto &[oper, newIterArgIdx] : + llvm::zip_equal(args, iterArgIdxMap)) { + if (newIterArgIdx != -1) + operands.push_back(oper); + } + + // For each of the init arg that we added additional Values in for loop, we + // need to add corresponding Values as yield operands. The loop below gathers + // BlockData for those values. + for (auto [i, v] : llvm::enumerate(args)) { + if (auto mappedV = rewriter.getRemappedValue(v)) { + // If this value is a tensor of pointers produced by AddPtrOp, + // we should have already converted to a ReinterpretCastOp without + // layout information for the normal cases + if (v.getDefiningOp() || + v.getDefiningOp() || + v.getDefiningOp()) { + if (auto castOp = mappedV.getDefiningOp()) { + v = castOp; + } else { + llvm_unreachable("mapped value defined by an unexpected op"); + } + } else { + // If this value is not a tensor of pointers, we will use the + // mapped value, and rely on the conversion will happen later + // automatically when we legalize loop body. + + // TODO: + // The scenario where a value is a tensor of pointers but not + // produced by AddPtrOp is not supported + if (isa(mappedV.getType()) && + isa( + dyn_cast(mappedV.getType()).getElementType())) + llvm_unreachable("unsupported scenario where a value is a tensor of " + "pointers but not produced by AddPtrOp"); + v = mappedV; + } + } + + if (blockArgIdxSet.find(i) == blockArgIdxSet.end()) + continue; + + auto reintCastOp = v.getDefiningOp(); + assert( + reintCastOp || + (isa(v.getType()) && + isa(dyn_cast(v.getType()).getElementType()))); + + BlockData state; + if (reintCastOp) { + parseReinterpretCast(reintCastOp, state, op.getLoc(), rewriter, known); + } else { + parse(v, state, op.getLoc(), rewriter, known); + } + initArgState.push_back(state); + } + + // For each of the BlockData recorded in the last step, extract value + // that correspond to offset and stride for each dimension and append + // them to yield operands. + for (auto state : initArgState) { + for (auto offset : state.getOffsetsRef()) { + // offsets can be IntAttr zeroes, since reinterpret_cast collapses + // them for the input memref, and the for loop may not update + // offsets other than offsets[0]. Create constants Values for those + // zeroes. + if (isa(offset)) { + auto constOffset = cast(offset); + assert(isa(constOffset) && + dyn_cast(constOffset).getInt() == 0 && + "attribute offsets should be zeroes"); + auto constOp = rewriter.create( + op.getLoc(), rewriter.getIndexAttr(0)); + operands.push_back(constOp.getResult()); + } else { + operands.push_back(cast(offset)); + } + } + + auto sizesRef = state.getSizesRef(); + size_t dimIdx = 0; + for (OpFoldResult stride : state.getStridesRef()) { + if (isa(stride)) { + auto constStride = cast(stride); + assert(isa(constStride) && + "attribute strides should be IntegerAttr"); + auto strideVal = dyn_cast(constStride).getInt(); + bool isSizeOne = + (dimIdx < sizesRef.size() && isa(sizesRef[dimIdx]) && + cast(cast(sizesRef[dimIdx])).getInt() == + 1); + assert((strideVal == 1 || (strideVal == 0 && isSizeOne)) && + "attribute strides should be ones"); + auto constOp = rewriter.create( + op.getLoc(), rewriter.getIndexAttr(1)); + operands.push_back(constOp.getResult()); + } else { + operands.push_back(cast(stride)); + } + dimIdx++; + } + } + + // Yield is a terminator op that must be at the end of the function + rewriter.setInsertionPointAfter(op); + Operation *newOp; + if constexpr (std::is_same_v) { + newOp = rewriter.replaceOpWithNewOp(op, operands); + } else { + newOp = rewriter.replaceOpWithNewOp(op, op.getCondition(), + operands); + } + + assert(op->getNumResults() == 0); + + LLVM_DEBUG({ + llvm::dbgs() << "new terminator: "; + newOp->print(llvm::dbgs(), OpPrintingFlags().printGenericOpForm()); + llvm::dbgs() << "\n"; + }); +} + +// This function is util function for rewriteLoopOp that +// check if given regionIterArg is used with given condition +bool isUsedWithCondition(Value v, std::function cond, + int depth = 0, + llvm::SmallSetVector *visited = nullptr) { + llvm::SmallSetVector localVisited; + if (!visited) { + visited = &localVisited; + } + + if (visited->contains(v)) { + return false; + } + visited->insert(v); + + for (auto &use : v.getUses()) { + auto *user = use.getOwner(); + if (user->hasAttr(ConverterUtils::discreteAttrName) || + isa(user)) + continue; + if (cond(&use)) + return true; + if (auto loopOp = dyn_cast(user); + loopOp && !loopOp->hasAttr("ExtractedLoadOrStore")) { + Value tiedArg = loopOp.getTiedLoopRegionIterArg(&use); + if (tiedArg && isUsedWithCondition(tiedArg, cond, depth + 1, visited)) + return true; + } else if (auto yieldOp = dyn_cast(user); + yieldOp && !isa(user->getParentOp())) { + if (depth && isUsedWithCondition(yieldOp->getParentOp()->getResult( + use.getOperandNumber()), + cond, depth - 1, visited)) + return true; + } else if (auto conditionOp = dyn_cast(user); + conditionOp && use.getOperandNumber() > 0) { + auto whileOp = cast(conditionOp->getParentOp()); + if (depth && + isUsedWithCondition(whileOp->getResult(use.getOperandNumber() - 1), + cond, depth - 1, visited)) + return true; + if (isUsedWithCondition( + whileOp.getAfterArguments()[use.getOperandNumber() - 1], cond, + depth, visited)) + return true; + } + for (auto res : user->getResults()) { + if (isUsedWithCondition(res, cond, depth, visited)) + return true; + } + } + return false; +} + +// This function is util function for rewriteLoopOp that create value from data. +// Assume data is structured, and from regionIterArg from LoopLikeOpInterface. +// +// For example, +// +// %7 = scf.for %arg2 = %c0_i32 to %c3_i32 step %c1_i32 iter_args(%arg3 = %4) -> +// (tensor<128xi32>) : i32 { +// %8 = tt.addptr %5, %arg3 : tensor<128x!tt.ptr>, tensor<128xi32> +// ... +// } +// +// is converted to +// +// %7 = scf.for %arg2 = %c0_i32 to %c3_i32 step %c1_i32 iter_args(%arg3 = %4, +// %arg4 = %5, %arg5 = %6) -> (tensor<128xi32>) : i32 { +// %scalarOffset = arith.index_cast %arg4 : index to i32 +// %scalarStride = arith.index_cast %arg5 : index to i32 +// ... +// %newRes = arith.addi %offset, %stride : tensor<128xi32> +// %8 = tt.addptr %5, %newRes : tensor<128x!tt.ptr>, tensor<128xi32> +// } +Value createFromData(RankedTensorType resType, const BlockData &data, + const Location &loc, OpBuilder &builder, + bool isMaskIterArg) { + auto resShape = resType.getShape(); + Value newRes = nullptr; + for (size_t i = 0; i < resShape.size(); i++) { + auto axisType = + RankedTensorType::get({resShape[i]}, resType.getElementType()); + auto axisI32Type = + RankedTensorType::get({resShape[i]}, builder.getIntegerType(32)); + Value axisValue = + builder.create(loc, axisI32Type, 0, resShape[i]); + if (axisType != axisI32Type) { + axisValue = builder.create(loc, axisType, axisValue); + } + Value offset = cast(data.getOffset(i)); + Value offsetValue = builder.create( + loc, resType.getElementType(), offset); + offsetValue = builder.create(loc, axisType, offsetValue); + Value stride = cast(data.getStride(i)); + if (!isMaskIterArg) { + Value strideValue = builder.create( + loc, resType.getElementType(), stride); + strideValue = builder.create(loc, axisType, strideValue); + axisValue = builder.create(loc, axisValue, strideValue); + } + axisValue = builder.create(loc, axisValue, offsetValue); + + for (size_t j = 0; j < resShape.size(); j++) { + if (i != j) + axisValue = builder.create(loc, axisValue, j); + } + axisValue = builder.create(loc, resType, axisValue); + if (newRes) { + newRes = builder.create(loc, newRes, axisValue); + } else { + newRes = axisValue; + } + } + return newRes; +} + +void BlockDataParser::rewriteLoopOp( + LoopLikeOpInterface op, ConversionPatternRewriter &rewriter, + llvm::SmallDenseMap &known) { + SmallVector newInitArgs; + SmallVector iterArgIdxMap; + SmallVector maskIterArgs; + int64_t argCnt = 0; + + SmallVector, 5> initArgIndexIfBlockData; + SmallVector, 5> knownPtrsTmp; + llvm::SmallDenseSet blockArgIdxSet; + + // Create a new list of init args + for (auto [i, arg] : llvm::enumerate(op.getInits())) { + auto mappedV = rewriter.getRemappedValue(arg); + memref::ReinterpretCastOp reintCastOp; + maskIterArgs.push_back(false); + + // If this init arg is supposed to be remapped, use the remapped + // value instead. + // In addition, if this init arg is a memref created by a reinterpret_cast + // or a tensor of index, there is a chance that it will be used in addptr. + // Create BlockData for each such init arg. + if (mappedV) { + // TODO: + // Passing a block argument pointer directly into a for loop not + // supported. + assert(!(isa(mappedV) && + isa(mappedV.getType())) && + "cannot take pointer block argument as init arg for for loop"); + if (auto reinterpretCastOp = + mappedV.getDefiningOp()) { + // Record memref::ReinterpretCastOp + reintCastOp = reinterpretCastOp; + newInitArgs.push_back(mappedV); + iterArgIdxMap.push_back(argCnt++); + } else { + newInitArgs.push_back(mappedV); + iterArgIdxMap.push_back(argCnt++); + } + } else { + newInitArgs.push_back(arg); + iterArgIdxMap.push_back(argCnt++); + } + + auto indexTensor = + isa(arg.getType()) && + isa(cast(arg.getType()).getElementType()) && + cast(cast(arg.getType()).getElementType()) + .getWidth() != 1 && + isUsedWithCondition(op.getRegionIterArgs()[i], [](OpOperand *use) { + auto *user = use->getOwner(); + return isa(user) || + (isa(user) && use->getOperandNumber() == 1) || + (isa(user) && use->getOperandNumber() == 2); + }); + + // Handle memref::ReinterpretCastOp and tensor specially + if (!reintCastOp && !indexTensor) + continue; + + BlockData data; + if (reintCastOp) { + parseReinterpretCast(reintCastOp, data, op.getLoc(), rewriter, + llvm::SmallDenseMap(0)); + } else { + parse(arg, data, op.getLoc(), rewriter, + llvm::SmallDenseMap(0)); + } + + maskIterArgs[i] = + indexTensor && + isUsedWithCondition(op.getRegionIterArgs()[i], [](OpOperand *use) { + auto *user = use->getOwner(); + return (isa(user) && use->getOperandNumber() == 1) || + (isa(user) && use->getOperandNumber() == 2); + }); + + if (indexTensor) { + newInitArgs.back() = nullptr; + iterArgIdxMap.back() = -1; + argCnt--; + } + + // Record the BlockData for later processing + initArgIndexIfBlockData.push_back(std::make_pair(i, data)); + } + + // Set insertion point to be before the for loop for new variables passed + // into the new loop. + auto origIp = rewriter.saveInsertionPoint(); + rewriter.setInsertionPoint(op); + + // For each of the BlockData recorded in the last step, insert new + // instructions to describe offset and stride for each dimension and append + // them to init args + for (auto [i, data] : initArgIndexIfBlockData) { + // For each dimension, if the corresponding offset and stride is an + // integer attribute, create a constant value and append them at the + // end of init arg list, which is prepared for calculate layout info with + // loop interation index + for (auto &dataOffset : data.getOffsetsRef()) { + if (isa(dataOffset)) { + auto constDataOffset = cast(dataOffset); + assert(isa(constDataOffset)); + auto constOp = rewriter.create( + op.getLoc(), rewriter.getIndexAttr( + dyn_cast(constDataOffset).getInt())); + newInitArgs.push_back(constOp.getResult()); + dataOffset = constOp.getResult(); + } else { + assert(isa(cast(dataOffset).getType())); + newInitArgs.push_back(cast(dataOffset)); + } + } + + for (auto &dataStride : data.getStridesRef()) { + if (isa(dataStride)) { + auto constDataStride = cast(dataStride); + assert(isa(constDataStride)); + auto constOp = rewriter.create( + op.getLoc(), rewriter.getIndexAttr( + dyn_cast(constDataStride).getInt())); + newInitArgs.push_back(constOp.getResult()); + dataStride = constOp.getResult(); + } else { + assert(isa(cast(dataStride).getType())); + newInitArgs.push_back(cast(dataStride)); + } + } + + // Note that we want the knownPtrs to be indexed by block arg, but we + // only have index for now. Also, the blockdata we record is the init + // arg, but want to to use newly created block arg. These block args + // are not created yet. We will translate this mapping later. + knownPtrsTmp.push_back(std::make_pair(i, data)); + blockArgIdxSet.insert(i); + + // If the original init arg is a memref produced by reinterpret_cast, + // create a new memref using new strides and offsets created above. + // This produces a canonicalized memref, which will match what the + // for loop generates if it modifies the memref. E.g., original + // reinterpret_cast can produce a memref with const stride: + // - memref<4x256xbf16, affine_map<(d0, d1)[s0, s1] -> (d0 * 256 + + // s0 + d1 + // * s1)>> + // The new reinterpret_cast will always have dynamic stride and + // offset: + // - memref<4x256xbf16, affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + // + s0 + d1 * s2)>> + if (newInitArgs[i] && + newInitArgs[i].getDefiningOp()) { + SmallVector resultShape; + for (auto size : data.getSizesRef()) { + auto constSize = getConstantIntValue(size); + assert(constSize && "expected constant size"); + resultShape.push_back(constSize.value()); + } + + // In current block data layout info, strides and offsets must be dynamic + // value + auto castOp = data.createCastOp(resultShape, op.getLoc(), rewriter); + if (resultShape.size() > 1) { + auto originalOffset = dyn_cast(data.getOffsetsRef()[0]); + for (auto &offsets : newInitArgs) { + if (offsets == originalOffset) { + offsets = castOp.getOffsets()[0]; + break; + } + } + data.getOffsetsRef()[0] = castOp.getOffsets()[0]; + } + + LLVM_DEBUG({ + llvm::dbgs() << "new reinterpret_cast with dynamic sizes " + "and offsets:"; + castOp->print(llvm::dbgs(), OpPrintingFlags().printGenericOpForm()); + llvm::dbgs() << "\n"; + }); + + newInitArgs[i] = castOp.getResult(); + } + } + + rewriter.restoreInsertionPoint(origIp); + IRMapping mapping; + + // Create a new LoopOp that uses updated init args and same loop body + LoopLikeOpInterface newOp; + auto newInits = to_vector( + make_filter_range(newInitArgs, [](Value v) { return v != nullptr; })); + auto commonBodyBuilder = [&](OpBuilder &b, Location loc, bool useInit, + ValueRange newRegionArgs, Region ®ion, + Block::BlockArgListType regionArgs, + ArrayRef isUsedForRegionArgs, + ArrayRef maskIterArgs) { + auto newArgIter = newRegionArgs.begin(); + for (const auto &[regionArg, isUsedForRegionArg] : + llvm::zip(regionArgs, isUsedForRegionArgs)) { + if (isUsedForRegionArg) { + mapping.map(regionArg, *newArgIter); + ++newArgIter; + } + } + + // Convert the book-keeping data structure to use the correct key and value. + // Key is converted from init arg index to newly created block arg, and + // Value's BlockData fields are converted from init arg to newly created + // block arg + + // TODO: remove (useInit = true) logic after supporting make_tensor_ptr + if (useInit) { + for (auto [i, data] : knownPtrsTmp) { + for (auto &offset : data.getOffsetsRef()) { + offset = *newArgIter; + ++newArgIter; + } + + for (auto &stride : data.getStridesRef()) { + stride = *newArgIter; + ++newArgIter; + } + + auto regionArg = regionArgs[i]; + auto key = mapping.lookupOrNull(regionArg); + if (!key) { + // Create IndexTensor regionArg from computed offset and stride data + key = createFromData(cast(regionArg.getType()), + data, op.getLoc(), rewriter, maskIterArgs[i]); + mapping.map(regionArg, key); + } + known.insert(std::make_pair(key, data)); + } + } else { + for (auto [i, isUsedForRegionArg] : + llvm::enumerate(isUsedForRegionArgs)) { + if (!isUsedForRegionArg) { + BlockData data; + auto regionArg = regionArgs[i]; + auto regionArgType = cast(regionArg.getType()); + data.getOffsetsRef().resize(regionArgType.getRank()); + data.getStridesRef().resize(regionArgType.getRank()); + for (auto &offset : data.getOffsetsRef()) { + offset = *newArgIter; + ++newArgIter; + } + for (auto &dim : regionArgType.getShape()) { + data.getSizesRef().push_back(rewriter.getIndexAttr(dim)); + } + for (auto &stride : data.getStridesRef()) { + stride = *newArgIter; + ++newArgIter; + } + + auto key = mapping.lookupOrNull(regionArg); + if (!key) { + // Create IndexTensor regionArg from computed offset and stride data + key = createFromData(regionArgType, data, op.getLoc(), rewriter, + maskIterArgs[i]); + mapping.map(regionArg, key); + } + known.insert(std::make_pair(key, data)); + } + } + } + + for (auto &bodyOp : region.getOps()) + b.clone(bodyOp, mapping); + }; + for (const auto &[initArg, newInitArg] : + llvm::zip(op.getInits(), newInitArgs)) { + if (newInitArg) { + mapping.map(initArg, newInitArg); + } + } + if (auto forOp = dyn_cast(op.getOperation())) { + SmallVector usedForRegionArgs; + for (auto newInitArg : newInitArgs) { + usedForRegionArgs.push_back(newInitArg ? true : false); + } + newOp = rewriter.create( + forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), + forOp.getStep(), newInits, + [&](OpBuilder &b, Location loc, Value iv, ValueRange args) { + mapping.map(forOp.getInductionVar(), iv); + commonBodyBuilder(b, loc, true, args, forOp.getRegion(), + op.getRegionIterArgs(), usedForRegionArgs, + maskIterArgs); + }); + + // Replace only the results that correspond to the original scf.for + auto newResultIter = newOp->result_begin(); + rewriter.setInsertionPointAfter(newOp); + for (const auto &[res, regionArg, newIterArgIdx, mask] : + llvm::zip_equal(op->getResults(), op.getRegionIterArgs(), + iterArgIdxMap, maskIterArgs)) { + if (newIterArgIdx != -1) { + rewriter.replaceAllUsesWith(res, *newResultIter); + ++newResultIter; + } else { + auto key = mapping.lookup(regionArg); + auto data = known.at(key); + for (auto &offset : data.getOffsetsRef()) + offset = + newOp.getTiedLoopResult(cast(cast(offset))); + for (auto &stride : data.getStridesRef()) + stride = + newOp.getTiedLoopResult(cast(cast(stride))); + auto newRes = + createFromData(cast(regionArg.getType()), data, + op.getLoc(), rewriter, mask); + rewriter.replaceAllUsesWith(res, newRes); + } + } + } else if (auto whileOp = dyn_cast(op.getOperation())) { + SmallVector resultTypes; + SmallVector usedForBeforeRegionArgs; + SmallVector usedForAfterRegionArgs; + llvm::SmallDenseSet blockArgIdxSetForAfter; + SmallVector iterArgIdxMapForAfter; + SmallVector maskIterArgsForAfter(whileOp->getNumResults()); + + int64_t indexCnt = 0; + + for (auto newInitArg : newInitArgs) { + usedForBeforeRegionArgs.push_back(newInitArg ? true : false); + } + for (size_t i = 0; i < whileOp->getNumResults(); i++) { + auto resType = whileOp->getResultTypes()[i]; + auto indexTensor = + isa(resType) && + isa(cast(resType).getElementType()) && + isUsedWithCondition(whileOp.getAfterArguments()[i], + [](OpOperand *use) { + auto *user = use->getOwner(); + return isa(user) || + (isa(user) && + use->getOperandNumber() == 1) || + (isa(user) && + use->getOperandNumber() == 2); + }); + if (indexTensor) { + indexCnt += 2 * cast(resType).getRank(); + usedForAfterRegionArgs.push_back(false); + iterArgIdxMapForAfter.push_back(-1); + maskIterArgsForAfter[i] = isUsedWithCondition( + whileOp.getAfterArguments()[i], [](OpOperand *use) { + auto *user = use->getOwner(); + return (isa(user) && + use->getOperandNumber() == 1) || + (isa(user) && + use->getOperandNumber() == 2); + }); + blockArgIdxSetForAfter.insert(i); + } else { + resultTypes.push_back(resType); + usedForAfterRegionArgs.push_back(true); + iterArgIdxMapForAfter.push_back(argCnt++); + } + } + resultTypes.append(indexCnt, rewriter.getIndexType()); + newOp = rewriter.create( + whileOp.getLoc(), resultTypes, newInits, + [&](OpBuilder &b, Location loc, ValueRange args) { + commonBodyBuilder(b, loc, true, args, whileOp.getBefore(), + whileOp.getBeforeArguments(), + usedForBeforeRegionArgs, maskIterArgs); + }, + [&](OpBuilder &b, Location loc, ValueRange args) { + commonBodyBuilder(b, loc, false, args, whileOp.getAfter(), + whileOp.getAfterArguments(), usedForAfterRegionArgs, + maskIterArgsForAfter); + }); + + auto newResultIter = newOp->result_begin(); + rewriter.setInsertionPointAfter(newOp); + for (const auto &[res, regionArg, newIterArgIdx, mask] : + llvm::zip_equal(op->getResults(), whileOp.getAfterArguments(), + iterArgIdxMapForAfter, maskIterArgsForAfter)) { + if (newIterArgIdx != -1) { + rewriter.replaceAllUsesWith(res, *newResultIter); + ++newResultIter; + } else { + auto key = mapping.lookup(regionArg); + auto data = known.at(key); + for (auto &offset : data.getOffsetsRef()) + offset = newOp->getResult( + cast(cast(offset)).getArgNumber()); + for (auto &stride : data.getStridesRef()) + stride = newOp->getResult( + cast(cast(stride)).getArgNumber()); + auto newRes = + createFromData(cast(regionArg.getType()), data, + op.getLoc(), rewriter, mask); + rewriter.replaceAllUsesWith(res, newRes); + } + } + + auto conditionOp = + cast(newOp.getOperation()).getConditionOp(); + rewriteTerminator(conditionOp, rewriter, blockArgIdxSetForAfter, + iterArgIdxMapForAfter, known); + } + + // Copy all attributes from op to newOp + newOp->setAttrs(op->getAttrs()); + rewriter.eraseOp(op); + + // Update the loop body. Manually invoke the rewrite logic on addptr and yield + // in the loop body, so we can take advantage of the states we built up + for (auto *region : newOp.getLoopRegions()) { + for (auto &bodyOp : region->getOps()) { + if (auto customOp = dyn_cast(bodyOp)) { + auto adaptor = hivm::CustomOp::Adaptor(customOp); + rewriteCustomOp(customOp, adaptor, rewriter, known); + } else if (auto addptrOp = dyn_cast(bodyOp)) { + // FIXME: Constructed adaptor here does not hold the transformed op + // info. + auto adaptor = triton::AddPtrOp::Adaptor(addptrOp); + rewriteAddPtr(addptrOp, adaptor, rewriter, known); + } else if (auto advanceOp = dyn_cast(bodyOp)) { + rewriteAdvanceOp(advanceOp, rewriter, known); + } else if (auto makeTensorPtrOp = + dyn_cast(bodyOp)) { + ConversionPatternRewriter::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(makeTensorPtrOp); + rewriteMakeTensorPtrOp( + makeTensorPtrOp, + rewriter.getRemappedValue(makeTensorPtrOp.getBase()), rewriter, + known); + } else if (auto loopOp = dyn_cast(bodyOp); + loopOp && !loopOp->hasAttr("ExtractedLoadOrStore")) { + ConversionPatternRewriter::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(loopOp); + // Remove UnhandledLoopOp attr before process + loopOp->removeAttr("UnhandledLoopOp"); + rewriteLoopOp(loopOp, rewriter, known); + } + } + } + + if (!op.getRegionIterArgs().empty()) { + auto yieldOp = cast( + newOp.getLoopRegions().back()->back().getTerminator()); + rewriteTerminator(yieldOp, rewriter, blockArgIdxSet, iterArgIdxMap, known); + } + + LLVM_DEBUG({ + llvm::dbgs() << "new loop\n"; + newOp.getOperation()->print(llvm::dbgs(), + OpPrintingFlags().printGenericOpForm()); + llvm::dbgs() << "\n"; + }); +} + +/// @brief Rewrite the triton::AddPtrOp to handle unstructured memory access. +/// @param op The triton::AddPtrOp to be rewritten. +/// @param adaptor The adaptor of the triton::AddPtrOp, used to get operands. +/// @param rewriter The pattern rewriter used to modify the IR. +/// @param data The BlockData containing information about the memory access. +void BlockDataParser::rewriteAddPtrToUnstrucMemAcc( + triton::AddPtrOp op, triton::AddPtrOp::Adaptor &adaptor, + ConversionPatternRewriter &rewriter, BlockData &data) { + auto loc = op.getLoc(); + auto &offsets = data.getOffsetsRef(); + auto &blockSizes = data.getSizesRef(); + auto &strides = data.getStridesRef(); + Value ptrOffset = adaptor.getOffset(); + Value zeroIdx = + rewriter.create(loc, rewriter.getIndexAttr(0)); + Value oneIdx = + rewriter.create(loc, rewriter.getIndexAttr(1)); + auto addptrRes = op.getResult(); + assert(addptrRes.hasOneUse() && "Invalid: tt.addptr has multiple users"); + auto loadOp = *(addptrRes.user_begin()); + + // Prepare empty tensor for loop based scalar load + // FIXME: We use cast here because addptr must return tensor>. + // True? + auto resTy = cast(addptrRes.getType()); + auto resEPtrTy = resTy.getElementType(); + auto resETy = cast(resEPtrTy).getPointeeType(); + Value loaded = rewriter.create(loc, blockSizes, resETy); + SmallVector initArgs; + initArgs.push_back(loaded); + + SmallVector forLBs; + SmallVector forUBs; + SmallVector forSteps; + for (auto &s : offsets) { + forLBs.push_back(zeroIdx); + } + for (auto &s : blockSizes) { + forUBs.push_back(getValueOrCreateConstantIndexOp(rewriter, loc, s)); + } + for (auto &s : strides) { + forSteps.push_back(oneIdx); + } + SmallVector ivs; + OpBuilder builder(op); + auto loop = createNestedLoops( + builder, loc, 0, blockSizes.size(), forLBs, forUBs, forSteps, ivs, + initArgs, + [&](OpBuilder &bB, Location bLoc, SmallVector &allIVs, + ValueRange iterArgs) { + OpBuilder::InsertionGuard g(bB); + bB.setInsertionPointToStart(bB.getBlock()); + + Value scalarOffsetRaw = + bB.create(bLoc, ptrOffset, allIVs); + Value scalarOffset = bB.create( + bLoc, bB.getIndexType(), scalarOffsetRaw); + OpFoldResult baseOffset = bB.getIndexAttr(0); + for (auto ofr : data.getOffsetsRef()) { + baseOffset = addOpFoldResult(baseOffset, ofr, bLoc, bB); + } + Value baseVal = getValueOrCreateConstantIndexOp(bB, bLoc, baseOffset); + Value combinedOffset = + bB.create(bLoc, baseVal, scalarOffset); + // Replace offset & size. Only single element. + data.getOffsetsRef().clear(); + data.getOffsetsRef().push_back(combinedOffset); + data.getSizesRef().clear(); + data.getSizesRef().push_back(bB.getIndexAttr(1)); + data.getStridesRef().clear(); + data.getStridesRef().push_back(bB.getIndexAttr(1)); + memref::ReinterpretCastOp castOp = data.createCastOp({1}, bLoc, bB); + rewriter.replaceOp(op, castOp); + // Move tt.load using this tt.addptr into this block + loadOp->moveAfter(castOp); + loadOp->setAttr("IndirectLoad", UnitAttr::get(op.getContext())); + bB.create(bLoc, iterArgs); + }); +} + +} // namespace triton +} // namespace mlir diff --git a/compiler/lib/TritonToLinalg/CMakeLists.txt b/compiler/lib/TritonToLinalg/CMakeLists.txt new file mode 100644 index 00000000..61c89f82 --- /dev/null +++ b/compiler/lib/TritonToLinalg/CMakeLists.txt @@ -0,0 +1,35 @@ +add_triton_library(TritonToLinalg + TritonToLinalgPass.cpp + LoadStoreConverter.cpp + FunctionConverter.cpp + ArgMinMaxConverter.cpp + TritonOpConverter.cpp + HoistBroadcast.cpp + BlockPtrAnalysis.cpp + MaskAnalysis.cpp + UseAnalysis.cpp + ImplicitPermute.cpp + DescriptorConverter.cpp + MarkTensorKindPass.cpp + AscendNPUIRLegalizePass.cpp + + DEPENDS + TritonToLinalgConversionPassIncGen + + LINK_LIBS PUBLIC + MLIRArithDialect + MLIRDialectUtils + MLIRIR + MLIRMathDialect + MLIRPass + MLIRTensorDialect + MLIRTransforms + MLIRSupport + TritonIR + TritonTransforms + TritonAnalysis + MLIRTritonNPUUtils + MLIRSCFTransforms + MLIRLinalgTransforms + BiShengIRHIVMDialect +) \ No newline at end of file diff --git a/compiler/lib/TritonToLinalg/DescriptorConverter.cpp b/compiler/lib/TritonToLinalg/DescriptorConverter.cpp new file mode 100644 index 00000000..10140575 --- /dev/null +++ b/compiler/lib/TritonToLinalg/DescriptorConverter.cpp @@ -0,0 +1,546 @@ + + +#include "dicp/TritonToLinalg/DescriptorConverter.h" +#include "dicp/TritonToLinalg/BlockPtrAnalysis.h" +#include "dicp/TritonToLinalg/MaskAnalysis.h" +#include "dicp/TritonToLinalg/TritonOpConverter.h" +#include "dicp/TritonToLinalg/TritonToLinalgPass.h" +#include "dicp/Utils/Utils.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/LogicalResult.h" +#include "llvm/Support/raw_ostream.h" +#include + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Transforms/DialectConversion.h" + +namespace DescriptorConverter { +using namespace mlir; +using namespace triton; + +bool hasATensorDescriptorType(mlir::TypeRange types) { + return llvm::any_of(types, [](mlir::Type t) { + return llvm::isa(t); + }); +} + +Descriptor unpackDescriptor(TensorDescType type, Value desc, + ConversionPatternRewriter &rewriter) { + auto makeDescOp = desc.getDefiningOp(); + assert(makeDescOp && "Descriptor must be defined by MakeTensorDescOp"); + + Descriptor res; + + res.base = makeDescOp.getBase(); + for (auto s : makeDescOp.getShape()) { + res.shape.push_back(rewriter.createOrFold( + makeDescOp.getLoc(), rewriter.getI64Type(), s)); + } + for (auto st : makeDescOp.getStrides()) { + res.strides.push_back(rewriter.createOrFold( + makeDescOp.getLoc(), rewriter.getI64Type(), st)); + } + res.padding = makeDescOp.getPaddingAttr(); + + return res; +} + +SmallVector computeOrder(ArrayRef shape) { + SmallVector order; + int rank = shape.size(); + order.reserve(rank); + // default by [dims - 1, ..., 0] + for (int i = rank - 1; i >= 0; --i) { + order.push_back(i); + } + return order; +} + +DenseI32ArrayAttr getFullBoundaryCheckAttr(ConversionPatternRewriter &rewriter, + ArrayRef shape) { + SmallVector boundaryCheck; + boundaryCheck.reserve(shape.size()); + for (int32_t dim = 0; dim < static_cast(shape.size()); ++dim) { + boundaryCheck.push_back(dim); + } + return rewriter.getDenseI32ArrayAttr(boundaryCheck); +} + +Value expandOffsets(OpBuilder &builder, Location loc, + ArrayRef blockShape, Value offsets, unsigned dim) { + Value expandedResult = offsets; + for (size_t j = 0; j < blockShape.size(); ++j) { + if (j == dim) { + continue; + } + expandedResult = + builder.create(loc, expandedResult, j); + } + + return expandedResult; +} + +Value getExpandedOffsetWithRange(OpBuilder &builder, const Location &loc, + ArrayRef blockShape, + Value offset, unsigned dim) { + // Add range + auto indexI32RowType = + RankedTensorType::get({blockShape[dim]}, builder.getI32Type()); + auto indexRowType = + RankedTensorType::get({blockShape[dim]}, builder.getI64Type()); + Value splatOffset = + builder.create(loc, indexRowType, offset); + Value range = builder.create(loc, indexI32RowType, 0, + blockShape[dim]); + Value i64Range = builder.create(loc, indexRowType, range); + + Value offsets = builder.create(loc, splatOffset, i64Range); + return expandOffsets(builder, loc, blockShape, offsets, dim); +} + +Value generatePtrFromOffsetRanges(OpBuilder &builder, Location loc, + ArrayRef blockShape, + Descriptor &desc, ValueRange offsets) { + assert(blockShape.size() == desc.shape.size()); + assert(blockShape.size() == offsets.size()); + auto indexTensorType = + RankedTensorType::get(blockShape, builder.getI64Type()); + auto ptrType = cast(desc.base.getType()); + auto ptrTensorType = RankedTensorType::get(blockShape, ptrType); + + // Generate offsets per dimension + Value ptr = builder.create(loc, ptrTensorType, desc.base); + for (unsigned i = 0; i < blockShape.size(); ++i) { + // We must splat strides into the expanded shape not a row for retaining + // the divisibility information given by strides + Value splatStride = builder.create( + loc, offsets[i].getType(), desc.strides[i]); + Value offsetWithStride = + builder.create(loc, offsets[i], splatStride); + Value broadcasted = builder.create( + loc, indexTensorType, offsetWithStride); + + // Add to the pointer + ptr = + builder.create(loc, ptrTensorType, ptr, broadcasted); + } + + return ptr; +} + +Value generatePtr(OpBuilder &builder, const Location &loc, + ArrayRef blockShape, Descriptor &desc, + ValueRange offsets) { + assert(blockShape.size() == desc.shape.size()); + assert(blockShape.size() == offsets.size()); + SmallVector offsetRanges; + for (unsigned i = 0; i < blockShape.size(); ++i) { + auto offsetWithRange = + getExpandedOffsetWithRange(builder, loc, blockShape, offsets[i], i); + offsetRanges.push_back(offsetWithRange); + } + + return generatePtrFromOffsetRanges(builder, loc, blockShape, desc, + offsetRanges); +} + +Value generateMaskFromOffsetRanges(OpBuilder &builder, const Location &loc, + ArrayRef blockShape, + Descriptor &desc, ValueRange offsetRanges) { + assert(blockShape.size() == desc.shape.size()); + assert(blockShape.size() == offsetRanges.size()); + + // Generate mask per dimension + auto maskTensorType = RankedTensorType::get(blockShape, builder.getI1Type()); + Value mask; + for (std::size_t i = 0; i < blockShape.size(); ++i) { + auto offsetWithRange = offsetRanges[i]; + + // Compare with lower bound + Value lowerBound = builder.create( + loc, builder.getI64Type(), 0); + Value splatLowerBound = builder.create( + loc, offsetWithRange.getType(), lowerBound); + Value cmpLower = builder.create( + loc, arith::CmpIPredicate::sge, offsetWithRange, splatLowerBound); + + // Compare with upper bound + Value splatUpperBound = builder.create( + loc, offsetWithRange.getType(), desc.shape[i]); + Value cmpUpper = builder.create( + loc, arith::CmpIPredicate::slt, offsetWithRange, splatUpperBound); + + // And and broadcast + Value andResult = builder.create(loc, cmpLower, cmpUpper); + Value broadcasted = + builder.create(loc, maskTensorType, andResult); + + // And up all results + if (!mask) { + mask = broadcasted; + } else { + mask = builder.create(loc, mask, broadcasted); + } + } + + return mask; +} + +Value generateMask(OpBuilder &builder, const Location &loc, + ArrayRef blockShape, Descriptor &desc, + ValueRange offsets) { + assert(blockShape.size() == desc.shape.size()); + assert(blockShape.size() == offsets.size()); + SmallVector offsetRanges; + for (unsigned i = 0; i < blockShape.size(); ++i) { + auto offsetWithRange = + getExpandedOffsetWithRange(builder, loc, blockShape, offsets[i], i); + offsetRanges.push_back(offsetWithRange); + } + + return generateMaskFromOffsetRanges(builder, loc, blockShape, desc, + offsetRanges); +} + +SmallVector castToI64(OpBuilder &builder, + mlir::ValueRange values) { + auto i64Type = builder.getI64Type(); + return llvm::map_to_vector(values, [&](mlir::Value v) { + return builder.createOrFold(v.getLoc(), i64Type, v); + }); +} + +LogicalResult DescriptorLoadConverter::matchAndRewrite( + triton::DescriptorLoadOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + const auto blockShape = op.getDesc().getType().getBlockType().getShape(); + auto descTy = op.getDesc().getType(); + auto indices = op.getIndices(); + + // 1. unpack descriptor + auto desc = unpackDescriptor(descTy, adaptor.getDesc(), rewriter); + + // 2. create make_tensor_ptr + SmallVector tensorShapeValues; + for (auto dim : blockShape) { + tensorShapeValues.push_back(static_cast(dim)); + } + Value tensorPtr = + rewriter.create(loc, + desc.base, // base + desc.shape, // shape + desc.strides, // strides + indices, // offset + tensorShapeValues, // tensorShape + computeOrder(blockShape) // order + ); + // 3. replace tt.load + auto boundaryCheck = getFullBoundaryCheckAttr(rewriter, blockShape); + triton::PaddingOptionAttr padding = desc.padding; + auto cache = triton::CacheModifierAttr::get(rewriter.getContext(), + triton::CacheModifier::NONE); + auto evict = triton::EvictionPolicyAttr::get(rewriter.getContext(), + triton::EvictionPolicy::NORMAL); + auto isVolatile = rewriter.getBoolAttr(false); + + if (auto a = op->getAttrOfType("cache")) + cache = a; + if (auto a = op->getAttrOfType("evict")) + evict = a; + if (auto a = op->getAttrOfType("isVolatile")) + isVolatile = a; + + auto newLoad = rewriter.create( + loc, descTy.getSignlessBlockType(), tensorPtr, + Value(), // mask + Value(), // other + boundaryCheck, padding, cache, evict, isVolatile); + + rewriter.replaceOp(op, newLoad.getResult()); + + return success(); +} + +LogicalResult DescriptorStoreConverter::matchAndRewrite( + triton::DescriptorStoreOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + const auto blockShape = op.getDesc().getType().getBlockType().getShape(); + auto descTy = op.getDesc().getType(); + auto indices = op.getIndices(); + + // 1. unpack descriptor + auto desc = unpackDescriptor(descTy, adaptor.getDesc(), rewriter); + + // 2. create make_tensor_ptr + SmallVector tensorShapeValues; + for (auto dim : blockShape) { + tensorShapeValues.push_back(static_cast(dim)); + } + Value tensorPtr = + rewriter.create(loc, + desc.base, // base + desc.shape, // shape + desc.strides, // strides + indices, // offset + tensorShapeValues, // tensorShape + computeOrder(blockShape) // order + ); + + // 3. replace tt.store + Value valueToStore = adaptor.getSrc(); + + auto maskType = RankedTensorType::get(blockShape, rewriter.getI1Type()); + rewriter.create(loc, + DenseElementsAttr::get(maskType, true)); + auto boundaryCheck = getFullBoundaryCheckAttr(rewriter, blockShape); + auto cacheModifier = triton::CacheModifierAttr::get( + rewriter.getContext(), triton::CacheModifier::NONE); + auto evictionPolicy = triton::EvictionPolicyAttr::get( + rewriter.getContext(), triton::EvictionPolicy::NORMAL); + + auto newStore = rewriter.create(loc, tensorPtr, valueToStore, + Value(), // mask + boundaryCheck, cacheModifier, + evictionPolicy); + + rewriter.eraseOp(op); + return success(); +} + +LogicalResult DescriptorScatterConverter::matchAndRewrite( + triton::DescriptorScatterOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto descTy = op.getDesc().getType(); + auto srcType = cast(op.getSrc().getType()); + const auto rowBlockShape = descTy.getSignlessBlockType().getShape(); + + auto desc = unpackDescriptor(descTy, adaptor.getDesc(), rewriter); + SmallVector tensorShapeValues; + tensorShapeValues.reserve(rowBlockShape.size()); + for (auto dim : rowBlockShape) { + tensorShapeValues.push_back(static_cast(dim)); + } + + auto zeroIndex = rewriter.create(loc, 0); + auto oneIndex = rewriter.create(loc, 1); + Value rowUpperBound; + if (srcType.isDynamicDim(0)) { + rowUpperBound = rewriter.create(loc, adaptor.getSrc(), 0); + } else { + rowUpperBound = + rewriter.create(loc, srcType.getShape()[0]); + } + auto rowBoundaryCheck = getFullBoundaryCheckAttr(rewriter, rowBlockShape); + auto cacheModifier = triton::CacheModifierAttr::get( + rewriter.getContext(), triton::CacheModifier::NONE); + auto evictionPolicy = triton::EvictionPolicyAttr::get( + rewriter.getContext(), triton::EvictionPolicy::NORMAL); + + auto loop = rewriter.create( + loc, zeroIndex, rowUpperBound, oneIndex, ValueRange{}, + [&](OpBuilder &nestedBuilder, Location nestedLoc, Value rowIv, + ValueRange) { + Value xOffset = nestedBuilder.create( + nestedLoc, adaptor.getXOffsets(), ValueRange{rowIv}); + Value tensorPtr = nestedBuilder.create( + nestedLoc, desc.base, desc.shape, desc.strides, + ValueRange{xOffset, adaptor.getYOffset()}, tensorShapeValues, + computeOrder(rowBlockShape)); + SmallVector extractOffsets{rowIv, + nestedBuilder.getIndexAttr(0)}; + SmallVector extractSizes{ + nestedBuilder.getIndexAttr(rowBlockShape[0]), + nestedBuilder.getIndexAttr(rowBlockShape[1])}; + SmallVector extractStrides{nestedBuilder.getIndexAttr(1), + nestedBuilder.getIndexAttr(1)}; + auto rowValue = nestedBuilder.create( + nestedLoc, adaptor.getSrc(), extractOffsets, extractSizes, + extractStrides); + auto rowStore = nestedBuilder.create( + nestedLoc, tensorPtr, rowValue.getResult(), Value(), + rowBoundaryCheck, cacheModifier, evictionPolicy); + rowStore->setAttr(ConverterUtils::discreteAttrName, + nestedBuilder.getUnitAttr()); + nestedBuilder.create(nestedLoc); + }); + loop->setAttr("ExtractedLoadOrStore", rewriter.getUnitAttr()); + loop->setAttr("hivm.parallel_loop", rewriter.getUnitAttr()); + + rewriter.eraseOp(op); + return success(); +} + +SmallVector filterSegmentSizes(ArrayRef attrs) { + SmallVector filteredAttrs; + llvm::copy_if(attrs, std::back_inserter(filteredAttrs), + [](const NamedAttribute &attr) { + return attr.getName().getValue() != "operandSegmentSizes"; + }); + return filteredAttrs; +} + +LogicalResult DescriptorGatherConverter::matchAndRewrite( + triton::DescriptorGatherOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto descTy = cast(op.getDesc().getType()); + auto resultType = cast(op.getResult().getType()); + const auto blockShape = resultType.getShape(); + const auto rowBlockShape = descTy.getSignlessBlockType().getShape(); + + auto desc = unpackDescriptor(descTy, adaptor.getDesc(), rewriter); + SmallVector tensorShapeValues; + tensorShapeValues.reserve(rowBlockShape.size()); + for (auto dim : rowBlockShape) { + tensorShapeValues.push_back(static_cast(dim)); + } + + auto zeroIndex = rewriter.create(loc, 0); + auto oneIndex = rewriter.create(loc, 1); + Value rowUpperBound = + rewriter.create(loc, adaptor.getXOffsets(), zeroIndex); + auto rowBoundaryCheck = getFullBoundaryCheckAttr(rewriter, rowBlockShape); + auto cache = triton::CacheModifierAttr::get(rewriter.getContext(), + triton::CacheModifier::NONE); + auto evict = triton::EvictionPolicyAttr::get(rewriter.getContext(), + triton::EvictionPolicy::NORMAL); + auto isVolatile = rewriter.getBoolAttr(false); + + if (auto attr = op->getAttrOfType("cache")) + cache = attr; + if (auto attr = op->getAttrOfType("evict")) + evict = attr; + if (auto attr = op->getAttrOfType("isVolatile")) + isVolatile = attr; + + SmallVector dynamicResultSizes; + dynamicResultSizes.reserve(resultType.getNumDynamicDims()); + for (const auto &[dim, size] : llvm::enumerate(blockShape)) { + if (!ShapedType::isDynamic(size)) + continue; + if (dim == 0) { + dynamicResultSizes.push_back(rowUpperBound); + continue; + } + dynamicResultSizes.push_back( + rewriter.create(loc, rowBlockShape[dim])); + } + + auto initialTensor = rewriter.create( + loc, blockShape, resultType.getElementType(), dynamicResultSizes); + auto loop = rewriter.create( + loc, zeroIndex, rowUpperBound, oneIndex, + ValueRange{initialTensor.getResult()}, + [&](OpBuilder &nestedBuilder, Location nestedLoc, Value rowIv, + ValueRange iterArgs) { + Value xOffset = nestedBuilder.create( + nestedLoc, adaptor.getXOffsets(), ValueRange{rowIv}); + Value tensorPtr = nestedBuilder.create( + nestedLoc, desc.base, desc.shape, desc.strides, + ValueRange{xOffset, adaptor.getYOffset()}, tensorShapeValues, + computeOrder(rowBlockShape)); + auto rowLoad = nestedBuilder.create( + nestedLoc, descTy.getSignlessBlockType(), tensorPtr, + Value(), // mask + Value(), // other + rowBoundaryCheck, desc.padding, cache, evict, isVolatile); + for (const auto &attr : filterSegmentSizes(op->getAttrs())) { + if (!rowLoad->hasAttr(attr.getName())) { + rowLoad->setAttr(attr.getName(), attr.getValue()); + } + } + rowLoad->setAttr(ConverterUtils::discreteAttrName, + nestedBuilder.getUnitAttr()); + + auto insertSlice = nestedBuilder.create( + nestedLoc, rowLoad.getResult(), iterArgs[0], + SmallVector{rowIv, nestedBuilder.getIndexAttr(0)}, + SmallVector{ + nestedBuilder.getIndexAttr(rowBlockShape[0]), + nestedBuilder.getIndexAttr(rowBlockShape[1])}, + SmallVector{nestedBuilder.getIndexAttr(1), + nestedBuilder.getIndexAttr(1)}); + insertSlice->setAttr(ConverterUtils::discreteAttrName, + nestedBuilder.getUnitAttr()); + nestedBuilder.create(nestedLoc, insertSlice.getResult()); + }); + loop->setAttr("ExtractedLoadOrStore", rewriter.getUnitAttr()); + + rewriter.replaceOp(op, loop.getResult(0)); + return success(); +} + +std::optional translateReduceKind(DescriptorReduceKind kind, + TensorDescType ty) { + auto scalarTy = ty.getBlockType().getElementType(); + switch (kind) { + case DescriptorReduceKind::ADD: + return scalarTy.isInteger() ? RMWOp::ADD : RMWOp::FADD; + case DescriptorReduceKind::MIN: + if (scalarTy.isUnsignedInteger()) { + return RMWOp::UMIN; + } else if (scalarTy.isSignedInteger()) { + return RMWOp::MIN; + } + return {}; + case DescriptorReduceKind::MAX: + if (scalarTy.isUnsignedInteger()) { + return RMWOp::UMAX; + } else if (scalarTy.isSignedInteger()) { + return RMWOp::MAX; + } + return {}; + case DescriptorReduceKind::AND: + return RMWOp::AND; + case DescriptorReduceKind::OR: + return RMWOp::OR; + case DescriptorReduceKind::XOR: + return RMWOp::XOR; + default: + break; + } + return {}; +} + +LogicalResult DescriptorReduceConverter::matchAndRewrite( + triton::DescriptorReduceOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto descTy = op.getDesc().getType(); + const auto blockShape = descTy.getBlockType().getShape(); + auto desc = unpackDescriptor(descTy, adaptor.getDesc(), rewriter); + auto offsets = castToI64(rewriter, op.getIndices()); + auto rmwOp = translateReduceKind(op.getKind(), descTy); + if (!rmwOp) { + std::string msgstring; + llvm::raw_string_ostream msg(msgstring); + msg << "Cannot fallback on descriptor atomic op, unsupported for type " + << descTy.getBlockType().getElementType(); + return op->emitError(msgstring); + } + + rewriter.create( + loc, descTy.getSignlessBlockType(), *rmwOp, + generatePtr(rewriter, loc, blockShape, desc, offsets), op.getSrc(), + generateMask(rewriter, loc, blockShape, desc, offsets), + MemSemantic::ACQUIRE_RELEASE, MemSyncScope::GPU); + op.erase(); + return success(); +} + +} // namespace DescriptorConverter diff --git a/compiler/lib/TritonToLinalg/FunctionConverter.cpp b/compiler/lib/TritonToLinalg/FunctionConverter.cpp new file mode 100644 index 00000000..ac2de14e --- /dev/null +++ b/compiler/lib/TritonToLinalg/FunctionConverter.cpp @@ -0,0 +1,36 @@ + + +#include "dicp/TritonToLinalg/FunctionConverter.h" + +namespace FunctionConverter { +using namespace mlir; +using namespace triton; + +LogicalResult GetProgramIDConverter::matchAndRewrite( + triton::GetProgramIdOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto axis = (uint32_t)op.getAxis(); + assert(axis < GetProgramIDConverter::LAUNCH_GRID_RANK && + "Invalid axis for GetProgramIdOp"); + auto func = op->getParentOfType(); + auto numArgs = func.getNumArguments(); + auto id = func.getArgument(numArgs - GetProgramIDConverter::LAUNCH_GRID_RANK + + axis); + rewriter.replaceOp(op, id); + return success(); +} + +LogicalResult GetNumProgramsConverter::matchAndRewrite( + triton::GetNumProgramsOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto axis = (uint32_t)op.getAxis(); + assert(axis < GetNumProgramsConverter::LAUNCH_GRID_RANK && + "Invalid axis for GetNumProgramsOp"); + auto func = op->getParentOfType(); + auto numArgs = func.getNumArguments(); + auto id = func.getArgument( + numArgs - GetNumProgramsConverter::LAUNCH_GRID_RANK * 2 + axis); + rewriter.replaceOp(op, id); + return success(); +} +} // namespace FunctionConverter diff --git a/compiler/lib/TritonToLinalg/HoistBroadcast.cpp b/compiler/lib/TritonToLinalg/HoistBroadcast.cpp new file mode 100644 index 00000000..e1baec43 --- /dev/null +++ b/compiler/lib/TritonToLinalg/HoistBroadcast.cpp @@ -0,0 +1,206 @@ + + +#include "dicp/TritonToLinalg/HoistBroadcast.h" +#include "dicp/TritonToLinalg/TritonToLinalgPass.h" +#include "dicp/Utils/Utils.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/LogicalResult.h" +#include "llvm/Support/raw_ostream.h" +#include + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/MemRef/Transforms/Passes.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/ValueRange.h" + +namespace HoistBroadcast { +using namespace mlir; +using namespace triton; + +LogicalResult +BroadcastConverter::matchAndRewrite(triton::BroadcastOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + assert(op->getNumResults() == 1 && "BroadcastOp assumes single result"); + + if (!isa(op.getType().getElementType())) { + return rewriter.notifyMatchFailure( + op, "only support hoist broadcast for tt.ptr tensor right now."); + } + auto loc = op.getLoc(); + BroadcastHoister hoister(op); + + if (!hoister.canBroadcast()) { + return failure(); + } + + if (hoister.parse(op.getSrc(), loc, rewriter).failed()) { + return failure(); + } + if (hoister.replaceBroadcastOp(op, rewriter).failed()) { + return failure(); + } + return success(); +} + +BroadcastHoister::BroadcastHoister(triton::BroadcastOp op) { + source = nullptr; + if (findSrc(op.getSrc()).failed()) { + LLVM_DEBUG({ llvm::dbgs() << "No legal source found for broadcast op\n"; }); + } + opToHoist = op; + auto resultType = dyn_cast(op.getType()); + for (size_t i = 0; i < resultType.getShape().size(); ++i) { + tensorSizes.push_back(resultType.getShape()[i]); + } +} + +LogicalResult BroadcastHoister::findSrc(Value operand) { + // ptr tensor can only be defined by AddPtrOp or SplatOp in this converter + // another broadcast to be complemented + if (auto op = operand.getDefiningOp()) { + return findSrc(op.getPtr()); + } else if (auto op = operand.getDefiningOp()) { + source = op.getSrc(); + return success(); + } else { + LLVM_DEBUG({ + llvm::dbgs() << "Unsupported operation in BroadcastHoister::findSrc: " + << *operand.getDefiningOp() << "\n"; + }); + return failure(); + } +} + +LogicalResult BroadcastHoister::parse(Value operand, const Location &loc, + ConversionPatternRewriter &rewriter) { + if (auto op = operand.getDefiningOp()) { + return parseAddptr(op, loc, rewriter); + } else if (auto op = operand.getDefiningOp()) { + return parseSplat(op, loc, rewriter); + } else if (auto op = operand.getDefiningOp()) { + return parseBroadcast(op, loc, rewriter); + } else { + // Handle other cases or throw an error + LLVM_DEBUG({ + llvm::dbgs() << "Unsupported operation in BroadcastHoister::parse: " + << *operand.getDefiningOp() << "\n"; + }); + return failure(); + } +} + +LogicalResult +BroadcastHoister::parseAddptr(triton::AddPtrOp addptrOp, const Location &loc, + ConversionPatternRewriter &rewriter) { + // Implementation for parsing AddptrOp + if (parse(addptrOp.getPtr(), loc, rewriter).failed()) { + return failure(); + } + auto broadcastedPtr = broadcastMap[addptrOp.getPtr()]; + + RankedTensorType offsetType = + dyn_cast(addptrOp.getOffset().getType()); + if (!offsetType || !offsetType.hasStaticShape()) { + LLVM_DEBUG({ + llvm::dbgs() << "Offset must be a ranked tensor with static shape.\n"; + }); + return failure(); + } + + auto elementType = offsetType.getElementType(); + auto broadcastType = RankedTensorType::get({tensorSizes}, elementType); + auto broadcastedOffset = rewriter.create( + loc, broadcastType, addptrOp.getOffset()); + + auto ptrType = dyn_cast(source.getType()); + auto ptrTensorType = RankedTensorType::get({tensorSizes}, ptrType); + + auto newAddPtrOp = rewriter.create( + loc, ptrTensorType, broadcastedPtr, broadcastedOffset); + + size_t hoistDim = -1; + for (size_t i = 0; i < offsetType.getShape().size(); ++i) { + if (offsetType.getShape()[i] == 1 && + offsetType.getShape()[i] != tensorSizes[i]) { + hoistDim = i; + break; + } + } + if (hoistDim == -1) { + LLVM_DEBUG({ + llvm::dbgs() << "No dimension to hoist found in AddPtrOp offset.\n"; + }); + } + newAddPtrOp->setAttr( + "hoist_dim", rewriter.getI64IntegerAttr(static_cast(hoistDim))); + broadcastMap[addptrOp.getResult()] = newAddPtrOp.getResult(); + return success(); +} + +LogicalResult +BroadcastHoister::parseSplat(triton::SplatOp splatOp, const Location &loc, + ConversionPatternRewriter &rewriter) { + // End of parse: splat for ptr + auto src = splatOp.getSrc(); + auto dst = splatOp.getResult(); + if (!isa(src.getType())) { + LLVM_DEBUG( + { llvm::dbgs() << "SplatOp source must be of pointer type.\n"; }); + return failure(); + } + + auto ptrType = dyn_cast(source.getType()); + auto ptrTensorType = RankedTensorType::get({tensorSizes}, ptrType); + auto newSplatOp = rewriter.create(loc, ptrTensorType, src); + broadcastMap[splatOp.getResult()] = newSplatOp.getResult(); + return success(); +} + +LogicalResult +BroadcastHoister::parseBroadcast(triton::BroadcastOp broadcastOp, + const Location &loc, + ConversionPatternRewriter &rewriter) { + // Another broadcast for ptr tensor + // To be fixed if needed in future + LLVM_DEBUG({ + llvm::dbgs() << "Now cannot handle multi broadcast for ptr tensor.\n"; + }); + return failure(); +} + +LogicalResult +BroadcastHoister::replaceBroadcastOp(triton::BroadcastOp op, + ConversionPatternRewriter &rewriter) { + auto newOp = broadcastMap[op.getSrc()]; + rewriter.replaceOp(op, newOp); + return success(); +} + +bool BroadcastHoister::canBroadcast() { + auto resultType = dyn_cast(opToHoist.getType()); + auto srcType = dyn_cast(opToHoist.getSrc().getType()); + int broadcastedDims = 0; + for (size_t i = 0; i < resultType.getShape().size(); ++i) { + if (srcType.getShape()[i] == 1 && resultType.getShape()[i] != 1) { + broadcastedDims++; + } + } + if (broadcastedDims != 1) { + LLVM_DEBUG({ + llvm::dbgs() << "Now cannot handle broadcast for ptr tensor with multi " + "broadcasted dimension.\n"; + }); + return false; + } + return source != nullptr && isa(source.getType()); +} +} // namespace HoistBroadcast diff --git a/compiler/lib/TritonToLinalg/ImplicitPermute.cpp b/compiler/lib/TritonToLinalg/ImplicitPermute.cpp new file mode 100644 index 00000000..27663323 --- /dev/null +++ b/compiler/lib/TritonToLinalg/ImplicitPermute.cpp @@ -0,0 +1,609 @@ + + +#include +#include +#include + +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Passes.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Location.h" +#include "mlir/IR/Matchers.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Value.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/MathExtras.h" + +#include "llvm/Support/Debug.h" + +#include "dicp/TritonToLinalg/ImplicitPermute.h" +#include "dicp/TritonToLinalg/TritonToLinalgPass.h" +#include "dicp/TritonToStructured/CannonicalizerConverter.h" +#include "dicp/TritonToStructured/MaskAnalysis.h" +#include "dicp/TritonToStructured/PtrAnalysis.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/Utils/Utils.h" +#include "dicp/Utils/InterleaveOptimization.h" +#include "dicp/Utils/Utils.h" + +#define DEBUG_TYPE "triton-to-linalg-implicit-permute" + +namespace ImplicitPermute { +using namespace mlir; +using namespace triton; +using namespace TritonToStructured; + +LogicalResult LoadConverter::matchAndRewrite(triton::LoadOp op, + PatternRewriter &rewriter) const { + // no need to analyze and rewrite + if (compileOn91095Flag && !existDotFlag) { + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "compileOn91095Flag :" << compileOn91095Flag << "\n"; + llvm::dbgs() << "!existDotFlag :" << !existDotFlag << "\n"; + llvm::dbgs() << "no need to analyze and rewrite Load" + << "\n"; + llvm::dbgs() << "----------------------------------------------\n"; + }); + return success(); + } + + auto loc = op.getLoc(); + auto oldPtr = op.getPtr(); + auto oldMask = op.getMask(); + auto oldOther = op.getOther(); + + MemOpTransformer tf(MemOpTransformer::MemType::load); + + Value newPtr = nullptr; + if (oldPtr.getDefiningOp()) { + newPtr = tf.createNewTensorPtr(oldPtr, loc, rewriter); + } else if (oldPtr.getDefiningOp()) { + newPtr = tf.createNewAdvancePtr(oldPtr, loc, rewriter); + } else if (oldPtr.getDefiningOp()) { + newPtr = tf.createNewAddPtr(oldPtr, loc, rewriter); + } else { + InFlightDiagnostic diag = emitWarning(loc) + << "PtrAnalysis: only MakeTensorPtrOp, " + "AdvanceOp, and AddPtrOp are supported."; + return success(); + } + if (!tf.ptrState.isPermuted) { + // no need to rewrite + return success(); + } + + auto newMask = tf.createNewMask(oldMask, loc, rewriter); + auto newOther = tf.createNewOther(oldOther, loc, rewriter); + + if (!newPtr) { + InFlightDiagnostic diag = emitWarning(loc) + << "PtrAnalysis: failed to analyze load pointer."; + return failure(); + } + + if (oldMask && !newMask) { + InFlightDiagnostic diag = emitWarning(loc) + << "MaskAnalysis: failed to analyze load mask."; + return failure(); + } + + auto newBoundaryCheck = tf.getBoundaryCheck(op.getBoundaryCheck()); + + auto loadOp = rewriter.create( + loc, newPtr, newMask, newOther, newBoundaryCheck, op.getPadding(), + op.getCache(), op.getEvict(), op.getIsVolatile()); + + auto permuteResult = + tf.materializeImplicitPermute(loadOp.getResult(), loc, rewriter); + + rewriter.replaceOp(op, permuteResult); + return success(); +} + +LogicalResult StoreConverter::matchAndRewrite(triton::StoreOp op, + PatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto oldPtr = op.getPtr(); + auto oldMask = op.getMask(); + auto oldValue = op.getValue(); + + MemOpTransformer tf(MemOpTransformer::MemType::store); + Value newPtr = nullptr; + if (oldPtr.getDefiningOp()) { + newPtr = tf.createNewTensorPtr(oldPtr, loc, rewriter); + } else if (oldPtr.getDefiningOp()) { + newPtr = tf.createNewAdvancePtr(oldPtr, loc, rewriter); + } else if (oldPtr.getDefiningOp()) { + newPtr = tf.createNewAddPtr(oldPtr, loc, rewriter); + } else { + InFlightDiagnostic diag = emitWarning(loc) + << "PtrAnalysis: only MakeTensorPtrOp, " + "AdvanceOp, and AddPtrOp are supported."; + return success(); + } + if (!tf.ptrState.isPermuted) { + // no need to rewrite + return success(); + } + auto newMask = tf.createNewMask(oldMask, loc, rewriter); + + if (!newPtr) { + InFlightDiagnostic diag = + emitWarning(loc) << "PtrAnalysis: failed to analyze store pointer."; + return failure(); + } + + if (oldMask && !newMask) { + InFlightDiagnostic diag = emitWarning(loc) + << "MaskAnalysis: failed to analyze store mask."; + return failure(); + } + + auto permuteResult = tf.materializeImplicitPermute(oldValue, loc, rewriter); + + auto newBoundaryCheck = tf.getBoundaryCheck(op.getBoundaryCheck()); + + auto storeOp = rewriter.create(loc, newPtr, permuteResult, + newMask, newBoundaryCheck, + op.getCache(), op.getEvict()); + + rewriter.eraseOp(op); + return success(); +} + +LogicalResult +AtomicRMWConverter::matchAndRewrite(triton::AtomicRMWOp op, + PatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto oldPtr = op.getPtr(); + auto oldMask = op.getMask(); + auto oldVal = op.getVal(); + + MemOpTransformer tf(MemOpTransformer::MemType::store); + + Value newPtr = nullptr; + if (oldPtr.getDefiningOp()) { + newPtr = tf.createNewAddPtr(oldPtr, loc, rewriter); + } else { + emitWarning(loc) << "PtrAnalysis: AtomicRMW only support AddPtrOp."; + return success(); + } + + Value newMask = tf.createNewMask(oldMask, loc, rewriter); + + if (!tf.ptrState.isPermuted) { + return success(); + } + if (!newPtr) { + emitWarning(loc) << "PtrAnalysis: failed to analyze atomic_rmw pointer."; + return failure(); + } + if (oldMask && !newMask) { + emitWarning(loc) << "MaskAnalysis: failed to analyze atomic_rmw mask."; + return failure(); + } + + Value newVal = tf.materializeImplicitPermute(oldVal, loc, rewriter); + + Type newAtomicResTy = newVal.getType(); + + auto newAtomic = rewriter.create( + loc, newAtomicResTy, op.getAtomicRmwOp(), newPtr, newVal, newMask, + op.getSem(), op.getScope()); + + // The returned old value should be in OLD layout for users => permute back + // (load-side). + MemOpTransformer tfLoad(MemOpTransformer::MemType::load); + tfLoad.ptrState = tf.ptrState; + Value permutedRes = + tfLoad.materializeImplicitPermute(newAtomic.getResult(), loc, rewriter); + + rewriter.replaceOp(op, permutedRes); + return success(); +} + +LogicalResult +AtomicCASConverter::matchAndRewrite(triton::AtomicCASOp op, + PatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto oldPtr = op.getPtr(); + auto oldCmp = op.getCmp(); + auto oldVal = op.getVal(); + + MemOpTransformer tf(MemOpTransformer::MemType::store); + + Value newPtr = nullptr; + if (oldPtr.getDefiningOp()) { + newPtr = tf.createNewAddPtr(oldPtr, loc, rewriter); + } else { + emitWarning(loc) << "PtrAnalysis: AtomicRMW only support AddPtrOp."; + return success(); + } + + if (!tf.ptrState.isPermuted) { + return success(); + } + if (!newPtr) { + emitWarning(loc) << "PtrAnalysis: failed to analyze atomic_cas pointer."; + return failure(); + } + + Value newCmp = tf.materializeImplicitPermute(oldCmp, loc, rewriter); + Value newVal = tf.materializeImplicitPermute(oldVal, loc, rewriter); + + // CAS result (old value) must have same shape as cmp/val operands. + Type newAtomicResTy = newVal.getType(); + + auto newAtomic = rewriter.create( + loc, newAtomicResTy, newPtr, newCmp, newVal, op.getSem(), op.getScope()); + + MemOpTransformer tfLoad(MemOpTransformer::MemType::load); + tfLoad.ptrState = tf.ptrState; + Value permutedRes = + tfLoad.materializeImplicitPermute(newAtomic.getResult(), loc, rewriter); + + rewriter.replaceOp(op, permutedRes); + return success(); +} + +Value MemOpTransformer::materializeImplicitPermute(Value srcTensor, + const Location loc, + PatternRewriter &rewriter) { + auto inTy = dyn_cast(srcTensor.getType()); + if (!inTy || !ptrState.isPermuted) + return srcTensor; + + auto inShape = inTy.getShape(); + auto orderSize = ptrState.sizes.size(); + SmallVector permuteOrder(orderSize); + for (size_t i = 0; i < orderSize; ++i) { + if (currentType == MemType::load) { + if (ptrState.isBlockPtr()) { + permuteOrder[i] = orderSize - 1 - ptrState.order[i]; + } else { + permuteOrder[ptrState.permuteIds[i]] = i; + } + } else { + if (ptrState.isBlockPtr()) { + permuteOrder[i] = ptrState.order[orderSize - 1 - i]; + } else { + permuteOrder[i] = ptrState.permuteIds[i]; + } + } + } + SmallVector outShape(permuteOrder.size()); + if (inShape.size() != outShape.size()) { + InFlightDiagnostic diag = emitWarning(loc) + << "PtrAnalysis: incompatible shape for permute"; + return srcTensor; + } + + for (size_t i = 0; i < outShape.size(); ++i) { + outShape[i] = inShape[permuteOrder[i]]; + } + + auto outTy = RankedTensorType::get(outShape, inTy.getElementType()); + auto transOp = + rewriter.create(loc, outTy, srcTensor, permuteOrder); + return transOp.getResult(); +} + +Value MemOpTransformer::createNewAddPtr(Value oldPtr, const Location loc, + PatternRewriter &rewriter) { + TritonToStructured::PtrAnalysis ptrAnalysis; + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "PtrAnalysis: analyzing load/store's ptr.\n"; + }); + + if (ptrAnalysis.visitOperand(oldPtr, ptrState, loc, rewriter).failed()) { + ptrState.shouldLinearize = false; + InFlightDiagnostic diag = + emitWarning(loc) << "PtranAlysis: failed to analyze load/store ptr."; + return oldPtr; + } + + // compute missing strides + // if stateinfo.shape is 1 and sizes[dimIndex] is 1, + // then the stride is the accumulated size of all dimensions on the right side + // ie. for shape [1, 128], sizes [1, 128], originally stride is [0, 1], + // after normalization, stride is [128, 1] + OpFoldResult maxStride = rewriter.getIndexAttr(1); + for (auto it = ptrState.stateInfo.rbegin(); it != ptrState.stateInfo.rend(); + ++it) { + if (TritonToStructured::isOne(it->shape) && isZero(it->stride)) { + it->stride = maxStride; + } + maxStride = maxOpFoldResult(maxStride, it->stride, loc, rewriter); + } + + ptrState.analyzePermute(); + return ptrState.createAddPtrOp(rewriter, loc); +} + +Value MemOpTransformer::createNewTensorPtr(Value oldPtr, const Location loc, + PatternRewriter &rewriter) { + TritonToStructured::PtrAnalysis ptrAnalysis; + auto makeTPtrOp = oldPtr.getDefiningOp(); + if (!makeTPtrOp) { + InFlightDiagnostic diag = emitWarning(loc) + << "PtrAnalysis: load pointer must originate " + "from 'make_tensor_ptr' operation"; + return oldPtr; + } + if (ptrAnalysis.visitOperandMakeTensorPtr(makeTPtrOp, ptrState, loc, rewriter) + .failed()) { + ptrState.isPermuted = false; + return oldPtr; + } + ptrState.analyzePermute(); + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "After ptrState.analyzePermute:\n"; + llvm::dbgs() << "compileOn91095Flag: " << compileOn91095Flag << "\n"; + ptrState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + return ptrState.createMakeTensorPtrOp(rewriter, loc); +} + +Value MemOpTransformer::createNewAdvancePtr(Value oldPtr, const Location loc, + PatternRewriter &rewriter) { + TritonToStructured::PtrAnalysis ptrAnalysis; + auto advOp = oldPtr.getDefiningOp(); + if (!advOp) { + emitWarning(loc) + << "PtrAnalysis: pointer must originate from 'advance' operation"; + return oldPtr; + } + + Value basePtr = advOp.getPtr(); + if (!basePtr || !basePtr.getDefiningOp()) { + emitWarning(loc) << "PtrAnalysis: advance base ptr must originate from " + "'make_tensor_ptr' operation"; + return oldPtr; + } + auto newBasePtr = createNewTensorPtr(basePtr, loc, rewriter); + if (!newBasePtr) + return oldPtr; + + if (!ptrState.isPermuted) + return oldPtr; + // 2) Rewrite advance offsets according to the new make_tensor_ptr layout. + auto oldOffsets = advOp.getOffsets(); + SmallVector newOffsets; + size_t rank = ptrState.order.size(); + // iterate reversed safely: i = rank-1, ..., 0 + for (size_t i = rank; i-- > 0;) { + newOffsets.push_back(oldOffsets[ptrState.order[i]]); + } + return rewriter.create(loc, newBasePtr.getType(), + newBasePtr, newOffsets); +} + +Value MemOpTransformer::createNewMask(Value oldMask, const Location loc, + PatternRewriter &rewriter) { + if (!oldMask) + return nullptr; + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "MaskAnalysis: analyzing load/store mask.\n"; + }); + + if (!oldMask || maskState.analysisMask(oldMask).failed()) { + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "MaskAnalysis: no mask or failed to analyze mask.\n"; + llvm::dbgs() << "oldMask:" << oldMask << "\n"; + maskState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + InFlightDiagnostic diag = + emitWarning(loc) << "MaskAnalysis: failed to analyze load/store mask."; + return nullptr; + } + + SmallVector newMaskInfo; + auto itPtr = ptrState.stateInfo.begin(); + auto itMask = maskState.stateInfo.begin(); + + // match and create new mask info + while (itPtr != ptrState.stateInfo.end() && + itMask != maskState.stateInfo.end()) { + // ptr'shape must be multiple of mask'shape or vice versa + if (!isMultiple(itMask->shape, itPtr->shape)) { + InFlightDiagnostic diag = + emitWarning(loc) + << "MaskAnalysis: incompatible shapes between ptr and mask."; + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + ptrState.dump(); + llvm::dbgs() << "oldMask:" << oldMask << "\n"; + maskState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + return nullptr; + } + + auto newShape = minOpFoldResult(itMask->shape, itPtr->shape, loc, rewriter); + if (isLess(newShape, itMask->shape) && !itMask->hasBroadCast) { + InFlightDiagnostic diag = + emitWarning(loc) + << "MaskAnalysis: the mask shape is incompatible with ptr shape."; + return nullptr; + } + + TritonToStructured::dimInfo newInfo(itMask->offset, newShape, + itMask->dimIndex, itMask->hasBroadCast, + itMask->currentType, itMask->rhs); + + if (!isZero(itPtr->stride)) { + newMaskInfo.emplace_back(newInfo); + } + + ++itPtr; + if (isEqual(itMask->shape, newShape)) { + ++itMask; + } + } + + if (itPtr != ptrState.stateInfo.end() || + itMask != maskState.stateInfo.end()) { + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "MaskAnalysis: failed to apply permute on mask.\n"; + ptrState.dump(); + llvm::dbgs() << "oldMask:" << oldMask << "\n"; + maskState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + InFlightDiagnostic diag = emitWarning(loc) + << "MaskAnalysis: incompatible number of " + "dimensions between ptr and mask."; + return nullptr; + } + + maskState.stateInfo = newMaskInfo; + + if (ptrState.isPermuted && !applyPermuteOnMask()) { + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "MaskAnalysis: failed to apply permute on mask.\n"; + ptrState.dump(); + llvm::dbgs() << "oldMask:" << oldMask << "\n"; + maskState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + InFlightDiagnostic diag = + emitWarning(loc) << "MaskAnalysis: failed to apply permute on mask."; + return nullptr; + } + + LLVM_DEBUG({ + llvm::dbgs() << "After matching MaskState: \n"; + for (auto info : newMaskInfo) { + info.dump(); + } + llvm::dbgs() << "----------------------------------------------\n"; + }); + + auto newMask = maskState.createNewMask(loc, rewriter); + return newMask; +} + +Value MemOpTransformer::createNewOther(Value oldOther, const Location loc, + PatternRewriter &rewriter) { + if (!oldOther || !maskState.newMask) + return nullptr; + + auto ptrType = dyn_cast(ptrState.source.getType()); + if (!ptrType) { + InFlightDiagnostic diag = + emitWarning(loc) + << "PtrAnalysis: source of ptrState is not a pointer type."; + return nullptr; + } + Type elementType = ptrType.getPointeeType(); + + SmallVector targetShape; + for (auto info : maskState.stateInfo) { + auto staticShape = getIntAttr(info.shape); + if (!staticShape.has_value()) { + InFlightDiagnostic diag = + emitWarning(loc) + << "MaskAnalysis: dynamic shape is not supported in reshape\n"; + return oldOther; + } + targetShape.emplace_back(staticShape.value()); + } + auto targetShapeAttr = DenseIntElementsAttr::get( + RankedTensorType::get({static_cast(targetShape.size())}, + rewriter.getI64Type()), + targetShape); + auto targetShapeType = RankedTensorType::get(targetShape, elementType); + auto targetShapeValue = + rewriter.create(loc, targetShapeAttr); + + auto reshapeOp = rewriter.create( + loc, targetShapeType, oldOther, targetShapeValue); + + return reshapeOp.getResult(); +} + +// Remap boundary_check for block ptr implicit permute. +// Formula +// - newAxis = rank - 1 - position(oldAxis in ptrState.order) +// Rules implemented +// - ptrState.order records original axes in memory-priority order. +// - createMakeTensorPtrOp rebuilds the new block ptr from +// reverse(ptrState.order), +// so the rebuilt ptr is canonicalized to descending order [rank-1, ..., 0]. +// - So each old boundary_check axis must be translated into the axis index +// of the rewritten block ptr before creating the new load/store. +// Examples +// - order=[1,0], boundary_check=[0] => [0] +// - order=[0,1], boundary_check=[0,1] => [1,0] +// - order=[2,0,1], boundary_check=[0,2] => [1,2] +SmallVector +MemOpTransformer::getBoundaryCheck(ArrayRef oldBoundaryCheck) const { + SmallVector newBoundaryCheck(oldBoundaryCheck.begin(), + oldBoundaryCheck.end()); + if (!ptrState.isPermuted || !ptrState.isBlockPtr() || + newBoundaryCheck.empty()) { + return newBoundaryCheck; + } + + int32_t rank = static_cast(ptrState.order.size()); + for (auto &boundaryAxis : newBoundaryCheck) { + auto pos = llvm::find(ptrState.order, static_cast(boundaryAxis)); + if (pos == ptrState.order.end()) { + continue; + } + boundaryAxis = + rank - 1 - + static_cast(std::distance(ptrState.order.begin(), pos)); + } + return newBoundaryCheck; +} + +bool MemOpTransformer::applyPermuteOnMask() { + if (!ptrState.isPermuted || maskState.isEmpty()) { + return true; + } + if (ptrState.permuteIds.size() != maskState.stateInfo.size()) { + return false; + } + SmallVector newMaskInfo; + for (auto id : ptrState.permuteIds) { + newMaskInfo.push_back(maskState.stateInfo[id]); + } + maskState.stateInfo = newMaskInfo; + return true; +} + +} // namespace ImplicitPermute diff --git a/compiler/lib/TritonToLinalg/LoadStoreConverter.cpp b/compiler/lib/TritonToLinalg/LoadStoreConverter.cpp new file mode 100644 index 00000000..e38fb4bd --- /dev/null +++ b/compiler/lib/TritonToLinalg/LoadStoreConverter.cpp @@ -0,0 +1,1242 @@ + + +#include "dicp/TritonToLinalg/LoadStoreConverter.h" +#include "dicp/TritonToLinalg/BlockPtrAnalysis.h" +#include "dicp/TritonToLinalg/MaskAnalysis.h" +#include "dicp/TritonToLinalg/TritonToLinalgPass.h" +#include "dicp/Utils/InterleaveOptimization.h" +#include "dicp/Utils/Utils.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HFusion/IR/HFusion.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" + +#include "dicp/Utils/Utils.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Passes.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Location.h" +#include "mlir/IR/Matchers.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Value.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/MathExtras.h" + +#include "llvm/Support/Debug.h" + +#include +#include +#include + +#define DEBUG_TYPE "triton-load-store-converter" + +namespace LoadStoreConverter { +using namespace mlir; +using namespace triton; + +const std::string MayImplicitTransposeWithLastAxisTAG = + "MayImplicitTransposeWithLastAxis"; + +LogicalResult +AddPtrConverter::matchAndRewrite(triton::AddPtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + llvm::SmallDenseMap known; + BlockDataParser::rewriteAddPtr(op, adaptor, rewriter, known); + return success(); +} + +LogicalResult LoadConverter::toTensorAndReplace( + triton::LoadOp &op, RankedTensorType &tensorType, Value localMem, + bool mayImplicitTransposeWithLastAxis, const Location &loc, + ConversionPatternRewriter &rewriter) const { + Value loadedTensor = rewriter.create( + loc, tensorType, localMem, true, true); + propagateWasBoolToInt8Attr(op.getOperation(), loadedTensor.getDefiningOp(), + rewriter); + + if (mayImplicitTransposeWithLastAxis) { + auto markOp = rewriter.create(loc, loadedTensor); + markOp->setAttr(MayImplicitTransposeWithLastAxisTAG, + UnitAttr::get(rewriter.getContext())); + } + rewriter.replaceOp(op, loadedTensor); + return success(); +} + +/// @brief Check whether the triton::LoadOp has been modified to the specified +/// state by the AddPtrConverter. +/// @param op The triton::LoadOp operation to be checked. +/// @return Return success if the operation conforms to the specified state; +/// otherwise, return failure. +LogicalResult +LoadConverter::checkModifiedByAddPtrConverter(triton::LoadOp &op) const { + if (!isa(op->getParentOp())) { + return failure(); + } + if (!op->hasAttr("IndirectLoad")) { + return failure(); + } + auto ptrOp = op.getPtr().getDefiningOp(); + auto ptrBlock = ptrOp->getBlock(); + auto opBlock = op->getBlock(); + if (ptrBlock == opBlock) { + return failure(); + } + + return success(); +} + +void LoadConverter::propagateWasBoolToInt8Attr( + Operation *srcLoadOp, Operation *dstOp, PatternRewriter &rewriter) const { + const std::string WasBoolToInt8TAG = "was_bool_to_int8"; + if (!srcLoadOp || !dstOp) + return; + if (srcLoadOp->hasAttr(WasBoolToInt8TAG)) { + dstOp->setAttr(WasBoolToInt8TAG, rewriter.getBoolAttr(true)); + } +} + +/// @brief Continue to modify the triton::LoadOp from the state modified by the +/// AddPtrConverter. +/// @param op The triton::LoadOp operation to be processed. +/// @param adaptor The adaptor for the operation, used to obtain operands. +/// @param rewriter The pattern rewriter used to rewrite the operation. +/// @return Return success if the operation is successful; otherwise, return +/// failure. +LogicalResult LoadConverter::continueModifyFromAddPtrConverter( + triton::LoadOp &op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto forOp = op->getParentOfType(); + Operation *firstOp = &forOp.getBody()->front(); + auto extractOp = cast(firstOp); + auto ivs = extractOp.getIndices(); + // Single iterArg which is inserted by AddPtrConverter. + auto iterArg = forOp.getRegionIterArg(0); + auto ptr = adaptor.getPtr(); + + rewriter.setInsertionPointAfter(op); + Value castVal = ptr.getDefiningOp(); + Value idxZero = + rewriter.create(loc, rewriter.getIndexAttr(0)); + Value loadVal = + rewriter.create(loc, castVal, ValueRange{idxZero}); + propagateWasBoolToInt8Attr(op.getOperation(), loadVal.getDefiningOp(), + rewriter); + Value insertedVal = + rewriter.create(loc, loadVal, iterArg, ValueRange{ivs}); + // a yield op is already created by AddPtrConverter. + // so we need to replace it with a new yield op. + Operation *terminator = forOp.getBody()->getTerminator(); + scf::YieldOp oldYieldOp = cast(terminator); + auto yieldOp = rewriter.create(loc, ValueRange{insertedVal}); + rewriter.replaceOp(oldYieldOp, yieldOp); + // Now the scf.for is complete, we can replace tt.load with it. + auto rank = cast(op.getResult().getType()).getShape().size(); + Operation *rootForOp = op; + while (rank != 0) { + rank--; + rootForOp = rootForOp->getParentOfType(); + } + rewriter.replaceOp(op, rootForOp); + LLVM_DEBUG({ llvm::dbgs() << *getModuleOpFromOperation(rootForOp) << "\n"; }); + return success(); +} + +void LoadConverter::fillTensorWithOtherForMaskScenario( + Value other, Value localMem, ArrayRef maskDim, + ConversionPatternRewriter &rewriter) const { + auto loc = localMem.getLoc(); + MemRefType originalType = cast(localMem.getType()); + assert(originalType.hasStaticShape() && "only support static shape"); + assert(originalType.getRank() == maskDim.size() && + "shape and mask must have same rank"); + + auto fillFlag = + rewriter.create(loc, rewriter.getBoolAttr(false)) + .getResult(); + + for (size_t i = 0; i < originalType.getShape().size(); ++i) { + // Use dynamic value to judge whether overstep boundary + auto shapeVal = rewriter.create( + loc, rewriter.getIndexAttr(originalType.getDimSize(i))); + + Value maskDimVal; + if (isa(maskDim[i])) + maskDimVal = rewriter.create( + loc, cast(cast(maskDim[i]))); + else + maskDimVal = cast(maskDim[i]); + + auto curCmp = rewriter.create(loc, arith::CmpIPredicate::slt, + maskDimVal, shapeVal); + + fillFlag = rewriter.create(loc, fillFlag, curCmp.getResult()) + .getResult(); + } + auto ifOp = rewriter.create(loc, fillFlag); + { + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front()); + rewriter.create(loc, ValueRange{other}, + ValueRange{localMem}); + } + ifOp->setAttr(rewriter.getStringAttr("hivm.unlikely_condition"), + UnitAttr::get(rewriter.getContext())); +} + +LoadConverter::LoadConverter(MLIRContext *context) + : OpConversionPattern(context) {} + +LogicalResult +LoadConverter::matchAndRewrite(triton::LoadOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + + // Check if tt.load is modified by AddPtrConverter to a specified state. + if (checkModifiedByAddPtrConverter(op).succeeded()) { + return continueModifyFromAddPtrConverter(op, adaptor, rewriter); + } + + auto ptr = adaptor.getPtr(); + auto mask = op.getMask(); + auto other = op.getOther(); + auto loc = op.getLoc(); + + // handling scalar + if (!isa(op.getResult().getType())) { + auto scalarMemref = + BlockDataParser::getScalarMemRef(op.getPtr(), ptr, loc, rewriter); + auto resTy = op.getResult().getType(); + auto idxZero = + rewriter.create(loc, rewriter.getIndexAttr(0)); + auto loadedValue = rewriter + .create(loc, resTy, scalarMemref, + idxZero.getResult()) + .getResult(); + propagateWasBoolToInt8Attr(op.getOperation(), loadedValue.getDefiningOp(), + rewriter); + if (mask && other) { + mask = rewriter.create( + loc, RankedTensorType::get({1}, mask.getType()), mask); + loadedValue = rewriter.create( + loc, RankedTensorType::get({1}, loadedValue.getType()), loadedValue); + other = rewriter.create( + loc, RankedTensorType::get({1}, other.getType()), other); + loadedValue = + rewriter.create(loc, mask, loadedValue, other); + rewriter.replaceOpWithNewOp(op, loadedValue, + ValueRange({idxZero})); + } else { + rewriter.replaceOp(op, loadedValue); + } + return success(); + } + + int64_t lastStride = -1; + if (isa(ptr)) { + auto u = ptr; + while (auto blkArg = dyn_cast(u)) { + if (auto forOp = dyn_cast(blkArg.getOwner()->getParentOp())) { + auto prt = forOp->getOperand(3 + blkArg.getArgNumber() - 1); + u = prt; + } else { + u = nullptr; + break; + } + } + if (u && isa(u.getDefiningOp())) { + auto ret = mlir::ConverterUtils::getLastStrideOfReinterpretCastOp( + dyn_cast(u.getDefiningOp())); + if (ret.has_value()) + lastStride = *ret; + } + } + + // handling no mask + auto memRefType = dyn_cast(ptr.getType()); + if (!memRefType) { + return rewriter.notifyMatchFailure( + op, "LoadOp expects a memref, not a memref of pointers"); + } + if (!op->hasAttr(ConverterUtils::GeneratedByMakeTensorPtrTAG)) { + auto memrefOp = dyn_cast(ptr.getDefiningOp()); + auto ret = mlir::ConverterUtils::getLastStrideOfReinterpretCastOp(memrefOp); + if (ret.has_value()) + lastStride = *ret; + } + bool mayImplicitTransposeWithLastAxis = + (existDotFlag) && + (!op->hasAttr(ConverterUtils::GeneratedByMakeTensorPtrTAG)) && + (lastStride != 1 && + mlir::ConverterUtils::isaPermutedMemRefType(memRefType)); + auto memRefShape = memRefType.getShape(); + auto memRefElementType = memRefType.getElementType(); + + Value allocOp; + Value allocOpTmp; + if (op->hasAttr(ConverterUtils::discreteAttrName)) { + Operation *loop = op->getParentOp(); + int extractedLoopCount = 1; + for (auto parentOp = loop->getParentOp(); + parentOp->hasAttr("ExtractedLoadOrStore"); + parentOp = parentOp->getParentOp()) { + loop = parentOp; + extractedLoopCount++; + } + rewriter.setInsertionPoint(loop); + auto loopOp = cast(loop); + auto fullMemRefShape = + cast(loopOp.getInitArgs()[0].getType()).getShape(); + auto fullMemRefType = MemRefType::get(fullMemRefShape, memRefElementType); + bool isIndexSelectScenario = + (extractedLoopCount == 1) && (fullMemRefShape.size() > 1u); + if (isIndexSelectScenario) + loopOp->setAttr("hivm.parallel_loop", rewriter.getUnitAttr()); + allocOp = rewriter.create(loc, fullMemRefType); + allocOpTmp = allocOp; + rewriter.setInsertionPointAfter(loop); + auto toTensorOp = rewriter.create( + loc, RankedTensorType::get(fullMemRefShape, memRefElementType), allocOp, + true, true); + rewriter.replaceAllUsesWith(loopOp->getResult(0), toTensorOp->getResult(0)); + tensor::InsertSliceOp insertSliceOp = nullptr; + for (auto *user : op->getUsers()) { + if (auto targetOp = dyn_cast(user)) { + insertSliceOp = targetOp; + break; + } + } + auto offsets = insertSliceOp.getMixedOffsets(); + auto sizes = insertSliceOp.getMixedSizes(); + auto strides = insertSliceOp.getMixedStrides(); + auto allocType = memref::SubViewOp::inferResultType(fullMemRefType, offsets, + sizes, strides); + rewriter.setInsertionPoint(op); + allocOp = rewriter.create( + loc, cast(allocType), allocOp, offsets, sizes, strides); + rewriter.replaceAllUsesExcept(insertSliceOp.getResult(), + insertSliceOp.getDest(), insertSliceOp); + rewriter.eraseOp(insertSliceOp); + } else { + allocOp = rewriter.create( + loc, MemRefType::get(memRefShape, memRefElementType)); + } + + auto tensorType = RankedTensorType::get(memRefShape, memRefElementType); + // boundary check + auto boundaryCheck = op.getBoundaryCheck(); + if (!boundaryCheck.empty()) { + auto makeTensorPtrOp = op.getPtr().getDefiningOp(); + auto boundarySizes = mlir::ConverterUtils::getBoundarySizes( + boundaryCheck, /*remapped*/ ptr, loc, rewriter); + // handle the padding + auto padding = op.getPadding(); + SmallVector srcOffsets(boundarySizes.size(), + rewriter.getIndexAttr(0)); + SmallVector dstOffsets; + if (makeTensorPtrOp) { + auto zeroVal = rewriter.createOrFold( + loc, rewriter.getI32IntegerAttr(0)); + for (auto [idx, offVal] : llvm::enumerate(makeTensorPtrOp.getOffsets())) { + if (llvm::find(boundaryCheck, idx) == boundaryCheck.end()) { + dstOffsets.push_back(srcOffsets[idx]); + continue; + } + Value offset = + rewriter.createOrFold(loc, zeroVal, offVal); + Value size = + getValueOrCreateConstantIndexOp(rewriter, loc, boundarySizes[idx]); + offset = rewriter.createOrFold(loc, offset, zeroVal); + offset = rewriter.createOrFold( + loc, rewriter.getIndexType(), offset); + OpFoldResult ofr; + if (auto constOp = offset.getDefiningOp()) { + ofr = constOp.getValue(); + } else { + ofr = offset; + } + ofr = minOpFoldResult(ofr, size, loc, rewriter); + boundarySizes[idx] = subOpFoldResult(size, ofr, loc, rewriter); + dstOffsets.push_back(ofr); + } + } else { + dstOffsets = srcOffsets; + } + if (padding.has_value()) { + TypedAttr padAttr = rewriter.getZeroAttr(memRefElementType); + // triton already ensure only NAN and ZERO are passed in + if (padding.value() == triton::PaddingOption::PAD_NAN) { + // FIXME: Why NaN requires elemTy to be non-int or non-index? + assert(!memRefElementType.isIntOrIndex()); + auto apNaN = llvm::APFloat::getNaN( + cast(padAttr).getValue().getSemantics()); + padAttr = rewriter.getFloatAttr(memRefElementType, apNaN); + } + auto padVal = rewriter.create(loc, padAttr); + + fillTensorWithOtherForMaskScenario(padVal, allocOp, boundarySizes, + rewriter); + } + auto srcSubView = mlir::ConverterUtils::makeSubViewOp( + ptr, srcOffsets, boundarySizes, loc, rewriter); + auto dstSubview = mlir::ConverterUtils::makeSubViewOp( + allocOp, dstOffsets, boundarySizes, loc, rewriter); + auto copyOp = rewriter.create(loc, srcSubView, dstSubview); + propagateWasBoolToInt8Attr(op.getOperation(), copyOp.getOperation(), + rewriter); + if (mayImplicitTransposeWithLastAxis) { + auto markOp = rewriter.create(loc, dstSubview); + markOp->setAttr(MayImplicitTransposeWithLastAxisTAG, + UnitAttr::get(rewriter.getContext())); + } + return this->toTensorAndReplace(op, tensorType, allocOp, + mayImplicitTransposeWithLastAxis, loc, + rewriter); + } + + if (!mask) { + assert(!other && "can not input 'other' when 'mask' is not set"); + if (auto unrealizedCastOp = + ptr.getDefiningOp()) { + // TODO : not support handle associate with "module" + // hint : can be handled in Linearize + op->emitError("meeting unexpected UCC in LoadConverter!"); + return failure(); + } else { + // If last dimension stride equals 2, try deinterleave optimization. + auto [ptrStrides, ptrOffsets] = memRefType.getStridesAndOffset(); + if (ptrStrides.back() == 2 && (memRefShape.back() % 2 == 0) && + mlir::triton::DeinterleaveStatusOptimization(op, adaptor, rewriter) + .succeeded()) { + return success(); + } + auto copyOp = rewriter.create(loc, ptr, allocOp); + propagateWasBoolToInt8Attr(op.getOperation(), copyOp.getOperation(), + rewriter); + if (mayImplicitTransposeWithLastAxis && + allocOp.getDefiningOp()) { + auto markOp = rewriter.create(loc, allocOp); + markOp->setAttr(MayImplicitTransposeWithLastAxisTAG, + UnitAttr::get(rewriter.getContext())); + } else if (mayImplicitTransposeWithLastAxis && + allocOp.getDefiningOp()) { + auto markOp = rewriter.create(loc, allocOpTmp); + markOp->setAttr(MayImplicitTransposeWithLastAxisTAG, + UnitAttr::get(rewriter.getContext())); + } + } + + return this->toTensorAndReplace(op, tensorType, allocOp, + mayImplicitTransposeWithLastAxis, loc, + rewriter); + } + + MaskState mstate; + auto isContMask = mstate.parse(mask, loc, rewriter); + if (isContMask.failed()) { + return rewriter.notifyMatchFailure( + op, "can not lower uncontinuout masked loads"); + } + + if (other) { + auto scalarOther = + mlir::ConverterUtils::getScalarValue(other, loc, rewriter); + assert( + scalarOther && + "other value used in masked load produced by unsupported instruction!"); + + fillTensorWithOtherForMaskScenario(scalarOther, allocOp, mstate.dims, + rewriter); + } + + // To enable deinterleave optimization with mask load, mask state along last + // dimension couldn't be split, which means `dims.back()` must be equal to + // origin type last dimension constant size and `offsets.back()` must be 0. + // + // The basis is that last dimension range comparison would generate + // unaccepted discontinuous mask. + if (mstate.getRank() == memRefType.getRank() && + isConstantIntValue(mstate.offsets.back(), 0) && + isConstantIntValue(mstate.dims.back(), memRefType.getShape().back())) { + auto [ptrStrides, ptrOffsets] = memRefType.getStridesAndOffset(); + if (ptrStrides.back() == 2 && (memRefType.getShape().back() % 2 == 0) && + DeinterleaveStatusWithMaskOptimization(op, adaptor, rewriter, mstate, + allocOp) + .succeeded()) { + return success(); + } + } + + if (auto unrealizedCastOp = ptr.getDefiningOp()) { + // TODO : not support handle associate with "module" + // hint : can be handled in Linearize + op->emitError("meeting unexpected UCC in LoadConverter!"); + return failure(); + } else { + if (mstate.isMemrefSubviewValid(ptr, rewriter)) { + memref::SubViewOp srcSubView = mstate.getSubview(ptr, loc, rewriter); + memref::SubViewOp dstSubView = mstate.getSubview(allocOp, loc, rewriter); + MemRefType dstSubViewType = mlir::cast(dstSubView.getType()); + + auto [srcStrides, srcOffset] = dstSubViewType.getStridesAndOffset(); + MemRefType castType = MemRefType::get( + dstSubViewType.getShape(), dstSubViewType.getElementType(), + makeStridedLinearLayoutMap(srcStrides, srcOffset, + rewriter.getContext())); + auto castOp = rewriter.create(loc, castType, dstSubView); + auto copyOp = rewriter.create(loc, srcSubView, castOp); + propagateWasBoolToInt8Attr(op.getOperation(), copyOp.getOperation(), + rewriter); + } + + if (mayImplicitTransposeWithLastAxis && + allocOp.getDefiningOp()) { + auto markOp = rewriter.create(loc, allocOp); + markOp->setAttr(MayImplicitTransposeWithLastAxisTAG, + UnitAttr::get(rewriter.getContext())); + } else if (mayImplicitTransposeWithLastAxis && + allocOp.getDefiningOp()) { + auto markOp = rewriter.create(loc, allocOpTmp); + markOp->setAttr(MayImplicitTransposeWithLastAxisTAG, + UnitAttr::get(rewriter.getContext())); + } + } + return this->toTensorAndReplace( + op, tensorType, allocOp, mayImplicitTransposeWithLastAxis, loc, rewriter); +} + +AtomicRMWConverter::AtomicRMWConverter(MLIRContext *context) + : OpConversionPattern(context) {} + +// lowering tt.atomicRMW to linalg.generic +// If atomic op's return value is used by other op as it's the old value stored +// at the ptrwe will use tt.load to get it +// +// example: +// input: +// %return_value = tt.atomic_rmw fadd, acq_rel, gpu, +// %output_memref, %input_tensor, %mask : +// (tensor<256x!tt.ptr>, tensor<256xf32>, tensor<256xi1>) +// -> tensor<256xf32> +// +// output: +// memref.copy %output_memref, %ub_buf : memref to memref +// %17 = bufferization.to_tensor %alloc_3 restrict writable : memref<256xf32> +// linalg.generic +// {indexing_maps = [#map, #map, #map], iterator_types = ["parallel"]} +// ins(%output_memref, %masked_input_memref : memref, memref) +// outs(%subview_2 : memref) +// attrs = {GenericAtomicRMW = "fadd", MemSemantic = "acq_rel", +// MemSyncScope = "gpu"} { +// ^bb0(%in: f32, %in_9: f32, %out: f32): +// %25 = arith.addf %in, %in_9 : f32 +// linalg.yield %25 : f32 +// } +LogicalResult +AtomicRMWConverter::matchAndRewrite(triton::AtomicRMWOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto ptr = adaptor.getPtr(); + auto val = op.getVal(); + auto loc = op.getLoc(); + auto mask = op.getMask(); + auto rmwOp = op.getAtomicRmwOp(); + auto resType = dyn_cast(op.getResult().getType()); + auto ptrType = dyn_cast(ptr.getType()); + + if (!resType) + return rewriter.notifyMatchFailure( + op, "atomicRMWConverter: scalar will be handled by " + "ScalarAtomicRMWCanonicalizer"); + if (!ptrType) + return rewriter.notifyMatchFailure( + op, "AtomicRMWOp expects a memref, not a memref of pointers"); + + const std::map atomicKindMap = { + {RMWOp::ADD, hivm::AtomicKind::ADD}, + {RMWOp::FADD, hivm::AtomicKind::ADD}, + {RMWOp::OR, hivm::AtomicKind::OR}, + {RMWOp::XOR, hivm::AtomicKind::XOR}, + {RMWOp::AND, hivm::AtomicKind::AND}, + {RMWOp::MIN, hivm::AtomicKind::MIN}, + {RMWOp::UMIN, hivm::AtomicKind::UMIN}, + {RMWOp::MAX, hivm::AtomicKind::MAX}, + {RMWOp::UMAX, hivm::AtomicKind::UMAX}, + {RMWOp::XCHG, hivm::AtomicKind::XCHG}, + }; + const std::map hfusionAtomicKindMap = { + {RMWOp::ADD, hfusion::AtomicKind::ADD}, + {RMWOp::FADD, hfusion::AtomicKind::ADD}, + {RMWOp::OR, hfusion::AtomicKind::OR}, + {RMWOp::XOR, hfusion::AtomicKind::XOR}, + {RMWOp::AND, hfusion::AtomicKind::AND}, + {RMWOp::MIN, hfusion::AtomicKind::MIN}, + {RMWOp::UMIN, hfusion::AtomicKind::UMIN}, + {RMWOp::MAX, hfusion::AtomicKind::MAX}, + {RMWOp::UMAX, hfusion::AtomicKind::UMAX}, + {RMWOp::XCHG, hfusion::AtomicKind::XCHG}, + }; + assert(atomicKindMap.find(rmwOp) != atomicKindMap.end()); + auto atomicKind = + hivm::AtomicKindAttr::get(rewriter.getContext(), atomicKindMap.at(rmwOp)); + assert(hfusionAtomicKindMap.find(rmwOp) != hfusionAtomicKindMap.end()); + auto hfusionAtomicKind = hfusion::AtomicKindAttr::get( + rewriter.getContext(), hfusionAtomicKindMap.at(rmwOp)); + + auto dstMemref = ptr; + Value inputVal = val; + + // Lazily materialize a memref view only when we truly need buffer + // semantics (e.g., mask subview or XCHG lowering). Otherwise keep tensor + // inputs to avoid redundant to_memref conversions before hivm.store. + auto getInputMemref = [&]() -> Value { + if (isa(inputVal.getType())) + return inputVal; + return rewriter.create(loc, ptrType, inputVal); + }; + auto inputMemref = getInputMemref(); + auto inputMemrefType = cast(inputMemref.getType()); + auto elementType = inputMemrefType.getElementType(); + auto isHardwareSupported = + (rmwOp == RMWOp::ADD || rmwOp == RMWOp::FADD || rmwOp == RMWOp::MAX || + rmwOp == RMWOp::MIN) && + (elementType.isF16() || elementType.isBF16() || elementType.isF32() || + elementType.isInteger(8) || elementType.isInteger(16) || + elementType.isInteger(32)); + + bool isDiscreteMask = false; + if (mask) { + auto constantMask = mask.getDefiningOp(); + if (constantMask && !isConstantMaskTrue(mask)) { + rewriter.eraseOp(op); + return success(); + } + MaskState mstate; + isDiscreteMask = mstate.parse(mask, loc, rewriter).failed(); + if (!constantMask && !isDiscreteMask) { + // For dstMemref (store output), use subview to maintain reference to + // original memref. For inputVal (store input), use tensor.extract_slice + // to keep tensor semantics. + dstMemref = mstate.getSubview(ptr, loc, rewriter); + if (isHardwareSupported) { + auto inputTensorType = RankedTensorType::get( + inputMemrefType.getShape(), inputMemrefType.getElementType()); + if (!isa(inputVal.getType())) + inputVal = rewriter.create( + loc, inputTensorType, inputMemref, true, true); + inputVal = mstate.getExtractSlice(inputVal, loc, rewriter); + } else { + inputMemref = mstate.getSubview(inputMemref, loc, rewriter); + } + } + } + + if (!op.getResult().use_empty()) { + auto tensorType = + RankedTensorType::get(ptrType.getShape(), ptrType.getElementType()); + auto alloc = rewriter.create( + loc, MemRefType::get(ptrType.getShape(), ptrType.getElementType())); + rewriter.create(loc, ptr, alloc); + Value tensorToReplace = rewriter.create( + loc, tensorType, alloc, true /* restrict */, true /* writable */); + rewriter.replaceOp(op, tensorToReplace); + } + + if (isDiscreteMask) { + if (rmwOp != RMWOp::XCHG) { + return op.emitError( + "Discrete mask is only expected for XCHG; other atomics " + "should be lowered without discrete masks"); + } + Value memrefMask = mask; + if (auto maskTypeT = dyn_cast(mask.getType())) { + MemRefType maskTypeM = + MemRefType::get(maskTypeT.getShape(), maskTypeT.getElementType()); + memrefMask = + rewriter.create(loc, maskTypeM, mask); + } + rewriter.create(op.getLoc(), TypeRange(), + inputMemref, dstMemref, memrefMask); + } else { + if (isHardwareSupported) + rewriter.create(op.getLoc(), TypeRange{}, inputVal, + dstMemref, atomicKind); + else if (rmwOp == RMWOp::XCHG) + rewriter.create(op.getLoc(), TypeRange(), + inputMemref, dstMemref); + else { + if (rmwOp == RMWOp::OR || rmwOp == RMWOp::XOR || rmwOp == RMWOp::AND) { + if (!elementType.isSignlessIntOrIndex()) { + return op->emitOpError() + << "must be signless-integer-like, but got " << elementType; + } + } + // Currently, for atomic kind and element type that is not supported by + // the hardware, we use software to simulate the computation. However, + // decompose now happens in both HFusion and HIVM, and is not consistent + // for 910B and 91095. Therefore, we convert to different atomic/store ops + // for now. This should be unified and refactored later. + if (compileOn91095Flag) { + rewriter.create( + op.getLoc(), TypeRange{}, ValueRange{inputMemref}, + ValueRange{dstMemref}, hfusionAtomicKind, + ArrayRef{}); + } else { + rewriter.create(op.getLoc(), TypeRange(), + inputMemref, dstMemref, + hfusionAtomicKind); + } + } + } + + if (op.getResult().use_empty()) { + rewriter.eraseOp(op); + } + return success(); +} + +LogicalResult +AtomicCASConverter::matchAndRewrite(triton::AtomicCASOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + // If the result of AtomicCASOp is not used, we don't need to load the old + // data stored at the ptr + auto ptr = adaptor.getPtr(); + auto cmp = op.getCmp(); + auto val = op.getVal(); + auto loc = op.getLoc(); + + auto resType = dyn_cast(op.getResult().getType()); + if (!resType) { + return rewriter.notifyMatchFailure( + op, "atomicCASConverter: scalar will be handled by " + "ScalarAtomicCASCanonicalizer"); + } + + // 1. Simple case where no mask is used. + auto type = dyn_cast(ptr.getType()); + if (!type) { + // Seen when implicit broadcasting is done late in a chain of + // operations. The workaround is to broadcast the pointers early in the + // address calculation. A proper fix is complicated, but at least we can + // provide a better error message. + return rewriter.notifyMatchFailure( + op, "AtomicCASOp expects a memref, not a memref of pointers"); + } + + auto dstMemref = ptr; + // Well, linalg structure op wouldn't support mixed tensor/buffer semantics + // any more in latest LLVM(triton LLVM dependency has involed this), so we + // need to convert tensor to buffer early. + auto dstOriType = cast(dstMemref.getType()); + MemRefType dstType = + MemRefType::get(dstOriType.getShape(), dstOriType.getElementType()); + Value inputMemref = + rewriter.create(loc, dstType, val); + + Value cmpMemref = + rewriter.create(loc, dstType, cmp); + + // create element-wise map + int64_t rank = type.getRank(); + SmallVector inputDims; + auto context = rewriter.getContext(); + + for (int i = 0; i < rank; i++) { + inputDims.push_back(getAffineDimExpr(i, context)); + } + + SmallVector indexingMaps; + // As mask has been erased for now + // the number of input must be 2 + // the input memref is also the output memref + // Thus, there are a total of four inputs and outputs. + // so here we have 4 map to create + for (int i = 0; i < 4; i++) { // 4: 3 input and 1 output + indexingMaps.push_back(AffineMap::get(rank, 0, inputDims, context)); + } + + if (!op.getResult().use_empty()) { + auto tensorType = + RankedTensorType::get(type.getShape(), type.getElementType()); + auto alloc = rewriter.create( + loc, MemRefType::get(type.getShape(), type.getElementType())); + + // For the return value, don't need to care about mask for now + // this op don't support other, so we best not fill it + rewriter.create(loc, ptr, alloc); + Value tensor = rewriter.create( + loc, tensorType, alloc, true /* restrict */, true /* writable */); + rewriter.replaceOp(op, tensor); + } + + auto linalgOp = rewriter.create( + loc, ValueRange{dstMemref, cmpMemref, inputMemref}, + mlir::ValueRange{dstMemref}, indexingMaps, + mlir::ConverterUtils::getNParallelLoopsAttrs(rank), + [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange blockArgs) { + Value lhs = blockArgs[0]; + Value rhs = blockArgs[1]; + Value setValue = blockArgs[2]; + Value cond; + if (mlir::isa(lhs.getType())) { + cond = nestedBuilder.create( + nestedLoc, arith::CmpFPredicate::UEQ, lhs, rhs); + } else { + cond = nestedBuilder.create( + nestedLoc, arith::CmpIPredicate::eq, lhs, rhs); + } + auto ifOp = nestedBuilder.create( + nestedLoc, TypeRange{setValue.getType()}, cond, true); + { + OpBuilder::InsertionGuard guard(nestedBuilder); + nestedBuilder.setInsertionPointToEnd(&ifOp.getThenRegion().front()); + nestedBuilder.create(nestedLoc, setValue); + } + { + OpBuilder::InsertionGuard guard(nestedBuilder); + nestedBuilder.setInsertionPointToEnd(&ifOp.getElseRegion().front()); + nestedBuilder.create(nestedLoc, lhs); + } + nestedBuilder.setInsertionPointToEnd(nestedBuilder.getBlock()); + nestedBuilder.create(nestedLoc, + ifOp.getResult(0)); + }); + + const StringRef genericAtomicRMW = "GenericAtomicRMW"; + const StringRef memSemantic = "MemSemantic"; + const StringRef memSyncScope = "MemSyncScope"; + auto attr = mlir::StringAttr::get(context, "cas"); + + linalgOp->setAttr(genericAtomicRMW, attr); + linalgOp->setAttr(memSemantic, + rewriter.getStringAttr(stringifyEnum(op.getSem()))); + linalgOp->setAttr(memSyncScope, + rewriter.getStringAttr(stringifyEnum(op.getScope()))); + + linalgOp->setAttr("Software", rewriter.getUnitAttr()); + + // tt.atomicRMW op has two part of feature + // 1. load the old data at the ptr + // 2. atomically store the data on ub to the ptr + // at the same time it perform the action it has been assigned + // So we lower this op to load + atomically store + // + // The first part is not necessary when the returned value of atomic op + // is not used, it will be deleted cause it's meaningless + // Here, we preemptively determine whether it will be used + // and decide whether it is necessary to create the load process based on + // this assessment. + // + // logic of handling is copied + if (op.getResult().use_empty()) { + rewriter.eraseOp(op); + } + return success(); +} + +LogicalResult +ScalarStoreCanonicalizer::matchAndRewrite(triton::StoreOp op, + PatternRewriter &rewriter) const { + if (!op.getValue().getType().isIntOrIndexOrFloat()) { + return rewriter.notifyMatchFailure( + op, "ScalarStoreCanonicalizer handles scalar store scene!"); + } + auto ptr = op.getPtr(); + auto mask = op.getMask(); + auto value = op.getValue(); + if (mask) { + rewriter.replaceOpWithNewOp( + op, mask, [&](OpBuilder &b, Location loc) { + b.create(loc, ptr, value, op.getCache(), + op.getEvict()); + b.create(loc); + }); + return success(); + } + + auto ptrTy = RankedTensorType::get({(int64_t)1}, ptr.getType()); + auto ptrSplat = rewriter.create(op.getLoc(), ptrTy, ptr); + auto valTy = RankedTensorType::get({(int64_t)1}, value.getType()); + auto valSplat = rewriter.create(op.getLoc(), valTy, value); + auto newStoreOp = rewriter.create( + op.getLoc(), ptrSplat, valSplat, op.getCache(), op.getEvict()); + rewriter.replaceOp(op, newStoreOp); + return success(); +} + +LogicalResult +ScalarAtomicRMWCanonicalizer::matchAndRewrite(triton::AtomicRMWOp op, + PatternRewriter &rewriter) const { + if (!op.getVal().getType().isIntOrIndexOrFloat()) { + return rewriter.notifyMatchFailure( + op, "ScalarAtomicRMWCanonicalizer handles scalar atomic rmw op scene!"); + } + + auto ptr = op.getPtr(); + auto ptrTy = RankedTensorType::get({(int64_t)1}, ptr.getType()); + auto ptrSplat = rewriter.create(op.getLoc(), ptrTy, ptr); + auto valTy = RankedTensorType::get({(int64_t)1}, op.getVal().getType()); + auto valSplat = + rewriter.create(op.getLoc(), valTy, op.getVal()); + auto maskTy = RankedTensorType::get({(int64_t)1}, op.getMask().getType()); + auto maskSplat = + rewriter.create(op.getLoc(), maskTy, op.getMask()); + + auto newAtomicOp = rewriter.create( + op.getLoc(), valTy, op.getAtomicRmwOp(), ptrSplat, valSplat, maskSplat, + op.getSem(), op.getScope()); + auto idxZero = + rewriter.create(op.getLoc(), rewriter.getIndexAttr(0)); + rewriter.replaceOpWithNewOp(op, newAtomicOp, + ValueRange({idxZero})); + return success(); +} + +LogicalResult +ScalarAtomicCASCanonicalizer::matchAndRewrite(triton::AtomicCASOp op, + PatternRewriter &rewriter) const { + if (!op.getVal().getType().isIntOrIndexOrFloat() && + !op.getCmp().getType().isIntOrIndexOrFloat()) { + return rewriter.notifyMatchFailure( + op, "ScalarAtomicCASCanonicalizer handles scalar atomic cas op scene!"); + } + + auto ptr = op.getPtr(); + auto ptrTy = RankedTensorType::get({(int64_t)1}, ptr.getType()); + auto ptrSplat = rewriter.create(op.getLoc(), ptrTy, ptr); + auto cmpTy = RankedTensorType::get({(int64_t)1}, op.getCmp().getType()); + auto cmpSplat = + rewriter.create(op.getLoc(), cmpTy, op.getCmp()); + auto valTy = RankedTensorType::get({(int64_t)1}, op.getVal().getType()); + auto valSplat = + rewriter.create(op.getLoc(), valTy, op.getVal()); + + auto newAtomicOp = rewriter.create( + op.getLoc(), valTy, ptrSplat, cmpSplat, valSplat, op.getSem(), + op.getScope()); + auto idxZero = + rewriter.create(op.getLoc(), rewriter.getIndexAttr(0)); + rewriter.replaceOpWithNewOp(op, newAtomicOp, + ValueRange({idxZero})); + return success(); +} + +// The atomic max op with float input will be devided into +// two atomic max ops with integer input +// One handles the part of the tensor greater than zero +// the other deals with the part less than zero +// It will lead to maskAnalysis failure +// So here we need to revert the procedures in semantics.py +// The triton IR is like +// +// %cst_0 = arith.constant dense<0.000000e+00> : tensor<1x256xf32> +// %1 = tt.bitcast %value : tensor<1x256xf32> -> tensor<1x256xi32> +// %2 = tt.bitcast %ptr : tensor<1x256x!tt.ptr> -> +// tensor<1x256x!tt.ptr> %3 = arith.cmpf oge, %1, %cst_0 %4 = arith.cmpf +// olt, %1, %cst_0 %5 = arith.andi %8, %3 %6 = tt.atomic_rmw max, acq_rel, gpu, +// %2, %1, %5 : +// (tensor<1x256x!tt.ptr>, tensor<1x256xi32>, tensor<1x256xi1>) -> +// tensor<1x256xi32> +// %7 = arith.andi %8, %4 +// %8 = tt.atomic_rmw umin, acq_rel, gpu, %2, %1, %7 : +// (tensor<1x256x!tt.ptr>, tensor<1x256xi32>, tensor<1x256xi1>) -> +// tensor<1x256xi32> +// +// it's hard to handle and meaningless complicated for our device +// so we revert it to +// %0 = tt.atomic_rmw max, acq_rel, gpu, %23, %21, %8 : +// (tensor<1x256x!tt.ptr>, tensor<1x256xf32>, tensor<1x256xi1>) -> +// tensor<1x256xf32> +LogicalResult +AtomicMaxMinCanonicalizer::matchAndRewrite(triton::AtomicRMWOp op, + PatternRewriter &rewriter) const { + // Revert the op to its original form + auto ptrBitcastOp = op.getPtr().getDefiningOp(); + auto valueBitcastOp = op.getVal().getDefiningOp(); + if (!ptrBitcastOp || !valueBitcastOp) { + return failure(); + } + + // We only need to handle the op when the element type is float + auto elementType = + dyn_cast(valueBitcastOp.getSrc().getType()).getElementType(); + if (!isa(elementType)) { + return failure(); + } + + auto rmwOp = op.getAtomicRmwOp(); + // here we know that atomic UMAX/UMIN + // is created by special logic of triton right now + // so we can simply delete it + if (rmwOp == triton::RMWOp::UMAX || rmwOp == triton::RMWOp::UMIN) { + // if the return value of op is used, we can't simply erase it + if (op.getResult().use_empty()) { + rewriter.eraseOp(op); + return success(); + } + return failure(); + } + + if (rmwOp != triton::RMWOp::MAX && rmwOp != triton::RMWOp::MIN) { + return failure(); + } + + // 1. Though semantic interpreter will generate full true tensor as original + // mask if atomicrmwOp don't have it, above float devision process will also + // generate positive and negative comparison mask, which will cause to fold + // true mask. + // 2. While if atomicrmwOp has original mask, there exists andiop between + // original mask and positive/negative comparison mask + // + // Here wanna extract original mask + Value originalMask = op.getMask(); + if (auto andOp = originalMask.getDefiningOp()) + // LHS is convention in semantic interpreter + originalMask = andOp.getLhs(); + else if (auto cmpOp = originalMask.getDefiningOp()) { + if (cmpOp.getPredicate() != mlir::arith::CmpFPredicate::OGE || + !matchPattern(cmpOp.getRhs(), + /*positive float zero matcher*/ m_PosZeroFloat())) + // Here recheck frontend interpreter generation in no manual mask state + return op->emitError("Illegal mask for atomicrmwOp of float type"); + // Restore original true mask + originalMask = rewriter.create( + op->getLoc(), + /*typed attr*/ DenseElementsAttr::get( + cast(originalMask.getType()), true)); + } else + return op->emitError("Illegal mask for atomicrmwOp of float type"); + + auto originAtomicOp = rewriter.create( + op.getLoc(), valueBitcastOp.getSrc().getType(), op.getAtomicRmwOp(), + ptrBitcastOp.getSrc(), valueBitcastOp.getSrc(), originalMask, op.getSem(), + op.getScope()); + + // if the return value of op is used + // we need to handle its usage + // In semantic.py, if the atomic Max/Min with float input is used + // It will use select + bitcast to get float value + // so here we need to revert it too + // + // For example: + // %0 = tt.atomic_rmw max, acq_rel, gpu, %gm, %input, %mask1 : + // (tensor<32x!tt.ptr>... %1 = tt.atomic_rmw umin, acq_rel, gpu, %gm, + // %input, %mask2 : (tensor<32x!tt.ptr>... %2 = arith.select + // %devidedMask, %0, %1 : tensor<32xi1>, tensor<32xi32> %3 = tt.bitcast %2 : + // tensor<32xi32> -> tensor<32xf32> tt.store %outputMemref, %3 : + // tensor<32x!tt.ptr> + // + // will be revert to: + // %0 = tt.atomic_rmw max, acq_rel, gpu, %gm, %input, %mask : + // (tensor<32x!tt.ptr>... tt.store %outputMemref, %0 : + // tensor<32x!tt.ptr> + // + if (!op.getResult().use_empty()) { + for (OpOperand &use : op->getUses()) { + auto selectOp = dyn_cast(use.getOwner()); + if (!selectOp) + continue; + + for (OpOperand &selectUse : selectOp->getUses()) { + if (auto bitcastOp = + dyn_cast(selectUse.getOwner())) { + bitcastOp.getResult().replaceAllUsesWith(originAtomicOp); + } + } + } + rewriter.replaceOp(op, originAtomicOp); + } else { + rewriter.eraseOp(op); + } + + return success(); +} + +StoreConverter::StoreConverter(MLIRContext *context) + : OpConversionPattern(context) {} + +LogicalResult +StoreConverter::matchAndRewrite(triton::StoreOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + + // triton store op basic + auto mask = op.getMask(); + auto loc = op.getLoc(); + auto ptr = adaptor.getPtr(); + auto val = adaptor.getValue(); + + // 1. boundary size check + auto boundaryCheck = op.getBoundaryCheck(); + if (!boundaryCheck.empty()) { + auto makeTensorPtrOp = op.getPtr().getDefiningOp(); + auto boundarySizes = mlir::ConverterUtils::getBoundarySizes( + boundaryCheck, /*remapped*/ ptr, loc, rewriter); + SmallVector srcOffsets; + SmallVector dstOffsets(boundarySizes.size(), + rewriter.getIndexAttr(0)); + if (makeTensorPtrOp) { + auto zeroVal = rewriter.createOrFold( + loc, rewriter.getI32IntegerAttr(0)); + for (auto [idx, offVal] : llvm::enumerate(makeTensorPtrOp.getOffsets())) { + if (llvm::find(boundaryCheck, idx) == boundaryCheck.end()) { + srcOffsets.push_back(dstOffsets[idx]); + continue; + } + Value offset = + rewriter.createOrFold(loc, zeroVal, offVal); + Value size = + getValueOrCreateConstantIndexOp(rewriter, loc, boundarySizes[idx]); + offset = rewriter.createOrFold(loc, offset, zeroVal); + offset = rewriter.createOrFold( + loc, rewriter.getIndexType(), offset); + OpFoldResult ofr; + if (auto constOp = offset.getDefiningOp()) { + ofr = constOp.getValue(); + } else { + ofr = offset; + } + ofr = minOpFoldResult(ofr, size, loc, rewriter); + boundarySizes[idx] = subOpFoldResult(size, ofr, loc, rewriter); + srcOffsets.push_back(ofr); + } + } else { + srcOffsets = dstOffsets; + } + auto srcSlice = mlir::ConverterUtils::makeExtractSliceOp( + val, srcOffsets, boundarySizes, loc, rewriter); + auto dstSubview = mlir::ConverterUtils::makeSubViewOp( + ptr, dstOffsets, boundarySizes, loc, rewriter); + auto storeOp = rewriter.create( + loc, srcSlice, dstSubview); + storeOp.setWritable(true); + rewriter.eraseOp(op); + return success(); + } + + // 2. Simple load with no mask + if (!mask) { + auto storeOp = rewriter.create( + loc, val, ptr); + storeOp.setWritable(true); + rewriter.eraseOp(op); + return success(); + } + + // 3. Continuous masked stores. + // Analyze the mask operand to determine at runtime the size of the data we + // are moving. + MaskState mstate; + auto isContMask = mstate.parse(mask, loc, rewriter); + + if (isContMask.failed()) { + return failure(); + } + LLVM_DEBUG({ llvm::dbgs() << *getModuleOpFromOperation(op) << "\n"; }); + auto srcSlice = mstate.getExtractSlice(val, loc, rewriter); + auto dstSubview = mstate.getSubview(ptr, loc, rewriter); + auto storeOp = rewriter.create( + loc, srcSlice, dstSubview); + storeOp.setWritable(true); + rewriter.eraseOp(op); + return success(); +} + +bool ReinterpretCastStrideCanonicalizer::hasFixableZeroStride( + memref::ReinterpretCastOp op) { + auto staticSizes = op.getStaticSizes(); + auto staticStrides = op.getStaticStrides(); + auto dynamicStrides = op.getStrides(); + + if (staticSizes.size() != staticStrides.size()) + return false; + + // now handle: size all static + if (llvm::any_of(staticSizes, ShapedType::isDynamic)) + return false; + + unsigned dynStrideIdx = 0; + for (unsigned i = 0; i < staticStrides.size(); ++i) { + // now handle: dynamic stride 0 with static size 1 + if (!ShapedType::isDynamic(staticStrides[i])) + continue; + + if (dynStrideIdx >= dynamicStrides.size()) + return false; + + Value st = dynamicStrides[dynStrideIdx]; + dynStrideIdx++; + if (staticSizes[i] == 1 && mlir::isZero(OpFoldResult(st))) + return true; + } + return false; +} + +LogicalResult ReinterpretCastStrideCanonicalizer::matchAndRewrite( + memref::ReinterpretCastOp op, PatternRewriter &rewriter) const { + if (!hasFixableZeroStride(op)) + return failure(); + + auto staticSizes = op.getStaticSizes(); + auto staticStrides = op.getStaticStrides(); + auto dynamicStrides = op.getStrides(); + + SmallVector newDynamicStrides; + newDynamicStrides.reserve(dynamicStrides.size()); + + unsigned dynStrideIdx = 0; + bool changed = false; + Value c1 = + rewriter.create(op.getLoc(), rewriter.getIndexAttr(1)); + + for (unsigned i = 0, e = staticStrides.size(); i < e; ++i) { + if (!ShapedType::isDynamic(staticStrides[i])) + continue; + + if (dynStrideIdx >= dynamicStrides.size()) + return failure(); + + Value oldStride = dynamicStrides[dynStrideIdx]; + dynStrideIdx++; + if (staticSizes[i] == 1 && mlir::isZero(OpFoldResult(oldStride))) { + newDynamicStrides.push_back(c1); + changed = true; + } else { + newDynamicStrides.push_back(oldStride); + } + } + + // all dynStride should be visited, and at least one should be changed + if (dynStrideIdx != dynamicStrides.size()) + return failure(); + if (!changed) + return failure(); + + auto newReinterpretCast = rewriter.create( + op.getLoc(), cast(op.getResult().getType()), op.getSource(), + op.getOffsets(), op.getSizes(), newDynamicStrides, op.getStaticOffsets(), + op.getStaticSizes(), op.getStaticStrides()); + + rewriter.replaceOp(op, newReinterpretCast.getResult()); + + return success(); +} + +} // namespace LoadStoreConverter diff --git a/compiler/lib/TritonToLinalg/MarkTensorKindPass.cpp b/compiler/lib/TritonToLinalg/MarkTensorKindPass.cpp new file mode 100644 index 00000000..ea66566f --- /dev/null +++ b/compiler/lib/TritonToLinalg/MarkTensorKindPass.cpp @@ -0,0 +1,148 @@ + + +#include "dicp/TritonToLinalg/MarkTensorKindPass.h" +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" + +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" + +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/ADT/SmallPtrSet.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "mark-tensor-kind" + +using namespace mlir; +using namespace triton; +const unsigned INT_BIT_WIDTH = 32; +const unsigned SET_INIT_SIZE = 16; + +template struct has_getPtr : std::false_type {}; +template +struct has_getPtr().getPtr())>> + : std::true_type {}; + +template struct has_getSrc : std::false_type {}; +template +struct has_getSrc().getSrc())>> + : std::true_type {}; + +template struct has_getBase : std::false_type {}; +template +struct has_getBase().getBase())>> + : std::true_type {}; + +template static Value extractPointer(OpTy op) { + if constexpr (has_getPtr::value) + return op.getPtr(); + else if constexpr (has_getSrc::value) + return op.getSrc(); + else if constexpr (has_getBase::value) + return op.getBase(); + else { + Operation *raw = op.getOperation(); + if (!raw || raw->getNumOperands() == 0) + return Value(); + return raw->getOperand(0); + } +} + +static void setBlockArgumentAttr(BlockArgument blockArg, triton::FuncOp func, + TensorKind tensorKind) { + unsigned argIdx = blockArg.getArgNumber(); + auto existingAttr = + func.getArgAttrOfType(argIdx, "tt.tensor_kind"); + TensorKind oldVal = existingAttr + ? static_cast(existingAttr.getInt()) + : TensorKind::NONE; + + TensorKind finalVal = tensorKind; + if ((oldVal == TensorKind::INPUT && tensorKind == TensorKind::OUTPUT) || + (oldVal == TensorKind::OUTPUT && tensorKind == TensorKind::INPUT)) { + finalVal = TensorKind::INPUT_OUTPUT; + } else if (oldVal == TensorKind::INPUT_OUTPUT) { + finalVal = oldVal; + } + + LLVM_DEBUG(llvm::dbgs() << "Setting tensor_kind for argument " << argIdx + << ": " << finalVal << "\n";); + + func.setArgAttr( + argIdx, "tt.tensor_kind", + IntegerAttr::get(IntegerType::get(func.getContext(), INT_BIT_WIDTH), + static_cast(finalVal))); +} + +template +static void addTensorKindToArguments(OpTy op, TensorKind tensorKind) { + Value ptr = extractPointer(op); + if (!ptr) + return; + + LLVM_DEBUG(llvm::dbgs() << "Processing op: " << *op.getOperation() << "\n";); + + Value cur = ptr; + llvm::SmallPtrSet visited; + while (visited.insert(cur).second) { + if (auto blockArg = dyn_cast(cur)) { + if (auto func = dyn_cast_or_null( + blockArg.getOwner()->getParentOp())) { + if (blockArg.getOwner() == &func.getBody().front() && + isa(blockArg.getType())) { + setBlockArgumentAttr(blockArg, func, tensorKind); + break; + } + } + } + + Operation *defOp = cur.getDefiningOp(); + if (!defOp || defOp->getNumOperands() == 0) + break; + cur = defOp->getOperand(0); + } +} + +template +struct MarkTensorKindPattern : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(OpTy op, + PatternRewriter &rewriter) const override { + addTensorKindToArguments(op, Kind); + return success(); + } +}; + +void MarkTensorKindPass::runOnOperation() { + RewritePatternSet patterns(&getContext()); + + // INPUT tensors + patterns.add< + MarkTensorKindPattern, + MarkTensorKindPattern, + MarkTensorKindPattern, + MarkTensorKindPattern>( + &getContext()); + + // OUTPUT tensors + patterns.add< + MarkTensorKindPattern, + MarkTensorKindPattern, + MarkTensorKindPattern, + MarkTensorKindPattern>( + &getContext()); + + // INPUT_OUTPUT tensors + patterns.add< + MarkTensorKindPattern, + MarkTensorKindPattern>( + &getContext()); + + (void)applyPatternsGreedily(getOperation(), std::move(patterns)); +} + +std::unique_ptr> triton::createMarkTensorKindPass() { + return std::make_unique(); +} \ No newline at end of file diff --git a/compiler/lib/Conversion/DiscreteMaskAccessConversion/MaskAnalysis.cpp b/compiler/lib/TritonToLinalg/MaskAnalysis.cpp similarity index 70% rename from compiler/lib/Conversion/DiscreteMaskAccessConversion/MaskAnalysis.cpp rename to compiler/lib/TritonToLinalg/MaskAnalysis.cpp index e9e9ea55..683451d7 100644 --- a/compiler/lib/Conversion/DiscreteMaskAccessConversion/MaskAnalysis.cpp +++ b/compiler/lib/TritonToLinalg/MaskAnalysis.cpp @@ -1,4 +1,7 @@ -#include "dicp/Conversion/DiscreteMaskAccessConversion/MaskAnalysis.h" + + +#include "dicp/TritonToLinalg/MaskAnalysis.h" +#include "dicp/Utils/Utils.h" #include "triton/Dialect/Triton/IR/Dialect.h" @@ -12,266 +15,49 @@ #include "llvm/ADT/TypeSwitch.h" #include "llvm/Support/Debug.h" #include -#include #include -#include -#include -#include -#include -#include -#include -#define DEBUG_TYPE "dicp-mask-analysis" +#define DEBUG_TYPE "mask-analysis" namespace mlir { -static Value createConstIndexValueOp(const Location &loc, OpBuilder &b, - int64_t value) { - return b.create(loc, b.getIndexAttr(value)).getResult(); -} - -static std::optional getConstantOfAttr(const OpFoldResult &arg) { - if (isa(arg)) { - return getConstantIntValue(arg); - } +namespace triton { - return std::nullopt; -} +namespace { -static bool isZeroIndex(OpFoldResult v) { - if (!v) - return false; - if (auto attr = dyn_cast(v)) { - IntegerAttr intAttr = dyn_cast(attr); - return intAttr && intAttr.getValue().isZero(); +template +std::optional runMaskAnalysisImpl(MemAccOpTy op, + OpBuilder &builder) { + auto mask = op.getMask(); + if (!mask) { + return std::nullopt; } - return false; -} - -OpFoldResult addOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b) { - auto lhsInt = getConstantOfAttr(lhs); - auto rhsInt = getConstantOfAttr(rhs); - - if (lhsInt && rhsInt) - return b.getIndexAttr(lhsInt.value() + rhsInt.value()); - - if (!lhsInt && rhsInt && rhsInt.value() == 0) - return lhs; - if (!rhsInt && lhsInt && lhsInt.value() == 0) - return rhs; - - auto lhsValue = dyn_cast(lhs); - if (lhsInt) - lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); - else - assert(isa(lhsValue.getType())); - - auto rhsValue = dyn_cast(rhs); - if (rhsInt) - rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); - else - assert(isa(rhsValue.getType())); - - return b.create(loc, lhsValue, rhsValue).getResult(); -} -OpFoldResult subOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b) { - auto lhsInt = getConstantOfAttr(lhs); - auto rhsInt = getConstantOfAttr(rhs); + PatternRewriter::InsertionGuard insertGuard(builder); + builder.setInsertionPoint(op); - if (lhsInt && rhsInt) - return b.getIndexAttr(lhsInt.value() - rhsInt.value()); - - if (!lhsInt && rhsInt && rhsInt.value() == 0) - return lhs; - - auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); - if (lhsInt) - lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); - else - assert(isa(lhsValue.getType())); - - if (rhsInt) - rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); - else - assert(isa(rhsValue.getType())); - - return b.create(loc, lhsValue, rhsValue).getResult(); -} - -OpFoldResult mulOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b) { - auto lhsInt = getConstantOfAttr(lhs); - auto rhsInt = getConstantOfAttr(rhs); - - if (lhsInt && rhsInt) - return b.getIndexAttr(lhsInt.value() * rhsInt.value()); - - if (lhsInt) { - if (lhsInt.value() == 0) - return lhs; - if (lhsInt.value() == 1) - return rhs; - } - if (rhsInt) { - if (rhsInt.value() == 0) - return rhs; - if (rhsInt.value() == 1) - return lhs; + MaskState mstate; + if (mstate.parse(mask, op.getLoc(), builder).failed()) { + return std::nullopt; } - - auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); - if (lhsInt) - lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); - else - assert(isa(lhsValue.getType())); - - if (rhsInt) - rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); - else - assert(isa(rhsValue.getType())); - - return b.create(loc, lhsValue, rhsValue).getResult(); + return mstate; } -OpFoldResult divOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b) { - auto lhsInt = getConstantOfAttr(lhs); - auto rhsInt = getConstantOfAttr(rhs); - - if (rhsInt && rhsInt.value() == 0) { - emitError(loc) << "cannot div 0!"; - return OpFoldResult(); - } +} // namespace - if (lhsInt && rhsInt) - return b.getIndexAttr(lhsInt.value() / rhsInt.value()); - - if (lhsInt) { - if (lhsInt.value() == 0) - return lhs; - } - - if (rhsInt) { - if (rhsInt.value() == 1) - return lhs; - } - - auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); - if (lhsInt) - lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); - else - assert(isa(lhsValue.getType())); - - if (rhsInt) - rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); - else - assert(isa(rhsValue.getType())); - - return b.create(loc, lhsValue, rhsValue).getResult(); -} - -OpFoldResult remOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b) { - auto lhsInt = getConstantOfAttr(lhs); - auto rhsInt = getConstantOfAttr(rhs); - - if (rhsInt && rhsInt.value() == 0) { - emitError(loc) << "cannot remainder by 0!"; - return OpFoldResult(); - } - - if (lhsInt && rhsInt) - return b.getIndexAttr(lhsInt.value() % rhsInt.value()); - - if (lhsInt) { - if (lhsInt.value() == 0) - return lhs; - } - - auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); - if (lhsInt) - lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); - else - assert(isa(lhsValue.getType())); - - if (rhsInt) - rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); - else - assert(isa(rhsValue.getType())); - - return b.create(loc, lhsValue, rhsValue).getResult(); -} - -OpFoldResult minOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b) { - auto lhsInt = getConstantOfAttr(lhs); - auto rhsInt = getConstantOfAttr(rhs); - if (lhsInt && rhsInt) - return b.getIndexAttr(std::min(lhsInt.value(), rhsInt.value())); - - auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); - if (lhsInt) - lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); - else - assert(isa(lhsValue.getType())); - - if (rhsInt) - rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); - else - assert(isa(rhsValue.getType())); - - return b.create(loc, lhsValue, rhsValue).getResult(); -} - -OpFoldResult maxOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, - const Location &loc, OpBuilder &b) { - auto lhsInt = getConstantOfAttr(lhs); - auto rhsInt = getConstantOfAttr(rhs); - if (lhsInt && rhsInt) - return b.getIndexAttr(std::max(lhsInt.value(), rhsInt.value())); - - auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); - if (lhsInt) - lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); - else - assert(isa(lhsValue.getType())); - - if (rhsInt) - rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); - else - assert(isa(rhsValue.getType())); - - return b.create(loc, lhsValue, rhsValue).getResult(); -} - -// Fold layout constant info to attr, otherwise convert to index type value -OpFoldResult getOpFoldResultOfLayoutInfo(Value value, OpBuilder &builder) { - OpFoldResult constantFold = getAsOpFoldResult(value); - if (llvm::isa(constantFold)) { - assert(isa(cast(constantFold))); - return constantFold; - } - - if (!isa(value.getType())) - llvm_unreachable("Illegal data type when parse block data layout info"); - - if (!isa(value.getType())) { - if (value.getType().isInteger(/*width*/ 1)) - value = builder.create( - value.getLoc(), builder.getIndexType(), value); - else - value = builder.create(value.getLoc(), - builder.getIndexType(), value); +OpFoldResult MaskState::clampToNonNegativeIndex(const OpFoldResult value, + const Location &loc, + OpBuilder &builder) const { + if (auto cst = getConstantIntValue(value)) { + return builder.getIndexAttr(std::max(0, *cst)); } + // For non-constant value, we could generate max(value, 0) to ensure the value + // is non-negative. But this caused error in atomic max/min ut test. We need + // to investigate more on this. return value; } -namespace dicp { - LogicalResult MaskState::parse(Value operand, const Location &loc, OpBuilder &builder) { if (isa(operand.getType())) { @@ -320,6 +106,8 @@ LogicalResult MaskState::parse(Value operand, const Location &loc, [&](auto op) { return this->parseDiv(op, loc, builder); }) .Case( [&](auto op) { return this->parseSel(op, loc, builder); }) + .Case( + [&](auto op) { return this->parseInsert(op, loc, builder); }) .Default([&](Operation *op) { return failure(); }); } @@ -336,6 +124,19 @@ tensor::ExtractSliceOp MaskState::getExtractSlice(Value source, dims, strides); } +tensor::ExtractSliceOp MaskState::getExtractSlice( + Value source, const Location &loc, OpBuilder &builder, + SmallVector offsets, SmallVector dims) const { + auto sourceRType = cast(source.getType()); + SmallVector strides(getRank(), builder.getIndexAttr(1)); + + auto dstRType = tensor::ExtractSliceOp::inferResultType(sourceRType, offsets, + dims, strides); + return builder.create(loc, dstRType, source, offsets, + dims, strides); +} + +// insertSlice tensor::InsertSliceOp MaskState::getInsertSlice(Value source, Value dest, const Location &loc, OpBuilder &builder) const { @@ -344,14 +145,76 @@ tensor::InsertSliceOp MaskState::getInsertSlice(Value source, Value dest, strides); } +tensor::InsertSliceOp +MaskState::getInsertSlice(Value source, Value dest, const Location &loc, + OpBuilder &builder, SmallVector offsets, + SmallVector dims) const { + SmallVector strides(getRank(), builder.getIndexAttr(1)); + return builder.create(loc, source, dest, offsets, dims, + strides); +} + memref::SubViewOp MaskState::getSubview(Value source, const Location &loc, OpBuilder &builder) const { auto sourceType = cast(source.getType()); - SmallVector strides(getRank(), builder.getIndexAttr(1)); - auto dstType = - memref::SubViewOp::inferResultType(sourceType, offsets, dims, strides); - return builder.create(loc, cast(dstType), - source, offsets, dims, strides); + int64_t rank = sourceType.getRank(); + SmallVector strides(rank, builder.getIndexAttr(1)); + + SmallVector fixedOffsets(offsets.begin(), offsets.end()); + SmallVector fixedDims(dims.begin(), dims.end()); + fixedOffsets.resize(rank, builder.getIndexAttr(0)); + fixedDims.resize(rank, builder.getIndexAttr(1)); + + auto dstType = memref::SubViewOp::inferResultType(sourceType, fixedOffsets, + fixedDims, strides); + return builder.create( + loc, cast(dstType), source, fixedOffsets, fixedDims, strides); +} + +// In llvm later version, for each dimension, assert that: +// 0 <= offset < dim_size +// 0 <= offset + (size - 1) *stride < dim_size +// To adatpt to this change, add this verification to avoid llvm assert error +// And currently, in the function calling scenario, the stride coefficient is +// always 1 +bool MaskState::isMemrefSubviewValid(Value source, OpBuilder &builder) const { + auto sourceType = cast(source.getType()); + int64_t rank = sourceType.getRank(); + + SmallVector fixedOffsets(offsets.begin(), offsets.end()); + SmallVector fixedDims(dims.begin(), dims.end()); + fixedOffsets.resize(rank, builder.getIndexAttr(0)); + fixedDims.resize(rank, builder.getIndexAttr(1)); + + for (int64_t i = 0; i < rank; ++i) { + int64_t sourceSize = sourceType.getDimSize(i); + if (ShapedType::isDynamic(sourceSize)) + continue; + std::optional offsetVal = + mlir::getConstantIntValue(fixedOffsets[i]); + std::optional dimVal = mlir::getConstantIntValue(fixedDims[i]); + if (offsetVal.has_value() && dimVal.has_value()) { + if (offsetVal.value() >= sourceSize || offsetVal.value() < 0) { + LLVM_DEBUG({ + llvm::dbgs() << "MemrefSubview offset check faied at dim " << i + << "sourceSize :" << sourceSize + << "offsetVal :" << offsetVal << "\n"; + }); + return false; + } + + int64_t computedEnd = offsetVal.value() + dimVal.value(); + if (computedEnd > sourceSize) { + LLVM_DEBUG({ + llvm::dbgs() << "MemrefSubview offset end check faied at dim " << i + << "dimVal :" << dimVal << "offsetVal :" << offsetVal + << "\n"; + }); + return false; + } + } + } + return true; } static memref::SubViewOp createSubview(Value src, const Location &loc, @@ -374,6 +237,14 @@ LogicalResult MaskState::addStateScalar(const MaskState &state, end = addOpFoldResult(state.end, scalar, loc, builder); dims = state.dims; offsets = state.offsets; + + bool allDimsOne = llvm::all_of(state.dims, [](OpFoldResult dim) { + return getConstantIntValue(dim).value() == std::optional(1); + }); + if (allDimsOne) { + this->scalar = this->start; + } + return success(); } @@ -415,7 +286,7 @@ LogicalResult MaskState::divStates(const MaskState &lhsState, const MaskState &rhsState, const Location &loc, OpBuilder &builder) { if (!lhsState.scalar && rhsState.scalar) { - if (isZeroIndex(rhsState.scalar)) { + if (isZeroInteger(rhsState.scalar)) { InFlightDiagnostic diag = emitError(loc) << "Unsupported scenario where rhs is zero constant in divide!"; @@ -450,9 +321,10 @@ LogicalResult MaskState::minStates(const MaskState &lhsState, auto rhsEnd = addOpFoldResult(rhsOffset, rhsDim, loc, builder); auto newEnd = minOpFoldResult(lhsEnd, rhsEnd, loc, builder); auto newDim = subOpFoldResult(newEnd, newOffset, loc, builder); + auto clampedNewDim = clampToNonNegativeIndex(newDim, loc, builder); offsets.push_back(newOffset); - dims.push_back(newDim); + dims.push_back(clampedNewDim); } return success(); } @@ -468,8 +340,19 @@ LogicalResult MaskState::parseConstant(arith::ConstantOp constOp, auto elementType = attr.getElementType(); assert(attr.isSplat() && isa(elementType) && "All elements must share a single integer constant value"); - this->scalar = builder.getIndexAttr( - attr.getSplatValue().getValue().getSExtValue()); + + if (elementType.isInteger(1) && + isa(constOp.getValue().getType())) { + auto shapedType = cast(constOp.getValue().getType()); + auto shape = shapedType.getShape(); + for (size_t i = 0; i < shape.size(); i++) { + this->dims.push_back(builder.getIndexAttr(shape[i])); + this->offsets.push_back(builder.getIndexAttr(0)); + } + } else { + this->scalar = builder.getIndexAttr( + attr.getSplatValue().getValue().getSExtValue()); + } } else { auto value = cast(constOp.getValue()).getInt(); this->scalar = builder.getIndexAttr(value); @@ -504,6 +387,7 @@ LogicalResult MaskState::parseAdd(arith::AddIOp addOp, const Location &loc, LogicalResult MaskState::parseDiv(arith::DivSIOp divOp, const Location &loc, OpBuilder &builder) { assert(this->isEmpty()); + return failure(); // temporarily disable parseDiv MaskState lhsState; if (failed(lhsState.parse(divOp.getLhs(), loc, builder))) { return failure(); @@ -645,8 +529,9 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location &loc, maxOpFoldResult(lhsState.start, rhsState.scalar, loc, builder); auto newEnd = minOpFoldResult(lhsState.end, realBound, loc, builder); auto newDim = subOpFoldResult(newEnd, lhsState.start, loc, builder); + auto clampedNewDim = clampToNonNegativeIndex(newDim, loc, builder); - this->dims[cmpDim] = newDim; + this->dims[cmpDim] = clampedNewDim; break; } case arith::CmpIPredicate::sle: { @@ -656,8 +541,9 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location &loc, auto realBound = maxOpFoldResult(lhsState.start, rhsPlusOne, loc, builder); auto newEnd = minOpFoldResult(lhsState.end, realBound, loc, builder); auto newDim = subOpFoldResult(newEnd, lhsState.start, loc, builder); + auto clampedNewDim = clampToNonNegativeIndex(newDim, loc, builder); - this->dims[cmpDim] = newDim; + this->dims[cmpDim] = clampedNewDim; break; } case arith::CmpIPredicate::sge: { @@ -666,9 +552,10 @@ LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location &loc, auto newStart = minOpFoldResult(lhsState.end, realBound, loc, builder); auto newOffset = subOpFoldResult(newStart, lhsState.start, loc, builder); auto newDim = subOpFoldResult(lhsState.end, newStart, loc, builder); + auto clampedNewDim = clampToNonNegativeIndex(newDim, loc, builder); this->offsets[cmpDim] = newOffset; - this->dims[cmpDim] = newDim; + this->dims[cmpDim] = clampedNewDim; break; } case arith::CmpIPredicate::eq: { @@ -800,6 +687,26 @@ LogicalResult MaskState::parseSplat(triton::SplatOp splatOp, return success(); } +LogicalResult MaskState::parseInsert(tensor::InsertOp insertOp, + const Location &loc, OpBuilder &builder) { + if (auto tensorType = dyn_cast(insertOp.getType())) { + if (llvm::all_of(tensorType.getShape(), + [](int64_t dim) { return dim == 1; })) { + SmallVector indices( + tensorType.getRank(), builder.create(loc, 0)); + auto scalarlizeTensor = + builder.create(loc, insertOp, indices).getResult(); + this->offsets = SmallVector(tensorType.getRank(), + builder.getIndexAttr(0)); + this->dims = SmallVector(tensorType.getRank(), + builder.getIndexAttr(1)); + this->scalar = getOpFoldResultOfLayoutInfo(scalarlizeTensor, builder); + return success(); + } + } + return failure(); +} + LogicalResult MaskState::parseExpandDims(triton::ExpandDimsOp expandDimsOp, const Location &loc, OpBuilder &builder) { @@ -824,25 +731,51 @@ void MaskState::eraseInsertedOps(Operation *rawOp, PatternRewriter &rewriter) { auto moduleOp = rawOp->getParentOfType(); SmallVector worklist; moduleOp->walk([&](Operation *op) { - if (isOpTriviallyDead(op)) + if (isOpTriviallyDead(op) && op->use_empty()) { worklist.push_back(op); + } }); while (!worklist.empty()) { Operation *op = worklist.pop_back_val(); - if (!isOpTriviallyDead(op)) + if (!isOpTriviallyDead(op) || !op->use_empty()) { + continue; + } + if (!op->getBlock()) { continue; + } + SmallVector operandDefOps; for (Value value : op->getOperands()) { - if (auto defOp = value.getDefiningOp()) - worklist.push_back(defOp); + if (auto defOp = value.getDefiningOp()) { + if (defOp->getBlock()) { + operandDefOps.push_back(defOp); + } + } } LLVM_DEBUG({ llvm::dbgs() << "[MaskState]==> inserted op: \n" << *op << "\n[MaskState]<== is removed\n"; }); + rewriter.eraseOp(op); + for (auto defOp : operandDefOps) { + worklist.push_back(defOp); + } + } +} + +std::optional runMaskAnalysis(Operation *op, OpBuilder &builder) { + if (auto loadOp = dyn_cast(op)) { + return runMaskAnalysisImpl(loadOp, builder); } + if (auto storeOp = dyn_cast(op)) { + return runMaskAnalysisImpl(storeOp, builder); + } + if (auto atomicRMWOp = dyn_cast(op)) { + return runMaskAnalysisImpl(atomicRMWOp, builder); + } + return std::nullopt; } -} // namespace dicp +} // namespace triton } // namespace mlir diff --git a/compiler/lib/TritonToLinalg/TritonOpConverter.cpp b/compiler/lib/TritonToLinalg/TritonOpConverter.cpp new file mode 100644 index 00000000..ec71ae70 --- /dev/null +++ b/compiler/lib/TritonToLinalg/TritonOpConverter.cpp @@ -0,0 +1,3165 @@ + + +#include "dicp/TritonToLinalg/TritonOpConverter.h" +#include "dicp/TritonToLinalg/BlockPtrAnalysis.h" +#include "dicp/TritonToLinalg/MaskAnalysis.h" +#include "dicp/TritonToLinalg/TritonToLinalgPass.h" +#include "dicp/Utils/Utils.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/StringRef.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/LogicalResult.h" +#include "llvm/Support/raw_ostream.h" +#include +#include +#include + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/MemRef/Transforms/Passes.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/Interfaces/CallInterfaces.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HFusion/IR/HFusion.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" + +namespace TTOpConverters { +using namespace mlir; +using namespace triton; + +static const llvm::SmallVector libdeviceOps = { + // Basic operations + "__hmf_pow_fp32", + "__hmf_div_rz_fp32", + "__hmf_fmod_fp32", + "__hmf_float_as_int_fp32", + "__hmf_trunc_fp32", + "__hmf_trunc_fp16", + "__hmf_nearbyint_fp32", + "__hmf_signbit_fp32", + "__hmf_signbit_fp16", + "__hmf_copysign_fp32", + "__hmf_log10_fp32", + // Trigonometric operations + "__hmf_tanh_fp32", + "__hmf_asin_fp32", + "__hmf_asin_fp16", + "__hmf_acos_fp32", + "__hmf_acos_fp16", + "__hmf_atan2_fp32", + "__hmf_atan2_fp16", + "__hmf_sinh_fp32", + "__hmf_sinh_fp16", + "__hmf_cosh_fp32", + "__hmf_cosh_fp16", + "__hmf_asinh_fp32", + "__hmf_asinh_fp16", + "__hmf_acosh_fp32", + "__hmf_acosh_fp16", + "__hmf_atanh_fp32", + "__hmf_atanh_fp16", + // Other operations + "__hmf_expm1_fp32", + "__hmf_expm1_fp16", + "__hmf_nextafter_fp32", + "__hmf_nextafter_fp16", + "__hmf_hypot_fp32", + "__hmf_hypot_fp16", + "__hmf_cyl_bessel_i0_fp32", + "__hmf_cyl_bessel_i0_fp16", + "__hmf_erfinv_fp32", + "__hmf_lgamma_fp32", +}; + +/** + * Retrieves a boolean environment variable. + * @param envVar The name of the environment variable. + * @param defaultValue The default value to return if the variable is not set or + * cannot be parsed. + * @return true if the environment variable exists and its value is parsed as + * "true", otherwise returns defaultValue. Parsing rules (case-insensitive): + * "true" values: any non-empty string not equal to "0", "false", "no", "off" is + * considered true. "false" values: an empty string or a string equal to any of + * the false literals is considered false. + */ +bool getEnvBool(const char *envVar, bool defaultValue) { + const char *val = std::getenv(envVar); + if (val == nullptr) { + return defaultValue; // variable not set + } + + std::string s(val); + // Convert to lowercase for easier comparison + std::transform(s.begin(), s.end(), s.begin(), + [](unsigned char c) { return std::tolower(c); }); + + // Common false literals + if (s.empty() || s == "0" || s == "false" || s == "no" || s == "off") { + return false; + } + // All other cases (including "1", "true", "yes", "on", etc.) are considered + // true + return true; +} + +static llvm::SmallString +generateUniqueFuncName(ModuleOp moduleOp, llvm::StringRef funcNameBase) { + llvm::SmallString funcName = funcNameBase; + int uniqueId = 0; + while (SymbolTable::lookupSymbolIn(moduleOp, funcName)) { + funcName = funcNameBase; + funcName += ("_" + std::to_string(uniqueId++)); + } + return funcName; +} + +LogicalResult +BitcastConverter::matchAndRewrite(triton::BitcastOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + Value result; + auto loc = op.getLoc(); + + if (auto dstPtrTy = dyn_cast(op.getType())) { + auto srcPtrTy = cast(op.getSrc().getType()); + auto resType = + MemRefType::get({ShapedType::kDynamic}, dstPtrTy.getPointeeType()); + + auto i1Ty = rewriter.getIntegerType(1); + auto i8Ty = rewriter.getIntegerType(8); + bool isI1toI8 = (srcPtrTy.getPointeeType() == i1Ty) && + (dstPtrTy.getPointeeType() == i8Ty); + // handling special case: ptr -> ptr, directly forward without + // arith.bitcast + if (isI1toI8) { + // TypeConverter has already converted i1 to i8 memref, + LLVM_DEBUG({ + llvm::dbgs() + << "[BitcastConverter] Special i1->i8 pointer bitcast. Forward " + "without arith.bitcast. srcConvertedTy=" + << adaptor.getSrc().getType() << "\n"; + }); + rewriter.replaceOp(op, adaptor.getSrc()); + return success(); + } + result = rewriter.create(loc, resType, adaptor.getSrc()); + } else { + // handling normal case: bitcast between tensors/memrefs + result = + rewriter.create(loc, op.getType(), adaptor.getSrc()); + } + rewriter.replaceOp(op, result); + return success(); +} + +LogicalResult +TransposeConverter::matchAndRewrite(triton::TransOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto src = adaptor.getSrc(); + auto res = ConverterUtils::getTransposedValue(src, op.getLoc(), rewriter, + op.getOrder()); + rewriter.replaceOp(op, res); + return success(); +} + +LogicalResult +YieldConverter::matchAndRewrite(scf::YieldOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + rewriter.replaceOpWithNewOp(op, adaptor.getOperands()); + return success(); +} + +LogicalResult +AdvanceConverter::matchAndRewrite(triton::AdvanceOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + llvm::SmallDenseMap known; + BlockDataParser::rewriteAdvanceOp(op, rewriter, known); + return success(); +} + +// ToDo: +// 1. Refactor MakeTensorPtrConverter and AdvanceConverter with +// memref::ReinterpretCastOp and memref::SubViewOp. +// Use recast to describe full shape of tensor, and use subview to represent +// current block tensor. +LogicalResult MakeTensorPtrConverter::matchAndRewrite( + triton::MakeTensorPtrOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + llvm::SmallDenseMap known; + BlockDataParser::rewriteMakeTensorPtrOp(op, adaptor.getBase(), rewriter, + known); + return success(); +} + +LogicalResult PreciseDivConverter::matchAndRewrite( + triton::PreciseDivFOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + Value opa = op.getX(); + Value opb = op.getY(); + auto loc = op.getLoc(); + + auto resType = dyn_cast(op.getResult().getType()); + auto divOp = rewriter.create(loc, resType, opa, opb); + + rewriter.replaceOp(op, divOp); + return success(); +} + +LogicalResult +SelectCanonicalizer::matchAndRewrite(arith::SelectOp op, + PatternRewriter &rewriter) const { + auto loc = op.getLoc(); + + // 0. Shortcut for scalars and bool type + auto type = dyn_cast(op.getResult().getType()); + if (!type) { + // do nothing non-tensor select + return failure(); + } + auto elementType = type.getElementType(); + if (elementType.isInteger(1)) { + // do nothing with bool type + return failure(); + } + auto tensorShape = type.getShape(); + auto mask = op.getCondition(); + if (!isa(mask.getType())) { + // do nothing for scalar mask + return failure(); + } + + // 1. Check for continuous masked loads. + // Analyze the mask operand to determine at runtime the size of the data we + // are moving. + MaskState mstate; + auto isContMask = mstate.parse(mask, loc, rewriter); + + if (isContMask.failed()) { + mstate.eraseInsertedOps(op, rewriter); + return rewriter.notifyMatchFailure( + op, "Cannot lower continuous masked selects"); + } + + // 2. Get mask position + MaskPosition maskPos = mstate.getMaskPosition(tensorShape); + LLVM_DEBUG({ + llvm::dbgs() << "[SelectAnalysis] MaskPosition detected: " + << (maskPos == MaskPosition::Head ? "Head" + : maskPos == MaskPosition::Tail ? "Tail" + : maskPos == MaskPosition::Middle ? "Middle" + : "Unknown") + << "\n"; + }); + + if (maskPos == MaskPosition::Unknown) { + mstate.eraseInsertedOps(op, rewriter); + return failure(); + } + auto trueTensor = op.getTrueValue(); + auto falseTensor = op.getFalseValue(); + + // 3. Slice and insert out the masked part + if (maskPos == MaskPosition::Head) { + // Slice out the masked part of true tensor + auto extractSliceOp = mstate.getExtractSlice(trueTensor, loc, rewriter); + + // Insert out the sliced true tensor into false tensor + auto insertSliceOp = + mstate.getInsertSlice(extractSliceOp, falseTensor, loc, rewriter); + + LLVM_DEBUG({ + llvm::dbgs() << " -> Created ExtractSlice: " + << *extractSliceOp.getOperation() << "\n" + << " -> Created InsertSlice: " + << *insertSliceOp.getOperation() << "\n"; + }); + rewriter.replaceOp(op, insertSliceOp); + return success(); + } + + // For Tail or Middle positions, we need to compute inverted dimensions + // to handle the masking logic + SmallVector invertOffsets; + SmallVector invertFalseDims; + SmallVector invertTrueDims; + OpFoldResult falseDimOp; + OpFoldResult trueDimOp; + int valDim = -1; + for (int i = 0; i < mstate.getRank(); ++i) { + const auto &offVal = mstate.offsets[i]; + const auto &dimVal = mstate.dims[i]; + auto constOffVal = getConstantIntValue(offVal); + invertOffsets.push_back(rewriter.getIndexAttr(0)); + if (constOffVal.has_value() && constOffVal.value() == 0) { + invertFalseDims.push_back(dimVal); + invertTrueDims.push_back(dimVal); + } else { + assert(valDim == -1 && + "The offset in only one dimension can be not zero."); + if (!constOffVal.has_value()) { + valDim = i; + falseDimOp = offVal; + } + + invertFalseDims.push_back(offVal); + trueDimOp = addOpFoldResult(offVal, dimVal, loc, rewriter); + invertTrueDims.push_back(trueDimOp); + } + } + + // Slice out the invert first masked part of false tensor + auto falseExtractSliceOp = mstate.getExtractSlice( + falseTensor, loc, rewriter, invertOffsets, invertFalseDims); + // Insert out the sliced false tensor into true tensor + auto trueInsertSliceOp = + mstate.getInsertSlice(falseExtractSliceOp, trueTensor, loc, rewriter, + invertOffsets, invertFalseDims); + // Slice out the invert first masked and masked part of inserted true tensor + auto extractSliceOp = mstate.getExtractSlice(trueInsertSliceOp, loc, rewriter, + invertOffsets, invertTrueDims); + // Insert out the sliced true tensor into false tensor + auto insertSliceOp = + mstate.getInsertSlice(extractSliceOp, falseTensor, loc, rewriter, + invertOffsets, invertTrueDims); + if (valDim != -1) { + rewriter.setInsertionPointAfter(trueInsertSliceOp); + assert(isa(falseDimOp) && + "Expected to be a runtime Value for dynamic dimension check."); + Value zeroIndex = rewriter.create(loc, 0); + Value isNegative = rewriter.create( + loc, arith::CmpIPredicate::slt, cast(falseDimOp), zeroIndex); + + Value sizeIndex = + rewriter.create(loc, tensorShape[valDim]); + Value isOutOfRange = rewriter.create( + loc, arith::CmpIPredicate::sge, cast(falseDimOp), sizeIndex); + auto orOp = rewriter.create(loc, isNegative, isOutOfRange); + auto ifOp = rewriter.create(loc, TypeRange{type}, + orOp.getResult(), true, true); + + Block *thenBlock = &ifOp.getThenRegion().front(); + rewriter.setInsertionPointToStart(thenBlock); + rewriter.create(loc, ValueRange{falseTensor}); + + Block *elseBlock = &ifOp.getElseRegion().front(); + rewriter.setInsertionPointToStart(elseBlock); + falseExtractSliceOp->moveBefore(elseBlock, elseBlock->begin()); + trueInsertSliceOp->moveAfter(falseExtractSliceOp); + extractSliceOp->moveAfter(trueInsertSliceOp); + insertSliceOp->moveAfter(extractSliceOp); + + rewriter.setInsertionPointAfter(insertSliceOp); + rewriter.create(loc, ValueRange{insertSliceOp.getResult()}); + rewriter.replaceOp(op, ifOp); + } else { // static offsets + rewriter.replaceOp(op, insertSliceOp); + } + LLVM_DEBUG({ + llvm::dbgs() << " -> [invert] Created false tensor extractSlice: " + << *falseExtractSliceOp.getOperation() << "\n" + << " -> [invert] Created true tensor insertSlice: " + << *trueInsertSliceOp.getOperation() << "\n" + << " -> [invert] Created ExtractSlice: " + << *extractSliceOp.getOperation() << "\n" + << " -> [invert] Created InsertSlice: " + << *insertSliceOp.getOperation() << "\n"; + }); + return success(); +} + +/* + * Move tt.bitcast to a previous location if tt.bitcast is not directly applied + * on function arguments + */ +LogicalResult +BitcastCanonicalizer::matchAndRewrite(triton::BitcastOp bitcastOp, + PatternRewriter &rewriter) const { + Value castSrc = bitcastOp.getSrc(); + Value castRes = bitcastOp.getResult(); + Type castSrcTy = castSrc.getType(); + Type castSrcPtrTy = isa(castSrcTy) + ? cast(castSrcTy).getElementType() + : castSrcTy; + if (!isa(castSrcPtrTy)) + return failure(); + + auto origBitwidth = getPointeeBitWidth(castSrc.getType()); + auto castBitwidth = getPointeeBitWidth(castRes.getType()); + + if (origBitwidth == 1) + origBitwidth = 8; + if (castBitwidth == 1) + castBitwidth = 8; + if (origBitwidth != castBitwidth) { + bitcastOp.emitError() << "Casting pointers with unmatched bitwidth!\n"; + return failure(); + } + + Operation *beforeCastOp = castSrc.getDefiningOp(); + if (beforeCastOp == nullptr) { + return failure(); + } + + auto newRes = + TypeSwitch>(beforeCastOp) + // before: addptr - bitcast - load/store + // after: bitcast - addptr - load/store + .Case([&](triton::AddPtrOp addptrOp) { + auto newCastOp = rewriter.create( + bitcastOp.getLoc(), castRes.getType(), addptrOp.getPtr()); + return rewriter.create( + bitcastOp.getLoc(), castRes.getType(), newCastOp.getResult(), + addptrOp.getOffset()); + }) + .Case([&](triton::SplatOp splatOp) { + Type newCastSrcTy = + cast(castRes.getType()).getElementType(); + + Value splatSrc = splatOp.getSrc(); + Type splatSrcTy = splatSrc.getType(); + if (auto splatSrcTensorTy = dyn_cast(splatSrcTy)) + newCastSrcTy = + splatSrcTensorTy.cloneWith(std::nullopt, newCastSrcTy); + auto newCastOp = rewriter.create( + bitcastOp.getLoc(), newCastSrcTy, splatSrc); + return rewriter.create( + bitcastOp.getLoc(), castRes.getType(), newCastOp); + }) + // before: bitcast - bitcast + // after(fusion optimization): bitcast + .Case([&](triton::BitcastOp prevCastOp) { + return rewriter.create( + bitcastOp.getLoc(), castRes.getType(), prevCastOp.getSrc()); + }) + .Default([&](Operation *op) { + return rewriter.notifyMatchFailure(bitcastOp, + "Unknown bitcast pattern"); + }); + if (succeeded(newRes)) { + rewriter.replaceOp(bitcastOp, newRes.value()); + if (beforeCastOp->use_empty()) { + rewriter.eraseOp(beforeCastOp); + } + return success(); + } + return failure(); +} + +LogicalResult +FpToFpCanonicalizer::matchAndRewrite(triton::FpToFpOp op, + PatternRewriter &rewriter) const { + auto loc = op.getLoc(); + Value input = op.getSrc(); + auto resultType = op.getResult().getType(); + + // Check if rounding mode is specified + auto roundingMode = op.getRounding(); + if (roundingMode.has_value() && + roundingMode.value() != triton::RoundingMode::RTNE) { + // Non-RTNE rounding modes (e.g., RTZ) should be handled by TritonToHFusion + // pass Return failure here so this pattern doesn't match + return failure(); + } + + // Handle RTNE (default) rounding mode with arith.truncf/extf + auto srcType = cast(input.getType()); + auto dstType = cast(resultType); + auto srcElemType = srcType.getElementType(); + auto dstElemType = dstType.getElementType(); + if (!isa(srcElemType) || !isa(dstElemType)) { + return op.emitError("FpToFp expects floating point types"); + } + + unsigned srcBitwidth = srcElemType.getIntOrFloatBitWidth(); + unsigned dstBitwidth = dstElemType.getIntOrFloatBitWidth(); + + // Create round_mode attribute (RINT for RTNE) + auto roundModeAttr = hfusion::RoundModeAttr::get(rewriter.getContext(), + hfusion::RoundMode::RINT); + + if (srcBitwidth > dstBitwidth) { + // Downcast: use arith.truncf with round_mode=rint + auto truncOp = rewriter.create(loc, dstType, input); + truncOp->setAttr("round_mode", roundModeAttr); + rewriter.replaceOp(op, truncOp.getResult()); + } else if (srcBitwidth < dstBitwidth) { + // Upcast: use arith.extf with round_mode=rint + auto extOp = rewriter.create(loc, dstType, input); + extOp->setAttr("round_mode", roundModeAttr); + rewriter.replaceOp(op, extOp.getResult()); + } else { + // Same bitwidth, should not happen but handle gracefully + rewriter.replaceOp(op, input); + } + + return success(); +} + +void rewriteUserWithNewOrder( + mlir::OpOperand *use, PatternRewriter &rewriter, + llvm::SmallVector &blkShapeI64, // 8: container size + mlir::Location &loc, llvm::ArrayRef &order, size_t &orderSize) { + Operation *user = use->getOwner(); + rewriter.setInsertionPointAfter(user); + if (auto loadOp = dyn_cast(user)) { + auto loadResTy = loadOp.getResult().getType(); + auto loadResShapedTy = cast(loadResTy); + auto newLoadTy = loadResShapedTy.cloneWith( + blkShapeI64, loadResShapedTy.getElementType()); + auto newLoadOp = rewriter.create( + loc, newLoadTy, loadOp->getOperands(), loadOp->getAttrs()); + newLoadOp->setAttr(ConverterUtils::GeneratedByMakeTensorPtrTAG, + UnitAttr::get(rewriter.getContext())); + rewriter.replaceOp(loadOp, newLoadOp); + // load contiguous data then permute. thus the permute order is as + // follows. + SmallVector permuteOrder; // 8: container size + for (auto [i, v] : llvm::enumerate(order)) { + permuteOrder.push_back(orderSize - 1 - order[i]); + } + auto permuteOp = rewriter.create( + loc, newLoadOp.getResult(), + DenseI32ArrayAttr::get(loadOp.getContext(), permuteOrder)); + newLoadOp.getResult().replaceAllUsesExcept(permuteOp.getResult(), + permuteOp); + } else if (auto storeOp = dyn_cast(user)) { + // permute to contiguous then store. thus the permute order is as follows. + SmallVector permuteOrder; // 8: container size + for (auto [i, v] : llvm::enumerate(order)) { + permuteOrder.push_back(order[orderSize - 1 - i]); + } + auto permuteOp = rewriter.create( + loc, storeOp.getValue(), + DenseI32ArrayAttr::get(storeOp.getContext(), permuteOrder)); + storeOp.getValue().replaceAllUsesExcept(permuteOp.getResult(), permuteOp); + auto newStoreOp = rewriter.create( + loc, storeOp.getPtr(), storeOp.getValue(), storeOp.getMask(), + storeOp.getBoundaryCheck(), storeOp.getCache(), storeOp.getEvict()); + rewriter.replaceOp(storeOp, newStoreOp); + } else if (auto advanceOp = dyn_cast(user)) { + auto advanceResPtrTy = + cast(advanceOp.getResult().getType()); + auto advanceResShapedTy = + cast(advanceResPtrTy.getPointeeType()); + auto newAdvanceResShapedTy = advanceResShapedTy.cloneWith( + blkShapeI64, advanceResShapedTy.getElementType()); + auto newAdvanceResPtrTy = triton::PointerType::get( + newAdvanceResShapedTy, advanceResPtrTy.getAddressSpace()); + auto advanceOffsets = advanceOp.getOffsets(); + llvm::SmallVector newAdvanceOffsets; // 8: container size + for (int i = orderSize - 1; i >= 0; i--) { + newAdvanceOffsets.push_back(advanceOffsets[order[i]]); + } + SmallVector resUses; + for (auto &use : advanceOp->getUses()) + resUses.push_back(&use); + auto newAdvanceOp = rewriter.create( + loc, newAdvanceResPtrTy, advanceOp.getPtr(), newAdvanceOffsets); + rewriter.replaceOp(advanceOp, newAdvanceOp); + for (auto resUse : resUses) + rewriteUserWithNewOrder(resUse, rewriter, blkShapeI64, loc, order, + orderSize); + } else if (auto loopOp = dyn_cast(user)) { + auto initArg = use->get(); + auto iterArg = loopOp.getTiedLoopRegionIterArg(use); + auto resultValue = loopOp.getTiedLoopResult(use); + iterArg.setType(initArg.getType()); + resultValue.setType(initArg.getType()); + for (auto &argUse : iterArg.getUses()) + rewriteUserWithNewOrder(&argUse, rewriter, blkShapeI64, loc, order, + orderSize); + for (auto &resUse : resultValue.getUses()) + rewriteUserWithNewOrder(&resUse, rewriter, blkShapeI64, loc, order, + orderSize); + } else if (isa(user)) { + return; + } else { + llvm_unreachable( + "[MakeTensorPtrCanonicalizer] tt.make_tensor_ptr's result is " + "not used by load/store/advance op"); + } +} + +void markLoadUsers(mlir::OpOperand *use, PatternRewriter &rewriter) { + Operation *user = use->getOwner(); + if (auto loadOp = dyn_cast(user)) { + loadOp->setAttr(ConverterUtils::GeneratedByMakeTensorPtrTAG, + UnitAttr::get(rewriter.getContext())); + } else if (auto storeOp = dyn_cast(user)) { + return; + } else if (auto advanceOp = dyn_cast(user)) { + SmallVector resUses; + for (auto &use : advanceOp->getUses()) + resUses.push_back(&use); + for (auto resUse : resUses) + markLoadUsers(resUse, rewriter); + } else if (auto loopOp = dyn_cast(user)) { + auto initArg = use->get(); + auto iterArg = loopOp.getTiedLoopRegionIterArg(use); + auto resultValue = loopOp.getTiedLoopResult(use); + iterArg.setType(initArg.getType()); + resultValue.setType(initArg.getType()); + for (auto &argUse : iterArg.getUses()) + markLoadUsers(&argUse, rewriter); + for (auto &resUse : resultValue.getUses()) + markLoadUsers(&resUse, rewriter); + } else if (isa(user)) { + return; + } else { + llvm_unreachable( + "[MakeTensorPtrCanonicalizer] tt.make_tensor_ptr's result is " + "not used by load/store/advance op"); + } +} + +LogicalResult +MakeTensorPtrCanonicalizer::matchAndRewrite(triton::MakeTensorPtrOp op, + PatternRewriter &rewriter) const { + auto order = op.getOrder(); + auto orderSize = order.size(); + if (orderSize == 1) { + return rewriter.notifyMatchFailure( + op, "make_tensor_ptr's order has single value."); + } + + bool isPermuted = false; + for (auto [first, second] : llvm::zip(order.slice(0, orderSize - 1), + order.slice(1, orderSize - 1))) { + if (first != second + 1) { + isPermuted = true; + break; + } + } + + auto loc = op.getLoc(); + auto base = op.getBase(); + auto shape = op.getShape(); + auto strides = op.getStrides(); + auto offsets = op.getOffsets(); + auto result = op.getResult(); + SmallVector opUses; + + for (auto &use : result.getUses()) + opUses.push_back(&use); + for (auto use : opUses) + markLoadUsers(use, rewriter); + + if (!isPermuted) { + return rewriter.notifyMatchFailure( + op, "make_tensor_ptr's order is contiguous."); + } + + llvm::SmallVector blkShapeI32; + llvm::SmallVector blkShapeI64; + auto resPtrType = cast(result.getType()); + if (auto resShapedTy = dyn_cast(resPtrType.getPointeeType())) { + auto resBlkShape = resShapedTy.getShape(); + for (auto [i, v] : llvm::enumerate(resBlkShape)) { + auto reverseI = orderSize - 1 - i; + blkShapeI32.push_back(resBlkShape[order[reverseI]]); + blkShapeI64.push_back(resBlkShape[order[reverseI]]); + } + } + + llvm::SmallVector newShape; + llvm::SmallVector newStrides; + llvm::SmallVector newOffsets; + for (int i = orderSize - 1; i >= 0; i--) { + newShape.push_back(shape[order[i]]); + newStrides.push_back(strides[order[i]]); + newOffsets.push_back(offsets[order[i]]); + } + + llvm::SmallVector contiguousOrder; + for (int i = orderSize - 1; i >= 0; i--) + contiguousOrder.push_back(i); + + rewriter.setInsertionPoint(op); + auto newMakeTensorPtrOp = rewriter.create( + loc, base, ValueRange(newShape), ValueRange(newStrides), + ValueRange(newOffsets), blkShapeI32, contiguousOrder); + rewriter.replaceOp(op, newMakeTensorPtrOp); + for (auto use : opUses) + rewriteUserWithNewOrder(use, rewriter, blkShapeI64, loc, order, orderSize); + return success(); +} + +LogicalResult +ReduceSingleCanonicalizer::matchAndRewrite(triton::ReduceOp reduceOp, + PatternRewriter &rewriter) const { + assert(reduceOp.getSrcs().size() <= 2 && + "Only reduce or reduce with index are supported"); + auto src = reduceOp.getSrcs()[0]; + auto srcType = cast(src.getType()); + auto srcShape = srcType.getShape(); + if (llvm::any_of(srcShape, [](auto s) { return s != 1; })) + return rewriter.notifyMatchFailure( + reduceOp, "reduce's srcs are not all with single element"); + auto loc = reduceOp->getLoc(); + + // Handle Reduce Value + auto res = reduceOp.getResult()[0]; + Value extracted; + if (srcType.getRank() == 1) { + auto zero = + rewriter.create(loc, rewriter.getIndexAttr(0)); + extracted = rewriter.create(loc, src, zero.getResult()) + .getResult(); + } else { + auto resShape = cast(res.getType()).getShape(); + auto collapseReassociationIndicesOptional = + getReassociationIndicesForCollapse(srcShape, resShape); + if (!collapseReassociationIndicesOptional.has_value()) { + return rewriter.notifyMatchFailure( + reduceOp, "Failure with getReassociationIndicesForCollapse call"); + } + auto collapseReassociationIndices = + collapseReassociationIndicesOptional.value(); + extracted = rewriter + .create( + loc, src, collapseReassociationIndices) + .getResult(); + } + res.replaceAllUsesWith(extracted); + + // Handle Reduce Index + if (reduceOp.getSrcs().size() == 1) + return success(); + + auto resIdx = reduceOp.getResult()[1]; + auto zeroI32 = + rewriter.create(loc, rewriter.getI32IntegerAttr(0)); + if (srcType.getRank() == 1) { + resIdx.replaceAllUsesWith(zeroI32); + } else { + auto resIdxShape = cast(resIdx.getType()).getShape(); + auto initTensor = rewriter.create(loc, resIdxShape, + rewriter.getI32Type()); + auto fillOp = rewriter.create(loc, ValueRange{zeroI32}, + ValueRange{initTensor}); + resIdx.replaceAllUsesWith(fillOp.getResult(0)); + } + + return success(); +} + +LogicalResult DenseConstantConverter::matchAndRewrite( + arith::ConstantOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto denseAttr = cast(op.getValue()); + auto loc = op.getLoc(); + auto constSplatOp = arith::ConstantOp::materialize( + rewriter, denseAttr.getSplatValue(), + denseAttr.getElementType(), loc); + auto emptyOp = rewriter.create( + loc, cast(op.getResult().getType()).getShape(), + denseAttr.getElementType()); + + rewriter.replaceOpWithNewOp(op, ValueRange{constSplatOp}, + ValueRange{emptyOp}); + + return success(); +} + +LogicalResult +MakeRangeConverter::matchAndRewrite(triton::MakeRangeOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto type = cast(op.getResult().getType()); + auto shape = type.getShape(); + auto elementType = type.getElementType(); + auto context = op.getContext(); + + assert(type.getShape().size() == 1 && + isa(type.getElementType()) && + type.getElementType().getIntOrFloatBitWidth() == 32 && + "make range can only return 1D int32 tensor"); + + SmallVector indexingMaps{AffineMap::get( + /* dimCount */ 1, /* symbolCount */ 0, + {mlir::getAffineDimExpr(0, context)}, context)}; + + auto init = rewriter.create(loc, shape, elementType); + + auto nestedBody = [&](OpBuilder &nestedBuilder, Location nestedLoc, + ValueRange blockArgs) { + Value index = nestedBuilder.create(loc, 0); + Value res = + nestedBuilder.create(loc, elementType, index); + nestedBuilder.create(loc, res); + }; + + auto linalgOp = rewriter.create( + loc, op->getResultTypes(), /* operands */ ValueRange{}, ValueRange{init}, + indexingMaps, ConverterUtils::getNParallelLoopsAttrs(1), nestedBody); + + linalgOp->setAttr("tt.from_make_range", mlir::UnitAttr::get(context)); + linalgOp->setAttr("tt.make_range_offset", + mlir::IntegerAttr::get(mlir::IndexType::get(context), 0)); + linalgOp->setAttr( + "tt.make_range_size", + mlir::IntegerAttr::get(mlir::IndexType::get(context), shape[0])); + + int32_t startVal = op.getStartAttr().getInt(); + if (startVal == 0) { + rewriter.replaceOp(op, linalgOp->getResults()); + return success(); + } + + // Apply start offset + Value startScaler = rewriter.create( + loc, rewriter.getI32IntegerAttr(static_cast(startVal))); + auto startInit = rewriter.create(loc, shape, elementType); + Value startTensor = rewriter + .create(loc, ValueRange{startScaler}, + ValueRange{startInit}) + .getResult(0); + auto addOp = + rewriter.create(loc, linalgOp->getResult(0), startTensor); + rewriter.replaceOp(op, addOp); + return success(); +} + +LogicalResult +SplatConverter::matchAndRewrite(triton::SplatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto shape = op.getType().getShape(); + auto init = rewriter.create(loc, shape, + op.getType().getElementType()); + if (llvm::all_of(shape, [](int64_t dim) { return dim == 1; })) { + SmallVector idx(shape.size(), rewriter.create( + loc, rewriter.getIndexAttr(0))); + rewriter.replaceOpWithNewOp(op, adaptor.getSrc(), init, + idx); + } else { + rewriter.replaceOpWithNewOp( + op, ValueRange{adaptor.getSrc()}, ValueRange{init}); + } + return success(); +} + +LogicalResult +UnsplatConverter::matchAndRewrite(triton::UnsplatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto src = adaptor.getSrc(); + auto srcType = cast(src.getType()); + auto shape = srcType.getShape(); + + // Create index constants for all dimensions (all zeros since we're extracting + // the single element) + SmallVector indices; + for (int64_t dim : shape) { + indices.push_back( + rewriter.create(loc, rewriter.getIndexAttr(0))); + } + + // Extract the scalar element from the tensor + auto elementType = srcType.getElementType(); + auto extractOp = + rewriter.create(loc, elementType, src, indices); + rewriter.replaceOp(op, extractOp.getResult()); + return success(); +} + +LogicalResult +ReshapeConverter::matchAndRewrite(triton::ReshapeOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto src = op.getSrc(); + auto dst = op.getResult(); + Value shape = rewriter.create( + loc, + rewriter.getI64TensorAttr(cast(dst.getType()).getShape())); + auto reshapeOp = + rewriter.create(loc, dst.getType(), src, shape); + rewriter.replaceOp(op, reshapeOp.getResult()); + return success(); +} + +LogicalResult ExpandDimsConverter::matchAndRewrite( + triton::ExpandDimsOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto src = op.getSrc(); + auto resShape = cast(op.getResult().getType()).getShape(); + auto axis = op.getAxis(); + + SmallVector reassociation; + + auto src_last_dim = resShape.size() - 2; + auto map_func = [&](unsigned i) -> ReassociationIndices { + if (i < axis) { + return i == src_last_dim ? ReassociationIndices{i, i + 1} + : ReassociationIndices{i}; + } + return i == axis ? ReassociationIndices{i, i + 1} + : ReassociationIndices{i + 1}; + }; + + reassociation = llvm::to_vector( + llvm::map_range(llvm::seq(0, src_last_dim + 1), map_func)); + + auto expandShapeOp = rewriter.create( + op.getLoc(), op.getResult().getType(), src, reassociation); + rewriter.replaceOp(op, expandShapeOp.getResult()); + return success(); +} + +LogicalResult +ClampFConverter::matchAndRewrite(triton::ClampFOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto input = adaptor.getX(); + auto min_para = adaptor.getMin(); + auto max_para = adaptor.getMax(); + auto propagateNan_para = adaptor.getPropagateNan(); + + if (auto input_type = dyn_cast(input.getType())) { + if (isa(min_para.getType())) { + auto minEmptyTensor = rewriter.create( + loc, input_type.getShape(), input_type.getElementType()); + min_para = rewriter + .create(loc, ValueRange{min_para}, + ValueRange{minEmptyTensor}) + .result(); + } + if (isa(max_para.getType())) { + auto maxEmptyTensor = rewriter.create( + loc, input_type.getShape(), input_type.getElementType()); + max_para = rewriter + .create(loc, ValueRange{max_para}, + ValueRange{maxEmptyTensor}) + .result(); + } + } + + if (propagateNan_para == PropagateNan::NONE) { + auto minOp = rewriter.create(loc, input, max_para); + auto maxOp = rewriter.create(loc, min_para, minOp); + rewriter.replaceOp(op, ValueRange{maxOp}); + } else if (propagateNan_para == PropagateNan::ALL) { + auto minOp = rewriter.create(loc, input, max_para); + auto maxOp = rewriter.create(loc, min_para, minOp); + rewriter.replaceOp(op, ValueRange{maxOp}); + } else { + return failure(); + } + + return success(); +} + +// Here convert tt.broadcast to linalg.broadcast +// +// before +// %out = tt.broadcast %in : tensor<1x4x8xf32> -> tensor<128x4x8xf32> +// +// after +// %collpased = tensor.collapse_shape %in [[0, 1], [2]] : +// tensor<1x4x8xf32> into tensor<4x8xf32> +// %out = linalg.broadcast ins(%collpased : tensor<4x8xf32>) +// outs(%empty : tensor<128x4x8xf32>) dimensions = [0] +LogicalResult +BroadcastConverter::matchAndRewrite(triton::BroadcastOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + assert(op->getNumResults() == 1 && "BroadcastOp assumes single result"); + + RankedTensorType sourceType = + cast(adaptor.getSrc().getType()); + RankedTensorType resultType = cast(op.getType()); + auto elementType = resultType.getElementType(); + auto loc = op.getLoc(); + + auto initEmpty = + rewriter.create(loc, resultType.getShape(), elementType); + + SmallVector broadcastDims = + ConverterUtils::getBroadcastDims(sourceType, resultType); + SmallVector unbroadcastDims = + ConverterUtils::getUnbroadcastDims(sourceType, resultType); + + SmallVector collapseReassociationIndices; + auto collapseReassociationIndicesOptional = + getReassociationIndicesForCollapse(sourceType.getShape(), + unbroadcastDims); + if (!collapseReassociationIndicesOptional.has_value()) { + return rewriter.notifyMatchFailure( + op, "Failure with getReassociationIndicesForCollapse call"); + } + collapseReassociationIndices = collapseReassociationIndicesOptional.value(); + + RankedTensorType collapseResultType = + RankedTensorType::get(unbroadcastDims, sourceType.getElementType()); + + auto collpasedOp = rewriter.create( + loc, collapseResultType, adaptor.getSrc(), collapseReassociationIndices); + + auto broadcastOp = rewriter.create( + loc, collpasedOp, initEmpty, + rewriter.getDenseI64ArrayAttr(broadcastDims)); + + rewriter.replaceOp(op, broadcastOp.getResults()); + return success(); +} + +// Reduce Converter +bool ReduceConverter::isReductionOpSupported(Operation *redOp) const { + return isa(redOp); +} + +bool ReduceConverter::isMultiReductionOpSupported(Operation *redOp) { + return isa( + redOp); +} + +Value ReduceConverter::cloneReduceOps(OpBuilder &builder, Value in, Value out, + Value opIns, Value opOuts, + triton::ReduceOp op) const { + auto ® = op->getRegion(0); + assert(reg.getBlocks().size() == 1); + auto &body = reg.getBlocks().front(); + auto numArguments = 2; + assert(body.getNumArguments() == numArguments); + + Value ttIn = body.getArgument(0); + Value ttOut = body.getArgument(1); + + IRMapping mapping; + mapping.map(ttIn, in); + mapping.map(ttOut, out); + + for (auto &op : body.without_terminator()) { + builder.clone(op, mapping); + } + auto yield = cast(body.getTerminator()); + return mapping.lookup(yield->getOperand(0)); +} + +void ReduceConverter::checkIsNotCallOp( + const llvm::SmallVector &reductionOps) const { + llvm::for_each(reductionOps, [](Operation *op) { + assert(!isa(op) && "tt.call ops expected to be inlined in " + "tt.reduce body in ttir building stage"); + }); +} + +bool ReduceConverter::isSCFOpReduce( + const llvm::SmallVector &reductionOps) const { + return (reductionOps.size() == 1 && + reductionOps.front()->getDialect()->getNamespace() == + scf::SCFDialect::getDialectNamespace()); +} + +bool ReduceConverter::isMultiOpReduce( + const llvm::SmallVector &reductionOps) const { + this->checkIsNotCallOp(reductionOps); + + return (reductionOps.size() > 1) || + (reductionOps.size() == 1 && + this->isMultiReductionOpSupported(reductionOps.front())) || + this->isSCFOpReduce(reductionOps); +} + +Value ReduceConverter::computeReduceResultWithCompileFlag( + OpBuilder &opBuilder, Location loc, Value lhs, Value rhs, Value source, + Value initTensor, triton::ReduceOp op, bool compileOn91095Flag) const { + // Original operation list (including all operations) + auto originalReductionOps = this->getReductionOps(op); + auto realReductionOps = this->getRealReductionOps(op); + + // If the size of the original operation list is greater than 1, + // there are additional operations such as type conversion, and these + // operations must be cloned. + bool needClone = compileOn91095Flag || originalReductionOps.size() > 1; + if (needClone) { + return this->cloneReduceOps(opBuilder, lhs, rhs, source, initTensor, op); + } else { + assert(realReductionOps.size() == 1); + auto rop = realReductionOps.front(); + return this->getReductionElement(lhs, rhs, loc, rop, opBuilder, false); + } +} + +LogicalResult +ReduceConverter::convertToTargetOp(triton::ReduceOp op, + typename triton::ReduceOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto source = adaptor.getOperands().front(); + auto sourceType = cast(source.getType()); + auto elemType = sourceType.getElementType(); + auto resType = op.getResult().front().getType(); + auto loc = op.getLoc(); + + // Actual operation list (filtering type conversion operations, leaving only + // actual reduce operations) + auto realReductionOps = this->getRealReductionOps(op); + + bool multiOpReduce = this->isMultiOpReduce(realReductionOps); + // Reduction of arbitrary operations isn't supported because using the first + // element across the reduction dimension requires us to iterate over a + // subview that skips over each first element. + if (!multiOpReduce && + !this->isReductionOpSupported(realReductionOps.front())) { + if (compileOn91095Flag) { + llvm_unreachable("All reduction cases expected to be covered"); + } + return rewriter.notifyMatchFailure( + op, "Only support lowering reduction with single op and limited types " + "of reduction"); + } + + auto rop = realReductionOps.front(); + auto ropLoc = rop->getLoc(); + auto axis = op.getAxis(); + auto isVectorReduce = sourceType.getRank() == 1; + + auto constantType = elemType; + + auto accBaseConstOp = + multiOpReduce + ? this->getMultiOpReductionBaseConstOp(rewriter, op, ropLoc, + constantType) + : this->getReductionBaseConstOp(rewriter, rop, constantType); + + Value initTensor; + if (isVectorReduce) { + auto holder = rewriter.create( + loc, RankedTensorType::get({}, constantType), ValueRange{}); + initTensor = rewriter + .create(loc, accBaseConstOp.getResult(), + holder.getResult()) + .getResult(0); + } else { + Value init = rewriter.create( + loc, cast(resType).getShape(), constantType); + initTensor = + rewriter.create(loc, accBaseConstOp.getResult(), init) + .getResult(0); + } + + Value finalResult = + rewriter + .create( + loc, ValueRange{source}, ValueRange{initTensor}, + SmallVector{axis}, + [&](OpBuilder &opBuilder, Location loc, ValueRange inputs) { + assert(inputs.size() == 2); + Value result = this->computeReduceResultWithCompileFlag( + opBuilder, loc, inputs[0], inputs[1], source, initTensor, + op, compileOn91095Flag); + opBuilder.create(loc, result); + }) + .getResult(0); + + if (sourceType.getRank() == 1) { + finalResult = + rewriter.create(loc, constantType, finalResult); + } + + rewriter.replaceOp(op, finalResult); + return success(); +} + +LogicalResult ReduceConverter::convertToTargetOpExtended( + triton::ReduceOp op, typename triton::ReduceOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto elemTypes = op.getElementTypes(); + + auto valueResultType = dyn_cast(op.getType(0)); + const auto isScalarReduce = valueResultType == nullptr; + + SmallVector outputs; + for (auto i = 0; i < op.getResult().size() && i < elemTypes.size(); i++) { + auto result = dyn_cast(op.getType(i)); + SmallVector resultShape{ + isScalarReduce ? SmallVector{} + : SmallVector(result.getShape())}; + outputs.push_back( + rewriter.create(loc, resultShape, elemTypes[i])); + } + + auto linalgOp = rewriter.create( + loc, adaptor.getOperands(), outputs, + SmallVector{adaptor.getAxis()}, + [&](OpBuilder &b, Location loc, ValueRange inputs) { + auto tritonReduceBlock = op.getBody(); + IRMapping mapping; + mapping.map(tritonReduceBlock->getArguments(), inputs); + + for (auto &op : tritonReduceBlock->without_terminator()) { + b.clone(op, mapping); + } + + auto tritonYield = tritonReduceBlock->getTerminator(); + auto results = + llvm::map_to_vector(tritonYield->getOperands(), + [&](Value val) { return mapping.lookup(val); }); + b.create(loc, results); + }); + + auto params = getReduceWithIndexParams(op); + if (failed(params)) { + return rewriter.notifyMatchFailure(op, "meaningless reduce operation"); + } else if (params->withIndexType != ReduceWithIndexType::None) { + addReduceWithIndexAttr(*params, rewriter, linalgOp); + } + + if (isScalarReduce) { + SmallVector reduceResults; + for (auto i = 0; i < linalgOp.getResults().size() && i < elemTypes.size(); + i++) { + reduceResults.push_back(rewriter.create( + loc, elemTypes[i], linalgOp.getResults()[i], ValueRange{})); + } + rewriter.replaceOp(op, reduceResults); + } else { + rewriter.replaceOp(op, linalgOp); + } + return success(); +} + +bool ScanConverter::isReductionOpSupported(Operation *reductionOp) const { + return isa( + reductionOp); +} + +LogicalResult +ScanConverter::convertToTargetOp(triton::ScanOp op, + typename triton::ScanOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto reductionOps = this->getReductionOps(op); + if (reductionOps.empty()) { + return rewriter.notifyMatchFailure(op, + "No reduction op found in scan body"); + } + + llvm::SmallString<64> funcName; + auto rop = reductionOps.front(); + if (this->isReductionOpSupported(reductionOps.front())) { + if (isa(rop)) { + funcName = "triton_cumsum"; + } else if (isa(rop)) { + funcName = "triton_cumprod"; + } + + auto moduleOp = op->getParentOfType(); + rewriter.setInsertionPoint(moduleOp.getBody(), + std::prev(moduleOp.getBody()->end())); + + auto loc = op.getLoc(); + auto src = adaptor.getOperands().front(); + auto resTy = op.getResult().front().getType(); + auto libFnType = rewriter.getFunctionType( + {src.getType(), rewriter.getI32Type(), rewriter.getI1Type()}, {resTy}); + auto funcOp = rewriter.create(loc, funcName.str(), libFnType); + + SymbolTable symTab(moduleOp); + auto maybePrintFuncNameAttr = symTab.renameToUnique(funcOp, {&symTab}); + if (failed(maybePrintFuncNameAttr)) { + return op->emitError( + "failed to create a unique func name for device_print"); + } + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + + rewriter.setInsertionPoint(op); + auto scanAxis = op.getAxis(); + auto scanReverse = op.getReverse(); + Value axis = rewriter.create(loc, scanAxis, 32); + Value reverseVal = + rewriter.create(loc, scanReverse, 1); + auto callOp = rewriter.create( + loc, funcOp.getSymNameAttr(), TypeRange({resTy}), + ValueRange({src, axis, reverseVal})); + + rewriter.replaceOp(op, callOp); + + return success(); + } else { + // This branch is the associative_scan op. + bool reverse = op.getReverse(); + + auto loc = op.getLoc(); + + Value scanInput = op.getOperand(0); + + auto srcType = mlir::dyn_cast(scanInput.getType()); + if (!srcType) { + return rewriter.notifyMatchFailure( + op, "Expected RankedTensorType input for associative_scan"); + } + + auto elementType = srcType.getElementType(); + auto shape = srcType.getShape(); + int rank = shape.size(); + int axis = op.getAxis(); + + if (axis < 0 || axis >= rank) { + return rewriter.notifyMatchFailure(op, "Invalid scan axis: " + + std::to_string(axis)); + } + + if (op->getNumRegions() < 1 || op->getRegion(0).empty()) { + return rewriter.notifyMatchFailure(op, "Missing combine region"); + } + + OpBuilder::InsertionGuard guard(rewriter); + + auto memrefType = MemRefType::get(shape, elementType); + Value inputMemRef = + rewriter.create(loc, memrefType, scanInput); + Value outputMemRef = rewriter.create(loc, memrefType); + + auto processDimension = [&](ArrayRef baseIdxsArray) { + auto startInd = rewriter.create(op.getLoc(), 0); + if (reverse) { + startInd = rewriter.create(op.getLoc(), + shape[axis] - 1); + } + llvm::SmallVector baseIdxs(baseIdxsArray.begin(), + baseIdxsArray.end()); + llvm::SmallVector firstIdx = baseIdxs; + if (axis <= firstIdx.size()) { + firstIdx.insert(firstIdx.begin() + axis, startInd); + } else { + firstIdx.push_back(startInd); + } + + Value firstVal = + rewriter.create(loc, inputMemRef, firstIdx); + rewriter.create(loc, firstVal, outputMemRef, firstIdx); + + Value axisSize = + rewriter.create(loc, inputMemRef, axis).getResult(); + Value one = rewriter.create(loc, 1); + + Value cmp = rewriter.create(loc, arith::CmpIPredicate::sgt, + axisSize, one); + auto ifOp = rewriter.create(loc, cmp, false); + + // Create a loop only when the axis size is greater than 1. + rewriter.setInsertionPointToStart(ifOp.thenBlock()); + + auto forOp = rewriter.create(loc, one, axisSize, one); + rewriter.setInsertionPointToStart(forOp.getBody()); + + Value k = forOp.getInductionVar(); + if (reverse) { + llvm::SmallVector fixInd; + fixInd.push_back( + rewriter + .create(op.getLoc(), shape[axis] - 1) + .getResult()); + fixInd.push_back(k); + auto fixIndVal = rewriter.create(op.getLoc(), fixInd); + k = fixIndVal.getResult(); + } + llvm::SmallVector currIdx = baseIdxs; + if (axis <= currIdx.size()) { + currIdx.insert(currIdx.begin() + axis, k); + } else { + currIdx.push_back(k); + } + + Value km1 = rewriter.create(loc, k, one); + if (reverse) { + km1 = rewriter.create(loc, k, one); + } + llvm::SmallVector prevIdx = baseIdxs; + if (axis <= prevIdx.size()) { + prevIdx.insert(prevIdx.begin() + axis, km1); + } else { + prevIdx.push_back(km1); + } + + Value currentVal = + rewriter.create(loc, inputMemRef, currIdx); + Value prevResult = + rewriter.create(loc, outputMemRef, prevIdx); + + Region &combineRegion = op->getRegion(0); + Block &combineBlock = combineRegion.front(); + IRMapping mapping; + mapping.map(combineBlock.getArgument(0), prevResult); + mapping.map(combineBlock.getArgument(1), currentVal); + + for (Operation &innerOp : combineBlock.without_terminator()) { + rewriter.clone(innerOp, mapping); + } + + Operation *yieldOp = combineBlock.getTerminator(); + Value resultVal = mapping.lookup(yieldOp->getOperand(0)); + + rewriter.create(loc, resultVal, outputMemRef, currIdx); + + rewriter.setInsertionPointAfter(ifOp); + }; + + // Constructing loops for non-scanning dimensions + llvm::SmallVector nonScanDims; + for (int i = 0; i < rank; ++i) { + if (i != axis) + nonScanDims.push_back(i); + } + + createSimpleNestedLoops(rewriter, loc, outputMemRef, nonScanDims, + processDimension); + + rewriter.setInsertionPointAfter(op); + + mlir::Type resultType = mlir::memref::getTensorTypeFromMemRefType( + dyn_cast(outputMemRef.getType())); + Value outputTensor = rewriter.create( + loc, resultType, outputMemRef, true); + rewriter.replaceOp(op, outputTensor); + return success(); + } +} + +LogicalResult ScanConverter::convertToTargetOpExtended( + triton::ScanOp op, typename triton::ScanOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + bool reverse = op.getReverse(); + + // 1. Extract all input tensors (supports multiple inputs) + auto operands = op->getOperands(); + if (operands.empty()) { + return rewriter.notifyMatchFailure(op, + "No input operands for extended scan"); + } + + // 2. Validate all inputs are of RankedTensorType + llvm::SmallVector inputTensTypes; + for (auto operand : operands) { + auto tensorTy = dyn_cast(operand.getType()); + if (!tensorTy) { + return rewriter.notifyMatchFailure(op, + "All inputs must be RankedTensorType"); + } + inputTensTypes.push_back(tensorTy); + } + + // 3. Validate all input tensors have the same shape (scan operation requires + // matching input dimensions) + auto baseShape = inputTensTypes[0].getShape(); + int rank = baseShape.size(); + int axis = op.getAxis(); + if (axis < 0 || axis >= rank) { + return rewriter.notifyMatchFailure(op, "Invalid scan axis: " + + std::to_string(axis)); + } + for (size_t i = 1; i < inputTensTypes.size(); ++i) { + if (inputTensTypes[i].getShape() != baseShape) { + return rewriter.notifyMatchFailure(op, + "All inputs must have the same shape"); + } + } + + // 4. Prepare MemRefs for multiple inputs/outputs + llvm::SmallVector inputMemRefs; + llvm::SmallVector outputMemRefs; + llvm::SmallVector memRefTypes; + for (size_t i = 0; i < inputTensTypes.size(); ++i) { + auto &tensorTy = inputTensTypes[i]; + auto memRefTy = + MemRefType::get(tensorTy.getShape(), tensorTy.getElementType()); + memRefTypes.push_back(memRefTy); + // Convert input tensors to MemRefs + inputMemRefs.push_back( + rewriter.create(loc, memRefTy, operands[i])); + // Allocate MemRefs for outputs + outputMemRefs.push_back(rewriter.create(loc, memRefTy)); + } + + // 5. Define scanning logic for multiple inputs/outputs + LogicalResult loopResult = success(); + auto processDimension = [&](ArrayRef baseIdxsArray) { + llvm::SmallVector baseIdxs(baseIdxsArray.begin(), + baseIdxsArray.end()); + + auto startInd = rewriter.create(op.getLoc(), 0); + if (reverse) { + startInd = rewriter.create(op.getLoc(), + baseShape[axis] - 1); + } + + llvm::SmallVector firstIdx = baseIdxs; + if (axis <= firstIdx.size()) { + firstIdx.insert(firstIdx.begin() + axis, startInd); + } else { + firstIdx.push_back(startInd); + } + + // 5.1 Process the first element: directly copy multiple inputs to multiple + // outputs (initialize cumulative results) + for (size_t i = 0; i < inputMemRefs.size(); ++i) { + Value firstVal = + rewriter.create(loc, inputMemRefs[i], firstIdx); + rewriter.create(loc, firstVal, outputMemRefs[i], + firstIdx); + } + + Value axisSize = + rewriter.create(loc, baseShape[axis]); + Value one = rewriter.create(loc, 1); + + Value cmp = rewriter.create(loc, arith::CmpIPredicate::sgt, + axisSize, one); + auto ifOp = rewriter.create(loc, cmp, false); + + // Create a loop only when the axis size is greater than 1. + rewriter.setInsertionPointToStart(ifOp.thenBlock()); + + // Use a forward loop, but handle reverse indexing inside the loop. + auto forOp = rewriter.create(loc, one, axisSize, one); + rewriter.setInsertionPointToStart(forOp.getBody()); + + Value k = forOp.getInductionVar(); + + if (reverse) { + // Reverse scanning: Convert the forward loop index to the actual reverse + // index. (axis_size - 1) - k + Value axisSizeVal = + rewriter.create(loc, baseShape[axis]); + Value axisSizeMinusOne = + rewriter.create(loc, axisSizeVal, one); + k = rewriter.create(loc, axisSizeMinusOne, k); + } + + llvm::SmallVector currIdx = baseIdxs; + if (axis <= currIdx.size()) { + currIdx.insert(currIdx.begin() + axis, k); + } else { + currIdx.push_back(k); + } + + Value prevIndex; + if (reverse) { + prevIndex = rewriter.create(loc, k, one); + } else { + prevIndex = rewriter.create(loc, k, one); + } + + llvm::SmallVector prevIdx = baseIdxs; + if (axis <= prevIdx.size()) { + prevIdx.insert(prevIdx.begin() + axis, prevIndex); + } else { + prevIdx.push_back(prevIndex); + } + + // 5.4 Load current elements and previous cumulative results + llvm::SmallVector currentVals; + llvm::SmallVector prevResults; + for (size_t i = 0; i < inputMemRefs.size(); ++i) { + currentVals.push_back( + rewriter.create(loc, inputMemRefs[i], currIdx)); + prevResults.push_back( + rewriter.create(loc, outputMemRefs[i], prevIdx)); + } + + // 5.5 Bind parameters for custom reduction logic + Region &combineRegion = op->getRegion(0); + if (combineRegion.empty()) { + op->emitError("Missing combine region in extended scan"); + loopResult = failure(); + return; + } + Block &combineBlock = combineRegion.front(); + // Validate that the number of reduction region arguments matches (number of + // previous results + number of current elements) + if (combineBlock.getNumArguments() != 2 * inputMemRefs.size()) { + op->emitError("Combine region arguments mismatch with input count"); + loopResult = failure(); + return; + } + IRMapping mapping; + for (size_t i = 0; i < inputMemRefs.size(); ++i) { + // Bind previous results (previous value of the i-th output) to the i-th + // argument of the reduction region + mapping.map(combineBlock.getArgument(i), prevResults[i]); + // Bind current elements (current value of the i-th input) to the i+N-th + // argument of the reduction region (N is the number of inputs) + mapping.map(combineBlock.getArgument(i + inputMemRefs.size()), + currentVals[i]); + } + + // 5.6 Clone all operations within the reduction region + for (Operation &innerOp : combineBlock.without_terminator()) { + rewriter.clone(innerOp, mapping); + } + + // 5.7 Extract reduction results and store them in outputMemRef + Operation *yieldOp = combineBlock.getTerminator(); + if (yieldOp->getNumOperands() != outputMemRefs.size()) { + op->emitError("Combine region returns mismatch with output count"); + loopResult = failure(); + return; + } + for (size_t i = 0; i < outputMemRefs.size(); ++i) { + Value resultVal = mapping.lookup(yieldOp->getOperand(i)); + rewriter.create(loc, resultVal, outputMemRefs[i], + currIdx); + } + + rewriter.setInsertionPointAfter(ifOp); + }; + + // 6. Generate nested loops for non-scan dimensions + llvm::SmallVector nonScanDims; + for (int i = 0; i < rank; ++i) { + if (i != axis) + nonScanDims.push_back(i); + } + createSimpleNestedLoops(rewriter, loc, outputMemRefs[0], nonScanDims, + processDimension); + + if (failed(loopResult)) { + return failure(); + } + + // 7. Convert multiple output MemRefs back to tensors and replace the original + // tt.scan operation + llvm::SmallVector outputTensors; + for (auto outputMemRef : outputMemRefs) { + mlir::Type resultType = mlir::memref::getTensorTypeFromMemRefType( + dyn_cast(outputMemRef.getType())); + outputTensors.push_back(rewriter.create( + loc, resultType, outputMemRef, true)); + } + rewriter.replaceOp(op, outputTensors); + + return success(); +} + +LogicalResult ExternElementwiseClOpConverter::matchAndRewrite( + triton::ExternElementwiseOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + if (!op.getPure()) { + op->emitWarning() << "impure elementwise op!"; + return failure(); + } + if (op.getSymbol().contains("__hmf_")) { + // 1. get or create the declaration of external elementwise function + Type dstTy = op.getResult().getType(); + bool isDstScalar = !isa(dstTy); + Type dstElemTy = + isDstScalar ? dstTy : cast(dstTy).getElementType(); + SmallVector srcElemTys; + SmallVector srcs; + for (auto src : op.getSrcs()) { + if (!isa(src.getType())) { + src = rewriter.create( + op.getLoc(), RankedTensorType::get({(int64_t)1}, src.getType()), + src); + } + srcs.push_back(src); + srcElemTys.push_back( + cast(src.getType()).getElementType()); + } + FunctionType elemFuncType = + FunctionType::get(rewriter.getContext(), srcElemTys, {dstElemTy}); + auto mod = SymbolTable::getNearestSymbolTable(op); + auto extFunc = dyn_cast_or_null( + SymbolTable::lookupSymbolIn(mod, op.getSymbol())); + // std::string symbol = op.getSymbol().str(); + bool is_libdevice = llvm::is_contained(libdeviceOps, op.getSymbol()) && + getEnvBool("TRITON_ENABLE_LIBDEVICE_SIMT", false); + if (!extFunc) { + OpBuilder::InsertionGuard guard(rewriter); + rewriter.setInsertionPointToStart(&mod->getRegion(0).front()); + extFunc = rewriter.create(rewriter.getUnknownLoc(), + op.getSymbol(), elemFuncType); + extFunc.setPrivate(); + extFunc->setAttr(LLVM::LLVMDialect::getReadnoneAttrName(), + UnitAttr::get(rewriter.getContext())); + // set coreType for external func, otherwise InferFuncCoreTypePass will + // fail + if (is_libdevice) { + hivm::TFuncCoreType e = hivm::TFuncCoreType::AIV; + extFunc->setAttr( + hivm::TFuncCoreTypeAttr::name, + hivm::TFuncCoreTypeAttr::get(extFunc->getContext(), e)); + } + } + assert(isa( + SymbolTable::lookupSymbolIn(mod, op.getSymbol()))); + // 2. prepare the output tensor + Value output; + if (isDstScalar) { + dstTy = RankedTensorType::get({(int64_t)1}, dstElemTy); + } + bool found = false; + for (Value v : srcs) { + if (v.getType() == dstTy) { + found = true; + output = v; + break; + } + } + if (!found) { + output = rewriter.create( + op.getLoc(), cast(dstTy).getShape(), dstElemTy); + } + + if (is_libdevice) { + auto srcType = cast(srcs[0].getType()); + SmallVector dimSizes; + int64_t rank = srcType.getRank(); + for (int i = 0; i < rank; ++i) { + if (srcType.isDynamicDim(i)) { + auto dimOp = rewriter.create(loc, srcs[0], i); + dimSizes.push_back(dimOp); + } else { + auto constOp = rewriter.create( + loc, srcType.getDimSize(i)); + dimSizes.push_back(constOp); + } + } + // building nested loops by recursion + std::function, Value)> + buildLoops = [&](OpBuilder &b, Location loc, + SmallVector indices, Value acc) -> Value { + int64_t dim = indices.size(); + if (dim == rank) { + // innermost loop + SmallVector elemVals; + for (auto src : srcs) { + auto extract = b.create(loc, src, indices); + elemVals.push_back(extract); + } + auto call = + b.create(loc, op.getSymbol(), dstElemTy, elemVals); + auto insert = + b.create(loc, call.getResult(0), acc, indices); + return insert; + } else { + Value lower = b.create(loc, 0); + Value upper = dimSizes[dim]; + Value step = b.create(loc, 1); + auto loop = + b.create(loc, lower, upper, step, ValueRange{acc}); + Block *body = loop.getBody(); + OpBuilder innerBuilder = OpBuilder::atBlockBegin(body); + SmallVector newIndices = indices; + newIndices.push_back(loop.getInductionVar()); + Value innerAcc = loop.getRegionIterArgs()[0]; + Value updatedAcc = + buildLoops(innerBuilder, loc, newIndices, innerAcc); + innerBuilder.create(loc, updatedAcc); + return loop.getResult(0); + } + }; + + Value result = buildLoops(rewriter, loc, {}, output); + if (isDstScalar) { + SmallVector zeroIndices( + rank, rewriter.create(loc, 0)); + auto extract = + rewriter.create(loc, result, zeroIndices); + rewriter.replaceOp(op, extract); + } else { + rewriter.replaceOp(op, result); + } + return success(); + } + // 3. create the linalg.map op + auto mapOp = rewriter.create( + loc, + /*inputs=*/srcs, + /*init=*/output, + /*bodyBuilder=*/ + [&](OpBuilder &builder, Location loc, ValueRange regionArgs) { + auto elemOp = builder.create(loc, + /*name=*/op.getSymbol(), + /*resultType=*/dstElemTy, + /*operands=*/regionArgs); + builder.create(loc, elemOp->getResults()); + }); + if (isDstScalar) { + // need to convert tensor back to scalar + auto indexType = rewriter.getIndexType(); + Value zeroConstant = rewriter.create( + loc, indexType, rewriter.getIntegerAttr(indexType, 0)); + auto extractOp = rewriter.create( + loc, mapOp.getResults()[0], zeroConstant); + rewriter.replaceOp(op, extractOp); + } else { + rewriter.replaceOp(op, mapOp); + } + return success(); + } + return failure(); +} + +LogicalResult UnrealizedCastConverter::matchAndRewrite( + UnrealizedConversionCastOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + rewriter.eraseOp(op); + return success(); +} + +LogicalResult +JoinConverter::matchAndRewrite(triton::JoinOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + Value opa = op.getLhs(); + Value opb = op.getRhs(); + auto loc = op.getLoc(); + + auto resType = dyn_cast(op.getResult().getType()); + Value emptyOp = rewriter.create(loc, resType.getShape(), + resType.getElementType()); + + auto shape = dyn_cast(opa.getType()).getShape(); + auto sizes = llvm::map_to_vector(shape, [&](int64_t t) { + return OpFoldResult(rewriter.getI64IntegerAttr(t)); + }); + sizes.push_back(rewriter.getI64IntegerAttr(1)); + + int64_t rank = resType.getRank(); + + // Set last dimension stride to 2 in layout + // As last dimension size is always 1, last dimension stride here could be + // either 1 or 2, while stride `2` could carry interleave trait and it's + // convenient for next lower. + SmallVector strides(rank, rewriter.getIndexAttr(1)); + strides.back() = rewriter.getIndexAttr(2); + + SmallVector offsets(rank, rewriter.getIndexAttr(0)); + + auto insert0 = rewriter.create( + loc, opa, emptyOp, offsets, sizes, strides); + + offsets.back() = rewriter.getIndexAttr(1); + auto insert1 = rewriter.create( + loc, opb, insert0, offsets, sizes, strides); + rewriter.replaceOp(op, insert1); + return success(); +} + +LogicalResult +CatConverter::matchAndRewrite(triton::CatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + Value opa = op.getLhs(); + Value opb = op.getRhs(); + auto loc = op.getLoc(); + + auto resType = dyn_cast(op.getResult().getType()); + if (!resType || resType.getRank() != 1) { + return rewriter.notifyMatchFailure(op, "only support 1D cat"); + } + + auto inputTypeA = dyn_cast(opa.getType()); + auto inputTypeB = dyn_cast(opb.getType()); + if (!inputTypeA || !inputTypeB || inputTypeA.getRank() != 1 || + inputTypeB.getRank() != 1) { + return rewriter.notifyMatchFailure(op, "inputs must be 1D tensors"); + } + + int64_t sizeA = inputTypeA.getShape()[0]; + int64_t sizeB = inputTypeB.getShape()[0]; + + // Only handle the case where both inputs have size 1 (i.e., scalar-like) + if (sizeA == 1 && sizeB == 1) { + // Use scalar extract + insert + auto emptyOp = rewriter.create(loc, resType.getShape(), + resType.getElementType()); + + Value zero = rewriter.create(loc, 0); + Value one = rewriter.create(loc, 1); + + Value scalarA = rewriter.create(loc, opa, zero); + Value scalarB = rewriter.create(loc, opb, zero); + + Value inserted0 = + rewriter.create(loc, scalarA, emptyOp, zero); + Value inserted1 = + rewriter.create(loc, scalarB, inserted0, one); + + rewriter.replaceOp(op, inserted1); + return success(); + } + + // General case: use tensor.insert_slice + auto emptyOp = rewriter.create(loc, resType.getShape(), + resType.getElementType()); + + auto rank = resType.getRank(); + SmallVector offsets(rank, rewriter.getIndexAttr(0)); + SmallVector strides(rank, rewriter.getIndexAttr(1)); + + auto inputType = dyn_cast(opa.getType()); + + SmallVector sizes = + llvm::map_to_vector(inputType.getShape(), [&](int64_t t) { + return OpFoldResult(rewriter.getI64IntegerAttr(t)); + }); + + auto insert0 = rewriter.create( + loc, opa, emptyOp, offsets, sizes, strides); + + offsets[0] = + rewriter.getIndexAttr(inputType.getRank() ? inputType.getShape()[0] : 1); + auto insert1 = rewriter.create( + loc, opb, insert0, offsets, sizes, strides); + + rewriter.replaceOp(op, insert1); + return success(); +} + +/// @brief Convert tt.gather to func.call. BiShengIR captures the func +/// with assumed semantics. +/// @param op The `triton::GatherOp` operation to be rewritten. +/// @param adaptor An adaptor for the operation's operands. +/// @param rewriter A pattern rewriter used to modify the IR. +/// @return A `LogicalResult` indicating whether the rewrite was successful. +LogicalResult +GatherConverter::matchAndRewrite(triton::GatherOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + Value src = adaptor.getSrc(); + Value idx = adaptor.getIndices(); + Value res = op.getResult(); + auto gatherAxis = op.getAxis(); + + auto moduleOp = op->getParentOfType(); + rewriter.setInsertionPoint(moduleOp.getBody(), + std::prev(moduleOp.getBody()->end())); + + llvm::SmallString funcName = gatherFuncNameBase; + int uniqueId = 0; + while (SymbolTable::lookupSymbolIn(moduleOp, funcName)) { + funcName = gatherFuncNameBase; + funcName += ("_" + std::to_string(uniqueId++)); + } + + auto resTy = res.getType(); + auto libFnType = rewriter.getFunctionType( + {src.getType(), idx.getType(), rewriter.getI32Type()}, {resTy}); + auto funcOp = rewriter.create(loc, funcName.str(), libFnType); + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + + rewriter.setInsertionPoint(op); + Value axis = rewriter.create(loc, gatherAxis, 32); + auto callOp = rewriter.create(loc, funcOp.getSymNameAttr(), + TypeRange({resTy}), + ValueRange({src, idx, axis})); + + rewriter.replaceOp(op, callOp); + + return success(); +} + +LogicalResult +SplitConverter::matchAndRewrite(triton::SplitOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + Value input = op.getSrc(); + auto loc = op.getLoc(); + auto inputType = cast(input.getType()); + + int64_t rank = inputType.getRank(); + SmallVector offsets(rank, rewriter.getIndexAttr(0)); + // Similar to JoinConverter, here adjust last dimension stride + SmallVector strides(rank, rewriter.getIndexAttr(1)); + strides.back() = rewriter.getIndexAttr(2); + + auto outType = dyn_cast(op.getOutLHS().getType()); + auto sizes = llvm::map_to_vector(outType.getShape(), [&](int64_t t) { + return OpFoldResult(rewriter.getIndexAttr(t)); + }); + sizes.push_back(rewriter.getIndexAttr(1)); + + auto slice0 = rewriter.create( + loc, outType, input, offsets, sizes, strides); + + offsets.back() = rewriter.getIndexAttr(1); + auto slice1 = rewriter.create( + loc, outType, input, offsets, sizes, strides); + + SmallVector slices = {slice0.getResult(), slice1.getResult()}; + rewriter.replaceOp(op, ValueRange(slices)); + return success(); +} + +/* +the element-wise most significant N bits of the 2N-bit product of x and y +%x:2 = arith.mulsi_extended %y, %z : tensor<4x?xi32> +*/ +LogicalResult TritonMulhiuiConverter::matchAndRewrite( + triton::MulhiUIOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + Value opl = op.getX(); + Value opr = op.getY(); + Value res = op.getResult(); + auto newMulOp = rewriter.create( + loc, res.getType(), res.getType(), opl, opr); + // triton only need the high value + rewriter.replaceOp(op, ValueRange{newMulOp.getHigh()}); + return success(); +} + +LogicalResult TritonPreciseSqrtConverter::matchAndRewrite( + triton::PreciseSqrtOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + rewriter.replaceOpWithNewOp(op, adaptor.getOperands()); + return success(); +} + +LogicalResult DevicePrintConverter::matchAndRewrite( + triton::PrintOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto moduleOp = op->getParentOfType(); + rewriter.setInsertionPoint(moduleOp.getBody(), + std::prev(moduleOp.getBody()->end())); + SmallVector inputTypes; + for (auto arg : op.getArgs()) { + inputTypes.push_back(arg.getType()); + } + auto libFnType = rewriter.getFunctionType(inputTypes, {}); + auto funcOp = + rewriter.create(op.getLoc(), printFuncNameBase, libFnType); + SymbolTable symTab(moduleOp); + auto maybePrintFuncNameAttr = symTab.renameToUnique(funcOp, {&symTab}); + if (failed(maybePrintFuncNameAttr)) { + return op->emitError( + "failed to create a unique func name for device_print"); + } + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + auto prefixAttr = op.getPrefixAttr(); + funcOp->setAttr(prefixAttrName, prefixAttr); + auto hexAttr = op.getHexAttr(); + funcOp->setAttr(hexAttrName, hexAttr); + + rewriter.setInsertionPoint(op); + rewriter.create(op.getLoc(), funcOp, op.getArgs()); + + rewriter.eraseOp(op); + return success(); +} + +LogicalResult DeviceAssertConverter::matchAndRewrite( + triton::AssertOp op, OpAdaptor adaptor, + mlir::ConversionPatternRewriter &rewriter) const { + auto msgAttr = op.getMessageAttr(); + // Filter out automatically inserted assert ops + if (auto strAttr = mlir::dyn_cast(msgAttr)) { + llvm::StringRef msg = strAttr.getValue(); + if (msg.contains("overflow detected for operation")) { + rewriter.eraseOp(op); + return success(); + } + } + + auto moduleOp = op->getParentOfType(); + rewriter.setInsertionPoint(moduleOp.getBody(), + std::prev(moduleOp.getBody()->end())); + auto conditionType = op.getCondition().getType(); + + auto libFnType = rewriter.getFunctionType({conditionType}, {}); + auto funcOp = + rewriter.create(op.getLoc(), printFuncNameBase, libFnType); + mlir::SymbolTable symTab(moduleOp); + auto maybePrintFuncNameAttr = symTab.renameToUnique(funcOp, {&symTab}); + if (failed(maybePrintFuncNameAttr)) { + return op->emitError( + "failed to create a unique func name for device_assert"); + } + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + funcOp->setAttr(msgAttrName, msgAttr); + + rewriter.setInsertionPoint(op); + rewriter.create(op.getLoc(), funcOp, + ValueRange{op.getCondition()}); + + rewriter.eraseOp(op); + return success(); +} + +LogicalResult +MatmulConverter::matchAndRewrite(triton::DotOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto opa = adaptor.getA(); + auto opb = adaptor.getB(); + auto opc = adaptor.getC(); + auto dstType = cast(op.getType()); + auto elemTy = dstType.getElementType(); + auto inputPrec = op.getInputPrecision(); + + auto createOp = [&](auto &&rewriter, ValueRange operands, + ValueRange results) -> Operation * { + if (dstType.getRank() == 2) + return rewriter.template create(op.getLoc(), operands, + results); + else if (dstType.getRank() == 3) + return rewriter.template create(op.getLoc(), + operands, results); + llvm_unreachable("Datatype of DotOp operands could only be 2D or 3D"); + }; + + auto replaceOp = [&](auto &&rewriter, ValueRange operands, + ValueRange results) -> Operation * { + if (dstType.getRank() == 2) + return rewriter.template replaceOpWithNewOp( + op, operands, results); + else if (dstType.getRank() == 3) + return rewriter.template replaceOpWithNewOp( + op, operands, results); + llvm_unreachable("Datatype of DotOp operands could only be 2D or 3D"); + }; + + Operation *matmulOp; + if (mlir::isa(elemTy) && !elemTy.isF32()) { + RankedTensorType opcFp32Ty = + RankedTensorType::get(dstType.getShape(), rewriter.getF32Type()); + Value opcFp32 = rewriter.create(op.getLoc(), opcFp32Ty, opc); + matmulOp = createOp(rewriter, ValueRange{opa, opb}, ValueRange{opcFp32}); + auto roundModeAttr = hfusion::RoundModeAttr::get(rewriter.getContext(), + hfusion::RoundMode::RINT); + auto truncOp = rewriter.replaceOpWithNewOp( + op, dstType, matmulOp->getResult(0)); + truncOp->setAttr("round_mode", roundModeAttr); + } else { + matmulOp = replaceOp(rewriter, ValueRange{opa, opb}, ValueRange{opc}); + } + matmulOp->setAttr("input_precision", + rewriter.getStringAttr(stringifyInputPrecision(inputPrec))); + return success(); +} + +LogicalResult +FlipOpConverter::matchAndRewrite(triton::dicp::FlipOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + Value src = adaptor.getSrc(); + auto rankedSrcTy = cast(src.getType()); + + MLIRContext *ctx = rewriter.getContext(); + + Type valuesTy = src.getType(); + Location loc = op.getLoc(); + + auto dimAttr = op->getAttrOfType("dim"); + if (!dimAttr) { + op->emitError("missing 'dim' attribute"); + return failure(); + } + + auto moduleOp = op->getParentOfType(); + if (!moduleOp) { + op->emitError("must be inside a module"); + return failure(); + } + + // Unique callee name: triton_flip, triton_flip_1, … + std::string funcName = baseFuncName.str(); + int uniqueId = 0; + while (SymbolTable::lookupSymbolIn(moduleOp, funcName)) + funcName = (baseFuncName + Twine("_") + Twine(uniqueId++)).str(); + + auto i64Ty = IntegerType::get(ctx, 64); + auto libFnType = + rewriter.getFunctionType({rankedSrcTy, i64Ty}, {rankedSrcTy}); + + // Declare the callee + auto moduleIP = rewriter.saveInsertionPoint(); + rewriter.setInsertionPointToEnd(moduleOp.getBody()); + auto funcOp = rewriter.create(loc, funcName, libFnType); + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + rewriter.restoreInsertionPoint(moduleIP); + + // dim constant + Value dimVal = + rewriter.create(loc, dimAttr.getInt(), 64); + + // Call the backend function + auto callee = SymbolRefAttr::get(ctx, funcOp.getSymName()); + auto callOp = rewriter.create( + loc, TypeRange({rankedSrcTy}), callee, ValueRange({src, dimVal})); + + Value finalValues = callOp.getResult(0); + + rewriter.replaceOp(op, {finalValues}); + return success(); +} + +LogicalResult +SortOpConverter::matchAndRewrite(triton::dicp::SortOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + Value src = adaptor.getSrc(); + auto rankedSrcTy = cast(src.getType()); + auto srcElemTy = rankedSrcTy.getElementType(); + auto srcShape = rankedSrcTy.getShape(); + auto srcEnc = rankedSrcTy.getEncoding(); + + MLIRContext *ctx = rewriter.getContext(); + + Type backendElemTy = srcElemTy; + if (srcElemTy.isInteger(8)) { + backendElemTy = Float16Type::get(ctx); // i8 -> f16 + } else if (srcElemTy.isInteger(16)) { + backendElemTy = Float32Type::get(ctx); // i16 -> f32 + } + Type backendTensorTy = RankedTensorType::get(srcShape, backendElemTy, srcEnc); + + Type valuesTy = src.getType(); + + Location loc = op.getLoc(); + auto dimAttr = op->getAttrOfType("dim"); + auto descAttr = op->getAttrOfType("descending"); + if (!dimAttr || !descAttr) { + op->emitError("missing 'dim' or 'descending' attribute"); + return failure(); + } + + auto moduleOp = op->getParentOfType(); + if (!moduleOp) { + op->emitError("must be inside a module"); + return failure(); + } + + llvm::SmallString<64> baseName("triton_sort"); + llvm::SmallString<64> funcName = baseName; + int uniqueId = 0; + while (SymbolTable::lookupSymbolIn(moduleOp, funcName)) { + funcName = baseName; + funcName += ("_" + std::to_string(uniqueId++)); + } + + auto i64Ty = IntegerType::get(ctx, 64); + auto i1Ty = IntegerType::get(ctx, 1); + auto libFnType = rewriter.getFunctionType({backendTensorTy, i64Ty, i1Ty}, + {backendTensorTy}); + + auto moduleIP = rewriter.saveInsertionPoint(); + rewriter.setInsertionPointToEnd(moduleOp.getBody()); + auto funcOp = rewriter.create(loc, funcName.str(), libFnType); + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + rewriter.restoreInsertionPoint(moduleIP); + + Value srcForCall = src; + if (backendElemTy != srcElemTy) { + srcForCall = rewriter.create(loc, backendTensorTy, src); + } + + Value dimVal = + rewriter.create(loc, dimAttr.getInt(), 64); + Value descVal = rewriter.create( + loc, descAttr.getValue() ? 1 : 0, 1); + + auto callee = SymbolRefAttr::get(ctx, funcOp.getSymName()); + auto callOp = + rewriter.create(loc, TypeRange({backendTensorTy}), callee, + ValueRange({srcForCall, dimVal, descVal})); + + Value valuesFloat = callOp.getResult(0); // tensor + + Value finalValues = valuesFloat; + if (backendElemTy != srcElemTy) { + finalValues = rewriter.create(loc, valuesTy, valuesFloat); + } + + rewriter.replaceOp(op, {finalValues}); + + return success(); +} + +LogicalResult +DotScaledConverter::matchAndRewrite(triton::DotScaledOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + Location loc = op.getLoc(); + + Value lhs = adaptor.getA(); + Value rhs = adaptor.getB(); + Value c = adaptor.getC(); + Value lhsScale = adaptor.getAScale(); + Value rhsScale = adaptor.getBScale(); + RankedTensorType dstType = cast(op.getType()); + + auto lhsElemType = op.getAElemType(); + auto rhsElemType = op.getBElemType(); + + bool isFP8Input = (lhsElemType == triton::ScaleDotElemType::E4M3 || + lhsElemType == triton::ScaleDotElemType::E5M2) && + (rhsElemType == triton::ScaleDotElemType::E4M3 || + rhsElemType == triton::ScaleDotElemType::E5M2); + bool isFP4Input = (lhsElemType == triton::ScaleDotElemType::E2M1) && + (rhsElemType == triton::ScaleDotElemType::E2M1); + if (isFP8Input || isFP4Input) { + if (!rhsScale) { + RankedTensorType defaultScaleTy = + RankedTensorType::get({1}, rewriter.getI8Type()); + Value defaultScaleVal = + rewriter.create(loc, rewriter.getI8IntegerAttr(1)); + Value defaultScaleEmpty = rewriter.create( + loc, defaultScaleTy.getShape(), defaultScaleTy.getElementType()); + rhsScale = rewriter + .create(loc, ValueRange{defaultScaleVal}, + ValueRange{defaultScaleEmpty}) + .getResult(0); + } + Value acc = c ? c + : rewriter.create(loc, dstType.getShape(), + dstType.getElementType()); + + auto convertFormat = + [&](triton::ScaleDotElemType ty) -> mlir::hfusion::DataformatAttr { + auto ctx = rewriter.getContext(); + switch (ty) { + case triton::ScaleDotElemType::E2M1: + return mlir::hfusion::DataformatAttr::get( + ctx, mlir::hfusion::Dataformat::FP4E2M1_T); + case triton::ScaleDotElemType::E4M3: + return mlir::hfusion::DataformatAttr::get( + ctx, mlir::hfusion::Dataformat::FP8E4M3_T); + case triton::ScaleDotElemType::E5M2: + return mlir::hfusion::DataformatAttr::get( + ctx, mlir::hfusion::Dataformat::FP8E5M2_T); + default: + llvm_unreachable("unsupported ScaleDotElemType"); + } + }; + + auto lhsFmt = convertFormat(lhsElemType); + auto rhsFmt = convertFormat(rhsElemType); + + Value matmulMxResult = rewriter.create( + loc, dstType, lhs, rhs, lhsScale, rhsScale, acc, lhsFmt, rhsFmt); + + Value finalResult = matmulMxResult; + if (dstType.getElementType().isBF16()) { + finalResult = + rewriter.create(loc, dstType, matmulMxResult); + } + rewriter.replaceOp(op, finalResult); + return success(); + } + + if (!lhsScale) { + return op.emitError("lhsScale is required for non-FP8 input"); + } + + RankedTensorType lhsTy = cast(lhs.getType()); + RankedTensorType lhsScaleTy = cast(lhsScale.getType()); + RankedTensorType rhsScaleTy = + rhsScale ? cast(rhsScale.getType()) : nullptr; + RankedTensorType rhsTy = cast(rhs.getType()); + + Value lhsScaleOut; + Value rhsScaleOut; + Value c127 = rewriter.create( + op.getLoc(), rewriter.getI16Type(), rewriter.getI16IntegerAttr(127)); + Value c7 = rewriter.create( + op.getLoc(), rewriter.getI16Type(), rewriter.getI16IntegerAttr(7)); + Type i16Ty = rewriter.getI16Type(); + Type bf16Ty = rewriter.getBF16Type(); + Type fp16Ty = rewriter.getF16Type(); + Type fp32Ty = rewriter.getF32Type(); + bool fastMath = op.getFastMath(); + + auto createNanSplat = [&](RankedTensorType tensorTy) -> Value { + auto floatTy = cast(tensorTy.getElementType()); + auto nanAttr = rewriter.getFloatAttr( + floatTy, APFloat::getNaN(floatTy.getFloatSemantics())); + Value empty = rewriter.create(loc, tensorTy.getShape(), + tensorTy.getElementType()); + return rewriter + .create( + loc, ValueRange{rewriter.create(loc, nanAttr)}, + ValueRange{empty}) + .getResult(0); + }; + + auto createNaNMask = [&](Value scaleTensor, + RankedTensorType scaleTy) -> Value { + if (scaleTy.getElementType().isIntOrIndex()) { + auto bitWidth = scaleTy.getElementTypeBitWidth(); + auto allOnes = APInt::getAllOnes(bitWidth); + auto sentinel = rewriter.create( + loc, scaleTy, DenseElementsAttr::get(scaleTy, allOnes)); + return rewriter + .create(loc, arith::CmpIPredicate::eq, scaleTensor, + sentinel) + .getResult(); + } + + return rewriter + .create(loc, arith::CmpFPredicate::UNO, scaleTensor, + scaleTensor) + .getResult(); + }; + + auto applyNaNMask = [&](Value valueTensor, Value maskTensor) -> Value { + auto valueTy = cast(valueTensor.getType()); + Value nanTensor = createNanSplat(valueTy); + return rewriter + .create(loc, maskTensor, nanTensor, valueTensor) + .getResult(); + }; + + if (lhsScaleTy.getElementType().isIntOrIndex()) { + RankedTensorType lhsScaleI16Ty = + RankedTensorType::get(lhsScaleTy.getShape(), i16Ty); + Value lhsScaleI16 = + rewriter.create(op.getLoc(), lhsScaleI16Ty, lhsScale); + + Value lhsShift127Empty = rewriter.create( + op.getLoc(), lhsScaleI16Ty.getShape(), i16Ty); + Value lhsShift127 = + rewriter + .create(op.getLoc(), ValueRange{c127}, + ValueRange{lhsShift127Empty}) + .getResult(0); + + Value lhsScaleI16Add127 = + rewriter.create(op.getLoc(), lhsScaleI16, lhsShift127); + + Value lhsShift7Empty = rewriter.create( + op.getLoc(), lhsScaleI16Ty.getShape(), i16Ty); + Value lhsShift7 = rewriter + .create(op.getLoc(), ValueRange{c7}, + ValueRange{lhsShift7Empty}) + .getResult(0); + Value lhsScaleI16Shifted = rewriter.create( + op.getLoc(), lhsScaleI16Add127, lhsShift7); + + RankedTensorType lhsScaleBF16Ty = + RankedTensorType::get(lhsScaleTy.getShape(), bf16Ty); + Value lhsScaleBF16 = rewriter.create( + op.getLoc(), lhsScaleBF16Ty, lhsScaleI16Shifted); + if (lhsTy.getElementType() == fp16Ty) { + RankedTensorType lhsScaleFp32Ty = + RankedTensorType::get(lhsScaleTy.getShape(), fp32Ty); + Value lhsScaleFp32 = rewriter.create( + op.getLoc(), lhsScaleFp32Ty, lhsScaleBF16); + RankedTensorType lhsScaleFp16Ty = + RankedTensorType::get(lhsScaleTy.getShape(), fp16Ty); + lhsScaleOut = rewriter.create( + op.getLoc(), lhsScaleFp16Ty, lhsScaleFp32); + } else { + lhsScaleOut = lhsScaleBF16; + } + } else { + lhsScaleOut = + rewriter + .create( + op.getLoc(), + RankedTensorType::get(lhsScaleTy.getShape(), fp32Ty), lhsScale) + .getResult(); + } + + if (rhsScale && rhsScaleTy.getElementType().isIntOrIndex()) { + if (rhsScaleTy.getRank() != 2) { + return op.emitError("rhsScale must be 2D for transpose"); + } + + SmallVector transposedShape = {rhsScaleTy.getShape()[1], + rhsScaleTy.getShape()[0]}; + RankedTensorType transposedRhsScaleTy = + RankedTensorType::get(transposedShape, rhsScaleTy.getElementType()); + + Value transposedRhsScale = rewriter.create( + op.getLoc(), transposedRhsScaleTy, rhsScale, + DenseI32ArrayAttr::get(rewriter.getContext(), ArrayRef{1, 0})); + RankedTensorType rhsScaleI16Ty = + RankedTensorType::get(transposedShape, i16Ty); + Value rhsScaleI16 = rewriter.create( + op.getLoc(), rhsScaleI16Ty, transposedRhsScale); + Value rhsShift127Empty = rewriter.create( + op.getLoc(), rhsScaleI16Ty.getShape(), i16Ty); + Value rhsShift127 = + rewriter + .create(op.getLoc(), ValueRange{c127}, + ValueRange{rhsShift127Empty}) + .getResult(0); + + Value rhsScaleI16Add127 = + rewriter.create(op.getLoc(), rhsScaleI16, rhsShift127); + Value rhsShift7Empty = rewriter.create( + op.getLoc(), rhsScaleI16Ty.getShape(), i16Ty); + Value rhsShift7 = rewriter + .create(op.getLoc(), ValueRange{c7}, + ValueRange{rhsShift7Empty}) + .getResult(0); + Value rhsScaleI16Shifted = rewriter.create( + op.getLoc(), rhsScaleI16Add127, rhsShift7); + + RankedTensorType rhsScaleBF16Ty = + RankedTensorType::get(transposedShape, bf16Ty); + Value rhsScaleBF16 = rewriter.create( + op.getLoc(), rhsScaleBF16Ty, rhsScaleI16Shifted); + + if (rhsTy.getElementType() == fp16Ty) { + RankedTensorType rhsScaleFp32Ty = + RankedTensorType::get(transposedShape, fp32Ty); + Value rhsScaleFp32 = rewriter.create( + op.getLoc(), rhsScaleFp32Ty, rhsScaleBF16); + RankedTensorType rhsScaleFp16Ty = + RankedTensorType::get(transposedShape, fp16Ty); + rhsScaleOut = rewriter.create( + op.getLoc(), rhsScaleFp16Ty, rhsScaleFp32); + } else { + rhsScaleOut = rhsScaleBF16; + } + int64_t rhsD0 = rhsScaleTy.getShape()[1]; + int64_t rhsD1 = rhsScaleTy.getShape()[0]; + SmallVector rhsExpandedShape1 = {rhsD0, rhsD1, 1}; + RankedTensorType rhsExpandedTy1 = + RankedTensorType::get(rhsExpandedShape1, rhsTy.getElementType()); + Value rhsExpanded1 = rewriter + .create( + op.getLoc(), rhsExpandedTy1, rhsScaleOut, + rewriter.getI32IntegerAttr(2)) + .getResult(); + + int64_t rhsDim1 = rhsTy.getShape()[0]; + if (rhsDim1 % rhsD0 != 0) { + return op.emitError( + "rhs dim0 must be an integer multiple of rhsScale dim0"); + } + int64_t rhsD2 = rhsDim1 / rhsD0; + SmallVector rhsBroadcastShape = {rhsD0, rhsD1, rhsD2}; + RankedTensorType rhsBroadcastTy = + RankedTensorType::get(rhsBroadcastShape, rhsTy.getElementType()); + Value rhsBroadcasted = rewriter + .create( + op.getLoc(), rhsBroadcastTy, rhsExpanded1) + .getResult(); + + SmallVector transposeOrder = {0, 2, 1}; + Value transposedBroadcasted = rewriter.create( + op.getLoc(), + RankedTensorType::get({rhsD0, rhsD2, rhsD1}, rhsTy.getElementType()), + rhsBroadcasted, + DenseI32ArrayAttr::get(rewriter.getContext(), transposeOrder)); + SmallVector rhsReassociation; + rhsReassociation.push_back({0, 1}); + rhsReassociation.push_back({2}); + + Value scaledRhs = rewriter + .create( + op.getLoc(), + RankedTensorType::get({rhsD0 * rhsD2, rhsD1}, + rhsTy.getElementType()), + transposedBroadcasted, rhsReassociation) + .getResult(); + + rhs = + rewriter.create(op.getLoc(), rhs, scaledRhs).getResult(); + + if (!fastMath) { + Value rhsScaleNaNMask = + createNaNMask(transposedRhsScale, transposedRhsScaleTy); + Value rhsExpandedMask = + rewriter + .create( + op.getLoc(), + RankedTensorType::get(rhsExpandedShape1, + rewriter.getI1Type()), + rhsScaleNaNMask, rewriter.getI32IntegerAttr(2)) + .getResult(); + Value rhsBroadcastMask = + rewriter + .create( + op.getLoc(), + RankedTensorType::get(rhsBroadcastShape, + rewriter.getI1Type()), + rhsExpandedMask) + .getResult(); + Value transposedBroadcastMask = + rewriter + .create( + op.getLoc(), + RankedTensorType::get({rhsD0, rhsD2, rhsD1}, + rewriter.getI1Type()), + rhsBroadcastMask, + DenseI32ArrayAttr::get(rewriter.getContext(), transposeOrder)) + .getResult(); + Value collapsedRhsMask = + rewriter + .create( + op.getLoc(), + RankedTensorType::get({rhsD0 * rhsD2, rhsD1}, + rewriter.getI1Type()), + transposedBroadcastMask, rhsReassociation) + .getResult(); + rhs = applyNaNMask(rhs, collapsedRhsMask); + } + } + + int64_t D0 = lhsScaleTy.getShape()[0]; + int64_t D1 = lhsScaleTy.getShape()[1]; + SmallVector expandedShape1 = {D0, D1, 1}; + RankedTensorType expandedTy1 = + RankedTensorType::get(expandedShape1, lhsTy.getElementType()); + Value expanded1 = + rewriter + .create(op.getLoc(), expandedTy1, lhsScaleOut, + rewriter.getI32IntegerAttr(2)) + .getResult(); + + int64_t lhsDim1 = lhsTy.getShape()[1]; + if (lhsDim1 % D1 != 0) { + return op.emitError( + "lhs dim1 must be an integer multiple of lhsScale dim1"); + } + int64_t D2 = lhsDim1 / D1; + SmallVector broadcastShape = {D0, D1, D2}; + RankedTensorType broadcastTy = + RankedTensorType::get(broadcastShape, lhsTy.getElementType()); + Value broadcasted = + rewriter.create(op.getLoc(), broadcastTy, expanded1) + .getResult(); + + SmallVector reassociation; + reassociation.push_back({0}); + reassociation.push_back({1, 2}); + + Value scaledLhs = + rewriter + .create( + op.getLoc(), + RankedTensorType::get({D0, D1 * D2}, lhsTy.getElementType()), + broadcasted, reassociation) + .getResult(); + + Value scaledLhsFinal = + rewriter.create(op.getLoc(), lhs, scaledLhs).getResult(); + + if (!fastMath) { + Value lhsScaleNaNMask = createNaNMask(lhsScale, lhsScaleTy); + Value lhsExpandedMask = + rewriter + .create( + op.getLoc(), + RankedTensorType::get(expandedShape1, rewriter.getI1Type()), + lhsScaleNaNMask, rewriter.getI32IntegerAttr(2)) + .getResult(); + Value lhsBroadcastMask = + rewriter + .create( + op.getLoc(), + RankedTensorType::get(broadcastShape, rewriter.getI1Type()), + lhsExpandedMask) + .getResult(); + Value collapsedLhsMask = + rewriter + .create( + op.getLoc(), + RankedTensorType::get({D0, D1 * D2}, rewriter.getI1Type()), + lhsBroadcastMask, reassociation) + .getResult(); + scaledLhsFinal = applyNaNMask(scaledLhsFinal, collapsedLhsMask); + } + + Operation *matmulOp; + if (dstType.getRank() == 2) { + matmulOp = rewriter.create( + op.getLoc(), ValueRange{scaledLhsFinal, rhs}, ValueRange{c}); + } else if (dstType.getRank() == 3) { + matmulOp = rewriter.create( + op.getLoc(), ValueRange{scaledLhsFinal, rhs}, ValueRange{c}); + } else { + return op.emitError("DotScaledOp only support 2D or 3D tensor"); + } + + rewriter.replaceOp(op, matmulOp->getResults()); + return success(); +} + +LogicalResult +PtrToIntConverter::matchAndRewrite(triton::PtrToIntOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + Value ptr = adaptor.getSrc(); + + if (!mlir::isa(ptr.getType())) { + return rewriter.notifyMatchFailure(op, "input is not a memref type"); + } + + auto resultType = op.getType(); + + // memref.extract_aligned_pointer_as_index is used to obtain the integer + // representation of the base address. + auto ptrToIndexOp = + rewriter.create(loc, ptr); + + Value intResult = + rewriter.create(loc, resultType, ptrToIndexOp); + + rewriter.replaceOp(op, intResult); + return success(); +} + +LogicalResult +IndexPutConverter::matchAndRewrite(triton::dicp::IndexPutOp op, + OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + + auto moduleOp = op->getParentOfType(); + rewriter.setInsertionPoint(moduleOp.getBody(), + std::prev(moduleOp.getBody()->end())); + + auto funcName = generateUniqueFuncName(moduleOp, funcNameBase); + + auto ptr = adaptor.getPtr(); + auto index = op.getIndex(); + auto value = op.getValue(); + auto dim = op.getDim(); + auto indexBoundary = op.getIndexBoundary(); + auto endOffset = op.getEndOffset(); + auto startOffset = op.getStartOffset(); + auto dstStride = adaptor.getDstStride(); + + // convert !tt.ptr to memref + auto ptrTy = dyn_cast(ptr.getType()); + if (!ptrTy) { + return rewriter.notifyMatchFailure(op, "expected MemRefType for ptr"); + } + SmallVector inputTypes({ptrTy, index.getType(), value.getType(), + dim.getType(), indexBoundary.getType()}); + inputTypes.append(endOffset.getTypes().begin(), endOffset.getTypes().end()); + inputTypes.append(startOffset.getTypes().begin(), + startOffset.getTypes().end()); + inputTypes.append(dstStride.getTypes().begin(), dstStride.getTypes().end()); + auto libFnType = rewriter.getFunctionType(inputTypes, {}); + auto funcOp = rewriter.create(loc, funcName.str(), libFnType); + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + + rewriter.setInsertionPoint(op); + SmallVector inputVals({ptr, index, value, dim, indexBoundary}); + inputVals.append(endOffset.begin(), endOffset.end()); + inputVals.append(startOffset.begin(), startOffset.end()); + inputVals.append(dstStride.begin(), dstStride.end()); + rewriter.create(loc, funcOp.getSymNameAttr(), TypeRange({}), + inputVals); + rewriter.eraseOp(op); + return success(); +} + +LogicalResult GatherOutToUbConverter::matchAndRewrite( + triton::dicp::GatherOutToUbOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + + auto moduleOp = op->getParentOfType(); + rewriter.setInsertionPoint(moduleOp.getBody(), + std::prev(moduleOp.getBody()->end())); + + auto funcName = generateUniqueFuncName(moduleOp, funcNameBase); + + auto src = adaptor.getSrc(); + auto index = op.getIndex(); + auto indexBoundary = op.getIndexBoundary(); + auto dim = op.getDim(); + auto srcStride = op.getSrcStride(); + auto endOffset = op.getEndOffset(); + auto startOffset = op.getStartOffset(); + auto other = op.getOther(); + + auto res = op.getResult(); + auto resTy = res.getType(); + + // convert !tt.ptr to memref + auto srcTy = dyn_cast(src.getType()); + if (!srcTy) { + return rewriter.notifyMatchFailure(op, "expected MemRefType for src"); + } + + SmallVector inputTypes( + {srcTy, index.getType(), indexBoundary.getType(), dim.getType()}); + inputTypes.append(srcStride.getTypes().begin(), srcStride.getTypes().end()); + inputTypes.append(endOffset.getTypes().begin(), endOffset.getTypes().end()); + inputTypes.append(startOffset.getTypes().begin(), + startOffset.getTypes().end()); + if (other) + inputTypes.push_back(other.getType()); + + auto libFnType = rewriter.getFunctionType(inputTypes, {resTy}); + auto funcOp = rewriter.create(loc, funcName.str(), libFnType); + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + + rewriter.setInsertionPoint(op); + SmallVector inputVals({src, index, indexBoundary, dim}); + inputVals.append(srcStride.begin(), srcStride.end()); + inputVals.append(endOffset.begin(), endOffset.end()); + inputVals.append(startOffset.begin(), startOffset.end()); + if (other) + inputVals.push_back(other); + auto callOp = rewriter.create(loc, funcOp.getSymNameAttr(), + TypeRange({resTy}), inputVals); + rewriter.replaceOp(op, callOp); + return success(); +} + +LogicalResult ScatterUbToOutConverter::matchAndRewrite( + triton::dicp::ScatterUbToOutOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + + auto moduleOp = op->getParentOfType(); + rewriter.setInsertionPoint(moduleOp.getBody(), + std::prev(moduleOp.getBody()->end())); + + auto funcName = generateUniqueFuncName(moduleOp, funcNameBase); + + auto ptr = adaptor.getPtr(); + auto value = op.getValue(); + auto index = op.getIndex(); + auto indexBoundary = op.getIndexBoundary(); + auto dim = op.getDim(); + auto dstStride = op.getDstStride(); + auto endOffset = op.getEndOffset(); + auto startOffset = op.getStartOffset(); + + // convert !tt.ptr to memref + auto ptrTy = dyn_cast(ptr.getType()); + if (!ptrTy) { + return rewriter.notifyMatchFailure(op, "expected MemRefType for ptr"); + } + + SmallVector inputTypes({ptrTy, value.getType(), index.getType(), + indexBoundary.getType(), dim.getType()}); + inputTypes.append(dstStride.getTypes().begin(), dstStride.getTypes().end()); + inputTypes.append(endOffset.getTypes().begin(), endOffset.getTypes().end()); + inputTypes.append(startOffset.getTypes().begin(), + startOffset.getTypes().end()); + + auto libFnType = rewriter.getFunctionType(inputTypes, {}); + auto funcOp = rewriter.create(loc, funcName.str(), libFnType); + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + + rewriter.setInsertionPoint(op); + SmallVector inputVals({ptr, value, index, indexBoundary, dim}); + inputVals.append(dstStride.begin(), dstStride.end()); + inputVals.append(endOffset.begin(), endOffset.end()); + inputVals.append(startOffset.begin(), startOffset.end()); + rewriter.create(loc, funcOp.getSymNameAttr(), TypeRange({}), + inputVals); + rewriter.eraseOp(op); + return success(); +} + +LogicalResult IndirectLoadConverter::matchAndRewrite( + triton::dicp::IndirectLoadOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + + auto moduleOp = op->getParentOfType(); + rewriter.setInsertionPoint(moduleOp.getBody(), + std::prev(moduleOp.getBody()->end())); + + auto funcName = generateUniqueFuncName(moduleOp, funcNameBase); + + auto src = adaptor.getSrc(); + auto offsets = op.getOffsets(); + auto mask = op.getMask(); + auto other = op.getOther(); + auto res = op.getResult(); + auto resTy = res.getType(); + + // convert !tt.ptr to memref + auto srcTy = dyn_cast(src.getType()); + if (!srcTy) { + return rewriter.notifyMatchFailure(op, "expected MemRefType for src"); + } + SmallVector inputTypes({srcTy, offsets.getType()}); + if (mask) + inputTypes.push_back(mask.getType()); + if (other) + inputTypes.push_back(other.getType()); + auto libFnType = rewriter.getFunctionType(inputTypes, {resTy}); + auto funcOp = rewriter.create(loc, funcName.str(), libFnType); + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + + rewriter.setInsertionPoint(op); + SmallVector inputVals({src, offsets}); + if (mask) + inputVals.push_back(mask); + if (other) + inputVals.push_back(other); + auto callOp = rewriter.create(loc, funcOp.getSymNameAttr(), + TypeRange({resTy}), inputVals); + rewriter.replaceOp(op, callOp); + return success(); +} + +LogicalResult IndirectStoreConverter::matchAndRewrite( + triton::dicp::IndirectStoreOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + + auto moduleOp = op->getParentOfType(); + rewriter.setInsertionPoint(moduleOp.getBody(), + std::prev(moduleOp.getBody()->end())); + + auto funcName = generateUniqueFuncName(moduleOp, funcNameBase); + + auto src = adaptor.getSrc(); + auto offsets = op.getOffsets(); + auto value = op.getValue(); + auto mask = op.getMask(); + + // convert !tt.ptr to memref + auto srcTy = dyn_cast(src.getType()); + if (!srcTy) { + return rewriter.notifyMatchFailure(op, "expected MemRefType for src"); + } + SmallVector inputTypes({srcTy, offsets.getType(), value.getType()}); + if (mask) + inputTypes.push_back(mask.getType()); + + auto libFnType = rewriter.getFunctionType(inputTypes, {}); + auto funcOp = rewriter.create(loc, funcName.str(), libFnType); + SymbolTable::setSymbolVisibility(funcOp, SymbolTable::Visibility::Private); + + rewriter.setInsertionPoint(op); + SmallVector inputVals({src, offsets, value}); + if (mask) + inputVals.push_back(mask); + rewriter.create(loc, funcOp.getSymNameAttr(), TypeRange({}), + inputVals); + rewriter.eraseOp(op); + return success(); +} + +IndexSelectSimdConverter::IndexSelectSimdConverter(MLIRContext *context) + : OpConversionPattern(context) {} + +LogicalResult IndexSelectSimdConverter::matchAndRewrite( + triton::dicp::IndexSelectSimdOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto loc = op.getLoc(); + + // Get converted operands + Value src = adaptor.getSrc(); + Value indexTensor = adaptor.getIndex(); + auto srcShapeVals = adaptor.getSrcShape(); + auto srcOffsetVals = adaptor.getSrcOffset(); + auto readShapeAttr = op.getReadShape(); + int32_t dim = op.getDim(); + + // Get result type + auto resultTensorType = cast(op.getResult().getType()); + auto elemType = resultTensorType.getElementType(); + auto resultShape = resultTensorType.getShape(); + + // Convert src (tt.ptr -> memref) to the correct memref shape + // src is now memref after type conversion, need to reinterpret to full + // shape + auto srcMemRefType = cast(src.getType()); + + // DenseI32ArrayAttr can be implicitly converted to ArrayRef + ArrayRef readShape = readShapeAttr; + + // Helper lambda to convert Value to Index type if needed + auto toIndexValue = [&](Value val) -> Value { + if (!val.getType().isIndex()) { + return rewriter.create(loc, rewriter.getIndexType(), + val); + } + return val; + }; + + // Build multi-dimensional memref type and determine static sizes + // Merge two passes: build fullSrcShape and staticSizes in one loop + SmallVector fullSrcShape; + SmallVector staticSizes; + SmallVector sizes; + + for (size_t i = 0; i < srcShapeVals.size(); ++i) { + bool isDynamicDim = (i == static_cast(dim) && readShape[i] == -1); + int64_t staticSize; + + if (isDynamicDim) { + // Dynamic dimension: readShape[i] == -1 indicates dynamic + staticSize = ShapedType::kDynamic; + fullSrcShape.push_back(ShapedType::kDynamic); + sizes.push_back(toIndexValue(srcShapeVals[i])); + } else if (auto constOp = + srcShapeVals[i].getDefiningOp()) { + // Static dimension: use constant value + staticSize = constOp.value(); + fullSrcShape.push_back(staticSize); + } else { + // Runtime value: must use dynamic + staticSize = ShapedType::kDynamic; + fullSrcShape.push_back(ShapedType::kDynamic); + sizes.push_back(toIndexValue(srcShapeVals[i])); + } + staticSizes.push_back(staticSize); + } + auto fullSrcMemRefType = MemRefType::get(fullSrcShape, elemType); + + // Build static offsets, sizes, and strides for ReinterpretCastOp + SmallVector offsets, strides; + SmallVector staticOffsets, staticStrides; + staticOffsets.push_back(0); // offsets are 0 + + // Calculate static strides: stride[i] = product of all dimensions after i + for (size_t i = 0; i < srcShapeVals.size(); ++i) { + int64_t staticStride = 1; + bool isDynamic = false; + + // Check if stride needs to be dynamic (any dimension after i is dynamic) + for (size_t j = i + 1; j < srcShapeVals.size(); ++j) { + if (staticSizes[j] == ShapedType::kDynamic) { + isDynamic = true; + break; + } + staticStride *= staticSizes[j]; + } + + if (isDynamic) { + staticStride = ShapedType::kDynamic; + // Compute stride dynamically: stride[i] = product of sizes after i + Value strideVal = rewriter.create(loc, 1); + for (size_t j = i + 1; j < srcShapeVals.size(); ++j) { + if (staticSizes[j] != ShapedType::kDynamic) { + strideVal = rewriter.create( + loc, strideVal, + rewriter.create(loc, staticSizes[j])); + } else { + // Dynamic dimension: use runtime value + strideVal = rewriter.create( + loc, strideVal, toIndexValue(srcShapeVals[j])); + } + } + strides.push_back(strideVal); + } + staticStrides.push_back(staticStride); + } + + auto srcMemRef = rewriter.create( + loc, fullSrcMemRefType, src, offsets, sizes, strides, staticOffsets, + staticSizes, staticStrides); + + // Allocate output buffer + auto resultMemRefType = MemRefType::get(resultShape, elemType); + auto outputBuffer = rewriter.create(loc, resultMemRefType); + + // Get indices tensor type for extracting + auto indicesTensorType = cast(indexTensor.getType()); + int64_t numIndices = indicesTensorType.getShape()[0]; + + // Create for loop + auto zeroIdx = rewriter.create(loc, 0); + auto numIndicesVal = rewriter.create(loc, numIndices); + auto stepOne = rewriter.create(loc, 1); + auto forOp = + rewriter.create(loc, zeroIdx, numIndicesVal, stepOne); + + // Mark as parallel loop + forOp->setAttr("hivm.parallel_loop", rewriter.getUnitAttr()); + + // Build loop body + Block *loopBody = forOp.getBody(); + auto savedInsertionPoint = rewriter.saveInsertionPoint(); + rewriter.setInsertionPointToStart(loopBody); + + // Remove the terminator temporarily + Operation *terminator = &loopBody->back(); + rewriter.setInsertionPoint(terminator); + + Value iv = forOp.getInductionVar(); + + // Extract index from indices tensor + Value selectedIdx = + rewriter.create(loc, indexTensor, ValueRange{iv}); + Value selectedIdxAsIndex = rewriter.create( + loc, rewriter.getIndexType(), selectedIdx); + + // Build source subview offsets/sizes/strides + SmallVector srcSubviewOffsets, srcSubviewSizes, + srcSubviewStrides; + + for (size_t i = 0; i < srcOffsetVals.size(); ++i) { + if (i == static_cast(dim)) { + // Use the selected index for this dimension + srcSubviewOffsets.push_back(selectedIdxAsIndex); + srcSubviewSizes.push_back(rewriter.getIndexAttr(1)); + } else { + // Use provided offset and read size for other dimensions + Value offsetVal = srcOffsetVals[i]; + if (!offsetVal.getType().isIndex()) { + offsetVal = rewriter.create( + loc, rewriter.getIndexType(), offsetVal); + } + srcSubviewOffsets.push_back(offsetVal); + srcSubviewSizes.push_back(rewriter.getIndexAttr(readShape[i])); + } + srcSubviewStrides.push_back(rewriter.getIndexAttr(1)); + } + + auto srcSubview = rewriter.create( + loc, srcMemRef, srcSubviewOffsets, srcSubviewSizes, srcSubviewStrides); + + // Build destination subview + SmallVector dstSubviewOffsets, dstSubviewSizes, + dstSubviewStrides; + for (size_t i = 0; i < resultShape.size(); ++i) { + if (i == static_cast(dim)) { + dstSubviewOffsets.push_back(iv); + dstSubviewSizes.push_back(rewriter.getIndexAttr(1)); + } else { + dstSubviewOffsets.push_back(rewriter.getIndexAttr(0)); + dstSubviewSizes.push_back(rewriter.getIndexAttr(readShape[i])); + } + dstSubviewStrides.push_back(rewriter.getIndexAttr(1)); + } + + auto dstSubview = rewriter.create( + loc, outputBuffer, dstSubviewOffsets, dstSubviewSizes, dstSubviewStrides); + + // Check if index_select is on the trailing axis (last dimension) + if (static_cast(dim) == fullSrcShape.size() - 1) { + // For index_select on the trailing axis, mark as discrete memory access + // This degrades to scalar read/write handling to avoid alignment issues + auto copyOp = rewriter.create(loc, srcSubview, dstSubview); + copyOp->setAttr(ConverterUtils::discreteAttrName, rewriter.getUnitAttr()); + } else { + // For index_select on non-trailing axes, add stride alignment annotation + // This tells the backend to handle address alignment for DMA operations + auto dstMarkOp = rewriter.create(loc, dstSubview); + dstMarkOp->setAttr( + "hfusion.stride_align_dims", + rewriter.getDenseI32ArrayAttr({static_cast(dim)})); + dstMarkOp->setAttr("hfusion.stride_align_value_in_byte", + rewriter.getDenseI32ArrayAttr({32})); + + // Copy from source to destination + rewriter.create(loc, srcSubview, dstSubview); + } + + // Restore insertion point + rewriter.restoreInsertionPoint(savedInsertionPoint); + + // Convert memref to tensor + auto resultTensor = rewriter.create( + loc, resultTensorType, outputBuffer, true, true); + + // Mark as index_select_simd + resultTensor->setAttr("index_select_simd", rewriter.getUnitAttr()); + + // Replace the original op + rewriter.replaceOp(op, resultTensor); + + return success(); +} + +} // namespace TTOpConverters diff --git a/compiler/lib/TritonToLinalg/TritonToLinalgPass.cpp b/compiler/lib/TritonToLinalg/TritonToLinalgPass.cpp new file mode 100644 index 00000000..3df8d9e0 --- /dev/null +++ b/compiler/lib/TritonToLinalg/TritonToLinalgPass.cpp @@ -0,0 +1,1160 @@ + + +#include "dicp/TritonToLinalg/TritonToLinalgPass.h" +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" +#include "dicp/TritonToLinalg/ArgMinMaxConverter.h" +#include "dicp/TritonToLinalg/DescriptorConverter.h" +#include "dicp/TritonToLinalg/FunctionConverter.h" +#include "dicp/TritonToLinalg/HoistBroadcast.h" +#include "dicp/TritonToLinalg/ImplicitPermute.h" +#include "dicp/TritonToLinalg/LoadStoreConverter.h" +#include "dicp/TritonToLinalg/MarkTensorKindPass.h" +#include "dicp/TritonToLinalg/TritonOpConverter.h" +#include "dicp/TritonToLinalg/UseAnalysis.h" +#include "dicp/TritonToStructured/CannonicalizerConverter.h" +#include "dicp/Utils/InterleaveOptimization.h" +#include "dicp/Utils/Utils.h" + +#include "bishengir/Dialect/HFusion/IR/HFusion.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Operation.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" + +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" + +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Transforms/Transforms.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Visitors.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/Passes.h" + +#include "llvm/ADT/BitVector.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/Twine.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/LogicalResult.h" + +#include +#include +#include + +#define DEBUG_TYPE "triton-to-linalg" + +using namespace mlir; +using namespace triton; + +int nd2nzFlag = 0; +bool compileOn91095Flag = false; +bool existDotFlag = false; + +// Convert CustomOp after operand type converted, +// for example tt.ptr converted to memref. +class CustomOpConverter : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(hivm::CustomOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + llvm::SmallDenseMap known; + BlockDataParser::rewriteCustomOp(op, adaptor, rewriter, known); + return success(); + } +}; + +static bool isSIMTOp(Operation *op) { + if (auto custom_op = dyn_cast(op)) { + return custom_op.getCoreType() == hivm::TCoreType::VECTOR && + custom_op.getVFMode() == hivm::VFMode::SIMT; + } + return isa(op); +} + +TritonTypeConverter::TritonTypeConverter() { + addConversion([](Type type) { return type; }); + + addConversion([](triton::PointerType ptrType) { + Type elem = ptrType.getPointeeType(); + // Handling special case: ptr -> memref + if (auto it = dyn_cast(elem); it && it.getWidth() == 1) { + elem = IntegerType::get(ptrType.getContext(), 8); + LLVM_DEBUG({ + llvm::dbgs() << "[TritonTypeConverter] Normalize i1 pointer to i8 " + "memref. elemType=" + << elem << "\n"; + }); + } + return MemRefType::get({ShapedType::kDynamic}, elem); + }); + + addConversion([](TensorType tensorType) -> Type { + auto elemType = tensorType.getElementType(); + if (auto ptrType = dyn_cast(elemType)) { + elemType = ptrType.getPointeeType(); + } + // Handling special case: tensor -> memref + if (auto it = dyn_cast(elemType); it && it.getWidth() == 1) { + elemType = IntegerType::get(tensorType.getContext(), 8); + LLVM_DEBUG({ + llvm::dbgs() << "[TritonTypeConverter] Normalize i1 tensor to i8 " + "memref. elemType=" + << elemType << "\n"; + }); + } + return MemRefType::get(tensorType.getShape(), elemType); + }); +} + +void TritonToLinalgPass::addProgramInfo(triton::FuncOp func, + bool globalKernel) { + OpBuilder b(func); + + auto origFuncType = func.getFunctionType(); + auto origInputTypes = origFuncType.getInputs(); + SmallVector newInputTypes(origInputTypes); + newInputTypes.append(TRITON_PROGRAM_INFO_ARG_COUNT, b.getI32Type()); + + auto newFuncType = + b.getFunctionType(newInputTypes, origFuncType.getResults()); + + func.setFunctionType(newFuncType); + + // If argument attributes exist, extend attribute list. + if (func.getAllArgAttrs()) { + SmallVector newArgAttrs; + func.getAllArgAttrs(newArgAttrs); + newArgAttrs.append(TRITON_PROGRAM_INFO_ARG_COUNT, DictionaryAttr()); + func.setAllArgAttrs(newArgAttrs); + } + + // Append the arguments to the entry block. + for (unsigned i = 0; i < TRITON_PROGRAM_INFO_ARG_COUNT; i++) { + func.getBody().front().addArgument(b.getI32Type(), func.getLoc()); + } + + if (globalKernel) { + func->setAttr(globalKernelAttr, b.getStringAttr("")); + } else { + func->setAttr(globalKernelAttr, b.getStringAttr("local")); + } +} + +LogicalResult +TritonToLinalgPass::convertMultipleBlockControlFlow(Operation *funcOp, + OpBuilder &builder) { + if (!isa(funcOp)) { + funcOp->emitError( + "convertMultipleBlockControlFlow can only process func::FuncOp!"); + return failure(); + } + + SmallVector candidate; + SmallVector eraseBlocks; + for (Block &block : dyn_cast(funcOp).getBody()) { + auto curTerminator = block.getTerminator(); + if (isa(curTerminator)) { + candidate.push_back(curTerminator); + } else if (isa(curTerminator)) { + if (candidate.empty()) { + curTerminator->emitError( + "funcOp has more than one Block but got an early 'tt.return' Op."); + return failure(); + } + } else if (!isa(curTerminator)) { + funcOp->emitError( + "funcOp has more than one Block but found unsupported Terminator: ") + << *curTerminator; + return failure(); + } + + if (!block.isEntryBlock()) + eraseBlocks.push_back(&block); + } + + LLVM_DEBUG({ + llvm::dbgs() << "Found " << candidate.size() + << " candidate cond_branch operations to convert.\n"; + }); + + if (candidate.empty()) { + funcOp->emitError("funcOp has more than one Block but no candidate " + "Terminator was found!"); + return failure(); + } + + llvm::BitVector visitFlag(candidate.size(), false); + + // Recursive function to convert all cf::CondBranchOp to scf::IfOp + std::function convertToSCF = + [&](Operation *op, Operation *insertPosOp) -> void { + auto condBranchOp = dyn_cast_if_present(op); + auto iter = llvm::find(candidate, condBranchOp); + if (!(condBranchOp && iter != candidate.end())) { + op->emitError( + "convertToSCF must process with condBranchOp in candidates!"); + return; + } + visitFlag.set(iter - candidate.begin()); + + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPointAfter(insertPosOp); + + // Well, here force to destory original control flow + builder.create( + condBranchOp->getLoc(), condBranchOp.getCondition(), + /*thenBuilder=*/ + [&](OpBuilder &builder, Location loc) { + SmallVector movedOps = llvm::map_to_vector( + condBranchOp.getTrueDest()->without_terminator(), + [](Operation &op) { return &op; }); + for (auto *innerOp : movedOps) { + innerOp->moveBefore(builder.getInsertionBlock(), + builder.getInsertionPoint()); + } + + auto blockTerm = condBranchOp.getTrueDest()->getTerminator(); + if (auto nextCond = dyn_cast(blockTerm)) { + if (movedOps.empty()) { + blockTerm->emitError("movedOps can not be empty before entering " + "convertToSCF (then)!"); + return; + } + convertToSCF(nextCond, movedOps.back()); + } else if (!isa(blockTerm)) { + blockTerm->emitError( + "Unsupported terminator in then branch after structuring"); + } + + builder.create(loc); + }, + /*elseBuilder=*/ + [&](OpBuilder &builder, Location loc) { + SmallVector movedOps = llvm::map_to_vector( + condBranchOp.getFalseDest()->without_terminator(), + [](Operation &op) { return &op; }); + for (auto *innerOp : movedOps) { + innerOp->moveBefore(builder.getInsertionBlock(), + builder.getInsertionPoint()); + } + + auto blockTerm = condBranchOp.getFalseDest()->getTerminator(); + if (auto nextCond = dyn_cast(blockTerm)) { + if (movedOps.empty()) { + blockTerm->emitError("movedOps can not be empty before entering " + "convertToSCF (else)!"); + return; + } + convertToSCF(nextCond, movedOps.back()); + } else if (!isa(blockTerm)) { + blockTerm->emitError( + "Unsupported terminator in else branch after structuring"); + } + builder.create(loc); + }); + }; + + Block::iterator insertOp(candidate.front()); + if (insertOp == candidate.front()->getBlock()->begin()) { + // if the first operation is a cond_branch, we need to insert before it + convertToSCF(candidate.front(), candidate.front()); + } else { + --insertOp; + convertToSCF(candidate.front(), &(*insertOp)); + } + + if (!visitFlag.all()) { + funcOp->emitError("Not all cf.cond_br converted!"); + return failure(); + } + + OpBuilder::InsertionGuard guard(builder); + builder.setInsertionPoint(candidate.front()); + builder.create(candidate.front()->getLoc()); + + for (Operation *eachTerm : candidate) + eachTerm->erase(); + for (Block *block : llvm::reverse(eraseBlocks)) + block->erase(); + + return success(); +} + +void TritonToLinalgPass::convertTTFunc(triton::FuncOp func, const bool existDot, + const bool existSIMTOp) { + OpBuilder builder(func); + + auto name = func.getName(); + auto type = func.getFunctionType(); + + SmallVector argAttrs, resAttrs; + func.getAllArgAttrs(argAttrs); + func.getAllResultAttrs(resAttrs); + + // Special handling for bit-casted tt.ptr arguments + SmallVector inputTypes{type.getInputs()}; + SmallVector retTypes{type.getResults()}; + if (func.getSymVisibility() == "public" && !func.isDeclaration()) { + for (size_t i = 0; i < func.getNumArguments(); ++i) { + auto arg = func.getArgument(i); + // Special method for i1 arg + if (!isa(arg.getType()) || + dyn_cast(arg.getType()).getElementTypeBitWidth() != + 1) { + continue; + } + + SmallVector argVaildUser{arg.getUsers()}; + llvm::erase_if(argVaildUser, [](Operation *op) -> bool { + return isOpTriviallyDead(op); + }); + + if (!argVaildUser.empty()) { + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << arg << " has users:\n"; + int cnt = 0; + for (auto it : argVaildUser) { + os << "users[" << cnt++ << "] = " << *it; + } + }); + if (llvm::all_of(argVaildUser, [](Operation *userOp) { + return isa(userOp); + })) { + auto castOp = cast(*argVaildUser.begin()); + if (castOp.getInputs().size() == 1 && + castOp.getOutputs().size() == 1) { + arg.setType(castOp.getOutputs()[0].getType()); + inputTypes[i] = arg.getType(); + } + } else { + func->emitError(Twine("Unsupported use of func arg at index ") + + Twine(i)); + } + } else { + // Process unused bool ptr type specially, which guarantees bool pointer + // argument's type is realistic and don't mislead backend compiler. + // realistic memory layout of bool pointer is 8 bit width + auto memType = dyn_cast(arg.getType()) + .cloneWith(std::nullopt, builder.getI8Type()); + arg.setType(memType); + inputTypes[i] = arg.getType(); + } + } + } + auto castType = FunctionType::get(func.getContext(), inputTypes, retTypes); + + auto funcFunc = builder.create(func.getLoc(), name, castType); + funcFunc.setAllArgAttrs(argAttrs); + funcFunc.setAllResultAttrs(resAttrs); + auto kernelAttr = func->getAttr(globalKernelAttr); + if (kernelAttr) { + funcFunc->setAttr(globalKernelAttr, kernelAttr); + } + std::string kernelMixMode = "aiv"; + if (existDot) { + // mix also works for pure cube kernel by using the same MAGIC_ELF keyword + kernelMixMode = "mix"; + } + // Set mix_mode in the func attrs so that the backend could know + // the mix_mode by parse the func attrs. + // The backend needs to know the mix_mode because the host wrapper + // needs to set the devbin.magic. Check npu_utils.cpp. + funcFunc->setAttr(kernelMixModeName, builder.getStringAttr(kernelMixMode)); + + std::string parallelMode = "simd"; + if (existSIMTOp) { + parallelMode = "mix_simd_simt"; + } + funcFunc->setAttr(kernelParallelModeName, + builder.getStringAttr(parallelMode)); + + auto autoBlockifyAttr = func->getAttr("auto_blockify_size"); + if (autoBlockifyAttr) + funcFunc->setAttr("auto_blockify_size", autoBlockifyAttr); + + auto &funcFuncBody = funcFunc.getBody(); + auto &funcBody = func.getBody(); + + IRMapping map; + funcBody.cloneInto(&funcFuncBody, map); + + if (!funcFuncBody.hasOneBlock()) { + if (failed(convertMultipleBlockControlFlow(funcFunc, builder))) { + llvm_unreachable("Encounter unsupported control flow"); + } + } + + for (Block &block : funcFuncBody.getBlocks()) { + auto term = block.getTerminator(); + builder.setInsertionPoint(term); + builder.create(func.getLoc(), term->getOperands()); + term->erase(); + } + func.erase(); +} + +void TritonToLinalgPass::addDynamicLegal( + ConversionTarget &target, TritonTypeConverter &tritonTypeConverter) { + target.addLegalDialect(); + + // add legal dialect on condition + target.addLegalOp(); + + // decide which ops need conversion based on uses + target.addDynamicallyLegalOp( + [](mlir::Operation *op) { + if (op->use_empty()) { + return false; + } else { + return true; + } + }); + + target.addDynamicallyLegalOp([&](triton::FuncOp op) { + return tritonTypeConverter.isSignatureLegal(op.getFunctionType()); + }); + + // For CustomOp, tt.ptr should be converted to memref. + target.addDynamicallyLegalOp([&](hivm::CustomOp op) { + return all_of(op->getOperandTypes(), [](Type t) { + if (isa(t)) { + return false; + } + if (auto shapedType = dyn_cast(t)) { + return !isa(shapedType.getElementType()); + } + return true; + }); + }); + + target.addDynamicallyLegalOp([](arith::ConstantOp op) { + auto res = op.getResult(); + if (!isa(res.getType())) { + return true; + } + + if (auto denseAttr = dyn_cast(op.getValue())) { + if (!denseAttr.isSplat() || + !isa(denseAttr.getElementType())) { + return true; + } + if (res.hasOneUse() && isa(*res.user_begin())) { + return true; + } + return false; + } + return true; + }); + + target.addDynamicallyLegalOp([](Operation *op) { + return llvm::all_of(op->getOperandTypes(), [](Type t) { + if (isa(t)) { + return false; + } + if (auto shapedType = dyn_cast(t)) { + return shapedType.getElementType().isIntOrFloat(); + } + assert(t.isIntOrIndexOrFloat()); + return true; + }); + }); + + target.addDynamicallyLegalDialect( + [this](Operation *op) { + if (op->hasAttr("MetaUse")) { + return false; + } + + if (isa(op)) { + return true; + } + + bool operateOnTensors = + llvm::all_of(op->getOperandTypes(), + [](Type type) { return isa(type); }); + + return this->namedOps || !operateOnTensors; + }); +} + +void TritonToLinalgPass::populateTritonToLinalgCanonicalizationPatterns( + RewritePatternSet &patterns) { + patterns.add, + LoadStoreConverter::LoadStoreCanonicalizer, + LoadStoreConverter::LoadStoreCanonicalizer, + LoadStoreConverter::LoadStoreCanonicalizer>( + patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add< + TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + // TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer< + arith:: + MulFOp>, // TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer, + TTOpConverters::ScalarMathCanonicalizer + // By test, the following ops do not need canonicalization. + // TTOpConverters::ScalarMathCanonicalizer + // TTOpConverters::ScalarMathCanonicalizer + // TTOpConverters::ScalarMathCanonicalizer + >(patterns.getContext()); + patterns.add( + patterns.getContext()); + if (this->enableSelectAnalysis) { + patterns.add(patterns.getContext()); + } +} + +void TritonToLinalgPass::populateTritonToLinalgConversionPatterns( + TypeConverter &typeConverter, RewritePatternSet &patterns, + unsigned int launchGridRank) { + nd2nzFlag = this->enableNd2nzOnVector; + populateFunctionOpInterfaceTypeConversionPattern( + patterns, typeConverter); + + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + // reduce converters + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + + patterns.add(patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add>( + patterns.getContext()); + patterns.add>( + patterns.getContext()); + patterns.add(patterns.getContext()); + + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + + // Add convert pattern for CustomOp. + patterns.add(patterns.getContext()); + + if (!this->namedOps) { + linalg::populateElementwiseToLinalgConversionPatterns(patterns); + } +} + +void TritonToLinalgPass::getDependentDialects(DialectRegistry ®istry) const { + registry.insert(); +} + +LogicalResult +TritonToLinalgPass::processDescriptorOperations(ModuleOp moduleOp) { + // --- ConversionTarget: dynamic legality checks --- + mlir::ConversionTarget target(getContext()); + target.addLegalDialect(); + + // Dialect-level dynamic legality: ops are legal if none of their + // operands/results use TensorDescType. + target.addDynamicallyLegalDialect< + mlir::arith::ArithDialect, mlir::scf::SCFDialect, triton::TritonDialect>( + [](mlir::Operation *op) { + return !DescriptorConverter::hasATensorDescriptorType( + op->getOperandTypes()) && + !DescriptorConverter::hasATensorDescriptorType( + op->getResultTypes()); + }); + // Function signature legality: Triton FuncOp is legal if its inputs/outputs + // contain no TensorDescType. + target.addDynamicallyLegalOp([](triton::FuncOp funcOp) { + return !DescriptorConverter::hasATensorDescriptorType( + funcOp.getFunctionType().getInputs()) && + !DescriptorConverter::hasATensorDescriptorType( + funcOp.getFunctionType().getResults()); + }); + target.addLegalOp(); + target.addIllegalOp(); + + // --- Patterns --- + mlir::RewritePatternSet patterns(&getContext()); + patterns.add( + patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add( + patterns.getContext()); + + mlir::ConversionConfig config; + config.buildMaterializations = true; + if (failed(applyPartialConversion(moduleOp, target, std::move(patterns), + config))) { + moduleOp->emitError("failed to convert tensor descriptor operations"); + return failure(); + } + + return success(); +} + +LogicalResult +TritonToLinalgPass::processPtrBroadcastOperations(ModuleOp moduleOp) { + // --- ConversionTarget: dynamic legality checks --- + mlir::ConversionTarget target(getContext()); + target.addLegalOp(); + target.addLegalOp(); + target.addDynamicallyLegalOp([](triton::BroadcastOp op) { + if (op->hasAttr("MetaUse")) { + return true; + } + auto resultType = dyn_cast(op.getType()); + HoistBroadcast::BroadcastHoister hoister(op); + return !(isa(resultType.getElementType()) && + hoister.canBroadcast()); + }); + + // --- Patterns --- + mlir::RewritePatternSet patterns(&getContext()); + patterns.add(patterns.getContext()); + + if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) { + moduleOp->emitError("failed to convert ptr broadcast operations"); + return failure(); + } + + return success(); +} + +LogicalResult +TritonToLinalgPass::processImplicitPermuteOperations(ModuleOp moduleOp) { + mlir::RewritePatternSet patterns(&getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); + patterns.add( + patterns.getContext()); + + if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { + LLVM_DEBUG({ llvm::dbgs() << "ImplicitPermute: rewrite MemOp failed\n"; }); + } + + mlir::PassManager pm(&getContext(), moduleOp.getOperationName()); + pm.addPass(createCSEPass()); + pm.addPass(createCanonicalizerPass()); + return runPipeline(pm, getOperation()); +} + +LogicalResult +TritonToLinalgPass::processLegalStrideOperations(ModuleOp moduleOp) { + mlir::ConversionTarget target(getContext()); + target.addLegalOp(); + target.addDynamicallyLegalOp( + [](memref::ReinterpretCastOp op) { + return !LoadStoreConverter::ReinterpretCastStrideCanonicalizer:: + hasFixableZeroStride(op); + }); + + mlir::RewritePatternSet patterns(&getContext()); + patterns.add( + patterns.getContext()); + + if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) { + moduleOp->emitError( + "failed to legalize reinterpret_cast dynamic stride(0) with size(1)"); + return failure(); + } + + return success(); +} + +void TritonToLinalgPass::runOnOperation() { + compileOn91095Flag = this->compileOn91095; + + auto moduleOp = getOperation(); + + // Check if the kernel contains tl.dot. Without tl.dot, + // the kernel would be pure AIV kernel. + bool existDot = false; + moduleOp.walk([&](triton::DotOp dotOp) { + existDot = true; + return WalkResult::interrupt(); + }); + moduleOp.walk([&](triton::DotScaledOp dotScaledOp) { + existDot = true; + return WalkResult::interrupt(); + }); + existDotFlag = existDot; + + bool existSIMTOp = false; + moduleOp.walk([&](Operation *op) { + if (isSIMTOp(op)) { + existSIMTOp = true; + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "Found SIMT op in function: "; + os << op->getName(); + os << "\n"; + }); + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + + // Execute tensor descriptor operations conversion + if (failed(processDescriptorOperations(moduleOp))) { + signalPassFailure(); + } + + // Execute implicit permute + if (failed(processImplicitPermuteOperations(moduleOp))) { + LLVM_DEBUG( + { llvm::dbgs() << "Failed to process implicit permute operations\n"; }); + signalPassFailure(); + } + + // 0. Annotate Memory-Related Triton FuncOps with tensor_kind (used by + // profiling). + { + PassManager pm(&getContext(), moduleOp.getOperationName()); + pm.addPass(triton::createMarkTensorKindPass()); + if (failed(runPipeline(pm, moduleOp))) { + moduleOp->emitError("failed to run LoopCanonicalizerPass"); + signalPassFailure(); + return; + } + } + + RewritePatternSet canonicalizerPatterns(&getContext()); + // 1. Canonicalize load/store related patterns. + this->populateTritonToLinalgCanonicalizationPatterns(canonicalizerPatterns); + if (failed( + applyPatternsGreedily(moduleOp, std::move(canonicalizerPatterns)))) { + moduleOp->emitError("failed to apply Canonicalizer Patterns"); + signalPassFailure(); + } + + // 2.1 Pre-clean dead control-flow before use analysis. + // This helps remove unreachable branches such as `scf.if %true` else-region, + // so runUseAnalysis won't walk dead ops with missing lattice states. + { + PassManager pm(&getContext(), moduleOp.getOperationName()); + pm.addPass(createCSEPass()); + pm.addPass(createCanonicalizerPass()); + if (failed(runPipeline(pm, moduleOp))) { + moduleOp->emitError( + "failed to pre-clean dead control-flow before use analysis"); + signalPassFailure(); + return; + } + } + + // 2. Perform use analysis on FuncOp. + moduleOp.walk([this](triton::FuncOp op) { + if (failed(runUseAnalysis(op))) { + signalPassFailure(); + } + }); + + RewritePatternSet patterns(&getContext()); + ConversionTarget target(getContext()); + TritonTypeConverter tritonTypeConverter{}; + + // 3. Mark legal dialects and operations. + this->addDynamicLegal(target, tritonTypeConverter); + + // 4. Mark ops that must be converted explicitly (e.g. tt.scan). + auto loopOpLegalFn = [](LoopLikeOpInterface op) { + return !op.getOperation()->hasAttr("UnhandledLoopOp"); + }; + + target.addIllegalOp(); + target.addDynamicallyLegalOp(loopOpLegalFn); + target.addDynamicallyLegalOp(loopOpLegalFn); + + // 5. Register converters for all illegal Triton ops. + // Execute ptr broadcast operations conversion + if (failed(processPtrBroadcastOperations(moduleOp))) { + signalPassFailure(); + } + this->populateTritonToLinalgConversionPatterns(tritonTypeConverter, patterns, + LAUNCH_GRID_RANK); + + // 6. Inject program id / number of programs arguments into each Triton kernel + // function. + for (auto func : getOperation().getOps()) { + addProgramInfo(func, globalKernel); + } + + moduleOp.walk([this](LoopLikeOpInterface loopOp) { + auto *op = loopOp.getOperation(); + if (!op->hasAttr("ExtractedLoadOrStore")) + op->setAttr("UnhandledLoopOp", UnitAttr::get(op->getContext())); + + for (auto res : loopOp->getResults()) { + if (auto tensorType = dyn_cast(res.getType()); + tensorType && + !isa(tensorType.getElementType())) { + IRRewriter rewriter(op->getContext()); + rewriter.setInsertionPointAfter(op); + auto newVal = + rewriter.create(op->getLoc(), res.getType(), res); + rewriter.replaceAllUsesExcept(res, newVal, newVal); + } + } + }); + + // 7. Convert ops. + if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) { + moduleOp->emitError("failed to apply Conversion Patterns"); + signalPassFailure(); + } + + // Execute legal stride operations conversion + if (failed(processLegalStrideOperations(moduleOp))) { + signalPassFailure(); + } + + // 8. Convert function prologue/epilogue. + moduleOp.walk([&](triton::FuncOp func) { + this->convertTTFunc(func, existDot, existSIMTOp); + }); + + // 9. Clean up dead code and simplify IR. + PassManager pm(&getContext(), moduleOp.getOperationName()); + pm.addPass(createCSEPass()); + pm.addPass(createCanonicalizerPass()); + if (failed(runPipeline(pm, getOperation()))) { + signalPassFailure(); + } + + // Calculate size of PointerCastOp precisely + SmallVector castOps; + + moduleOp.walk([&](hivm::PointerCastOp op) { castOps.push_back(op); }); + + for (auto op : castOps) { + SmallVector userOps(op->getUsers().begin(), + op->getUsers().end()); + IRRewriter rewriter(&getContext()); + rewriter.setInsertionPointAfter(op); + Value addr = op.getAddrs()[0]; + auto elementType = + cast(op.getResult().getType()).getElementType(); + Value elementTypeSize; + if (auto intType = dyn_cast(elementType)) { + elementTypeSize = rewriter.create( + op.getLoc(), + rewriter.getIntegerAttr(addr.getType(), intType.getWidth() / 8)); + } else if (auto floatType = dyn_cast(elementType)) { + elementTypeSize = rewriter.create( + op.getLoc(), + rewriter.getIntegerAttr(addr.getType(), floatType.getWidth() / 8)); + } else { + llvm_unreachable("Cannot get memory size"); + } + + for (auto userOp : userOps) { + auto reinterpretCastOp = cast(userOp); + auto sizes = reinterpretCastOp.getStaticSizes(); + auto staticStrides = reinterpretCastOp.getStaticStrides(); + auto strides = reinterpretCastOp.getStrides(); + if (reinterpretCastOp.getStaticOffsets().size() != 1) + userOp->emitError("IntToPtrOp must converted to PointerCastOp of " + "memref type"); + int64_t castOpSize = 0; + SmallVector dynamicSizes; + for (const auto &[size, stride] : llvm::zip_equal(sizes, staticStrides)) { + assert(!ShapedType::isDynamic(size)); + if (ShapedType::isDynamic(stride)) + dynamicSizes.push_back(size); + else + castOpSize = size * stride; + } + rewriter.setInsertionPoint(reinterpretCastOp); + Value dynamicSize = rewriter.create( + op.getLoc(), rewriter.getIndexAttr(castOpSize)); + for (const auto &[size, stride] : + llvm::zip_equal(dynamicSizes, strides)) { + Value axisSize = rewriter.create( + op.getLoc(), rewriter.getIndexAttr(size)); + axisSize = + rewriter.create(op.getLoc(), stride, axisSize); + dynamicSize = + rewriter.create(op.getLoc(), dynamicSize, axisSize); + } + Value offsetValue; + auto staticOffset = reinterpretCastOp.getStaticOffsets()[0]; + if (ShapedType::isDynamic(staticOffset)) { + offsetValue = reinterpretCastOp.getOffsets()[0]; + if (offsetValue.getType() != addr.getType()) + offsetValue = rewriter.create( + op.getLoc(), addr.getType(), offsetValue); + } else { + offsetValue = rewriter.create( + op.getLoc(), rewriter.getIntegerAttr(addr.getType(), staticOffset)); + } + offsetValue = rewriter.create(op.getLoc(), offsetValue, + elementTypeSize); + Value realAddr = + rewriter.create(op.getLoc(), addr, offsetValue); + auto memrefType = MemRefType::get({ShapedType::kDynamic}, elementType); + auto newCastOp = rewriter.create( + op.getLoc(), memrefType, realAddr, dynamicSize); + auto markOp = rewriter.create(op.getLoc(), + newCastOp.getResult()); + markOp->setAttr(hivm::AddressSpaceAttr::getMnemonic(), + {hivm::AddressSpaceAttr::get(rewriter.getContext(), + hivm::AddressSpace::GM)}); + auto oldMemrefType = + cast(reinterpretCastOp.getResult().getType()); + auto newMemrefType = MemRefType::get( + oldMemrefType.getShape(), oldMemrefType.getElementType(), + StridedLayoutAttr::get(&getContext(), ShapedType::kDynamic, + staticStrides)); + rewriter.replaceOpWithNewOp( + reinterpretCastOp, newMemrefType, newCastOp, ValueRange({}), + reinterpretCastOp.getSizes(), reinterpretCastOp.getStrides(), + SmallVector({0}), reinterpretCastOp.getStaticSizes(), + reinterpretCastOp.getStaticStrides()); + } + rewriter.eraseOp(op); + } + + // Try interleave optimization + llvm::DenseMap> interleaveCandidate; + llvm::DenseMap> + interleaveCandidateWithMask; + moduleOp.walk([&](bufferization::MaterializeInDestinationOp materializeOp) { + if (auto reinterpretCastOp = + materializeOp.getDest() + .getDefiningOp()) { + if (llvm::isa(reinterpretCastOp.getSource()) && + reinterpretCastOp.getStaticStrides().back() == 2) { + interleaveCandidate[llvm::cast( + reinterpretCastOp.getSource())] + .push_back(materializeOp); + } + } + + // Difference is that converted op chain of store with mask has + // `memref::SubViewOp` + if (auto subviewOp = + materializeOp.getDest().getDefiningOp()) { + if (!llvm::isa( + materializeOp.getSource().getDefiningOp())) + return WalkResult::advance(); + + if (auto reinterpretCastOp = + subviewOp.getSource() + .getDefiningOp()) { + if (llvm::isa(reinterpretCastOp.getSource()) && + reinterpretCastOp.getStaticStrides().back() == 2) { + interleaveCandidateWithMask[llvm::cast( + reinterpretCastOp.getSource())] + .push_back(materializeOp); + } + } + } + + return WalkResult::advance(); + }); + + for (auto [blockArg, materializeVec] : interleaveCandidate) { + // Just enable optimization where exists double materializeOp with same + // block argument destination. + if (materializeVec.size() != 2) + continue; + auto result = InterleaveStatusOptimization(materializeVec); + } + + for (auto [blockArg, materializeVec] : interleaveCandidateWithMask) { + if (materializeVec.size() != 2) + continue; + auto result = InterleaveStatusWithMaskOptimization(materializeVec); + } + + // Force to add an argument at the beginning of function arguments, which + // represents stub arg for workspace. Default type is memref + for (auto func : getOperation().getOps()) { + if (!func->hasAttr("global_kernel")) + continue; + + auto context = func.getContext(); + constexpr int64_t syncBlockLockArgIdx = 0; + NamedAttribute syncBlockLockArgAttr( + StringAttr::get(context, "syncBlockLock"), UnitAttr::get(context)); + MemRefType syncBlockLockArgType = + MemRefType::get(SmallVector(1, ShapedType::kDynamic), + IntegerType::get(context, 8)); + llvm::LogicalResult syncBlockLockArg = + func.insertArgument(syncBlockLockArgIdx, // argIndex + syncBlockLockArgType, // argType + nullptr, func->getLoc()); // dicAttr + func->setAttr("SyncBlockLockArgIdx", + IntegerAttr::get(IntegerType::get(&getContext(), 64), + 0)); // 64: 64位整型 + + constexpr int64_t workspaceArgIdx = 1; + MemRefType workspaceArgType = + MemRefType::get(SmallVector(1, ShapedType::kDynamic), + IntegerType::get(context, 8)); + NamedAttribute workspaceArgAttr(StringAttr::get(context, "workspace"), + UnitAttr::get(context)); + + llvm::LogicalResult workspaceArg = + func.insertArgument(/*argIndex*/ workspaceArgIdx, + /*argType*/ workspaceArgType, + /*dicAttr*/ nullptr, func->getLoc()); + func->setAttr("WorkspaceArgIdx", + IntegerAttr::get(IntegerType::get(&getContext(), 64), + 1)); // 64: 64位整型 + } + + // Fix the Location info + moduleOp.walk([&](Operation *op) { + auto loc = op->getLoc(); + if (isa(loc)) { + llvm::SmallPtrSet stopOps; + traverseForwardUpdateUserChainIf( + op, + /*conditionFn*/ + [](Operation *curOp) { return false; }, + /*stopFn*/ + [](Operation *curOp) { return !isa(curOp->getLoc()); }, + /*actionFn*/ + nullptr, stopOps); + if (stopOps.empty()) { + op->emitWarning() << *op << " and its users all have no location!"; + } else { + Operation *goodOp = *stopOps.begin(); + op->setLoc(goodOp->getLoc()); + } + } + return WalkResult::advance(); + }); +} + +std::unique_ptr> triton::createTritonToLinalgPass( + bool globalKernel, bool namedOps, bool enableNd2nzOnVector, + bool enableSelectAnalysis, bool compileOn91095) { + return std::make_unique( + globalKernel, namedOps, enableNd2nzOnVector, enableSelectAnalysis, + compileOn91095); +} + +std::unique_ptr> triton::createTritonToLinalgPass() { + return std::make_unique(); +} diff --git a/compiler/lib/TritonToLinalg/UseAnalysis.cpp b/compiler/lib/TritonToLinalg/UseAnalysis.cpp new file mode 100644 index 00000000..6d42d83a --- /dev/null +++ b/compiler/lib/TritonToLinalg/UseAnalysis.cpp @@ -0,0 +1,560 @@ + + +#include "dicp/TritonToLinalg/UseAnalysis.h" +#include "dicp/Utils/Utils.h" + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "mlir/Analysis/DataFlow/ConstantPropagationAnalysis.h" +#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h" + +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Debug.h" + +using namespace mlir; +using namespace triton; +using namespace dataflow; + +#define DEBUG_TYPE "triton-use-analysis" + +std::string stringifyUseType(UseType useTy) { + std::string ret; + if (useTy == UseType::MetaUse) { + ret = "MetaUse"; + } else if (useTy == UseType::DataUse) { + ret = "DataUse"; + } else if (useTy == UseType::MixUse) { + ret = "MixUse"; + } else if (useTy == UseType::Undefined) { + ret = "Undefined"; + } + return ret; +} + +#if LLVM_VERSION_MAJOR >= 20 +LogicalResult +triton::UseAnalysis::visitOperation(Operation *op, ArrayRef operands, + ArrayRef results) { +#else +void triton::UseAnalysis::visitOperation(Operation *op, + ArrayRef operands, + ArrayRef results) { +#endif + + if (op->getResults().size() == 1) { + auto resultType = dyn_cast(op->getResult(0).getType()); + if (resultType && isa(resultType.getElementType())) { + for (auto opnd : operands) { + propagateUse(opnd, UseType::MetaUse); + } + } + } + + TypeSwitch(op) + .Case([&](auto load) { + propagateUse(operands[0], UseType::MetaUse); + auto mask = load.getMask(); + auto other = load.getOther(); + if (mask) { + assert(mask != other && "mask and other cannot be the same"); + propagateUse(operands[1], UseType::MetaUse); + } + if (other) { + propagateUse(operands[2], UseType::MetaUse); + } + }) + .Case( + [&](auto print) { propagateUse(operands[0], UseType::DataUse); }) + .Case( + [&](auto assert) { propagateUse(operands[0], UseType::DataUse); }) + .Case([&](auto store) { + propagateUse(operands[0], UseType::MetaUse); + propagateUse(operands[1], UseType::DataUse); + auto value = store.getValue(); + auto mask = store.getMask(); + if (mask) { + assert(mask != value && "mask and data cannot be the same"); + propagateUse(operands[2], UseType::MetaUse); + } + }) + .Case([&](auto store) { + propagateUse(operands[0], UseType::MetaUse); + propagateUse(operands[1], UseType::MetaUse); + propagateUse(operands[2], UseType::DataUse); + auto value = store.getValue(); + auto mask = store.getMask(); + if (mask) { + assert(mask != value && "mask and data cannot be the same"); + propagateUse(operands[3], UseType::MetaUse); + } + }) + // Consider triton::AtomicRMWOp as store operation + .Case([&](auto atomicOp) { + propagateUse(operands[0], UseType::MixUse); + propagateUse(operands[1], UseType::DataUse); + auto value = atomicOp.getVal(); + auto mask = atomicOp.getMask(); + if (mask) { + assert(mask != value && "mask and data cannot be the same"); + propagateUse(operands[2], UseType::MetaUse); + } + }) + .Case([&](auto atomicOp) { + propagateUse(operands[0], UseType::MetaUse); + propagateUse(operands[1], UseType::DataUse); + propagateUse(operands[2], UseType::DataUse); + auto value = atomicOp.getVal(); + }) + .Case([&](auto dot) { + propagateResults(operands[0], results); + propagateResults(operands[1], results); + + auto opc = dot.getC(); + triton::SplatOp splat; + if (opc) { + splat = opc.template getDefiningOp(); + } + + if (opc && splat && splat.getSrc().getDefiningOp()) { + propagateUse(operands[2], UseType::MetaUse); + } else { + propagateUse(operands[2], UseType::DataUse); + } + }) + .Case([&](auto loopOp) { + for (const auto &[yield, init, result] : llvm::zip_equal( + loopOp.getYieldedValues(), loopOp.getInits(), results)) { + propagateResults(getLatticeElement(yield), {result}); + propagateResults(getLatticeElement(init), {result}); + } + }) + .Case([&](auto reduceOp) { + for (auto operand : operands) { + propagateUse(operand, UseType::DataUse); + } + }) + .Case( + [&](auto fixpipeOp) { propagateUse(operands[0], UseType::DataUse); }) + .Case( + [&](auto copyOp) { propagateUse(operands[0], UseType::DataUse); }) + .Default([&](Operation *op) { + // this condition account for tt.addptr + for (auto operand : operands) { + propagateResults(operand, results); + } + }); +#if LLVM_VERSION_MAJOR >= 20 + return success(); +#endif +} + +void setMixUseRecursively(Operation *rootOp, bool applyRoot = true) { + traverseBackwardUpdateOperandChainIf( + rootOp, + // ConditionFn + [rootOp, applyRoot](Operation *curOp) { + for (auto res : curOp->getResults()) { + auto tensorType = dyn_cast(res.getType()); + if (tensorType && + isa(tensorType.getElementType())) + return false; + } + return isMetaUse(curOp) && (curOp != rootOp || applyRoot); + }, + // StopFn + [rootOp](Operation *curOp) { + return isa(curOp) && curOp != rootOp; + }, + // ActionFn + [](OpBuilder &b, Operation *op) { + LLVM_DEBUG({ op->setAttr("MixUse", UnitAttr::get(b.getContext())); }); + op->removeAttr("MetaUse"); + }); +} + +static void setMixUseFromValue(Value v) { + if (auto *defOp = v.getDefiningOp()) { + setMixUseRecursively(defOp); + return; + } + + auto blockArg = dyn_cast(v); + if (!blockArg) { + return; + } + + auto *parentOp = blockArg.getOwner()->getParentOp(); + auto loopLikeOp = dyn_cast_or_null(parentOp); + if (!loopLikeOp) { + return; + } + + if (OpOperand *init = loopLikeOp.getTiedLoopInit(blockArg)) { + if (auto *initDefOp = init->get().getDefiningOp()) + setMixUseRecursively(initDefOp); + } + + if (OpOperand *yielded = loopLikeOp.getTiedLoopYieldedValue(blockArg)) { + if (auto *yieldDefOp = yielded->get().getDefiningOp()) + setMixUseRecursively(yieldDefOp); + } +} + +std::optional isIterArgMixUse(Value v, Value target, + const DataFlowSolver &solver) { + auto defOp = v.getDefiningOp(); + auto *use = solver.lookupState(v); + if ((use && use->type == UseType::DataUse) || + isa_and_nonnull(defOp)) + return true; + if (v == target) + return false; + if (!defOp) + return std::nullopt; + for (auto oper : defOp->getOperands()) { + auto res = isIterArgMixUse(oper, target, solver); + if (res.has_value()) + return res.value() || !isMetaUse(defOp); + } + return std::nullopt; +} + +void postProcessWhileOp(scf::WhileOp op, const DataFlowSolver &solver) { + for (const auto &[res, arg] : + llvm::zip_equal(op->getResults(), op.getConditionOp().getArgs())) { + auto *defOp = arg.getDefiningOp(); + if (!defOp) + continue; + auto *use = solver.lookupState(res); + if (use && use->type == UseType::DataUse) + setMixUseRecursively(defOp); + } + for (const auto &[yield, regionArg] : llvm::zip_equal( + op.getYieldOp().getOperands(), op.getBeforeArguments())) { + auto *defOp = yield.getDefiningOp(); + if (!defOp) + continue; + if (isIterArgMixUse(yield, regionArg, solver).value_or(false)) + setMixUseRecursively(defOp); + } +} + +void postProcessLoopOp(LoopLikeOpInterface loopOp, + const DataFlowSolver &solver) { + if (auto whileOp = dyn_cast(loopOp.getOperation())) { + postProcessWhileOp(whileOp, solver); + return; + } + for (const auto &[res, yield, regionArg] : + llvm::zip_equal(loopOp->getResults(), loopOp.getYieldedValues(), + loopOp.getRegionIterArgs())) { + auto *defOp = yield.getDefiningOp(); + if (!defOp) + continue; + auto *use = solver.lookupState(res); + if ((use && use->type == UseType::DataUse) || + isIterArgMixUse(yield, regionArg, solver).value_or(false)) + setMixUseRecursively(defOp); + } +} + +LogicalResult triton::runUseAnalysis(triton::FuncOp &funcOp) { + MLIRContext *context = funcOp.getContext(); + SymbolTableCollection symbolTable; + + DataFlowSolver solver; + solver.load(); + solver.load(); + solver.load(symbolTable); + if (failed(solver.initializeAndRun(funcOp))) { + return failure(); + } + auto &os = llvm::dbgs(); + // Walk the func op, convert tags on operands to tags on operations + funcOp.walk([&](Operation *op) { + LLVM_DEBUG({ os << "[UseAnalysis] op is " << *op << "\n"; }); + UseType useType = UseType::Undefined; + for (auto result : op->getResults()) { + LLVM_DEBUG({ os << "[UseAnalysis] ===> result is " << result << "\n"; }); + auto use = solver.lookupState(result); + assert(use && "Lattice value not found"); + auto thisUseType = use->type; + LLVM_DEBUG({ + os << "[UseAnalysis] ==========> useType is " + << stringifyUseType(thisUseType) << "\n"; + }); + if (thisUseType == UseType::Undefined) { + continue; + } + if (useType == UseType::Undefined) { + useType = thisUseType; + } + if (thisUseType == UseType::MixUse || thisUseType != useType) { + useType = UseType::MixUse; + break; + } + } + + if (useType == UseType::Undefined) { + LLVM_DEBUG({ op->setAttr("Undefined", UnitAttr::get(context)); }); + return; + } else if (useType == UseType::MetaUse) { + auto memEffect = dyn_cast(op); + if (memEffect) { + if (isa(op)) { + LLVM_DEBUG( + { os << "force protecting side-effect op:" << *op << "\n"; }); + op->setAttr("DataUse", UnitAttr::get(context)); + return; + } + } + if (!isa(op)) { + assert(op->getNumResults() == 1 && + "Ops used for meta computation are expected to have one result"); + } + for (auto it = 0; it < op->getNumResults(); ++it) { + // Only set the tag if the operation uses tensors + if (isa(op->getResult(it).getType()) || + (isa(op) && + op->hasAttr(ConverterUtils::discreteAttrName)) || + (isa(op) && + isa(op->getResult(it).getType()))) { + // Setting tag for erasing op later + op->setAttr("MetaUse", UnitAttr::get(context)); + } + } + return; + } else if (useType == UseType::DataUse) { + LLVM_DEBUG({ op->setAttr("DataUse", UnitAttr::get(context)); }); + return; + } + + assert(useType == UseType::MixUse); + + // If the operation only produces scalars, no need to clone it + bool shapedResult = true; + for (auto result : op->getResults()) + shapedResult &= isa(result.getType()); + if (!shapedResult || + isa(op)) { + LLVM_DEBUG({ op->setAttr("MixUse", UnitAttr::get(context)); }); + return; + } + llvm::SetVector metaUsers; + for (auto result : op->getResults()) { + for (auto user : result.getUsers()) { + TypeSwitch(user) + .Case([&](auto load) { + auto ptr = load.getPtr(); + auto mask = load.getMask(); + auto other = load.getOther(); + if (result == ptr || result == mask || result == other) { + metaUsers.insert(user); + } + }) + .Case([&](auto store) { + auto ptr = store.getPtr(); + auto mask = store.getMask(); + if (result == ptr || result == mask) { + metaUsers.insert(user); + } + }) + .Case([&](auto indirectstore) { + auto src = indirectstore.getSrc(); + auto offset = indirectstore.getOffsets(); + auto mask = indirectstore.getMask(); + if (result == src || result == offset || result == mask) { + metaUsers.insert(user); + } + }) + .Case([&](auto atomicOp) { + auto ptr = atomicOp.getPtr(); + auto mask = atomicOp.getMask(); + if (result == ptr || result == mask) + metaUsers.insert(user); + }) + .Case([&](auto atomicOp) { + auto ptr = atomicOp.getPtr(); + if (result == ptr) + metaUsers.insert(user); + }) + .Case([&](auto dot) { + auto opc = dot.getC(); + triton::SplatOp splat; + if (opc) { + splat = opc.template getDefiningOp(); + } + + if (opc && splat && + splat.getSrc().getDefiningOp()) { + metaUsers.insert(user); + } + }) + .Case([&](auto print) {}) + .Default([&](Operation *op) { + bool allMeta = true; + for (auto res : op->getResults()) { + auto resUse = solver.lookupState(res); + if (resUse->type != UseType::MetaUse) { + allMeta = false; + break; + } + } + if (allMeta) { + metaUsers.insert(user); + } + }); + } + } + + // If the operation doesn't have direct meta users, no need to clone it + if (metaUsers.empty()) { + LLVM_DEBUG({ op->setAttr("MixUse", UnitAttr::get(context)); }); + return; + } + + if (isa(op)) + return; + + if (isa(op)) + return; + + // Clone the operation; switch all meta users to use the clone + OpBuilder builder(op); + auto clone = builder.clone(*op); + LLVM_DEBUG({ op->setAttr("MixUse", UnitAttr::get(context)); }); + + // Setting tag for erasing op later + clone->setAttr("MetaUse", UnitAttr::get(context)); + + for (auto [res_i, result] : llvm::enumerate(op->getResults())) { + for (auto user : metaUsers) { + for (auto &operand : user->getOpOperands()) { + if (operand.get() == result) { + operand.set(clone->getResult(res_i)); + } + } + } + } + }); + LLVM_DEBUG({ + os << "[UseAnalysis] Before post-process, funcOp is " << *funcOp << "\n"; + }); + // Post-process + funcOp.walk([&](Operation *op) { + // Handle indirect load and store case. + // For example, load(1st) -> computeOp -> load(2nd), + // or load -> computeOp -> store + // The first load is IndirectLoadInterfaceOp. + // Do not inplace replace MetaUse by MixUse. Because the condition checking + // depends on that the op has the attr of MetaUse. + // Handle the indirect load interface op + // We first trace from the 1st load to the 2nd load with the ops between + // them marked as MixUse. Then we traceback from the 2nd load to mark defs + // MixUse. + if (opIsIndirectLoad(op) || opIsIndirectCalc(op) || + isa(op)) { + LLVM_DEBUG({ + os << "[UseAnalysis] Found indirect load interface op: " << *op << "\n"; + }); + llvm::SmallPtrSet stopOps; + // Modify the users of this op's result. + traverseForwardUpdateUserChainIf( + op, + /*conditionFn*/ + [op](Operation *curOp) { return isMetaUse(curOp) && curOp != op; }, + /*stopFn*/ + [&](Operation *curOp) { + // triton::LoadOp or triton::StoreOp without MetaUse means + // it is an indirect load or store + // instead of the load providing the offset. + // The pattern is as follows, + // load -> ops -> load + // load -> ops -> store + // We need to ensure the intermediate ops are marked MixUse + // so that they will be replaced instead of be erased without + // conversion. + return (isa(curOp) || isa(curOp) || + isa(curOp)) && + !isMetaUse(curOp); + }, + /*actionFn*/ + [](OpBuilder &b, Operation *op) { setMixUseRecursively(op); }, + stopOps); + LLVM_DEBUG({ + os << "[UseAnalysis] stopOps are \n"; + for (auto [idx, stopOp] : llvm::enumerate(stopOps)) + os << idx << ": " << *stopOp << "\n"; + }); + LLVM_DEBUG({ + os << "[UseAnalysis] After trace, funcOp is " << *funcOp << "\n"; + }); + for (auto *stopOp : stopOps) + setMixUseRecursively(stopOp, /*applyRoot=*/false); + LLVM_DEBUG({ + os << "[UseAnalysis] After traceback of stopOp, funcOp is " << *funcOp + << "\n"; + }); + // Modify this op. + LLVM_DEBUG({ op->setAttr("MixUse", UnitAttr::get(context)); }); + op->removeAttr("MetaUse"); + } + if (op->hasAttr(ConverterUtils::discreteAttrName)) + setMixUseRecursively(op); + if (auto loopOp = dyn_cast(op)) { + postProcessLoopOp(loopOp, solver); + } else if (auto ifOp = dyn_cast(op)) { + SmallVector yields(ifOp.thenYield().getOperands()); + if (!ifOp.getElseRegion().empty()) + yields.append(llvm::to_vector(ifOp.elseYield().getOperands())); + for (auto yield : yields) { + setMixUseFromValue(yield); + } + } else if (auto atomicRmwOp = dyn_cast(op)) { + auto mask = atomicRmwOp.getMask(); + if (mask && op->hasAttr(ConverterUtils::discreteMaskAttrName)) + setMixUseRecursively(mask.getDefiningOp()); + } + }); + // Remove MetaUse in case of MixUse existing in the op + funcOp.walk([&](Operation *op) { + if (isMetaUse(op) && isMixUse(op)) { + op->removeAttr("MetaUse"); + } + }); + // hivm.custom present library call, shouldn't be metause + funcOp.walk([&](hivm::CustomOp op) { + if (isMetaUse(op)) { + op->removeAttr("MetaUse"); + } + }); + LLVM_DEBUG({ + os << "[UseAnalysis] After post-process, funcOp is " << *funcOp << "\n"; + }); + return success(); +} + +MetaUseEraser::MetaUseEraser(MLIRContext *context) + : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/10, context) {} + +LogicalResult MetaUseEraser::matchAndRewrite(Operation *op, + PatternRewriter &rewriter) const { + LLVM_DEBUG({ + int64_t count = 0; + for (auto result : op->getResults()) { + count += std::distance(result.use_begin(), result.use_end()); + } + llvm::dbgs() << "Number of user: " << count << "\n"; + }); + if (isa(op)) { + return rewriter.notifyMatchFailure(op, + "AddPtrOp will be handled separately"); + } + if (isMetaUse(op)) { + rewriter.eraseOp(op); + return success(); + } + return rewriter.notifyMatchFailure(op, "requires meta ops"); +} diff --git a/compiler/lib/TritonToStructured/CMakeLists.txt b/compiler/lib/TritonToStructured/CMakeLists.txt new file mode 100644 index 00000000..7c2def92 --- /dev/null +++ b/compiler/lib/TritonToStructured/CMakeLists.txt @@ -0,0 +1,27 @@ +add_triton_library(TritonToStructured + TritonToStructuredPass.cpp + PtrAnalysis.cpp + CannonicalizerConverter.cpp + MemOpConverter.cpp + MaskAnalysis.cpp + + DEPENDS + TritonToStructuredConversionPassIncGen + + LINK_LIBS PUBLIC + MLIRArithDialect + MLIRDialectUtils + MLIRIR + MLIRMathDialect + MLIRPass + MLIRTensorDialect + MLIRTransforms + MLIRSupport + TritonIR + TritonTransforms + TritonAnalysis + MLIRTritonNPUUtils + MLIRSCFTransforms + MLIRLinalgTransforms + BiShengIRHIVMDialect +) diff --git a/compiler/lib/TritonToStructured/CannonicalizerConverter.cpp b/compiler/lib/TritonToStructured/CannonicalizerConverter.cpp new file mode 100644 index 00000000..2c8d8233 --- /dev/null +++ b/compiler/lib/TritonToStructured/CannonicalizerConverter.cpp @@ -0,0 +1,2594 @@ + + +#include "dicp/TritonToStructured/CannonicalizerConverter.h" + +#include +#include +#include + +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Passes.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Location.h" +#include "mlir/IR/Matchers.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Value.h" +#include "mlir/Support/LLVM.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/MathExtras.h" + +#include "llvm/Support/Debug.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "dicp/TritonToStructured/PtrAnalysis.h" +#include "dicp/TritonToStructured/TritonToStructuredPass.h" +#include "dicp/Utils/InterleaveOptimization.h" +#include "dicp/Utils/Utils.h" + +#define DEBUG_TYPE "triton-cannonicalizer-converter" + +namespace CannonicalizerConverter { +using namespace mlir; +using namespace triton; +constexpr int INT_TYPE_BIT_WIDTH = 32; + +// Match and rewrite pattern for optimizing cmp.ne (select(cond, 1, 0), 0) -> +// cond This pattern transforms: +// %select = arith.select %cond, %true_val, %false_val +// %cmp = arith.cmpi ne, %select, %zero +// Where: +// - %true_val is a constant splat tensor of 1s +// - %false_val is a constant splat tensor of 0s +// - %zero is a constant splat tensor of 0s +// Into: +// %cond (directly replace the cmp with the select condition) +// +// This optimization is valid because: +// select(cond, 1, 0) != 0 +// is equivalent to: cond != 0 +// Since the result of select is either 1 or 0, the only way it's not equal to +// 0 is when it's 1, which happens exactly when cond is true. +// +// Example: +// Input IR: +// %39 = arith.cmpi slt, %15, %cst_14 : tensor<128xi32> +// %40 = arith.select %39, %cst_13, %cst_12 : tensor<128xi1>, +// tensor<128xi32> %41 = arith.cmpi ne, %40, %cst_12 : tensor<128xi32> +// Where cst_13 is constant dense<1> and cst_12 is constant dense<0> +// Output IR: +// %39 = arith.cmpi slt, %15, %cst_14 : tensor<128xi32> +LogicalResult CmpConverter::matchAndRewrite(arith::CmpIOp cmpOp, + PatternRewriter &rewriter) const { + // Only handle "not equal" comparison + auto cmpType = cmpOp.getPredicate(); + if (cmpType != arith::CmpIPredicate::ne) { + return failure(); + } + + Value rhs = cmpOp.getRhs(); + Value lhs = cmpOp.getLhs(); + + // 1. Check if RHS is a constant zero + APInt rhsValue; + if (!matchPattern(rhs, m_ConstantInt(&rhsValue))) { + return failure(); // RHS is not a constant + } + + if (!rhsValue.isZero()) { + return failure(); // RHS is not zero + } + + // 2. Check if LHS is defined by a select operation + auto selectOp = lhs.getDefiningOp(); + if (!selectOp) { + return failure(); + } + + // 3. Check if select's true and false values are constants + DenseElementsAttr trueAttr; + DenseElementsAttr falseAttr; + if (!matchPattern(selectOp.getTrueValue(), m_Constant(&trueAttr)) || + !matchPattern(selectOp.getFalseValue(), m_Constant(&falseAttr))) { + return failure(); // Either true or false value is not constant + } + + // 4. Check if true value is all 1s and false value is all 0s + if (!trueAttr.isSplat() || !trueAttr.getSplatValue().isOne() || + !falseAttr.isSplat() || !falseAttr.getSplatValue().isZero()) { + return failure(); + } + + // 5. Optimization matched, replace cmp with select's condition + rewriter.replaceOp(cmpOp, selectOp.getCondition()); + return success(); +} + +// Detect when both operands of the cmpOp are triton::SplatOp. If so, +// replace the original comparison by comparing the underlying scalar values +// and then splatting (broadcasting) the scalar comparison result back to the +// original tensor shape. +// Example: +// Input IR: +// %splat_lhs = tt.splat %val1 : tensor<128xi32> +// %splat_rhs = tt.splat %val2 : tensor<128xi32> +// %cmp = arith.cmpi slt, %splat_lhs, %splat_rhs : tensor<128xi32> +// Output IR: +// %cmp_scalar = arith.cmpi slt, %val1, %val2 +// %splat_cmp = tt.splat %cmp_scalar : tensor<128xi1> +LogicalResult +SplatCmpConverter::matchAndRewrite(arith::CmpIOp cmpOp, + PatternRewriter &rewriter) const { + auto lhs = cmpOp.getLhs(); + auto rhs = cmpOp.getRhs(); + auto lhsSplatOp = lhs.getDefiningOp(); + auto rhsSplatOp = rhs.getDefiningOp(); + if (!lhsSplatOp || !rhsSplatOp) { + return failure(); + } + auto lhsSrc = lhsSplatOp.getSrc(); + auto rhsSrc = rhsSplatOp.getSrc(); + auto newCmpOp = rewriter.create( + cmpOp.getLoc(), cmpOp.getPredicate(), lhsSrc, rhsSrc); + auto cmpType = dyn_cast(cmpOp.getType()); + if (!cmpType) { + return failure(); + } + auto splatType = + RankedTensorType::get(cmpType.getShape(), newCmpOp.getType()); + auto splatOp = rewriter.create(cmpOp.getLoc(), splatType, + newCmpOp.getResult()); + rewriter.replaceOp(cmpOp, splatOp.getResult()); + return success(); +} + +// Convert AddPtr when ptr is produced by tt.splat and offset is produced by +// tt.broadcast of a smaller-shaped value. Transformation: +// %ptr_splat = tt.splat %ptr_src : tensor> +// %offset_b = tt.broadcast %offset_src : tensor +// %res = tt.addptr %ptr_splat, %offset_b +// => +// %ptr_small_splat = tt.splat %ptr_src : tensor> +// %add_small = tt.addptr %ptr_small_splat, %offset_src +// %res = tt.broadcast %add_small : tensor> +LogicalResult +AddPtrSplatConverter::matchAndRewrite(triton::AddPtrOp addPtrOp, + PatternRewriter &rewriter) const { + // Match ptr produced by splat + auto ptr = addPtrOp.getPtr(); + auto ptrSplatOp = ptr.getDefiningOp(); + if (!ptrSplatOp) + return failure(); + + // Match offset produced by broadcast + auto offset = addPtrOp.getOffset(); + auto offsetBroadcastOp = offset.getDefiningOp(); + if (!offsetBroadcastOp) + return failure(); + + // Extract sources + auto ptrSrc = ptrSplatOp.getSrc(); + auto offsetSrc = offsetBroadcastOp.getSrc(); + + // Types must be ranked to reason about shapes + auto resultType = dyn_cast(addPtrOp.getResult().getType()); + auto offsetSrcType = dyn_cast(offsetSrc.getType()); + if (!resultType || !offsetSrcType) + return failure(); + + // If offset source already has same shape as result, nothing to do + if (resultType.getShape() == offsetSrcType.getShape()) + return failure(); + + auto loc = addPtrOp.getLoc(); + + // Build pointer tensor type corresponding to offset's (smaller) shape + // ptrSrc should be a pointer type (element type) + auto ptrElementType = dyn_cast(ptrSrc.getType()); + auto smallPtrTensorType = + RankedTensorType::get(offsetSrcType.getShape(), ptrElementType); + + // Create splat of ptrSrc to the smaller shape + auto smallPtrSplat = + rewriter.create(loc, smallPtrTensorType, ptrSrc); + + // Create addptr on the smaller shape + auto smallAdd = rewriter.create( + loc, smallPtrTensorType, smallPtrSplat.getResult(), offsetSrc); + + // Broadcast the small add result back to the original (larger) shape + auto broadcasted = rewriter.create(loc, resultType, + smallAdd.getResult()); + + rewriter.replaceOp(addPtrOp, broadcasted.getResult()); + return success(); +} + +// Move load before broadcast when possible: +// If load.ptr is a triton::BroadcastOp and +// - load.mask is null; or +// - load.mask is a triton::BroadcastOp and the broadcast sources (before +// broadcast) +// of ptr and mask have identical shapes, +// then replace +// %ptr_b = tt.broadcast %ptr_src +// %mask_b = tt.broadcast %mask_src? (optional) +// %v = tt.load %ptr_b, %mask_b +// with +// %v_small = tt.load %ptr_src, %mask_src? +// %v = tt.broadcast %v_small +LogicalResult +LoadBroadcastConverter::matchAndRewrite(triton::LoadOp loadOp, + PatternRewriter &rewriter) const { + // Match when ptr is defined by BroadcastOp + Value ptr = loadOp.getPtr(); + auto ptrBroadcast = ptr.getDefiningOp(); + if (!ptrBroadcast) + return failure(); + + auto ptrSrcType = dyn_cast(ptrBroadcast.getSrc().getType()); + if (!ptrSrcType) + return failure(); + + // mask can be null + Value mask = loadOp.getMask(); + Value maskSrc = nullptr; + if (mask) { + auto maskBroadcast = mask.getDefiningOp(); + if (!maskBroadcast) + return failure(); + maskSrc = maskBroadcast.getSrc(); + // shapes of ptrBroadcast.src and maskBroadcast.src must match + auto ptrSrcType = + dyn_cast(ptrBroadcast.getSrc().getType()); + auto maskSrcType = dyn_cast(maskSrc.getType()); + if (!ptrSrcType || !maskSrcType) + return failure(); + if (ptrSrcType.getShape() != maskSrcType.getShape()) + return failure(); + } + + // Prepare the smaller load: load from ptrBroadcast.src with maskSrc (or null) + Location loc = loadOp.getLoc(); + Value smallPtr = ptrBroadcast.getSrc(); + + // Preserve other operands of load (like 'other', cache, evict, isVolatile) + Value other = loadOp.getOther(); + auto cache = loadOp.getCache(); + auto evict = loadOp.getEvict(); + auto isVolatile = loadOp.getIsVolatile(); + + // If 'other' exists, it must be a constant DenseElementsAttr so we can + // construct a smaller-shaped constant to feed the small load. If it's + // non-constant, abort the transformation. + Value newOther = other; + if (other) { + Attribute otherAttr; + if (!matchPattern(other, m_Constant(&otherAttr))) + return failure(); + auto denseOther = dyn_cast(otherAttr); + if (!denseOther) + return failure(); + + // Build a small-shaped tensor type that uses the ptr src's shape but the + // element type of the 'other' constant + auto ptrSrcRT = dyn_cast(ptrBroadcast.getSrc().getType()); + if (!ptrSrcRT) + return failure(); + + auto elemType = denseOther.getType().getElementType(); + if (!elemType) + return failure(); + + auto smallType = RankedTensorType::get(ptrSrcRT.getShape(), elemType); + + DenseElementsAttr newDense; + if (denseOther.isSplat()) { + // Reuse the splat value; assume element/value types are compatible. + newDense = DenseElementsAttr::get(smallType, + denseOther.getSplatValue()); + } else { + // Multi-element dense constant cannot be safely reshaped here + return failure(); + } + + auto constOp = rewriter.create(loc, newDense); + newOther = constOp.getResult(); + } + + auto newLoad = rewriter.create( + loc, smallPtr, maskSrc, newOther, cache, evict, isVolatile); + + // Broadcast result back to original result type + auto resultType = dyn_cast(loadOp.getResult().getType()); + if (!resultType) + return failure(); + auto broadcasted = rewriter.create(loc, resultType, + newLoad.getResult()); + + rewriter.replaceOp(loadOp, broadcasted.getResult()); + return success(); +} + +LogicalResult PromotePointerIterArgsPattern::matchAndRewrite( + scf::ForOp forOp, PatternRewriter &rewriter) const { + // 1. Check if the loop meets transformation conditions + if (failed(matchLoop(forOp))) { + return failure(); + } + + // 2. Collect pointer iteration arguments to be processed + if (!failed(matchAndRewriteForAddPtr(forOp, rewriter))) { + return success(); + } + return matchAndRewriteAdvancePtr(forOp, rewriter); +} + +// Transform a for loop that uses pointer iteration arguments into one that uses +// integer offsets instead. This pattern handles the specific case where: +// 1. The loop has pointer iteration arguments of type like +// tensor<1024x!tt.ptr> +// 2. Each pointer is used in a load/store operation and then incremented by +// a constant offset via tt.addptr +// 3. The updated pointer (from addptr) is yielded back as the next iteration +// value +// +// The transformation converts: +// scf.for iter_args(%ptr = %base_ptr) { +// %val = tt.load %ptr +// tt.store %other_ptr, %val +// %new_ptr = tt.addptr %ptr, %offset +// scf.yield %new_ptr +// } +// +// Into: +// scf.for iter_args(%offset_int = 0) { +// %splat_offset = tt.splat %offset_int +// %current_ptr = tt.addptr %base_ptr, %splat_offset +// %val = tt.load %current_ptr +// tt.store %other_ptr, %val +// %new_offset = arith.addi %offset_int, %const_offset +// scf.yield %new_offset +// } +// +LogicalResult PromotePointerIterArgsPattern::matchAndRewriteForAddPtr( + scf::ForOp forOp, PatternRewriter &rewriter) const { + // 1. Collect pointer iteration arguments to be processed + auto pointerArgsInfo = collectPointerIterArgs(forOp); + if (pointerArgsInfo.empty()) { + return failure(); + } + // 2. Create new iteration argument types and initial values + auto [newInitArgs, newIterArgTypes, indexMap] = + createNewIterArgs(forOp, pointerArgsInfo, rewriter); + // 3. Create the new for loop + auto newForOp = + createNewForLoop(forOp, newInitArgs, newIterArgTypes, rewriter); + // 4. Rewrite the loop body + if (failed(rewriteLoopBody(forOp, newForOp, pointerArgsInfo, indexMap, + rewriter))) { + return failure(); + } + // 5. Replace original loop results + return replaceResults(forOp, newForOp, pointerArgsInfo, indexMap, rewriter); +} + +// Transform a for loop that uses pointer iteration arguments into one that uses +// integer offsets instead. This pattern handles the specific case where: +// 1. The loop has pointer iteration arguments of type like +// !tt.ptr> +// 2. Each pointer is used in a load/store operation and then incremented by +// a constant offset via tt.advanceptr +// 3. The updated pointer (from advanceptr) is yielded back as the next +// iteration value +// +// The transformation converts: +// scf.for iter_args(%ptr1 = %base_ptr1, %ptr2 = %base_ptr2 ) { +// %val = tt.load %ptr1 +// tt.store %base_ptr2, %val +// %new_ptr1 = tt.adcanceptr %ptr1, %offset +// %new_ptr2 = tt.adcanceptr %ptr2, %offset +// scf.yield %new_ptr1,%new_ptr2 +// } +// +// Into: +// The transformation converts: +// scf.for iter_args(%offset_int1 = 0,%offset_int2 = 0) { +// %new_ptr1 = tt.adcanceptr %ptr1, %offset_int1 +// %new_ptr2 = tt.adcanceptr %ptr2, %offset_int1 +// %val = tt.load %new_ptr1 +// tt.store %new_ptr2, %val +// %new_offset1 = arith.addi %offset_int, %const_offset +// %new_offset2 = arith.addi %offset_int, %const_offset +// scf.yield %new_offset1,%new_offset2 +// } +// +LogicalResult PromotePointerIterArgsPattern::matchAndRewriteAdvancePtr( + scf::ForOp forOp, PatternRewriter &rewriter) const { + // 1. Check if the loop meets transformation conditions + if (failed(matchLoop(forOp))) { + return failure(); + } + // 2. Collect pointer iteration arguments to be processed + auto pointerArgsInfo = collectPointerIterArgsForAdvancePtr(forOp); + if (pointerArgsInfo.size() == 0) { + return failure(); + } + // 3. Create new iteration argument types and initial values + auto [newInitArgs, newIterArgTypes, indexMap] = + createNewIterArgsForAdvancePtr(forOp, pointerArgsInfo, rewriter); + + // 4. Create the new for loop + auto newForOp = + createNewForLoop(forOp, newInitArgs, newIterArgTypes, rewriter); + // 5. Rewrite the loop body + if (failed(rewriteLoopBodyForAdvancePtr(forOp, newForOp, pointerArgsInfo, + indexMap, rewriter))) { + return failure(); + } + rewriter.replaceOp(forOp, newForOp); + return success(); +} + +LogicalResult PromotePointerIterArgsPattern::matchLoop(scf::ForOp forOp) const { + auto lowerBound = forOp.getLowerBound(); + auto upperBound = forOp.getUpperBound(); + auto step = forOp.getStep(); + if (!matchPattern(lowerBound, m_Constant()) || + !matchPattern(upperBound, m_Constant()) || + !matchPattern(step, m_Constant())) { + return failure(); + } + return success(); +} + +SmallVector +PromotePointerIterArgsPattern::collectPointerIterArgs(scf::ForOp forOp) const { + SmallVector result; + auto &loopBody = *forOp.getBody(); + for (auto [idx, iterArg] : llvm::enumerate(forOp.getRegionIterArgs())) { + if (isPointerIterArg(iterArg)) { + auto info = analyzePointerIterArg(iterArg, loopBody); + if (info.has_value()) { + info->oldIndex = static_cast(idx), + info->basePointer = forOp.getInitArgs()[idx], + result.push_back(info.value()); + } + } + } + return result; +} + +SmallVector +PromotePointerIterArgsPattern::collectPointerIterArgsForAdvancePtr( + scf::ForOp forOp) const { + SmallVector result; + auto &loopBody = *forOp.getBody(); + + for (auto [idx, iterArg] : llvm::enumerate(forOp.getRegionIterArgs())) { + if (isa(iterArg.getType())) { + auto info = analyzePointerIterArgForAdvancePtr(iterArg, loopBody); + if (info.has_value()) { + info->oldIndex = static_cast(idx); + info->basePointer = forOp.getInitArgs()[idx]; + result.push_back(info.value()); + } + } + } + return result; +} + +bool PromotePointerIterArgsPattern::isPointerIterArg(Value iterArg) const { + auto ptrType = dyn_cast(iterArg.getType()); + return ptrType && isa(ptrType.getElementType()); +} + +std::optional +PromotePointerIterArgsPattern::analyzePointerIterArg(Value iterArg, + Block &loopBody) const { + int memCount = + 0; // Count of memory operations (load/store) using this pointer + int addPtrCount = 0; // Count of addptr operations on this pointer + Value addPtrResult = nullptr; // Result of the addptr operation + Value offset = nullptr; // Offset value used in addptr + Value addPtrValue = nullptr; // The addptr operation result value + + for (auto &op : loopBody) { + TypeSwitch(&op) + .Case([&](auto memoryOp) { + // Check if this memory operation uses the pointer we're analyzing + if (memoryOp.getPtr() == iterArg) + ++memCount; + }) + .Case([&](auto addPtrOp) { + // Check if this addptr operation updates the pointer we're analyzing + if (addPtrOp.getPtr() == iterArg) { + ++addPtrCount; + addPtrResult = addPtrOp.getResult(); + offset = addPtrOp.getOffset(); + addPtrValue = addPtrOp.getResult(); + } + }) + .Default([](auto) {}); // Ignore other operations + } + + // Check the terminator to see if the addptr result is yielded + auto yieldOp = dyn_cast(loopBody.getTerminator()); + if (!yieldOp) + return std::nullopt; + + bool isYielded = false; + for (auto operand : yieldOp.getOperands()) { + if (operand == addPtrResult) { + isYielded = true; + break; + } + } + + // Pattern matched if: + // 1. Exactly one addptr operation on this pointer + // 2. At least one memory operation using this pointer + // 3. The addptr result is yielded + if (addPtrCount == 1 && memCount >= 1 && isYielded) { + return PointerArgInfo{ + .oldIndex = 0, + .basePointer = nullptr, // Will be set in collectPointerIterArgs + .offsetValue = offset, + .newIterArg = nullptr, // Will be set in createNewIterArgs + .addPtrValue = addPtrValue}; + } + return std::nullopt; +} + +std::optional +PromotePointerIterArgsPattern::analyzePointerIterArgForAdvancePtr( + Value iterArg, Block &loopBody) const { + int memCount = + 0; // Count of memory operations (load/store) using this pointer + int advancePtrCount = 0; // Count of advanceptr operations on this pointer + Value advancePtrResult = nullptr; // Result of the addptr operation + SmallVector offsetValues; + Value advancePtrValue = nullptr; // The addptr operation result value + int nonZeroConstant = 0; // Number of non-zero constants + for (auto &op : loopBody) { + TypeSwitch(&op) + .Case([&](auto memoryOp) { + // Check if this memory operation uses the pointer we're analyzing + if (memoryOp.getPtr() == iterArg) + ++memCount; + }) + .Case([&](triton::AdvanceOp advancePtrOp) { + // Check if this addptr operation updates the pointer we're analyzing + if (advancePtrOp.getPtr() == iterArg) { + ++advancePtrCount; + advancePtrResult = advancePtrOp.getResult(); + for (auto offsetVal : advancePtrOp.getOffsets()) { + if (auto offsetInt = getConstantIntValue(offsetVal)) { + if (*offsetInt != 0) { + ++nonZeroConstant; + } + offsetValues.push_back(offsetVal); + } + } + advancePtrValue = advancePtrOp.getResult(); + } + }) + .Default([](auto) {}); // Ignore other operations + } + // Check the terminator to see if the addptr result is yielded + auto yieldOp = dyn_cast(loopBody.getTerminator()); + if (!yieldOp) + return std::nullopt; + bool isYielded = false; + for (auto operand : yieldOp.getOperands()) { + if (operand == advancePtrResult) { + isYielded = true; + break; + } + } + // Pattern matched if: + // 1. Exactly one addptr operation on this pointer + // 2. At least one memory operation using this pointer + // 3. The addptr result is yielded + if (advancePtrCount == 1 && nonZeroConstant == 1 && memCount >= 1 && + isYielded) { + return PointerArgInfo{ + .oldIndex = 0, + .basePointer = nullptr, // Will be set in collectPointerIterArgs + .offsetValue = nullptr, + .newIterArg = nullptr, // Will be set in createNewIterArgs + .addPtrValue = advancePtrValue, + .offsetValues = offsetValues}; + } + return std::nullopt; +} + +std::tuple, SmallVector, DenseMap> +PromotePointerIterArgsPattern::createNewIterArgs( + scf::ForOp forOp, ArrayRef pointerArgs, + PatternRewriter &rewriter) const { + SmallVector newInitArgs; + SmallVector newIterArgTypes; + DenseMap indexMap; + + for (unsigned i = 0; i < forOp.getInitArgs().size(); ++i) { + if (isPointerArgIndex(pointerArgs, i)) { + // Replace pointer with integer offset (initialized to 0) + Value zero = rewriter.create(forOp.getLoc(), 0, + INT_TYPE_BIT_WIDTH); + newInitArgs.push_back(zero); + newIterArgTypes.push_back(rewriter.getIntegerType(INT_TYPE_BIT_WIDTH)); + } else { + // Preserve original argument unchanged + newInitArgs.push_back(forOp.getInitArgs()[i]); + newIterArgTypes.push_back(forOp.getInitArgs()[i].getType()); + } + // Identity mapping: argument count and order unchanged， + // may change in future + indexMap[i] = i; + } + + return {newInitArgs, newIterArgTypes, indexMap}; +} + +std::tuple, SmallVector, DenseMap> +PromotePointerIterArgsPattern::createNewIterArgsForAdvancePtr( + scf::ForOp forOp, SmallVector &pointerArgs, + PatternRewriter &rewriter) const { + DenseMap indexMap; + SmallVector newInitArgs; + SmallVector newIterArgTypes; + SmallVector newInitArgsTemp; + SmallVector newIterArgTypesTemp; + + for (unsigned i = 0; i < forOp.getInitArgs().size(); ++i) { + PointerArgInfo *infoTemp = nullptr; + bool isMatch = false; + + // find the matching pointerArg + for (size_t k = 0; k < pointerArgs.size(); ++k) { + auto &info = pointerArgs[k]; + if (info.oldIndex == i) { + infoTemp = &info; + isMatch = true; + break; + } + } + + if (!isMatch) { + // Preserve original argument unchanged + newInitArgsTemp.push_back(forOp.getInitArgs()[i]); + newIterArgTypesTemp.push_back(forOp.getInitArgs()[i].getType()); + } else if (infoTemp != nullptr) { + // Replace pointer with integer offset (initialized to 0) + SmallVector newInitArgs; + SmallVector newIterArgTypes; + + if (infoTemp->offsetValues.empty()) { + Value zero = + rewriter.create(forOp.getLoc(), 0, 32); + newInitArgs.push_back(zero); + newIterArgTypes.push_back(rewriter.getIntegerType(INT_TYPE_BIT_WIDTH)); + newInitArgsTemp.push_back(zero); + newIterArgTypesTemp.push_back( + rewriter.getIntegerType(INT_TYPE_BIT_WIDTH)); + } else { + Value zero = + rewriter.create(forOp.getLoc(), 0, 32); + for (size_t j = 0; j < infoTemp->offsetValues.size(); ++j) { + Value offset = infoTemp->offsetValues[j]; + if (auto maskInt = getConstantIntValue(offset)) { + newInitArgs.push_back(zero); + newIterArgTypes.push_back( + rewriter.getIntegerType(INT_TYPE_BIT_WIDTH)); + if (*maskInt == 0) { + continue; + } + newInitArgsTemp.push_back(zero); + newIterArgTypesTemp.push_back( + rewriter.getIntegerType(INT_TYPE_BIT_WIDTH)); + } + } + } + + infoTemp->newInitArgs = newInitArgs; + infoTemp->newIterArgTypes = newIterArgTypes; + } else { + newInitArgsTemp.push_back(forOp.getInitArgs()[i]); + newIterArgTypesTemp.push_back(forOp.getInitArgs()[i].getType()); + } + // may change in future + indexMap[i] = i; + } + return {newInitArgsTemp, newIterArgTypesTemp, indexMap}; +} + +scf::ForOp PromotePointerIterArgsPattern::createNewForLoop( + scf::ForOp forOp, ArrayRef newInitArgs, + ArrayRef newIterArgTypes, PatternRewriter &rewriter) const { + return rewriter.create(forOp.getLoc(), forOp.getLowerBound(), + forOp.getUpperBound(), forOp.getStep(), + newInitArgs); +} + +LogicalResult PromotePointerIterArgsPattern::rewriteLoopBody( + scf::ForOp oldForOp, scf::ForOp newForOp, + SmallVector &pointerArgs, + DenseMap &indexMap, PatternRewriter &rewriter) const { + Block &oldBody = *oldForOp.getBody(); + Block &newBody = *newForOp.getBody(); + + rewriter.setInsertionPointToStart(&newBody); + + // Create IR mapping that maps original values to their transformed + // equivalents + IRMapping mapping = + createIRMapping(oldForOp, newForOp, pointerArgs, indexMap, rewriter); + // Clone instructions from original loop body, applying the mapping + return cloneInstructions(oldBody, newBody, pointerArgs, indexMap, mapping, + rewriter); +} + +LogicalResult PromotePointerIterArgsPattern::rewriteLoopBodyForAdvancePtr( + scf::ForOp oldForOp, scf::ForOp newForOp, + SmallVector &pointerArgs, + DenseMap &indexMap, PatternRewriter &rewriter) const { + Block &oldBody = *oldForOp.getBody(); + Block &newBody = *newForOp.getBody(); + rewriter.setInsertionPointToStart(&newBody); + // Create IR mapping that maps original values to their transformed + // equivalents + IRMapping mapping = createIRMappingForAdvancePtr( + oldForOp, newForOp, pointerArgs, indexMap, rewriter); + // Clone instructions from original loop body, applying the mapping + return cloneInstructionsForAdvancePtr(oldBody, newBody, pointerArgs, indexMap, + mapping, rewriter); +} + +IRMapping PromotePointerIterArgsPattern::createIRMapping( + scf::ForOp oldForOp, scf::ForOp newForOp, + SmallVector &pointerArgs, + DenseMap &indexMap, PatternRewriter &rewriter) const { + IRMapping mapping; + mapping.map(oldForOp.getInductionVar(), newForOp.getInductionVar()); + + // Process iteration arguments + for (unsigned i = 0; i < oldForOp.getRegionIterArgs().size(); ++i) { + Value oldIterArg = oldForOp.getRegionIterArgs()[i]; + Value newIterArg = newForOp.getRegionIterArgs()[indexMap[i]]; + + if (isPointerArgIndex(pointerArgs, i)) { + // Update the PointerArgInfo with the new integer iteration argument + for (auto &info : pointerArgs) { + if (info.oldIndex == i) { + info.newIterArg = newIterArg; + break; + } + } + + // Map original pointer argument to a reconstructed pointer + mapping.map(oldIterArg, + rebuildPointer(oldForOp, pointerArgs, i, rewriter)); + } else { + // Direct mapping for non-pointer arguments + mapping.map(oldIterArg, newIterArg); + } + } + return mapping; +} + +IRMapping PromotePointerIterArgsPattern::createIRMappingForAdvancePtr( + scf::ForOp oldForOp, scf::ForOp newForOp, + SmallVector &pointerArgs, + DenseMap &indexMap, PatternRewriter &rewriter) const { + IRMapping mapping; + mapping.map(oldForOp.getInductionVar(), newForOp.getInductionVar()); + // Process iteration arguments + for (unsigned i = 0; i < oldForOp.getRegionIterArgs().size(); ++i) { + Value oldIterArg = oldForOp.getRegionIterArgs()[i]; + Value newIterArg = newForOp.getRegionIterArgs()[indexMap[i]]; + if (isPointerArgIndex(pointerArgs, i)) { + // Update the PointerArgInfo with the new integer iteration argument + for (auto &info : pointerArgs) { + if (info.oldIndex == i) { + info.newIterArg = newIterArg; + break; + } + } + Value newIterArgTemp = + rebuildPointerForAdvancePtr(oldForOp, pointerArgs, i, rewriter); + // Map original pointer argument to a reconstructed pointer + mapping.map(oldIterArg, newIterArgTemp); + } else { + // Direct mapping for non-pointer arguments + mapping.map(oldIterArg, newIterArg); + } + } + return mapping; +} + +bool PromotePointerIterArgsPattern::isPointerArgIndex( + ArrayRef pointerArgs, unsigned idx) const { + for (auto &info : pointerArgs) { + if (info.oldIndex == idx) + return true; + } + return false; +} + +Value PromotePointerIterArgsPattern::rebuildPointer( + scf::ForOp forOp, ArrayRef pointerArgs, unsigned idx, + PatternRewriter &rewriter) const { + const PointerArgInfo *info = nullptr; + for (auto &argInfo : pointerArgs) { + if (argInfo.oldIndex == idx) { + info = &argInfo; + break; + } + } + if (!info) + return nullptr; + // Create splat operation to broadcast integer offset to tensor shape + auto baseType = info->basePointer.getType(); + Value splatOffset = nullptr; + if (auto rankedType = dyn_cast(baseType)) { + // Get the shape of the original tensor + auto shape = rankedType.getShape(); + + splatOffset = rewriter.create( + forOp.getLoc(), RankedTensorType::get(shape, rewriter.getI32Type()), + info->newIterArg); + } else { + return nullptr; + } + // Create addptr operation: base pointer + splatted offset + return rewriter.create(forOp.getLoc(), + info->basePointer.getType(), + info->basePointer, splatOffset); +} + +Value PromotePointerIterArgsPattern::rebuildPointerForAdvancePtr( + scf::ForOp forOp, ArrayRef pointerArgs, unsigned idx, + PatternRewriter &rewriter) const { + const PointerArgInfo *info = nullptr; + for (auto &argInfo : pointerArgs) { + if (argInfo.oldIndex == idx) { + info = &argInfo; + break; + } + } + if (!info) { + return nullptr; + } + SmallVector advanceOpOld; + advanceOpOld.push_back(info->basePointer); + for (size_t i = 0; i < info->offsetValues.size(); ++i) { + if (auto maskInt = getConstantIntValue(info->offsetValues[i])) { + if (*maskInt != 0) { + advanceOpOld.push_back(info->newIterArg); + continue; + } + } + advanceOpOld.push_back(info->offsetValues[i]); + } + auto advanceOp = rewriter.create( + forOp.getLoc(), info->basePointer.getType(), advanceOpOld); + return advanceOp; +} + +LogicalResult PromotePointerIterArgsPattern::cloneInstructions( + Block &oldBody, Block &newBody, ArrayRef pointerArgs, + DenseMap &indexMap, IRMapping &mapping, + PatternRewriter &rewriter) const { + // Collect all operations from the old loop body except the terminator + SmallVector toClone; + for (auto &op : oldBody.without_terminator()) { + toClone.push_back(&op); + } + // Build a set of addptr operations to skip (those that update pointer + // iteration arguments) + DenseSet addPtrOpsToSkip; + for (const auto &info : pointerArgs) { + if (info.addPtrValue) { + addPtrOpsToSkip.insert(info.addPtrValue); + } + } + // Clone all operations except the skipped addptr operations + for (auto *op : toClone) { + // Only skip addptr operations that are updating pointer iteration arguments + if (auto addPtrOp = dyn_cast(op)) { + if (addPtrOpsToSkip.contains(addPtrOp.getResult())) { + continue; + } + } + rewriter.clone(*op, mapping); + } + // Handle the yield terminator separately + auto yieldOp = dyn_cast(oldBody.getTerminator()); + if (!yieldOp) { + return failure(); + } + return cloneYieldOp(yieldOp, pointerArgs, indexMap, mapping, rewriter); +} + +LogicalResult PromotePointerIterArgsPattern::cloneInstructionsForAdvancePtr( + Block &oldBody, Block &newBody, ArrayRef pointerArgs, + DenseMap &indexMap, IRMapping &mapping, + PatternRewriter &rewriter) const { + // Collect all operations from the old loop body except the terminator + SmallVector toClone; + for (auto &op : oldBody.without_terminator()) { + toClone.push_back(&op); + } + // Build a set of addptr operations to skip (those that update pointer + // iteration arguments) + DenseSet advancePtrOpsToSkip; + for (const auto &info : pointerArgs) { + if (info.addPtrValue) { + advancePtrOpsToSkip.insert(info.addPtrValue); + } + } + // Clone all operations except the skipped addptr operations + for (auto *op : toClone) { + if (auto advancePtrOp = dyn_cast(op)) { + if (advancePtrOpsToSkip.contains(advancePtrOp.getResult())) { + continue; + } + } + auto clonedOp = rewriter.clone(*op, mapping); + } + // Handle the yield terminator separately + auto yieldOp = dyn_cast(oldBody.getTerminator()); + if (!yieldOp) { + return failure(); + } + return cloneYieldOpForAdvancePtr(yieldOp, pointerArgs, indexMap, mapping, + rewriter); +} + +LogicalResult PromotePointerIterArgsPattern::cloneYieldOpForAdvancePtr( + scf::YieldOp yieldOp, ArrayRef pointerArgs, + DenseMap &indexMap, IRMapping &mapping, + PatternRewriter &rewriter) const { + SmallVector newOperands; + // Process each operand of the original yield operation + for (unsigned i = 0; i < yieldOp.getNumOperands(); ++i) { + auto operand = yieldOp.getOperand(i); + if (isPointerArgIndex(pointerArgs, i)) { + // For pointer arguments being promoted: create integer addition + SmallVector intResult = + createOffsetsForAdvancePtr(i, pointerArgs, indexMap, rewriter); + for (size_t i = 0; i < intResult.size(); ++i) { + newOperands.push_back(intResult[i]); + } + } else { + // For other arguments: use the value from the IR mapping + Value mappedValue = mapping.lookupOrDefault(operand); + newOperands.push_back(mappedValue); + } + } + // Validate that all new operands are non-null + for (size_t i = 0; i < newOperands.size(); ++i) { + auto v = newOperands[i]; + if (!v) { + return failure(); + } + } + // Create the new yield operation in the transformed loop + auto newYieldOp = + rewriter.create(yieldOp.getLoc(), newOperands); + return success(); +} + +LogicalResult PromotePointerIterArgsPattern::cloneYieldOp( + scf::YieldOp yieldOp, ArrayRef pointerArgs, + DenseMap &indexMap, IRMapping &mapping, + PatternRewriter &rewriter) const { + SmallVector newOperands; + // Process each operand of the original yield operation + for (unsigned i = 0; i < yieldOp.getNumOperands(); ++i) { + if (isPointerArgIndex(pointerArgs, i)) { + // For pointer arguments being promoted: create integer addition + Value intResult = createIntegerAdd(i, pointerArgs, indexMap, rewriter); + newOperands.push_back(intResult); + } else { + // For other arguments: use the value from the IR mapping + newOperands.push_back(mapping.lookupOrDefault(yieldOp.getOperand(i))); + } + } + // Validate that all new operands are non-null + for (auto v : newOperands) { + if (!v) { + return failure(); + } + } + // Create the new yield operation in the transformed loop + rewriter.create(yieldOp.getLoc(), newOperands); + return success(); +} + +SmallVector PromotePointerIterArgsPattern::createOffsetsForAdvancePtr( + unsigned idx, ArrayRef pointerArgs, + DenseMap &indexMap, PatternRewriter &rewriter) const { + const PointerArgInfo *info = nullptr; + for (auto &argInfo : pointerArgs) { + if (argInfo.oldIndex == idx) { + info = &argInfo; + break; + } + } + if (!info) { + return SmallVector(); + } + + if (info->offsetValues.empty()) { + return SmallVector(); + } + SmallVector offsets; + // Process all offset values in the collection + for (size_t i = 0; i < info->offsetValues.size(); ++i) { + Value offset = info->offsetValues[i]; + Value newArg = info->newInitArgs[i]; + // Get a location for creating constants + Location loc = offset.getLoc(); + // Check if this offset is a constant + Attribute offsetAttr; + if (matchPattern(offset, m_Constant(&offsetAttr))) { + // Case 1: Integer attribute (scalar constant) + if (auto intAttr = dyn_cast(offsetAttr)) { + if (intAttr.getInt() == 0) { + continue; + } + + Value result = + rewriter.create(loc, info->newIterArg, offset); + offsets.push_back(result); + } + } + } + return offsets; +} + +Value PromotePointerIterArgsPattern::createIntegerAdd( + unsigned idx, ArrayRef pointerArgs, + DenseMap &indexMap, PatternRewriter &rewriter) const { + const PointerArgInfo *info = nullptr; + for (auto &argInfo : pointerArgs) { + if (argInfo.oldIndex == idx) { + info = &argInfo; + break; + } + } + if (!info) + return nullptr; + // Try to extract constant offset value + Attribute offsetAttr; + if (matchPattern(info->offsetValue, m_Constant(&offsetAttr))) { + Location loc = info->offsetValue.getLoc(); + // Case 1: Integer attribute (scalar constant) + if (auto intAttr = dyn_cast(offsetAttr)) { + Value constOffset = rewriter.create( + loc, intAttr.getInt(), INT_TYPE_BIT_WIDTH); + return rewriter.create(loc, info->newIterArg, constOffset); + } + // Case 2: DenseElementsAttr (tensor constant) + if (auto denseAttr = dyn_cast(offsetAttr)) { + // Check if it's a splat (all elements are the same) + if (denseAttr.isSplat()) { + // For integer-type DenseElementsAttr + if (denseAttr.getElementType().isInteger(INT_TYPE_BIT_WIDTH)) { + auto splatValue = denseAttr.getSplatValue(); + Value constOffset = rewriter.create( + loc, splatValue.getZExtValue(), INT_TYPE_BIT_WIDTH); + return rewriter.create(loc, info->newIterArg, + constOffset); + } + } else { + // If not a splat, but has only one element, we can still handle it + if (denseAttr.getNumElements() == 1) { + auto firstElement = *denseAttr.getValues().begin(); + Value constOffset = rewriter.create( + loc, firstElement.getZExtValue(), INT_TYPE_BIT_WIDTH); + return rewriter.create(loc, info->newIterArg, + constOffset); + } + } + } + } + + // Return nullptr if offset is not a constant (pattern only handles constant + // offsets) + return nullptr; +} + +LogicalResult PromotePointerIterArgsPattern::replaceResults( + scf::ForOp oldForOp, scf::ForOp newForOp, + ArrayRef pointerArgs, + DenseMap &indexMap, PatternRewriter &rewriter) const { + SmallVector newResults; + + for (unsigned i = 0; i < oldForOp.getNumResults(); ++i) { + if (isPointerArgIndex(pointerArgs, i)) { + Value ptrResult = reconstructPointer( + oldForOp, i, newForOp.getResult(indexMap[i]), pointerArgs, rewriter); + newResults.push_back(ptrResult); + } else { + newResults.push_back(newForOp.getResult(indexMap[i])); + } + } + + for (auto v : newResults) { + if (!v) { + return failure(); + } + } + rewriter.replaceOp(oldForOp, newResults); + return success(); +} + +SmallVector PromotePointerIterArgsPattern::reconstructPointerForAdvance( + scf::ForOp forOp, unsigned idx, Value intResult, + ArrayRef pointerArgs, PatternRewriter &rewriter) const { + const PointerArgInfo *info = nullptr; + for (auto &argInfo : pointerArgs) { + if (argInfo.oldIndex == idx) { + info = &argInfo; + break; + } + } + if (!info) { + return SmallVector(); + } + // Create splat operation to broadcast integer result to tensor shape + auto baseType = info->basePointer.getType(); + SmallVector addPtr; + // Create a tensor with the same shape, where all elements are the integer + // result + SmallVector newInitArgs = info->newInitArgs; + for (size_t i = 0; i < newInitArgs.size(); ++i) { + if (auto maskInt = getConstantIntValue(newInitArgs[i])) { + if (*maskInt == 0) { + continue; + } + Value constOffset = rewriter.create( + forOp.getLoc(), *maskInt, INT_TYPE_BIT_WIDTH); + Value addPtrResult = rewriter.create( + forOp.getLoc(), newInitArgs[i], constOffset); + addPtr.push_back(addPtrResult); + } + } + return addPtr; +} + +Value PromotePointerIterArgsPattern::reconstructPointer( + scf::ForOp forOp, unsigned idx, Value intResult, + ArrayRef pointerArgs, PatternRewriter &rewriter) const { + const PointerArgInfo *info = nullptr; + for (auto &argInfo : pointerArgs) { + if (argInfo.oldIndex == idx) { + info = &argInfo; + break; + } + } + if (!info) + return nullptr; + + // Create splat operation to broadcast integer result to tensor shape + auto baseType = info->basePointer.getType(); + Value splatOffset = nullptr; + if (auto rankedType = dyn_cast(baseType)) { + // Get the shape of the original tensor + auto shape = rankedType.getShape(); + + splatOffset = rewriter.create( + forOp.getLoc(), RankedTensorType::get(shape, rewriter.getI32Type()), + intResult); + } else { + return nullptr; + } + + // Create a tensor with the same shape, where all elements are the integer + // result + return rewriter.create(forOp.getLoc(), + info->basePointer.getType(), + info->basePointer, splatOffset); +} + +void SimplifyTensorIterArgsPattern::ShapeChainInfo::dump() const { + LLVM_DEBUG({ + llvm::dbgs() << "ShapeChainInfo:\n"; + llvm::dbgs() << " base: " << base << "\n"; + llvm::dbgs() << " chain:\n"; + for (Operation *op : chain) { + llvm::dbgs() << " " << *op << "\n"; + } + }); +} + +void SimplifyTensorIterArgsPattern::CandidateInfo::dump() const { + LLVM_DEBUG({ + llvm::dbgs() << "CandidateInfo:\n"; + llvm::dbgs() << " idx: " << idx << "\n"; + llvm::dbgs() << " ShapeInfo: \n"; + shapeInfo.dump(); + for (Operation *op : arithOps) { + llvm::dbgs() << " ArithOp: " << *op << "\n"; + } + }); +} + +Value SimplifyTensorIterArgsPattern::cloneShapeChain( + Location loc, Value base, ArrayRef chain, + PatternRewriter &rewriter) const { + Value cur = base; + for (Operation *op : chain) { + if (auto splat = dyn_cast(op)) { + auto dstTy = cast(splat.getType()); + cur = rewriter.create(loc, dstTy, cur); + continue; + } + if (auto bcast = dyn_cast(op)) { + auto dstTy = cast(bcast.getType()); + cur = rewriter.create(loc, dstTy, cur); + continue; + } + if (auto expand = dyn_cast(op)) { + auto dstTy = cast(expand.getType()); + cur = rewriter.create(loc, dstTy, cur, + expand.getAxis()); + continue; + } + return Value(); + } + return cur; +} + +bool SimplifyTensorIterArgsPattern::isBlockArgumentFromAnotherForLoop( + Value v) const { + auto barg = dyn_cast(v); + if (!barg) { + return false; + } + + Block *owner = barg.getOwner(); + if (!owner) { + return false; + } + + auto parentFor = dyn_cast_or_null(owner->getParentOp()); + if (!parentFor) { + return false; + } + + // Only handle block args that belong to the body block of scf.for. + if (&parentFor.getRegion().front() != owner) { + return false; + } + + // scf.for body block args layout: + // arg#0 : induction variable + // arg#1..arg#N : region iter args + unsigned argNo = barg.getArgNumber(); + if (argNo == 0) { + // IV: not an iter arg init source. + return false; + } + + unsigned iterIdx = argNo - 1; + if (iterIdx >= parentFor.getInitArgs().size()) { + return false; + } + + return true; +} + +Value SimplifyTensorIterArgsPattern::normalizeInitArgForShapePeel( + Value v) const { + if (!isBlockArgumentFromAnotherForLoop(v)) { + // Not a block argument from another for loop, return as is. + return v; + } + auto barg = dyn_cast(v); + if (!barg) { + return v; + } + + Block *owner = barg.getOwner(); + if (!owner) { + return v; + } + + auto parentFor = dyn_cast_or_null(owner->getParentOp()); + if (!parentFor) { + return v; + } + + unsigned argNo = barg.getArgNumber(); + if (argNo == 0) { + return v; + } + + unsigned iterIdx = argNo - 1; + if (iterIdx >= parentFor.getInitArgs().size()) { + return v; + } + // For the simple nested relay case: + // inner.initArg = outer.regionIterArg ==> resolve to outer.initArg + return parentFor.getInitArgs()[iterIdx]; +} + +std::optional +SimplifyTensorIterArgsPattern::getRelayMapM1(scf::ForOp innerFor, + scf::ForOp outerFor, + unsigned innerIdx) const { + if (innerIdx >= innerFor.getInitArgs().size() || + innerIdx >= innerFor.getNumResults()) { + return std::nullopt; + } + + // 1) inner initArg -> outer iterArg idx + Value innerInit = innerFor.getInitArgs()[innerIdx]; + if (!isBlockArgumentFromAnotherForLoop(innerInit)) { + return std::nullopt; + } + + auto barg = dyn_cast(innerInit); + if (!barg) { + return std::nullopt; + } + + unsigned outerInitIdx = barg.getArgNumber() - 1; + if (outerInitIdx >= outerFor.getRegionIterArgs().size()) { + return std::nullopt; + } + + // NEW: ensure outer iterArg lane is used only as the mapped inner initArg. + // i.e. no extra users in outer body. + Value outerIterArg = outerFor.getRegionIterArgs()[outerInitIdx]; + if (!outerIterArg.hasOneUse()) { + return std::nullopt; + } + OpOperand &onlyUse = *outerIterArg.getUses().begin(); + if (onlyUse.getOwner() != innerFor.getOperation() || + onlyUse.get() != innerFor.getInitArgs()[innerIdx]) { + return std::nullopt; + } + + // 2) inner result -> outer yield slot + if (!outerFor.getBody() || !outerFor.getBody()->mightHaveTerminator()) { + return std::nullopt; + } + auto outerYield = dyn_cast(outerFor.getBody()->getTerminator()); + if (!outerYield) { + return std::nullopt; + } + + // only handle the case where inner result directly feeds into outer yield + // operand, without any intermediate use/def. + Value innerRes = innerFor.getResult(innerIdx); + std::optional outerYieldIdx; + for (unsigned k = 0; k < outerYield.getNumOperands(); ++k) { + if (outerYield.getOperand(k) == innerRes) { + outerYieldIdx = k; + break; + } + } + if (!outerYieldIdx.has_value()) { + return std::nullopt; + } + + if (outerInitIdx != *outerYieldIdx) { + // The position of iterArg and yield operand should be the same in this + // simple relay case. + return std::nullopt; + } + + return RelayMapM1{ + .innerIdx = innerIdx, + .outerInitIdx = outerInitIdx, + .outerYieldIdx = *outerYieldIdx, + }; +} + +void SimplifyTensorIterArgsPattern::splitCandidatesByRelay( + SmallVector all, + SmallVector &locals, + SmallVector &relays) const { + for (const auto &c : all) { + if (c.relayMap.has_value()) { + relays.push_back(c); + } else { + locals.push_back(c); + } + } +} + +std::optional +SimplifyTensorIterArgsPattern::peelShapeChain(Value v) const { + ShapeChainInfo info; + Value cur = v; + SmallVector rev; + while (Operation *def = cur.getDefiningOp()) { + if (isa(def)) { + rev.push_back(def); + cur = def->getOperand(0); + continue; + } + break; + } + if (rev.empty()) { + return std::nullopt; + } + std::reverse(rev.begin(), rev.end()); + info.base = cur; + info.chain = std::move(rev); + return info; +} + +bool SimplifyTensorIterArgsPattern::isArithWithConst(Operation *op, + Value curVal, + Value &nextVal, + Value &constVal) const { + nextVal = Value(); + constVal = Value(); + + Value lhs; + Value rhs; + if (!extractBinaryArithOperands(op, lhs, rhs)) { + return false; + } + + bool lhsIsConst = matchPattern(lhs, m_Constant()); + bool rhsIsConst = matchPattern(rhs, m_Constant()); + if (lhsIsConst == rhsIsConst) { + return false; + } + + if (lhs == curVal && rhsIsConst) { + nextVal = op->getResult(0); + constVal = rhs; + return true; + } + if (rhs == curVal && lhsIsConst) { + nextVal = op->getResult(0); + constVal = lhs; + return true; + } + + return false; +} + +Value SimplifyTensorIterArgsPattern::getNewConstLikeOperand( + Value cst, Type targetTy, PatternRewriter &rewriter) const { + Attribute attr; + if (!matchPattern(cst, m_Constant(&attr))) { + return Value(); + } + auto tensorTy = dyn_cast(targetTy); + if (!tensorTy) { + return Value(); + } + // Only support arith.constant dense right now. + auto dense = dyn_cast(attr); + if (!dense || !dense.isSplat()) { + return Value(); + } + auto splatAttr = + DenseElementsAttr::get(tensorTy, dense.getSplatValue()); + return rewriter.create(cst.getLoc(), tensorTy, splatAttr); +} + +bool SimplifyTensorIterArgsPattern::canBuildConstLikeOperand( + Value cst, Type targetTy) const { + Attribute attr; + if (!matchPattern(cst, m_Constant(&attr))) { + return false; + } + auto tensorTy = dyn_cast(targetTy); + if (!tensorTy) { + return false; + } + auto dense = dyn_cast(attr); + return dense && dense.isSplat(); +} + +LogicalResult SimplifyTensorIterArgsPattern::collectReverseLinearYieldPath( + Value yielded, Value iterArg, + SmallVectorImpl &opsInExecOrder) const { + SmallVector revOps; + Value cur = yielded; + + while (cur != iterArg) { + Operation *def = cur.getDefiningOp(); + if (!def || isa(def)) { + return failure(); + } + + Value lhs; + Value rhs; + if (!extractBinaryArithOperands(def, lhs, rhs)) { + return failure(); + } + + bool lhsIsConst = matchPattern(lhs, m_Constant()); + bool rhsIsConst = matchPattern(rhs, m_Constant()); + if (lhsIsConst == rhsIsConst) { + return failure(); + } + + Value upstream = lhsIsConst ? rhs : lhs; + if (!upstream) { + return failure(); + } + + // Safety: only accept a strict linear chain. + // The current def result must be consumed exclusively by `cur`. + // If there is any extra user, rewriting this lane may break semantics. + if (def->getNumResults() != 1) { + return failure(); + } + Value defRes = def->getResult(0); + if (!defRes.hasOneUse()) { + return failure(); + } + OpOperand &onlyUse = *defRes.getUses().begin(); + if (onlyUse.get() != cur) { + return failure(); + } + + revOps.push_back(def); + cur = upstream; + } + + opsInExecOrder.assign(revOps.rbegin(), revOps.rend()); + return success(); +} + +LogicalResult SimplifyTensorIterArgsPattern::matchAndRewrite( + scf::ForOp forOp, PatternRewriter &rewriter) const { + LLVM_DEBUG({ + llvm::dbgs() << "Now Handling For Op: \n"; + forOp.dump(); + }); + + // if you want only simplify iter args once to avoid infinite pattern + // application, return failure when meeting done label + if (forOp->hasAttr(kSimplifiedAttr)) { + LLVM_DEBUG({ + llvm::dbgs() << "This For Op has been simplified before.\n"; + forOp.dump(); + }); + } + // If the forOp has been marked as failed before, it means we attempted to + // simplify it but couldn't, so we should not try again. + if (forOp->hasAttr(kFailedAttr) || forOp->hasAttr(kIncompleteAttr)) { + return failure(); + } + + Block &oldBody = *forOp.getBody(); + if (!oldBody.mightHaveTerminator()) { + return failure(); + } + auto oldYield = dyn_cast(oldBody.getTerminator()); + if (!oldYield) { + return failure(); + } + + SmallVector candidates; + auto regionIterArgs = forOp.getRegionIterArgs(); + auto initArgs = forOp.getInitArgs(); + + for (unsigned i = 0; i < regionIterArgs.size(); ++i) { + Value iterArg = regionIterArgs[i]; + Value initArg = initArgs[i]; + Value yielded = oldYield.getOperand(i); + + auto iterTy = dyn_cast(iterArg.getType()); + if (!iterTy || !iterTy.hasStaticShape()) { + continue; + } + + // if initArg comes from another iterArg of an outer loop, get the ultimate + // source initArg. This handles the common nested relay pattern where inner + // loop iter arg is directly yielded from outer loop iter arg without + // modification. multiple for-loop levels of relay may require more complex + // data flow analysis to resolve the ultimate source initArg. + Value normalizedInitArg = normalizeInitArgForShapePeel(initArg); + + auto shapeInfoOpt = peelShapeChain(normalizedInitArg); + if (!shapeInfoOpt.has_value()) { + continue; + } + auto shapeInfo = *shapeInfoOpt; + + SmallVector chainOps; + if (failed(collectReverseLinearYieldPath(yielded, iterArg, chainOps))) { + continue; + } + + std::optional relayMap; + if (auto outerFor = dyn_cast_or_null(forOp->getParentOp())) { + relayMap = getRelayMapM1(forOp, outerFor, i); + } + + if (!relayMap.has_value() && isBlockArgumentFromAnotherForLoop(initArg)) { + LLVM_DEBUG({ + llvm::dbgs() + << "Init arg is a block argument from another for loop, but does " + "not form a simple relay pattern. Skipping candidate.\n"; + }); + continue; + } + + candidates.push_back(CandidateInfo{ + .idx = i, + .shapeInfo = std::move(shapeInfo), + .arithOps = std::move(chainOps), + .relayMap = relayMap, + }); + } + + if (candidates.empty()) { + return failure(); + } + + // Pre-validate all candidate chains before mutating IR. + // This avoids creating partial new IR and then returning failure(). + for (auto &c : candidates) { + auto baseTensorTy = dyn_cast(c.shapeInfo.base.getType()); + if (!baseTensorTy) { + return failure(); + } + + Value iterCur = regionIterArgs[c.idx]; + for (Operation *oldOp : c.arithOps) { + Value nextVal, oldConst; + if (!isArithWithConst(oldOp, iterCur, nextVal, oldConst)) { + return failure(); + } + if (!canBuildConstLikeOperand(oldConst, baseTensorTy)) { + return failure(); + } + iterCur = nextVal; + } + } + + if (!isSafeToRewriteLanesByResultUses(forOp, candidates)) { + return failure(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "Now Simplify For Op: \n"; + forOp.dump(); + llvm::dbgs() << "Found " << candidates.size() + << " candidate iter args to simplify in for loop at " + << forOp.getLoc() << "\n"; + for (auto &c : candidates) { + c.dump(); + } + }); + + SmallVector localCandidates; + SmallVector relayCandidates; + splitCandidatesByRelay(candidates, localCandidates, relayCandidates); + + FailureOr newForRes = rewriteForWithLocalCandidates( + forOp, localCandidates, /*outerCaptureMap=*/nullptr, rewriter); + if (failed(newForRes)) { + return failure(); + } + auto newFor = *newForRes; + + // case A: has relay -> continue relay pipeline + // If there are relay candidates, we need to rewrite the innerFor and outerFor + // together to maintain the relay relationship. The innerFor needs to be + // rewritten because the iterArg shape is changed after peeling, and the + // outerFor needs to be rewritten together to maintain the relay relationship + // between inner and outer iter args. + if (!relayCandidates.empty()) { + LLVM_DEBUG({ + llvm::dbgs() << "Nested For-loop args need to be rewritten to maintain " + "relay relationship\n"; + }); + scf::ForOp oldFor = forOp; + if (failed(rewriteForWithRelayCandidates(newFor, oldFor, relayCandidates, + rewriter))) { + newFor->setAttr(kFailedAttr, rewriter.getUnitAttr()); + return failure(); + } + return success(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "Only local candidates found, committing local rewrite\n"; + }); + // case B: no relay -> commit local rewrite now + if (newFor->hasAttr(kIncompleteAttr)) { + newFor->removeAttr(kIncompleteAttr); + } + newFor->setAttr(kSimplifiedAttr, rewriter.getUnitAttr()); + LLVM_DEBUG({ + llvm::dbgs() << "Successfully rewrote ForOp to: \n"; + newFor.dump(); + }); + rewriter.replaceOp(forOp, newFor.getResults()); + return success(); +} + +bool SimplifyTensorIterArgsPattern::extractBinaryArithOperands( + Operation *op, Value &lhs, Value &rhs) const { + if (auto v = dyn_cast(op)) { + lhs = v.getLhs(); + rhs = v.getRhs(); + return true; + } + if (auto v = dyn_cast(op)) { + lhs = v.getLhs(); + rhs = v.getRhs(); + return true; + } + if (auto v = dyn_cast(op)) { + lhs = v.getLhs(); + rhs = v.getRhs(); + return true; + } + if (auto v = dyn_cast(op)) { + lhs = v.getLhs(); + rhs = v.getRhs(); + return true; + } + + if (auto v = dyn_cast(op)) { + lhs = v.getLhs(); + rhs = v.getRhs(); + return true; + } + if (auto v = dyn_cast(op)) { + lhs = v.getLhs(); + rhs = v.getRhs(); + return true; + } + if (auto v = dyn_cast(op)) { + lhs = v.getLhs(); + rhs = v.getRhs(); + return true; + } + if (auto v = dyn_cast(op)) { + lhs = v.getLhs(); + rhs = v.getRhs(); + return true; + } + + return false; +} + +Value SimplifyTensorIterArgsPattern::createSameBinaryArithOp( + Operation *oldOp, Location loc, Value lhs, Value rhs, + PatternRewriter &rewriter) const { + if (isa(oldOp)) + return rewriter.create(loc, lhs, rhs).getResult(); + if (isa(oldOp)) + return rewriter.create(loc, lhs, rhs).getResult(); + if (isa(oldOp)) + return rewriter.create(loc, lhs, rhs).getResult(); + if (isa(oldOp)) + return rewriter.create(loc, lhs, rhs).getResult(); + + if (isa(oldOp)) + return rewriter.create(loc, lhs, rhs).getResult(); + if (isa(oldOp)) + return rewriter.create(loc, lhs, rhs).getResult(); + if (isa(oldOp)) + return rewriter.create(loc, lhs, rhs).getResult(); + if (isa(oldOp)) + return rewriter.create(loc, lhs, rhs).getResult(); + + return Value(); +} + +FailureOr +SimplifyTensorIterArgsPattern::rewriteForWithLocalCandidates( + scf::ForOp forOp, + ArrayRef candidates, + const IRMapping *outerCaptureMap, PatternRewriter &rewriter) const { + // This function is used to rewrite the forOp with local candidate. + // The caller should have already validated that the candidate's shape chain + // can be peeled and the arithmetic chain can be rebuilt with constants. + if (candidates.empty()) { + return forOp; + } + + Block &oldBody = *forOp.getBody(); + if (!oldBody.mightHaveTerminator()) { + return failure(); + } + auto oldYield = dyn_cast(oldBody.getTerminator()); + if (!oldYield) { + return failure(); + } + + auto regionIterArgs = forOp.getRegionIterArgs(); + auto initArgs = forOp.getInitArgs(); + + DenseMap candMap; + for (auto &c : candidates) { + candMap[c.idx] = &c; + } + + DenseMap arithOpToCandIdx; + DenseSet candidateArithSet; + for (auto &c : candidates) { + for (Operation *op : c.arithOps) { + candidateArithSet.insert(op); + arithOpToCandIdx[op] = c.idx; + } + } + + // Build new init args (candidate iter arg uses base) + SmallVector newInitArgs; + newInitArgs.reserve(initArgs.size()); + for (unsigned i = 0; i < initArgs.size(); ++i) { + if (auto it = candMap.find(i); it != candMap.end()) { + newInitArgs.push_back(it->second->shapeInfo.base); + continue; + } + Value v = initArgs[i]; + if (outerCaptureMap) { + if (Value mapped = outerCaptureMap->lookupOrNull(v)) { + v = mapped; + } + } + newInitArgs.push_back(v); + } + + auto newFor = rewriter.create( + forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), + forOp.getStep(), newInitArgs); + newFor->setAttr(kIncompleteAttr, rewriter.getUnitAttr()); + + auto failAfterCreate = [&]() -> LogicalResult { + newFor->setAttr(kFailedAttr, rewriter.getUnitAttr()); + return failure(); + }; + + IRMapping mapping; + mapping.map(forOp.getInductionVar(), newFor.getInductionVar()); + + // map region args + for (unsigned i = 0; i < regionIterArgs.size(); ++i) { + Value oldArg = regionIterArgs[i]; + Value newArg = newFor.getRegionIterArgs()[i]; + mapping.map(oldArg, newArg); + } + + rewriter.setInsertionPointToStart(newFor.getBody()); + + // Materialize candidate iter-shape values and remap old iterArg semantic + // values. + for (auto &c : candidates) { + Value baseArg = newFor.getRegionIterArgs()[c.idx]; + Value mat = + cloneShapeChain(forOp.getLoc(), baseArg, c.shapeInfo.chain, rewriter); + if (!mat) { + return failAfterCreate(); + } + mapping.map(regionIterArgs[c.idx], mat); + } + + // Track current base-domain running value for each candidate lane. + DenseMap baseCurByCand; + for (auto &c : candidates) { + baseCurByCand[c.idx] = newFor.getRegionIterArgs()[c.idx]; + } + + // Replay old body in-order: + // - candidate arith op: rebuild now and map old result -> new result + // - other op: clone with mapping + for (Operation &op : oldBody.without_terminator()) { + if (candidateArithSet.contains(&op)) { + auto itIdx = arithOpToCandIdx.find(&op); + if (itIdx == arithOpToCandIdx.end()) { + return failAfterCreate(); + } + unsigned candIdx = itIdx->second; + + Value baseCur = baseCurByCand[candIdx]; + auto baseTensorTy = dyn_cast(baseCur.getType()); + if (!baseTensorTy) { + return failAfterCreate(); + } + + Value lhs; + Value rhs; + if (!extractBinaryArithOperands(&op, lhs, rhs)) { + return failAfterCreate(); + } + + bool lhsIsCst = matchPattern(lhs, m_Constant()); + bool rhsIsCst = matchPattern(rhs, m_Constant()); + if (lhsIsCst == rhsIsCst) { + return failAfterCreate(); + } + + Value oldConst = lhsIsCst ? lhs : rhs; + Value baseConst = + getNewConstLikeOperand(oldConst, baseTensorTy, rewriter); + if (!baseConst) { + return failAfterCreate(); + } + + Value newRes = createSameBinaryArithOp(&op, op.getLoc(), baseCur, + baseConst, rewriter); + if (!newRes) { + return failAfterCreate(); + } + + mapping.map(op.getResult(0), newRes); + baseCurByCand[candIdx] = newRes; + continue; + } + + // Non-candidate op: clone as-is with mapping. + if (outerCaptureMap) { + for (Value operand : op.getOperands()) { + if (mapping.lookupOrNull(operand)) { + continue; // already mapped (iv/iterArg/local rebuilt values) + } + if (Value mappedOuter = outerCaptureMap->lookupOrNull(operand)) { + mapping.map(operand, mappedOuter); + } + } + } + rewriter.clone(op, mapping); + } + + // Rebuild yield. + SmallVector newYieldOperands; + newYieldOperands.reserve(oldYield.getNumOperands()); + for (unsigned i = 0; i < oldYield.getNumOperands(); ++i) { + if (candMap.count(i)) { + auto it = baseCurByCand.find(i); + if (it == baseCurByCand.end() || !it->second) { + return failAfterCreate(); + } + newYieldOperands.push_back(it->second); + } else { + Value mapped = mapping.lookupOrDefault(oldYield.getOperand(i)); + if (!mapped) { + return failAfterCreate(); + } + newYieldOperands.push_back(mapped); + } + } + rewriter.create(oldYield.getLoc(), newYieldOperands); + if (newFor->hasAttr(kIncompleteAttr)) { + newFor->removeAttr(kIncompleteAttr); + } + newFor->setAttr(kSimplifiedAttr, rewriter.getUnitAttr()); + return newFor; +} + +bool SimplifyTensorIterArgsPattern::isSafeToRewriteLanesByResultUses( + scf::ForOp forOp, + ArrayRef candidates) const { + if (candidates.empty()) { + return false; + } + + SmallVector laneIdxs; + laneIdxs.reserve(candidates.size()); + for (const auto &c : candidates) { + laneIdxs.push_back(c.idx); + } + + auto parentOuter = dyn_cast_or_null(forOp->getParentOp()); + if (!parentOuter || !parentOuter.getBody() || + !parentOuter.getBody()->mightHaveTerminator()) { + // No parent loop or parent loop body is malformed, be conservative and + // return false to avoid unsafe rewrite. + return false; + } + auto parentOuterYield = + parentOuter + ? dyn_cast(parentOuter.getBody()->getTerminator()) + : nullptr; + + for (const auto &c : candidates) { + unsigned idx = c.idx; + if (idx >= forOp.getNumResults()) { + return false; + } + + Value r = forOp.getResult(idx); + if (r.use_empty()) { + continue; + } + + // relay case: the rewritten lane can only be used by the corresponding + // outer yield operand in the parent loop, and cannot have any other uses. + if (c.relayMap.has_value()) { + if (!parentOuter || !parentOuterYield) { + return false; + } + unsigned outerYieldIdx = c.relayMap->outerYieldIdx; + if (outerYieldIdx >= parentOuterYield.getNumOperands()) { + return false; + } + + for (OpOperand &use : r.getUses()) { + auto y = dyn_cast(use.getOwner()); + if (!y || y != parentOuterYield) { + return false; + } + if (use.getOperandNumber() != outerYieldIdx) { + return false; + } + } + + // outer result cannot have any uses outside of the parent loop's yield + // operand + unsigned outerLane = c.relayMap->outerYieldIdx; + if (outerLane >= parentOuter.getNumResults()) { + return false; + } + if (!parentOuter.getResult(outerLane).use_empty()) { + return false; + } + + continue; + } + + // local case: the rewritten lane cannot have any uses outside of the forOp + // results + return false; + } + + return true; +} + +LogicalResult SimplifyTensorIterArgsPattern::precheckRelayCandidates( + scf::ForOp innerFor, + ArrayRef relayCandidates, + scf::ForOp &outerForOut) const { + if (relayCandidates.empty()) { + return failure(); + } + + auto outerFor = dyn_cast_or_null(innerFor->getParentOp()); + if (!outerFor) { + return failure(); + } + + DenseSet usedOuterInitIdx; + DenseSet usedOuterYieldIdx; + if (!outerFor.getBody() || !outerFor.getBody()->mightHaveTerminator()) { + return failure(); + } + auto outerYield = dyn_cast(outerFor.getBody()->getTerminator()); + if (!outerYield) { + return failure(); + } + + for (const auto &c : relayCandidates) { + if (!c.relayMap.has_value()) { + return failure(); + } + const auto &m = *c.relayMap; + + // must be the same level mapping for the current innerFor + if (m.innerIdx != c.idx) { + return failure(); + } + + // M1 strict lane + if (m.outerInitIdx != m.outerYieldIdx) { + return failure(); + } + + if (m.outerInitIdx >= outerFor.getInitArgs().size()) { + return failure(); + } + if (m.outerYieldIdx >= outerYield.getNumOperands()) { + return failure(); + } + + // Conflict check: outer lane cannot be occupied multiple times + // different inner iter args mapped to the same outer iter arg or yield + // operand is not supported in this simple relay case, as it may require + // more complex data flow analysis to resolve the ultimate source init arg + // and final yield operand. + if (!usedOuterInitIdx.insert(m.outerInitIdx).second) { + return failure(); + } + if (!usedOuterYieldIdx.insert(m.outerYieldIdx).second) { + return failure(); + } + } + + outerForOut = outerFor; + return success(); +} + +FailureOr +SimplifyTensorIterArgsPattern::rewriteInnerForWithRelayCandidates( + scf::ForOp innerFor, + ArrayRef relayCandidates, + const IRMapping *outerCaptureMap, PatternRewriter &rewriter) const { + if (!innerFor || relayCandidates.empty()) { + return failure(); + } + + // Reuse local-like rebuild on the already-correct inner init args. + // IMPORTANT: For relay mode, caller must ensure innerFor init args on relay + // lanes are already wired to new outer iter args. + return rewriteForWithLocalCandidates(innerFor, relayCandidates, + outerCaptureMap, rewriter); +} + +FailureOr +SimplifyTensorIterArgsPattern::rewriteOuterForWithRelayCandidates( + scf::ForOp innerFor, scf::ForOp oldInnerFor, scf::ForOp outerFor, + ArrayRef relayCandidates, + PatternRewriter &rewriter) const { + if (!outerFor || !innerFor || relayCandidates.empty()) { + return failure(); + } + + constexpr llvm::StringLiteral kFailedAttr = + "tts.simplify_tensor_iter_args.failed"; + + Block &oldOuterBody = *outerFor.getBody(); + if (!oldOuterBody.mightHaveTerminator()) { + return failure(); + } + auto oldOuterYield = dyn_cast(oldOuterBody.getTerminator()); + if (!oldOuterYield) { + return failure(); + } + + // Build new outer init args: relay lanes switch to base. + SmallVector newOuterInitArgs(outerFor.getInitArgs().begin(), + outerFor.getInitArgs().end()); + for (const auto &c : relayCandidates) { + if (!c.relayMap.has_value()) { + return failure(); + } + const auto &m = *c.relayMap; + if (m.outerInitIdx >= newOuterInitArgs.size()) { + return failure(); + } + newOuterInitArgs[m.outerInitIdx] = c.shapeInfo.base; + } + + rewriter.setInsertionPoint(outerFor); + auto newOuterFor = rewriter.create( + outerFor.getLoc(), outerFor.getLowerBound(), outerFor.getUpperBound(), + outerFor.getStep(), newOuterInitArgs); + newOuterFor->setAttr(kIncompleteAttr, rewriter.getUnitAttr()); + + auto failAfterCreate = [&]() -> FailureOr { + newOuterFor->setAttr(kFailedAttr, rewriter.getUnitAttr()); + return failure(); + }; + + IRMapping outerMap; + outerMap.map(outerFor.getInductionVar(), newOuterFor.getInductionVar()); + for (unsigned i = 0; i < outerFor.getRegionIterArgs().size(); ++i) { + outerMap.map(outerFor.getRegionIterArgs()[i], + newOuterFor.getRegionIterArgs()[i]); + } + + rewriter.setInsertionPointToStart(newOuterFor.getBody()); + + scf::ForOp rewrittenInnerInNewOuter; + + // Clone outer body with special handling for old inner. + for (Operation &op : oldOuterBody.without_terminator()) { + // Fast path: not the anchor inner op we want to rebuild. + if (&op != innerFor.getOperation()) { + if (&op == oldInnerFor.getOperation()) { + // Drop old inner only when it is different from anchor inner. + continue; + } + rewriter.clone(op, outerMap); + continue; + } + + // Build relay-wired inner init args: + // relay lane: inner init <- new outer iterArg[outerInitIdx] + // non-relay lane: mapped old inner init + SmallVector relayWiredInnerInitArgs; + relayWiredInnerInitArgs.reserve(innerFor.getInitArgs().size()); + + DenseMap innerIdxToOuterInitIdx; + for (const auto &c : relayCandidates) { + const auto &m = *c.relayMap; + innerIdxToOuterInitIdx[m.innerIdx] = m.outerInitIdx; + } + + for (unsigned i = 0; i < innerFor.getInitArgs().size(); ++i) { + auto it = innerIdxToOuterInitIdx.find(i); + if (it != innerIdxToOuterInitIdx.end()) { + unsigned outerInitIdx = it->second; + if (outerInitIdx >= newOuterFor.getRegionIterArgs().size()) { + return failAfterCreate(); + } + relayWiredInnerInitArgs.push_back( + newOuterFor.getRegionIterArgs()[outerInitIdx]); + } else { + relayWiredInnerInitArgs.push_back( + outerMap.lookupOrDefault(innerFor.getInitArgs()[i])); + } + } + + // Build inner relay candidates whose base comes from new outer iter args. + SmallVector relayCandidatesForInner; + relayCandidatesForInner.reserve(relayCandidates.size()); + for (const auto &c : relayCandidates) { + CandidateInfo cc = c; // copy + + if (!cc.relayMap.has_value()) { + return failAfterCreate(); + } + unsigned outerInitIdx = cc.relayMap->outerInitIdx; + if (outerInitIdx >= newOuterFor.getRegionIterArgs().size()) { + return failAfterCreate(); + } + + // relay semantics: inner base is new outer iter arg on mapped lane + cc.shapeInfo.base = newOuterFor.getRegionIterArgs()[outerInitIdx]; + relayCandidatesForInner.push_back(std::move(cc)); + } + + FailureOr rewrittenInnerRes = failure(); + { + PatternRewriter::InsertionGuard guard(rewriter); + rewrittenInnerRes = rewriteInnerForWithRelayCandidates( + innerFor, relayCandidatesForInner, &outerMap, rewriter); + } + if (failed(rewrittenInnerRes)) { + return failAfterCreate(); + } + rewrittenInnerInNewOuter = *rewrittenInnerRes; + + // Map old inner results to rewritten inner results for outer cloning/yield + // mapping. + if (innerFor.getNumResults() != rewrittenInnerInNewOuter.getNumResults()) { + return failAfterCreate(); + } + for (unsigned r = 0; r < innerFor.getNumResults(); ++r) { + outerMap.map(innerFor.getResult(r), + rewrittenInnerInNewOuter.getResult(r)); + } + } + + if (!rewrittenInnerInNewOuter) { + return failAfterCreate(); + } + + // Rebuild outer yield; relay lanes forced from rewritten inner results. + SmallVector newOuterYieldOps; + newOuterYieldOps.reserve(oldOuterYield.getNumOperands()); + DenseMap outerYieldToInnerIdx; + for (const auto &c : relayCandidates) { + const auto &m = *c.relayMap; + outerYieldToInnerIdx[m.outerYieldIdx] = m.innerIdx; + } + + for (unsigned i = 0; i < oldOuterYield.getNumOperands(); ++i) { + auto it = outerYieldToInnerIdx.find(i); + if (it != outerYieldToInnerIdx.end()) { + unsigned innerIdx = it->second; + if (innerIdx >= rewrittenInnerInNewOuter.getNumResults()) { + return failAfterCreate(); + } + newOuterYieldOps.push_back(rewrittenInnerInNewOuter.getResult(innerIdx)); + } else { + newOuterYieldOps.push_back( + outerMap.lookupOrDefault(oldOuterYield.getOperand(i))); + } + } + + rewriter.setInsertionPointToEnd(newOuterFor.getBody()); + rewriter.create(oldOuterYield.getLoc(), newOuterYieldOps); + + if (newOuterFor->hasAttr(kIncompleteAttr)) { + newOuterFor->removeAttr(kIncompleteAttr); + } + newOuterFor->setAttr(kSimplifiedAttr, rewriter.getUnitAttr()); + return newOuterFor; +} + +LogicalResult SimplifyTensorIterArgsPattern::rewriteForWithRelayCandidates( + scf::ForOp newfor, scf::ForOp oldFor, + ArrayRef relayCandidates, PatternRewriter &rewriter) const { + // forOp is innerFor (already local-rewritten or to-be-relay-rewritten) + scf::ForOp innerFor = newfor; + scf::ForOp outerFor; + if (failed(precheckRelayCandidates(innerFor, relayCandidates, outerFor))) { + return failure(); + } + + FailureOr newOuterRes = rewriteOuterForWithRelayCandidates( + innerFor, oldFor, outerFor, relayCandidates, rewriter); + if (failed(newOuterRes)) { + return failure(); + } + + scf::ForOp newOuterFor = *newOuterRes; + LLVM_DEBUG({ + llvm::dbgs() << "Successfully rewrote outer ForOp to: \n"; + newOuterFor.dump(); + }); + + // Only replace outer; old inner is nested under old outer and will be removed + // together. + rewriter.replaceOp(outerFor, newOuterFor.getResults()); + return success(); +} + +bool IfYieldAddHoistConverter::isSupportedTensorResultType(Type type) const { + auto tensorType = dyn_cast(type); + return tensorType && !isa(tensorType.getElementType()); +} + +bool IfYieldAddHoistConverter::isDefinedOutsideIf(Value value, + scf::IfOp ifOp) const { + if (auto blockArg = dyn_cast(value)) { + Block *owner = blockArg.getOwner(); + return owner && owner->getParentOp() != ifOp.getOperation(); + } + + Operation *defOp = value.getDefiningOp(); + return defOp && !ifOp->isAncestor(defOp); +} + +bool IfYieldAddHoistConverter::extractAddendFromAddExpr( + Value maybeAddExpr, Value baseValue, Value &addendOut) const { + if (auto addi = maybeAddExpr.getDefiningOp()) { + if (addi.getLhs() == baseValue) { + addendOut = addi.getRhs(); + return true; + } + if (addi.getRhs() == baseValue) { + addendOut = addi.getLhs(); + return true; + } + return false; + } + + if (auto addf = maybeAddExpr.getDefiningOp()) { + if (addf.getLhs() == baseValue) { + addendOut = addf.getRhs(); + return true; + } + if (addf.getRhs() == baseValue) { + addendOut = addf.getLhs(); + return true; + } + return false; + } + + return false; +} + +Value IfYieldAddHoistConverter::buildZeroTensorLikeType( + Type laneType, Location loc, PatternRewriter &rewriter) const { + auto tensorType = dyn_cast(laneType); + if (!tensorType) + return Value(); + + Type elemType = tensorType.getElementType(); + Attribute zeroElemAttr; + if (isa(elemType)) { + zeroElemAttr = rewriter.getFloatAttr(elemType, 0.0); + } else if (elemType.isIntOrIndex()) { + zeroElemAttr = rewriter.getIntegerAttr(elemType, 0); + } else { + return Value(); + } + + auto zeroTensorAttr = DenseElementsAttr::get(tensorType, zeroElemAttr); + return rewriter.create(loc, laneType, zeroTensorAttr); +} + +bool IfYieldAddHoistConverter::tryRewriteSingleLane( + unsigned laneIdx, Value baseBranchYield, Value addExprBranchYield, + bool baseInThenBranch, Type laneType, scf::IfOp ifOp, + PatternRewriter &rewriter, SmallVectorImpl &updatedThenYieldOperands, + SmallVectorImpl &updatedElseYieldOperands, + SmallVectorImpl &hoistedBasePerLane, + SmallVectorImpl &laneRewrittenFlags) const { + if (!isDefinedOutsideIf(baseBranchYield, ifOp)) + return false; + + Value addendValue; + if (!extractAddendFromAddExpr(addExprBranchYield, baseBranchYield, + addendValue)) + return false; + if (!addendValue || addendValue.getType() != laneType) + return false; + + Value zeroTensor = buildZeroTensorLikeType(laneType, ifOp.getLoc(), rewriter); + if (!zeroTensor) + return false; + + if (baseInThenBranch) { + updatedThenYieldOperands[laneIdx] = zeroTensor; + updatedElseYieldOperands[laneIdx] = addendValue; + } else { + updatedThenYieldOperands[laneIdx] = addendValue; + updatedElseYieldOperands[laneIdx] = zeroTensor; + } + + hoistedBasePerLane[laneIdx] = baseBranchYield; + laneRewrittenFlags[laneIdx] = true; + return true; +} + +// scf.if lane rewrite: +// one branch yields A, the other yields A + B +// => if yields 0 / B, then outside do A + if_result +LogicalResult +IfYieldAddHoistConverter::matchAndRewrite(scf::IfOp ifOp, + PatternRewriter &rewriter) const { + if (ifOp.getNumResults() == 0) { + return failure(); + } + + // Robust guard for malformed / emptied regions. + Block *thenBlockPtr = ifOp.thenBlock(); + Block *elseBlockPtr = ifOp.elseBlock(); + if (!thenBlockPtr || !elseBlockPtr) { + return failure(); + } + if (!thenBlockPtr->mightHaveTerminator() || + !elseBlockPtr->mightHaveTerminator()) { + return failure(); + } + + auto thenYieldOp = dyn_cast(ifOp.thenBlock()->getTerminator()); + auto elseYieldOp = dyn_cast(ifOp.elseBlock()->getTerminator()); + if (!thenYieldOp || !elseYieldOp) { + return failure(); + } + + Location loc = ifOp.getLoc(); + bool anyLaneRewritten = false; + + SmallVector updatedThenYieldOperands(thenYieldOp.getOperands().begin(), + thenYieldOp.getOperands().end()); + SmallVector updatedElseYieldOperands(elseYieldOp.getOperands().begin(), + elseYieldOp.getOperands().end()); + + // for each lane, record whether we need post-if add with baseYield + SmallVector hoistedBasePerLane(ifOp.getNumResults(), Value()); + SmallVector laneRewrittenFlags(ifOp.getNumResults(), false); + + for (unsigned laneIdx = 0; laneIdx < ifOp.getNumResults(); ++laneIdx) { + Type laneType = ifOp.getResultTypes()[laneIdx]; + if (!isSupportedTensorResultType(laneType)) { + continue; + } + + Value thenYieldVal = thenYieldOp.getOperand(laneIdx); + Value elseYieldVal = elseYieldOp.getOperand(laneIdx); + + // Assume that only one of the two branches can have the add pattern, and + // try both ways to find a rewrite opportunity. + bool rewritten = + tryRewriteSingleLane( + laneIdx, thenYieldVal, elseYieldVal, /* baseInThenBranch= */ true, + laneType, ifOp, rewriter, updatedThenYieldOperands, + updatedElseYieldOperands, hoistedBasePerLane, laneRewrittenFlags) || + tryRewriteSingleLane( + laneIdx, elseYieldVal, thenYieldVal, /* baseInThenBranch= */ false, + laneType, ifOp, rewriter, updatedThenYieldOperands, + updatedElseYieldOperands, hoistedBasePerLane, laneRewrittenFlags); + + anyLaneRewritten |= rewritten; + } + + if (!anyLaneRewritten) { + return failure(); + } + + rewriter.setInsertionPoint(ifOp); + auto rewrittenIfOp = rewriter.create(loc, ifOp.getResultTypes(), + ifOp.getCondition(), + /* withElseRegion= */ true); + + { + IRMapping thenMapping; + Block *oldThenBlock = ifOp.thenBlock(); + Block *newThenBlock = rewrittenIfOp.thenBlock(); + + rewriter.setInsertionPointToStart(newThenBlock); + for (Operation &op : oldThenBlock->without_terminator()) { + rewriter.clone(op, thenMapping); + } + + SmallVector mappedThenYieldOperands; + mappedThenYieldOperands.reserve(updatedThenYieldOperands.size()); + for (Value v : updatedThenYieldOperands) { + mappedThenYieldOperands.push_back(thenMapping.lookupOrDefault(v)); + } + + rewriter.create(loc, mappedThenYieldOperands); + } + + { + IRMapping elseMapping; + Block *oldElseBlock = ifOp.elseBlock(); + Block *newElseBlock = rewrittenIfOp.elseBlock(); + + rewriter.setInsertionPointToStart(newElseBlock); + for (Operation &op : oldElseBlock->without_terminator()) { + rewriter.clone(op, elseMapping); + } + + SmallVector mappedElseYieldOperands; + mappedElseYieldOperands.reserve(updatedElseYieldOperands.size()); + for (Value v : updatedElseYieldOperands) { + mappedElseYieldOperands.push_back(elseMapping.lookupOrDefault(v)); + } + + rewriter.create(loc, mappedElseYieldOperands); + } + + rewriter.setInsertionPointAfter(rewrittenIfOp); + SmallVector finalResults; + finalResults.reserve(ifOp.getNumResults()); + + for (unsigned laneIdx = 0; laneIdx < ifOp.getNumResults(); ++laneIdx) { + if (!laneRewrittenFlags[laneIdx]) { + finalResults.push_back(rewrittenIfOp.getResult(laneIdx)); + continue; + } + + Value hoistedBase = hoistedBasePerLane[laneIdx]; + Value laneDelta = rewrittenIfOp.getResult(laneIdx); + + auto laneTensorType = dyn_cast(laneDelta.getType()); + if (!laneTensorType) + return failure(); + + Type elemType = laneTensorType.getElementType(); + Value reconstructedLane; + if (isa(elemType)) { + reconstructedLane = + rewriter.create(loc, hoistedBase, laneDelta); + } else if (elemType.isIntOrIndex()) { + reconstructedLane = + rewriter.create(loc, hoistedBase, laneDelta); + } else { + return failure(); + } + + finalResults.push_back(reconstructedLane); + } + + rewriter.replaceOp(ifOp, finalResults); + return success(); +} + +} // namespace CannonicalizerConverter diff --git a/compiler/lib/TritonToStructured/MaskAnalysis.cpp b/compiler/lib/TritonToStructured/MaskAnalysis.cpp new file mode 100644 index 00000000..5a57e7b6 --- /dev/null +++ b/compiler/lib/TritonToStructured/MaskAnalysis.cpp @@ -0,0 +1,992 @@ + + +#include "dicp/TritonToStructured/MaskAnalysis.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/LogicalResult.h" + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/IR/Visitors.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Support/LogicalResult.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" + +#include "dicp/TritonToStructured/PtrAnalysis.h" +#include "dicp/Utils/Utils.h" + +#define DEBUG_TYPE "triton-to-structured-mask-analysis" + +namespace TritonToStructured { +using namespace mlir; +using namespace triton; + +bool dimInfo::setType(arith::CmpIPredicate Type) { + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Setting compare type for dimIndex " << dimIndex << "\n"; + llvm::dbgs() << "Type: " << Type << "\n"; + llvm::dbgs() << "----------------------------------------------\n"; + }); + + switch (Type) { + case arith::CmpIPredicate::slt: + this->currentType = dimInfo::CompareType::slt; + break; + case arith::CmpIPredicate::ult: + this->currentType = dimInfo::CompareType::ult; + break; + case arith::CmpIPredicate::sge: + this->currentType = dimInfo::CompareType::sge; + break; + case arith::CmpIPredicate::uge: + this->currentType = dimInfo::CompareType::uge; + break; + default: + return false; + } + return true; +} + +bool dimInfo::compareTypeIsLess() const { + return this->currentType == dimInfo::CompareType::slt || + this->currentType == dimInfo::CompareType::ult; +} + +void dimInfo::dump() const { + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "MaskDimInfo: \n"; + llvm::dbgs() << "offset = " << offset << "\n"; + llvm::dbgs() << "shape = " << shape << "\n"; + llvm::dbgs() << "rhs = " << rhs << "\n"; + llvm::dbgs() << "isLessMode = " << compareTypeIsLess() << "\n"; + llvm::dbgs() << "hasBroadCast = " << hasBroadCast << "\n"; + llvm::dbgs() << "----------------------------------------------\n"; +}; + +void MaskState::dump() const { + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "MaskState :\n"; + llvm::dbgs() << "scalar = " << scalar << "\n"; + llvm::dbgs() << "stateInfo.size = " << stateInfo.size() << "\n"; + for (auto info : stateInfo) + info.dump(); + llvm::dbgs() << "----------------------------------------------\n"; +}; + +LogicalResult MaskState::parse(Value operand, const Location loc, + OpBuilder &builder) { + if (isa(operand.getType()) && + operand.getType().getIntOrFloatBitWidth() != 1) { + return this->parseIntScalar(operand, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseConstant(op, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseAdd(op, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseAnd(op, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseCmp(op, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseMakeRange(op, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseBroadcast(op, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseSplat(op, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseExpandDims(op, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseExtSI(op, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseRem(op, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return this->parseDiv(op, loc, builder); + } + LLVM_DEBUG({ + InFlightDiagnostic diag = emitWarning(loc) + << "MaskAnalysis: compare operand produced by an " + "unsupported operation\n"; + }); + return failure(); +} + +LogicalResult MaskState::parseConstant(arith::ConstantOp constOp, + const Location loc, OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + constOp.emitError( + "MaskAnalysis: MaskState should be empty when visiting constant"); + }); + return failure(); + } + + if (auto intType = dyn_cast(constOp.getType())) { + if (intType.getWidth() == 1) { + LLVM_DEBUG({ + constOp.emitWarning("MaskAnalysis: Unsupported constant for int1"); + }); + return failure(); + } + } + + if (isa(constOp.getValue())) { + auto attr = cast(constOp.getValue()); + auto elementType = attr.getElementType(); + if (!attr.isSplat() || !isa(elementType)) { + LLVM_DEBUG({ + constOp.emitError("MaskAnalysis: only support splat integer constant"); + }); + return failure(); + } + auto values = attr.getValues(); + auto value = values[0].getValue(); + auto constAttr = builder.getIndexAttr(value.getSExtValue()); + auto op = arith::ConstantOp::materialize(builder, constAttr, + builder.getIndexType(), loc); + this->scalar = op.getValue(); + } else { + auto value = cast(constOp.getValue()).getInt(); + this->scalar = builder.getIndexAttr(value); + } + return success(); +} + +LogicalResult MaskState::parseIntScalar(Value scalar, const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitError(loc) << "MaskAnalysis: MaskState should be empty when " + "visiting integer scalar"; + }); + return failure(); + } + auto castOp = + builder.create(loc, builder.getIndexType(), scalar); + this->scalar = castOp.getResult(); + return success(); +} + +LogicalResult MaskState::parseMakeRange(triton::MakeRangeOp rangeOp, + const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + rangeOp.emitError( + "MaskAnalysis: MaskState should be empty when visiting make_range"); + }); + return failure(); + } + + auto shape = cast(rangeOp.getType()).getShape(); + auto start = rangeOp.getStart(); + auto end = rangeOp.getEnd(); + auto stride = (end - start + shape[0] - 1) / shape[0]; + + if (stride != 1) { + LLVM_DEBUG( + { + InFlightDiagnostic diag = + emitError(loc) + << "stride must be 1 for make_range whose result is used " + "as load or store masks"; + }); + return failure(); + } + + stateInfo.emplace_back(builder.getIndexAttr(start), + builder.getIndexAttr(shape[0])); + return success(); +} + +LogicalResult MaskState::parseExtSI(arith::ExtSIOp op, const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + op->emitError( + "MaskAnalysis: MaskState should be empty when visiting extsi"); + }); + return failure(); + } + return parse(op.getIn(), loc, builder); +} + +LogicalResult MaskState::parseSplat(triton::SplatOp splatOp, const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + splatOp.emitError( + "MaskAnalysis: MaskState should be empty when visiting splat"); + }); + return failure(); + } + + auto src = splatOp.getSrc(); + auto dst = splatOp.getResult(); + auto dstShape = cast(dst.getType()).getShape(); + + if (!isa(src.getType())) { + LLVM_DEBUG( + { + splatOp.emitError() + << "splat source must be an integer scalar for load/store masks"; + }); + return failure(); + } + + if (failed(this->parse(src, loc, builder))) + return failure(); + + if (stateInfo.size() > 1 || + (stateInfo.size() == 1 && !isOne(stateInfo.back().shape))) { + LLVM_DEBUG({ + splatOp.emitError() << "splat from a non-scalar source is not supported, " + "unless it's state size and shape are 1"; + }); + return failure(); + } + + auto zeroAttr = builder.getIndexAttr(0); + SmallVector newStateInfo; + for (auto [i, shape] : llvm::enumerate(dstShape)) { + auto shapeAttr = builder.getIndexAttr(shape); + newStateInfo.emplace_back(zeroAttr, shapeAttr, i, true); + if (!stateInfo.empty() && shape == 1) { + newStateInfo.back() = stateInfo.back(); + newStateInfo.back().dimIndex = i; + } + } + this->stateInfo = newStateInfo; + return success(); +} + +LogicalResult MaskState::parseExpandDims(triton::ExpandDimsOp expandDimsOp, + const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + expandDimsOp.emitError( + "MaskAnalysis: MaskState should be empty when visiting expand_dims"); + }); + return failure(); + } + + auto zeroAttr = builder.getIndexAttr(0); + auto defaultShape = builder.getIndexAttr(1); + + if (failed(this->parse(expandDimsOp.getSrc(), loc, builder))) + return failure(); + + auto dstShape = + cast(expandDimsOp.getResult().getType()).getShape(); + auto axis = expandDimsOp.getAxis(); + if (dstShape[axis] != 1) { + LLVM_DEBUG({ + expandDimsOp.emitError( + "MaskAnalysis: unexpected dimension size in expand_dims"); + }); + return failure(); + } + + size_t insertPos = 0; + for (auto &info : stateInfo) { + if (info.dimIndex >= axis) + ++info.dimIndex; + if (info.dimIndex < axis) + ++insertPos; + } + + dimInfo insertInfo(zeroAttr, defaultShape, axis, true); + stateInfo.insert(stateInfo.begin() + insertPos, insertInfo); + + return success(); +} + +LogicalResult MaskState::parseAdd(arith::AddIOp addOp, const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + addOp.emitError( + "MaskAnalysis: MaskState should be empty when visiting add"); + }); + return failure(); + } + + MaskState lhsState; + if (failed(lhsState.parse(addOp.getLhs(), loc, builder))) + return failure(); + + MaskState rhsState; + if (failed(rhsState.parse(addOp.getRhs(), loc, builder))) + return failure(); + + return this->addStates(lhsState, rhsState, loc, builder); +} + +LogicalResult MaskState::addStates(const MaskState &lhsState, + const MaskState &rhsState, Location loc, + OpBuilder &builder) { + if (lhsState.scalar && rhsState.scalar) { + LLVM_DEBUG( + { + InFlightDiagnostic diag = + emitWarning(loc) + << "Unexpected case where both lhs and rhs are scalars"; + }); + return failure(); + } + + if (!lhsState.scalar && !rhsState.scalar) { + LLVM_DEBUG( + { + InFlightDiagnostic diag = + emitWarning(loc) + << "Unsupported scenario where neither lhs nor rhs is a scalar"; + }); + return failure(); + } + + if (lhsState.scalar) + return addStateScalar(rhsState, lhsState.scalar, loc, builder); + else + return addStateScalar(lhsState, rhsState.scalar, loc, builder); +} + +LogicalResult MaskState::addStateScalar(const MaskState &state, + const OpFoldResult scalar, Location loc, + OpBuilder &builder) { + for (auto info : state.stateInfo) { + info.offset = addOpFoldResult(info.offset, scalar, loc, builder); + this->stateInfo.emplace_back(info); + } + return success(); +} + +LogicalResult MaskState::parseBroadcast(triton::BroadcastOp broadcastOp, + const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + broadcastOp.emitError( + "MaskAnalysis: MaskState should be empty when visiting broadcast"); + }); + return failure(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Parsing BROADCAST operation: " << broadcastOp << "\n"; + }); + + auto src = broadcastOp.getSrc(); + auto dst = broadcastOp.getResult(); + if (!isa(dst.getType())) { + LLVM_DEBUG({ + broadcastOp.emitError( + "MaskAnalysis: broadcast dst should be a shaped type"); + }); + return failure(); + } + + auto srcShape = cast(src.getType()).getShape(); + auto dstShape = cast(dst.getType()).getShape(); + if (srcShape.size() != dstShape.size()) { + LLVM_DEBUG({ + broadcastOp.emitError( + "MaskAnalysis: broadcast src and dst should have the same rank"); + }); + return failure(); + } + + if (failed(parse(src, loc, builder))) + return failure(); + + LLVM_DEBUG({ + llvm::dbgs() << "Before BROADCAST MaskState: \n"; + this->dump(); + }); + + for (size_t i = 0; i < srcShape.size(); ++i) { + if (srcShape[i] == dstShape[i]) + continue; + else if (srcShape[i] < dstShape[i] && srcShape[i] == 1) { + for (auto &info : stateInfo) { + if (info.dimIndex != i) + continue; + info.shape = builder.getIndexAttr(dstShape[i]); + info.hasBroadCast = true; + } + } else { + LLVM_DEBUG({ + broadcastOp.emitError( + "MaskAnalysis: unexpected dimensions used in broadcast"); + }); + return failure(); + } + } + + LLVM_DEBUG({ + llvm::dbgs() << "After BROADCAST MaskState: \n"; + this->dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + + return success(); +} + +LogicalResult MaskState::parseCmp(arith::CmpIOp cmpOp, const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + cmpOp.emitError( + "MaskAnalysis: MaskState should be empty when visiting cmpi"); + }); + return failure(); + } + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Parsing CMP operation: " << cmpOp << "\n"; + }); + + if (isa(cmpOp.getLhs().getType()) && + (cmpOp.getLhs().getDefiningOp() || + cmpOp.getRhs().getDefiningOp())) { + LLVM_DEBUG({ + cmpOp.emitWarning( + "MaskAnalysis: Unsupported nested cmpi scenario for int1"); + }); + return failure(); + } + + MaskState lhsState; + if (failed(lhsState.parse(cmpOp.getLhs(), loc, builder))) + return failure(); + + MaskState rhsState; + if (failed(rhsState.parse(cmpOp.getRhs(), loc, builder))) + return failure(); + + if (lhsState.scalar) { + if (lhsState.stateInfo.empty()) { + lhsState.stateInfo.emplace_back(builder.getIndexAttr(0), + builder.getIndexAttr(1)); + } + for (auto &info : lhsState.stateInfo) { + if (isOne(info.shape)) { + info.offset = lhsState.scalar; + } + } + } + + // lhs must be a Value and rhs must be scalar + if (!rhsState.scalar) { + LLVM_DEBUG( + { cmpOp.emitWarning("MaskAnalysis: Unsupported cmpi scenario"); }); + return failure(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "LHS MaskState: \n"; + lhsState.dump(); + llvm::dbgs() << "RHS MaskState: \n"; + rhsState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + + // In the case where the values we are loading are entirely masked off like + // the following: + // + // ---|-------|-----------| + // ^ ^ ^ + // scalar start end + // + // newEnd = min(end, scalar) = scalar + // Now scalar < start, so simply doing dim = newEnd - start is incorrect. + // + // The correct formula is to optionally move `newDim` back to `start` using + // max(newEnd, start). + auto cmpType = cmpOp.getPredicate(); + for (auto &info : lhsState.stateInfo) { + if (info.hasBroadCast) + continue; + if (!info.setType(cmpType)) { + LLVM_DEBUG({ cmpOp.emitWarning("MaskAnalysis: Unsupported cmpi type"); }); + return failure(); + } + info.rhs = rhsState.scalar; + } + this->stateInfo = lhsState.stateInfo; + + LLVM_DEBUG({ + llvm::dbgs() << "After CMP MaskState: \n"; + this->dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + return success(); +} + +LogicalResult MaskState::parseRem(arith::RemSIOp remOp, const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + remOp.emitError( + "MaskAnalysis: MaskState should be empty when visiting REMSI"); + }); + return failure(); + } + + MaskState lhsState; + if (failed(lhsState.parse(remOp.getLhs(), loc, builder))) + return failure(); + + MaskState rhsState; + if (failed(rhsState.parse(remOp.getRhs(), loc, builder))) + return failure(); + + if (lhsState.scalar || !rhsState.scalar) { + LLVM_DEBUG( + { remOp.emitRemark("MaskAnalysis: Unsupported REMSI scenario"); }); + return failure(); + } + + auto divisorAttr = rhsState.scalar; + + if (!getIntAttr(divisorAttr).has_value()) { + LLVM_DEBUG({ + remOp.emitError("MaskAnalysis: do not support dynamic divisor in REMSI."); + }); + return failure(); + } + + SmallVector newStateInfo; + auto zeroAttr = builder.getIndexAttr(0); + for (auto info : lhsState.stateInfo) { + if (info.hasBroadCast) { + newStateInfo.emplace_back(info); + continue; + } + if (!isMultiple(divisorAttr, info.shape) && + !isMultiple(info.shape, divisorAttr)) { + LLVM_DEBUG({ + remOp.emitError( + "MaskAnalysis: do not support dynamic stride before REMSI."); + }); + return failure(); + } + + auto contiguousSize = + minOpFoldResult(divisorAttr, info.shape, loc, builder); + auto nonContiguousSize = + divOpFoldResult(info.shape, contiguousSize, loc, builder); + + auto staticNonContiguousSize = getIntAttr(nonContiguousSize); + if (!staticNonContiguousSize.has_value()) { + LLVM_DEBUG({ + remOp.emitError( + "MaskAnalysis: do not support dynamic size before REMSI."); + }); + return failure(); + } + + if (staticNonContiguousSize.value() != 0) + newStateInfo.emplace_back(zeroAttr, nonContiguousSize, info.dimIndex, + true); + + auto newOffset = remOpFoldResult(info.offset, divisorAttr, loc, builder); + newStateInfo.emplace_back(newOffset, nonContiguousSize, info.dimIndex); + } + + this->stateInfo = newStateInfo; + + return success(); +} + +LogicalResult MaskState::parseDiv(arith::DivSIOp divOp, const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + divOp.emitError( + "MaskAnalysis: MaskState should be empty when visiting DIVSI"); + }); + return failure(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Parsing DIV operation: " << divOp << "\n"; + }); + + MaskState lhsState; + if (failed(lhsState.parse(divOp.getLhs(), loc, builder))) + return failure(); + + MaskState rhsState; + if (failed(rhsState.parse(divOp.getRhs(), loc, builder))) + return failure(); + + LLVM_DEBUG({ + llvm::dbgs() << "LHS MaskState: \n"; + lhsState.dump(); + llvm::dbgs() << "RHS MaskState: \n"; + rhsState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + + if (lhsState.scalar || !rhsState.scalar) { + LLVM_DEBUG( + { divOp.emitRemark("MaskAnalysis: Unsupported DIVSI scenario"); }); + return failure(); + } + + auto divisorAttr = rhsState.scalar; + + if (!getIntAttr(divisorAttr).has_value()) { + LLVM_DEBUG({ + divOp.emitError("MaskAnalysis: do not support dynamix divisor in DIVSI."); + }); + return failure(); + } + + SmallVector newStateInfo; + auto zeroAttr = builder.getIndexAttr(0); + for (auto info : lhsState.stateInfo) { + if (info.hasBroadCast) { + newStateInfo.emplace_back(info); + continue; + } + if (!isMultiple(divisorAttr, info.shape) && + !isMultiple(info.shape, divisorAttr)) { + LLVM_DEBUG({ + divOp.emitError( + "MaskAnalysis: do not support dynamix stride before DIVSI."); + }); + return failure(); + } + + auto nonContiguousSize = + minOpFoldResult(divisorAttr, info.shape, loc, builder); + auto contiguousSize = + divOpFoldResult(info.shape, nonContiguousSize, loc, builder); + + auto staticContiguousSize = getIntAttr(contiguousSize); + if (!staticContiguousSize.has_value()) { + LLVM_DEBUG({ + divOp.emitError( + "MaskAnalysis: do not support dynamix size before DIVSI."); + }); + return failure(); + } + + if (staticContiguousSize.value() != 0) { + auto newOffset = divOpFoldResult(info.offset, divisorAttr, loc, builder); + newStateInfo.emplace_back(newOffset, contiguousSize, info.dimIndex); + } + + newStateInfo.emplace_back(zeroAttr, nonContiguousSize, info.dimIndex, true); + } + + this->stateInfo = newStateInfo; + + LLVM_DEBUG({ + llvm::dbgs() << "After DIV MaskState: \n"; + this->dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + return success(); +} + +LogicalResult MaskState::parseAnd(arith::AndIOp andOp, const Location loc, + OpBuilder &builder) { + if (!this->isEmpty()) { + LLVM_DEBUG({ + andOp.emitError( + "MaskAnalysis: MaskState should be empty when visiting and"); + }); + return failure(); + } + auto zeroAttr = builder.getIndexAttr(0); + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Parsing AND operation: " << andOp << "\n"; + }); + + if (isa(andOp.getLhs().getType())) { + LLVM_DEBUG({ + andOp.emitWarning("MaskAnalysis: Unsupported andi scenario for int1"); + }); + return failure(); + } + + MaskState lhsState; + if (failed(lhsState.parse(andOp.getLhs(), loc, builder))) + return failure(); + MaskState rhsState; + if (failed(rhsState.parse(andOp.getRhs(), loc, builder))) + return failure(); + + LLVM_DEBUG({ + llvm::dbgs() << "LHS MaskState: \n"; + lhsState.dump(); + llvm::dbgs() << "RHS MaskState: \n"; + rhsState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + + SmallVector newStateInfo; + auto lIt = lhsState.stateInfo.begin(); + auto rIt = rhsState.stateInfo.begin(); + + while (lIt != lhsState.stateInfo.end() && rIt != rhsState.stateInfo.end()) { + if (lIt->dimIndex != rIt->dimIndex) { + auto newInfo = lIt->dimIndex < rIt->dimIndex ? *lIt++ : *rIt++; + newStateInfo.emplace_back(newInfo); + continue; + } + + if (!isMultiple(lIt->shape, rIt->shape) && + !isMultiple(rIt->shape, lIt->shape)) { + LLVM_DEBUG({ + llvm::dbgs() << "LHS MaskState: \n"; + lhsState.dump(); + llvm::dbgs() << "RHS MaskState: \n"; + rhsState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + LLVM_DEBUG({ + andOp.emitError( + "MaskAnalysis: the add operation have incompatible sizes"); + }); + return failure(); + } + + dimInfo newInfo; + newInfo.dimIndex = lIt->dimIndex; + newInfo.shape = minOpFoldResult(lIt->shape, rIt->shape, loc, builder); + if ((isLess(newInfo.shape, lIt->shape) && !lIt->hasBroadCast || + isLess(newInfo.shape, rIt->shape) && !rIt->hasBroadCast)) { + LLVM_DEBUG({ + llvm::dbgs() << "LHS MaskState: \n"; + lhsState.dump(); + llvm::dbgs() << "RHS MaskState: \n"; + rhsState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + LLVM_DEBUG({ + andOp.emitError( + "MaskAnalysis: the add operation have incompatible sizes." + "Valid dimensions are split."); + }); + return failure(); + } + newInfo.currentType = + lIt->hasBroadCast ? rIt->currentType : lIt->currentType; + if (lIt->currentType != dimInfo::CompareType::deafaultType && + rIt->currentType != dimInfo::CompareType::deafaultType && + lIt->currentType != rIt->currentType) { + LLVM_DEBUG({ + andOp.emitError( + "MaskAnalysis: do not suppport different compare mode within" + "the same dimension."); + }); + return failure(); + } + + if (lIt->hasBroadCast) { + newInfo.offset = rIt->offset; + newInfo.rhs = rIt->rhs; + newInfo.hasBroadCast = rIt->hasBroadCast; + } else if (rIt->hasBroadCast) { + newInfo.offset = lIt->offset; + newInfo.rhs = lIt->rhs; + newInfo.hasBroadCast = lIt->hasBroadCast; + } else { + LLVM_DEBUG({ + andOp.emitError("MaskAnalysis: do not suppport " + "and in the same dimension."); + }); + return failure(); + } + + newStateInfo.emplace_back(newInfo); + + if (isEqual(lIt->shape, newInfo.shape)) + ++lIt; + else + lIt->shape = divOpFoldResult(lIt->shape, newInfo.shape, loc, builder); + if (isEqual(rIt->shape, newInfo.shape)) + ++rIt; + else + rIt->shape = divOpFoldResult(rIt->shape, newInfo.shape, loc, builder); + } + + while (rIt != rhsState.stateInfo.end()) { + newStateInfo.push_back(*rIt++); + } + while (lIt != lhsState.stateInfo.end()) { + newStateInfo.push_back(*lIt++); + } + + this->stateInfo = newStateInfo; + + LLVM_DEBUG({ + llvm::dbgs() << "After AND MaskState: \n"; + this->dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + return success(); +} + +LogicalResult MaskState::analysisMask(Value operand) { + auto op = operand.getDefiningOp(); + if (!op) { + return failure(); + } + auto loc = op->getLoc(); + OpBuilder builder(op); + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Analyzing mask: " << operand << "\n"; + }); + + if (this->parse(operand, loc, builder).failed() || this->isEmpty()) { + return failure(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "Mask analysis result: \n"; + this->dump(); + llvm::dbgs() << "MaskAnalysis: successfully analyzed mask.\n"; + llvm::dbgs() << "----------------------------------------------\n"; + }); + return success(); +} + +Value MaskState::createNewMask(const Location loc, OpBuilder &builder) { + if (isEmpty()) + return nullptr; + + SmallVector shape; + for (auto info : stateInfo) { + auto staticShape = getIntAttr(info.shape); + if (!staticShape.has_value()) { + LLVM_DEBUG( + { + InFlightDiagnostic diag = + emitError(loc) + << "MaskAnalysis: dynamic shape is not supported in mask " + "generation\n"; + }); + return nullptr; + } + shape.emplace_back(staticShape.value()); + } + SmallVector cacheResults; + auto maskShape = RankedTensorType::get(shape, builder.getI1Type()); + + auto createRhsValue = [&](OpFoldResult rhs) -> Value { + if (auto rhsInt = getIntAttr(rhs)) { + auto rhsAttr = + builder.getI32IntegerAttr(static_cast(rhsInt.value())); + return builder.create(loc, rhsAttr).getResult(); + } + Value rhsValue = dyn_cast(rhs); + if (rhsValue.getType().isIndex()) { + rhsValue = builder.create(loc, builder.getI32Type(), + rhsValue); + } + return rhsValue; + }; + for (size_t i = 0; i < stateInfo.size(); ++i) { + auto info = stateInfo[i]; + if (info.hasBroadCast) { + continue; + } + auto indexI32RowType = + RankedTensorType::get(shape[i], builder.getI32Type()); + Value newMask = + builder.create(loc, indexI32RowType, 0, shape[i]); + + Value newOffset = createRhsValue(info.offset); + if (newOffset.getType().isIndex()) { + newOffset = builder.create(loc, builder.getI32Type(), + newOffset); + } + Value splatRhs = + builder.create(loc, indexI32RowType, newOffset); + newMask = builder.create(loc, newMask, splatRhs); + + auto rhsValue = createRhsValue(info.rhs); + auto splatOp = + builder.create(loc, indexI32RowType, rhsValue); + + if (info.currentType == dimInfo::CompareType::deafaultType) { + LLVM_DEBUG({ + InFlightDiagnostic diag = emitError(loc) + << "MaskAnalysis: cannot generate mask when " + "compare type is not set\n"; + }); + return nullptr; + } + auto cmpOp = builder.create(loc, + info.compareTypeIsLess() + ? arith::CmpIPredicate::slt + : arith::CmpIPredicate::sge, + newMask, splatOp.getResult()); + + auto expandValue = cmpOp.getResult(); + for (size_t j = 0; j < stateInfo.size(); ++j) { + if (j == i) + continue; + expandValue = builder.create(loc, expandValue, j); + } + + auto broadcastValue = + builder.create(loc, maskShape, expandValue); + + cacheResults.push_back(broadcastValue); + } + + if (cacheResults.empty()) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) << "MaskAnalysis: cannot generate mask when all " + "dimensions are broadcasted"; + }); + return nullptr; + } + newMask = cacheResults[0]; + for (size_t i = 1; i < cacheResults.size(); ++i) { + newMask = builder.create(loc, newMask, cacheResults[i]); + } + return newMask; +} + +} // namespace TritonToStructured diff --git a/compiler/lib/TritonToStructured/MemOpConverter.cpp b/compiler/lib/TritonToStructured/MemOpConverter.cpp new file mode 100644 index 00000000..88782d55 --- /dev/null +++ b/compiler/lib/TritonToStructured/MemOpConverter.cpp @@ -0,0 +1,554 @@ + + +#include "dicp/TritonToStructured/MemOpConverter.h" + +#include +#include +#include + +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Passes.h" +#include "mlir/Dialect/Utils/ReshapeOpsUtils.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Location.h" +#include "mlir/IR/Matchers.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/IR/Value.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/MathExtras.h" + +#include "llvm/Support/Debug.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HIVM/Utils/Utils.h" +#include "dicp/TritonToStructured/CannonicalizerConverter.h" +#include "dicp/TritonToStructured/MaskAnalysis.h" +#include "dicp/TritonToStructured/PtrAnalysis.h" +#include "dicp/TritonToStructured/TritonToStructuredPass.h" +#include "dicp/Utils/InterleaveOptimization.h" +#include "dicp/Utils/Utils.h" + +#define DEBUG_TYPE "triton-mem-op-converter" + +namespace MemOpConverter { +using namespace mlir; +using namespace triton; +using namespace TritonToStructured; + +LogicalResult LoadConverter::matchAndRewrite(triton::LoadOp op, + PatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto oldPtr = op.getPtr(); + auto oldMask = op.getMask(); + auto oldOther = op.getOther(); + + MemOpTransformer tf(MemOpTransformer::MemType::load, optimizeDynamicOffset); + + auto newPtr = tf.createNewPtr(oldPtr, loc, rewriter); + auto newMask = tf.createNewMask(oldMask, loc, rewriter); + auto newOther = tf.createNewOther(oldOther, loc, rewriter); + + if (!tf.ptrState.shouldLinearize) { + // no need to rewrite + return failure(); + } + + if (!newPtr) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) << "PtrAnalysis: failed to analyze load pointer."; + }); + return failure(); + } + + if (!enableMaskFallbackConversion && oldMask && !newMask) { + LLVM_DEBUG({ + InFlightDiagnostic diag = emitWarning(loc) + << "MaskAnalysis: failed to analyze load mask."; + }); + return failure(); + } + + auto loadOp = rewriter.create(loc, newPtr, newMask, newOther, + op.getCache(), op.getEvict(), + op.getIsVolatile()); + + // insert implicit ops + auto broadCastResult = + tf.materializeImplicitBroadcast(loadOp.getResult(), loc, rewriter); + auto permuteResult = + tf.materializeImplicitPermute(broadCastResult, loc, rewriter); + auto reshapeResult = + tf.materializeImplicitReshape(permuteResult, loc, rewriter); + auto selectResult = tf.materializeImplicitSelect(reshapeResult, oldMask, + oldOther, loc, rewriter); + + rewriter.replaceOp(op, selectResult); + return success(); +} + +LogicalResult StoreConverter::matchAndRewrite(triton::StoreOp op, + PatternRewriter &rewriter) const { + auto loc = op.getLoc(); + auto oldPtr = op.getPtr(); + auto oldMask = op.getMask(); + auto oldValue = op.getValue(); + + MemOpTransformer tf(MemOpTransformer::MemType::store, optimizeDynamicOffset); + + auto newPtr = tf.createNewPtr(oldPtr, loc, rewriter); + auto newMask = tf.createNewMask(oldMask, loc, rewriter); + + if (!tf.ptrState.shouldLinearize) { + // no need to rewrite + return failure(); + } + + if (!newPtr) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) << "PtrAnalysis: failed to analyze store pointer."; + }); + return failure(); + } + + if (!enableMaskFallbackConversion && oldMask && !newMask) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) << "MaskAnalysis: failed to analyze store mask."; + }); + return failure(); + } + + // insert sync_block_lock + auto lockVar = createSyncBlockLockVar(rewriter, loc); + if (oldMask && !newMask) { + rewriter.create(loc, lockVar); + } + + auto selectResult = + tf.materializeImplicitSelect(oldValue, oldMask, oldPtr, loc, rewriter); + auto reshapeResult = + tf.materializeImplicitReshape(selectResult, loc, rewriter); + auto permuteResult = + tf.materializeImplicitPermute(reshapeResult, loc, rewriter); + + auto storeOp = rewriter.create( + loc, newPtr, permuteResult, newMask, op.getBoundaryCheck(), op.getCache(), + op.getEvict()); + + // insert sync_block_unlock + if (oldMask && !newMask) { + rewriter.create(loc, lockVar); + } + rewriter.eraseOp(op); + return success(); +} + +Value MemOpTransformer::materializeImplicitBroadcast( + Value srcTensor, const Location loc, PatternRewriter &rewriter) { + SmallVector broadCastIndex; + SmallVector broadCastShape; + for (auto [i, info] : llvm::enumerate(ptrState.stateInfo)) { + if (isZero(info.stride)) { + broadCastIndex.emplace_back(i); + } + auto staticShape = getIntAttr(info.shape); + if (!staticShape.has_value()) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) + << "PtrAnalysis: dynamic shape is not supported in broadcast\n"; + }); + return srcTensor; + } + broadCastShape.emplace_back(staticShape.value()); + } + + if (broadCastIndex.empty()) + return srcTensor; + + // when load is a scalar, we need to use splat to broadcast + auto srcType = srcTensor.getType(); + if (srcType.isIntOrFloat()) { + auto broadCastType = RankedTensorType::get(broadCastShape, srcType); + auto splatOp = + rewriter.create(loc, broadCastType, srcTensor); + return splatOp.getResult(); + } + + auto init = rewriter.create( + loc, broadCastShape, + cast(srcTensor.getType()).getElementType()); + + auto broadCastOp = rewriter.create(loc, srcTensor, init, + broadCastIndex); + + return broadCastOp->getResult(0); +} + +Value MemOpTransformer::materializeImplicitReshape(Value srcTensor, + const Location loc, + PatternRewriter &rewriter) { + if (ptrState.sizes.size() == ptrState.stateInfo.size()) + return srcTensor; + SmallVector targetShape; + if (currentType == MemType::load) { + for (auto size : ptrState.sizes) { + auto staticShape = getIntAttr(size); + if (!staticShape.has_value()) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) + << "PtrAnalysis: dynamic shape is not supported in reshape\n"; + }); + return srcTensor; + } + targetShape.emplace_back(staticShape.value()); + } + } else { + for (auto info : ptrState.stateInfo) { + auto staticShape = getIntAttr(info.shape); + if (!staticShape.has_value()) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) + << "PtrAnalysis: dynamic shape is not supported in reshape\n"; + }); + return srcTensor; + } + targetShape.emplace_back(staticShape.value()); + } + } + + auto targetShapeAttr = DenseIntElementsAttr::get( + RankedTensorType::get({static_cast(targetShape.size())}, + rewriter.getI64Type()), + targetShape); + auto targetShapeType = RankedTensorType::get( + targetShape, cast(srcTensor.getType()).getElementType()); + auto targetShapeValue = + rewriter.create(loc, targetShapeAttr); + auto reshapeOp = rewriter.create( + loc, targetShapeType, srcTensor, targetShapeValue); + return reshapeOp.getResult(); +} + +Value MemOpTransformer::materializeImplicitSelect(Value srcTensor, Value mask, + Value other, + const Location loc, + PatternRewriter &rewriter) { + if (!mask || maskState.newMask) + return srcTensor; + auto TensorType = cast(srcTensor.getType()); + if (cast(mask.getType()).getShape() != TensorType.getShape()) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) << "MaskAnalysis: mask shape is not same as Value"; + }); + return srcTensor; + } + + if (currentType == MemType::store) { + auto loadOp = rewriter.create(loc, other, nullptr, nullptr, + ArrayRef(), nullptr); + other = loadOp.getResult(); + } + + if (!other) { + auto elementType = TensorType.getElementType(); + auto emptyOp = rewriter.create(loc, TensorType.getShape(), + elementType); + other = emptyOp.getResult(); + } + auto selectOp = rewriter.create(loc, mask, srcTensor, other); + return selectOp->getResult(0); +} + +Value MemOpTransformer::materializeImplicitPermute(Value srcTensor, + const Location loc, + PatternRewriter &rewriter) { + auto inTy = dyn_cast(srcTensor.getType()); + if (!inTy || !ptrState.isPermuted) + return srcTensor; + + auto inShape = inTy.getShape(); + SmallVector order(ptrState.permuteIds.size()); + for (size_t i = 0; i < ptrState.permuteIds.size(); ++i) { + if (currentType == MemType::load) { + order[ptrState.permuteIds[i]] = i; + } else { + order[i] = ptrState.permuteIds[i]; + } + } + SmallVector outShape(order.size()); + if (inShape.size() != outShape.size()) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) << "PtrAnalysis: incompatible shape for permute"; + }); + return srcTensor; + } + + for (size_t i = 0; i < outShape.size(); ++i) { + outShape[i] = inShape[order[i]]; + } + + auto outTy = RankedTensorType::get(outShape, inTy.getElementType()); + auto transOp = rewriter.create(loc, outTy, srcTensor, order); + return transOp.getResult(); +} + +Value MemOpTransformer::createNewPtr(Value oldPtr, const Location loc, + PatternRewriter &rewriter) { + TritonToStructured::PtrAnalysis ptrAnalysis(optimizeDynamicOffset); + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "PtrAnalysis: analyzing load/store's ptr.\n"; + }); + + if (ptrAnalysis.visitOperand(oldPtr, ptrState, loc, rewriter).failed()) { + ptrState.shouldLinearize = false; + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) << "PtranAlysis: failed to analyze load/store ptr."; + }); + return oldPtr; + } + + // compute missing strides + // if stateinfo.shape is 1 and sizes[dimIndex] is 1, + // then the stride is the accumulated size of all dimensions on the right side + // ie. for shape [1, 128], sizes [1, 128], originally stride is [0, 1], + // after normalization, stride is [128, 1] + OpFoldResult maxStride = rewriter.getIndexAttr(1); + for (auto it = ptrState.stateInfo.rbegin(); it != ptrState.stateInfo.rend(); + ++it) { + if (TritonToStructured::isOne(it->shape) && isZero(it->stride)) { + it->stride = maxStride; + } + maxStride = maxOpFoldResult(maxStride, it->stride, loc, rewriter); + } + + for (auto it = ptrState.stateInfo.rbegin(); it != ptrState.stateInfo.rend(); + ++it) { + if (isZero(it->stride)) { + ptrState.shouldLinearize = true; + } + } + + ptrState.generateOriginPermuteIds(); + + return ptrState.createAddPtrOp(rewriter, loc); +} + +Value MemOpTransformer::createNewMask(Value oldMask, const Location loc, + PatternRewriter &rewriter) { + if (!oldMask) + return nullptr; + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "MaskAnalysis: analyzing load/store mask.\n"; + }); + + if (!oldMask || maskState.analysisMask(oldMask).failed()) { + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "MaskAnalysis: no mask or failed to analyze mask.\n"; + llvm::dbgs() << "oldMask:" << oldMask << "\n"; + maskState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + LLVM_DEBUG( + { + InFlightDiagnostic diag = + emitWarning(loc) + << "MaskAnalysis: failed to analyze load/store mask."; + }); + return nullptr; + } + + SmallVector newMaskInfo; + auto itPtr = ptrState.stateInfo.begin(); + auto itMask = maskState.stateInfo.begin(); + + // match and create new mask info + while (itPtr != ptrState.stateInfo.end() && + itMask != maskState.stateInfo.end()) { + // ptr'shape must be multiple of mask'shape or vice versa + if (!isMultiple(itMask->shape, itPtr->shape)) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) + << "MaskAnalysis: incompatible shapes between ptr and mask."; + llvm::dbgs() << "----------------------------------------------\n"; + ptrState.dump(); + llvm::dbgs() << "oldMask:" << oldMask << "\n"; + maskState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + return nullptr; + } + + auto newShape = minOpFoldResult(itMask->shape, itPtr->shape, loc, rewriter); + if (isLess(newShape, itMask->shape) && !itMask->hasBroadCast) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) + << "MaskAnalysis: the mask shape is incompatible with ptr shape."; + }); + return nullptr; + } + + TritonToStructured::dimInfo newInfo(itMask->offset, newShape, + itMask->dimIndex, itMask->hasBroadCast, + itMask->currentType, itMask->rhs); + + if (!isZero(itPtr->stride)) { + newMaskInfo.emplace_back(newInfo); + } + + ++itPtr; + if (isEqual(itMask->shape, newShape)) { + ++itMask; + } + } + + if (itPtr != ptrState.stateInfo.end() || + itMask != maskState.stateInfo.end()) { + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "MaskAnalysis: failed to apply permute on mask.\n"; + ptrState.dump(); + llvm::dbgs() << "oldMask:" << oldMask << "\n"; + maskState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + LLVM_DEBUG({ + InFlightDiagnostic diag = emitWarning(loc) + << "MaskAnalysis: incompatible number of " + "dimensions between ptr and mask."; + }); + return nullptr; + } + + maskState.stateInfo = newMaskInfo; + + if (ptrState.isPermuted && !applyPermuteOnMask()) { + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "MaskAnalysis: failed to apply permute on mask.\n"; + ptrState.dump(); + llvm::dbgs() << "oldMask:" << oldMask << "\n"; + maskState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + InFlightDiagnostic diag = + emitWarning(loc) << "MaskAnalysis: failed to apply permute on mask."; + }); + return nullptr; + } + + LLVM_DEBUG({ + llvm::dbgs() << "After matching MaskState: \n"; + for (auto info : newMaskInfo) { + info.dump(); + } + llvm::dbgs() << "----------------------------------------------\n"; + }); + + auto newMask = maskState.createNewMask(loc, rewriter); + return newMask; +} + +Value MemOpTransformer::createNewOther(Value oldOther, const Location loc, + PatternRewriter &rewriter) { + if (!oldOther || !maskState.newMask) + return nullptr; + + auto ptrType = dyn_cast(ptrState.source.getType()); + if (!ptrType) { + LLVM_DEBUG( + { + InFlightDiagnostic diag = + emitWarning(loc) + << "PtrAnalysis: source of ptrState is not a pointer type."; + }); + return nullptr; + } + Type elementType = ptrType.getPointeeType(); + + SmallVector targetShape; + for (auto info : maskState.stateInfo) { + auto staticShape = getIntAttr(info.shape); + if (!staticShape.has_value()) { + LLVM_DEBUG({ + InFlightDiagnostic diag = + emitWarning(loc) + << "MaskAnalysis: dynamic shape is not supported in reshape\n"; + }); + return oldOther; + } + targetShape.emplace_back(staticShape.value()); + } + auto targetShapeAttr = DenseIntElementsAttr::get( + RankedTensorType::get({static_cast(targetShape.size())}, + rewriter.getI64Type()), + targetShape); + auto targetShapeType = RankedTensorType::get(targetShape, elementType); + auto targetShapeValue = + rewriter.create(loc, targetShapeAttr); + + auto reshapeOp = rewriter.create( + loc, targetShapeType, oldOther, targetShapeValue); + + return reshapeOp.getResult(); +} + +bool MemOpTransformer::applyPermuteOnMask() { + if (!ptrState.isPermuted || maskState.isEmpty()) { + return true; + } + if (ptrState.permuteIds.size() != maskState.stateInfo.size()) { + return false; + } + SmallVector newMaskInfo; + for (auto id : ptrState.permuteIds) { + newMaskInfo.push_back(maskState.stateInfo[id]); + } + maskState.stateInfo = newMaskInfo; + return true; +} + +hivm::CreateSyncBlockLockOp createSyncBlockLockVar(OpBuilder &builder, + Location loc) { + SmallVector shape = {1}; + auto elementType = builder.getI64Type(); + Type memrefType = MemRefType::get(shape, elementType); + + auto createSyncBlockLockOp = + builder.create(loc, memrefType, Value()); + return createSyncBlockLockOp; +} +} // namespace MemOpConverter diff --git a/compiler/lib/TritonToStructured/PtrAnalysis.cpp b/compiler/lib/TritonToStructured/PtrAnalysis.cpp new file mode 100644 index 00000000..6951ff57 --- /dev/null +++ b/compiler/lib/TritonToStructured/PtrAnalysis.cpp @@ -0,0 +1,1557 @@ + + +#include "dicp/TritonToStructured/PtrAnalysis.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Value.h" +#include "mlir/IR/ValueRange.h" +#include "mlir/IR/Visitors.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Support/LogicalResult.h" + +#include "mlir/IR/IRMapping.h" +#include "mlir/Transforms/DialectConversion.h" + +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" + +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/LogicalResult.h" + +#include "dicp/Utils/Utils.h" + +#define DEBUG_TYPE "triton-to-structured-ptr-analysis" + +namespace TritonToStructured { +using namespace mlir; +using namespace triton; + +bool isMultiple(const OpFoldResult ÷nd, const OpFoldResult &divisor) { + auto staticDividend = getIntAttr(dividend); + auto staticDivisor = getIntAttr(divisor); + if (!staticDividend || !staticDivisor) { + return false; + } + return staticDividend.value() % staticDivisor.value() == 0; +} + +bool isEqual(const OpFoldResult &ofr1, const OpFoldResult &ofr2) { + auto staticOfr1 = getIntAttr(ofr1); + auto staticOfr2 = getIntAttr(ofr2); + return staticOfr1 == staticOfr2; +} + +bool isLess(const OpFoldResult &ofr1, const OpFoldResult &ofr2) { + auto staticOfr1 = getIntAttr(ofr1); + auto staticOfr2 = getIntAttr(ofr2); + // When sorting for permute, the value determined at runtime + // is greater than the value determined at compile time. + if (!staticOfr1 && !staticOfr2) { + return false; // keep relative order (stable_sort) + } + if (!staticOfr1 && staticOfr2) { + return false; // dynamic > static + } + if (staticOfr1 && !staticOfr2) { + return true; // static < dynamic + } + return staticOfr1.value() < staticOfr2.value(); +} + +bool isGreater(const OpFoldResult &ofr1, const OpFoldResult &ofr2) { + auto staticOfr1 = getIntAttr(ofr1); + auto staticOfr2 = getIntAttr(ofr2); + // When sorting for permute, the value determined at runtime + // is greater than the value determined at compile time. + if (!staticOfr1 && !staticOfr2) { + return false; // keep relative order (stable_sort) + } + if (!staticOfr1 && staticOfr2) { + return true; // dynamic > static + } + if (staticOfr1 && !staticOfr2) { + return false; // static < dynamic + } + return staticOfr1.value() > staticOfr2.value(); +} + +void StateInfo::dump() const { + llvm::dbgs() << "StateInfo: \n"; + llvm::dbgs() << "dimIndex = " << dimIndex << "\n"; + llvm::dbgs() << "shape = " << shape << "\n"; + llvm::dbgs() << "stride = " << stride << "\n"; +} + +void PtrState::dump() const { + llvm::dbgs() << "PtrState: \n"; + llvm::dbgs() << "source:" << source << "\n"; + llvm::dbgs() << "scalar:" << offset << "\n"; + llvm::dbgs() << "size: ["; + for (auto size : sizes) + llvm::dbgs() << size << ", "; + llvm::dbgs() << "]\n"; + llvm::dbgs() << "shouldLinearize: " << shouldLinearize << "\n"; + llvm::dbgs() << "isPermuted: " << isPermuted << "\n"; + llvm::dbgs() << "isBlockPtr: " << isBlockPtr() << "\n"; + + llvm::dbgs() << "permuteIds: ["; + for (auto id : permuteIds) + llvm::dbgs() << id << ", "; + llvm::dbgs() << "]\n"; + llvm::dbgs() << "order: ["; + for (auto id : order) + llvm::dbgs() << id << ", "; + llvm::dbgs() << "]\n"; + + llvm::dbgs() << "stateInfo:\n"; + llvm::dbgs() << "\n"; + for (auto info : stateInfo) { + llvm::dbgs() << "-----------------------------------------\n"; + info.dump(); + llvm::dbgs() << "-----------------------------------------\n"; + } +} + +bool PtrState::isEmpty() const { + return (stateInfo.empty() && !source && !offset); +} + +bool PtrState::isScalar() const { + bool scalar = true; + for (auto info : stateInfo) { + auto staticStride = getIntAttr(info.stride); + if (!staticStride.has_value() || staticStride.value() != 0) + scalar = false; + } + return scalar && (offset || source); +} + +bool PtrState::hasSource() const { return source != nullptr; } + +bool PtrState::isBlockPtr() const { return !order.empty(); } + +bool PtrState::isSameSizeAs(const PtrState &x) const { + if (this->sizes.size() != x.sizes.size()) + return false; + + for (size_t i = 0; i < this->sizes.size(); ++i) { + if (this->sizes[i] != x.sizes[i]) + return false; + } + return true; +} + +void PtrState::updatePtrState(SmallVector stateInfo, + SmallVector sizes, Value source, + OpFoldResult offset, const Location loc, + OpBuilder &builder, bool shouldLinearize) { + this->stateInfo = stateInfo; + this->sizes = sizes; + this->source = source; + this->offset = offset; + this->shouldLinearize = shouldLinearize; + this->normalizeState(loc, builder); +} + +void PtrState::normalizeState(const Location loc, OpBuilder &builder) { + SmallVector newStateInfo; + auto zeroAttr = builder.getIndexAttr(0); + + // merge continuous zero strides + // e.g., stride [0, 0, 1] shape [4, 32, 16] --> stride [0, 1] shape [128, 16] + for (auto it = this->stateInfo.begin(); it != this->stateInfo.end(); ++it) { + while (it != this->stateInfo.end() && isZero(it->stride)) { + auto newShape = it->shape; + auto dimIndex = it->dimIndex; + for (++it; it != this->stateInfo.end() && isZero(it->stride) && + it->dimIndex == dimIndex; + ++it) { + newShape = mulOpFoldResult(newShape, it->shape, loc, builder); + } + newStateInfo.emplace_back(zeroAttr, newShape, dimIndex); + } + if (it == this->stateInfo.end()) + break; + // if the info is the only one with oriSize 1 in this dimension, skip it + // e.g., stride [0, 1] shape [1, 128] sizes [1, 128] do not delete the first + // info + if (isOne(it->shape) && !isOne(sizes[it->dimIndex])) + continue; + newStateInfo.emplace_back(*it); + } + + this->stateInfo = newStateInfo; +} + +LogicalResult PtrAnalysis::visitOperandAddptr(triton::AddPtrOp addptrOp, + PtrState &state, + const Location loc, + OpBuilder &builder) { + if (!state.isEmpty()) { + LLVM_DEBUG({ + addptrOp.emitError( + "PtrAnalysis: PtrState should be empty when visiting addptr"); + }); + return failure(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Visit addptr operation: " << addptrOp << "\n"; + }); + + PtrState ptrState; + if (visitOperand(addptrOp.getPtr(), ptrState, addptrOp.getLoc(), builder) + .failed()) { + return failure(); + } + + PtrState offsetState; + if (visitOperand(addptrOp.getOffset(), offsetState, addptrOp.getLoc(), + builder) + .failed()) { + return failure(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Before visiting addptr operands: \n"; + llvm::dbgs() << "PtrState: \n"; + ptrState.dump(); + llvm::dbgs() << "OffsetState: \n"; + offsetState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + + if (!ptrState.source) { + LLVM_DEBUG({ + addptrOp.emitError("ptr field should provide source / base pointer"); + }); + return failure(); + } + return state.addState(ptrState, offsetState, addptrOp, builder); +} + +LogicalResult +PtrAnalysis::visitOperandMakeTensorPtr(triton::MakeTensorPtrOp makeTPtrOp, + PtrState &state, const Location loc, + OpBuilder &builder) { + if (!state.isEmpty()) { + makeTPtrOp.emitError( + "PtrAnalysis: PtrState should be empty when visiting make_tensor_ptr"); + return failure(); + } + if (makeTPtrOp.getOrder().empty()) { + LLVM_DEBUG(makeTPtrOp->emitRemark( + "PtrAnalysis: expect tt.make_tensor_ptr to have order field set")); + return failure(); + } + + // Build: + // - stateInfo: per-dimension (stride, shape, dimIndex) of the parent tensor + // - sizes: original tensor shape of the block + // - dimOffsets: the offset to the block in the parent tensor + state.source = makeTPtrOp.getBase(); + state.dimOffsets = makeTPtrOp.getOffsets(); + state.order = SmallVector(makeTPtrOp.getOrder()); + + auto resType = cast(makeTPtrOp.getResult().getType()); + auto pointeeType = cast(resType.getPointeeType()); + auto pointeeShape = pointeeType.getShape(); + const int64_t rank = pointeeType.getRank(); + + SmallVector newStateInfo; + for (int64_t i = 0; i < rank; i++) { + state.sizes.push_back(builder.getIndexAttr(pointeeShape[i])); + newStateInfo.emplace_back(makeTPtrOp.getStrides()[i], + makeTPtrOp.getShape()[i], i); + } + state.stateInfo = newStateInfo; + + assert(state.isBlockPtr() && + "tt.make_tensor_ptr pointer state should describe a block pointer"); + + return success(); +} + +bool PtrAnalysis::operandIsScalar(Value operand) { + auto tensorType = dyn_cast(operand.getType()); + auto elementType = + tensorType ? tensorType.getElementType() : operand.getType(); + bool isScalar = true; + if (tensorType) { + for (size_t i = 0; i < tensorType.getRank() && isScalar; ++i) { + isScalar = tensorType.getDimSize(i) == 1; + } + } + return isScalar && + (isa(elementType) || isa(elementType)); +} + +LogicalResult PtrAnalysis::initStateByScalar(Value operand, PtrState &state, + const Location loc, + OpBuilder &builder) { + OpFoldResult newOffset; + SmallVector newSizes; + SmallVector newStateInfo; + if (isa(operand.getType())) { + OpBuilder::InsertionGuard guard(builder); + if (!isa(operand) && operand.getDefiningOp()) { + builder.setInsertionPointAfter(operand.getDefiningOp()); + } + auto castOp = builder.create( + loc, builder.getIndexType(), operand); + newOffset = castOp.getResult(); + } else if (isa(operand.getType())) { + newOffset = operand; + } else { + auto tensorType = dyn_cast(operand.getType()); + auto index = builder.create(loc, 0); + auto zeroAttr = builder.getIndexAttr(0); + auto oneAttr = builder.getIndexAttr(1); + SmallVector indices; + for (size_t i = 0; i < tensorType.getRank(); ++i) { + indices.push_back(index); + newSizes.emplace_back(oneAttr); + newStateInfo.emplace_back(zeroAttr, oneAttr, i); + } + auto extractedElement = + builder.create(loc, operand, indices); + newOffset = extractedElement.getResult(); + } + state.updatePtrState(newStateInfo, newSizes, nullptr, newOffset, loc, + builder); + return success(); +} + +LogicalResult PtrAnalysis::initStateByPointer(Value operand, PtrState &state, + const Location loc, + OpBuilder &builder) { + Value newSource; + SmallVector newSizes; + SmallVector newStateInfo; + + if (auto op = operand.getDefiningOp()) { + if (auto addPtrOp = dyn_cast(op)) { + return visitOperandAddptr(cast(op), state, loc, + builder); + } else if (auto bitCastOp = dyn_cast(op)) { + newSource = operand; + } else if (auto makeTensorOp = dyn_cast(op)) { + LLVM_DEBUG({ + op->emitWarning("Unexpected operand defining operation tts.make_tptr."); + }); + return failure(); + } else if (auto intToPtrOp = dyn_cast(op)) { + newSource = operand; + } else { + LLVM_DEBUG({ op->emitWarning("PtrAnalysis: Unexpected operand."); }); + return failure(); + } + } else { + newSource = operand; + } + auto newOffset = builder.getIndexAttr(0); + state.updatePtrState(newStateInfo, newSizes, newSource, newOffset, loc, + builder); + return success(); +} + +LogicalResult PtrState::mulState(const PtrState &lhsState, + const PtrState &rhsState, Operation *op, + OpBuilder &builder) { + auto loc = op->getLoc(); + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "mulState: " << op << "\n"; + }); + + if (!isEmpty()) { + LLVM_DEBUG({ + op->emitError("PtrAnalysis: PtrState should be empty when multiplying"); + }); + return failure(); + } + + // neither lhs nor rhs should have source, since multiplying base pointer + // does not make sense + if (lhsState.hasSource() || rhsState.hasSource()) { + LLVM_DEBUG({ + op->emitError("PtrAnalysis: do not support base inters in multiplying"); + }); + return failure(); + } else if (!lhsState.isScalar() && !rhsState.isScalar()) { + // do not support both tensors are effectively non-scalar + LLVM_DEBUG({ + op->emitError( + "PtrAnalysis: only support multiplying pointer states when one of " + "them represent a scalar"); + }); + return failure(); + } + + PtrState const *lhs = &lhsState; + PtrState const *rhs = &rhsState; + + if (!rhs->isScalar() && lhs->isScalar()) { + std::swap(lhs, rhs); + } + + SmallVector newStateInfo; + for (auto info : lhs->stateInfo) { + OpFoldResult newStride = + mulOpFoldResult(info.stride, rhs->offset, loc, builder); + newStateInfo.emplace_back(newStride, info.shape, info.dimIndex); + } + + auto newOffset = + mulOpFoldResult(lhsState.offset, rhsState.offset, loc, builder); + updatePtrState(newStateInfo, lhs->sizes, lhs->source, newOffset, loc, builder, + lhs->shouldLinearize); + + LLVM_DEBUG({ + llvm::dbgs() << "After mulState: \n"; + this->dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + + return success(); +} + +LogicalResult PtrState::subState(const PtrState &lhsState, + const PtrState &rhsState, Operation *op, + OpBuilder &builder) { + auto loc = op->getLoc(); + if (!isEmpty()) { + LLVM_DEBUG({ + op->emitError("PtrAnalysis: PtrState should be empty when subtracting"); + }); + return failure(); + } + + if (lhsState.hasSource() && rhsState.hasSource()) { + LLVM_DEBUG({ + op->emitError( + "PtrAnalysis: do not support both sides have base pointers in sub"); + }); + return failure(); + } + + if (!rhsState.isScalar()) { + LLVM_DEBUG({ + op->emitError("PtrAnalysis: only support sub when one of " + "them represents a scalar"); + }); + return failure(); + } + + auto newOffset = + subOpFoldResult(lhsState.offset, rhsState.offset, loc, builder); + updatePtrState(lhsState.stateInfo, lhsState.sizes, lhsState.source, newOffset, + loc, builder, lhsState.shouldLinearize); + + return success(); +} + +LogicalResult PtrState::addState(PtrState &lhsState, PtrState &rhsState, + Operation *op, OpBuilder &builder) { + auto loc = op->getLoc(); + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "addState: " << op << "\n"; + }); + + if (!isEmpty()) { + LLVM_DEBUG({ + op->emitError("PtrAnalysis: PtrState should be empty when adding"); + }); + return failure(); + } + if (!lhsState.isSameSizeAs(rhsState)) { + LLVM_DEBUG({ + op->emitError( + "PtrAnalysis: The original size of the addition should be the same"); + }); + return failure(); + } + + SmallVector newStateInfo; + auto lIt = lhsState.stateInfo.begin(); + auto rIt = rhsState.stateInfo.begin(); + while (lIt != lhsState.stateInfo.end() && rIt != rhsState.stateInfo.end()) { + if (lIt->dimIndex != rIt->dimIndex) { + auto newInfo = lIt->dimIndex < rIt->dimIndex ? *lIt++ : *rIt++; + newStateInfo.emplace_back(newInfo); + continue; + } + if (!isMultiple(lIt->shape, rIt->shape) && + !isMultiple(rIt->shape, lIt->shape)) { + LLVM_DEBUG({ + llvm::dbgs() << "LHS PtrState: \n"; + lhsState.dump(); + llvm::dbgs() << "RHS PtrState: \n"; + rhsState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + LLVM_DEBUG({ + op->emitError("PtrAnalysis: the add operation have incompatible sizes"); + }); + return failure(); + } + + auto newShape = minOpFoldResult(lIt->shape, rIt->shape, loc, builder); + if ((isLess(newShape, lIt->shape) && !isZero(lIt->stride) || + isLess(newShape, rIt->shape) && !isZero(rIt->stride))) { + LLVM_DEBUG({ + llvm::dbgs() << "LHS PtrState: \n"; + lhsState.dump(); + llvm::dbgs() << "RHS PtrState: \n"; + rhsState.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + LLVM_DEBUG({ + op->emitError("PtrAnalysis: the add operation have incompatible sizes." + "Valid dimensions are split."); + }); + return failure(); + } + + auto newStride = addOpFoldResult(lIt->stride, rIt->stride, loc, builder); + newStateInfo.emplace_back(newStride, newShape, lIt->dimIndex); + + if (isEqual(lIt->shape, newShape)) + ++lIt; + else + lIt->shape = divOpFoldResult(lIt->shape, newShape, loc, builder); + if (isEqual(rIt->shape, newShape)) + ++rIt; + else + rIt->shape = divOpFoldResult(rIt->shape, newShape, loc, builder); + } + + while (rIt != rhsState.stateInfo.end()) { + newStateInfo.push_back(*rIt++); + } + while (lIt != lhsState.stateInfo.end()) { + newStateInfo.push_back(*lIt++); + } + + auto newSource = source = lhsState.source ? lhsState.source : rhsState.source; + auto newOffset = + addOpFoldResult(lhsState.offset, rhsState.offset, loc, builder); + auto newShouldLinearize = + lhsState.shouldLinearize || rhsState.shouldLinearize; + auto newSizes = lhsState.sizes; + + updatePtrState(newStateInfo, newSizes, newSource, newOffset, loc, builder, + newShouldLinearize); + + LLVM_DEBUG({ + llvm::dbgs() << "After addState: \n"; + this->dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + + return success(); +} + +triton::AddPtrOp PtrState::createAddPtrOp(OpBuilder &builder, Location loc) { + SmallVector tensorSizes; + SmallVector tensorStrides; + + auto zeroAttr = builder.getIndexAttr(0); + auto oneAttr = builder.getIndexAttr(1); + + for (auto id : permuteIds) { + auto info = stateInfo[id]; + if (isZero(info.stride)) + continue; + tensorStrides.emplace_back(info.stride); + tensorSizes.emplace_back(getIntAttr(info.shape).value()); + } + + // load a scalar pointer + if (tensorSizes.empty()) { + Value offsetValue = materializeValue(builder, loc, offset); + if (offsetValue.getType().isIndex()) { + offsetValue = builder.create( + loc, builder.getI32Type(), offsetValue); + } + auto addptrOp = builder.create(loc, source.getType(), + source, offsetValue); + return addptrOp; + } + + SmallVector cachedRange; + auto ptrType = cast(source.getType()); + auto ptrTensorType = RankedTensorType::get({tensorSizes}, ptrType); + auto broadCastType = + RankedTensorType::get({tensorSizes}, builder.getI32Type()); + + if (tensorSizes.size() != tensorStrides.size()) { + LLVM_DEBUG( + { + InFlightDiagnostic diag = + emitError(loc) + << "PtrAnalysis: inconsistent tensor sizes and strides"; + }); + return nullptr; + } + for (size_t i = 0; i < tensorSizes.size(); ++i) { + // make range + auto indexI32RowType = + RankedTensorType::get({tensorSizes[i]}, builder.getI32Type()); + Value makeRangeOp = builder.create( + loc, indexI32RowType, 0, tensorSizes[i]); + + // multiply stride + Value strideValue = materializeValue(builder, loc, tensorStrides[i]); + if (strideValue.getType().isIndex()) { + strideValue = builder.create( + loc, builder.getI32Type(), strideValue); + } + Value splatStride = + builder.create(loc, indexI32RowType, strideValue); + auto rangeAfterMul = + builder.create(loc, makeRangeOp, splatStride); + + // reshape + Value expandedValue = rangeAfterMul; + for (size_t j = 0; j < tensorSizes.size(); ++j) { + if (j == i) + continue; + expandedValue = + builder.create(loc, expandedValue, j); + } + + // broadcast + auto broadcastValue = + builder.create(loc, broadCastType, expandedValue); + cachedRange.push_back(broadcastValue); + } + + // combine the cachedRange + Value rangeAfterCombine = cachedRange[0]; + for (size_t i = 1; i < cachedRange.size(); ++i) { + rangeAfterCombine = + builder.create(loc, rangeAfterCombine, cachedRange[i]); + } + + // addOffset + Value addValue = materializeValue(builder, loc, offset); + if (addValue.getType().isIndex()) { + addValue = + builder.create(loc, builder.getI32Type(), addValue); + } + Value splatOffset = + builder.create(loc, broadCastType, addValue); + auto rangeAfterAdd = + builder.create(loc, rangeAfterCombine, splatOffset); + + // addPtr + Value splatPtr = builder.create(loc, ptrTensorType, source); + auto addptrOp = builder.create(loc, ptrTensorType, splatPtr, + rangeAfterAdd); + return addptrOp; +} + +triton::MakeTensorPtrOp PtrState::createMakeTensorPtrOp(OpBuilder &builder, + Location loc) { + SmallVector newShape; + SmallVector newStrides; + SmallVector newOffsets; + SmallVector newBlkShape; + SmallVector newOrder; // must be int32_t for MakeTensorPtrOp builder + + const size_t rank = order.size(); + if (rank == 0) { + emitError(loc) << "PtrAnalysis: empty order in createMakeTensorPtrOp"; + return nullptr; + } + + // iterate reversed safely: i = rank-1, ..., 0 + for (size_t i = rank; i-- > 0;) { + size_t dim = order[i]; + + if (dim >= stateInfo.size() || dim >= dimOffsets.size() || + dim >= sizes.size()) { + emitError(loc) + << "PtrAnalysis: invalid dim index in createMakeTensorPtrOp"; + return nullptr; + } + + auto info = stateInfo[dim]; + + newShape.push_back(materializeValue(builder, loc, info.shape)); + newStrides.push_back(materializeValue(builder, loc, info.stride)); + newOffsets.push_back(materializeValue(builder, loc, dimOffsets[dim])); + + auto blkSzOpt = getIntAttr(sizes[dim]); + if (!blkSzOpt.has_value()) { + emitError(loc) << "PtrAnalysis: dynamic block_shape is not supported for " + "tt.make_tensor_ptr"; + return nullptr; + } + newBlkShape.push_back(static_cast(blkSzOpt.value())); + newOrder.push_back(static_cast(i)); + } + + return builder.create( + loc, source, ValueRange(newShape), ValueRange(newStrides), + ValueRange(newOffsets), newBlkShape, newOrder); +} + +LogicalResult PtrAnalysis::visitOperandMul(arith::MulIOp mulOp, PtrState &state, + const Location loc, + OpBuilder &builder) { + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Visit Mul operation: " << mulOp << "\n"; + }); + + PtrState lhsState; + if (visitOperand(mulOp.getLhs(), lhsState, loc, builder).failed()) { + return failure(); + } + + PtrState rhsState; + if (visitOperand(mulOp.getRhs(), rhsState, loc, builder).failed()) { + return failure(); + } + + return state.mulState(lhsState, rhsState, mulOp, builder); +} + +LogicalResult PtrAnalysis::visitOperandSub(arith::SubIOp subOp, PtrState &state, + const Location loc, + OpBuilder &builder) { + PtrState lhsState; + if (visitOperand(subOp.getLhs(), lhsState, loc, builder).failed()) { + return failure(); + } + + PtrState rhsState; + if (visitOperand(subOp.getRhs(), rhsState, loc, builder).failed()) { + return failure(); + } + + return state.subState(lhsState, rhsState, subOp, builder); +} + +LogicalResult PtrAnalysis::visitOperandMakeRange(triton::MakeRangeOp rangeOp, + PtrState &state, Location loc, + OpBuilder &builder) { + if (!state.isEmpty()) { + LLVM_DEBUG({ + rangeOp.emitError( + "PtrAnalysis: PtrState should be empty when visiting make_range"); + }); + return failure(); + } + + auto shape = cast(rangeOp.getType()).getShape(); + + auto start = rangeOp.getStart(); + auto end = rangeOp.getEnd(); + auto stride = (end - start + shape[0] - 1) / shape[0]; + if (stride != 1) { + LLVM_DEBUG({ + rangeOp.emitError( + "PtrAnalysis: make_range op with stride != 1 is not supported"); + }); + return failure(); + } + + auto infoStride = builder.getIndexAttr(stride); + auto size = builder.getIndexAttr(shape[0]); + auto offset = builder.getIndexAttr(start); + + SmallVector stateInfo; + SmallVector sizes; + stateInfo.emplace_back(infoStride, size); + sizes.emplace_back(size); + + state.updatePtrState(stateInfo, sizes, nullptr, offset, loc, builder); + return success(); +} + +LogicalResult +PtrAnalysis::visitOperandBroadcast(triton::BroadcastOp broadcastOp, + PtrState &state, const Location loc, + OpBuilder &builder) { + if (!state.isEmpty()) { + LLVM_DEBUG({ + broadcastOp.emitError( + "PtrAnalysis: PtrState should be empty when visiting broadcast"); + }); + return failure(); + } + + auto src = broadcastOp.getSrc(); + auto dst = broadcastOp.getResult(); + if (!isa(dst.getType())) { + LLVM_DEBUG({ + broadcastOp.emitRemark( + "PtrAnalysis: broadcast dst should be a shaped type"); + }); + return failure(); + } + + auto srcShape = cast(src.getType()).getShape(); + auto dstShape = cast(dst.getType()).getShape(); + if (srcShape.size() != dstShape.size()) { + LLVM_DEBUG({ + broadcastOp.emitRemark( + "PtrAnalysis: broadcast src and dst should have the same rank"); + }); + return failure(); + } + if (visitOperand(src, state, loc, builder).failed()) { + return failure(); + } + + if (state.sizes.size() != dstShape.size()) { + llvm::dbgs() << broadcastOp << "\n"; + state.dump(); + llvm::dbgs() << "dst.size = " << dstShape.size() << "\n"; + for (auto x : dstShape) + llvm::dbgs() << x << ", "; + llvm::dbgs() << "\n"; + } + + SmallVector newStateInfo(state.stateInfo); + SmallVector newSizes; + if (srcShape.size() != dstShape.size()) { + LLVM_DEBUG({ + broadcastOp.emitRemark( + "PtrAnalysis: unexpected state info size in broadcast"); + }); + return failure(); + } + for (size_t i = 0; i < dstShape.size(); ++i) { + newSizes.emplace_back(builder.getIndexAttr(dstShape[i])); + if (srcShape[i] == dstShape[i]) { + continue; + } else if (srcShape[i] < dstShape[i] && srcShape[i] == 1) { + for (auto &info : newStateInfo) { + if (info.dimIndex != i) + continue; + info.shape = builder.getIndexAttr(dstShape[i]); + } + } else { + LLVM_DEBUG({ + broadcastOp.emitRemark("unexpected dimensions used in broadcast"); + }); + return failure(); + } + } + state.updatePtrState(newStateInfo, newSizes, state.source, state.offset, loc, + builder, state.shouldLinearize); + return success(); +} + +LogicalResult PtrAnalysis::visitOperandSplat(triton::SplatOp splatOp, + PtrState &state, + const Location loc, + OpBuilder &builder) { + if (!state.isEmpty()) { + LLVM_DEBUG({ + splatOp.emitError( + "PtrAnalysis: PtrState should be empty when visiting splat"); + }); + return failure(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Visit SPLAT operation: " << splatOp << "\n"; + }); + + auto src = splatOp.getSrc(); + auto dst = splatOp.getResult(); + auto dstShape = cast(dst.getType()).getShape(); + + if (visitOperand(src, state, loc, builder).failed()) { + return failure(); + } + + LLVM_DEBUG({ + llvm::dbgs() << "splat ptrState: \n"; + state.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + + if (!state.isScalar()) { + LLVM_DEBUG( + { splatOp.emitRemark("PtrAnalysis: splat source should be scalar"); }); + return failure(); + } + + SmallVector newStateInfo; + SmallVector newSizes; + auto zeroAttr = builder.getIndexAttr(0); + if (isa(src.getType())) { + for (size_t i = 0; i < dstShape.size(); ++i) { + auto currentSize = builder.getIndexAttr(dstShape[i]); + newSizes.emplace_back(currentSize); + newStateInfo.emplace_back(zeroAttr, currentSize, i); + } + } else { + LLVM_DEBUG( + { splatOp.emitRemark("PtrAnalysis: unsupported splat pattern"); }); + return failure(); + } + state.updatePtrState(newStateInfo, newSizes, state.source, state.offset, loc, + builder, state.shouldLinearize); + + LLVM_DEBUG({ + llvm::dbgs() << "After SPLAT ptrState: \n"; + state.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + return success(); +} + +LogicalResult +PtrAnalysis::visitOperandExpandDims(triton::ExpandDimsOp expandDimsOp, + PtrState &state, const Location loc, + OpBuilder &builder) { + if (!state.isEmpty()) { + LLVM_DEBUG({ + expandDimsOp.emitError( + "PtrAnalysis: PtrState should be empty when visiting expand_dims"); + }); + return failure(); + } + + if (visitOperand(expandDimsOp.getSrc(), state, loc, builder).failed()) { + return failure(); + } + + auto dstShape = + cast(expandDimsOp.getResult().getType()).getShape(); + auto axis = expandDimsOp.getAxis(); + + SmallVector newStateInfo(state.stateInfo); + SmallVector newSizes(state.sizes); + size_t insertPos = 0; + for (auto &info : newStateInfo) { + if (info.dimIndex >= axis) + ++info.dimIndex; + if (info.dimIndex < axis) + ++insertPos; + } + auto zeroAttr = builder.getIndexAttr(0); + auto oneAttr = builder.getIndexAttr(1); + StateInfo insertInfo(zeroAttr, oneAttr, axis); + + newStateInfo.insert(newStateInfo.begin() + insertPos, insertInfo); + newSizes.insert(newSizes.begin() + axis, oneAttr); + + state.updatePtrState(newStateInfo, newSizes, state.source, state.offset, loc, + builder, state.shouldLinearize); + return success(); +} + +LogicalResult PtrAnalysis::visitOperandConstSplat(arith::ConstantOp op, + PtrState &state, + const Location loc, + OpBuilder &builder) { + if (!state.isEmpty()) { + LLVM_DEBUG({ + op->emitError( + "PtrAnalysis: PtrState should be empty when visiting const_splat"); + }); + return failure(); + } + + auto attr = cast(op.getValue()); + auto elementType = attr.getElementType(); + if (!attr.isSplat() || !isa(elementType)) { + LLVM_DEBUG( + { op->emitError("PtrAnalysis: only support splat integer constant"); }); + return failure(); + } + + auto value = attr.getValues()[0].getValue(); + auto constAttr = builder.getIndexAttr(value.getSExtValue()); + + auto resultShape = cast(op.getResult().getType()).getShape(); + + SmallVector sizes; + SmallVector stateInfo; + auto defaultAttr = builder.getIndexAttr(0); + + for (auto [i, shape] : llvm::enumerate(resultShape)) { + auto shapeAttr = builder.getIndexAttr(shape); + sizes.emplace_back(shapeAttr); + stateInfo.emplace_back(defaultAttr, shapeAttr, i); + } + + state.updatePtrState(stateInfo, sizes, nullptr, constAttr, loc, builder); + return success(); +} + +LogicalResult PtrAnalysis::visitOperandExtSI(arith::ExtSIOp extOp, + PtrState &state, + const Location loc, + OpBuilder &builder) { + if (!state.isEmpty()) { + LLVM_DEBUG({ + extOp.emitError( + "PtrAnalysis: PtrState should be empty when visiting extsi"); + }); + return failure(); + } + + if (visitOperand(extOp.getIn(), state, loc, builder).failed()) { + return failure(); + } + + return success(); +} + +LogicalResult PtrAnalysis::visitOperandRem(arith::RemSIOp remOp, + PtrState &state, const Location loc, + OpBuilder &builder) { + if (!state.isEmpty()) { + LLVM_DEBUG({ + remOp.emitError( + "PtrAnalysis: PtrState should be empty when visiting remsi"); + }); + return failure(); + } + LLVM_DEBUG({ + llvm::dbgs() << "before VisitRemOperands \n"; + state.dump(); + }); + + PtrState rhsState; + if (visitOperand(remOp.getRhs(), rhsState, loc, builder).failed()) { + return failure(); + } + + if (!rhsState.isScalar() || rhsState.hasSource()) { + LLVM_DEBUG({ + remOp.emitRemark("PtrAnalysis: only support cases when rhs of remainder " + "contains scalar"); + }); + return failure(); + } + + if (visitOperand(remOp.getLhs(), state, loc, builder).failed()) { + return failure(); + } + + bool hasAnnotation = optimizeDynamicOffset; + + auto zeroAttr = builder.getIndexAttr(0); + auto oneAttr = builder.getIndexAttr(1); + auto divisorAttr = rhsState.offset; + + auto staticOffset = getIntAttr(state.offset); + if ((!staticOffset.has_value() || !isMultiple(state.offset, divisorAttr)) && + !hasAnnotation) { + LLVM_DEBUG({ + remOp.emitRemark( + "PtrAnalysis: dynamic offset before REMSI, adding annotation"); + }); + return failure(); + } + + if (!getIntAttr(divisorAttr).has_value()) { + LLVM_DEBUG({ + remOp.emitError("PtrAnalysis: do not support dynamix divisor in REMSI."); + }); + return failure(); + } + + SmallVector newStateInfo; + for (auto info : state.stateInfo) { + if (!getIntAttr(info.stride).has_value()) { + LLVM_DEBUG({ + remOp.emitError( + "PtrAnalysis: do not support dynamix stride before REMSI."); + }); + return failure(); + } + + if (!isMultiple(divisorAttr, info.shape) && + !isMultiple(info.shape, divisorAttr)) { + LLVM_DEBUG({ + remOp.emitError( + "PtrAnalysis: do not support dynamix stride before REMSI."); + }); + } + + if (isMultiple(info.stride, divisorAttr)) { + newStateInfo.emplace_back(zeroAttr, info.shape, info.dimIndex); + } else if (isMultiple(divisorAttr, info.stride)) { + auto contiguousSize = + divOpFoldResult(divisorAttr, info.stride, loc, builder); + contiguousSize = + minOpFoldResult(contiguousSize, info.shape, loc, builder); + auto nonContiguousSize = + divOpFoldResult(info.shape, contiguousSize, loc, builder); + + auto staticNonContiguousSize = getIntAttr(nonContiguousSize); + if (!staticNonContiguousSize.has_value()) { + LLVM_DEBUG({ + remOp.emitError( + "PtrAnalysis: do not support dynamix size before REMSI."); + }); + return failure(); + } + + if (staticNonContiguousSize.value() > 1) + newStateInfo.emplace_back(zeroAttr, nonContiguousSize, info.dimIndex); + + newStateInfo.emplace_back(info.stride, contiguousSize, info.dimIndex); + } else { + LLVM_DEBUG({ + remOp.emitError("PtrAnalysis: stride that are not divisible by REMSI " + "are not allowed " + "to precede REMSI"); + }); + return failure(); + } + } + + auto newOffset = remOpFoldResult(state.offset, divisorAttr, loc, builder); + state.updatePtrState(newStateInfo, state.sizes, state.source, newOffset, loc, + builder, true); + + LLVM_DEBUG({ + llvm::dbgs() << "after VisitRemOperands \n"; + state.dump(); + }); + + return success(); +} + +LogicalResult PtrAnalysis::visitOperandDiv(arith::DivSIOp divOp, + PtrState &state, const Location loc, + OpBuilder &builder) { + if (!state.isEmpty()) { + LLVM_DEBUG({ + divOp.emitError( + "PtrAnalysis: PtrState should be empty when visiting divsi"); + }); + return failure(); + } + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "Visit DIVSI operation: " << divOp << "\n"; + }); + + PtrState rhsState; + if (visitOperand(divOp.getRhs(), rhsState, loc, builder).failed()) { + return failure(); + } + + if (!rhsState.isScalar() || rhsState.hasSource()) { + LLVM_DEBUG({ + divOp.emitRemark("PtrAnalysis: only support cases when rhs of remainder " + "contains scalar"); + }); + return failure(); + } + + if (visitOperand(divOp.getLhs(), state, loc, builder).failed()) { + return failure(); + } + + bool hasAnnotation = optimizeDynamicOffset; + + auto staticMultipleOf = extractDivisibilityFromOpFoldResult(state.offset); + if (!hasAnnotation && staticMultipleOf.has_value()) { + auto attr = builder.getIndexAttr(staticMultipleOf.value()); + hasAnnotation = isMultiple(attr, rhsState.offset); + } + + // add divState + auto zeroAttr = builder.getIndexAttr(0); + auto oneAttr = builder.getIndexAttr(1); + auto divisorAttr = rhsState.offset; + + auto staticOffset = getIntAttr(state.offset); + if ((!staticOffset.has_value() || !isMultiple(state.offset, divisorAttr)) && + !hasAnnotation) { + LLVM_DEBUG({ + divOp.emitRemark( + "PtrAnalysis: dynamic offset before DIVSI, adding annotation"); + }); + return failure(); + } + + if (!getIntAttr(divisorAttr).has_value()) { + LLVM_DEBUG({ + divOp.emitError("PtrAnalysis: do not support dynamix divisor in DIVSI."); + }); + return failure(); + } + + SmallVector newStateInfo; + for (auto info : state.stateInfo) { + auto staticStride = getIntAttr(info.stride); + if (!staticStride.has_value()) { + LLVM_DEBUG({ + divOp.emitError( + "PtrAnalysis: do not support dynamix stride before DIVSI."); + }); + return failure(); + } + + if (!isMultiple(divisorAttr, info.shape) && + !isMultiple(info.shape, divisorAttr)) { + LLVM_DEBUG({ + divOp.emitError( + "PtrAnalysis: do not support dynamix stride before DivSI."); + }); + } + + if (isMultiple(info.stride, divisorAttr)) { + auto newStride = divOpFoldResult(info.stride, divisorAttr, loc, builder); + newStateInfo.emplace_back(newStride, info.shape, info.dimIndex); + } else if (isMultiple(divisorAttr, info.stride)) { + auto nonContiguousSize = + divOpFoldResult(divisorAttr, info.stride, loc, builder); + nonContiguousSize = + minOpFoldResult(nonContiguousSize, info.shape, loc, builder); + auto contiguousSize = + divOpFoldResult(info.shape, nonContiguousSize, loc, builder); + + auto staticContiguousSize = getIntAttr(contiguousSize); + if (!staticContiguousSize.has_value()) { + LLVM_DEBUG({ + divOp.emitError( + "PtrAnalysis: do not support dynamix size before DIVSI."); + }); + return failure(); + } + + if (staticContiguousSize.value() != 0) + newStateInfo.emplace_back(oneAttr, contiguousSize, info.dimIndex); + + newStateInfo.emplace_back(zeroAttr, nonContiguousSize, info.dimIndex); + } else { + LLVM_DEBUG({ + divOp.emitError("PtrAnalysis: stride that are not divisible by DIVSI " + "are not allowed " + "to precede DIVSI"); + }); + return failure(); + } + } + + auto newOffset = divOpFoldResult(state.offset, divisorAttr, loc, builder); + state.updatePtrState(newStateInfo, state.sizes, state.source, newOffset, loc, + builder, true); + + LLVM_DEBUG({ + llvm::dbgs() << "after VisitDivOperands \n"; + state.dump(); + }); + return success(); +} + +LogicalResult PtrAnalysis::visitOperandAdd(arith::AddIOp addOp, PtrState &state, + const Location loc, + OpBuilder &builder) { + PtrState lhsState; + if (visitOperand(addOp.getLhs(), lhsState, loc, builder).failed()) { + return failure(); + } + + PtrState rhsState; + if (visitOperand(addOp.getRhs(), rhsState, loc, builder).failed()) + return failure(); + return state.addState(lhsState, rhsState, addOp, builder); +} + +LogicalResult PtrAnalysis::visitOperand(Value operand, PtrState &state, + const Location loc, + OpBuilder &builder) { + if (knownPtrs.find(operand) != knownPtrs.end()) { + state = knownPtrs.lookup(operand); + return success(); + } + + if (operandIsScalar(operand)) { + return initStateByScalar(operand, state, loc, builder); + } + + if (isa(operand.getType())) { + return initStateByPointer(operand, state, loc, builder); + } + + if (auto op = operand.getDefiningOp()) { + return visitOperandAdd(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandMul(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandSub(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandMakeRange(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandBroadcast(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandSplat(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandExpandDims(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandAddptr(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandConstSplat(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandRem(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandDiv(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + return visitOperandExtSI(op, state, loc, builder); + } else if (auto op = operand.getDefiningOp()) { + LLVM_DEBUG({ + op.emitRemark("TritonToStructured: Invalid dynamic offset" + "The load operation's offset cannot be derived from " + "another load result."); + }); + return failure(); + } else if (auto op = operand.getDefiningOp()) { + LLVM_DEBUG({ + op.emitWarning("IllegalTypeConversionInAddressCalculation" + "float-to-int precision conversion is not supported " + "during address computation."); + llvm::dbgs() << "Operand: \n"; + operand.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + return failure(); + } else if (!operand.getDefiningOp()) { + if (!knownPtrs.contains(operand)) { + LLVM_DEBUG({ + llvm::dbgs() << "TritonToStructured: Pointer analysis is not supported " + "for input parameters\n"; + }); + return failure(); + } + + // This operand must be an iter-arg of an inner-loop in a multiple-level + // nested loop, which means its PtrState must have already been populated + // during rewriteForOp of the parent loop. + state = knownPtrs[operand]; + return success(); + } else { + auto op = operand.getDefiningOp(); + LLVM_DEBUG({ + op->emitWarning("TritonToStructured: encountered addptr operand produced " + "by an unsupported operation"); + llvm::dbgs() << "Operand: \n"; + operand.dump(); + llvm::dbgs() << "----------------------------------------------\n"; + }); + return failure(); + } + return success(); +} + +LogicalResult PtrAnalysis::rewriteAddptrOp(triton::AddPtrOp op) { + OpBuilder builder(op); + auto loc = op.getLoc(); + + PtrState state; + if (visitOperandAddptr(op, state, op.getLoc(), builder).failed()) { + return failure(); + } + + auto maketptrOp = state.createAddPtrOp(builder, op.getLoc()); + knownPtrs[op.getResult()] = state; + ptrMap.map(op.getResult(), maketptrOp.getResult()); + + return success(); +} + +void PtrState::analyzePermute() { + const size_t n = stateInfo.size(); + generateOriginPermuteIds(); + + if (n <= 1) + return; + + // ============================================================ + // === 1. make_tensor_ptr (block ptr): order-based === + // Rule: permute if order is NOT canonical (i.e. strictly decreasing, + // representing inner-to-outer memory layout priority, e.g. [n-1, ..., 0]) + // + // NOTE: order describes memory layout priority, not axis permutation. + // Do NOT translate order into permuteIds. + // ============================================================ + if (isBlockPtr()) { + for (size_t i = 0; i + 1 < n; ++i) { + if (order[i] <= order[i + 1]) { + isPermuted = true; + break; + } + } + return; + } + + // ============================================================ + // === 2. addptr: stride-based permuteIds + contiguous-axis === + // Rule: permute if contiguous axes increased + // + // analyze constraints: must have at least one static stride + // ============================================================ + bool hasStatic = false; + for (auto &s : stateInfo) { + if (getIntAttr(s.stride).has_value()) { + hasStatic = true; + break; + } + } + if (!hasStatic) { + return; + } + auto isIdentity = [&](ArrayRef perm) -> bool { + for (size_t i = 0; i < perm.size(); ++i) + if (perm[i] != i) + return false; + return true; + }; + + std::stable_sort(permuteIds.begin(), permuteIds.end(), + [&](size_t a, size_t b) { + return isGreater(stateInfo[a].stride, stateInfo[b].stride); + }); + + // If already in canonical axis order, do not permute. + if (isIdentity(permuteIds)) { + return; + } + + // Tail axis must be physically contiguous (stride == 1), + // otherwise addptr-based permutation is invalid. + auto tailStride = getIntAttr(stateInfo[permuteIds.back()].stride); + if (!tailStride.has_value() || tailStride.value() != 1) { + generateOriginPermuteIds(); + return; + } + + // Compute new contiguous axes count using the permuted stateInfo. + SmallVector newStateInfo; + newStateInfo.reserve(n); + for (size_t id : permuteIds) { + newStateInfo.push_back(stateInfo[id]); + } + const size_t oldContig = countContiguousAxes(stateInfo); + const size_t newContig = countContiguousAxes(newStateInfo); + LLVM_DEBUG({ + llvm::dbgs() << "----------------------------------------------\n"; + llvm::dbgs() << "after analyzePermute:\n" + << "oldContig: " << oldContig << "\n" + << "newContig: " << newContig << "\n"; + dump(); + }); + // only permute if contiguous axes increased + if (newContig > oldContig) { + isPermuted = true; + return; + } + + // otherwise: no permute + generateOriginPermuteIds(); +} + +std::optional +extractDivisibilityFromOpFoldResult(mlir::OpFoldResult ofr) { + auto value = dyn_cast(ofr); + if (!value) { + return std::nullopt; + } + auto defOp = value.getDefiningOp(); + if (!defOp) { + return std::nullopt; + } + + auto divisibilityAttr = defOp->getAttr("tt.divisibility"); + if (!divisibilityAttr) { + return std::nullopt; + } + + auto denseAttr = dyn_cast(divisibilityAttr); + if (!denseAttr || denseAttr.empty()) { + return std::nullopt; + } + + return denseAttr.getValues()[0]; +} + +void PtrState::generateOriginPermuteIds() { + permuteIds.clear(); + isPermuted = false; + for (size_t i = 0; i < stateInfo.size(); ++i) { + permuteIds.emplace_back(i); + } + return; +} + +// Formula of "contiguous axes" +// - axis i is contiguous if stride[i] == product(shape[0..i-1]). +// We only prove contiguity from the rightmost axis outward. +// Rules implemented +// - Start from the rightmost axis. It is contiguous only if stride == 1 +// (static). +// - Then move left; expectedStride multiplies by the static shape of the axis +// to the right. +// - If we encounter any dynamic stride or dynamic shape needed for +// expectedStride, +// stop (cannot prove further) and return the count accumulated so far. +// Examples +// - shape=[2,3,4,5], stride=[A,B,2,C] => rightmost stride is dynamic => +// count=0 +// - shape=[2,3,4,5], stride=[A,B,C,1] => rightmost stride=1 => count=1 +// - shape=[2,3,4,5], stride=[60,20,5,1] => count=4 +size_t PtrState::countContiguousAxes(SmallVector stateInfo) const { + if (stateInfo.empty()) + return 0; + int64_t expected = 1; + size_t cnt = 0; + // iterate reversed safely: i = stateInfo.size()-1, ..., 0 + for (size_t i = stateInfo.size(); i-- > 0;) { + auto stride = getIntAttr(stateInfo[i].stride); + if (!stride.has_value()) + break; + if (stride.value() != expected) + break; + ++cnt; + // Update expected for the next (outer) axis: expected *= shape[i] + auto shape = getIntAttr(stateInfo[i].shape); + if (!shape.has_value()) + break; + expected *= shape.value(); + } + return cnt; +} +} // namespace TritonToStructured diff --git a/compiler/lib/TritonToStructured/TritonToStructuredPass.cpp b/compiler/lib/TritonToStructured/TritonToStructuredPass.cpp new file mode 100644 index 00000000..4f71cacf --- /dev/null +++ b/compiler/lib/TritonToStructured/TritonToStructuredPass.cpp @@ -0,0 +1,147 @@ + + +#include "dicp/TritonToStructured/TritonToStructuredPass.h" + +#include +#include +#include + +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "dicp/Utils/InterleaveOptimization.h" +#include "dicp/Utils/Utils.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/Operation.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/Dialect/Linalg/Transforms/Transforms.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Visitors.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/Passes.h" + +#include "dicp/TritonToStructured/CannonicalizerConverter.h" +#include "dicp/TritonToStructured/MemOpConverter.h" +#include "dicp/TritonToStructured/PtrAnalysis.h" +#include "llvm/ADT/BitVector.h" +#include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/ADT/SmallVectorExtras.h" +#include "llvm/ADT/Twine.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/LogicalResult.h" + +#define DEBUG_TYPE "triton-to-structured" + +using namespace mlir; +using namespace triton; + +void TritonToStructuredPass::getDependentDialects( + DialectRegistry ®istry) const { + registry.insert(); +} + +void TritonToStructuredPass::populateTritonToStructuredCanonicalizationPatterns( + RewritePatternSet &patterns) { + // TODO enable this optimization after fixing the bisheng bug it causes in + // current version + // patterns.add(patterns.getContext()); + patterns.add( + patterns.getContext()); + patterns.add( + patterns.getContext()); + // Add addptr splat->broadcast hoisting converter + patterns.add( + patterns.getContext()); + // Move loads before broadcasts when safe + patterns.add( + patterns.getContext()); +} + +void TritonToStructuredPass::populateTritonToStructuredPatterns( + RewritePatternSet &patterns, bool optimizeDynamicOffset, + bool enableMaskFallbackConversion) { + patterns.add(patterns.getContext(), + optimizeDynamicOffset, + enableMaskFallbackConversion); + patterns.add(patterns.getContext(), + optimizeDynamicOffset, + enableMaskFallbackConversion); +} + +LogicalResult +TritonToStructuredPass::processSplatBinaryOperations(ModuleOp moduleOp) { + mlir::RewritePatternSet patterns(&getContext()); + patterns.add( + patterns.getContext()); + if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { + moduleOp.emitWarning("Splat binary op processing failed"); + return failure(); + } + return success(); +} + +void TritonToStructuredPass::runOnOperation() { + auto moduleOp = getOperation(); + ConversionTarget target(getContext()); + RewritePatternSet canonicalizerPatterns(&getContext()); + + this->populateTritonToStructuredCanonicalizationPatterns( + canonicalizerPatterns); + if (failed( + applyPatternsGreedily(moduleOp, std::move(canonicalizerPatterns)))) { + moduleOp.emitWarning("Canonicalize failed"); + } + + RewritePatternSet tritonToStructuredPatterns(&getContext()); + populateTritonToStructuredPatterns(tritonToStructuredPatterns, + optimizeDynamicOffset, + enableMaskFallbackConversion); + + if (failed(applyPatternsGreedily(moduleOp, + std::move(tritonToStructuredPatterns)))) { + LLVM_DEBUG({ moduleOp->emitRemark("PtrAnalysis: rewrite MemOp failed"); }); + } + + if (failed(processSplatBinaryOperations(moduleOp))) { + moduleOp.emitWarning("Splat binary op processing failed"); + } + + PassManager pm(&getContext(), moduleOp.getOperationName()); + pm.addPass(createCSEPass()); + pm.addPass(createCanonicalizerPass()); + if (failed(runPipeline(pm, getOperation()))) { + moduleOp->emitWarning("Canonicalize failed"); + } +} + +std::unique_ptr> +triton::createTritonToStructuredPass() { + return std::make_unique(); +} + +std::unique_ptr> +triton::createTritonToStructuredPass(bool enableMaskFallbackConversion, + bool optimizeDynamicOffset) { + return std::make_unique(enableMaskFallbackConversion, + optimizeDynamicOffset); +} diff --git a/compiler/lib/Conversion/TritonToUnstructure/BubbleUpOperation.cpp b/compiler/lib/TritonToUnstructure/BubbleUpOperation.cpp similarity index 78% rename from compiler/lib/Conversion/TritonToUnstructure/BubbleUpOperation.cpp rename to compiler/lib/TritonToUnstructure/BubbleUpOperation.cpp index 43677f9e..8037d5b9 100644 --- a/compiler/lib/Conversion/TritonToUnstructure/BubbleUpOperation.cpp +++ b/compiler/lib/TritonToUnstructure/BubbleUpOperation.cpp @@ -1,4 +1,6 @@ -#include "dicp/Conversion/TritonToUnstructure/BubbleUpOperation.h" + + +#include "dicp/TritonToUnstructure/BubbleUpOperation.h" #include "dicp/Utils/Utils.h" #include "mlir/Pass/PassManager.h" @@ -7,63 +9,6 @@ #define DEBUG_TYPE "triton-bubble-up-operation" -template -class BubbleUpExtract : public OpRewritePattern { - static_assert(std::is_same_v || - std::is_same_v); - -public: - using OpRewritePattern::OpRewritePattern; - BubbleUpExtract(MLIRContext *context, bool enableAggressiveMode); - LogicalResult matchAndRewrite(ExtractOpTy op, - PatternRewriter &rewriter) const override; - -private: - Value createExtractOp(ExtractOpTy op, Value value, Location loc, - PatternRewriter &rewriter) const; - template - void bubbleUpIntBinaryOp(ExtractOpTy op, BinOpTy binOp, Location loc, - PatternRewriter &rewriter) const; - template - void bubbleUpFloatBinaryOp(ExtractOpTy op, BinOpTy binOp, Location loc, - PatternRewriter &rewriter) const; - - void bubbleUpOperation(ExtractOpTy op, arith::ExtSIOp parentOp, Location loc, - PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, arith::CmpIOp parentOp, Location loc, - PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, arith::TruncFOp parentOp, Location loc, - PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, arith::ExtFOp parentOp, Location loc, - PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, arith::FPToSIOp parentOp, Location loc, - PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, arith::SIToFPOp parentOp, Location loc, - PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, triton::ClampFOp parentOp, - Location loc, PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, arith::CmpFOp parentOp, Location loc, - PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, triton::BroadcastOp parentOp, - Location loc, PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, triton::ExpandDimsOp parentOp, - Location loc, PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, triton::SplatOp parentOp, Location loc, - PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, triton::MakeRangeOp parentOp, - Location loc, PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, triton::AddPtrOp parentOp, - Location loc, PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, math::FloorOp parentOp, Location loc, - PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, math::CeilOp parentOp, Location loc, - PatternRewriter &rewriter) const; - void bubbleUpOperation(ExtractOpTy op, tensor::ExtractSliceOp parentOp, - Location loc, PatternRewriter &rewriter) const; - - bool enableAggressiveMode; -}; - template BubbleUpExtract::BubbleUpExtract(MLIRContext *context, bool enableAggressiveMode) @@ -193,7 +138,7 @@ Value BubbleUpExtract::createExtractOp( PatternRewriter &rewriter) const { auto extractedOp = rewriter.create(loc, value, op.getIndices()); - extractedOp->setAttr(mlir::dicp::discreteAttrName, + extractedOp->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); return extractedOp; } @@ -202,10 +147,12 @@ template <> Value BubbleUpExtract::createExtractOp( tensor::ExtractSliceOp op, Value value, Location loc, PatternRewriter &rewriter) const { + auto extractedType = getExtractSlicedType( + op.getMixedSizes(), op.getDroppedDims(), getElementTypeOrSelf(value)); auto extractedOp = rewriter.create( - loc, value, op.getMixedOffsets(), op.getMixedSizes(), + loc, extractedType, value, op.getMixedOffsets(), op.getMixedSizes(), op.getMixedStrides()); - extractedOp->setAttr(mlir::dicp::discreteAttrName, + extractedOp->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); return extractedOp; } @@ -269,7 +216,7 @@ void BubbleUpExtract::bubbleUpOperation( } } auto extractedOp = rewriter.create(loc, src, newIndices); - extractedOp->setAttr(mlir::dicp::discreteAttrName, + extractedOp->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); rewriter.replaceOp(op, extractedOp); } @@ -295,13 +242,15 @@ void BubbleUpExtract::bubbleUpOperation( if (getConstantIntValue(newSizes.back()).value_or(-1) != 1) isScalarLikeSrc = false; } + auto extractedType = getExtractSlicedType(newSizes, op.getDroppedDims(), + getElementTypeOrSelf(src)); auto extractedOp = rewriter.create( - loc, src, newOffsets, newSizes, op.getMixedStrides()); - extractedOp->setAttr(mlir::dicp::discreteAttrName, + loc, extractedType, src, newOffsets, newSizes, op.getMixedStrides()); + extractedOp->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); if (isScalarLikeSrc) { SmallVector indices( - srcShape.size(), + extractedType.getRank(), rewriter.create(loc, rewriter.getIndexAttr(0))); auto extractedValue = rewriter.create(loc, extractedOp, indices); @@ -324,7 +273,7 @@ void BubbleUpExtract::bubbleUpOperation( newIndices.push_back(index.value()); } auto extractedOp = rewriter.create(loc, src, newIndices); - extractedOp->setAttr(mlir::dicp::discreteAttrName, + extractedOp->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); rewriter.replaceOp(op, extractedOp); } @@ -335,22 +284,38 @@ void BubbleUpExtract::bubbleUpOperation( PatternRewriter &rewriter) const { auto src = parentOp.getSrc(); auto srcShape = cast(src.getType()).getShape(); + auto offsets = op.getMixedOffsets(); + auto sizes = op.getMixedSizes(); + auto strides = op.getMixedStrides(); + auto axis = parentOp.getAxis(); + int64_t axisOffset = 0; SmallVector newOffsets; SmallVector newSizes; SmallVector newStrides; + auto droppedDim = op.getDroppedDims(); + llvm::SmallBitVector srcDroppedDims; for (size_t i = 0; i <= srcShape.size(); i++) { - if (i != parentOp.getAxis()) { - newOffsets.push_back(op.getMixedOffsets()[i]); - newSizes.push_back(op.getMixedSizes()[i]); - newStrides.push_back(op.getMixedStrides()[i]); + if (i != axis) { + newOffsets.push_back(offsets[i]); + newSizes.push_back(sizes[i]); + newStrides.push_back(strides[i]); + srcDroppedDims.push_back(droppedDim[i]); + if (i < axis && droppedDim[i]) + axisOffset++; } } - auto extractedOp = rewriter.create( - loc, src, newOffsets, newSizes, newStrides); - extractedOp->setAttr(mlir::dicp::discreteAttrName, - UnitAttr::get(rewriter.getContext())); - rewriter.replaceOpWithNewOp(op, extractedOp, - parentOp.getAxisAttr()); + auto extractedType = + getExtractSlicedType(newSizes, srcDroppedDims, getElementTypeOrSelf(src)); + if (extractedType == src.getType()) { + rewriter.replaceOp(op, src); + } else { + auto extractedOp = rewriter.create( + loc, extractedType, src, newOffsets, newSizes, newStrides); + extractedOp->setAttr(ConverterUtils::discreteAttrName, + UnitAttr::get(rewriter.getContext())); + rewriter.replaceOpWithNewOp(op, extractedOp, + axis - axisOffset); + } } template <> @@ -375,8 +340,17 @@ void BubbleUpExtract::bubbleUpOperation( tensor::ExtractOp op, triton::MakeRangeOp parentOp, Location loc, PatternRewriter &rewriter) const { auto resultType = cast(parentOp.getResult().getType()); - rewriter.replaceOpWithNewOp( - op, resultType.getElementType(), op.getIndices()[0]); + int32_t start = parentOp.getStart(); + Value idx = op.getIndices()[0]; + Value result = rewriter.create( + op.getLoc(), resultType.getElementType(), idx); + if (start != 0) { + Value startVal = rewriter.create( + op.getLoc(), + rewriter.getIntegerAttr(resultType.getElementType(), start)); + result = rewriter.create(op.getLoc(), result, startVal); + } + rewriter.replaceOp(op, result); } template <> @@ -384,14 +358,8 @@ void BubbleUpExtract::bubbleUpOperation( tensor::ExtractSliceOp op, triton::MakeRangeOp parentOp, Location loc, PatternRewriter &rewriter) const { auto resultType = cast(parentOp.getResult().getType()); - Value idx; - if (auto offsetVal = dyn_cast(op.getMixedOffsets()[0])) { - idx = offsetVal; - } else { - idx = rewriter.create( - op.getLoc(), rewriter.getIndexAttr( - getConstantIntValue(op.getMixedOffsets()[0]).value())); - } + auto idxOfr = op.getMixedOffsets()[0]; + Value idx = getValueOrCreateConstantIndexOp(rewriter, op.getLoc(), idxOfr); idx = rewriter.create(op.getLoc(), resultType.getElementType(), idx); rewriter.replaceOpWithNewOp(op, op.getResult().getType(), @@ -485,23 +453,31 @@ template <> void BubbleUpExtract::bubbleUpOperation( tensor::ExtractOp op, tensor::ExtractSliceOp parentOp, Location loc, PatternRewriter &rewriter) const { + SmallVector indices; SmallVector newIndices; - for (const auto &[offset, index] : - llvm::zip_equal(parentOp.getMixedOffsets(), op.getIndices())) { - Value offsetVal; - if (auto v = dyn_cast(offset)) { - offsetVal = v; + auto indiceIter = op.getIndices().begin(); + auto droppedDims = parentOp.getDroppedDims(); + for (auto [idx, offset] : llvm::enumerate(parentOp.getMixedOffsets())) { + if (droppedDims[idx]) { + auto zeroIdx = rewriter.create( + op.getLoc(), rewriter.getIndexAttr(0)); + indices.push_back(zeroIdx); } else { - offsetVal = rewriter.create( - op.getLoc(), rewriter.getIndexAttr(*getConstantIntValue(offset))); + indices.push_back(*indiceIter); + ++indiceIter; } + } + for (const auto &[offset, index] : + llvm::zip_equal(parentOp.getMixedOffsets(), indices)) { + Value offsetVal = + getValueOrCreateConstantIndexOp(rewriter, op.getLoc(), offset); newIndices.push_back( rewriter.create(op.getLoc(), offsetVal, index)); } rewriter .replaceOpWithNewOp(op, parentOp.getSource(), newIndices) - ->setAttr(mlir::dicp::discreteAttrName, + ->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); } diff --git a/compiler/lib/Conversion/TritonToUnstructure/CMakeLists.txt b/compiler/lib/TritonToUnstructure/CMakeLists.txt similarity index 87% rename from compiler/lib/Conversion/TritonToUnstructure/CMakeLists.txt rename to compiler/lib/TritonToUnstructure/CMakeLists.txt index 802f5130..7761d1f7 100644 --- a/compiler/lib/Conversion/TritonToUnstructure/CMakeLists.txt +++ b/compiler/lib/TritonToUnstructure/CMakeLists.txt @@ -2,6 +2,7 @@ add_triton_library(TritonToUnstructure UnstructureConversionPass.cpp OffsetAnalysis.cpp BubbleUpOperation.cpp + ReplaceArguments.cpp DEPENDS TritonToUnstructureConversionPassIncGen @@ -17,4 +18,5 @@ add_triton_library(TritonToUnstructure TritonIR TritonAnalysis MLIRSCFTransforms + BiShengIRHIVMDialect ) \ No newline at end of file diff --git a/compiler/lib/Conversion/TritonToUnstructure/OffsetAnalysis.cpp b/compiler/lib/TritonToUnstructure/OffsetAnalysis.cpp similarity index 69% rename from compiler/lib/Conversion/TritonToUnstructure/OffsetAnalysis.cpp rename to compiler/lib/TritonToUnstructure/OffsetAnalysis.cpp index cf66046b..b07020b1 100644 --- a/compiler/lib/Conversion/TritonToUnstructure/OffsetAnalysis.cpp +++ b/compiler/lib/TritonToUnstructure/OffsetAnalysis.cpp @@ -1,8 +1,14 @@ -#include "dicp/Conversion/TritonToUnstructure/OffsetAnalysis.h" + + +#include "dicp/TritonToUnstructure/OffsetAnalysis.h" #include "dicp/Utils/Utils.h" +#include "bishengir/Dialect/Annotation/IR/Annotation.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "llvm/Support/Casting.h" #include "llvm/Support/Debug.h" #define DEBUG_TYPE "triton-offset-analysis" @@ -16,31 +22,31 @@ PtrOffsetInfo::PtrOffsetInfo(const PtrOffsetInfo &other) { *this = other; } PtrOffsetInfo::PtrOffsetInfo(const Value &ptr) : ptr(ptr) { setZeroOffset(); } -PtrOffsetInfo::PtrOffsetInfo(ArrayRef structured) +PtrOffsetInfo::PtrOffsetInfo(ArrayRef structured) : ptr(nullptr), offset(nullptr) { setStructured(structured); } -PtrOffsetInfo::PtrOffsetInfo(const Value &ptr, bool structured) : ptr(ptr) { +PtrOffsetInfo::PtrOffsetInfo(const Value &ptr, AxisInfo structured) : ptr(ptr) { setZeroOffset(); if (auto tensorType = dyn_cast(ptr.getType())) this->structured.resize(tensorType.getRank(), structured); } -PtrOffsetInfo::PtrOffsetInfo(const Value &ptr, ArrayRef structured) +PtrOffsetInfo::PtrOffsetInfo(const Value &ptr, ArrayRef structured) : ptr(ptr) { setStructured(structured); } PtrOffsetInfo::PtrOffsetInfo(const Value &ptr, const Value &offset, - bool structured) + AxisInfo structured) : ptr(ptr), offset(offset) { if (auto tensorType = dyn_cast(ptr.getType())) this->structured.resize(tensorType.getRank(), structured); } PtrOffsetInfo::PtrOffsetInfo(const Value &ptr, const Value &offset, - ArrayRef structured) + ArrayRef structured) : ptr(ptr), offset(offset) { setStructured(structured); } @@ -63,10 +69,11 @@ SmallVector &PtrOffsetInfo::getOffsetsRef() { return this->tptOffsets; } bool PtrOffsetInfo::isScalarLike() const { return this->scalarLike; } -SmallVector &PtrOffsetInfo::getStructuredRef() { +SmallVector &PtrOffsetInfo::getStructuredRef() { return this->structured; } -const SmallVector &PtrOffsetInfo::getStructured() const { +const SmallVector & +PtrOffsetInfo::getStructured() const { return this->structured; } @@ -85,27 +92,32 @@ void PtrOffsetInfo::setStructured() { assert(ptr && "ptr Should be to infer rank"); this->structured.clear(); if (auto tensorType = dyn_cast(ptr.getType())) - this->structured.resize(tensorType.getRank(), true); + this->structured.resize(tensorType.getRank(), AxisInfo::structured); } void PtrOffsetInfo::setStructured(int rank) { this->structured.clear(); - this->structured.resize(rank, true); + this->structured.resize(rank, AxisInfo::structured); +} + +void PtrOffsetInfo::setStructured(int rank, AxisInfo info) { + this->structured.clear(); + this->structured.resize(rank, info); } void PtrOffsetInfo::setUnstructured() { assert(ptr && "ptr Should be to infer rank"); this->structured.clear(); if (auto tensorType = dyn_cast(ptr.getType())) - this->structured.resize(tensorType.getRank(), false); + this->structured.resize(tensorType.getRank(), AxisInfo::unstructured); } void PtrOffsetInfo::setUnstructured(int rank) { this->structured.clear(); - this->structured.resize(rank, false); + this->structured.resize(rank, AxisInfo::unstructured); } -void PtrOffsetInfo::setStructured(ArrayRef structured) { +void PtrOffsetInfo::setStructured(ArrayRef structured) { this->structured.resize(structured.size()); for (size_t i = 0; i < structured.size(); i++) this->structured[i] = structured[i]; @@ -120,16 +132,26 @@ void PtrOffsetInfo::setScalarLike(bool scalarLike) { } bool PtrOffsetInfo::isStructured(int dim) const { - return this->scalarLike || structured[dim]; + return this->scalarLike || structured[dim] == AxisInfo::structured || + structured[dim] == AxisInfo::scalar; } bool PtrOffsetInfo::isStructured() const { - return this->scalarLike || - llvm::all_of(structured, [](auto dim) { return dim; }); + return this->scalarLike || llvm::all_of(structured, [](auto dim) { + return dim == AxisInfo::structured || dim == AxisInfo::scalar; + }); } bool PtrOffsetInfo::isUnstructured() const { - return llvm::all_of(structured, [](auto dim) { return !dim; }); + return llvm::all_of(structured, + [](auto dim) { return dim == AxisInfo::unstructured; }); +} + +bool PtrOffsetInfo::isUnstructuredOrScalarlike() const { + return llvm::all_of(structured, [](auto dim) { + return dim == AxisInfo::unstructured || dim == AxisInfo::scalarlike || + dim == AxisInfo::scalar; + }); } void PtrOffsetInfo::setZeroOffset() { @@ -156,10 +178,12 @@ PtrOffsetInfo combineInfo(const PtrOffsetInfo &lhs, const PtrOffsetInfo &rhs) { assert(lhs.getRank() == rhs.getRank() && "Rank must be same to be combined"); info.setScalarLike(lhs.isScalarLike() && rhs.isScalarLike()); - SmallVector &structuredRef = info.getStructuredRef(); + auto &structuredRef = info.getStructuredRef(); + auto lhsStructured = lhs.getStructured(); + auto rhsStructured = rhs.getStructured(); structuredRef.resize(lhs.getRank()); for (size_t i = 0; i < structuredRef.size(); i++) - structuredRef[i] = lhs.isStructured(i) && rhs.isStructured(i); + structuredRef[i] = std::min(lhsStructured[i], rhsStructured[i]); return info; } @@ -177,7 +201,6 @@ void parse(Value operand, const Location &loc, RewriterBase &rewriter, auto &os = llvm::dbgs(); os << "parse\n" << operand << '\n'; }); - LLVM_DEBUG({}); if (auto *defOp = operand.getDefiningOp()) { if (isa(defOp->getDialect())) { @@ -193,6 +216,18 @@ void parse(Value operand, const Location &loc, RewriterBase &rewriter, parseLoopOp(loopOp, loc, rewriter, offsetMap, operand); } else if (auto extractOp = dyn_cast(defOp)) { parseExtract(extractOp, loc, rewriter, offsetMap); + } else if (auto insertOp = dyn_cast(defOp)) { + parseInsert(insertOp, loc, rewriter, offsetMap); + } else if (auto extractSliceOp = + dyn_cast(defOp)) { + parseExtractSlice(extractSliceOp, loc, rewriter, offsetMap); + } else if (auto insertSliceOp = dyn_cast(defOp)) { + parseInsertSlice(insertSliceOp, loc, rewriter, offsetMap); + } else if (auto customOp = dyn_cast(defOp)) { + auto opResult = dyn_cast(operand); + assert(opResult && "Expected operand to be an OpResult"); + unsigned resultIdx = opResult.getResultNumber(); + parseCustomOp(customOp, loc, rewriter, offsetMap, resultIdx); } } } else if (auto blockArgument = dyn_cast(operand)) { @@ -203,7 +238,8 @@ void parse(Value operand, const Location &loc, RewriterBase &rewriter, }); if (isa(parentOp)) { if (auto ptrType = dyn_cast(operand.getType())) { - offsetMap[operand] = PtrOffsetInfo(operand, true); + offsetMap[operand] = + PtrOffsetInfo(operand, PtrOffsetInfo::AxisInfo::scalar); } else { offsetMap[operand] = PtrOffsetInfo(); } @@ -225,9 +261,15 @@ void parse(Value operand, const Location &loc, RewriterBase &rewriter, os << "finish parse\n" << operand << '\n'; auto data = offsetMap.at(operand); for (auto s : data.getStructuredRef()) - os << s; + os << static_cast(s); os << "\n"; }); + + if (auto tensorType = dyn_cast(operand.getType()); + tensorType && isa(tensorType.getElementType())) { + auto data = offsetMap.at(operand); + assert(data.getPtr() && "pointer type should be parsed"); + } } void parseLoopRegionIterArg(LoopLikeOpInterface loopOp, const Location &loc, @@ -239,7 +281,8 @@ void parseLoopRegionIterArg(LoopLikeOpInterface loopOp, const Location &loc, auto argNum = regionIterArg.getArgNumber(); auto conditionArg = whileOp.getConditionOp().getArgs()[argNum]; parse(conditionArg, loc, rewriter, offsetMap); - offsetMap[regionIterArg] = offsetMap[conditionArg]; + auto tmp = offsetMap[conditionArg]; + offsetMap[regionIterArg] = tmp; return; } OpOperand *initArgOperand = loopOp.getTiedLoopInit(regionIterArg); @@ -247,7 +290,8 @@ void parseLoopRegionIterArg(LoopLikeOpInterface loopOp, const Location &loc, return; Value initArg = initArgOperand->get(); parse(initArg, loc, rewriter, offsetMap); - offsetMap[regionIterArg] = offsetMap[initArg]; + auto tmp = offsetMap[initArg]; + offsetMap[regionIterArg] = tmp; } void parseArithOp(Operation *arithOp, const Location &loc, @@ -330,11 +374,13 @@ void parseTritonOp(Operation *tritonOp, const Location &loc, parseExpandDims(expandDimsOp, loc, rewriter, offsetMap); } else if (auto clampFOp = dyn_cast(tritonOp)) { parseClampF(clampFOp, loc, rewriter, offsetMap); - } else if (auto makeTensorDescOp = - dyn_cast(tritonOp)) { - parseMakeTensorDesc(makeTensorDescOp, loc, rewriter, offsetMap); - } else if (auto makeTensorPtrOp = - dyn_cast(tritonOp)) { + } + // FIXME:Z|wait triton version upgrade to 3.5 + // else if (auto makeTensorDescOp = + // dyn_cast(tritonOp)) { + // parseMakeTensorDesc(makeTensorDescOp, loc, rewriter, offsetMap); + // } + else if (auto makeTensorPtrOp = dyn_cast(tritonOp)) { parseMakeTensorPtr(makeTensorPtrOp, loc, rewriter, offsetMap); } else if (auto reduceOp = dyn_cast(tritonOp)) { parseReduce(reduceOp, loc, rewriter, offsetMap); @@ -396,15 +442,15 @@ void parseAddPtr(triton::AddPtrOp op, const Location &loc, offsetMap[dst] = dstOffsetInfo; LLVM_DEBUG({ auto &os = llvm::dbgs(); - SmallVector &ptrStructured = ptrOffsetInfo.getStructuredRef(); - SmallVector &offsetStructured = offsetOffsetInfo.getStructuredRef(); + auto &ptrStructured = ptrOffsetInfo.getStructuredRef(); + auto &offsetStructured = offsetOffsetInfo.getStructuredRef(); os << "[parseAddPtr] ptrStructured: "; for (size_t i = 0; i < ptrStructured.size(); i++) - os << ptrStructured[i]; + os << static_cast(ptrStructured[i]); os << "\n"; os << "[parseAddPtr] offsetStructured: "; for (size_t i = 0; i < offsetStructured.size(); i++) - os << offsetStructured[i]; + os << static_cast(offsetStructured[i]); os << "\n"; }); } @@ -434,8 +480,10 @@ void parseSplat(triton::SplatOp op, const Location &loc, RewriterBase &rewriter, dstOffsetInfo.setOffset(offset); } // Set addPtr offset map - - dstOffsetInfo.setStructured(dstType.getRank()); + auto &dstStructured = dstOffsetInfo.getStructuredRef(); + for (auto dim : dstType.getShape()) + dstStructured.push_back(dim == 1 ? PtrOffsetInfo::AxisInfo::scalar + : PtrOffsetInfo::AxisInfo::scalarlike); dstOffsetInfo.setScalarLike(true); offsetMap[dst] = dstOffsetInfo; } @@ -446,17 +494,18 @@ void parseBinaryOp(BinOpTy op, const Location &loc, RewriterBase &rewriter, auto lhs = op.getLhs(); parse(lhs, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo lhsOffsetInfo = offsetMap.at(lhs); - SmallVector &lhsStructured = lhsOffsetInfo.getStructuredRef(); + auto &lhsStructured = lhsOffsetInfo.getStructuredRef(); auto rhs = op.getRhs(); parse(rhs, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo rhsOffsetInfo = offsetMap.at(rhs); - SmallVector &rhsStructured = rhsOffsetInfo.getStructuredRef(); + auto &rhsStructured = rhsOffsetInfo.getStructuredRef(); auto dst = op->getResult(0); PtrOffsetInfo dstOffsetInfo; dstOffsetInfo.setScalarLike(lhsOffsetInfo.isScalarLike() && rhsOffsetInfo.isScalarLike()); if (dstOffsetInfo.isScalarLike()) - dstOffsetInfo.setStructured(lhsStructured.size()); + dstOffsetInfo.setStructured(lhsStructured.size(), + PtrOffsetInfo::AxisInfo::scalarlike); else dstOffsetInfo.setUnstructured(lhsStructured.size()); offsetMap[dst] = dstOffsetInfo; @@ -503,7 +552,8 @@ void parseIndexCast(arith::IndexCastOp op, const Location &loc, parse(src, op.getLoc(), rewriter, offsetMap); // Set indexCast offset map auto dst = op.getOut(); - offsetMap[dst] = offsetMap.at(src); + auto srcOffsetInfo = offsetMap.at(src); + offsetMap[dst] = srcOffsetInfo; } template @@ -512,8 +562,13 @@ void parseConstantOp(ConstOpTy dst, const Location &loc, RewriterBase &rewriter, // Set constant offset map offsetMap[dst] = PtrOffsetInfo(); offsetMap[dst].setScalarLike(true); - if (auto tensorType = dyn_cast(dst->getResult(0).getType())) - offsetMap[dst].setStructured(tensorType.getRank()); + if (auto tensorType = + dyn_cast(dst->getResult(0).getType())) { + auto &dstStructured = offsetMap[dst].getStructuredRef(); + for (auto dim : tensorType.getShape()) + dstStructured.push_back(dim == 1 ? PtrOffsetInfo::AxisInfo::scalar + : PtrOffsetInfo::AxisInfo::scalarlike); + } } void parseMakeRange(triton::MakeRangeOp op, const Location &loc, @@ -532,7 +587,8 @@ void parseExtSI(arith::ExtSIOp op, const Location &loc, RewriterBase &rewriter, parse(src, op.getLoc(), rewriter, offsetMap); // Set extSI offset map auto dst = op.getOut(); - offsetMap[dst] = offsetMap.at(src); + auto srcOffsetInfo = offsetMap.at(src); + offsetMap[dst] = srcOffsetInfo; } void parseBitcast(triton::BitcastOp op, const Location &loc, @@ -542,7 +598,7 @@ void parseBitcast(triton::BitcastOp op, const Location &loc, auto src = op.getSrc(); parse(src, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo srcOffsetInfo = offsetMap.at(src); - SmallVector &srcStructured = srcOffsetInfo.getStructuredRef(); + auto &srcStructured = srcOffsetInfo.getStructuredRef(); // Set extSI offset map auto dst = op.getResult(); if (auto ptr = srcOffsetInfo.getPtr()) { @@ -580,20 +636,20 @@ void parseMulI(arith::MulIOp op, const Location &loc, RewriterBase &rewriter, auto lhs = op.getLhs(); parse(lhs, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo lhsOffsetInfo = offsetMap.at(lhs); - SmallVector &lhsStructured = lhsOffsetInfo.getStructuredRef(); + auto &lhsStructured = lhsOffsetInfo.getStructuredRef(); bool lhsScalarLike = lhsOffsetInfo.isScalarLike(); // Get muli rhs auto rhs = op.getRhs(); parse(rhs, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo rhsOffsetInfo = offsetMap.at(rhs); - SmallVector &rhsStructured = rhsOffsetInfo.getStructuredRef(); + auto &rhsStructured = rhsOffsetInfo.getStructuredRef(); bool rhsScalarLike = rhsOffsetInfo.isScalarLike(); // Set muli offset map size_t maxSize = std::max(lhsStructured.size(), rhsStructured.size()); auto dst = op.getResult(); offsetMap[dst] = PtrOffsetInfo(); offsetMap[dst].setScalarLike(lhsScalarLike && rhsScalarLike); - SmallVector &dstStructured = offsetMap[dst].getStructuredRef(); + auto &dstStructured = offsetMap[dst].getStructuredRef(); dstStructured.resize(maxSize); for (size_t i = 0; i < maxSize; i++) if (lhsScalarLike) @@ -601,7 +657,7 @@ void parseMulI(arith::MulIOp op, const Location &loc, RewriterBase &rewriter, else if (rhsScalarLike) dstStructured[i] = lhsStructured[i]; else - dstStructured[i] = false; + dstStructured[i] = PtrOffsetInfo::AxisInfo::unstructured; } void parseBroadcast(triton::BroadcastOp op, const Location &loc, @@ -611,7 +667,7 @@ void parseBroadcast(triton::BroadcastOp op, const Location &loc, auto src = op.getSrcMutable().get(); parse(src, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo srcOffsetInfo = offsetMap.at(src); - SmallVector &srcStructured = srcOffsetInfo.getStructuredRef(); + auto &srcStructured = srcOffsetInfo.getStructuredRef(); // Get broadcast dim auto dst = op.getResult(); assert(isa(src.getType()) && @@ -620,7 +676,7 @@ void parseBroadcast(triton::BroadcastOp op, const Location &loc, auto dstType = cast(dst.getType()); assert(srcType.getRank() == dstType.getRank() && "rank of source shoule be equal to destnation"); - auto broadcastDim = mlir::dicp::getBroadcastDims(srcType, dstType); + auto broadcastDim = ConverterUtils::getBroadcastDims(srcType, dstType); // Set broadcast offset map offsetMap[dst] = PtrOffsetInfo(srcOffsetInfo.getPtr()); offsetMap[dst].setScalarLike(srcOffsetInfo.isScalarLike()); @@ -637,23 +693,25 @@ void parseBroadcast(triton::BroadcastOp op, const Location &loc, offsetMap[dst].setOffset(offset); } - SmallVector &dstStructured = offsetMap[dst].getStructuredRef(); + auto &dstStructured = offsetMap[dst].getStructuredRef(); + auto dstShape = dstType.getShape(); dstStructured.resize(srcStructured.size()); for (size_t i = 0; i < dstStructured.size(); i++) - if (llvm::find(broadcastDim, i) != broadcastDim.end()) - dstStructured[i] = true; - else + if (llvm::find(broadcastDim, i) != broadcastDim.end() && dstShape[i] != 1) { + dstStructured[i] = PtrOffsetInfo::AxisInfo::scalarlike; + } else { dstStructured[i] = srcStructured[i]; + } } void parseExpandDims(triton::ExpandDimsOp op, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap) { // Get expandDims src - auto src = op.getSrcMutable().get(); + auto src = op.getSrc(); parse(src, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo srcOffsetInfo = offsetMap.at(src); - SmallVector &srcStructured = srcOffsetInfo.getStructuredRef(); + auto &srcStructured = srcOffsetInfo.getStructuredRef(); // Set expandDims offset map auto dst = op.getResult(); offsetMap[dst] = PtrOffsetInfo(srcOffsetInfo.getPtr()); @@ -662,17 +720,17 @@ void parseExpandDims(triton::ExpandDimsOp op, const Location &loc, RewriterBase::InsertionGuard guard(rewriter); rewriter.setInsertionPoint(op); Value valueOffset = srcOffsetInfo.getOffset(); - Value offset = rewriter.create(loc, valueOffset, - op.getAxisAttr()); + Value offset = + rewriter.create(loc, valueOffset, op.getAxis()); offsetMap[dst].setOffset(offset); } - SmallVector &dstStructured = offsetMap[dst].getStructuredRef(); + auto &dstStructured = offsetMap[dst].getStructuredRef(); dstStructured.resize(srcStructured.size() + 1); size_t j = 0; for (size_t i = 0; i < dstStructured.size(); i++) if (i == op.getAxis()) { - dstStructured[i] = true; + dstStructured[i] = PtrOffsetInfo::AxisInfo::scalar; } else { dstStructured[i] = srcStructured[j]; j++; @@ -718,15 +776,13 @@ void parseSelect(arith::SelectOp op, const Location &loc, auto trueValue = op.getTrueValue(); parse(trueValue, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo trueValueOffsetInfo = offsetMap.at(trueValue); - SmallVector &trueValueStructured = - trueValueOffsetInfo.getStructuredRef(); + auto &trueValueStructured = trueValueOffsetInfo.getStructuredRef(); bool trueValueScalarLike = trueValueOffsetInfo.isScalarLike(); // Get select falseValue auto falseValue = op.getFalseValue(); parse(falseValue, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo falseValueOffsetInfo = offsetMap.at(falseValue); - SmallVector &falseValueStructured = - falseValueOffsetInfo.getStructuredRef(); + auto &falseValueStructured = falseValueOffsetInfo.getStructuredRef(); bool falseValueScalarLike = falseValueOffsetInfo.isScalarLike(); // Set select offset map auto dst = op.getResult(); @@ -734,7 +790,30 @@ void parseSelect(arith::SelectOp op, const Location &loc, auto dstType = dyn_cast(dst.getType()); if (!dstType) return; - offsetMap[dst].setUnstructured(dstType.getRank()); + + // recognize "all dims size == 1", which cannot be handled in linalg pass's + // rewrite loop right now fix rewrite loop in linalg pass and remove this + // special handling + bool dstAllDimsAreOne = false; + if (auto rankedDstType = dyn_cast(dstType)) { + dstAllDimsAreOne = llvm::all_of(rankedDstType.getShape(), + [](int64_t dim) { return dim == 1; }); + } + + if (dstAllDimsAreOne) { + offsetMap[dst].setUnstructured(dstType.getRank()); + return; + } + + auto dstIsScalar = + trueValueScalarLike && falseValueScalarLike && conditionScalarLike; + offsetMap[dst].setScalarLike(dstIsScalar); + + auto &dstStructured = offsetMap[dst].getStructuredRef(); + dstStructured.resize(trueValueStructured.size()); + for (size_t i = 0; i < dstStructured.size(); i++) + dstStructured[i] = (dstIsScalar) ? PtrOffsetInfo::AxisInfo::scalarlike + : PtrOffsetInfo::AxisInfo::unstructured; } void parseFPToSI(arith::FPToSIOp op, const Location &loc, @@ -752,7 +831,8 @@ void parseFPToSI(arith::FPToSIOp op, const Location &loc, if (!dstType) return; if (offsetMap[dst].isScalarLike()) - offsetMap[dst].setStructured(dstType.getRank()); + offsetMap[dst].setStructured(dstType.getRank(), + PtrOffsetInfo::AxisInfo::scalarlike); else offsetMap[dst].setUnstructured(dstType.getRank()); } @@ -772,22 +852,24 @@ void parseSIToFP(arith::SIToFPOp op, const Location &loc, if (!dstType) return; if (offsetMap[dst].isScalarLike()) - offsetMap[dst].setStructured(dstType.getRank()); + offsetMap[dst].setStructured(dstType.getRank(), + PtrOffsetInfo::AxisInfo::scalarlike); else offsetMap[dst].setUnstructured(dstType.getRank()); } -void parseMakeTensorDesc(triton::MakeTensorDescOp op, const Location &loc, - RewriterBase &rewriter, - llvm::DenseMap &offsetMap) { - // Set MakeTensorDesc offset map - auto dst = op.getResult(); - offsetMap[dst] = PtrOffsetInfo(); - auto dstType = dyn_cast(dst.getType()); - if (!dstType) - return; - offsetMap[dst].setStructured(dstType.getRank()); -} +// FIXME:Z|wait triton version upgrade to 3.5 +// void parseMakeTensorDesc(triton::MakeTensorDescOp op, const Location &loc, +// RewriterBase &rewriter, +// llvm::DenseMap &offsetMap) { +// // Set MakeTensorDesc offset map +// auto dst = op.getResult(); +// offsetMap[dst] = PtrOffsetInfo(); +// auto dstType = dyn_cast(dst.getType()); +// if (!dstType) +// return; +// offsetMap[dst].setStructured(dstType.getRank()); +// } void parseMakeTensorPtr(triton::MakeTensorPtrOp op, const Location &loc, RewriterBase &rewriter, @@ -810,7 +892,8 @@ void parseAdvance(triton::AdvanceOp op, const Location &loc, auto ptr = op.getPtr(); parse(ptr, op.getLoc(), rewriter, offsetMap); auto dst = op.getResult(); - offsetMap[dst] = offsetMap.at(ptr); + auto ptrOffsetInfo = offsetMap.at(ptr); + offsetMap[dst] = ptrOffsetInfo; auto dstType = dyn_cast( cast(dst.getType()).getPointeeType()); if (!dstType) @@ -833,7 +916,7 @@ void parseReduce(triton::ReduceOp op, const Location &loc, Value src = op->getOperand(0); parse(src, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo srcOffsetInfo = offsetMap.at(src); - SmallVector &srcStructured = srcOffsetInfo.getStructuredRef(); + auto &srcStructured = srcOffsetInfo.getStructuredRef(); // Set reduce offset map Value dst = op->getResult(0); auto dstType = dyn_cast(dst.getType()); @@ -841,12 +924,12 @@ void parseReduce(triton::ReduceOp op, const Location &loc, offsetMap[dst].setScalarLike(srcOffsetInfo.isScalarLike()); if (!dstType) return; - SmallVector &dstStructured = offsetMap[dst].getStructuredRef(); + auto &dstStructured = offsetMap[dst].getStructuredRef(); auto dstShape = dstType.getShape(); dstStructured.resize(dstShape.size()); for (size_t i = 0; i < dstStructured.size(); i++) if (dstShape[i] == 1) - dstStructured[i] = true; + dstStructured[i] = PtrOffsetInfo::AxisInfo::scalar; else dstStructured[i] = srcStructured[i]; } @@ -858,7 +941,7 @@ void parseReduceReturn(triton::ReduceReturnOp op, const Location &loc, Value src = op->getOperand(0); parse(src, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo srcOffsetInfo = offsetMap.at(src); - SmallVector &srcStructured = srcOffsetInfo.getStructuredRef(); + auto &srcStructured = srcOffsetInfo.getStructuredRef(); // Set reduce offset map Value dst = op->getResult(0); auto dstType = dyn_cast(dst.getType()); @@ -866,12 +949,12 @@ void parseReduceReturn(triton::ReduceReturnOp op, const Location &loc, offsetMap[dst].setScalarLike(srcOffsetInfo.isScalarLike()); if (!dstType) return; - SmallVector &dstStructured = offsetMap[dst].getStructuredRef(); + auto &dstStructured = offsetMap[dst].getStructuredRef(); auto dstShape = dstType.getShape(); dstStructured.resize(dstShape.size()); for (size_t i = 0; i < dstStructured.size(); i++) if (dstShape[i] == 1) - dstStructured[i] = true; + dstStructured[i] = PtrOffsetInfo::AxisInfo::scalar; else dstStructured[i] = srcStructured[i]; } @@ -884,10 +967,11 @@ void parseIf(scf::IfOp op, const Location &loc, RewriterBase &rewriter, Value thenYieldedValue = thenBlock.getTerminator()->getOperand(index); parse(thenYieldedValue, op.getLoc(), rewriter, offsetMap); PtrOffsetInfo thenOffsetInfo = offsetMap.at(thenYieldedValue); - SmallVector &thenStructured = thenOffsetInfo.getStructuredRef(); + auto &thenStructured = thenOffsetInfo.getStructuredRef(); + auto thenSrcPtr = thenOffsetInfo.getPtr(); // Get if else region bool dstIsScalar = thenOffsetInfo.isScalarLike(); - SmallVector elseStructured; + SmallVector elseStructured; if (op.elseBlock()) { Block &elseBlock = op.getElseRegion().front(); Value elseYieldedValue = elseBlock.getTerminator()->getOperand(index); @@ -895,17 +979,31 @@ void parseIf(scf::IfOp op, const Location &loc, RewriterBase &rewriter, PtrOffsetInfo elseOffsetInfo = offsetMap.at(elseYieldedValue); elseStructured = elseOffsetInfo.getStructuredRef(); dstIsScalar = dstIsScalar && elseOffsetInfo.isScalarLike(); + if (thenSrcPtr != elseOffsetInfo.getPtr()) { + emitError(loc) + << "Currently ptr type from different source not supported"; + } } + // Set if offset map offsetMap[dst] = PtrOffsetInfo(); + offsetMap[dst].setPtr(thenSrcPtr); offsetMap[dst].setScalarLike(dstIsScalar); - SmallVector &dstStructured = offsetMap[dst].getStructuredRef(); + auto &dstStructured = offsetMap[dst].getStructuredRef(); dstStructured.resize(thenStructured.size()); for (size_t i = 0; i < dstStructured.size(); i++) if (op.elseBlock()) - dstStructured[i] = thenStructured[i] && elseStructured[i]; + dstStructured[i] = (dstIsScalar) ? PtrOffsetInfo::AxisInfo::scalarlike + : PtrOffsetInfo::AxisInfo::unstructured; else dstStructured[i] = thenStructured[i]; + SmallVector dstOffsets(thenOffsetInfo.getOffsetsRef().size()); + if (!dstOffsets.empty()) { + // Assumes ifOp is already rewritten + for (size_t i = 0; i < dstOffsets.size(); i++) + dstOffsets[i] = op->getResult(index + i); + offsetMap[dst].setOffsets(dstOffsets); + } } void parseYield(scf::YieldOp op, const Location &loc, RewriterBase &rewriter, @@ -926,18 +1024,81 @@ void parseLoopOp(LoopLikeOpInterface op, const Location &loc, yieldedValue = op.getYieldedValues()[resNum]; } parse(yieldedValue, op.getLoc(), rewriter, offsetMap); - offsetMap[dst] = offsetMap.at(yieldedValue); + auto yieldOffsetInfo = offsetMap.at(yieldedValue); + offsetMap[dst] = yieldOffsetInfo; } void parseExtractSlice(tensor::ExtractSliceOp op, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap) { // Get extractSlice src - auto src = op.getOperand(0); + auto src = op.getSource(); parse(src, op.getLoc(), rewriter, offsetMap); // Set extractSlice offset map auto dst = op.getResult(); - offsetMap[dst] = offsetMap.at(src); + auto srcPtrInfo = offsetMap.at(src); + auto srcPtr = srcPtrInfo.getPtr(); + auto srcOffset = srcPtrInfo.getOffset(); + auto srcStructured = srcPtrInfo.getStructured(); + auto droppedDims = op.getDroppedDims(); + if (srcOffset) { + RewriterBase::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(op); + auto offsetType = getExtractSlicedType(op.getMixedSizes(), droppedDims, + getElementTypeOrSelf(srcOffset)); + srcOffset = rewriter.create( + op.getLoc(), offsetType, srcOffset, op.getMixedOffsets(), + op.getMixedSizes(), op.getMixedStrides()); + } + SmallVector dstStructured; + for (size_t i = 0; i < srcStructured.size(); i++) { + if (!droppedDims[i]) + dstStructured.push_back(srcStructured[i]); + } + offsetMap[dst] = PtrOffsetInfo(srcPtr, srcOffset, dstStructured); +} + +void parseInsertSlice(tensor::InsertSliceOp op, const Location &loc, + RewriterBase &rewriter, + llvm::DenseMap &offsetMap) { + // Get insertSlice src and dst + auto src = op.getSource(); + parse(src, op.getLoc(), rewriter, offsetMap); + auto dst = op.getDest(); + parse(dst, op.getLoc(), rewriter, offsetMap); + // Set insertSlice offset map + auto res = op.getResult(); + auto srcPtrInfo = offsetMap.at(src); + auto dstPtrInfo = offsetMap.at(dst); + PtrOffsetInfo resPtrInfo; + if (auto srcOffset = srcPtrInfo.getOffset()) { + RewriterBase::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(op); + auto resOffset = rewriter.create( + op.getLoc(), srcOffset, dstPtrInfo.getOffset(), op.getMixedOffsets(), + op.getMixedSizes(), op.getMixedStrides()); + auto srcPtr = srcPtrInfo.getPtr(); + auto dstPtr = dstPtrInfo.getPtr(); + assert(srcPtr == dstPtr && "ptrInfo for insert slice should be consistent"); + resPtrInfo.setPtr(srcPtr); + resPtrInfo.setOffset(resOffset); + } + auto droppedDims = op.getDroppedDims(); + auto srcStructuredIter = srcPtrInfo.getStructured().begin(); + SmallVector resStructured; + auto srcShape = op.getStaticSizes(); + auto dstShape = cast(dst.getType()).getShape(); + for (size_t i = 0; i < dstShape.size(); i++) { + if (!ShapedType::isDynamic(srcShape[i]) && srcShape[i] == dstShape[i]) { + resStructured.push_back(*srcStructuredIter); + } else { + resStructured.push_back(PtrOffsetInfo::AxisInfo::unstructured); + } + if (!droppedDims[i]) + ++srcStructuredIter; + } + resPtrInfo.setStructured(resStructured); + offsetMap[res] = resPtrInfo; } void parseExtract(tensor::ExtractOp op, const Location &loc, @@ -949,10 +1110,38 @@ void parseExtract(tensor::ExtractOp op, const Location &loc, offsetMap[dst] = PtrOffsetInfo(); if (isa(dst.getType())) { offsetMap[dst].setPtr(dst); + offsetMap[dst].setZeroOffset(); } offsetMap[dst].setScalarLike(true); } +void parseInsert(tensor::InsertOp op, const Location &loc, + RewriterBase &rewriter, + llvm::DenseMap &offsetMap) { + auto src = op.getScalar(); + parse(src, op.getLoc(), rewriter, offsetMap); + auto dst = op.getDest(); + parse(dst, op.getLoc(), rewriter, offsetMap); + + auto res = op.getResult(); + auto srcPtrInfo = offsetMap.at(src); + auto dstPtrInfo = offsetMap.at(dst); + + PtrOffsetInfo resPtrInfo; + if (auto srcOffset = srcPtrInfo.getOffset()) { + RewriterBase::InsertionGuard guard(rewriter); + rewriter.setInsertionPoint(op); + auto resOffset = rewriter.create( + op.getLoc(), srcOffset, dstPtrInfo.getOffset(), op.getIndices()); + auto srcPtr = srcPtrInfo.getPtr(); + auto dstPtr = dstPtrInfo.getPtr(); + resPtrInfo.setPtr(srcPtr); + resPtrInfo.setOffset(resOffset); + } + resPtrInfo.setUnstructured(dstPtrInfo.getRank()); + offsetMap[res] = resPtrInfo; +} + void parseIntToPtr(triton::IntToPtrOp op, const Location &loc, RewriterBase &rewriter, llvm::DenseMap &offsetMap) { @@ -961,5 +1150,47 @@ void parseIntToPtr(triton::IntToPtrOp op, const Location &loc, offsetMap[dst].setScalarLike(true); } +void parseCustomOp(hivm::CustomOp op, const Location &loc, + RewriterBase &rewriter, + llvm::DenseMap &offsetMap, + unsigned int resultIdx) { + for (auto operand : op.getInputs()) { + parse(operand, op->getLoc(), rewriter, offsetMap); + } + auto dst = op->getResult(resultIdx); + offsetMap[dst] = PtrOffsetInfo(); + auto tensorType = dyn_cast(dst.getType()); + if (!tensorType) { + if (isa(dst.getType())) { + offsetMap[dst].setPtr(dst); + offsetMap[dst].setZeroOffset(); + } else if (isa(dst.getType())) { + offsetMap[dst].setOffset(dst); + } else { + emitError(loc) << "Unsupported return type for hivm.custom: " + << dst.getType(); + } + return; + } + if (llvm::isa(tensorType.getElementType())) { + if (checkStructureAnnotated(op, rewriter)) { + auto srcValArrayAttr = op->getAttrOfType( + ConverterUtils::customSrcPtrIndexAttrName); + assert(srcValArrayAttr && + "structure hivm.custom op should present src tensor"); + auto srcValArray = srcValArrayAttr.asArrayRef(); + assert(srcValArray[resultIdx] != -1 && + "tensor result should map to src tensor"); + auto srcOffsetInfo = offsetMap[op->getOperand(srcValArray[resultIdx])]; + offsetMap[dst] = srcOffsetInfo; + return; + } + emitError(loc) << "Unsupported return unstructure RankedTensor of tt.ptr " + "for hivm.custom: " + << dst; + } + offsetMap[dst].setUnstructured(tensorType.getRank()); +} + } // namespace triton } // namespace mlir diff --git a/compiler/lib/TritonToUnstructure/ReplaceArguments.cpp b/compiler/lib/TritonToUnstructure/ReplaceArguments.cpp new file mode 100644 index 00000000..9bef0734 --- /dev/null +++ b/compiler/lib/TritonToUnstructure/ReplaceArguments.cpp @@ -0,0 +1,301 @@ + + +#include "dicp/TritonToUnstructure/UnstructureConversionPass.h" +#include "dicp/Utils/Utils.h" + +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "triton-replace-arguments" + +using namespace mlir; +using namespace triton; + +void replaceOperands(MutableArrayRef oprs, RewriterBase &rewriter, + llvm::DenseMap &offsetMap) { + for (auto it = oprs.begin(); it != oprs.end(); ++it) { + auto &opr = *it; + auto operand = opr.get(); + if (auto tensorType = dyn_cast(operand.getType()); + tensorType && isa(tensorType.getElementType())) { + parse(operand, operand.getLoc(), rewriter, offsetMap); + opr.set(offsetMap.at(operand).getOffset()); + } else if (auto ptrType = + dyn_cast(operand.getType())) { + parse(operand, operand.getLoc(), rewriter, offsetMap); + if (auto tensorType = + dyn_cast(ptrType.getPointeeType())) { + for (auto offset : offsetMap.at(operand).getOffsets()) { + it->set(offset); + ++it; + } + --it; + } else { + opr.set(offsetMap.at(operand).getOffset()); + } + } + } +} + +void replaceArgs(ValueRange args, RewriterBase &rewriter, + llvm::DenseMap &offsetMap) { + for (auto it = args.begin(); it != args.end(); ++it) { + auto arg = *it; + if (auto tensorType = dyn_cast(arg.getType()); + tensorType && isa(tensorType.getElementType())) { + RewriterBase::InsertionGuard guard(rewriter); + if (auto blockArg = dyn_cast(arg)) { + rewriter.setInsertionPointToStart(blockArg.getOwner()); + } else { + rewriter.setInsertionPointAfterValue(arg); + } + auto tempVar = rewriter + .create( + arg.getLoc(), arg.getType(), ValueRange({})) + ->getResult(0); + parse(arg, arg.getLoc(), rewriter, offsetMap); + auto src = offsetMap.at(arg).getPtr(); + rewriter.replaceAllUsesWith(arg, tempVar); + arg.setType(RankedTensorType::get(tensorType.getShape(), + rewriter.getIntegerType(64))); + src = rewriter.create(arg.getLoc(), tempVar.getType(), + src); + rewriter.replaceOpWithNewOp( + tempVar.getDefiningOp(), tempVar.getType(), src, arg); + } else if (auto ptrType = dyn_cast(arg.getType())) { + RewriterBase::InsertionGuard guard(rewriter); + if (auto blockArg = dyn_cast(arg)) { + rewriter.setInsertionPointToStart(blockArg.getOwner()); + } else { + rewriter.setInsertionPointAfterValue(arg); + } + auto tempVar = rewriter + .create( + arg.getLoc(), arg.getType(), ValueRange({})) + ->getResult(0); + parse(arg, arg.getLoc(), rewriter, offsetMap); + rewriter.replaceAllUsesWith(arg, tempVar); + if (auto tensorType = + dyn_cast(ptrType.getPointeeType())) { + auto srcOp = + offsetMap.at(arg).getPtr().getDefiningOp(); + arg.setType(rewriter.getIntegerType(32)); + SmallVector newOffsets; + for (auto offset : offsetMap.at(arg).getOffsets()) { + newOffsets.push_back(*it); + ++it; + } + --it; + rewriter.replaceOpWithNewOp( + tempVar.getDefiningOp(), tempVar.getType(), srcOp.getBase(), + srcOp.getShape(), srcOp.getStrides(), newOffsets, srcOp.getOrder()); + } else { + auto src = offsetMap.at(arg).getPtr(); + arg.setType(rewriter.getIntegerType(64)); + rewriter.replaceOpWithNewOp( + tempVar.getDefiningOp(), tempVar.getType(), src, arg); + } + } + } +} + +void convertTensorPtrPre(Operation *op, RewriterBase &rewriter, + llvm::DenseMap &offsetMap) { + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "[convertTensorPtr]: Preorder start\n" << *op << "\n"; + }); + if (auto whileOp = dyn_cast(op)) { + replaceArgs(whileOp.getBeforeArguments(), rewriter, offsetMap); + replaceOperands(whileOp.getInitsMutable(), rewriter, offsetMap); + replaceArgs(whileOp.getAfterArguments(), rewriter, offsetMap); + replaceArgs(whileOp->getResults(), rewriter, offsetMap); + replaceOperands(whileOp.getConditionOp().getArgsMutable(), rewriter, + offsetMap); + } else if (auto loopOp = dyn_cast(op)) { + replaceArgs(loopOp.getRegionIterArgs(), rewriter, offsetMap); + replaceOperands(loopOp.getInitsMutable(), rewriter, offsetMap); + } else if (auto ifOp = dyn_cast(op)) { + replaceArgs(ifOp->getResults(), rewriter, offsetMap); + replaceOperands(ifOp.thenYield().getResultsMutable(), rewriter, offsetMap); + replaceOperands(ifOp.elseYield().getResultsMutable(), rewriter, offsetMap); + } + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "[convertTensorPtr]: Preorder end\n" << *op << "\n"; + }); +} + +void convertTensorPtrPost(Operation *op, RewriterBase &rewriter, + llvm::DenseMap &offsetMap) { + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "[convertTensorPtr]: Postorder start\n" << *op << "\n"; + }); + if (auto whileOp = dyn_cast(op)) { + replaceOperands(whileOp.getYieldOp()->getOpOperands(), rewriter, offsetMap); + } else if (auto loopOp = dyn_cast(op)) { + replaceArgs(loopOp->getResults(), rewriter, offsetMap); + replaceOperands(*loopOp.getYieldedValuesMutable(), rewriter, offsetMap); + } + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "[convertTensorPtr]: Postorder end\n" << *op << "\n"; + }); +} + +int getPtrTensorRank(Type type) { + if (auto ptrType = dyn_cast(type)) { + if (auto tensorType = + dyn_cast(ptrType.getPointeeType())) { + return tensorType.getRank(); + } + } + return 0; +} + +SmallVector constructOperands(ValueRange operands, Value tempVar, + IRMapping mapping) { + SmallVector newOperands; + for (auto opr : operands) { + opr = mapping.lookupOrDefault(opr); + newOperands.push_back(opr); + auto numAppend = getPtrTensorRank(opr.getType()) - 1; + if (numAppend > 0) + newOperands.append(numAppend, tempVar); + } + return newOperands; +} + +SmallVector constructTypes(TypeRange types) { + SmallVector newTypes; + for (auto type : types) { + newTypes.push_back(type); + if (auto ptrType = dyn_cast(type)) { + if (auto tensorType = + dyn_cast(ptrType.getPointeeType())) { + if (tensorType.getRank() > 0) + newTypes.append(tensorType.getRank() - 1, + IntegerType::get(type.getContext(), 32)); + } + } + } + return newTypes; +} + +void replacePtrArguments(triton::FuncOp funcOp, + llvm::DenseMap &offsetMap) { + IRRewriter rewriter(funcOp.getContext()); + rewriter.setInsertionPointToStart(&funcOp.getBody().front()); + Value tempVar = rewriter + .create( + funcOp.getLoc(), rewriter.getI32Type(), ValueRange{}) + ->getResult(0); + std::function convertTensorPtr = [&](Operation *op) { + IRMapping mapping; + Operation *newOp = nullptr; + rewriter.setInsertionPointAfter(op); + if (auto forOp = dyn_cast(op)) { + newOp = rewriter.create( + forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), + forOp.getStep(), + constructOperands(forOp.getInitArgs(), tempVar, mapping), + [&](OpBuilder &b, Location loc, Value iv, ValueRange args) { + mapping.map(forOp.getInductionVar(), iv); + auto newArgIter = args.begin(); + for (auto oldArg : forOp.getRegionIterArgs()) { + mapping.map(oldArg, *newArgIter); + std::advance(newArgIter, + std::max(getPtrTensorRank(oldArg.getType()), 1)); + } + for (auto &bodyOp : forOp.getBody()->without_terminator()) { + b.clone(bodyOp, mapping); + } + auto yieldOp = cast(forOp.getBody()->getTerminator()); + b.create( + yieldOp.getLoc(), + constructOperands(yieldOp.getOperands(), tempVar, mapping)); + }); + } else if (auto whileOp = dyn_cast(op)) { + newOp = rewriter.create( + whileOp.getLoc(), constructTypes(whileOp->getResultTypes()), + constructOperands(whileOp.getInits(), tempVar, mapping), + [&](OpBuilder &b, Location loc, ValueRange args) { + auto newArgIter = args.begin(); + for (auto oldArg : whileOp.getBeforeArguments()) { + mapping.map(oldArg, *newArgIter); + std::advance(newArgIter, + std::max(getPtrTensorRank(oldArg.getType()), 1)); + } + for (auto &bodyOp : whileOp.getBeforeBody()->without_terminator()) { + b.clone(bodyOp, mapping); + } + auto conditionOp = whileOp.getConditionOp(); + b.create( + conditionOp.getLoc(), + mapping.lookup(conditionOp.getCondition()), + constructOperands(conditionOp.getArgs(), tempVar, mapping)); + }, + [&](OpBuilder &b, Location loc, ValueRange args) { + auto newArgIter = args.begin(); + for (auto oldArg : whileOp.getAfterArguments()) { + mapping.map(oldArg, *newArgIter); + std::advance(newArgIter, + std::max(getPtrTensorRank(oldArg.getType()), 1)); + } + for (auto &bodyOp : whileOp.getAfterBody()->without_terminator()) { + b.clone(bodyOp, mapping); + } + auto yieldOp = whileOp.getYieldOp(); + b.create( + yieldOp.getLoc(), + constructOperands(yieldOp.getOperands(), tempVar, mapping)); + }); + } else if (auto ifOp = dyn_cast(op); + ifOp && ifOp->getNumResults() > 0) { + newOp = rewriter.create( + ifOp.getLoc(), ifOp.getCondition(), + [&](OpBuilder &b, Location loc) { + for (auto &bodyOp : ifOp.thenBlock()->without_terminator()) { + b.clone(bodyOp, mapping); + } + auto yieldOp = ifOp.thenYield(); + b.create( + yieldOp.getLoc(), + constructOperands(yieldOp.getOperands(), tempVar, mapping)); + }, + [&](OpBuilder &b, Location loc) { + for (auto &bodyOp : ifOp.elseBlock()->without_terminator()) { + b.clone(bodyOp, mapping); + } + auto yieldOp = ifOp.elseYield(); + b.create( + yieldOp.getLoc(), + constructOperands(yieldOp.getOperands(), tempVar, mapping)); + }); + } else if (auto loopOp = dyn_cast(op)) { + llvm_unreachable("Unsupported loop op"); + } + if (newOp) { + newOp->setAttrs(op->getAttrs()); + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "Converting\n" << *op << "\nto\n" << *newOp << "\n"; + }); + auto resIter = newOp->result_begin(); + for (auto res : op->getResults()) { + rewriter.replaceAllUsesWith(res, *resIter); + std::advance(resIter, std::max(getPtrTensorRank(res.getType()), 1)); + } + rewriter.eraseOp(op); + op = newOp; + convertTensorPtrPre(op, rewriter, offsetMap); + for (auto ®ion : op->getRegions()) + region.walk(convertTensorPtr); + convertTensorPtrPost(op, rewriter, offsetMap); + return WalkResult::skip(); + } + return WalkResult::advance(); + }; + + funcOp->walk(convertTensorPtr); +} diff --git a/compiler/lib/Conversion/TritonToUnstructure/UnstructureConversionPass.cpp b/compiler/lib/TritonToUnstructure/UnstructureConversionPass.cpp similarity index 63% rename from compiler/lib/Conversion/TritonToUnstructure/UnstructureConversionPass.cpp rename to compiler/lib/TritonToUnstructure/UnstructureConversionPass.cpp index 3f69b378..89279d1a 100644 --- a/compiler/lib/Conversion/TritonToUnstructure/UnstructureConversionPass.cpp +++ b/compiler/lib/TritonToUnstructure/UnstructureConversionPass.cpp @@ -1,4 +1,8 @@ -#include "dicp/Conversion/TritonToUnstructure/UnstructureConversionPass.h" + + +#include "dicp/TritonToUnstructure/UnstructureConversionPass.h" +#include "dicp/TritonToLinalg/MaskAnalysis.h" +#include "dicp/TritonToStructured/CannonicalizerConverter.h" #include "dicp/Utils/Utils.h" #include "triton/Dialect/Triton/IR/Dialect.h" @@ -13,8 +17,6 @@ #include "llvm/ADT/STLExtras.h" -#include - #define DEBUG_TYPE "triton-unstructure-converter" using namespace mlir; @@ -22,6 +24,8 @@ using namespace triton; #include "llvm/Support/Debug.h" +bool forceSimtTemplateFlag = false; + template bool UnstructuredMemAccessConverter::checkUnstructureAnnotated( MemAccOpTy op, PatternRewriter &rewriter) const { @@ -55,17 +59,12 @@ Value UnstructuredMemAccessConverter::createExtractOp( if (!value) return value; SmallVector indices; - for (auto idx : iterIdx) { - if (auto val = dyn_cast(idx)) { - indices.push_back(val); - } else { - auto idxVal = rewriter.create( - loc, rewriter.getIndexAttr(*getConstantIntValue(idx))); - indices.push_back(idxVal); - } + for (auto idxOfr : iterIdx) { + auto idx = getValueOrCreateConstantIndexOp(rewriter, loc, idxOfr); + indices.push_back(idx); } auto extractedOp = rewriter.create(loc, value, indices); - extractedOp->setAttr(mlir::dicp::discreteAttrName, + extractedOp->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); return extractedOp; } @@ -84,7 +83,7 @@ Value UnstructuredMemAccessConverter::createExtractOp( }); auto extractedOp = rewriter.create( loc, value, offsets, sizes, strides); - extractedOp->setAttr(mlir::dicp::discreteAttrName, + extractedOp->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); return extractedOp; } @@ -221,6 +220,7 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( auto ptr = op.getPtr(); auto ptrType = dyn_cast(ptr.getType()); + auto isDiscreteMask = op->hasAttr("is_discrete_mask"); if (auto ptrPtrType = dyn_cast(ptr.getType())) { if (auto ptrTensorType = @@ -228,7 +228,7 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( ptrType = ptrTensorType; } - if (!ptrType || op->hasAttr(mlir::dicp::discreteAttrName)) + if (!ptrType || op->hasAttr(ConverterUtils::discreteAttrName)) return failure(); if (!offsetMap.contains(ptr)) return op.emitError() << "PtrOffsetInfo should be computed\n" << ptr; @@ -238,7 +238,7 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( if (checkUnstructureAnnotated(op, rewriter)) ptrOffsetInfo.setUnstructured(ptrOffsetInfo.getRank()); - if (ptrOffsetInfo.isStructured() && + if (ptrOffsetInfo.isStructured() && !isDiscreteMask && (!ptrOffsetInfo.isScalarLike() || llvm::all_of(ptrType.getShape(), [](int64_t dim) { return dim == 1; }))) return failure(); @@ -247,17 +247,22 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( auto &os = llvm::dbgs(); os << "Converting " << op->getName() << "\n"; os << op << "\n"; - os << ptrOffsetInfo.isStructured() << "\n"; + for (auto structured : ptrOffsetInfo.getStructuredRef()) + os << static_cast(structured); + os << "\n"; os << ptrOffsetInfo.isScalarLike() << "\n"; }); - if constexpr (std::is_same_v) + if constexpr (std::is_same_v) { if (ptrOffsetInfo.isScalarLike()) { splatAndLoadScenario(op, ptrOffsetInfo.getRank(), rewriter); return success(); } + } + + std::optional mstate = runMaskAnalysis(op, rewriter); - if (op->hasAttr(mlir::dicp::discreteMaskAttrName)) { + if (op->hasAttr(ConverterUtils::discreteMaskAttrName)) { if constexpr (std::is_same_v) { auto selectOp = op.getValue().template getDefiningOp(); op = rewriter.replaceOpWithNewOp( @@ -266,11 +271,13 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( rewriter.setInsertionPoint(op); ptrOffsetInfo.setUnstructured(ptrOffsetInfo.getRank()); } else if constexpr (std::is_same_v) { - auto selectOp = op.getVal().template getDefiningOp(); - op = rewriter.replaceOpWithNewOp( - op, op.getType(), op.getAtomicRmwOp(), op.getPtr(), - selectOp.getTrueValue(), selectOp.getCondition(), op.getSem(), - op.getScope()); + if (auto selectOp = + op.getVal().template getDefiningOp()) { + op = rewriter.replaceOpWithNewOp( + op, op.getType(), op.getAtomicRmwOp(), op.getPtr(), + selectOp.getTrueValue(), selectOp.getCondition(), op.getSem(), + op.getScope()); + } } rewriter.setInsertionPoint(op); ptrOffsetInfo.setUnstructured(ptrOffsetInfo.getRank()); @@ -322,8 +329,45 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( os << "UnStructured Flag check:\n"; os << "ptrOffsetInfo.isStructured: " << ptrOffsetInfo.isStructured() << "\n"; + os << "compileOn91095Flag: " << compileOn91095Flag << "\n"; + os << "forceSimtTemplateFlag: " << forceSimtTemplateFlag << "\n"; }); + // Fast path on A5: rewrite tt.load/store to tt.indirect_load/store directly. + if (compileOn91095Flag && forceSimtTemplateFlag && + (ptrOffsetInfo.isUnstructuredOrScalarlike() || isDiscreteMask)) { + if constexpr (std::is_same_v) { + assert(isa(srcPtr.getType()) && + "src must be ptr type"); + Value mask = op.getMask(); + Value other = op.getOther(); + auto resultType = op.getType(); + auto indirect = rewriter.create( + loc, resultType, srcPtr, ptrOffset, mask, other); + rewriter.replaceOp(op, indirect.getResult()); + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "Rewriting tt.load to tt.indirect_load\n"; + os << indirect << "\n"; + }); + return success(); + } else if constexpr (std::is_same_v) { + assert(isa(srcPtr.getType()) && + "src must be ptr type"); + Value value = op.getValue(); + Value mask = op.getMask(); + auto indirect = rewriter.create( + loc, srcPtr, ptrOffset, value, mask); + rewriter.eraseOp(op); + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "Rewriting tt.store to tt.indirect_store\n"; + os << indirect << "\n"; + }); + return success(); + } + } + Value iterArg = nullptr; // Only load case @@ -342,7 +386,8 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( for (size_t i = 0; i < resultShape.size(); i++) { auto size = resultShape[i]; - auto structured = ptrOffsetInfo.getStructuredRef()[i]; + auto structured = ptrOffsetInfo.getStructuredRef()[i] == + PtrOffsetInfo::AxisInfo::structured; // handle indirect dimension strides.push_back(rewriter.getIndexAttr(1)); Value sizeVal = @@ -362,8 +407,25 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( } sizeVal = rewriter.create(loc, sizeVal, tptShape); } + + Value loopLower = zeroIdx; + Value loopUpper = sizeVal; + if (mstate && i < mstate->dims.size() && i < mstate->offsets.size()) { + Value maskOffset = + getValueOrCreateConstantIndexOp(rewriter, loc, mstate->offsets[i]); + maskOffset = rewriter.create(loc, maskOffset, zeroIdx); + maskOffset = rewriter.create(loc, maskOffset, sizeVal); + loopLower = maskOffset; + + Value maskDim = + getValueOrCreateConstantIndexOp(rewriter, loc, mstate->dims[i]); + maskDim = rewriter.create(loc, maskOffset, maskDim); + maskDim = rewriter.create(loc, maskDim, sizeVal); + loopUpper = maskDim; + } + if (isLoadLike) { - forOp = rewriter.create(loc, zeroIdx, sizeVal, oneIdx, + forOp = rewriter.create(loc, loopLower, loopUpper, oneIdx, ValueRange({iterArg})); if (!newOpResult) { newOpResult = forOp->getResult(0); @@ -372,7 +434,7 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( } iterArg = forOp.getRegionIterArg(0); } else { - forOp = rewriter.create(loc, zeroIdx, sizeVal, oneIdx); + forOp = rewriter.create(loc, loopLower, loopUpper, oneIdx); } sizes.push_back(rewriter.getIndexAttr(1)); offsets.push_back(forOp.getInductionVar()); @@ -383,7 +445,7 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( } } - bool fullyUnstructured = ptrOffsetInfo.isUnstructured(); + bool fullyUnstructured = ptrOffsetInfo.isUnstructuredOrScalarlike(); auto extractedType = RankedTensorType::get(extractedShape, resultElementType); Value extractedOffset; @@ -437,7 +499,7 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( createMemAccOp(op, ptrToAccess, loc, rewriter, offsets, sizes, strides); } - accessedOp->setAttr(mlir::dicp::discreteAttrName, + accessedOp->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); if (isLoadLike) { @@ -452,14 +514,9 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( } if (!isa(value.getType())) { SmallVector indices; - for (auto idx : offsets) { - if (auto val = dyn_cast(idx)) { - indices.push_back(val); - } else { - auto idxVal = rewriter.create( - loc, rewriter.getIndexAttr(*getConstantIntValue(idx))); - indices.push_back(idxVal); - } + for (auto idxOfr : offsets) { + auto idx = getValueOrCreateConstantIndexOp(rewriter, loc, idxOfr); + indices.push_back(idx); } result = rewriter.create(loc, value, iterArg, indices); } else { @@ -467,7 +524,7 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( offsets, sizes, strides); } rewriter.create(loc, result) - ->setAttr(mlir::dicp::discreteAttrName, + ->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); rewriter.restoreInsertionPoint(insertPoint); if constexpr (std::is_same_v) { @@ -475,7 +532,7 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( rewriter .replaceOpWithNewOp(op, op.getMask(), newOpResult, op.getOther()) - ->setAttr(mlir::dicp::discreteAttrName, + ->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); } else { rewriter.replaceOp(op, newOpResult); @@ -495,7 +552,7 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( loc, accessedOp.getType(), accessedOp.getAtomicRmwOp(), accessedOp.getPtr(), accessedOp.getVal(), nullptr, accessedOp.getSem(), accessedOp.getScope()) - ->setAttr(mlir::dicp::discreteAttrName, + ->setAttr(ConverterUtils::discreteAttrName, UnitAttr::get(rewriter.getContext())); b.create(loc); }); @@ -514,250 +571,6 @@ LogicalResult UnstructuredMemAccessConverter::matchAndRewrite( return success(); } -void replaceOperands(MutableArrayRef oprs, RewriterBase &rewriter, - llvm::DenseMap &offsetMap) { - for (auto it = oprs.begin(); it != oprs.end(); ++it) { - auto &opr = *it; - auto operand = opr.get(); - if (auto tensorType = dyn_cast(operand.getType()); - tensorType && isa(tensorType.getElementType())) { - parse(operand, operand.getLoc(), rewriter, offsetMap); - opr.set(offsetMap.at(operand).getOffset()); - } else if (auto ptrType = - dyn_cast(operand.getType())) { - parse(operand, operand.getLoc(), rewriter, offsetMap); - if (auto tensorType = - dyn_cast(ptrType.getPointeeType())) { - for (auto offset : offsetMap.at(operand).getOffsets()) { - it->set(offset); - ++it; - } - --it; - } else { - opr.set(offsetMap.at(operand).getOffset()); - } - } - } -} - -void replaceArgs(ValueRange args, RewriterBase &rewriter, - llvm::DenseMap &offsetMap) { - for (auto it = args.begin(); it != args.end(); ++it) { - auto arg = *it; - if (auto tensorType = dyn_cast(arg.getType()); - tensorType && isa(tensorType.getElementType())) { - RewriterBase::InsertionGuard guard(rewriter); - if (auto blockArg = dyn_cast(arg)) { - rewriter.setInsertionPointToStart(blockArg.getOwner()); - } else { - rewriter.setInsertionPointAfterValue(arg); - } - auto tempVar = rewriter - .create( - arg.getLoc(), arg.getType(), ValueRange({})) - ->getResult(0); - parse(arg, arg.getLoc(), rewriter, offsetMap); - auto src = offsetMap.at(arg).getPtr(); - rewriter.replaceAllUsesWith(arg, tempVar); - arg.setType(RankedTensorType::get(tensorType.getShape(), - rewriter.getIntegerType(64))); - src = rewriter.create(arg.getLoc(), tempVar.getType(), - src); - rewriter.replaceOpWithNewOp( - tempVar.getDefiningOp(), tempVar.getType(), src, arg); - } else if (auto ptrType = dyn_cast(arg.getType())) { - RewriterBase::InsertionGuard guard(rewriter); - if (auto blockArg = dyn_cast(arg)) { - rewriter.setInsertionPointToStart(blockArg.getOwner()); - } else { - rewriter.setInsertionPointAfterValue(arg); - } - auto tempVar = rewriter - .create( - arg.getLoc(), arg.getType(), ValueRange({})) - ->getResult(0); - parse(arg, arg.getLoc(), rewriter, offsetMap); - rewriter.replaceAllUsesWith(arg, tempVar); - if (auto tensorType = - dyn_cast(ptrType.getPointeeType())) { - auto srcOp = - offsetMap.at(arg).getPtr().getDefiningOp(); - arg.setType(rewriter.getIntegerType(32)); - SmallVector newOffsets; - for (auto offset : offsetMap.at(arg).getOffsets()) { - newOffsets.push_back(*it); - ++it; - } - --it; - rewriter.replaceOpWithNewOp( - tempVar.getDefiningOp(), tempVar.getType(), srcOp.getBase(), - srcOp.getShape(), srcOp.getStrides(), newOffsets, srcOp.getOrder()); - } else { - auto src = offsetMap.at(arg).getPtr(); - arg.setType(rewriter.getIntegerType(64)); - rewriter.replaceOpWithNewOp( - tempVar.getDefiningOp(), tempVar.getType(), src, arg); - } - } - } -} - -void convertTensorPtrPre(LoopLikeOpInterface op, RewriterBase &rewriter, - llvm::DenseMap &offsetMap) { - if (auto whileOp = dyn_cast(op.getOperation())) { - replaceArgs(whileOp.getBeforeArguments(), rewriter, offsetMap); - replaceOperands(whileOp.getInitsMutable(), rewriter, offsetMap); - replaceArgs(whileOp.getAfterArguments(), rewriter, offsetMap); - replaceArgs(whileOp->getResults(), rewriter, offsetMap); - replaceOperands(whileOp.getConditionOp().getArgsMutable(), rewriter, - offsetMap); - } else { - replaceArgs(op.getRegionIterArgs(), rewriter, offsetMap); - replaceOperands(op.getInitsMutable(), rewriter, offsetMap); - } -} - -void convertTensorPtrPost(LoopLikeOpInterface op, RewriterBase &rewriter, - llvm::DenseMap &offsetMap) { - if (auto whileOp = dyn_cast(op.getOperation())) { - replaceOperands(whileOp.getYieldOp()->getOpOperands(), rewriter, offsetMap); - } else { - replaceArgs(op->getResults(), rewriter, offsetMap); - replaceOperands(*op.getYieldedValuesMutable(), rewriter, offsetMap); - } -} - -int getPtrTensorRank(Type type) { - if (auto ptrType = dyn_cast(type)) { - if (auto tensorType = - dyn_cast(ptrType.getPointeeType())) { - return tensorType.getRank(); - } - } - return 0; -} - -SmallVector constructOperands(ValueRange operands, Value tempVar, - IRMapping mapping) { - SmallVector newOperands; - for (auto opr : operands) { - opr = mapping.lookupOrDefault(opr); - newOperands.push_back(opr); - auto numAppend = getPtrTensorRank(opr.getType()) - 1; - if (numAppend > 0) - newOperands.append(numAppend, tempVar); - } - return newOperands; -} - -SmallVector constructTypes(TypeRange types) { - SmallVector newTypes; - for (auto type : types) { - newTypes.push_back(type); - if (auto ptrType = dyn_cast(type)) { - if (auto tensorType = - dyn_cast(ptrType.getPointeeType())) { - if (tensorType.getRank() > 0) - newTypes.append(tensorType.getRank() - 1, - IntegerType::get(type.getContext(), 32)); - } - } - } - return newTypes; -} - -void replacePtrLoopArguments(Operation *rootOp, - llvm::DenseMap &offsetMap) { - std::function convertTensorPtr = - [&](LoopLikeOpInterface op) { - IRRewriter rewriter(op.getContext()); - IRMapping mapping; - LoopLikeOpInterface newOp; - rewriter.setInsertionPointAfter(op); - Value tempVar = - rewriter - .create( - op.getLoc(), rewriter.getI32Type(), ValueRange({})) - ->getResult(0); - if (auto forOp = dyn_cast(op.getOperation())) { - newOp = rewriter.create( - forOp.getLoc(), forOp.getLowerBound(), forOp.getUpperBound(), - forOp.getStep(), - constructOperands(forOp.getInitArgs(), tempVar, mapping), - [&](OpBuilder &b, Location loc, Value iv, ValueRange args) { - mapping.map(forOp.getInductionVar(), iv); - auto newArgIter = args.begin(); - for (auto oldArg : forOp.getRegionIterArgs()) { - mapping.map(oldArg, *newArgIter); - std::advance(newArgIter, - std::max(getPtrTensorRank(oldArg.getType()), 1)); - } - for (auto &bodyOp : forOp.getBody()->without_terminator()) { - b.clone(bodyOp, mapping); - } - auto yieldOp = - cast(forOp.getBody()->getTerminator()); - b.create( - yieldOp.getLoc(), - constructOperands(yieldOp.getOperands(), tempVar, mapping)); - }); - } else if (auto whileOp = dyn_cast(op.getOperation())) { - newOp = rewriter.create( - whileOp.getLoc(), constructTypes(whileOp->getResultTypes()), - constructOperands(whileOp.getInits(), tempVar, mapping), - [&](OpBuilder &b, Location loc, ValueRange args) { - auto newArgIter = args.begin(); - for (auto oldArg : whileOp.getBeforeArguments()) { - mapping.map(oldArg, *newArgIter); - std::advance(newArgIter, - std::max(getPtrTensorRank(oldArg.getType()), 1)); - } - for (auto &bodyOp : - whileOp.getBeforeBody()->without_terminator()) { - b.clone(bodyOp, mapping); - } - auto conditionOp = whileOp.getConditionOp(); - b.create( - conditionOp.getLoc(), - mapping.lookup(conditionOp.getCondition()), - constructOperands(conditionOp.getArgs(), tempVar, mapping)); - }, - [&](OpBuilder &b, Location loc, ValueRange args) { - auto newArgIter = args.begin(); - for (auto oldArg : whileOp.getAfterArguments()) { - mapping.map(oldArg, *newArgIter); - std::advance(newArgIter, - std::max(getPtrTensorRank(oldArg.getType()), 1)); - } - for (auto &bodyOp : - whileOp.getAfterBody()->without_terminator()) { - b.clone(bodyOp, mapping); - } - auto yieldOp = whileOp.getYieldOp(); - b.create( - yieldOp.getLoc(), - constructOperands(yieldOp.getOperands(), tempVar, mapping)); - }); - } else { - llvm_unreachable("Unsupported loop op"); - } - auto resIter = newOp->result_begin(); - for (auto res : op->getResults()) { - rewriter.replaceAllUsesWith(res, *resIter); - std::advance(resIter, std::max(getPtrTensorRank(res.getType()), 1)); - } - rewriter.eraseOp(op); - op = newOp; - convertTensorPtrPre(op, rewriter, offsetMap); - for (auto *region : op.getLoopRegions()) - region->walk(convertTensorPtr); - convertTensorPtrPost(op, rewriter, offsetMap); - return WalkResult::skip(); - }; - - rootOp->walk(convertTensorPtr); -} - void TritonToUnstructurePass::runPreparse(LoopLikeOpInterface op) { IRRewriter rewriter(&getContext()); auto loc = op.getLoc(); @@ -785,7 +598,7 @@ void TritonToUnstructurePass::runPreparse(LoopLikeOpInterface op) { auto &os = llvm::dbgs(); os << "Pre-parsing result of\n" << arg << "\nis "; for (auto structured : offsetMap[arg].getStructuredRef()) - os << structured; + os << static_cast(structured); os << '\n'; }); } @@ -822,13 +635,45 @@ void TritonToUnstructurePass::runParse(MemAccOpTy op) { isFromTensorArg(op.getPtr(), fromTensorArg); } +LogicalResult +TritonToUnstructurePass::processIfYieldAddHoistOperations(ModuleOp moduleOp) { + mlir::RewritePatternSet patterns(&getContext()); + patterns.add( + patterns.getContext()); + if (failed(applyPatternsGreedily(moduleOp, std::move(patterns)))) { + moduleOp.emitWarning("IfYieldAddHoist processing failed"); + return failure(); + } + return success(); +} + +TritonToUnstructurePass::TritonToUnstructurePass( + const TritonToUnstructureOptions &options) + : TritonToUnstructureBase(options) {} + void TritonToUnstructurePass::runOnOperation() { + compileOn91095Flag = this->compileOn91095; + forceSimtTemplateFlag = this->forceSimtTemplate; + + LLVM_DEBUG({ + auto &os = llvm::dbgs(); + os << "TritonToUnstructurePass started with options:\n"; + os << " compileOn91095: " << compileOn91095Flag << "\n"; + os << " forceSimtTemplate: " << forceSimtTemplateFlag << "\n"; + }); ModuleOp moduleOp = getOperation(); MLIRContext *ctx = &getContext(); - replacePtrLoopArguments(moduleOp, offsetMapForLoopArgs); + moduleOp->walk([this](triton::FuncOp funcOp) { + replacePtrArguments(funcOp, offsetMapForLoopArgs); + }); offsetMapForLoopArgs.clear(); + + if (failed(processIfYieldAddHoistOperations(moduleOp))) { + moduleOp.emitWarning("Failed to process IfYieldAddHoist operations"); + } + moduleOp->walk([this](LoopLikeOpInterface op) { runPreparse(op); }); moduleOp->walk([this](Operation *op) { if (auto loadOp = dyn_cast(op)) { @@ -873,10 +718,10 @@ void TritonToUnstructurePass::getDependentDialects( registry.insert(); + triton::TritonDialect, triton::dicp::TritonDicpDialect>(); } -std::unique_ptr> -triton::createTritonToUnstructurePass() { - return std::make_unique(); +std::unique_ptr> triton::createTritonToUnstructurePass( + const TritonToUnstructureOptions &options) { + return std::make_unique(options); } diff --git a/compiler/lib/Utils/CMakeLists.txt b/compiler/lib/Utils/CMakeLists.txt index 9315bf5c..59417471 100644 --- a/compiler/lib/Utils/CMakeLists.txt +++ b/compiler/lib/Utils/CMakeLists.txt @@ -1,7 +1,9 @@ -add_triton_library(DICPUtils - Utils.cpp +add_triton_library(MLIRTritonNPUUtils + Utils.cpp + InterleaveOptimization.cpp LINK_LIBS PUBLIC MLIRIR TritonIR + BiShengIRHIVMDialect ) \ No newline at end of file diff --git a/compiler/lib/Utils/InterleaveOptimization.cpp b/compiler/lib/Utils/InterleaveOptimization.cpp new file mode 100644 index 00000000..4fe29524 --- /dev/null +++ b/compiler/lib/Utils/InterleaveOptimization.cpp @@ -0,0 +1,705 @@ + + +#include "dicp/Utils/InterleaveOptimization.h" +#include "dicp/Utils/Utils.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "mlir/Dialect/MemRef/IR/MemRef.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" +#include "mlir/Dialect/Utils/StaticValueUtils.h" +#include "mlir/IR/Attributes.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/OpDefinition.h" +#include "mlir/Interfaces/ViewLikeInterface.h" +#include "mlir/Support/LogicalResult.h" + +#include "mlir/IR/Operation.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Casting.h" +#include "llvm/Support/Debug.h" +#include +#include + +namespace mlir { +namespace triton { +// For origin MemRefType of ReinterpretCastOp under interleave state, here wanna +// adjust its shape info by expanding last dimension double. +MemRefType expandInterleaveMemRefType(MemRefType originType) { + // Double the last dimension shape + SmallVector shape(originType.getShape()); + shape.back() = shape.back() * 2; + + // Adjuest layout attribute + StridedLayoutAttr originLayout = + llvm::dyn_cast(originType.getLayout()); + // If offset is static, just reset it to 0 + auto offset = originLayout.getOffset() == ShapedType::kDynamic + ? originLayout.getOffset() + : 0; + // Set last dimension stride to 1 + SmallVector stride(originLayout.getStrides()); + stride.back() = 1; + + return MemRefType::get( + shape, originType.getElementType(), + StridedLayoutAttr::get(originType.getContext(), offset, stride)); +} + +// ********************* +// ** NOTE ** +// ********************* +// How to determine new offset is a little tricky and specific +// Here just consider this state in triton language: +// +// dim_range = tl.arange(0, BLOCK // 2) +// last_dim_even_range = dim_range * 2 +// last_dim_odd_range = dim_range * 2 + 1 +// +// Here `multiply two` represents that last dimension stride is 2, and +// `add constant one` represents whether it's odd index part of +// deinterleave result. +// +// Therefore, how to distinguish interleave/deinterleave on even index or odd +// index is whether last dimension range explicitly `add constant one` without +// any other operation. In IR it's shown that whether defining op of +// `castOffset` is an arith::addOp, as this arith::addOp would contain above +// `add constant one` opeartion after LegacyAddPtrConverter. +// +// Well, index mode should be passed to interleave/deinterleave, in other words, +// `add constant one` should work on offset of next insert_slice/extract_slic. +// The new reinterpretcast just wanna describe whole tensor, so new castOffset +// is just from non-last diemsnion accumulation and remove `add constant one` +bool checkIsCaseOffsetValid(OpFoldResult originOffset) { + // If offset is constant int(IndexAttr), the int value could only be 0 or 1 + // if offset is a value from add constant operation and not from `add constant + // one` operation, it's invalid. + if (llvm::isa(originOffset)) { + int64_t intOffset = getConstantIntValue(originOffset).value(); + return intOffset == 0 || intOffset == 1; + } else if (llvm::isa(originOffset)) { + auto op = cast(originOffset).getDefiningOp(); + if (op && llvm::isa(op)) { + if (auto addOp = dyn_cast(op)) { + if (auto constLHS = addOp.getLhs().getDefiningOp()) { + return dyn_cast(constLHS.getValueAttr()).getInt() == 1; + } + if (auto constRHS = addOp.getRhs().getDefiningOp()) { + return dyn_cast(constRHS.getValueAttr()).getInt() == 1; + } + } + } + } + return true; +} + +std::pair +recountReinterpretCastOffset(OpFoldResult originOffset, Builder &builder) { + // To trace value type offset + std::function traceOffset = [&](Operation *op) -> bool { + // Consider constant one in `add constant one` operation + if (llvm::isa(op)) + return false; + + if (llvm::isa(op)) { + auto addOp = llvm::cast(op); + if (auto constLHS = addOp.getLhs().getDefiningOp()) { + assert(dyn_cast(constLHS.getValueAttr()).getInt() == 1 && + "Arith::constant value of addi's operand must be 1 when " + "calculate deinterleave offset"); + return false; + } + if (auto constRHS = addOp.getRhs().getDefiningOp()) { + assert(dyn_cast(constRHS.getValueAttr()).getInt() == 1 && + "Arith::constant value of addi's operand must be 1 when " + "calculate deinterleave offset"); + return false; + } + } + return true; + }; + + IndexMode evenOrOdd = IndexMode::EVEN_MODE; + // Reuse origin offset if there's no 'add constant one' + OpFoldResult newOffset = originOffset; + if (llvm::isa(originOffset)) { + // If offset is constant int(IndexAttr), + // the int value could only be 0 or 1 + int64_t intOffset = getConstantIntValue(originOffset).value(); + assert((intOffset == 0 || intOffset == 1)); + if (intOffset == 1) { + evenOrOdd = IndexMode::ODD_MODE; + newOffset = builder.getIndexAttr(0); + } + } else if (llvm::isa(originOffset)) { + if (!traceOffset(cast(originOffset).getDefiningOp())) { + evenOrOdd = IndexMode::ODD_MODE; + Operation *traceResult = findFirstMatchingOperandDef( + cast(originOffset).getDefiningOp(), traceOffset); + assert(traceResult->getNumResults() == 1 && + "Offset defining operation must have one result"); + newOffset = traceResult->getResult(0); + } + } + + return {newOffset, evenOrOdd}; +} + +LogicalResult +DeinterleaveStatusOptimization(triton::LoadOp op, + triton::LoadOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) { + auto ptr = adaptor.getPtr(); + if (auto reinterpretCast = ptr.getDefiningOp()) { + auto loc = op.getLoc(); + + // 1. Get new source memref type + auto srcType = expandInterleaveMemRefType(reinterpretCast.getType()); + + // 2. Create new ReinterpretCastOp + auto originCastOffset = reinterpretCast.getConstifiedMixedOffset(); + auto castSize = reinterpretCast.getConstifiedMixedSizes(); + auto castStride = reinterpretCast.getConstifiedMixedStrides(); + // Actually, `castSize` is always constant value as `MemRefType` result + if (auto lastDimSize = getConstantIntValue(castSize.back())) { + castSize.back() = rewriter.getIndexAttr(lastDimSize.value() * 2); + } else { + return failure(); + } + // Last element of castStride is also constant value as prerequisite + // is that last dimension stride of casted memref type is always 2. + castStride.back() = rewriter.getIndexAttr(1); + if (!checkIsCaseOffsetValid(originCastOffset)) { + return failure(); + } + auto [castOffset, indexMode] = + recountReinterpretCastOffset(originCastOffset, rewriter); + auto newCastOp = rewriter.create( + loc, srcType, reinterpretCast.getViewSource(), castOffset, castSize, + castStride); + + // 3. Create new memref allocOp + auto newAllocOp = rewriter.create( + loc, MemRefType::get(srcType.getShape(), srcType.getElementType())); + + // 4. Implement memref copy and bufferization back to tensor + rewriter.create(loc, newCastOp.getResult(), newAllocOp); + Value newTensor = rewriter.create( + loc, + RankedTensorType::get(srcType.getShape(), srcType.getElementType()), + newAllocOp, true /* restrict */, true /* writable */); + + // 5. Implement tensor extract_slice to represent deinterleave + // Here use `castOffset` to determine whether even index deinterleave or + // odd index. + SmallVector extractOffsets(srcType.getRank(), + rewriter.getIndexAttr(0)); + SmallVector extractStrides(srcType.getRank(), + rewriter.getIndexAttr(1)); + SmallVector extractSizes = llvm::to_vector( + llvm::map_range(srcType.getShape(), [&](int64_t dim) -> OpFoldResult { + return rewriter.getIndexAttr(dim); + })); + + // Adjust extract_slice shape + switch (indexMode) { + case IndexMode::EVEN_MODE: + extractOffsets.back() = rewriter.getIndexAttr(0); + break; + case IndexMode::ODD_MODE: + extractOffsets.back() = rewriter.getIndexAttr(1); + break; + } + extractStrides.back() = rewriter.getIndexAttr(2); + extractSizes.back() = rewriter.getIndexAttr(srcType.getShape().back() / 2); + + Value deinterleaveSlice = rewriter.create( + loc, newTensor, extractOffsets, extractSizes, extractStrides); + + rewriter.replaceOp(op, deinterleaveSlice); + return success(); + } + + return failure(); +} + +LogicalResult DeinterleaveStatusWithMaskOptimization( + triton::LoadOp op, triton::LoadOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter, MaskState &mstate, Value localMem) { + auto ptr = adaptor.getPtr(); + if (auto reinterpretCast = ptr.getDefiningOp()) { + auto loc = op.getLoc(); + + // 1. Get new source memref type + auto srcType = expandInterleaveMemRefType(reinterpretCast.getType()); + + // 2. Create new ReinterpretCastOp + auto originCastOffset = reinterpretCast.getConstifiedMixedOffset(); + auto castSize = reinterpretCast.getConstifiedMixedSizes(); + auto castStride = reinterpretCast.getConstifiedMixedStrides(); + + if (auto lastDimSize = getConstantIntValue(castSize.back())) { + castSize.back() = rewriter.getIndexAttr(lastDimSize.value() * 2); + } else { + return failure(); + } + castStride.back() = rewriter.getIndexAttr(1); + if (!checkIsCaseOffsetValid(originCastOffset)) { + return failure(); + } + auto [castOffset, indexMode] = + recountReinterpretCastOffset(originCastOffset, rewriter); + + auto newCastOp = rewriter.create( + loc, srcType, reinterpretCast.getViewSource(), castOffset, castSize, + castStride); + + // 3. Create new memref allocOp + // To reuse existing linalg::fill, here need to change insertion point + auto savedInsertPoint = rewriter.saveInsertionPoint(); + rewriter.setInsertionPointAfterValue(localMem); + auto newAllocOp = rewriter.create( + loc, MemRefType::get(srcType.getShape(), srcType.getElementType())); + rewriter.restoreInsertionPoint(savedInsertPoint); + + // 4. Broadcast other value by linalg.fill if necessary + auto other = op.getOther(); + // While deinterleave optimization will just adjust last dimension info + // and origin mask state wouldn't involve last dimension. Therefore in + // current `scf.if + linalg.fill` combination, condition of `if` could be + // kept and just replace linalg.fill' + if (other) { + assert(localMem.hasOneUse() && + llvm::isa(*(localMem.getUsers().begin()))); + auto originFillOp = + llvm::dyn_cast(*(localMem.getUsers().begin())); + + assert(llvm::isa(originFillOp->getParentOp())); + auto ifOp = llvm::dyn_cast(originFillOp->getParentOp()); + + auto newFillOp = ifOp.getThenBodyBuilder().create( + originFillOp.getLoc(), originFillOp.getInputs(), + ValueRange{newAllocOp}); + rewriter.replaceOp(originFillOp, newFillOp); + } + + // 5. Implement new subview, memref copy and bufferization back to tensor + SmallVector subviewStrides(srcType.getRank(), + rewriter.getIndexAttr(1)); + SmallVector subviewOffsets = mstate.offsets; + SmallVector subviewSizes = mstate.dims; + // Just adjust last dimension size to double + std::optional originSubviewLastDim = + getConstantIntValue(subviewSizes.back()); + assert(originSubviewLastDim.has_value()); + subviewSizes.back() = + rewriter.getIndexAttr(originSubviewLastDim.value() * 2); + + auto argSubviewType = memref::SubViewOp::inferResultType( + srcType, subviewOffsets, subviewSizes, subviewStrides); + // alloca subview type doesn't carry layout attribute + auto allocSubviewType = memref::SubViewOp::inferResultType( + newAllocOp.getType(), subviewOffsets, subviewSizes, subviewStrides); + + memref::SubViewOp srcSubview = rewriter.create( + loc, llvm::cast(argSubviewType), newCastOp, subviewOffsets, + subviewSizes, subviewStrides); + memref::SubViewOp dstSubview = rewriter.create( + loc, llvm::cast(allocSubviewType), newAllocOp, + subviewOffsets, subviewSizes, subviewStrides); + rewriter.create(loc, srcSubview, dstSubview); + Value newTensor = rewriter.create( + loc, + RankedTensorType::get(srcType.getShape(), srcType.getElementType()), + newAllocOp, true /* restrict */, true /* writable */); + + // 6. Implement tensor extract_slice to represent deinterleave + // Here use `castOffset` to determine whether even index deinterleave or + // odd index. + SmallVector extractOffsets(srcType.getRank(), + rewriter.getIndexAttr(0)); + SmallVector extractStrides(srcType.getRank(), + rewriter.getIndexAttr(1)); + SmallVector extractSizes = llvm::to_vector( + llvm::map_range(srcType.getShape(), [&](int64_t dim) -> OpFoldResult { + return rewriter.getIndexAttr(dim); + })); + + switch (indexMode) { + case IndexMode::EVEN_MODE: + extractOffsets.back() = rewriter.getIndexAttr(0); + break; + case IndexMode::ODD_MODE: + extractOffsets.back() = rewriter.getIndexAttr(1); + break; + } + extractStrides.back() = rewriter.getIndexAttr(2); + extractSizes.back() = rewriter.getIndexAttr(srcType.getShape().back() / 2); + + Value deinterleaveSlice = rewriter.create( + loc, newTensor, extractOffsets, extractSizes, extractStrides); + + rewriter.replaceOp(op, deinterleaveSlice); + return success(); + } + return failure(); +} + +LogicalResult +InterleaveStatusOptimization(SmallVector materializeVec) { + OpBuilder builder(materializeVec[1]); + auto loc = materializeVec[1]->getLoc(); + + auto firstReinterpretCastOp = + llvm::dyn_cast( + materializeVec[0]) + .getDest() + .getDefiningOp(); + auto secondReinterpretCastOp = + llvm::dyn_cast( + materializeVec[1]) + .getDest() + .getDefiningOp(); + + assert(firstReinterpretCastOp && secondReinterpretCastOp); + + // Judge whether two `ReinterpretCastOp` shape satisfy interleave state + // a. both size are equal + if (!isEqualConstantIntOrValueArray( + firstReinterpretCastOp.getConstifiedMixedSizes(), + secondReinterpretCastOp.getConstifiedMixedSizes())) { + return failure(); + } + // b. both strides are equal + if (!isEqualConstantIntOrValueArray( + firstReinterpretCastOp.getConstifiedMixedStrides(), + secondReinterpretCastOp.getConstifiedMixedStrides())) { + return failure(); + } + // c. both offsets should satisfy tricky rule + auto firstOriginCastOffset = + firstReinterpretCastOp.getConstifiedMixedOffset(); + auto secondOriginCastOffset = + secondReinterpretCastOp.getConstifiedMixedOffset(); + if (!checkIsCaseOffsetValid(firstOriginCastOffset) || + !checkIsCaseOffsetValid(secondOriginCastOffset)) { + return failure(); + } + + std::pair indexModeRecord; + OpFoldResult newCastOffset; + if (llvm::isa(firstOriginCastOffset) && + llvm::isa(secondOriginCastOffset)) { + auto [firstCastOffset, firstIndexMode] = + recountReinterpretCastOffset(firstOriginCastOffset, builder); + auto [secondCastOffset, secondIndexMode] = + recountReinterpretCastOffset(secondOriginCastOffset, builder); + + if (!(static_cast(firstIndexMode) ^ static_cast(secondIndexMode))) + return failure(); + newCastOffset = builder.getIndexAttr(0); + indexModeRecord = {firstIndexMode, secondIndexMode}; + + } else if (llvm::isa(firstOriginCastOffset) && + llvm::isa(secondOriginCastOffset)) { + auto [firstCastOffset, firstIndexMode] = + recountReinterpretCastOffset(firstOriginCastOffset, builder); + auto [secondCastOffset, secondIndexMode] = + recountReinterpretCastOffset(secondOriginCastOffset, builder); + + if (!(static_cast(firstIndexMode) ^ + static_cast(secondIndexMode)) || + (llvm::dyn_cast(firstCastOffset) != + llvm::dyn_cast(secondCastOffset))) + return failure(); + + if (firstIndexMode == IndexMode::EVEN_MODE) { + newCastOffset = llvm::dyn_cast(firstCastOffset); + } + if (secondIndexMode == IndexMode::EVEN_MODE) { + newCastOffset = llvm::dyn_cast(secondCastOffset); + } + indexModeRecord = {firstIndexMode, secondIndexMode}; + + } else { + return failure(); + } + + // Create new op + // 1. Get new destination memref type + auto dstType = expandInterleaveMemRefType(firstReinterpretCastOp.getType()); + + // 2. New tensor::EmptyOp + auto emptyTensor = builder.create(loc, dstType.getShape(), + dstType.getElementType()); + + // 3. New insert_slice from materialization source into new empty tensor + SmallVector insertOffsets(dstType.getRank(), + builder.getIndexAttr(0)); + SmallVector insertStrides(dstType.getRank(), + builder.getIndexAttr(1)); + SmallVector insertSizes = llvm::to_vector( + llvm::map_range(dstType.getShape(), [&](int64_t dim) -> OpFoldResult { + return builder.getIndexAttr(dim); + })); + insertStrides.back() = builder.getIndexAttr(2); + insertSizes.back() = builder.getIndexAttr(dstType.getShape().back() / 2); + if (indexModeRecord.first == IndexMode::ODD_MODE) { + insertOffsets.back() = builder.getIndexAttr(1); + } else { + insertOffsets.back() = builder.getIndexAttr(0); + } + auto insertFirst = builder.create( + loc, + llvm::dyn_cast( + materializeVec[0]) + .getSource(), + emptyTensor.getResult(), insertOffsets, insertSizes, insertStrides); + + if (indexModeRecord.second == IndexMode::ODD_MODE) { + insertOffsets.back() = builder.getIndexAttr(1); + } else { + insertOffsets.back() = builder.getIndexAttr(0); + } + auto insertSecond = builder.create( + loc, + llvm::dyn_cast( + materializeVec[1]) + .getSource(), + insertFirst.getResult(), insertOffsets, insertSizes, insertStrides); + + // 4. Reinterpret_cast block arg + auto newCastSize = firstReinterpretCastOp.getConstifiedMixedSizes(); + auto newCastStride = firstReinterpretCastOp.getConstifiedMixedStrides(); + newCastSize.back() = builder.getIndexAttr(dstType.getShape().back()); + newCastStride.back() = builder.getIndexAttr(1); + auto newCastOp = builder.create( + loc, dstType, firstReinterpretCastOp.getViewSource(), newCastOffset, + newCastSize, newCastStride); + + // 5. Create new bufferization::MaterializeInDestinationOp + auto newStoreOp = builder.create( + loc, insertSecond.getResult(), newCastOp.getResult()); + // Setting writable is necessary as dst is memref type + newStoreOp.setWritable(true); + + // 6. Erase origin materialization + materializeVec[0]->erase(); + materializeVec[1]->erase(); + + return success(); +} + +LogicalResult +InterleaveStatusWithMaskOptimization(SmallVector materializeVec) { + OpBuilder builder(materializeVec[1]); + + auto firstSubviewOpOfReCast = + llvm::dyn_cast( + materializeVec[0]) + .getDest() + .getDefiningOp(); + auto firstSrcExtractSlice = + llvm::dyn_cast( + materializeVec[0]) + .getSource() + .getDefiningOp(); + auto firstReinterpretCastOp = firstSubviewOpOfReCast.getSource() + .getDefiningOp(); + + auto secondSubviewOpOfReCast = + llvm::dyn_cast( + materializeVec[1]) + .getDest() + .getDefiningOp(); + auto secondSrcExtractSlice = + llvm::dyn_cast( + materializeVec[1]) + .getSource() + .getDefiningOp(); + auto secondReinterpretCastOp = + secondSubviewOpOfReCast.getSource() + .getDefiningOp(); + + // 1. Both source shapes of subview and extract_slice are equal + if (firstSubviewOpOfReCast.getSourceType().getShape() != + firstSrcExtractSlice.getSourceType().getShape()) + return failure(); + if (secondSubviewOpOfReCast.getSourceType().getShape() != + secondSrcExtractSlice.getSourceType().getShape()) + return failure(); + if (firstSubviewOpOfReCast.getSourceType().getShape() != + secondSubviewOpOfReCast.getSourceType().getShape()) + return failure(); + + // 2. both mask state are equal + std::function cmpFunc = + mlir::isEqualConstantIntOrValue; + if (!mlir::detail::sameOffsetsSizesAndStrides(firstSubviewOpOfReCast, + firstSrcExtractSlice, cmpFunc)) + return failure(); + if (!mlir::detail::sameOffsetsSizesAndStrides(secondSubviewOpOfReCast, + secondSrcExtractSlice, cmpFunc)) + return failure(); + if (!mlir::detail::sameOffsetsSizesAndStrides( + firstSubviewOpOfReCast, secondSubviewOpOfReCast, cmpFunc)) + return failure(); + + // 3. Still judge whether two `ReinterpretCastOp` shape satisfy request + // a. both size are equal + if (!isEqualConstantIntOrValueArray( + firstReinterpretCastOp.getConstifiedMixedSizes(), + secondReinterpretCastOp.getConstifiedMixedSizes())) + return failure(); + // b. both strides are equal + if (!isEqualConstantIntOrValueArray( + firstReinterpretCastOp.getConstifiedMixedStrides(), + secondReinterpretCastOp.getConstifiedMixedStrides())) + return failure(); + // c. both offsets should satisfy tricky rule + auto firstOriginCastOffset = + firstReinterpretCastOp.getConstifiedMixedOffset(); + auto secondOriginCastOffset = + secondReinterpretCastOp.getConstifiedMixedOffset(); + if (!checkIsCaseOffsetValid(firstOriginCastOffset) || + !checkIsCaseOffsetValid(secondOriginCastOffset)) { + return failure(); + } + + std::pair indexModeRecord; + OpFoldResult newCastOffset; + if (llvm::isa(firstOriginCastOffset) && + llvm::isa(secondOriginCastOffset)) { + auto [firstCastOffset, firstIndexMode] = + recountReinterpretCastOffset(firstOriginCastOffset, builder); + auto [secondCastOffset, secondIndexMode] = + recountReinterpretCastOffset(secondOriginCastOffset, builder); + + if (!(static_cast(firstIndexMode) ^ static_cast(secondIndexMode))) + return failure(); + newCastOffset = builder.getIndexAttr(0); + indexModeRecord = {firstIndexMode, secondIndexMode}; + + } else if (llvm::isa(firstOriginCastOffset) && + llvm::isa(secondOriginCastOffset)) { + auto [firstCastOffset, firstIndexMode] = + recountReinterpretCastOffset(firstOriginCastOffset, builder); + auto [secondCastOffset, secondIndexMode] = + recountReinterpretCastOffset(secondOriginCastOffset, builder); + + if (!(static_cast(firstIndexMode) ^ + static_cast(secondIndexMode)) || + (llvm::dyn_cast(firstCastOffset) != + llvm::dyn_cast(secondCastOffset))) + return failure(); + + if (firstIndexMode == IndexMode::EVEN_MODE) { + newCastOffset = llvm::dyn_cast(firstCastOffset); + } + if (secondIndexMode == IndexMode::EVEN_MODE) { + newCastOffset = llvm::dyn_cast(secondCastOffset); + } + indexModeRecord = {firstIndexMode, secondIndexMode}; + + } else { + return failure(); + } + auto loc = materializeVec[1]->getLoc(); + + // Create new op + // 1. Get new destination memref type + auto dstType = expandInterleaveMemRefType(firstReinterpretCastOp.getType()); + + // 2. New tensor::EmptyOp + auto emptyTensor = builder.create(loc, dstType.getShape(), + dstType.getElementType()); + + // 3. New insert_slice from extract_slice source into new empty tensor + SmallVector insertOffsets(dstType.getRank(), + builder.getIndexAttr(0)); + SmallVector insertStrides(dstType.getRank(), + builder.getIndexAttr(1)); + SmallVector insertSizes = llvm::to_vector( + llvm::map_range(dstType.getShape(), [&](int64_t dim) -> OpFoldResult { + return builder.getIndexAttr(dim); + })); + insertStrides.back() = builder.getIndexAttr(2); + insertSizes.back() = builder.getIndexAttr(dstType.getShape().back() / 2); + if (indexModeRecord.first == IndexMode::ODD_MODE) { + insertOffsets.back() = builder.getIndexAttr(1); + } else { + insertOffsets.back() = builder.getIndexAttr(0); + } + auto insertFirst = builder.create( + loc, firstSrcExtractSlice.getSource(), emptyTensor.getResult(), + insertOffsets, insertSizes, insertStrides); + + if (indexModeRecord.second == IndexMode::ODD_MODE) { + insertOffsets.back() = builder.getIndexAttr(1); + } else { + insertOffsets.back() = builder.getIndexAttr(0); + } + auto insertSecond = builder.create( + loc, secondSrcExtractSlice.getSource(), insertFirst.getResult(), + insertOffsets, insertSizes, insertStrides); + + // 4. To enable store with mask, create new extract_slice + SmallVector extractOffsets = + firstSrcExtractSlice.getMixedOffsets(); + SmallVector extractStrides = + firstSrcExtractSlice.getMixedStrides(); + SmallVector extractSizes = firstSrcExtractSlice.getMixedSizes(); + if (!llvm::isa(extractSizes.back())) { + return failure(); + } + extractSizes.back() = builder.getIndexAttr( + getConstantIntValue(extractSizes.back()).value() * 2); + auto newSrcExtractSlice = builder.create( + loc, insertSecond.getResult(), extractOffsets, extractSizes, + extractStrides); + + // 5. Reinterpret_cast block arg + auto newCastSize = firstReinterpretCastOp.getConstifiedMixedSizes(); + auto newCastStride = firstReinterpretCastOp.getConstifiedMixedStrides(); + newCastSize.back() = builder.getIndexAttr(dstType.getShape().back()); + newCastStride.back() = builder.getIndexAttr(1); + auto newCastOp = builder.create( + loc, dstType, firstReinterpretCastOp.getViewSource(), newCastOffset, + newCastSize, newCastStride); + + // 6. Create new memref::SubViewOp of above new reinterpret_cast + // Here could reuse shape info of new extract_slice + auto dstSubviewType = memref::SubViewOp::inferResultType( + dstType, extractOffsets, extractSizes, extractStrides); + auto newSubviewOpOfReCast = builder.create( + loc, llvm::cast(dstSubviewType), newCastOp, extractOffsets, + extractSizes, extractStrides); + + // 7. Create new bufferization::MaterializeInDestinationOp + auto newStoreOp = builder.create( + loc, newSrcExtractSlice.getResult(), newSubviewOpOfReCast.getResult()); + // Setting writable is necessary as dst is memref type + newStoreOp.setWritable(true); + + // 8. Erase origin operation + materializeVec[0]->erase(); + materializeVec[1]->erase(); + if (firstSubviewOpOfReCast->use_empty()) { + firstSubviewOpOfReCast->erase(); + } + if (firstSrcExtractSlice->use_empty()) { + firstSrcExtractSlice->erase(); + } + if (secondSubviewOpOfReCast->use_empty()) { + secondSubviewOpOfReCast->erase(); + } + if (secondSrcExtractSlice->use_empty()) { + secondSrcExtractSlice->erase(); + } + + return success(); +} + +} // namespace triton +} // namespace mlir diff --git a/compiler/lib/Utils/Utils.cpp b/compiler/lib/Utils/Utils.cpp index cced4713..326da453 100644 --- a/compiler/lib/Utils/Utils.cpp +++ b/compiler/lib/Utils/Utils.cpp @@ -1,5 +1,8 @@ + + #include "dicp/Utils/Utils.h" +#include "bishengir/Dialect/Annotation/IR/Annotation.h" #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" @@ -12,6 +15,7 @@ #include "mlir/IR/Diagnostics.h" #include "mlir/IR/OpDefinition.h" #include "mlir/IR/Operation.h" +#include "mlir/IR/PatternMatch.h" #include "mlir/IR/Value.h" #include "mlir/Transforms/DialectConversion.h" @@ -34,24 +38,12 @@ #include #include #include +#include #include -#define DEBUG_TYPE "Dicp-Utils" -using namespace mlir; - -namespace mlir::dicp { - -llvm::StringRef getBackend(ModuleOp module) { - if (!module) - return llvm::StringRef(); +#define DEBUG_TYPE "DicpUtils" - if (auto strAttr = module->getAttrOfType("dicp.backend")) - return strAttr.getValue(); // StringRef，适配 StringSwitch - - return llvm::StringRef(); // 空字符串 -} - -bool isAscendBackend(ModuleOp module) { return getBackend(module) == "ascend"; } +namespace mlir { static Value createConstIndexValueOp(const Location &loc, OpBuilder &b, int64_t value) { @@ -66,6 +58,8 @@ static std::optional getConstantOfAttr(const OpFoldResult &arg) { return std::nullopt; } +namespace ConverterUtils { + std::optional getLastStrideOfReinterpretCastOp(memref::ReinterpretCastOp op) { SmallVector mixedStrides = op.getMixedStrides(); @@ -75,9 +69,28 @@ getLastStrideOfReinterpretCastOp(memref::ReinterpretCastOp op) { } OpFoldResult lastStride = mixedStrides.back(); - if (auto attr = dyn_cast(lastStride)) { + + if (op.getStaticStrides().back() > 0) { + return op.getStaticStrides().back(); + } else if (isa(op.getStrides().back())) { + auto u = op.getStrides().back(); + while (auto blkArg = dyn_cast(u)) { + if (auto forOp = dyn_cast(blkArg.getOwner()->getParentOp())) { + auto prt = forOp->getOperand(3 + blkArg.getArgNumber() - 1); + u = prt; + } else { + u = nullptr; + break; + } + } + if (!u) + return std::nullopt; + lastStride = u; + } + + if (auto attr = lastStride.dyn_cast()) { return getConstantOfAttr(lastStride); - } else if (auto value = dyn_cast(lastStride)) { + } else if (auto value = lastStride.dyn_cast()) { auto defOp = value.getDefiningOp(); if (auto constIndexOp = dyn_cast(defOp)) { int64_t constValue = constIndexOp.value(); @@ -90,6 +103,27 @@ getLastStrideOfReinterpretCastOp(memref::ReinterpretCastOp op) { return std::nullopt; } +bool isaPermutedMemRefType(MemRefType memRefType) { + auto [ptrStrides, ptrOffsets] = memRefType.getStridesAndOffset(); + LLVM_DEBUG({ + llvm::dbgs() << "---------- [BEG] ptrStrides ----------\n"; + for (auto stride : ptrStrides) + llvm::dbgs() << stride << " "; + llvm::dbgs() << "\n"; + llvm::dbgs() << "---------- [END] ptrStrides ----------\n"; + }); + + switch (ptrStrides.size()) { + case 0: + return false; + case 1: + return false; + default: { + return ptrStrides[ptrStrides.size() - 1] != 1; + } + } +} + Value getTransposedValue(Value source, const Location loc, ConversionPatternRewriter &rewriter, llvm::ArrayRef order) { @@ -183,6 +217,190 @@ Value getScalarValue(Value operand, Location loc, return nullptr; } +memref::SubViewOp makeSubViewOp(Value src, + const llvm::SmallVector &offsets, + const llvm::SmallVector &sizes, + const Location &loc, + ConversionPatternRewriter &rewriter) { + auto srcType = cast(src.getType()); + SmallVector strides(srcType.getRank(), + rewriter.getIndexAttr(1)); + auto dstType = + memref::SubViewOp::inferResultType(srcType, offsets, sizes, strides); + return rewriter.create(loc, dyn_cast(dstType), + src, offsets, sizes, strides); +} + +tensor::ExtractSliceOp +makeExtractSliceOp(Value src, const llvm::SmallVector &offsets, + const llvm::SmallVector &sizes, + const Location &loc, ConversionPatternRewriter &rewriter) { + auto srcType = cast(src.getType()); + SmallVector strides(srcType.getRank(), + rewriter.getIndexAttr(1)); + auto dstType = + tensor::ExtractSliceOp::inferResultType(srcType, offsets, sizes, strides); + return rewriter.create(loc, dstType, src, offsets, + sizes, strides); +} + +std::optional getFullShapeOp(Value val, + ConversionPatternRewriter &rewriter) { + while (true) { + if (isa(val)) { + auto blockArg = dyn_cast(val); + Operation *parentOp = blockArg.getOwner()->getParentOp(); + + // When BlockArgument is from the scf.for loop, trace its initial value + if (auto forOp = dyn_cast_or_null(parentOp)) { + auto init = forOp.getTiedLoopInit(blockArg); + if (!init) + return std::nullopt; + val = init->get(); + continue; + } + + // When BlockArgument is not scf::ForOp but FuncOp, it means that shape + // information can no longer be tracked. In this case, std::nullopt is + // returned, and getBoundarySizes() is called to return the current shape + // as the boundary. + if (isa(parentOp)) { + return std::nullopt; + } + + emitWarning(val.getLoc()) + << "getFullShapeOp() only support ReinterpretCastOp, " + "UnrealizedConversionCastOp " + "and scf.for's block argument, but got : " + << val << "\n"; + return std::nullopt; + } + + Operation *defOp = val.getDefiningOp(); + if (!defOp) + return std::nullopt; + + if (auto castOp = dyn_cast(defOp)) { + if (castOp.getInputs().size() == 1) { + val = castOp.getInputs()[0]; + continue; + } + return std::nullopt; + } + + if (!isa(val.getType())) + return std::nullopt; + + if (auto reCastOp = dyn_cast(defOp)) { + if (reCastOp->hasAttr("tensor_ptr_full_shape")) + return reCastOp; + val = reCastOp.getSource(); + continue; + } + + emitWarning(val.getLoc()) + << "getFullShapeOp() only support ReinterpretCastOp, " + "UnrealizedConversionCastOp " + "and scf.for's block argument, but got : " + << val << "\n"; + return std::nullopt; + } +} + +SmallVector +getBoundarySizes(llvm::ArrayRef boundaryCheck, Value ptr, + const Location &loc, ConversionPatternRewriter &rewriter) { + if (isa(ptr.getType())) { + ptr = rewriter.getRemappedValue(ptr); + } + + auto shapedType = dyn_cast_if_present(ptr.getType()); + if (!shapedType) { + LLVM_DEBUG(llvm::dbgs() << "ptr is not a ShapedType.\n";); + return {}; + } + + if (!shapedType.hasStaticShape()) { + LLVM_DEBUG(llvm::dbgs() << "shapedType does not have a static shape\n";); + return {}; + } + + auto fullShapeOp = getFullShapeOp(ptr, rewriter); + if (!fullShapeOp.has_value()) { + // If fullShapeOp has no value, the current shape is returned as the + // boundary. + SmallVector boundarySize = + getAsIndexOpFoldResult(rewriter.getContext(), shapedType.getShape()); + + return boundarySize; + } + + SmallVector boundarySize = + getAsIndexOpFoldResult(rewriter.getContext(), shapedType.getShape()); + + auto fullShapeReCast = + dyn_cast(fullShapeOp.value()); + if (!fullShapeReCast) { + return getAsIndexOpFoldResult(rewriter.getContext(), shapedType.getShape()); + } + + OpFoldResult curPtrOffset; + if (auto curReCast = ptr.getDefiningOp()) { + curPtrOffset = curReCast.getConstifiedMixedOffset(); + } else if (isa(ptr) && + isa(ptr.getParentBlock()->getParentOp())) { + // Here's to process loop state where ptr is just from loop interator. + // Following assertion corresponds to conversion result from `rewriteFor` + auto blockArg = dyn_cast(ptr); + auto forOp = dyn_cast(ptr.getParentBlock()->getParentOp()); + auto initReCastOfLoop = forOp.getTiedLoopInit(blockArg) + ->get() + .getDefiningOp(); + assert(initReCastOfLoop && initReCastOfLoop.getOffsets().size() == 1); + Value initReCastOffset = initReCastOfLoop.getOffsets()[0]; + + for (OpOperand &use : initReCastOffset.getUses()) { + if (use.getOwner() == initReCastOfLoop) + continue; + else if (isa(use.getOwner())) + continue; + else if (use.getOwner() == forOp) + curPtrOffset = OpFoldResult(forOp.getTiedLoopRegionIterArg(&use)); + else + llvm_unreachable("Illegal interation offset after rewriteFor"); + } + } else { + llvm_unreachable("Unsupported state when check tensor_ptr boundary"); + } + + assert(curPtrOffset); + + OpFoldResult offsetShift = subOpFoldResult( + curPtrOffset, fullShapeReCast.getConstifiedMixedOffset(), loc, rewriter); + + for (int i = 0; i < shapedType.getRank(); ++i) { + if (llvm::find(boundaryCheck, i) != boundaryCheck.end()) { + auto fullShape = fullShapeReCast.getConstifiedMixedSizes()[i]; + + OpFoldResult curOffset = divOpFoldResult( + offsetShift, fullShapeReCast.getConstifiedMixedStrides()[i], loc, + rewriter); + OpFoldResult curLeftSize = + maxOpFoldResult(subOpFoldResult(fullShape, curOffset, loc, rewriter), + rewriter.getIndexAttr(0), loc, rewriter); + + boundarySize[i] = + minOpFoldResult(boundarySize[i], curLeftSize, loc, rewriter); + + offsetShift = remOpFoldResult( + offsetShift, fullShapeReCast.getConstifiedMixedStrides()[i], loc, + rewriter); + } + } + + return boundarySize; +} + SmallVector getBroadcastDims(RankedTensorType src, RankedTensorType dst) { SmallVector broadcastDims; @@ -215,6 +433,212 @@ SmallVector getUnbroadcastDims(RankedTensorType src, return unbroadcastDims; } +} // namespace ConverterUtils + +namespace triton { + +mlir::Operation * +findFirstMatchingOperandDef(mlir::Operation *rootOp, + const std::function &condFn) { + LLVM_DEBUG(llvm::dbgs() << "[findFirstMatchingOperandDef] Current op: " + << *rootOp << "\n"); + mlir::Value lhs = nullptr; + mlir::Value rhs = nullptr; + if (auto op = dyn_cast(rootOp)) { + lhs = op.getPtr(); + rhs = op.getOffset(); + } else if (auto op = dyn_cast(rootOp)) { + lhs = op.getLhs(); + rhs = op.getRhs(); + } else if (auto op = dyn_cast(rootOp)) { + lhs = op.getLhs(); + rhs = op.getRhs(); + } else if (auto op = dyn_cast(rootOp)) { + lhs = op.getLhs(); + rhs = op.getRhs(); + } else if (auto op = dyn_cast(rootOp)) { + lhs = op.getLhs(); + rhs = op.getRhs(); + } else if (auto op = dyn_cast(rootOp)) { + lhs = op.getLhs(); + rhs = op.getRhs(); + } else if (auto op = dyn_cast(rootOp)) { + lhs = op.getSrc(); + } else if (auto op = dyn_cast(rootOp)) { + } else { + rootOp->emitRemark("Backtracing encounters unsupported Operation"); + return nullptr; + } + // Backtrace operands + if (!lhs) { + return nullptr; + } + auto lhsDef = lhs.getDefiningOp(); + mlir::Operation *targetOp; + if (lhsDef) { + if (condFn(lhsDef)) { + targetOp = lhsDef; + } else { + targetOp = findFirstMatchingOperandDef(lhsDef, condFn); + } + if (targetOp) { + return targetOp; + } + } + if (!rhs) { + return nullptr; + } + auto rhsDef = rhs.getDefiningOp(); + if (rhsDef) { + if (condFn(rhsDef)) { + targetOp = rhsDef; + } else { + targetOp = findFirstMatchingOperandDef(rhsDef, condFn); + } + if (targetOp) { + return targetOp; + } + } + return nullptr; +} + +void traverseBackwardUpdateOperandChainIf( + Operation *op, std::function conditionFn, + std::function stopFn, + std::function actionFn, OpBuilder &builder, + DenseSet &handledOperation) { + + if (!op || handledOperation.contains(op)) + return; + + handledOperation.insert(op); + + if (stopFn(op)) + return; + + if (conditionFn(op)) + actionFn(builder, op); + + DenseSet handledOperand; + + std::function handler = [&](Value operand) { + if (handledOperand.contains(operand)) + return; + handledOperand.insert(operand); + if (Operation *defOp = operand.getDefiningOp()) { + traverseBackwardUpdateOperandChainIf(defOp, conditionFn, stopFn, actionFn, + builder, handledOperation); + } else { + auto blockArgument = cast(operand); + auto parentOp = blockArgument.getOwner()->getParentOp(); + if (auto whileOp = dyn_cast(parentOp); + whileOp && whileOp.getAfterBody() == blockArgument.getOwner()) { + auto argNum = blockArgument.getArgNumber(); + auto conditionArg = whileOp.getConditionOp().getArgs()[argNum]; + handler(conditionArg); + } else if (auto loopOp = dyn_cast(parentOp)) { + OpOperand *initArgOperand = loopOp.getTiedLoopInit(blockArgument); + if (!initArgOperand) + return; + Value initArg = initArgOperand->get(); + handler(initArg); + Value yieldedValue = + loopOp.getTiedLoopYieldedValue(blockArgument)->get(); + if (yieldedValue != blockArgument) + handler(yieldedValue); + } + } + }; + + for (Value operand : op->getOperands()) { + handler(operand); + } + + if (auto loopOp = dyn_cast(op)) { + for (auto yieldedValue : loopOp.getYieldedValues()) + handler(yieldedValue); + } +} + +// Note: rootOp will also be processed. +void traverseBackwardUpdateOperandChainIf( + Operation *rootOp, std::function conditionFn, + std::function stopFn, + std::function actionFn) { + + OpBuilder builder(rootOp->getContext()); + DenseSet handledOperation; + + traverseBackwardUpdateOperandChainIf(rootOp, conditionFn, stopFn, actionFn, + builder, handledOperation); +} + +void traverseForwardUpdateUserChainIf( + Operation *op, std::function conditionFn, + std::function stopFn, + std::function actionFn, OpBuilder &builder, + llvm::SmallPtrSet &stopOps) { + + if (!op) { + return; + } + + if (stopFn(op)) { + stopOps.insert(op); + return; + } + + if (conditionFn(op)) { + actionFn(builder, op); + } + + for (auto res : op->getResults()) { + for (auto userOp : res.getUsers()) { + traverseForwardUpdateUserChainIf(userOp, conditionFn, stopFn, actionFn, + builder, stopOps); + } + } +} + +// Note: rootOp will also be processed. +void traverseForwardUpdateUserChainIf( + Operation *rootOp, std::function conditionFn, + std::function stopFn, + std::function actionFn, + llvm::SmallPtrSet &stopOps) { + + OpBuilder builder(rootOp->getContext()); + + traverseForwardUpdateUserChainIf(rootOp, conditionFn, stopFn, actionFn, + builder, stopOps); +} + +bool isMetaUse(Operation *op) { return op->hasAttr("MetaUse"); } + +bool isMixUse(Operation *op) { return op->hasAttr("MixUse"); } + +IndirectLoadInterfaceOpType getIndirectLoadInterfaceOpType(Operation *op) { + auto ty = IndirectLoadInterfaceOpType::Undefined; + if (isMetaUse(op)) { + if (isa(op)) { + ty = IndirectLoadInterfaceOpType::Load; + } else if (isa(op)) { + ty = IndirectLoadInterfaceOpType::Calc; + } + } + return ty; +} + +bool opIsIndirectLoad(Operation *op) { + auto opType = getIndirectLoadInterfaceOpType(op); + return opType == IndirectLoadInterfaceOpType::Load; +} + +bool opIsIndirectCalc(Operation *op) { + auto opType = getIndirectLoadInterfaceOpType(op); + return opType == IndirectLoadInterfaceOpType::Calc; +} + scf::ForOp createNestedLoops( OpBuilder &builder, Location loc, unsigned currentDim, unsigned totalDims, ValueRange LBs, ValueRange UBs, ValueRange steps, SmallVector &ivs, @@ -243,229 +667,553 @@ scf::ForOp createNestedLoops( return loop; } +ModuleOp getModuleOpFromOperation(Operation *op) { + Operation *parent = op; + while (parent != nullptr && !isa(parent)) { + parent = parent->getParentOp(); // 向上查找 + } + return cast(parent); // 如果没找到会抛出异常 +} + +bool isTensorPtrType(Type type) { + auto ptrType = dyn_cast(type); + if (!ptrType) + return false; + return isa(ptrType.getPointeeType()); +} + +} // namespace triton + +// TODO: imply these function below +OpFoldResult addOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b) { + auto lhsInt = getConstantOfAttr(lhs); + auto rhsInt = getConstantOfAttr(rhs); + + if (lhsInt && rhsInt) + return b.getIndexAttr(lhsInt.value() + rhsInt.value()); + + if (!lhsInt && rhsInt && rhsInt.value() == 0) + return lhs; + if (!rhsInt && lhsInt && lhsInt.value() == 0) + return rhs; + + auto lhsValue = dyn_cast(lhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + auto rhsValue = dyn_cast(rhs); + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + return b.create(loc, lhsValue, rhsValue).getResult(); +} + +OpFoldResult subOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b) { + auto lhsInt = getConstantOfAttr(lhs); + auto rhsInt = getConstantOfAttr(rhs); + + if (lhsInt && rhsInt) + return b.getIndexAttr(lhsInt.value() - rhsInt.value()); + + if (!lhsInt && rhsInt && rhsInt.value() == 0) + return lhs; + + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + return b.create(loc, lhsValue, rhsValue).getResult(); +} + +OpFoldResult mulOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b) { + auto lhsInt = getConstantOfAttr(lhs); + auto rhsInt = getConstantOfAttr(rhs); + + if (lhsInt && rhsInt) + return b.getIndexAttr(lhsInt.value() * rhsInt.value()); + + if (lhsInt) { + if (lhsInt.value() == 0) + return lhs; + if (lhsInt.value() == 1) + return rhs; + } + if (rhsInt) { + if (rhsInt.value() == 0) + return rhs; + if (rhsInt.value() == 1) + return lhs; + } + + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + return b.create(loc, lhsValue, rhsValue).getResult(); +} + +OpFoldResult divOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b) { + auto lhsInt = getConstantOfAttr(lhs); + auto rhsInt = getConstantOfAttr(rhs); + + if (rhsInt && rhsInt.value() == 0) { + emitError(loc) << "cannot div 0!"; + return OpFoldResult(); + } + + if (lhsInt && rhsInt) + return b.getIndexAttr(lhsInt.value() / rhsInt.value()); + + if (lhsInt) { + if (lhsInt.value() == 0) + return lhs; + } + + if (rhsInt) { + if (rhsInt.value() == 1) + return lhs; + } + + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + return b.create(loc, lhsValue, rhsValue).getResult(); +} + +OpFoldResult remOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b) { + auto lhsInt = getConstantOfAttr(lhs); + auto rhsInt = getConstantOfAttr(rhs); + + if (rhsInt && rhsInt.value() == 0) { + emitError(loc) << "cannot remainder by 0!"; + return OpFoldResult(); + } + + if (lhsInt && rhsInt) + return b.getIndexAttr(lhsInt.value() % rhsInt.value()); + + if (lhsInt) { + if (lhsInt.value() == 0) + return lhs; + } + + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + return b.create(loc, lhsValue, rhsValue).getResult(); +} + +OpFoldResult minOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b) { + auto lhsInt = getConstantOfAttr(lhs); + auto rhsInt = getConstantOfAttr(rhs); + if (lhsInt && rhsInt) + return b.getIndexAttr(std::min(lhsInt.value(), rhsInt.value())); + + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + return b.create(loc, lhsValue, rhsValue).getResult(); +} + +OpFoldResult maxOpFoldResult(const OpFoldResult &lhs, const OpFoldResult &rhs, + const Location &loc, OpBuilder &b) { + auto lhsInt = getConstantOfAttr(lhs); + auto rhsInt = getConstantOfAttr(rhs); + if (lhsInt && rhsInt) + return b.getIndexAttr(std::max(lhsInt.value(), rhsInt.value())); + + auto lhsValue = dyn_cast(lhs), rhsValue = dyn_cast(rhs); + if (lhsInt) { + lhsValue = createConstIndexValueOp(loc, b, lhsInt.value()); + } else { + lhsValue = convertToIndexIfNeeded(lhsValue, loc, b); + assert(isa(lhsValue.getType())); + } + + if (rhsInt) { + rhsValue = createConstIndexValueOp(loc, b, rhsInt.value()); + } else { + rhsValue = convertToIndexIfNeeded(rhsValue, loc, b); + assert(isa(rhsValue.getType())); + } + + return b.create(loc, lhsValue, rhsValue).getResult(); +} + +void addReduceWithIndexAttr(ReduceWithIndexParams params, + ConversionPatternRewriter &rewriter, + linalg::ReduceOp reduceOp) { + const StringRef reduceRef = "reduce_mode"; + const StringRef tieBreakLeftRef = "tie_break_left"; + const StringRef unsignedSrcRef = "unsigned_src"; + + const StringRef tieBreakStr = + params.tieBreakType == TieBreakType::LEFT ? "true" : "false"; + const StringRef withIndexStr = + params.withIndexType == ReduceWithIndexType::MAX ? "max_with_index" + : "min_with_index"; + const StringRef unsignedSrcStr = params.isUnsignedSrc ? "true" : "false"; + + reduceOp->setAttr(reduceRef, rewriter.getStringAttr(withIndexStr)); + reduceOp->setAttr(tieBreakLeftRef, rewriter.getStringAttr(tieBreakStr)); + reduceOp->setAttr(unsignedSrcRef, rewriter.getStringAttr(unsignedSrcStr)); +} + +llvm::FailureOr +getReduceWithIndexParams(triton::ReduceOp op) { + auto tritonReduceBlock = op.getBody(); + auto *tritonYield = tritonReduceBlock->getTerminator(); + auto yieldValues = tritonYield->getOperands(); + constexpr int yieldValuesNum = 2; + if (yieldValues.empty()) { + return llvm::failure(); + } + if (yieldValues.size() != yieldValuesNum) { + return ReduceWithIndexParams{}; + } + + // Unify signed/unsigned and int/float predicate + enum class Predicate { Undefined = 0, lt = 1, gt = 2, eq = 3 }; + enum class Signedness { NotApplicable = 0, Signed = 1, Unsigned = 2 }; + auto unifyPredicateI = + [](arith::CmpIPredicate p) -> std::pair { + switch (p) { + case arith::CmpIPredicate::slt: + return {Predicate::lt, Signedness::Signed}; + case arith::CmpIPredicate::ult: + return {Predicate::lt, Signedness::Unsigned}; + case arith::CmpIPredicate::sgt: + return {Predicate::gt, Signedness::Signed}; + case arith::CmpIPredicate::ugt: + return {Predicate::gt, Signedness::Unsigned}; + case arith::CmpIPredicate::eq: + return {Predicate::eq, Signedness::NotApplicable}; + default: + return {Predicate::Undefined, Signedness::NotApplicable}; + } + }; + auto unifyPredicateF = + [](arith::CmpFPredicate p) -> std::pair { + switch (p) { + case arith::CmpFPredicate::OLT: + return {Predicate::lt, Signedness::Signed}; + case arith::CmpFPredicate::ULT: + return {Predicate::lt, Signedness::Unsigned}; + case arith::CmpFPredicate::OGT: + return {Predicate::gt, Signedness::Signed}; + case arith::CmpFPredicate::UGT: + return {Predicate::gt, Signedness::Unsigned}; + case arith::CmpFPredicate::OEQ: + return {Predicate::eq, Signedness::Signed}; + case arith::CmpFPredicate::UEQ: + return {Predicate::eq, Signedness::Unsigned}; + default: + return {Predicate::Undefined, Signedness::NotApplicable}; + } + }; + + // Composite predicate to pick index of min (or max) element have to be + // written in following form: (v means value and i means index) + // For leftmost element: + // (new_v == old_v and new_i < old_i) or new_v < old_v + // new_v < old_v or (new_v == old_v and new_i < old_i) + // new_v < old_v // python3.11 ttir + // (new_v == old_v and new_i < old_i) or new_v > old_v + // new_v > old_v or (new_v == old_v and new_i < old_i) + // new_v > old_v // python3.11 ttir + // For rightmost element: + // (new_v == old_v and new_i > old_i) or new_v < old_v + // new_v < old_v or (new_v == old_v and new_i > old_i) + // (new_v == old_v and new_i > old_i) or new_v > old_v + // new_v > old_v or (new_v == old_v and new_i > old_i) + + std::map, std::pair> + m{ + // leftmost + {{Predicate::eq, Predicate::lt, Predicate::lt}, + {ReduceWithIndexType::MIN, TieBreakType::LEFT}}, + {{Predicate::lt, Predicate::eq, Predicate::lt}, + {ReduceWithIndexType::MIN, TieBreakType::LEFT}}, + {{Predicate::lt}, {ReduceWithIndexType::MIN, TieBreakType::LEFT}}, + {{Predicate::eq, Predicate::lt, Predicate::gt}, + {ReduceWithIndexType::MAX, TieBreakType::LEFT}}, + {{Predicate::gt, Predicate::eq, Predicate::lt}, + {ReduceWithIndexType::MAX, TieBreakType::LEFT}}, + {{Predicate::gt}, {ReduceWithIndexType::MAX, TieBreakType::LEFT}}, + // rightmost + {{Predicate::eq, Predicate::gt, Predicate::lt}, + {ReduceWithIndexType::MIN, TieBreakType::RIGHT}}, + {{Predicate::lt, Predicate::eq, Predicate::gt}, + {ReduceWithIndexType::MIN, TieBreakType::RIGHT}}, + {{Predicate::eq, Predicate::gt, Predicate::gt}, + {ReduceWithIndexType::MAX, TieBreakType::RIGHT}}, + {{Predicate::gt, Predicate::eq, Predicate::gt}, + {ReduceWithIndexType::MAX, TieBreakType::RIGHT}}, + }; + + std::vector preds; + std::vector signednesses; + // A better way is to trace the arith.select + // Checking the operations one by one is hacky :( + for (auto &op : tritonReduceBlock->without_terminator()) { + Predicate pred = Predicate::Undefined; + Signedness signedness = Signedness::NotApplicable; + if (auto cmpiOp = dyn_cast(op)) { + auto predi = cmpiOp.getPredicate(); + std::tie(pred, signedness) = unifyPredicateI(predi); + } + if (auto cmpfOp = dyn_cast(op)) { + auto predf = cmpfOp.getPredicate(); + std::tie(pred, signedness) = unifyPredicateF(predf); + } + if (pred != Predicate::Undefined) { + preds.push_back(pred); + signednesses.push_back(signedness); + } + } + + // check if sequence of predicates matches any sequence for min/max + // leftmost/rightmost + if (m.find(preds) == m.end()) { + return llvm::failure(); + } + + assert(!signednesses.empty()); + const bool isUnsignedSrc = + signednesses[0] == Signedness::Unsigned || + signednesses[signednesses.size() - 1] == Signedness::Unsigned; + return ReduceWithIndexParams{.withIndexType = m.at(preds).first, + .tieBreakType = m.at(preds).second, + .isUnsignedSrc = isUnsignedSrc}; +} + +// Fold layout constant info to attr, otherwise convert to index type value +OpFoldResult getOpFoldResultOfLayoutInfo(Value value, OpBuilder &builder) { + OpFoldResult constantFold = getAsOpFoldResult(value); + if (llvm::isa(constantFold)) { + assert(isa(cast(constantFold))); + return constantFold; + } + + if (!isa(value.getType())) + llvm_unreachable("Illegal data type when parse block data layout info"); + + if (!isa(value.getType())) { + if (value.getType().isInteger(/*width*/ 1)) + value = builder.create( + value.getLoc(), builder.getIndexType(), value); + else + value = builder.create(value.getLoc(), + builder.getIndexType(), value); + } + + return value; +} + +// Specialize the Typeless Value (Zero, Min, Max) into a mlir TypedAttr FailureOr specializeTypelessValueToAttr(TypelessValue value, Type type, OpBuilder &b) { - // Common float and integer MLIR types used as map keys. mlir::Type f16Ty = Float16Type::get(b.getContext()); - mlir::Type f32Ty = Float32Type::get(b.getContext()); mlir::Type bf16Ty = BFloat16Type::get(b.getContext()); - + mlir::Type f32Ty = Float32Type::get(b.getContext()); mlir::Type i8TySL = IntegerType::get( b.getContext(), 8, IntegerType::SignednessSemantics::Signless); mlir::Type i8TyS = IntegerType::get(b.getContext(), 8, IntegerType::SignednessSemantics::Signed); mlir::Type i8TyU = IntegerType::get( b.getContext(), 8, IntegerType::SignednessSemantics::Unsigned); - mlir::Type i16TySL = IntegerType::get( b.getContext(), 16, IntegerType::SignednessSemantics::Signless); mlir::Type i16TyS = IntegerType::get( b.getContext(), 16, IntegerType::SignednessSemantics::Signed); mlir::Type i16TyU = IntegerType::get( b.getContext(), 16, IntegerType::SignednessSemantics::Unsigned); - mlir::Type i32TySL = IntegerType::get( b.getContext(), 32, IntegerType::SignednessSemantics::Signless); mlir::Type i32TyS = IntegerType::get( b.getContext(), 32, IntegerType::SignednessSemantics::Signed); mlir::Type i32TyU = IntegerType::get( b.getContext(), 32, IntegerType::SignednessSemantics::Unsigned); - mlir::Type i64TySL = IntegerType::get( b.getContext(), 64, IntegerType::SignednessSemantics::Signless); mlir::Type i64TyS = IntegerType::get( b.getContext(), 64, IntegerType::SignednessSemantics::Signed); mlir::Type i64TyU = IntegerType::get( b.getContext(), 64, IntegerType::SignednessSemantics::Unsigned); - - // Create APFloat values for float semantics (half, single, bfloat). llvm::APFloat halfZero = llvm::APFloat::getZero(llvm::APFloat::IEEEhalf()); llvm::APFloat halfOne(llvm::APFloat::IEEEhalf(), 1); llvm::APFloat halfMax = llvm::APFloat::getInf(llvm::APFloat::IEEEhalf()); llvm::APFloat halfMin = - llvm::APFloat::getInf(llvm::APFloat::IEEEhalf(), /*Negative=*/true); - - llvm::APFloat floatZero = llvm::APFloat::getZero(llvm::APFloat::IEEEsingle()); - llvm::APFloat floatOne(llvm::APFloat::IEEEsingle(), 1); - llvm::APFloat floatMax = llvm::APFloat::getInf(llvm::APFloat::IEEEsingle()); - llvm::APFloat floatMin = - llvm::APFloat::getInf(llvm::APFloat::IEEEsingle(), /*Negative=*/true); - - // BF16 (bfloat16) semantics via APFloat. + llvm::APFloat::getInf(llvm::APFloat::IEEEhalf(), true); llvm::APFloat bfloatZero = llvm::APFloat::getZero(llvm::APFloat::BFloat()); llvm::APFloat bfloatOne(llvm::APFloat::BFloat(), 1); llvm::APFloat bfloatMax = llvm::APFloat::getInf(llvm::APFloat::BFloat()); llvm::APFloat bfloatMin = - llvm::APFloat::getInf(llvm::APFloat::BFloat(), /*Negative=*/true); - - // Helper to use the opaque pointer of a Type as a stable key. + llvm::APFloat::getInf(llvm::APFloat::BFloat(), true); + llvm::APFloat floatZero = llvm::APFloat::getZero(llvm::APFloat::IEEEsingle()); + llvm::APFloat floatOne(llvm::APFloat::IEEEsingle(), 1); + llvm::APFloat floatMax = llvm::APFloat::getInf(llvm::APFloat::IEEEsingle()); + llvm::APFloat floatMin = + llvm::APFloat::getInf(llvm::APFloat::IEEEsingle(), true); auto toPtr = [](mlir::Type ty) { return ty.getAsOpaquePointer(); }; - // Store initialization values. Use signed and unsigned integer variants to - // avoid narrowing/overflow problems. - using InitValVariant = - std::variant; - - std::map, InitValVariant> initMap = { - // Zero values (floats and integers). - {{TypelessValue::Zero, toPtr(f16Ty)}, halfZero}, - {{TypelessValue::Zero, toPtr(f32Ty)}, floatZero}, - {{TypelessValue::Zero, toPtr(bf16Ty)}, bfloatZero}, - - {{TypelessValue::Zero, toPtr(i8TySL)}, (int8_t)0}, - {{TypelessValue::Zero, toPtr(i8TyS)}, (int8_t)0}, - {{TypelessValue::Zero, toPtr(i8TyU)}, (uint8_t)0}, - - {{TypelessValue::Zero, toPtr(i16TySL)}, (int16_t)0}, - {{TypelessValue::Zero, toPtr(i16TyS)}, (int16_t)0}, - {{TypelessValue::Zero, toPtr(i16TyU)}, (uint16_t)0}, - - {{TypelessValue::Zero, toPtr(i32TySL)}, (int32_t)0}, - {{TypelessValue::Zero, toPtr(i32TyS)}, (int32_t)0}, - {{TypelessValue::Zero, toPtr(i32TyU)}, (uint32_t)0}, - - {{TypelessValue::Zero, toPtr(i64TySL)}, (int64_t)0}, - {{TypelessValue::Zero, toPtr(i64TyS)}, (int64_t)0}, - {{TypelessValue::Zero, toPtr(i64TyU)}, (uint64_t)0}, - - // Min values (floats and integers). - {{TypelessValue::Min, toPtr(f16Ty)}, halfMin}, - {{TypelessValue::Min, toPtr(f32Ty)}, floatMin}, - {{TypelessValue::Min, toPtr(bf16Ty)}, bfloatMin}, - - {{TypelessValue::Min, toPtr(i8TySL)}, std::numeric_limits::min()}, - {{TypelessValue::Min, toPtr(i8TyS)}, std::numeric_limits::min()}, - {{TypelessValue::Min, toPtr(i8TyU)}, std::numeric_limits::min()}, - - {{TypelessValue::Min, toPtr(i16TySL)}, - std::numeric_limits::min()}, - {{TypelessValue::Min, toPtr(i16TyS)}, - std::numeric_limits::min()}, - {{TypelessValue::Min, toPtr(i16TyU)}, - std::numeric_limits::min()}, - - {{TypelessValue::Min, toPtr(i32TySL)}, - std::numeric_limits::min()}, - {{TypelessValue::Min, toPtr(i32TyS)}, - std::numeric_limits::min()}, - {{TypelessValue::Min, toPtr(i32TyU)}, - std::numeric_limits::min()}, - - {{TypelessValue::Min, toPtr(i64TySL)}, - std::numeric_limits::min()}, - {{TypelessValue::Min, toPtr(i64TyS)}, - std::numeric_limits::min()}, - {{TypelessValue::Min, toPtr(i64TyU)}, - std::numeric_limits::min()}, // 0 - - // Max values (floats and integers). - {{TypelessValue::Max, toPtr(f16Ty)}, halfMax}, - {{TypelessValue::Max, toPtr(f32Ty)}, floatMax}, - {{TypelessValue::Max, toPtr(bf16Ty)}, bfloatMax}, - - {{TypelessValue::Max, toPtr(i8TySL)}, std::numeric_limits::max()}, - {{TypelessValue::Max, toPtr(i8TyS)}, std::numeric_limits::max()}, - {{TypelessValue::Max, toPtr(i8TyU)}, std::numeric_limits::max()}, - - {{TypelessValue::Max, toPtr(i16TySL)}, - std::numeric_limits::max()}, - {{TypelessValue::Max, toPtr(i16TyS)}, - std::numeric_limits::max()}, - {{TypelessValue::Max, toPtr(i16TyU)}, - std::numeric_limits::max()}, - - {{TypelessValue::Max, toPtr(i32TySL)}, - std::numeric_limits::max()}, - {{TypelessValue::Max, toPtr(i32TyS)}, - std::numeric_limits::max()}, - {{TypelessValue::Max, toPtr(i32TyU)}, - std::numeric_limits::max()}, - - {{TypelessValue::Max, toPtr(i64TySL)}, - std::numeric_limits::max()}, - {{TypelessValue::Max, toPtr(i64TyS)}, - std::numeric_limits::max()}, - {{TypelessValue::Max, toPtr(i64TyU)}, - std::numeric_limits::max()}, - }; + std::map, + std::variant> + initMap = { + {{TypelessValue::Zero, toPtr(f16Ty)}, halfZero}, + {{TypelessValue::Zero, toPtr(bf16Ty)}, bfloatZero}, + {{TypelessValue::Zero, toPtr(f32Ty)}, floatZero}, + {{TypelessValue::Zero, toPtr(i16TySL)}, (int16_t)0}, + {{TypelessValue::Zero, toPtr(i16TyS)}, (int16_t)0}, + {{TypelessValue::Zero, toPtr(i16TyU)}, (int16_t)0}, + {{TypelessValue::Zero, toPtr(i32TySL)}, 0}, + {{TypelessValue::Zero, toPtr(i32TyS)}, 0}, + {{TypelessValue::Zero, toPtr(i32TyU)}, 0}, + {{TypelessValue::Zero, toPtr(i64TySL)}, (int64_t)0}, + {{TypelessValue::Zero, toPtr(i64TyS)}, (int64_t)0}, + {{TypelessValue::Zero, toPtr(i64TyU)}, (int64_t)0}, + {{TypelessValue::Min, toPtr(f16Ty)}, halfMin}, + {{TypelessValue::Min, toPtr(bf16Ty)}, bfloatMin}, + {{TypelessValue::Min, toPtr(f32Ty)}, floatMin}, + {{TypelessValue::Min, toPtr(i16TySL)}, + std::numeric_limits::min()}, + {{TypelessValue::Min, toPtr(i16TyS)}, + std::numeric_limits::min()}, + {{TypelessValue::Min, toPtr(i16TyU)}, + std::numeric_limits::min()}, + {{TypelessValue::Min, toPtr(i32TySL)}, + std::numeric_limits::min()}, + {{TypelessValue::Min, toPtr(i32TyS)}, + std::numeric_limits::min()}, + {{TypelessValue::Min, toPtr(i32TyU)}, + std::numeric_limits::min()}, + {{TypelessValue::Min, toPtr(i64TySL)}, + std::numeric_limits::min()}, + {{TypelessValue::Min, toPtr(i64TyS)}, + std::numeric_limits::min()}, + {{TypelessValue::Min, toPtr(i64TyU)}, + std::numeric_limits::min()}, + {{TypelessValue::Max, toPtr(f16Ty)}, halfMax}, + {{TypelessValue::Max, toPtr(bf16Ty)}, bfloatMax}, + {{TypelessValue::Max, toPtr(f32Ty)}, floatMax}, + {{TypelessValue::Max, toPtr(i16TySL)}, + std::numeric_limits::max()}, + {{TypelessValue::Max, toPtr(i16TyS)}, + std::numeric_limits::max()}, + {{TypelessValue::Max, toPtr(i16TyU)}, + std::numeric_limits::max()}, + {{TypelessValue::Max, toPtr(i32TySL)}, + std::numeric_limits::max()}, + {{TypelessValue::Max, toPtr(i32TyS)}, + std::numeric_limits::max()}, + {{TypelessValue::Max, toPtr(i32TyU)}, + std::numeric_limits::max()}, + {{TypelessValue::Max, toPtr(i64TySL)}, + std::numeric_limits::max()}, + {{TypelessValue::Max, toPtr(i64TyS)}, + std::numeric_limits::max()}, + {{TypelessValue::Max, toPtr(i64TyU)}, + std::numeric_limits::max()}, + }; - // Lookup key for the requested typeless value + concrete type. std::pair key = std::make_pair(value, toPtr(type)); - auto it = initMap.find(key); - if (it == initMap.end()) + if (initMap.find(key) == initMap.end()) return failure(); - - // Integer handling: prefer using the provided 'type' for IntegerAttr so - // signedness/width are preserved. - if (type.isInteger(8) || type.isInteger(16) || type.isInteger(32) || - type.isInteger(64)) { - unsigned bitWidth = type.getIntOrFloatBitWidth(); - - // Signed integers: extract signed variant and create IntegerAttr directly. - if (type.isSignedInteger(bitWidth)) { - switch (bitWidth) { - case 8: - return success(IntegerAttr::get(type, std::get(it->second))); - case 16: - return success(IntegerAttr::get(type, std::get(it->second))); - case 32: - return success(IntegerAttr::get(type, std::get(it->second))); - case 64: - return success(IntegerAttr::get(type, std::get(it->second))); - default: - return failure(); - } - } - - // Unsigned integers: extract unsigned variant. For 64-bit unsigned use - // APInt to avoid overflow of signed int64_t. - if (type.isUnsignedInteger(bitWidth)) { - switch (bitWidth) { - case 8: - return success(IntegerAttr::get( - type, static_cast(std::get(it->second)))); - case 16: - return success(IntegerAttr::get( - type, static_cast(std::get(it->second)))); - case 32: - return success(IntegerAttr::get( - type, static_cast(std::get(it->second)))); - case 64: { - uint64_t uval = std::get(it->second); - llvm::APInt apv(/*numBits=*/64, uval, /*isSigned=*/false); - return success(IntegerAttr::get(type, apv)); - } - default: - return failure(); - } - } - - // Signless integers: treat as signless using the signed variants (original - // code used signless integers everywhere for constants). - switch (bitWidth) { - case 8: - return success(IntegerAttr::get(type, std::get(it->second))); - case 16: - return success(IntegerAttr::get(type, std::get(it->second))); - case 32: - return success(IntegerAttr::get(type, std::get(it->second))); - case 64: - return success(IntegerAttr::get(type, std::get(it->second))); - default: - return failure(); - } - } - - // Floating-point handling (half, bf16, single). + if (type.isInteger(8)) + return success(IntegerAttr::get(IntegerType::get(b.getContext(), 8), + std::get(initMap.at(key)))); + if (type.isInteger(16)) + return success(IntegerAttr::get(IntegerType::get(b.getContext(), 16), + std::get(initMap.at(key)))); + if (type.isInteger(32)) + return success(IntegerAttr::get(IntegerType::get(b.getContext(), 32), + std::get(initMap.at(key)))); + if (type.isInteger(64)) + return success(IntegerAttr::get(IntegerType::get(b.getContext(), 64), + std::get(initMap.at(key)))); if (isa(type)) - return success(FloatAttr::get(f16Ty, std::get(it->second))); - if (isa(type)) - return success(FloatAttr::get(f32Ty, std::get(it->second))); + return success( + FloatAttr::get(f16Ty, std::get(initMap.at(key)))); if (isa(type)) - return success(FloatAttr::get(bf16Ty, std::get(it->second))); - + return success( + FloatAttr::get(bf16Ty, std::get(initMap.at(key)))); + if (isa(type)) + return success( + FloatAttr::get(f32Ty, std::get(initMap.at(key)))); return failure(); } @@ -504,4 +1252,98 @@ FailureOr specializeTypelessValueToConstant(TypelessValue value, return failure(); } +std::optional getIntAttr(const OpFoldResult ofr) { + Attribute attr; + if (auto val = dyn_cast(ofr)) { + if (!val.getDefiningOp()) + return std::nullopt; + attr = cast(val.getDefiningOp().getValue()); + } else { + attr = dyn_cast(ofr); + } + if (attr && isa(attr)) + return dyn_cast(attr).getInt(); + return std::nullopt; +} + +Value materializeValue(OpBuilder &builder, Location loc, OpFoldResult ofr) { + if (auto val = ofr.dyn_cast()) { + return val; + } + + auto intVal = getIntAttr(ofr); + if (intVal.has_value()) { + return builder.create( + loc, builder.getI32IntegerAttr(intVal.value())); + } + assert(intVal.has_value()); + return Value(); + + // return builder.create( + // loc, dyn_cast(attr).getInt()); +} + +bool isZero(const OpFoldResult ofr) { + auto staticOfr = getIntAttr(ofr); + return staticOfr.has_value() && staticOfr.value() == 0; +} + +bool isOne(const OpFoldResult ofr) { + auto staticOfr = getIntAttr(ofr); + return staticOfr.has_value() && staticOfr.value() == 1; +} + +Value convertToIndexIfNeeded(Value input, const Location &loc, OpBuilder &b) { + auto inputType = input.getType(); + if (auto intType = dyn_cast(inputType)) { + if (intType.isInteger(32) || intType.isInteger(64)) { + return b.create(loc, b.getIndexType(), input); + } + } + return input; +} + +RankedTensorType getExtractSlicedType(ArrayRef shape, + const llvm::SmallBitVector &droppedDims, + Type elemType) { + SmallVector targetShape; + for (auto [idx, dimOfr] : llvm::enumerate(shape)) { + if (!droppedDims[idx]) { + if (auto dim = getConstantIntValue(dimOfr)) { + targetShape.push_back(dim.value()); + } else { + targetShape.push_back(ShapedType::kDynamic); + } + } + } + return RankedTensorType::get(targetShape, elemType); +} + +bool checkStructureAnnotated(Operation *op, RewriterBase &rewriter) { + return llvm::any_of(op->getUsers(), [&rewriter](Operation *user) { + auto annotationOp = dyn_cast(user); + if (annotationOp && + annotationOp->hasAttr(ConverterUtils::continuousAttrName)) { + rewriter.eraseOp(annotationOp); + return true; + } + return false; + }); +} +} // namespace mlir + +namespace mlir::dicp { + +llvm::StringRef getBackend(ModuleOp module) { + if (!module) + return llvm::StringRef(); + + if (auto strAttr = module->getAttrOfType("dicp.backend")) + return strAttr.getValue(); + + return llvm::StringRef(); +} + +bool isDicpBackend(ModuleOp module) { return getBackend(module) == "dicp"; } + } // namespace mlir::dicp diff --git a/conda.sh b/conda.sh new file mode 100644 index 00000000..512a0108 --- /dev/null +++ b/conda.sh @@ -0,0 +1,79 @@ +#!/bin/bash +set -e + +# Conda 安装目录（可自定义） +CONDA_DIR="${CONDA_DIR:-/opt/conda}" + +echo "==> 1. 下载并安装 Miniconda" +if [ -d "$CONDA_DIR" ]; then + echo " 已存在 $CONDA_DIR，跳过安装" +else + ARCH=$(uname -m) + MINI_ARCH="$ARCH" + [ "$ARCH" = "aarch64" ] && MINI_ARCH="aarch64" + [ "$ARCH" = "x86_64" ] && MINI_ARCH="x86_64" + + wget -q "https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/Miniconda3-latest-Linux-${MINI_ARCH}.sh" -O /tmp/miniconda.sh + bash /tmp/miniconda.sh -b -p "$CONDA_DIR" + rm -f /tmp/miniconda.sh +fi + +export PATH="$CONDA_DIR/bin:$PATH" + +echo "==> 2. 配置 Conda 使用清华源" +# 系统级 .condarc 中的 defaults 会导致访问 repo.anaconda.com 并要求接受 ToS， +# 直接覆盖系统级与用户级 .condarc，只保留清华镜像 +cat > "$CONDA_DIR/.condarc" <<'EOF' +channels: + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge +show_channel_urls: true +EOF +cat > "$HOME/.condarc" <<'EOF' +channels: + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge +show_channel_urls: true +EOF + +echo "==> 3. 创建 dlcompiler 环境 (Python 3.10)" +if conda env list | awk '{print $1}' | grep -qx "dlcompiler"; then + echo " 已存在 dlcompiler 环境，跳过创建" +else + conda create -n dlcompiler python=3.10 -y +fi + +# 激活环境 +source "$CONDA_DIR/etc/profile.d/conda.sh" +conda activate dlcompiler + +echo "==> 4. 配置 pip 使用清华源" +pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple +pip config set global.trusted-host pypi.tuna.tsinghua.edu.cn + +echo "==> 5. 安装 Triton 构建依赖" +pip install --no-cache-dir \ + autopep8 isort numpy pytest pytest-forked pytest-xdist \ + "scipy>=1.7.1" llnl-hatchet expecttest \ + "setuptools>=40.8.0" wheel "cmake>=3.20,<4.0" \ + "ninja>=1.11.1" "pybind11>=2.13.1" lit nanobind + +echo "==> 6. 安装 PyTorch (CPU 版) 与 torch_npu (通过 pip 指定版本，不下载 wheel)" +# 安装 PyTorch CPU 版本，使用官方索引 +pip install torch==2.9.0+cpu --index-url https://download.pytorch.org/whl/cpu + +# 安装 torch_npu，默认从 PyPI 获取（若需 Ascend 专用源可加 --extra-index-url） +pip install torch-npu==2.9.0.post2 + +echo "==> 7. 安装额外 DLCompiler 依赖 (requirements.txt)" +if [ -f requirements.txt ]; then + pip install --no-cache-dir -r requirements.txt +else + echo "警告: 未找到 requirements.txt 文件，跳过该步骤" +fi + +echo "==> 环境安装完成！请执行以下命令激活 dlcompiler 环境：" +echo " source $CONDA_DIR/etc/profile.d/conda.sh" +echo " conda activate dlcompiler" \ No newline at end of file diff --git a/dicp_triton.cc b/dicp_triton.cc deleted file mode 100644 index e7a286c7..00000000 --- a/dicp_triton.cc +++ /dev/null @@ -1,5 +0,0 @@ -#include - -namespace py = pybind11; - -void init_triton_dicp_triton(py::module &&m) {} diff --git a/docker/Dockerfile b/docker/Dockerfile index a4199669..4e175a59 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,152 +1,142 @@ -FROM ccr.ccs.tencentyun.com/library/ubuntu:22.04 - -# FROM ubuntu:22.04 +ARG TARGETPLATFORM=linux/arm64 +FROM --platform=${TARGETPLATFORM} swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:9.0.0-910b-ubuntu22.04-py3.11-devel SHELL ["/bin/bash", "-c"] ENV DEBIAN_FRONTEND=noninteractive -ENV TZ="Asia/shanghai" - -RUN apt update && \ - apt install --yes --no-install-recommends --no-install-suggests \ - bash \ - ca-certificates \ - curl \ - git \ - gnupg \ - make \ - sudo \ - unzip \ - vim \ - wget && \ - apt clean && \ - rm -rf /var/lib/apt/lists/* && \ - update-ca-certificates - -RUN echo "deb http://ppa.launchpad.net/deadsnakes/ppa/ubuntu jammy main" > /etc/apt/sources.list.d/deadsnakes-ppa.list && \ - apt-key adv --keyserver keyserver.ubuntu.com --recv-keys BA6932366A755776 && \ - apt-get update - -RUN apt update && \ - apt install --yes --no-install-recommends --no-install-suggests \ - clang \ - clang-format \ - cmake \ - lld \ - ninja-build \ - python3-dev \ - python3-venv \ - python3.11 \ - python3.11-dev \ - python3.11-venv \ - python3.11-distutils \ - python3.9 \ - python3.9-dev \ - python3.9-venv \ - python3.9-distutils \ - zlib1g-dev && \ - apt clean && \ - rm -rf /var/lib/apt/lists/* - -RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.9 && \ - curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10 && \ - curl -sS https://bootstrap.pypa.io/get-pip.py | python3.11 && \ - update-alternatives --install /usr/bin/python python /usr/bin/python3.9 1 \ - --slave /usr/bin/python3 python3 /usr/bin/python3.9 \ - --slave /usr/bin/pip pip /usr/local/bin/pip3.9 \ - --slave /usr/bin/pip3 pip3 /usr/local/bin/pip3.9 && \ - update-alternatives --install /usr/bin/python python /usr/bin/python3.10 2 \ - --slave /usr/bin/python3 python3 /usr/bin/python3.10 \ - --slave /usr/bin/pip pip /usr/local/bin/pip3.10 \ - --slave /usr/bin/pip3 pip3 /usr/local/bin/pip3.10 && \ - update-alternatives --install /usr/bin/python python /usr/bin/python3.11 3 \ - --slave /usr/bin/python3 python3 /usr/bin/python3.11 \ - --slave /usr/bin/pip pip /usr/local/bin/pip3.11 \ - --slave /usr/bin/pip3 pip3 /usr/local/bin/pip3.11 && \ - rm -f /usr/local/bin/pip /usr/local/bin/pip3 && \ - update-alternatives --set python /usr/bin/python3.10 - -USER root -WORKDIR /root/workspace - -RUN set -ex && \ - apt update && \ - apt install -y zlib1g-dev clang-15 lld-15 && \ - apt install -y ccache && \ - update-alternatives --install /usr/bin/clang clang /usr/bin/clang-15 20 && \ - update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-15 20 && \ - update-alternatives --install /usr/bin/lld lld /usr/bin/lld-15 20 - -# 添加Miniconda安装（Python 3.10版本） -ENV CONDA_DIR /opt/conda -RUN wget https://repo.anaconda.com/miniconda/Miniconda3-py310_23.1.0-1-Linux-aarch64.sh -O /tmp/miniconda.sh && \ +ENV TZ=Asia/Shanghai +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +# ============================================================ +# 1. 替换 Ubuntu 源为阿里源 +# ============================================================ +RUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \ + cat > /etc/apt/sources.list << 'EOF' +deb http://mirrors.aliyun.com/ubuntu/ jammy main restricted universe multiverse +deb http://mirrors.aliyun.com/ubuntu/ jammy-security main restricted universe multiverse +deb http://mirrors.aliyun.com/ubuntu/ jammy-updates main restricted universe multiverse +deb http://mirrors.aliyun.com/ubuntu/ jammy-backports main restricted universe multiverse +deb-src http://mirrors.aliyun.com/ubuntu/ jammy main restricted universe multiverse +deb-src http://mirrors.aliyun.com/ubuntu/ jammy-security main restricted universe multiverse +deb-src http://mirrors.aliyun.com/ubuntu/ jammy-updates main restricted universe multiverse +deb-src http://mirrors.aliyun.com/ubuntu/ jammy-backports main restricted universe multiverse +EOF +RUN apt-get update + +# ============================================================ +# 2. 安装基础依赖 +# ============================================================ +RUN apt-get install -y --no-install-recommends \ + wget curl git build-essential software-properties-common \ + apt-transport-https ca-certificates gnupg lsb-release cmake +# ============================================================ +# 5. 安装其他构建依赖 +# ============================================================ +RUN apt-get install -y --no-install-recommends \ + ninja-build python3-pip libssl-dev libffi-dev \ + bash make sudo unzip vim zlib1g-dev +RUN apt-get clean && rm -rf /var/lib/apt/lists/* + +# ============================================================ +# 3. 安装 LLVM/Clang (>=15)，固定为 15 版本 +# 对应原脚本步骤3，但只装15且使用阿里云源 +# ============================================================ +RUN wget -O - https://apt.llvm.org/llvm-snapshot.gpg.key | apt-key add - && \ + cat > /etc/apt/sources.list.d/llvm-aliyun.list << 'EOF' +deb http://mirrors.aliyun.com/llvm/apt/jammy/ llvm-toolchain-jammy-15 main +EOF +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + clang-15 clangd-15 lld-15 \ + clang-format-15 clang-tidy-15 \ + llvm-15-dev libclang-15-dev +RUN update-alternatives --install /usr/bin/clang clang /usr/bin/clang-15 100 && \ + update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-15 100 && \ + update-alternatives --install /usr/bin/clangd clangd /usr/bin/clangd-15 100 && \ + update-alternatives --install /usr/bin/lld lld /usr/bin/lld-15 100 && \ + update-alternatives --install /usr/bin/ld.lld ld.lld /usr/bin/ld.lld-15 100 && \ + update-alternatives --install /usr/bin/clang-format clang-format /usr/bin/clang-format-15 100 && \ + update-alternatives --install /usr/bin/clang-tidy clang-tidy /usr/bin/clang-tidy-15 100 + + + + +# ============================================================ +# 6. 安装 Miniconda +# ============================================================ +ENV CONDA_DIR=/opt/conda +RUN ARCH=$(uname -m) && \ + MINI_ARCH=$ARCH && \ + [ "$ARCH" = "aarch64" ] && MINI_ARCH="aarch64" && \ + [ "$ARCH" = "x86_64" ] && MINI_ARCH="x86_64" && \ + wget -q "https://mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/Miniconda3-latest-Linux-${MINI_ARCH}.sh" -O /tmp/miniconda.sh && \ bash /tmp/miniconda.sh -b -p $CONDA_DIR && \ rm -f /tmp/miniconda.sh && \ $CONDA_DIR/bin/conda clean -tipy && \ - ln -s $CONDA_DIR/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \ - echo ". $CONDA_DIR/etc/profile.d/conda.sh" >> /etc/bash.bashrc && \ - echo "conda activate base" >> /etc/bash.bashrc - -ENV PATH $CONDA_DIR/bin:$PATH - -# 初始化 conda -RUN conda init bash - -# 创建新的 conda 环境 -RUN set -ex \ - && conda create -n dlcompiler python=3.10 -ENV LD_LIBRARY_PATH=/usr/lib/aarch64-linux-gnu:/usr/lib:/lib:$CONDA_DIR/envs/dlcompiler/lib:$LD_LIBRARY_PATH - -RUN echo "source activate dlcompiler" > ~/.bashrc -RUN set -ex && \ - apt-get update && apt-get install -y openssh-server && \ - sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config && \ - service ssh start - -COPY cann cann -RUN set -ex && \ - cd /root/workspace/cann && \ - chmod +x Ascend-cann-toolkit_8.3.RC1_linux-aarch64.run && \ - chmod +x Ascend-cann-kernels-910b_8.3.RC1_linux-aarch64.run && \ - ./Ascend-cann-toolkit_8.3.RC1_linux-aarch64.run --quiet --install --install-for-all && \ - ./Ascend-cann-kernels-910b_8.3.RC1_linux-aarch64.run --quiet --install --install-for-all && \ - echo "alias ll='ls -alh'" >> ~/.bashrc && \ - source ~/.bashrc && \ - echo "source /usr/local/Ascend/ascend-toolkit/set_env.sh" >> ~/.bashrc - -RUN set -ex && \ - source ~/.bashrc && \ - conda activate dlcompiler && \ - pip install ninja cmake wheel pybind11 -i https://mirrors.huaweicloud.com/repository/pypi/simple && \ - pip install attrs==24.2.0 numpy==1.26.4 scipy==1.13.1 decorator==5.1.1 psutil==6.0.0 pytest==8.3.2 pytest-xdist==3.6.1 pyyaml nanobind==2.9.2 torch==2.6.0 torch-npu==2.6.0 -i https://mirrors.huaweicloud.com/repository/pypi/simple - -COPY requirements.txt . -RUN set -ex && \ - source ~/.bashrc && \ - conda activate dlcompiler && \ - pip install -r requirements.txt -i https://mirrors.huaweicloud.com/repository/pypi/simple - -ENV LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/driver:$LD_LIBRARY_PATH - -RUN mkdir -p /root/.triton -COPY nvidia /root/.triton/nvidia - -COPY code . - -RUN ln -sf /usr/lib/aarch64-linux-gnu/libstdc++.so.6 $CONDA_DIR/envs/dlcompiler/lib/libstdc++.so.6 && \ - ldconfig - + echo ". $CONDA_DIR/etc/profile.d/conda.sh" >> ~/.bashrc && \ + echo "conda activate dlcompiler" >> ~/.bashrc +ENV PATH="$CONDA_DIR/bin:$PATH" +# 直接覆盖 .condarc，避免保留默认的 defaults channel +# （新版 conda 访问 repo.anaconda.com 时会要求接受 ToS，导致 conda create 失败） +RUN cat > $CONDA_DIR/.condarc <<'EOF' && cp $CONDA_DIR/.condarc /root/.condarc +channels: + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free + - https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge +show_channel_urls: true +EOF + +# ============================================================ +# 7. 创建并激活 dlcompiler 环境（python=3.10），pip 改为清华源 +# ============================================================ +RUN conda create -n dlcompiler python=3.10 -y + +RUN source $CONDA_DIR/etc/profile.d/conda.sh && conda activate dlcompiler && \ + pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple && \ + pip config set global.trusted-host pypi.tuna.tsinghua.edu.cn + +# ============================================================ +# 8. 安装 Triton 构建依赖 +# ============================================================ +RUN source $CONDA_DIR/etc/profile.d/conda.sh && conda activate dlcompiler && \ + pip install --no-cache-dir \ + autopep8 isort numpy pytest pytest-forked pytest-xdist \ + "scipy>=1.7.1" llnl-hatchet expecttest \ + "setuptools>=40.8.0" wheel "cmake>=3.20,<4.0" \ + "ninja>=1.11.1" "pybind11>=2.13.1" lit nanobind + +# ============================================================ +# 9. 安装 PyTorch + torch_npu +# ============================================================ +RUN source $CONDA_DIR/etc/profile.d/conda.sh && conda activate dlcompiler && \ + ARCH=$(uname -m) && \ + if [ "$ARCH" = "aarch64" ]; then \ + wget -q -P /tmp https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp311-cp311-manylinux_2_28_aarch64.whl && \ + pip install /tmp/torch-2.9.0+cpu-cp311-cp311-manylinux_2_28_aarch64.whl && \ + rm -f /tmp/torch-2.9.0+cpu-cp311-cp311-manylinux_2_28_aarch64.whl && \ + wget -q -P /tmp https://gitcode.com/Ascend/pytorch/releases/download/v26.0.0-pytorch2.9.0/torch_npu-2.9.0.post2-cp311-cp311-manylinux_2_28_aarch64.whl && \ + pip install /tmp/torch_npu-2.9.0.post2-cp311-cp311-manylinux_2_28_aarch64.whl && \ + rm -f /tmp/torch_npu-2.9.0.post2-cp311-cp311-manylinux_2_28_aarch64.whl; \ + elif [ "$ARCH" = "x86_64" ]; then \ + wget -q -P /tmp https://download.pytorch.org/whl/cpu/torch-2.9.0%2Bcpu-cp311-cp311-manylinux_2_28_x86_64.whl && \ + pip install /tmp/torch-2.9.0+cpu-cp311-cp311-manylinux_2_28_x86_64.whl && \ + rm -f /tmp/torch-2.9.0+cpu-cp311-cp311-manylinux_2_28_x86_64.whl && \ + wget -q -P /tmp https://gitcode.com/Ascend/pytorch/releases/download/v26.0.0-pytorch2.9.0/torch_npu-2.9.0.post2-cp311-cp311-manylinux_2_28_x86_64.whl && \ + pip install /tmp/torch_npu-2.9.0.post2-cp311-cp311-manylinux_2_28_x86_64.whl && \ + rm -f /tmp/torch_npu-2.9.0.post2-cp311-cp311-manylinux_2_28_x86_64.whl; \ + else \ + echo "Unsupported architecture: $ARCH" && exit 1; \ + fi + +# ============================================================ +# 10. 设置工作目录和 CANN 环境 +# ============================================================ +WORKDIR /root/workspace RUN echo "alias ll='ls -alh --color=auto'" >> ~/.bashrc && \ echo "alias ls='ls --color=auto'" >> ~/.bashrc && \ - echo "export LS_COLORS='rs=0:di=01;34:ln=01;36:mh=00:pi=40;33:so=01;35:do=01;35:bd=40;33;01:cd=40;33;01:or=40;31;01:mi=00:su=37;41:sg=30;43:ca=30;41:tw=30;42:ow=34;42:st=37;44:ex=01;32'" >> ~/.bashrc + echo "source /usr/local/Ascend/ascend-toolkit/set_env.sh" >> ~/.bashrc && \ + echo "export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/driver:\$LD_LIBRARY_PATH" >> ~/.bashrc -RUN set -ex && \ - source ~/.bashrc && \ - conda activate dlcompiler && \ - cd /root/workspace && \ - export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/driver:$LD_LIBRARY_PATH && \ - ln -sf /usr/lib/aarch64-linux-gnu/libstdc++.so.6 $CONDA_DIR/envs/dlcompiler/lib/libstdc++.so.6 && \ - ldconfig && \ - echo "当前LD_LIBRARY_PATH: $LD_LIBRARY_PATH" && \ - echo "Python路径: $(which python)" && \ - echo "PIP路径: $(which pip)" && \ - echo y | bash compile.sh +# 额外 DLCompiler 依赖（原 Dockerfile 保留） +COPY requirements.txt . +RUN source $CONDA_DIR/etc/profile.d/conda.sh && conda activate dlcompiler && \ + pip install --no-cache-dir -r requirements.txt diff --git a/format.sh b/format.sh index 28ac6000..12060139 100644 --- a/format.sh +++ b/format.sh @@ -5,4 +5,71 @@ set -e script_dir=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) cd "$script_dir" -find tools/ compiler/ -regex '.*\.\(h\|cpp\|cc\|c\)' -print0 | xargs -0 clang-format -i + +# Determine the base ref: if on a branch with origin/main, diff against main; otherwise diff against HEAD~1 +if git rev-parse --verify origin/main &> /dev/null; then + base_ref="origin/main" +else + base_ref="HEAD~1" +fi + +# Collect changed files, excluding third_party and .git (per format.yaml ignore rules) +# Include: committed changes vs base, staged changes, and unstaged working tree changes +filter_exclude() { + grep -v '^third_party/' | grep -v '^\.git/' || true +} + +changed_cpp_files=$({ + git diff --name-only --diff-filter=ACMR "$base_ref"...HEAD + git diff --cached --name-only --diff-filter=ACMR + git diff --name-only --diff-filter=ACMR +} | sort -u | grep -E '\.(h|cpp|cc|c)$' | filter_exclude) +changed_md_files=$({ + git diff --name-only --diff-filter=ACMR "$base_ref"...HEAD + git diff --cached --name-only --diff-filter=ACMR + git diff --name-only --diff-filter=ACMR +} | sort -u | grep '\.md$' | filter_exclude) +changed_py_files=$({ + git diff --name-only --diff-filter=ACMR "$base_ref"...HEAD + git diff --cached --name-only --diff-filter=ACMR + git diff --name-only --diff-filter=ACMR +} | sort -u | grep '\.py$' | filter_exclude) + +# clang-format +if [ -n "$changed_cpp_files" ]; then + if command -v clang-format &> /dev/null; then + echo "clang-format on changed C/C++ files:" + echo "$changed_cpp_files" + echo "$changed_cpp_files" | xargs clang-format -i + else + echo "clang-format not found, skipping C/C++ formatting" + fi +else + echo "No C/C++ files modified, skipping clang-format" +fi + +# markdownlint +if [ -n "$changed_md_files" ]; then + if command -v markdownlint-cli2 &> /dev/null; then + echo "markdownlint on changed Markdown files:" + echo "$changed_md_files" + echo "$changed_md_files" | xargs markdownlint-cli2 --fix + else + echo "markdownlint-cli2 not found, skipping markdown linting" + fi +else + echo "No Markdown files modified, skipping markdownlint" +fi + +# python-black +if [ -n "$changed_py_files" ]; then + if command -v black &> /dev/null; then + echo "black on changed Python files:" + echo "$changed_py_files" + echo "$changed_py_files" | xargs black + else + echo "black not found, skipping python formatting" + fi +else + echo "No Python files modified, skipping black" +fi diff --git a/language/deeplink/__init__.py b/language/deeplink/__init__.py index b4f5ddca..1bd67bbf 100644 --- a/language/deeplink/__init__.py +++ b/language/deeplink/__init__.py @@ -1,46 +1,55 @@ -from triton.backends.dicp_triton.utils import init_dicp_driver -from . import libdevice +from triton.language import math + +from .cann import libdevice +from .cann import extension from .async_task import async_task -from .core import ( - insert_slice, - extract_slice, - sync_block_all, - set_cross_flag, - wait_cross_flag, - parallel, - inline_lambda, - alloc, - compile_hint, - ND, - NZ, - fragment, - UB, - L1, - L0A, - L0B, - L0C, - SyncFlag, -) -from .custom_op import ( - custom, - custom_semantic, - register_custom_op, - CORE, - PIPE, - MODE, -) +from .core import inline_lambda + +# --------------------------------------------------------------------------- +# Glue layer: delegate standard math functions to triton.language.math +# (aligned with triton-ascend cann/__init__.py) +# --------------------------------------------------------------------------- + +libdevice.umulhi = math.umulhi +libdevice.exp = math.exp +libdevice.exp2 = math.exp2 +libdevice.log = math.log +libdevice.log2 = math.log2 +libdevice.cos = math.cos +libdevice.sin = math.sin +libdevice.sqrt = math.sqrt +libdevice.sqrt_rn = math.sqrt_rn +libdevice.rsqrt = math.rsqrt +libdevice.div_rn = math.div_rn +libdevice.erf = math.erf +libdevice.floor = math.floor +libdevice.ceil = math.ceil +libdevice.fdiv = math.fdiv +libdevice.fma = math.fma +libdevice.abs = math.abs + +# Reverse override: libdevice's tanh supports bf16 via cast, replace math.tanh +math.tanh = libdevice.tanh __all__ = [ "libdevice", + "extension", + "async_task", + "inline_lambda", +] + +_EXTENSION_ATTRS = { "insert_slice", "extract_slice", "sync_block_all", + "sync_block_set", + "sync_block_wait", "set_cross_flag", "wait_cross_flag", "parallel", - "inline_lambda", "alloc", "compile_hint", + "multibuffer", "ND", "NZ", "fragment", @@ -50,13 +59,37 @@ "L0B", "L0C", "SyncFlag", - "async_task", "custom", "custom_semantic", "register_custom_op", "CORE", "PIPE", "MODE", -] + "scope", + "layout", + "builtin", + "is_builtin", + "get_element", + "sort", + "flip", + "gather", + "index_put", + "gather_out_to_ub", + "scatter_ub_to_out", + "index_select_simd", +} + + +def __getattr__(name): + if name in _EXTENSION_ATTRS: + return getattr(extension, name) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def ensure_driver_initialized(): + from triton.runtime.driver import driver + + if driver._active is None: + from triton.backends.dicp_triton.utils import init_dicp_driver -init_dicp_driver() + init_dicp_driver() diff --git a/language/deeplink/cann/__init__.py b/language/deeplink/cann/__init__.py new file mode 100644 index 00000000..31c390af --- /dev/null +++ b/language/deeplink/cann/__init__.py @@ -0,0 +1,24 @@ +# 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. + +from . import extension +from . import libdevice + +__all__ = ["extension", "libdevice"] diff --git a/language/deeplink/cann/buffer/__init__.py b/language/deeplink/cann/buffer/__init__.py new file mode 100644 index 00000000..22371aaa --- /dev/null +++ b/language/deeplink/cann/buffer/__init__.py @@ -0,0 +1,43 @@ +# 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. + +__all__ = [ + "builtin", + "is_builtin", + "address_space", + "buffer_type", + "buffer", + "alloc", + "to_buffer", + "to_tensor", + "subview", +] + +from .core import ( + builtin, + is_builtin, + address_space, + buffer_type, + buffer, + alloc, + to_buffer, + to_tensor, + subview, +) diff --git a/language/deeplink/cann/buffer/builder.py b/language/deeplink/cann/buffer/builder.py new file mode 100644 index 00000000..fd360a2a --- /dev/null +++ b/language/deeplink/cann/buffer/builder.py @@ -0,0 +1,82 @@ +# 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 + ) diff --git a/language/deeplink/cann/buffer/core.py b/language/deeplink/cann/buffer/core.py new file mode 100644 index 00000000..c06addae --- /dev/null +++ b/language/deeplink/cann/buffer/core.py @@ -0,0 +1,396 @@ +# 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__ = [ + "address_space", + "buffer_type", + "subview", + "alloc", + "buffer", + "to_buffer", + "to_tensor", +] + +import importlib +from typing import TypeVar, List +from functools import wraps + +from triton._C.libtriton import ir +import triton.language.core as tl + + +T = TypeVar("T") + +TRITON_BUILTIN = "__triton_builtin__" +BUFFER_BUILTIN = "__buffer_builtin__" + + +def builtin(fn: T) -> T: + """Mark a function as a buffer language builtin.""" + assert callable(fn) + + @wraps(fn) + def wrapper(*args, **kwargs): + if "_semantic" not in kwargs or kwargs["_semantic"] is None: + raise ValueError( + "Did you forget to add @triton.jit ? " + "(`_semantic` argument must be provided outside of JIT functions.)" + ) + return fn(*args, **kwargs) + + # also set triton_builtin to true so that CodeGenerator will recognize this function + setattr(wrapper, TRITON_BUILTIN, True) + setattr(wrapper, BUFFER_BUILTIN, True) + + return wrapper + + +def is_builtin(fn) -> bool: + """Is this a registered buffer language builtin function?""" + return getattr(fn, BUFFER_BUILTIN, False) + + +class address_space: + """Represents a buffer's address space. + + The :code:`address_space` of a buffer is a target-specific attribute. + """ + + def to_ir(self, builder: ir.builder) -> ir.type: + raise NotImplementedError("Abstract address_space cannot be converted to ir") + + +class buffer_type(tl.dtype): + def __init__( + self, + element_ty: tl.dtype, + shape: List, + space: address_space = None, + strides: List = None, + ): + self.element_ty = element_ty + self.shape = shape if isinstance(shape, list) else list(shape) + self.space = space + self.strides = strides if strides is not None else [] + self.name = self._make_name() + + def _make_name(self): + res = ( + "" + + def to_ir(self, builder: ir.builder) -> ir.type: + element_ty_ir = self.element_ty.to_ir(builder) + addr_space_attr = ( + self.space.to_ir(builder) if self.space else builder.get_null_attr() + ) + + # use the method with strides if strides is not empty + if self.strides: + return builder.get_buffer_ty_with_strides( + self.shape, element_ty_ir, self.strides, addr_space_attr + ) + else: + return builder.get_buffer_ty(self.shape, element_ty_ir, addr_space_attr) + + def __str__(self): + return self.name + + def __repr__(self): + return self.__str__() + + def __eq__(self, other) -> bool: + if not isinstance(other, buffer_type): + return False + return ( + self.element_ty == other.element_ty + and self.shape == other.shape + and self.space == other.space + and self.strides == other.strides + ) + + def __ne__(self, other) -> bool: + return not self.__eq__(other) + + @property + def scalar(self): + return self.element_ty + + def mangle(self) -> str: + elt = self.element_ty.mangle() + shape = "_".join(map(str, self.shape)) + return f"B{elt}S{shape}S" + + def _unflatten_ir(self, handles: List[ir.value], cursor: int): + return buffer(handles[cursor], self), cursor + 1 + + +# ----------------------- +# buffer +# ----------------------- + + +class buffer(tl.base_value): + """Represents a region of memory. + + :code:`buffer` is the fundamental data structure for Triton programs using + the buffer language extension. Most functions in + :py:mod:`deeplink.cann.buffer` operate on and return buffers. + + Most of the named member functions here are duplicates of the free functions + in :code:`triton.language`. For example, :code:`triton.language.sqrt(x)` is + equivalent to :code:`x.sqrt()`. + + .. rubric:: Constructors + .. + For some reason Sphinx includes __init__ before printing the full table + of methods. Not what I want, but I can't figure out how to fix it. Give + it its own section so it looks intentional. :) + """ + + def __init__(self, handle, buffer_ty: buffer_type): + """Not called by user code.""" + super().__init__() + self.handle = handle + self.type = buffer_ty + self.dtype = buffer_ty.element_ty.scalar + self.shape = buffer_ty.shape + self.space = buffer_ty.space + self.strides = buffer_ty.strides + + def _flatten_ir(self, handles: List[ir.value]) -> None: + handles.append(self.handle) + + def __str__(self) -> str: + # ex. "<16x32xfloat32, address_space>" + res = "<" + "x".join(str(s) for s in self.shape) + "x" + str(self.dtype) + if self.space: + res += ", " + str(self.space) + return res + ">" + + @builtin + def subview( + self, + offsets: List[tl.constexpr], + sizes: List[tl.constexpr], + strides: List[tl.constexpr], + _semantic=None, + ) -> "buffer": + return subview(self, offsets, sizes, strides, _semantic=_semantic) + + @builtin + def to_tensor(self, writable=True, target_shape=None, _semantic=None): + """Convert this buffer to a tl.tensor""" + return to_tensor( + self, writable=writable, target_shape=target_shape, _semantic=_semantic + ) + + +semantic = importlib.import_module(".semantic", package=__package__) + + +@builtin +def alloc( + etype: tl.dtype, + shape: List[tl.constexpr], + _address_space: address_space = None, + is_mem_unique: bool = False, + _semantic=None, +) -> buffer: + """ + Allocates a region of local memory with the specified shape and type. + + :param etype: the element type of the buffer. + :type etype: tl.dtype + :param shape: A list of non-negative integers representing the shape of the buffer. + :type shape: List[tl.constexpr] + :param _address_space: (Optional) backend-specific local memory address space + :type _address_space: bl.address_space + """ + return semantic.alloc( + etype, shape, _address_space, is_mem_unique, _semantic.builder + ) + + +@builtin +def to_buffer( + tensor: tl.tensor, + space: address_space = None, + bind_buffer: buffer = None, + _semantic=None, +) -> buffer: + """ + Convert a tensor to a buffer. + + :param tensor: the tensor to convert. + :type tensor: tl.tensor + :param space: the address space for the buffer (optional). + :type space: address_space + """ + return semantic.to_buffer(tensor, space, bind_buffer, _semantic.builder) + + +@builtin +def to_tensor( + memref: buffer, writable: bool = True, target_shape=None, _semantic=None +) -> tl.tensor: + """ + Create a tl.tensor from a bl.buffer. + + :param memref: the input bl.buffer object. + :memref type: bl.buffer + :param writable: If set true, the resultant tensor is considered "writable" during bufferization. + :type writable: bool + """ + return semantic.to_tensor( + memref, writable, _semantic.builder, target_shape=target_shape + ) + + +def check_subview(src, offsets, sizes, strides): + """ + Check data of subview methods which the data length and the offset value must be 32-byte aligned. + + The conditions for checking data are as follows: + 1. offset value must be 32-bytes aligned. + 2. all strides must be 1. + 3. the first point's offset in the second row of the last dimension must be 32-byte aligned. + + For instance, the following example fails to satisfy the specified criteria. + %subview = memref.subview %arg0[1, 1][4, 4][2, 2] + : memref<8x8xf32, strided<[8, 1], offset: 0>> to + memref<4x4xf32, strided<[16, 2], offset: 9>> + offsets = [8, 8] | sizes = [4, 4] | strides = [2, 2] + result_offset = 9 + second_row_start_offset = 25 + The scene will be go wrong because the follow conditions are not meet. + 1) result_offset is not 32-bytes aligned. + 2) strides = [2, 2], not all strides are equal to 1. + 3) second_row_start_offset are not 32-bytes aligned. + """ + bytes_per_block = 32 + bits_per_byte = 8 + base_byte = bytes_per_block // (src.dtype.primitive_bitwidth // bits_per_byte) + result_strides = [] + result_offset = 0 + second_row_start_offset = 0 + length = len(strides) + src_strides = [1] * length + if length == 1: + if offset[0] % base_byte != 0: + raise TypeError( + f"all strides should be 1 and the offset value should be 32-bytes aligned." + ) + return + for i in range(length - 2, -1, -1): + src_strides[i] = src_strides[i + 1] * src.shape[i + 1] + for i in range(0, length): + if isinstance(offsets[i], tl.tensor): + return + result_strides.append(src_strides[i] * strides[i]) + result_offset = result_offset + offsets[i] * src_strides[i] + second_row_start_offset = result_offset + src_strides[-2] * strides[-2] + is_unaligned = False + if sizes[1] > 1: + is_unaligned = second_row_start_offset % base_byte != 0 + stride_1 = all(s == 1 for s in strides) + is_unaligned = result_offset % base_byte != 0 or is_unaligned or not stride_1 + if is_unaligned: + raise TypeError( + f"all strides should be 1 and the offset value should be 32-bytes aligned." + ) + + +@builtin +def subview( + src: buffer, + offsets: List[tl.constexpr], + sizes: List[tl.constexpr], + strides: List[tl.constexpr], + _semantic=None, +) -> buffer: + """ + Creates a subview of the source buffer with the specified offsets, sizes, and strides. + + :param src: The source buffer to create a subview from. + :type src: buffer + :param offsets: A list of non-negative integers representing the offsets in each dimension. + :type offsets: List[tl.constexpr] + :param sizes: A list of non-negative integers representing the sizes in each dimension. + :type sizes: List[tl.constexpr] + :param strides: A list of non-negative integers representing the strides in each dimension. + :type strides: List[tl.constexpr] + :return: A new buffer representing the subview of the source buffer. + :rtype: buffer + """ + # Validate that sizes and strides contain only constexpr values + new_sizes = [] + for i, size in enumerate(sizes): + if isinstance(size, int): + # Convert regular integers to constexpr + new_sizes.append(tl.constexpr(size)) + elif isinstance(size, tl.constexpr): + new_sizes.append(size) + else: + raise TypeError(f"sizes[{i}] must be constexpr, got {type(size).__name__}") + + new_strides = [] + for i, stride in enumerate(strides): + if isinstance(stride, int): + # Convert regular integers to constexpr + new_strides.append(tl.constexpr(stride)) + elif isinstance(stride, tl.constexpr): + new_strides.append(stride) + else: + raise TypeError( + f"strides[{i}] must be constexpr, got {type(stride).__name__}" + ) + + check_offsets = [] + new_offsets = [] + for offset in offsets: + if isinstance(offset, tl.constexpr): + # Check that constexpr offset values cannot be negative + if offset < 0: + raise ValueError(f"Offset value must be non-negative, got {offset}") + new_offsets.append(_semantic.to_tensor(offset)) + check_offsets.append(offset) + elif isinstance(offset, int): + # Convert regular integers to constexpr and then to tensor + if offset < 0: + raise ValueError(f"Offset value must be non-negative, got {offset}") + new_offsets.append(_semantic.to_tensor(tl.constexpr(offset))) + check_offsets.append(tl.constexpr(offset)) + else: + # Assume it's already a tensor + new_offsets.append(offset) + check_offsets.append(offset) + + check_subview(src, check_offsets, new_sizes, new_strides) + return semantic.subview(src, new_offsets, new_sizes, new_strides, _semantic.builder) diff --git a/language/deeplink/cann/buffer/semantic.py b/language/deeplink/cann/buffer/semantic.py new file mode 100644 index 00000000..70e529ab --- /dev/null +++ b/language/deeplink/cann/buffer/semantic.py @@ -0,0 +1,158 @@ +# 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. + +from typing import TypeVar, List + +from triton._C.libtriton import ir +import triton.language.core as tl + +from . import core as bl + + +T = TypeVar("T") + + +def alloc( + etype: tl.dtype, + shape: List[tl.constexpr], + address_space: bl.address_space, + is_mem_unique, + builder: ir.builder, +) -> bl.buffer: + shape = tl._unwrap_shape(shape) + if etype == tl.int1: + raise TypeError("Unsupported alloc int1 type") + if not isinstance(shape, (tl.tuple, list)): + raise TypeError("shape must be list/tuple") + etype = tl._unwrap_if_constexpr(etype) + address_space = tl._unwrap_if_constexpr(address_space) + element_ty_ir = etype.to_ir(builder) + addr_space_attr = ( + address_space.to_ir(builder) if address_space else builder.get_null_attr() + ) + memref_ty = builder.get_buffer_ty(shape, element_ty_ir, addr_space_attr) + handle = builder.alloc(memref_ty) + if is_mem_unique: + builder.create_annotation_mark(handle, "mem_unique", builder.get_unit_attr()) + builder.create_annotation_mark( + handle, "effects", builder.get_str_array_attr(["write", "read"]) + ) + + buffer_ty = bl.buffer_type(element_ty=etype, shape=shape, space=address_space) + return bl.buffer(handle, buffer_ty) + + +def to_buffer( + tensor: tl.tensor, + address_space: bl.address_space, + bind_buffer: bl.buffer, + builder: ir.builder, +) -> bl.buffer: + if not isinstance(tensor.shape, (tl.tuple, list)) or not tensor.shape: + raise TypeError("scalar type cannot be converted to buffer") + if isinstance(bind_buffer, bl.buffer): + builder.create_bind_buffer(tensor.handle, bind_buffer.handle) + return bind_buffer + if not (bind_buffer is None): + raise ValueError("bind_buffer must be a buffer or None") + address_space = tl._unwrap_if_constexpr(address_space) + addr_space_attr = ( + address_space.to_ir(builder) if address_space else builder.get_null_attr() + ) + handle = builder.to_buffer(tensor.handle, addr_space_attr) + buffer_ty = bl.buffer_type( + element_ty=tensor.dtype, shape=tensor.shape, space=address_space + ) + return bl.buffer(handle, buffer_ty) + + +def to_tensor( + memref: bl.buffer, writable: bool, builder: ir.builder, target_shape=None +) -> tl.tensor: + if not isinstance(memref, bl.buffer): + raise TypeError("memref must be bl.buffer") + + need_convert_layout = False + shape = memref.shape + if target_shape: + need_convert_layout = True + shape = tl._unwrap_shape(target_shape) + assert shape != memref.shape, "target shape is the same as source shape" + if not isinstance(shape, (tl.tuple, list)): + raise TypeError("shape must be list/tuple") + tensor_type = tl.block_type(memref.dtype, shape) + + memref_value = memref.handle + if need_convert_layout: + buffer_ty = bl.buffer_type( + element_ty=memref.dtype, + shape=shape, + space=memref.space, + ) + memref_value = builder.create_convert_layout( + memref_value, buffer_ty.to_ir(builder) + ) + + return tl.tensor(builder.to_tensor(memref_value, writable), tensor_type) + + +def subview( + src: bl.buffer, + offsets: List[tl.tensor], + sizes: List[tl.constexpr], + strides: List[tl.constexpr], + builder: ir.builder, +) -> bl.buffer: + + new_offsets = [offset.handle for offset in offsets] + sizes_int = tl._unwrap_shape(sizes) + strides_int = tl._unwrap_shape(strides) + + result_handle = builder.subview(src.handle, new_offsets, sizes_int, strides_int) + + # calculate the memory layout strides of the source buffer + if src.strides: + # use the strides of the source buffer + src_memory_strides = src.strides + else: + # calculate the default row-major strides + src_memory_strides = [] + stride = 1 + for dim_size in reversed(src.shape): + if dim_size < 0: + raise ValueError( + "Cannot compute strides for buffer with dynamic dimensions" + ) + src_memory_strides.insert(0, stride) + stride *= dim_size + + result_memory_strides = [] + for src_stride, subview_stride in zip(src_memory_strides, strides_int): + result_memory_strides.append(src_stride * subview_stride) + + # create buffer_type with strides + buffer_ty = bl.buffer_type( + element_ty=src.dtype, + shape=sizes_int, + space=src.space, + strides=result_memory_strides, + ) + return bl.buffer(result_handle, buffer_ty) diff --git a/language/deeplink/cann/extension/__init__.py b/language/deeplink/cann/extension/__init__.py new file mode 100644 index 00000000..2624fd01 --- /dev/null +++ b/language/deeplink/cann/extension/__init__.py @@ -0,0 +1,195 @@ +# 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. + +from triton._C.libtriton import ir, dicp_triton + +# MLIR affine bindings (same objects as triton._C.libtriton.ascend.ir). +affine_expr = dicp_triton.ir.affine_expr +affine_constant_expr = dicp_triton.ir.affine_constant_expr +affine_dim_expr = dicp_triton.ir.affine_dim_expr +affine_symbol_expr = dicp_triton.ir.affine_symbol_expr +affine_binary_op_expr = dicp_triton.ir.affine_binary_op_expr +affine_map = dicp_triton.ir.affine_map + +AffineExpr = affine_expr +AffineConstantExpr = affine_constant_expr +AffineDimExpr = affine_dim_expr +AffineSymbolExpr = affine_symbol_expr +AffineBinaryOpExpr = affine_binary_op_expr +AffineMap = affine_map + +from .core import ( + ascend_address_space, + builtin, + CORE, + copy_from_ub_to_l1, + copy, + debug_barrier, + fixpipe, + FixpipeDMAMode, + FixpipeDualDstMode, + FixpipePreQuantMode, + FixpipePreReluMode, + int64, + is_builtin, + MODE, + PIPE, + IteratorType, + sub_vec_id, + sub_vec_num, + sync_block_all, + sync_block_set, + sync_block_wait, + alloc, + SyncFlag, + set_cross_flag, + wait_cross_flag, + SYNC_IN_VF, +) + +from .scope import scope + +from .layout import ( + layout, + ND, + NZ, + fragment, + UB, + L1, + L0A, + L0B, + L0C, +) + +from .custom_op import ( + custom, + custom_semantic, + register_custom_op, +) + +from .math_ops import ( + atan2, + isfinited, + finitef, +) + +from .aux_ops import ( + parallel, + compile_hint, + multibuffer, +) + +from .vec_ops import ( + insert_slice, + extract_slice, + get_element, + sort, + flip, +) + +from .mem_ops import ( + index_put, + gather_out_to_ub, + scatter_ub_to_out, + index_select_simd, +) + +# gather is a standard triton op; re-export for backward compat. +from triton.language.core import gather + +__all__ = [ + # core + "builtin", + "copy_from_ub_to_l1", + "copy", + "CORE", + "debug_barrier", + "fixpipe", + "FixpipeDMAMode", + "FixpipeDualDstMode", + "FixpipePreQuantMode", + "FixpipePreReluMode", + "int64", + "is_builtin", + "MODE", + "PIPE", + "IteratorType", + "sub_vec_id", + "sub_vec_num", + "sync_block_all", + "sync_block_set", + "sync_block_wait", + "alloc", + "SyncFlag", + "set_cross_flag", + "wait_cross_flag", + "SYNC_IN_VF", + # address space + "ascend_address_space", + # MLIR affine + "affine_expr", + "affine_constant_expr", + "affine_dim_expr", + "affine_symbol_expr", + "affine_binary_op_expr", + "affine_map", + "AffineExpr", + "AffineConstantExpr", + "AffineDimExpr", + "AffineSymbolExpr", + "AffineBinaryOpExpr", + "AffineMap", + # scope + "scope", + "layout", + "ND", + "NZ", + "fragment", + "UB", + "L1", + "L0A", + "L0B", + "L0C", + # custom op + "custom", + "custom_semantic", + "register_custom_op", + # math ops + "atan2", + "isfinited", + "finitef", + # aux ops + "parallel", + "compile_hint", + "multibuffer", + # vec ops + "insert_slice", + "extract_slice", + "get_element", + "sort", + "flip", + # mem ops + "index_put", + "gather_out_to_ub", + "scatter_ub_to_out", + "index_select_simd", + # standard + "gather", +] diff --git a/language/deeplink/cann/extension/aux_ops.py b/language/deeplink/cann/extension/aux_ops.py new file mode 100644 index 00000000..b5cfdb74 --- /dev/null +++ b/language/deeplink/cann/extension/aux_ops.py @@ -0,0 +1,72 @@ +# 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. + +from triton.language.core import _unwrap_if_constexpr, range, tensor +import triton.language.core as tl_core + +from .core import builtin +from . import semantic + +__all__ = ["parallel", "compile_hint", "multibuffer"] + + +class parallel(range): + """ + Iterator that counts upward with parallel execution semantics. + """ + + def __init__( + self, + arg1, + arg2=None, + step=None, + num_stages=None, + loop_unroll_factor=None, + bind_sub_block: bool = False, + ): + super().__init__(arg1, arg2, step, num_stages, loop_unroll_factor) + self.bind_sub_block = bind_sub_block + + +@builtin +def compile_hint(ptr, hint_name, hint_val=None, _semantic=None): + def _unwrap(val): + return _unwrap_if_constexpr(val) if val else val + + hint_name = _unwrap_if_constexpr(hint_name) + assert isinstance(hint_name, str), f"hint name: {hint_name} is not string" + if isinstance(hint_val, (list, tl_core.tuple)): + hint_val = [_unwrap(val) for val in hint_val] + else: + hint_val = _unwrap(hint_val) + hint_val = _unwrap_if_constexpr(hint_val) if hint_val else hint_val + semantic.compile_hint(ptr, hint_name, hint_val, _semantic.builder) + + +@builtin +def multibuffer(src: tensor, size, _semantic=None): + """ + Set multi_buffer for an existing tensor. + """ + buffer_size = _unwrap_if_constexpr(size) + assert ( + isinstance(buffer_size, int) and buffer_size == 2 + ), "only support bufferize equals 2" + semantic.compile_hint(src, "multi_buffer", buffer_size, _semantic.builder) diff --git a/language/deeplink/cann/extension/builder.py b/language/deeplink/cann/extension/builder.py new file mode 100644 index 00000000..32a74c79 --- /dev/null +++ b/language/deeplink/cann/extension/builder.py @@ -0,0 +1,73 @@ +""" +DICP NPU builder utilities for unified builder interface. +""" + +__all__ = [ + "create_builder_method_wrapper", + "attach_builder_methods", + "setup_unified_builder", +] + + +def create_builder_method_wrapper(main_builder, delegate_builder, method_name): + 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(main_builder, delegate_builder, method_names): + for method_name in method_names: + wrapper = create_builder_method_wrapper( + main_builder, delegate_builder, method_name + ) + setattr(main_builder, method_name, wrapper) + + +def setup_unified_builder(main_builder, dicp_builder): + main_builder._dicp_builder = dicp_builder + dicp_methods = [ + "create_scope_op", + "scope_return", + "get_t_core_type_attr_name", + "get_t_core_type_cube_attr", + "get_t_core_type_vector_attr", + "get_target_attribute", + "create_get_sub_vec_id", + "create_copy_buffer", + "create_copy_tensor", + "create_fixpipe", + "create_bind_buffer", + "create_debug_barrier", + "is_910_95", + "sync_block_set", + "sync_block_wait", + "create_convert_layout", + "sync_block_all", + "get_int_attr", + "get_str_array_attr", + "get_i64_array_attr", + "get_core_type_attr", + "get_pipe_attr", + "get_vf_mode_attr", + "get_iterator_types_attr", + "parse_attr", + "get_affine_map_attr", + "get_affine_map_array_attr", + "get_buffer_ty_with_affine_map", + "create_custom_op", + ] + attach_builder_methods(main_builder, dicp_builder, dicp_methods) diff --git a/language/deeplink/cann/extension/code_generator.py b/language/deeplink/cann/extension/code_generator.py new file mode 100644 index 00000000..5b4c6af6 --- /dev/null +++ b/language/deeplink/cann/extension/code_generator.py @@ -0,0 +1,162 @@ +""" +Ascend-specific code generation handlers for 'with' statement context managers. +""" + +__all__ = ["handle_scope_with", "mangle_ty"] +import ast + + +def mangle_ty(ty): + from triton import language + from ...buffer.core import buffer_type + + if isinstance(ty, buffer_type): + elt = mangle_ty(ty.element_ty) + shape = "_".join(map(str, ty.shape)) + return f"B{elt}S{shape}S" + + if ty.is_ptr(): + return "P" + mangle_ty(ty.element_ty) + if ty.is_int(): + SIGNED = language.dtype.SIGNEDNESS.SIGNED + prefix = "i" if ty.int_signedness == SIGNED else "u" + return prefix + str(ty.int_bitwidth) + if ty.is_floating(): + return str(ty) + if ty.is_block(): + elt = mangle_ty(ty.scalar) + shape = "_".join(map(str, ty.shape)) + return f"{elt}S{shape}S" + if ty.is_void(): + return "V" + raise TypeError(f"Unsupported type {ty}") + + +def _extract_scope_attributes(context_expr): + scope_attrs = {} + for keyword in context_expr.keywords: + if isinstance(keyword.value, ast.Constant): + scope_attrs[keyword.arg] = keyword.value.value + return scope_attrs + + +def _py_value_to_mlir_attr(builder, value): + attr_creators = { + str: lambda v: builder.get_string_attr(v), + bool: lambda v: builder.get_bool_attr(v), + int: lambda v: builder.get_int32_attr(v), + list: lambda v: builder.get_i64_array_attr(v), + } + creator = attr_creators.get(type(value)) + return creator(value) if creator else value + + +def _handle_core_mode_attr(builder, core_mode): + if core_mode not in ("cube", "vector"): + return {} + return { + builder.get_t_core_type_attr_name(): ( + builder.get_t_core_type_cube_attr() + if core_mode == "cube" + else builder.get_t_core_type_vector_attr() + ) + } + + +def _build_mlir_attrs_from_scope_attrs(builder, scope_attrs): + mlir_attrs = {"noinline": builder.get_unit_attr()} + for k, v in scope_attrs.items(): + if k == "core_mode": + mlir_attrs.update(_handle_core_mode_attr(builder, v)) + elif k == "noinline": + if not v: + mlir_attrs.pop("noinline") + elif k == "disable_auto_sync": + if v: + mlir_attrs["hivm.disable_auto_sync"] = _py_value_to_mlir_attr( + builder, v + ) + else: + mlir_attrs[k] = _py_value_to_mlir_attr(builder, v) + return mlir_attrs + + +def _verify_loop_carried_variable( + _is_triton_value, _is_triton_tensor, name, loop_val, live_val +): + assert _is_triton_value(loop_val), f"cannot reassign constexpr {name} in the loop" + assert _is_triton_value(live_val), f"cannot reassign constexpr {name} in the loop" + assert type(loop_val) == type( + live_val + ), f"Loop carried variable {name} changed type" + assert not _is_triton_tensor(loop_val) or loop_val.type == live_val.type, ( + f"Loop-carried variable {name} has initial type {live_val.type} " + f"but is re-assigned to {loop_val.type} in loop! " + f"Please make sure that the type stays consistent." + ) + + +def _reconstruct_value_from_ir(language, entry_block_arg, ret_type): + return language.core.tensor(entry_block_arg, ret_type) + + +def handle_scope_with(generator, node): + from triton import language + from triton.compiler.code_generator import ( + enter_sub_region, + _is_triton_value, + _is_triton_tensor, + ) + + context_expr = node.items[0].context_expr + scope_attrs = _extract_scope_attributes(context_expr) + + with enter_sub_region(generator) as sr: + liveins, _ = sr + ip, last_loc = generator._get_insertion_point_and_loc() + + dummy = generator.builder.create_block() + generator.builder.set_insertion_point_to_start(dummy) + generator.visit_compound_statement(node.body) + scope_defs = generator.local_defs + dummy.erase() + + names = [] + ret_types = [] + for name in scope_defs: + scope_val = scope_defs[name] + ret_types.append(scope_val.type) + names.append(name) + if name in liveins: + live_val = liveins[name] + _verify_loop_carried_variable( + _is_triton_value, _is_triton_tensor, name, scope_val, live_val + ) + + mlir_attrs = _build_mlir_attrs_from_scope_attrs(generator.builder, scope_attrs) + + generator._set_insertion_point_and_loc(ip, last_loc) + scope_op = generator.builder.create_scope_op( + mlir_attrs, [ty.to_ir(generator.builder) for ty in ret_types] + ) + + entry_block = generator.builder.create_block_with_parent( + scope_op.get_region(0), [] + ) + generator.builder.set_insertion_point_to_start(entry_block) + + generator.lscope = liveins.copy() + generator.visit_compound_statement(node.body) + generator.builder.set_insertion_point_to_end(entry_block) + reconstructed_values = [] + + for i in range(len(names)): + reconstructed_values.append(generator.lscope[names[i]].handle) + generator.builder.scope_return(reconstructed_values) + + for i, name in enumerate(names): + generator.set_value( + name, + _reconstruct_value_from_ir(language, scope_op.get_result(i), ret_types[i]), + ) + return None diff --git a/language/deeplink/cann/extension/core.py b/language/deeplink/cann/extension/core.py new file mode 100644 index 00000000..76b006be --- /dev/null +++ b/language/deeplink/cann/extension/core.py @@ -0,0 +1,405 @@ +# 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. + +__all__ = [ + "ascend_address_space", + "builtin", + "CORE", + "copy_from_ub_to_l1", + "copy", + "debug_barrier", + "fixpipe", + "FixpipeDMAMode", + "FixpipeDualDstMode", + "FixpipePreQuantMode", + "FixpipePreReluMode", + "int64", + "is_builtin", + "MODE", + "PIPE", + "IteratorType", + "sub_vec_id", + "sub_vec_num", + "sync_block_all", + "sync_block_set", + "sync_block_wait", + "alloc", + "SyncFlag", + "set_cross_flag", + "wait_cross_flag", + "SYNC_IN_VF", +] + +import enum +from typing import TypeVar, List, Union +from functools import wraps + +from triton._C.libtriton import ir, dicp_triton +import triton.language.core as tl +from triton.language.core import _shape_check_impl, _unwrap_if_constexpr + +from triton.backends.dicp_triton.npu_driver import NPUUtils + +from . import semantic as semantic + + +T = TypeVar("T") + +TRITON_BUILTIN = "__triton_builtin__" +ASCEND_BUILTIN = "__ascend_builtin__" + + +def _constexpr_to_value(v): + if isinstance(v, tl.constexpr): + return v.value + return v + + +def builtin(fn: T) -> T: + """Mark a function as a buffer language builtin.""" + assert callable(fn) + + @wraps(fn) + def wrapper(*args, **kwargs): + if "_semantic" not in kwargs or kwargs["_semantic"] is None: + raise ValueError( + "Did you forget to add @triton.jit ? " + "(`_semantic` argument must be provided outside of JIT functions.)" + ) + return fn(*args, **kwargs) + + setattr(wrapper, TRITON_BUILTIN, True) + setattr(wrapper, ASCEND_BUILTIN, True) + + return wrapper + + +def is_builtin(fn) -> bool: + """Is this a registered ascend language builtin function?""" + return getattr(fn, ASCEND_BUILTIN, False) + + +class int64(int): + def __new__(cls, value): + obj = int.__new__(cls, value) + obj.type = tl.int64 + return obj + + +class CORE(enum.Enum): + VECTOR = dicp_triton.ir.CoreType.VECTOR + CUBE = dicp_triton.ir.CoreType.CUBE + CUBE_OR_VECTOR = dicp_triton.ir.CoreType.CUBE_OR_VECTOR + CUBE_AND_VECTOR = dicp_triton.ir.CoreType.CUBE_AND_VECTOR + + +class PIPE(enum.Enum): + PIPE_S = dicp_triton.ir.PIPE.PIPE_S + PIPE_V = dicp_triton.ir.PIPE.PIPE_V + PIPE_M = dicp_triton.ir.PIPE.PIPE_M + PIPE_MTE1 = dicp_triton.ir.PIPE.PIPE_MTE1 + PIPE_MTE2 = dicp_triton.ir.PIPE.PIPE_MTE2 + PIPE_MTE3 = dicp_triton.ir.PIPE.PIPE_MTE3 + PIPE_ALL = dicp_triton.ir.PIPE.PIPE_ALL + PIPE_FIX = dicp_triton.ir.PIPE.PIPE_FIX + + +class MODE(enum.Enum): + SIMD = dicp_triton.ir.MODE.SIMD + SIMT = dicp_triton.ir.MODE.SIMT + MIX = dicp_triton.ir.MODE.MIX + + +class IteratorType(enum.Enum): + Parallel = dicp_triton.ir.IteratorType.Parallel + Broadcast = dicp_triton.ir.IteratorType.Broadcast + Transpose = dicp_triton.ir.IteratorType.Transpose + Reduction = dicp_triton.ir.IteratorType.Reduction + Interleave = dicp_triton.ir.IteratorType.Interleave + Deinterleave = dicp_triton.ir.IteratorType.Deinterleave + Inverse = dicp_triton.ir.IteratorType.Inverse + Pad = dicp_triton.ir.IteratorType.Pad + Concat = dicp_triton.ir.IteratorType.Concat + Gather = dicp_triton.ir.IteratorType.Gather + Cumulative = dicp_triton.ir.IteratorType.Cumulative + Opaque = dicp_triton.ir.IteratorType.Opaque + + +class FixpipeDMAMode(enum.Enum): + NZ2DN = dicp_triton.ir.FixpipeDMAMode.NZ2DN + NZ2ND = dicp_triton.ir.FixpipeDMAMode.NZ2ND + NZ2NZ = dicp_triton.ir.FixpipeDMAMode.NZ2NZ + + +class FixpipeDualDstMode(enum.Enum): + NO_DUAL = dicp_triton.ir.FixpipeDualDstMode.NO_DUAL + COLUMN_SPLIT = dicp_triton.ir.FixpipeDualDstMode.COLUMN_SPLIT + ROW_SPLIT = dicp_triton.ir.FixpipeDualDstMode.ROW_SPLIT + + +class FixpipePreQuantMode(enum.Enum): + NO_QUANT = dicp_triton.ir.FixpipePreQuantMode.NO_QUANT + F322BF16 = dicp_triton.ir.FixpipePreQuantMode.F322BF16 + F322F16 = dicp_triton.ir.FixpipePreQuantMode.F322F16 + S322I8 = dicp_triton.ir.FixpipePreQuantMode.S322I8 + + +class FixpipePreReluMode(enum.Enum): + LEAKY_RELU = dicp_triton.ir.FixpipePreReluMode.LEAKY_RELU + NO_RELU = dicp_triton.ir.FixpipePreReluMode.NO_RELU + NORMAL_RELU = dicp_triton.ir.FixpipePreReluMode.NORMAL_RELU + P_RELU = dicp_triton.ir.FixpipePreReluMode.P_RELU + + +class ascend_address_space_base: + def __init__(self, address_space_value): + self.real_address_space = address_space_value + + def to_ir(self, builder: ir.builder) -> ir.attribute: + return builder.get_target_attribute(self.real_address_space) + + +class ascend_address_space: + def __init__(self): + for k, v in { + k: v + for k, v in dicp_triton.ir.AddressSpace.__dict__.items() + if isinstance(v, dicp_triton.ir.AddressSpace) + }.items(): + setattr(self, k, ascend_address_space_base(v)) + + +ascend_address_space = ascend_address_space() + + +@builtin +def sub_vec_id(_semantic=None) -> tl.tensor: + return semantic.sub_vec_id(_semantic) + + +@builtin +def copy_from_ub_to_l1(src, dst, _semantic=None): + from warnings import warn + + warn("copy_from_ub_to_l1 is deprecated, please use copy instead.") + return semantic.copy_from_ub_to_l1(src, dst, _semantic) + + +@builtin +def copy(src, dst, _semantic=None): + return semantic.copy(src, dst, _semantic) + + +def create_sync_block( + sender, + receiver, + event_id, + is_set: bool, + sender_pipe=None, + receiver_pipe=None, + _semantic=None, +): + sender = _unwrap_if_constexpr(sender) + receiver = _unwrap_if_constexpr(receiver) + assert isinstance(sender, str) and sender in ( + "cube", + "vector", + ), f"ERROR: sender = {sender}" + assert isinstance(receiver, str) and receiver in ( + "cube", + "vector", + ), f"ERROR: receiver = {receiver}" + if isinstance(event_id, int): + assert 0 <= event_id < 16, f"event_id: {event_id} should be 0 ~ 15" + if sender == receiver: + raise ValueError(f"Unexpected pair: {sender} -> {receiver}") + if sender_pipe is None and receiver_pipe is None: + if sender == "cube": + sender_pipe = PIPE.PIPE_FIX + receiver_pipe = PIPE.PIPE_MTE2 + if sender == "vector": + sender_pipe = PIPE.PIPE_MTE3 + receiver_pipe = PIPE.PIPE_MTE2 + if not isinstance(sender_pipe, PIPE) or not isinstance(receiver_pipe, PIPE): + raise TypeError("sender_pipe and receiver_pipe must be instances of PIPE enum") + if is_set: + return semantic.create_sync_block_set( + sender, receiver, event_id, sender_pipe, receiver_pipe, _semantic + ) + return semantic.create_sync_block_wait( + sender, receiver, event_id, sender_pipe, receiver_pipe, _semantic + ) + + +@builtin +def sync_block_set( + sender, receiver, event_id, sender_pipe=None, receiver_pipe=None, _semantic=None +): + return create_sync_block( + sender, receiver, event_id, True, sender_pipe, receiver_pipe, _semantic + ) + + +@builtin +def sync_block_wait( + sender, receiver, event_id, sender_pipe=None, receiver_pipe=None, _semantic=None +): + return create_sync_block( + sender, receiver, event_id, False, sender_pipe, receiver_pipe, _semantic + ) + + +@builtin +def sync_block_all(mode, event_id, _semantic=None): + mode = _unwrap_if_constexpr(mode) + event_id = _unwrap_if_constexpr(event_id) + assert isinstance(mode, str), f"mode: {mode} is not string" + assert ( + isinstance(event_id, int) and 0 <= event_id < 16 + ), f"event_id: {event_id} should be 0 ~ 15" + assert mode in ("all_cube", "all_vector", "all"), f"ERROR: mode = {mode}" + semantic.custom_sync_op( + _semantic.builder, "sync_block_all", mode=mode, event_id=event_id + ) + + +@builtin +def alloc(shape, value, dtype, layout=None, scope=None, _semantic=None): + """ + Returns a tensor filled with the scalar value for the given shape and dtype. + """ + shape = _shape_check_impl(shape) + value = _constexpr_to_value(value) + dtype = _constexpr_to_value(dtype) + layout = _constexpr_to_value(layout) + scope = _constexpr_to_value(scope) + return semantic.alloc(shape, value, dtype, layout, scope, _semantic.builder) + + +class SyncFlagType: + ASCEND = ["cube_to_vector", "vector_to_cube"] + + def __init__(self, name): + name = _unwrap_if_constexpr(name) + self.name = name + assert name in SyncFlagType.ASCEND, name + + def __str__(self): + return self.name + + def codegen_name(self): + return self.name + + def sender(self): + if self.name == "cube_to_vector": + return "cube" + if self.name == "vector_to_cube": + return "vector" + assert self.name in SyncFlagType.ASCEND + + @property + def cache_key_part(self) -> str: + return self.name + + def __repr__(self): + return f"triton.language.{self.codegen_name()}" + + +class SyncFlag: + C2V = SyncFlagType("cube_to_vector") + V2C = SyncFlagType("vector_to_cube") + + +def _get_cross_flag_pipes(sender): + if sender == "cube": + return "vector", PIPE.PIPE_FIX, PIPE.PIPE_MTE2 + if sender == "vector": + return "cube", PIPE.PIPE_MTE3, PIPE.PIPE_MTE2 + raise AssertionError(f"Unexpected sender: {sender}") + + +@builtin +def set_cross_flag(sync_flag_type: SyncFlagType, event_id: int, _semantic=None): + sender = _unwrap_if_constexpr(sync_flag_type.sender()) + event_id = _unwrap_if_constexpr(event_id) + assert ( + isinstance(event_id, int) and 0 <= event_id < 16 + ), f"event_id: {event_id} should be 0 ~ 15" + receiver, sender_pipe, receiver_pipe = _get_cross_flag_pipes(sender) + return sync_block_set( + sender, receiver, event_id, sender_pipe, receiver_pipe, _semantic=_semantic + ) + + +@builtin +def wait_cross_flag(sync_flag_type: SyncFlagType, event_id: int, _semantic=None): + sender = _unwrap_if_constexpr(sync_flag_type.sender()) + event_id = _unwrap_if_constexpr(event_id) + assert ( + isinstance(event_id, int) and 0 <= event_id < 16 + ), f"event_id: {event_id} should be 0 ~ 15" + receiver, sender_pipe, receiver_pipe = _get_cross_flag_pipes(sender) + return sync_block_wait( + sender, receiver, event_id, sender_pipe, receiver_pipe, _semantic=_semantic + ) + + +@builtin +def fixpipe( + src, + dst=None, + dma_mode: FixpipeDMAMode = FixpipeDMAMode.NZ2ND, + dual_dst_mode: FixpipeDualDstMode = FixpipeDualDstMode.NO_DUAL, + _semantic=None, +): + pre_quant_mode = FixpipePreQuantMode.NO_QUANT + pre_relu_mode = FixpipePreReluMode.NO_RELU + return semantic.fixpipe( + src, dst, dma_mode, dual_dst_mode, pre_quant_mode, pre_relu_mode, _semantic + ) + + +class SYNC_IN_VF(enum.Enum): + VV_ALL = enum.auto() + VST_VLD = enum.auto() + VLD_VST = enum.auto() + VST_VST = enum.auto() + VS_ALL = enum.auto() + VST_LD = enum.auto() + VLD_ST = enum.auto() + VST_ST = enum.auto() + SV_ALL = enum.auto() + ST_VLD = enum.auto() + LD_VST = enum.auto() + ST_VST = enum.auto() + + +@builtin +def debug_barrier(sync_mode: SYNC_IN_VF, _semantic=None): + return semantic.debug_barrier(sync_mode.name, _semantic) + + +@builtin +def sub_vec_num(_semantic=None) -> tl.constexpr: + npuUtils = NPUUtils() + cube_num = npuUtils.get_aivector_core_num() + vector_num = npuUtils.get_aicore_num() + const_val = cube_num // vector_num + return tl.constexpr(const_val) diff --git a/language/deeplink/custom_op.py b/language/deeplink/cann/extension/custom_op.py similarity index 50% rename from language/deeplink/custom_op.py rename to language/deeplink/cann/extension/custom_op.py index 53ea97b8..794488b8 100644 --- a/language/deeplink/custom_op.py +++ b/language/deeplink/cann/extension/custom_op.py @@ -1,55 +1,53 @@ +# Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved. +# 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__ = ["custom", "custom_semantic", "register_custom_op"] + import inspect import types import typing import itertools -import enum -from triton.language import core, semantic -import triton.language as tl +import triton.language.core as tl +from .core import CORE, PIPE, MODE -__all__ = ["custom", "custom_semantic", "register_custom_op", "CORE", "PIPE", "MODE"] +# Registry for custom op, mapping name to its configuration. _custom_op_registry = {} -class CORE(enum.Enum): - CUBE = "CUBE" - VECTOR = "VECTOR" - CUBE_OR_VECTOR = "CUBE_OR_VECTOR" - CUBE_AND_VECTOR = "CUBE_AND_VECTOR" - - -class PIPE(enum.Enum): - PIPE_S = "PIPE_S" - PIPE_V = "PIPE_V" - PIPE_M = "PIPE_M" - PIPE_MTE1 = "PIPE_MTE1" - PIPE_MTE2 = "PIPE_MTE2" - PIPE_MTE3 = "PIPE_MTE3" - PIPE_ALL = "PIPE_ALL" - PIPE_FIX = "PIPE_FIX" - - -class MODE(enum.Enum): - SIMD = "SIMD" - SIMT = "SIMT" - MIX = "MIX" - - def _get_op_class(name): # Try to get op class in _custom_op_registry. op_class = _custom_op_registry.get(name) if op_class is None: - # Allow bulitin custom ops used without registry. + # Allow builtin custom ops used without registry. assert name.startswith("__builtin_"), f"Custom Op '{name}' not registered." - # Return a dummy op class for builtin custom op. op_class = type( "_builtin_custom_op", (object,), { "name": name, - "core": core.CORE.VECTOR, - "pipe": core.PIPE.PIPE_V, - "mode": core.MODE.SIMT, + "core": CORE.VECTOR, + "pipe": PIPE.PIPE_V, + "mode": MODE.SIMT, "signature": inspect.signature(object), }, ) @@ -59,7 +57,7 @@ def _get_op_class(name): def _unwrap_constexpr(arg): if isinstance(arg, tl.constexpr): return arg.value - if isinstance(arg, tuple): + if isinstance(arg, (tuple, tl.tuple)): return tuple(_unwrap_constexpr(x) for x in arg) if isinstance(arg, list): return [_unwrap_constexpr(x) for x in arg] @@ -68,64 +66,62 @@ def _unwrap_constexpr(arg): return arg -def _to_value(value, builder, ty=None): +def _to_value(value, _semantic=None, ty=None): # Try to use 'type' attribute if ty not set. ty = getattr(value, "type", ty) if ty is None else ty if isinstance(value, tl.tensor): if not value.type.is_block() and isinstance(ty, tl.dtype) and value.type != ty: - # For a scalar variable, if its type is not the expected one - # that specified by type hint 'ty', insert a cast for it. - return tl.semantic.cast(value, ty, builder).handle + return _semantic.cast(value, ty).handle return value.handle if isinstance(value, bool): - return builder.get_int1(value) + return _semantic.builder.get_int1(value) if isinstance(value, int): if isinstance(ty, tl.dtype): if ty.is_int64(): - return builder.get_int64(value) + return _semantic.builder.get_int64(value) if ty.is_uint64(): - return builder.get_uint64(value) + return _semantic.builder.get_uint64(value) if ty.is_int32(): - return builder.get_int32(value) + return _semantic.builder.get_int32(value) if ty.is_uint32(): - return builder.get_uint32(value) + return _semantic.builder.get_uint32(value) if ty.is_int16(): - return builder.get_int16(value) + return _semantic.builder.get_int16(value) if ty.is_uint16(): - return builder.get_uint16(value) + return _semantic.builder.get_uint16(value) if ty.is_int8(): - return builder.get_int8(value) + return _semantic.builder.get_int8(value) if ty.is_uint8(): - return builder.get_uint8(value) + return _semantic.builder.get_uint8(value) # default int32 - return builder.get_int32(value) + return _semantic.builder.get_int32(value) if isinstance(value, float): if isinstance(ty, tl.dtype): if ty.is_fp64(): - return builder.get_fp64(value) + return _semantic.builder.get_fp64(value) if ty.is_fp32(): - return builder.get_fp32(value) + return _semantic.builder.get_fp32(value) if ty.is_fp16(): - return builder.get_fp16(value) + return _semantic.builder.get_fp16(value) if ty.is_bf16(): - return builder.get_bf16(value) + return _semantic.builder.get_bf16(value) # default float32 - return builder.get_fp32(value) + return _semantic.builder.get_fp32(value) if isinstance(value, tl.constexpr): - return _to_value(value.value, builder) + return _to_value(value.value, _semantic) raise TypeError(f"Unsupported argument type {value} : {type(value)}") -def _to_operands(args, builder): +def _to_operands(args, _semantic=None): operands = [] for value in args: if value is None: continue - if isinstance(value, (list, tuple)): + if isinstance(value, (list, tuple, tl.tuple)): for item in value: - operands.append(_to_value(item, builder)) + operands.append(_to_value(item, _semantic)) else: - operands.append(_to_value(value, builder)) + operands.append(_to_value(value, _semantic)) return operands @@ -135,40 +131,60 @@ def _get_element_type(ty): return ty -def _args_to_operands(op, builder, args, kwargs): +def _args_to_operands(op, _semantic, args, kwargs): if not op.signature.parameters: # Without parameters in signature, use the actual parameter order. - return _to_operands(itertools.chain(args, kwargs.values()), builder) + return _to_operands(itertools.chain(args, kwargs.values()), _semantic) # Convert arguments to operands according the signature. operands = [] bind = op.signature.bind(*args, **kwargs) for param in op.signature.parameters.values(): - value = bind.arguments.get(param.name, None) + value = bind.arguments.get(param.name) if value is None: continue ty = op.arg_type.get(param.name, param.annotation) - if isinstance(value, (list, tuple)): + if isinstance(value, (list, tuple, tl.tuple)): ty = _get_element_type(ty) for item in value: - operands.append(_to_value(item, builder, ty)) + operands.append(_to_value(item, _semantic, ty)) else: - operands.append(_to_value(value, builder, ty)) + operands.append(_to_value(value, _semantic, ty)) return operands -def _add_optional_attr(op, name, builder, attrs): - if hasattr(op, name): - attrs[name] = getattr(op, name) +def _bind_op_arguments(op, args, kwargs): + if not op.signature.parameters: + return None + return op.signature.bind(*args, **kwargs) -def _add_bitcode_attr(op, builder, attrs): - if not hasattr(op, "bitcode"): +def _make_align_dim_attrs(op, builder, arg_attrs): + name = "align_dim" + if not hasattr(op, name): return - from pathlib import Path - bitcode_path = _resolve_bitcode_path(getattr(op, "bitcode")) - attrs["bitcode"] = bitcode_path + align_arg_indices = {} + if hasattr(op, "signature"): + param_names = list(op.signature.parameters.keys()) + for arg_name in op.align_dim.keys(): + if arg_name in param_names: + align_arg_indices[arg_name] = param_names.index(arg_name) + + for arg, align_val in op.align_dim.items(): + if isinstance(arg, str) and arg in align_arg_indices: + arg_attrs[align_arg_indices[arg]] = {name: builder.get_int_attr(align_val)} + elif isinstance(arg, int): + arg_attrs[arg] = {name: builder.get_int_attr(align_val)} + else: + assert False, f"{name}'s keys should be string or int" + + +def _make_arg_attrs(op, builder): + num_args = len(op.signature.parameters) if hasattr(op, "signature") else 0 + arg_attrs = [{} for _ in range(num_args)] + _make_align_dim_attrs(op, builder, arg_attrs) + return arg_attrs def _get_bitcode_search_paths(): @@ -188,7 +204,7 @@ def _get_bitcode_search_paths(): # 2. Relative to this file: language/deeplink/bitcode/bc/ # Works both in source tree and after pip install. - local_bc = Path(__file__).parent / "bitcode" / "bc" + local_bc = Path(__file__).parent.parent.parent / "bitcode" / "bc" if local_bc.is_dir(): paths.append(str(local_bc)) @@ -249,28 +265,83 @@ def _resolve_bitcode_path(bitcode_ref): ) +def _add_optional_attr(op, name, builder, attrs): + if hasattr(op, name): + attrs[name] = builder.get_string_attr(getattr(op, name)) + + +def _add_bitcode_attr(op, builder, attrs): + name = "bitcode" + if not hasattr(op, name): + return + + from pathlib import Path + + bitcode_path = _resolve_bitcode_path(getattr(op, name)) + attrs[name] = builder.get_string_attr(bitcode_path) + + +def _add_optional_extra_buffer_attr(op, builder, attrs): + name = "extra_buffers" + if not hasattr(op, name): + return + + extra_buffers = getattr(op, name) + if isinstance(extra_buffers, tuple): + extra_buffers = [extra_buffers] + + extra_buffer_types, extra_buffer_sizes = zip(*extra_buffers) + if hasattr(builder, "get_type_array_attr"): + attrs[name + "_types"] = builder.get_type_array_attr( + [ty.to_ir(builder) for ty in extra_buffer_types] + ) + attrs[name + "_sizes"] = builder.get_i64_array_attr(list(extra_buffer_sizes)) + else: + type_strs = [str(ty.to_ir(builder)) for ty in extra_buffer_types] + attrs[name + "_types"] = builder.parse_attr("[" + ", ".join(type_strs) + "]") + size_strs = [str(s) for s in extra_buffer_sizes] + attrs[name + "_sizes"] = builder.parse_attr("[" + ", ".join(size_strs) + "]") + + +def _add_optional_indexing_map_attr(op, builder, attrs): + name = "indexing_map" + if not hasattr(op, name): + return + indexing_map = getattr(op, name) + attrs[name] = builder.get_affine_map_array_attr(indexing_map) + + +def _add_optional_iterator_types_attr(op, builder, attrs): + name = "iterator_types" + if not hasattr(op, name): + return + attrs[name] = builder.get_iterator_types_attr( + [iterator_type.value for iterator_type in getattr(op, name)] + ) + + def _make_attrs(op, builder): attrs = { - "hivm.tcore_type": f"#hivm.tcore_type<{op.core.value}>", - "hivm.pipe": f"#hivm.pipe<{op.pipe.value}>", - "hivm.vf_mode": f"#hivm.vf_mode<{op.mode.value}>", + "hivm.tcore_type": builder.get_core_type_attr(op.core.value), + "hivm.pipe": builder.get_pipe_attr(op.pipe.value), + "hivm.vf_mode": builder.get_vf_mode_attr(op.mode.value), } if not op.name.startswith("__builtin_"): - assert hasattr(op, "symbol"), f"Non builtin custom op, symbol is required." + assert hasattr(op, "symbol"), "Non builtin custom op, symbol is required." assert hasattr( op, "bitcode" - ), f"Non builtin custom op, bitcode path is required." + ), "Non builtin custom op, bitcode path is required." - # Add bit code path attribute, formalize to abosulte path. _add_bitcode_attr(op, builder, attrs) - + _add_optional_indexing_map_attr(op, builder, attrs) + _add_optional_iterator_types_attr(op, builder, attrs) + _add_optional_extra_buffer_attr(op, builder, attrs) _add_optional_attr(op, "symbol", builder, attrs) _add_optional_attr(op, "source", builder, attrs) _add_optional_attr(op, "compile", builder, attrs) - # Extra attributes can be added here, such as op.extra_attr="attr_a=xx" _add_optional_attr(op, "extra_attr", builder, attrs) - _add_optional_attr(op, "iterator_types", builder, attrs) + return attrs @@ -281,99 +352,85 @@ def _to_result(res, res_types): return None if n_res == 1: return tl.tensor(res[0], res_types[0]) - return tuple(tl.tensor(res[i], res_types[i]) for i in range(n_res)) + return tl.tuple(tl.tensor(res[i], res_types[i]) for i in range(n_res)) def _init_op(op_class, *args, **kwargs): op = op_class.__new__(op_class) - # Add arg_type dict to support dynamic argument type specifying. setattr(op, "arg_type", {}) if op_class.signature.parameters: - # Init with arguments validate. op_class.__init__(op, *args, **kwargs) return op def custom_semantic(name: str, *args, _semantic=None, **kwargs): - _builder = _semantic.builder name = _unwrap_constexpr(name) - assert name in _custom_op_registry, f"Custom op '{name}' not found." - # Get op class according the name. op_class = _get_op_class(name) - # Convert constexpr to value in arguments. args = _unwrap_constexpr(args) kwargs = _unwrap_constexpr(kwargs) - # Create op instance from op class with the arguments. op = _init_op(op_class, *args, **kwargs) - # Prepare inputs and outputs operands. out = kwargs.pop("out", []) - outs = out if isinstance(out, (list, tuple)) else [out] - outputs = _to_operands(outs, _builder) - inputs = _args_to_operands(op, _builder, args, kwargs) - # Setup attributes. - attrs = _make_attrs(op, _builder) - # Build IR for the custom op. - res = _builder.create_custom_op(name, attrs, inputs, outputs) - # Results with same types as outputs. + outs = out if isinstance(out, (list, tuple, tl.tuple)) else [out] + outputs = _to_operands(outs, _semantic) + inputs = _args_to_operands(op, _semantic, args, kwargs) + builder = _semantic.builder + attrs = _make_attrs(op, builder) + arg_attrs = _make_arg_attrs(op, builder) + res = builder.create_custom_op(name, attrs, inputs, outputs, arg_attrs) res_types = [out.type for out in outs] return _to_result(res, res_types) -@core.builtin +@tl.builtin def custom(name: str, *args, _semantic=None, **kwargs): """Invoke a custom operation with the given name and arguments.""" return custom_semantic(name, *args, _semantic=_semantic, **kwargs) def register_custom_op(op): - """Register a custom operation so that we can invoke it using al.custom().""" + """Register a custom operation so that we can invoke it using dl.custom().""" assert inspect.isclass(op), "@register_custom_op should decorate on a class." - # Use class name if name not set. if not hasattr(op, "name"): setattr(op, "name", op.__name__) - # The op name should not be used. assert ( op.name not in _custom_op_registry ), f"Custom op name '{op.name}' already used." - # Check required core, pipe, mode fields. + assert hasattr(op, "core"), "'core' field is required." assert hasattr(op, "pipe"), "'pipe' field is required." assert hasattr(op, "mode"), "'mode' field is required." assert isinstance(op.core, CORE), "Invalid 'core' field, CORE type is required." assert isinstance(op.pipe, PIPE), "Invalid 'pipe' field, PIPE type is required." assert isinstance(op.mode, MODE), "Invalid 'mode' field, MODE type is required." - # Retrieve arguments signature from __init__ method and save it. signature = inspect.signature(op) setattr(op, "signature", signature) - # Register the custom op configuration. _custom_op_registry[op.name] = op return op -# ----------------------- -# SPMD Programming Model -# ----------------------- -def _constexpr_to_value(v): - if isinstance(v, tl.constexpr): - return v.value - return v - - -def _is_int_like_elem(x) -> bool: - """Accept int / tl.constexpr(int) / tl.tensor(int*).""" - if isinstance(x, int): - return True - if isinstance(x, tl.constexpr): - # constexpr value should be python int - return isinstance(x.value, int) - if isinstance(x, tl.tensor): - # Offsets/strides must be integer typed (i32/i64 etc.) - return x.dtype.is_int() - return False - - -def _assert_int_like_tuple(name: str, xs): - assert isinstance( - xs, (tuple, list) - ), f"{name} should be a tuple/list, but got {type(xs)}" - assert all(_is_int_like_elem(x) for x in xs), f"{name} should be integer" +_dtype_cname_dict = { + "int1": "bool", + "int8": "int8_t", + "int16": "int16_t", + "int32": "int32_t", + "int64": "int64_t", + "uint8": "uint8_t", + "uint16": "uint16_t", + "uint32": "uint32_t", + "uint64": "uint64_t", + "fp16": "half", + "bf16": "bfloat16_t", + "fp32": "float", + "fp64": "double", + "fp8e5": "float8_e5m2_t", + "fp8e4nv": "float8_e4m3_t", +} + + +def _cname(self): + """Return the corresponding C name of the given tl.dtype""" + return _dtype_cname_dict.get(self.name, self.name) + + +# Add 'cname' property to tl.dtype class. +tl.dtype.cname = property(_cname, None) diff --git a/language/deeplink/cann/extension/dispatch.py b/language/deeplink/cann/extension/dispatch.py new file mode 100644 index 00000000..08985222 --- /dev/null +++ b/language/deeplink/cann/extension/dispatch.py @@ -0,0 +1,13 @@ +""" +Dispatch table for Ascend-specific 'with' statement context managers. +""" + +from .scope import scope +from .code_generator import handle_scope_with, mangle_ty + +__all__ = ["ASCEND_WITH_DISPATCH"] + +ASCEND_WITH_DISPATCH = { + scope: handle_scope_with, + "mangle_ty": mangle_ty, +} diff --git a/language/deeplink/cann/extension/layout.py b/language/deeplink/cann/extension/layout.py new file mode 100644 index 00000000..3789a342 --- /dev/null +++ b/language/deeplink/cann/extension/layout.py @@ -0,0 +1,90 @@ +# 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. + +from triton.language.core import _unwrap_if_constexpr + +__all__ = [ + "layout", + "ND", + "NZ", + "fragment", + "UB", + "L1", + "L0A", + "L0B", + "L0C", +] + + +class layout: + ASCEND = ["ND", "NZ"] + + def __init__(self, name): + name = _unwrap_if_constexpr(name) + self.name = name + assert name in layout.ASCEND, name + + def __str__(self): + return self.name + + def codegen_name(self): + return self.name + + @property + def cache_key_part(self) -> str: + return self.name + + def __repr__(self): + return f"triton.language.{self.codegen_name()}" + + +ND = layout("ND") +NZ = layout("NZ") + + +class _memory_scope: + GPU = ["fragment"] + ASCEND = ["UB", "L1", "L0A", "L0B", "L0C"] + + def __init__(self, name): + name = _unwrap_if_constexpr(name) + self.name = name + assert name in _memory_scope.ASCEND + _memory_scope.GPU, name + + def __str__(self): + return self.name + + def codegen_name(self): + return self.name + + @property + def cache_key_part(self) -> str: + return self.name + + def __repr__(self): + return f"triton.language.{self.codegen_name()}" + + +fragment = _memory_scope("fragment") +UB = _memory_scope("UB") +L1 = _memory_scope("L1") +L0A = _memory_scope("L0A") +L0B = _memory_scope("L0B") +L0C = _memory_scope("L0C") diff --git a/language/deeplink/cann/extension/math_ops.py b/language/deeplink/cann/extension/math_ops.py new file mode 100644 index 00000000..29b9aea9 --- /dev/null +++ b/language/deeplink/cann/extension/math_ops.py @@ -0,0 +1,24 @@ +# 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. + +# Re-export math functions from libdevice for alignment with triton-ascend layout. +from ..libdevice import atan2, isfinited, finitef + +__all__ = ["atan2", "isfinited", "finitef"] diff --git a/language/deeplink/cann/extension/mem_ops.py b/language/deeplink/cann/extension/mem_ops.py new file mode 100644 index 00000000..64d24684 --- /dev/null +++ b/language/deeplink/cann/extension/mem_ops.py @@ -0,0 +1,182 @@ +# 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. + +from triton.language.core import ( + _unwrap_if_constexpr, + builtin, + tensor, +) + +import triton.language.core as tl + + +__all__ = [ + "index_select_simd", + "gather_out_to_ub", + "scatter_ub_to_out", + "index_put", +] + + +@builtin +def index_select_simd( + src, dim, index, src_shape, src_offset, read_shape, _semantic=None +): + dim = _unwrap_if_constexpr(dim) + newsrc_shape = [ + _semantic.to_tensor(o) if isinstance(o, tl.constexpr) else o for o in src_shape + ] + newsrc_offset = [ + _semantic.to_tensor(o) if isinstance(o, tl.constexpr) else o for o in src_offset + ] + assert len(index.shape) == 1, "index must be a 1D tensor" + + ndim = len(newsrc_shape) + return_shape = [index.shape[0] if i == dim else read_shape[i] for i in range(ndim)] + element_ty = src.type.element_ty + output_ty = tl.block_type(element_ty, return_shape) + + newsrc_shape_handles = [ + s.handle if isinstance(s, tensor) else s for s in newsrc_shape + ] + newsrc_offset_handles = [ + s.handle if isinstance(s, tensor) else s for s in newsrc_offset + ] + out = _semantic.builder.create_index_select_simd( + src.handle, + index.handle, + dim, + newsrc_shape_handles, + newsrc_offset_handles, + read_shape, + return_shape, + ) + return tl.tensor(out, output_ty) + + +@builtin +def gather_out_to_ub( + src, + index, + index_boundary, + dim, + src_stride, + end_offset, + start_offset, + other=None, + _semantic=None, +): + dim = _unwrap_if_constexpr(dim) + index_boundary = _unwrap_if_constexpr(index_boundary) + + src_stride_handles = [s.handle if isinstance(s, tensor) else s for s in src_stride] + end_offset_handles = [s.handle if isinstance(s, tensor) else s for s in end_offset] + start_offset_handles = [ + s.handle if isinstance(s, tensor) else s for s in start_offset + ] + other_handle = None + if other is not None: + other = _semantic.cast(other, src.dtype.element_ty) + other_handle = other.handle if isinstance(other, tensor) else other + + ret = _semantic.builder.create_gather_out_to_ub( + src.handle, + index.handle, + index_boundary, + dim, + src_stride_handles, + end_offset_handles, + start_offset_handles, + other_handle, + ) + ret_shape = [_unwrap_if_constexpr(s) for s in index.shape] + return _semantic.wrap_tensor(ret, src.dtype.element_ty, ret_shape) + + +@builtin +def scatter_ub_to_out( + ptr, + value, + index, + index_boundary, + dim, + dst_stride, + end_offset, + start_offset, + _semantic=None, +): + dim = _unwrap_if_constexpr(dim) + index_boundary = _unwrap_if_constexpr(index_boundary) + + dst_stride_handles = [s.handle if isinstance(s, tensor) else s for s in dst_stride] + end_offset_handles = [s.handle if isinstance(s, tensor) else s for s in end_offset] + start_offset_handles = [ + s.handle if isinstance(s, tensor) else s for s in start_offset + ] + + return tl.tensor( + _semantic.builder.create_scatter_ub_to_out( + ptr.handle, + value.handle, + index.handle, + index_boundary, + dim, + dst_stride_handles, + end_offset_handles, + start_offset_handles, + ), + tl.void, + ) + + +@builtin +def index_put( + ptr, + index, + value, + dim, + index_boundary, + end_offset, + start_offset, + dst_stride, + _semantic=None, +): + dim = _unwrap_if_constexpr(dim) + index_boundary = _unwrap_if_constexpr(index_boundary) + + end_offset_handles = [s.handle if isinstance(s, tensor) else s for s in end_offset] + start_offset_handles = [ + s.handle if isinstance(s, tensor) else s for s in start_offset + ] + dst_stride_handles = [s.handle if isinstance(s, tensor) else s for s in dst_stride] + + return tl.tensor( + _semantic.builder.create_index_put( + ptr.handle, + index.handle, + value.handle, + dim, + index_boundary, + end_offset_handles, + start_offset_handles, + dst_stride_handles, + ), + tl.void, + ) diff --git a/language/deeplink/cann/extension/scope.py b/language/deeplink/cann/extension/scope.py new file mode 100644 index 00000000..0b1ab966 --- /dev/null +++ b/language/deeplink/cann/extension/scope.py @@ -0,0 +1,43 @@ +# 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. + +from triton.language.core import _unwrap_if_constexpr + + +class scope: + def __init__(self, core_mode: str, _builder=None, _semantic=None, **kwargs): + self.core_mode = ( + _unwrap_if_constexpr(core_mode) if _builder is None else core_mode + ) + self._builder = _builder + self._semantic = _semantic + self.disable_auto_sync = kwargs.get("disable_auto_sync", False) + if self.core_mode not in ("cube", "vector"): + raise ValueError( + f'core_mode must be "cube" or "vector", got {self.core_mode}' + ) + + def __enter__(self): + if self._builder is None: + raise RuntimeError("scope can only be used inside a Triton kernel") + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + return False diff --git a/language/deeplink/cann/extension/semantic.py b/language/deeplink/cann/extension/semantic.py new file mode 100644 index 00000000..477d4da5 --- /dev/null +++ b/language/deeplink/cann/extension/semantic.py @@ -0,0 +1,318 @@ +# 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. + +from typing import List +from contextlib import contextmanager +from triton.language import core as tl +from triton.language import semantic as tl_semantic +from triton._C.libtriton import ir, dicp_triton + +_dry_run = False + + +@contextmanager +def dry_run_context(): + global _dry_run + _dry_run = True + try: + yield + finally: + _dry_run = False + + +_SENDER_RECEIVER_MAP = { + "cube": ("vector", dicp_triton.ir.PIPE.PIPE_FIX, dicp_triton.ir.PIPE.PIPE_MTE2), + "vector": ("cube", dicp_triton.ir.PIPE.PIPE_MTE3, dicp_triton.ir.PIPE.PIPE_MTE2), +} + + +def create_address_space( + address_space: dicp_triton.ir.AddressSpace, + builder, +) -> ir.attribute: + return builder.get_target_attribute(address_space) + + +class PIPE: + PIPE_S = dicp_triton.ir.PIPE.PIPE_S + PIPE_V = dicp_triton.ir.PIPE.PIPE_V + PIPE_M = dicp_triton.ir.PIPE.PIPE_M + PIPE_MTE1 = dicp_triton.ir.PIPE.PIPE_MTE1 + PIPE_MTE2 = dicp_triton.ir.PIPE.PIPE_MTE2 + PIPE_MTE3 = dicp_triton.ir.PIPE.PIPE_MTE3 + PIPE_ALL = dicp_triton.ir.PIPE.PIPE_ALL + PIPE_FIX = dicp_triton.ir.PIPE.PIPE_FIX + + +def insert_slice( + ful: tl.tensor, + sub: tl.tensor, + offsets: List[tl.tensor], + sizes: List[int], + strides: List[int], + _semantic: tl_semantic.TritonSemantic, +) -> tl.tensor: + assert len(ful.shape) == len(offsets) + assert len(ful.shape) == len(sizes) + assert len(ful.shape) == len(strides) + assert all([s >= 1 for s in sizes]) + assert all([s >= 0 for s in strides]) + new_offsets = [o.handle for o in offsets] + ret_type = tl.block_type(ful.type.scalar, ful.shape) + out = _semantic.builder.create_insert_slice( + ful.handle, sub.handle, new_offsets, sizes, strides + ) + return tl.tensor(out, ret_type) + + +def extract_slice( + ful: tl.tensor, + offsets: List[tl.tensor], + sizes: List[int], + strides: List[int], + _semantic: tl_semantic.TritonSemantic, +) -> tl.tensor: + assert len(ful.shape) == len(offsets) + assert len(ful.shape) == len(sizes) + assert len(ful.shape) == len(strides) + assert all([s >= 1 for s in sizes]) + assert all([s >= 0 for s in strides]) + new_offsets = [o.handle for o in offsets] + ret_type = tl.block_type(ful.type.scalar, sizes) + out = _semantic.builder.create_extract_slice( + ful.handle, new_offsets, sizes, strides + ) + return tl.tensor(out, ret_type) + + +def compile_hint(ptr: tl.tensor, hint_name: str, hint_val, builder: ir.builder): + if isinstance(hint_val, bool): + hint_val = builder.get_bool_attr(hint_val) + elif not hint_val: + hint_val = builder.get_unit_attr() + elif isinstance(hint_val, int): + hint_val = builder.get_int32_attr(hint_val) + elif isinstance(hint_val, tl.constexpr): + hint_val = builder.get_string_attr(hint_val.value) + elif isinstance(hint_val, (list, tl.tuple)): + hint_val = builder.get_i64_array_attr(hint_val) + else: + raise ValueError(f"Unsupported hint value type: {type(hint_val)}") + builder.create_annotation_mark(ptr.handle, hint_name, hint_val) + + +def alloc( + shape: List[int], value, dtype: tl.dtype, layout, scope, builder: ir.builder +) -> tl.tensor: + if isinstance(value, tl.tensor): + assert value.numel.value == 1, "only accepts size-1 tensor" + value = tl_semantic.cast(value, dtype, builder) + else: + if dtype is None: + raise ValueError("dtype must be specified when value is not a tensor") + if value == 0: + value = builder.get_null_value(dtype.to_ir(builder)) + else: + get_value_fn = getattr(builder, f"get_{dtype.name}") + value = get_value_fn(value) + value = tl.tensor(value, dtype) + if len(shape) == 0: + return value + ret_ty = tl.block_type(value.dtype, shape) + x = tl.tensor(builder.create_splat(value.handle, shape), ret_ty) + if layout is not None: + builder.create_annotation_mark( + x.handle, "layout", builder.get_string_attr(str(layout)) + ) + if scope is not None: + builder.create_annotation_mark( + x.handle, "scope", builder.get_string_attr(str(scope)) + ) + return x + + +def custom_sync_op(builder: ir.builder, op_name: str, **kwargs): + if _dry_run: + return None + if op_name == "sync_block_all": + return builder.sync_block_all(kwargs["mode"], kwargs["event_id"]) + elif op_name == "sync_block_set": + sender = kwargs["sender"] + receiver, sender_pipe, receiver_pipe = _SENDER_RECEIVER_MAP[sender] + event_id = kwargs["event_id"] + id_value = builder.get_int64(event_id) + return builder.sync_block_set( + sender, receiver, id_value, sender_pipe, receiver_pipe + ) + elif op_name == "sync_block_wait": + sender = kwargs["sender"] + receiver, sender_pipe, receiver_pipe = _SENDER_RECEIVER_MAP[sender] + event_id = kwargs["event_id"] + id_value = builder.get_int64(event_id) + return builder.sync_block_wait( + sender, receiver, id_value, sender_pipe, receiver_pipe + ) + raise ValueError(f"Unsupported custom op: {op_name}") + + +def create_sync_block_set( + sender, receiver, event_id, sender_pipe, receiver_pipe, _semantic=None +): + if isinstance(event_id, int): + _semantic.builder.sync_block_set( + sender, + receiver, + _semantic.to_tensor(tl.constexpr(event_id)).handle, + sender_pipe.value, + receiver_pipe.value, + ) + elif isinstance(event_id, tl.constexpr): + _semantic.builder.sync_block_set( + sender, + receiver, + _semantic.to_tensor(event_id).handle, + sender_pipe.value, + receiver_pipe.value, + ) + else: + _semantic.builder.sync_block_set( + sender, receiver, event_id.handle, sender_pipe.value, receiver_pipe.value + ) + + +def create_sync_block_wait( + sender, receiver, event_id, sender_pipe, receiver_pipe, _semantic=None +): + if isinstance(event_id, int): + _semantic.builder.sync_block_wait( + sender, + receiver, + _semantic.to_tensor(tl.constexpr(event_id)).handle, + sender_pipe.value, + receiver_pipe.value, + ) + elif isinstance(event_id, tl.constexpr): + _semantic.builder.sync_block_wait( + sender, + receiver, + _semantic.to_tensor(event_id).handle, + sender_pipe.value, + receiver_pipe.value, + ) + else: + _semantic.builder.sync_block_wait( + sender, receiver, event_id.handle, sender_pipe.value, receiver_pipe.value + ) + + +def sub_vec_id(_semantic=None): + return tl.tensor(_semantic.builder.create_get_sub_vec_id(), tl.int64) + + +def copy_from_ub_to_l1(src, dst, _semantic=None): + from ..buffer.core import buffer as bl_buffer + from . import core as _core + + if isinstance(src, tl.tensor) or isinstance(dst, tl.tensor): + raise TypeError("tensor not support yet") + if src.shape != dst.shape: + raise TypeError("src and dst must have same shape") + if src.dtype != dst.dtype: + raise TypeError("src and dst need to have the same type") + if isinstance(src, bl_buffer) and isinstance(dst, bl_buffer): + if src.space != _core.ascend_address_space.UB: + raise TypeError("src's AddressSpace must be UB") + if dst.space != _core.ascend_address_space.L1: + raise TypeError("dst's AddressSpace must be L1") + _semantic.builder.create_copy_buffer(src.handle, dst.handle) + else: + raise TypeError("src and dst must be tl.tensor or bl.buffer") + + +def copy(src, dst, _semantic=None): + from ..buffer.core import buffer as bl_buffer + from . import core as _core + + if isinstance(src, tl.tensor) or isinstance(dst, tl.tensor): + raise TypeError("tensor not support yet") + if src.shape != dst.shape: + raise TypeError("src and dst must have same shape") + if src.dtype != dst.dtype: + raise TypeError("src and dst need to have the same type") + if isinstance(src, bl_buffer) and isinstance(dst, bl_buffer): + if src.space != _core.ascend_address_space.UB: + raise TypeError("src's AddressSpace must be UB") + if dst.space not in ( + _core.ascend_address_space.L1, + _core.ascend_address_space.UB, + ): + raise TypeError("dst's AddressSpace must be UB or L1") + _semantic.builder.create_copy_buffer(src.handle, dst.handle) + else: + raise TypeError("src and dst must be tl.tensor or bl.buffer") + + +def fixpipe( + src, + dst, + dma_mode, + dual_dst_mode, + pre_quant_mode, + pre_relu_mode, + _semantic=None, +): + if dst is None: + result = _semantic.builder.create_fixpipe( + src.handle, + None, + dma_mode.value, + dual_dst_mode.value, + pre_quant_mode.value, + pre_relu_mode.value, + ) + if dual_dst_mode.value.name == "ROW_SPLIT": + new_shape = list(src.type.shape) + if len(new_shape) >= 1 and new_shape[0] > 0: + new_shape[0] = new_shape[0] // 2 + new_type = tl.block_type(src.type.element_ty, new_shape) + return tl.tensor(result, new_type) + elif dual_dst_mode.value.name == "COLUMN_SPLIT": + new_shape = list(src.type.shape) + if len(new_shape) >= 2 and new_shape[1] > 0: + new_shape[1] = new_shape[1] // 2 + new_type = tl.block_type(src.type.element_ty, new_shape) + return tl.tensor(result, new_type) + else: + return tl.tensor(result, src.type) + else: + _semantic.builder.create_fixpipe( + src.handle, + dst.handle, + dma_mode.value, + dual_dst_mode.value, + pre_quant_mode.value, + pre_relu_mode.value, + ) + + +def debug_barrier(sync_mode: str, _semantic=None): + target = tl.tensor(_semantic.builder.get_int64(0), tl.int64) + attr = _semantic.builder.get_string_attr(sync_mode) + _semantic.builder.create_debug_barrier(target.handle, "SYNC_IN_VF", attr) diff --git a/language/deeplink/cann/extension/vec_ops.py b/language/deeplink/cann/extension/vec_ops.py new file mode 100644 index 00000000..f9b2ac8f --- /dev/null +++ b/language/deeplink/cann/extension/vec_ops.py @@ -0,0 +1,195 @@ +# 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. + +import builtins + +from triton.language.core import ( + _tensor_member_fn, + _unwrap_if_constexpr, + builtin, + constexpr, + slice, + tensor, +) + +import triton.language.core as tl +from . import semantic as dl_semantic + + +__all__ = ["insert_slice", "extract_slice", "get_element", "sort", "flip"] + + +def _constexpr_to_value(v): + if isinstance(v, constexpr): + return v.value + return v + + +def _extract_slice(sl: slice, shape: constexpr): + def constexpr_or_none_to_value(v, default: int): + if v is None: + return default + assert isinstance( + v, (constexpr, int) + ), f"slice only can be constexpr or int, got: {v}" + return _constexpr_to_value(v) + + start = constexpr_or_none_to_value(sl.start, 0) + stop = constexpr_or_none_to_value(sl.stop, _constexpr_to_value(shape)) + step = constexpr_or_none_to_value(sl.step, 1) + size = (stop - start + step - 1) // step + assert ( + start >= 0 and stop >= 0 and step >= 0 and size >= 0 + ), "slice should be greater than 0" + return start, size, step + + +@_tensor_member_fn +@builtin +def __getitem__(self, slices, _semantic=None): + if ( + isinstance(slices, (builtins.slice, slice, constexpr, tensor, int)) + or slices is None + ): + slices = [slices] + if isinstance(slices, tuple): + slices = slices.values + ret = self + offsets = [] + sizes = [] + strides = [] + dst_shape = [] + need_extract_slice = False + for dim, sl in enumerate(slices): + if sl is None or isinstance(sl, constexpr) and sl.value is None: + ret = _semantic.expand_dims(ret, dim) + offsets.append(_semantic.builder.get_int32(0)) + dst_shape.append(constexpr(1)) + sizes.append(constexpr(1)) + strides.append(constexpr(1)) + elif ( + isinstance(sl, slice) + and sl.start is None + and sl.stop is None + and sl.step is None + ): + pass + elif isinstance(sl, constexpr) and sl.value is not None: + offsets.append(_semantic.builder.get_int32(_constexpr_to_value(sl))) + need_extract_slice = True + sizes.append(constexpr(1)) + strides.append(constexpr(1)) + elif isinstance(sl, int): + offsets.append(_semantic.builder.get_int32(sl)) + need_extract_slice = True + sizes.append(constexpr(1)) + strides.append(constexpr(1)) + elif isinstance(sl, tensor): + offsets.append(sl.handle) + sizes.append(constexpr(1)) + strides.append(constexpr(1)) + need_extract_slice = True + elif isinstance(sl, (slice, builtins.slice)): + start, size, step = _extract_slice(sl, ret.shape[dim]) + offsets.append(start) + strides.append(constexpr(step)) + sizes.append(constexpr(size)) + dst_shape.append(constexpr(size)) + need_extract_slice = True + else: + raise ValueError(f"unsupported tensor index: {sl}") + + if need_extract_slice: + new_offsets = [ + (_semantic.to_tensor(o) if not isinstance(o, tensor) else o) + for o in offsets + ] + ret = dl_semantic.extract_slice( + self, new_offsets, sizes, strides, _semantic=_semantic + ) + return ret + + +@_tensor_member_fn +@builtin +def insert_slice( + ful, sub, offsets, sizes, strides, _builder=None, _generator=None, _semantic=None +) -> tensor: + """ + Insert a tensor to another tensor as specified by the operation's offsets, sizes and strides arguments. + """ + assert len(ful.shape) > 0 + assert len(ful.shape) == len(sub.shape) + new_offsets = [ + _semantic.to_tensor(o) if isinstance(o, constexpr) else o for o in offsets + ] + return dl_semantic.insert_slice( + ful, sub, new_offsets, sizes, strides, _semantic=_semantic + ) + + +@_tensor_member_fn +@builtin +def extract_slice( + ful, offsets, sizes, strides, _generator=None, _semantic=None +) -> tensor: + """ + Extract a tensor from another tensor as specified by the operation's offsets, sizes and strides arguments. + """ + assert len(ful.shape) > 0 + new_offsets = [ + _semantic.to_tensor(o) if isinstance(o, constexpr) else o for o in offsets + ] + return dl_semantic.extract_slice( + ful, new_offsets, sizes, strides, _semantic=_semantic + ) + + +@builtin +def get_element(src, indice, _semantic=None, _generator=None): + assert len(src.shape) > 0 + new_indice = [ + _semantic.to_tensor(i) if isinstance(i, constexpr) else i for i in indice + ] + new_indice_handles = [] + for i in new_indice: + if isinstance(i, tensor): + new_indice_handles.append(i.handle) + elif isinstance(i, int): + new_indice_handles.append(i) + else: + new_indice_handles.append(i.handle if hasattr(i, "handle") else i) + result = _semantic.builder.create_extract_scalar(src.handle, new_indice_handles) + return _semantic.wrap_tensor(result, src.type.scalar, None) + + +@builtin +def sort(ptr, dim=-1, descending=False, _semantic=None): + dim = _unwrap_if_constexpr(dim) + if hasattr(descending, "value"): + descending = bool(descending.value) + else: + descending = bool(descending) + sorted_vals = _semantic.builder.create_sort(ptr.handle, dim, descending) + return tensor(sorted_vals, type=ptr.type) + + +# flip is defined in libdevice; re-export here for alignment with triton-ascend +from ..libdevice import flip diff --git a/language/deeplink/cann/libdevice.py b/language/deeplink/cann/libdevice.py new file mode 100644 index 00000000..f44d53e5 --- /dev/null +++ b/language/deeplink/cann/libdevice.py @@ -0,0 +1,1056 @@ +from math import pi as math_pi +from triton.language import core, math, semantic +from triton._C.libtriton import ir +from triton.runtime.jit import jit +from triton.backends.dicp_triton.utils import get_ascend_arch_from_env +import triton + + +# TODO :triton-ascend 拆分为 extension/math_ops.py、extension/vec_ops.py，DLCompiler 直接放在 libdevice.py 中 + +# --------------------------------------------------------------------------- +# Pure extern functions (bitcode __hmf_* symbols) +# --------------------------------------------------------------------------- + + +@core.extern +def reciprocal(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_recipf", core.dtype("fp32")), + (core.dtype("fp16"),): ("__hmf_recipDh", core.dtype("fp16")), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def log1p(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_log1pf", core.dtype("fp32")), + (core.dtype("fp16"),): ("__hmf_log1pDh", core.dtype("fp16")), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def relu(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_reluf", core.dtype("fp32")), + (core.dtype("fp16"),): ("__hmf_reluDh", core.dtype("fp16")), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def isinf(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_isinf", core.dtype("int1")), + (core.dtype("fp16"),): ("__hmf_isinf", core.dtype("int1")), + (core.dtype("bf16"),): ("__hmf_isinf", core.dtype("int1")), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def tan(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_tanf", core.dtype("fp32")), + (core.dtype("fp16"),): ("__hmf_tanDh", core.dtype("fp16")), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def atan(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_atanf", core.dtype("fp32")), + (core.dtype("fp16"),): ("__hmf_atanDh", core.dtype("fp16")), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def tanh(arg0, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + original_dtype = arg0.dtype + if original_dtype == core.dtype("bf16"): + arg0 = _semantic.cast(arg0, core.float32) + + dispatch = { + (core.dtype("fp32"),): ("__hmf_tanhf", core.dtype("fp32")), + (core.dtype("fp16"),): ("__hmf_tanhDh", core.dtype("fp16")), + } + res = core.extern_elementwise( + "", "", [arg0], dispatch, is_pure=True, _semantic=_semantic + ) + if original_dtype == core.dtype("bf16"): + return _semantic.cast(res, core.dtype("bf16")) + return res + + +@core.extern +def ilogb(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_ilogbf", core.dtype("fp32")), + (core.dtype("fp16"),): ("__hmf_ilogbDh", core.dtype("fp16")), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def ldexp(arg0, arg1, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0, arg1], + { + (core.dtype("fp32"), core.dtype("int32")): ( + "__hmf_ldexpf", + core.dtype("fp32"), + ), + (core.dtype("fp16"), core.dtype("int32")): ( + "__hmf_ldexpDh", + core.dtype("fp16"), + ), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def pow(arg0, arg1, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0, arg1], + { + (core.dtype("fp32"), core.dtype("fp32")): ( + "__hmf_powf", + core.dtype("fp32"), + ), + (core.dtype("fp16"), core.dtype("fp16")): ( + "__hmf_powDh", + core.dtype("fp16"), + ), + (core.dtype("bf16"), core.dtype("bf16")): ( + "__hmf_powDb", + core.dtype("bf16"), + ), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def isnan(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_isnan", core.dtype("int1")), + (core.dtype("fp16"),): ("__hmf_isnan", core.dtype("int1")), + (core.dtype("bf16"),): ("__hmf_isnan", core.dtype("int1")), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def div_rz(arg0, arg1, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0, arg1], + { + (core.dtype("fp32"), core.dtype("fp32")): ( + "__hmf_div_rz_fp32", + core.dtype("fp32"), + ), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.builtin +def fast_dividef(arg0, arg1, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + arg1 = _semantic.to_tensor(arg1) + ret = _semantic.fdiv(arg0, arg1, False) + return ret + + +@core.builtin +def fast_expf(arg0, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + ret = core.tensor(_semantic.builder.create_exp(arg0.handle), arg0.type) + return ret + + +@core.extern +def fmod(arg0, arg1, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0, arg1], + { + (core.dtype("fp32"), core.dtype("fp32")): ( + "__hmf_fmod_fp32", + core.dtype("fp32"), + ), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def float_as_int(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_float_as_int_fp32", core.dtype("int32")), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def atan2(arg0, arg1, _semantic=None): + if arg0.dtype == core.dtype("bf16") or arg1.dtype == core.dtype("bf16"): + core.static_print( + "extern libdevice.atan2 for dtype bf16 is unsupported for now." + ) + core.static_assert(False) + return core.extern_elementwise( + "", + "", + [arg0, arg1], + { + (core.dtype("fp16"), core.dtype("fp16")): ( + "__hmf_atan2_fp16", + core.dtype("fp16"), + ), + (core.dtype("fp32"), core.dtype("fp32")): ( + "__hmf_atan2_fp32", + core.dtype("fp32"), + ), + }, + is_pure=True, + _semantic=_semantic, + ) + + +@core.extern +def round(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_roundf", core.dtype("fp32")), + }, + is_pure=True, + _semantic=_semantic, + ) + + +# --------------------------------------------------------------------------- +# IR builder math implementations (SIMD path) +# --------------------------------------------------------------------------- + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("arcsine") +def acos(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + pi_half = 1.5707963268 + abs_x = math.abs(arg0, _semantic=_semantic) + arg0_2 = _semantic.mul(arg0, arg0, True) + arg0_4 = _semantic.mul(arg0_2, arg0_2, True) + arg0_6 = _semantic.mul(arg0_4, arg0_2, True) + arg0_8 = _semantic.mul(arg0_6, arg0_2, True) + arg0_10 = _semantic.mul(arg0_8, arg0_2, True) + poly = _semantic.add(1.0, _semantic.mul(0.166667, arg0_2, True), True) + poly = _semantic.add(poly, _semantic.mul(0.075, arg0_4, True), True) + poly = _semantic.add(poly, _semantic.mul(0.044643, arg0_6, True), True) + poly = _semantic.add(poly, _semantic.mul(0.030380, arg0_8, True), True) + poly = _semantic.add(poly, _semantic.mul(0.022372, arg0_10, True), True) + acos_center = _semantic.sub(pi_half, _semantic.mul(arg0, poly, True), True) + + numerator_mid = _semantic.sub(1.0, abs_x, True) + denom_mid = _semantic.add(1.0, abs_x, True) + div_mid = _semantic.truediv(numerator_mid, denom_mid) + t_mid = math.sqrt(div_mid, _semantic=_semantic) + t2_mid = _semantic.mul(t_mid, t_mid, True) + t4_mid = _semantic.mul(t2_mid, t2_mid, True) + t6_mid = _semantic.mul(t4_mid, t2_mid, True) + poly_mid1 = _semantic.mul(0.1065976, t2_mid, True) + poly_mid2 = _semantic.add(-0.1420890, poly_mid1, True) + poly_mid3 = _semantic.mul(poly_mid2, t2_mid, True) + poly_mid4 = _semantic.add(0.1999341, poly_mid3, True) + poly_mid5 = _semantic.mul(poly_mid4, t2_mid, True) + poly_mid6 = _semantic.add(-0.3333310, poly_mid5, True) + poly_mid = _semantic.add(1.0, _semantic.mul(poly_mid6, t2_mid, True), True) + arctan_t = _semantic.mul(t_mid, poly_mid, True) + acos_mid = _semantic.mul(2.0, arctan_t, True) + is_neg_mid = _semantic.less_than(arg0, 0.0) + acos_mid_signed = _semantic.where( + is_neg_mid, _semantic.sub(3.1415926536, acos_mid, True), acos_mid + ) + is_center = _semantic.less_than(abs_x, 0.6) + return _semantic.where(is_center, acos_center, acos_mid_signed) + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("sinh") +def sinh(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + exp0 = core.tensor(_semantic.builder.create_exp(arg0.handle), arg0.type) + exp1 = _semantic.truediv(1.0, exp0) + tmp = _semantic.sub(exp0, exp1, True) + ret = _semantic.truediv(tmp, 2.0) + return ret + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("cosh") +def cosh(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + exp0 = core.tensor(_semantic.builder.create_exp(arg0.handle), arg0.type) + exp1 = _semantic.truediv(1.0, exp0) + tmp = _semantic.add(exp0, exp1, True) + ret = _semantic.truediv(tmp, 2.0) + return ret + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("acosh") +def acosh(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + tmp = _semantic.sub(_semantic.mul(arg0, arg0, True), 1.0, True) + sqrt_res = core.tensor(_semantic.builder.create_sqrt(tmp.handle), tmp.type) + sum_res = _semantic.add(arg0, sqrt_res, True) + return core.tensor(_semantic.builder.create_log(sum_res.handle), sum_res.type) + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("asinh") +def asinh(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + tmp = _semantic.add(_semantic.mul(arg0, arg0, True), 1.0, True) + sqrt_res = core.tensor(_semantic.builder.create_sqrt(tmp.handle), tmp.type) + sum_res = _semantic.add(arg0, sqrt_res, True) + return core.tensor(_semantic.builder.create_log(sum_res.handle), sum_res.type) + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("atanh") +def atanh(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + a = _semantic.add(1.0, arg0, True) + b = _semantic.sub(1.0, arg0, True) + lna = core.tensor(_semantic.builder.create_log(a.handle), a.type) + lnb = core.tensor(_semantic.builder.create_log(b.handle), b.type) + tmp = _semantic.sub(lna, lnb, True) + return _semantic.mul(tmp, 0.5, True) + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_1arg_docstr("expm1") +def expm1(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + tmp = core.tensor(_semantic.builder.create_exp(arg0.handle), arg0.type) + return _semantic.sub(tmp, 1, True) + + +@core.builtin +@math._check_dtype(dtypes=["fp16", "fp32"]) +@math._add_math_2arg_docstr("nextafter") +def nextafter(arg0: core.tensor, arg1: core.tensor, _semantic=None): + x = _semantic.to_tensor(arg0) + y = _semantic.to_tensor(arg1) + dtype_map = { + "bf16": core.int16, + "fp16": core.int16, + "fp32": core.int32, + } + min_pos_bit = { + "bf16": 0x0001, + "fp16": 0x0001, + "fp32": 0x00000001, + } + max_neg_bit = { + "bf16": 0x8001, + "fp16": 0x8001, + "fp32": 0x80000001, + } + int_type = dtype_map[x.type.scalar.name] + x_eq_y = _semantic.equal(x, y) + x_gt_0 = _semantic.greater_than(x, 0) + y_gt_x = _semantic.greater_than(y, x) + next_neg = _semantic.xor_(x_gt_0, y_gt_x) + next_pos = _semantic.not_(next_neg) + + p1 = _semantic.full(x.shape, 1, int_type) + n1 = _semantic.full(x.shape, -1, int_type) + dir_xy = _semantic.where(next_pos, p1, n1) + x_abs = math.abs(x, _semantic=_semantic) + x_is_0 = _semantic.equal(x_abs, 0) + + min_pos = _semantic.full(x.shape, min_pos_bit[x.type.scalar.name], int_type) + max_neg = _semantic.full(x.shape, max_neg_bit[x.type.scalar.name], int_type) + min_pos = _semantic.bitcast(min_pos, x.dtype) + max_neg = _semantic.bitcast(max_neg, x.dtype) + bits_x = _semantic.bitcast(x, int_type) + bits_next = _semantic.add(bits_x, dir_xy, True) + next_val = _semantic.bitcast(bits_next, x.dtype) + + need_min_pos = _semantic.logical_and(x_is_0, next_pos) + need_max_neg = _semantic.logical_and(x_is_0, next_neg) + next_val = _semantic.where(need_min_pos, min_pos, next_val) + next_val = _semantic.where(need_max_neg, max_neg, next_val) + return _semantic.where(x_eq_y, x, next_val) + + +@core.builtin +@math._check_dtype(dtypes=["bf16", "fp16", "fp32"]) +@math._add_math_2arg_docstr("hypot(Euclidean Distance)") +def hypot(arg0: core.tensor, arg1: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + arg1 = _semantic.to_tensor(arg1) + x2 = _semantic.mul(arg0, arg0, True) + y2 = _semantic.mul(arg1, arg1, True) + sum_res = _semantic.add(x2, y2, True) + return core.tensor(_semantic.builder.create_sqrt(sum_res.handle), sum_res.type) + + +# Derived from the Cephes Math Library release 2.8: June, 2000 +# https://netlib.org/cephes/ +@core.builtin +@math._check_dtype(dtypes=["fp16", "fp32"]) +@math._add_math_2arg_docstr( + "besseli0 (Modified Bessel function of the first kind, order 0)." +) +def cyl_bessel_i0(arg0: core.tensor, _semantic=None): + param1 = [ + -4.41534164647933937950e-18, + +3.33079451882223809783e-17, + -2.43127984654795469359e-16, + +1.71539128555513303061e-15, + -1.16853328779934516808e-14, + +7.67618549860493561688e-14, + -4.85644678311192946090e-13, + +2.95505266312963983461e-12, + -1.72682629144155570723e-11, + +9.67580903537323691224e-11, + -5.18979560163526290666e-10, + +2.65982372468238665035e-09, + -1.30002500998624804212e-08, + +6.04699502254191894932e-08, + -2.67079385394061173391e-07, + +1.11738753912010371815e-06, + -4.41673835845875056359e-06, + +1.64484480707288970893e-05, + -5.75419501008210370398e-05, + +1.88502885095841655729e-04, + -5.76375574538582365885e-04, + +1.63947561694133579842e-03, + -4.32430999505057594430e-03, + +1.05464603945949983183e-02, + -2.37374148058994688156e-02, + +4.93052842396707084878e-02, + -9.49010970480476444210e-02, + +1.71620901522208775349e-01, + -3.04682672343198398683e-01, + +6.76795274409476084995e-01, + ] + param2 = [ + -7.23318048787475395456e-18, + -4.83050448594418207126e-18, + +4.46562142029675999901e-17, + +3.46122286769746109310e-17, + -2.82762398051658348494e-16, + -3.42548561967721913462e-16, + +1.77256013305652638360e-15, + +3.81168066935262242075e-15, + -9.55484669882830764870e-15, + -4.15056934728722208663e-14, + +1.54008621752140982691e-14, + +3.85277838274214270114e-13, + +7.18012445138366623367e-13, + -1.79417853150680611778e-12, + -1.32158118404477131188e-11, + -3.14991652796324136454e-11, + +1.18891471078464383424e-11, + +4.94060238822496958910e-10, + +3.39623202570838634515e-09, + +2.26666899049817806459e-08, + +2.04891858946906374183e-07, + +2.89137052083475648297e-06, + +6.88975834691682398426e-05, + +3.36911647825569408990e-03, + +8.04490411014108831608e-01, + ] + arg0 = _semantic.to_tensor(arg0) + abs_x = core.tensor(_semantic.builder.create_fabs(arg0.handle), arg0.type) + x_a = _semantic.sub(_semantic.mul(abs_x, 0.5, True), 2.0, True) + a_n_2 = 0 + a_n_1 = 0 + a_n = param1[0] + for i in range(1, 30): + a_n_2 = a_n_1 + a_n_1 = a_n + a_n = _semantic.sub(_semantic.mul(x_a, a_n_1, True), a_n_2, True) + a_n = _semantic.add(a_n, param1[i], True) + + f_32 = _semantic.full(abs_x.shape, 32.0, abs_x.type.scalar) + x_b = _semantic.sub(_semantic.fdiv(f_32, abs_x, True), 2.0, True) + b_n_2 = 0 + b_n_1 = 0 + b_n = param2[0] + for i in range(1, 25): + b_n_2 = b_n_1 + b_n_1 = b_n + b_n = _semantic.sub(_semantic.mul(x_b, b_n_1, True), b_n_2, True) + b_n = _semantic.add(b_n, param2[i], True) + + half_exp = _semantic.mul( + core.tensor(_semantic.builder.create_exp(abs_x.handle), abs_x.type), 0.5, True + ) + res_a = _semantic.mul(half_exp, _semantic.sub(a_n, a_n_2, True), True) + res_b = _semantic.fdiv( + _semantic.mul(half_exp, _semantic.sub(b_n, b_n_2, True), True), + core.tensor(_semantic.builder.create_sqrt(abs_x.handle), abs_x.type), + True, + ) + cond = _semantic.less_equal(abs_x, 8.0) + return _semantic.where(cond, res_a, res_b) + + +@core.extern +@math._check_dtype(dtypes=["fp16", "fp32"]) +def signbit(arg0, _semantic=None): + arg0_scalar_ty = arg0.type.scalar + if arg0_scalar_ty == core.float32: + int_ty = core.int32 + else: + int_ty = core.int16 + + arg0 = _semantic.to_tensor(arg0) + int_tensor = _semantic.bitcast(arg0, int_ty) + if int_ty == core.int32: + shift = 31 + elif int_ty == core.int16: + shift = 15 + + shift = _semantic.full(arg0.shape, shift, int_ty) + sign_bit_tensor = _semantic.lshr(int_tensor, shift) + sign_bit_tensor = _semantic.and_( + sign_bit_tensor, _semantic.full(arg0.shape, 1, int_ty) + ) + return _semantic.equal(sign_bit_tensor, 1) + + +@core.extern +@math._check_dtype(dtypes=["fp32"]) +def erfinv(arg0, _semantic=None): + arg0_scalar_ty = arg0.type.scalar + arg0 = _semantic.to_tensor(arg0) + + inv_sqrt_pi_times_2 = _semantic.full(arg0.shape, 1.128379167, arg0_scalar_ty).handle + coeff_low_numerator = [-0.140543331, 0.914624893, -1.645349621, 0.886226899] + coeff_low_denominator = [0.012229801, -0.329097515, 1.442710462, -2.118377725, 1.0] + coeff_high_numerator = [1.641345311, 3.429567803, -1.624906493, -1.970840454] + coeff_high_denominator = [1.6370678, 3.5438892, 1.0] + + arg0_squared = _semantic.builder.create_fmul(arg0.handle, arg0.handle) + numerator_low_range = _semantic.full( + arg0.shape, coeff_low_numerator[0], arg0_scalar_ty + ).handle + for i in range(1, len(coeff_low_numerator)): + numerator_low_range = _semantic.builder.create_fma( + numerator_low_range, + arg0_squared, + _semantic.full(arg0.shape, coeff_low_numerator[i], arg0_scalar_ty).handle, + ) + + denominator_low_range = _semantic.full( + arg0.shape, coeff_low_denominator[0], arg0_scalar_ty + ).handle + for i in range(1, len(coeff_low_denominator)): + denominator_low_range = _semantic.builder.create_fma( + denominator_low_range, + arg0_squared, + _semantic.full(arg0.shape, coeff_low_denominator[i], arg0_scalar_ty).handle, + ) + + low_res = _semantic.builder.create_fmul( + arg0.handle, + _semantic.builder.create_fdiv(numerator_low_range, denominator_low_range), + ) + + arg0_erf_trans = _semantic.builder.create_sqrt( + _semantic.builder.create_fmul( + _semantic.full(arg0.shape, -1, arg0_scalar_ty).handle, + _semantic.builder.create_log( + _semantic.builder.create_fdiv( + _semantic.builder.create_fsub( + _semantic.full(arg0.shape, 1, arg0_scalar_ty).handle, + _semantic.builder.create_fabs(arg0.handle), + ), + _semantic.full(arg0.shape, 2, arg0_scalar_ty).handle, + ) + ), + ) + ) + numerator_high_range = _semantic.full( + arg0.shape, coeff_high_numerator[0], arg0_scalar_ty + ).handle + for i in range(1, len(coeff_high_numerator)): + numerator_high_range = _semantic.builder.create_fma( + numerator_high_range, + arg0_erf_trans, + _semantic.full(arg0.shape, coeff_high_numerator[i], arg0_scalar_ty).handle, + ) + + denominator_high_range = _semantic.full( + arg0.shape, coeff_high_denominator[0], arg0_scalar_ty + ).handle + for i in range(1, len(coeff_high_denominator)): + denominator_high_range = _semantic.builder.create_fma( + denominator_high_range, + arg0_erf_trans, + _semantic.full( + arg0.shape, coeff_high_denominator[i], arg0_scalar_ty + ).handle, + ) + + high_res = _semantic.builder.create_fdiv( + numerator_high_range, denominator_high_range + ) + high_res = _semantic.mul( + _semantic.where( + signbit(arg0, _semantic=_semantic), + _semantic.full(arg0.shape, -1, arg0_scalar_ty), + _semantic.full(arg0.shape, 1, arg0_scalar_ty), + ), + core.tensor(high_res, arg0.type), + True, + ).handle + + for _ in range(2): + low_res = _semantic.builder.create_fsub( + low_res, + _semantic.builder.create_fdiv( + _semantic.builder.create_fsub( + _semantic.builder.create_erf(low_res), arg0.handle + ), + _semantic.builder.create_fmul( + inv_sqrt_pi_times_2, + _semantic.builder.create_exp( + _semantic.builder.create_fmul( + _semantic.full(arg0.shape, -1, arg0_scalar_ty).handle, + _semantic.builder.create_fmul(low_res, low_res), + ) + ), + ), + ), + ) + + high_res = _semantic.builder.create_fsub( + high_res, + _semantic.builder.create_fdiv( + _semantic.builder.create_fsub( + _semantic.builder.create_erf(high_res), arg0.handle + ), + _semantic.builder.create_fmul( + inv_sqrt_pi_times_2, + _semantic.builder.create_exp( + _semantic.builder.create_fmul( + _semantic.full(arg0.shape, -1, arg0_scalar_ty).handle, + _semantic.builder.create_fmul(high_res, high_res), + ) + ), + ), + ), + ) + + arg0_abs = core.tensor(_semantic.builder.create_fabs(arg0.handle), arg0.type) + arg0_over = _semantic.greater_than( + arg0_abs, _semantic.full(arg0.shape, 1, arg0_scalar_ty) + ) + nan_tensor = _semantic.full(arg0.shape, float("nan"), arg0_scalar_ty) + arg0_equal1 = _semantic.equal( + arg0_abs, _semantic.full(arg0.shape, 1, arg0_scalar_ty) + ) + pos_inf_tensor = _semantic.full(arg0.shape, float("inf"), arg0_scalar_ty) + neg_inf_tensor = _semantic.full(arg0.shape, float("-inf"), arg0_scalar_ty) + inf_res = _semantic.where( + signbit(arg0, _semantic=_semantic), neg_inf_tensor, pos_inf_tensor + ) + arg0_high = _semantic.greater_equal( + arg0_abs, _semantic.full(arg0.shape, 0.7, arg0_scalar_ty) + ) + + return _semantic.where( + arg0_equal1, + inf_res, + _semantic.where( + arg0_over, + nan_tensor, + _semantic.where( + arg0_high, + core.tensor(high_res, arg0.type), + core.tensor(low_res, arg0.type), + ), + ), + ) + + +@core.extern +@math._check_dtype(dtypes=["fp32"]) +def gamma(arg0, _semantic=None): + arg0_scalar_ty = arg0.type.scalar + arg0 = _semantic.to_tensor(arg0) + pi_tensor = _semantic.full(arg0.shape, math_pi, arg0_scalar_ty).handle + sqrt_2pi_tensor = _semantic.full(arg0.shape, 2.506628275, arg0_scalar_ty).handle + lanczos_coeff = [ + 676.5203681218851, + -1259.1392167224028, + 771.32342877765313, + -176.61502916214059, + 12.507343278686905, + -0.13857109526572012, + 9.9843695780195716e-6, + 1.5056327351493116e-7, + ] + condition = _semantic.less_than(arg0, 0.5) + reflect_arg0 = _semantic.where(condition, _semantic.sub(1, arg0, True), arg0) + + x = _semantic.full(arg0.shape, 0.99999999999980993, arg0_scalar_ty) + for i in range(0, len(lanczos_coeff)): + x = _semantic.add( + x, + _semantic.fdiv( + _semantic.full(arg0.shape, lanczos_coeff[i], arg0_scalar_ty), + _semantic.add(reflect_arg0, i, True), + True, + ), + True, + ) + t = _semantic.add(reflect_arg0, 6.5, True) + + gamma_res = _semantic.builder.create_fmul( + _semantic.builder.create_fmul( + sqrt_2pi_tensor, + pow(t, _semantic.sub(reflect_arg0, 0.5, True), _semantic=_semantic).handle, + ), + _semantic.builder.create_fmul( + x.handle, + _semantic.builder.create_exp( + _semantic.builder.create_fmul( + t.handle, _semantic.full(arg0.shape, -1, arg0_scalar_ty).handle + ) + ), + ), + ) + + gamma_res_reflect = _semantic.builder.create_fdiv( + _semantic.builder.create_fdiv(pi_tensor, gamma_res), + _semantic.builder.create_sin( + _semantic.builder.create_fmul(pi_tensor, arg0.handle) + ), + ) + + is_neg_int = _semantic.logical_and( + _semantic.equal(math.floor(arg0, _semantic=_semantic), arg0), + _semantic.less_than(arg0, 0), + ) + pos_inf_tensor = _semantic.full(arg0.shape, float("inf"), arg0_scalar_ty) + neg_inf_tensor = _semantic.full(arg0.shape, float("-inf"), arg0_scalar_ty) + gamma_res_reflect = _semantic.where( + is_neg_int, pos_inf_tensor, core.tensor(gamma_res_reflect, arg0.type) + ) + + res = _semantic.where( + condition, gamma_res_reflect, core.tensor(gamma_res, arg0.type) + ) + is_pos_inf_input = _semantic.equal(arg0, pos_inf_tensor) + is_neg_inf_input = _semantic.equal(arg0, neg_inf_tensor) + + return _semantic.where( + is_pos_inf_input, + pos_inf_tensor, + _semantic.where(is_neg_inf_input, neg_inf_tensor, res), + ) + + +@core.extern +@math._check_dtype(dtypes=["fp32"]) +def lgamma(arg0, _semantic=None): + arg0_scalar_ty = arg0.type.scalar + arg0 = _semantic.to_tensor(arg0) + + inf_tensor = _semantic.full(arg0.shape, float("inf"), arg0_scalar_ty) + is_inf = _semantic.equal( + core.tensor(_semantic.builder.create_fabs(arg0.handle), arg0.type), inf_tensor + ) + gamma_res = _semantic.builder.create_fabs(gamma(arg0, _semantic=_semantic).handle) + lgamma_res = _semantic.builder.create_log(gamma_res) + + return _semantic.where(is_inf, inf_tensor, core.tensor(lgamma_res, arg0.type)) + + +@core.builtin +@math._check_dtype(dtypes=["fp32"]) +@math._add_math_1arg_docstr("trunc") +def trunc(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + zero = _semantic.full(arg0.shape, 0.0, arg0.type.scalar) + condition = _semantic.greater_equal(arg0, zero) + floor_result = core.tensor(_semantic.builder.create_floor(arg0.handle), arg0.type) + ceil_result = core.tensor(_semantic.builder.create_ceil(arg0.handle), arg0.type) + return _semantic.where(condition, floor_result, ceil_result) + + +@core.builtin +@math._check_dtype(dtypes=["fp32"]) +@math._add_math_1arg_docstr("nearbyint") +def nearbyint(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + half = _semantic.full(arg0.shape, 0.5, arg0.type.scalar) + positive_adjust = _semantic.add(arg0, half, True) + negative_adjust = _semantic.sub(arg0, half, True) + positive_result = core.tensor( + _semantic.builder.create_floor(positive_adjust.handle), arg0.type + ) + negative_result = core.tensor( + _semantic.builder.create_ceil(negative_adjust.handle), arg0.type + ) + zero = _semantic.full(arg0.shape, 0.0, arg0.type.scalar) + is_positive = _semantic.greater_equal(arg0, zero) + basic_round = _semantic.where(is_positive, positive_result, negative_result) + + fractional = _semantic.sub(arg0, basic_round, True) + abs_fractional = core.tensor( + _semantic.builder.create_fabs(fractional.handle), fractional.type + ) + is_half = _semantic.equal(abs_fractional, half) + two = _semantic.full(arg0.shape, 2.0, arg0.type.scalar) + half_value = math.fdiv(basic_round, two, _semantic=_semantic) + half_floor = core.tensor( + _semantic.builder.create_floor(half_value.handle), half_value.type + ) + double_half = _semantic.mul(half_floor, two, True) + is_even = _semantic.equal(basic_round, double_half) + + adjustment = _semantic.where( + is_positive, + _semantic.full(arg0.shape, -1.0, arg0.type.scalar), + _semantic.full(arg0.shape, 1.0, arg0.type.scalar), + ) + banker_result = _semantic.where( + is_even, basic_round, _semantic.add(basic_round, adjustment, True) + ) + return _semantic.where(is_half, banker_result, basic_round) + + +@core.builtin +@math._check_dtype(dtypes=["fp32"]) +@math._add_math_1arg_docstr("arcsine") +def asin(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + half_pi = _semantic.full(arg0.shape, 1.5707963267948966, arg0.type.scalar) + acos_val = acos(arg0, _semantic=_semantic) + return _semantic.sub(half_pi, acos_val, True) + + +@core.builtin +@math._check_dtype(dtypes=["fp32"]) +@math._add_math_1arg_docstr("base-10 logarithm") +def log10(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + log_val = math.log(arg0, _semantic=_semantic) + log10_const = _semantic.full(arg0.shape, 2.302585092994046, arg0.type.scalar) + return math.fdiv(log_val, log10_const, _semantic=_semantic) + + +@core.builtin +@math._check_dtype(dtypes=["fp32"]) +@math._add_math_2arg_docstr("copysign") +def copysign(arg0: core.tensor, arg1: core.tensor, _semantic=None): + x = _semantic.to_tensor(arg0) + y = _semantic.to_tensor(arg1) + magnitude = core.tensor(_semantic.builder.create_fabs(x.handle), x.type) + zero = _semantic.full(y.shape, 0.0, y.type.scalar) + one = _semantic.full(y.shape, 1.0, y.type.scalar) + is_zero = _semantic.equal(y, zero) + y_reciprocal = math.fdiv(one, y, _semantic=_semantic) + is_negative_reciprocal = _semantic.less_than(y_reciprocal, zero) + is_negative_zero = _semantic.and_(is_zero, is_negative_reciprocal) + is_negative_nonzero = _semantic.less_than(y, zero) + is_negative = _semantic.or_(is_negative_zero, is_negative_nonzero) + neg_magnitude = _semantic.mul( + magnitude, _semantic.full(magnitude.shape, -1.0, magnitude.type.scalar), True + ) + return _semantic.where(is_negative, neg_magnitude, magnitude) + + +if get_ascend_arch_from_env() == "Ascend910_9589": + + @core.extern + def rint(arg0, _semantic=None): + return core.extern_elementwise( + "", + "", + [arg0], + { + (core.dtype("fp32"),): ("__hmf_rint", core.dtype("fp32")), + (core.dtype("fp16"),): ("__hmf_rint", core.dtype("fp16")), + (core.dtype("bf16"),): ("__hmf_rint", core.dtype("bf16")), + }, + is_pure=True, + _semantic=_semantic, + ) + +else: + + @core.builtin + @math._check_dtype(dtypes=["fp16", "fp32", "bf16"]) + @math._add_math_1arg_docstr("rint") + def rint(arg0: core.tensor, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + floor_x = math.floor(arg0, _semantic=_semantic) + fractional = _semantic.sub(arg0, floor_x, True) + half = _semantic.full(arg0.shape, 0.5, arg0.type.scalar) + eps = _semantic.full(arg0.shape, 1e-8, arg0.type.scalar) + is_half = _semantic.less_than( + math.abs(_semantic.sub(fractional, half, True), _semantic=_semantic), eps + ) + floor_int = ( + floor_x.to(core.int32, _semantic=_semantic) + if hasattr(floor_x, "to") + else _semantic.cast(floor_x, core.int32) + ) + two_i32 = _semantic.full(arg0.shape, 2, core.int32) + is_even = _semantic.equal( + _semantic.mod(floor_int, two_i32), _semantic.full(arg0.shape, 0, core.int32) + ) + zero = _semantic.full(arg0.shape, 0.0, arg0.type.scalar) + is_pos = _semantic.greater_equal(arg0, zero) + round_pos = math.floor(_semantic.add(arg0, half, True), _semantic=_semantic) + round_neg = math.ceil(_semantic.sub(arg0, half, True), _semantic=_semantic) + normal_round = _semantic.where(is_pos, round_pos, round_neg) + half_round = _semantic.where( + is_even, floor_x, _semantic.add(floor_x, 1.0, True) + ) + return _semantic.where(is_half, half_round, normal_round) + + +@core.builtin +def flip(arg0, dim=-1, _semantic=None): + arg0 = _semantic.to_tensor(arg0) + shape = arg0.shape + rank = len(shape) + if dim < 0: + dim = dim + rank + if not (0 <= dim < rank): + raise ValueError(f"flip got invalid dim={dim} for shape {tuple(shape)}") + flipped = _semantic.builder.create_flip(arg0.handle, dim) + return core.tensor(flipped, arg0.type) + + +# --------------------------------------------------------------------------- +# JIT-style functions using libdevice primitives +# --------------------------------------------------------------------------- + + +@core._tensor_member_fn +@jit +@math._add_math_1arg_docstr("isfinited") +def isfinited(x): + _is_int8_type: core.constexpr = x.dtype.is_int8() + core.static_assert( + not _is_int8_type, f"Expected dtype fp16/fp32/bf16, but got int8 or int1" + ) + _is_floating_type: core.constexpr = x.dtype.is_floating() + core.static_assert( + _is_floating_type == True, + f"Expected dtype fp16/fp32/bf16, but got {core.constexpr(x.dtype)}", + ) + nan_mask = isnan(x) + inf_mask = isinf(x) + return (~nan_mask & ~inf_mask).to(core.int1) + + +@core._tensor_member_fn +@jit +@math._add_math_1arg_docstr("finitef") +def finitef(x): + _is_int8_type: core.constexpr = x.dtype.is_int8() + core.static_assert( + not _is_int8_type, f"finitef only supports float32, but got int8 or int1" + ) + core.static_assert( + x.dtype == core.float32, + f"finitef only supports float32, but got {core.constexpr(x.dtype)}", + ) + nan_mask = isnan(x) + inf_mask = isinf(x) + return (~nan_mask & ~inf_mask).to(core.int1) diff --git a/language/deeplink/core.py b/language/deeplink/core.py index 6cb1557a..458bee0b 100644 --- a/language/deeplink/core.py +++ b/language/deeplink/core.py @@ -1,363 +1,6 @@ -import numpy as np -from triton.language import semantic as tl_semantic -from triton.language.core import ( - _tensor_member_fn, - _shape_check_impl, - _unwrap_if_constexpr, - builtin, - constexpr, - tensor, - range, - slice, -) -import builtins -from . import semantic as dl_semantic - - -def _constexpr_to_value(v): - if isinstance(v, constexpr): - return v.value - return v - - -class layout: - ASCEND = ["ND", "NZ"] - - def __init__(self, name): - name = _unwrap_if_constexpr(name) - self.name = name - assert name in layout.ASCEND, name - - def __str__(self): - return self.name - - def codegen_name(self): - return self.name - - @property - def cache_key_part(self) -> str: - """See cache_key_part() in triton.cc.""" - return self.name - - def __repr__(self): - """Output of repr needs to be an evaluatable expression""" - return f"triton.language.{self.codegen_name()}" - - -ND = layout("ND") -NZ = layout("NZ") - - -class scope: - GPU = ["fragment"] - ASCEND = ["UB", "L1", "L0A", "L0B", "L0C"] - - def __init__(self, name): - name = _unwrap_if_constexpr(name) - self.name = name - assert name in scope.ASCEND + scope.GPU, name - - def __str__(self): - return self.name - - def codegen_name(self): - return self.name - - @property - def cache_key_part(self) -> str: - """See cache_key_part() in triton.cc.""" - return self.name - - def __repr__(self): - """Output of repr needs to be an evaluatable expression""" - return f"triton.language.{self.codegen_name()}" - - -fragment = scope("fragment") -UB = scope("UB") -L1 = scope("L1") -L0A = scope("L0A") -L0B = scope("L0B") -L0C = scope("L0C") - - -def _extract_slice(sl: slice, shape: constexpr): - def constexpr_or_none_to_value(v, default: int): - if v is None: - return default - assert isinstance( - v, (constexpr, int) - ), f"slice only can be constexpr or int, got: {v}" - return _constexpr_to_value(v) - - start = constexpr_or_none_to_value(sl.start, 0) - stop = constexpr_or_none_to_value(sl.stop, _constexpr_to_value(shape)) - step = constexpr_or_none_to_value(sl.step, 1) - size = (stop - start + step - 1) // step - assert ( - start >= 0 and stop >= 0 and step >= 0 and size >= 0 - ), f"slice should be greater than 0" - return start, size, step - - -@_tensor_member_fn -@builtin -def __getitem__(self, slices, _semantic=None): - if isinstance(slices, (builtins.slice, slice, constexpr, tensor)) or slices is None: - slices = [slices] - if isinstance(slices, tuple): - slices = slices.values - ret = self - offsets = [] - sizes = [] - strides = [] - dst_shape = [] - need_extract_slice = False - for dim, sl in enumerate(slices): - if sl is None or isinstance(sl, constexpr) and sl.value is None: - ret = _semantic.expand_dims(ret, dim) - offsets.append(_semantic.builder.get_int32(0)) - dst_shape.append(constexpr(1)) - sizes.append(constexpr(1)) - strides.append(constexpr(1)) - elif ( - isinstance(sl, slice) - and sl.start is None - and sl.stop is None - and sl.step is None - ): - pass - elif sl is None or isinstance(sl, (constexpr, int)) and sl.value is not None: - offsets.append(_semantic.builder.get_int32(_constexpr_to_value(sl))) - need_extract_slice = True - sizes.append(constexpr(1)) - strides.append(constexpr(1)) - elif isinstance(sl, tensor): - offsets.append(sl.handle) - sizes.append(constexpr(1)) - strides.append(constexpr(1)) - need_extract_slice = True - elif isinstance(sl, (slice, builtins.slice)): - start, size, step = _extract_slice(sl, ret.shape[dim]) - offsets.append(start) - strides.append(constexpr(step)) - sizes.append(constexpr(size)) - dst_shape.append(constexpr(size)) - need_extract_slice = True - else: - raise ValueError(f"unsupported tensor index: {sl}") - - if need_extract_slice: - new_offsets = [ - (_semantic.to_tensor(o) if not isinstance(o, tensor) else o) - for o in offsets - ] - ret = dl_semantic.extract_slice( - self, new_offsets, sizes, strides, _semantic=_semantic - ) - return ret - - -@builtin -def insert_slice( - ful, sub, offsets, sizes, strides, _builder=None, _generator=None, _semantic=None -) -> tensor: - """ - Insert a tensor to another tensor as specified by the operation’s offsets, sizes and strides arguments. - - :param ful: The tensor to receive tensor. - :type ful: Tensor - :param sub: The tensor to be inserted. - :type sub: Tensor - :param offsets: - :type offsets: tuple of ints - :param sizes: - :type sizes: tuple of ints - :param strides: - :type strides: tuple of ints - """ - assert len(ful.shape) > 0 - assert len(ful.shape) == len(sub.shape) - new_offsets = [ - _semantic.to_tensor(o) if isinstance(o, constexpr) else o for o in offsets - ] - out = dl_semantic.insert_slice( - ful, sub, new_offsets, sizes, strides, _semantic=_semantic - ) - return out - - -@builtin -def extract_slice( - ful, offsets, sizes, strides, _generator=None, _semantic=None -) -> tensor: - """ - Extract a tensor from another tensor as specified by the operation’s offsets, sizes and strides arguments. - - :param ful: The tensor to split. - :type ful: Tensor - :param offsets: - :type offsets: tuple of ints - :param sizes: - :type sizes: tuple of ints - :param strides: - :type strides: tuple of ints - """ - assert len(ful.shape) > 0 - new_offsets = [ - _semantic.to_tensor(o) if isinstance(o, constexpr) else o for o in offsets - ] - sub = dl_semantic.extract_slice( - ful, new_offsets, sizes, strides, _semantic=_semantic - ) - return sub - - -@builtin -def compile_hint(ptr, hint_name, hint_val=None, _semantic=None): - hint_name = _constexpr_to_value(hint_name) - assert isinstance(hint_name, str), f"hint name: {hint_name} is not string" - hint_val = _unwrap_if_constexpr(hint_val) if hint_val else hint_val - dl_semantic.compile_hint(ptr, hint_name, hint_val, _semantic.builder) - - -@builtin -def alloc(shape, value, dtype, layout=None, scope=None, _builder=None): - """ - Returns a tensor filled with the scalar value for the given :code:`shape` and :code:`dtype`. - - :param shape: Shape of the new array, e.g., (8, 16) or (8, ) - :type shape: tuple of ints - :param value: A scalar value to fill the array with - :type value: scalar - :param dtype: Data type of the new array, e.g., :code:`tl.float16` - :type dtype: tl.dtype - """ - shape = _shape_check_impl(shape) - value = _constexpr_to_value(value) - dtype = _constexpr_to_value(dtype) - layout = _constexpr_to_value(layout) - scope = _constexpr_to_value(scope) - return dl_semantic.alloc(shape, value, dtype, layout, scope, _builder) - - -@builtin -def multibuffer(src: tensor, size, _semantic=None): - """ - Set multi_buffer for an existing tensor - :src: tensor set to bufferize multiple time - :size: number of copies - """ - buffer_size = _constexpr_to_value(size) - assert ( - isinstance(buffer_size, int) and buffer_size == 2 - ), f"only support bufferize equals 2" - dl_semantic.compile_hint(src, "multi_buffer", buffer_size, _semantic.builder) - - -@builtin -def sync_block_all(mode, event_id, _builder=None): - mode = _constexpr_to_value(mode) - event_id = _constexpr_to_value(event_id) - assert isinstance(mode, str), f"mode: {mode} is not string" - assert ( - isinstance(event_id, int) and (event_id >= 0) and (event_id < 16) - ), f"event_id: {event_id} should be 0 ~ 15" - assert ( - mode == "all_cube" or mode == "all_vector" or mode == "all" - ), f"ERROR: mode = {mode}, only supports all_cube/all_vector/all" - dl_semantic.custom_sync_op(_builder, "sync_block_all", mode=mode, event_id=event_id) - - -class SyncFlagType: - ASCEND = ["cube_to_vector", "vector_to_cube"] - - def __init__(self, name): - name = _unwrap_if_constexpr(name) - self.name = name - assert name in SyncFlagType.ASCEND, name - - def __str__(self): - return self.name - - def codegen_name(self): - return self.name - - def sender(self): - if self.name == "cube_to_vector": - return "cube" - elif self.name == "vector_to_cube": - return "vector" - else: - assert self.name in SyncFlagType.ASCEND - - @property - def cache_key_part(self) -> str: - """See cache_key_part() in triton.cc.""" - return self.name - - def __repr__(self): - """Output of repr needs to be an evaluatable expression""" - return f"triton.language.{self.codegen_name()}" - - -class SyncFlag: - C2V = SyncFlagType("cube_to_vector") - V2C = SyncFlagType("vector_to_cube") - - -@builtin -def set_cross_flag(sync_flag_type: SyncFlagType, event_id: int, _semantic=None): - sender = _constexpr_to_value(sync_flag_type.sender()) - event_id = _constexpr_to_value(event_id) - assert ( - isinstance(event_id, int) and (event_id >= 0) and (event_id < 16) - ), f"event_id: {event_id} should be 0 ~ 15" - dl_semantic.custom_sync_op( - _semantic.builder, "sync_block_set", sender=sender, event_id=event_id - ) - - -@builtin -def wait_cross_flag(sync_flag_type: SyncFlagType, event_id: int, _semantic=None): - sender = _constexpr_to_value(sync_flag_type.sender()) - event_id = _constexpr_to_value(event_id) - assert ( - isinstance(event_id, int) and (event_id >= 0) and (event_id < 16) - ), f"event_id: {event_id} should be 0 ~ 15" - dl_semantic.custom_sync_op( - _semantic.builder, "sync_block_wait", sender=sender, event_id=event_id - ) - - -class parallel(range): - """ - Iterator that counts upward forever, with parallel execution semantics. - - This is a special iterator used to implement similar semantics to Python's :code:`range` in the context of - :code:`triton.jit` functions. In addition, it allows user to pass extra attributes to the compiler. - :param bind_sub_block: Tells the compiler if multiple vector cores participate in the loop. - This is used in the mixed cube-vector kernel on 910B. The number of vector cores is determined by the number of - iteration in this loop. Currently on 910B, max 2 vector cores could be used. - """ - - def __init__( - self, - arg1, - arg2=None, - step=None, - num_stages=None, - loop_unroll_factor=None, - bind_sub_block: bool = False, - ): - super().__init__(arg1, arg2, step, num_stages, loop_unroll_factor) - self.bind_sub_block = bind_sub_block - - class inline_lambda: """ Inline a lambda function into the current block. - This class is used to inline a lambda function into the current block. """ def __init__(self, node, closure_values, closure_names, arg_names): @@ -370,29 +13,23 @@ def __call__(self, *args, generator=None): if generator is None: raise RuntimeError("Generator must be provided for Lambda inlining") - # save old state old_lscope = generator.lscope.copy() old_local_defs = generator.local_defs.copy() old_insert_block = generator.builder.get_insertion_block() try: - # create closure parameters map closure_map = {} for name, value in zip(self.closure_names, self.closure_values): closure_map[name] = value - # create parameter map param_map = {} for name, value in zip(self.arg_names, args): param_map[name] = value - # merge closure and parameter maps generator.lscope = {**old_lscope, **closure_map, **param_map} generator.local_defs = {**old_local_defs, **closure_map, **param_map} - # visit the lambda body return generator.visit(self.node.body) finally: - # restore old state generator.lscope = old_lscope generator.local_defs = old_local_defs if old_insert_block: diff --git a/language/deeplink/extension.py b/language/deeplink/extension.py new file mode 100644 index 00000000..da989a78 --- /dev/null +++ b/language/deeplink/extension.py @@ -0,0 +1,7 @@ +# Backward compatibility: deeplink.extension is now a thin wrapper around cann.extension. +# The canonical Ascend extension APIs have moved to deeplink.cann.extension to align +# with triton-ascend's directory structure. +from .cann.extension import * +from .cann.extension import __all__ as _cann_all + +__all__ = _cann_all diff --git a/language/deeplink/libdevice.py b/language/deeplink/libdevice.py deleted file mode 100644 index 5516b5c5..00000000 --- a/language/deeplink/libdevice.py +++ /dev/null @@ -1,294 +0,0 @@ -from triton.language import core - - -@core.extern -def reciprocal(arg0, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("fp32"),): ("__hmf_recipf", core.dtype("fp32")), - (core.dtype("fp16"),): ("__hmf_recipDh", core.dtype("fp16")), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def log1p(arg0, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("fp32"),): ("__hmf_log1pf", core.dtype("fp32")), - (core.dtype("fp16"),): ("__hmf_log1pDh", core.dtype("fp16")), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def relu(arg0, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("fp32"),): ("__hmf_reluf", core.dtype("fp32")), - (core.dtype("fp16"),): ("__hmf_reluDh", core.dtype("fp16")), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def isinf(arg0, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("fp32"),): ("__hmf_isinf", core.dtype("int1")), - (core.dtype("fp16"),): ("__hmf_isinf", core.dtype("int1")), - (core.dtype("bf16"),): ("__hmf_isinf", core.dtype("int1")), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def tan(arg0, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("fp32"),): ("__hmf_tanf", core.dtype("fp32")), - (core.dtype("fp16"),): ("__hmf_tanDh", core.dtype("fp16")), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def atan(arg0, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("fp32"),): ("__hmf_atanf", core.dtype("fp32")), - (core.dtype("fp16"),): ("__hmf_atanDh", core.dtype("fp16")), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def tanh(arg0, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("fp32"),): ("__hmf_tanhf", core.dtype("fp32")), - (core.dtype("fp16"),): ("__hmf_tanhDh", core.dtype("fp16")), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def ilogb(arg0, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("fp32"),): ("__hmf_ilogbf", core.dtype("fp32")), - (core.dtype("fp16"),): ("__hmf_ilogbDh", core.dtype("fp16")), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def ldexp(arg0, arg1, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0, arg1], - { - (core.dtype("fp32"), core.dtype("fp32")): ( - "__hmf_ldexpf", - core.dtype("fp32"), - ), - (core.dtype("fp16"), core.dtype("fp16")): ( - "__hmf_ldexpDh", - core.dtype("fp16"), - ), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def pow(arg0, arg1, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0, arg1], - { - (core.dtype("fp32"), core.dtype("fp32")): ( - "__hmf_powf", - core.dtype("fp32"), - ), - (core.dtype("fp16"), core.dtype("fp16")): ( - "__hmf_powf", - core.dtype("fp16"), - ), - (core.dtype("bf16"), core.dtype("bf16")): ( - "__hmf_powf", - core.dtype("bf16"), - ), - (core.dtype("fp16"), core.dtype("fp32")): ( - "__hmf_powf", - core.dtype("fp16"), - ), - (core.dtype("bf16"), core.dtype("fp32")): ( - "__hmf_powf", - core.dtype("bf16"), - ), - (core.dtype("fp32"), core.dtype("fp16")): ( - "__hmf_powf", - core.dtype("fp32"), - ), - (core.dtype("fp32"), core.dtype("bf16")): ( - "__hmf_powf", - core.dtype("fp32"), - ), - (core.dtype("int64"), core.dtype("int64")): ( - "__hmf_powi", - core.dtype("int64"), - ), - (core.dtype("int32"), core.dtype("int32")): ( - "__hmf_powi", - core.dtype("int32"), - ), - (core.dtype("int16"), core.dtype("int16")): ( - "__hmf_powi", - core.dtype("int16"), - ), - (core.dtype("int8"), core.dtype("int8")): ( - "__hmf_powi", - core.dtype("int8"), - ), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def isnan(arg0, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("fp32"),): ("__hmf_isnan", core.dtype("int1")), - (core.dtype("fp16"),): ("__hmf_isnan", core.dtype("int1")), - (core.dtype("bf16"),): ("__hmf_isnan", core.dtype("int1")), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def flip(arg0, arg1=None, _semantic=None): - if arg1 == None: - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("bf16"),): ("__hmf_flipDhb", core.dtype("bf16")), - (core.dtype("fp16"),): ("__hmf_flipDh", core.dtype("fp16")), - (core.dtype("fp32"),): ("__hmf_flipf", core.dtype("fp32")), - (core.dtype("int8"),): ("__hmf_flipi8", core.dtype("int8")), - (core.dtype("int16"),): ("__hmf_flipi16", core.dtype("int16")), - (core.dtype("int32"),): ("__hmf_flipi32", core.dtype("int32")), - (core.dtype("uint32"),): ("__hmf_flipui32", core.dtype("uint32")), - (core.dtype("int64"),): ("__hmf_flipi64", core.dtype("int64")), - }, - is_pure=True, - _semantic=_semantic, - ) - - return core.extern_elementwise( - "", - "", - [arg0, arg1], - { - (core.dtype("bf16"), core.dtype("int32")): ( - "__hmf_flipDhb", - core.dtype("bf16"), - ), - (core.dtype("fp16"), core.dtype("int32")): ( - "__hmf_flipDh", - core.dtype("fp16"), - ), - (core.dtype("fp32"), core.dtype("int32")): ( - "__hmf_flipf", - core.dtype("fp32"), - ), - (core.dtype("int8"), core.dtype("int32")): ( - "__hmf_flipi8", - core.dtype("int8"), - ), - (core.dtype("int16"), core.dtype("int32")): ( - "__hmf_flipi16", - core.dtype("int16"), - ), - (core.dtype("int32"), core.dtype("int32")): ( - "__hmf_flipi32", - core.dtype("int32"), - ), - (core.dtype("uint32"), core.dtype("int32")): ( - "__hmf_flipui32", - core.dtype("uint32"), - ), - (core.dtype("int64"), core.dtype("int32")): ( - "__hmf_flipi64", - core.dtype("int64"), - ), - }, - is_pure=True, - _semantic=_semantic, - ) - - -@core.extern -def round(arg0, _semantic=None): - return core.extern_elementwise( - "", - "", - [arg0], - { - (core.dtype("fp32"),): ("__hmf_roundf", core.dtype("fp32")), - }, - is_pure=True, - _semantic=_semantic, - ) diff --git a/language/deeplink/runtime/__init__.py b/language/deeplink/runtime/__init__.py new file mode 100644 index 00000000..c7cf39b8 --- /dev/null +++ b/language/deeplink/runtime/__init__.py @@ -0,0 +1,3 @@ +from .libentry import LibEntry, LibTuner, libentry + +__all__ = ["LibEntry", "LibTuner", "libentry"] diff --git a/language/deeplink/runtime/code_cache.py b/language/deeplink/runtime/code_cache.py new file mode 100644 index 00000000..563d46c8 --- /dev/null +++ b/language/deeplink/runtime/code_cache.py @@ -0,0 +1,68 @@ +# Copyright (c) FlagOpen contributors +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# Copyright © 2024 BAAI. 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. + +# Modifications: +# - 2025-06-03: +# - init version: e9c7aa71832eb2f897a49ce787e42d5377404a72 +# + +import functools +import os +import shutil +from pathlib import Path + + +@functools.lru_cache(maxsize=None) # this is the same as functools.cache in Python 3.9+ +def cache_dir_path() -> Path: + """Return the cache directory for generated files in flaggems.""" + _cache_dir = os.environ.get("FLAGGEMS_CACHE_DIR") + if _cache_dir is None: + _cache_dir = Path.home() / ".flaggems" + else: + _cache_dir = Path(_cache_dir) + return _cache_dir + + +def cache_dir() -> Path: + """Return cache directory for generated files in flaggems. Create it if it does not exist.""" + _cache_dir = cache_dir_path() + os.makedirs(_cache_dir, exist_ok=True) + return _cache_dir + + +def code_cache_dir() -> Path: + _code_cache_dir = cache_dir() / "code_cache" + os.makedirs(_code_cache_dir, exist_ok=True) + return _code_cache_dir + + +def config_cache_dir() -> Path: + _config_cache_dir = cache_dir() / "config_cache" + os.makedirs(_config_cache_dir, exist_ok=True) + return _config_cache_dir + + +def clear_cache(): + """Clear the cache directory for code cache.""" + _cache_dir = cache_dir_path() + shutil.rmtree(_cache_dir) diff --git a/language/deeplink/runtime/libentry.py b/language/deeplink/runtime/libentry.py new file mode 100644 index 00000000..7391a047 --- /dev/null +++ b/language/deeplink/runtime/libentry.py @@ -0,0 +1,403 @@ +# Copyright 2018-2020 Philippe Tillet +# Copyright 2020-2022 OpenAI +# Copyright © 2024 BAAI. 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. + +# Modifications: +# - aligned with triton-ascend libentry for DLCompiler package integration. + +import ast +import inspect +import sqlite3 +import threading +import weakref +from collections import OrderedDict +from typing import Dict, Optional + +import torch +import triton +from language.deeplink.runtime.code_cache import config_cache_dir + + +# Prefer NPU device function in DLCompiler runtime; keep CUDA as fallback for envs +# without NPU extension. +torch_device_fn = getattr(torch, "npu", torch.cuda) +DEVICE_COUNT = ( + torch_device_fn.device_count() if hasattr(torch_device_fn, "device_count") else 0 +) +version = triton.__version__.split(".") +major_version, minor_version = eval(version[0]), eval(version[1]) + + +def quote_identifier(name: str) -> str: + if not name: + raise ValueError("empty identifier") + allowed = set("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789_") + if not (name[0].isalpha() or name[0] == "_"): + raise ValueError("identifier must start with letter or _") + if not all(ch in allowed for ch in name): + raise ValueError("identifier contains illegal char") + return '"' + name.replace('"', '""') + '"' + + +class LibTuner(triton.runtime.Autotuner): + def __init__( + self, + fn, + arg_names, + configs, + key, + reset_to_zero, + restore_value, + pre_hook=None, + post_hook=None, + prune_configs_by: Optional[Dict] = None, + warmup=None, + rep=None, + use_cuda_graph=False, + ): + if major_version == 2 or (major_version == 3 and minor_version <= 1): + if warmup is None: + warmup = 25 + if rep is None: + rep = 100 + self.base_fn = fn + while not inspect.isfunction(self.base_fn): + self.base_fn = self.base_fn.fn + if major_version == 2: + super().__init__( + fn, + arg_names, + configs, + key, + reset_to_zero, + restore_value, + prune_configs_by, + warmup, + rep, + ) + else: + super().__init__( + fn, + arg_names, + configs, + key, + reset_to_zero, + restore_value, + pre_hook, + post_hook, + prune_configs_by, + warmup, + rep, + use_cuda_graph, + ) + self.__name__ = self.base_fn.__name__ + self.table_name = quote_identifier(self.__name__) + self.cache_path = config_cache_dir() / "TunedConfig.db" + self.preload() + weakref.finalize(self, self.store) + + def preload(self): + connect = sqlite3.connect(self.cache_path) + c = connect.cursor() + c.execute( + f"CREATE TABLE IF NOT EXISTS {self.table_name} (key TEXT PRIMARY KEY, config TEXT)" + ) + cursor = c.execute(f"SELECT key, config from {self.table_name}") + + for row in cursor: + key_str, config_str = row + key = [ast.literal_eval(k) for k in key_str[1:-1].split(", ")] + + cfg_ls = [item.split(": ") for item in config_str.split(", ")] + config = triton.Config({}) + attrs = -5 if major_version == 2 else -4 + for k, v in cfg_ls[:attrs]: + config.kwargs[k] = ast.literal_eval(v) + config.num_warps = ast.literal_eval(cfg_ls[attrs][1]) + config.num_ctas = ast.literal_eval(cfg_ls[attrs + 1][1]) + config.num_stages = ast.literal_eval(cfg_ls[attrs + 2][1]) + if major_version == 2: + config.enable_warp_specialization = ast.literal_eval( + cfg_ls[attrs + 3][1] + ) + config.enable_persistent = ast.literal_eval(cfg_ls[attrs + 4][1]) + else: + config.maxnreg = ast.literal_eval(cfg_ls[attrs + 3][1]) + + self.cache[tuple(key)] = config + + connect.close() + self.volumn = len(self.cache) + + def store(self): + if len(self.cache) == self.volumn: + return + connect = sqlite3.connect(self.cache_path) + c = connect.cursor() + c.execute( + f"CREATE TABLE IF NOT EXISTS {self.table_name} (key TEXT PRIMARY KEY, config TEXT)" + ) + for key, config in self.cache.items(): + c.execute( + f"INSERT OR IGNORE INTO {self.table_name} (key, config) VALUES (?, ?)", + (str(key), config.__str__()), + ) + + connect.commit() + connect.close() + + +def libtuner( + configs, + key, + prune_configs_by=None, + reset_to_zero=None, + restore_value=None, + pre_hook=None, + post_hook=None, + warmup=25, + rep=100, + use_cuda_graph=False, +): + """ + Decorator for triton library autotuner. + """ + + def decorator(fn): + return LibTuner( + fn, + fn.arg_names, + configs, + key, + reset_to_zero, + restore_value, + pre_hook=pre_hook, + post_hook=post_hook, + prune_configs_by=prune_configs_by, + warmup=warmup, + rep=rep, + use_cuda_graph=use_cuda_graph, + ) + + return decorator + + +class LibEntry(triton.KernelInterface): + def __init__( + self, + fn, + ): + self.fn = fn + self.arg_names = fn.arg_names + self.divisibility = 16 + self.kernel_cache = tuple(dict() for _ in range(DEVICE_COUNT)) + + while not isinstance(fn, triton.runtime.JITFunction): + fn = fn.fn + self.jit_function: triton.runtime.JITFunction = fn + self.specialize_indices = [ + p.num + for p in self.jit_function.params + if not p.is_constexpr and not p.do_not_specialize + ] + self.do_not_specialize_indices = [ + p.num + for p in self.jit_function.params + if not p.is_constexpr and p.do_not_specialize + ] + self.lock = threading.Lock() + self.signature = fn.signature + + def key(self, spec_args, dns_args, const_args): + def spec_arg(arg): + if hasattr(arg, "data_ptr"): + return (arg.dtype, arg.data_ptr() % self.divisibility == 0) + return (type(arg), arg) + + def dns_arg(arg): + if hasattr(arg, "data_ptr"): + return arg.dtype + if not isinstance(arg, int): + return type(arg) + if -(2**31) <= arg and arg <= 2**31 - 1: + return "i32" + if 2**63 <= arg and arg <= 2**64 - 1: + return "u64" + return "i64" + + spec_key = [spec_arg(arg) for arg in spec_args] + dns_key = [dns_arg(arg) for arg in dns_args] + # const args passed by position + return tuple(spec_key + dns_key + const_args) + + def run(self, *args, **kwargs): + grid = kwargs["grid"] + + # collect all the arguments + spec_args = [] # specialize arguments + dns_args = [] # do not specialize arguments + const_args = [] # constexpr arguments + k_args = OrderedDict() + param_names = list(self.signature.parameters.keys()) + for i, arg in enumerate(args): + hashable_arg = arg + if ( + hasattr(arg, "__class__") + and arg.__class__.__name__ == "TensorDescriptor" + ): + # Create a hashable representation of TensorDescriptor + hashable_arg = ( + "TensorDescriptor", + tuple(arg.shape) if hasattr(arg, "shape") else None, + tuple(arg.strides) if hasattr(arg, "strides") else None, + tuple(arg.block_shape) if hasattr(arg, "block_shape") else None, + arg.padding if hasattr(arg, "padding") else None, + # Add other relevant attributes + ) + if i in self.specialize_indices: + k_args[param_names[i]] = arg + spec_args.append(hashable_arg) + elif i in self.do_not_specialize_indices: + k_args[param_names[i]] = arg + dns_args.append(hashable_arg) + else: + if major_version == 3 and 3 <= minor_version <= 6: + k_args[param_names[i]] = arg + const_args.append(hashable_arg) + for p in self.jit_function.params[len(args) :]: + if p.name in kwargs: + val = kwargs[p.name] + elif p.default is inspect._empty: + continue + else: + val = p.default + + if p.is_constexpr: + const_args.append(val) + if major_version == 3 and 3 <= minor_version <= 6: + k_args[p.name] = val + elif p.do_not_specialize: + dns_args.append(val) + k_args[p.name] = val + else: + spec_args.append(val) + k_args[p.name] = val + + entry_key = self.key(spec_args, dns_args, const_args) + device = torch_device_fn.current_device() + cache = self.kernel_cache[device] + while entry_key not in cache: + # NOTE: we serialize the first run of a jit function regardless of which device to run on + # because Triton runtime is currently not threadsafe. + with self.lock: + if entry_key in cache: + break + kernel = self.fn.run(*args, **kwargs) + fn = self.fn + # collect constexpr arguments for grid computation + constexprs = {} + tune_constexprs = {} + heur_constexprs = {} + while not isinstance(fn, triton.runtime.JITFunction): + if isinstance(fn, triton.runtime.Autotuner): + config = fn.best_config + constexprs["num_warps"] = config.num_warps + constexprs["num_stages"] = config.num_stages + constexprs["num_ctas"] = config.num_ctas + constexprs = {**constexprs, **config.kwargs} + tune_constexprs = {**tune_constexprs, **config.kwargs} + elif isinstance(fn, triton.runtime.Heuristics): + for v, heur in fn.values.items(): + heur_constexprs[v] = heur( + { + **dict(zip(fn.arg_names, args)), + **kwargs, + **constexprs, + } + ) + constexprs[v] = heur_constexprs[v] + else: + raise RuntimeError("Invalid Runtime Function") + fn = fn.fn + for p in self.jit_function.params: + if ( + p.is_constexpr + and p.name not in constexprs + and (p.default is not inspect._empty) + ): + constexprs[p.name] = p.default + cache[entry_key] = ( + kernel, + constexprs, + tune_constexprs, + heur_constexprs, + ) + return kernel, constexprs + + kernel, constexprs, tune_constexprs, heur_constexprs = cache[entry_key] + + if callable(grid): + # collect all arguments to the grid fn, ie: + # 1. args, + # 2. kwargs, + # 3. all all other captured arguments in CompiledKernel from Autotunner & Heuristics + # when kwargs & captured args conflict, captured args have higher priority + meta = {**dict(zip(self.arg_names, args)), **kwargs, **constexprs} + grid = grid(meta) + grid = grid + (1, 1) + + if major_version == 3 and 3 <= minor_version <= 6: + all_args = [] + missing_keys = [] + for key in list(self.signature.parameters.keys()): + if key in k_args: + all_args.append(k_args[key]) + elif key in tune_constexprs: + all_args.append(tune_constexprs[key]) + elif key in heur_constexprs: + all_args.append(heur_constexprs[key]) + elif key in constexprs: + all_args.append(constexprs[key]) + else: + missing_keys.append(key) + if len(missing_keys): + raise RuntimeError( + f"[libentry]: probably a bug, the following kernel params where not captured: {missing_keys}" + ) + kernel[grid[0:3]](*all_args) + else: + kernel[grid[0:3]](*k_args.values()) + return kernel, constexprs + + +def libentry(): + """ + Decorator for triton library entries. + """ + + def decorator(fn): + from triton.runtime.interpreter import InterpretedFunction + + if isinstance(fn, InterpretedFunction): + return fn + return LibEntry(fn) + + return decorator diff --git a/language/deeplink/semantic.py b/language/deeplink/semantic.py deleted file mode 100644 index 9516b2e6..00000000 --- a/language/deeplink/semantic.py +++ /dev/null @@ -1,107 +0,0 @@ -from typing import List -from triton.language import core as tl -from triton.language import semantic as tl_semantic -from triton._C.libtriton import ir - - -def insert_slice( - ful: tl.tensor, - sub: tl.tensor, - offsets: List[tl.tensor], - sizes: List[int], - strides: List[int], - _semantic: tl_semantic.TritonSemantic, -) -> tl.tensor: - assert len(ful.shape) == len(offsets) - assert len(ful.shape) == len(sizes) - assert len(ful.shape) == len(strides) - assert all([s >= 1 for s in sizes]) - assert all([s >= 0 for s in strides]) - new_offsets = [o.handle for o in offsets] - ret_type = tl.block_type(ful.type.scalar, ful.shape) - out = _semantic.builder.create_insert_slice( - ful.handle, sub.handle, new_offsets, sizes, strides - ) - return tl.tensor(out, ret_type) - - -def extract_slice( - ful: tl.tensor, - offsets: List[tl.tensor], - sizes: List[int], - strides: List[int], - _semantic: tl_semantic.TritonSemantic, -) -> tl.tensor: - assert len(ful.shape) == len(offsets) - assert len(ful.shape) == len(sizes) - assert len(ful.shape) == len(strides) - assert all([s >= 1 for s in sizes]) - assert all([s >= 0 for s in strides]) - new_offsets = [o.handle for o in offsets] - ret_type = tl.block_type(ful.type.scalar, sizes) - out = _semantic.builder.create_extract_slice( - ful.handle, new_offsets, sizes, strides - ) - return tl.tensor(out, ret_type) - - -def compile_hint(ptr: tl.tensor, hint_name: str, hint_val, builder: ir.builder): - if not hint_val: - hint_val = builder.get_unit_attr() - elif isinstance(hint_val, bool): - hint_val = builder.get_bool_attr(hint_val) - elif isinstance(hint_val, int): - hint_val = builder.get_int32_attr(hint_val) - else: - raise ValueError(f"Unsupported hint value type: {type(hint_val)}") - builder.create_annotation(ptr.handle, hint_name, hint_val) - - -def alloc( - shape: List[int], value, dtype: tl.dtype, layout, scope, builder: ir.builder -) -> tl.tensor: - if isinstance(value, tl.tensor): - assert value.numel.value == 1, "only accepts size-1 tensor" - value = tl_semantic.cast(value, dtype, builder) - else: - # scalar - if dtype is None: - raise ValueError("dtype must be specified when value is not a tensor") - if value == 0: - value = builder.get_null_value(dtype.to_ir(builder)) - else: - get_value_fn = getattr(builder, f"get_{dtype.name}") - value = get_value_fn(value) - value = tl.tensor(value, dtype) - if len(shape) == 0: - return value - ret_ty = tl.block_type(value.dtype, shape) - x = tl.tensor(builder.create_splat(value.handle, shape), ret_ty) - if layout is not None: - builder.create_annotation( - x.handle, "layout", builder.get_string_attr(str(layout)) - ) - if scope is not None: - builder.create_annotation( - x.handle, "scope", builder.get_string_attr(str(scope)) - ) - return x - - -def custom_sync_op(builder: ir.builder, op_name: str, **kwargs): - if op_name == "sync_block_all": - return builder.create_custom_op_for_inter_core_sync( - op_name, kwargs["mode"], kwargs["event_id"] - ) - - elif op_name == "sync_block_set": - return builder.create_custom_op_for_inter_core_sync( - op_name, kwargs["sender"], kwargs["event_id"] - ) - - elif op_name == "sync_block_wait": - return builder.create_custom_op_for_inter_core_sync( - op_name, kwargs["sender"], kwargs["event_id"] - ) - - raise ValueError(f"Unsupported custom op: {op_name}") diff --git a/patch/ascendnpu-ir.patch b/patch/ascendnpu-ir.patch deleted file mode 100644 index 2b428979..00000000 --- a/patch/ascendnpu-ir.patch +++ /dev/null @@ -1,789 +0,0 @@ -diff --git a/bishengir/include/bishengir/Dialect/HFusion/IR/HFusion.h b/bishengir/include/bishengir/Dialect/HFusion/IR/HFusion.h -index 01814d0..dde227f 100644 ---- a/bishengir/include/bishengir/Dialect/HFusion/IR/HFusion.h -+++ b/bishengir/include/bishengir/Dialect/HFusion/IR/HFusion.h -@@ -18,7 +18,7 @@ - #ifndef BISHENGIR_DIALECT_HFUSION_IR_HFUSION_H - #define BISHENGIR_DIALECT_HFUSION_IR_HFUSION_H - --#include "mlir/Dialect/Mesh/IR/MeshDialect.h" -+#include "mlir/Dialect/Shard/IR/ShardDialect.h" - #include "bishengir/Dialect/Symbol/IR/Symbol.h" - #include "bishengir/Interfaces/AggregatedOpInterface.h" - #include "mlir/Bytecode/BytecodeOpInterface.h" -diff --git a/bishengir/include/bishengir/Dialect/HFusion/IR/HFusionBase.td b/bishengir/include/bishengir/Dialect/HFusion/IR/HFusionBase.td -index 4257dc2..15f6b2c 100644 ---- a/bishengir/include/bishengir/Dialect/HFusion/IR/HFusionBase.td -+++ b/bishengir/include/bishengir/Dialect/HFusion/IR/HFusionBase.td -@@ -39,7 +39,7 @@ def HFusion_Dialect : Dialect { - #endif - "linalg::LinalgDialect", - "mathExt::MathExtDialect", -- "mesh::MeshDialect", -+ "shard::ShardDialect", - "symbol::SymbolDialect" - ]; - let hasCanonicalizer = 1; -diff --git a/bishengir/include/bishengir/Dialect/HFusion/IR/HFusionStructuredOps.td b/bishengir/include/bishengir/Dialect/HFusion/IR/HFusionStructuredOps.td -index ce55e7f..089f75e 100644 ---- a/bishengir/include/bishengir/Dialect/HFusion/IR/HFusionStructuredOps.td -+++ b/bishengir/include/bishengir/Dialect/HFusion/IR/HFusionStructuredOps.td -@@ -114,8 +114,11 @@ def ReduceWithIndexOp : HFusionStructuredBase_Op<"reduce_with_index", - // Declare functions necessary for LinalgStructuredInterface. - SmallVector getIteratorTypesArray(); - ArrayAttr getIndexingMaps(); -+ // static std::function)> - static std::function)> -+ mlir::ArrayRef, -+ function_ref)> - getRegionBuilder(); - }]; - } -@@ -154,8 +157,11 @@ def ArangeOp : HFusionStructuredBase_Op<"arange", [AttrSizedOperandSegments, - // Declare functions necessary for LinalgStructuredInterface. - SmallVector getIteratorTypesArray(); - ArrayAttr getIndexingMaps(); -+ // static std::function)> - static std::function)> -+ mlir::ArrayRef, -+ function_ref)> - getRegionBuilder(); - /// Precondition: `val` must be of type ShapedType - static void getStridesFromValue(OpBuilder & builder, Location loc, -@@ -226,8 +232,11 @@ def GatherOp - // Declare functions necessary for LinalgStructuredInterface. - SmallVector getIteratorTypesArray(); - ArrayAttr getIndexingMaps(); -+ // static std::function)> - static std::function)> -+ mlir::ArrayRef, -+ function_ref)> - getRegionBuilder(); - - // Used for AggregateOpInterface to decompose into legal operations -diff --git a/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMAttrs.td b/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMAttrs.td -index f245841..47db011 100644 ---- a/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMAttrs.td -+++ b/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMAttrs.td -@@ -783,6 +783,23 @@ def HIVM_StorageAligned : HIVM_Attr<"StorageAligned", "storage_aligned"> { - // Misc. - //===----------------------------------------------------------------------===// - -+//SH -+def HIVM_VF_SIMD : I32EnumAttrCase<"SIMD", 0>; -+def HIVM_VF_SIMT : I32EnumAttrCase<"SIMT", 1>; -+def HIVM_VF_MIX : I32EnumAttrCase<"MIX", 2>; -+ -+def HIVM_VFModeEnum -+ : HIVM_I32Enum<"VFMode", -+ "HIVM VF Mode", [HIVM_VF_SIMD, HIVM_VF_SIMT, HIVM_VF_MIX]>; -+ -+def HIVM_VFModeAttr : HIVM_I32EnumAttr<"vf_mode", HIVM_VFModeEnum> { -+ let description = [{ -+ HIVM VF mode attribute. -+ }]; -+} -+//SH -+ -+ - def HIVM_MultiBufferAttr : HIVM_Attr<"MultiBuffer", "multi_buffer"> { - let description = [{ - HIVM multi-buffer attribute. -diff --git a/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMImpl.h b/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMImpl.h -index cfd9f28..945c6b8 100644 ---- a/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMImpl.h -+++ b/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMImpl.h -@@ -58,8 +58,10 @@ std::optional traceDefOp(Value v, bool isSingleChain = false) { - return traceDefOp(tensorCollapseShape.getSrc(), isSingleChain); - } else if (auto subViewOp = v.getDefiningOp()) { - return traceDefOp(subViewOp.getViewSource(), isSingleChain); -- } else if (auto toMemrefOp = v.getDefiningOp()) { -- return traceDefOp(toMemrefOp.getOperand(), isSingleChain); -+ // } else if (auto toMemrefOp = v.getDefiningOp()) { -+ // return traceDefOp(toMemrefOp.getOperand(), isSingleChain); -+} else if (auto toBufferOp = v.getDefiningOp()) { -+ return traceDefOp(toBufferOp.getOperand(), isSingleChain); - } else if (auto toTensorOp = v.getDefiningOp()) { - return traceDefOp(toTensorOp.getOperand(), isSingleChain); - } else if (auto viewOp = v.getDefiningOp()) { -diff --git a/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMInterfaces.td b/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMInterfaces.td -index f039870..5271c3c 100644 ---- a/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMInterfaces.td -+++ b/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMInterfaces.td -@@ -522,7 +522,8 @@ def HIVMStructuredOpInterface : OpInterface<"HIVMStructuredOp", - /*methodBody=*/"", - /*defaultImplementation=*/[{ - auto maps = $_op.getIndexingMapsArray(); -- return concatAffineMaps(maps); -+ // return concatAffineMaps(maps); -+ return concatAffineMaps(maps, $_op.getContext()); - }] - >, - InterfaceMethod< -diff --git a/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMOps.td b/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMOps.td -index 9c16c24..aea8554 100644 ---- a/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMOps.td -+++ b/bishengir/include/bishengir/Dialect/HIVM/IR/HIVMOps.td -@@ -219,4 +219,101 @@ def FinishDebugOp : HIVM_Op<"finish_debug", [CubeVectorCoreTypeTrait]> { - }]; - } - -+//===----------------------------------------------------------------------===// -+// CustomOp -+//===----------------------------------------------------------------------===// -+ -+def CustomOp -+ : HIVM_StructuredOp< -+ "custom", [AttrSizedOperandSegments, -+ MemoryEffects<[MemRead, MemWrite]>, SinglePipeOpTrait, -+ DeclareOpInterfaceMethods< -+ HIVMInferCoreTypeInterface, ["inferCoreType"]>]> { -+ let summary = [{ -+ Custom operation is a generic op interface for users to write their own custom implementation. -+ -+ Scenarios: -+ 1. Existing operations could not fulfill the desired functionality. -+ 2. Existing operations could fulfill the functionality, but overall performance is not optimal. -+ 3. Desire for private operation. -+ }]; -+ -+ let description = [{ -+ General interface for custom op, where: -+ - name : unique op name. -+ -+ Note : there are names reserved for builtins, usually starts with "__builtin". -+ Compiler will link these builtins to self-contained template library, -+ which comes together within bishengir-compile. -+ -+ For normal names/cases, user needs to specify implementation location/compilation commands (TODO), -+ and all ther necessary informations. -+ -+ Available builtin names: -+ "__builtin_gather_load" -+ -+ - inputs : input parameters. -+ - outputs : output results, designated "init" operands, which act as initial values for the results -+ of the operation or the init locations to which the results of the op will be written. -+ -+ In order to adapt to future enhancements quickly and dynamically, custom op relies on attributes -+ to retreive necessary information, required informations are: -+ - CoreType : which core type to execute on, refer to TCoreTypeAttr. -+ - Pipe : which pipe to execute on, refer to PipeAttr. -+ - VFMode : which mode to run on vector units, refer to VFModeAttr. -+ this attribute is ignored when core type is cube. -+ -+ Note : for builtins, user could specify these informations or not, -+ compiler will help to check the correctness and canonicalize. -+ -+ TODO: -+ - Impl : user provided implementation. -+ - Multi Pipe : custom op wants to use multiple pipes, which is a MacroOp in HIVM's context. -+ }]; -+ -+ let arguments = (ins StrAttr:$name, Variadic:$inputs, -+ Variadic:$outputs); -+ -+ let results = (outs Variadic:$results); -+ -+ let extraClassDeclaration = [{ -+ // TODO: Customize -+ static int getOpLibraryMaxRankImpl() { return 5; } -+ -+ PIPE getPipe(); -+ -+ ::mlir::MutableOperandRange getDpsInitsMutable() { -+ return getOutputsMutable(); -+ } -+ -+ // Helper functions -+ void setPipe(PIPE); -+ -+ std::optional getCoreType(); -+ void setCoreType(TCoreType); -+ -+ std::optional getVFMode(); -+ void setVFMode(VFMode); -+ -+ bool isBuiltin(); -+ -+ // Builtins helpers -+ struct BuiltinInfo { -+ TCoreType coreType; -+ PIPE pipe; -+ VFMode vfMode; -+ }; -+ -+ // Map <-> {CORE_TYPE, PIPE, VF_MODE} -+ static const DenseMap kBuiltins; -+ }]; -+ -+ let hasVerifier = 1; -+ let hasCustomAssemblyFormat = 1; -+ let hasCanonicalizer = 1; -+} -+ -+// TODO: Add CustomMacroOp -+//SH -+ - #endif // BISHENGIR_DIALECT_HIVM_IR_HIVMOPS_TD -diff --git a/bishengir/lib/Dialect/HFusion/IR/HFusionOps.cpp b/bishengir/lib/Dialect/HFusion/IR/HFusionOps.cpp -index 9309262..cdd81b9 100644 ---- a/bishengir/lib/Dialect/HFusion/IR/HFusionOps.cpp -+++ b/bishengir/lib/Dialect/HFusion/IR/HFusionOps.cpp -@@ -68,9 +68,14 @@ using namespace mlir::hfusion; - // Support for named HFusion ops defined in ods-gen. - //===----------------------------------------------------------------------===// - -+// using RegionBuilderFn = llvm::function_ref)>; -+ // static std::function, -+ // function_ref)> - using RegionBuilderFn = llvm::function_ref)>; -- -+ ArrayRef, -+ function_ref)>; - /// Fills the region of a structured operation using the provided - /// `regionBuilder`. The method is used by both named structured ops created by - /// ods-gen and by manually defined C++ ops. It is called by both builders and -@@ -103,7 +108,10 @@ static void fillStructuredOpRegion(OpBuilder &opBuilder, Region ®ion, - - opBuilder.setInsertionPointToStart(body); - ImplicitLocOpBuilder b(opBuilder.getUnknownLoc(), opBuilder); -- regionBuilder(b, *body, attrs); -+ // regionBuilder(b, *body, attrs); -+ regionBuilder(b, *body, attrs, []() -> InFlightDiagnostic { -+ llvm_unreachable("diagnostic function should not be called"); -+ }); - - // indexing_maps is an auto-generated method. - -@@ -996,10 +1004,18 @@ void codeGenWithIndexDispatch(OpBuilder &builder, Block &block, Type elemType, - } - } - --std::function)> -+// RegionBuilderFn ReduceWithIndexOp::getRegionBuilder() { -+// return [](ImplicitLocOpBuilder &b, Block &block, -+// ArrayRef attrs, -+// llvm::function_ref emitError) { -+// using RegionBuilderFn = llvm::function_ref, -+// function_ref)>; -+std::function, llvm::function_ref)> - ReduceWithIndexOp::getRegionBuilder() { - return [](ImplicitLocOpBuilder &b, Block &block, -- ArrayRef attrs) { -+ ArrayRef attrs, -+ llvm::function_ref emitError) { - // check numArgs - constexpr int kNumArgsWithoutIndex = 3; - auto numArgs = block.getNumArguments(); -@@ -1833,7 +1849,9 @@ ParseResult ArangeOp::parse(OpAsmParser &parser, OperationState &result) { - ImplicitLocOpBuilder builder(unknownLoc, parser.getContext()); - builder.setInsertionPointToStart(&block); - // Build the region -- getRegionBuilder()(builder, block, result.attributes.getAttrs()); -+ getRegionBuilder()(builder, block, result.attributes.getAttrs(), []() -> InFlightDiagnostic { -+ return InFlightDiagnostic(); -+ }); - - return success(); - } -@@ -1858,10 +1876,11 @@ ArrayAttr ArangeOp::getIndexingMaps() { - return builder.getAffineMapArrayAttr(maps); - } - --std::function)> -+std::function, llvm::function_ref)> - ArangeOp::getRegionBuilder() { - return [](ImplicitLocOpBuilder &builder, Block &block, -- ArrayRef attrs) { -+ ArrayRef attrs, -+ llvm::function_ref emitError) { - OpBuilder::InsertionGuard guard(builder); - - auto segmentSizes = cast_or_null( -@@ -2045,10 +2064,11 @@ void GatherOp::build(OpBuilder &odsBuilder, OperationState &odsState, Value src, - /// %cmp = arith.cmpi eq, , %iter - /// %sel = arith.select %cmp, , - /// linalg.yield %sel --std::function)> -+std::function, llvm::function_ref)> - GatherOp::getRegionBuilder() { - return [](ImplicitLocOpBuilder &builder, Block &block, -- ArrayRef attrs) { -+ ArrayRef attrs, -+ llvm::function_ref emitError) { - assert(block.getNumArguments() == 3 && - "GatherOp expecting 3 block arguments"); - Value srcVal = block.getArgument(0); -diff --git a/bishengir/lib/Dialect/HFusion/Transforms/OpFusion/FusibleHelper.cpp b/bishengir/lib/Dialect/HFusion/Transforms/OpFusion/FusibleHelper.cpp -index d3ad5f2..f190405 100644 ---- a/bishengir/lib/Dialect/HFusion/Transforms/OpFusion/FusibleHelper.cpp -+++ b/bishengir/lib/Dialect/HFusion/Transforms/OpFusion/FusibleHelper.cpp -@@ -506,11 +506,11 @@ OpPattern FusibleHelper::getOpPattern(Operation *op) { - [](auto) -> OpPattern { return OpPattern::kMidFusionAuxiliary; }) - .Case( - [](auto) -> OpPattern { return OpPattern::kMidFusionImportantAux; }) -- .Case( -+ .Case( - [](auto) -> OpPattern { return OpPattern::kAllReduce; }) -- .Case( -+ .Case( - [](auto) -> OpPattern { return OpPattern::kAllGather; }) -- .Case( -+ .Case( - [](auto) -> OpPattern { return OpPattern::kReduceScatter; }) - .Case( - [](auto) -> OpPattern { return OpPattern::kInterleave; }) -diff --git a/bishengir/lib/Dialect/HIVM/IR/HIVMCanonicalizations.cpp b/bishengir/lib/Dialect/HIVM/IR/HIVMCanonicalizations.cpp -index 968b75d..9dce8dd 100644 ---- a/bishengir/lib/Dialect/HIVM/IR/HIVMCanonicalizations.cpp -+++ b/bishengir/lib/Dialect/HIVM/IR/HIVMCanonicalizations.cpp -@@ -558,4 +558,41 @@ LogicalResult StoreOp::fold(hivm::StoreOp::FoldAdaptor adaptor, - void mlir::hivm::HIVMDialect::getCanonicalizationPatterns( - ::mlir::RewritePatternSet &results) const { - results.add(getContext()); --} -\ No newline at end of file -+} -+ -+//SH -+struct CustomOpCanonicalizer : public OpRewritePattern { -+ using OpRewritePattern::OpRewritePattern; -+ -+ LogicalResult matchAndRewrite(CustomOp customOp, -+ PatternRewriter &rewriter) const final { -+ if (!customOp.isBuiltin()) -+ return failure(); -+ -+ const auto &builtinInfo = CustomOp::kBuiltins.at(customOp.getName()); -+ const auto &coreType = customOp.getCoreType(); -+ if (!coreType || *coreType != builtinInfo.coreType) { -+ customOp.setCoreType(builtinInfo.coreType); -+ return success(); -+ } -+ -+ if (customOp.getPipe() != builtinInfo.pipe) { -+ customOp.setPipe(builtinInfo.pipe); -+ return success(); -+ } -+ -+ const auto &vfMode = customOp.getVFMode(); -+ if (!vfMode || *vfMode != builtinInfo.vfMode) { -+ customOp.setVFMode(builtinInfo.vfMode); -+ return success(); -+ } -+ -+ return failure(); -+ } -+}; -+ -+void CustomOp::getCanonicalizationPatterns(::mlir::RewritePatternSet &results, -+ ::mlir::MLIRContext *context) { -+ results.add(context); -+} -+//SH -diff --git a/bishengir/lib/Dialect/HIVM/IR/HIVMImpl.cpp b/bishengir/lib/Dialect/HIVM/IR/HIVMImpl.cpp -index b3d7807..2e8b56c 100644 ---- a/bishengir/lib/Dialect/HIVM/IR/HIVMImpl.cpp -+++ b/bishengir/lib/Dialect/HIVM/IR/HIVMImpl.cpp -@@ -302,7 +302,7 @@ Type getAnnotationMarkByteAlignment(Value value) { - auto memrefType = cast(shapedType); - bool isAlreadyAligned = true; - -- auto [strides, offset] = getStridesAndOffset(memrefType); -+ auto [strides, offset] = memrefType.getStridesAndOffset(); - llvm::SmallVector alignedStrides(rank, 1); - for (int64_t i = 0; i < rank; i++) { - if (strideAlignElems[i] == 1) { -diff --git a/bishengir/lib/Dialect/HIVM/IR/HIVMOps.cpp b/bishengir/lib/Dialect/HIVM/IR/HIVMOps.cpp -index e663a2c..827ac4f 100644 ---- a/bishengir/lib/Dialect/HIVM/IR/HIVMOps.cpp -+++ b/bishengir/lib/Dialect/HIVM/IR/HIVMOps.cpp -@@ -412,3 +412,204 @@ std::string hivm::detail::getTypeName(Location loc, Type type) { - return unknown; - } - -+ -+//SH===----------------------------------------------------------------------===// -+// CustomOp -+//===----------------------------------------------------------------------===// -+ -+// Helper functions -+void CustomOp::setPipe(PIPE pipe) { -+ getOperation()->setAttr(PipeAttr::name, PipeAttr::get(getContext(), pipe)); -+} -+ -+std::optional CustomOp::getCoreType() { -+ if (const auto coreTypeAttr = -+ getOperation()->template getAttrOfType( -+ TCoreTypeAttr::name)) { -+ return coreTypeAttr.getTcoretype(); -+ } -+ -+ return {}; -+} -+ -+void CustomOp::setCoreType(TCoreType coreType) { -+ getOperation()->setAttr(TCoreTypeAttr::name, -+ TCoreTypeAttr::get(getContext(), coreType)); -+} -+ -+std::optional CustomOp::getVFMode() { -+ if (const auto vfModeAttr = -+ getOperation()->template getAttrOfType( -+ VFModeAttr::name)) { -+ return vfModeAttr.getValue(); -+ } -+ -+ return {}; -+} -+ -+void CustomOp::setVFMode(VFMode vfMode) { -+ getOperation()->setAttr(VFModeAttr::name, -+ VFModeAttr::get(getContext(), vfMode)); -+} -+ -+bool CustomOp::isBuiltin() { return kBuiltins.contains(getName()); } -+ -+ParseResult CustomOp::parse(OpAsmParser &parser, OperationState &result) { -+ if (succeeded(parser.parseOptionalLess())) { -+ if (parser.parseAttribute(result.propertiesAttr) || parser.parseGreater()) -+ return failure(); -+ } -+ -+ // Parse attributes -+ SMLoc attrsLoc = parser.getCurrentLocation(); -+ if (parser.parseOptionalAttrDict(result.attributes)) -+ return failure(); -+ -+ { // Parse name -+ std::string name{}; -+ if (parser.parseString(&name)) -+ return failure(); -+ -+ result.addAttribute("name", parser.getBuilder().getStringAttr(name)); -+ } -+ -+ { // Parse variadic args -+ SmallVector variadicArgsSizes; -+ auto parseVariadicArgs = [&](const std::string &nameHint) { -+ SMLoc loc; -+ SmallVector types; -+ SmallVector operands; -+ -+ if (succeeded(parser.parseOptionalKeyword(nameHint))) { -+ loc = parser.getCurrentLocation(); -+ if (parser.parseLParen() || parser.parseOperandList(operands) || -+ parser.parseColonTypeList(types) || parser.parseRParen()) -+ return failure(); -+ } -+ -+ if (parser.resolveOperands(operands, types, loc, result.operands)) { -+ return failure(); -+ } -+ -+ variadicArgsSizes.push_back(static_cast(operands.size())); -+ return success(); -+ }; -+ -+ if (failed(parseVariadicArgs("ins")) || failed(parseVariadicArgs("outs"))) { -+ return failure(); -+ } -+ -+ // Update operandSegmentSizes attribute -+ const auto operandSegmentSizesAttr = -+ parser.getBuilder().getDenseI32ArrayAttr(variadicArgsSizes); -+ // This is a bit complex because we're trying to be backward compatible with -+ // operation syntax that mix the inherent attributes and the discardable -+ // ones in the same dictionary. If the properties are used, we append the -+ // operandSegmentSizes there directly. Otherwise we append it to the -+ // discardable attributes dictionary where it is handled by the generic -+ // Operation::create(...) method. -+ if (result.propertiesAttr) { -+ NamedAttrList attrs = llvm::cast(result.propertiesAttr); -+ attrs.append("operandSegmentSizes", operandSegmentSizesAttr); -+ result.propertiesAttr = attrs.getDictionary(parser.getContext()); -+ } else { -+ result.addAttribute("operandSegmentSizes", operandSegmentSizesAttr); -+ std::optional info = -+ result.name.getRegisteredInfo(); -+ if (info) { -+ if (failed(info->verifyInherentAttrs(result.attributes, [&]() { -+ return parser.emitError(attrsLoc) -+ << "'" << result.name.getStringRef() << "' op "; -+ }))) -+ return failure(); -+ } -+ } -+ } -+ -+ { // Parse result types -+ SmallVector resultTypes; -+ if (parser.parseOptionalArrowTypeList(resultTypes)) { -+ return failure(); -+ } -+ result.addTypes(resultTypes); -+ } -+ -+ return success(); -+} -+ -+void CustomOp::print(OpAsmPrinter &p) { -+ p.printOptionalAttrDict(getOperation()->getAttrs(), -+ /*elidedAttrs=*/{"operandSegmentSizes", "name"}); -+ -+ p << " "; -+ p.printString(getName()); -+ -+ auto printVariadicArgs = [&](const auto &args, const std::string &nameHint) { -+ if (!args.empty()) -+ p << " " << nameHint << "(" << args << " : " << args.getTypes() << ")"; -+ }; -+ -+ printVariadicArgs(getInputs(), "ins"); -+ printVariadicArgs(getOutputs(), "outs"); -+ -+ if (!getResults().empty()) -+ p.printOptionalArrowTypeList(getResultTypes()); -+} -+ -+static LogicalResult verifyBuiltins(CustomOp op) { -+ const auto &builtinInfo = CustomOp::kBuiltins.at(op.getName()); -+ -+ const auto &coreType = op.getCoreType(); -+ if (coreType && *coreType != builtinInfo.coreType) -+ return op.emitOpError() << "Specified core type conflict with " -+ << op.getName() << "'s core type."; -+ -+ const auto &pipe = op.getPipe(); -+ if (pipe != PIPE::PIPE_UNASSIGNED && pipe != builtinInfo.pipe) -+ return op.emitOpError() -+ << "Specified pipe conflict with " << op.getName() << "'s pipe."; -+ -+ const auto &vfMode = op.getVFMode(); -+ if (vfMode && *vfMode != builtinInfo.vfMode) -+ return op.emitOpError() << "Specified vf mode conflict with " -+ << op.getName() << "'s vf mode."; -+ -+ return success(); -+} -+ -+LogicalResult CustomOp::verify() { -+ // Check builtins -+ // if (isBuiltin()) -+ // return verifyBuiltins(*this); -+ -+ // // Check core type attribute -+ // const auto coreType = getCoreType(); -+ // if (!coreType) -+ // return emitOpError() << "Missing core type information"; -+ -+ // // Check pipe attribute -+ // if (getPipe() == PIPE::PIPE_UNASSIGNED) -+ // return emitOpError() << "Missing pipe information"; -+ -+ // // Check VF mode attribute -+ // if (*coreType != TCoreType::CUBE) { -+ // if (!getVFMode()) -+ // return emitOpError() << "Missing vf mode information"; -+ // } else { // Pure cube -+ // // Cube function ignores vf mode information -+ // } -+ -+ return success(); -+} -+ -+PIPE CustomOp::getPipe() { -+ if (auto pipAttr = -+ getOperation()->template getAttrOfType(PipeAttr::name)) -+ return pipAttr.getPipe(); -+ -+ return PIPE::PIPE_UNASSIGNED; -+} -+ -+const DenseMap CustomOp::kBuiltins{ -+ {"__builtin_gather_load", {TCoreType::VECTOR, PIPE::PIPE_V, VFMode::SIMT}}}; -+//SH -diff --git a/bishengir/lib/Dialect/HIVM/IR/HIVMSynchronizationOps.cpp b/bishengir/lib/Dialect/HIVM/IR/HIVMSynchronizationOps.cpp -index 8c61377..0d2d7d8 100644 ---- a/bishengir/lib/Dialect/HIVM/IR/HIVMSynchronizationOps.cpp -+++ b/bishengir/lib/Dialect/HIVM/IR/HIVMSynchronizationOps.cpp -@@ -151,7 +151,9 @@ void SyncBlockSetOp::build(OpBuilder &odsBuilder, OperationState &odsState, - /*tsync_instr_mode=*/{}); - } else { - build(odsBuilder, odsState, tcore_type, tpipe, pipe, nullptr, -- flag_id.get(), nullptr, /*tsync_instr_mode=*/{}); -+ // flag_id.get(), nullptr, /*tsync_instr_mode=*/{}); -+ // flag_id.cast(), nullptr, /*tsync_instr_mode=*/{}); -+ cast(flag_id), nullptr, /*tsync_instr_mode=*/{}); - } - } - -@@ -165,7 +167,8 @@ void SyncBlockSetOp::build(OpBuilder &odsBuilder, OperationState &odsState, - cast(attr), nullptr, ffts_base_addr, tsync_instr_mode); - } else { - build(odsBuilder, odsState, tcore_type, tpipe, pipe, nullptr, -- flag_id.get(), ffts_base_addr, tsync_instr_mode); -+ // flag_id.get(), ffts_base_addr, tsync_instr_mode); -+ cast(flag_id), ffts_base_addr, tsync_instr_mode); - } - } - -@@ -201,7 +204,8 @@ void SyncBlockWaitOp::build(OpBuilder &odsBuilder, OperationState &odsState, - cast(attr), nullptr); - } else { - build(odsBuilder, odsState, tcore_type, tpipe, pipe, nullptr, -- flag_id.get()); -+ // flag_id.get()); -+ cast(flag_id)); - } - } - -diff --git a/bishengir/lib/Dialect/HIVM/IR/HIVMTraits.cpp b/bishengir/lib/Dialect/HIVM/IR/HIVMTraits.cpp -index 7f9c39f..9775dca 100644 ---- a/bishengir/lib/Dialect/HIVM/IR/HIVMTraits.cpp -+++ b/bishengir/lib/Dialect/HIVM/IR/HIVMTraits.cpp -@@ -19,6 +19,8 @@ - - #include "mlir/IR/BuiltinTypeInterfaces.h" - #include "mlir/IR/TypeUtilities.h" -+// #include "mlir/Analysis/Strides.h" -+#include "mlir/IR/BuiltinTypes.h" - - #include "llvm/ADT/SmallSet.h" - #include "llvm/ADT/SmallVector.h" -@@ -43,7 +45,8 @@ inline int64_t getRank(const Type &type) { - - // Returns stride of shaped memref type. - inline SmallVector getStride(const Type &type) { -- auto [strides, offset] = getStridesAndOffset(cast(type)); -+ // auto [strides, offset] = getStridesAndOffset(cast(type)); -+ auto [strides, offset] = cast(type).getStridesAndOffset(); - return strides; - } - -diff --git a/bishengir/lib/Dialect/HIVM/IR/InferCoreTypeInterface/InferCoreType.cpp b/bishengir/lib/Dialect/HIVM/IR/InferCoreTypeInterface/InferCoreType.cpp -index 39a1231..36a95e6 100644 ---- a/bishengir/lib/Dialect/HIVM/IR/InferCoreTypeInterface/InferCoreType.cpp -+++ b/bishengir/lib/Dialect/HIVM/IR/InferCoreTypeInterface/InferCoreType.cpp -@@ -111,6 +111,18 @@ inferCoreTypeForGlobalMixMatmulOps(GlobalMixMatmulTy *mixMatmulOp) { - // HIVM Ops - //===----------------------------------------------------------------------===// - -+//SH -+std::optional CustomOp::inferCoreType() { -+ if (auto coreTypeAttr = getOperation()->template getAttrOfType( -+ TCoreTypeAttr::name)) { -+ return coreTypeAttr.getTcoretype(); -+ } -+ -+ return {}; -+ } -+//SH -+ -+ - std::optional ConvertLayoutOp::inferCoreType() { - BaseMemRefType srcMemRefTy = getSource().getType(); - hivm::AddressSpace addrSpace = -diff --git a/bishengir/lib/Dialect/Utils/Util.cpp b/bishengir/lib/Dialect/Utils/Util.cpp -index cd4ccbe..2a2f5ac 100644 ---- a/bishengir/lib/Dialect/Utils/Util.cpp -+++ b/bishengir/lib/Dialect/Utils/Util.cpp -@@ -364,7 +364,9 @@ getTensorOrMemrefDynSizes(OpBuilder &builder, Location loc, Value source, - - inline bool isPureStatic(ArrayRef mixedValues) { - return llvm::all_of(mixedValues, -- [](OpFoldResult x) { return x.is(); }); -+ // [](OpFoldResult x) { return x.is(); }); -+ [](OpFoldResult x) { return isa(x); }); -+ - } - - inline void markDynShapeAlloc(OpBuilder &builder, Value source, -@@ -957,10 +959,24 @@ bool isElementwiseOp(Operation *op) { - return true; - } - -+bool isElementwiseLinalgOp(Operation* op) { -+ if (auto linalgOp = dyn_cast(op)) { -+ auto iteratorTypes = linalgOp.getIteratorTypesArray(); -+ return llvm::all_of(iteratorTypes, -+ [](mlir::utils::IteratorType iterType) { -+ return iterType == mlir::utils::IteratorType::parallel; -+ }); -+ } -+ return false; -+} -+ -+// https://discourse.llvm.org/t/rfc-deprecate-linalg-elemwise-unary-and-elemwise-binary/87144/2 - bool isMarkedAsElementwiseOp(Operation *op) { - // This would handle scalar as well -- return isa_and_present(op); -+ // return isa_and_present(op); -+ return (isElementwiseLinalgOp(op) || -+ isa(op)); - } - - bool isZeroDimensionOp(Operation *op) { -@@ -975,7 +991,9 @@ bool isZeroDimensionOp(Operation *op) { - - bool isMarkedAsElementwiseUnaryOp(Operation *op) { - // This would handle scalar as well -- return isa_and_present(op); -+ // return isa_and_present(op); -+ return (isElementwiseLinalgOp(op) || -+ isa(op)); - } - - bool isAllParallelOp(Operation *op) { -@@ -990,8 +1008,8 @@ bool isAllParallelOp(Operation *op) { - - // TODO: Need to refactor this. - bool isLegalOp(Operation *op) { -- if (isa(op)) { -diff --git a/bishengir/tools/bishengir-hfusion-ods-gen/bishengir-hfusion-ods-yaml-gen.cpp b/bishengir/tools/bishengir-hfusion-ods-gen/bishengir-hfusion-ods-yaml-gen.cpp -index 3907aed..3e5fb17 100644 ---- a/bishengir/tools/bishengir-hfusion-ods-gen/bishengir-hfusion-ods-yaml-gen.cpp -+++ b/bishengir/tools/bishengir-hfusion-ods-gen/bishengir-hfusion-ods-yaml-gen.cpp -@@ -608,9 +608,11 @@ def {0} : HFusionStructuredBase_Op<"{1}", !listconcat([AttrSizedOperandSegments] - SmallVector getIteratorTypesArray(); - ArrayAttr getIndexingMaps(); - static void regionBuilder(ImplicitLocOpBuilder &b, -- Block &block, ArrayRef attrs); -+ Block &block, ArrayRef attrs, -+ function_ref diag); - static std::function)> -+ Block &, ArrayRef, -+ function_ref)> - getRegionBuilder() {{ - return regionBuilder; - } -@@ -1055,7 +1057,8 @@ LogicalResult {0}::verifyIndexingMapRequiredAttributes() {{ - // {3}: Statements - static const char structuredOpRegionBuilderFormat[] = R"FMT( - void {0}::regionBuilder(ImplicitLocOpBuilder &b, -- Block &block, ArrayRef attrs) {{ -+ Block &block, ArrayRef attrs, -+ function_ref diag) {{ - assert({1} > 0 && block.getNumArguments() == {1} && - "{0} regionBuilder expects {1} (>=0) args"); - RegionBuilderHelper helper(block.getArgument(0).getContext(), block); diff --git a/patch/triton/CMakeLists_txt.patch b/patch/triton/CMakeLists_txt.patch new file mode 100644 index 00000000..c426276d --- /dev/null +++ b/patch/triton/CMakeLists_txt.patch @@ -0,0 +1,19 @@ +diff --git a/CMakeLists.txt b/CMakeLists.txt +index 47a0f3b17..fc7729ce3 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -83,6 +83,14 @@ endif() + + # Compiler flags + include_directories(${CMAKE_CURRENT_SOURCE_DIR}/include) ++ ++# Ascend NPU: 检测 npu-smi，禁用 AddPtrOp fold 以避免与 LoadStoreCanonicalizer 死循环 ++find_program(NPU_SMI_CMD npu-smi) ++if(NPU_SMI_CMD) ++ add_definitions(-DWITH_ASCEND=1) ++ message(STATUS "WITH_ASCEND=1: npu-smi found, AddPtrOp::fold disabled") ++endif() ++ + if(NOT MSVC) + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -D__STDC_FORMAT_MACROS -fPIC -std=gnu++17") + else() diff --git a/patch/triton/include_triton_Dialect_Triton_IR_TritonOps_td.patch b/patch/triton/include_triton_Dialect_Triton_IR_TritonOps_td.patch deleted file mode 100644 index 050e5818..00000000 --- a/patch/triton/include_triton_Dialect_Triton_IR_TritonOps_td.patch +++ /dev/null @@ -1,88 +0,0 @@ -diff --git a/include/triton/Dialect/Triton/IR/TritonOps.td b/include/triton/Dialect/Triton/IR/TritonOps.td -index e9f892da0..07b83751f 100644 ---- a/include/triton/Dialect/Triton/IR/TritonOps.td -+++ b/include/triton/Dialect/Triton/IR/TritonOps.td -@@ -1410,5 +1410,83 @@ def TT_DescriptorScatterOp : TT_Op<"descriptor_scatter", [TT_DescriptorStoreLike - let hasVerifier = 1; - } - -+def TT_AnnotationOp : TT_Op<"annotation", [Pure, MemoryEffects<[MemWrite]>]> { -+ let summary = "Annotate a tensor with key-value attribute pairs"; -+ let description = [{ -+ `tt.annotation` operation can be used to annotate a tensor with -+ key-value attribute pairs. -+ -+ Example: -+ ```mlir -+ tt.annotation %target {key : val} -+ ``` -+ }]; -+ let arguments = (ins TT_Tensor:$src); -+ let assemblyFormat = [{ -+ $src attr-dict `:` type($src) -+ }]; -+} -+ -+def TT_CustomSyncOp : TT_Op<"custom_sync", [Pure, MemoryEffects<[MemWrite]>]> { -+ let summary = "self-defined custom sync operation"; -+ let description = [{ -+ `tt.custom_sync` triton custom sync op is designed to pass self-defined custom sync operation. -+ }]; -+ let arguments = (ins StrAttr:$op_name, StrAttr:$mode_or_sender, I32Attr:$id); -+ let assemblyFormat = "$op_name attr-dict"; -+} -+ -+//===----------------------------------------------------------------------===// -+// CustomOp -+//===----------------------------------------------------------------------===// -+def TT_CustomOp : TT_Op<"custom_op", [AttrSizedOperandSegments, MemoryEffects<[MemRead, MemWrite]>]> { -+ let summary = [{ -+ Custom operation is a generic op interface for users to write their own custom implementation. -+ -+ Scenarios: -+ 1. Existing operations could not fulfill the desired functionality. -+ 2. Existing operations could fulfill the functionality, but overall performance is not optimal. -+ 3. Desire for private operation. -+ }]; -+ -+ let description = [{ -+ General interface for custom op, where: -+ - name : unique op name. -+ -+ Note : there are names reserved for builtins, usually starts with "__builtin". -+ Compiler will link these builtins to self-contained template library, -+ which comes together within bishengir-compile. -+ -+ For normal names/cases, user needs to specify implementation location/compilation commands (TODO), -+ and all ther necessary informations. -+ -+ Available builtin names: -+ "__builtin_gather_load" -+ -+ - inputs : input parameters. -+ - outputs : output results, designated "init" operands, which act as initial values for the results -+ of the operation or the init locations to which the results of the op will be written. -+ -+ In order to adapt to future enhancements quickly and dynamically, custom op relies on attributes -+ to retreive necessary information, required informations are: -+ - CoreType : which core type to execute on, refer to TCoreTypeAttr. -+ - Pipe : which pipe to execute on, refer to PipeAttr. -+ - VFMode : which mode to run on vector units, refer to VFModeAttr. -+ this attribute is ignored when core type is cube. -+ -+ Note : for builtins, user could specify these informations or not, -+ compiler will help to check the correctness and canonicalize. -+ -+ TODO: -+ - Impl : user provided implementation. -+ - Multi Pipe : custom op wants to use multiple pipes, which is a MacroOp in HIVM's context. -+ }]; -+ -+ let arguments = (ins StrAttr:$name, Variadic:$inputs, -+ Variadic:$outputs); -+ -+ let results = (outs Variadic:$results); -+ -+} - - #endif // Triton_OPS diff --git a/patch/triton/lib_Dialect_Triton_IR_Ops_cpp.patch b/patch/triton/lib_Dialect_Triton_IR_Ops_cpp.patch new file mode 100644 index 00000000..7ef956c2 --- /dev/null +++ b/patch/triton/lib_Dialect_Triton_IR_Ops_cpp.patch @@ -0,0 +1,22 @@ +diff --git a/lib/Dialect/Triton/IR/Ops.cpp b/lib/Dialect/Triton/IR/Ops.cpp +index 06082cf9d..48e08e526 100644 +--- a/lib/Dialect/Triton/IR/Ops.cpp ++++ b/lib/Dialect/Triton/IR/Ops.cpp +@@ -995,6 +995,9 @@ void MakeTensorPtrOp::build(OpBuilder &builder, OperationState &state, + } + + //-- AddPtrOp -- ++#ifdef WITH_ASCEND ++OpFoldResult AddPtrOp::fold(FoldAdaptor adaptor) { return {}; } ++#else + OpFoldResult AddPtrOp::fold(FoldAdaptor adaptor) { + // addptr(ptr, 0) -> ptr + if (matchPattern(adaptor.getOffset(), m_Zero())) { +@@ -1002,6 +1005,7 @@ OpFoldResult AddPtrOp::fold(FoldAdaptor adaptor) { + } + return {}; + } ++#endif + + //-- AdvanceOp -- + OpFoldResult AdvanceOp::fold(FoldAdaptor adaptor) { diff --git a/patch/triton/python_src_ir_cc.patch b/patch/triton/python_src_ir_cc.patch index 557435a0..c03ee15d 100644 --- a/patch/triton/python_src_ir_cc.patch +++ b/patch/triton/python_src_ir_cc.patch @@ -1,8 +1,22 @@ diff --git a/python/src/ir.cc b/python/src/ir.cc -index 4c8a4233b..b88942c30 100644 +index d79a9e70f..2ea7c4db3 100644 --- a/python/src/ir.cc +++ b/python/src/ir.cc -@@ -339,6 +339,10 @@ void init_triton_ir(py::module &&m) { +@@ -245,6 +245,13 @@ py::list getTensorDescMetadata(ModuleOp &mod) { + + } // anonymous namespace + ++namespace ir { ++ ++static py::class_ *builderClassPtr = nullptr; ++py::class_ *getBuilderClass() { return builderClassPtr; } ++ ++} // namespace ir ++ + /*****************************************************************************/ + /* Python bindings for ir */ + /*****************************************************************************/ +@@ -390,6 +397,10 @@ void init_triton_ir(py::module &&m) { .def("param_types", [](FunctionType &self) { return std::vector(self.getInputs().begin(), self.getInputs().end()); @@ -13,7 +27,7 @@ index 4c8a4233b..b88942c30 100644 }); py::class_(m, "location", py::module_local()) -@@ -464,6 +468,7 @@ void init_triton_ir(py::module &&m) { +@@ -515,6 +526,7 @@ void init_triton_ir(py::module &&m) { py::class_(m, "integer_attr", py::module_local()); py::class_(m, "bool_attr", py::module_local()); py::class_(m, "unit_attr", py::module_local()); @@ -21,114 +35,16 @@ index 4c8a4233b..b88942c30 100644 // Ops py::class_(m, "OpState", py::module_local()) -@@ -1740,6 +1745,110 @@ void init_triton_ir(py::module &&m) { - .def("create_gather", - [](TritonOpBuilder &self, Value src, Value indices, int axis) - -> Value { return self.create(src, indices, axis); }) -+ .def("create_extract_slice", -+ [](TritonOpBuilder &self, Value &ful, std::vector &offs_vec, -+ std::vector &sizs_vec, std::vector &strd_vec) -> Value { -+ self.getContext()->getOrLoadDialect(); -+ llvm::SmallVector offsets; -+ for (const auto &o : offs_vec) { -+ auto oTy = o.getType(); -+ if (!oTy.isIndex()) { -+ auto v = self.create( -+ self.getBuilder().getIndexType(), o); -+ offsets.push_back(v); -+ } else { -+ offsets.push_back(o); -+ } -+ } -+ llvm::SmallVector sizes; -+ llvm::SmallVector retSizes; -+ for (const auto &s : sizs_vec) { -+ auto v = self.create(s); -+ sizes.push_back(v); -+ retSizes.push_back(s); -+ } -+ llvm::SmallVector strides; -+ for (const auto &s : strd_vec) { -+ auto v = self.create(s); -+ strides.push_back(v); -+ } -+ auto retTy = RankedTensorType::get( -+ retSizes, -+ cast(ful.getType()).getElementType()); -+ auto ret = self.create(retTy, ful, offsets, -+ sizes, strides); -+ return ret; -+ }) -+ .def("create_insert_slice", -+ [](TritonOpBuilder &self, Value &ful, Value &sub, -+ std::vector &offs_vec, std::vector &sizs_vec, -+ std::vector &strd_vec) -> Value { -+ self.getContext()->getOrLoadDialect(); -+ llvm::SmallVector offsets; -+ for (const auto &o : offs_vec) { -+ auto oTy = o.getType(); -+ if (!oTy.isIndex()) { -+ auto v = self.create( -+ self.getBuilder().getIndexType(), o); -+ offsets.push_back(v); -+ } else { -+ offsets.push_back(o); -+ } -+ } -+ llvm::SmallVector sizes; -+ llvm::SmallVector retSizes; -+ for (const auto &s : sizs_vec) { -+ auto v = self.create(s); -+ sizes.push_back(v); -+ retSizes.push_back(s); -+ } -+ llvm::SmallVector strides; -+ for (const auto &s : strd_vec) { -+ auto v = self.create(s); -+ strides.push_back(v); -+ } -+ auto retTy = RankedTensorType::get( -+ retSizes, -+ cast(ful.getType()).getElementType()); -+ auto ret = self.create(sub, ful, offsets, -+ sizes, strides); -+ return ret; -+ }) -+ .def("create_annotation", -+ [](TritonOpBuilder &self, Value &ptr, const std::string &attrKey, -+ Attribute &attrVal) { -+ auto annotationOp = self.create(ptr); -+ annotationOp->setAttr(self.getBuilder().getStringAttr(attrKey), -+ attrVal); -+ }) -+ .def("create_custom_op_for_inter_core_sync", -+ [](TritonOpBuilder &self, std::string &op_name, -+ std::string &mode_or_sender, int id) -> void { -+ self.create( -+ self.getBuilder().getStringAttr(op_name), -+ self.getBuilder().getStringAttr(mode_or_sender), -+ self.getBuilder().getI32IntegerAttr(id)); -+ }) -+ .def("create_custom_op", -+ [](TritonOpBuilder &self, -+ const std::string &name, -+ const py::dict &attrs, -+ const std::vector &ins, -+ const std::vector &outs) -> std::vector { -+ ValueRange inputs{ins}; -+ ValueRange outputs{outs}; -+ TypeRange res_types{outputs}; -+ auto op = self.create(res_types, name, inputs, outputs); -+ for (auto &attr : attrs) { -+ std::string attr_name = py::cast(attr.first); -+ std::string attr_value = py::cast(attr.second); -+ op->setAttr(attr_name, self.getBuilder().getStringAttr(attr_value)); -+ // Attribute attr_value = py::cast(attr.second); -+ // op->setAttr(attr_name, attr_value); -+ } -+ auto results = op->getResults(); -+ return std::vector(results.begin(), results.end()); -+ }) - // Force GPU barrier - .def("create_barrier", - [](TritonOpBuilder &self) { self.create(); }) +@@ -777,8 +789,10 @@ void init_triton_ir(py::module &&m) { + + py::class_(m, "InsertPoint", py::module_local()); + +- py::class_(m, "builder", py::module_local(), +- py::dynamic_attr()) ++ static py::class_ builderClass( ++ m, "builder", py::module_local(), py::dynamic_attr()); ++ ir::builderClassPtr = &builderClass; ++ builderClass + .def(py::init()) + .def("get_op_builder", &TritonOpBuilder::getBuilder, ret::reference) + // getters diff --git a/patch/triton/python_src_ir_h.patch b/patch/triton/python_src_ir_h.patch new file mode 100644 index 00000000..f00325e2 --- /dev/null +++ b/patch/triton/python_src_ir_h.patch @@ -0,0 +1,23 @@ +diff --git a/python/src/ir.h b/python/src/ir.h +index e1f9ce848..66af85b9f 100644 +--- a/python/src/ir.h ++++ b/python/src/ir.h +@@ -3,6 +3,8 @@ + #include "triton/Tools/Sys/GetEnv.hpp" + #include "llvm/ADT/ArrayRef.h" + #include ++#include ++#include + + typedef int AsyncTaskId; + void setAsyncTaskIds(mlir::Operation *op, +@@ -103,3 +105,9 @@ private: + bool lineInfoEnabled = + !mlir::triton::tools::getBoolEnv("TRITON_DISABLE_LINE_INFO"); + }; ++ ++namespace py = pybind11; ++ ++namespace ir { ++extern py::class_ *getBuilderClass(); ++} // namespace ir diff --git a/patch/triton/python_triton_utils_py.patch b/patch/triton/python_triton__utils_py.patch similarity index 100% rename from patch/triton/python_triton_utils_py.patch rename to patch/triton/python_triton__utils_py.patch diff --git a/patch/triton/python_triton_compiler_code_generator_py.patch b/patch/triton/python_triton_compiler_code_generator_py.patch index c2f30a91..b51aeb14 100644 --- a/patch/triton/python_triton_compiler_code_generator_py.patch +++ b/patch/triton/python_triton_compiler_code_generator_py.patch @@ -1,8 +1,17 @@ diff --git a/python/triton/compiler/code_generator.py b/python/triton/compiler/code_generator.py -index df09b3198..3dac7c740 100644 +index 176b6b515..be2b5aff3 100644 --- a/python/triton/compiler/code_generator.py +++ b/python/triton/compiler/code_generator.py -@@ -20,7 +20,6 @@ from .._utils import find_paths_if, get_iterable_path, set_iterable_path +@@ -12,7 +12,7 @@ from types import ModuleType + from typing import Any, Callable, Dict, Optional, Tuple, Type, Union, Iterable, List + + from .. import knobs, language +-from .._C.libtriton import ir, gluon_ir ++from .._C.libtriton import ir, gluon_ir, dicp_triton + from ..language import constexpr, str_to_ty, tensor, tuple as tl_tuple + from ..language.core import _unwrap_if_constexpr, base_value, base_type + # ideally we wouldn't need any runtime component +@@ -21,7 +21,6 @@ from .._utils import find_paths_if, get_iterable_path, set_iterable_path from .errors import (CompilationError, CompileTimeAssertionFailure, UnsupportedLanguageConstruct) @@ -10,7 +19,76 @@ index df09b3198..3dac7c740 100644 def check_identifier_legality(name, type): pattern = r'^[a-zA-Z_][a-zA-Z0-9_]*$' if not re.match(pattern, name): -@@ -578,6 +577,39 @@ class CodeGenerator(ast.NodeVisitor): +@@ -29,6 +28,17 @@ def check_identifier_legality(name, type): + return name + + ++# Central registry for 'with' statement handlers ++WITH_DISPATCH = {} ++ ++# Import and register deeplink extension dispatch handlers ++from triton.language.extra.deeplink.cann.extension.dispatch import ASCEND_WITH_DISPATCH as DEEPLINK_WITH_DISPATCH ++from triton.language.extra.deeplink.cann.extension.builder import setup_unified_builder ++from triton.language.extra.deeplink.cann.buffer.builder import setup_unified_builder_with_buffer_builder ++from triton.language.extra.deeplink.cann.buffer.core import is_builtin as is_dicp_builtin ++WITH_DISPATCH.update(DEEPLINK_WITH_DISPATCH) ++ ++ + def mangle_fn(name, arg_tys, constants, caller_context): + # doesn't mangle ret type, which must be a function of arg tys + mangled_arg_names = '_'.join([ty.mangle() for ty in arg_tys]) +@@ -308,6 +318,14 @@ class CodeGenerator(ast.NodeVisitor): + self.builder = ir.builder(context) + self.semantic = TritonSemantic(self.builder) + ++ self.dicp_builder = dicp_triton.ir.dicp_npu_ir_builder(context, getattr(options, "arch", "")) ++ self.dicp_builder.set_loc(file_name, begin_line, 0) ++ setup_unified_builder(self.builder, self.dicp_builder) ++ ++ self.buffer_builder = dicp_triton.ir.buffer_builder(context) ++ self.buffer_builder.set_loc(file_name, begin_line, 0) ++ setup_unified_builder_with_buffer_builder(self.builder, self.buffer_builder) ++ + self.name_loc_as_prefix = None + self.file_name = file_name + # node.lineno starts from 1, so we need to subtract 1 +@@ -456,17 +474,19 @@ class CodeGenerator(ast.NodeVisitor): + self.lscope[name] = value + self.local_defs[name] = value + +- def _get_insertion_point_and_loc(self): ++ def _get_insertion_point_and_loc(self, builder=None): + # XXX: this is a hack to get the location of the insertion point. + # The insertion point's location could be invalid sometimes, + # so we need to explicitly set the location +- loc = self.builder.get_loc() +- ip = self.builder.get_insertion_point() ++ _builder = builder if builder else self.builder ++ loc = _builder.get_loc() ++ ip = _builder.get_insertion_point() + return ip, loc + +- def _set_insertion_point_and_loc(self, ip, loc): +- self.builder.restore_insertion_point(ip) +- self.builder.set_loc(loc) ++ def _set_insertion_point_and_loc(self, ip, loc, builder=None): ++ _builder = builder if builder else self.builder ++ _builder.restore_insertion_point(ip) ++ _builder.set_loc(loc) + + def _find_carries(self, node, liveins): + # create loop body block +@@ -474,7 +494,8 @@ class CodeGenerator(ast.NodeVisitor): + self.builder.set_insertion_point_to_start(block) + # dry visit loop body + self.scf_stack.append(node) +- self.visit_compound_statement(node.body) ++ with language.extra.deeplink.cann.extension.semantic.dry_run_context(): ++ self.visit_compound_statement(node.body) + self.scf_stack.pop() + block.erase() + +@@ -579,6 +600,39 @@ class CodeGenerator(ast.NodeVisitor): assert isinstance(args, language.core.tuple) return args.values @@ -50,9 +128,39 @@ index df09b3198..3dac7c740 100644 def visit_FunctionDef(self, node): arg_names, kwarg_names = self.visit(node.args) if self.fn: -@@ -1044,6 +1076,10 @@ class CodeGenerator(ast.NodeVisitor): - f'but is re-assigned to {loop_val.type} in loop! '\ - f'Please make sure that the type stays consistent.' +@@ -972,16 +1026,20 @@ class CodeGenerator(ast.NodeVisitor): + return self.visit(node.orelse) + + def visit_With(self, node): +- # Lower `with` statements by constructing context managers and calling their enter/exit hooks +- # Instantiate each context manager with builder injection +- if len(node.items) == 1: # Handle async_task ++ # Try dispatch mechanism for deeplink context managers (e.g., scope) ++ if len(node.items) == 1: + context = node.items[0].context_expr +- withitemClass = self.visit(context.func) +- if withitemClass == language.async_task: +- args = [self.visit(arg) for arg in context.args] +- with withitemClass(*args, _builder=self.builder): +- self.visit_compound_statement(node.body) +- return ++ if isinstance(context, ast.Call): ++ withitemClass = self.visit(context.func) ++ handler = WITH_DISPATCH.get(withitemClass) ++ if handler: ++ return handler(self, node) ++ # Handle async_task ++ if withitemClass == language.async_task: ++ args = [self.visit(arg) for arg in context.args] ++ with withitemClass(*args, _builder=self.builder): ++ self.visit_compound_statement(node.body) ++ return + + cm_list = [] + for item in node.items: +@@ -1058,6 +1116,10 @@ class CodeGenerator(ast.NodeVisitor): + def visit_withitem(self, node): + return self.visit(node.context_expr) + def visit_With(self, node): + assert len(node.items) == 1 @@ -61,36 +169,47 @@ index df09b3198..3dac7c740 100644 def visit_While(self, node): with enter_sub_region(self) as sr: liveins, insert_block = sr -@@ -1137,7 +1173,8 @@ class CodeGenerator(ast.NodeVisitor): +@@ -1151,7 +1213,7 @@ class CodeGenerator(ast.NodeVisitor): flatten = False warp_specialize = False disable_licm = False - if IteratorClass is language.range: -+ bind_sub_block = None + if IteratorClass in [language.range, language.extra.deeplink.parallel]: iterator = IteratorClass(*iter_args, **iter_kwargs) # visit iterator arguments # note: only `range` iterator is supported now -@@ -1151,6 +1188,8 @@ class CodeGenerator(ast.NodeVisitor): - flatten = iterator.flatten - warp_specialize = iterator.warp_specialize - disable_licm = iterator.disable_licm -+ if (IteratorClass is language.extra.deeplink.parallel): -+ bind_sub_block = iterator.bind_sub_block - elif IteratorClass is range: - # visit iterator arguments - # note: only `range` iterator is supported now -@@ -1210,6 +1249,9 @@ class CodeGenerator(ast.NodeVisitor): +@@ -1224,6 +1286,8 @@ class CodeGenerator(ast.NodeVisitor): if disable_licm: for_op.set_attr("llvm.loop_annotation", self.builder.get_disable_loop_licm_attr()) -+ if (bind_sub_block is not None) and bind_sub_block: -+ for_op.set_attr("bind_sub_block", self.builder.get_bool_attr(bind_sub_block)) -+ ++ if (IteratorClass is language.extra.deeplink.parallel): ++ for_op.set_attr("hivm.parallel_loop", self.builder.get_unit_attr()) self.scf_stack.append(node) for_op_body = for_op.get_body(0) self.builder.set_insertion_point_to_start(for_op_body) -@@ -1363,6 +1405,13 @@ class CodeGenerator(ast.NodeVisitor): +@@ -1338,6 +1402,11 @@ class CodeGenerator(ast.NodeVisitor): + return self.call_JitFunction(fn, args, kws) + if (hasattr(fn, '__self__') and _is_triton_value(fn.__self__)) or language.core.is_builtin(fn) or isinstance( + fn, ConstexprFunction): ++ # Copy builder's location and insertion point. ++ ip, last_loc = self._get_insertion_point_and_loc() ++ # Use dicp_builder if this function is a DICP builtin extension operation. ++ _builder = getattr(self, 'dicp_builder', self.builder) if is_dicp_builtin(fn) else self.builder ++ self._set_insertion_point_and_loc(ip, last_loc, _builder) + extra_kwargs = dict() + + if isinstance(fn, ConstexprFunction): +@@ -1353,6 +1422,9 @@ class CodeGenerator(ast.NodeVisitor): + # builtin functions return plain tuples for readability + if isinstance(ret, tuple): + ret = language.tuple(ret) ++ # Sync the builder's location before return. ++ ip, last_loc = self._get_insertion_point_and_loc(_builder) ++ self._set_insertion_point_and_loc(ip, last_loc) + return ret + except Exception as e: + if knobs.compilation.front_end_debugging: +@@ -1385,6 +1457,13 @@ class CodeGenerator(ast.NodeVisitor): def visit_Call(self, node): fn = _unwrap_if_constexpr(self.visit(node.func)) diff --git a/patch/triton/python_triton_compiler_compiler_py.patch b/patch/triton/python_triton_compiler_compiler_py.patch index c652bf96..fdd9e1b0 100644 --- a/patch/triton/python_triton_compiler_compiler_py.patch +++ b/patch/triton/python_triton_compiler_compiler_py.patch @@ -1,23 +1,23 @@ diff --git a/python/triton/compiler/compiler.py b/python/triton/compiler/compiler.py -index 81d7b3267..a838359de 100644 +index 1f38e9ddf..5585f4707 100644 --- a/python/triton/compiler/compiler.py +++ b/python/triton/compiler/compiler.py -@@ -15,6 +15,7 @@ import re - import functools +@@ -16,6 +16,7 @@ import functools import os import time + import copy +import sys # - ^\s*tt\.func\s+ : match the start of the string, any leading whitespace, the keyword func, # and any following whitespace -@@ -463,8 +464,10 @@ class CompiledKernel: +@@ -470,8 +471,9 @@ class CompiledKernel: 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 = driver.active.utils.load_binary( -+ # self.module, self.function, self.n_regs, self.n_spills, self.n_max_threads = driver.active.utils.load_binary( +- self.name, self.kernel, self.metadata.shared, device) + self.module, self.function, self.n_regs, self.n_spills = driver.active.utils.load_binary( - self.name, self.kernel, self.metadata.shared, device) ++ self.metadata.kernel_name, self.kernel, self.metadata.shared, device, self.metadata.mix_mode) + self.n_max_threads = sys.maxsize # sys.maxsize warp_size = driver.active.get_current_target().warp_size if self.metadata.num_warps * warp_size > self.n_max_threads: diff --git a/patch/triton/python_triton_language_semantic_py.patch b/patch/triton/python_triton_language_semantic_py.patch index 7fab092d..1c3ca20a 100644 --- a/patch/triton/python_triton_language_semantic_py.patch +++ b/patch/triton/python_triton_language_semantic_py.patch @@ -1,5 +1,5 @@ diff --git a/python/triton/language/semantic.py b/python/triton/language/semantic.py -index 8341c2c4a..1fca39d66 100644 +index 5fc9cf67e..c8d1217a3 100644 --- a/python/triton/language/semantic.py +++ b/python/triton/language/semantic.py @@ -115,6 +115,7 @@ class TritonSemantic(Generic[TensorTy]): diff --git a/patch/triton/setup_py.patch b/patch/triton/setup_py.patch index f2f302ef..b123a108 100644 --- a/patch/triton/setup_py.patch +++ b/patch/triton/setup_py.patch @@ -1,18 +1,18 @@ diff --git a/setup.py b/setup.py -index dffc45ab1..a4fe0b8a8 100644 +index 8a3007ce4..68a1d9447 100644 --- a/setup.py +++ b/setup.py @@ -182,6 +182,13 @@ class Package: # json def get_json_package_info(): url = "https://github.com/nlohmann/json/releases/download/v3.11.3/include.zip" -+ local_json_path = os.environ.get("JSON_PATH34", "") ++ local_json_path = os.environ.get("JSON_PATH35", "") + if local_json_path != "": + if os.path.exists(local_json_path): + url = "file://" + local_json_path -+ print(f"JSON_PATH34 {local_json_path} exists. url={url}") ++ print(f"JSON_PATH35 {local_json_path} exists. url={url}") + else: -+ print(f"JSON_PATH34 {local_json_path} does not exist.") ++ print(f"JSON_PATH35 {local_json_path} does not exist.") return Package("json", "", url, "JSON_INCLUDE_DIR", "", "JSON_SYSPATH") @@ -20,35 +20,27 @@ index dffc45ab1..a4fe0b8a8 100644 # Create a stable symlink that doesn't include revision sym_name = f"llvm-{system_suffix}" url = f"https://oaitriton.blob.core.windows.net/public/llvm-builds/{name}.tar.gz" -+ local_llvm_tgz_path = os.environ.get("LLVM_TGZ_PATH34", "") ++ local_llvm_tgz_path = os.environ.get("LLVM_TGZ_PATH35", "") + if local_llvm_tgz_path != "": + if os.path.exists(local_llvm_tgz_path): + url = "file://" + local_llvm_tgz_path -+ print(f"LLVM_TGZ_PATH34 {local_llvm_tgz_path} exists. url={url}") ++ print(f"LLVM_TGZ_PATH35 {local_llvm_tgz_path} exists. url={url}") + else: -+ print(f"LLVM_TGZ_PATH34 {local_llvm_tgz_path} does not exist.") ++ print(f"LLVM_TGZ_PATH35 {local_llvm_tgz_path} does not exist.") return Package("llvm", name, url, "LLVM_INCLUDE_DIRS", "LLVM_LIBRARY_DIR", "LLVM_SYSPATH", sym_name=sym_name) -@@ -528,6 +542,37 @@ class CMakeBuild(build_ext): +@@ -528,6 +542,29 @@ class CMakeBuild(build_ext): subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=cmake_dir) subprocess.check_call(["cmake", "--build", ".", "--target", "mlir-doc"], cwd=cmake_dir) + if(check_env_flag("IS_NOT_PUBLISH","1")): + cmake_build_dir = get_cmake_dir() -+ tools_src_dir = os.path.join( -+ cmake_build_dir, -+ "third_party/dicp_triton/third_party/triton_shared/tools/triton-shared-opt/" -+ ) + dicp_opt_src = os.path.join( + cmake_build_dir, + "third_party/dicp_triton/tools/dicp_triton_opt/" + ) + shutil.copy( -+ os.path.join(tools_src_dir, "triton-shared-opt"), -+ os.path.join(extdir, "triton-shared-opt-v3_4") -+ ) -+ shutil.copy( + os.path.join(dicp_opt_src, "dicp_opt"), + extdir + ) @@ -68,7 +60,7 @@ index dffc45ab1..a4fe0b8a8 100644 def download_and_copy_dependencies(): nvidia_version_path = os.path.join(get_base_dir(), "cmake", "nvidia-toolchain-version.json") -@@ -649,6 +694,9 @@ def get_packages(): +@@ -649,6 +686,9 @@ def get_packages(): # Install the contents of each backend's `tools` directory into # `triton.tools.extra`. for x in os.listdir(backend.tools_dir): @@ -78,7 +70,7 @@ index dffc45ab1..a4fe0b8a8 100644 yield f"triton.tools.extra.{x}" if check_env_flag("TRITON_BUILD_PROTON", "ON"): # Default ON -@@ -795,15 +843,15 @@ PYTHON_CLASSIFIERS = [ +@@ -795,17 +835,18 @@ PYTHON_CLASSIFIERS = [ CLASSIFIERS = BASE_CLASSIFIERS + PYTHON_CLASSIFIERS setup( @@ -91,13 +83,16 @@ index dffc45ab1..a4fe0b8a8 100644 description="A language and compiler for custom Deep Learning operations", long_description="", install_requires=[ - "setuptools>=40.8.0", "importlib-metadata; python_version < '3.10'", + "torch_npu>=2.6.0", ], packages=list(get_packages()), package_dir=dict(get_package_dirs()), -@@ -824,7 +872,7 @@ setup( ++ package_data={"triton.backends.dicp_triton": ["npu_utils.cpp"]}, + entry_points=get_entry_points(), + include_package_data=True, + ext_modules=[CMakeExtension("triton", "triton/_C/")], +@@ -823,7 +864,7 @@ setup( zip_safe=False, # for PyPI keywords=["Compiler", "Deep Learning"], diff --git a/patch/triton/unittest_googletest_cmake.patch b/patch/triton/unittest_googletest_cmake.patch index c65cc6ed..e9b2a775 100644 --- a/patch/triton/unittest_googletest_cmake.patch +++ b/patch/triton/unittest_googletest_cmake.patch @@ -6,8 +6,8 @@ index 064dc8860..4ba22f32a 100644 include(FetchContent) -set(GOOGLETEST_DIR "" CACHE STRING "Location of local GoogleTest repo to build against") -+if(DEFINED ENV{GOOGLETEST_DIR34}) -+ set(GOOGLETEST_DIR $ENV{GOOGLETEST_DIR34}) ++if(DEFINED ENV{GOOGLETEST_DIR35}) ++ set(GOOGLETEST_DIR $ENV{GOOGLETEST_DIR35}) +else() + set(GOOGLETEST_DIR "") +endif() diff --git a/patch/ttshared/triton_shared.patch b/patch/ttshared/triton_shared.patch deleted file mode 100644 index b2f6c863..00000000 --- a/patch/ttshared/triton_shared.patch +++ /dev/null @@ -1,116 +0,0 @@ -diff --git a/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td b/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td -index 17c3ce9..30b88eb 100644 ---- a/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td -+++ b/include/triton-shared/Dialect/TritonStructured/IR/TritonStructuredDialect.td -@@ -233,8 +233,8 @@ def TTS_GetStructuredStateOp : TTS_Op<"get_structured_state", [AttrSizedResultSe - let summary = "Placeholder for the structured pointer states computed during PtrAnalysis."; - let description = "Used to pass the offsets and strides to scf.for op to simplify IR rewrites."; - -- let arguments = (ins AnyTypeOf<[TT_PtrLike, I32Tensor]>:$input); -- let results = (outs AnyTypeOf<[TT_PtrLike, I32Tensor]>:$structured, Variadic:$offsets, Variadic:$strides); -+ let arguments = (ins AnyTypeOf<[TT_PtrLike, TT_IndexTensorLike]>:$input); -+ let results = (outs AnyTypeOf<[TT_PtrLike, TT_IndexTensorLike]>:$structured, Variadic:$offsets, Variadic:$strides); - - let builders = [ - OpBuilder<(ins "Value":$input)>, -diff --git a/lib/Analysis/UseAnalysis.cpp b/lib/Analysis/UseAnalysis.cpp -index 62e4508..2d81db4 100644 ---- a/lib/Analysis/UseAnalysis.cpp -+++ b/lib/Analysis/UseAnalysis.cpp -@@ -120,12 +120,19 @@ LogicalResult triton::runUseAnalysis(triton::FuncOp &funcOp) { - LLVM_DEBUG({ op->setAttr("Undefined", UnitAttr::get(context)); }); - return; - } else if (useType == UseType::MetaUse) { -- assert(op->getNumResults() == 1 && -- "Ops used for meta computation are expected to have one result"); -- // Only set the tag if the operation uses tensors -- if (isa(op->getResult(0).getType())) { -- // Setting tag for erasing op later -- op->setAttr("MetaUse", UnitAttr::get(context)); -+ auto opName = op->getName().getStringRef(); -+ // if (!isa(op)) { -+ // assert(op->getNumResults() == 1 && -+ // "Ops used for meta computation are expected to have one result"); -+ // } -+ for (auto it = 0; it < op->getNumResults(); ++it) { -+ // Only set the tag if the operation uses tensors -+ if (isa(op->getResult(it).getType()) || -+ (isa(op) && -+ isa(op->getResult(it).getType()))) { -+ // Setting tag for erasing op later -+ op->setAttr("MetaUse", UnitAttr::get(context)); -+ } - } - return; - } else if (useType == UseType::DataUse) { -diff --git a/lib/Conversion/TritonArithToLinalg/CMakeLists.txt b/lib/Conversion/TritonArithToLinalg/CMakeLists.txt -index 432dee7..f357483 100644 ---- a/lib/Conversion/TritonArithToLinalg/CMakeLists.txt -+++ b/lib/Conversion/TritonArithToLinalg/CMakeLists.txt -@@ -1,6 +1,7 @@ - add_triton_library(TritonArithToLinalg - TritonArithToLinalg.cpp - TritonArithToLinalgPass.cpp -+ ${TRTION_SHARED_NPU_SPECIFIC_SOURCES}/TritonArithToLinalg/NPUSpecific.cpp - - DEPENDS - TritonArithToLinalgConversionPassIncGen -@@ -20,4 +21,6 @@ add_triton_library(TritonArithToLinalg - TritonTilingExtIR - TritonStructuredIR - TritonSharedUtils -+ -+ DICPUtils - ) -diff --git a/lib/Conversion/TritonArithToLinalg/TritonArithToLinalgPass.cpp b/lib/Conversion/TritonArithToLinalg/TritonArithToLinalgPass.cpp -index 8c86fe2..7a05c59 100644 ---- a/lib/Conversion/TritonArithToLinalg/TritonArithToLinalgPass.cpp -+++ b/lib/Conversion/TritonArithToLinalg/TritonArithToLinalgPass.cpp -@@ -4,6 +4,8 @@ - // Licensed under the MIT license. - // - //===----------------------------------------------------------------------===// -+#include "dicp/Conversion/TritonToLinalgNPU/TritonArithToLinalg/NPUSpecific.h" -+#include "dicp/Utils/Utils.h" - - #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" - #include "mlir/IR/BuiltinTypeInterfaces.h" -@@ -26,7 +28,7 @@ - - using namespace mlir; - using namespace triton; -- -+using namespace mlir::dicp; - namespace mlir { - namespace triton { - #define GEN_PASS_DEF_TRITONARITHTOLINALG -@@ -124,7 +126,10 @@ public: - target.addLegalOp(); - - target.addDynamicallyLegalDialect( -- [](Operation *op) { -+ [&](Operation *op) { -+ if (isAscendBackend(moduleOp)) -+ return linked::isLegalConstantAndTensorArithmeticOpForNPU(op); -+ - // Lower dense constant to linalg.fill - if (auto constOp = dyn_cast(op)) { - if (!isa(constOp.getResult().getType())) { -@@ -185,9 +190,14 @@ public: - target.addLegalOp(); - } - -- triton::populateTritonArithToLinalgConversionPatterns( -- pidsToFuncArgs, addptrToLinalg, assertToCf, transposeReduceToRank0, -- patterns); -+ DISPATCH_BACKEND_CONVERSION_PATTERNS( -+ mlir::dicp::getBackend(moduleOp), -+ linked::populateTritonArithToLinalgNPUConversionPatterns( -+ pidsToFuncArgs, addptrToLinalg, assertToCf, -+ transposeReduceToRank0, patterns), -+ triton::populateTritonArithToLinalgConversionPatterns( -+ pidsToFuncArgs, addptrToLinalg, assertToCf, -+ transposeReduceToRank0, patterns)); - - if (pidsToFuncArgs) { - for (auto func : getOperation().getOps()) { diff --git a/requirements.txt b/requirements.txt index ddeb35e1..ffced462 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,90 +1,38 @@ -absl-py==2.3.1 -Arpeggio==2.0.3 -attrs==24.2.0 -auditwheel==6.4.2 -autopep8==2.3.2 -backports.tarfile==1.2.0 -caliper-reader==0.4.1 -certifi==2025.10.5 -cffi==2.0.0 -charset-normalizer==3.4.4 -cmake==4.1.0 -contourpy==1.3.2 -cryptography==46.0.3 -cycler==0.12.1 -Cython==3.1.4 -decorator==5.1.1 -dill==0.4.0 -docutils==0.22.2 -einops==0.8.1 -exceptiongroup==1.3.0 -execnet==2.1.1 -expecttest==0.3.0 -filelock==3.20.0 -fonttools==4.60.1 -fsspec==2025.9.0 -id==1.5.0 -idna==3.11 -importlib_metadata==8.7.0 -iniconfig==2.1.0 -isort==7.0.0 -jaraco.classes==3.4.0 -jaraco.context==6.0.1 -jaraco.functools==4.3.0 -jeepney==0.9.0 -Jinja2==3.1.6 -keyring==25.6.0 -kiwisolver==1.4.9 -llnl-hatchet==2025.1.0 -markdown-it-py==4.0.0 -MarkupSafe==3.0.3 -matplotlib==3.10.7 -mdurl==0.1.2 -more-itertools==10.8.0 -mpmath==1.3.0 -multiprocess==0.70.18 -nanobind==2.9.2 -networkx==3.4.2 -nh3==0.3.1 -ninja==1.13.0 -numpy==1.26.4 -packaging==25.0 -pandas==2.3.3 -patchelf==0.18.0.0 -pathlib2==2.3.7.post1 -pillow==12.0.0 -pluggy==1.6.0 -protobuf==3.20.0 -psutil==6.0.0 -py==1.11.0 -pybind11==3.0.1 -pycodestyle==2.14.0 -pycparser==2.23 -pydot==4.0.1 -pyelftools==0.32 -Pygments==2.19.2 -pyparsing==3.2.5 -pytest==8.3.2 -pytest-forked==1.6.0 -pytest-xdist==3.6.1 -python-dateutil==2.9.0.post0 -pytz==2025.2 -PyYAML==6.0.3 -readme_renderer==44.0 -requests==2.32.5 -requests-toolbelt==1.0.0 -rfc3986==2.0.0 -rich==14.2.0 -scipy==1.13.1 -SecretStorage==3.4.0 -six==1.17.0 -sympy==1.13.1 -textX==4.2.3 -tomli==2.3.0 -torch==2.6.0 -torch_npu==2.6.0 -twine==6.2.0 -typing_extensions==4.15.0 -tzdata==2025.2 -urllib3==2.5.0 -zipp==3.23.0 +setuptools>=40.8.0 +wheel +cmake>=3.20,<4.0 +ninja>=1.11.1 +pybind11>=2.13.1 +lit +tabulate +sphinx +matplotlib +myst_parser +sphinx-rtd-theme +pandas +pytest +sphinx-gallery +sphinx-multiversion +llnl-hatchet +attrs +cython +numpy<2 +numpoly +decorator +sympy +cffi +pyyaml +pathlib2 +psutil +protobuf==3.20 +scipy +requests +absl-py +auditwheel +einops +expecttest +patchelf +pytest-xdist +twine +nanobind>=2.4 +pytest-timeout diff --git a/scripts/ci_python_black.sh b/scripts/ci_python_black.sh index d76a98c6..b7cc51a5 100644 --- a/scripts/ci_python_black.sh +++ b/scripts/ci_python_black.sh @@ -11,6 +11,5 @@ fi echo "Found modified Python files for PR:" echo "$modified_py_files" | tr '\n' ' ' # 显示修改的文件列表 -# 使用 Black 检查这些文件（只检查，不修改） -pip install black +pip install "black==25.*" black --check --diff $modified_py_files \ No newline at end of file diff --git a/test/ascend/autotune/01-vector-add.py b/test/ascend/autotune/01-vector-add.py new file mode 100644 index 00000000..1486d0c7 --- /dev/null +++ b/test/ascend/autotune/01-vector-add.py @@ -0,0 +1,48 @@ +import os + +import torch +import torch_npu +import triton +import triton.language as tl +from backend.testing import do_bench_npu +import triton.backends.dicp_triton.ascend_autotune_hooks + + +@triton.autotune(configs=[], key=["n_elements"]) +@triton.jit +def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + tl.store(output_ptr + offsets, output, mask=mask) + + +def add_torch(x, y): + return x + y + + +def add_autotune(x, y): + output = torch.empty_like(x) + n_elements = output.numel() + add_kernel[lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)]( + x, y, output, n_elements + ) + return output + + +def test_add(size: int): + x = torch.rand(size, device="npu") + y = torch.rand(size, device="npu") + + output_torch = add_torch(x, y) + output_triton = add_autotune(x, y) + assert torch.allclose(output_triton, output_torch) + print(f"Vector Add {size} PASSED!") + + +if __name__ == "__main__": + test_add(98432) diff --git a/test/ascend/autotune/02-fused-softmax.py b/test/ascend/autotune/02-fused-softmax.py new file mode 100644 index 00000000..659fea80 --- /dev/null +++ b/test/ascend/autotune/02-fused-softmax.py @@ -0,0 +1,73 @@ +import os + +import torch +import torch_npu +import triton +import triton.language as tl +from backend.testing import do_bench_npu +import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 — install proxy before @triton.autotune + + +@triton.autotune( + configs=[], + key=["n_rows", "n_cols"], +) +@triton.jit +def softmax_kernel( + output_ptr, + input_ptr, + input_row_stride, + output_row_stride, + n_rows, + n_cols, + BLOCK_SIZE: tl.constexpr, + XBLOCK: tl.constexpr, + XBLOCK_SUB: tl.constexpr, +): + row_start = tl.program_id(0) * XBLOCK + for row_idx in tl.range(0, XBLOCK, XBLOCK_SUB): + row_offsets = row_start + row_idx + tl.arange(0, XBLOCK_SUB)[:, None] + col_offsets = tl.arange(0, BLOCK_SIZE)[None, :] + xmask = row_offsets < n_rows + ymask = col_offsets < n_cols + mask = xmask & ymask + input_ptrs = input_ptr + (row_offsets * input_row_stride + col_offsets) + row = tl.load(input_ptrs, mask=mask, other=-float("inf")) + row_minus_max = row - tl.max(row, axis=1).reshape(XBLOCK_SUB, 1).broadcast_to( + XBLOCK_SUB, BLOCK_SIZE + ) + numerator = tl.exp(row_minus_max) + denominator = ( + tl.sum(numerator, axis=1) + .reshape(XBLOCK_SUB, 1) + .broadcast_to(XBLOCK_SUB, BLOCK_SIZE) + ) + softmax_output = numerator / denominator + output_ptrs = output_ptr + (row_offsets * output_row_stride + col_offsets) + tl.store(output_ptrs, softmax_output, mask=mask) + + +def softmax_torch(x): + return torch.softmax(x, axis=-1) + + +def softmax_autotune(x): + n_rows, n_cols = x.shape + BLOCK_SIZE = n_cols + y = torch.empty_like(x) + softmax_kernel[lambda meta: (triton.cdiv(n_rows, meta["XBLOCK"]), 1, 1)]( + y, x, x.stride(0), y.stride(0), n_rows, n_cols, BLOCK_SIZE=BLOCK_SIZE + ) + return y + + +def test_softmax(shape, dtype): + x = torch.randn(shape, dtype=dtype, device="npu") + y_torch = softmax_torch(x) + y_triton = softmax_autotune(x) + assert torch.allclose(y_triton, y_torch) + print(f"Fused Softmax {shape} {dtype} PASSED!") + + +if __name__ == "__main__": + test_softmax((16896, 1024), torch.float32) diff --git a/test/ascend/autotune/03-layer-norm.py b/test/ascend/autotune/03-layer-norm.py new file mode 100644 index 00000000..f5568afe --- /dev/null +++ b/test/ascend/autotune/03-layer-norm.py @@ -0,0 +1,105 @@ +import os + +import torch +import torch_npu +import triton +import triton.language as tl +from backend.testing import do_bench_npu +import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 — install proxy before @triton.autotune + + +@triton.autotune( + configs=[], + key=["M", "N"], +) +@triton.jit +def _layer_norm_fwd_fused( + X, + Y, + W, + B, + Mean, + Rstd, + stride, + N, + M, + eps, + XBLOCK_SIZE: tl.constexpr, + RBLOCK_SIZE: tl.constexpr, +): + row_begin = tl.program_id(0) * XBLOCK_SIZE + row_idx = row_begin + tl.arange(0, XBLOCK_SIZE) + row_mask = row_idx < M + row_offsets = row_idx[:, None] * stride + _mean = tl.zeros((XBLOCK_SIZE, RBLOCK_SIZE), dtype=tl.float32) + for off in range(0, N, RBLOCK_SIZE): + col_idx = off + tl.arange(0, RBLOCK_SIZE) + col_mask = col_idx < N + mask = row_mask[:, None] & col_mask[None, :] + a = tl.load(X + row_offsets + col_idx[None, :], mask=mask, other=0.0).to( + tl.float32 + ) + _mean += a + mean = tl.sum(_mean, axis=1, keep_dims=True) / N + _var = tl.zeros((XBLOCK_SIZE, RBLOCK_SIZE), dtype=tl.float32) + for off in range(0, N, RBLOCK_SIZE): + col_idx = off + tl.arange(0, RBLOCK_SIZE) + col_mask = col_idx < N + mask = row_mask[:, None] & col_mask[None, :] + x = tl.load(X + row_offsets + col_idx[None, :], mask=mask, other=0.0).to( + tl.float32 + ) + x = tl.where(mask, x - mean, 0.0) + _var += x * x + var = tl.sum(_var, axis=1, keep_dims=True) / N + rstd = 1 / tl.sqrt(var + eps) + tl.store(Mean + row_idx[:, None], mean, mask=row_mask[:, None]) + tl.store(Rstd + row_idx[:, None], rstd, mask=row_mask[:, None]) + for off in range(0, N, RBLOCK_SIZE): + col_idx = off + tl.arange(0, RBLOCK_SIZE) + col_mask = col_idx < N + mask = row_mask[:, None] & col_mask[None, :] + w = tl.load(W + col_idx, mask=col_mask).reshape((1, RBLOCK_SIZE)) + b = tl.load(B + col_idx, mask=col_mask).reshape((1, RBLOCK_SIZE)) + x = tl.load(X + row_offsets + col_idx[None, :], mask=mask, other=0.0).to( + tl.float32 + ) + x_hat = (x - mean) * rstd + y = x_hat * w + b + tl.store(Y + row_offsets + col_idx[None, :], y, mask=mask) + + +def layer_norm_torch(args): + x, w_shape, weight, bias, eps, dtype = args + return torch.nn.functional.layer_norm(x, w_shape, weight, bias, eps).to(dtype) + + +def layer_norm_autotune(args): + x, weight, bias, eps = args + y = torch.empty_like(x) + x_arg = x.reshape(-1, x.shape[-1]) + M, N = x_arg.shape + mean = torch.empty((M,), dtype=torch.float32, device=x.device) + rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + _layer_norm_fwd_fused[lambda meta: (triton.cdiv(M, meta["XBLOCK_SIZE"]), 1, 1)]( + x_arg, y, weight, bias, mean, rstd, x_arg.stride(0), N, M, eps + ) + return y + + +def test_layer_norm(shape, dtype, eps=1e-5): + M, N = shape + device = "npu" + x_shape = shape + w_shape = (x_shape[-1],) + weight = torch.rand(w_shape, dtype=dtype, device=device) + bias = torch.rand(w_shape, dtype=dtype, device=device) + x = -2.3 + 0.5 * torch.randn(x_shape, dtype=dtype, device=device) + y_torch = layer_norm_torch((x, w_shape, weight, bias, eps, dtype)) + y_triton = layer_norm_autotune((x, weight, bias, eps)) + assert torch.allclose(y_triton, y_torch, atol=1e-2, rtol=0) + print(f"Layer Normalization {M},{N} {dtype} PASSED!") + + +if __name__ == "__main__": + test_layer_norm((128, 32), torch.float16) diff --git a/test/ascend/autotune/04-libentry.py b/test/ascend/autotune/04-libentry.py new file mode 100644 index 00000000..a229bd0b --- /dev/null +++ b/test/ascend/autotune/04-libentry.py @@ -0,0 +1,59 @@ +import os + +import torch +import torch_npu +import triton +import triton.language as tl + +from language.deeplink.runtime import libentry + +from backend.testing import do_bench_npu +import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 — install proxy before @triton.autotune + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 1 * 1024, "multibuffer": True}), + triton.Config({"BLOCK_SIZE": 12 * 1024, "multibuffer": True}), + triton.Config({"BLOCK_SIZE": 12 * 1024, "multibuffer": False}), + triton.Config({"BLOCK_SIZE": 8 * 1024, "multibuffer": True}), + ], + key=["n_elements"], +) +@libentry() +@triton.jit +def add_kernel(x_ptr, y_ptr, output_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + tl.store(output_ptr + offsets, output, mask=mask) + + +def add_torch(x, y): + return x + y + + +def add_autotune(x, y): + output = torch.empty_like(x) + n_elements = output.numel() + add_kernel[lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)]( + x, y, output, n_elements + ) + return output + + +def test_add(size: int): + x = torch.rand(size, device="npu") + y = torch.rand(size, device="npu") + output_torch = add_torch(x, y) + output_triton = add_autotune(x, y) + assert torch.allclose(output_triton, output_torch) + print(f"Vector Add {size} with libentry PASSED!") + + +if __name__ == "__main__": + test_add(98432) diff --git a/test/ascend/autotune/conftest.py b/test/ascend/autotune/conftest.py new file mode 100644 index 00000000..a2e61e64 --- /dev/null +++ b/test/ascend/autotune/conftest.py @@ -0,0 +1,4 @@ +# The proxy replaces triton.autotune / triton.max_autotune at import time. +# We import this early so every test module sees the proxy (which auto-detects +# ascend via triton.runtime.driver.active.target at call time). +import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 — side-effect import diff --git a/test/ascend/autotune/demo_backend_runtime_autotune.py b/test/ascend/autotune/demo_backend_runtime_autotune.py new file mode 100644 index 00000000..401aab40 --- /dev/null +++ b/test/ascend/autotune/demo_backend_runtime_autotune.py @@ -0,0 +1,48 @@ +import torch +import torch_npu +import triton +import triton.language as tl + +from backend.ascend_autotune_runtime import autotune as ascend_autotune + + +@ascend_autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 256}, num_warps=4), + triton.Config({"BLOCK_SIZE": 512}, num_warps=4), + triton.Config({"BLOCK_SIZE": 1024}, num_warps=8), + ], + key=["n_elements"], + hints={"compile_options": "vector"}, +) +@triton.jit +def add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + tl.store(out_ptr + offsets, x + y, mask=mask) + + +def add(x, y): + out = torch.empty_like(x) + n_elements = out.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + add_kernel[grid](x, y, out, n_elements) + return out + + +def main(): + n_elements = 98432 + x = torch.rand(n_elements, device="npu", dtype=torch.float32) + y = torch.rand(n_elements, device="npu", dtype=torch.float32) + + out = add(x, y) + torch.npu.synchronize() + torch.testing.assert_close(out, x + y, rtol=1e-3, atol=1e-3) + print("backend.ascend_autotune_runtime autotune demo PASSED") + + +if __name__ == "__main__": + main() diff --git a/test/ascend/autotune/test_autotune_doc_e2e.py b/test/ascend/autotune/test_autotune_doc_e2e.py new file mode 100644 index 00000000..4150b82e --- /dev/null +++ b/test/ascend/autotune/test_autotune_doc_e2e.py @@ -0,0 +1,275 @@ +import os + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + +import triton.backends.dicp_triton.ascend_autotune_hooks # noqa: F401 - install proxy before decorators + +os.environ.setdefault("TRITON_AUTOTUNE_PARALLEL_COMPILE", "0") + + +@triton.autotune( + configs=[ + triton.Config({"XS": 128, "multibuffer": True}), + triton.Config({"XS": 1024, "multibuffer": True}), + triton.Config({"XS": 1024, "multibuffer": False}), + ], + key=["numel"], +) +@triton.jit +def _explicit_config_exp_add_kernel(out_ptr, x_ptr, y_ptr, numel, XS: tl.constexpr): + offsets = tl.program_id(0) * XS + tl.arange(0, XS) + mask = offsets < numel + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + y = tl.load(y_ptr + offsets, mask=mask, other=0.0) + out = tl.full((XS,), 0.0, tl.float32) + for i in range(8): + out = tl.exp(x) + y + i + tl.store(out_ptr + offsets, out, mask=mask) + + +def _explicit_config_exp_add(x, y): + out = torch.empty_like(x) + numel = out.numel() + grid = lambda meta: (triton.cdiv(numel, meta["XS"]), 1, 1) + _explicit_config_exp_add_kernel[grid](out, x, y, numel) + return out + + +@triton.autotune(configs=[], key=["n_elements"]) +@triton.jit +def _auto_tiling_add_kernel( + x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + y = tl.load(y_ptr + offsets, mask=mask, other=0.0) + tl.store(out_ptr + offsets, x + y, mask=mask) + + +def _auto_tiling_add(x, y): + out = torch.empty_like(x) + n_elements = out.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) + _auto_tiling_add_kernel[grid](x, y, out, n_elements) + return out + + +@triton.autotune( + configs=[], + key=["n_elements"], + hints={"compile_options": "vector"}, +) +@triton.jit +def _auto_tiling_compile_options_add_kernel( + x_ptr, + y_ptr, + out_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + y = tl.load(y_ptr + offsets, mask=mask, other=0.0) + tl.store(out_ptr + offsets, x + y * 3.0, mask=mask) + + +def _auto_tiling_compile_options_add(x, y): + out = torch.empty_like(x) + n_elements = out.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) + _auto_tiling_compile_options_add_kernel[grid](x, y, out, n_elements) + return out + + +@triton.autotune( + configs=[], + key={"x": "n_elements"}, + hints={ + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + }, +) +@triton.jit +def _hinted_tiling_add_kernel( + x_ptr, + y_ptr, + out_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, + BLOCK_SIZE_SUB: tl.constexpr, +): + block_start = tl.program_id(0) * BLOCK_SIZE + sub_offsets = tl.arange(0, BLOCK_SIZE_SUB) + loops = (BLOCK_SIZE + BLOCK_SIZE_SUB - 1) // BLOCK_SIZE_SUB + for loop in range(loops): + offsets = block_start + loop * BLOCK_SIZE_SUB + sub_offsets + mask = offsets < min(block_start + BLOCK_SIZE, n_elements) + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + y = tl.load(y_ptr + offsets, mask=mask, other=0.0) + tl.store(out_ptr + offsets, x + y, mask=mask) + + +def _hinted_tiling_add(x, y): + out = torch.empty_like(x) + n_elements = out.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) + _hinted_tiling_add_kernel[grid](x, y, out, n_elements) + return out + + +@triton.autotune( + configs=[triton.Config({"BLOCK_SIZE": 128, "multibuffer": False})], + key=["n_elements"], + hints={"auto_gen_config": True}, +) +@triton.jit +def _auto_and_user_config_add_kernel( + x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + y = tl.load(y_ptr + offsets, mask=mask, other=0.0) + tl.store(out_ptr + offsets, x - y, mask=mask) + + +def _auto_and_user_config_add(x, y): + out = torch.empty_like(x) + n_elements = out.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) + _auto_and_user_config_add_kernel[grid](x, y, out, n_elements) + return out + + +@triton.max_autotune( + configs=[triton.Config({"BLOCK_SIZE": 256})], + key=["n_elements"], + kernel_type="vector", + num_stages=[1, 2], + enable_ubuf_saving=[True, False], +) +@triton.jit +def _max_autotune_vector_kernel( + x_ptr, y_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + y = tl.load(y_ptr + offsets, mask=mask, other=0.0) + tl.store(out_ptr + offsets, x * 2.0 + y, mask=mask) + + +def _max_autotune_vector(x, y): + out = torch.empty_like(x) + n_elements = out.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) + _max_autotune_vector_kernel[grid](x, y, out, n_elements) + return out + + +@triton.max_autotune( + configs=[triton.Config({"BLOCK_SIZE": 128})], + key=["n_elements"], + kernel_type="vector", + enable_ubuf_saving=[True, False], +) +@triton.jit +def _max_autotune_default_stage_kernel( + x_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr +): + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + tl.store(out_ptr + offsets, x + 1.0, mask=mask) + + +def _max_autotune_default_stage(x): + out = torch.empty_like(x) + n_elements = out.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]), 1, 1) + _max_autotune_default_stage_kernel[grid](x, out, n_elements) + return out + + +@pytest.mark.autotune +def test_community_autotune_explicit_configs_e2e(): + x = torch.randn(4096, dtype=torch.float32, device="npu") + y = torch.randn(4096, dtype=torch.float32, device="npu") + + actual = _explicit_config_exp_add(x, y) + expected = torch.exp(x) + y + 7 + + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + + +@pytest.mark.autotune +def test_advanced_autotune_empty_configs_auto_tiling_e2e(): + x = torch.randn(4096, dtype=torch.float32, device="npu") + y = torch.randn(4096, dtype=torch.float32, device="npu") + + actual = _auto_tiling_add(x, y) + expected = x + y + + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + + +@pytest.mark.autotune +def test_advanced_autotune_auto_tiling_compile_options_e2e(): + x = torch.randn(4096, dtype=torch.float32, device="npu") + y = torch.randn(4096, dtype=torch.float32, device="npu") + + actual = _auto_tiling_compile_options_add(x, y) + expected = x + y * 3.0 + + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + + +@pytest.mark.autotune +def test_advanced_autotune_hints_dict_key_e2e(): + x = torch.randn(4096, dtype=torch.float32, device="npu") + y = torch.randn(4096, dtype=torch.float32, device="npu") + + actual = _hinted_tiling_add(x, y) + expected = x + y + + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + + +@pytest.mark.autotune +def test_advanced_autotune_user_configs_merge_auto_configs_e2e(): + x = torch.randn(4096, dtype=torch.float32, device="npu") + y = torch.randn(4096, dtype=torch.float32, device="npu") + + actual = _auto_and_user_config_add(x, y) + expected = x - y + + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + + +@pytest.mark.autotune +def test_max_autotune_vector_expanded_configs_e2e(): + x = torch.randn(4096, dtype=torch.float32, device="npu") + y = torch.randn(4096, dtype=torch.float32, device="npu") + + actual = _max_autotune_vector(x, y) + expected = x * 2.0 + y + + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + + +@pytest.mark.autotune +def test_max_autotune_uses_ascend_default_num_stages_e2e(): + x = torch.randn(4096, dtype=torch.float32, device="npu") + + actual = _max_autotune_default_stage(x) + expected = x + 1.0 + + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) diff --git a/test/ascend/autotune/test_autotune_param_valid.py b/test/ascend/autotune/test_autotune_param_valid.py new file mode 100644 index 00000000..e4fdca42 --- /dev/null +++ b/test/ascend/autotune/test_autotune_param_valid.py @@ -0,0 +1,171 @@ +import os + +import pytest +import torch +import torch_npu +import triton +import triton.language as tl + + +@triton.autotune( + configs=[], + key={"x": "n_elements"}, + hints={ + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + }, +) +@triton.jit +def add_kernel( + x_ptr, + y_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, + BLOCK_SIZE_SUB: tl.constexpr, +): + offset = tl.program_id(0) * BLOCK_SIZE + loops1 = (BLOCK_SIZE + BLOCK_SIZE_SUB - 1) // BLOCK_SIZE_SUB + for loop in range(0, loops1): + x0 = offset + loop * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE_SUB) + mask = x0 < n_elements + x = tl.load(x_ptr + x0, mask) + y = tl.load(y_ptr + x0, mask) + output = x + y + tl.store(output_ptr + x0, output) + + +def add_torch(x, y): + return x + y + + +def add_autotune(x, y): + output = torch.empty_like(x) + n_elements = output.numel() + add_kernel[lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)]( + x, y, output, n_elements + ) + return output + + +@pytest.mark.autotune +@pytest.mark.parametrize( + "size", + [ + 2048, + ], +) +def test_add(size: int): + x = torch.rand(size, device="npu") + y = torch.rand(size, device="npu") + + output_torch = add_torch(x, y) + output_triton = add_autotune(x, y) + assert torch.allclose(output_triton, output_torch) + + +@pytest.mark.autotune +def test_add_no_reduction_axes(): + try: + + @triton.autotune( + configs=[], + key={"x": "n_elements"}, + hints={ + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "low_dim_axes": ["x"], + }, + ) + @triton.jit + def add_kernel_exception(): + pass + + except ValueError as e: + assert "reduction_axes must be a list" in str(e) + + +@pytest.mark.autotune +def test_add_no_low_dim_axes(): + try: + + @triton.autotune( + configs=[], + key={"x": "n_elements"}, + hints={ + "split_params": {"x": "BLOCK_SIZE"}, + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "reduction_axes": [], + }, + ) + @triton.jit + def add_kernel_exception(): + pass + + except ValueError as e: + assert "low_dim_axes must be a list" in str(e) + + +@pytest.mark.autotune +def test_add_no_tiling_params(): + try: + + @triton.autotune( + configs=[], + key={"x": "n_elements"}, + hints={ + "split_params": {"x": "BLOCK_SIZE"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + }, + ) + @triton.jit + def add_kernel_exception(): + pass + + except ValueError as e: + assert "tiling_params must be a dict" in str(e) + + +@pytest.mark.autotune +def test_add_no_split_params(): + try: + + @triton.autotune( + configs=[], + key={"x": "n_elements"}, + hints={ + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + }, + ) + @triton.jit + def add_kernel_exception(): + pass + + except ValueError as e: + assert "split_params must be a dict" in str(e) + + +@pytest.mark.autotune +def test_add_no_keyname(): + try: + + @triton.autotune( + configs=[], + key={"x0": "n_elements"}, + hints={ + "tiling_params": {"x": "BLOCK_SIZE_SUB"}, + "low_dim_axes": ["x"], + "reduction_axes": [], + }, + ) + @triton.jit + def add_kernel_exception(): + pass + + except ValueError as e: + assert "All keys in 'key' must be valid axis names" in str(e) diff --git a/test/ascend/autotune/test_common.py b/test/ascend/autotune/test_common.py new file mode 100644 index 00000000..350ba874 --- /dev/null +++ b/test/ascend/autotune/test_common.py @@ -0,0 +1,83 @@ +import unittest.mock as mock +import pytest +import torch + + +def MockAutoTilingTunerRun(self, *args, **kwargs): + self.nargs = dict(zip(self.arg_names, args)) + + # generate key + all_args = {**self.nargs, **kwargs} + try: + self._autoparse_axis_params(all_args) + except ValueError as e: + if "Missing required arguments" in str(e): + pass + else: + raise + return { + "keys": self.keys, + "split_params": self.split_params, + "tiling_params": self.tiling_params, + "low_dim_axes": self.low_dim_axes, + "reduction_axes": self.reduction_axes, + "persistent_reduction": self.persistent_reduction, + } + + +def check_axes_parse_res(act: dict, ref: dict): + ref_keys = ref["keys"] + act_keys = act["keys"] + + assert set(ref_keys.values()) == set( + act_keys.values() + ), f"Semantic dimensions mismatch: ref={set(ref_keys.values())}, act={set(act_keys.values())}" + + def normalize_param_dict(param_dict: dict, sym_to_sem: dict) -> dict: + return {sym_to_sem[sym]: value for sym, value in param_dict.items()} + + ref_split = normalize_param_dict(ref["split_params"], ref_keys) + act_split = normalize_param_dict(act["split_params"], act_keys) + + ref_tiling = normalize_param_dict(ref["tiling_params"], ref_keys) + act_tiling = normalize_param_dict(act["tiling_params"], act_keys) + + def normalize_axis_list(axis_list: list, sym_to_sem: dict) -> list: + return sorted(sym_to_sem[sym] for sym in axis_list) + + ref_low = normalize_axis_list(ref["low_dim_axes"], ref_keys) + act_low = normalize_axis_list(act["low_dim_axes"], act_keys) + + ref_red = normalize_axis_list(ref["reduction_axes"], ref_keys) + act_red = normalize_axis_list(act["reduction_axes"], act_keys) + + assert ref_split == act_split, f"split_params mismatch: {ref_split} vs {act_split}" + assert ( + ref_tiling == act_tiling + ), f"tiling_params mismatch: {ref_tiling} vs {act_tiling}" + assert ref_low == act_low, f"low_dim_axes mismatch: {ref_low} vs {act_low}" + assert ref_red == act_red, f"reduction_axes mismatch: {ref_red} vs {act_red}" + + +@pytest.fixture +def mock_autotuner(): + with mock.patch( + "triton.backends.dicp_triton.ascend_autotune_runtime.autotuner.AutoTilingTuner.run", + new=MockAutoTilingTunerRun, + ): + yield + + +def generate_tensor(shape, dtype): + if dtype == "float32" or dtype == "float16" or dtype == "bfloat16": + return torch.randn(size=shape, dtype=eval("torch." + dtype)) + elif dtype == "int32" or dtype == "int64" or dtype == "int16": + return torch.randint(low=0, high=2000, size=shape, dtype=eval("torch." + dtype)) + elif dtype == "int8": + return torch.randint(low=0, high=127, size=shape, dtype=eval("torch." + dtype)) + elif dtype == "bool": + return torch.randint(low=0, high=2, size=shape).bool() + elif dtype == "uint8": + return torch.randint(low=0, high=255, size=shape, dtype=torch.uint8) + else: + raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) diff --git a/test/ascend/autotune/test_compile_options.py b/test/ascend/autotune/test_compile_options.py new file mode 100644 index 00000000..9a3d18eb --- /dev/null +++ b/test/ascend/autotune/test_compile_options.py @@ -0,0 +1,406 @@ +import pytest +import triton + +from backend.ascend_autotune_runtime.compile_options import ( + expand_compile_option_configs, + format_compile_option_result, + parse_compile_options_hint, +) + + +def test_compile_options_string_hint_expands_vector_defaults(): + spec = parse_compile_options_hint("vector") + configs = expand_compile_option_configs( + [triton.Config({"BLOCK_SIZE": 1024})], + spec, + generated_tiling=True, + ) + + assert len(configs) == 4 + assert {cfg.num_stages for cfg in configs} == {1, 2} + assert {cfg.kwargs["enable_ubuf_saving"] for cfg in configs} == {True, False} + assert {cfg.kwargs["enable_tuning_mode"] for cfg in configs} == {True} + assert {cfg.kwargs["BLOCK_SIZE"] for cfg in configs} == {1024} + + +def test_compile_options_string_hint_expands_mixcv_auto_search_space_without_limit(): + spec = parse_compile_options_hint("mixcv") + configs = expand_compile_option_configs( + [triton.Config({"BLOCK_SIZE": 1024})], + spec, + generated_tiling=True, + ) + + assert spec.max_configs is None + assert len(configs) == 156 + assert {cfg.num_stages for cfg in configs} == {1, 2} + assert {cfg.kwargs["enable_tuning_mode"] for cfg in configs} == {True} + assert {cfg.kwargs["unit_flag"] for cfg in configs} == {False, True} + assert {cfg.kwargs["enable_hivm_auto_cv_balance"] for cfg in configs} == {True} + assert {cfg.kwargs["enable_ubuf_saving"] for cfg in configs} == {False, True} + + stage1_configs = [cfg for cfg in configs if cfg.num_stages == 1] + assert len(stage1_configs) == 4 + assert all("multibuffer" not in cfg.kwargs for cfg in stage1_configs) + assert all( + "limit_auto_multi_buffer_only_for_local_buffer" not in cfg.kwargs + for cfg in stage1_configs + ) + assert all( + "limit_auto_multi_buffer_of_local_buffer" not in cfg.kwargs + for cfg in stage1_configs + ) + assert all("set_workspace_multibuffer" not in cfg.kwargs for cfg in stage1_configs) + assert all("tile_mix_vector_loop" not in cfg.kwargs for cfg in stage1_configs) + assert all("tile_mix_cube_loop" not in cfg.kwargs for cfg in stage1_configs) + assert all("enable_preload" not in cfg.kwargs for cfg in stage1_configs) + assert {cfg.kwargs["enable_auto_bind_sub_block"] for cfg in stage1_configs} == { + True + } + + stage2_configs = [cfg for cfg in configs if cfg.num_stages == 2] + assert len(stage2_configs) == 152 + assert all("multibuffer" not in cfg.kwargs for cfg in stage2_configs) + assert { + cfg.kwargs["limit_auto_multi_buffer_only_for_local_buffer"] + for cfg in stage2_configs + } == {False, True} + assert { + cfg.kwargs["limit_auto_multi_buffer_of_local_buffer"] for cfg in stage2_configs + } == {"no-limit", "no-l0c"} + assert { + cfg.kwargs["set_workspace_multibuffer"] + for cfg in stage2_configs + if "set_workspace_multibuffer" in cfg.kwargs + } == {2, 4} + assert { + cfg.kwargs["tile_mix_vector_loop"] + for cfg in stage2_configs + if "tile_mix_vector_loop" in cfg.kwargs + } == {1, 2, 4} + assert { + cfg.kwargs["tile_mix_cube_loop"] + for cfg in stage2_configs + if "tile_mix_cube_loop" in cfg.kwargs + } == {1, 2, 4} + assert {cfg.kwargs["enable_auto_bind_sub_block"] for cfg in stage2_configs} == { + True + } + assert any( + "set_workspace_multibuffer" in cfg.kwargs + and "tile_mix_vector_loop" in cfg.kwargs + and "tile_mix_cube_loop" in cfg.kwargs + for cfg in stage2_configs + ) + assert any( + "set_workspace_multibuffer" not in cfg.kwargs + and "tile_mix_vector_loop" not in cfg.kwargs + and "tile_mix_cube_loop" not in cfg.kwargs + for cfg in stage2_configs + ) + workspace_configs = [ + cfg for cfg in stage2_configs if "set_workspace_multibuffer" in cfg.kwargs + ] + assert len(workspace_configs) == 144 + assert {cfg.num_stages for cfg in workspace_configs} == {2} + assert { + cfg.kwargs["limit_auto_multi_buffer_only_for_local_buffer"] + for cfg in workspace_configs + } == {False} + + +def test_compile_options_workspace_pruned_when_auto_multibuffer_disabled(): + spec = parse_compile_options_hint("mixcv") + configs = expand_compile_option_configs( + [triton.Config({"BLOCK_SIZE": 1024})], + spec, + generated_tiling=True, + fixed_options={"multibuffer": False, "num_stages": 2}, + ) + + assert configs + assert all("set_workspace_multibuffer" not in cfg.kwargs for cfg in configs) + + +def test_compile_options_workspace_pruned_when_workspace_limit_enabled(): + spec = parse_compile_options_hint("mixcv") + configs = expand_compile_option_configs( + [triton.Config({"BLOCK_SIZE": 1024})], + spec, + generated_tiling=True, + fixed_options={ + "limit_auto_multi_buffer_only_for_local_buffer": True, + "num_stages": 2, + }, + ) + + assert configs + assert all("set_workspace_multibuffer" not in cfg.kwargs for cfg in configs) + + +def test_compile_options_explicit_mixcv_values_are_not_restricted_by_auto_search_space(): + spec = parse_compile_options_hint( + { + "kernel_type": "mixcv", + "num_stages": [2], + "multibuffer": [False], + "unit_flag": [True], + "limit_auto_multi_buffer_only_for_local_buffer": [True], + "limit_auto_multi_buffer_of_local_buffer": ["no-limit"], + "set_workspace_multibuffer": [4], + "enable_hivm_auto_cv_balance": [False], + "tile_mix_vector_loop": [8], + "tile_mix_cube_loop": [8], + "enable_ubuf_saving": [False], + "enable_auto_bind_sub_block": [False], + } + ) + configs = expand_compile_option_configs( + [triton.Config({"BLOCK_SIZE": 1024})], + spec, + generated_tiling=True, + ) + + assert len(configs) == 1 + assert configs[0].num_stages == 2 + assert configs[0].kwargs["enable_tuning_mode"] is True + assert configs[0].kwargs["multibuffer"] is False + assert configs[0].kwargs["unit_flag"] is True + assert configs[0].kwargs["enable_hivm_auto_cv_balance"] is False + assert configs[0].kwargs["enable_auto_bind_sub_block"] is False + assert "limit_auto_multi_buffer_only_for_local_buffer" not in configs[0].kwargs + assert "limit_auto_multi_buffer_of_local_buffer" not in configs[0].kwargs + assert "set_workspace_multibuffer" not in configs[0].kwargs + assert "tile_mix_vector_loop" not in configs[0].kwargs + assert "tile_mix_cube_loop" not in configs[0].kwargs + + +def test_compile_options_explicit_mixcv_stage1_values_are_preserved(): + spec = parse_compile_options_hint( + { + "kernel_type": "mixcv", + "num_stages": [1], + "multibuffer": [True], + "limit_auto_multi_buffer_only_for_local_buffer": [False], + "limit_auto_multi_buffer_of_local_buffer": ["no-l0c"], + "set_workspace_multibuffer": [2], + "tile_mix_vector_loop": [4], + "tile_mix_cube_loop": [4], + "enable_auto_bind_sub_block": [True], + } + ) + configs = expand_compile_option_configs( + [triton.Config({"BLOCK_SIZE": 1024})], + spec, + generated_tiling=True, + ) + + assert len(configs) == 4 + assert {cfg.num_stages for cfg in configs} == {1} + assert {cfg.kwargs["enable_tuning_mode"] for cfg in configs} == {True} + assert {cfg.kwargs["unit_flag"] for cfg in configs} == {False, True} + assert {cfg.kwargs["enable_ubuf_saving"] for cfg in configs} == {False, True} + assert all(cfg.kwargs["multibuffer"] is True for cfg in configs) + assert all("set_workspace_multibuffer" not in cfg.kwargs for cfg in configs) + assert all("tile_mix_vector_loop" not in cfg.kwargs for cfg in configs) + assert all("tile_mix_cube_loop" not in cfg.kwargs for cfg in configs) + + +def test_compile_options_auto_search_overrides_base_compile_options(): + spec = parse_compile_options_hint("mixcv") + configs = expand_compile_option_configs( + [ + triton.Config( + { + "BLOCK_SIZE": 1024, + "enable_tuning_mode": False, + "unit_flag": True, + "enable_ubuf_saving": True, + "set_workspace_multibuffer": 4, + "tile_mix_vector_loop": 4, + "tile_mix_cube_loop": 4, + }, + num_stages=1, + ) + ], + spec, + generated_tiling=False, + ) + + assert len(configs) == 156 + assert {cfg.kwargs["enable_tuning_mode"] for cfg in configs} == {True} + assert {cfg.kwargs["unit_flag"] for cfg in configs} == {False, True} + assert {cfg.kwargs["enable_ubuf_saving"] for cfg in configs} == {False, True} + assert all( + "multibuffer" not in cfg.kwargs for cfg in configs if cfg.num_stages == 1 + ) + assert all( + "limit_auto_multi_buffer_only_for_local_buffer" not in cfg.kwargs + for cfg in configs + if cfg.num_stages == 1 + ) + assert all( + "limit_auto_multi_buffer_of_local_buffer" not in cfg.kwargs + for cfg in configs + if cfg.num_stages == 1 + ) + assert all( + "set_workspace_multibuffer" not in cfg.kwargs + for cfg in configs + if cfg.num_stages == 1 + ) + assert all( + "tile_mix_vector_loop" not in cfg.kwargs + for cfg in configs + if cfg.num_stages == 1 + ) + assert all( + "tile_mix_cube_loop" not in cfg.kwargs for cfg in configs if cfg.num_stages == 1 + ) + assert all( + "enable_preload" not in cfg.kwargs for cfg in configs if cfg.num_stages == 1 + ) + assert { + cfg.kwargs["enable_auto_bind_sub_block"] + for cfg in configs + if cfg.num_stages == 1 + } == {True} + + +def test_compile_options_runtime_fixed_options_are_not_redefined(): + spec = parse_compile_options_hint("mixcv") + configs = expand_compile_option_configs( + [ + triton.Config( + { + "BLOCK_SIZE": 1024, + "multibuffer": False, + }, + num_stages=2, + ) + ], + spec, + generated_tiling=False, + fixed_options={ + "multibuffer": True, + }, + ) + + assert configs + assert all("multibuffer" not in cfg.kwargs for cfg in configs) + assert {cfg.num_stages for cfg in configs} == {1, 2} + + +def test_compile_options_runtime_fixed_num_stages_limits_search(): + spec = parse_compile_options_hint("mixcv") + configs = expand_compile_option_configs( + [triton.Config({"BLOCK_SIZE": 1024})], + spec, + generated_tiling=False, + fixed_options={"num_stages": 2}, + ) + + assert len(configs) == 152 + assert {cfg.num_stages for cfg in configs} == {2} + + +def test_compile_options_format_stage1_effective_options(): + spec = parse_compile_options_hint("mixcv") + config = triton.Config( + { + "BLOCK_M": 32, + "BLOCK_N": 32, + "enable_tuning_mode": True, + "enable_ubuf_saving": True, + "enable_hivm_auto_cv_balance": True, + "unit_flag": True, + }, + num_stages=1, + ) + + text = format_compile_option_result(config, spec) + + assert "selected_meta: BLOCK_M=32, BLOCK_N=32, num_stages=1" in text + assert "enable_auto_multi_buffer=False" in text + assert "set_workspace_multibuffer=" in text + assert "tile_mix_vector_loop=" in text + assert "tile_mix_cube_loop=" in text + assert "enable_preload" not in text + + +def test_compile_options_format_stage2_effective_options(): + spec = parse_compile_options_hint("mixcv") + config = triton.Config( + { + "BLOCK_M": 64, + "BLOCK_N": 128, + "enable_tuning_mode": True, + "enable_ubuf_saving": False, + "enable_hivm_auto_cv_balance": True, + "limit_auto_multi_buffer_only_for_local_buffer": False, + "limit_auto_multi_buffer_of_local_buffer": "no-l0c", + "set_workspace_multibuffer": 4, + "tile_mix_vector_loop": 2, + "tile_mix_cube_loop": 4, + "unit_flag": False, + }, + num_stages=2, + ) + + text = format_compile_option_result(config, spec) + + assert "selected_meta: BLOCK_M=64, BLOCK_N=128, num_stages=2" in text + assert "enable_auto_multi_buffer=True" in text + assert "set_workspace_multibuffer=4" in text + assert "tile_mix_vector_loop=2" in text + assert "tile_mix_cube_loop=4" in text + assert " {tt.divisibility = 16 : i32} , %arg1: !tt.ptr {tt.divisibility = 16 : i32, tt.shape_1 = 0 : i32} , %arg2: i32 {tt.divisibility = 16 : i32} ) attributes {noinline = false} { diff --git a/test/ascend/mlir/linalg_multi_assign.mlir b/test/ascend/mlir/linalg_multi_assign.mlir index 380242f0..3db45b1f 100644 --- a/test/ascend/mlir/linalg_multi_assign.mlir +++ b/test/ascend/mlir/linalg_multi_assign.mlir @@ -1,4 +1,4 @@ -// RUN: %triton-shared-opt-v3_4 %s --triton-to-linalg | %FileCheck %s +// RUN: %dicp_opt %s --triton-to-linalg | %FileCheck %s module { tt.func public @gcd_kernel(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: i32 {tt.divisibility = 16 : i32}) attributes {noinline = false} { diff --git a/test/ascend/passed_tests/test_atan.py b/test/ascend/passed_tests/test_atan.py index 595a8cf8..9949b9df 100644 --- a/test/ascend/passed_tests/test_atan.py +++ b/test/ascend/passed_tests/test_atan.py @@ -27,7 +27,8 @@ import torch import torch_npu -import triton.language.extra.deeplink.libdevice as libdevice +import triton.language.extra.deeplink.cann.libdevice as libdevice + def standard_unary(x0, dtype): res = torch.atan(x0) @@ -48,7 +49,9 @@ def triton_elementwise_unary(in_ptr0, out_ptr0, N: tl.constexpr, NUMEL: tl.const @triton.jit -def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr): +def triton_elementwise_binary( + in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL: tl.constexpr +): idx_block = tl.arange(0, NUMEL) x = tl.load(in_ptr0 + idx_block, mask=idx_block < N) y = tl.load(in_ptr1 + idx_block, mask=idx_block < N) @@ -57,8 +60,8 @@ def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL types = [ - (torch.float32, 'float32'), - (torch.float16, 'float16'), + (torch.float32, "float32"), + (torch.float16, "float16"), # (torch.bfloat16, 'bfloat16'), # (torch.int8, 'int8'), # (torch.int16, 'int16'), @@ -77,8 +80,8 @@ def triton_elementwise_binary(in_ptr0, in_ptr1, out_ptr0, N: tl.constexpr, NUMEL map_for_64_t = {37: 31} -@pytest.mark.parametrize('dtype,sigtype', types) -@pytest.mark.parametrize('N,NUMEL', shapes) +@pytest.mark.parametrize("dtype,sigtype", types) +@pytest.mark.parametrize("N,NUMEL", shapes) def test_elementwsie_common(dtype, sigtype, N, NUMEL): N = (-N) // torch.tensor(0, dtype=dtype).element_size() if N < 0 else N @@ -97,4 +100,4 @@ def test_elementwsie_common(dtype, sigtype, N, NUMEL): triton_elementwise_unary[1, 1, 1](x0, out, N=N, NUMEL=NUMEL, debug=True) print(out) - test_common.validate_cmp(sigtype, out, ans) \ No newline at end of file + test_common.validate_cmp(sigtype, out, ans) diff --git a/test/ascend/passed_tests/test_common.py b/test/ascend/passed_tests/test_common.py index 3fe03fca..b45b3f95 100644 --- a/test/ascend/passed_tests/test_common.py +++ b/test/ascend/passed_tests/test_common.py @@ -24,122 +24,177 @@ import pytest import functools import re +import numpy as np -_float_dtypes = [ - 'float32', 'float16', 'bfloat16' -] -_int_dtypes = [ - 'int32', 'int64', 'int16', 'int8' -] +_float_dtypes = ["float32", "float16", "bfloat16"] +_int_dtypes = ["int32", "int64", "int16", "int8"] +_uint_dtypes = ["uint8", "uint16", "uint32", "uint64"] _all_dtypes_no_bool = _float_dtypes + _int_dtypes -_all_dtypes = _all_dtypes_no_bool + ['bool'] -_32bit_dtypes = ['float32', 'int32'] -_16bit_dtypes = ['float16', 'bfloat16', 'int16'] +_all_dtypes = _all_dtypes_no_bool + ["bool"] +_32bit_dtypes = ["float32", "int32"] +_16bit_dtypes = ["float16", "bfloat16", "int16"] + + +def generate_numpy(shape, dtype, low=None, high=None): + if dtype in _int_dtypes + _uint_dtypes: + iinfo = np.iinfo(getattr(np, dtype)) + low = iinfo.min if low is None else max(low, iinfo.min) + high = iinfo.max if high is None else min(high, iinfo.max) + dty = getattr(np, dtype) + return np.random.randint(low, high, shape, dtype=dty) + elif dtype == "float16" or dtype == "float32": + return np.random.normal(0, 1, shape).astype(dtype) + elif dtype == "bfloat16": + return ( + np.random.normal(0, 1, shape).astype("float32").view("uint32") + & np.uint32(0xFFFF0000) + ).view("float32") + elif dtype == "bool": + return np.random.randint(low=0, high=2, size=shape).astype(bool) + else: + raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) + def generate_tensor(shape, dtype): - if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16': - return torch.randn(size=shape, dtype=eval('torch.' + dtype)) - elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16': - return torch.randint(low=0, high=2000, size=shape, dtype=eval('torch.' + dtype)) - elif dtype == 'int8': - return torch.randint(low=0, high=127, size=shape, dtype=eval('torch.' + dtype)) - elif dtype == 'bool': + if dtype == "float32" or dtype == "float16" or dtype == "bfloat16": + return torch.randn(size=shape, dtype=eval("torch." + dtype)) + elif dtype == "int32" or dtype == "int64" or dtype == "int16": + return torch.randint(low=0, high=2000, size=shape, dtype=eval("torch." + dtype)) + elif dtype == "int8": + return torch.randint(low=0, high=127, size=shape, dtype=eval("torch." + dtype)) + elif dtype == "bool": return torch.randint(low=0, high=2, size=shape).bool() + elif dtype == "uint8": + return torch.randint(low=0, high=255, size=shape, dtype=torch.uint8) else: - raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) def get_triton_sig_typename(dtype): - if dtype == 'float32': + if dtype == "float32": tyname = "*fp32" - elif dtype == 'int32': + elif dtype == "int32": tyname = "*i32" - elif dtype == 'int64': + elif dtype == "int64": tyname = "*i64" - elif dtype == 'float16': + elif dtype == "float16": tyname = "*fp16" - elif dtype == 'int16': + elif dtype == "int16": tyname = "*i16" - elif dtype == 'int8': + elif dtype == "int8": tyname = "*i8" - elif dtype == 'bool': + elif dtype == "bool": tyname = "*i1" else: - raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) return tyname + # Relative error: abs(x_ref - x_cal) / abs(x_ref) # Absolute error: abs(x_ref - x_cal) + # calculation type operators require different error range # It is a stricter verification and not satisfied now, save it here def validate_cal(dtype, y_cal, y_ref): - if dtype == 'float16': + if dtype == "float16": if torch.mean(y_ref) < 0.001: - assert torch.abs(y_cal - y_ref) < 0.001, "|y_cal - y_ref| < 0.001 is required !" + assert ( + torch.abs(y_cal - y_ref) < 0.001 + ), "|y_cal - y_ref| < 0.001 is required !" else: diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.001 # all true assert diff.all(), "Relative error is less than 0.001 !" - if dtype == 'float32': + if dtype == "float32": if torch.mean(y_ref) < 0.0001: - assert torch.abs(y_cal - y_ref) < 0.0001, "|y_cal - y_ref| < 0.0001 is required !" + assert ( + torch.abs(y_cal - y_ref) < 0.0001 + ), "|y_cal - y_ref| < 0.0001 is required !" else: diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.0001 assert diff.all(), "Relative error is less than 0.001 !" - elif dtype == 'bfloat16': + elif dtype == "bfloat16": diff = torch.div(torch.abs(y_cal - y_ref), torch.abs(y_cal)) < 0.001 assert diff.all(), "Relative error is less than 0.001 !" - elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16': + elif dtype == "int32" or dtype == "int64" or dtype == "int16" or dtype == "int8": + assert torch.equal(y_cal, y_ref) + elif ( + dtype == "uint8" or dtype == "uint16" or dtype == "uint32" or dtype == "uint64" + ): assert torch.equal(y_cal, y_ref) - elif dtype == 'bool': + elif dtype == "bool": assert torch.equal(y_cal, y_ref) else: - raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) + # moving and comparison ops require no precision error def validate_cmp(dtype, y_cal, y_ref, overflow_mode: Optional[str] = None): - y_cal=y_cal.npu() - y_ref=y_ref.npu() + y_cal = y_cal.npu() + y_ref = y_ref.npu() if overflow_mode == "saturate": - if dtype in ['float32', 'float16']: + if dtype in ["float32", "float16"]: min_value = -torch.finfo(dtype).min max_value = torch.finfo(dtype).max - elif dtype in ['int32', 'int16', 'int8']: + elif dtype in ["int32", "int16", "int8"]: min_value = torch.iinfo(dtype).min max_value = torch.iinfo(dtype).max - elif dtype == 'bool': + elif dtype == "bool": min_value = 0 max_value = 1 else: raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) y_ref = torch.clamp(y_ref, min=min_value, max=max_value) - if dtype == 'float16': - torch.testing.assert_close(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) - elif dtype == 'bfloat16': - torch.testing.assert_close(y_ref.to(torch.float32), y_cal.to(torch.float32), rtol=1e-03, atol=1e-03, equal_nan=True) - elif dtype == 'float32': - torch.testing.assert_close(y_ref, y_cal, rtol=1e-04, atol=1e-04, equal_nan=True) - elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8': + if dtype == "float16": + torch.testing.assert_close(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + elif dtype == "bfloat16": + torch.testing.assert_close( + y_ref.to(torch.float32), + y_cal.to(torch.float32), + rtol=1e-03, + atol=1e-03, + equal_nan=True, + ) + elif dtype == "float32": + torch.testing.assert_close(y_ref, y_cal, rtol=1e-04, atol=1e-04, equal_nan=True) + elif dtype == "int32" or dtype == "int64" or dtype == "int16" or dtype == "int8": + assert torch.equal(y_cal, y_ref) + elif ( + dtype == "uint8" or dtype == "uint16" or dtype == "uint32" or dtype == "uint64" + ): assert torch.equal(y_cal, y_ref) - elif dtype == 'bool': + elif dtype == "bool": assert torch.equal(y_cal, y_ref) else: - raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) + def validate_cmp_with_expection(dtype, y_cal, y_ref, expect): - if dtype == 'float32' or dtype == 'float16' or dtype == 'bfloat16': + if dtype == "float32" or dtype == "float16" or dtype == "bfloat16": if expect: - assert torch.allclose(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) + assert torch.allclose(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) else: - assert not torch.allclose(y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True) - elif dtype == 'int32' or dtype == 'int64' or dtype == 'int16' or dtype == 'int8': + assert not torch.allclose( + y_ref, y_cal, rtol=1e-03, atol=1e-03, equal_nan=True + ) + elif ( + dtype == "int32" + or dtype == "int64" + or dtype == "int16" + or dtype == "int8" + or dtype == "uint8" + or dtype == "uint16" + or dtype == "uint32" + or dtype == "uint64" + ): if expect: assert torch.equal(y_cal, y_ref) else: assert not torch.equal(y_cal, y_ref) else: - raise ValueError('Invalid parameter \"dtype\" is found : {}'.format(dtype)) + raise ValueError('Invalid parameter "dtype" is found : {}'.format(dtype)) + # Use the following pytest fixture to run one test case by only single worker. # Refer to https://pytest-xdist.readthedocs.io/en/stable/how-to.html#making-session-scoped-fixtures-execute-only-once @@ -149,36 +204,44 @@ def pytest_runonce(worker_id, request, cache): cache.set(request.node.nodeid, worker_id) else: file_name = f"pytest_{worker_id}.txt" - with open(file_name, 'a') as file: + with open(file_name, "a") as file: file.write(f"{request.node.nodeid} is already processed by {worker_id}") return True yield True cache.set(request.node.nodeid, "none") + def raises_with_match(expected_exception, match_pattern): def decorator(test_func): @functools.wraps(test_func) def wrapper(*args, **kwargs): with pytest.raises(expected_exception, match=match_pattern): return test_func(*args, **kwargs) + return wrapper + return decorator + def capture_output(expected_output): def decorator(test_func): @functools.wraps(test_func) def wrapper(*args, **kwargs): - capsys = kwargs.pop('capsys', None) + capsys = kwargs.pop("capsys", None) if capsys is None: try: capsys = pytest.fixture(capsys)() except: - raise RuntimeError("This decorator requires pytest's capsys fixture") + raise RuntimeError( + "This decorator requires pytest's capsys fixture" + ) test_func(capsys, *args, **kwargs) captured = capsys.readouterr() # pybind11::scoped_ostream_redirect captures std::cout with \x00 inserted # for now, no idea how to eliminate \x00 from C++ side. cleaned = re.sub(r"\x00", "", captured.out) assert expected_output in cleaned + return wrapper + return decorator diff --git a/test/ascend/passed_tests/test_compile_hint.py b/test/ascend/passed_tests/test_compile_hint.py index babcd860..4d75b42a 100644 --- a/test/ascend/passed_tests/test_compile_hint.py +++ b/test/ascend/passed_tests/test_compile_hint.py @@ -2,7 +2,7 @@ import triton import triton.language as tl import pytest -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl # eg: pytest -v test_compile_hint.py::test_compile_hint ############################# diff --git a/test/ascend/passed_tests/test_cv_flash_attention.py b/test/ascend/passed_tests/test_cv_flash_attention.py index e0e6dcc3..24ccd8d3 100644 --- a/test/ascend/passed_tests/test_cv_flash_attention.py +++ b/test/ascend/passed_tests/test_cv_flash_attention.py @@ -3,7 +3,8 @@ import torch_npu import triton import triton.language as tl -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl +from triton.language.extra.deeplink import async_task DEVICE = "npu" @@ -258,7 +259,6 @@ def _attn_fwd_base( class AttentionBase(torch.autograd.Function): - @staticmethod def forward(ctx, q, k, v, sm_scale, BM, BN): """ @@ -447,7 +447,7 @@ def _attn_fwd_split_cv( block_shape=(BLOCK_M, HEAD_DIM), order=(1, 0), ) - with dl.async_task(scope=dl.async_task.cube): + with async_task(scope=async_task.cube): q = tl.load(Q_block_ptr) lo, hi = 0, N_CTX # Process the entire context K_block_ptr = tl.advance( @@ -477,7 +477,7 @@ def _attn_fwd_split_cv( V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) K_block_ptr = tl.advance(K_block_ptr, (BLOCK_N, 0)) - with dl.async_task(scope=dl.async_task.vector): + with async_task(scope=async_task.vector): offs_m = task_m_idx * BLOCK_M + tl.arange(0, BLOCK_M) m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 @@ -517,7 +517,6 @@ def _attn_fwd_split_cv( class AttentionSplitCV(torch.autograd.Function): - @staticmethod def forward(ctx, q, k, v, sm_scale, BM, BN): """ diff --git a/test/ascend/passed_tests/test_extract_slice.py b/test/ascend/passed_tests/test_extract_slice.py index 508b0dfb..d47ba969 100644 --- a/test/ascend/passed_tests/test_extract_slice.py +++ b/test/ascend/passed_tests/test_extract_slice.py @@ -2,7 +2,7 @@ import triton import triton.language as tl -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl @triton.jit diff --git a/test/ascend/passed_tests/test_fused_rms_norm_rope.py b/test/ascend/passed_tests/test_fused_rms_norm_rope.py index 74cf64ea..227bf867 100644 --- a/test/ascend/passed_tests/test_fused_rms_norm_rope.py +++ b/test/ascend/passed_tests/test_fused_rms_norm_rope.py @@ -1,7 +1,7 @@ import torch import triton import triton.language as tl -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl import pytest diff --git a/test/ascend/passed_tests/test_fusedattention.py b/test/ascend/passed_tests/test_fusedattention.py index 5fe7ab9a..c06fec9b 100644 --- a/test/ascend/passed_tests/test_fusedattention.py +++ b/test/ascend/passed_tests/test_fusedattention.py @@ -38,7 +38,7 @@ import torch_npu import triton import triton.language as tl -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl DEVICE = "npu" @@ -334,7 +334,6 @@ def _attn_fwd( class _attention(torch.autograd.Function): - @staticmethod def forward(ctx, q, k, v, causal, sm_scale, BM, BN): """ diff --git a/test/ascend/passed_tests/test_insert_slice.py b/test/ascend/passed_tests/test_insert_slice.py index 56ad364d..fe7ef998 100644 --- a/test/ascend/passed_tests/test_insert_slice.py +++ b/test/ascend/passed_tests/test_insert_slice.py @@ -1,7 +1,7 @@ import torch import triton import triton.language as tl -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl @triton.jit diff --git a/test/ascend/passed_tests/test_isnan.py b/test/ascend/passed_tests/test_isnan.py index 58a2a2b9..b4d5d538 100644 --- a/test/ascend/passed_tests/test_isnan.py +++ b/test/ascend/passed_tests/test_isnan.py @@ -39,7 +39,6 @@ @pytest.mark.parametrize("sigtype", types) @pytest.mark.parametrize("N", shapes) def test_isnan(sigtype, N): - def torch_func(x0): res = torch.isnan(x0) return res @@ -48,7 +47,7 @@ def torch_func(x0): def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = tl.arange(0, N) x0 = tl.load(in_ptr0 + idx) - ret = tl.extra.deeplink.libdevice.isnan(x0) + ret = tl.extra.deeplink.cann.libdevice.isnan(x0) tl.store(out_ptr0 + idx, ret) def triton_func(x0, N): @@ -83,7 +82,7 @@ def triton_func(x0, N): # def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): # idx = tl.arange(0, N) # x0 = tl.load(in_ptr0 + idx) -# ret = tl.extra.deeplink.libdevice.isnan(x0) +# ret = tl.extra.deeplink.cann.libdevice.isnan(x0) # tl.store(out_ptr0 + idx, ret) diff --git a/test/ascend/passed_tests/test_log1p.py b/test/ascend/passed_tests/test_log1p.py index a589c497..65030290 100644 --- a/test/ascend/passed_tests/test_log1p.py +++ b/test/ascend/passed_tests/test_log1p.py @@ -23,7 +23,7 @@ import triton.language as tl import torch import test_common -import triton.language.extra.deeplink.libdevice as libdevice +import triton.language.extra.deeplink.cann.libdevice as libdevice def torch_log1p(x0, x1): @@ -32,7 +32,9 @@ def torch_log1p(x0, x1): @triton.jit -def triton_log1p(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): +def triton_log1p( + in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr +): xoffset = tl.program_id(0) * XBLOCK for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] @@ -43,10 +45,12 @@ def triton_log1p(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOC tl.store(out_ptr0 + x_index, tmp2, xmask) -@pytest.mark.parametrize('param_list', - [ - ['float32', (2, 4096, 8), 2, 32768, 1024], - ]) +@pytest.mark.parametrize( + "param_list", + [ + ["float32", (2, 4096, 8), 2, 32768, 1024], + ], +) def test_log1p(param_list): # 生成数据 dtype, shape, ncore, xblock, xblock_sub = param_list diff --git a/test/ascend/passed_tests/test_matrix_multiplication_optimized.py b/test/ascend/passed_tests/test_matrix_multiplication_optimized.py index 0f98b1bc..86f33747 100644 --- a/test/ascend/passed_tests/test_matrix_multiplication_optimized.py +++ b/test/ascend/passed_tests/test_matrix_multiplication_optimized.py @@ -4,7 +4,7 @@ import triton import triton.language as tl import triton.runtime.driver as driver -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl # get device properties of npu diff --git a/test/ascend/passed_tests/test_mod.py b/test/ascend/passed_tests/test_mod.py index 2d1d0d06..5c304c57 100644 --- a/test/ascend/passed_tests/test_mod.py +++ b/test/ascend/passed_tests/test_mod.py @@ -26,8 +26,15 @@ import test_common -def torch_pointwise(x0, x1): - res = x0 % x1 +def torch_pointwise(x0, x1, dtype): + if dtype == "float16": + x0 = x0.to(torch.float32) + x1 = x1.to(torch.float32) + elif dtype == "float32": + x0 = x0.to(torch.float64) + x1 = x1.to(torch.float64) + res = torch.div(x0, x1, rounding_mode="trunc") + res = x0 - x1 * res return res @@ -67,8 +74,9 @@ def test_case(param_list): else: x0 = test_common.generate_tensor(shape, dtype).npu() x1 = test_common.generate_tensor(shape, dtype).npu() - y_ref = torch_pointwise(x0.cpu(), x1.cpu()) - y_ref = y_ref.npu() + y_ref = torch_pointwise(x0, x1, dtype) + if dtype == "float16": + y_ref = y_ref.to(torch.float16) y_cal = torch.zeros(shape, dtype=eval("torch." + dtype)).npu() triton_mod[ncore, 1, 1](x0, x1, y_cal, xblock, xblock_sub) # test_common.validate_cmp(dtype, y_cal, y_ref.npu()) diff --git a/test/ascend/passed_tests/test_multi_return.py b/test/ascend/passed_tests/test_multi_return.py index 3bfad8c2..2c02009d 100644 --- a/test/ascend/passed_tests/test_multi_return.py +++ b/test/ascend/passed_tests/test_multi_return.py @@ -24,7 +24,7 @@ import torch import torch_npu import triton -from triton.language.extra.deeplink.libdevice import tanh +from triton.language.extra.deeplink.cann.libdevice import tanh import triton.language as tl device = "npu" diff --git a/test/ascend/passed_tests/test_pow.py b/test/ascend/passed_tests/test_pow.py index ac0a6b23..2e5e1617 100644 --- a/test/ascend/passed_tests/test_pow.py +++ b/test/ascend/passed_tests/test_pow.py @@ -20,7 +20,7 @@ import triton import triton.language as tl -from triton.language.extra.deeplink.libdevice import pow +from triton.language.extra.deeplink.cann.libdevice import pow import torch import torch_npu import pytest @@ -28,20 +28,23 @@ types = [ "float32", - "float16", - "bfloat16", + # "float16", + # "bfloat16", ] shapes = [ - 16, - 256, + # 3, + # 32, + 37, + # 256, + # 781, ] -# @pytest.mark.skip(reason="waiting for bishengir-compile to support") + +@pytest.mark.skip(reason="waiting for bishengir-compile to support") @pytest.mark.parametrize("sigtype", types) @pytest.mark.parametrize("N", shapes) def test_pow_vv(sigtype, N): - def torch_func(x0, x1): res = torch.pow(x0, x1) return res @@ -66,11 +69,11 @@ def triton_func(x0, x1, N): torch_ref = torch_func(x0, x1) test_common.validate_cmp(sigtype, triton_cal, torch_ref) + @pytest.mark.skip(reason="waiting for bishengir-compile to support") @pytest.mark.parametrize("sigtype", types) @pytest.mark.parametrize("N", shapes) def test_pow_vs_dynamic(sigtype, N): - def torch_func(x0, x1): res = torch.pow(x0, x1) return res @@ -95,11 +98,11 @@ def triton_func(x0, x1, N): torch_ref = torch_func(x0, x1) test_common.validate_cmp(sigtype, triton_cal, torch_ref) + # @pytest.mark.skip(reason="waiting for bishengir-compile to support") @pytest.mark.parametrize("sigtype", types) @pytest.mark.parametrize("N", shapes) def test_pow_vs_const(sigtype, N): - def torch_func(x0, x1): res = torch.pow(x0, x1) return res @@ -121,4 +124,4 @@ def triton_func(x0, x1, N): triton_cal = triton_func(x0, x1, N) torch_ref = torch_func(x0, x1) - test_common.validate_cmp(sigtype, triton_cal, torch_ref) \ No newline at end of file + test_common.validate_cmp(sigtype, triton_cal, torch_ref) diff --git a/test/ascend/passed_tests/test_relu.py b/test/ascend/passed_tests/test_relu.py index adf676e0..8a9753ee 100644 --- a/test/ascend/passed_tests/test_relu.py +++ b/test/ascend/passed_tests/test_relu.py @@ -23,7 +23,7 @@ import torch import pytest import test_common -import triton.language.extra.deeplink.libdevice as libdevice +import triton.language.extra.deeplink.cann.libdevice as libdevice def torch_relu(x0, x1): @@ -32,7 +32,9 @@ def torch_relu(x0, x1): @triton.jit -def triton_relu(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): +def triton_relu( + in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr +): xoffset = tl.program_id(0) * XBLOCK for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): x_index = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] @@ -43,11 +45,13 @@ def triton_relu(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK tl.store(out_ptr0 + x_index, tmp2, xmask) -@pytest.mark.parametrize('param_list', - [ - ['float32', (2, 4096, 8), 2, 32768, 512], - ['float16', (2, 4096, 8), 2, 32768, 512], - ]) +@pytest.mark.parametrize( + "param_list", + [ + ["float32", (2, 4096, 8), 2, 32768, 512], + ["float16", (2, 4096, 8), 2, 32768, 512], + ], +) def test_relu(param_list): # 生成数据 dtype, shape, ncore, xblock, xblock_sub = param_list diff --git a/test/ascend/passed_tests/test_scalar_calc.py b/test/ascend/passed_tests/test_scalar_calc.py index 76e44169..9b1a200b 100644 --- a/test/ascend/passed_tests/test_scalar_calc.py +++ b/test/ascend/passed_tests/test_scalar_calc.py @@ -22,7 +22,7 @@ import torch_npu import triton import triton.language as tl -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.libdevice as libdevice import pytest import test_common @@ -30,7 +30,6 @@ ### add @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_add_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -41,7 +40,7 @@ def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): def torch_func(x0): y = x0[0] y = y + 2.0 - return y + return torch.tensor(y) dtype, N = param_list x0 = test_common.generate_tensor((N,), dtype).npu() @@ -54,7 +53,6 @@ def torch_func(x0): ### sub @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_sub_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -65,7 +63,7 @@ def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): def torch_func(x0): y = x0[0] y = y - 2.0 - return y + return torch.tensor(y) dtype, N = param_list x0 = test_common.generate_tensor((N,), dtype).npu() @@ -78,7 +76,6 @@ def torch_func(x0): ### mul @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_mul_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -89,7 +86,7 @@ def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): def torch_func(x0): y = x0[0] y = y * 2.0 - return y + return torch.tensor(y) dtype, N = param_list x0 = test_common.generate_tensor((N,), dtype).npu() @@ -102,7 +99,6 @@ def torch_func(x0): ### div @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_div_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -113,7 +109,7 @@ def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): def torch_func(x0): y = x0[0] y = y / 2.0 - return y + return torch.tensor(y) dtype, N = param_list x0 = test_common.generate_tensor((N,), dtype).npu() @@ -126,7 +122,6 @@ def torch_func(x0): ### remf @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_remf_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -136,8 +131,8 @@ def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): def torch_func(x0): y = x0[0] - y = y % 2.0 - return y + y = y - 2.0 * torch.div(y, 2.0, rounding_mode="trunc") + return torch.tensor(y) dtype, N = param_list x0 = test_common.generate_tensor((N,), dtype).npu() @@ -150,7 +145,6 @@ def torch_func(x0): ### negf @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_negf_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -174,7 +168,6 @@ def torch_func(x0): ### cmpf @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_cmpf_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -198,7 +191,6 @@ def torch_func(x0): ### ceil @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_ceil_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -222,7 +214,6 @@ def torch_func(x0): ### floor @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_floor_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -247,7 +238,6 @@ def torch_func(x0): # setting propagate_nan=tl.PropagateNan.ALL to generate arith::MaximumFOp @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_maximum_nanall_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): tl.static_assert(N > 1) @@ -274,7 +264,6 @@ def torch_func(x0): # setting propagate_nan=tl.PropagateNan.NONE to generate arith::MaxNumFOp @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_maximum_nannone_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): tl.static_assert(N > 1) @@ -301,7 +290,6 @@ def torch_func(x0): # setting propagate_nan=tl.PropagateNan.ALL to generate arith::MinimumFOp @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_minimum_nanall_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): tl.static_assert(N > 1) @@ -328,7 +316,6 @@ def torch_func(x0): # setting propagate_nan=tl.PropagateNan.NONE to generate arith::MinNumFOp @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_minimum_nannone_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): tl.static_assert(N > 1) @@ -354,7 +341,6 @@ def torch_func(x0): ### extf @pytest.mark.parametrize("param_list", [["float16", "float32", 16]]) def test_scalar_extf_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -378,7 +364,6 @@ def torch_func(x0): ### truncf @pytest.mark.parametrize("param_list", [["float32", "float16", 16]]) def test_scalar_truncf_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -402,7 +387,6 @@ def torch_func(x0): ### exp @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_exp_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -426,7 +410,6 @@ def torch_func(x0): ### exp2 @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_exp_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -450,7 +433,6 @@ def torch_func(x0): ### log @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_log_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -476,7 +458,6 @@ def torch_func(x0): ### log2 @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_log2_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -502,7 +483,6 @@ def torch_func(x0): ### sin @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_sin_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -526,7 +506,6 @@ def torch_func(x0): ### cos @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_cos_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -550,7 +529,6 @@ def torch_func(x0): ### abs @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_abs_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -574,7 +552,6 @@ def torch_func(x0): ### erf @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_erf_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -598,7 +575,6 @@ def torch_func(x0): ### sqrt @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_sqrt_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -624,7 +600,6 @@ def torch_func(x0): ### rsqrt @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_rsqrt_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 @@ -650,12 +625,11 @@ def torch_func(x0): ### tanh @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_tanh_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): idx = 0 tmp0 = tl.load(in_ptr0 + idx) - tmp1 = dl.libdevice.tanh(tmp0) + tmp1 = libdevice.tanh(tmp0) tl.store(out_ptr0 + idx, tmp1) def torch_func(x0): @@ -674,7 +648,6 @@ def torch_func(x0): ### sum @pytest.mark.parametrize("param_list", [["float32", 16]]) def test_scalar_sum_calc(param_list): - @triton.jit def triton_kernel(out_ptr0, in_ptr0, N: tl.constexpr): tmp0 = tl.load(in_ptr0 + tl.arange(0, N)) diff --git a/test/ascend/passed_tests/test_sync_block.py b/test/ascend/passed_tests/test_sync_block.py index fc7306f3..846c8d15 100644 --- a/test/ascend/passed_tests/test_sync_block.py +++ b/test/ascend/passed_tests/test_sync_block.py @@ -1,7 +1,7 @@ import torch import triton import triton.language as tl -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl import test_common diff --git a/test/ascend/passed_tests/test_zeros.py b/test/ascend/passed_tests/test_zeros.py index 64631738..6e07e9a8 100644 --- a/test/ascend/passed_tests/test_zeros.py +++ b/test/ascend/passed_tests/test_zeros.py @@ -39,7 +39,9 @@ def fn_npu_f32(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.con ret = tl.zeros((XB, YB, ZB), dtype=tl.float32) - oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + oidx = ( + xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + ) tl.store(output_ptr + oidx, ret) @@ -56,7 +58,9 @@ def fn_npu_f16(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.con ret = tl.zeros((XB, YB, ZB), dtype=tl.float16) - oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + oidx = ( + xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + ) tl.store(output_ptr + oidx, ret) @@ -73,29 +77,32 @@ def fn_npu_i8(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.cons ret = tl.zeros((XB, YB, ZB), dtype=tl.int8) - oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + oidx = ( + xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + ) tl.store(output_ptr + oidx, ret) -@pytest.mark.parametrize('param_list', - [ - ['float32', (2, 256, 16), 1, 2, 256, 16], - ['float32', (8, 8, 4), 1, 8, 8, 4], - ['float16', (2, 256, 16), 1, 2, 256, 16], - ['float16', (8, 8, 4), 1, 8, 8, 4], - ['int8', (2, 256, 16), 1, 2, 256, 16], - ['int8', (8, 8, 4), 1, 8, 8, 4], - ] - ) +@pytest.mark.parametrize( + "param_list", + [ + ["float32", (2, 256, 16), 1, 2, 256, 16], + ["float32", (8, 8, 4), 1, 8, 8, 4], + ["float16", (2, 256, 16), 1, 2, 256, 16], + ["float16", (8, 8, 4), 1, 8, 8, 4], + ["int8", (2, 256, 16), 1, 2, 256, 16], + ["int8", (8, 8, 4), 1, 8, 8, 4], + ], +) def test_case(param_list): dtype, shape, ncore, XB, YB, ZB = param_list - x0 = test_common.generate_tensor(shape, dtype) + x0 = test_common.generate_tensor(shape, dtype).npu() - y_ref = torch.full((XB, YB, ZB), 0, dtype=eval('torch.' + dtype)).npu() + y_ref = torch.full((XB, YB, ZB), 0, dtype=eval("torch." + dtype)).npu() print(f"y_ref = {y_ref[0, 0, 0:4]}") - y_cal = torch.randint(1, (XB, YB, ZB), dtype=eval('torch.' + dtype)).npu() + y_cal = torch.randint(1, (XB, YB, ZB), dtype=eval("torch." + dtype)).npu() if dtype == "float32": fn_npu_f32[ncore, 1, 1](y_cal, x0, XB, YB, ZB) elif dtype == "float16": diff --git a/test/ascend/passed_tests/test_zeroslike.py b/test/ascend/passed_tests/test_zeroslike.py index b2835993..0f2e4a2e 100644 --- a/test/ascend/passed_tests/test_zeroslike.py +++ b/test/ascend/passed_tests/test_zeroslike.py @@ -39,27 +39,30 @@ def fn_npu_(output_ptr, x_ptr, XB: tl.constexpr, YB: tl.constexpr, ZB: tl.conste ret = tl.zeros_like(X) - oidx = xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + oidx = ( + xidx[:, None, None] * YB * ZB + yidx[None, :, None] * ZB + zidx[None, None, :] + ) tl.store(output_ptr + oidx, ret) -@pytest.mark.parametrize('param_list', - [ - ['float32', (2, 256, 16), 1, 2, 256, 16], - ['float32', (8, 8, 4), 1, 8, 8, 4], - ['float16', (2, 256, 16), 1, 2, 256, 16], - ['float16', (8, 8, 4), 1, 8, 8, 4], - ['int8', (2, 256, 16), 1, 2, 256, 16], - ['int8', (8, 8, 4), 1, 8, 8, 4], - ] - ) +@pytest.mark.parametrize( + "param_list", + [ + ["float32", (2, 256, 16), 1, 2, 256, 16], + ["float32", (8, 8, 4), 1, 8, 8, 4], + ["float16", (2, 256, 16), 1, 2, 256, 16], + ["float16", (8, 8, 4), 1, 8, 8, 4], + ["int8", (2, 256, 16), 1, 2, 256, 16], + ["int8", (8, 8, 4), 1, 8, 8, 4], + ], +) def test_case(param_list): dtype, shape, ncore, XB, YB, ZB = param_list - x0 = test_common.generate_tensor(shape, dtype) - y_ref = torch.zeros_like(x0, dtype=eval('torch.' + dtype)).npu() + x0 = test_common.generate_tensor(shape, dtype).npu() + y_ref = torch.zeros_like(x0, dtype=eval("torch." + dtype)).npu() print(f"y_ref = {y_ref[0, 0, 0:4]}") - y_cal = torch.zeros(shape, dtype=eval('torch.' + dtype)).npu() + y_cal = torch.zeros(shape, dtype=eval("torch." + dtype)).npu() fn_npu_[ncore, 1, 1](y_cal, x0, XB, YB, ZB) print(f"y_cal = {y_cal[0, 0, 0:4]}") diff --git a/test/ascend/run_tests.sh b/test/ascend/run_tests.sh index 72237d9f..22832da4 100644 --- a/test/ascend/run_tests.sh +++ b/test/ascend/run_tests.sh @@ -5,6 +5,9 @@ script=$(readlink -f "$0") script_dir=$(dirname "$script") function run_pytestcases() { + PRIMARY_ASCEND_DEVICE="${ASCEND_RT_VISIBLE_DEVICES:-}" + FALLBACK_ASCEND_DEVICE="${ASCEND_FALLBACK_RT_VISIBLE_DEVICES:-5}" + if [ -d ${HOME}/.triton/dump ]; then rm -rf ${HOME}/.triton/dump fi @@ -16,9 +19,9 @@ function run_pytestcases() { TARGET_DIR="$1" cd ${TARGET_DIR} - echo "[Phase 1] Run tests in parallel" + echo "[Phase 1] Run tests in parallel on Ascend device ${PRIMARY_ASCEND_DEVICE}" set +e - timeout --signal=TERM 40m pytest . -n 8 --dist=loadscope --reruns 5 --reruns-delay 5 + timeout --signal=TERM 20m pytest . -n 8 --dist=loadscope --reruns 1 --reruns-delay 2 parallel_rc=$? set -e @@ -26,19 +29,31 @@ function run_pytestcases() { echo "[SUCCESS] All tests passed" return 0 fi + if [ "${parallel_rc}" -eq 124 ]; then + echo "[ERROR] Parallel run timed out" + return "${parallel_rc}" + fi + + echo "[INFO] Failed cases collected from device ${PRIMARY_ASCEND_DEVICE}:" + if [ -f .pytest_cache/v/cache/lastfailed ]; then + sed -n '1,200p' .pytest_cache/v/cache/lastfailed + else + echo "[WARN] .pytest_cache/v/cache/lastfailed not found" + fi - echo "[Phase 2] Parallel run failed, rerun failed cases serially" + echo "[Phase 2] Rerun failed cases serially on Ascend device ${FALLBACK_ASCEND_DEVICE}" + export ASCEND_RT_VISIBLE_DEVICES="${FALLBACK_ASCEND_DEVICE}" set +e - pytest --lf --last-failed-no-failures=none -n 0 -v . + pytest --lf --last-failed-no-failures=none -n 0 -v . --reruns 0 serial_rc=$? set -e if [ "${serial_rc}" -ne 0 ]; then - echo "[ERROR] Serial rerun still has failures" + echo "[ERROR] Serial rerun still has failures on Ascend device ${FALLBACK_ASCEND_DEVICE}" return 1 fi - echo "[SUCCESS] Cases run passed in serial rerun" + echo "[SUCCESS] Cases run passed on Ascend device ${FALLBACK_ASCEND_DEVICE}" return 0 } diff --git a/test/ascend/test_custom_op.py b/test/ascend/test_custom_op.py index 8f411ee3..75cc9a28 100644 --- a/test/ascend/test_custom_op.py +++ b/test/ascend/test_custom_op.py @@ -1,33 +1,31 @@ -#!/usr/bin/env python3 -import subprocess +""" +Triton add kernel using dl.custom() with bitcode auto-resolution. + +Verifies the custom op `add` produces correct results on Ascend NPU. +""" + import os + +# Ensure bishengir tools are in PATH before triton imports read BISHENG_INSTALL_PATH. +_BISHENG_INSTALL = ( + "/mnt/data01/zmz/workspace/04ttshared/fordlc/ascendnpu-ir-0514/build/install/bin/" +) +if os.path.isdir(_BISHENG_INSTALL): + os.environ.setdefault("BISHENG_INSTALL_PATH", _BISHENG_INSTALL) + if _BISHENG_INSTALL not in os.environ.get("PATH", ""): + os.environ["PATH"] = _BISHENG_INSTALL + os.pathsep + os.environ.get("PATH", "") + +import pytest +import torch +import torch_npu # noqa: F401 import triton import triton.language as tl -import triton.language.extra.deeplink as dl - -from triton.compiler.compiler import ASTSource -from triton.compiler.code_generator import ast_to_ttir -from triton._C.libtriton import ir -import hashlib -from triton.backends.dicp_triton.npu import NPUUtils -from triton.backends.compiler import GPUTarget -from triton.compiler.compiler import make_backend, IRSource, filter_traceback -from triton import __version__, knobs -from triton.runtime.cache import ( - get_cache_manager, - get_dump_manager, - get_override_manager, - get_cache_key, -) +import triton.language.extra.deeplink.cann.extension as dl -from triton.backends.dicp_triton.npu import ( - make_ttir, - ttir_to_linalg, - ttir_to_ttsharedir_ascend, - ttsharedir_to_linkedir, - linalg_to_bin_enable_npu_compile, - NPUOptions, -) + +# ====================================================================== +# DSL custom op registration — bitcode auto-resolved by name +# ====================================================================== @dl.register_custom_op @@ -37,203 +35,45 @@ class add: mode = dl.MODE.SIMD def __init__(self, a, b, out=None): - assert out, "out is required" + assert out is not None, "dl.custom() requires out= parameter" self.symbol = "custom_add_" + str(a.dtype) - # bitcode name, auto-resolved via dlcompiler/bitcode/bc/ or - # DLCOMPILER_BITCODE_PATH environment variable - self.bitcode = "add" + self.bitcode = "add" # auto-resolved to add.aiv.bc + + +# ====================================================================== +# Triton kernel +# ====================================================================== @triton.jit -def triton_custom_add(output_ptr, a_ptr, b_ptr, L: tl.constexpr): +def custom_add_kernel(output_ptr, a_ptr, b_ptr, L: tl.constexpr): idx = tl.arange(0, L) - a = tl.load(a_ptr + idx) b = tl.load(b_ptr + idx) - buf = tl.full([L], 0, a.dtype) res = dl.custom("add", a, b, out=buf) - tl.store(output_ptr + idx, res) -def compile(src, target=None, options=None, _env_vars=None): - compilation_listener = knobs.compilation.listener - if compilation_listener: - timer = CompileTimer() - - if target is None: - target = driver.active.get_current_target() - assert isinstance(target, GPUTarget), "target must be of GPUTarget type" - backend = make_backend(target) - ir_source = not isinstance(src, ASTSource) - # create backend - if ir_source: - assert isinstance(src, str), "source must be either AST or a filepath" - context = ir.context() - src = IRSource(src, context, backend) - - extra_options = src.parse_options() - options = backend.parse_options(dict(options or dict(), **extra_options)) - # create cache manager - env_vars = get_cache_invalidating_env_vars() if _env_vars is None else _env_vars - key = get_cache_key(src, backend, options, env_vars=env_vars) - hash = hashlib.sha256(key.encode("utf-8")).hexdigest() - fn_cache_manager = get_cache_manager(hash) - # For dumping/overriding only hash the source as we want it to be independent of triton - # core changes to make it easier to track kernels by hash. - enable_override = knobs.compilation.override - enable_ir_dump = knobs.compilation.dump_ir - store_only_binary = knobs.compilation.store_binary_only - fn_override_manager = get_override_manager(src.hash()) if enable_override else None - fn_dump_manager = get_dump_manager(src.hash()) if enable_ir_dump else None - # Pre-truncate the file name here to avoid hitting the 255 character limit on common platforms. - # The final file name in the cache will have a format of f"{filename}.{ext}.tmp.pid_{pid}_{uuid}". - # A PID string can be 5-character long. A UUID string has typically 36 characters. Let's truncate - # the file name to 150 characters to be safe. - file_name = src.name[:150] - metadata_filename = f"{file_name}.json" - metadata_group = fn_cache_manager.get_group(metadata_filename) or {} - metadata_path = metadata_group.get(metadata_filename) - always_compile = knobs.compilation.always_compile - if not always_compile and metadata_path is not None: - # cache hit! - res = CompiledKernel(src, metadata_group, hash) - if compilation_listener: - compilation_listener( - src=src, - metadata=res.metadata._asdict(), - metadata_group=metadata_group, - times=timer.end(), - cache_hit=True, - ) - return res - - # initialize metadata - metadata = { - "hash": hash, - "target": target, - **options.__dict__, - **env_vars, - } - metadata["triton_version"] = __version__ - # run compilation pipeline and populate metadata - stages = dict() - backend.add_stages(stages, options, src.language) - first_stage = list(stages.keys()).index(src.ext) - # when the source is an IR file, don't apply the passes related to this stage. This makes it easier to write IR level tests. - if ir_source: - first_stage += 1 - - # For IRSource, we have already grabbed the context + called both - # ir.load_dialects and backend.load_dialects. - if not isinstance(src, IRSource): - context = ir.context() - ir.load_dialects(context) - backend.load_dialects(context) - - codegen_fns = backend.get_codegen_implementation(options) - module_map = backend.get_module_map() - try: - module = src.make_ir(target, options, codegen_fns, module_map, context) - except Exception as e: - filter_traceback(e) - raise - - if ir_source: - ir_filename = f"{file_name}.{src.ext}" - metadata_group[ir_filename] = fn_cache_manager.put(module, ir_filename) - else: - ir_filename = f"{file_name}.source" - metadata_group[ir_filename] = fn_cache_manager.put(module, ir_filename) - - use_ir_loc = knobs.compilation.use_ir_loc - if ir_source and use_ir_loc: - module.create_location_snapshot(src.path) - print(f"Creating new locations for {src.path}") - - if compilation_listener: - timer.finished_ir_initialization() - - if "npubin" in stages.keys(): - del stages["npubin"] - - for ext, compile_ir in list(stages.items())[first_stage:]: - next_module = compile_ir(module, metadata) - ir_filename = f"{file_name}.{ext}" - if fn_override_manager is None: - # Users can override kernels at scale by setting `ir_override` in autotune config - # without TRITON_KERNEL_OVERRIDE - if ( - ir_override := metadata.get("ir_override", None) - ) and ir_override.endswith(f".{ext}"): - next_module = parse(ir_override, ext, context) - elif full_name := fn_override_manager.get_file(ir_filename): - print(f"\nOverriding kernel with file {full_name}") - next_module = parse(full_name, ext, context) - # If TRITON_STORE_BINARY_ONLY is 1, only store cubin/hsaco/json - if (not store_only_binary) or (ext in ("cubin", "hsaco", "json")): - metadata_group[ir_filename] = fn_cache_manager.put(next_module, ir_filename) - if fn_dump_manager is not None: - fn_dump_manager.put(next_module, ir_filename) - if ext == "cubin": - sass = get_sass(next_module) - fn_dump_manager.put(sass, file_name + ".sass") - # use an env variable to parse ir from file - if use_ir_loc == ext: - ir_full_name = fn_cache_manager.get_file(ir_filename) - next_module.create_location_snapshot(ir_full_name) - print(f"Creating new locations for {ir_full_name}") - module = next_module - if compilation_listener: - timer.stage_finished(ext) - return module +# ====================================================================== +# Tests +# ====================================================================== + + +@pytest.mark.parametrize("L", [32, 128, 1024]) +def test_custom_add_int32(L): + a = torch.randint(0, 1000, (L,), dtype=torch.int32).npu() + b = torch.randint(0, 1000, (L,), dtype=torch.int32).npu() + out = torch.empty(L, dtype=torch.int32).npu() + + custom_add_kernel[1, 1, 1](out, a, b, L=L) + + ref = a.cpu() + b.cpu() + assert torch.equal(out.cpu(), ref), f"L={L}: out={out.cpu()}, ref={ref}" if __name__ == "__main__": - npuutiles = NPUUtils() - src = ASTSource( - triton_custom_add, - {"output_ptr": "*i32", "a_ptr": "*i32", "b_ptr": "*i32"}, - {"L": 32}, - ) - target = GPUTarget(backend="ascend", arch=npuutiles.get_arch(), warp_size=0) - options = { - "debug": False, - "sanitize_overflow": False, - "llvm_version": 15, - "kernel_name": "triton_", - "cluster_dims": (1, 1, 1), - "num_warps": -1, - "num_ctas": -1, - "num_stages": 2, - "num_buffers_warp_spec": 0, - "num_consumer_groups": 0, - "reg_dec_producer": 0, - "reg_inc_consumer": 0, - "enable_warp_specialization": False, - "enable_nd2nz_on_vector": False, - "enable_persistent": False, - "optimize_epilogue": False, - "enable_fp_fusion": True, - "allow_fp8e4nv": False, - "allowed_dot_input_precisions": ("ieee", "hf32"), - "enable_npu_compile": True, - "max_num_imprecise_acc_default": None, - "extern_libs": None, - "multibuffer": True, - "inject_barrier_all": False, - "disable_auto_inject_block_sync": False, - "unit_flag": False, - "disable_auto_cv_work_space_manage": False, - "enable_auto_bind_sub_block": True, - "tile_mix_vector_loop": None, - "tile_mix_cube_loop": None, - "limit_auto_multi_buffer_only_for_local_buffer": None, - "set_workspace_multibuffer": None, - "stream": None, - } - linkedir = compile(src, target, options, {}) - - print("=== MLIR (linkedir) ===") - print(linkedir) + for L in [32, 128, 1024]: + test_custom_add_int32(L) + print(f"[PASS] L={L}") + print("Done.") diff --git a/test/ascend/test_mlir.sh b/test/ascend/test_mlir.sh index a10c0676..1d4996fb 100644 --- a/test/ascend/test_mlir.sh +++ b/test/ascend/test_mlir.sh @@ -8,7 +8,7 @@ TRITON_PATH=$(python -c "import triton; import os; print(os.path.dirname(triton. export PATH="$TRITON_PATH:$TRITON_PATH/_C:$PATH" # 检查必要工具 -for tool in dicp_opt "triton-shared-opt-v3_4" FileCheck; do +for tool in dicp_opt FileCheck; do if ! command -v "$tool" &> /dev/null; then echo "Error: $tool is not available in PATH" >&2 exit 1 @@ -56,7 +56,6 @@ run_test() { # 替换占位符 local cmd=$(echo "$run_line" | sed "s|%s|$mlir_file|g" | \ sed 's|%dicp_opt|dicp_opt|g' | \ - sed 's|%triton-shared-opt-v3_4|triton-shared-opt-v3_4|g' | \ sed 's|%FileCheck|FileCheck|g') echo "TEST: $filename" diff --git a/test/ascend/test_softmax_custom_op.py b/test/ascend/test_softmax_custom_op.py index f674edf4..dec51ab5 100644 --- a/test/ascend/test_softmax_custom_op.py +++ b/test/ascend/test_softmax_custom_op.py @@ -22,7 +22,7 @@ import torch import triton import triton.language as tl -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl # ====================================================================== # DSL custom op registration — bitcode auto-resolved by name diff --git a/test/commonir/run_tests.sh b/test/commonir/run_tests.sh index bd8b91d4..a51f8771 100644 --- a/test/commonir/run_tests.sh +++ b/test/commonir/run_tests.sh @@ -15,7 +15,7 @@ function run_pytestcases() { cd ${script_dir} TARGET_DIR="$1" cd ${TARGET_DIR} - pytest -n 8 --dist=load . || { exit 1 ; } + pytest -n 8 --dist=load --timeout=60 . || { exit 1 ; } } diff --git a/test/dsl/test_alloc.py b/test/dsl/test_alloc.py index 89ddf0c0..a0a7bfd8 100644 --- a/test/dsl/test_alloc.py +++ b/test/dsl/test_alloc.py @@ -1,15 +1,17 @@ import torch import triton import triton.language as tl -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl import pytest + @triton.jit -def custom_func_kernel(x_ptr, # *Pointer* to first input vector. - output_ptr, # *Pointer* to output vector. - n_elements, # Size of the vector. - BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. - ): +def custom_func_kernel( + x_ptr, # *Pointer* to first input vector. + output_ptr, # *Pointer* to output vector. + n_elements, # Size of the vector. + BLOCK_SIZE: tl.constexpr, # Number of elements each program should process. +): pid = tl.program_id(axis=0) block_start = pid * BLOCK_SIZE offsets = block_start + tl.arange(0, BLOCK_SIZE) @@ -19,23 +21,28 @@ def custom_func_kernel(x_ptr, # *Pointer* to first input vector. output = x + y tl.store(output_ptr + offsets, output, mask=mask) + def custom_func(x: torch.Tensor): output = torch.empty_like(x) n_elements = output.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) custom_func_kernel[grid](x, output, n_elements, BLOCK_SIZE=128) return output + def test_add(): torch.manual_seed(0) size = 1024 - x = torch.rand(size, device='npu') + x = torch.rand(size, device="npu") output_torch = x + 1.68 output_triton = custom_func(x) - + assert torch.allclose(output_torch, output_triton, atol=1e-5) - print(f'The maximum difference between torch and triton is ' - f'{torch.max(torch.abs(output_torch - output_triton))}') + print( + f"The maximum difference between torch and triton is " + f"{torch.max(torch.abs(output_torch - output_triton))}" + ) + if __name__ == "__main__": test_add() diff --git a/test/dsl/test_compile_hint.py b/test/dsl/test_compile_hint.py index d0386d09..4d75b42a 100644 --- a/test/dsl/test_compile_hint.py +++ b/test/dsl/test_compile_hint.py @@ -2,14 +2,16 @@ import triton import triton.language as tl import pytest -import triton.language.extra.deeplink as dl +import triton.language.extra.deeplink.cann.extension as dl # eg: pytest -v test_compile_hint.py::test_compile_hint ############################# @triton.jit -def triton_compile_hint(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr): +def triton_compile_hint( + in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_SUB: tl.constexpr +): xoffset = tl.program_id(0) * XBLOCK for xoffset_sub in range(0, XBLOCK, XBLOCK_SUB): xindex = xoffset + xoffset_sub + tl.arange(0, XBLOCK_SUB)[:] @@ -23,20 +25,23 @@ def triton_compile_hint(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr, XBLOCK_ tl.store(out_ptr0 + (xindex), tmp2, xmask) -@pytest.mark.parametrize('param_list', - [ - ['float32', (2, 4096, 8), 2, 32768, 1024], - ] - ) +@pytest.mark.parametrize( + "param_list", + [ + ["float32", (2, 4096, 8), 2, 32768, 1024], + ], +) def test_compile_hint(param_list): dtype_str, shape, ncore, xblock, xblock_sub = param_list dtype = getattr(torch, dtype_str) x0 = torch.rand(shape, dtype=dtype).npu() y_ref = x0 y_cal = torch.rand(shape, dtype=dtype).npu() - triton_compile_hint[(ncore, )](x0, y_cal, x0.numel(), xblock, xblock_sub) + triton_compile_hint[(ncore,)](x0, y_cal, x0.numel(), xblock, xblock_sub) assert torch.allclose(y_cal, y_ref) - print(f'The maximum difference between torch and triton is ' - f'{torch.max(torch.abs(y_cal - y_ref))}') + print( + f"The maximum difference between torch and triton is " + f"{torch.max(torch.abs(y_cal - y_ref))}" + ) assert y_cal.dtype == y_ref.dtype print(f"dtype is same.") diff --git a/third_party/ascendnpu-ir b/third_party/ascendnpu-ir index af5499b3..ef9139b3 160000 --- a/third_party/ascendnpu-ir +++ b/third_party/ascendnpu-ir @@ -1 +1 @@ -Subproject commit af5499b3b9f3dbab50b2834bcfff5da5c2a1d920 +Subproject commit ef9139b323e25e8dae0a812ac585f3b47ab5d955 diff --git a/third_party/triton b/third_party/triton index e44bd1c8..c3c476f3 160000 --- a/third_party/triton +++ b/third_party/triton @@ -1 +1 @@ -Subproject commit e44bd1c83c1c3e8deac7c4f02683cfb3cc395c8b +Subproject commit c3c476f357f1e9768ea4e45aa5c17528449ab9ef diff --git a/third_party/triton_shared b/third_party/triton_shared deleted file mode 160000 index 2b728ad9..00000000 --- a/third_party/triton_shared +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 2b728ad97bc02af821a0805b09075838911d4c19 diff --git a/tools/dicp_triton_opt/CMakeLists.txt b/tools/dicp_triton_opt/CMakeLists.txt index bd00d88a..1dd64186 100644 --- a/tools/dicp_triton_opt/CMakeLists.txt +++ b/tools/dicp_triton_opt/CMakeLists.txt @@ -1,5 +1,7 @@ get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS) get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS) +get_property(translation_libs GLOBAL PROPERTY MLIR_TRANSLATION_LIBS) +get_property(extension_libs GLOBAL PROPERTY MLIR_EXTENSION_LIBS) add_llvm_executable(dicp_opt dicp_triton_opt.cpp PARTIAL_SOURCES_INTENDED) @@ -7,25 +9,31 @@ llvm_update_compile_flags(dicp_opt) target_link_libraries(dicp_opt PRIVATE TritonAnalysis TritonTransforms - TritonGPUTransforms - TritonSharedAnalysis ${dialect_libs} ${translation_libs} ${conversion_libs} ${extension_libs} MLIRRegisterAllPasses - DICPNPU - LinalgToNPU - LinalgToLinked - TritonExtTransforms - LinalgExtTransforms - LinkedToHIVM - DICPLinalgExt - DiscreteMaskAccessConversion + # DICP NPU pass libraries TritonToLinalg - TritonTilingExtIR - TritonToLinalgNPUCoversion + TritonToStructured + TritonToUnstructure + TritonToHIVM + TritonToAnnotation + TritonToHFusion + TritonToLLVM + TritonToGraph + AutoBlockify + AscendLegalize + CommonIRTransforms + DiscreteMaskAccessConversion + MLIRTritonNPUUtils + TritonAffinityOpt + TritonDicpIR + TritonStructuredIR + TritonNvidiaGPUTransforms + MLIROptLib MLIRPass MLIRTransforms diff --git a/tools/dicp_triton_opt/dicp_triton_opt.cpp b/tools/dicp_triton_opt/dicp_triton_opt.cpp index bd2a0c15..676b7677 100644 --- a/tools/dicp_triton_opt/dicp_triton_opt.cpp +++ b/tools/dicp_triton_opt/dicp_triton_opt.cpp @@ -1,14 +1,17 @@ -#include "dicp/Conversion/DiscreteMaskAccessConversion/Passes.h" -#include "dicp/Conversion/LinalgToLinked/Passes.h" -#include "dicp/Conversion/LinalgToNPU/Passes.h" -#include "dicp/Conversion/LinkedToHIVM/Passes.h" -#include "dicp/Conversion/TritonToLinalgNPU/MemRefCopyGatherToTensorInsert/Passes.h" -#include "dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/Passes.h" -#include "dicp/Conversion/TritonToUnstructure/Passes.h" -#include "dicp/Dialect/LinalgExt/IR/LinalgExtOps.h" -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h" -#include "dicp/Dialect/NPU/IR/NPUDialect.h" -#include "dicp/Dialect/TritonExt/Transforms/Passes.h" +#include "bishengir/InitAllDialects.h" +#include "dicp/AscendLegalize/Passes.h" +#include "dicp/AutoBlockify/Passes.h" +#include "dicp/Dialect/CommonIR/Passes.h" +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" +#include "dicp/DiscreteMaskAccessConversion/Passes.h" +#include "dicp/TritonAffinityOpt/Passes.h" +#include "dicp/TritonToAnnotation/Passes.h" +#include "dicp/TritonToHFusion/Passes.h" +#include "dicp/TritonToHIVM/Passes.h" +#include "dicp/TritonToLLVM/Passes.h" +#include "dicp/TritonToLinalg/Passes.h" +#include "dicp/TritonToStructured/Passes.h" +#include "dicp/TritonToUnstructure/Passes.h" #include "mlir/Conversion/ArithToEmitC/ArithToEmitC.h" #include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h" @@ -80,43 +83,39 @@ #include "mlir/Target/LLVMIR/Dialect/XeVM/XeVMToLLVMIRTranslation.h" #include "mlir/Tools/mlir-opt/MlirOptMain.h" -#include "triton-shared/Conversion/TritonToLinalgExperimental/Passes.h.inc" -#include "triton-shared/Dialect/TritonTilingExt/IR/TritonTilingExtDialect.h" - using namespace mlir; inline void registerDICPDialects(mlir::DialectRegistry ®istry) { mlir::registerAllPasses(); mlir::registerLinalgPasses(); - mlir::triton::registerDiscreteMaskAccessConversionPass(); - mlir::triton::registerTritonToUnstructurePass(); - mlir::triton::registerBubbleUpOperationPass(); - - dicp::npu::registerLinalgToNPUPass(); - dicp::linked::registerLinalgToLinkedPass(); - dicp::trtion_ext::registerCanonicalizeTritonIRAscendPass(); - dicp::trtion_ext::registerCanonicalizeCmpiPass(); - dicp::linked::registerLinkedToHIVMPass(); - dicp::linked::registerTritonToLinalgNPUCoversionPass(); - dicp::linked::registerMemRefCopyGatherToTensorInsertPass(); - dicp::linked::registerDebugCPUVerifyPass(); + // triton-dicp pass registrations + triton::registerAutoBlockifyPass(); + triton::registerAscendLegalizePass(); + mlir::dicp::CommonIR::registerCommonIRPasses(); + triton::registerTritonToStructuredPass(); + triton::registerDiscreteMaskAccessConversionPass(); + triton::registerTritonToAnnotationPass(); + triton::registerTritonToUnstructurePass(); + triton::registerTritonToHIVMPass(); + triton::registerTritonToHFusionPass(); + triton::registerTritonToLLVMPass(); + triton::registerBubbleUpOperationPass(); + triton::registerTritonToLinalgPass(); + triton::registerDAGSyncPass(); + triton::registerDAGScopePass(); + triton::registerDAGSSBufferPass(); - dicp::LinalgExt::registerLinalgIfToSelectPass(); - dicp::LinalgExt::registerLinalgGenericToSCFPass(); - dicp::LinalgExt::registerScalarTo1DTensorPass(); - dicp::LinalgExt::registerNormalizeSliceOpsPass(); - dicp::LinalgExt::registerVectorizeParallelLoopPass(); + registry + .insert(); - registry.insert(); + bishengir::registerAllDialects(registry); } int main(int argc, char **argv) { diff --git a/triton_dicp_triton.cc b/triton_dicp_triton.cc index 5c39d4f0..a2a1c884 100644 --- a/triton_dicp_triton.cc +++ b/triton_dicp_triton.cc @@ -1,22 +1,36 @@ -#include "dicp/Conversion/DiscreteMaskAccessConversion/Passes.h" -#include "dicp/Conversion/LinalgToLinked/LinalgToLinked.h" -#include "dicp/Conversion/LinalgToLinked/Passes.h" -#include "dicp/Conversion/LinalgToNPU/Passes.h" -#include "dicp/Conversion/LinkedToHIVM/Passes.h" -#include "dicp/Conversion/TritonToLinalgNPU/TritonToLinalgNPUCoversion/Passes.h" -#include "dicp/Conversion/TritonToUnstructure/BubbleUpOperation.h" -#include "dicp/Conversion/TritonToUnstructure/UnstructureConversionPass.h" -#include "dicp/Dialect/LinalgExt/Transforms/Passes.h" -#include "dicp/Dialect/TritonExt/Transforms/Passes.h" - -#include "triton-shared/Dialect/TritonTilingExt/IR/TritonTilingExtDialect.h" +#include "ir.h" #include "triton/Dialect/Triton/IR/Dialect.h" +#include "dicp/AscendLegalize/Passes.h" +#include "dicp/AutoBlockify/Passes.h" +#include "dicp/Dialect/CommonIR/Passes.h" +#include "dicp/Dialect/TritonDicp/IR/TritonDicpDialect.h" +#include "dicp/DiscreteMaskAccessConversion/Passes.h" +#include "dicp/TritonAffinityOpt/Passes.h" +#include "dicp/TritonToAnnotation/Passes.h" +#include "dicp/TritonToHFusion/Passes.h" +#include "dicp/TritonToHIVM/Passes.h" +#include "dicp/TritonToLLVM/Passes.h" +#include "dicp/TritonToLinalg/Passes.h" +#include "dicp/TritonToStructured/Passes.h" +#include "dicp/TritonToUnstructure/Passes.h" + +#include "bishengir/Dialect/Annotation/IR/Annotation.h" +#include "bishengir/Dialect/HACC/IR/HACC.h" +#include "bishengir/Dialect/HIVM/IR/HIVM.h" +#include "bishengir/Dialect/Scope/IR/Scope.h" + +#include "mlir/AsmParser/AsmParser.h" #include "mlir/Conversion/AffineToStandard/AffineToStandard.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Arith/Utils/Utils.h" +#include "mlir/Dialect/Bufferization/IR/Bufferization.h" #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" +#include "mlir/Dialect/Func/Extensions/InlinerExtension.h" #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/Dialect/Index/IR/IndexDialect.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/Linalg/Passes.h" #include "mlir/Dialect/Math/IR/Math.h" @@ -24,98 +38,1257 @@ #include "mlir/Dialect/SCF/Transforms/Passes.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/Dialect/Transform/IR/TransformDialect.h" -#include "mlir/InitAllDialects.h" -#include "mlir/InitAllExtensions.h" +#include "mlir/IR/AffineExpr.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/Types.h" #include "mlir/InitAllPasses.h" #include "mlir/Pass/PassManager.h" #include "mlir/Pass/PassOptions.h" +#include "mlir/Support/LLVM.h" #include "mlir/Transforms/Passes.h" -#include "llvm/IR/Constants.h" +#include "llvm/IR/Instructions.h" -#include "passes.h" +#include #include #include #include namespace py = pybind11; using namespace mlir; +using namespace mlir::triton::dicp; + +// ============================================================================= +// DICPNPUIROpBuilder +// ============================================================================= + +struct DICPNPUIROpBuilder : public TritonOpBuilder { + std::string target; + static constexpr char kTarget910_95[] = "Ascend910_95"; + static constexpr char kTarget950[] = "Ascend950"; + + explicit DICPNPUIROpBuilder(MLIRContext *context, std::string target = "") + : TritonOpBuilder(context), target(target) {} + + bool is_910_95() const { + constexpr size_t kLen910 = sizeof(kTarget910_95) - 1; + bool match_910 = target.size() >= kLen910 && + target.compare(0, kLen910, kTarget910_95) == 0; + + constexpr size_t kLen950 = sizeof(kTarget950) - 1; + bool match_950 = + target.size() >= kLen950 && target.compare(0, kLen950, kTarget950) == 0; + + return match_910 || match_950; + } +}; + +namespace { + +MLIRContext *gDefaultDICPContext = nullptr; -void init_triton_dicp_triton_pass_triton_shared_ascend(py::module &&m) { - ADD_PASS_WRAPPER_0("add_canonicalize_cmpi", - dicp::trtion_ext::createCanonicalizeCmpiPass); - ADD_PASS_WRAPPER_0("add_canonicalize_triton_ir_ascend", - dicp::trtion_ext::createCanonicalizeTritonIRAscendPass); - ADD_PASS_WRAPPER_0("add_triton_to_linalg_npu", - dicp::linked::createTritonToLinalgNPUCoversionPass); - ADD_PASS_OPTION_WRAPPER_2("add_discrete_mask_access_conversion", - triton::createDiscreteMaskAccessConversionPass, - bool, bool); - ADD_PASS_WRAPPER_0("add_triton_to_unstructure", - triton::createTritonToUnstructurePass); - ADD_PASS_WRAPPER_0("add_bubble_up_operation", - triton::createBubbleUpOperationPass); +MLIRContext *resolveContext(const py::object &contextObj) { + if (!contextObj.is_none()) { + return &py::cast(contextObj); + } + if (gDefaultDICPContext) { + return gDefaultDICPContext; + } + throw std::invalid_argument( + "No default MLIR context. Pass context explicitly or call " + "dicp_ir.load_dialects(context) first."); } -void init_triton_dicp_triton_pass_linked_npu(py::module &&m) { - ADD_PASS_WRAPPER_0("add_lower_affine", createLowerAffinePass); - m.def("add_normalize_slice_ops", [](mlir::PassManager &pm) { - pm.addNestedPass( - dicp::LinalgExt::createNormalizeSliceOpsPass()); +struct ModeAndPipes { + hivm::SyncBlockModeAttr modeAttr = {}; + hivm::PipeAttr cubePipe = {}; + hivm::PipeAttr vectorPipe = {}; +}; + +hivm::TCoreTypeAttr GetCore(MLIRContext *ctx, llvm::StringRef opName, + llvm::StringRef sender) { + hivm::TCoreTypeAttr core; + if (sender == "cube") { + if (opName == "sync_block_set") + core = hivm::TCoreTypeAttr::get(ctx, hivm::TCoreType::CUBE); + else + core = hivm::TCoreTypeAttr::get(ctx, hivm::TCoreType::VECTOR); + } else { + if (sender != "vector") { + throw std::runtime_error( + "sync_block_set/wait only supports 'cube' or 'vector' as sender"); + } + if (opName == "sync_block_set") + core = hivm::TCoreTypeAttr::get(ctx, hivm::TCoreType::VECTOR); + else + core = hivm::TCoreTypeAttr::get(ctx, hivm::TCoreType::CUBE); + } + return core; +} + +void buildSyncBlockOp(DICPNPUIROpBuilder &self, const std::string &opNameSnake, + std::string &sender, std::string &receiver, Value id, + hivm::PIPE senderPipe, hivm::PIPE receiverPipe) { + auto *ctx = self.getBuilder().getContext(); + hivm::TCoreTypeAttr coreAttr = GetCore(ctx, opNameSnake, sender); + hivm::PipeAttr prodPipe = hivm::PipeAttr::get(ctx, senderPipe); + hivm::PipeAttr consPipe = hivm::PipeAttr::get(ctx, receiverPipe); + const size_t I64 = 64; + auto i64Ty = IntegerType::get(ctx, I64); + Value idI64 = id; + if (!id.getType().isInteger(I64)) { + idI64 = mlir::convertScalarToDtype(self.getBuilder(), id.getLoc(), id, + i64Ty, true); + } + if (opNameSnake == "sync_block_set") { + self.create(coreAttr, prodPipe, consPipe, idI64); + } else if (opNameSnake == "sync_block_wait") { + self.create(coreAttr, prodPipe, consPipe, idI64); + } else { + throw std::runtime_error("Unsupported operation name for SyncBlockOp"); + } +} + +ModeAndPipes GetSyncBlockModeAndPipes(MLIRContext *ctx, + const std::string &mode) { + hivm::SyncBlockModeAttr modeAttr = {}; + hivm::PipeAttr cubePipe = {}; + hivm::PipeAttr vectorPipe = {}; + + if (mode == "all_cube") { + modeAttr = hivm::SyncBlockModeAttr::get(ctx, hivm::SyncBlockMode::ALL_CUBE); + cubePipe = hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_ALL); + vectorPipe = hivm::PipeAttr{}; + } else if (mode == "all_vector") { + modeAttr = + hivm::SyncBlockModeAttr::get(ctx, hivm::SyncBlockMode::ALL_VECTOR); + cubePipe = hivm::PipeAttr{}; + vectorPipe = hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_ALL); + } else if (mode == "all") { + modeAttr = hivm::SyncBlockModeAttr::get(ctx, hivm::SyncBlockMode::ALL); + cubePipe = hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_ALL); + vectorPipe = hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_ALL); + } else if (mode == "all_sub_vector") { + modeAttr = + hivm::SyncBlockModeAttr::get(ctx, hivm::SyncBlockMode::ALL_SUB_VECTOR); + cubePipe = hivm::PipeAttr{}; + vectorPipe = hivm::PipeAttr::get(ctx, hivm::PIPE::PIPE_ALL); + } else { + llvm::report_fatal_error( + llvm::StringRef("Invalid sync-block mode: " + mode)); + } + return {modeAttr, cubePipe, vectorPipe}; +} + +} // namespace + +// ============================================================================= +// init_dicp_ir: IR builder bindings (merged) +// ============================================================================= + +void init_dicp_ir(py::module &&m) { + // --- AffineExpr bindings --- + auto affineExprClass = + py::class_(m, "affine_expr", py::module_local()); + affineExprClass + .def("__str__", + [](AffineExpr self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return os.str(); + }) + .def("__repr__", + [](AffineExpr self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return ""; + }) + .def("is_symbolic_or_constant", &AffineExpr::isSymbolicOrConstant) + .def("is_pure_affine", &AffineExpr::isPureAffine) + .def("is_function_of_dim", &AffineExpr::isFunctionOfDim) + .def("compose", + [](AffineExpr self, AffineMap map) { return self.compose(map); }) + .def("get_largest_known_divisor", &AffineExpr::getLargestKnownDivisor) + .def("floordiv", [](AffineExpr self, + AffineExpr other) { return self.floorDiv(other); }) + .def("ceildiv", [](AffineExpr self, + AffineExpr other) { return self.ceilDiv(other); }) + .def("mod", + [](AffineExpr self, AffineExpr other) { return self % other; }) + .def("__hash__", + [](AffineExpr self) { + return py::int_(static_cast(mlir::hash_value(self))); + }) + .def("__eq__", [](AffineExpr lhs, AffineExpr rhs) { return lhs == rhs; }) + .def(py::self + py::self) + .def(py::self - py::self) + .def(py::self * py::self) + .def(py::self % py::self); + affineExprClass + .def_static( + "get_constant", + [](int64_t val, py::object contextObj) { + auto *context = resolveContext(contextObj); + return getAffineConstantExpr(val, context); + }, + py::arg("value"), py::arg("context") = py::none()) + .def_static( + "get_dim", + [](uint32_t pos, py::object contextObj) { + auto *context = resolveContext(contextObj); + return getAffineDimExpr(pos, context); + }, + py::arg("pos"), py::arg("context") = py::none()) + .def_static( + "get_symbol", + [](uint32_t pos, py::object contextObj) { + auto *context = resolveContext(contextObj); + return getAffineSymbolExpr(pos, context); + }, + py::arg("pos"), py::arg("context") = py::none()); + + py::class_(m, "affine_constant_expr", + py::module_local()) + .def("get_value", &AffineConstantExpr::getValue); + py::class_(m, "affine_dim_expr", + py::module_local()) + .def("get_position", &AffineDimExpr::getPosition); + py::class_(m, "affine_symbol_expr", + py::module_local()) + .def("get_position", &AffineSymbolExpr::getPosition); + py::class_(m, "affine_binary_op_expr", + py::module_local()) + .def("get_lhs", &AffineBinaryOpExpr::getLHS) + .def("get_rhs", &AffineBinaryOpExpr::getRHS); + + // --- AffineMap bindings --- + auto affineMapClass = + py::class_(m, "affine_map", py::module_local()); + affineMapClass + .def("__str__", + [](AffineMap &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return os.str(); + }) + .def("__repr__", + [](AffineMap &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return ""; + }) + .def("is_identity", &AffineMap::isIdentity) + .def("is_permutation", &AffineMap::isPermutation) + .def("get_num_dims", &AffineMap::getNumDims) + .def("get_num_symbols", &AffineMap::getNumSymbols) + .def("get_num_results", &AffineMap::getNumResults) + .def("is_empty", &AffineMap::isEmpty) + .def("is_single_constant", &AffineMap::isSingleConstant) + .def("is_constant", &AffineMap::isConstant) + .def("get_constant_result", + [](AffineMap &self) -> int64_t { + if (!self.isSingleConstant()) { + throw std::runtime_error( + "affine map is not a single constant map"); + } + return self.getSingleConstantResult(); + }) + .def("get_result", + [](AffineMap &self, uint32_t pos) { + if (pos >= self.getNumResults()) { + throw py::index_error("result index out of range"); + } + return self.getResult(pos); + }) + .def("get_sub_map", + [](AffineMap &self, const std::vector &resultPos) { + return self.getSubMap(resultPos); + }) + .def("replace", + [](AffineMap &self, AffineExpr expr, AffineExpr replacement, + uint32_t numResultDims, uint32_t numResultSymbols) { + return self.replace(expr, replacement, numResultDims, + numResultSymbols); + }) + .def("compose", + [](AffineMap &self, AffineMap map) { return self.compose(map); }) + .def("get_results", + [](AffineMap &self) -> std::vector { + auto results = self.getResults(); + return std::vector(results.begin(), results.end()); + }) + .def("__hash__", + [](AffineMap &self) { + return py::int_(static_cast(mlir::hash_value(self))); + }) + .def("__eq__", [](AffineMap &lhs, AffineMap &rhs) { return lhs == rhs; }) + .def("inverse_permutation", + [](AffineMap &self) -> py::object { + if (!self.isPermutation()) { + throw py::value_error( + "AffineMap must be a valid permutation to compute inverse"); + } + AffineMap inverse = mlir::inversePermutation(self); + if (!inverse) { + throw py::value_error("Failed to compute inverse permutation"); + } + return py::cast(inverse); + }) + .def("to_dict", [](AffineMap &self) -> py::dict { + py::list results; + for (AffineExpr result : self.getResults()) { + if (auto dimExpr = dyn_cast(result)) { + results.append(dimExpr.getPosition()); + } else { + std::string exprStr; + llvm::raw_string_ostream os(exprStr); + result.print(os); + results.append(py::str(exprStr)); + } + } + py::dict ret; + ret["num_dims"] = self.getNumDims(); + ret["num_symbols"] = self.getNumSymbols(); + ret["results"] = std::move(results); + return ret; + }); + affineMapClass + .def_static( + "get", + [](int64_t numDims, int64_t numSymbols, const py::iterable &resultsIn, + py::object contextObj) -> AffineMap { + MLIRContext *context = nullptr; + if (numDims < 0 || numSymbols < 0) { + throw std::invalid_argument( + "num_dims and num_symbols must be non-negative"); + } + llvm::SmallVector results; + for (const auto &item : resultsIn) { + if (py::isinstance(item)) { + auto expr = py::cast(item); + if (!context) { + context = expr.getContext(); + } + results.push_back(expr); + continue; + } + if (py::isinstance(item)) { + if (!context) { + context = resolveContext(contextObj); + } + int64_t pos = py::cast(item); + if (pos < 0 || pos >= numDims) { + throw std::invalid_argument( + "result dim index is out of range for num_dims"); + } + results.push_back(getAffineDimExpr(pos, context)); + continue; + } + throw std::invalid_argument( + "results must contain affine_expr or int dim indices"); + } + if (!context) { + context = resolveContext(contextObj); + } + return AffineMap::get(numDims, numSymbols, results, context); + }, + py::arg("num_dims"), py::arg("num_symbols"), py::arg("result_dims"), + py::arg("context") = py::none()) + .def_static( + "get_identity", + [](int64_t numDims, py::object contextObj) -> AffineMap { + auto *context = resolveContext(contextObj); + if (numDims < 0) { + throw std::invalid_argument("num_dims must be non-negative"); + } + return AffineMap::getMultiDimIdentityMap(numDims, context); + }, + py::arg("num_dims"), py::arg("context") = py::none()) + .def_static( + "get_minor_identity", + [](int64_t dims, int64_t results, py::object contextObj) { + auto *context = resolveContext(contextObj); + if (dims < 0 || results < 0) { + throw std::invalid_argument("dims/results must be non-negative"); + } + return AffineMap::getMinorIdentityMap(dims, results, context); + }, + py::arg("dims"), py::arg("results"), py::arg("context") = py::none()) + .def_static( + "get_empty", + [](py::object contextObj) { + auto *context = resolveContext(contextObj); + return AffineMap::get(0, 0, {}, context); + }, + py::arg("context") = py::none()) + .def_static( + "get_permutation", + [](const std::vector &permutation, py::object contextObj) { + auto *context = resolveContext(contextObj); + return AffineMap::getPermutationMap(permutation, context); + }, + py::arg("permutation"), py::arg("context") = py::none()) + .def_static( + "get_constant", + [](int64_t value, py::object contextObj) { + auto *context = resolveContext(contextObj); + return AffineMap::getConstantMap(value, context); + }, + py::arg("value"), py::arg("context") = py::none()); + + // --- hivm enums --- + py::enum_(m, "AddressSpace", py::module_local()) + .value("L1", hivm::AddressSpace::L1) + .value("UB", hivm::AddressSpace::UB) + .value("L0A", hivm::AddressSpace::L0A) + .value("L0B", hivm::AddressSpace::L0B) + .value("L0C", hivm::AddressSpace::L0C) + .export_values(); + + py::enum_(m, "CoreType", py::module_local()) + .value("CUBE", hivm::TCoreType::CUBE) + .value("VECTOR", hivm::TCoreType::VECTOR) + .value("CUBE_OR_VECTOR", hivm::TCoreType::CUBE_OR_VECTOR) + .value("CUBE_AND_VECTOR", hivm::TCoreType::CUBE_AND_VECTOR) + .export_values(); + + py::enum_(m, "PIPE", py::module_local()) + .value("PIPE_S", hivm::PIPE::PIPE_S) + .value("PIPE_V", hivm::PIPE::PIPE_V) + .value("PIPE_M", hivm::PIPE::PIPE_M) + .value("PIPE_MTE1", hivm::PIPE::PIPE_MTE1) + .value("PIPE_MTE2", hivm::PIPE::PIPE_MTE2) + .value("PIPE_MTE3", hivm::PIPE::PIPE_MTE3) + .value("PIPE_ALL", hivm::PIPE::PIPE_ALL) + .value("PIPE_FIX", hivm::PIPE::PIPE_FIX) + .export_values(); + + py::enum_(m, "MODE", py::module_local()) + .value("SIMD", hivm::VFMode::SIMD) + .value("SIMT", hivm::VFMode::SIMT) + .value("MIX", hivm::VFMode::MIX) + .export_values(); + + py::enum_(m, "IteratorType", py::module_local()) + .value("Parallel", hivm::IteratorType::kParallel) + .value("Broadcast", hivm::IteratorType::kBroadcast) + .value("Transpose", hivm::IteratorType::kTranspose) + .value("Reduction", hivm::IteratorType::kReduction) + .value("Interleave", hivm::IteratorType::kInterleave) + .value("Deinterleave", hivm::IteratorType::kDeinterleave) + .value("Inverse", hivm::IteratorType::kInverse) + .value("Pad", hivm::IteratorType::kPad) + .value("Concat", hivm::IteratorType::kConcat) + .value("Gather", hivm::IteratorType::kGather) + .value("Cumulative", hivm::IteratorType::kCumulative) + .value("Opaque", hivm::IteratorType::kOpaque) + .export_values(); + + py::enum_(m, "FixpipeDMAMode", py::module_local()) + .value("NZ2DN", hivm::FixpipeDMAMode::NZ2DN) + .value("NZ2ND", hivm::FixpipeDMAMode::NZ2ND) + .value("NZ2NZ", hivm::FixpipeDMAMode::NZ2NZ) + .export_values(); + + py::enum_(m, "FixpipeDualDstMode", + py::module_local()) + .value("NO_DUAL", hivm::FixpipeDualDstMode::NO_DUAL) + .value("COLUMN_SPLIT", hivm::FixpipeDualDstMode::COLUMN_SPLIT) + .value("ROW_SPLIT", hivm::FixpipeDualDstMode::ROW_SPLIT) + .export_values(); + + py::enum_(m, "FixpipePreQuantMode", + py::module_local()) + .value("NO_QUANT", hivm::FixpipePreQuantMode::NO_QUANT) + .value("F322BF16", hivm::FixpipePreQuantMode::F322BF16) + .value("F322F16", hivm::FixpipePreQuantMode::F322F16) + .value("S322I8", hivm::FixpipePreQuantMode::S322I8) + .export_values(); + + py::enum_(m, "FixpipePreReluMode", + py::module_local()) + .value("LEAKY_RELU", hivm::FixpipePreReluMode::LEAKY_RELU) + .value("NO_RELU", hivm::FixpipePreReluMode::NO_RELU) + .value("NORMAL_RELU", hivm::FixpipePreReluMode::NORMAL_RELU) + .value("P_RELU", hivm::FixpipePreReluMode::P_RELU) + .export_values(); + + py::enum_(m, "DataLayout", py::module_local()) + .value("nZ", hivm::DataLayout::nZ) + .value("zN", hivm::DataLayout::zN) + .export_values(); + + m.def("load_dialects", [](MLIRContext &context) { + gDefaultDICPContext = &context; + DialectRegistry registry; + registry.insert(); + mlir::func::registerInlinerExtension(registry); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); }); - ADD_PASS_WRAPPER_0("add_linalg_if_to_select", - dicp::LinalgExt::createLinalgIfToSelectPass); - ADD_PASS_WRAPPER_0("add_linalg_generic_to_scf", - dicp::LinalgExt::createLinalgGenericToSCFPass); - m.def("add_scalar_to_1d_tensor", [](mlir::PassManager &pm) { + + // --- dicp_npu_ir_builder class --- + py::class_( + m, "dicp_npu_ir_builder", py::module_local(), py::dynamic_attr()) + .def(py::init(), py::arg("context"), + py::arg("target") = "") + .def("get_int_attr", + [](DICPNPUIROpBuilder &self, int64_t value) -> Attribute { + return IntegerAttr::get(self.getBuilder().getI64Type(), value); + }) + .def("get_str_array_attr", + [](DICPNPUIROpBuilder &self, + const std::vector &values) -> Attribute { + auto *ctx = self.getBuilder().getContext(); + llvm::SmallVector attrs; + attrs.reserve(values.size()); + for (const auto &v : values) + attrs.push_back(self.getBuilder().getStringAttr(v)); + return ArrayAttr::get(ctx, attrs); + }) + .def("get_i64_array_attr", + [](DICPNPUIROpBuilder &self, + const std::vector &values) -> Attribute { + return self.getBuilder().getI64ArrayAttr(values); + }) + .def( + "get_core_type_attr", + [](DICPNPUIROpBuilder &self, hivm::TCoreType core_type) -> Attribute { + return self.getBuilder().getAttr(core_type); + }) + .def("get_pipe_attr", + [](DICPNPUIROpBuilder &self, hivm::PIPE pipe) -> Attribute { + return self.getBuilder().getAttr(pipe); + }) + .def("get_vf_mode_attr", + [](DICPNPUIROpBuilder &self, hivm::VFMode mode) -> Attribute { + return self.getBuilder().getAttr(mode); + }) + .def("get_iterator_types_attr", + [](DICPNPUIROpBuilder &self, + const std::vector &array) { + auto attrs = llvm::to_vector( + llvm::map_range(array, [&self](hivm::IteratorType type) { + return cast( + self.getBuilder().getAttr(type)); + })); + return self.getBuilder().getArrayAttr(attrs); + }) + .def("get_t_core_type_attr_name", + [](DICPNPUIROpBuilder &self) -> std::string { + return hivm::TCoreTypeAttr::name.str(); + }) + .def("get_t_core_type_cube_attr", + [](DICPNPUIROpBuilder &self) -> Attribute { + return hivm::TCoreTypeAttr::get(self.getBuilder().getContext(), + hivm::TCoreType::CUBE); + }) + .def("get_t_core_type_vector_attr", + [](DICPNPUIROpBuilder &self) -> Attribute { + return hivm::TCoreTypeAttr::get(self.getBuilder().getContext(), + hivm::TCoreType::VECTOR); + }) + .def("parse_attr", + [](TritonOpBuilder &self, std::string value) -> Attribute { + auto *ctx = self.getBuilder().getContext(); + ctx->allowUnregisteredDialects(); + return mlir::parseAttribute(value, ctx); + }) + .def("get_affine_map_attr", + [](DICPNPUIROpBuilder &self, AffineMap affineMap) -> Attribute { + return AffineMapAttr::get(affineMap); + }) + .def("get_affine_map_array_attr", + [](DICPNPUIROpBuilder &self, + const std::vector &affineMaps) -> Attribute { + auto *ctx = self.getBuilder().getContext(); + llvm::SmallVector attrs; + attrs.reserve(affineMaps.size()); + for (const auto &map : affineMaps) { + attrs.push_back(AffineMapAttr::get(map)); + } + return ArrayAttr::get(ctx, attrs); + }) + .def("get_buffer_ty_with_affine_map", + [](DICPNPUIROpBuilder &self, std::vector &shape, + Type &elementType, AffineMap affineMap, + const Attribute &memorySpace) -> Type { + auto layout = AffineMapAttr::get(affineMap); + return MemRefType::get(shape, elementType, layout, memorySpace); + }) + .def("create_fixpipe", + [](DICPNPUIROpBuilder &self, Value src, py::object dst_obj, + hivm::FixpipeDMAMode dma_mode, + hivm::FixpipeDualDstMode dual_dst_mode, + hivm::FixpipePreQuantMode pre_quant_mode, + hivm::FixpipePreReluMode pre_relu_mode) -> py::object { + if (!dyn_cast(src.getType())) { + llvm_unreachable("src is not of RankedTensorType"); + } + auto *ctx = self.getBuilder().getContext(); + auto loc = self.getLastLoc(); + Value dstValue; + bool needCreateDst = dst_obj.is_none(); + if (needCreateDst) { + auto srcType = dyn_cast(src.getType()); + auto srcShape = srcType.getShape(); + llvm::SmallVector dstShape(srcShape.begin(), + srcShape.end()); + if (dual_dst_mode == hivm::FixpipeDualDstMode::ROW_SPLIT) { + if (dstShape.size() >= 1 && dstShape[0] > 0) { + dstShape[0] = dstShape[0] / 2; + } + } else if (dual_dst_mode == + hivm::FixpipeDualDstMode::COLUMN_SPLIT) { + if (dstShape.size() >= 2 && dstShape[1] > 0) { + dstShape[1] = dstShape[1] / 2; + } + } + auto dstType = + RankedTensorType::get(dstShape, srcType.getElementType()); + auto emptyTensor = self.create( + dstType.getShape(), dstType.getElementType()); + dstValue = emptyTensor.getResult(); + } else { + dstValue = py::cast(dst_obj); + if (!dyn_cast(dstValue.getType())) { + llvm_unreachable("dst is not of ShapedType"); + } + } + auto dma_mode_attr = + mlir::hivm::FixpipeDMAModeAttr::get(ctx, dma_mode); + auto dual_dst_mode_attr = + mlir::hivm::FixpipeDualDstModeAttr::get(ctx, dual_dst_mode); + auto pre_quant_mode_attr = + mlir::hivm::FixpipePreQuantModeAttr::get(ctx, pre_quant_mode); + auto pre_relu_mode_attr = + mlir::hivm::FixpipePreReluModeAttr::get(ctx, pre_relu_mode); + auto channel_split = BoolAttr::get(ctx, false); + if (needCreateDst) { + return py::cast( + self.create( + mlir::TypeRange{dstValue.getType()}, src, dstValue, + dma_mode_attr, dual_dst_mode_attr, + pre_quant_mode_attr, pre_relu_mode_attr, + channel_split) + .getResult(0)); + } else { + self.create(mlir::TypeRange{}, src, dstValue, + dma_mode_attr, dual_dst_mode_attr, + pre_quant_mode_attr, + pre_relu_mode_attr, channel_split); + return py::none(); + } + }) + .def("create_bind_buffer", + [](TritonOpBuilder &self, Value &src, Value &alloc) -> void { + auto ctx = self.getBuilder().getContext(); + auto bind = StringAttr::get(ctx, "bind_buffer"); + self.create(src, ValueRange{alloc}, + ArrayAttr::get(ctx, bind)); + }) + .def("create_debug_barrier", + [](TritonOpBuilder &self, Value &ptr, const std::string &attrKey, + Attribute &attrVal) { + auto annotationOp = self.create(ptr); + annotationOp->setAttr(self.getBuilder().getStringAttr(attrKey), + attrVal); + }) + .def("create_custom_op", + [](DICPNPUIROpBuilder &self, const std::string &name, + const py::dict &attrs, const std::vector &ins, + const std::vector &outs, + const std::vector &arg_attrs) -> std::vector { + ValueRange inputs{ins}; + ValueRange outputs{outs}; + ValueRange temp_buffers{}; + TypeRange res_types{outputs}; + auto op = self.create(res_types, name, inputs, + outputs, temp_buffers); + for (auto &attr : attrs) { + std::string attr_name = py::cast(attr.first); + Attribute attr_value = py::cast(attr.second); + op->setAttr(attr_name, attr_value); + } + SmallVector dictAttrs(arg_attrs.size()); + Attribute emptyDict = self.getBuilder().getDictionaryAttr({}); + for (const auto &[idx, attrs] : llvm::enumerate(arg_attrs)) { + if (idx >= op.getNumOperands()) + continue; + if (attrs.is_none()) { + dictAttrs[idx] = emptyDict; + continue; + } + llvm::SmallVector namedAttrs; + for (const auto &attr : attrs) { + std::string attr_name = py::cast(attr.first); + Attribute attr_value = py::cast(attr.second); + namedAttrs.push_back(NamedAttribute( + self.getBuilder().getStringAttr(attr_name), attr_value)); + } + dictAttrs[idx] = self.getBuilder().getDictionaryAttr(namedAttrs); + } + ArrayAttr arg_attrs_array = + self.getBuilder().getArrayAttr(dictAttrs); + op->setAttr("arg_attrs", arg_attrs_array); + auto results = op->getResults(); + return std::vector(results.begin(), results.end()); + }) + .def("create_scope_op", + [](DICPNPUIROpBuilder &self, py::dict &scopeAttrs, + std::vector resultTypes) -> OpState { + llvm::SmallVector attrs; + for (auto item : scopeAttrs) { + std::string key = py::cast(item.first); + Attribute value = py::cast(item.second); + attrs.push_back( + NamedAttribute(self.getBuilder().getStringAttr(key), value)); + } + auto scopeOp = self.create(TypeRange(resultTypes)); + scopeOp->setAttrs(attrs); + return OpState(scopeOp); + }) + .def( + "scope_return", + [](DICPNPUIROpBuilder &self, std::vector operands) -> OpState { + return self.create(ValueRange(operands)); + }) + .def("sync_block_set", + [](DICPNPUIROpBuilder &self, std::string &sender, + std::string &receiver, Value id, hivm::PIPE senderPipe, + hivm::PIPE receiverPipe) -> void { + buildSyncBlockOp(self, "sync_block_set", sender, receiver, id, + senderPipe, receiverPipe); + }) + .def("sync_block_wait", + [](DICPNPUIROpBuilder &self, std::string &sender, + std::string &receiver, Value id, hivm::PIPE senderPipe, + hivm::PIPE receiverPipe) -> void { + buildSyncBlockOp(self, "sync_block_wait", sender, receiver, id, + senderPipe, receiverPipe); + }) + .def("get_target_attribute", + [](DICPNPUIROpBuilder &self, + hivm::AddressSpace &addressSpace) -> Attribute { + return hivm::AddressSpaceAttr::get(self.getBuilder().getContext(), + addressSpace); + }) + .def("create_get_sub_vec_id", + [](DICPNPUIROpBuilder &self) -> Value { + auto subBlockIdxOp = self.create(); + auto moduleOp = subBlockIdxOp->getParentOfType(); + auto *ctx = self.getBuilder().getContext(); + moduleOp->setAttr("hivm.disable_auto_tile_and_bind_subblock", + mlir::UnitAttr::get(ctx)); + return subBlockIdxOp; + }) + .def("sync_block_all", + [](DICPNPUIROpBuilder &self, std::string &mode, int id) -> void { + auto *ctx = self.getBuilder().getContext(); + auto [modeAttr, cubePipe, vectorPipe] = + GetSyncBlockModeAndPipes(ctx, mode); + mlir::IndexType indexType = mlir::IndexType::get(ctx); + mlir::IntegerAttr indexAttribute = + mlir::IntegerAttr::get(indexType, static_cast(id)); + self.create( + modeAttr, indexAttribute, mlir::Value{}, cubePipe, vectorPipe); + }) + .def("is_910_95", + [](DICPNPUIROpBuilder &self) -> bool { return self.is_910_95(); }) + .def("create_copy_buffer", + [](DICPNPUIROpBuilder &self, Value src, Value dst) { + self.create(mlir::TypeRange{}, src, dst); + }) + .def("create_copy_tensor", + [](DICPNPUIROpBuilder &self, Value src, Value dst) { + return self + .create(mlir::TypeRange{dst.getType()}, src, dst) + .getResult(0); + }) + .def("create_convert_layout", + [](DICPNPUIROpBuilder &self, Value src, Type memrefType) -> Value { + auto *ctx = self.getBuilder().getContext(); + return self + .create( + memrefType, src, + hivm::DataLayoutAttr::get(ctx, hivm::DataLayout::ND), + hivm::DataLayoutAttr::get(ctx, hivm::DataLayout::ND)) + .getResult(); + }); + + // --- DICP extension methods on TritonOpBuilder (via getBuilderClass) --- + auto *builder_cls = ir::getBuilderClass(); + if (builder_cls) { + builder_cls + ->def("create_extract_slice", + [](TritonOpBuilder &self, Value &ful, + std::vector &offs_vec, std::vector &sizs_vec, + std::vector &strd_vec) -> Value { + self.getContext() + ->getOrLoadDialect(); + llvm::SmallVector offsets; + llvm::SmallVector staticOffsets; + for (const auto &o : offs_vec) { + auto oTy = o.getType(); + if (!oTy.isIndex()) { + auto v = self.create( + self.getBuilder().getIndexType(), o); + offsets.push_back(v); + } else { + offsets.push_back(o); + } + staticOffsets.push_back(ShapedType::kDynamic); + } + llvm::SmallVector sizes; + llvm::SmallVector staticSizes; + llvm::SmallVector retSizes; + for (const auto &s : sizs_vec) { + staticSizes.push_back(s); + retSizes.push_back(s); + } + llvm::SmallVector strides; + llvm::SmallVector staticStrides; + for (const auto &s : strd_vec) { + auto v = self.create(s); + strides.push_back(v); + staticStrides.push_back(ShapedType::kDynamic); + } + auto retTy = RankedTensorType::get( + retSizes, + cast(ful.getType()).getElementType()); + return self.create( + retTy, ful, offsets, sizes, strides, staticOffsets, + staticSizes, staticStrides); + }) + .def("create_insert_slice", + [](TritonOpBuilder &self, Value &ful, Value &sub, + std::vector &offs_vec, std::vector &sizs_vec, + std::vector &strd_vec) -> Value { + self.getContext() + ->getOrLoadDialect(); + llvm::SmallVector offsets; + llvm::SmallVector staticOffsets; + for (const auto &o : offs_vec) { + auto oTy = o.getType(); + if (!oTy.isIndex()) { + auto v = self.create( + self.getBuilder().getIndexType(), o); + offsets.push_back(v); + } else { + offsets.push_back(o); + } + staticOffsets.push_back(ShapedType::kDynamic); + } + llvm::SmallVector sizes; + llvm::SmallVector staticSizes; + llvm::SmallVector retSizes; + for (const auto &s : sizs_vec) { + staticSizes.push_back(s); + retSizes.push_back(s); + } + llvm::SmallVector strides; + llvm::SmallVector staticStrides; + for (const auto &s : strd_vec) { + auto v = self.create(s); + strides.push_back(v); + staticStrides.push_back(ShapedType::kDynamic); + } + auto retTy = RankedTensorType::get( + retSizes, + cast(ful.getType()).getElementType()); + auto ret = self.create( + sub, ful, offsets, sizes, strides, staticOffsets, + staticSizes, staticStrides); + return ret; + }) + .def("create_annotation_mark", + [](TritonOpBuilder &self, Value &ptr, const std::string &attrKey, + Attribute &attrVal) { + self.getContext() + ->getOrLoadDialect(); + auto annotationOp = self.create(ptr); + annotationOp->setAttr(self.getBuilder().getStringAttr(attrKey), + attrVal); + }) + .def("create_extract_scalar", + [](TritonOpBuilder &self, Value &src, + std::vector &indices) -> Value { + llvm::SmallVector arg_indices; + for (const auto &i : indices) { + if (!i.getType().isIndex()) { + arg_indices.push_back(self.create( + self.getBuilder().getIndexType(), i)); + } else { + arg_indices.push_back(i); + } + } + return self.create(src, arg_indices); + }) + .def("create_index_select_simd", + [](TritonOpBuilder &self, Value &src, Value &index, int32_t dim, + std::vector &srcShape, std::vector &srcOffset, + std::vector &readShape, + std::vector &returnShape) -> Value { + auto &builder = self.getBuilder(); + auto loc = self.getLastLoc(); + Type elemType; + if (auto ptrTy = dyn_cast(src.getType())) { + elemType = ptrTy.getPointeeType(); + } else { + llvm::report_fatal_error( + "index_select_simd: src must be pointer type"); + } + llvm::SmallVector retShape; + for (const auto &s : returnShape) + retShape.push_back(s); + auto retTensorType = RankedTensorType::get(retShape, elemType); + llvm::SmallVector srcShapeIndex; + for (auto val : srcShape) { + if (!val.getType().isIndex()) + val = self.create(builder.getIndexType(), + val); + srcShapeIndex.push_back(val); + } + llvm::SmallVector srcOffsetIndex; + for (auto val : srcOffset) { + if (!val.getType().isIndex()) + val = self.create(builder.getIndexType(), + val); + srcOffsetIndex.push_back(val); + } + auto op = builder.create( + loc, retTensorType, src, index, + builder.getI32IntegerAttr(dim), srcShapeIndex, + srcOffsetIndex, builder.getDenseI32ArrayAttr(readShape)); + return op.getResult(); + }) + .def("create_index_put", + [](TritonOpBuilder &self, Value &ptr, Value &index, Value &value, + const int32_t dim, const int64_t indexBoundary, + std::vector &endOffset, std::vector &startOffset, + std::vector &dstStride) -> void { + auto dim_val = self.create( + self.getBuilder().getI32Type(), static_cast(dim)); + auto bound_val = self.create( + self.getBuilder().getI64Type(), + static_cast(indexBoundary)); + self.create(ptr, index, value, dim_val, + bound_val, endOffset, + startOffset, dstStride); + }) + .def("create_gather_out_to_ub", + [](TritonOpBuilder &self, Value &src, Value &index, + const int64_t indexBoundary, const int32_t dim, + std::vector &srcStride, std::vector &endOffset, + std::vector &startOffset, + std::optional &other) -> Value { + auto elemTy = cast(src.getType()).getPointeeType(); + auto idxShape = + cast(index.getType()).getShape(); + std::vector retShape(idxShape.begin(), idxShape.end()); + auto resType = RankedTensorType::get(retShape, elemTy); + auto bound_val = self.create( + self.getBuilder().getI64Type(), + static_cast(indexBoundary)); + auto dim_val = self.create( + self.getBuilder().getI32Type(), static_cast(dim)); + return self.create( + resType, src, index, bound_val, dim_val, srcStride, + endOffset, startOffset, other.value_or(Value())); + }) + .def("create_scatter_ub_to_out", + [](TritonOpBuilder &self, Value &ptr, Value &value, Value &index, + const int64_t indexBoundary, const int32_t dim, + std::vector &dstStride, std::vector &endOffset, + std::vector &startOffset) -> void { + auto bound_val = self.create( + self.getBuilder().getI64Type(), + static_cast(indexBoundary)); + auto dim_val = self.create( + self.getBuilder().getI32Type(), static_cast(dim)); + self.create( + ptr, value, index, bound_val, dim_val, dstStride, endOffset, + startOffset); + }) + .def("create_sort", + [](TritonOpBuilder &self, Value src, int64_t dim, + bool descending) -> Value { + auto &builder = self.getBuilder(); + auto op = builder.create( + self.getLastLoc(), src, builder.getI64IntegerAttr(dim), + builder.getBoolAttr(descending)); + return op->getResult(0); + }) + .def("create_flip", + [](TritonOpBuilder &self, Value src, int64_t dim) -> Value { + auto op = self.getBuilder().create( + self.getLastLoc(), src, + self.getBuilder().getI64IntegerAttr(dim)); + return op->getResult(0); + }); + + // --- buffer_builder class --- + struct BufferOpBuilder : public TritonOpBuilder {}; + + m.def("load_buffer_dialects", [](MLIRContext &context) { + DialectRegistry registry; + registry + .insert(); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + }); + + py::class_( + m, "buffer_builder", py::module_local(), py::dynamic_attr()) + .def(py::init()) + .def("get_null_attr", + [](BufferOpBuilder &self) -> Attribute { return Attribute(); }) + .def("get_str_array_attr", + [](BufferOpBuilder &self, + const std::vector &array) -> ArrayAttr { + auto strRefVec = to_vector(llvm::map_range( + array, [](const auto &s) { return llvm::StringRef(s); })); + return self.getBuilder().getStrArrayAttr( + llvm::ArrayRef{strRefVec}); + }) + .def("alloc", + [](BufferOpBuilder &self, Type memrefType) -> Value { + return self.create( + mlir::cast(memrefType)); + }) + .def("to_buffer", + [](BufferOpBuilder &self, Value &src, + const Attribute &addressSpace) -> Value { + auto tensorType = dyn_cast(src.getType()); + if (!tensorType) { + llvm::report_fatal_error("to_buffer: src must be tensor type"); + } + auto memrefType = MemRefType::get(tensorType.getShape(), + tensorType.getElementType(), + MemRefLayoutAttrInterface{}); + Operation *memref = + self.create(memrefType, src); + if (addressSpace) { + memref = self.create( + MemRefType::get(memrefType.getShape(), + memrefType.getElementType(), + memrefType.getLayout(), addressSpace), + memref->getResult(0)); + } + return memref->getResult(0); + }) + .def("to_tensor", + [](BufferOpBuilder &self, Value &src, bool writable) -> Value { + const auto &memrefType = mlir::cast(src.getType()); + auto tensorType = mlir::RankedTensorType::get( + memrefType.getShape(), memrefType.getElementType()); + auto hasAddressSpace = memrefType.getMemorySpace(); + if (hasAddressSpace) { + MemRefType targetType = MemRefType::get( + memrefType.getShape(), memrefType.getElementType(), + memrefType.getLayout()); + return self.create( + tensorType, + self.create(targetType, src), + mlir::UnitAttr::get(self.getContext()), + writable ? mlir::UnitAttr::get(self.getContext()) + : nullptr); + } + return self.create( + tensorType, src, mlir::UnitAttr::get(self.getContext()), + writable ? mlir::UnitAttr::get(self.getContext()) : nullptr); + }) + .def("subview", + [](BufferOpBuilder &self, Value source, + std::vector &offsets, const std::vector &sizes, + const std::vector &strides) -> Value { + SmallVector mixedOffsets; + auto *context = self.getBuilder().getContext(); + auto &builder = self.getBuilder(); + auto sourceType = mlir::cast(source.getType()); + int64_t rank = sourceType.getRank(); + if (offsets.size() != rank || sizes.size() != rank || + strides.size() != rank) { + throw std::runtime_error( + "Number of offsets, sizes, and strides " + "must match memref rank"); + } + for (const auto &offset : offsets) { + auto indexType = builder.getIndexType(); + if (offset.getType() != indexType) { + Value offset_val = + self.create(indexType, offset); + mixedOffsets.push_back(offset_val); + } else { + mixedOffsets.push_back(offset); + } + } + constexpr unsigned kIntegerAttrBitWidth = 64; + SmallVector mixedSizes; + SmallVector mixedStrides; + for (int64_t i = 0; i < rank; ++i) { + int64_t size = sizes[i]; + int64_t stride = strides[i]; + int64_t srcDim = sourceType.getDimSize(i); + if (size <= 0) { + throw std::runtime_error("Expected sizes to be positive"); + } + if (stride <= 0) { + throw std::runtime_error("Expected strides to be positive"); + } + if (!ShapedType::isDynamic(srcDim)) { + if (size > srcDim) { + throw std::runtime_error( + "Subview size cannot exceed source dimension size"); + } + if (stride > srcDim) { + throw std::runtime_error( + "Stride cannot exceed source dimension size"); + } + } + mixedSizes.push_back(IntegerAttr::get( + IntegerType::get(context, kIntegerAttrBitWidth), size)); + mixedStrides.push_back(IntegerAttr::get( + IntegerType::get(context, kIntegerAttrBitWidth), stride)); + } + return self.create(source, mixedOffsets, + mixedSizes, mixedStrides); + }); + } +} + +// ============================================================================= +// Pass pipeline bindings (from triton_dicp_triton.cc original) +// ============================================================================= + +void init_triton_dicp_passes_commonir(py::module &&m) { + m.def("add_vectorize_parallel_loop", [](mlir::PassManager &pm) { pm.addNestedPass( - dicp::LinalgExt::createScalarTo1DTensorPass()); + mlir::dicp::CommonIR::createVectorizeParallelLoopPass()); }); - m.def("add_linalg_to_linked", [](mlir::PassManager &pm, bool globalKernel, - bool namedOps, bool cpuVerify) { - pm.addPass(mlir::dicp::linked::createLinalgToLinkedPass( - globalKernel, namedOps, cpuVerify)); + m.def("add_annotate_kernel_attrs", [](mlir::PassManager &pm) { + pm.addPass(mlir::dicp::CommonIR::createAnnotateKernelAttrsPass()); }); - ADD_PASS_WRAPPER_0("add_linked_to_hivm", - dicp::linked::createLinkedToHIVMPass); - ADD_PASS_WRAPPER_0("add_debug_cpu_verify", - dicp::linked::createDebugCPUVerifyPass); - m.def("add_vectorize_parallel_loop", [](mlir::PassManager &pm) { - pm.addNestedPass( - dicp::LinalgExt::createVectorizeParallelLoopPass()); +} + +void init_triton_dicp_passes_ttir(py::module &&m) { + m.def("add_auto_blockify", [](mlir::PassManager &pm, int autoBlockifySize) { + AutoBlockifyOptions opts; + opts.autoBlockifySize = autoBlockifySize; + pm.addPass(mlir::triton::createAutoBlockifyPass(opts)); + }); + + m.def("add_ascend_legalize", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createAscendLegalizePass()); + }); + + m.def("add_triton_to_structure", + [](mlir::PassManager &pm, bool enableMaskFallbackConversion, + bool optimizeDynamicOffset) { + pm.addPass(mlir::triton::createTritonToStructuredPass( + enableMaskFallbackConversion, optimizeDynamicOffset)); + }); + + m.def("add_discrete_mask_access_conversion", [](mlir::PassManager &pm, + bool compileOn91095, + bool forceSimtTemplate, + bool enableSyncBlockLock) { + DiscreteMaskAccessConversionOptions opts; + opts.compileOn91095 = compileOn91095; + opts.forceSimtTemplate = forceSimtTemplate; + opts.enableSyncBlockLock = enableSyncBlockLock; + pm.addPass(mlir::triton::createDiscreteMaskAccessConversionPass(opts)); + }); + + m.def("add_triton_to_annotation", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createTritonToAnnotationPass()); + }); + + m.def("add_triton_to_unstructure", + [](mlir::PassManager &pm, bool compileOn91095, bool forceSimtTemplate) { + TritonToUnstructureOptions opts; + opts.compileOn91095 = compileOn91095; + opts.forceSimtTemplate = forceSimtTemplate; + pm.addPass(mlir::triton::createTritonToUnstructurePass(opts)); + }); + + m.def("add_triton_to_hivm", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createTritonToHIVMPass()); + }); + + m.def("add_triton_to_hfusion", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createTritonToHFusionPass()); + }); + + m.def("add_triton_to_llvm", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createTritonToLLVMPass()); + }); + + m.def("add_bubble_up_operation", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createBubbleUpOperationPass()); + }); + + m.def("add_triton_to_linalg", + [](mlir::PassManager &pm, bool globalKernel, bool namedOps, + bool enableNd2nzOnVector, bool enableSelectAnalysis, + bool compileOn91095) { + pm.addPass(mlir::triton::createTritonToLinalgPass( + globalKernel, namedOps, enableNd2nzOnVector, enableSelectAnalysis, + compileOn91095)); + }); + + m.def("add_ascend_npu_ir_legalize", + [](mlir::PassManager &pm, bool unsafeMode) { + AscendNPUIRLegalizeOptions opts; + opts.unsafeMode = unsafeMode; + pm.addPass(mlir::triton::createAscendNPUIRLegalizePass(opts)); + }); + + m.def("add_dag_sync", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createDAGSyncPass()); + }); + + m.def("add_dag_scope", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createDAGScopePass()); + }); + + m.def("add_dag_ssbuffer", [](mlir::PassManager &pm) { + pm.addPass(mlir::triton::createDAGSSBufferPass()); }); } +// ============================================================================= +// Top-level init: init_triton_dicp_triton +// ============================================================================= + void init_triton_dicp_triton(py::module &&m) { - m.doc() = "Python bindings to the Deeplink Triton backend"; + m.doc() = "Python bindings to the DICP Triton backend (Ascend NPU)"; + auto passes = m.def_submodule("passes"); - init_triton_dicp_triton_pass_triton_shared_ascend( - passes.def_submodule("triton_shared_ascend")); - init_triton_dicp_triton_pass_linked_npu(passes.def_submodule("linked_npu")); + init_triton_dicp_passes_commonir(passes.def_submodule("commonir")); + init_triton_dicp_passes_ttir(passes.def_submodule("ttir")); + + // DICP NPU IR builder, affine types, hivm enums + init_dicp_ir(m.def_submodule("ir")); - // load dialects + // Load core dialects m.def("load_dialects", [](MLIRContext &context) { DialectRegistry registry; - registry.insert(); - - dicp::trtion_ext::registerCanonicalizeTritonIRAscendPass(); - dicp::trtion_ext::registerCanonicalizeCmpiPass(); - - dicp::linked::registerLinalgToLinkedPass(); - dicp::linked::registerLinkedToHIVMPass(); - dicp::linked::registerTritonToLinalgNPUCoversionPass(); - - dicp::LinalgExt::registerLinalgIfToSelectPass(); - dicp::LinalgExt::registerLinalgGenericToSCFPass(); - dicp::LinalgExt::registerScalarTo1DTensorPass(); - dicp::LinalgExt::registerNormalizeSliceOpsPass(); - dicp::LinalgExt::registerVectorizeParallelLoopPass(); - + registry + .insert(); + mlir::func::registerInlinerExtension(registry); context.appendDialectRegistry(registry); context.loadAllAvailableDialects(); }); -} \ No newline at end of file +}