From 9241e29683d865aecfed61a6bfbc1dc9f361064e Mon Sep 17 00:00:00 2001 From: lixingda Date: Fri, 26 Jun 2026 11:50:17 +0800 Subject: [PATCH 1/6] [BACKEND] add sunrise backend with tle. --- .github/workflows/sunrise-build-and-test.yml | 61 + .gitignore | 3 + CMakeLists.txt | 99 +- .../Dialect/TritonGPU/IR/CMakeLists.txt | 2 +- .../TritonToTritonGPU/CMakeLists.txt | 3 + .../TritonToTritonGPUPass.cpp | 8 + lib/Dialect/Triton/Transforms/CMakeLists.txt | 3 + .../Transforms/RewriteTensorPointer.cpp | 8 + lib/Dialect/TritonGPU/IR/CMakeLists.txt | 3 + lib/Dialect/TritonGPU/IR/Dialect.cpp | 8 + .../TritonGPU/IR/LinearLayoutConversions.cpp | 8 + .../TritonGPU/Transforms/CMakeLists.txt | 3 + lib/Dialect/TritonGPU/Transforms/Prefetch.cpp | 7 + .../Transforms/RemoveLayoutConversions.cpp | 8 + lib/Dialect/TritonGPU/Transforms/Utility.cpp | 8 + python/setup_tools/setup_helper.py | 82 +- python/setup_tools/utils/tools.py | 1 + python/triton/compiler/hint_manager.py | 12 +- python/triton/runtime/adjust_kernel_param.py | 16 + setup.py | 6 +- third_party/sunrise/CMakeLists.txt | 44 + third_party/sunrise/backend/__init__.py | 0 third_party/sunrise/backend/compiler.py | 445 ++ third_party/sunrise/backend/driver.c | 172 + third_party/sunrise/backend/driver.py | 564 +++ third_party/sunrise/backend/include/ptml.h | 925 ++++ third_party/sunrise/backend/include/tang.h | 2321 +++++++++ .../backend/include/tang_compiler_api.h | 223 + .../backend/include/tang_rt/device_assert.h | 43 + .../include/tang_rt/device_functions.h | 8 + .../backend/include/tang_rt/driver_types.h | 1281 +++++ .../sunrise/backend/include/tang_rt/fmt.hpp | 1097 +++++ .../backend/include/tang_rt/host_defines.h | 101 + .../backend/include/tang_rt/vector_types.h | 35 + .../sunrise/backend/include/tang_rt/version.h | 30 + .../sunrise/backend/include/tang_runtime.h | 32 + .../backend/include/tang_runtime_api.h | 1871 ++++++++ .../sunrise/backend/include/tapti/tapti.h | 11 + .../backend/include/tapti/tapti_activity.h | 956 ++++ .../backend/include/tapti/tapti_callbacks.h | 109 + .../backend/include/tapti/tapti_driver_cbid.h | 203 + .../backend/include/tapti/tapti_result.h | 248 + .../include/tapti/tapti_runtime_cbid.h | 147 + .../backend/include/tapti/tapti_version.h | 39 + third_party/sunrise/backend/lib/.empty | 0 third_party/sunrise/backend/spec/__init__.py | 0 .../backend/spec/include/flagtree_spec.h | 12 + .../sunrise_TritonToTritonGPUPass.h | 6 + .../Transforms/sunrise_RewriteTensorPointer.h | 6 + .../Dialect/TritonGPU/IR/TritonGPUAttrDefs.td | 1533 ++++++ .../Dialect/TritonGPU/IR/sunrise_Dialect.h | 6 + .../IR/sunrise_LinearLayoutConversions.h | 6 + .../TritonGPU/Transforms/sunrise_Prefetch.h | 6 + .../sunrise_RemoveLayoutConversions.h | 6 + .../TritonGPU/Transforms/sunrise_Utility.h | 6 + .../sunrise/backend/spec/lib/CMakeLists.txt | 2 + .../spec/lib/Conversion/CMakeLists.txt | 1 + .../TritonToTritonGPU/CMakeLists.txt | 14 + .../TritonToTritonGPUPass.cpp | 1043 ++++ .../backend/spec/lib/Dialect/CMakeLists.txt | 2 + .../spec/lib/Dialect/Triton/CMakeLists.txt | 1 + .../Dialect/Triton/Transforms/CMakeLists.txt | 13 + .../Transforms/RewriteTensorPointer.cpp | 619 +++ .../spec/lib/Dialect/TritonGPU/CMakeLists.txt | 2 + .../lib/Dialect/TritonGPU/IR/CMakeLists.txt | 14 + .../spec/lib/Dialect/TritonGPU/IR/Dialect.cpp | 4182 +++++++++++++++++ .../TritonGPU/IR/LinearLayoutConversions.cpp | 1827 +++++++ .../TritonGPU/Transforms/CMakeLists.txt | 9 + .../Dialect/TritonGPU/Transforms/Prefetch.cpp | 459 ++ .../Transforms/RemoveLayoutConversions.cpp | 1894 ++++++++ .../Dialect/TritonGPU/Transforms/Utility.cpp | 1777 +++++++ .../sunrise/backend/sunrise_hint_handler.py | 54 + third_party/sunrise/language/tang/__init__.py | 3 + .../sunrise/language/tang/libdevice.py | 203 + third_party/sunrise/plugin/.empty | 0 third_party/sunrise/python/src/gluon_ir.cc | 984 ++++ third_party/sunrise/python/src/interpreter.cc | 740 +++ third_party/sunrise/python/src/ir.cc | 2110 +++++++++ third_party/sunrise/python/src/ir.h | 113 + .../sunrise/python/src/linear_layout.cc | 223 + third_party/sunrise/python/src/llvm.cc | 826 ++++ third_party/sunrise/python/src/main.cc | 62 + third_party/sunrise/python/src/passes.cc | 135 + third_party/sunrise/python/src/passes.h | 43 + third_party/sunrise/python/src/specialize.cc | 584 +++ .../sunrise/python/test/01-vector-add.py | 50 + .../sunrise/python/test/02-tle-local_ptr.py | 110 + .../python/test/03-tle-extract_insert_tile.py | 261 + .../python/test/04-tle-hint-shared-memory.py | 55 + .../python/test/05-tle-copy-normcopy.py | 144 + .../sunrise/python/test/06-tle-cumsum.py | 82 + .../sunrise/python/test/07-tle-load.py | 78 + .../python/test/08-tle-dot-local-ptr.py | 94 + .../python/test/09-tle-negative-contract.py | 113 + .../TlePromoteLocalStoreStaging.cpp | 5 +- 95 files changed, 31756 insertions(+), 16 deletions(-) create mode 100644 .github/workflows/sunrise-build-and-test.yml create mode 100644 third_party/sunrise/CMakeLists.txt create mode 100644 third_party/sunrise/backend/__init__.py create mode 100644 third_party/sunrise/backend/compiler.py create mode 100644 third_party/sunrise/backend/driver.c create mode 100644 third_party/sunrise/backend/driver.py create mode 100755 third_party/sunrise/backend/include/ptml.h create mode 100755 third_party/sunrise/backend/include/tang.h create mode 100755 third_party/sunrise/backend/include/tang_compiler_api.h create mode 100755 third_party/sunrise/backend/include/tang_rt/device_assert.h create mode 100755 third_party/sunrise/backend/include/tang_rt/device_functions.h create mode 100755 third_party/sunrise/backend/include/tang_rt/driver_types.h create mode 100755 third_party/sunrise/backend/include/tang_rt/fmt.hpp create mode 100755 third_party/sunrise/backend/include/tang_rt/host_defines.h create mode 100755 third_party/sunrise/backend/include/tang_rt/vector_types.h create mode 100755 third_party/sunrise/backend/include/tang_rt/version.h create mode 100755 third_party/sunrise/backend/include/tang_runtime.h create mode 100755 third_party/sunrise/backend/include/tang_runtime_api.h create mode 100755 third_party/sunrise/backend/include/tapti/tapti.h create mode 100755 third_party/sunrise/backend/include/tapti/tapti_activity.h create mode 100755 third_party/sunrise/backend/include/tapti/tapti_callbacks.h create mode 100755 third_party/sunrise/backend/include/tapti/tapti_driver_cbid.h create mode 100755 third_party/sunrise/backend/include/tapti/tapti_result.h create mode 100755 third_party/sunrise/backend/include/tapti/tapti_runtime_cbid.h create mode 100755 third_party/sunrise/backend/include/tapti/tapti_version.h create mode 100644 third_party/sunrise/backend/lib/.empty create mode 100644 third_party/sunrise/backend/spec/__init__.py create mode 100644 third_party/sunrise/backend/spec/include/flagtree_spec.h create mode 100644 third_party/sunrise/backend/spec/include/triton/Conversion/TritonToTritonGPU/sunrise_TritonToTritonGPUPass.h create mode 100644 third_party/sunrise/backend/spec/include/triton/Dialect/Triton/Transforms/sunrise_RewriteTensorPointer.h create mode 100644 third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td create mode 100644 third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_Dialect.h create mode 100644 third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_LinearLayoutConversions.h create mode 100644 third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Prefetch.h create mode 100644 third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_RemoveLayoutConversions.h create mode 100644 third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Utility.h create mode 100644 third_party/sunrise/backend/spec/lib/CMakeLists.txt create mode 100644 third_party/sunrise/backend/spec/lib/Conversion/CMakeLists.txt create mode 100644 third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/CMakeLists.txt create mode 100644 third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/CMakeLists.txt create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/Triton/CMakeLists.txt create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/CMakeLists.txt create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/CMakeLists.txt create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/CMakeLists.txt create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp create mode 100644 third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Utility.cpp create mode 100644 third_party/sunrise/backend/sunrise_hint_handler.py create mode 100644 third_party/sunrise/language/tang/__init__.py create mode 100644 third_party/sunrise/language/tang/libdevice.py create mode 100644 third_party/sunrise/plugin/.empty create mode 100644 third_party/sunrise/python/src/gluon_ir.cc create mode 100644 third_party/sunrise/python/src/interpreter.cc create mode 100644 third_party/sunrise/python/src/ir.cc create mode 100644 third_party/sunrise/python/src/ir.h create mode 100644 third_party/sunrise/python/src/linear_layout.cc create mode 100644 third_party/sunrise/python/src/llvm.cc create mode 100644 third_party/sunrise/python/src/main.cc create mode 100644 third_party/sunrise/python/src/passes.cc create mode 100644 third_party/sunrise/python/src/passes.h create mode 100644 third_party/sunrise/python/src/specialize.cc create mode 100644 third_party/sunrise/python/test/01-vector-add.py create mode 100644 third_party/sunrise/python/test/02-tle-local_ptr.py create mode 100644 third_party/sunrise/python/test/03-tle-extract_insert_tile.py create mode 100644 third_party/sunrise/python/test/04-tle-hint-shared-memory.py create mode 100644 third_party/sunrise/python/test/05-tle-copy-normcopy.py create mode 100644 third_party/sunrise/python/test/06-tle-cumsum.py create mode 100644 third_party/sunrise/python/test/07-tle-load.py create mode 100644 third_party/sunrise/python/test/08-tle-dot-local-ptr.py create mode 100644 third_party/sunrise/python/test/09-tle-negative-contract.py diff --git a/.github/workflows/sunrise-build-and-test.yml b/.github/workflows/sunrise-build-and-test.yml new file mode 100644 index 0000000000..6aded30788 --- /dev/null +++ b/.github/workflows/sunrise-build-and-test.yml @@ -0,0 +1,61 @@ +name: Sunrise-Build-And-Test + +on: + push: + branches: [ "triton_v3.6.x" ] + pull_request: + branches: [ "triton_v3.6.x" ] + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + sunrise-build-and-test: + runs-on: flagtree-sunrise + if: ${{ github.repository == 'FlagTree/flagtree' || github.repository == 'flagos-ai/flagtree' }} + steps: + - name: Setup environment + shell: bash + run: | + # source ~/env.sh + env | grep -E '^(http_proxy|https_proxy|all_proxy|no_proxy)=' >> $GITHUB_ENV || true + + - name: Smart Checkout + uses: flagos-ai/FlagTree/.github/actions/smart-checkout@main + with: + checkout_version: 'v6' + + - name: Check if backend-relevant files changed + id: check_backend + uses: flagos-ai/FlagTree/.github/actions/check-backend-changed@main + with: + backend: sunrise + + - name: FlagTree Build on Sunrise + if: steps.check_backend.outputs.should_skip != 'true' + shell: bash + run: | + set -x + pip uninstall -y triton + export TRITON_BUILD_WITH_CLANG_LLD=1 + export TRITON_OFFLINE_BUILD=1 + export TRITON_BUILD_PROTON=OFF + export FLAGTREE_BACKEND=sunrise + MAX_JOBS=32 python3 -m pip install . --no-build-isolation + + - name: FlagTree Test on Sunrise + if: steps.check_backend.outputs.should_skip != 'true' + shell: bash + run: | + set -x + export LD_LIBRARY_PATH=/usr/local/tangrt/lib/linux-x86_64/stub:${LD_LIBRARY_PATH} + python3 -m pytest third_party/sunrise/python/test/01-vector-add.py + python3 -m pytest third_party/sunrise/python/test/02-tle-local_ptr.py + python3 -m pytest third_party/sunrise/python/test/03-tle-extract_insert_tile.py + python3 -m pytest third_party/sunrise/python/test/04-tle-hint-shared-memory.py + python3 -m pytest third_party/sunrise/python/test/05-tle-copy-normcopy.py + python3 -m pytest third_party/sunrise/python/test/06-tle-cumsum.py + python3 -m pytest third_party/sunrise/python/test/07-tle-load.py + python3 -m pytest third_party/sunrise/python/test/08-tle-dot-local-ptr.py + python3 -m pytest third_party/sunrise/python/test/09-tle-negative-contract.py \ No newline at end of file diff --git a/.gitignore b/.gitignore index c40bfc2436..772b539c3c 100644 --- a/.gitignore +++ b/.gitignore @@ -87,6 +87,9 @@ third_party/nvidia/backend/include third_party/nvidia/backend/lib/cupti third_party/tle/third_party +# third_party sunrise +third_party/sunrise/backend/lib/*.bc + # Docs docs/_build/ docs/python-api/generated/ diff --git a/CMakeLists.txt b/CMakeLists.txt index 35ce33293a..7fc372dfdc 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -10,6 +10,23 @@ set(CMAKE_CXX_STANDARD 17) set(CMAKE_INCLUDE_CURRENT_DIR ON) +project(triton CXX C) +include(CTest) + +list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") + +# Options +option(TRITON_BUILD_PYTHON_MODULE "Build Python Triton bindings" OFF) +if(FLAGTREE_BACKEND) + option(TRITON_BUILD_PROTON "Build the Triton Proton profiler" OFF) + option(TRITON_BUILD_UT "Build C++ Triton Unit Tests" OFF) +else() + option(TRITON_BUILD_PROTON "Build the Triton Proton profiler" ON) + option(TRITON_BUILD_UT "Build C++ Triton Unit Tests" ON) +endif() +option(TRITON_BUILD_WITH_CCACHE "Build with ccache (if available)" ON) +set(TRITON_CODEGEN_BACKENDS "" CACHE STRING "Enable different codegen backends") + # FLAGTREE Options set(FLAGTREE_BACKEND "$ENV{FLAGTREE_BACKEND}") if(FLAGTREE_BACKEND) @@ -50,14 +67,22 @@ elseif(FLAGTREE_BACKEND STREQUAL "metax") set(FLAGTREE_TLE OFF) remove_definitions(-D__TLE__) list(REMOVE_ITEM LLVM_TABLEGEN_FLAGS -D__TLE__) +elseif(FLAGTREE_BACKEND STREQUAL "sunrise") + find_package(Python3 3.10 REQUIRED COMPONENTS Development.Module Interpreter) endif() set(FLAGTREE_PLUGIN "$ENV{FLAGTREE_PLUGIN}") if(FLAGTREE_PLUGIN) add_definitions(-D__FLAGTREE_PLUGIN__) endif() -project(triton CXX C) -include(CTest) +# FLAGTREE SPEC LIB GET FUNC +function(get_flagtree_backend_lib lib_name output_lib) + set(ret FlagTree_${FLAGTREE_BACKEND}_${lib_name}) + if(NOT TARGET ${ret}) + set(ret "") + endif() + set(${output_lib} ${ret} PARENT_SCOPE) +endfunction() list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") @@ -77,6 +102,18 @@ if(FLAGTREE_BACKEND AND FLAGTREE_BACKEND STREQUAL "metax") list(APPEND TRITON_CODEGEN_BACKENDS "amd") endif() +# FLAGTREE SPEC TD FILE GET FUNC +function(set_flagtree_backend_td output_td td_filename) + set(ret ${td_filename}) + file(RELATIVE_PATH relative_path "${PROJECT_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}") + get_filename_component(BACKEND_SPEC_ROOT "${BACKEND_SPEC_INCLUDE_DIR}" DIRECTORY) + set(BACKEND_SPEC_TD ${BACKEND_SPEC_ROOT}/${relative_path}/${td_filename}) + if(EXISTS ${BACKEND_SPEC_TD}) + set(ret ${BACKEND_SPEC_TD}) + endif() + set(${output_td} ${ret} PARENT_SCOPE) +endfunction() + if(TRITON_BUILD_WITH_CCACHE) find_program(CCACHE_PROGRAM ccache) if(CCACHE_PROGRAM) @@ -131,7 +168,14 @@ if(TRITON_BUILD_UT) endif() # Compiler flags -set(BACKEND_INCLUDE_DIR ${CMAKE_CURRENT_SOURCE_DIR}/third_party/${FLAGTREE_BACKEND}/include) +set(FLAGTREE_BACKEND_DIR ${PROJECT_SOURCE_DIR}/third_party/${FLAGTREE_BACKEND}) +## flagtree spec include dir +set(BACKEND_SPEC_INCLUDE_DIR ${FLAGTREE_BACKEND_DIR}/backend/spec/include) +if(FLAGTREE_BACKEND AND EXISTS ${BACKEND_SPEC_INCLUDE_DIR}) + include_directories(${BACKEND_SPEC_INCLUDE_DIR}) +endif() +## flagtree third_party include dir +set(BACKEND_INCLUDE_DIR ${FLAGTREE_BACKEND_DIR}/include) if(FLAGTREE_BACKEND AND EXISTS "${BACKEND_INCLUDE_DIR}") include_directories(${BACKEND_INCLUDE_DIR}) else() @@ -464,6 +508,49 @@ if(TRITON_BUILD_PYTHON_MODULE) elseif(FLAGTREE_BACKEND STREQUAL "hcu") list(APPEND TRITON_PLUGIN_NAMES "distributed") add_subdirectory(test) + elseif(FLAGTREE_BACKEND STREQUAL "sunrise") + set(TRITON_LIBRARIES + ${triton_libs} + ${triton_plugins} + # mlir + # MLIRAMDGPUDialect + # MLIRNVVMDialect + MLIRSTVMDialect # STVM + MLIRNVVMToLLVMIRTranslation + MLIRSTVMToLLVMIRTranslation + MLIRGPUToNVVMTransforms + MLIRGPUToSTVMTransforms + MLIRGPUToGPURuntimeTransforms + MLIRGPUTransforms + MLIRIR + MLIRControlFlowToLLVM + MLIRBytecodeWriter + MLIRPass + MLIRTransforms + MLIRLLVMDialect + MLIRSupport + MLIRTargetLLVMIRExport + MLIRMathToLLVM + # MLIRROCDLToLLVMIRTranslation + MLIRGPUDialect + MLIRSCFToControlFlow + MLIRIndexToLLVM + MLIRGPUToROCDLTransforms + MLIRUBToLLVM + # LLVM + LLVMPasses + # LLVMNVPTXCodeGen + # LLVMAMDGPUCodeGen + # LLVMAMDGPUAsmParser + LLVMSTCUCodeGen + LLVMSTCUAsmParser + LLVMAArch64CodeGen + LLVMAArch64AsmParser + LLVMRISCVCodeGen + LLVMRISCVAsmParser + Python3::Module + pybind11::headers + ) endif() if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64" OR # Linux arm64 @@ -508,7 +595,7 @@ if(TRITON_BUILD_PYTHON_MODULE) ${PYTHON_SRC_PATH}/specialize.cc) # Link triton with its dependencies - target_link_libraries(triton PRIVATE ${TRITON_LIBRARIES}) + target_link_libraries(triton PUBLIC ${TRITON_LIBRARIES}) if(WIN32) target_link_libraries(triton PRIVATE ${CMAKE_DL_LIBS}) set_target_properties(triton PROPERTIES SUFFIX ".pyd") @@ -540,6 +627,8 @@ if (UNIX AND NOT APPLE) get_filename_component(_flagtree_llvm_archive_name "${_flagtree_llvm_archive}" NAME) set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,--exclude-libs,${_flagtree_llvm_archive_name}") endforeach() + elseif(FLAGTREE_BACKEND STREQUAL "sunrise") + set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,--export-dynamic") else() set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_SHARED_LINKER_FLAGS} -Wl,--exclude-libs,ALL") endif() @@ -572,7 +661,7 @@ find_package(Threads REQUIRED) add_subdirectory(third_party/f2reduce) -if(NOT FLAGTREE_BACKEND OR FLAGTREE_BACKEND MATCHES "^(aipu|tsingmicro|enflame|rpu|thrive|metax)$") +if(NOT FLAGTREE_BACKEND OR FLAGTREE_BACKEND MATCHES "^(aipu|tsingmicro|enflame|rpu|thrive|metax|sunrise)$") add_subdirectory(bin) if(FLAGTREE_TLE) flagtree_add_tle_generated_header_dependencies() diff --git a/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt b/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt index 436bbdc830..10b1b1b05b 100644 --- a/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt +++ b/include/triton/Dialect/TritonGPU/IR/CMakeLists.txt @@ -11,7 +11,7 @@ add_mlir_doc(TritonGPUDialect TritonGPUDialect dialects/ -gen-dialect-doc) add_mlir_doc(TritonGPUOps TritonGPUOps dialects/ -gen-op-doc) add_public_tablegen_target(TritonGPUTableGen) -set(LLVM_TARGET_DEFINITIONS TritonGPUAttrDefs.td) +set_flagtree_backend_td(LLVM_TARGET_DEFINITIONS TritonGPUAttrDefs.td) mlir_tablegen(AttrInterfaces.h.inc -gen-attr-interface-decls) mlir_tablegen(AttrInterfaces.cpp.inc -gen-attr-interface-defs) mlir_tablegen(AttrDefs.h.inc -gen-attrdef-decls) diff --git a/lib/Conversion/TritonToTritonGPU/CMakeLists.txt b/lib/Conversion/TritonToTritonGPU/CMakeLists.txt index 3770e406d5..dedd4baf08 100644 --- a/lib/Conversion/TritonToTritonGPU/CMakeLists.txt +++ b/lib/Conversion/TritonToTritonGPU/CMakeLists.txt @@ -4,6 +4,8 @@ else() set(_TLE_IR_LIBS "") endif() +get_flagtree_backend_lib("TritonToTritonGPU" _EXTRA_LINK_LIBS) + add_triton_library(TritonToTritonGPU RelayoutTritonGPU.cpp TritonGPUConversion.cpp @@ -20,4 +22,5 @@ add_triton_library(TritonToTritonGPU ProtonIR TritonGPUIR ${_TLE_IR_LIBS} + ${_EXTRA_LINK_LIBS} ) diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp index bf509ffe52..defd11073a 100644 --- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp +++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp @@ -1,3 +1,9 @@ +#if __has_include("flagtree_spec.h") +#include "flagtree_spec.h" +#endif + +#ifndef FLAGTREE_SPEC_triton_Conversion_TritonToTritonGPU_sunrise_TritonToTritonGPUPass + #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" #include "mlir/Dialect/UB/IR/UBOps.h" @@ -1002,3 +1008,5 @@ class ConvertTritonToTritonGPU }; } // namespace + +#endif diff --git a/lib/Dialect/Triton/Transforms/CMakeLists.txt b/lib/Dialect/Triton/Transforms/CMakeLists.txt index 8be846f589..c39797d88b 100644 --- a/lib/Dialect/Triton/Transforms/CMakeLists.txt +++ b/lib/Dialect/Triton/Transforms/CMakeLists.txt @@ -2,6 +2,8 @@ set(LLVM_TARGET_DEFINITIONS Combine.td) mlir_tablegen(TritonCombine.inc -gen-rewriters) add_public_tablegen_target(TritonCombineIncGen) +get_flagtree_backend_lib("TritonTransforms" _EXTRA_LINK_LIBS) + add_triton_library(TritonTransforms Combine.cpp LoopAwareCSE.cpp @@ -24,4 +26,5 @@ add_triton_library(TritonTransforms MLIRTransforms MLIRSCFToControlFlow TritonIR + ${_EXTRA_LINK_LIBS} ) diff --git a/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp b/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp index 24a66d3a7d..3a745727ed 100644 --- a/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp +++ b/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp @@ -1,3 +1,9 @@ +#if __has_include("flagtree_spec.h") +#include "flagtree_spec.h" +#endif + +#ifndef FLAGTREE_SPEC_triton_Dialect_Triton_Transforms_sunrise_RewriteTensorPointer + #include #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" @@ -565,3 +571,5 @@ class RewriteTensorPointerPass }; } // namespace mlir::triton + +#endif diff --git a/lib/Dialect/TritonGPU/IR/CMakeLists.txt b/lib/Dialect/TritonGPU/IR/CMakeLists.txt index af8d918502..d9dd096e4a 100644 --- a/lib/Dialect/TritonGPU/IR/CMakeLists.txt +++ b/lib/Dialect/TritonGPU/IR/CMakeLists.txt @@ -1,3 +1,5 @@ +get_flagtree_backend_lib("TritonGPUIR" _EXTRA_LINK_LIBS) + add_triton_library(TritonGPUIR Dialect.cpp LinearLayoutConversions.cpp @@ -15,4 +17,5 @@ add_triton_library(TritonGPUIR MLIRGPUDialect TritonIR TritonTools + ${_EXTRA_LINK_LIBS} ) diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp index e7fe5b7e2c..b8c6e2716c 100644 --- a/lib/Dialect/TritonGPU/IR/Dialect.cpp +++ b/lib/Dialect/TritonGPU/IR/Dialect.cpp @@ -1,3 +1,9 @@ +#if __has_include("flagtree_spec.h") +#include "flagtree_spec.h" +#endif + +#ifndef FLAGTREE_SPEC_triton_Dialect_TritonGPU_IR_sunrise_Dialect + #include "triton/Dialect/Triton/IR/Dialect.h" #include @@ -4011,3 +4017,5 @@ std::optional triton::gpu::getWarpSpecializeTag(Operation *op) { } return std::nullopt; } + +#endif diff --git a/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp b/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp index 6f8a5aceee..7e45fe9125 100644 --- a/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp +++ b/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp @@ -1,3 +1,9 @@ +#if __has_include("flagtree_spec.h") +#include "flagtree_spec.h" +#endif + +#ifndef FLAGTREE_SPEC_triton_Dialect_TritonGPU_IR_sunrise_LinearLayoutConversion + #include #include "triton/Dialect/Triton/IR/Utility.h" @@ -1689,3 +1695,5 @@ chooseMfmaLikeStoreLayout(RankedTensorType valType) { } } // namespace mlir::triton::gpu + +#endif diff --git a/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt b/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt index 287e17864f..6cdd688943 100644 --- a/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt +++ b/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt @@ -6,6 +6,8 @@ if(FLAGTREE_TLE) ) endif() +get_flagtree_backend_lib("TritonGPUTransforms" _EXTRA_LINK_LIBS) + add_triton_library(TritonGPUTransforms AccelerateMatmul.cpp ProcessSharedMemoryHint.cpp @@ -62,4 +64,5 @@ add_triton_library(TritonGPUTransforms TritonToTritonGPU TritonInstrumentIR MLIRTransformUtils + ${_EXTRA_LINK_LIBS} ) diff --git a/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp b/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp index 9d8e42eb77..cff02296f6 100644 --- a/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp +++ b/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp @@ -25,6 +25,11 @@ // scf.yield %next_a, ..., %a_prefetch_next // } //===----------------------------------------------------------------------===// +#if __has_include("flagtree_spec.h") +#include "flagtree_spec.h" +#endif + +#ifndef FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_Prefetch #include "mlir/IR/IRMapping.h" #include "mlir/Support/LLVM.h" @@ -455,3 +460,5 @@ struct PrefetchPass : public impl::TritonGPUPrefetchBase { } // namespace gpu } // namespace triton } // namespace mlir + +#endif diff --git a/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp b/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp index 7f20ec16b6..fd5766b625 100644 --- a/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp +++ b/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp @@ -1,3 +1,9 @@ +#if __has_include("flagtree_spec.h") +#include "flagtree_spec.h" +#endif + +#ifndef FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_RemoveLayoutConversion + #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Analysis/TopologicalSortUtils.h" #include "mlir/Dialect/SCF/IR/SCF.h" @@ -1879,3 +1885,5 @@ class TritonGPURemoveLayoutConversionsPass }; } // namespace mlir::triton::gpu + +#endif diff --git a/lib/Dialect/TritonGPU/Transforms/Utility.cpp b/lib/Dialect/TritonGPU/Transforms/Utility.cpp index 315d252012..50ee6955ee 100644 --- a/lib/Dialect/TritonGPU/Transforms/Utility.cpp +++ b/lib/Dialect/TritonGPU/Transforms/Utility.cpp @@ -1,3 +1,9 @@ +#if __has_include("flagtree_spec.h") +#include "flagtree_spec.h" +#endif + +#ifndef FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_Utility + #include "triton/Analysis/Utility.h" #include @@ -1773,3 +1779,5 @@ LogicalResult verifyBarrierType(Operation *op, } } // namespace mlir::triton + +#endif diff --git a/python/setup_tools/setup_helper.py b/python/setup_tools/setup_helper.py index d6919b8762..12a579bbac 100644 --- a/python/setup_tools/setup_helper.py +++ b/python/setup_tools/setup_helper.py @@ -1,6 +1,7 @@ import os import shutil import sys +import sysconfig import functools from pathlib import Path import hashlib @@ -32,10 +33,18 @@ def install_extension(*args, **kargs): def get_backend_cmake_args(*args, **kargs): + if "editable_wheel" in sys.argv: + editable = True + else: + editable = False + handle_plugin_backend(editable) try: - return configs.activated_module.get_backend_cmake_args(*args, **kargs) + cmake_args = configs.activated_module.get_backend_cmake_args(*args, **kargs) except Exception: - return [] + cmake_args = [] + if editable: + cmake_args += ["-DEDITABLE_MODE=ON"] + return cmake_args def get_device_name(): @@ -234,6 +243,7 @@ def store(self, file=None, condition=None, url=None, copy_src_path=None, copy_ds def get(self, file_name) -> Path: return self.cache_files[file_name] +cache = FlagTreeCache() # -----flagtree-tle-raw-----flagtree-mlir--- @@ -347,7 +357,7 @@ def get_package_dir(packages): package_dict = {} if flagtree_backend and flagtree_backend not in configs.plugin_backends: connection = [] - backend_triton_path = f"../third_party/{flagtree_backend}/python/" + backend_triton_path = f"./third_party/{flagtree_backend}/python/" for package in packages: if CommonUtils.skip_package_dir(package): continue @@ -367,7 +377,33 @@ def handle_flagtree_backend(): print(f"\033[1;32m[INFO] FlagtreeBackend is {flagtree_backend}\033[0m") configs.extend_backends.append(flagtree_backend) if "editable_wheel" in sys.argv and flagtree_backend not in configs.plugin_backends: - ext_sourcedir = os.path.abspath(f"../third_party/{flagtree_backend}/python/{configs.ext_sourcedir}") + "/" + ext_sourcedir = os.path.abspath(f"./third_party/{flagtree_backend}/python/{configs.ext_sourcedir}") + "/" + + +def handle_plugin_backend(editable): + plugin_mode = os.getenv("FLAGTREE_PLUGIN") + if (plugin_mode and plugin_mode.upper() not in ["0", "OFF"]) or not flagtree_backend: + return + flagtree_backend_dir = cache.sub_dirs[flagtree_backend] + flagtree_plugin_so = flagtree_backend + "TritonPlugin.so" + src_build_plugin_path = flagtree_backend_dir / flagtree_plugin_so + if not src_build_plugin_path.exists(): + return + if flagtree_backend in ["iluvatar", "mthreads", "sunrise"]: + if editable is False: + dst_build_plugin_dir = Path(sysconfig.get_path("purelib")) / "triton" / "_C" + if not os.path.exists(dst_build_plugin_dir): + os.makedirs(dst_build_plugin_dir) + dst_build_plugin_path = dst_build_plugin_dir / flagtree_plugin_so + shutil.copy(src_build_plugin_path, dst_build_plugin_path) + if flagtree_backend in ("mthreads",): + dst_install_plugin_dir = Path( + __file__).resolve().parent.parent.parent / "third_party" / flagtree_backend / "python" / "triton" / "_C" + else: + dst_install_plugin_dir = Path(__file__).resolve().parent.parent / "triton" / "_C" + if not os.path.exists(dst_install_plugin_dir): + os.makedirs(dst_install_plugin_dir) + shutil.copy(src_build_plugin_path, dst_install_plugin_dir) def set_env(env_dict: dict): @@ -558,3 +594,41 @@ def uninstall_triton(): pre_hook=lambda: check_env('LLVM_SYSPATH'), post_hook=set_llvm_env, ) + +# sunrise +def sunrise_cp_bc_files(path): + # mkdir -p third_party/sunrise/backend/lib + lib_dir = Path("third_party/sunrise/backend/lib") + os.makedirs(lib_dir, exist_ok=True) + # cp ${LLVM_SYSPATH}/stpu/bitcode/*.bc third_party/sunrise/backend/lib + bc_dir = Path(path) / "stpu" / "bitcode" + for bc_file in bc_dir.glob("*.bc"): + shutil.copy(bc_file, lib_dir) + + +def sunrise_set_llvm_env(path): + set_llvm_env(path) + sunrise_cp_bc_files(path) + + +def sunrise_pre_llvm_env(): + llvm_path = os.environ.get('LLVM_SYSPATH', '') + ret = llvm_path != '' + if ret: + sunrise_cp_bc_files(llvm_path) + return ret + + +cache.store( + file="sunrise_llvm22_dev_release", + condition=("sunrise" == flagtree_backend), + url= None, + pre_hook=sunrise_pre_llvm_env, + post_hook=sunrise_set_llvm_env, +) + +cache.store( + file="sunriseTritonPlugin.so", + condition=("sunrise" == flagtree_backend) and (not configs.flagtree_plugin), + url=None, + copy_dst_path=f"third_party/{flagtree_backend}") diff --git a/python/setup_tools/utils/tools.py b/python/setup_tools/utils/tools.py index 8443cee7da..75bd2a6ad4 100644 --- a/python/setup_tools/utils/tools.py +++ b/python/setup_tools/utils/tools.py @@ -40,6 +40,7 @@ class FlagtreeConfigs: "cambricon": "mlu", "thrive": "thrive", "metax": "metax", + "sunrise": "sunrise", })) def __post_init__(self): diff --git a/python/triton/compiler/hint_manager.py b/python/triton/compiler/hint_manager.py index eabd631c23..36c8ccdc9b 100644 --- a/python/triton/compiler/hint_manager.py +++ b/python/triton/compiler/hint_manager.py @@ -61,12 +61,19 @@ def _load_handler(self, backend): except ImportError: # print(f"[FlagTree] Warning: Failed to load Nvidia Hint Handler: {e}", file=sys.stderr) return BaseHintHandler() + elif backend == 'sunrise': + try: + module = importlib.import_module("triton.backends.sunrise.sunrise_hint_handler") + return module.SunriseHintHandler() + except ImportError as e: + print(f"[FlagTree] Warning: Failed to load Sunrise Hint Handler: {e}", file=sys.stderr) + return BaseHintHandler() else: return BaseHintHandler() # supported backend with matched version -SUPPORTED_BACKENDS = ["aipu", "npu", "cuda"] +SUPPORTED_BACKENDS = ["aipu", "npu", "cuda", "sunrise"] # TODO : npu will have conflicts if more backend involved # mapping name @@ -74,6 +81,9 @@ def _load_handler(self, backend): "ascend": "npu", "huawei": "npu", "nvidia": "cuda", + # sunrise: GPUTarget backend name is "tang", torch device type is "ptpu". + "tang": "sunrise", + "ptpu": "sunrise", } diff --git a/python/triton/runtime/adjust_kernel_param.py b/python/triton/runtime/adjust_kernel_param.py index 7f97d32eba..02ae776bb3 100644 --- a/python/triton/runtime/adjust_kernel_param.py +++ b/python/triton/runtime/adjust_kernel_param.py @@ -1121,6 +1121,11 @@ def analyze_kernel_dependencies(jit_fn, pre_hook_fn: object | None = None) -> tu tma_m_map, tma_k_map = analyzer.analyze_tma_dot_dim(tma_map) # BS_M: {A}, BS_K: {A, B} # Analyzer 4: tl.dot M/N/K via load_map ge_m_map, ge_k_map, ge_n_map = analyzer.analyze_general_dot_dim(load_map) # BS_M: {A}, BS_K: {A, B}, BS_N: {B} + # clear load_map without any adjustment while dot analyzing encounters a special case, + # or the later adjustment of load_map may lead to the BS_X below the dot lower bound. + # for example: tl.dot(tl.load(A + offset0), tl.load(A + offset1)) + if analyzer.dot_calls and not ge_k_map: + load_map.clear() # cache _analysis_cache[cache_key] = (load_map, tma_map, tma_m_map, tma_k_map, ge_m_map, ge_k_map, ge_n_map) @@ -1319,6 +1324,17 @@ def auto_adjust_block_sizes(nargs, fn, configs, current, config): print("[AABS] 4. adjust bs in tl.dot with general tl.load") adjust_block_size_general_dot_mn_dim(nargs, current, config, ge_m_map, 16) adjust_block_size_general_dot_mn_dim(nargs, current, config, ge_n_map, 16) + if FLAGTREE_BACKEND == "sunrise": + # sunrise min_dot_size = (M=8, N=8, K=16/4) (see sunrise compiler.py + # min_dot_size). The tl.load shrink path above can lower a BLOCK that + # also feeds tl.dot below the dot lower bound; bump M/N/K back up to + # at least min_dot_size so AABS never produces a config that violates + # the dot constraint checked in semantic.py. + if knobs.autotuning.print: + print("[AABS] 4. adjust bs in tl.dot with general tl.load") + adjust_block_size_general_dot_mn_dim(nargs, current, config, ge_k_map, 16) + adjust_block_size_general_dot_mn_dim(nargs, current, config, ge_m_map, 8) + adjust_block_size_general_dot_mn_dim(nargs, current, config, ge_n_map, 8) if knobs.autotuning.print: nargs_str = '' diff --git a/setup.py b/setup.py index 3c3870e472..8cb4b2777e 100644 --- a/setup.py +++ b/setup.py @@ -16,7 +16,7 @@ from pathlib import Path from typing import Optional -from setuptools import Extension, find_packages, setup +from setuptools import Extension, find_namespace_packages, setup from setuptools.command.build_ext import build_ext from setuptools.command.build_py import build_py from setuptools.command.develop import develop @@ -644,7 +644,7 @@ def download_and_copy_dependencies(): if helper.flagtree_backend: - if helper.flagtree_backend in ("aipu", "tsingmicro", "enflame", "rpu", "thrive"): + if helper.flagtree_backend in ("aipu", "tsingmicro", "enflame", "rpu", "thrive", "sunrise"): backends = [ *BackendInstaller.copy(helper.configs.default_backends + tuple(helper.configs.extend_backends)), *BackendInstaller.copy_externals(), @@ -686,7 +686,7 @@ def get_package_dirs(): def get_packages(): - yield from find_packages(where="python", include=["triton", "triton.*"]) + yield from find_namespace_packages(where="python", include=["triton", "triton.*"]) for backend in backends: yield f"triton.backends.{backend.name}" diff --git a/third_party/sunrise/CMakeLists.txt b/third_party/sunrise/CMakeLists.txt new file mode 100644 index 0000000000..0aea96bd52 --- /dev/null +++ b/third_party/sunrise/CMakeLists.txt @@ -0,0 +1,44 @@ +add_compile_options("-Wno-deprecated-declarations") +add_compile_options("-Wno-error=deprecated-declarations") + +option(EDITABLE_MODE "Build in developer (editable) mode" OFF) +set(SUNRISE_PLUGIN_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}) + +if(TRITON_BUILD_PYTHON_MODULE) + if(FLAGTREE_PLUGIN) + add_subdirectory(plugin) + add_triton_plugin(TritonSunrise + SHARED_LIB sunriseTritonPlugin + ) + else() + if(EDITABLE_MODE) + find_library(sunriseTritonPluginLib + NAMES + sunriseTritonPlugin.so + PATHS + ${SUNRISE_PLUGIN_DIR} + REQUIRED + ) + add_triton_plugin(TritonSunrise + SHARED_LIB ${sunriseTritonPluginLib} + ) + else() + find_library(sunriseTritonPluginLib + NAMES + sunriseTritonPlugin.so + PATHS + ${SUNRISE_PLUGIN_DIR} + REQUIRED + ) + add_triton_plugin(TritonSunrise + SHARED_LIB ${sunriseTritonPluginLib} + ) + endif() + endif() +endif() + +add_subdirectory(backend/spec/lib) +add_subdirectory(${PROJECT_SOURCE_DIR}/include + ${PROJECT_BINARY_DIR}/include) +add_subdirectory(${PROJECT_SOURCE_DIR}/lib + ${PROJECT_BINARY_DIR}/lib) diff --git a/third_party/sunrise/backend/__init__.py b/third_party/sunrise/backend/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/third_party/sunrise/backend/compiler.py b/third_party/sunrise/backend/compiler.py new file mode 100644 index 0000000000..a80f9f8ecf --- /dev/null +++ b/third_party/sunrise/backend/compiler.py @@ -0,0 +1,445 @@ +from triton.backends.compiler import BaseBackend, GPUTarget, Language +from triton._C.libtriton import ir, passes, llvm, sunrise +from triton._C.libtriton import tle +from triton import knobs +from triton.knobs import base_knobs, env_opt_str, env_int, env_bool +from dataclasses import dataclass +import functools +from typing import Any, Dict, Tuple +from types import ModuleType +import hashlib +import platform +import re +import tempfile +import os +import subprocess +from pathlib import Path + + +class sunrise_knobs(base_knobs): + lld_path: env_opt_str = env_opt_str("TRITON_SUNRISE_LLD_PATH") + triple: env_opt_str = env_opt_str("TRITON_SUNRISE_TRANSLATE_TRIPLE") + flag: env_opt_str = env_opt_str("TRITON_SUNRISE_OPTION_FLAG") + libtang_path: env_opt_str = env_opt_str("TRITON_LIBTANG_PATH") + opt_level: env_int = env_int("TRITON_SUNRISE_OPTIMIZATION_LEVEL", 3) + + +knobs_sunrise = sunrise_knobs() + + +def min_dot_size(target: GPUTarget): + return lambda lhsType, rhsType: (8, 8, 16) if lhsType.is_int8() else (8, 8, 4) + +@functools.lru_cache(None) +def file_hash(path): + with open(path, "rb") as f: + return hashlib.sha256(f.read()).hexdigest() + +@dataclass(frozen=True) +class SunriseOptions: + num_warps: int = 4 + num_ctas: int = 1 + num_stages: int = 3 + warp_size: int = 32 + enable_fp_fusion: bool = True + supported_fp8_dtypes: Tuple[str] = ("fp8e5", ) + deprecated_fp8_dot_operand_dtypes: Tuple[str] = () + default_dot_input_precision: str = "ieee" + allowed_dot_input_precisions: Tuple[str] = ("ieee", ) + max_num_imprecise_acc_default: bool = None + extern_libs: dict = None + debug: bool = False + backend_name: str = 'tang' + sanitize_overflow: bool = True + arch: str = None + instrumentation_mode: str = "" + + # libdivice库,需要区分S2/S3； + def __post_init__(self): + warp_size = 32 + object.__setattr__(self, 'warp_size', warp_size) + default_libdir = Path(__file__).parent / 'lib' + extern_libs ={} if self.extern_libs is None else dict(self.extern_libs) + lib_ver = 'S2' + if knobs_sunrise.triple and ('stcuv2' in knobs_sunrise.triple): + lib_ver = 'S3' + for lib in ["ocml", "ockl"]: + # extern_libs[lib] = str(default_libdir / f'{lib}_{lib_ver}.bc') + lib_path = str(default_libdir / f'{lib}_{lib_ver}.bc') + if os.path.exists(lib_path): + extern_libs[lib] = lib_path + else: + extern_libs[lib] = str(default_libdir / f'{lib}.bc') + if lib_ver == 'S3': + for lib in ["builtin_math_fdiv", "cccl"]: + lib_path = default_libdir / f'{lib}_S3.bc' + if lib_path.exists(): + extern_libs[lib] = str(lib_path) + object.__setattr__(self, 'extern_libs', tuple(extern_libs.items())) + assert self.num_warps > 0 and (self.num_warps & (self.num_warps - 1)) == 0, \ + "num_warps must be a power of 2" + + def hash(self): + hash_dict = dict(self.__dict__) + hash_dict["extern_libs"] = tuple((k, file_hash(v)) for k, v in sorted(hash_dict["extern_libs"])) + key = "_".join([f"{name}-{val}" for name, val in sorted(hash_dict.items())]) + return hashlib.sha256(key.encode("utf-8")).hexdigest() + + +class SunriseBackend(BaseBackend): + instrumentation = None + supports_native_tensor_specialization = False + stage_stcu_dump_dir_name = '0' # 用于保存 .asm 等文件 + + @staticmethod + def supports_target(target: GPUTarget): + return target.backend == 'tang' + + def get_target_name(self, options) -> str: + return f"tang:{options.arch}" # tang:S2 + + def __init__(self, target: GPUTarget) -> None: + super().__init__(target) + self.binary_ext = 'stcu' + + def parse_options(self, opts) -> Any: + args = {'arch': knobs.runtime.override_arch or self.target.arch} + if "enable_fp_fusion" not in opts: + args["enable_fp_fusion"] = knobs.language.default_fp_fusion + args["max_num_imprecise_acc_default"] = 0 # TODO + args.update({k: opts[k] for k in SunriseOptions.__dataclass_fields__.keys() \ + if k in opts and opts[k] is not None}) + return SunriseOptions(**args) + + def pack_metadata(self, metadata): + return ( + metadata.num_warps, + metadata.num_ctas, + metadata.shared, + ) + + def get_codegen_implementation(self, options): + codegen_fns = { + "min_dot_size": min_dot_size(self.target) + } + return codegen_fns + + def get_module_map(self) -> Dict[str, ModuleType]: + from triton.language.extra.tang import libdevice + return {"triton.language.extra.libdevice": libdevice} + + def load_dialects(self, ctx): + sunrise.load_dialects(ctx) + if SunriseBackend.instrumentation: + SunriseBackend.instrumentation.load_dialects(ctx) + + def path_to_clang_offload_bundler(): + lld_env_path = knobs_sunrise.lld_path + if lld_env_path is not None: + lld = Path(lld_env_path) + if lld.is_file(): + return lld + arch = platform.machine() + lld = Path(f"/usr/local/tangrt/toolchains/llvm/prebuilt/linux-{arch}/bin/clang-offload-bundler") + if lld.is_file(): + return lld + raise Exception("clang-offload-bundler not found. Set 'TRITON_SUNRISE_LLD_PATH' to its path.") + + @staticmethod + def get_triple(): + triple = knobs_sunrise.triple + if triple is None or triple == '': + return "stcu-unknown-tang" + return triple + + @staticmethod + def get_flag(metadata, opt): + flag = knobs_sunrise.flag + if flag is None or flag == []: + flag = ['enable-predicate'] + if isinstance(flag, str): + flag = flag.split() + if metadata["num_warps"] > 16: + flag.append('thread-regfile-size=64') + for name, path in opt.extern_libs: + if name == "ockl": + flag.append('ocklPath='+path) + return flag + + @staticmethod + def get_optimization_level(llvm): + opt = knobs_sunrise.opt_level + if int(opt) == 0: + return llvm.OPTIMIZE_O0 + if int(opt) == 1: + return llvm.OPTIMIZE_O1 + if int(opt) == 2: + return llvm.OPTIMIZE_O2 + return llvm.OPTIMIZE_O3 + + @staticmethod + def make_ttir(mod, metadata, opt): + pm = ir.pass_manager(mod.context) + pm.enable_debug() + passes.common.add_inliner(pm) + passes.ttir.add_rewrite_tensor_pointer(pm) + passes.ttir.add_rewrite_tensor_descriptor_to_pointer(pm) + passes.common.add_canonicalizer(pm) + passes.ttir.add_combine(pm) + passes.ttir.add_reorder_broadcast(pm) + passes.common.add_cse(pm) + passes.ttir.add_triton_licm(pm) + passes.common.add_symbol_dce(pm) + passes.ttir.add_loop_unroll(pm) + pm.run(mod, 'make_ttir') + return mod + + @staticmethod + def make_ttgir(mod, metadata, opt): + num_stages = opt.num_stages if opt.num_stages <= 3 else 3 + # TTIR -> TTGIR + pm = ir.pass_manager(mod.context) + pm.enable_debug() + emuTF32 = False + passes.ttir.add_convert_to_ttgpuir(pm, f"tang:{opt.arch}", opt.num_warps, 32, opt.num_ctas) + # flagtree tle raw: convert tle.dsl_region tensor args/results to shared + # MemDesc (sunrise variant: gpu.barrier instead of NVVM::Barrier0Op). + sunrise.passes.ttgpuir.add_tle_convert_arg_to_memdesc(pm) + # flagtree tle raw: drop redundant copies before TTGPU lowering obscures + # the load->store pattern (target-independent). + tle.raw_passes.add_tle_remove_redundant_copy(pm) + # flagtree tle: canonicalize static gmem->local_ptr(shared) chunk copies + # into direct async copies before local-pointer lowering hides the + # load->store pattern (target-independent). + tle.passes.add_optimize_local_pointer_async_stores(pm) + # optimize TTGIR + passes.ttgpuir.add_coalesce(pm) + passes.ttgpuir.add_process_shared_memory_hint(pm) # flagtree hints (#@hint: shared_memory) + passes.ttgpuir.add_f32_dot_tc(pm, emuTF32) + passes.ttgpuir.add_remove_layout_conversions(pm) + passes.ttgpuir.add_optimize_thread_locality(pm) + # begin flagtree tle: local pointer (tle.gpu.local_ptr / alloc) support. + # These passes are target-independent (live in third_party/tle) and lower + # tle.local_pointers + tl.load/tl.store on shared memory into ttg.local_*. + tle.passes.add_early_assign_memory_space(pm) + tle.passes.add_select_encodings(pm) + tle.passes.add_insert_local_pointer_barriers(pm) + tle.passes.add_optimize_local_pointer_loads(pm) + tle.passes.add_optimize_local_pointer_stores(pm) + # flagtree tle: promote transient local_store staging to pipelineable + # local_alloc (target-independent perf optimization). + tle.passes.add_promote_local_store_staging(pm) + # end flagtree tle + sunrise.passes.ttgpuir.add_combine_optimize(pm) + + sunrise.passes.ttgpuir.add_optimize_column_order_dot_operands(pm) + if os.getenv('OFF_MMA', '0') == '1': + print('not run accelerate_matmul pass') + else: + sunrise.passes.ttgpuir.add_accelerate_matmul(pm, 1, 0) # 版本：1.0 + # sunrise.passes.ttgpuir.add_mma_direct_store(pm) + passes.ttgpuir.add_remove_layout_conversions(pm) + passes.ttgpuir.add_optimize_dot_operands(pm, True) + + passes.common.add_cse(pm) + if os.getenv('DFT_PP', '0') == '1': + if os.getenv('OFF_ASYNC', '0') == '0': + passes.ttgpuir.add_assign_latencies(pm, num_stages) + passes.ttgpuir.add_schedule_loops(pm) + passes.ttgpuir.add_pipeline(pm, num_stages, True ) + if os.getenv('OFF_PREF', '0') == '0': + passes.ttir.add_loop_aware_cse(pm) + passes.common.add_canonicalizer(pm) + passes.ttir.add_loop_aware_cse(pm) + passes.ttgpuir.add_prefetch(pm) + else: + if os.getenv('OFF_ASYNC', '0') == '0': + passes.ttgpuir.add_assign_latencies(pm, num_stages) + passes.ttgpuir.add_schedule_loops(pm) + sunrise.passes.ttgpuir.add_pipeline(pm, num_stages, 1, 0) # 版本：1.0 + if os.getenv('OFF_PREF', '0') == '0': + passes.ttir.add_loop_aware_cse(pm) + passes.common.add_canonicalizer(pm) + passes.ttir.add_loop_aware_cse(pm) + passes.ttgpuir.add_prefetch(pm) + if os.getenv('OFF_MMA', '0') == '1': + print('not run accelerate_matmul pass') + else: + sunrise.passes.ttgpuir.add_accelerate_matmul(pm, 1, 0) # 版本：1.0 + sunrise.passes.ttgpuir.add_mma_direct_store(pm) + passes.ttgpuir.add_remove_layout_conversions(pm) + # sunrise.passes.ttgpuir.add_optimize_dot_operands(pm) + # passes.ttgpuir.add_remove_layout_conversions(pm) + passes.ttgpuir.add_reduce_data_duplication(pm) + passes.ttgpuir.add_reorder_instructions(pm) + # flagtree tle: downgrade async copies that cannot be lowered (e.g. width + # constraints) back to plain load+local_store, before async-load lowering. + tle.passes.add_downgrade_invalid_async_copy(pm) + # flagtree tle: Lowering load with tt.load.async attribute + tle.passes.add_lower_async_load(pm) + passes.common.add_cse(pm) + passes.common.add_symbol_dce(pm) + sunrise.passes.ttgpuir.add_split_dot(pm, 1, 0) + passes.common.add_canonicalizer(pm) + + pm.run(mod, 'make_ttgir') + metadata["tensordesc_meta"] = mod.get_tensordesc_metadata() + return mod + + def ttgir_opt(self, src, metadata, options): + mod = src + pm = ir.pass_manager(mod.context) + pm.enable_debug() + + passes.gluon.add_inliner(pm) + passes.gluon.add_resolve_auto_encodings(pm) + passes.common.add_sccp(pm) + passes.ttir.add_loop_aware_cse(pm) + passes.gluon.add_canonicalizer(pm) + passes.ttgpuir.add_combine_tensor_select_and_if(pm) + + pm.run(mod, 'gluon_to_ttgir') + metadata["tensordesc_meta"] = mod.get_tensordesc_metadata() + return mod + + @staticmethod + def make_llir(src, metadata, options): + mod = src + # TritonGPU -> LLVM-IR (MLIR) + pm = ir.pass_manager(mod.context) + pm.enable_debug() + passes.ttgpuir.add_combine_tensor_select_and_if(pm) + passes.convert.add_scf_to_cf(pm) + passes.gluon.add_inliner(pm) + passes.convert.add_index_to_llvmir(pm) + + # flagtree tle: TLE-aware shared-memory allocation so tle.extract_tile / + # tle.insert_tile SMEM relay gets an allocation.offset. Falls back to + # default sizing for every other op. + sunrise.passes.ttgpuir.add_allocate_shared_memory(pm) + # instrumentation point here so we can override IRs above (e.g., ttir and ttgir) + if SunriseBackend.instrumentation: + SunriseBackend.instrumentation.patch("ttgpuir_to_llvmir", pm, mod.context) + + # flagtree tle raw: inline tle.dsl_region regions before TritonGPU->LLVM + # so no tle.dsl_region op survives into conversion (target-independent). + tle.raw_passes.add_tle_dsl_region_inline(pm) + sunrise.passes.ttgpuir.add_to_llvmir(pm, options.arch) + sunrise.passes.ttgpuir.add_remove_repeated_fence(pm) + passes.convert.add_scf_to_cf(pm) + passes.convert.add_cf_to_llvmir(pm) + passes.convert.add_arith_to_llvmir(pm) + passes.common.add_canonicalizer(pm) + passes.common.add_cse(pm) + passes.common.add_symbol_dce(pm) + if not knobs.compilation.disable_line_info and not knobs.compilation.dump_ir_extract_di_local_variables: + passes.llvmir.add_di_scope(pm) + + # This can not be moved below the di_scope pass + if SunriseBackend.instrumentation: + SunriseBackend.instrumentation.patch("llvmir_to_llvm", pm, mod.context) + + pm.run(mod, 'make_llir') + + if knobs.compilation.dump_ir_extract_di_local_variables: + # comments below on why separate it + if not knobs.compilation.disable_line_info: + pm = ir.pass_manager(mod.context) + pm.enable_debug() + passes.llvmir.add_di_scope(pm) + pm.run(mod, 'make_llir.disable_line_info') + + # insert dbg intrinsic with several DI Attribute including source + # var name and type info note: unknown reason for now, but this + # pass and add_di_scope has to be run separately, otherwise if we + # put them into previous pipline, it trigger a segmentfault without + # any error message; could be due to a bug in mlir or pybind11 + pm = ir.pass_manager(mod.context) + pm.enable_debug() + passes.llvmir.add_di_local_variable(pm) + pm.run(mod, 'make_llir.dump_ir_extract_di_local_variables') + + # LLVM-IR (MLIR) -> LLVM-IR (LLVM) + llvm.init_targets() + context = llvm.context() + llvm_mod = llvm.to_module(mod, context) + triple = SunriseBackend.get_triple() + llvm.attach_datalayout(llvm_mod, triple, '', '') + # sunrise.set_nvvm_reflect_ftz(llvm_mod) # 属性设置，可以不需要 + if options.extern_libs: + for name, path in options.extern_libs: + # OCKL is still materialized by the backend through -ocklPath. + if name != "ockl": + llvm.link_extern_libs(llvm_mod, [path]) + llvm.optimize_module(llvm_mod, SunriseBackend.get_optimization_level(llvm)) + + # Get some metadata + total_num_warps = src.get_int_attr("ttg.total-num-warps") + if total_num_warps is not None: + metadata["num_warps"] = total_num_warps + metadata["shared"] = src.get_int_attr("ttg.shared") + metadata["profile_scratch_size"] = src.get_int_attr("ttg.profile_scratch_memory_size") or 0 + metadata["profile_scratch_align"] = src.get_int_attr("ttg.profile_scratch_memory_alignment") or 1 + ret = str(llvm_mod) + ret = ret.replace("define void @", "define dso_local cc200 void @") + del llvm_mod + del context + return ret + + @staticmethod + def make_stcu(src, metadata, opt): + names = re.findall(r"define dso_local cc200 void @([a-zA-Z_][a-zA-Z0-9_]*)", src) + + assert len(names) == 1 + metadata["name"] = names[0] + proc = '' + + triple = SunriseBackend.get_triple() + flag = SunriseBackend.get_flag(metadata, opt) + asm = llvm.translate_to_asm(src, triple, proc, '', flag, opt.enable_fp_fusion, True) + + bundler = SunriseBackend.path_to_clang_offload_bundler() + + major = 0 + try: + output = subprocess.check_output([bundler, "--version"], stderr=subprocess.STDOUT) + version_str = output.decode("utf-8").strip() + match = re.search(r"version\s+(\d+)\.(\d+)\.(\d+)", version_str) + if match: + major = int(match.group(1)) + else: + print("Cannot parse clang-offload-bundler version\n") + except Exception as e: + print("Error getting version:", e) + + arch = platform.machine() + + with tempfile.NamedTemporaryFile() as tmp_out: + with tempfile.NamedTemporaryFile() as tmp_in: + with open(tmp_in.name, 'wb') as fd_in: + fd_in.write(asm) + try: + cmd = f'{bundler} -type=o -targets=host-{arch}-unknown-linux,tang-{triple} -input=/dev/null -input={tmp_in.name} -output={tmp_out.name}' + subprocess.run(cmd, shell=True, check=True) + except subprocess.CalledProcessError as e: + print(" run error\n") + + with open(tmp_out.name, 'rb') as fd_out: + ret = fd_out.read() + return ret + + def add_stages(self, stages, options, language): + if language == Language.TRITON: + stages["ttir"] = lambda src, metadata: self.make_ttir(src, metadata, options) + stages["ttgir"] = lambda src, metadata: self.make_ttgir(src, metadata, options) + elif language == Language.GLUON: + stages["ttgir"] = lambda src, metadata: self.ttgir_opt(src, metadata, options) + stages["llir"] = lambda src, metadata: self.make_llir(src, metadata, options) + stages["stcu"] = lambda src, metadata: self.make_stcu(src, metadata, options) + if knobs.runtime.add_stages_inspection_hook is not None: + knobs.runtime.add_stages_inspection_hook(self, stages, options, language, None) + + @functools.lru_cache() + def hash(self): + version = subprocess.check_output([SunriseBackend.path_to_clang_offload_bundler(), "--version"], encoding='utf-8') + return f'{version}-{self.target.arch}' diff --git a/third_party/sunrise/backend/driver.c b/third_party/sunrise/backend/driver.c new file mode 100644 index 0000000000..336dfd0102 --- /dev/null +++ b/third_party/sunrise/backend/driver.c @@ -0,0 +1,172 @@ +#include "tang.h" +#include +#include +#include +#include +#define PY_SSIZE_T_CLEAN +#include + +// Raises a Python exception and returns false if code is not TANG_SUCCESS. +static bool gpuAssert(TAresult code, const char *file, int line) { + if (code == TANG_SUCCESS) + return true; + + const char *prefix = "Triton Error [TANG]: "; + const char *str; + taGetErrorString(code, &str); + char err[1024] = {0}; + strcat(err, prefix); + strcat(err, str); + PyGILState_STATE gil_state; + gil_state = PyGILState_Ensure(); + PyErr_SetString(PyExc_RuntimeError, err); + PyGILState_Release(gil_state); + return false; +} + +// To be used only *outside* a Py_{BEGIN,END}_ALLOW_THREADS block. +#define TANG_CHECK_AND_RETURN_NULL(ans) \ + do { \ + if (!gpuAssert((ans), __FILE__, __LINE__)) \ + return NULL; \ + } while (0) + +// To be used inside a Py_{BEGIN,END}_ALLOW_THREADS block. +#define TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(ans) \ + do { \ + if (!gpuAssert((ans), __FILE__, __LINE__)) { \ + PyEval_RestoreThread(_save); \ + return NULL; \ + } \ + } while (0) + +static PyObject *getDeviceProperties(PyObject *self, PyObject *args) { + int device_id; + if (!PyArg_ParseTuple(args, "i", &device_id)) + return NULL; + // Get device handle + TAdevice device; + taDeviceGet(&device, device_id); + + // create a struct to hold device properties + int max_shared_mem; + int max_num_regs; + int multiprocessor_count; + int warp_size; + int sm_clock_rate; + int mem_clock_rate; + int mem_bus_width; + TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( + &max_shared_mem, TA_DEV_ATTR_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, + device)); + TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( + &max_num_regs, TA_DEV_ATTR_REGS_PER_BLOCK, device)); + TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( + &multiprocessor_count, TA_DEV_ATTR_MULTIPROCESSOR_COUNT, device)); + TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( + &warp_size, TA_DEV_ATTR_WARP_SIZE, device)); + TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( + &sm_clock_rate, TA_DEV_ATTR_CLOCK_RATE, device)); + TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( + &mem_clock_rate, TA_DEV_ATTR_MEMORY_CLOCK_RATE, device)); + TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( + &mem_bus_width, TA_DEV_ATTR_MEMORY_BUS_WIDTH, device)); + + return Py_BuildValue("{s:i, s:i, s:i, s:i, s:i, s:i, s:i}", "max_shared_mem", + max_shared_mem, "max_num_regs", max_num_regs, + "multiprocessor_count", multiprocessor_count, "warpSize", + warp_size, "sm_clock_rate", sm_clock_rate, + "mem_clock_rate", mem_clock_rate, "mem_bus_width", + mem_bus_width); +} + +static PyObject *loadBinary(PyObject *self, PyObject *args) { + const char *name; + const char *data; + Py_ssize_t data_size; + int shared; + int device; + if (!PyArg_ParseTuple(args, "ss#ii", &name, &data, &data_size, &shared, + &device)) { + return NULL; + } + TAfunction fun; + TAmodule mod; + int32_t n_regs = 0; + int32_t n_spills = 0; + int32_t n_max_threads = 0; + // create driver handles + TAcontext pctx = 0; + TAdevice device_hd; + taDeviceGet(&device_hd, device); + + Py_BEGIN_ALLOW_THREADS; + TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(taCtxGetCurrent(&pctx)); + if (!pctx) { + TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( + taDevicePrimaryCtxRetain(&pctx, device_hd)); + TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(taCtxSetCurrent(pctx)); + } + + TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(taModuleLoadData(&mod, data, (size_t)data_size)); + TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( + taModuleGetFunction(&fun, mod, name)); + // get allocated registers and spilled registers from the function + /* 不支持属性获取, 按照默认0处理 */ + // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( + // taFuncGetAttribute(&n_regs, TA_FUNC_ATTRIBUTE_NUM_REGS, fun)); + // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( + // taFuncGetAttribute(&n_spills, CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, fun)); + // n_spills /= 4; + TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(taFuncGetAttribute( + &n_max_threads, TA_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, fun)); + // set dynamic shared memory if necessary + // int shared_optin; + // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(cuDeviceGetAttribute( + // &shared_optin, TA_DEV_ATTR_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, + // device)); + // if (shared > 49152 && shared_optin > 49152) { + // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( + // cuFuncSetCacheConfig(fun, CU_FUNC_CACHE_PREFER_SHARED)); + // int shared_total, shared_static; + // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(cuDeviceGetAttribute( + // &shared_total, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_MULTIPROCESSOR, + // device)); + // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(cuFuncGetAttribute( + // &shared_static, CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, fun)); + // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( + // cuFuncSetAttribute(fun, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + // shared_optin - shared_static)); + // } + Py_END_ALLOW_THREADS; + + if (PyErr_Occurred()) { + return NULL; + } + return Py_BuildValue("(KKiii)", (uint64_t)mod, (uint64_t)fun, n_regs, + n_spills, n_max_threads); +} + +static PyMethodDef ModuleMethods[] = { + {"load_binary", loadBinary, METH_VARARGS, + "Load provided cubin into TANG driver"}, + {"get_device_properties", getDeviceProperties, METH_VARARGS, + "Get the properties for a given device"}, + {NULL, NULL, 0, NULL} // sentinel +}; + +static struct PyModuleDef ModuleDef = {PyModuleDef_HEAD_INIT, "tang_utils", + NULL, // documentation + -1, // size + ModuleMethods}; + +PyMODINIT_FUNC PyInit_tang_utils(void) { + PyObject *m = PyModule_Create(&ModuleDef); + if (m == NULL) { + return NULL; + } + + PyModule_AddFunctions(m, ModuleMethods); + + return m; +} diff --git a/third_party/sunrise/backend/driver.py b/third_party/sunrise/backend/driver.py new file mode 100644 index 0000000000..583fd61e90 --- /dev/null +++ b/third_party/sunrise/backend/driver.py @@ -0,0 +1,564 @@ +import functools +import os +import platform +import subprocess +import re +import triton +from pathlib import Path +from triton import knobs +from triton.runtime.build import compile_module_from_src +from triton.runtime import _allocation +from triton.backends.compiler import GPUTarget +from triton.backends.driver import GPUDriver + +dirname = os.path.dirname(os.path.realpath(__file__)) +include_dirs = [os.path.join(dirname, "include")] +libdevice_dir = os.path.join(dirname, "lib") +libraries = ['tang', 'tangrt_shared'] +arch = platform.machine() + +@functools.lru_cache() +def libtang_dirs(): + if env_libtang_path := knobs.env_opt_str("TRITON_LIBTANG_PATH"): + return [env_libtang_path] + + libs = subprocess.check_output(["/sbin/ldconfig", "-p"]).decode(errors="ignore") + # each line looks like the following: + # libtang.so.1 (libc6,x86-64) => /lib/x86_64-linux-gnu/libtang.so.1 + locs = [line.split()[-1] for line in libs.splitlines() if "libtang.so" in line] + dirs = [os.path.dirname(loc) for loc in locs] + env_ld_library_path = os.getenv("LD_LIBRARY_PATH") + if env_ld_library_path and not dirs: + dirs = [dir for dir in env_ld_library_path.split(":") if os.path.exists(os.path.join(dir, "libtang.so"))] + if not dirs: + dirs = [f'/usr/local/tangrt/lib/linux-{arch}/stub/'] + msg = 'libtang.so cannot found!\n' + if locs: + msg += 'Possible files are located at %s.' % str(locs) + msg += 'Please create a symlink of libtang.so to any of the files.' + else: + msg += 'Please make sure GPU is set up and then run "/sbin/ldconfig"' + msg += ' (requires sudo) to refresh the linker cache.' + assert any(os.path.exists(os.path.join(path, 'libtang.so')) for path in dirs), msg + return dirs + + +@functools.lru_cache() +def library_dirs(): + return [libdevice_dir, *libtang_dirs(), f"/usr/local/tangrt/lib/linux-{arch}"] + + +# ------------------------ +# Utils +# ------------------------ + + +class SunriseUtils(object): + + def __new__(cls): + if not hasattr(cls, "instance"): + cls.instance = super(SunriseUtils, cls).__new__(cls) + return cls.instance + + def __init__(self): + mod = compile_module_from_src( + src=Path(os.path.join(dirname, "driver.c")).read_text(), + name="tang_utils", + library_dirs=library_dirs(), + include_dirs=include_dirs, + libraries=libraries, + ) + self.load_binary = mod.load_binary + self.get_device_properties = mod.get_device_properties + + +# ------------------------ +# Launcher +# ------------------------ + + +def ty_to_cpp(ty): + if ty[0] == '*': + return "TAdeviceptr" + return { + "i1": "int_t", + "i8": "int8_t", + "i16": "int16_t", + "i32": "int32_t", + "i64": "int64_t", + "u1": "uint8_t", + "u8": "uint8_t", + "u16": "uint16_t", + "u32": "uint32_t", + "u64": "uint64_t", + "fp16": "double", + "bf16": "double", + "fp32": "double", + "f32": "double", + "fp64": "double", + }[ty] + +FLOAT_STORAGE_TYPE = { + "fp16": "uint16_t", + "bf16": "uint16_t", + "fp32": "uint32_t", + "f32": "uint32_t", + "fp64": "uint64_t", +} +FLOAT_PACK_FUNCTION = { + "fp16": "pack_fp16", + "bf16": "pack_bf16", + "fp32": "pack_fp32", + "f32": "pack_fp32", + "fp64": "pack_fp64", +} + +_BASE_ARGS_FORMAT = "iiiKKOOOOO" + +def make_launcher(constants, signature, warp_size, tensordesc_meta): + def _expand_signature(signature): + output = [] + # Expand tensor descriptor arguments into base pointer, shape, and + # strides + for sig in signature: + if isinstance(sig, str) and sig.startswith("tensordesc"): + ndim = sig.count(",") + 1 + dtype = re.match("tensordesc<([^[>]*)", sig).group() + + output.append("*" + dtype) + # Currently the host side tensor descriptors get passed in as a + # tensor desc, shape, and strides. We have no way to use these + # shape and strides when processing tensor descriptors which is + # why we provide our own decomposition above. Sadly this means + # we have to pass the shape and strides twice. + for _ in range(2 * ndim): + output.append("i64") + output.append("i1") + + for _ in range(ndim): + output.append("i32") + for _ in range(ndim): + output.append("i64") + else: + output.append(sig) + + return output + + def _flatten_signature(sig, output): + # Flatten tuples + if isinstance(sig, tuple): + for x in sig: + _flatten_signature(x, output) + else: + output.append(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" + return { + "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)] + + expand_signature = _expand_signature(signature.values()) + signature = {i: s for i, s in enumerate(expand_signature)} + + args_format = ''.join([format_of(ty) for ty in signature.values()]) + format = _BASE_ARGS_FORMAT + args_format + + flat_signature = [] + for sig in signature.values(): + _flatten_signature(sig, flat_signature) + signature = {i: s for i, s in enumerate(flat_signature)} + args_list = ', ' + ', '.join(f"&_arg{i}" for i, ty in signature.items()) if len(signature) > 0 else '' + # Record the end of regular arguments; + # subsequent arguments are architecture-specific descriptors, such as tensor descriptors for CUDA. + arg_decl_list = [] + for i, ty in signature.items(): + if ty == "constexpr": + continue + if ty in FLOAT_STORAGE_TYPE: + arg_decl_list.append(f"{FLOAT_STORAGE_TYPE[ty]} arg{i}") + else: + arg_decl_list.append(f"{ty_to_cpp(ty)} arg{i}") + arg_decls = ', '.join(arg_decl_list) + internal_args_list = [] + for i, ty in signature.items(): + if ty[0] == "*": + internal_args_list.append(f"ptr_info{i}.dev_ptr") + elif ty in FLOAT_STORAGE_TYPE: + internal_args_list.append(f"_arg{i}_storage") + elif ty != "constexpr": + internal_args_list.append(f"_arg{i}") + + # generate glue code + 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] == "*" + ] + float_storage_decls = [ + f"{FLOAT_STORAGE_TYPE[ty]} _arg{i}_storage = {FLOAT_PACK_FUNCTION[ty]}(_arg{i});" + for i, ty in signature.items() + if ty in FLOAT_STORAGE_TYPE + ] + params = [f"&arg{i}" for i, ty in signature.items() if ty != "constexpr"] + params.append("&global_scratch") + params.append("&profile_scratch") + src = f""" +#include \"tang.h\" +#include \"tang_runtime.h\" +#include +#include +#include + +static inline void gpuAssert(TAresult code, const char *file, int line) +{{ + if (code != TANG_SUCCESS) + {{ + const char* prefix = "Triton Error [TANG]: "; + const char* str; + taGetErrorString(code, &str); + char err[1024] = {{0}}; + strcat(err, prefix); + strcat(err, str); + PyGILState_STATE gil_state; + gil_state = PyGILState_Ensure(); + PyErr_SetString(PyExc_RuntimeError, err); + PyGILState_Release(gil_state); + }} +}} + +#define TANG_CHECK(ans) {{ gpuAssert((ans), __FILE__, __LINE__); }} + +static void _launch(int gridX, int gridY, int gridZ, int num_warps, int num_ctas, int shared_memory, TAstream stream, TAfunction function, TAdeviceptr profile_scratch{', ' + arg_decls if len(arg_decls) > 0 else ''}) {{ + TAdeviceptr global_scratch = 0; + void *params[] = {{ {', '.join(params)} }}; + if (gridX*gridY*gridZ > 0) {{ + TANG_CHECK(taLaunchKernel(function, gridX, gridY, gridZ, {warp_size}*num_warps, 1, 1, shared_memory, stream, params, 0)); + }} +}} + +typedef struct _DevicePtrInfo {{ + TAdeviceptr 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 = 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 = PyLong_AsUnsignedLongLong(ret); + + if(!ptr_info.dev_ptr) {{ + return ptr_info; + }} + + // 暂时使用PyTorch接口的方案, 后续taPointerGetAttribute支持使用指针切换后，还是使用它 + // 获取 device 属性 + PyObject* device_obj = PyObject_GetAttrString(obj, "device"); + if (device_obj && device_obj != Py_None) {{ + // 获取 device.index + PyObject* index_obj = PyObject_GetAttrString(device_obj, "index"); + if (index_obj && PyLong_Check(index_obj)) {{ + int dev = PyLong_AsLong(index_obj); + // printf("[DEBUG] Switching to tensor device (device.index): %d\\n", dev); + tangSetDevice(dev); + }} + Py_XDECREF(index_obj); + }} + Py_XDECREF(device_obj); + + uint64_t dev_ptr; + int status = taPointerGetAttribute(&dev_ptr, TA_POINTER_ATTRIBUTE_DEVICE_POINTER, ptr_info.dev_ptr); + if (status == TANG_ERROR_INVALID_VALUE) {{ + PyErr_Format(PyExc_ValueError, + "Pointer argument (at %d) cannot be accessed from Triton (cpu tensor?)", idx); + ptr_info.valid = false; + }} else if (status != TANG_SUCCESS) {{ + TANG_CHECK(status); // Catch any other TANG API errors + ptr_info.valid = false; + }} + ptr_info.dev_ptr = dev_ptr; + Py_DECREF(ret); // Thanks ChatGPT! + 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; +}} + +static uint16_t pack_fp16(double f) {{ + uint16_t result; + // from https://github.com/python/pythoncapi-compat +#if 0x030600B1 <= PY_VERSION_HEX && PY_VERSION_HEX <= 0x030B00A1 && !defined(PYPY_VERSION) + _PyFloat_Pack2(f, (unsigned char*)&result, 1); +#else + PyFloat_Pack2(f, (unsigned char*)&result, 1); +#endif + return result; +}} + +static uint16_t pack_bf16(double f) {{ + float f32 = (float)f; + uint32_t u32 = *(uint32_t*)&f32; + return (uint16_t)(u32 >> 16); +}} + +static uint32_t pack_fp32(double f) {{ + float f32 = (float)f; + return *(uint32_t*)&f32; +}} + +static uint64_t pack_fp64(double f) {{ + return *(uint64_t*)&f; +}} + +static PyObject* launch(PyObject* self, PyObject* args) {{ + int gridX, gridY, gridZ; + uint64_t _stream; + uint64_t _function; + PyObject *launch_enter_hook = NULL; + PyObject *launch_exit_hook = NULL; + PyObject *kernel_metadata = NULL; + PyObject *launch_metadata = NULL; + PyObject *profile_scratch_obj = NULL; + {' '.join([f"{_extracted_type(ty)} _arg{i}; " for i, ty in signature.items()])} + if(!PyArg_ParseTuple(args, \"{format}\", &gridX, &gridY, &gridZ, &_stream, &_function, &profile_scratch_obj, + &kernel_metadata, &launch_metadata, + &launch_enter_hook, &launch_exit_hook {args_list})) {{ + PyErr_SetString(PyExc_TypeError, "get input data error"); + return NULL; + }} + + {' '.join(float_storage_decls)} + + int num_warps, num_ctas, shared_memory; + if (!PyArg_ParseTuple(kernel_metadata, \"iii\", &num_warps, &num_ctas, &shared_memory)) {{ + PyErr_SetString(PyExc_TypeError, "kernel_metadata must be a tuple"); + return NULL; + }} + + // extract launch metadata + 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; + Py_DECREF(ret); + }} + + TAdeviceptr profile_scratch = 0; + if (profile_scratch_obj != Py_None) {{ + DevicePtrInfo profile_scratch_info = getPointer(profile_scratch_obj, -1); + if (!profile_scratch_info.valid) {{ + return NULL; + }} + profile_scratch = profile_scratch_info.dev_ptr; + }} + + // 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(gridX, gridY, gridZ, num_warps, num_ctas, shared_memory, (TAstream)_stream, (TAfunction)_function, profile_scratch{', ' + ', '.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_DECREF(ret); + }} + + 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); + return m; +}} +""" + return src + +def wrap_handle_tensordesc(launcher, signature, tensordesc_meta): + has_tensor_desc_arg = any(isinstance(sig, str) and sig.startswith("tensordesc") for sig in signature.values()) + if not has_tensor_desc_arg: + return launcher + + from triton.tools.tensor_descriptor import TensorDescriptor + + def inner(*args): + meta_args = args[:len(_BASE_ARGS_FORMAT)] + raw_kernel_args = args[len(_BASE_ARGS_FORMAT):] + final_args = [] + for arg in raw_kernel_args: + if isinstance(arg, TensorDescriptor): + # Currently the host side tensor descriptors get decomposed in + # the frontend to tensor desc, shape, and strides. We have no + # way to use these shape and strides when processing tensor + # descriptors which is why we provide our own decomposition + # above. Sadly this means we have to pass the shape and strides + # twice. + final_args.extend([arg.base, *arg.shape, *arg.strides, *arg.shape, *arg.strides]) + else: + final_args.append(arg) + return launcher(*meta_args, *final_args) + + return inner + +class SunriseLauncher(object): + + def __init__(self, src, metadata): + constants = src.constants if hasattr(src, "constants") else dict() + arg_idx = lambda x: (src.fn.arg_names.index(x), ) if isinstance(x, str) else x + constants = {arg_idx(idx): value for idx, value in constants.items()} + signature = {idx: value for idx, value in src.signature.items()} + tensordesc_meta = getattr(metadata, "tensordesc_meta", None) + src = make_launcher(constants, signature, metadata.warp_size, tensordesc_meta) + mod = compile_module_from_src( + src=src, + name="__triton_launcher", + library_dirs=library_dirs(), + include_dirs=include_dirs, + libraries=libraries, + ) + has_tensor_desc_arg = any(isinstance(sig, str) and sig.startswith("tensordesc") for sig in signature.values()) + self.launch = wrap_handle_tensordesc(mod.launch) if has_tensor_desc_arg else mod.launch + self.profile_scratch_size = metadata.profile_scratch_size + self.profile_scratch_align = metadata.profile_scratch_align + + def __call__(self, gridX, gridY, gridZ, stream, function, *args): + def allocate_scratch(size, align, allocator): + if size > 0: + grid_size = gridX * gridY * gridZ + alloc_size = grid_size * self.num_ctas * size + alloc_fn = allocator.get() + return alloc_fn(alloc_size, align, stream) + return None + + profile_scratch = allocate_scratch(self.profile_scratch_size, self.profile_scratch_align, + _allocation._profile_allocator) + self.launch(gridX, gridY, gridZ, stream, function, profile_scratch,*args) + + +class SunriseDriver(GPUDriver): + + def __init__(self): + self.utils = SunriseUtils() # TODO: make static + self.launcher_cls = SunriseLauncher + import torch + if not hasattr(torch, "ptpu"): + raise RuntimeError("torch.ptpu is not available") + self.get_device_capability = torch.ptpu.get_device_capability + self.get_current_stream = lambda dev_idx: torch.ptpu.current_stream(dev_idx).ptpu_stream + self.get_current_device = torch.ptpu.current_device + self.set_current_device = torch.ptpu.set_device + + + def get_current_target(self): + arch = knobs.runtime.override_arch + if not arch: + arch = "S2" + warp_size = 32 + return GPUTarget("tang", arch, warp_size) + + def get_active_torch_device(self): + import torch + return torch.device("ptpu", self.get_current_device()) + + def get_device_interface(self): + import torch + return torch.ptpu + + @staticmethod + def is_active(): + # 默认用pts后端 + import torch + return True + return torch.cuda.is_available() and (torch.version.hip is None) and (torch.version.cuda is None) + + def map_python_to_cpp_type(self, ty: str) -> str: + return ty_to_cpp(ty) + + def get_benchmarker(self): + from triton.testing import do_bench + return do_bench + + def get_empty_cache_for_benchmark(self): + import torch + import torch_ptpu + + # We maintain a buffer of 256 MB that we clear + # before each kernel call to make sure that the L2 cache + # doesn't contain any input data before the run + cache_size = 256 * 1024 * 1024 + return torch.empty(int(cache_size // 4), dtype=torch.int, device='ptpu') + + def clear_cache(self, cache): + cache.zero_() diff --git a/third_party/sunrise/backend/include/ptml.h b/third_party/sunrise/backend/include/ptml.h new file mode 100755 index 0000000000..2b5d7b3e8b --- /dev/null +++ b/third_party/sunrise/backend/include/ptml.h @@ -0,0 +1,925 @@ +//////////////////////////////////////////////////////// +// @file ptml.h +// ptml api +// ptmlDevice_t represents the type of device index +//////////////////////////////////////////////////////// + +#ifndef _PT_ML_H_ +#define _PT_ML_H_ + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif //! __cplusplus + +#include + +#ifndef TA_PT_NUM_MAX +#define TA_PT_NUM_MAX 128 +#endif //! TA_PT_NUM_MAX + +#if defined(_MSC_VER) +#define PTML_DEPRECATED __declspec(deprecated) +#define PTML_API_EXPORT __declspec(dllexport) +#define PTML_API_IMPORT __declspec(dllimport) +#elif defined(__GNUC__) || defined(__clang__) +#define PTML_DEPRECATED __attribute__((deprecated)) +#define PTML_API_EXPORT __attribute__((visibility("default"))) +#define PTML_API_IMPORT __attribute__((visibility("default"))) +#else +#define PTML_DEPRECATED +#define PTML_API_EXPORT +#define PTML_API_IMPORT +#endif //! UNKNOWN COMPILER + +#if defined(ptml_EXPORTS) +#define PTML_API PTML_API_EXPORT +#else +#define PTML_API PTML_API_IMPORT +#endif //! For user + +#define PAGE_SIZE 4096 +#ifndef ALIGN +#define __ALIGN_KERNEL_MASK(x, mask) (((x) + (mask)) & ~(mask)) +#define __ALIGN_KERNEL(x, a) __ALIGN_KERNEL_MASK(x, (__typeof__(x))(a)-1) +#define ALIGN(x, a) __ALIGN_KERNEL((x), (a)) +#endif // ALIGN +#define PAGE_ALIGN(addr) ALIGN(addr, PAGE_SIZE) +#define CLK_NAME_MAX 16 +#define NUMBER_OF_CYCLES_IN_1_SEC 0x38400000 // 900M +#define am_interval_1s (1800000000) // 1800M stands for 1 second + +/*ioctl parm cmd*/ +#define CM3 0x10 +#define MLP 0x11 +#define LINUX 0x12 + +#define CMD_PMIC (0xa8) +#define MOD_TEMP_TYPE_NUM (1) +#define TEMP_IPID_CNT (8) +#define C2C_INFO_CNT (10) + +#define CMD_TEMPERATURE 0xab +#define CMD_HBM_TEMPERATURE 0xb2 +#define CMD_GET_GPIO_STATUS 0xb7 +#define CMD_DUMP_MEM 0xb5 +#define CMD_GET_CPLD_VERSION 0xbc +#define CMD_GET_MAX_POWER 0xbd +#define CMD_GET_EXCEPTION 0xb6 + +/*linux cmd*/ +#define LINUX_CMD_PTUTILI (0x11) +#define LINUX_CMD_PCIERELINK (0x13) +#define LINUX_CMD_HBMBWUTILI (0x1A) +#define LINUX_CMD_HBMUTILI (0x1B) +#define LINUX_CMD_C2CREVDB (0x14) +#define LINUX_CMD_C2CTRANSDB (0x15) +#define LINUX_CMD_PCIEREVDB (0x16) +#define LINUX_CMD_PCIETRANSDB (0x17) +#define LINUX_CMD_TUUTILI (0x18) +#define LINUX_CMD_THREADUTILI (0x19) + +/** + * @brief Return val for ptml API + */ +typedef enum ptmlReturn_enum { + PTML_SUCCESS = 0, //!< APT returns ok + PTML_ERROR_UNINITIALIZED, //!< ptmlInit is not called now + PTML_ERROR_INVALID_ARGUMENT, //!< invalid argument + PTML_ERROR_ALREADY_INITIALIZED, //!< ptmlInit is already called + PTML_ERROR_INSUFFICIENT_SIZE, //!< An input argument is not large enough + PTML_ERROR_IN_USE, //!< PT is in use + PTML_ERROR_DRIVER_NOT_LOADED, //!< driver is not loaded + PTML_ERROR_DEVICE_NOT_FOUND, //!< device is not found + PTML_ERROR_EVENT_TIMEOUT, //!< device is not found + PTML_ERROR_UNKNOWN, //!< An internal driver error occurred +} ptmlReturn_t; + + +typedef struct { + enum ptmlReturn_enum errorCode; + const char *errorMessage; +} ErrorDescriptor; + +typedef enum ptmlClockType { + //!< PTML_CLOCK_GRAPHICS = 0, + //!< PTML_CLOCK_SM, + PTML_CLOCK_PT = 0, + PTML_CLOCK_MEM, + //!< PTML_CLOCK_VIDEO, +} ptmlClockType_t; + +/** + * @brief Device Handle type + * + ********************************************/ +typedef int ptmlDevice_t; + +typedef struct ptMemory { + size_t total; //!< total memory + size_t used; //!< used memory + size_t free; //!< free memory +} ptMemory_t; + +#define PTML_DEVICE_PCI_BUS_ID_BUFFER_SIZE 32 + +typedef struct ptPciInfo { + int domain; //!< domain number + int bus; //!< bus number + int device; //!< dev && func number + int vendor; //!< vendor number + int pciSubSystemId; //!< subsystem Id + char busId[PTML_DEVICE_PCI_BUS_ID_BUFFER_SIZE]; //!< "domain:bus:device:0" + unsigned int max_link_speed; //!< max link speed(MT/s) + unsigned int max_link_width; //!< max link width + unsigned int max_bandwidth; //!< max bandwidth(MB/s) + unsigned int curr_link_speed; //!< current link speed(MT/s) + unsigned int curr_link_width; //!< current link width + unsigned int curr_bandwidth; //!< current bandwidth(MB/s) +} ptPciInfo_t; + +typedef struct ptProcessInfo { + char name[256]; + char pidStr[16]; +}ptProcessInfo_t; + +struct barInfo { + int barIdx; //!< 0-5 + void* addr; //!< The mapped vritual address. + uint64_t paddr; + size_t size; //!< The size of the mapped vritaul address space. +}; + +struct fw_param { + int cmd; + int len; + int data[]; +}; +#define _S2_IOC_SMI_INFO _IOWR(_S2_IOC_MAGIC, 160, struct fw_param) + +struct st_rpmsg_cmd { + unsigned int cmd; + unsigned char data[]; +} __attribute__((packed)); + +struct st_rpmsg_i2c_cmd { + unsigned char rx_data_len; + unsigned char tx_data_len; + unsigned char slave_len; + unsigned char trx_flag; + unsigned char pyload[]; // slave addr + tx data +} __attribute__((packed)); + +struct st_rpmsg_i2c_response { + unsigned char error_code; + unsigned char rx_data_len; + unsigned char slave_len; + unsigned char trx_flag; + unsigned char pyload[]; // slave addr + rx data +} __attribute__((packed)); + +struct st_rpmsg_pmic_cmd { + unsigned char pmic_cmd; + unsigned int pmic_param; +} __attribute__((packed)); + +struct st_rpmsg_pmic_response { + unsigned char pmic_cmd; + unsigned char error_code; + unsigned char pmic_data[4]; +} __attribute__((packed)); + +/* + * RPmsg Clock Command IDs + */ +enum rpmsg_clk_cmd_id { + RPMSG_CLK_GET_STATE, + RPMSG_CLK_GET_NAME, + RPMSG_CLK_GET_RATE, + RPMSG_CLK_ENABLE, + RPMSG_CLK_DISABLE, + RPMSG_CLK_CMD_COUNT, +}; + +enum plat_clock_idx { + MOD_CLOCK_G0_0, + MOD_CLOCK_G0_1, + MOD_CLOCK_G1_0, + MOD_CLOCK_G2_0, + MOD_CLOCK_G3_0, + MOD_CLOCK_G4_0, + MOD_CLOCK_G5_0, + MOD_CLOCK_G6_0, + MOD_CLOCK_G7_0, + MOD_CLOCK_G8_0, + MOD_CLOCK_G9_0, + MOD_CLOCK_G10_0, + MOD_CLOCK_G11_0, + MOD_CLOCK_G12_0, + MOD_CLOCK_G13_0, + MOD_CLOCK_G13_1, + MOD_CLOCK_G13_2, + MOD_CLOCK_G14_0, + MOD_CLOCK_G14_1, + + MOD_CLOCK_L0_CLK2000CLK, + MOD_CLOCK_L0_CLK1000CLK, + MOD_CLOCK_L0_CLK500CLK, + MOD_CLOCK_L0_CLK250CLK, + MOD_CLOCK_L0_CLK125CLK, + MOD_CLOCK_L0_CLK62P5CLK, + MOD_CLOCK_L0_CLK31P25CLK, + MOD_CLOCK_L0_SMB_MELESCLK, + MOD_CLOCK_L0_MELS_REF_CLK, + MOD_CLOCK_L0_SMB_32KCLK, + MOD_CLOCK_L0_PLL0_CLK, + MOD_CLOCK_L1_CORE_CLK_L, + MOD_CLOCK_L1_NOC_CLK0, + MOD_CLOCK_L1_PLL1_CLK, + MOD_CLOCK_L2_CORE_CLK_H, + MOD_CLOCK_L2_NOC_CLK1, + MOD_CLOCK_L2_PLL2_CLK, + MOD_CLOCK_L3_APBCLK, + MOD_CLOCK_L3_PLL3_CLK, + MOD_CLOCK_L4_VIDEO_CLK, + MOD_CLOCK_L4_PLL4_CLK, + MOD_CLOCK_L5_DMA_CLK, + MOD_CLOCK_L5_TIGER_CLK, + MOD_CLOCK_L5_PLL5_CLK, + MOD_CLOCK_L6_AXI_CLOCK0, + MOD_CLOCK_L6_PLL6_CLK, + MOD_CLOCK_L7_AXI_CLOCK1, + MOD_CLOCK_L7_PLL7_CLK, + MOD_CLOCK_L8_JPEG_CLK, + MOD_CLOCK_L8_PLL8_CLK, + MOD_CLOCK_L9_PLLREFCLK0, + MOD_CLOCK_L9_DFICLK0, + MOD_CLOCK_L9_DFIHDRCLK0, + MOD_CLOCK_L9_PLL9_CLK, + MOD_CLOCK_L10_PLLREFCLK1, + MOD_CLOCK_L10_DFICLK1, + MOD_CLOCK_L10_DFIHDRCLK1, + MOD_CLOCK_L10_PLL10_CLK, + MOD_CLOCK_L11_PLLREFCLK2, + MOD_CLOCK_L11_DFICLK2, + MOD_CLOCK_L11_DFIHDRCLK2, + MOD_CLOCK_L11_PLL11_CLK, + MOD_CLOCK_L12_PLLREFCLK3, + MOD_CLOCK_L12_DFICLK3, + MOD_CLOCK_L12_DFIHDRCLK3, + MOD_CLOCK_L12_PLL12_CLK, + MOD_CLOCK_L13_ACLK0, + MOD_CLOCK_L13_PLL13_CLK, + MOD_CLOCK_L15_AUX_CLK0, + MOD_CLOCK_L16_ACLK1, + MOD_CLOCK_L16_PLL16_CLK, + MOD_CLOCK_L17_AUX_CLK1, + + MOD_CLOCK_IDX_COUNT, +}; + +enum c2c_port_index { + C2C0_0 = 0, + C2C0_1, + C2C1_0, + C2C1_1, + C2C2_0, + C2C2_1, + C2C3_0, + C2C3_1, + C2C4_0, + C2C4_1, + PCIE, +}; +/* + * struct c2h_clk_msg - Response payload for RPMSG_CLK_ATTRIBUTES_DUMP command + * @status: Command status + * @state: Clock state(on or off) + * @rate: Clock rate in Hz, + * rate[0] 32bit lsb clock rate + * rate[1] 32bit hsb clock rate + * @name: Clock name + */ +struct c2h_clk_msg { + int status; + uint32_t state; + uint32_t rate[2]; + char name[16]; +}; + +struct c2h_temp_msg { + int status; + int temp[TEMP_IPID_CNT]; +}; + +#define ALLPORTS 0x3ff +typedef struct ptPhyTopo { + unsigned char local_chipid; + unsigned char local_port; + unsigned char remote_chipid; + unsigned char remote_port; + unsigned char link_status; + unsigned char isBif; + unsigned char max_speed; + unsigned char cur_speed; + unsigned char max_bandwidth; + unsigned char cur_bandwidth; +} ptPhyTopo_t; + +/** + * @brief Init ptml module + * + * @return int + * @note must called before using any ptml API + ********************************************/ +ptmlReturn_t PTML_API ptmlInit(void); +ptmlReturn_t PTML_API __ptmlUinit(void); +void PTML_API ptmlUninit(void); + +/** + * @brief Get system driver version 0.1.0 + * + * @param length is driver version's length + * @note driver Version is Equivalent to dirver version, specific information + *in + * "/sys/module/pt/version" + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlSystemGetDriverVersion(char * version, + unsigned int length); + +/** + * @brief Get system tang version 0.1.0 + * + * @note tang Version is Equivalent to cuda version + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlSystemGetTangVersion(int *version); +ptmlReturn_t PTML_API ptmlSystemGetTangVersionForSmi(char * version, + unsigned int length); + +/** + * @brief Get dev base info: ptType and memInfo + * + * @param device device handles + * @param ptTypeOut pointer to ptTypeOut + * @param memInfo pointer to dev memInfo + * @note ptTypeOut is pt200, device info is total mem size + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetBaseInfo(ptmlDevice_t device, + char * ptTypeOut, + unsigned int *memInfo); + +/** + * @brief Get PT board count + * + * @param devCount pointer to devCount + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetBoardCount(unsigned int *devCount); + +/** + * @brief Get PT type :pt200 + * + * @param device device handles + * @param ptTypeOut pointer to pt type + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetBrand(ptmlDevice_t device, char *ptTypeOut); + +/** + * @brief Get device Capacity + * + * @param device device handles + * @param value pointer to Capacity + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetCapacity(ptmlDevice_t device, float *value); + +/** + * @brief Get PT count + * + * @param devCount pointer to devCcount + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetCount(int *devCount); + +/** + * @brief Get device c2c rev DB + * + * @param device device handles + * @param c2cIndex is port num 0-10 + * @param interval is from 1-60(s) + * @param revdb pointer is c2c revdb + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetC2CRevDB(ptmlDevice_t device, + int c2cIndex, + int interval, + unsigned int *revdb); + +/** + * @brief Get device c2c trans DB + * + * @param device device handles + * @param c2cIndex is port num 0-10 + * @param interval is from 1-60(s) + * @param transdb pointer is c2c transdb + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetC2CTransDB(ptmlDevice_t device, + int c2cIndex, + int interval, + unsigned int *transdb); + +/** + * @brief Get device fw version + * + * @param device device handles + * @param version pointer to fw version + * @param length is version's length + * @note fw is cm3 linux mlp mix version Internal use only Not provided + *externally + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetFirmwareVersion(ptmlDevice_t device, + char * version, + unsigned int length); + +/** + * @brief Get Device Handle by idx + * + * @param idx idx of the target PT + * @param device pointer to the handle of target PT + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetHanldeByIdx(unsigned int idx, + ptmlDevice_t *device); + +/** + * @brief Get Device Handle by PCI + * + * @param idx idx of the target PT + * @param device pointer to the handle of target PT + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetHanldeByPciBusId(const char * busId, + ptmlDevice_t *device); + +/** + * @brief Get device mem BW Utilization + * + * @param device device handles + * @param interval is from 1-60(s) + * @param utilization pointer is memBWUtilization 0-100% + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetMemBWUtilizationRates(ptmlDevice_t device, + int interval, + float *utilization); + +/** + * @brief ptmlDeviceGetMemClockFrequency + * + * @param device device handles + * @param clock pointer to memClockFreq + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetMemClockFrequency(ptmlDevice_t device, + unsigned int *clock); + +/** + * @brief Get device memory information + * + * @param device device handle + * @param memInfo pointer to memory information + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetMemoryInfo(ptmlDevice_t device, + ptMemory_t * memInfo); +ptmlReturn_t PTML_API ptmlDeviceGetMemoryUsedInfo(ptmlDevice_t device, + unsigned int *usedInfo); +ptmlReturn_t PTML_API ptmlDeviceGetMemoryFreeInfo(ptmlDevice_t device, + unsigned int *usedInfo); + +/** + * @brief Get device mem Temperature + * + * @param device device handles + * @param temp pointer to ptTemperature + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetMemTemperature(ptmlDevice_t device, + unsigned int *temp); + +/** + * @brief Get device mem Utilization + * + * @param device device handles + * @param interval is from 1-60(s) + * @param utilization pointer is memUtilization 0-100% + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t ptmlDeviceGetMemUtilizationRates(ptmlDevice_t device, + int interval, + float *utilization); + +/** + * @brief Get PT node path /dev/ptpux + * + * @param device device handle + * @param path pointer to devNodePath + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetNodePath(ptmlDevice_t device, char *path); + +/** + * @brief Get device pcie relink times + * + * @param device device handles + * @param poniter to pcie relink times + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetPcieRelinkTime(ptmlDevice_t device, + int * count); + +/** + * @brief Get device pcie rev DB + * + * @param device device handles + * @param interval is from 1-60(s) + * @param revdb pointer is pcie revdb + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetPcieRevDB(ptmlDevice_t device, + int interval, + unsigned int *revdb); + +/** + * @brief Get device pcie trans DB + * + * @param device device handles + * @param interval is from 1-60(s) + * @param transdb pointer is pcie transdb + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetPcieTransDB(ptmlDevice_t device, + int interval, + unsigned int *transdb); + +/** + * @brief Get device pci information + * + * @param device device handle + * @param pciInfo pointer to pci information + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetPciInfo(ptmlDevice_t device, + ptPciInfo_t *pciInfo); + +/** + * @brief Reboot device + * + * @param device device handles + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceRebootBootloader(ptmlDevice_t device); + +/** + * @brief Get PT status 0:invalid 1:valid + * + * @param device device handles + * @param status pointer to ptstatus + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetStatus(ptmlDevice_t device, int *status); + +/** + * @brief ptmlDeviceGetPtClockFrequency + * + * @param device device handles + * @param clock pointer to ptClockFreq + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetPtClockFrequency(ptmlDevice_t device, + unsigned int *clock); + +/** + * @brief Get PTCTRL major and minor + * + * @param device device handles + * @param major pointer to ptpuctrl major + * @param minor pointer to ptpuctrl minor + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetPtCtrlMajorAndMinor(int *major, + int *minor); + +/** + * @brief Get PT major and minor /dev/ptpux + * + * @param device device handles + * @param major pointer to devmajor + * @param minor pointer to devminor + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetPtMajorAndMinor(ptmlDevice_t device, + int * major, + int * minor); + +/** + * @brief Get device pt Temperature + * + * @param device device handles + * @param temp pointer to ptTemperature + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetPtTemperature(ptmlDevice_t device, + unsigned int *temp); + +/** + * @brief Get device pt Utilization + * + * @param device device handles + * @param interval is from 1-60(s) + * @param utilization pointer to ptUtilization 0-100% + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetPtUtilizationRates(ptmlDevice_t device, + int interval, + float *utilization); + +/** + * @brief Get compute capability + * + * @param device device handles + * @param major pointer to major + * @param minor pointer to minor + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetTangComputeCapability(ptmlDevice_t device, + int * major, + int * minor); + +/** + * @brief Get device thread Utilization + * + * @param device device handles + * @param interval is from 1-60(s) + * @param utilization pointer to threadUtilization 0-100% + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetThreadUtilizationRates(ptmlDevice_t device, + int interval, + float *utilization); + +/** + * @brief Get device subcore Utilization + * + * @param device device handles + * @param interval is from 1-60(s) + * @param utilization pointer to subcoreUtilization 0-100% + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetTUUtilizationRates(ptmlDevice_t device, + int interval, + float * utilization); + +/** + * @brief Get device uuid + * + * @param device device handles + * @param uuid pointer to device uuid + * @param length device uuid length + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetUUID(ptmlDevice_t device, + char * uuid, + unsigned int length); + +/** + * @brief ptmlDeviceSetCUFrequency + * + * @param device device handles + * @param CU freq + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceSetCUFrequency(ptmlDevice_t device, + unsigned int freq) ; + +/** + * @brief Get device Mem Temperature + * + * @param device device handles + * @param temp pointer to ptTemperature + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetMemTemperature(ptmlDevice_t device, + unsigned int *temp); + +/** + * @brief ptmlDeviceGetGPIOStatus + * + * @param device device handles + * @param GPIO status + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetGPIOStatus(ptmlDevice_t device, + unsigned int *status); +/** + * @brief ptmlDeviceDumpCM3Regs + * + * @param device device handles + * @param addr is CM3 reg addr + * @param len is dump regs' length + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceDumpCM3Regs(ptmlDevice_t device, + unsigned int addr, + unsigned int len, + unsigned int *value); +/** + * @brief ptmlDeviceSetDumpTempSwitch + * + * @param device device handles + * @param switch + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceSetDumpTempSwitch(ptmlDevice_t device, + unsigned int switchFlag); + +int PTML_API ptmlDeviceGetBarInfo(ptmlDevice_t device, + unsigned int *size); + +/** + * @brief ptmlDeviceGetCPLDVersion + * + * @param device device handles + * @param large version + * @param small version + * @return ptmlReturn_t + ********************************************/ +ptmlReturn_t PTML_API ptmlDeviceGetCPLDVersion(ptmlDevice_t device, + int *lver, int *sver); + +ptmlReturn_t PTML_API ptmlDeviceGetProcessInfo(ptmlDevice_t device, + int processNum, struct ptProcessInfo *procInfo); +ptmlReturn_t PTML_API ptmlDeviceGetMaxPower(ptmlDevice_t device, + int *maxPower); +ptmlReturn_t PTML_API ptmlDeviceGetException(ptmlDevice_t device, + unsigned int *num); +/** + * @brief ptmlPtlinkEnableAll + * + * @return ptmlReturn_t + */ +ptmlReturn_t PTML_API ptmlPtlinkEnableAll(void); + +/** + * @brief ptmlPtlinkDisableAll + * + * @return ptmlReturn_t + */ +ptmlReturn_t PTML_API ptmlPtlinkDisableAll(void); + +/** + * @brief ptmlPtlinkPortControl + * + * @device device id + * @port port number + * @ ops operation: en, disable ... + * @return ptmlReturn_t + */ +ptmlReturn_t PTML_API ptmlPtlinkPortControl(ptmlDevice_t device, + uint32_t port, + uint32_t ops); + +/** + * @brief ptmlPtlinkPhytopoDetect + * + * @device device id + * @size memory size + * @buffer user buffer + * @return ptmlReturn_t + */ +ptmlReturn_t PTML_API ptmlPtlinkPhytopoDetect(ptmlDevice_t device, + uint32_t size, + void * buffer); + +/** + * @brief ptmlEngineCollAssign + * + * @device device id + * @coll_type collective type + * @buffer user buffer + * @size memory size + * @return ptmlReturn_t + */ +ptmlReturn_t PTML_API ptmlEngineCollAssign(ptmlDevice_t device, + uint32_t coll_type, + void * buffer, + uint32_t size); +/** + * @brief ptmlPtlinkGetConnectRelation + * + * @device1 device id + * @device2 device id + * @status the relationship between devices + * @return ptmlReturn_t + * @note before use this api please use ptmlPtlinkEnableAll + */ +ptmlReturn_t PTML_API ptmlPtlinkGetConnectRelation(ptmlDevice_t device1, + ptmlDevice_t device2, + int *status) ; + +/** + * @brief ptmlGetPtlinkStatus + * + * @device device id + * @port port id + * @status the status of the port + * @return ptmlReturn_t + * @note before use this api please use ptmlPtlinkEnableAll + */ +ptmlReturn_t PTML_API ptmlGetPtlinkStatus(ptmlDevice_t device, + int port, + int *status); + +/** + * @brief ptmlGetPtlinkRemoteDevicePciInfo + * + * @device device id + * @port port id + * @pciInfo pciInfo of remote device + * @return ptmlReturn_t + * @note before use this api please use ptmlPtlinkEnableAll + */ +ptmlReturn_t PTML_API ptmlGetPtlinkRemoteDevicePciInfo(ptmlDevice_t device, + int port, + ptPciInfo_t *pciInfo); + +/** + * @brief ptmlGetErrorCodeToDescription + * + * @errCode errCode + * @return errorDescription + */ +PTML_API const char *ptmlGetErrorCodeToDescription(int errorCode); + +typedef enum ptmlEventType_enum { + PTML_EVENT_TYPE_PSTATE, + PTML_EVENT_TYPE_ALL, +} ptmlEventType_t; + +typedef enum ptmlDeviceStateChange { + PTMLDEVICE_GOOD_TO_BAD, + PTMLDEVICE_BAD_TO_GOOD, +} ptmlDeviceStateChange_t; + +typedef enum ptmlDeviceState { + PTMLDEVICE_BAD, + PTMLDEVICE_GOOD, +} ptmlDeviceState_t; + +typedef enum ptmlEventStrategy { + PTMLEVENT_UN_MONITOR, + PTMLEVENT_MONITOR, +} ptmlEventStrategy_t; + +typedef struct ptmlEventData { + ptmlDevice_t device; + unsigned long eventType; + unsigned long eventData; +} ptmlEventData_t; + +typedef struct ptmlEvent { + ptmlEventStrategy_t strategy; + ptmlDeviceState_t state; +} ptmlEvent_t; + +typedef struct ptmlEventSet { + ptmlEvent_t deviceEvent[TA_PT_NUM_MAX][PTML_EVENT_TYPE_ALL]; +} ptmlEventSet_t; + +ptmlReturn_t PTML_API ptmlEventSetCreate(ptmlEventSet_t **set); +ptmlReturn_t PTML_API ptmlEventSetFree(ptmlEventSet_t *set); +ptmlReturn_t PTML_API ptmlDeviceRegisterEvents(ptmlDevice_t device, + unsigned long eventTypes, + ptmlEventSet_t *set); + +ptmlReturn_t PTML_API ptmlEventSetWait_v2(ptmlEventSet_t * set, + ptmlEventData_t *data, + unsigned int timeoutms); + +typedef int* ptmlDeviceErrorCode; +PTML_API int ptmlGetDeviceLastError(ptmlDevice_t device, + ptmlDeviceErrorCode DeviceErrorCode); + +#ifdef __cplusplus +} +#endif //! __cplusplus + +#endif //! _S2_ML_H_ diff --git a/third_party/sunrise/backend/include/tang.h b/third_party/sunrise/backend/include/tang.h new file mode 100755 index 0000000000..6037823193 --- /dev/null +++ b/third_party/sunrise/backend/include/tang.h @@ -0,0 +1,2321 @@ +//////////////////////////////////////////////////////// +// @file tang.h +// tang DRIVER INTERFACE +// @author linan +//////////////////////////////////////////////////////// + +#ifndef _TANG_H_ +#define _TANG_H_ +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif // __cplusplus + +#define TA_VERSION_MAJOR 0 +#define TA_VERSION_MINOR 13 +#define TA_VERSION_PATCH 0 + +#define TA_VERSION \ + ((TA_VERSION_MAJOR * 1000) + (TA_VERSION_MINOR * 10) + TA_VERSION_PATCH) + +#if defined(_MSC_VER) +#define TA_DEPRECATED __declspec(deprecated) +#define TA_API_EXPORT __declspec(dllexport) +#define TA_API_IMPORT __declspec(dllimport) +#elif defined(__GNUC__) || defined(__clang__) +#define TA_DEPRECATED __attribute__((deprecated)) +#define TA_API_EXPORT __attribute__((visibility("default"))) +#define TA_API_IMPORT __attribute__((visibility("default"))) +#else +#define TA_DEPRECATED +#define TA_API_EXPORT +#define TA_API_IMPORT +#endif //! UNKNOWN COMPILER + +#if defined(tang_EXPORTS) +#define TA_API TA_API_EXPORT +#else +#define TA_API TA_API_IMPORT +#endif //! For user + +#define COMMAND_MAGIC 0x100d58ba +enum COMMAND_SA { + COMMAND_ASYNC = 0, + COMMAND_SYNC, +}; + +enum MODE_TYPE { + PT_MGR_TYPE = 1, + PT_COLL_TYPE, + OTHER_TYPE, +}; + +enum OPERATIONS { + OPS_LINK_EN = 1, + OPS_LINK_DIS, + OPS_LINK_DETECT, + OPS_LINK_PORTADDR = 4, + OPS_LINK_BIF, + OPS_LINK_P2P_ATTR, + OPS_PEER_ACCESS_CAN = 7, + OPS_PEER_ACCESS_EN, + OPS_PEER_ACCESS_DIS, + OPS_LINK_INIT, + OPS_LINK_PORT_INIT, +}; + +enum COLL_OPS { + COLL_BROADCAST = 0, + COLL_REDUCE, + COLL_ALLGATHER, + COLL_REDUCESCATTER, + COLL_ALLREDUCE, + COLL_MAX_OPS +}; + +enum COMMANDS { + CMD_MAGIC = 0, + CMD_SYNC, + MODE_TYPE, + CMD_ID, + PORT, + DEV_ID, + RDEV_ID, + MSGIN_LEN, + MSGOUT_LEN, +}; + +struct scp_msg_ack { + int retval; // FW irq return value + int status; // simple status value + char payload[0]; // complex struct return +}; + +#define C2CSCP_MSG_HEAD (8) + +typedef uint64_t TAdeviceptr; //!< TANG device pointer + +#define TAdevice_nullptr (TAdeviceptr)0 + +typedef struct TAdevice_s* TAdevice; //!< TANG device +typedef struct TActx_s* TAcontext; //!< TANG context +typedef struct TAfunc_s* TAfunction; //!< TANG function handle +typedef struct TAevent_s* TAevent; //!< TANG event handle +typedef struct TAstream_s* TAstream; //!< TANG stream handle +typedef struct TAmodule_s* TAmodule; //!< TANG module handle +typedef struct TAvariable_s* TAvariable; //!< TANG variable +typedef struct TAgraph_s* TAgraph; //!< TANG graph handle +typedef struct TAgraphExec_s* TAgraphExec; //!< TANG graph exec handle +typedef struct TAgraphNode_s* TAgraphNode; //!< TANG graph node + +typedef struct TAdsoWrapper_s { + uintptr_t data[20]; +} TAdsoWrapper_t; + +/** + * @brief Stream flags. + * @sa __s2StreamFlags. + */ +typedef enum TAstream_flags_e { + TA_STREAM_DEFAULT = 0x0, //!< The default stream creation flag. + TA_STREAM_NON_BLOCKING = 0x1, //!< The non blocking stream creation flag. + //! TA_STREAM_LEGACY = 0x2, //!< The legacy stream creation flag. + //! //!< This flag can only be used internally. + //! //!< User use this flag will cause + //! //!< ::TANG_ERROR_INVALID_VALUE error. +} TAstream_flags; + +typedef enum TAstreamCaptureMode_e { + TA_STREAM_CAPTURE_MODE_GLOBAL = 0, + TA_STREAM_CAPTURE_MODE_THREAD_LOCAL = 1, + TA_STREAM_CAPTURE_MODE_RELAXED = 2, +} TAstreamCaptureMode; + +typedef enum TAstreamCaptureStatus_e { + TA_STREAM_CAPTURE_STATUS_NONE = 0, + TA_STREAM_CAPTURE_STATUS_ACTIVE = 1, + TA_STREAM_CAPTURE_STATUS_INVALIDATED = 2, +} TAstreamCaptureStatus; + +typedef struct TAgraphInfo_s { + int nr_nodes; +} TAgraphInfo; + +typedef struct TAeventTimestamp_s { + uint64_t comp; + uint64_t comp_sw; + uint64_t create; + uint64_t enqueue; + uint64_t writeq_beg; + uint64_t writeq_end; +} TAeventTimestamp; + +typedef enum TAevent_record_flags_e { + //!< The default record flag + TA_EVENT_RECORD_DEFAULT = 0, + + //!< Require hardware event + TA_EVENT_RECORD_HW = 0x0100, + + //!< Require software event + TA_EVENT_RECORD_SW = 0x0200, + + //!< Allow waiting while allocating hardware event. + TA_EVENT_RECORD_ALLOW_BLOCKING = 0x0400, +} TAevent_record_flags; + +typedef enum TAevent_flags_e { + TA_EVENT_DISABLE_TIMING = 0x02, + TA_EVENT_INTERPROCESS = 0x04, +} TAevent_flags; + +typedef enum TAevent_sync_flags_e { + //!< The default synchronization behaviour. + //!< 1. If the event has not been recorded, + //!< taEventSynchronize will return imediately. + //!< 2. If the event has been recorded, + //!< taEventSynchronize will block until the + //!< event is done. + TA_EVENT_SYNC_DEFAULT = 0x00, + + //!< Block until the event is recorded and done. + TA_EVENT_SYNC_RECORDED_AND_DONE = 0x01, +} TAevent_sync_flags; + +#define TA_IPC_HANDLE_SIZE 64U + +struct TAipcMemHandle_s { + unsigned long reserved[TA_IPC_HANDLE_SIZE / sizeof(unsigned long)]; +}; +typedef struct TAipcMemHandle_s TAipcMemHandle; + +struct TAipcEventHandle_s { + unsigned long reserved[TA_IPC_HANDLE_SIZE / sizeof(unsigned long)]; +}; +typedef struct TAipcEventHandle_s TAipcEventHandle; + +enum TAipcMem_flags_e { + TA_IPC_MEM_LAZY_ENABLE_PEER_ACCESS = 0x01, +}; +typedef enum TAipcMem_flags_e TAipcMem_flags; + +#define TA_LAUNCH_PARAM_END 0 +#define TA_LAUNCH_PARAM_INVALIDATE_L1P5 1 +#define TA_LAUNCH_PARAM_ICACHE_FLUSH 2 +#define TA_LAUNCH_PARAM_WORK_MODE 3 +#define TA_LAUNCH_PARAM_MAX_ACTIVE_BLOCK_COUNT_PER_CU 4 +#define TA_LAUNCH_PARAM_SHARE_MEM_MIRROR 5 +#define TA_LAUNCH_PARAM_CLST_DIMX 6 +#define TA_LAUNCH_PARAM_CLST_DIMY 7 +#define TA_LAUNCH_PARAM_CLST_DIMZ 8 + +struct TAextraLaunchParam_s { + unsigned long type; + union { + unsigned long val; + void *ptr; + }; +}; +typedef struct TAextraLaunchParam_s TAextraLaunchParam; + +typedef enum TAmemorytype_e { + TA_MEMORYTYPE_HOST = 0x01, + TA_MEMORYTYPE_DEVICE = 0x02, + TA_MEMORYTYPE_ARRAY = 0x04, + TA_MEMORYTYPE_UNIFIED = 0x05, +} TAmemorytype; + +typedef enum TApointer_attribute_e { + /**< The ::TAcontext on which a pointer is allocated and registered */ + TA_POINTER_ATTRIBUTE_CONTEXT = 1, + + /**< The ::TAmemorytype describing the physical location of a pointer */ + TA_POINTER_ATTRIBUTE_MEMORY_TYPE = 2, + + TA_POINTER_ATTRIBUTE_DEVICE_POINTER = 3, + TA_POINTER_ATTRIBUTE_HOST_POINTER = 4, + TA_POINTER_ATTRIBUTE_DEVICE_ORDINAL = 9, +} TApointer_attribute; + +typedef enum TAfunction_attribute_e { + // The maximum number of threads per block, beyond which + // a lanuch of the function would fail. + // This value depends on the function and the device. + TA_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 0, + + // The number of bytes statically allocated shared memory. + TA_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES = 1, + + // The number of bytes of user allocated constant memory. + // This attribute is not implemented in pt200 + TA_FUNC_ATTRIBUTE_CONST_SIZE_BYTES = 2, + + // The number of bytes of local memory used by each thread of the function. + TA_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES = 3, + + // The number of registers used by each thread of this function. + TA_FUNC_ATTRIBUTE_NUM_REGS = 4, + + // The maximum size of dynamically allocated shared memory that + // can be used by this function. + TA_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES = 8, +} TAfunction_attribute; + +typedef enum TAgraphNodeType_enum { + TA_GRAPH_NODE_TYPE_KERNEL = 0, + TA_GRAPH_NODE_TYPE_MEMCPY = 1, + TA_GRAPH_NODE_TYPE_MEMSET = 2, + TA_GRAPH_NODE_TYPE_HOST = 3, + TA_GRAPH_NODE_TYPE_GRAPH = 4, + TA_GRAPH_NODE_TYPE_EMPTY = 5, + TA_GRAPH_NODE_TYPE_WAIT_EVENT = 6, + TA_GRAPH_NODE_TYPE_EVENT_RECORD = 7, +} TAgraphNodeType; + +typedef enum TAmoduleSymbolType_enum { + TA_MODULE_SYMBOL_FUNCTION = 1, + TA_MODULE_SYMBOL_VARIABLE = 2, +} TAmoduleSymbolType; + +typedef struct TAmoduleSymbolHandle_s { + union { + TAfunction function; + TAvariable variable; + }; +#ifdef __cplusplus + TAmoduleSymbolHandle_s() + : function(nullptr) {} + + explicit TAmoduleSymbolHandle_s(TAfunction func) + : function(func) {} + + explicit TAmoduleSymbolHandle_s(TAvariable var) + : variable(var) {} +#endif // __cplusplus +} TAmoduleSymbolHandle; + +/** + * @brief TAmodule symbols iteration call back function type. + * @note If the function returns true, the iteration will stop. + * Always returns false to iterate all symbols. + */ +// typedef bool (*TAmoduleSymbolIterateFn)(TAmoduleSymbolType symbolType, +// const char* symbolName, +// TAmoduleSymbolHandle symbolHandle, +// void* userData); + +typedef void (*TAhostFn)(void* userData); + +typedef struct TANG_HOST_NODE_PARAMS { + TAhostFn fn; + void* userData; +} TANG_HOST_NODE_PARAMS; + +typedef struct TANG_KERNEL_NODE_PARAMS_s { + TAfunction func; + + unsigned int gridDimX; + unsigned int gridDimY; + unsigned int gridDimZ; + unsigned int blockDimX; + unsigned int blockDimY; + unsigned int blockDimZ; + unsigned int sharedMemBytes; + + void** kernelParams; + void** extra; +} TANG_KERNEL_NODE_PARAMS; + +#define L2_CACHE_MP_CNT 4 +#define L2_CACHE_MPX_SUBCNT 16 +#define WARP_SIZE 32 +#define BUFFER_SIZE 32 +#define IPATH_PERF_CNT 10 +#define EXT_PERF_CNT 13 +#define SHM_PERF_CNT 15 +#define CLUSTER_CNT 12 +#define SUBCORE_CNT 8 +#define CLUSTER_CNT_HALF 6 +#define CLUSTER_CNT_3 3 + +typedef struct ptProfileInfo_s { + uint32_t blockDim_x; + uint32_t blockDim_y; + uint32_t blockDim_z; + uint32_t gridDim_x; + uint32_t gridDim_y; + uint32_t gridDim_z; + int __s2Stream_t; + int __s2Context_t; + int regs; + int device; + int SSMem; + int DSMem; + uint32_t max_bkcnt; + uint32_t thread_cnt; + float waves_per_sm; + uint32_t max_warps_per_sched; + float perf_blk_occupation; + uint32_t achieved_active_warps_per_sm; + uint32_t blk_limit_reg; + uint32_t blk_limit_shared_mem; + uint32_t local_memory_size; + uint32_t blk_shm_size; + char funcName[512]; + uint64_t knlTime; + uint32_t knlSubTime[96]; + uint64_t ipathcnt[38]; //!< ipath-10 ext-13 shm-15 perf + uint64_t l2cache[6]; //!< [1]Monitor Read Hit Counter [Offset: 0x610] + //!< [2]Monitor Cacheable Read Request Counter [Offset: 0x608] + //!< [3]Monitor Read Request Counter [Offset: 0x600] + //!< [4]Monitor Cacheable Write Request Counter[Offset: 0x60c] + //!< [5]Monitor Write Request Counter [Offset: 0x604] + //!< [6]Monitor Write Hit Counter [Offset: 0x614] + uint64_t l1p5cache[6]; //!< [1]Monitor Read Hit Counter [Offset: 0x610] + //!< [2]Monitor Cacheable Read Request Counter [Offset: 0x608] + //!< [3]Monitor Read Request Counter [Offset: 0x600] + //!< [4]Monitor Cacheable Write Request Counter[Offset: 0x60c] + //!< [5]Monitor Write Request Counter [Offset: 0x604] + //!< [6]Monitor Write Hit Counter [Offset: 0x614] + uint32_t clock; + uint32_t knlInfo[3]; //!< blk_cnt\knl_rcvd_cnt\knl_cmpl_cnt + uint32_t l2HitRate[L2_CACHE_MP_CNT*L2_CACHE_MPX_SUBCNT][4]; + uint32_t l1p5HitRate[L2_CACHE_MPX_SUBCNT*(CLUSTER_CNT)][4]; + uint64_t extDetail[CLUSTER_CNT*SUBCORE_CNT][EXT_PERF_CNT]; + uint64_t shmDetail[CLUSTER_CNT*SUBCORE_CNT][SHM_PERF_CNT]; + float perf_achieved_warp_occupation; + uint32_t l1invalid; + uint32_t warp_regfile_size; +} ptProfileInfo; + +#ifndef TA_STREAM_LEGACY +#define TA_STREAM_LEGACY ((TAstream)0x01) +#endif //! TA_STREAM_LEGACY + +#ifndef TA_STREAM_PER_THREAD +#define TA_STREAM_PER_THREAD ((TAstream)0x02) +#endif //! TA_STREAM_PER_THREAD + +/** + * If set, host memory is page locked. + * Flag for ::taMemHostAlloc() + */ +#define TA_MEMHOSTALLOC_DEFAULT 0x00 + +/** + * If set, host memory is portable between TANG contexts. + * Flag for ::taMemHostAlloc() + */ +#define TA_MEMHOSTALLOC_PORTABLE 0x01 + +/** + * If set, host memory is mapped into TANG address space and + * ::taMemHostGetDevicePointer() may be called on the host pointer. + * Flag for ::taMemHostAlloc() + */ +#define TA_MEMHOSTALLOC_DEVICEMAP 0x02 + +/** + * If set, host memory is allocated as write-combined - fast to write, + * faster to DMA, slow to read except via SSE4 streaming load instruction + * (MOVNTDQA). + * Flag for ::taMemHostAlloc() + */ +#define TA_MEMHOSTALLOC_WRITECOMBINED 0x04 + +/** + * If set allocate memory from device side and map it to the user space. + * + */ +#define TA_MEMHOSTALLOC_MAP_DEVICE_MEMORY 0x100 + +/** + * If set, host memory is page locked. + * Flag for ::taMemHostRegister() + */ +#define TA_MEMHOSTREGISTER_DEFAULT 0x00 + +/** + * If set, host memory is portable between TANG contexts. + * Flag for ::taMemHostRegister() + */ +#define TA_MEMHOSTREGISTER_PORTABLE 0x01 + +/** + * If set, host memory is mapped into TANG address space and + * ::taMemHostGetDevicePointer() may be called on the host pointer. + * Flag for ::taMemHostRegister() + */ +#define TA_MEMHOSTREGISTER_DEVICEMAP 0x02 + +/** + * If set, the passed memory pointer is treated as pointing to some + * memory-mapped I/O space, e.g. belonging to a third-party PCIe device. + * Flag for ::taMemHostRegister() + */ +#define TA_MEMHOSTREGISTER_IOMEMORY 0x04 + +/** + * If set, the passed memory pointer is treated as pointing to memory that is + * considered read-only by the device. + * Flag for ::taMemHostRegister() + */ +#define TA_MEMHOSTREGISTER_READ_ONLY 0x08 + +/** + * @ingroup PT_ERROR error handling + * @{ + ************************************************/ +/** + * @brief Driver API error codes. + ************************************************/ +typedef enum TAresult_e { + /** + * @brief The API call returned with no errors. + * @note For asynchronous operations, \c TANG_SUCCESS + * just means the operation is ququed on the \c stream + * successfully. + */ + TANG_SUCCESS = 0, + + /** + * @brief This indicates one or more invalid parameters + * are passed to the API call. + */ + TANG_ERROR_INVALID_VALUE = 1, + + /** + * @brief This indicates the API call failed because + * it can not allocate enough memory to perform the requested + * operation. + */ + TANG_ERROR_OUT_OF_MEMORY = 2, + + /** + * @brief This indicates that the PT dirver has not been initialized + * with ::__taInit or thar initialization has failed. + */ + TANG_ERROR_NOT_INITIALIZED = 3, + + /** + * @brief This indicates that the PT driver is int the process of shutting + * down. + */ + TANG_ERROR_DEINITIALIZED = 4, + + //!< The device is remove for some reason. + //!< echo "1" > /sys/../remove + TANG_ERROR_DEVICE_REMOVED = 5, + + //!< The device is reseted. + //!< Example: enable or disable SR-IOV + TANG_ERROR_DEVICE_RESET = 6, + + //!< The operation is not allowed. + TANG_ERROR_NOT_PERMITTED = 7, + + //!< No such file or directroy + TANG_ERROR_NO_SUCH_FILE = 8, + + /** + * This indicates that a kernel launch is requesting resources that can + * never be satisfied by the current device. + */ + TANG_ERROR_INVALID_CONFIGURATION = 9, + + //!< Null pointer is passed as argument but it is not allowed. + TANG_ERROR_NULL_POINTER = 10, + + //!< The kernel mode driver is not compatible with current runtime. + TANG_ERROR_INCOMPATIBLE_DRIVER = 11, + + //!< Can allocate enough resources to perform the requested operation. + TANG_ERROR_OUT_OF_RESOURCES = 12, + + TANG_ERROR_TIMEOUT = 13, + + /** + * @brief This indicates the API call is not implemented + * and just a stub or for the given parameter(s) the function + * has not been implemented yet. + */ + TANG_ERROR_NOT_IMPLEMENTED = 99, + + /** + * @brief No available PT devices. + */ + TANG_ERROR_NO_DEVICE = 100, + + /** + * @brief Invalid device. + */ + TANG_ERROR_INVALID_DEVICE = 101, + + //!< Bad file descriptor. + TANG_ERROR_BAD_FD = 102, + + //!< Normal indicate some invariant are broken + TANG_ERROR_UNREACHABLE_CODE = 103, + + //!< More than one function use the same symbol name. + TANG_ERROR_DUPLICATE_FUNC_NAME = 198, + + //!< More than one global value use the same symbol name. + TANG_ERROR_DUPLICATE_VAR_NAME = 199, + + /** + * @brief + */ + TANG_ERROR_INVALID_IMAGE = 200, + + /** + * @brief This most frequently indicates there is + * no context bound to the current thread. + * This error code is also returned when an invalid + * context is passed to API call. + */ + TANG_ERROR_INVALID_CONTEXT = 201, + + /** + * @brief No context is bound to the calling thread. + */ + TANG_ERROR_NO_CONTEXT_BOUND = 202, + + /** + * @brief Invalid host address encountered. + */ + TANG_ERROR_ILLEGAL_HOST_ADDRESS = 203, + + //!< Context mismatch + TANG_ERROR_CONTEXT_MISMATCH = 204, + + /** + * This indicates that the ::TAlimit passed to the API call is not + * supported by the active device. + */ + TANG_ERROR_UNSUPPORTED_LIMIT = 215, + + /** + * @brief The key is not found! + * + */ + TANG_ERROR_NOT_FOUND = 301, + + // This indicates that a resource required by the API call + // is not in a valid state to perform the requested operation. + TANG_ERROR_ILLEGAL_STATE = 302, + + // This error indicates that the operation is not permitted + // then the stream is capturing. + TANG_ERROR_STREAM_CAPTURE_UNSUPPORTED = 303, + + // This error indicates that the current capture + // sequence on the stream has been invalidated + // due a previous error. + TANG_ERROR_STREAM_CAPTURE_INVALIDATED = 304, + + // This error indicates that the operation whould + // have resulted in a merge of two independent capture + // sequences. + TANG_ERROR_STREAM_CAPTURE_MERGE = 305, + + // This error indicates that the capture was not initiated in this stream. + TANG_ERROR_STREAM_CAPTURE_UNMATCHED = 306, + + // This error indicates that the capture sequence contains a fork that was + // not joined to the primary stream. + TANG_ERROR_STREAM_CAPTURE_UNJOINED = 307, + + // This error indicates that a dependency would have been created which + // crossed the capture sequence boundary. Only implicit in-stream ordering + // dependencies are allowed to cross the boundary. + TANG_ERROR_STREAM_CAPTURE_ISOLATION = 308, + + // This error indicates a disallowed implicit dependency on a current + // capture sequence from TA_STREAM_LEGACY. + TANG_ERROR_STREAM_CAPTURE_IMPLICIT = 309, + + // A stream capture sequence not initiated with the + // ::TA_STREAM_CAPTURE_MODE_RELAXED argument to taStreamBeginCapture was + // passed to ::cuStreamEndCapture in a different thread. + TANG_ERROR_STREAM_CAPTURE_WRONG_THREAD = 310, + + // This error indicates that the operation is not permitted on an event + // which was last recorded in a capturing stream. + TANG_ERROR_CAPTURED_EVENT = 311, + + /** + * @brief This indicates an invalid resource handle + * passed to a API call. + * In general, resource handles are opaque type like + * ::TAstream and ::TAcontext. + */ + TANG_ERROR_INVALID_HANDLE = 400, + + /** + * @brief This error code indicates asynchronous operations issued previously + * have not been completed yet. + */ + TANG_ERROR_NOT_READY = 600, + + /** + * @brief A load or store instruction on an invalid + * memory address occured when the device executing a + * kernel. + * This error makes the process is an inconsitant state. + * The process should be terminated and relanuched. + */ + TANG_ERROR_ILLEGAL_ADDRESS = 700, + + /** + * @brief resouce is not enougn for the kernel + * + */ + TANG_ERROR_LAUNCH_OUT_OF_RESOURCES = 701, + + /** + * @brief lanch kernel timeout + * + */ + TANG_ERROR_LAUNCH_TIMEOUT = 702, + + TANG_ERROR_PEER_ACCESS_ALREADY_ENABLED = 704, + TANG_ERROR_PEER_ACCESS_NOT_ENABLED = 705, + + /** + * @brief The premary for a context has been + * initialized. + */ + TANG_ERROR_PRIMARY_CONTEXT_ACTIVE = 708, + + /** + * @brief + * + */ + TANG_ERROR_NOT_SUPPORTED = 801, + + /** + * @brief This indicates that an unknown internal error has occurred. + */ + TANG_ERROR_UNKNOWN = 999, + + /** + * @brief context is destroyed or in destroying in kernel + * + */ + TANG_ERROR_CONTEXT_IS_DESTROYED = 3000, + + /** + * @brief context is not valid in kernel + * + */ + TANG_ERROR_CONTEXT_INVALID = 3001, + + /** + * @brief stream is destroyed or in destroying in kernel + * + */ + TANG_ERROR_STREAM_IS_DESTROYED = 3002, + + /** + * @brief stream is not valid in kernel + * + */ + TANG_ERROR_STREAM_INVALID = 3003, + + /** + * @brief event is destroyed or in destroying in kernel + * + */ + TANG_ERROR_EVENT_IS_DESTROYED = 3004, + + /** + * @brief event is not valid in kernel + * + */ + TANG_ERROR_EVENT_INVALID = 3005, + + /** + * @brief device memory is not enough for current operation + * + */ + TANG_ERROR_DEVICE_OUT_OF_MEMORY = 3006, + + /** + * @brief device memory is not found + * + */ + TANG_ERROR_DEVICE_MEMORY_NOT_FOUND = 3007, + + /** + * @brief pcie fatal error occured + * + */ + TANG_ERROR_PCIE_FATAL = 3012, + + /** + * @brief pcie non-fatal unrecovered error occured + * + */ + TANG_ERROR_PCIE_NON_FATAL_UNRECOVERED = 3013, + + /** + * @brief no more event exist + * + */ + TANG_ERROR_SCP_EVENT_NOT_EXIST = 3014, + + /** + * @brief record event failed + * + */ + TANG_ERROR_SCP_EVENT_RECORD_FAILED = 3015, + + /** + * @brief scp packet crc check failed + * + */ + TANG_ERROR_SCP_PACKET_CRC_FAILED = 3016, + + /** + * @brief scp dispatch send failed + * + */ + TANG_ERROR_SCP_DISP_SEND_FAILED = 3017, + + /** + * @brief sq write sequence error + * + */ + TANG_ERROR_SCP_SQ_WRITE_INVALID = 3018, + + /** + * @brief udrc pcie xdma packet invalid + * + */ + TANG_ERROR_UDRC_PCIE_DMA_CMD_PACKET_INVALID = 3019, + + /** + * @brief udrc mp dma packet invalid + * + */ + TANG_ERROR_UDRC_MP_DMA_CMD_PACKET_INVALID = 3020, + + /** + * @brief udrc reg packet invalid + * + */ + TANG_ERROR_UDRC_REG_CMD_PACKET_INVALID = 3021, + + /** + * @brief udrc reg access invalid + * + */ + TANG_ERROR_UDRC_REG_ACCESS_INVALID = 3022, + + /** + * @brief aiss cluster is not configured + * + */ + TANG_ERROR_AISS_VF_CTRL_CLUST_USR_NOT_ALLOCATED = 3023, + + /** + * @brief barrier is destroyed or in destroying in kernel + * + */ + TANG_ERROR_BARRIER_IS_DESTROYED = 3024, + + /** + * @brief barrier is not valid in kernel + * + */ + TANG_ERROR_BARRIER_INVALID = 3025, + + /** + * @brief one obj is destroyed or in destroying in kernel + * + */ + TANG_ERROR_IS_DESTROYED = 3026, + + /** + * @brief xdma C2H align mismath + * + */ + TANG_ERROR_XDMA_C2H_ALIGN_MISMATCH = 3300, + + /** + * @brief xdma C2H invalid magic stopped + * + */ + TANG_ERROR_XDMA_C2H_INVALID_MAGIC_STOPPED = 3301, + + /** + * @brief xdma C2H invalid Len + * + */ + TANG_ERROR_XDMA_C2H_INVALID_LEN = 3302, + + /** + * @brief xdma C2H decode error + * + */ + TANG_ERROR_XDMA_C2H_DECODE = 3303, + + /** + * @brief xdma C2H slave + * + */ + TANG_ERROR_XDMA_C2H_SLAVE = 3304, + + /** + * @brief xdma C2H descriptor unsupport request + * + */ + TANG_ERROR_XDMA_C2H_DESC_UNSUPPORT_REQUEST = 3305, + + /** + * @brief xdma C2H descriptor completer abort + * + */ + TANG_ERROR_XDMA_C2H_DESC_COMPLETER_ABORT = 3306, + + /** + * @brief xdma C2H descriptor parity + * + */ + TANG_ERROR_XDMA_C2H_DESC_PARITY = 3307, + + /** + * @brief xdma C2H descriptor header ep + * + */ + TANG_ERROR_XDMA_C2H_DESC_HEADER_EP = 3308, + + /** + * @brief xdma C2H descriptor unexpected comp + * + */ + TANG_ERROR_XDMA_C2H_DESC_UNEXPECTED_COMP = 3309, + + /** + * @brief xdma C2H timeout + * + */ + TANG_ERROR_XDMA_C2H_TIMEOUT = 3310, + + /** + * @brief xdma C2H unknown + * + */ + TANG_ERROR_XDMA_C2H_UNKNOWN = 3311, + + /** + * @brief xdma H2C align mismatch + * + */ + TANG_ERROR_XDMA_H2C_ALIGN_MISMATCH = 3350, + + /** + * @brief xdma H2C invalid magic stopped + * + */ + TANG_ERROR_XDMA_H2C_INVALID_MAGIC_STOPPED = 3351, + + /** + * @brief xdma H2C invalid len + * + */ + TANG_ERROR_XDMA_H2C_INVALID_LEN = 3352, + + /** + * @brief xdma H2C read unsupport request + * + */ + TANG_ERROR_XDMA_H2C_READ_UNSUPPORT_REQUEST = 3353, + + /** + * @brief xdma H2C read completer abort + * + */ + TANG_ERROR_XDMA_H2C_READ_COMPLETER_ABORT = 3354, + + /** + * @brief xdma H2C read parity + * + */ + TANG_ERROR_XDMA_H2C_READ_PARITY = 3355, + + /** + * @brief xdma H2C read header ep + * + */ + TANG_ERROR_XDMA_H2C_READ_HEADER_EP = 3356, + + /** + * @brief xdma H2C read unexpected comp + * + */ + TANG_ERROR_XDMA_H2C_READ_UNEXPECTED_COMP = 3357, + + /** + * @brief xdma H2C decode error + * + */ + TANG_ERROR_XDMA_H2C_DECODE = 3358, + + /** + * @brief xdma H2C slave + * + */ + TANG_ERROR_XDMA_H2C_SLAVE = 3359, + + /** + * @brief xdma H2C descriptor unsupport request + * + */ + TANG_ERROR_XDMA_H2C_DESC_UNSUPPORT_REQUEST = 3360, + + /** + * @brief xdma H2C descriptor completer abort + * + */ + TANG_ERROR_XDMA_H2C_DESC_COMPLETER_ABORT = 3361, + + /** + * @brief xdma H2C descriptor parity + * + */ + TANG_ERROR_XDMA_H2C_DESC_PARITY = 3362, + + /** + * @brief xdma H2C descriptor header ep + * + */ + TANG_ERROR_XDMA_H2C_DESC_HEADER_EP = 3363, + + /** + * @brief xdma H2C descriptor unexpected com + * + */ + TANG_ERROR_XDMA_H2C_DESC_UNEXPECTED_COMP = 3364, + + /** + * @brief xdma H2C timeout + * + */ + TANG_ERROR_XDMA_H2C_TIMEOUT = 3365, + + /** + * @brief xdma H2C unknown + * + */ + TANG_ERROR_XDMA_H2C_UNKNOWN = 3366, + /** + * @brief gpu profling share mem out of size + * + */ + TANG_ERROR_GPU_PROFLING_SHAMEM_OUT_OF_SIZE = 3367, + + /** + * @brief The requested ipc mem is destroied. + * @sa taIpcOpenMemHandle + */ + TANG_ERROR_IPC_MEM_DESTROIED = 3368, +} TAresult; + +/** + * @brief Get the TANG SDK Runtime version. + * + * @param runtimeVersion - Returned runtime version number. + * @return int + * ::TANG_SUCCESS - \p runtimeVersion is a non-null poiner. + * ::TANG_ERROR_INVALID_VALUE - \p runtimeVersion is a null pointer. + * @deleted Do not use this function. + ******************************************************/ +TAresult TA_API taRuntimeGetVersion(int* runtimeVersion); + +/** + * @brief Get the TANG SDK Driver version. + * + * @param driverVersion - Returned driver version number. + * @return int + * ::TANG_SUCCESS - \p driverVersion is a non-null poiner. + * ::TANG_ERROR_INVALID_VALUE - \p driverVersion is a null pointer. + ******************************************************/ +TAresult TA_API taDriverGetVersion(int* driverVersion); + +/** + * @brief Get the kernel mode driver version. + * + * @param kernelDriverVersion + * @return int + */ +TAresult TA_API taKernelDriverGetVersion(int* kernelDriverVersion); + +/** + * @brief Get error description. + * + * @param error + * @param ppstr - Returned Null-terminated string. + * @return int + * ::TANG_SUCCESS - \p error is a valid error code. + * ::TANG_ERROR_INVALID_VALUE - \p error is an invalid + * error code. + ******************************************************/ +TAresult TA_API taGetErrorString(TAresult error, char const** ppstr); + +/** + * @brief Get the string representation of an error code. + * + * @param error + * @param ppstr - Returned Null-terminated string. + * @return int + * ::TANG_SUCCESS - \p error is a valid error code. + * ::TANG_ERROR_INVALID_VALUE - \p error is an invalid + * error code. + ******************************************************/ +TAresult TA_API taGetErrorName(TAresult error, char const** ppstr); +/** @} PT_ERROR */ + +/** + * @brief Initilaize driver module. + * + * This function initialize driver module. + * @param flags Initialization flags for driver API + * @return + * ::TANG_SUCCESS + *******************************************************/ +TAresult TA_API taInit(unsigned int flags); +// TAresult TA_API taDeinit(void); + +/** + * @brief __taDeviceGet + * Get a handle to a compute device + * @param device Pointer to a device handle. + * @param ordinal The device number to get handle for + * @return int + ********************************************************/ +TAresult TA_API taDeviceGet(TAdevice* device, int ordinal); + +/** + * @brief Wait for all work completed + * + * @param device + * @return int + ********************************************************/ +TAresult TA_API taDeviceSynchronize(TAdevice device); + +/** + * @brief Synchronize with the current device of the calling thread. + * + * @warning Not a public API, may change in the future. + * @return int + * ::TANG_SUCCESS + */ +TAresult TA_API __taDeviceSynchronizeCurrent(void); + +/** + * @brief Reset the current device. + * Only the calling process is impacted. + * @return int + * @warning This is a dangerous API. The caller must + * ensure that all resources allocated from the current + * device will not be used again. The most difficult to + * handle is TAcontext which may be pushed onto thread's context + * stack. + * BE CAREFUL. + **********************************************************/ +TAresult TA_API taDeviceReset(void); + +/** + * @brief Returns a handle to a compute device. + * + * @param device - device handle + * @param pciBusId - PCI Bus ID + * @return int + * ::TANG_SUCCESS - \p error is a valid error code. + * ::TANG_ERROR_INVALID_DEVICE - \p error is an invalid + * error code. + * ::TANG_ERROR_INVALID_VALUE - \p error is an invalid + * error code. + */ +TAresult TA_API taDeviceGetByPCIBusId(TAdevice* device, const char* pciBusId); + +/** + * @brief Returns a PCI Bus Id string for the device. + * + * @param pciBusId - PCI Bus ID + * @param len - Maximum length of pciBusId name string + * @param device - device handle + * @return int + * ::TANG_SUCCESS - \p error is a valid error code. + * ::TANG_ERROR_INVALID_DEVICE - \p error is an invalid + * error code. + * ::TANG_ERROR_INVALID_VALUE - \p error is an invalid + * error code. + */ +TAresult TA_API taDeviceGetPCIBusId(char* pciBusId, int len, TAdevice device); + +/** + * Limits + */ +typedef enum TAlimit_enum { + TA_LIMIT_STACK_SIZE = 0x00, /**< GPU thread stack size */ + TA_LIMIT_PRINTF_FIFO_SIZE = 0x01, /**< GPU printf FIFO size */ + TA_LIMIT_MALLOC_HEAP_SIZE = 0x02, /**< GPU malloc heap size */ + TA_LIMIT_DEV_RUNTIME_SYNC_DEPTH = + 0x03, /**< GPU device runtime launch synchronize depth */ + TA_LIMIT_DEV_RUNTIME_PENDING_LAUNCH_COUNT = + 0x04, /**< GPU device runtime pending launch count */ + TA_LIMIT_MAX_L2_FETCH_GRANULARITY = + 0x05, /**< A value between 0 and 128 that indicates the maximum fetch + granularity of L2 (in Bytes). This is a hint */ + TA_LIMIT_MAX +} TAlimit; + +/** + * @brief Get Resource limits of current context + * + * @param [out] pValue + * @param [in] limit + * @return int + * ::TANG_SUCCESS - \p error is a valid error code. + * ::TANG_ERROR_INVALID_VALUE - \p error is an invalid + * error code. + * ::TANG_ERROR_UNSUPPORTED_LIMIT - \p error is an invalid + * error code. + * + */ +TAresult TA_API taCtxGetLimit(size_t* pValue, TAlimit limit); +TAresult TA_API taCtxSetLimit(TAlimit limit, size_t value); + +/** + * @brief Query detail limit information. + * Not a plublic interface. + * + * @param context + * @param limit + * @param pCurrent The current value of the \p limit + * @param pLimit The limit of the \p limit + * @return int + * ::TANG_SUCCESS + * ::TANG_ERROR_INVALID_VALUE + */ +TAresult TA_API __taCtxQueryLimit(TAcontext context, + TAlimit limit, + size_t* pCurrent, + size_t* pLimit); + +/** + * TA device attributes + */ +typedef enum taDeviceAttr { + TA_DEV_ATTR_SHARED_MEM_PER_BLOCK = 0, //!< sharedMemPerBlock + TA_DEV_ATTR_REGS_PER_BLOCK, //!< regsPerBlock + TA_DEV_ATTR_WARP_SIZE, //!< warpSize + TA_DEV_ATTR_MEM_PITCH, //!< memPitch + TA_DEV_ATTR_MAX_THREADS_PER_BLOCK, //!< maxThreadsPerBlock + TA_DEV_ATTR_MAX_THREADS_DIM_X, //!< maxThreadsDimX + TA_DEV_ATTR_MAX_THREADS_DIM_Y, //!< maxThreadsDimY + TA_DEV_ATTR_MAX_THREADS_DIM_Z, //!< maxThreadsDimZ + TA_DEV_ATTR_MAX_GRID_SIZE_X, //!< maxGridSizeX + TA_DEV_ATTR_MAX_GRID_SIZE_Y, //!< maxGridSizeY + TA_DEV_ATTR_MAX_GRID_SIZE_Z, //!< maxGridSizeZ + TA_DEV_ATTR_CLOCK_RATE, //!< clockRate + TA_DEV_ATTR_TOTAL_CONST_MEM, //!< totalConstMem + TA_DEV_ATTR_MULTIPROCESSOR_COUNT, //!< multiProcessorCount + TA_DEV_ATTR_MAX_BLOCKS_PER_MULTIPROCESSOR, //!< maxBlocksPerMultiProcessor + TA_DEV_ATTR_ASYNC_ENGINE_COUNT, //!< asyncEngineCount + TA_DEV_ATTR_MEMORY_CLOCK_RATE, //!< memoryClockRate + TA_DEV_ATTR_MEMORY_BUS_WIDTH, //!< memoryBusWidth + TA_DEV_ATTR_L2_CACHE_SIZE, //!< l2CacheSize + TA_DEV_ATTR_MAX_THREADS_PER_MULTIPROCESSOR, //!< maxThreadsPerMultiProcessor + TA_DEV_ATTR_GLOBAL_L1_CACHE_SUPPORTED, //!< globalL1CacheSupported + TA_DEV_ATTR_LOCAL_L1_CACHE_SUPPORTED, //!< localL1CacheSupported + TA_DEV_ATTR_SHARED_MEM_PER_MULTIPROCESSOR, //!< sharedMemPerMultiprocessor + TA_DEV_ATTR_REGS_PER_MULTIPROCESSOR, //!< regsPerMultiprocessor + TA_DEV_ATTR_STREAM_PRIORITIES_SUPPORTED, //!< streamPrioritiesSupported + TA_DEV_ATTR_CONCURRENT_KERNELS, //!< concurrentKernels + TA_DEV_ATTR_COMPUTE_PREEMPTION_SUPPORTED, //!< computePreemptionSupported + TA_DEV_ATTR_KERNEL_EXEC_TIMEOUT_ENABLED, //!< kernelExecTimeoutEnabled + TA_DEV_ATTR_ECC_ENABLED, //!< ECCEnabled + TA_DEV_ATTR_ACCESS_POLICY_MAX_WINDOW_SIZE, //!< accessPolicyMaxWindowSize + TA_DEV_ATTR_TCC_DRIVER, //!< tccDriver + TA_DEV_ATTR_SINGLE_TO_DOUBLE_PRECISION_PER_RATIO, //!< singleToDoublePrecisionPerfRatio + TA_DEV_ATTR_COOPERATIVE_LAUNCH, //!< cooperativeLaunch + TA_DEV_ATTR_COOPERATIVE_MULTI_DEVICE_LAUNCH, //!< cooperativeMultiDeviceLaunch + TA_DEV_ATTR_PERSISTING_L2_CACHE_MAX_SIZE, //!< persistingL2CacheMaxSize + TA_DEV_ATTR_CAN_MAP_HOST_MEMORY, //!< canMapHostMemory + TA_DEV_ATTR_UNIFIED_ADDRESSING, //!< unifiedAddressing + TA_DEV_ATTR_MANAGED_MEMORY, //!< managedMemory + TA_DEV_ATTR_CONCURRENT_MANAGED_ACCESS, //!< concurrentManagedAccess + TA_DEV_ATTR_DIRECT_MANAGED_MEM_ACCESS_FROM_HOST, //!< directManagedMemAccessFromHost + TA_DEV_ATTR_PAGEABLE_MEMORY_ACCESS, //!< pageableMemoryAccess + TA_DEV_ATTR_PAGEABLE_MEMORY_ACCESS_USES_HOST_PAGE_TABLES, //!< pageableMemoryAccessUsesHostPageTables + TA_DEV_ATTR_CAN_USE_HOST_POINTER_FOR_REGISTERED_MEM, //!< canUseHostPointerForRegisteredMem + TA_DEV_ATTR_HOST_NATIVE_ATOMIC_SUPPORTED, //!< hostNativeAtomicSupported + TA_DEV_ATTR_CAN_FLUSH_REMOTE_WRITES, //!< canFlushRemoteWrites + TA_DEV_ATTR_GPU_OVERLAP, //!< gpuOverlap + TA_DEV_ATTR_INTEGRATED, //!< integrated + TA_DEV_ATTR_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, //!< maxSharedMemoryPerBlockOptin + TA_DEV_ATTR_GPU_DIRECT_RDMA_SUPPORTED, //!< gpuDirectRDMASupported + TA_DEV_ATTR_GPU_DIRECT_RDMA_FLUSH_WRITES_OPTIONS, //!< gpuDirectRDMAFlushWritesOptions + TA_DEV_ATTR_GPU_DIRECT_RDMA_WRITES_ORDERING, //!< gpuDirectRDMAWritesOrdering + TA_DEV_ATTR_MAJOR, //!< major + TA_DEV_ATTR_MINOR, //!< minor + TA_DEV_ATTR_PCI_BUS_ID, //!< pciBusID + TA_DEV_ATTR_PCI_DEVICE_ID, //!< pciDeviceID + TA_DEV_ATTR_PCI_DOMAIN_ID, //!< pciDomainID + TA_DEV_ATTR_IS_MULTI_GPU_BOARD, //!< isMultiGpuBoard + TA_DEV_ATTR_GPU_BOARD_GROUP_ID, //!< multiGpuBoardGroupID + TA_DEV_ATTR_COMPUTE_MODE, //!< computeMode + TA_DEV_ATTR_RESERVED_SHARED_MEMORY_PER_BLOCK, //!< reservedSharedMemoryPerBlock + TA_DEV_ATTR_SPARSE_TANG_ARRAY_SUPPORTED, //!< sparseTangArraySupported + TA_DEV_ATTR_HOST_REGISTER_SUPPORTED, //!< hostRegisterSupported + TA_DEV_ATTR_HOST_REGISTER_READ_ONLY_SUPPORTED, //!< hostRegisterReadOnlySupported + TA_DEV_ATTR_MEMORY_POOLS_SUPPORTED, //!< memoryPoolsSupported + TA_DEV_ATTR_MEMORY_POOL_SUPPORTED_HANDLE_TYPES, //!< memoryPoolSupportedHandleTypes + TA_DEV_ATTR_MAX +} taDeviceAttr; + +TAresult TA_API taDeviceGetAttribute(int* value, taDeviceAttr attr, TAdevice dev); + +TAresult TA_API taDeviceGetName(char* name, int len, TAdevice dev); +TAresult TA_API taDeviceGetUuid(char* uuid, TAdevice dev); +TAresult TA_API taDeviceTotalMem(size_t* bytes, TAdevice dev); + +/** + * @brief Get the number of available devices. + * + * @param count Returns the count of available devices. + * @return int + * ::TANG_SUCCESS + * ::TANG_ERROR_NOT_INITIALIZED + ********************************************************/ +TAresult TA_API taDeviceGetCount(int* count); + +/** + * @brief Gets free and total memory of the current device. + * + * @param free - Returned free memory in bytes + * @param total - Returned total memory in bytes + * @return int + */ +TAresult TA_API taMemGetInfo(size_t* free, size_t* total); + +/** + * @brief Gets free and total memory of \p device. + * + * @param device + * @param free + * @param total + * @return int + */ +TAresult TA_API taDeviceMemGetInfo(TAdevice device, size_t* free, size_t* total); + +/** + * @brief Create a context. + * + * @param pctx - Returned newly created context. + * @param flags - Flags for creating the context. + * @param dev - The device ID. + * @return int + * ::TANG_SUCCESS + * ::TANG_ERROR_OUT_OF_MEMORY + * @note When the context is no longer used, the caller + * should call \c taCtxDestroy to destroy the context. + * @sa taCtxDestroy + ********************************************************/ +TAresult TA_API taCtxCreate(TAcontext* pctx, unsigned int flags, TAdevice dev); + +/** + * @brief Destroy the context \p ctx. + * + * @param ctx + * @return int + * ::TANG_SUCCESS + * ::TANG_ERROR_INVALID_CONTEXT + * @note It's the caller's responsibility ensure that the + * context \p ctx is no longer referenced by other objects. + * @taCtxCreate + ********************************************************/ +TAresult TA_API taCtxDestroy(TAcontext ctx); + +/** + * @brief Bind the context \p ctx to the calling thread. + * + * @param ctx - The context to be bound to the calling thread. + * @return int + * @note If \p ctx is NULL, this function just POPs the context + * stack of the calling thread if the context stack is not + * empty. + * If \p ctx is no NULL, this function replace the top of + * the stack if the context stack is no empty. + *********************************************************/ +TAresult TA_API taCtxSetCurrent(TAcontext ctx); + +/** + * @brief Get the current context bound to the calling + * thread. + * + * @param pctx - Returned context bound to the calling thread. + * @return int + * ::TANG_SUCCESS + * ::TANG_ERROR_NO_CONTEXT_BOUND + ********************************************************/ +TAresult TA_API taCtxGetCurrent(TAcontext* pctx); + +/** + * @brief Query the current context bound to the calling + * thread. + * + * @param pCtx + * @return int + */ +TAresult TA_API taCtxQueryCurrent(TAcontext* pCtx); + +/** + * @brief Get the device ID of the current context bound + * to the call thread. + * + * @param dev - Returned device ID of the current context. + * @return int + * ::TANG_SUCCESS + * ::TANG_ERROR_NO_CONTEXT_BOUND + ********************************************************/ +TAresult TA_API taCtxGetCurrentDevice(TAdevice* dev); + +TAresult TA_API taCtxGetDevice(TAcontext ctx, TAdevice* dev); +TAresult TA_API taCtxGetOrdinal(TAcontext ctx, int* ordinal); + +TAresult TA_API taCtxGetFunction(TAcontext ctx, + const void* hostFunc, + TAfunction* phFunc); + +TAresult TA_API taCtxRegisterFunction(TAcontext ctx, + const void* hostFunc, + TAfunction hFunc); + +TAresult TA_API taCtxGetBuiltInFunction(TAcontext ctx, + const char* funcName, + TAfunction* phFunc); + +TAresult TA_API taCtxRegisterBuiltInFunction(TAcontext ctx, + const char* funcName, + TAfunction hFunc); + +TAresult TA_API taCtxGetVariable(TAcontext ctx, + const void* hostVar, + TAvariable* hVar); + +TAresult TA_API taCtxRegisterVariable(TAcontext ctx, + const void* hostVar, + TAvariable hVar); + +TAresult TA_API taCtxPushCurrent(TAcontext ctx); +TAresult TA_API taCtxPopCurrent(TAcontext* ctx); +TAresult TA_API taCtxSynchronize(TAcontext ctx); + +/** + * @ingroup Primary Context Management. + * @{ + */ +/** + * @brief Retain the primary context of \p dev. + * + * @param pctx Pointer to receive the primary context handle. + * @param dev + * @return int + ********************************************************/ +TAresult TA_API taDevicePrimaryCtxRetain(TAcontext* pctx, TAdevice dev); + +/** + * @brief Release the primary context of \p dev. + * + * @param dev + * @return int + ********************************************************/ +TAresult TA_API taDevicePrimaryCtxRelease(TAdevice dev); + +/** + * @brief Reset the primary context of \p dev. + * + * @param dev + * @return int + ********************************************************/ +TAresult TA_API taDevicePrimaryCtxReset(TAdevice dev); + +/** + * @brief Get the state of primary context of device \p dev. + * + * @param dev Device to get primary context's state for. + * @param flags Pointer to receive the flags. + * @param active + * @return int + **********************************************************/ +TAresult TA_API taDevicePrimaryCtxGetState(TAdevice dev, + unsigned int* flags, + int* active); + +/** + * @brief Set flags for the primary context of the device \p dev. + * + * @param dev + * @param flags + * @return int + */ +TAresult TA_API taDevicePrimaryCtxSetFlags(TAdevice dev, unsigned int flags); +/** }@ */ + +TAresult TA_API taDeviceGetOrdinal(TAdevice dev, int* ordinal); + +/** + * @brief Allocate a block of memory in the device. + * + * @param dptr Receives the allocated device memory block handle. + * @param size + * @return int + * On success, zero is returned. + * @sa taMemFree + ***********************************************************/ +TAresult TA_API taMemAlloc(TAdeviceptr* dptr, size_t size); +TAresult TA_API taMemAlloc(TAdeviceptr* dptr, size_t size); +TAresult TA_API taMemAllocAsync(TAdeviceptr* dptr, size_t size, TAstream stream); +TAresult TA_API taMemAllocAsync_ptsz(TAdeviceptr* dptr, + size_t size, + TAstream stream); + +/** + * @brief Free a device memory block. + * + * @param dptr Device memory block handle. + * @return int + * On success, zero is returned. + * @sa taMemAlloc + ************************************************************/ +TAresult TA_API taMemFree(TAdeviceptr dptr); +TAresult TA_API taMemFreeAsync(TAdeviceptr dptr, TAstream hStream); +TAresult TA_API taMemFreeAsync_ptsz(TAdeviceptr dptr, TAstream hStream); + +/** + * @brief Allocate page locked host memory. + * + * @param hptr Pointer to the allocated page locked host memory + * @param sizeBytes Requested memory size + * @return int + * On success, zero is returned. + * @sa taMemFreeHost + ***********************************************************/ +TAresult TA_API taMemAllocHost(void** hptr, size_t sizeBytes); + +/** + * @brief Allocate page locked host memory. + * + * @param hptr Pointer to the allocated page locked host memory + * @param sizeBytes Requested memory size + * @param flags See below. + * flags: + * - #TA_MEMHOSTALLOC_PORTABLE Memory is considered registered by all + *contexts. + * - #TA_MEMHOSTALLOC_DEVICEMAP Map the allocation into the address space + *for the current device. + * - #TA_MEMHOSTALLOC_WRITECOMBINED Allocates the memory as write-combined (WC). + * TANG does not support IOMMU on device side, so flags of + *TA_MEMHOSTALLOC_DEVICEMAP and TA_MEMHOSTALLOC_WRITECOMBINED will always return + *false. + * @return int + * On success, zero is returned. + * @sa taMemFreeHost + ***********************************************************/ +TAresult TA_API taMemHostAlloc(void** hptr, size_t sizeBytes, unsigned int flags); +TAresult TA_API taMemHostGetDevicePointer(TAdeviceptr* pdptr, + void* pHost, + unsigned int flags); + +/** + * @brief Get the flags that were used for allocation. + * + * @param pFlags + * @param p + * @return int + */ +TAresult TA_API taMemHostGetFlags(unsigned int* pFlags, void* p); + +/** + * @brief Free page locked host memory. + * + * @param hptr Pointer to memory to be freed + * @return int + * On success, zero is returned. + * @sa taMemAllocHost, taMemHostAlloc + ************************************************************/ +TAresult TA_API taMemFreeHost(void* hptr); + +/** + * @brief Register host memory as page locked memory. + * + * @param hptr Pointer to host memory to be registered. + * @param sizeBytes Requested memory size + * @param flags See below. + * flags: + * - #TA_MEMHOSTREGISTER_PORTABLE Memory is considered registered by all + *contexts. + * - #TA_MEMHOSTREGISTER_DEVICEMAP Map the allocation into the address space for + *the current device. + * - #TA_MEMHOSTREGISTER_IOMEMORY The passed memory pointer is treated as + *pointing to some memory-mapped I/O space. + * - #TA_MEMHOSTREGISTER_READ_ONLY The passed memory pointer is treated as + *pointing to memory that is considered read-only by the device. TANG does not + *support IOMMU on device side, so flags of TA_MEMHOSTREGISTER_DEVICEMAP and + *TA_MEMHOSTREGISTER_IOMEMORY and TA_MEMHOSTREGISTER_READ_ONLY will always + *return false. + * @return int + * On success, zero is returned. + * @sa taMemHostUnregister + ***********************************************************/ +TAresult TA_API taMemHostRegister(void* hptr, size_t sizeBytes, unsigned int flags); + +/** + * @brief Un-register host pointer + * + * @param hptr Host pointer previously registered + * @return int + * On success, zero is returned. + * @sa taMemHostRegister + ************************************************************/ +TAresult TA_API taMemHostUnregister(void* hptr); + +/** + * @brief Returns information about a pointer. + * + * @param data - Pointer to the returned attribute value. + * @param attr - Pointer attribute to query + * @param ptr - Pointer to be queried. + * @return int + * ::TANG_SUCCESS + * ::TANG_ERROR_INVALID_VALUE + */ +TAresult TA_API taPointerGetAttribute(void* data, + TApointer_attribute attr, + TAdeviceptr ptr); + +TAresult TA_API taMemset(TAdeviceptr dptr, int value, size_t size); +TAresult TA_API taMemset_ptds(TAdeviceptr dptr, int value, size_t size); +TAresult TA_API taMemsetAsync(TAdeviceptr dptr, + int value, + size_t size, + TAstream stream); +TAresult TA_API taMemsetAsync_ptsz(TAdeviceptr dptr, + int value, + size_t size, + TAstream stream); + +/** + * @brief Copy data from host memory to host memory. + * + * @param dstHost - Host destination data address. + * @param srcHost - Host source data address. + * @param size - The size in bytes of data to be copied. + * @return int + * ::TANG_SUCCESS + ************************************************************/ +TAresult TA_API taMemcpyH2H(void* dstHost, const void* srcHost, size_t size); + +/** + * @brief Copy data from host memory to host memory. + * + * @param dstHost - Host destination data address. + * @param srcHost - Host source data address + * @param size The size in bytes of data to be copied. + * @return int + ************************************************************/ +TAresult TA_API taMemcpyH2H_ptds(void* dstHost, + const void* srcHost, + size_t size); + + +/** + * @brief Copy data from host memory to device memory. + * + * @param dstDevice - Device destination data address. + * @param srcHost - Host source data address. + * @param size - The size in bytes of data to be copied. + * @return int + * ::TANG_SUCCESS + ************************************************************/ +TAresult TA_API taMemcpyH2D(TAdeviceptr dstDevice, const void* srcHost, size_t size); + +/** + * @brief Copy data from host memory to device memory. + * + * @param dstDevice Device destination data address + * @param srcHost Host source data address + * @param size The size in bytes of data to be copied. + * @return int + ************************************************************/ +TAresult TA_API taMemcpyH2D_ptds(TAdeviceptr dstDevice, + const void* srcHost, + size_t size); + +/** + * @brief Copy data from device to host. + * + * @param dstHost + * @param srcDevice + * @param size + * @return int + * ::TANG_SUCCESS + *************************************************************/ +TAresult TA_API taMemcpyD2H(void* dstHost, TAdeviceptr srcDevice, size_t size); + +/** + * @brief + * + * @param dstHost + * @param srcDevice + * @param size + * @return int + *************************************************************/ +TAresult TA_API taMemcpyD2H_ptds(void* dstHost, TAdeviceptr srcDevice, size_t size); + +/** + * @brief + * + * @param dstDevice + * @param srcDevice + * @param size + * @return int + * ::TANG_ERROR_NOT_IMPLEMENTED. + *************************************************************/ +TAresult TA_API taMemcpyD2D(TAdeviceptr dstDevice, + TAdeviceptr srcDevice, + size_t size); + +/** + * @brief + * + * @param dstDevice + * @param srcDevice + * @param size + * @return int + *************************************************************/ +TAresult TA_API taMemcpyD2D_ptds(TAdeviceptr dstDevice, + TAdeviceptr srcDevice, + size_t size); + +TAresult TA_API taMemcpyH2HAsync(void* dstHost, + const void* srcHost, + size_t size, + TAstream stream); + +TAresult TA_API taMemcpyH2HAsync_ptsz(void* dstHost, + const void* srcHost, + size_t size, + TAstream stream); + +TAresult TA_API taMemcpyH2DAsync(TAdeviceptr dstDevice, + const void* srcHost, + size_t size, + TAstream stream); + +TAresult TA_API taMemcpyH2DAsync_ptsz(TAdeviceptr dstDevice, + const void* srcHost, + size_t size, + TAstream stream); + +TAresult TA_API taMemcpyD2HAsync(void* dstHost, + TAdeviceptr srcDevice, + size_t size, + TAstream stream); + +TAresult TA_API taMemcpyD2HAsync_ptsz(void* dstHost, + TAdeviceptr srcDevice, + size_t size, + TAstream stream); + +TAresult TA_API taMemcpyD2DAsync(TAdeviceptr dstDevice, + TAdeviceptr srcDevice, + size_t size, + TAstream stream); + +TAresult TA_API taMemcpyD2DAsync_ptsz(TAdeviceptr dstDevice, + TAdeviceptr srcDevice, + size_t size, + TAstream stream); + +TAresult TA_API taStreamCreate(TAstream* pStream, unsigned int flags); +TAresult TA_API taStreamCreateWithPriority(TAstream* pstream, + unsigned int flags, + int priority); +TAresult TA_API taStreamGetPriority(TAstream hstream, int* priority); +TAresult TA_API taStreamGetPriority_ptsz(TAstream hstream, int* priority); +TAresult TA_API taStreamGetFlags(TAstream hstream, unsigned int* priority); +TAresult TA_API taStreamGetFlags_ptsz(TAstream hstream, unsigned int* priority); +TAresult TA_API taStreamGetId(TAstream hstream, int* pId); +TAresult TA_API taStreamGetId_ptsz(TAstream hstream, int* pId); +TAresult TA_API taStreamDestroy(TAstream hStream); + +/********************************************************* + ********************************************************/ +#ifndef TANGRT_DEVICE_P2P_ATTR_ENUM +#define TANGRT_DEVICE_P2P_ATTR_ENUM +/** + * TANG Device P2P attributes + * TODO: This is design bug fix this. Remove this + */ +typedef enum tangDeviceP2PAttr { + tangDevP2PAttrPerformanceRank = 1, + tangDevP2PAttrAccessSupported = 2, + tangDevP2PAttrNativeAtomicSupported = 3, + tangDevP2PAttrTangArrayAccessSupported = 4, +} tangDeviceP2PAttr; +#endif // TANGRT_DEVICE_P2P_ATTR_ENUM + +TAresult TA_API taStreamC2Ctransfers(TAstream hStream, + uint32_t* cmd, + uint32_t cmd_count, + uint64_t device_addr, + uint32_t mem_size); + +TAresult TA_API taDeviceGetP2PAttribute(int* value, + tangDeviceP2PAttr attr, + int srcDevice, + int dstDevice); +TAresult TA_API taDeviceGetPeerPointer(int device, + int port, + void* peerAddr, + void** accessAddr); +TAresult TA_API taDeviceCanAccessPeer(int* canAccessPeer, + int device, + int peerDevice); +TAresult TA_API taDeviceEnablePeerAccess(int peerDevice, unsigned int flags); +TAresult TA_API taDeviceDisablePeerAccess(int peerDevice); +TAresult TA_API taMemcpyPeer(void* dst, + int dstDevice, + const void* src, + int srcDevice, + size_t count); + +TAresult TA_API taMemcpyPeerAsync(void* dst, + int dstDevice, + const void* src, + int srcDevice, + size_t count, + TAstream stream); + +//!< @ingroup Memory between peers. +//!< Copy data by HDMA +//!< @{{{ +TAresult TA_API taMemcpyPeer_v2(TAdeviceptr dst, + TAdevice dstDevice, + TAdeviceptr src, + TAdevice srcDevice, + size_t size); + +TAresult TA_API taMemcpyPeer_v2_ptds(TAdeviceptr dst, + TAdevice dstDevice, + TAdeviceptr src, + TAdevice srcDevice, + size_t size); + +TAresult TA_API taMemcpyPeerAsync_v2(TAdeviceptr dst, + TAdevice dstDevice, + TAdeviceptr src, + TAdevice srcDevice, + size_t size, + TAstream hStream); + +TAresult TA_API taMemcpyPeerAsync_v2_ptsz(TAdeviceptr dst, + TAdevice dstDevice, + TAdeviceptr src, + TAdevice srcDevice, + size_t size, + TAstream hStream); + +TAresult TA_API taMemcpyFromPeerAsync(TAdeviceptr dst, + TAdeviceptr src, + TAdevice srcDevice, + size_t size, + TAstream stream); + +TAresult TA_API taMemcpyToPeerAsync(TAdeviceptr dst, + TAdevice dstDevice, + TAdeviceptr src, + size_t size, + TAstream stream); + +TAresult TA_API taMemcpyFromPeerAsync_ptsz(TAdeviceptr dst, + TAdeviceptr src, + TAdevice srcDevice, + size_t size, + TAstream stream); + +TAresult TA_API taMemcpyToPeerAsync_ptsz(TAdeviceptr dst, + TAdevice dstDevice, + TAdeviceptr src, + size_t size, + TAstream stream); +//!< }}}@ + +TAresult TA_API taStreamWaitEvent(TAstream hStream, + TAevent hEvent, + unsigned int flags); +TAresult TA_API taStreamWaitEvent_ptsz(TAstream hStream, + TAevent hEvent, + unsigned int flags); +TAresult TA_API taStreamSynchronize(TAstream hStream); +TAresult TA_API taStreamSynchronize_ptsz(TAstream hStream); +TAresult TA_API taStreamQuery(TAstream hStream); +TAresult TA_API taStreamQuery_ptsz(TAstream hStream); + +TAresult TA_API taStreamBeginCapture(TAstream hStream, TAstreamCaptureMode mode); +TAresult TA_API taStreamBeginCapture_ptsz(TAstream hStream, + TAstreamCaptureMode mode); +TAresult TA_API taThreadExchangeStreamCaptureMode(TAstreamCaptureMode* mode); +TAresult TA_API taStreamEndCapture(TAstream hStream, TAgraph* phGraph); +TAresult TA_API taStreamEndCapture_ptsz(TAstream hStream, TAgraph* phGraph); +TAresult TA_API taStreamIsCapturing(TAstream hStream, + TAstreamCaptureStatus* captureStatus); +TAresult TA_API taStreamIsCapturing_ptsz(TAstream hStream, + TAstreamCaptureStatus* captureStatus); +TAresult TA_API taStreamGetCaptureInfo(TAstream hStream, + TAstreamCaptureStatus* pStatus, + unsigned long long* pId, + TAgraph* pGraph, + const TAgraphNode** deps, + size_t* numDeps); +TAresult TA_API taStreamGetCaptureInfo_ptsz(TAstream hStream, + TAstreamCaptureStatus* pStatus, + unsigned long long* pId, + TAgraph* pGraph, + const TAgraphNode** deps, + size_t* numDeps); + +TAresult TA_API taGraphInstantiateWithFlags(TAgraphExec* phGraphExec, + TAgraph hGraph, + unsigned long long flags); +TAresult TA_API taGraphLaunch(TAgraphExec hGraphExec, TAstream hStream); +TAresult TA_API taGraphLaunch_ptsz(TAgraphExec hGraphExec, TAstream hStream); +TAresult TA_API taGraphDestroy(TAgraph hGraph); +TAresult TA_API taGraphExecDestroy(TAgraphExec hGraphExec); +TAresult TA_API taGraphGetInfo(TAgraph hGraph, TAgraphInfo* pInfo); +TAresult TA_API taGraphCreate(TAgraph* phGraph, unsigned int flags); + +TAresult TA_API taGraphAddHostNode(TAgraphNode* phGraphNode, + TAgraph hGraph, + const TAgraphNode* dependencies, + size_t numDependencies, + const TANG_HOST_NODE_PARAMS* nodeParams); + +TAresult TA_API taGraphAddKernelNode(TAgraphNode* phGraphNode, + TAgraph hGraph, + const TAgraphNode* dependencies, + size_t numDependencies, + const TANG_KERNEL_NODE_PARAMS* nodeParams); + +TAresult TA_API taEventCreate(TAevent* phEvent, unsigned int flags); +TAresult TA_API taEventDestroy(TAevent hEvent); +TAresult TA_API taEventRecord(TAevent hEvent, TAstream hStream); +TAresult TA_API taEventRecord_ptsz(TAevent hEvent, TAstream hStream); +TAresult TA_API taEventRecordWithFlags(TAevent hEvent, + TAstream hStream, + unsigned int flags); +TAresult TA_API taEventRecordWithFlags_ptsz(TAevent hEvent, + TAstream hStream, + unsigned int flags); +TAresult TA_API taEventSynchronize(TAevent hEvent); +TAresult TA_API taEventElapsedTime(float* pMilliseconds, + TAevent hStart, + TAevent hEnd); +TAresult TA_API taEventQuery(TAevent hEvent); +TAresult TA_API taEventQueryTimestamp(TAevent hEvent, TAeventTimestamp* pTs); + +TAresult TA_API taEventSynchronizeWithFlags(TAevent hEvent, unsigned int flags); + +void TA_API taDsoWrapperInit(TAdsoWrapper_t* dso); +void TA_API taDsoWrapperDeinit(TAdsoWrapper_t* dso); + +TAresult TA_API taGetBuiltinModule(TAmodule* phModule, const char* name); + +TAresult TA_API taModuleLoad(TAmodule* phModule, const char* filename); +TAresult TA_API taModuleLoadData(TAmodule* phModule, const void* image, size_t size); +TAresult TA_API taModuleUnload(TAmodule hModule); + +TAresult TA_API taModuleLoadFatBinaryManaged(TAmodule* phModule, + const void* fatbin, + const char* fatbinInfo, + TAdsoWrapper_t* dso); + +TAresult TA_API taModuleUnloadManaged(TAmodule hModule); + +/** + * @brief Get the module symbol type name. + * + * @param name The pointer to receive the name of the type. + * @param type The symbol type. + * @return TAresult + * ::TANG_SUCCESS if the type is a valid value. + * ::TANG_ERROR_INVALID_VALUE if the type is not a valid type. + */ +TAresult TA_API taModuleSymbolTypeGetName(char const** name, TAmoduleSymbolType type); + +/** + * @brief Iterate symbols in \p hmod. + * + * @param hmod + * @param fn The call back function. Return true will cause the + * iteration to stop. + * @param userData + * @return TAresult + */ +// TAresult TA_API taModuleIterateSymbols(TAmodule hmod, +// TAmoduleSymbolIterateFn fn, +// void* userData); + +TAresult TA_API taModuleGetFunction(TAfunction* hfunc, + TAmodule hmod, + const char* name); + +TAresult TA_API taFuncGetAttribute(int* pi, + TAfunction_attribute attr, + TAfunction hfunc); + +TAresult TA_API taFuncGetModule(TAmodule* hmod, TAfunction func); + +TAresult TA_API taFunctionGetAddress(TAfunction func, TAdeviceptr* address); + +TAresult TA_API taFunctionGetNumArgs(TAfunction func, size_t* numArgs); + +TAresult TA_API taFunctionGetInfo(TAfunction func, + TAdeviceptr* address, + size_t* lenArgs, + size_t* thread_regfile_size, + size_t* shared_regfile_base, + size_t* shared_regfile_size, + size_t* warp_regfile_size, + size_t* local_memory_size, + size_t* static_shared_mem_size, + size_t* shared_memory_mirror, + size_t* max_threads_per_block, + size_t* max_dynamic_shared_mem_size_per_block, + size_t* max_block_count); + +TAresult TA_API taModuleGetVariable(TAvariable* hVar, + TAmodule hMod, + const char* varName); + +TAresult TA_API taVariableGetInfo(TAvariable hVar, + TAdeviceptr* address, + size_t* size); + +TAresult TA_API taVariableCopyFromDevice(TAvariable hVar, + TAdeviceptr src, + size_t size, + size_t offset); + +TAresult TA_API taVariableCopyFromDevice_ptds(TAvariable hVar, + TAdeviceptr src, + size_t size, + size_t offset); + +TAresult TA_API taVariableCopyFromDeviceAsync(TAvariable hVar, + TAdeviceptr src, + size_t size, + size_t offset, + TAstream stream); + +TAresult TA_API taVariableCopyFromDeviceAsync_ptsz(TAvariable hVar, + TAdeviceptr src, + size_t size, + size_t offset, + TAstream stream); + +TAresult TA_API taVariableCopyFromHost(TAvariable hVar, + const void* src, + size_t size, + size_t offset); + +TAresult TA_API taVariableCopyFromHost_ptds(TAvariable hVar, + const void* src, + size_t size, + size_t offset); + +TAresult TA_API taVariableCopyFromHostAsync(TAvariable hVar, + const void* src, + size_t size, + size_t offset, + TAstream stream); + +TAresult TA_API taVariableCopyFromHostAsync_ptsz(TAvariable hVar, + const void* src, + size_t size, + size_t offset, + TAstream stream); + +TAresult TA_API taVariableCopyToDevice(TAdeviceptr dst, + TAvariable hVar, + size_t size, + size_t offset); + +TAresult TA_API taVariableCopyToDevice_ptds(TAdeviceptr dst, + TAvariable hVar, + size_t size, + size_t offset); + +TAresult TA_API taVariableCopyToDeviceAsync(TAdeviceptr dst, + TAvariable hVar, + size_t size, + size_t offset, + TAstream stream); + +TAresult TA_API taVariableCopyToDeviceAsync_ptsz(TAdeviceptr dst, + TAvariable hVar, + size_t size, + size_t offset, + TAstream stream); + +TAresult TA_API taVariableCopyToHost(void* dst, + TAvariable hVar, + size_t size, + size_t offset); + +TAresult TA_API taVariableCopyToHost_ptds(void* dst, + TAvariable hVar, + size_t size, + size_t offset); + +TAresult TA_API taVariableCopyToHostAsync(void* dst, + TAvariable hVar, + size_t size, + size_t offset, + TAstream stream); + +TAresult TA_API taVariableCopyToHostAsync_ptsz(void* dst, + TAvariable hVar, + size_t size, + size_t offset, + TAstream stream); + +/** + * @brief Enqueue A raw SCP command packet onto stream. + * + * @param stream The stream. + * @param regs The SCP command packet. + * @param size The size of the command packet in byte. + * @return int + * @warning A raw SCP command packet needs four bytes aligned. + * The \p size must be integral multiple of 4. + ****************************************************************/ +TAresult TA_API taEnqueueCommand(TAstream stream, void* regs, size_t size); + +TAresult TA_API taEnqueueCommand_ptsz(TAstream stream, void* regs, size_t size); + +/** + * @brief Launch a kernel function. + * + * @param func + * @param gridX + * @param gridY + * @param gridZ + * @param blockX + * @param blockY + * @param blockZ + * @param sharedMemBytes + * @param stream + * @param funcParams + * @param extra + * @code {.cpp} + * + * @endcode + * @param extra + * @return int + * ::TANG_SUCCESS + * ::TANG_ERROR_OUT_OF_MEMORY + * ::TANG_ERROR_NOT_IMPLEMENTED + * @sa taModuleGetFunction + * @sa TA_STREAM_LEGACY + ********************************************************/ +TAresult TA_API taLaunchFunction(TAfunction func, + unsigned int gridX, + unsigned int gridY, + unsigned int gridZ, + unsigned int blockX, + unsigned int blockY, + unsigned int blockZ, + unsigned int sharedMemBytes, + TAstream stream, + void** funcParams, + void** extra); + +TAresult TA_API taLaunchFunction_ptsz(TAfunction func, + unsigned int gridX, + unsigned int gridY, + unsigned int gridZ, + unsigned int blockX, + unsigned int blockY, + unsigned int blockZ, + unsigned int sharedMemBytes, + TAstream stream, + void** funcParams, + void** extra); + +TAresult TA_API taLaunchKernel(TAfunction func, + unsigned int gridX, + unsigned int gridY, + unsigned int gridZ, + unsigned int blockX, + unsigned int blockY, + unsigned int blockZ, + unsigned int sharedMemBytes, + TAstream stream, + void** funcParams, + void** extra); + +TAresult TA_API taLaunchKernel_ptsz(TAfunction func, + unsigned int gridX, + unsigned int gridY, + unsigned int gridZ, + unsigned int blockX, + unsigned int blockY, + unsigned int blockZ, + unsigned int sharedMemBytes, + TAstream stream, + void** funcParams, + void** extra); + +//!< fn(usrData) +TAresult TA_API taLaunchHostFunc(TAstream hStream, TAhostFn fn, void* userData); +TAresult TA_API taLaunchHostFunc_ptsz(TAstream hStream, TAhostFn fn, void* userData); + +typedef void (*TAstreamCallback)(TAstream hStream, + TAresult status, + void* userData); + +//!< callback(hStream, status, userData); +TAresult TA_API taStreamAddCallback(TAstream hStream, + TAstreamCallback callback, + void* userData, + unsigned int flags); + +TAresult TA_API taStreamAddCallback_ptsz(TAstream hStream, + TAstreamCallback callback, + void* userData, + unsigned int flags); + +//!< proxy(proxy_data, func, func_data, error) +TAresult TA_API taLaunchHostFuncProxy(TAstream stream, + void* proxy, + void* proxy_data, + void* func, + void* func_data); + +TAresult TA_API taLaunchHostFuncProxy_ptsz(TAstream stream, + void* proxy, + void* proxy_data, + void* func, + void* func_data); + +TAresult TA_API taOccupancyMaxActiveBlocksPerMultiprocessor(int* numBlocks, + TAfunction func, + int blockSize, + size_t dynamicSMemSize); +TAresult TA_API taProfilerStart(); +TAresult TA_API taProfilerStop(); + +/** + * @brief Gets an interprocess communication memory handle from device memory + * allocated by tangMalloc or cuMemAlloc + * + * @param pHandle + * @param dptr + * @return + * ::TANG_SUCCESS + * ::TANG_ERROR_INVALID + * ::TANG_ERROR_OUT_OF_MEMORY + * @sa + * ::taIpcOpenMemHandle + * ::taIpcCloseMemHandle + */ +TAresult TA_API taIpcGetMemHandle(TAipcMemHandle* pHandle, TAdeviceptr dptr); + +/** + * @brief Opens an interprocess communication memory handle exported from another + * process and map it into the current context and returns a device pointer. + * + * @param pdptr + * @param handle + * @param flags + * @return + * ::TANG_SUCCESS + * ::TANG_ERROR_IPC_MEM_DESTROIED + * ::TANG_ERROR_OUT_OF_MEMORY + * @sa + * ::taIpcGetMemHandle + * ::taIpcCloseMemHandle + */ +TAresult TA_API taIpcOpenMemHandle(TAdeviceptr *pdptr, + TAipcMemHandle handle, + unsigned int flags); + +/** + * @brief Unmap the memory got from taIpcOpenMemHandle. + * + * @param dptr + * @return int + */ +TAresult TA_API taIpcCloseMemHandle(TAdeviceptr dptr); + +/** + * @brief Gets an interprocess event handle. The event must be created with + * ::TA_EVENT_INTERPROCESS flag set. + * + * @param pHandle + * @param event + * @return int + */ +TAresult TA_API taIpcGetEventHandle(TAipcEventHandle* pHandle, TAevent event); + +/** + * @brief Opens an interprocess event handle for the calling process. + * + * Opens an interprocess event handle exported from another process with + * ::taIpcGetEventHandle. + * + * Use ::taEventDestroy to free the event. + * @param phEvent + * @param handle + * @return int + * ::TANG_SUCCESS + * ::TANG_ERROR_OUT_MEMORY + */ +TAresult TA_API taIpcOpenEventHandle(TAevent* phEvent, TAipcEventHandle handle); + +/** + * @brief launch engine collectives witch stream. + * + * @param devId which ptpu + * @param collType engine collectives type + * @param devAddr HBM address + * @param size params length + * @param stream for sync + * @return int + */ +TAresult TA_API taStreamEngineCollAssign(int devId, int collType, uint64_t devAddr, int size, TAstream stream); + +// enum TAqueryInfoType_enum { +// TA_QUERY_INFO_MEMORY_USAGE = 0x01, +// }; +// typedef enum TAqueryInfoType_enum TAqueryInfoType; +// +// TAresult TA_API taCtxQueryMemoryUsage(int64_t* pTotal, TAcontext context); + +TAresult TA_API taGetExportTable(void** pTable, void* args); + +#ifdef __cplusplus +} +#endif // __cplusplus + +#endif // _TANG_H_ + diff --git a/third_party/sunrise/backend/include/tang_compiler_api.h b/third_party/sunrise/backend/include/tang_compiler_api.h new file mode 100755 index 0000000000..2c2f4ecd4d --- /dev/null +++ b/third_party/sunrise/backend/include/tang_compiler_api.h @@ -0,0 +1,223 @@ +/* +Copyright declaration. +*/ + +#ifndef _TANG_RT_TANG_COMPILER_API_H +#define _TANG_RT_TANG_COMPILER_API_H + +#include "tang_rt/driver_types.h" +#include "tang_rt/host_defines.h" +#include "tang_rt/vector_types.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * @brief C compliant register fatbinary API + * + * @param [in] fatTangbin - pointer of tang fatbin data + * + * @returns tang fatbin handle registered in runtime + * + */ +TANGRT_API_PUBLIC void** __tangRegisterFatBinary(void* fatTangbin); + +/** + * @brief + * + * @param fatbinHandle + */ +TANGRT_API_PUBLIC void __tangRegisterFatBinaryEnd(void** fatbinHandle); + +/** + * @brief C compliant unregister fatbinary API + * + * @param [in] fatTangbinHandle - tang fatbin handle registered in runtime + * + * @returns void + * + */ +TANGRT_API_PUBLIC void __tangUnregisterFatBinary(void** fatTangbinHandle); + +/** + * @brief Register fatbinary + * + * @param fatbinWrapper + * @return void** + * @example + * @code{.cpp} + * //!< Generated by compiler + * const unsigned long fatbinText[] = { 0x12344567, ..., 0xabcd1235 }; + * static __tangFatbinaryWrapper fatbinWrapper = { + * .version = 0, + * .fatbin = &fatbinText[0], + * .size = sizeof(fatbinText), + * }; + * static void** fatbinHandle; + * static void __tang_RegisterAll(void) __attribute__((constructor)) + * static void __tang_UnregisterAll(void); + * static void __tang_RegisterAll(void) { + * fatbinHandle = __tangRegisterFatBinary_v2(&fatbinWrapper); + * __tangRegisterFunction((const void*)vecadd, + * (const char*)"_Z6vecaddPiS_S_", + * fatbinHandle); + * __tangRegisterFatBinaryEnd(fatbinHandle); + * atexit(__tang_UnregisterAll); + * } + * static void __tang_UnregisterAll(void) { + * __tangUnregisterFatBinary_v2(fatbinHandle); + * } + * @endcode + */ +TANGRT_API_PUBLIC void** __tangRegisterFatBinary_v2(void* fatbinWrapper); + +/** + * @brief Unregister fatbinary + * + * @param fatTangbinHandle + * @return TANGRT_API_PUBLIC + */ +TANGRT_API_PUBLIC void __tangUnregisterFatBinary_v2(void** fatTangbinHandle); + +TANGRT_API_PUBLIC int __tangInitModule(void** fatTangbinHandle); + +/** + * @brief C compliant set fatbinary info API + * + * @param [in] fatbinHandle - tang fatbin handle registered in runtime + * @param [in] info - pointer of tang fatbin info + * + * @returns void + * + */ +TANGRT_API_PUBLIC void __tangSetFatBinaryInfo(void** fatbinHandle, + const char* info); + +/** + * @brief C compliant register function API + * + * @param [in] hostFunc - the pointer of the host stub function + * @param [in] deviceFuncName - the name string of device function. + * @param [in] fatbinHandle - tang fatbin handle registered in runtime + * + * @returns void + * + */ +TANGRT_API_PUBLIC void __tangRegisterFunction(const void* hostFunc, + const char* deviceFuncName, + void** fatbinHandle); + +/** + * @brief Register device variable. + * + * @param fatbinHandle The fatbin handle returned by \c __tangRegisterFatBinary + * @param hostVar The corresponding host variable address used as the key. + * @param deviceVarAddress The device variable address. + * @param deviceVarName The symbol name of the devcie variable. + * @param ext + * @param size The size of variable in bytes. + * @param constant If the variable is in const memory ? + * @param global From the host's point of view, device variable is always local. + * Thus, param "global" is always 0. + * @remark + * For code `__device__ int xyz;` the compiler should generate the following + * code. + * @code{.cpp} + * static int xyz; + * __tangRegisterVariable(fatbinHandle, + * &xyz, (char*)"xyz", "xyz", + * 0, sizeof(int), 0, 0); + * @endcode + */ +TANGRT_API_PUBLIC void __tangRegisterVariable(void** fatbinHandle, + const void* hostVar, + char* deviceVarAddress, + const char* deviceVarName, + int ext, + size_t size, + int constant, + int global); + +/** + * @brief C compliant push call configuration API + * + * @param [in] gridDim - number of blocks in a grid + * @param [in] blockDim - number of threads in a block + * @param [in] sharedMemBytes - Amount of dynamic shared memory to allocate for + * this kernel. The Kernel can access this with TANG_DYNAMIC_SHARED. + * @param [in] stream - Stream where the kernel should be dispatched. May be 0, + * in which case the default stream is used with associated synchronization + * rules. + * + * @returns #tangSuccess + * + */ +TANGRT_API_PUBLIC tangError_t +__tangPushCallConfiguration(dim3 gridDim, + dim3 blockDim, + size_t sharedMemBytes __dparm(0), + tangStream_t stream __dparm(0)); + +/** + * @brief C compliant pop call configuration API + * + * @param [out] gridDim - number of blocks in a grid + * @param [out] blockDim - number of threads in a block + * @param [out] sharedMemBytes - Amount of dynamic shared memory to allocate for + * this kernel. The Kernel can access this with TANG_DYNAMIC_SHARED. + * @param [out] stream - Stream where the kernel should be dispatched. May be + * 0, in which case the default stream is used with associated synchronization + * rules. + * + * @returns #tangSuccess + * + */ +TANGRT_API_PUBLIC tangError_t __tangPopCallConfiguration(dim3* gridDim, + dim3* blockDim, + size_t* sharedMemBytes, + tangStream_t* stream); + +/** + * @brief C compliant kernel launch API + * + * @param [in] hostFunc - the pointer of the host stub function + * @param [in] gridDim - number of blocks in a grid + * @param [in] blockDim - number of threads in a block + * @param [in] args - kernel arguments + * @param [in] numArgs - number of kernel arguments + * @param [in] sharedMemBytes - Amount of dynamic shared memory to allocate for + * this kernel. The Kernel can access this with TANG_DYNAMIC_SHARED. + * @param [in] stream - Stream where the kernel should be dispatched. May be 0, + * in which case th default stream is used with associated synchronization + * rules. + * + * @returns #tangSuccess, #tangErrorInvalidValue, tangInvalidDevice + * + */ +TANGRT_API_PUBLIC tangError_t tangLaunchKernel(const void* hostFunc, + dim3 gridDim, + dim3 blockDim, + void** args, + size_t numArgs, + size_t sharedMemBytes __dparm(0), + tangStream_t stream __dparm(0)); + +TANGRT_API_PUBLIC tangError_t +tangLaunchKernel_ptsz(const void* hostFunc, + dim3 gridDim, + dim3 blockDim, + void** args, + size_t numArgs, + size_t sharedMemBytes __dparm(0), + tangStream_t stream __dparm(0)); + +#ifdef __cplusplus +} +#endif + +#if defined(__TANGRT_API_PER_THREAD_DEFAULT_STREAM) +#define tangLaunchKernel __TANGRT_API_PTSZ(tangLaunchKernel) +#endif //! __TANGRT_API_PER_THREAD_DEFAULT_STREAM + +#endif //! _TANG_RT_TANG_COMPILER_API_H diff --git a/third_party/sunrise/backend/include/tang_rt/device_assert.h b/third_party/sunrise/backend/include/tang_rt/device_assert.h new file mode 100755 index 0000000000..691fc8f5b3 --- /dev/null +++ b/third_party/sunrise/backend/include/tang_rt/device_assert.h @@ -0,0 +1,43 @@ +#ifndef _TANGRT_DEVICE_ASSERT_H_ +#define _TANGRT_DEVICE_ASSERT_H_ + +#include + +#include + +#include "tang_rt/device_functions.h" + +extern "C" { +// #pragma push_macro("size_t") +// #define size_t unsigned +__device__ void __assertfail(const char *__message, + const char *__file, + unsigned __line, + const char *__function, + unsigned __charSize) +//__attribute__((noreturn)) +{ + __pt_printf("%d: block: [%d,%d,%d], thread: [%d,%d,%d] Assertion failed.\n", + __line, + blockIdx.x, + blockIdx.y, + blockIdx.z, + threadIdx.x, + threadIdx.y, + threadIdx.z); + asm volatile("exit\n\t" ::: "memory"); +} +// #undef size_t +// #pragma pop_macro("size_t") + +// In order for standard assert() macro on linux to work we need to +// provide device-side __assert_fail() +__device__ static inline void __assert_fail(const char *__message, + const char *__file, + unsigned __line, + const char *__function) { + __assertfail(__message, __file, __line, __function, sizeof(char)); +} +} // end extern "C" + +#endif diff --git a/third_party/sunrise/backend/include/tang_rt/device_functions.h b/third_party/sunrise/backend/include/tang_rt/device_functions.h new file mode 100755 index 0000000000..f95c0112bd --- /dev/null +++ b/third_party/sunrise/backend/include/tang_rt/device_functions.h @@ -0,0 +1,8 @@ +#ifndef _TANGRT_DEVICE_FUNCTIONS_H_ +#define _TANGRT_DEVICE_FUNCTIONS_H_ +#include + +#include "tang_rt/fmt.hpp" + +#endif //!< _TANGRT_DEVICE_FUNCTIONS_H_ + diff --git a/third_party/sunrise/backend/include/tang_rt/driver_types.h b/third_party/sunrise/backend/include/tang_rt/driver_types.h new file mode 100755 index 0000000000..0b76cc20ce --- /dev/null +++ b/third_party/sunrise/backend/include/tang_rt/driver_types.h @@ -0,0 +1,1281 @@ +/* +Copyright declaration. +*/ +#ifndef _TANG_RT_DRIVER_TYPES_H_ +#define _TANG_RT_DRIVER_TYPES_H_ +#include "tang_rt/vector_types.h" + +#define tangHostAllocDefault 0x00 /**< Default page-locked allocation flag */ +#define tangHostAllocPortable \ + 0x01 /**< Pinned memory accessible by all TANG contexts */ +#define tangHostAllocMapped 0x02 /**< Map allocation into device space */ +#define tangHostAllocWriteCombined 0x04 /**< Write-combined memory */ +#define tangHostAllocMapDeviceMemory \ + 0x100 /**< Allocate device memory and map it to userspace */ + +#define tangHostRegisterDefault \ + 0x00 /**< Default host memory registration flag */ +#define tangHostRegisterPortable \ + 0x01 /**< Pinned memory accessible by all TANG contexts */ +#define tangHostRegisterMapped \ + 0x02 /**< Map registered memory into device space */ +#define tangHostRegisterIoMemory 0x04 /**< Memory-mapped I/O space */ +#define tangHostRegisterReadOnly 0x08 /**< Memory-mapped read-only */ + +/**< Default behavior */ +#define tangOccupancyDefault 0x00 +/**< Assume global caching is enabled and cannot be automatically turned off */ +#define tangOccupancyDisableCachingOverride 0x01 +/* + * TANG error types + */ +// Developer note - when updating these, update the tangGetErrorString +// functions. + +/** + * Not consider Cooperative Launch, Peer Access, Pitch, Texture, Graph, JIT, + * Managed Memory, Multi-thread(compute mode) right now + */ + +enum tangError { + /** + * Some of the tangError are not supported, + * such as tangErrorInvalidPitchValue, tangErrorInvalidHostPointer, etc + * If we develop tools to convert tang program to tang, + * tangTangErrorToTangError should redirect them to tangErrorUnknown. + * Need check. + */ + + /** + * The API call returned with no errors. In the case of query calls, this + * also means that the operation being queried is complete (see + * ::tangEventQuery() and ::tangStreamQuery()). + * Need check. + */ + tangSuccess = 0, + + /** + * This indicates that one or more of the parameters passed to the API call + * is not within an acceptable range of values. + */ + tangErrorInvalidValue = 1, + + /** + * The API call failed because it was unable to allocate enough memory to + * perform the requested operation. + * In cuda: cudaErrorMemoryAllocation + * In hip: hipErrorOutOfMemory + */ + tangErrorMemoryAllocation = 2, + + /** + * The API call failed because the TANG driver and runtime could not be + * initialized. + * In cuda: cudaErrorInitializationError + * In hip: hipErrorNotInitialized + */ + tangErrorInitializationError = 3, + + /** + * This indicates that a TANG Runtime API call cannot be executed because + * it is being called during process shut down, at a point in time after + * Tang runtime has been deinitialized. + * In cuda: cudaErrorCudartUnloading + * In hip: hipErrorDeinitialized + * hip declares it but not use. + * Maybe hip will not meet such a point. + * Need check. + */ + tangErrorDeinitialized = 4, + + /** + * This indicates that the device is removed. + */ + tangErrorDeviceRemoved = 5, + + /** + * This indicated that the device is reset. + */ + tangErrorDeviceReset = 6, + + /** + * This operation is not allowed + */ + tangErrorNotPermitted = 7, + + /** + * The specified file or directory is not found + */ + tangErrorNoSuchFile = 8, + + /** + * This indicates that a kernel launch is requesting resources that can + * never be satisfied by the current device. Requesting more shared memory + * per block than the device supports will trigger this error, as will + * requesting too many threads or blocks. See ::tangDeviceProp for more + * device limitations. + */ + tangErrorInvalidConfiguration = 9, + + /** + * @brief Null pointer is passed as argument but it is disallowed. + */ + tangErrorNullPointer = 10, + + tangErrorOutOfResources = 11, + + /** + * This indicates that the symbol name/identifier passed to the API call + * is not a valid name or identifier. + */ + tangErrorInvalidSymbol = 13, + + /** + * This indicates that at least one device pointer passed to the API call is + * not a valid device pointer. + * Note: This error is deprecated from CUDA 10.1, + * but hip still use it + */ + tangErrorInvalidDevicePointer = 17, + + /** + * This indicates that the direction of the memcpy passed to the API call is + * not one of the types specified by ::tangMemcpyKind. + * hip declares it but not use. + * But it seems useful. + * Need check. + */ + tangErrorInvalidMemcpyDirection = 21, + + /** + * This indicates that the installed TANG driver version is mismatch with the + * runtime version. + * hip declares it but not use. + * But it seems useful. + * Need check. + */ + tangErrorInsufficientDriver = 35, + + /** + * The device function being invoked (usually via ::tangLaunchKernel()) was + * not previously configured via the ::tangConfigureCall() function + * nor provided with ".kernel_info" in ELF. + */ + tangErrorMissingConfiguration = 52, + + /** + * The requested device function does not exist or is not compiled for the + * proper device architecture. + */ + tangErrorInvalidDeviceFunction = 98, + + /** + * This error indicates the attempted operation is not implemented. + */ + tangErrorNotImplemented = 99, + + /** + * This indicates that no TANG-capable devices were detected by the installed + * TANG driver. Call to tangGetDeviceCount returned 0 devices. + */ + tangErrorNoDevice = 100, + + /** + * This indicates that the device ordinal supplied by the user does not + * correspond to a valid TANG device. + * DeviceID must be in range 0...#compute-devices. + */ + tangErrorInvalidDevice = 101, + + /** + * This indicates that the device kernel image is invalid. + * In cuda: cudaErrorInvalidKernelImage + * In hip: hipErrorInvalidImage + */ + tangErrorInvalidKernelImage = 200, + + /** + * This most frequently indicates that there is no context bound to the + * current thread. This can also be returned if the context passed to an + * API call is not a valid handle (such as a context that has had + * ::taCtxDestroy() invoked on it). + * In cuda: cudaErrorDeviceUninitialized + * In hip：hipErrorInvalidContext + */ + tangErrorInvalidContext = 201, + + /** + * This indicates that there is no kernel image available that is suitable + * for the device. This can occur when a user specifies code generation + * options for a particular TANG source file that do not include the + * corresponding device configuration. + * In cuda: cudaErrorNoKernelImageForDevice + * In hip: hipErrorNoBinaryForGpu + * Need check. + */ + tangErrorNoKernelImageForDevice = 209, + + /** + * This indicates that the ::tangLimit passed to the API call is not + * supported by the active device. + */ + tangErrorUnsupportedLimit = 215, + + /** + * This indicates that a call tried to access an exclusive-thread device that + * is already in use by a different thread. + * In cuda: cudaErrorDeviceAlreadyInUse + * In hip: hipErrorContextAlreadyInUse + */ + tangErrorContextAlreadyInUse = 216, + + /** + * A PTX compilation failed. The runtime may fall back to compiling PTX if + * an application does not contain a suitable binary for the current device. + * In cuda: cudaErrorInvalidPtx + * In hip: hipErrorInvalidKernelFile + */ + tangErrorInvalidKernelFile = 218, + + /** + * When launch kernel, unable to find the corresponding fatbinary. + */ + tangErrorFatBinaryNotFound = 300, + + /** + * This indicates that the file specified was not found. + */ + tangErrorFileNotFound = 301, + + /** + * This indicates that a link to a shared object failed to resolve. + */ + tangErrorSharedObjectSymbolNotFound = 302, + + /** + * This indicates that initialization of a shared object failed. + */ + tangErrorSharedObjectInitFailed = 303, + + /** + * This error indicates that an OS call failed. + * hip declares it but not use. + * But it seems useful. + * Need check. + */ + tangErrorOperatingSystem = 304, + + tangErrorIllegalState = 305, + + tangErrorStreamCaptureUnsupported = 306, + + tangErrorStreamCaptureInvalidated = 307, + + tangErrorStreamCaptureMerge = 308, + + tangErrorStreamCaptureUnmatched = 309, + + tangErrorStreamCaptureUnjoined = 310, + + tangErrorStreamCaptureIsolation = 311, + + tangErrorStreamCaptureImplicit = 312, + + tangErrorStreamCaptureWrongThread = 313, + + tangErrorCapturedEvent = 314, + + /** + * This indicates that a resource handle passed to the API call was not + * valid. Resource handles are opaque types like ::tangStream_t and + * ::tangEvent_t. + * In cuda: cudaErrorInvalidResourceHandle + * In hip: hipErrorInvalidHandle + */ + tangErrorInvalidResourceHandle = 400, + + /** + * This indicates that a named symbol was not found. Examples of symbols + * are global/constant variable names, texture names, and surface names. + * In cuda: cudaErrorSymbolNotFound + * In hip: tangErrorNotFound + */ + tangErrorSymbolNotFound = 500, + + /** + * This indicates that asynchronous operations issued previously have not + * completed yet. This result is not actually an error, but must be indicated + * differently than ::tangSuccess (which indicates completion). Calls that + * may return this value include ::tangEventQuery() and ::tangStreamQuery(). + */ + tangErrorNotReady = 600, + + /** + * The device encountered a load or store instruction on an invalid memory + * address. This leaves the process in an inconsistent state and any further + * TANG work will return the same error. To continue using TANG, the process + * must be terminated and relaunched. hip declares it but not use. But it + * seems useful. Need check. + */ + tangErrorIllegalAddress = 700, + + /** + * This indicates that a launch did not occur because it did not have + * appropriate resources. Although this error is similar to + * ::tangErrorInvalidConfiguration, this error usually indicates that the + * user has attempted to pass too many arguments to the device kernel, or the + * kernel launch specifies too many threads for the kernel's register count. + */ + tangErrorLaunchOutOfResources = 701, + + /** + * This indicates that the device kernel took too long to execute. This can + * only occur if timeouts are enabled - see the device property + * \ref ::tangDeviceProp::kernelExecTimeoutEnabled "kernelExecTimeoutEnabled" + * for more information. + * This leaves the process in an inconsistent state and any further TANG work + * will return the same error. To continue using TANG, the process must be + * terminated and relaunched. hip declares it but not use. But it seems + * useful. Need check. + */ + tangErrorLaunchTimeOut = 702, + + tangErrorPeerAccessAlreadyEnabled = 704, + tangErrorPeerAccessNotEnabled = 705, + + /** + * This error indicates that the memory range passed to ::tangHostRegister() + * has already been registered. + * hip declares it but not use. + * But it seems useful. + * Need check. + */ + tangErrorHostMemoryAlreadyRegistered = 712, + + /** + * This error indicates that the pointer passed to ::tangHostUnregister() + * does not correspond to any currently registered memory region. + * hip declares it but not use. + * But it seems useful. + * Need check. + */ + tangErrorHostMemoryNotRegistered = 713, + + /** + * An exception occurred on the device while executing a kernel. Common + * causes include dereferencing an invalid device pointer and accessing + * out of bounds shared memory. Less common cases can be system specific - + * more information about these cases can be found in the system specific user + * guide. This leaves the process in an inconsistent state and any further + * TANG work will return the same error. To continue using TANG, the process + * must be terminated and relaunched. + */ + tangErrorLaunchFailure = 719, + + /** + * This error indicates the attempted operation is not supported + * on the current system or device. + */ + tangErrorNotSupported = 801, + + /** + * This indicates that an unknown internal error has occurred. + */ + tangErrorUnknown = 999, + + /** + * @brief context is destroyed or in destroying in kernel + * + */ + tangErrorContextIsDestroyed = 3000, + + /** + * @brief context is not valid in kernel + * + */ + tangErrorContextInvalid = 3001, + + /** + * @brief stream is destroyed or in destroying in kernel + * + */ + tangErrorStreamIsDestroyed = 3002, + + /** + * @brief stream is not valid in kernel + * + */ + tangErrorStreamInvalid = 3003, + + /** + * @brief event is destroyed or in destroying in kernel + * + */ + tangErrorEventIsDestroyed = 3004, + + /** + * @brief event is not valid in kernel + * + */ + tangErrorEventInvalid = 3005, + + /** + * @brief device memory is not enough for current operation + * + */ + tangErrorDeviceOutOfMemory = 3006, + + /** + * @brief device memory is not found + * + */ + tangErrorDeviceMemoryNotFound = 3007, + + /** + * @brief pcie fatal error occured + * + */ + tangErrorPcieFatal = 3012, + + /** + * @brief pcie non-fatal unrecovered error occured + * + */ + tangErrorPcieNonFatalUnrecovered = 3013, + + /** + * @brief no more event exist + * + */ + tangErrorScpEventNotExist = 3014, + + /** + * @brief record event failed + * + */ + tangErrorSCPEventRecordFailed = 3015, + + /** + * @brief scp packet crc check failed + * + */ + tangErrorSCPCrcPacketFailed = 3016, + + /** + * @brief scp dispatch send failed + * + */ + tangErrorSCPDispSendFailed = 3017, + + /** + * @brief sq write sequence error + * + */ + tangErrorSCPSqWriteFailed = 3018, + + /** + * @brief udrc pcie xdma packet invalid + * + */ + tangErrorUdrcPcieDmaPacketInvalid = 3019, + + /** + * @brief udrc mp dma packet invalid + * + */ + tangErrorUdrcMpDmaPacketInvalid = 3020, + + /** + * @brief udrc reg packet invalid + * + */ + tangErrorUdrcRegPacketInvalid = 3021, + + /** + * @brief udrc reg access invalid + * + */ + tangErrorUdrcRegAcessInvalid = 3022, + + /** + * @brief aiss cluster is not configured + * + */ + tangErrorAissClusterUsrNotAllocated = 3023, + + /** + * @brief barrier is destroyed or in destroying in kernel + * + */ + tangErrorBarrierDestroyed = 3024, + + /** + * @brief barrier is not valid in kernel + * + */ + tangErrorBarrierInvalid = 3025, + + /** + * @brief one obj is destroyed or in destroying in kernel + * + */ + tangErrorDestroyed = 3026, + + /** + * @brief xdma C2H align mismath + * + */ + tangErrorXdmaC2HAlignMismatch = 3300, + + /** + * @brief xdma C2H invalid magic stopped + * + */ + tangErrorXdmaC2HInvalidMagicStopped = 3301, + + /** + * @brief xdma C2H invalid Len + * + */ + tangErrorXdmaC2HInvalidLen = 3302, + + /** + * @brief xdma C2H decode error + * + */ + tangErrorXdmaC2HDecode = 3303, + + /** + * @brief xdma C2H slave + * + */ + tangErrorXdmaC2HSlave = 3304, + + /** + * @brief xdma C2H descriptor unsupport request + * + */ + tangErrorXdmaC2HDescUnsupportRequest = 3305, + + /** + * @brief xdma C2H descriptor completer abort + * + */ + tangErrorXdmaC2HDescCompleterAbort = 3306, + + /** + * @brief xdma C2H descriptor parity + * + */ + tangErrorXdmaC2HDescParity = 3307, + + /** + * @brief xdma C2H descriptor header ep + * + */ + tangErrorXdmaC2HDescHeaderEp = 3308, + + /** + * @brief xdma C2H descriptor unexpected comp + * + */ + tangErrorXdmaC2HDescUnexpectedComp = 3309, + + /** + * @brief xdma C2H timeout + * + */ + tangErrorXdmaC2HTimeout = 3310, + + /** + * @brief xdma C2H unknown + * + */ + tangErrorXdmaC2HUnknown = 3311, + + /** + * @brief xdma H2C align mismatch + * + */ + tangErrorXdmaH2CAlignMismatch = 3350, + + /** + * @brief xdma H2C invalid magic stopped + * + */ + tangErrorXdmaH2CInvaildMagicStopped = 3351, + + /** + * @brief xdma H2C invalid len + * + */ + tangErrorXdmaH2CInvalidLen = 3352, + + /** + * @brief xdma H2C read unsupport request + * + */ + tangErrorXdmaH2CReadUnSupportRequest = 3353, + + /** + * @brief xdma H2C read completer abort + * + */ + tangErrorXdmaH2CReadCompleterAbort = 3354, + + /** + * @brief xdma H2C read parity + * + */ + tangErrorXdmaH2CReadParity = 3355, + + /** + * @brief xdma H2C read header ep + * + */ + tangErrorXdmaH2CReadHeaderEp = 3356, + + /** + * @brief xdma H2C read unexpected comp + * + */ + tangErrorXdmaH2CReadUnExpectedComp = 3357, + + /** + * @brief xdma H2C decode error + * + */ + tangErrorXdmaH2CDecode = 3358, + + /** + * @brief xdma H2C slave + * + */ + tangErrorXdmaH2CSlave = 3359, + + /** + * @brief xdma H2C descriptor unsupport request + * + */ + tangErrorXdmaH2CDescUnSupportRequest = 3360, + + /** + * @brief xdma H2C descriptor completer abort + * + */ + tangErrorXdmaH2CDescCompleterAbort = 3361, + + /** + * @brief xdma H2C descriptor parity + * + */ + tangErrorXdmaH2CDescParity = 3362, + + /** + * @brief xdma H2C descriptor header ep + * + */ + tangErrorXdmaH2CDescHeaderEp = 3363, + + /** + * @brief xdma H2C descriptor unexpected com + * + */ + tangErrorXdmaH2CDescUnExpectedComp = 3364, + + /** + * @brief xdma H2C timeout + * + */ + tangErrorXdmaH2CTimeout = 3365, + + /** + * @brief xdma H2C unknown + * + */ + tangErrorXdmaH2CUnknown = 3366, + + /** + * This indicates that the IOCTL operation of TANG driver is failed. + * Added by TANG. hip and tang do not use. + * Need check, avoid to use the save number with cuda and hip. + */ + tangErrorDriverIoctlFailed = 10000, +}; + +/** + * TANG memory copy types + */ +typedef enum tangMemcpyKind { + tangMemcpyHostToHost = 0, /**< Host -> Host */ + tangMemcpyHostToDevice = 1, /**< Host -> Device */ + tangMemcpyDeviceToHost = 2, /**< Device -> Host */ + tangMemcpyDeviceToDevice = 3, /**< Device -> Device */ + /** + * Direction of the transfer is inferred from the pointer values. + * Requires unified virtual addressing, thus tang doesn't support. + */ + // tangMemcpyDefault = 4 +} tangMemcpyKind; + +enum tangMemoryType { + tangMemoryTypeUnregistered = 0, /**< Unregistered memory */ + tangMemoryTypeHost = 1, /**< Host memory */ + tangMemoryTypeDevice = 2, /**< Device memory */ + tangMemoryTypeManaged = 3, /**< Managed memory */ +}; + +struct tangPointerAttributes { + enum tangMemoryType type; + + int device; + + void* devicePointer; + + void* hostPointer; +}; + +/** + * TANG function attributes + */ +typedef struct tangFuncAttributes { + /** + * The size in bytes of statically-allocated shared memory per block + * required by this function. This does not include dynamically-allocated + * shared memory requested by the user at runtime. + */ + size_t sharedSizeBytes; + + /** + * The size in bytes of user-allocated constant memory required by this + * function. + */ + // PT devices use global memory to perform as constant memory, + // and constant memory belongs to the module, rather than the function + size_t constSizeBytes; + + /** + * The size in bytes of local memory used by each thread of this function. + */ + size_t localSizeBytes; + + /** + * The maximum number of threads per block, beyond which a launch of the + * function would fail. This number depends on both the function and the + * device on which the function is currently loaded. + */ + int maxThreadsPerBlock; + + /** + * The number of registers used by each thread of this function. + */ + int numRegs; + + /** + * The PTX virtual architecture version for which the function was + * compiled. This value is the major PTX version * 10 + the minor PTX + * version, so a PTX version 1.3 function would return the value 13. + */ + int ptxVersion; + + /** + * The binary architecture version for which the function was compiled. + * This value is the major binary version * 10 + the minor binary version, + * so a binary version 1.3 function would return the value 13. + */ + int binaryVersion; + + /** + * The attribute to indicate whether the function has been compiled with + * user specified option "-Xptxas --dlcm=ca" set. + */ + int cacheModeCA; + + /** + * The maximum size in bytes of dynamic shared memory per block for + * this function. Any launch must have a dynamic shared memory size + * smaller than this value. + */ + int maxDynamicSharedSizeBytes; + + /** + * On devices where the L1 cache and shared memory use the same hardware + * resources, this sets the shared memory carveout preference, in percent of + * the maximum shared memory. Refer to + * ::tangDevAttrSHARED_MEM_PER_MULTIPROCESSOR. This is only a hint, and the + * driver can choose a different ratio if required to execute the function. + * See ::tangFuncSetAttribute + * + * PT devices do not suppport to config L1 cache/shared memory. + */ + int preferredShmemCarveout; +} tangFuncAttributes; + +/** + * TANG function attributes that can be set using ::tangFuncSetAttribute + */ +typedef enum tangFuncAttribute { + tangFuncAttributeMaxDynamicSharedMemorySize = + 8, /**< Maximum dynamic shared memory size */ + tangFuncAttributePreferredSharedMemoryCarveout = + 9, /**< Preferred shared memory-L1 cache split */ + tangFuncAttributeMax +} tangFuncAttribute; + +/** + * TANG function cache configurations + * @warning On PT2 devices, L1 cache and shared memory are separated, + * thus these hints and controls are ignored. + */ +typedef enum tangFuncCache { + tangFuncCachePreferNone, ///< no preference for shared memory or L1 (default) + tangFuncCachePreferShared, ///< prefer larger shared memory and smaller L1 + ///< cache + tangFuncCachePreferL1, ///< prefer larger L1 cache and smaller shared memory + tangFuncCachePreferEqual, ///< prefer equal size L1 cache and shared memory +} tangFuncCache; + +/** + * TANG shared memory configuration + * @warning On PT2 devices, shard memory bank size is fix to 4-bytes, + * thus these hints and controls are ignored. + */ +typedef enum tangSharedMemConfig { + tangSharedMemBankSizeDefault, ///< The compiler selects a device-specific + ///< value for the banking. + tangSharedMemBankSizeFourByte, ///< Shared mem is banked at 4-bytes intervals + ///< and performs best when adjacent threads + ///< access data 4 bytes apart. + tangSharedMemBankSizeEightByte ///< Shared mem is banked at 8-byte intervals + ///< and performs best when adjacent threads + ///< access data 4 bytes apart. +} tangSharedMemConfig; + +/** + * TANG Limits + */ +enum tangLimit { + tangLimitStackSize = 0x00, /**< GPU thread stack size */ + tangLimitPrintfFifoSize = 0x01, /**< GPU printf FIFO size */ + tangLimitMallocHeapSize = 0x02, /**< GPU malloc heap size */ + tangLimitDevRuntimeSyncDepth = + 0x03, /**< GPU device runtime synchronize depth */ + tangLimitDevRuntimePendingLaunchCount = + 0x04, /**< GPU device runtime pending launch count */ + tangLimitMaxL2FetchGranularity = + 0x05 /**< A value between 0 and 128 that indicates the +// maximum fetch granularity of L2 (in Bytes). This is a hint */ +}; + +enum tangEventFlags_e { + tangEventDefault = 0x00, + tangEventDisableTiming = 0x02, + tangEventInterprocess = 0x04, +}; + +enum tangEventSyncFlags_e { + tangEventSyncDefault = 0x00, + + //!< Block until the event is recorded and done. + tangEventSyncRecordedAndDone = 0x01, +}; + +enum tangEventRecordFlags_e { + //!< The default recording mode. + tangEventRecordDefault = 0x00, + + //!< Always use the hardware event. + tangEventRecordHW = 0x0100, + + //!< Always use the software event. + tangEventRecordSW = 0x0200, + + //!< Allow to blockinig the calling thread is resource is not available. + tangEventRecordBlockingAllowed = 0x0400, +}; + +/** + * TANG Memory Advise values + */ +// enum tangMemoryAdvise {}; + +/** + * TANG range attributes + */ +// enum tangMemRangeAttribute {}; + +typedef enum tangStreamCaptureMode_e { + tangStreamCaptureModeGlobal = 0, + tangStreamCaptureModeThreadLocal = 1, + tangStreamCaptureModeRelaxed = 2, +} tangStreamCaptureMode; + +typedef enum tangStreamCaptureStatus_e { + tangStreamCaptureStatusNone = 0, + tangStreamCaptureStatusActive = 1, + tangStreamCaptureStatusInvalidated = 2, +} tangStreamCaptureStatus; + +/** + * TANG device attribute enum + */ +typedef enum tangDeviceAttr { + tangDevAttrMaxSharedMemPerBlock = 0, //!< sharedMemPerBlock + tangDevAttrMaxRegsPerBlock, //!< regsPerBlock + tangDevAttrWarpSize, //!< warpSize + tangDevAttrMemPitch, //!< memPitch + tangDevAttrMaxThreadsPerBlock, //!< maxThreadsPerBlock + tangDevAttrMaxBlockDimX, //!< maxThreadsDimX + tangDevAttrMaxBlockDimY, //!< maxThreadsDimY + tangDevAttrMaxBlockDimZ, //!< maxThreadsDimZ + tangDevAttrMaxGridDimX, //!< maxGridSizeX + tangDevAttrMaxGridDimY, //!< maxGridSizeY + tangDevAttrMaxGridDimZ, //!< maxGridSizeZ + tangDevAttrClockRate, //!< clockRate + tangDevAttrTotalConstantMemory, //!< totalConstMem + tangDevAttrMultiProcessorCount, //!< multiProcessorCount + tangDevAttrMaxBlocksPerMultiProcessor, //!< maxBlocksPerMultiProcessor + tangDevAttrAsyncEngineCount, //!< asyncEngineCount + tangDevAttrMemoryClockRate, //!< memoryClockRate + tangDevAttrGlobalMemoryBusWidth, //!< memoryBusWidth + tangDevAttrL2CacheSize, //!< l2CacheSize + tangDevAttrMaxThreadsPerMultiProcessor, //!< maxThreadsPerMultiProcessor + tangDevAttrGlobalL1CacheSupported, //!< globalL1CacheSupported + tangDevAttrLocalL1CacheSupported, //!< localL1CacheSupported + tangDevAttrMaxSharedMemoryPerMultiprocessor,//!< sharedMemPerMultiprocessor + tangDevAttrMaxRegistersPerMultiprocessor, //!< regsPerMultiprocessor + tangDevAttrStreamPrioritiesSupported, //!< streamPrioritiesSupported + tangDevAttrConcurrentKernels, //!< concurrentKernels + tangDevAttrComputePreemptionSupported, //!< computePreemptionSupported + tangDevAttrKernelExecTimeout, //!< kernelExecTimeoutEnabled + tangDevAttrEccEnabled, //!< ECCEnabled + tangDevAttrMaxAccessPolicyWindowSize, //!< accessPolicyMaxWindowSize + tangDevAttrTccDriver, //!< tccDriver + tangDevAttrSingleToDoublePrecisionPerfRatio,//!< singleToDoublePrecisionPerfRatio + tangDevAttrCooperativeLaunch, //!< cooperativeLaunch + tangDevAttrCooperativeMultiDeviceLaunch, //!< cooperativeMultiDeviceLaunch + tangDevAttrMaxPersistingL2CacheSize, //!< persistingL2CacheMaxSize + tangDevAttrCanMapHostMemory, //!< canMapHostMemory + tangDevAttrUnifiedAddressing, //!< unifiedAddressing + tangDevAttrManagedMemory, //!< managedMemory + tangDevAttrConcurrentManagedAccess, //!< concurrentManagedAccess + tangDevAttrDirectManagedMemAccessFromHost, //!< directManagedMemAccessFromHost + tangDevAttrPageableMemoryAccess, //!< pageableMemoryAccess + tangDevAttrPageableMemoryAccessUsesHostPageTables, //!< pageableMemoryAccessUsesHostPageTables + tangDevAttrCanUseHostPointerForRegisteredMem, //!< canUseHostPointerForRegisteredMem + tangDevAttrHostNativeAtomicSupported, //!< hostNativeAtomicSupported + tangDevAttrCanFlushRemoteWrites, //!< canFlushRemoteWrites + tangDevAttrGpuOverlap, //!< gpuOverlap + tangDevAttrIntegrated, //!< integrated + tangDevAttrMaxSharedMemoryPerBlockOptin, //!< maxSharedMemoryPerBlockOptin + tangDevAttrGPUDirectRDMASupported, //!< gpuDirectRDMASupported + tangDevAttrGPUDirectRDMAFlushWritesOptions, //!< gpuDirectRDMAFlushWritesOptions + tangDevAttrGPUDirectRDMAWritesOrdering, //!< gpuDirectRDMAWritesOrdering + tangDevAttrComputeCapabilityMajor, //!< major + tangDevAttrComputeCapabilityMinor, //!< minor + tangDevAttrPciBusId, //!< pciBusID + tangDevAttrPciDeviceId, //!< pciDeviceID + tangDevAttrPciDomainId, //!< pciDomainID + tangDevAttrIsMultiGpuBoard, //!< isMultiGpuBoard + tangDevAttrMultiGpuBoardGroupID, //!< multiGpuBoardGroupID + tangDevAttrComputeMode, //!< computeMode + tangDevAttrReservedSharedMemoryPerBlock, //!< reservedSharedMemoryPerBlock + tangDevAttrSparseTangArraySupported, //!< sparseTangArraySupported + tangDevAttrHostRegisterSupported, //!< hostRegisterSupported + tangDevAttrHostRegisterReadOnlySupported, //!< hostRegisterReadOnlySupported + tangDevAttrMemoryPoolsSupported, //!< memoryPoolsSupported + tangDevAttrMemoryPoolSupportedHandleTypes, //!< memoryPoolSupportedHandleTypes + tangDevAttrMax +} tangDeviceAttr; + +#ifndef TANGRT_DEVICE_P2P_ATTR_ENUM +#define TANGRT_DEVICE_P2P_ATTR_ENUM +/** + * TANG Device P2P attributes + */ +enum tangDeviceP2PAttr { + tangDevP2PAttrPerformanceRank = 1, + tangDevP2PAttrAccessSupported = 2, + tangDevP2PAttrNativeAtomicSupported = 3, + tangDevP2PAttrTangArrayAccessSupported = 4, +}; +#endif // TANGRT_DEVICE_P2P_ATTR_ENUM + +/** + * + * TANG device properties + * Inconsistent with cudaDeviceProp + * Most field unused. + * Need check. + */ +typedef struct tangDeviceProp { + char name[256]; ///< Device name. + char uuid[16]; ///< a 16-byte unique identifier + uint64_t totalGlobalMem; ///< size of global memory region (in bytes). + int sharedMemPerBlock; ///< the maximum amount of shared memory + ///< available to a thread block in bytes. + int regsPerBlock; ///< the maximum number of 32-bit registers available to a + ///< thread block. + int warpSize; ///< the warp size in threads. + int memPitch; ///< the maximum pitch in bytes allowed by the memory copy + ///< functions + int maxThreadsPerBlock; ///< the maximum number of threads per block. + int maxThreadsDim[3]; ///< Max number of threads in each dimension (XYZ) of a + ///< block. + int maxGridSize[3]; ///< Max grid dimensions (XYZ). + int clockRate; ///< Max clock frequency of the multiProcessors in khz. + int totalConstMem; ///< the total amount of constant memory available on + ///< the device in bytes. + int multiProcessorCount; ///< Number of multi-processors (compute units). + int maxBlocksPerMultiProcessor; ///< the number of multiprocessors on the + ///< device + int asyncEngineCount; ///< + int memoryClockRate; ///< Max global memory clock frequency in khz. + int memoryBusWidth; ///< Global memory bus width in bits. + int l2CacheSize; ///< L2 cache size. + int maxThreadsPerMultiProcessor; ///< Maximum resident threads per + ///< multi-processor. + int globalL1CacheSupported; ///< whether the device supports caching of + ///< globals in L1 cache + int localL1CacheSupported; ///< whether the device supports caching of locals + ///< in L1 cache + int sharedMemPerMultiprocessor; ///< Maximum Shared Memory Per Multiprocessor. + int regsPerMultiprocessor; ///< the maximum amount of shared memory available + ///< to a multiprocessor in bytes + int streamPrioritiesSupported; ///< whether the device supports stream + ///< priorities + int concurrentKernels; ///< Device can possibly execute multiple kernels + ///< concurrently. + int computePreemptionSupported; ///< whether the device supports Compute + ///< Preemption + int kernelExecTimeoutEnabled; ///< Run time limit for kernels executed on the + ///< device + int ECCEnabled; ///< Device has ECC support enabled + int accessPolicyMaxWindowSize; ///< the maximum value of + ///< tangAccessPolicyWindow::num_bytes + int tccDriver; ///< whether device is a Tesla device using TCC driver + int singleToDoublePrecisionPerfRatio; ///< the ratio of single precision + ///< performance (in floating-point + ///< operations per second) to double + ///< precision performance + int cooperativeLaunch; ///< whether the device supports launching cooperative + ///< kernels via tangLaunchCooperativeKernel + int cooperativeMultiDeviceLaunch; ///< whether the device supports launching + ///< cooperative kernels via + ///< tangLaunchCooperativeKernelMultiDevice + int persistingL2CacheMaxSize; ///< L2 cache's maximum persisting lines size + ///< in bytes + int canMapHostMemory; ///< Check whether TANG can map host memory + int unifiedAddressing; ///< whether the device shares a unified address space + ///< with the host and 0 otherwise + int managedMemory; ///< whether the device supports allocating managed memory + ///< on this system, or 0 if it is not supported + int concurrentManagedAccess; ///< whether the device can coherently access + ///< managed memory concurrently with the CPU + int directManagedMemAccessFromHost; ///< whether the host can directly access + ///< managed memory on the device without + ///< migration + int pageableMemoryAccess; ///< whether the device supports coherently + ///< accessing pageable memory without calling + ///< tangHostRegister on it + int pageableMemoryAccessUsesHostPageTables; ///< whether the device accesses + ///< pageable memory via the + ///< host's page tables + int canUseHostPointerForRegisteredMem; ///< whether the device can access + ///< host registered memory at the + ///< same virtual address as the CPU + + int hostNativeAtomicSupported; ///< Link between the device and the host + ///< supports native atomic operations + int canFlushRemoteWrites; ///< Device supports flushing of outstanding remote + ///< writes + int gpuOverlap; ///< Device can possibly copy memory and execute a kernel + ///< concurrently + int integrated; ///< Device is integrated with host memory + int maxSharedMemoryPerBlockOptin; ///< The maximum optin shared memory per + ///< block. This value may vary by chip. + ///< See ::tangFuncSetAttribute + int gpuDirectRDMASupported; ///< Device supports GPUDirect RDMA APIs + int gpuDirectRDMAFlushWritesOptions; ///< The returned attribute shall be + ///< interpreted as a bitmask, where the + ///< individual bits are listed in the + ///< ::tangFlushGPUDirectRDMAWritesOptions + ///< enum + int gpuDirectRDMAWritesOrdering; ///< GPUDirect RDMA writes to the device do + ///< not need to be flushed for consumers + ///< within the scope indicated by the + ///< returned attribute. See + ///< ::tangGPUDirectRDMAWritesOrdering for + ///< the numerical values returned here. + int major; ///< the major revision numbers defining the device's compute + ///< capability + int minor; ///< the minor revision numbers defining the device's compute + ///< capability + int pciBusID; ///< PCI Bus ID. + int pciDeviceID; ///< PCI Device ID. + int pciDomainID; ///< PCI Domain ID + int isMultiGpuBoard; ///< whether device is on a multi-GPU board. + int multiGpuBoardGroupID; ///< a unique identifier for a group of devices + ///< associated with the same board + int computeMode; ///< the compute mode that the device is currently in + int reservedSharedMemoryPerBlock; ///< Shared memory reserved by TANG driver + ///< per block in bytes + int sparseTangArraySupported; ///< Device supports sparse arrays and sparse + ///< mipmapped arrays + int hostRegisterSupported; ///< Device supports host memory registration via + ///< ::tangHostRegister + int hostRegisterReadOnlySupported; ///< Device supports using the + ///< ::tangHostRegister flag + ///< tangHostRegisterReadOnly to register + ///< memory that must be mapped as + ///< read-only to the GPU + int memoryPoolsSupported; ///< Device supports using the ::tangMallocAsync + ///< and ::tangMemPool family of APIs + int memoryPoolSupportedHandleTypes; ///< Handle types supported with mempool + ///< based IPC +} __attribute__((packed)) tangDeviceProp; + +/** + * TANG launch parameters + */ +/* +struct __device_builtin__ tangLaunchParams { + void* func; ///< Device function symbol + dim3 gridDim; ///< Grid dimentions + dim3 blockDim; ///< Block dimentions + void **args; ///< Arguments + int sharedMem; ///< Shared memory + tangStream_t stream; ///< Stream identifier +}; +*/ +/******************************************************************************* + * * + * SHORTHAND TYPE DEFINITION USED BY RUNTIME API * + * * + *******************************************************************************/ + +/** + * TANG Error types + */ +typedef enum tangError tangError_t; + +/** + * @brief TANG Device + * @sa TAdevice + */ +typedef struct TAdevice_s* tangDevice_t; + +/** + * @brief TANG context + * @sa TAcontext + */ +typedef struct TActx_s* tangContext_t; + +/** + * @brief TANG stream + * @sa TAstream + */ +typedef struct TAstream_s* tangStream_t; + +/** + * @brief TANG event + * @sa TAevent + */ +typedef struct TAevent_s* tangEvent_t; + +/** + * @brief TANG function + * @sa TAfunction + */ +typedef struct TAfunc_s* tangFunction_t; + +/** + * @brief TANG graph & executable graph handle + * @sa tangStreamBeginCapture + * @sa tangStreamEndCapture + * @sa tangGraphLaunch + * @sa tangGraphInstantiate + */ +typedef struct TAgraph_s* tangGraph_t; +typedef struct TAgraphExec_s* tangGraphExec_t; +typedef struct TAgraphNode_s* tangGraphNode_t; + +typedef void (*tangHostFn_t)(void* userData); + +typedef struct tangHostNodeParams_s { + tangHostFn_t fn; + void* userData; +} tangHostNodeParams; + +typedef struct tangKernelNodeParams_s { + void* func; /**< Kernel to launch */ + dim3 gridDim; /**< Grid dimensions */ + dim3 blockDim; /**< Block dimensions */ + + /**< Dynamic shared memory size per thread block in bytes */ + unsigned int sharedMemBytes; + + /**< Kernel parameters */ + void** kernelParams; + void** extra; +} tangKernelNodeParams; + +typedef struct tangGraphInfo_s { + int nr_nodes; +} tangGraphInfo; + +typedef struct tangEventTimestamp_s { + uint64_t comp; + uint64_t comp_sw; + uint64_t create; + uint64_t enqueue; + uint64_t writeq_beg; + uint64_t writeq_end; +} tangEventTimestamp; + +struct tangLanuchParams { + void* func; + dim3 gridDim; + dim3 blockDim; + void** args; + size_t sharedMemBytes; + tangStream_t stream; +}; + +/** + * @brief tangFatbinaryWrapper + * TANGCC will provides the following asm code on x86_64 platform. + * @code{.s} + * __tang_fatbin_wrapper: + * .long 0 # 0x0, version, 4 bytes + * .zero 4 # padding space, 4 bytes + * .quad .L_Z9vectorAddPfS_S_.11 # fatbin + * .zero 160 # data[20] + * .size __tang_fatbin_wrapper, 176 + * @endcode + */ +struct __tangFatbinaryWrapper { + int version; + const void* fatbin; + // The TANGCC does not reserve space for size. + // The size will be parsed from fatbin + // unsigned long size; + struct { + uintptr_t data[20]; + } dso; +}; + +#define TANG_IPC_HANDLE_SIZE 64U +#define TANG_IPC_MEM_HANDLE_SIZE 64U + +#define tangIpcMemLazyEnablePeerAccess 0x01 + +typedef struct tangIpcMemHandle_s { + unsigned long reserved[TANG_IPC_MEM_HANDLE_SIZE / sizeof(unsigned long)]; +} tangIpcMemHandle_t; + +typedef struct tangIpcEventHandle_s { + unsigned long reserved[TANG_IPC_HANDLE_SIZE / sizeof(unsigned long)]; +} tangIpcEventHandle_t; + +#endif //! _TANG_RT_DRIVER_TYPES_H_ diff --git a/third_party/sunrise/backend/include/tang_rt/fmt.hpp b/third_party/sunrise/backend/include/tang_rt/fmt.hpp new file mode 100755 index 0000000000..1928ae095c --- /dev/null +++ b/third_party/sunrise/backend/include/tang_rt/fmt.hpp @@ -0,0 +1,1097 @@ +#ifndef _TANGRT_FMT_HPP_ +#define _TANGRT_FMT_HPP_ + +#include + +#include +#include +#include + +// #define PT_PRINTF_ENDMARKER +#define PT_PRINTF_READY + +#if !defined(__TANGC_MAJOR__) && !defined(__device__) +#define __device__ +#endif //!< __TANGC_MAJOR__ + +#ifndef unlikely +#define unlikely(x) __builtin_expected(!!(x), 0) +#endif //!< unlikely + +#ifndef lower_32_bit +#define lower_32_bit(x) (((uint64_t)x) & 0xffffffff) +#endif + +#ifndef upper_32_bit +#define upper_32_bit(x) (((uint64_t)x) >> 32) +#endif + +namespace tangrt { +namespace fmt { + +static constexpr unsigned kArgAlignment = 4; +static constexpr unsigned kArgAlignmentMask = ~(kArgAlignment - 1); + +enum ArgId { + ArgId_None = 0, + + ArgId_char = 1, + ArgId_schar = 2, + ArgId_uchar = 3, + + ArgId_short = 5, + ArgId_ushort = 6, + + ArgId_int = 7, + ArgId_uint = 8, + + ArgId_long = 9, + ArgId_ulong = 10, + + ArgId_longlong = 11, + ArgId_ulonglong = 12, + + ArgId_float = 13, + ArgId_double = 14, + ArgId_long_double = 15, + + //!< nullptr + ArgId_nullptr = 20, + + //!< generic pointer type. + ArgId_pointer = 21, + + //!< char* ptr = nullptr; + ArgId_char_nullptr = 22, + ArgId_schar_nullptr = 23, + ArgId_uchar_nullptr = 24, + + ArgId_char_pointer = 25, + ArgId_schar_pointer = 26, + ArgId_uchar_pointer = 27, + + //!< char[] + ArgId_char_array = 28, + ArgId_schar_array = 29, + ArgId_uchar_array = 30, +}; + +static inline const char* GetArgIdName(const int id) { +#define _case(x, str) \ + case x: \ + return str + + switch (id) { + _case(ArgId_char, "char"); + _case(ArgId_schar, "signed char"); + _case(ArgId_uchar, "unsigned char"); + + _case(ArgId_short, "short"); + _case(ArgId_ushort, "unsigned short"); + _case(ArgId_int, "int"); + _case(ArgId_uint, "unsigned int"); + _case(ArgId_long, "long"); + _case(ArgId_ulong, "unsigned long"); + _case(ArgId_longlong, "long long"); + _case(ArgId_ulonglong, "unsigned long long"); + + _case(ArgId_float, "float"); + _case(ArgId_double, "double"); + _case(ArgId_long_double, "long double"); + + _case(ArgId_nullptr, "nullptr"); + _case(ArgId_pointer, "pointer"); + + _case(ArgId_char_nullptr, "char nullptr"); + _case(ArgId_schar_nullptr, "schar nullptr"); + _case(ArgId_uchar_nullptr, "uchar nullptr"); + + _case(ArgId_char_pointer, "char pointer"); + _case(ArgId_schar_pointer, "schar pointer"); + _case(ArgId_uchar_pointer, "uchar pointer"); + + _case(ArgId_char_array, "char array"); + _case(ArgId_schar_array, "schar array"); + _case(ArgId_uchar_array, "uchar array"); + default: + return "None"; + } +#undef _case +} + +static __device__ inline unsigned ArgAlign(unsigned int x) { + return (x + kArgAlignment - 1) & ~(kArgAlignment - 1); +} + +static __device__ inline unsigned int ArgTraitsStrLen(const char* s) { + const char* p = s; + while (*p) { + ++p; + } + return p - s; +} + +namespace detail { + +union u32c4_u { + char c[4]; + uint32_t u32; + + __device__ u32c4_u(char ch0, char ch1 = 0, char ch2 = 0, char ch3 = 0) + : c{ch0, ch1, ch2, ch3} {} +}; + +template +__device__ void FundamentalFillImpl(const uint32_t id, + const T t, + uint32_t* buf, + uint32_t& pos, + const uint32_t mask) { + static_assert(N == 1 || N == 2 || N == 4, ""); + static_assert(!std::is_same::value, + "This helper function is not suitable for char*."); + + union { + uint32_t u[N]; + T t; + } x; + x.t = t; + buf[pos++ & mask] = (sizeof(T) << 16) | id; + for (unsigned int i = 0; i < N; ++i) { + buf[pos++ & mask] = x.u[i]; + } +} + +template +__device__ void FundamentalFill(const uint32_t id, + const T t, + uint32_t* buf, + uint32_t& pos, + const uint32_t mask) { + FundamentalFillImpl(id, t, buf, pos, mask); +} + +#if 0 +template <> +__device__ void FundamentalFill(const uint32_t id, const float t, uint32_t* buf, + uint32_t& pos, const uint32_t mask) +{ + buf[pos++ & mask] = (sizeof(float) << 16) | id; + union + { + uint32_t u[2]; + float f; + } x; + x.f = t; + buf[pos++ & mask] = x.u[0]; +} + +//!< to avoid strict aliasing +//!< -fno-strict-aliasing +template <> +__device__ void FundamentalFill(const uint32_t id, const double t, uint32_t* buf, + uint32_t& pos, const uint32_t mask) +{ + buf[pos++ & mask] = (sizeof(double) << 16) | id; + union + { + uint32_t u[2]; + double d; + } x; + x.d = t; + buf[pos++ & mask] = x.u[0]; + buf[pos++ & mask] = x.u[1]; +} +#endif + +static __device__ inline void StringFillData(const char* s, + const uint32_t sizeBytes, + uint32_t* const buf, + uint32_t& pos, + const uint32_t mask) { + const uint32_t* s32 = (const uint32_t*)s; + for (unsigned int i = 0; i < sizeBytes / sizeof(uint32_t); ++i) { + buf[pos++ & mask] = s32[i]; + } + switch (sizeBytes & (4 - 1)) { + case 1: { + detail::u32c4_u x(s[sizeBytes - 1]); + buf[pos++ & mask] = x.u32; + break; + }; + case 2: { + detail::u32c4_u x(s[sizeBytes - 2], s[sizeBytes - 1]); + buf[pos++ & mask] = x.u32; + break; + }; + case 3: { + detail::u32c4_u x(s[sizeBytes - 3], s[sizeBytes - 2], s[sizeBytes - 1]); + buf[pos++ & mask] = x.u32; + break; + } + } +} // StringFillData function + +} // namespace detail + +template +struct FmtTraits; + +template +struct FmtTraits { + static __device__ unsigned int FmtLength(const char*) { + return 4 + ArgAlign(N); + } + + static __device__ void Fill(const char* s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + //!< Fill header + buf[pos++ & mask] = (N << 16) | ArgId_schar_array; + detail::StringFillData(s, N, buf, pos, mask); + } +}; + +template +struct FmtTraits { + static __device__ unsigned int FmtLength(const char*) { + return 4 + ArgAlign(N); + } + + static __device__ void Fill(const char* s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + //!< Fill header + buf[pos++ & mask] = (N << 16) | ArgId_schar_array; + detail::StringFillData(s, N, buf, pos, mask); + } +}; + +template <> +struct FmtTraits { + static __device__ unsigned int FmtLength(const char* s) { + return 4 + (s ? ArgAlign(ArgTraitsStrLen(s)) : 0); + } + + static __device__ void Fill(const char* s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + if (!s) { + buf[pos++ & mask] = ArgId_char_nullptr; + return; + } + const uint32_t N = ArgTraitsStrLen(s); + detail::StringFillData(s, N, buf, pos, mask); + } +}; + +template <> +struct FmtTraits { + static __device__ unsigned int FmtLength(const char* s) { + return 4 + (s ? ArgAlign(ArgTraitsStrLen(s)) : 0); + } + + static __device__ void Fill(const char* s, + uint32_t* buf, + uint32_t pos, + uint32_t mask) { + if (!s) { + buf[pos++ & mask] = ArgId_char_nullptr; + return; + } + const uint32_t N = ArgTraitsStrLen(s); + detail::StringFillData(s, N, buf, pos, mask); + } +}; + +template +struct ArgTraits; + +// template +// struct ArgTraits { +// static const int id = ArgId_None; +// +// static unsigned int ArgLength(T&& t) { +// return 4 + sizeof(T); +// } +//}; + +template <> +struct ArgTraits { + static const int id = ArgId_char; + + static __device__ unsigned int ArgLength(const signed char t) { + return ArgAlign(4); + } + + static __device__ void Fill(const char ch, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + buf[pos++ & mask] = ch << 16 | id; + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_schar; + + static __device__ unsigned int ArgLength(const signed char t) { + return ArgAlign(sizeof(t)); + } + + static __device__ void Fill(const signed char ch, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + buf[pos++ & mask] = ch << 16 | id; + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_uchar; + + static __device__ unsigned int ArgLength(const unsigned char t) { + return ArgAlign(sizeof(t)); + } + + static __device__ void Fill(const unsigned char ch, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + buf[pos++ & mask] = ch << 16 | id; + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_short; + + static __device__ unsigned int ArgLength(const short int s) { + return ArgAlign(sizeof(s)); + } + + static __device__ void Fill(const short int ch, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + buf[pos++ & mask] = (ch << 16) | id; + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_ushort; + + static __device__ unsigned int ArgLength(const unsigned short int s) { + return ArgAlign(sizeof(s)); + } + + static __device__ void Fill(const unsigned short ch, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + buf[pos++ & mask] = (ch << 16) | id; + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_int; + + static __device__ unsigned int ArgLength(const int i) { + return 4 + sizeof(i); + } + + static __device__ void Fill(const int s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_int, s, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_uint; + + static __device__ unsigned int ArgLength(const unsigned int i) { + return 4 + sizeof(i); + } + + static __device__ void Fill(const unsigned int s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_uint, s, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const long id = ArgId_long; + + static __device__ unsigned int ArgLength(const long l) { + return 4 + sizeof(l); + } + + static __device__ void Fill(const long s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_long, s, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_ulong; + + static __device__ unsigned int ArgLength(const unsigned long l) { + return 4 + sizeof(l); + } + + static __device__ void Fill(const unsigned long s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_ulong, s, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const long id = ArgId_longlong; + + static __device__ unsigned int ArgLength(const long long l) { + return 4 + sizeof(l); + } + + static __device__ void Fill(const long long s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_longlong, s, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_ulonglong; + + static __device__ unsigned int ArgLength(const unsigned long long l) { + return sizeof(l) + 4; + } + + static __device__ void Fill(const unsigned long long s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_ulonglong, s, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_float; + + static __device__ unsigned int ArgLength(const float l) { + return sizeof(l) + 4; + } + + static __device__ void Fill(const float s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_float, s, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_double; + + static __device__ unsigned int ArgLength(const double d) { + return sizeof(d) + 4; + } + + static __device__ void Fill(const double s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_double, s, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_long_double; + + static __device__ unsigned int ArgLength(const long double d) { + return sizeof(d) + 4; + } + + static __device__ void Fill(const long double s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_long_double, s, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_nullptr; + + static __device__ unsigned int ArgLength(...) { return 4; } + + static __device__ void Fill(const std::nullptr_t, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + buf[pos++ & mask] = ArgId_nullptr; + } +}; + +template +struct ArgTraits { + static const int id = ArgId_pointer; + + static __device__ unsigned int ArgLength(...) { return 4 + sizeof(T*); } + + static __device__ void Fill(T* const s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_pointer, s, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_char_pointer; + + static __device__ unsigned int ArgLength(const char* s) { + return 4 + (s ? ArgAlign(ArgTraitsStrLen(s)) + 8 : 0); + } + + static __device__ void Fill(const char* s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + if (!s) { + buf[pos++ & mask] = ArgId_char_nullptr; + return; + } + const uint32_t N = ArgTraitsStrLen(s); + buf[pos++ & mask] = (N << 16) | id; + buf[pos++ & mask] = lower_32_bit(s); + buf[pos++ & mask] = upper_32_bit(s); + detail::StringFillData(s, N, buf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_char_pointer; + + static __device__ unsigned int ArgLength(const char* s) { + return 4 + (s ? ArgAlign(ArgTraitsStrLen(s)) + 8 : 0); + } + + static __device__ void Fill(const char* s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + if (!s) { + buf[pos++ & mask] = ArgId_char_nullptr; + return; + } + const uint32_t N = ArgTraitsStrLen(s); + buf[pos++ & mask] = (N << 16) | id; + buf[pos++ & mask] = lower_32_bit(s); + buf[pos++ & mask] = upper_32_bit(s); + detail::StringFillData(s, N, buf, pos, mask); + } +}; + +template +struct ArgTraits { + static const int id = ArgId_char_array; + + static __device__ unsigned int ArgLength(const char* s) { + return 4 + ArgAlign(N) + 8; + } + + static __device__ void Fill(const char* s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + buf[pos++ & mask] = (N << 16) | id; + buf[pos++ & mask] = lower_32_bit(s); + buf[pos++ & mask] = upper_32_bit(s); + detail::StringFillData(s, N, buf, pos, mask); + } +}; + +template +struct ArgTraits { + static const int id = ArgId_schar_array; + + static __device__ unsigned int ArgLength(const char* s) { + return 4 + ArgAlign(N) + 8; + } + + static __device__ void Fill(const char* s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + buf[pos++ & mask] = (N << 16) | id; + buf[pos++ & mask] = lower_32_bit(s); + buf[pos++ & mask] = upper_32_bit(s); + detail::StringFillData(s, N, buf, pos, mask); + } +}; + +template +struct ArgTraits { + static const int id = ArgId_uchar_array; + + static __device__ unsigned int ArgLength(const char* s) { + return 4 + ArgAlign(N) + 8; + } + + static __device__ void Fill(const char* s, + uint32_t* buf, + uint32_t& pos, + uint32_t mask) { + buf[pos++ & mask] = (N << 16) | id; + buf[pos++ & mask] = lower_32_bit(s); + buf[pos++ & mask] = upper_32_bit(s); + detail::StringFillData(s, N, buf, pos, mask); + } +}; + +// template +// struct ArgTraits { +// using type = T; +// +// static const int id = ArgId_any_array; +//}; + +// template +// struct ArgTraits { +// static const int id = ArgId_uchar_array; +//}; + +namespace detail { +template +__device__ unsigned SumArgsLength(Args&&... args); + +template <> +__device__ constexpr unsigned SumArgsLength() { + return 0; +} + +template +__device__ unsigned SumArgsLength(T&& t, Args&&... args) { + typedef typename std::remove_reference::type _type; + typedef typename std::remove_cv<_type>::type type; + + return ArgTraits::ArgLength(std::forward(t)) + + SumArgsLength(std::forward(args)...); +} + +struct FifoInfo { + uint32_t get; + uint32_t put; + + //!< num words + uint32_t size; + + uint32_t fifoSize; + uint64_t fifoAddress; +}; + +template +struct ArgTraitsHasFillFifoInfo { + template + static auto Check(int) -> decltype(&U::FillFifoInfo); + + template + static void Check(...); + + static const bool value = !std::is_same(0)), void>::value; +}; + +template +struct ArgTraitsFillProxy; + +template +struct ArgTraitsFillProxy { + template + __device__ static void Fill(U&& u, + uint32_t* fifobuf, + const FifoInfo& msgInfo, + uint32_t& pos, + uint32_t& mask) { + ArgTraits::FillFifoInfo(std::forward(u), fifobuf, msgInfo, pos, mask); + } +}; + +template +struct ArgTraitsFillProxy { + template + __device__ static void Fill(U&& u, + uint32_t* fifobuf, + const FifoInfo& msgInfo, + uint32_t& pos, + uint32_t& mask) { + ArgTraits::Fill(std::forward(u), fifobuf, pos, mask); + } +}; + +template +__device__ void FillArgs(uint32_t* fifobuf, + const FifoInfo& msgInfo, + uint32_t& pos, + uint32_t mask, + Args&&... args); + +template <> +__device__ inline void FillArgs(uint32_t* fifobuf, + const FifoInfo& msgInfo, + uint32_t& pos, + uint32_t mask) {} + +template +__device__ void FillArgs(uint32_t* fifobuf, + const FifoInfo& msgInfo, + uint32_t& pos, + uint32_t mask, + T&& t, + Args&&... args) { + typedef typename std::remove_reference::type _type; + typedef typename std::remove_cv<_type>::type type; + + ArgTraitsFillProxy>::value>:: + Fill(t, fifobuf, msgInfo, pos, mask); + //ArgTraits::Fill(t, fifobuf, pos, mask); + FillArgs(fifobuf, msgInfo, pos, mask, std::forward(args)...); +} + +template +struct CountOfArgs; + +template <> +struct CountOfArgs<> { + static const int value = 0; +}; + +template +struct CountOfArgs { + static const int value = CountOfArgs::value + 1; +}; + +} // namespace detail + +namespace debug { + +//!< The get the __pt_printf load. +struct MsgGet {}; + +//!< The begin position of the current __pt_printf. +struct MsgBeg {}; + +//!< The number words the current __pt_printf will consume. +struct MsgSize {}; + +//!< Print the begin address of the current __pt_printf fifo. +struct FifoAddress {}; + +//!< Print the size of the print fifo. +struct FifoSize {}; + +} // namespace debug + +template <> +struct ArgTraits { + static const int id = ArgId_uint; + + __device__ static unsigned int ArgLength(debug::MsgBeg put) { return 8; } + + __device__ static void FillFifoInfo(debug::MsgBeg, + uint32_t* fifobuf, + const detail::FifoInfo& msgInfo, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_uint, msgInfo.put, fifobuf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_uint; + + __device__ static unsigned int ArgLength(debug::MsgGet put) { return 8; } + + __device__ static void FillFifoInfo(debug::MsgGet, + uint32_t* fifobuf, + const detail::FifoInfo& msgInfo, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_uint, msgInfo.get, fifobuf, pos, mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_uint; + + __device__ static unsigned int ArgLength(debug::MsgSize s) { return 8; } + + __device__ static void FillFifoInfo(debug::MsgSize, + uint32_t* fifobuf, + const detail::FifoInfo& msgInfo, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_uint, + msgInfo.size, + fifobuf, + pos, + mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_uint; + + __device__ static unsigned int ArgLength(debug::FifoSize s) { return 8; } + + __device__ static void FillFifoInfo(debug::FifoSize, + uint32_t* fifobuf, + const detail::FifoInfo& msgInfo, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_uint, + msgInfo.fifoSize, + fifobuf, + pos, + mask); + } +}; + +template <> +struct ArgTraits { + static const int id = ArgId_pointer; + + __device__ static unsigned int ArgLength(debug::FifoAddress s) { return 12; } + + __device__ static void FillFifoInfo(debug::FifoAddress, + uint32_t* fifobuf, + const detail::FifoInfo& msgInfo, + uint32_t& pos, + uint32_t mask) { + detail::FundamentalFill(ArgId_pointer, + msgInfo.fifoAddress, + fifobuf, + pos, + mask); + } +}; + +struct fifo { + unsigned int put __attribute__((aligned(128))); + unsigned int mask; + uint32_t* data; + + unsigned int ready __attribute__((aligned(128))); + + unsigned int get __attribute__((aligned(128))); +}; + +extern "C" struct fifo __ptPrintfFifo; + +inline __device__ struct fifo* __ptSelectPrintfFifo(void) { + //unsigned int bidx = threadIdx.z * (blockDim.x * blockDim.y) + + // threadIdx.y * blockDim.x + threadIdx.x; + //return &__ptPrintfFifo + (widx / 32) & 0x01; +#ifdef __TANGC_MAJOR__ + return &__ptPrintfFifo + (__phywarpid() & 0x01); +#else + return &__ptPrintfFifo; +#endif //!< __TANGC_MAJOR__ +} + +inline __device__ uint32_t __ptPrintfFifoAlloc(struct fifo* fifo, + uint32_t n, + uint32_t* pGet) { + [[maybe_unused]] unsigned int tmp; + unsigned int newPut; + unsigned int oldPut; + unsigned int avail; + unsigned int get; + unsigned int mask = fifo->mask; + + if (n >= mask) { + return std::numeric_limits::max(); + } + do { +#ifdef __TANGC_MAJOR__ + oldPut = *((volatile unsigned int*)&fifo->put); + // oldPut = __ldcg(&fifo->put); +#else + oldPut = __atomic_load_n(&fifo->put, __ATOMIC_RELAXED); +#endif //!< __TANGC_MAJOR__ + +//#ifdef __TANGC_MAJOR__ +// __threadfence_memory(); +//#endif //!< __TANGC_MAJOR__ + +#ifdef __TANGC_MAJOR__ + // get = *((volatile unsigned int*)&fifo->get); + get = __ldcg(&fifo->get); +#else + get = __atomic_load_n(&fifo->get, __ATOMIC_RELAXED); +#endif //!< __TANGC_MAJOR__ + + // avail = mask - ((oldPut - get) & mask); + avail = (get - oldPut - 1) & mask; + if (avail < n) { + return std::numeric_limits::max(); + } + newPut = oldPut + n; + // newPut = (oldPut + n) & mask; +#ifdef __TANGC_MAJOR__ + tmp = atomicCAS(&fifo->put, oldPut, newPut); + // Compiler group provides this solution. +# if 1 + asm volatile("loop 1, 0, 1, 500000\n\tnop"); +# else + __stvm_bar_sync0(); +# endif + } while (oldPut != tmp); +#else + } while (!__atomic_compare_exchange_n(&fifo->put, + &oldPut, + newPut, + true, + __ATOMIC_RELAXED, + __ATOMIC_RELAXED)); +#endif //!< __TANGC_MAJOR__ + + *pGet = get; + return oldPut; +} + +inline __device__ void __ptPrintfFifoUpdateReady(struct fifo* fifo, + uint32_t orig_pos, + uint32_t pos) { + [[maybe_unused]] unsigned int tmp; + unsigned int oldReady = orig_pos; +#ifdef __TANGC_MAJOR__ + //! Make sure all writes before this call happens before + //! all writes after this call. + //! __syncthreads(); + __threadfence_memory(); +#endif //!< __TANGC_MAJOR__ + do { +#ifdef __TANGC_MAJOR__ + tmp = atomicCAS(&fifo->ready, oldReady, pos); +# if 1 + asm volatile("loop 1, 0, 1, 500000\n\tnop"); +# else + __stvm_bar_sync0(); +# endif + } while (oldReady != tmp); +#else + } while (!__atomic_compare_exchange_n(&fifo->ready, + &oldReady, + pos, + true, + __ATOMIC_RELEASE, + __ATOMIC_RELAXED)); +#endif //!< __TANGC_MAJOR__ +} + +template +__device__ void __pt_printf(Fmt&& fmt, Args&&... args) { + typedef typename std::remove_reference::type _fmt_type; + typedef typename std::remove_cv<_fmt_type>::type fmt_type; + + auto fifo = __ptSelectPrintfFifo(); + + // struct fifo = __ptPrintfFifo; + // numWords: the number bytes of the message; + // countOfArgs: the number of args + // fmt: fmt data + // arg[numArgs]: arg data + // endMarker: maybe not required + + unsigned int numFmtWords = FmtTraits::FmtLength(fmt) / 4; + unsigned int numArgWords = + detail::SumArgsLength(std::forward(args)...) / 4; + + static_assert(detail::CountOfArgs::value == sizeof...(Args), ""); + + // unsigned int countOfArgs = detail::CountOfArgs::value; + unsigned int countOfArgs = sizeof...(Args); + +#ifdef PT_PRINTF_ENDMARKER + unsigned int numMsgWords = 3 + numFmtWords + numArgWords; +#else + unsigned int numMsgWords = 2 + numFmtWords + numArgWords; +#endif //!< PT_PRINTF_ENDMARKER + + // align message to 128byte, 32uint32_t boundary. +#if 0 + numMsgWords = (numMsgWords + 31) & ~31; +#endif + + uint32_t get; + uint32_t pos = __ptPrintfFifoAlloc(fifo, numMsgWords, &get); + if (pos == std::numeric_limits::max()) { + return; + } + + const detail::FifoInfo fifoInfo = { + .get = get, + .put = pos, + .size = numMsgWords, + .fifoSize = fifo->mask + 1, + .fifoAddress = (uint64_t)fifo->data, + }; + + uint32_t const orig_pos = pos; + uint32_t const mask = fifo->mask; + uint32_t* const fifobuf = (uint32_t*)fifo->data; + + fifobuf[pos++ & mask] = numMsgWords; + fifobuf[pos++ & mask] = countOfArgs; + + FmtTraits::Fill(std::forward(fmt), fifobuf, pos, mask); + detail::FillArgs(fifobuf, fifoInfo, pos, mask, std::forward(args)...); + +#ifdef PT_PRINTF_ENDMARKER + fifobuf[pos++ & mask] = numMsgWords; +#endif //!< PT_PRINTF_ENDMARKER +#if 0 + while ((pos - orig_pos) < numMsgWords) { + fifobuf[pos++ & mask] = orig_pos; + } +#endif + + __ptPrintfFifoUpdateReady(fifo, orig_pos, (orig_pos + numMsgWords)); +} + +} // namespace fmt +} // namespace tangrt + +using tangrt::fmt::__pt_printf; + +#endif //!< _TANGRT_FMT_HPP_ diff --git a/third_party/sunrise/backend/include/tang_rt/host_defines.h b/third_party/sunrise/backend/include/tang_rt/host_defines.h new file mode 100755 index 0000000000..f7639fb98b --- /dev/null +++ b/third_party/sunrise/backend/include/tang_rt/host_defines.h @@ -0,0 +1,101 @@ +/* +Copyright declaration. +*/ + +#ifndef _TANG_RT_INCLUDE_HOST_DEFINES_H_ +#define _TANG_RT_INCLUDE_HOST_DEFINES_H_ +#include +#include + +#ifdef __cplusplus +#define __dparm(x) = x +#else +#define __dparm(x) +#endif + +#ifdef __TANGC_MAJOR__ + +#ifndef __device__ +#define __device__ __Tdevice__ +#define __Tdevice__ __attribute__((Tdevice)) +#endif + +#ifndef __global__ +#define __global__ __Tglobal__ +#define __Tglobal__ __attribute__((Tglobal)) +#endif + +#ifndef __constant__ +#define __constant__ __Tconstant__ +#define __Tconstant__ __attribute__((Tconstant)) +#endif + +#ifndef __host__ +#define __host__ __Thost__ +#define __Thost__ __attribute__((Thost)) +#endif + +#ifndef __shared__ +#define __shared__ __Tshared__ +#define __Tshared__ __attribute__((Tshared)) +#endif + +#ifndef __forceinline__ +#define __forceinline__ __inline__ __attribute__((always_inline)) +#endif + +#endif + +#if defined(_MSC_VER) +#define TANGRT_DEPRECATED __declspec(deprecated) +#define TANGRT_API_EXPORT __declspec(dllexport) +#define TANGRT_API_IMPORT __declspec(dllimport) +#elif defined(__GNUC__) || defined(__clang__) +#define TANG_DEPRECATED __attribute__((deprecated)) +#define TANG_API_EXPORT __attribute__((visibility("default"))) +#define TANG_API_IMPORT __attribute__((visibility("default"))) +#else +#define TANG_DEPRECATED +#define TANG_API_EXPORT +#define TANG_API_IMPORT +#endif // unknown compiler, may needs extra care. + +#if defined(tangrt_shared_EXPORTS) +#define TANGRT_API_PUBLIC TANG_API_EXPORT +#elif !defined(__TANGRT_API_VERSION_INTERNAL) +#define TANGRT_API_PUBLIC TANG_API_IMPORT +#else +#define TANGRT_API_PUBLIC +#endif + +/************************************************** + * _ptds: Per-Thread-Default-Stream API use ptds to + * run commands. + * _ptsz suffix: Per-Thread-Stream-Zero API use ptds to + * run commands when the given stream is null. + * See the following code for details: + * @code + * tangError_t tangMemcpyAsync_ptsz(..., tangStream_t stream) { + * return tangMemcpyAsyncImpl(..., stream ? stream : TA_STREAM_PER_THREAD); + * } + * tangError_t tangMemcpyAsync(..., tangStream_t stream) { + * return tangMemcpyAsyncImpl(..., stream ? stream : TA_STREAM_LEGACY); + * } + * @endcode + **************************************************/ +#if defined(__TANGRT_API_PER_THREAD_DEFAULT_STREAM) +#define __TANGRT_API_PTDS(api) api##_ptds +#define __TANGRT_API_PTSZ(api) api##_ptsz +#else +#define __TANGRT_API_PTDS(api) api +#define __TANGRT_API_PTSZ(api) api +#endif //! __TANGRT_API_PER_THREAD_DEFAULT_STREAM + +#if defined(__TANGRT_API_VERSION_INTERNAL) +#undef __TANGRT_API_PTDS +#undef __TANGRT_API_PTSZ +#define __TANGRT_API_PTDS(api) api +#define __TANGRT_API_PTSZ(api) api +#endif // __TANGRT_API_VERSION_INTERNAL + +#endif //! _TANG_RT_INCLUDE_HOST_DEFINES_H_ diff --git a/third_party/sunrise/backend/include/tang_rt/vector_types.h b/third_party/sunrise/backend/include/tang_rt/vector_types.h new file mode 100755 index 0000000000..e4aca338ce --- /dev/null +++ b/third_party/sunrise/backend/include/tang_rt/vector_types.h @@ -0,0 +1,35 @@ +/* +Copyright declaration. +*/ + +// cuda/include/vector_types.h + +#ifndef _TANG_RT_INCLUDE_VECTOR_TYPES_H_ +#define _TANG_RT_INCLUDE_VECTOR_TYPES_H_ + +#include "tang_rt/host_defines.h" + +/** + * Struct for data in 3D + */ + +#if defined(__DIM3_TYPE__) +typedef dim3 __DIM3_TYPE__; +#else +typedef struct dim3 { + unsigned x; ///< x + unsigned y; ///< y + unsigned z; ///< z +#ifdef __cplusplus +#if __cplusplus >= 201103L + constexpr dim3(unsigned _x = 1, unsigned _y = 1, unsigned _z = 1) + : x(_x), y(_y), z(_z) {} +#else + dim3(unsigned _x = 1, unsigned _y = 1, unsigned _z = 1) + : x(_x), y(_y), z(_z) {} +#endif //! __cplusplus >= 201103 +#endif //! __cplusplus +} dim3; +#endif //! no __DIM3_TYPE__ + +#endif //! _TANG_RT_INCLUDE_VECTOR_TYPES_H_ diff --git a/third_party/sunrise/backend/include/tang_rt/version.h b/third_party/sunrise/backend/include/tang_rt/version.h new file mode 100755 index 0000000000..3ac846e765 --- /dev/null +++ b/third_party/sunrise/backend/include/tang_rt/version.h @@ -0,0 +1,30 @@ +#ifndef _TANG_RUNTIME_VERSION_H_ +#define _TANG_RUNTIME_VERSION_H_ +#define TANG_VERSION_MAJOR 0 +#define TANG_VERSION_MINOR 13 +#define TANG_VERSION_PATCH 0 + +#define TANG_VERSION_GIT_SHA "" + +///////////////////////////////////////////////////////// + +#define TANGRT_VERSION_MAJOR 0 +#define TANGRT_VERSION_MINOR 13 +#define TANGRT_VERSION_PATCH 0 + +#define TANGRT_VERSION_GIT_SHA "04137493 Merge branch 'ln/bugfix/taStreamIsCapturing' into 'master'" + +///////////////////////////////////////////////////////// +#define TANGRT_TANGCC_VERSION_MAJOR 2 +#define TANGRT_TANGCC_VERSION_MINOR 2 + +#ifdef __TANGC_MAJOR__ +# if (TANGRT_TANGCC_VERSION_MAJOR <= 1) && (__TANGC_MAJOR__ >= 2) +#warning "the ptcc used is not compatible with the tang runtime library\nptcc less than 2.0.0 is required." +//#error "the ptcc used is not compatible with the tang runtime library\nptcc less than 2.0.0 is required." +//# elif (TANGRT_TANGCC_VERSION_MAJOR >= 2) && (__TANGC_MAJOR__ <= 1) +//#error "the ptcc used is not compatible with the tang runtime library\nptcc 2.0.0 or later is required." +# endif +#endif // __TANGC_MAJOR__ + +#endif //! _TANG_RUNTIME_VERSION_H_ diff --git a/third_party/sunrise/backend/include/tang_runtime.h b/third_party/sunrise/backend/include/tang_runtime.h new file mode 100755 index 0000000000..3a58f7a022 --- /dev/null +++ b/third_party/sunrise/backend/include/tang_runtime.h @@ -0,0 +1,32 @@ +#ifndef _TANG_RUNTIME_H_ +#define _TANG_RUNTIME_H_ +#include "tang_rt/version.h" +#include "tang_rt/driver_types.h" +#include "tang_rt/vector_types.h" +#include "tang_runtime_api.h" + +#ifndef TA_STREAM_LEGACY +#define TA_STREAM_LEGACY ((tangStream_t)0x01) +#endif //! TA_STREAM_LEGACY + +#ifndef TA_STREAM_PER_THREAD +#define TA_STREAM_PER_THREAD ((tangStream_t)0x02) +#endif //! TA_STREAM_PER_THREAD + +#ifndef tangStreamLegacy +#define tangStreamLegacy ((tangStream_t)0x01) +#endif //! tangStreamLegacy + +#ifndef tangStreamPerThread +#define tangStreamPerThread ((tangStream_t)0x02) +#endif //! tangStreamPerThread + +#ifdef __cplusplus +extern "C" { +#endif //! __cplusplus + +#ifdef __cplusplus +} +#endif //! __cplusplus + +#endif //! _TANG_RUNTIME_H_ diff --git a/third_party/sunrise/backend/include/tang_runtime_api.h b/third_party/sunrise/backend/include/tang_runtime_api.h new file mode 100755 index 0000000000..106c1b2930 --- /dev/null +++ b/third_party/sunrise/backend/include/tang_runtime_api.h @@ -0,0 +1,1871 @@ +/* +Copyright declaration. +*/ +#ifndef _TANG_RUNTIME_API_H_ +#define _TANG_RUNTIME_API_H_ + +#include "tang_rt/driver_types.h" +#include "tang_rt/host_defines.h" +#include "tang_rt/vector_types.h" + +/** + * @brief Flags that can be used with tangStreamCreateWithFlags + * @{ + */ +#define tangStreamDefault 0x00 ///< Default stream creation flags +#define tangStreamNonBlocking \ + 0x01 ///< Stream does not implicitly synchronize with null stream + +//! Flags that can be used with tangEventCreateWithFlags: +#define tangEventDefault 0x0 ///< Default flags +#define tangEventBlockingSync \ + 0x1 ///< Waiting will yield CPU. Power-friendly and usage-friendly but may + ///< increase latency. +#define tangEventDisableTiming \ + 0x2 ///< Disable event's capability to record timing information. May + ///< improve performance. +#define tangEventInterprocess \ + 0x4 ///< Event can support IPC. @warning - not supported right now. + +//! Flags that can be used with tangStreamWaitEvent: +#define tangEventWaitDefault 0x00 ///< Default stream creation flags +#define tangEventWaitExternal \ + 0x01 ///< Event is captured in the graph as an external event node when + ///< performing stream capture. @warning - not supported right now. + +/** + * @brief enum values that can be used with tangStreamCreateWithPriority and + * tangDeviceGetStreamPriorityRange + * @{ + */ +enum stream_priority { + priority_high = -2, + priority_middle = -1, + priority_normal = 0, + priority_low = 1 +}; + +#ifdef __cplusplus +extern "C" { +#endif //! __cplusplus + +/** + *------------------------------------------------------------------------------------------------- + *------------------------------------------------------------------------------------------------- + * @defgroup Version Management + * @{ + */ + +/** + * @brief Returns the TANG SDK Runtime version. + * + * @param [out] runtimeVersion + * + * @returns #tangSuccess, #tangErrorInavlidValue + * + * @warning The TANG SDK runtime version does not correspond to an exact CUDA + * SDK runtime revision. + * + * @see tangDriverGetVersion + */ +tangError_t TANGRT_API_PUBLIC tangRuntimeGetVersion(int* runtimeVersion); + +/** + * @brief Returns the TANG SDK Driver version. + * + * @param [out] driverVersion + * + * @returns #tangSuccess, #tangErrorInavlidValue + * + * @warning The TANG SDK driver veriosn does not correspond to an exact CUDA SDK + * driver revision. + * + * @see tangRuntimeGetVersion + */ +tangError_t TANGRT_API_PUBLIC tangDriverGetVersion(int* driverVersion); + +// end doxygen Error +/** + * @} + */ + +/** + *------------------------------------------------------------------------------------------------- + *------------------------------------------------------------------------------------------------- + * @defgroup Error Handling + * @{ + */ + +/** + * @brief Return last error returned by any TANG runtime API call and resets the + * stored error code to #tangSuccess + * + * @returns return code from last TANG called from the active host thread + * + * Returns the last error that has been returned by any of the runtime calls in + * the same host thread, and then resets the saved error to #tangSuccess. + * + * @see tangGetErrorString, tangGetLastError, tangPeakAtLastError, tangError_t + */ +tangError_t TANGRT_API_PUBLIC tangGetLastError(void); + +/** + * @brief Return last error returned by any TANG runtime API call. + * + * @return #tangSuccess + * + * Returns the last error that has been returned by any of the runtime calls in + * the same host thread. Unlike tangGetLastError, this function does not reset + * the saved error code. + * + * @see tangGetErrorString, tangGetLastError, tangPeakAtLastError, tangError_t + */ +tangError_t TANGRT_API_PUBLIC tangPeekAtLastError(void); + +/** + * @brief Return name of the specified error code in text form. + * + * @param tang_error Error code to convert to name. + * @return const char pointer to the NULL-terminated error name + * + * @see tangGetErrorString, tangGetLastError, tangPeakAtLastError, tangError_t + */ +TANGRT_API_PUBLIC const char* tangGetErrorName(tangError_t tang_error); + +/** + * @brief Return handy text string message to explain the error which occurred + * + * @param tangError Error code to convert to string. + * @return const char pointer to the NULL-terminated error string + * + * @warning : on HCC, this function returns the name of the error (same as + * tangGetErrorName) + * + * @see tangGetErrorName, tangGetLastError, tangPeakAtLastError, tangError_t + */ +TANGRT_API_PUBLIC const char* tangGetErrorString(tangError_t tangError); + +// end doxygen Error +/** + * @} + */ + +/** + *------------------------------------------------------------------------------------------------- + *------------------------------------------------------------------------------------------------- + * @defgroup Device Management + * @{ + */ + +/** + * @brief Return number of compute-capable devices. + * + * @param [output] count Returns number of compute-capable devices. + * + * @returns #tangSuccess, #tangErrorNoDevice + * + * + * Returns in @p *count the number of devices that have ability to run compute + * commands. If there are no such devices, then @ref tangGetDeviceCount will + * return #tangErrorNoDevice. If 1 or more devices can be found, then + * tangGetDeviceCount returns #tangSuccess. + */ +tangError_t TANGRT_API_PUBLIC tangGetDeviceCount(int* count); + +/** + * @brief Set default device to be used for subsequent tang API calls from this + * thread. + * + * @param[in] deviceId Valid device in range 0...tangGetDeviceCount(). + * + * Sets @p device as the default device for the calling host thread. Valid + * device id's are 0... (tangGetDeviceCount()-1). + * + * Many TANG APIs implicitly use the "default device" : + * + * - Any device memory subsequently allocated from this host thread (using + * tangMalloc) will be allocated on device. + * - Any streams or events created from this host thread will be associated with + * device. + * - Any kernels launched from this host thread (using tangLaunchKernel) will be + * executed on device (unless a specific stream is specified, in which case the + * device associated with that stream will be used). + * + * This function may be called from any host thread. Multiple host threads may + * use the same device. This function does no synchronization with the previous + * or new device, and has very little runtime overhead. Applications can use + * tangSetDevice to quickly switch the default device before making a TANG + * runtime call which uses the default device. + * + * The default device is stored in thread-local-storage for each thread. + * Thread-pool implementations may inherit the default device of the previous + * thread. A good practice is to always call tangSetDevice at the start of TANG + * coding sequency to establish a known standard device. + * + * @returns #tangSuccess, #tangErrorInvalidDevice, #tangErrorDeviceAlreadyInUse + * + * @see tangGetDevice, tangGetDeviceCount + */ +tangError_t TANGRT_API_PUBLIC tangSetDevice(int deviceId); + +/** + * @brief Return the default device id for the calling host thread. + * + * @param [out] device *device is written with the default device + * + * TANG maintains an default device for each thread using thread-local-storage. + * This device is used implicitly for TANG runtime APIs called by this thread. + * tangGetDevice returns in * @p device the default device for the calling host + * thread. + * + * @returns #tangSuccess, #tangErrorInvalidDevice, #tangErrorInvalidValue + * + * @see tangSetDevice, tangGetDevicesizeBytes + */ +tangError_t TANGRT_API_PUBLIC tangGetDevice(int* deviceId); + +/** + * @brief Waits on all active streams on current device + * + * When this command is invoked, the host thread gets blocked until all the + * commands associated with streams associated with the device. TANG does not + * support multiple blocking modes (yet!). + * + * @returns #tangSuccess + * + * @see tangSetDevice, tangDeviceReset + */ +tangError_t TANGRT_API_PUBLIC tangDeviceSynchronize(void); + +/** + * @brief Returns device properties. + * + * @param [out] props written with device properties + * @param [in] deviceId which device to query for information + * + * @return #tangSuccess, #tangErrorInvalidDevice + * @bug HCC always returns 0 for maxThreadsPerMultiProcessor + * @bug HCC always returns 0 for regsPerBlock + * @bug HCC always returns 0 for l2CacheSize + * + * Populates tangGetDeviceProperties with information for the specified device. + */ +TANGRT_API_PUBLIC tangError_t tangGetDeviceProperties(tangDeviceProp* props, + int deviceId); + +/** + * @brief Query for a specific device attribute. + * + * @param [out] value pointer to value to return + * @param [in] attr attribute to query + * @param [in] deviceId which device to query for information + * + * @returns #tangSuccess, #tangErrorInvalidDevice, #tangErrorInvalidValue + */ +TANGRT_API_PUBLIC tangError_t tangDeviceGetAttribute(int* value, + tangDeviceAttr attr, + int deviceId); + +/** + * @brief tangDeviceGetPeerPointer. + * @via port to convert addr access a peer device's memory. + * + * @param[in ] s2 device index + * @param[in ] ptlink used port index. + * @param[in ] memory address alloc in peer device; + * @param[out ] the pointer conver peerAddr to accessAddr; + * @return #tangSuccess, #tangErrorInvalidValue + */ +TANGRT_API_PUBLIC tangError_t tangDeviceGetPeerPointer(int srcDevice, + int port, + void* peerAddr, + void** accessAddr); + +/** + * @brief tangDeviceGetP2PAttribute. + * @Queries attributes of the link between two devices. + * + * @param[out ] returned value of the requested attribute + * @param[in ] the supported attributes. + * @param[in ] source device of the target link. + * @param[in ] destination device of the target link. + * @return #tangSuccess, #tangErrorInvalidValue, #tangErrorInvalidDevice + * + */ +TANGRT_API_PUBLIC tangError_t tangDeviceGetP2PAttribute(int* value, + tangDeviceP2PAttr attr, + int srcDevice, + int dstDevice); + +/** + * @brief tangDeviceCanAccessPeer. + * @Queries if a device may directly access a peer device's memory. + * + * @param[out ] canAccessPeer return value, 1: success, 0: false. + * @param[in ] device local device id. + * @param[in ] peerDevice remote device id; + * @return #tangSuccess, #tangErrorInvalidDevice + * + */ +TANGRT_API_PUBLIC tangError_t tangDeviceCanAccessPeer(int* canAccessPeer, + int device, + int peerDevice); + +/** + * @brief tangDeviceEnablePeerAccess. + * @Enables direct access to memory allocations on a peer device. + * + * @param[in ] peerDevice remote device id. + * @param[in ] flags set 0; + * @return #tangSuccess, #tangErrorInvalidValue, #tangErrorInvalidDevice, + * #tangErrorPeerAccessAlreadyEnabled + * + */ +TANGRT_API_PUBLIC tangError_t tangDeviceEnablePeerAccess(int peerDevice, + unsigned int flags); + +/** + * @brief tangMemcpyPeer. + * @memory copy from a device to a peer device. + * + * @param[in ] dst, dst device memory point; + * @param[in ] dstDevice, dst device id; + * @param[in ] src, src device memory point; + * @param[in ] srcDevice, dst device id; + * @param[in ] size of memory copy in bytes; + * @return #tangSuccess + * + */ +TANGRT_API_PUBLIC tangError_t tangMemcpyPeer(void* dst, + int dstDevice, + const void* src, + int srcDevice, + size_t count); + +/** + * @brief tangMemcpyPeerAsync. + * @memory copy from a device to a peer device. + * + * @param[in ] dst, dst device memory point; + * @param[in ] dstDevice, dst device id; + * @param[in ] src, src device memory point; + * @param[in ] srcDevice, dst device id; + * @param[in ] size of memory copy in bytes; + * @param[in ] stream, used stream; + * @return #tangSuccess + * + */ +TANGRT_API_PUBLIC tangError_t tangMemcpyPeerAsync(void* dst, + int dstDevice, + const void* src, + int srcDevice, + size_t count, + tangStream_t stream); + +TANGRT_API_PUBLIC tangError_t tangMemcpyPeer_v2(void* dst, + int dstDevice, + const void* src, + int srcDevice, + size_t count); + +TANGRT_API_PUBLIC tangError_t tangMemcpyPeer_v2_ptds(void* dst, + int dstDevice, + const void* src, + int srcDevice, + size_t count); + +TANGRT_API_PUBLIC tangError_t tangMemcpyPeerAsync_v2(void* dst, + int dstDevice, + const void* src, + int srcDevice, + size_t count, + tangStream_t stream); + +TANGRT_API_PUBLIC tangError_t tangMemcpyPeerAsync_v2_ptsz(void* dst, + int dstDevice, + const void* src, + int srcDevice, + size_t count, + tangStream_t stream); + +/** + * @brief tangDeviceDisablePeerAccess. + * @Disables direct access to memory allocations on a peer device. + * + * @param[in ] peerDevice remote device id; + * @return #tangSuccess, #tangErrorInvalidDevice, #tangErrorPeerAccessNotEnabled + * + */ +TANGRT_API_PUBLIC tangError_t tangDeviceDisablePeerAccess(int peerDevice); + +/** + * @brief Get Resource limits of current device + * + * @param [out] pValue + * @param [in] limit + * + * @returns #tangSuccess, #tangErrorUnsupportedLimit, #tangErrorInvalidValue + * Note: Currently, only tangLimitMallocHeapSize is available + * + */ +TANGRT_API_PUBLIC tangError_t tangDeviceGetLimit(size_t* pValue, + enum tangLimit limit); + +TANGRT_API_PUBLIC tangError_t tangDeviceSetLimit(enum tangLimit limit, + size_t value); + +TANGRT_API_PUBLIC tangError_t tangDeviceReset(void); + +/** + * @brief Returns a handle to a compute device. + * @param [out] device handle + * @param [in] PCI Bus ID + * + * @returns #tangSuccess, #tangErrorInavlidDevice, #tangErrorInvalidValue + */ +TANGRT_API_PUBLIC tangError_t tangDeviceGetByPCIBusId(int* device, + const char* pciBusId); + +/** + * @brief Returns a PCI Bus Id string for the device. + * + * @param [out] pciBusId - PCI Bus ID + * @param [in] len - Maximum length of pciBusId name string + * @param [in] deviceId - device handle + * @returns #tangSuccess, #tangErrorInavlidDevice, #tangErrorInvalidValue + */ +TANGRT_API_PUBLIC tangError_t tangDeviceGetPCIBusId(char* pciBusId, + int len, + int deviceId); + +/** + * @brief Set L1/Shared cache partition. + * + * @param [in] config + * + * @returns #tangSuccess, #tangErrorNotInitialized + * Note: On PT2 devices, L1 cache and shared memory are separated. + * Thus these hints and controls are ignored on those architectures. + * + */ +TANGRT_API_PUBLIC tangError_t tangDeviceSetCacheConfig(tangFuncCache config); + +/** + * @brief Set Cache configuration for a specific function + * + * @param [in] config + * + * @returns #tangSuccess, #tangErrorNotInitialized + * Note: On PT2 devices, L1 cache and shared memory are separated. + * Thus these hints and controls are ignored on those architectures. + * + */ +TANGRT_API_PUBLIC tangError_t tangDeviceGetCacheConfig(tangFuncCache* config); + +/** + * @brief The bank width of shared memory on current device is set + * + * @param [in] config + * + * @returns #tangSuccess, #tangErrorInvalidValue, #tangErrorNotInitialized + * + * Note: On PT2 devices, shard memory bank size is fix to 4-bytes. + * Thus these hints and controls are ignored on those architectures. + * + */ +TANGRT_API_PUBLIC tangError_t +tangDeviceSetSharedMemConfig(tangSharedMemConfig config); + +/** + * @brief Returns bank width of shared memory for current device + * + * @param [out] config + * + * @returns #tangSuccess, #tangErrorInvalidValue, #tangErrorNotInitialized + * + * Note: On PT2 devices, shard memory bank size is fix to 4-bytes. + * Thus these hints and controls are ignored on those architectures. + * + */ +TANGRT_API_PUBLIC tangError_t +tangDeviceGetSharedMemConfig(tangSharedMemConfig* config); + +/** + * @brief Returns numerical values that correspond to the least and greatest + * stream priority. + * + * @param[out] leastPriority pointer in which value corresponding to least + * priority is returned. + * @param[out] greatestPriority pointer in which value corresponding to greatest + * priority is returned. + * + * Returns in *leastPriority and *greatestPriority the numerical values that + * correspond to the least and greatest stream priority respectively. Stream + * priorities follow a convention where lower numbers imply greater priorities. + * The range of meaningful stream priorities is given by * [*greatestPriority, + * *leastPriority]. If the user attempts to create a stream with a priority + * value that is outside the the meaningful range as specified by this API, the + * priority is automatically clamped to within the valid range. + */ +TANGRT_API_PUBLIC tangError_t +tangDeviceGetStreamPriorityRange(int* leastPriority, int* greatestPriority); + +/** + * @brief Set a list of devices that can be used for TANG. + * + * @param[in] List of devices to try. + * @param[in] Number of devices in specified list. + * + * Sets a list of devices for TANG execution in priority order using device_arr. + * The parameter len specifies the number of elements in the list. TANG will try + * devices from the list sequentially until it finds one that works. If this + * function is not called, or if it is called with a len of 0, then TANG will go + * back to its default behavior of trying devices sequentially from a default + * list containing all of the available TANG devices in the system. If a + * specified device ID in the list does not exist, this function will return + * tangErrorInvalidDevice. If len is not 0 and device_arr is NULL or if len + * exceeds the number of devices in the system, then tangErrorInvalidValue is + * returned. + * + * @return #tangSuccess, #tangErrorInvalidValue, #tangErrorInvalidDevice + * + */ +TANGRT_API_PUBLIC tangError_t tangSetValidDevices(int* device_arr, int len); + +/** + * @brief Select compute-device which best matches criteria. + * + * @param[out] device Device with best match + * @param[in] properties Desired device properties + * + * @return #tangSuccess, #tangErrorInvalidValue + * + */ +TANGRT_API_PUBLIC tangError_t +tangChooseDevice(int* device, const tangDeviceProp* properties); + +// end doxygen Device +/** + * @} + */ + +/** + *------------------------------------------------------------------------------------------------- + *------------------------------------------------------------------------------------------------- + * @defgroup Memory Management + * @{ + * + */ + +/** + * @brief Allocate memory on the default accelerator + * + * @param[out] pptr Pointer to the allocated memory + * @param[in] size Requested memory size + * + * If size is 0, no memory is allocated, *ptr returns nullptr, and tangSuccess + * is returned. + * + * @return #tangSuccess, #tangErrorOutOfMemory, #tangErrorInvalidValue (bad + * context, null *ptr) + * + * @see tangMallocPitch, tangFree, tangMallocArray, tangFreeArray, + * tangMalloc3D, tangMalloc3DArray, tangHostFree, tangHostMalloc + */ +tangError_t TANGRT_API_PUBLIC tangMalloc(void** pptr, size_t sizeBytes); + +/** + * @brief Allocate memory. + * + * @param pptr + * @param sizeBytes + * @param hStream + * @return tangError_t + */ +tangError_t TANGRT_API_PUBLIC tangMallocAsync(void** pptr, + size_t sizeBytes, + tangStream_t hStream); + +tangError_t TANGRT_API_PUBLIC tangMallocAsync_ptsz(void** pptr, + size_t sizeBytes, + tangStream_t hStream); + +/** + * @brief Free memory allocated by the hcc tang memory allocation API. + * This API performs an implicit tangDeviceSynchronize() call. + * If pointer is NULL, the tang runtime is initialized and tangSuccess is + * returned. + * + * @param[in] ptr Pointer to memory to be freed + * @return #tangSuccess + * @return #tangErrorInvalidDevicePointer (if pointer is invalid, including + * host pointers allocated with tangHostMalloc) + * + * @see tangMalloc, tangMallocPitch, tangMallocArray, tangFreeArray, + * tangHostFree, tangMalloc3D, tangMalloc3DArray, tangHostMalloc + */ +tangError_t TANGRT_API_PUBLIC tangFree(void* ptr); + +/** + * @brief Free memory block async. + * + * @param ptr + * @param hStream + * @return tangError_t + */ +tangError_t TANGRT_API_PUBLIC tangFreeAsync(void* ptr, tangStream_t hStream); +tangError_t TANGRT_API_PUBLIC tangFreeAsync_ptsz(void* ptr, + tangStream_t hStream); + +/** + * @brief Allocate page locked host memory + * + * @param[out] pptr Pointer to the allocated page locked host memory + * @param[in] sizeBytes Requested memory size + * + * If size is 0, no memory is allocated, *ptr returns nullptr, and tangSuccess + * is returned. + * + * @return #tangSuccess, #tangErrorOutOfMemory + * + */ +tangError_t TANGRT_API_PUBLIC tangMallocHost(void** pptr, size_t sizeBytes); + +/** + * @brief Allocate page locked host memory + * + * @param[out] pptr Pointer to the allocated page locked host memory + * @param[in] sizeBytes Requested memory size + * @param[in] flags See below. + * + * flags: + * - #tangHostAllocDefault Memory is page locked. + * - #tangHostAllocPortable Memory is considered registered by all + * contexts. + * - #tangHostAllocMapped Map the allocation into the address space for + * the current device. + * - #tangHostAllocWriteCombined Allocates the memory as write-combined (WC). + * TANG does not support IOMMU on device side, so flags of tangHostAllocMapped + * and tangHostAllocWriteCombined will always return false. + * + * If size is 0, no memory is allocated, *ptr returns nullptr, and tangSuccess + * is returned. + * + * @return #tangSuccess, #tangErrorOutOfMemory + */ +tangError_t TANGRT_API_PUBLIC tangHostAlloc(void** pptr, + size_t sizeBytes, + unsigned int flags); + +/** + * @brief Passes back the device pointer corresponding to the mapped, pinned + * host buffer allocated by tangHostAlloc(). Note: on PT2 devices, device + * pointer of mapped host memory requires 4-byte aligned access (because of the + * PCIE access mode). The start address assigned by tangHostAlloc is 4-byte + * aligned by default, but further use of the "offset" over the device pointer + * of mapped host memory should be careful. Access that is not 4-byte aligned + * may result in incorrect calculations. + * + * @param[out] pDevice Returned device pointer for mapped memory + * @param[in] pHost Requested host pointer mapping + * @param[in] flags Flags for extensions (must be 0 for now) + * + * @return #tangSuccess, #tangErrorInvalidValue, #tangErrorMemoryAllocation + */ +tangError_t TANGRT_API_PUBLIC tangHostGetDevicePointer(void** pDevice, + void* pHost, + unsigned int flags); + +/** + * @brief Passes back flags used to allocate pinned host memory allocated by + * tangHostAlloc. + * + * @param[out] pFlags Returned flags word + * @param[in] pHost Host pointer + * + * pFlags: + * - #tangHostAllocDefault Memory is page locked. + * - #tangHostAllocPortable Memory is considered registered by all + * contexts. + * - #tangHostAllocMapped Map the allocation into the address space for + * the current device. + * - #tangHostAllocWriteCombined Allocates the memory as write-combined (WC). + * TANG does not support IOMMU on device side, so flags of tangHostAllocMapped + * and tangHostAllocWriteCombined are not supported. + * + * @return #tangSuccess, #tangErrorOutOfMemory + */ +TANGRT_API_PUBLIC tangError_t tangHostGetFlags(unsigned int* pFlags, + void* pHost); + +/** + * @brief Free the page locked host memory allocated by the tang host memory + allocation API. + * + * @param[in] ptr Pointer to memory to be freed + * + * @return #tangSuccess, + * #tangErrorInvalidValue (if pointer is invalid, including device + pointers allocated with tangMalloc) + */ +tangError_t TANGRT_API_PUBLIC tangFreeHost(void* ptr); + +/** + * @brief Register host memory as page locked memory. + * + * @param[out] ptr Pointer to host memory to be registered. + * @param[in] sizeBytes Size of the host memory + * @param[in] flags See below. + * + * flags: + * - #tangHostRegisterDefault Memory is page locked. + * - #tangHostRegisterPortable Memory is considered registered by all contexts. + * - #tangHostRegisterMapped Map the allocation into the address space for + * the current device. + * - #tangHostRegisterIoMemory The passed memory pointer is treated as pointing + * to some memory-mapped I/O space. + * - #tangHostRegisterReadOnly The passed memory pointer is treated as pointing + * to memory that is considered read-only by the device. TANG does not support + * IOMMU on device side, so flags of tangHostRegisterMapped and + * tangHostRegisterIoMemory and tangHostRegisterReadOnly will always return + * false. + * + * @return #tangSuccess, #tangErrorOutOfMemory + * + * @see tangHostUnregister, tangHostGetFlags, tangHostGetDevicePointer + */ +tangError_t TANGRT_API_PUBLIC tangHostRegister(void* ptr, + size_t sizeBytes, + unsigned int flags); + +/** + * @brief Un-register host pointer + * + * @param[in] ptr Host pointer previously registered + * @return Error code + * + * @see tangHostRegister + */ +tangError_t TANGRT_API_PUBLIC tangHostUnregister(void* ptr); + +/** + * @brief Copy data from src to dst. + * + * It supports memory from host to device, + * device to host, device to device and host to host + * The src and dst must not overlap. + * + * For tangMemcpy, the copy is always performed by the current device (set by +tangSetDevice). + * For multi-gpu or peer-to-peer configurations, it is recommended to set the +current device to the + * device where the src data is physically located. For optimal peer-to-peer +copies, the copy + * device must be able to access the src and dst pointers (by calling +tangDeviceEnablePeerAccess + * with copy agent as the current device and src/dest as the peerDevice +argument. if this is not + * done, the tangMemcpy will still work, but will perform the copy using a +staging buffer on the + * host. Calling tangMemcpy with dst and src pointers that do not match the +tangMemcpyKind results + * in undefined behavior. + * + * @param[out] dst Data being copy to + * @param[in] src Data being copy from + * @param[in] sizeBytes Data size in bytes + * @param[in] kind Memory copy type + * @return #tangSuccess, #tangErrorInvalidValue, +#tangErrorInvalidMemcpyDirection , #tangErrorDriverIoctlFailed + * + * @see tangArrayCreate, tangArrayDestroy, tangArrayGetDescriptor, +tangMemAlloc, tangMemAllocHost, + * tangMemAllocPitch, tangMemcpy2D, tangMemcpy2DAsync, tangMemcpy2DUnaligned, +tangMemcpyAtoA, + * tangMemcpyAtoD, tangMemcpyAtoH, tangMemcpyAtoHAsync, tangMemcpyDtoA, +tangMemcpyDtoD, + * tangMemcpyDtoDAsync, tangMemcpyDtoH, tangMemcpyDtoHAsync, tangMemcpyHtoA, +tangMemcpyHtoAAsync, + * tangMemcpyHtoDAsync, tangMemFree, tangMemFreeHost, tangMemGetAddressRange, +tangMemGetInfo, + * tangMemHostAlloc, tangMemHostGetDevicePointer + */ +tangError_t TANGRT_API_PUBLIC tangMemcpy(void* dst, + const void* src, + size_t sizeBytes, + tangMemcpyKind kind); + +tangError_t TANGRT_API_PUBLIC tangMemcpy_ptds(void* dst, + const void* src, + size_t sizeBytes, + tangMemcpyKind kind); + +/** + * @brief Copy data from src to dst asynchronously. + * + * It supports memory from host to device, + * device to host, device to device and host to host + * The src and dst must not overlap. + * + * For tangMemcpyAsync, the copy is always performed by the current device (set + * by tangSetDevice). For multi-gpu or peer-to-peer configurations, it is + * recommended to set the current device to the device where the src data is + * physically located. For optimal peer-to-peer copies, the copy device must be + * able to access the src and dst pointers (by calling + * tangDeviceEnablePeerAccess with copy agent as the current device and src/dest + * as the peerDevice argument. if this is not done, the tangMemcpyAsync will + * still work, but will perform the copy using a staging buffer on the host. + * Calling tangMemcpy with dst and src pointers that do not match the + * tangMemcpyKind results in undefined behavior. + * + * @param[out] dst Data being copy to + * @param[in] src Data being copy from + * @param[in] sizeBytes Data size in bytes + * @param[in] kind Memory copy type + * @param[in] stream Stream to execute copy on + * @return #tangSuccess, #tangErrorInvalidValue, + * #tangErrorInvalidMemcpyDirection, #tangErrorDriverIoctlFailed + * + * @see tangArrayCreate, tangArrayDestroy, tangArrayGetDescriptor, + * tangMemAlloc, tangMemAllocHost, tangMemAllocPitch, tangMemcpy2D, + * tangMemcpy2DAsync, tangMemcpy2DUnaligned, tangMemcpyAtoA, tangMemcpyAtoD, + * tangMemcpyAtoH, tangMemcpyAtoHAsync, tangMemcpyDtoA, tangMemcpyDtoD, + * tangMemcpyDtoDAsync, tangMemcpyDtoH, tangMemcpyDtoHAsync, tangMemcpyHtoA, + * tangMemcpyHtoAAsync, tangMemcpyHtoDAsync, tangMemFree, tangMemFreeHost, + * tangMemGetAddressRange, tangMemGetInfo, tangMemHostAlloc, + * tangMemHostGetDevicePointer + */ +tangError_t TANGRT_API_PUBLIC +tangMemcpyAsync(void* dst, + const void* src, + size_t sizeBytes, + tangMemcpyKind kind, + tangStream_t strem __dparm(nullptr)); + +tangError_t TANGRT_API_PUBLIC +tangMemcpyAsync_ptsz(void* dst, + const void* src, + size_t sizeBytes, + tangMemcpyKind kind, + tangStream_t strem __dparm(nullptr)); + +/** + * @brief Fills the first sizeBytes bytes of the memory area pointed to by dest + * with the constant byte value value. + * + * @param[out] dst Data being filled + * @param[in] constant value to be set + * @param[in] sizeBytes Data size in bytes + * @return #tangSuccess, #tangErrorInvalidValue, #tangErrorNotInitialized + */ +tangError_t TANGRT_API_PUBLIC tangMemset(void* dst, + int value, + size_t sizeBytes); + +tangError_t TANGRT_API_PUBLIC tangMemset_ptds(void* dst, + int value, + size_t sizeBytes); + +/** + * @brief Fills the first sizeBytes bytes of the memory area pointed to by dev + * with the constant byte value value. + * + * tangMemsetAsync() is asynchronous with respect to the host, so the call may + * return before the memset is complete. The operation can optionally be + * associated to a stream by passing a non-zero stream argument. If stream is + * non-zero, the operation may overlap with operations in other streams. + * + * @param[out] dst Pointer to device memory + * @param[in] value - Value to set for each byte of specified memory + * @param[in] sizeBytes - Size in bytes to set + * @param[in] stream - Stream identifier + * @return #tangSuccess, #tangErrorInvalidValue, #tangErrorMemoryFree + */ +tangError_t TANGRT_API_PUBLIC +tangMemsetAsync(void* dst, + int value, + size_t sizeBytes, + tangStream_t stream __dparm(nullptr)); + +tangError_t TANGRT_API_PUBLIC +tangMemsetAsync_ptsz(void* dst, + int value, + size_t sizeBytes, + tangStream_t strem __dparm(nullptr)); + +tangError_t TANGRT_API_PUBLIC +tangMemcpyFromSymbol(void* dst, + const void* symbol, + size_t count, + size_t offset __dparm(0), + tangMemcpyKind kind __dparm(tangMemcpyDeviceToHost)); + +tangError_t TANGRT_API_PUBLIC +tangMemcpyFromSymbol_ptds(void* dst, + const void* symbol, + size_t count, + size_t offset __dparm(0), + tangMemcpyKind kind __dparm(tangMemcpyDeviceToHost)); + +tangError_t TANGRT_API_PUBLIC +tangMemcpyToSymbol(const void* symbol, + const void* src, + size_t count, + size_t offset __dparm(0), + tangMemcpyKind kind __dparm(tangMemcpyHostToDevice)); + +tangError_t TANGRT_API_PUBLIC +tangMemcpyToSymbol_ptds(const void* symbol, + const void* src, + size_t count, + size_t offset __dparm(0), + tangMemcpyKind kind __dparm(tangMemcpyHostToDevice)); + +tangError_t TANGRT_API_PUBLIC +tangMemcpyToSymbolAsync(const void* symbol, + const void* src, + size_t count, + size_t offset, + tangMemcpyKind kind, + tangStream_t stream __dparm(nullptr)); + +tangError_t TANGRT_API_PUBLIC +tangMemcpyToSymbolAsync_ptsz(const void* symbol, + const void* src, + size_t count, + size_t offset, + tangMemcpyKind kind, + tangStream_t stream __dparm(nullptr)); + +tangError_t TANGRT_API_PUBLIC +tangMemcpyFromSymbolAsync(void* dst, + const void* symbol, + size_t count, + size_t offset, + tangMemcpyKind kind, + tangStream_t stream __dparm(nullptr)); + +tangError_t TANGRT_API_PUBLIC +tangMemcpyFromSymbolAsync_ptsz(void* dst, + const void* symbol, + size_t count, + size_t offset, + tangMemcpyKind kind, + tangStream_t stream __dparm(nullptr)); + +/** + * @brief Query memory info. + * Return snapshot of free memory, and total allocatable memory on the device. + * + * Returns in *free a snapshot of the current free memory. + * @returns #tangSuccess, #tangErrorInvalidDevice, #tangErrorInvalidValue + * @warning On HCC, the free memory only accounts for memory allocated by this + *process and may be optimistic. + **/ +tangError_t TANGRT_API_PUBLIC tangMemGetInfo(size_t* free, size_t* total); + +/** + * @brief Finds the address associated with a TANG symbol. + * + * @param[out] devPtr Device pointer associated with symbol + * @param[in] symbol Device symbol address + * @return #tangSuccess, #tangErrorInvalidValue + **/ +tangError_t TANGRT_API_PUBLIC tangGetSymbolAddress(void** devPtr, + const void* symbol); + +/** + * @brief Finds the size of the object associated with a TANG symbol. + * + * @param[out] size Size of object associated with symbol + * @param[in] symbol Device symbol address + * @return #tangSuccess, #tangErrorInvalidValue + **/ +tangError_t TANGRT_API_PUBLIC tangGetSymbolSize(size_t* size, + const void* symbol); + +// doxygen end Memory +/** + * @} + */ + +/** + *------------------------------------------------------------------------------------------------- + *------------------------------------------------------------------------------------------------- + * @defgroup Stream Management + * @{ + */ + +/** + * @brief Create an asynchronous stream. + * + * @param[in, out] stream Valid pointer to tangStream_t. This function writes + * the memory with the newly created stream. + * @return #tangSuccess, #tangErrorInvalidValue + * + * Create a new asynchronous stream. @p stream returns an opaque handle that + * can be used to reference the newly created stream in subsequent tangStream* + * commands. The stream is allocated on the heap and will remain allocated even + * if the handle goes out-of-scope. To release the memory used by the stream, + * applicaiton must call tangStreamDestroy. + * + * @return #tangSuccess, #tangErrorInvalidValue + * + * @see tangStreamCreateWithFlags, tangStreamCreateWithPriority, + * tangStreamSynchronize, tangStreamWaitEvent, tangStreamDestroy + */ +tangError_t TANGRT_API_PUBLIC tangStreamCreate(tangStream_t* stream); + +/** + * @brief communicate to c2c. + * + * @param[in, out] stream Pointer to new stream + * @param[in ] cmd is that command packets + * @param[in ] cmd_count is that command count + * @param[in ] device_addr is that hbm addr + * @return #tangSuccess, #tangErrorInvalidValue + * + * send command to c2c module, it can be used by ptlink. + * + */ +tangError_t TANGRT_API_PUBLIC tangStreamC2Ctransfers(tangStream_t stream, + uint32_t* cmd, + uint32_t cmd_count, + uint64_t device_addr, + uint32_t mem_size); + +/** + * @brief Create an asynchronous stream. + * + * @param[in, out] stream Pointer to new stream + * @param[in ] flags to control stream creation. + * @return #tangSuccess, #tangErrorInvalidValue + * + * Create a new asynchronous stream. @p stream returns an opaque handle that + * can be used to reference the newly created stream in subsequent tangStream* + * commands. The stream is allocated on the heap and will remain allocated even + * if the handle goes out-of-scope. To release the memory used by the stream, + * applicaiton must call tangStreamDestroy. Flags controls behavior of the + * stream. See #tangStreamDefault, #tangStreamNonBlocking. + * + * @see tangStreamCreate, tangStreamCreateWithPriority, tangStreamSynchronize, + * tangStreamWaitEvent, tangStreamDestroy + */ +tangError_t TANGRT_API_PUBLIC tangStreamCreateWithFlags(tangStream_t* stream, + unsigned int flags __dparm(tangStreamDefault)); + +/** + * @brief Create an asynchronous stream with the specified priority. + * + * @param[in, out] stream Pointer to new stream + * @param[in ] flags to control stream creation. + * @param[in ] priority of the stream. Lower numbers represent higher + * priorities. + * @return #tangSuccess, #tangErrorInvalidValue + * + * Create a new asynchronous stream with the specified priority. @p stream + * returns an opaque handle that can be used to reference the newly created + * stream in subsequent tangStream* commands. The stream is allocated on the + * heap and will remain allocated even if the handle goes out-of-scope. To + * release the memory used by the stream, applicaiton must call + * tangStreamDestroy. Flags controls behavior of the stream. See + * #tangStreamDefault, #tangStreamNonBlocking. + * + * @see tangStreamCreate, tangStreamSynchronize, tangStreamWaitEvent, + * tangStreamDestroy + */ +tangError_t TANGRT_API_PUBLIC tangStreamCreateWithPriority(tangStream_t* stream, + unsigned int flags __dparm(tangStreamDefault), + int priority __dparm(priority_normal)); + +/** + * @brief Returns numerical values that correspond to the least and greatest + * stream priority. + * + * @param[in, out] leastPriority pointer in which value corresponding to least + * priority is returned. + * @param[in, out] greatestPriority pointer in which value corresponding to + * greatest priority is returned. + * + * Returns in *leastPriority and *greatestPriority the numerical values that + * correspond to the least and greatest stream priority respectively. Stream + * priorities follow a convention where lower numbers imply greater priorities. + * The range of meaningful stream priorities is given by + * [*greatestPriority, *leastPriority]. If the user attempts to create a stream + * with a priority value that is outside the the meaningful range as specified + * by this API, the priority is automatically clamped to within the valid range. + */ +tangError_t tangDeviceGetStreamPriorityRange(int* leastPriority, + int* greatestPriority); + +/** + * @brief Query the priority of a stream. + * + * @param[in] hStream stream to be queried + * @param[in,out] priority Pointer to an unsigned integer in which the stream's + * priority is returned + * @return #tangSuccess, #tangErrorInvalidValue + * + * Query the priority of a stream. The priority is returned in in priority. + * + * @see tangStreamCreateWithFlags + */ +tangError_t TANGRT_API_PUBLIC tangStreamGetPriority(tangStream_t hStream, + int* priority); + +tangError_t TANGRT_API_PUBLIC tangStreamGetPriority_ptsz(tangStream_t stream, + int* priority); + +/** + * @brief Destroys the specified stream. + * + * @param[in, out] stream Valid pointer to tangStream_t. This function writes + * the memory with the newly created stream. + * @return #tangSuccess #tangErrorInvalidHandle + * + * Destroys the specified stream. + * + * If commands are still executing on the specified stream, some may complete + * execution before the queue is deleted. + * + * The queue may be destroyed while some commands are still inflight, or may + * wait for all commands queued to the stream before destroying it. + * + * @see tangStreamCreate, tangStreamCreateWithFlags, + * tangStreamCreateWithPriority, tangStreamQuery, tangStreamWaitEvent, + * tangStreamSynchronize + */ +tangError_t TANGRT_API_PUBLIC tangStreamDestroy(tangStream_t stream); + +/** + * @brief Wait for all commands in stream to complete. + * + * @param[in] stream stream identifier. + * + * @return #tangSuccess, #tangErrorInvalidHandle + * + * This command is host-synchronous : the host will block until the specified + * stream is empty. + * + * This command follows standard null-stream semantics. Specifically, + * specifying the null stream will cause the command to wait for other streams + * on the same device to complete all pending operations. + * + * This command honors the tangDeviceLaunchBlocking flag, which controls whether + * the wait is active or blocking. + * + * @see tangStreamCreate, tangStreamCreateWithFlags, + * tangStreamCreateWithPriority, tangStreamWaitEvent, tangStreamDestroy + * + */ +tangError_t TANGRT_API_PUBLIC tangStreamSynchronize(tangStream_t stream); + +tangError_t TANGRT_API_PUBLIC tangStreamSynchronize_ptsz(tangStream_t stream); + +/** + * @brief Check if a stream has completed all its commands + * + * @param stream + * @return tangError_t + * tangSuccess + * tangErrorNotReady + */ +tangError_t TANGRT_API_PUBLIC tangStreamQuery(tangStream_t stream); +tangError_t TANGRT_API_PUBLIC tangStreamQuery_ptsz(tangStream_t stream); + +tangError_t TANGRT_API_PUBLIC +tangStreamBeginCapture(tangStream_t stream, tangStreamCaptureMode mode); + +tangError_t TANGRT_API_PUBLIC +tangStreamBeginCapture_ptsz(tangStream_t stream, + tangStreamCaptureMode mode); + +tangError_t TANGRT_API_PUBLIC tangStreamEndCapture(tangStream_t stream, + tangGraph_t* pGraph); + +tangError_t TANGRT_API_PUBLIC tangStreamEndCapture_ptsz(tangStream_t stream, + tangGraph_t* pGraph); + +tangError_t TANGRT_API_PUBLIC +tangStreamIsCapturing(tangStream_t stream, + tangStreamCaptureStatus* pStatus); + +tangError_t TANGRT_API_PUBLIC +tangStreamIsCapturing_ptsz(tangStream_t stream, + tangStreamCaptureStatus* pStatus); + +tangError_t TANGRT_API_PUBLIC +tangStreamGetCaptureInfo(tangStream_t hStream, + tangStreamCaptureStatus* pStatus, + unsigned long long* pId __dparm(0), + tangGraph_t* pGraph __dparm(0), + const tangGraphNode_t** deps __dparm(0), + size_t* numDeps __dparm(0)); + +tangError_t TANGRT_API_PUBLIC +tangStreamGetCaptureInfo_ptsz(tangStream_t hStream, + tangStreamCaptureStatus* pStatus, + unsigned long long* pId __dparm(0), + tangGraph_t* pGraph __dparm(0), + const tangGraphNode_t** deps __dparm(0), + size_t* numDeps __dparm(0)); + +tangError_t TANGRT_API_PUBLIC +tangThreadExchangeStreamCaptureMode(tangStreamCaptureMode* mode); + +tangError_t TANGRT_API_PUBLIC tangGraphInstantiate(tangGraphExec_t* pGraphExec, + tangGraph_t graph, + void*, + void*, + unsigned long long); + +tangError_t TANGRT_API_PUBLIC tangGraphLaunch(tangGraphExec_t graphExec, + tangStream_t stream); + +tangError_t TANGRT_API_PUBLIC tangGraphLaunch_ptsz(tangGraphExec_t graphExec, + tangStream_t stream); +tangError_t TANGRT_API_PUBLIC +tangGraphInstantiateWithFlags(tangGraphExec_t* pGraphExec, + tangGraph_t graph, + unsigned long long flags); + +tangError_t TANGRT_API_PUBLIC tangGraphDestroy(tangGraph_t graph); +tangError_t TANGRT_API_PUBLIC tangGraphExecDestroy(tangGraphExec_t graphExec); + +tangError_t TANGRT_API_PUBLIC tangGraphGetInfo(tangGraph_t graph, + tangGraphInfo* pInfo); + +tangError_t TANGRT_API_PUBLIC tangGraphCreate(tangGraph_t* pGraph, + unsigned int flags); + +tangError_t TANGRT_API_PUBLIC +tangGraphAddHostNode(tangGraphNode_t* pGraphNode, + tangGraph_t graph, + const tangGraphNode_t* dependencies, + size_t numDependencies, + const tangHostNodeParams* nodeParams); + +tangError_t TANGRT_API_PUBLIC +tangGraphAddKernelNode(tangGraphNode_t* pGraphNode, + tangGraph_t graph, + const tangGraphNode_t* dependencies, + size_t numDependencies, + const tangKernelNodeParams* nodeParams); + +/** + * @brief Make the specified compute stream wait for an event + * + * @param[in] stream stream to make wait. + * @param[in] event event to wait on + * @param[in] flag control operation + * + * @return #tangSuccess, #tangErrorInvalidHandle + * + * This function inserts a wait operation into the specified stream. + * All future work submitted to @p stream will wait until @p event reports + * completion before beginning execution. + * + * This function only waits for commands in the current stream to complete. + * Notably,, this function does not impliciy wait for commands in the default + * stream to complete, even if the specified stream is created with + * tangStreamNonBlocking = 0. + * + * @see tangStreamCreate, tangStreamCreateWithFlags, + * tangStreamCreateWithPriority, tangStreamSynchronize, tangStreamDestroy + */ +tangError_t TANGRT_API_PUBLIC tangStreamWaitEvent(tangStream_t stream, + tangEvent_t event, + unsigned int flag __dparm(0)); + +tangError_t TANGRT_API_PUBLIC +tangStreamWaitEvent_ptsz(tangStream_t stream, + tangEvent_t event, + unsigned int flag __dparm(0)); +/** + * @brief Return flags associated with this stream. + * + * @param[in] stream stream to be queried + * @param[in,out] flag Pointer to an unsigned integer in which the stream's + * flags are returned + * @return #tangSuccess, #tangErrorInvalidValue, #tangErrorInvalidHandle + * + * @returns #tangSuccess #tangErrorInvalidValue #tangErrorInvalidHandle + * + * Return flags associated with this stream in *@p flag. + * + * @see tangStreamCreateWithFlags + */ +tangError_t TANGRT_API_PUBLIC tangStreamGetFlags(tangStream_t stream, + unsigned int* flags); + +tangError_t TANGRT_API_PUBLIC tangStreamGetFlags_ptsz(tangStream_t stream, + unsigned int* flags); + +tangError_t TANGRT_API_PUBLIC tangStreamGetId(tangStream_t stream, + int* pId); + +tangError_t TANGRT_API_PUBLIC tangStreamGetId_ptsz(tangStream_t stream, + int* pId); + +typedef void (*tangStreamCallback_t)(tangStream_t stream, + tangError_t status, + void* userData); + +tangError_t TANGRT_API_PUBLIC +tangStreamAddCallback(tangStream_t stream, + tangStreamCallback_t callback, + void* userData, + unsigned int flags); + +tangError_t TANGRT_API_PUBLIC +tangStreamAddCallback_ptsz(tangStream_t stream, + tangStreamCallback_t callback, + void* userData, + unsigned int flags); + +tangError_t TANGRT_API_PUBLIC tangLaunchHostFunc(tangStream_t stream, + tangHostFn_t fn, + void* userData); + +tangError_t TANGRT_API_PUBLIC tangLaunchHostFunc_ptsz(tangStream_t stream, + tangHostFn_t fn, + void* userData); + +tangError_t TANGRT_API_PUBLIC tangProfilerStart(); +tangError_t TANGRT_API_PUBLIC tangProfilerStop(); +// end doxygen Stream +/** + * @} + */ + +tangError_t TANGRT_API_PUBLIC tangEngineCollAssign(int devId, + int collType, + uint64_t devAddr, + int length, + tangStream_t stream); + + +/** + *------------------------------------------------------------------------------------------------- + *------------------------------------------------------------------------------------------------- + * @defgroup Event Management + * @{ + */ + +/** + * @brief Create an event with the specified flags + * + * @param[in,out] event Returns the newly created event. + * @param[in] flag Flag to control event behavior. Valid values are + #tangEventDefault, #tangEventBlockingSync, #tangEventDisableTiming, + #tangEventInterprocess + + * #tangEventDefault : Default flag. The event will use active synchronization + and will support timing. Blocking synchronization provides lowest possible + latency at the expense of dedicating a CPU to poll on the event. + * #tangEventBlockingSync : The event will use blocking synchronization : if + tangEventSynchronize is called on this event, the thread will block until the + event completes. This can increase latency for the synchroniation but can + result in lower power and more resources for other CPU threads. + * #tangEventDisableTiming : Disable recording of timing information. On ROCM + platform, timing information is always recorded and this flag has no + performance benefit. + + * @warning tangEventInterprocess support is under development. Use of this + flag will return an error. + * + * @returns #tangSuccess, #tangErrorNotInitialized, #tangErrorInvalidValue, + #tangErrorLaunchFailure, #tangErrorOutOfMemory + * + * @see tangEventCreate, tangEventSynchronize, tangEventDestroy, + tangEventElapsedTime + */ +tangError_t TANGRT_API_PUBLIC tangEventCreateWithFlags(tangEvent_t* event, + unsigned flag); + +/** + * Create an event + * + * @param[in,out] event Returns the newly created event. + * + * @returns #tangSuccess, #tangErrorNotInitialized, #tangErrorInvalidValue, + * #tangErrorLaunchFailure, #tangErrorOutOfMemory + * + * @see tangEventCreateWithFlags, tangEventRecord, tangEventQuery, + * tangEventSynchronize, tangEventDestroy, tangEventElapsedTime + */ +tangError_t TANGRT_API_PUBLIC tangEventCreate(tangEvent_t* event); + +/** + * @brief Record an event in the specified stream. + * + * @param[in] event event to record. + * @param[in] stream stream in which to record event. + * @returns #tangSuccess, #tangErrorInvalidValue, #tangErrorNotInitialized, + * #tangErrorInvalidHandle, #tangErrorLaunchFailure + * + * tangEventQuery() or tangEventSynchronize() must be used to determine when the + * event transitions from "recording" (after tangEventRecord() is called) to + * "recorded" (when timestamps are set, if requested). + * + * Events which are recorded in a non-NULL stream will transition to + * from recording to "recorded" state when they reach the head of + * the specified stream, after all previous + * commands in that stream have completed executing. + * + * If tangEventRecord() has been previously called on this event, then this call + * will overwrite any existing state in event. + * + * If this function is called on an event that is currently being recorded, + * results are undefined + * - either outstanding recording may save state into the event, and the order + * is not guaranteed. + * + * @see tangEventCreate, tangEventCreateWithFlags, tangEventQuery, + * tangEventSynchronize, tangEventDestroy, tangEventElapsedTime + * + */ +tangError_t TANGRT_API_PUBLIC +tangEventRecord(tangEvent_t event, tangStream_t stream __dparm(nullptr)); + +tangError_t TANGRT_API_PUBLIC +tangEventRecord_ptsz(tangEvent_t event, tangStream_t stream __dparm(nullptr)); + +tangError_t TANGRT_API_PUBLIC +tangEventRecordWithFlags(tangEvent_t event, + tangStream_t stream __dparm(nullptr), + unsigned int flags __dparm(0)); + +tangError_t TANGRT_API_PUBLIC +tangEventRecordWithFlags_ptsz(tangEvent_t event, + tangStream_t stream __dparm(nullptr), + unsigned int flags __dparm(0)); +/** + * @brief Destroy the specified event. + * + * @param[in] event Event to destroy. + * @returns #tangSuccess, #tangErrorNotInitialized, #tangErrorInvalidValue, + * #tangErrorLaunchFailure + * + * Releases memory associated with the event. If the event is recording but + * has not completed recording when tangEventDestroy() is called, the function + * will return immediately and the completion_future resources will be released + * later, when the tangDevice is synchronized. + * + * @see tangEventCreate, tangEventCreateWithFlags, tangEventQuery, + * tangEventSynchronize, tangEventRecord, tangEventElapsedTime + * + * @returns #tangSuccess + */ +tangError_t TANGRT_API_PUBLIC tangEventDestroy(tangEvent_t event); + +/** + * @brief Wait for an event to complete. + * + * This function will block until the event is ready, waiting for all previous + * work in the stream specified when event was recorded with tangEventRecord(). + * + * If tangEventRecord() has not been called on @p event, this function returns + * immediately. + * + * TODO-hcc - This function needs to support tangEventBlockingSync parameter. + * + * @param[in] event Event on which to wait. + * @returns #tangSuccess, #tangErrorInvalidValue, #tangErrorNotInitialized, + * #tangErrorInvalidHandle, #tangErrorLaunchFailure + * + * @see tangEventCreate, tangEventCreateWithFlags, tangEventQuery, + * tangEventDestroy, tangEventRecord, tangEventElapsedTime + */ +tangError_t TANGRT_API_PUBLIC tangEventSynchronize(tangEvent_t event); +tangError_t TANGRT_API_PUBLIC tangEventSynchronizeWithFlags(tangEvent_t event, + unsigned int flags); + +/** + * @brief Return the elapsed time between two events. + * + * @param[out] ms : Return time between start and stop in ms. + * @param[in] start : Start event. + * @param[in] stop : Stop event. + * @returns #tangSuccess, #tangErrorInvalidValue, #tangErrorNotReady, + * #tangErrorInvalidHandle, #tangErrorNotInitialized, #tangErrorLaunchFailure + * + * Computes the elapsed time between two events. Time is computed in ms, with + * a resolution of approximately 1 us. + * + * Events which are recorded in a NULL stream will block until all commands + * on all other streams complete execution, and then record the timestamp. + * + * Events which are recorded in a non-NULL stream will record their timestamp + * when they reach the head of the specified stream, after all previous + * commands in that stream have completed executing. Thus the time that + * the event recorded may be significantly after the host calls + * tangEventRecord(). + * + * If tangEventRecord() has not been called on either event, then + * #tangErrorInvalidHandle is returned. If tangEventRecord() has been called on + * both events, but the timestamp has not yet been recorded on one or both + * events (that is, tangEventQuery() would return #tangErrorNotReady on at least + * one of the events), then #tangErrorNotReady is returned. + * + * @see tangEventCreate, tangEventCreateWithFlags, tangEventQuery, + * tangEventDestroy, tangEventRecord, tangEventSynchronize + */ +tangError_t TANGRT_API_PUBLIC tangEventElapsedTime(float* ms, + tangEvent_t start, + tangEvent_t stop); + +/** + * @brief Query event status + * + * @param[in] event Event to query. + * @returns #tangSuccess, #tangErrorNotReady, #tangErrorInvalidHandle, + * #tangErrorInvalidValue, #tangErrorNotInitialized, #tangErrorLaunchFailure + * + * Query the status of the specified event. This function will return + * #tangErrorNotReady if all commands in the appropriate stream (specified to + * tangEventRecord()) have completed. If that work has not completed, or if + * tangEventRecord() was not called on the event, then #tangSuccess is returned. + * + * @see tangEventCreate, tangEventCreateWithFlags, tangEventRecord, + * tangEventDestroy, tangEventSynchronize, tangEventElapsedTime + */ +tangError_t TANGRT_API_PUBLIC tangEventQuery(tangEvent_t event); + +tangError_t TANGRT_API_PUBLIC tangEventQueryTimestamp(tangEvent_t event, + tangEventTimestamp* ts); +// end doxygen Event +/** + * @} + */ + +/** + *------------------------------------------------------------------------------------------------- + *------------------------------------------------------------------------------------------------- + * @defgroup Excution Control + * @{ + */ + +/** + * @brief Find out attributes for a given function. + * + * @param [out] attr + * @param [in] func + * + * @returns tangSuccess, tangErrorInvalidValue, tangErrorInvalidDeviceFunction + * + * NOTE: runtime only has tangFuncGetAttributes API, has no tangFuncGetAttribute + * API. while user mode driver only has taFuncGetAttribute API, has no + * taFuncGetAttributes API. + */ + +tangError_t TANGRT_API_PUBLIC tangFuncGetAttributes(tangFuncAttributes* attr, + const void* func); + +tangError_t TANGRT_API_PUBLIC tangGetFuncBySymbol(tangFunction_t *hFunc, + const void *symbol); + +tangError_t TANGRT_API_PUBLIC +tangPointerGetAttributes(struct tangPointerAttributes* attributes, + const void* ptr); + +/** + * @brief Set attribute for a specific function + * + * @param [in] func; + * @param [in] attr; + * @param [in] value; + * + * @returns #tangSuccess, #tangErrorInvalidDeviceFunction, + * #tangErrorInvalidValue + * + * Note: PT devices do not support shared cache banking, and the hint is + * ignored. + * + * NOTE: runtime tangFuncSetAttribute API only supports two types of + * tangFuncAttribute. while user mode driver taFuncSetAttribute API supports + * more types of TAfunction_attribute. + * + */ +tangError_t TANGRT_API_PUBLIC tangFuncSetAttribute(const void* func, + tangFuncAttribute attr, + int value); + +/** + * @brief Set Cache configuration for a specific function + * + * @param [in] func; + * @param [in] config; + * + * @returns #tangSuccess, #tangErrorNotInitialized + * Note: PT devices do not support reconfigurable cache. This hint is ignored. + * + */ +tangError_t TANGRT_API_PUBLIC tangFuncSetCacheConfig(const void* func, + tangFuncCache config); + +/** + * @brief Set shared memory configuation for a specific function + * + * @param [in] func + * @param [in] config + * + * @returns #tangSuccess, #tangErrorInvalidValue, + * #tangErrorInvalidDeviceFunction + * + * Note: PT devices do not support shared cache banking, and the hint is + * ignored. + * + */ +tangError_t TANGRT_API_PUBLIC +tangFuncSetSharedMemConfig(const void* func, tangSharedMemConfig config); + +/** + * @brief Converts a double argument to be executed on a device. + * + * @param[in][out] d Double to convert. + * @returns #tangSuccess, #tangErrorInvalidValue + * + * Note: PT2 devices do not support double both on hardware and software + * simulation. + * + */ +TANGRT_API_PUBLIC tangError_t tangSetDoubleForDevice(double* d); + +/** + * @brief Converts a double argument after execution on a device. + * + * @param[in][out] d Double to convert. + * @returns #tangSuccess, #tangErrorInvalidValue + * + * Note: PT2 devices do not support double both on hardware and software + * simulation. + * + */ +TANGRT_API_PUBLIC tangError_t tangSetDoubleForHost(double* d); + +// doxygen end Execution Control +/** + * @} + */ + +/** + *------------------------------------------------------------------------------------------------- + *------------------------------------------------------------------------------------------------- + * @defgroup Occupancy + * @{ + * + */ + +/** + * @brief Returns occupancy for a device function. + * + * @param[out] numBlocks Returned occupancy + * @param[in] func Kernel function for which occupancy is calculated + * @param[in] blockSize Block size the kernel is intended to be launched with + * @param[in] dynamicSMemSize Per-block dynamic shared memory usage intended, in + * bytes + * @returns #tangSuccess, #tangErrorInvalidValue + * + */ +TANGRT_API_PUBLIC tangError_t +tangOccupancyMaxActiveBlocksPerMultiprocessor(int* numBlocks, + const void* func, + int blockSize, + size_t dynamicSMemSize); + +/** + * @brief Returns occupancy for a device function with the specified flags. + * + * @param[out] numBlocks Returned occupancy + * @param[in] func Kernel function for which occupancy is calculated + * @param[in] blockSize Block size the kernel is intended to be launched with + * @param[in] dynamicSMemSize Per-block dynamic shared memory usage intended, + * in bytes + * @param[in] flags Requested behavior for the occupancy calculator + * @returns #tangSuccess, #tangErrorInvalidValue + * + */ +TANGRT_API_PUBLIC tangError_t +tangOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(int* numBlocks, + const void* func, + int blockSize, + size_t dynamicSMemSize, + unsigned int flags); + +TANGRT_API_PUBLIC tangError_t tangIpcGetMemHandle(tangIpcMemHandle_t* pHandle, + void* devPtr); + +TANGRT_API_PUBLIC tangError_t tangIpcOpenMemHandle(void** devPtr, + tangIpcMemHandle_t handle, + unsigned int flags); + +TANGRT_API_PUBLIC tangError_t tangIpcCloseMemHandle(void* devPtr); + +TANGRT_API_PUBLIC tangError_t +tangIpcGetEventHandle(tangIpcEventHandle_t* pHandle, tangEvent_t event); + +TANGRT_API_PUBLIC tangError_t +tangIpcOpenEventHandle(tangEvent_t* phEvent, tangIpcEventHandle_t handle); + +// private api +TANGRT_API_PUBLIC tangError_t tangGetExportTable(void** pExportedTable, + void* args); + +// end doxygen Occupancy +/** + * @} + */ + +#ifdef __cplusplus +} +#endif //! __cplusplus + +#ifdef __cplusplus +template +inline tangError_t tangMalloc(T** pptr, size_t sizeBytes) { + return ::tangMalloc((void**)pptr, sizeBytes); +} +template +inline tangError_t tangMallocHost(T** pptr, size_t sizeBytes) { + return ::tangMallocHost((void**)pptr, sizeBytes); +} +template +inline tangError_t tangMallocAsync(T** pptr, size_t sizeBytes, tangStream_t stream) { + return ::tangMallocAsync((void**)pptr, sizeBytes, stream); +} +template +inline tangError_t tangMallocAsync_ptsz(T** pptr, size_t sizeBytes, tangStream_t stream) { + return ::tangMallocAsync_ptsz((void**)pptr, sizeBytes, stream); +} +template +inline tangError_t tangHostAlloc(T** pptr, + size_t sizeBytes, + unsigned int flags) { + return ::tangHostAlloc((void**)pptr, sizeBytes, flags); +} +template +inline tangError_t tangHostGetDevicePointer(T** pDevice, + void* pHost, + unsigned int flags) { + return ::tangHostGetDevicePointer((void**)pDevice, pHost, flags); +} +template +inline tangError_t tangIpcOpenMemHandle(T** pDevPtr, + tangIpcMemHandle_t handle, + unsigned int flags) { + return ::tangIpcOpenMemHandle((void**)pDevPtr, handle, flags); +} +#endif // __cplusplus + +#if defined(__TANGRT_API_PER_THREAD_DEFAULT_STREAM) +#define tangMemset __TANGRT_API_PTDS(tangMemset) +#define tangMemsetAsync __TANGRT_API_PTSZ(tangMemsetAsync) +#define tangMemcpy __TANGRT_API_PTDS(tangMemcpy) +#define tangMemcpyAsync __TANGRT_API_PTSZ(tangMemcpyAsync) +#define tangMallocAsync __TANGRT_API_PTSZ(tangMallocAsync) +#define tangFreeAsync __TANGRT_API_PTSZ(tangFreeAsync) +#define tangStreamSynchronize __TANGRT_API_PTSZ(tangStreamSynchronize) +#define tangStreamQuery __TANGRT_API_PTSZ(tangStreamQuery) +#define tangStreamWaitEvent __TANGRT_API_PTSZ(tangStreamWaitEvent) +#define tangStreamC2Ctransfers __TANGRT_API_PTSZ(tangStreamC2Ctransfers) +#define tangStreamGetFlags __TANGRT_API_PTSZ(tangStreamGetFlags) +#define tangStreamGetId __TANGRT_API_PTSZ(tangStreamGetId) +#define tangStreamGetPriority __TANGRT_API_PTSZ(tangStreamGetPriority) +#define tangStreamAddCallback __TANGRT_API_PTSZ(tangStreamAddCallback) +#define tangStreamBeginCapture __TANGRT_API_PTSZ(tangStreamBeginCapture) +#define tangStreamEndCapture __TANGRT_API_PTSZ(tangStreamEndCapture) +#define tangStreamIsCapturing __TANGRT_API_PTSZ(tangStreamIsCapturing) +#define tangStreamGetCaptureInfo __TANGRT_API_PTSZ(tangStreamGetCaptureInfo) +#define tangGraphLaunch __TANGRT_API_PTSZ(tangGraphLaunch) +#define tangLaunchHostFunc __TANGRT_API_PTSZ(tangLaunchHostFunc) +#define tangEventRecord __TANGRT_API_PTSZ(tangEventRecord) +#define tangEventRecordWithFlags __TANGRT_API_PTSZ(tangEventRecordWithFlags) +#define tangMemcpyFromSymbol __TANGRT_API_PTDS(tangMemcpyFromSymbol) +#define tangMemcpyFromSymbolAsync __TANGRT_API_PTSZ(tangMemcpyFromSymbolAsync) +#define tangMemcpyToSymbol __TANGRT_API_PTDS(tangMemcpyToSymbol) +#define tangMemcpyToSymbolAsync __TANGRT_API_PTSZ(tangMemcpyToSymbolAsync) +//#define tangDeviceCanAccessPeer __TANGRT_API_PTSZ(tangDeviceCanAccessPeer) +//#define tangDeviceEnablePeerAccess +//__TANGRT_API_PTSZ(tangDeviceEnablePeerAccess) #define +//tangDeviceDisablePeerAccess __TANGRT_API_PTSZ(tangDeviceDisablePeerAccess) +#endif //! __TANGRT_API_PER_THREAD_DEFAULT_STREAM +// end doxygen Events +/** + * @} + */ + +#ifdef __cplusplus + +// template is only available in cxx. +// And these template APIs may cause ambiguous. +// If you don't want to use these api, please define TANGRT_DISABLE_SYMBOL_TEMPLATE_API +// before #include or #include +#if defined(__cplusplus) && !defined(TANGRT_DISABLE_SYMBOL_TEMPLATE_API) + +template +inline tangError_t tangGetSymbolAddress(void** devPtr, const T& symbol) { + return ::tangGetSymbolAddress((void**)devPtr, (const void*)&symbol); +} + +template +inline tangError_t tangGetSymbolSize(size_t* size, const T& symbol) { + return ::tangGetSymbolSize(size, (const void*)&symbol); +} + +template +inline tangError_t tangMemcpyToSymbol( + const T& symbol, + const void* src, + size_t count, + size_t offset = 0, + enum tangMemcpyKind kind = tangMemcpyHostToDevice) { + return ::tangMemcpyToSymbol((const void*)&symbol, src, count, offset, kind); +} + +template +inline tangError_t tangMemcpyToSymbolAsync( + const T& symbol, + const void* src, + size_t count, + size_t offset = 0, + enum tangMemcpyKind kind = tangMemcpyHostToDevice, + tangStream_t stream = 0) { + return ::tangMemcpyToSymbolAsync((const void*)&symbol, + src, + count, + offset, + kind, + stream); +} + +template +inline tangError_t tangMemcpyFromSymbol( + void* dst, + const T& symbol, + size_t count, + size_t offset = 0, + enum tangMemcpyKind kind = tangMemcpyDeviceToHost) { + return ::tangMemcpyFromSymbol(dst, (const void*)&symbol, count, offset, kind); +} + +template +inline tangError_t tangMemcpyFromSymbolAsync( + void* dst, + const T& symbol, + size_t count, + size_t offset = 0, + enum tangMemcpyKind kind = tangMemcpyDeviceToHost, + tangStream_t stream = 0) { + return ::tangMemcpyFromSymbolAsync(dst, + (const void*)&symbol, + count, + offset, + kind, + stream); +} +#endif //!< TANGRT_SYMBOL_FORCE_CXX_API + +#endif // __cplusplus + +#endif //! _TANG_RUNTIME_API_H_ diff --git a/third_party/sunrise/backend/include/tapti/tapti.h b/third_party/sunrise/backend/include/tapti/tapti.h new file mode 100755 index 0000000000..923cc2b57f --- /dev/null +++ b/third_party/sunrise/backend/include/tapti/tapti.h @@ -0,0 +1,11 @@ +#ifndef _TAPTI_HPP_ +#define _TAPTI_HPP_ + +#include "tang_runtime.h" +#include "tang.h" +#include "tapti_activity.h" +#include "tapti_version.h" +#include "tapti_callbacks.h" +#include "tapti_result.h" + +#endif // _TAPTI_HPP_ diff --git a/third_party/sunrise/backend/include/tapti/tapti_activity.h b/third_party/sunrise/backend/include/tapti/tapti_activity.h new file mode 100755 index 0000000000..f61765c7a3 --- /dev/null +++ b/third_party/sunrise/backend/include/tapti/tapti_activity.h @@ -0,0 +1,956 @@ +#ifndef _TAPTI_ACTIVITY_HPP_ +#define _TAPTI_ACTIVITY_HPP_ + +#include +#include + +#include "tapti_callbacks.h" + +/** + * \brief The kind of a memory copy, indicating the source and + * destination targets of the copy. + * + * Each kind represents the source and destination targets of a memory + * copy. Targets are host, device, and array. + */ +typedef enum { + /** + * The memory copy kind is not known. + */ + TAPTI_ACTIVITY_MEMCPY_KIND_UNKNOWN = 0, + + /** + * A host to device memory copy. + */ + TAPTI_ACTIVITY_MEMCPY_KIND_HTOD = 1, + + /** + * A device to host memory copy. + */ + TAPTI_ACTIVITY_MEMCPY_KIND_DTOH = 2, + + /** + * A device to device memory copy on the same device. + */ + TAPTI_ACTIVITY_MEMCPY_KIND_DTOD = 3, + + /** + * A host to host memory copy. + */ + TAPTI_ACTIVITY_MEMCPY_KIND_HTOH = 4, + + /** + * A peer to peer memory copy across different devices. + */ + TAPTI_ACTIVITY_MEMCPY_KIND_PTOP = 5, + + TAPTI_ACTIVITY_MEMCPY_KIND_FORCE_INT = 0x7fffffff +} TApti_ActivityMemcpyKind; + +typedef enum { + TAPTI_EXTERNAL_CORRELATION_KIND_INVALID = 0, + TAPTI_EXTERNAL_CORRELATION_KIND_UNKNOWN = 1, + TAPTI_EXTERNAL_CORRELATION_KIND_OPENACC = 2, + TAPTI_EXTERNAL_CORRELATION_KIND_CUSTOM0 = 3, + TAPTI_EXTERNAL_CORRELATION_KIND_CUSTOM1 = 4, + TAPTI_EXTERNAL_CORRELATION_KIND_CUSTOM2 = 5, + TAPTI_EXTERNAL_CORRELATION_KIND_SIZE, + TAPTI_EXTERNAL_CORRELATION_KIND_FORCE_INT = 0x7fffffff +} TApti_ExternalCorrelationKind; + +#define ACTIVITY_RECORD_ALIGNMENT 8 +#define PACKED_ALIGNMENT __attribute__ ((__packed__)) __attribute__ ((aligned (ACTIVITY_RECORD_ALIGNMENT))) + +/** + * \brief The kinds of activity records. + * + * Each activity record kind represents information about a GPU or an + * activity occurring on a CPU or GPU. Each kind is associated with a + * activity record structure that holds the information associated + * with the kind. + * \see TApti_Activity + * \see TApti_ActivityAPI + * \see TApti_ActivityExternalCorrelation + * \see TApti_ActivityKernel + * \see TApti_ActivityMemcpy + * \see TApti_ActivityMemcpyPtoP + * \see TApti_ActivityMemset + */ + typedef enum { + /** + * The activity record is invalid. + */ + TAPTI_ACTIVITY_KIND_INVALID = 0, + + /** + * A host<->host, host<->device, or device<->device memory copy. The + * corresponding activity record structure is \ref + */ + TAPTI_ACTIVITY_KIND_MEMCPY = 1, + + /** + * A memory set executing on the GPU. The corresponding activity + * record structure is \ref TApti_ActivityMemset. + */ + TAPTI_ACTIVITY_KIND_MEMSET = 2, + + /** + * A kernel executing on the GPU. This activity kind may significantly change + * the overall performance characteristics of the application because all + * kernel executions are serialized on the GPU. Other activity kind for kernel + * TAPTI_ACTIVITY_KIND_CONCURRENT_KERNEL doesn't break kernel concurrency. + * The corresponding activity record structure is \ref TApti_ActivityKernel. + */ + TAPTI_ACTIVITY_KIND_KERNEL = 3, + + /** + * A TANG driver API function execution. The corresponding activity + * record structure is \ref TApti_ActivityAPI. + */ + TAPTI_ACTIVITY_KIND_DRIVER = 4, + + /** + * A TANG runtime API function execution. The corresponding activity + * record structure is \ref TApti_ActivityAPI. + */ + TAPTI_ACTIVITY_KIND_RUNTIME = 5, + + /** + * Records for correlation of different programming APIs. The + * corresponding activity record structure is \ref + * TApti_ActivityExternalCorrelation. + */ + TAPTI_ACTIVITY_KIND_EXTERNAL_CORRELATION = 6, + + TAPTI_ACTIVITY_KIND_COUNT, + + TAPTI_ACTIVITY_KIND_FORCE_INT = 0x7fffffff +} TApti_ActivityKind; + +/** + * \brief The kinds of memory accessed by a memory operation/copy. + * + * Each kind represents the type of the memory + * accessed by a memory operation/copy. + */ +typedef enum { + /** + * The memory kind is unknown. + */ + TAPTI_ACTIVITY_MEMORY_KIND_UNKNOWN = 0, + + /** + * The memory is pageable. + */ + TAPTI_ACTIVITY_MEMORY_KIND_PAGEABLE = 1, + + /** + * The memory is pinned. + */ + TAPTI_ACTIVITY_MEMORY_KIND_PINNED = 2, + + /** + * The memory is on the device. + */ + TAPTI_ACTIVITY_MEMORY_KIND_DEVICE = 3, + + /** + * The memory is an array. + */ + TAPTI_ACTIVITY_MEMORY_KIND_ARRAY = 4, + + /** + * The memory is managed + */ + TAPTI_ACTIVITY_MEMORY_KIND_MANAGED = 5, + + /** + * The memory is device static + */ + TAPTI_ACTIVITY_MEMORY_KIND_DEVICE_STATIC = 6, + + /** + * The memory is managed static + */ + TAPTI_ACTIVITY_MEMORY_KIND_MANAGED_STATIC = 7, + + TAPTI_ACTIVITY_MEMORY_KIND_FORCE_INT = 0x7fffffff +} TApti_ActivityMemoryKind; + + +/** + * \brief The base activity record. + * + * The activity API uses a TApti_Activity as a generic representation + * for any activity. The 'kind' field is used to determine the + * specific activity kind, and from that the TAPTI_Activity object can + * be cast to the specific activity record type appropriate for that kind. + * + * Note that all activity record types are padded and aligned to + * ensure that each member of the record is naturally aligned. + * + * \see TApti_ActivityKind + */ +typedef struct PACKED_ALIGNMENT { + /** + * The kind of this activity. + */ + TApti_ActivityKind kind; +} TApti_Activity; + +/** + * \brief The activity record for correlation with external records + * + * This activity record correlates native TANG records (e.g. TANG Driver API, + * kernels, memcpys, ...) with records from external APIs such as OpenACC. + * (TAPTI_ACTIVITY_KIND_EXTERNAL_CORRELATION). + * + * \see TApti_ActivityKind + */ +typedef struct PACKED_ALIGNMENT { + TApti_ActivityKind kind; + + /** + * The kind of external API this record correlated to. + */ + TApti_ExternalCorrelationKind externalKind; + + /** + * The correlation ID of the associated non-TANG API record. + * The exact field in the associated external record depends + * on that record's activity kind (\see externalKind). + */ + uint64_t externalId; + + /** + * The correlation ID of the associated TANG driver or runtime API record. + */ + uint32_t correlationId; + /** + * Undefined. Reserved for internal use. + */ + uint32_t reserved; +} TApti_ActivityExternalCorrelation; + +/** + * \brief The activity record for a driver or runtime API invocation. + * + * This activity record represents an invocation of a driver or + * runtime API (TAPTI_ACTIVITY_KIND_DRIVER and + * TAPTI_ACTIVITY_KIND_RUNTIME). + */ + typedef struct PACKED_ALIGNMENT { + TApti_ActivityKind kind; + /** + * The ID of the driver or runtime function. + */ + TApti_CallbackId cbid; + /** + * The start timestamp for the function, in ns. A value of 0 for + * both the start and end timestamps indicates that timestamp + * information could not be collected for the function. + */ + uint64_t start; + + /** + * The end timestamp for the function, in ns. A value of 0 for both + * the start and end timestamps indicates that timestamp information + * could not be collected for the function. + */ + uint64_t end; + + /** + * The ID of the process where the driver or runtime TANG function + * is executing. + */ + uint32_t processId; + + /** + * The ID of the thread where the driver or runtime TANG function is + * executing. + */ + uint32_t threadId; + + /** + * The correlation ID of the driver or runtime TANG function. Each + * function invocation is assigned a unique correlation ID that is + * identical to the correlation ID in the memcpy, memset, or kernel + * activity record that is associated with this function. + */ + uint32_t correlationId; + + /** + * The return value for the function. For a TANG driver function + * with will be a TAresult value, and for a TANG runtime function + * this will be a tangError_t value. + */ + uint32_t returnValue; +} TApti_ActivityAPI; + +/** + * \brief The activity record for kernel. (deprecated) + * + * This activity record represents a kernel execution + * (TAPTI_ACTIVITY_KIND_KERNEL and + * TAPTI_ACTIVITY_KIND_CONCURRENT_KERNEL) but is no longer generated + * by TAPTI. Kernel activities are now reported using the + * TApti_ActivityKernel9 activity record. + */ + typedef struct PACKED_ALIGNMENT { + /** + * The activity record kind, must be TAPTI_ACTIVITY_KIND_KERNEL + * or TAPTI_ACTIVITY_KIND_CONCURRENT_KERNEL. + */ + TApti_ActivityKind kind; + + /** + * The cache configuration requested by the kernel. The value is one + * of the CUfunc_cache enumeration values from cuda.h. + */ + uint8_t cacheConfigRequested; + + /** + * The cache configuration used for the kernel. The value is one of + * the CUfunc_cache enumeration values from cuda.h. + */ + uint8_t cacheConfigExecuted; + + /** + * The number of registers required for each thread executing the + * kernel. + */ + uint16_t registersPerThread; + + /** + * The start timestamp for the kernel execution, in ns. A value of 0 + * for both the start and end timestamps indicates that timestamp + * information could not be collected for the kernel. + */ + uint64_t start; + + /** + * The end timestamp for the kernel execution, in ns. A value of 0 + * for both the start and end timestamps indicates that timestamp + * information could not be collected for the kernel. + */ + uint64_t end; + + /** + * The ID of the device where the kernel is executing. + */ + uint32_t deviceId; + + /** + * The ID of the context where the kernel is executing. + */ + uint32_t contextId; + + /** + * The ID of the stream where the kernel is executing. + */ + uint32_t streamId; + + /** + * The X-dimension grid size for the kernel. + */ + int32_t gridX; + + /** + * The Y-dimension grid size for the kernel. + */ + int32_t gridY; + + /** + * The Z-dimension grid size for the kernel. + */ + int32_t gridZ; + + /** + * The X-dimension block size for the kernel. + */ + int32_t blockX; + + /** + * The Y-dimension block size for the kernel. + */ + int32_t blockY; + + /** + * The Z-dimension grid size for the kernel. + */ + int32_t blockZ; + + /** + * The static shared memory allocated for the kernel, in bytes. + */ + int32_t staticSharedMemory; + + /** + * The dynamic shared memory reserved for the kernel, in bytes. + */ + int32_t dynamicSharedMemory; + + /** + * The amount of local memory reserved for each thread, in bytes. + */ + uint32_t localMemoryPerThread; + + /** + * The total amount of local memory reserved for the kernel, in + * bytes. + */ + uint32_t localMemoryTotal; + + /** + * The correlation ID of the kernel. Each kernel execution is + * assigned a unique correlation ID that is identical to the + * correlation ID in the driver API activity record that launched + * the kernel. + */ + uint32_t correlationId; + + /** + * The runtime correlation ID of the kernel. Each kernel execution + * is assigned a unique runtime correlation ID that is identical to + * the correlation ID in the runtime API activity record that + * launched the kernel. + */ + uint32_t runtimeCorrelationId; + + /** + * Undefined. Reserved for internal use. + */ + uint32_t pad; + + /** + * The name of the kernel. This name is shared across all activity + * records representing the same kernel, and so should not be + * modified. + */ + const char *name; + + /** + * Undefined. Reserved for internal use. + */ + void *reserved0; +} TApti_ActivityKernel; + +/** + * \brief The activity record for memory copies. (deprecated) + * + * This activity record represents a memory copy + * (TAPTI_ACTIVITY_KIND_MEMCPY). + */ + typedef struct PACKED_ALIGNMENT { + /** + * The activity record kind, must be TAPTI_ACTIVITY_KIND_MEMCPY. + */ + TApti_ActivityKind kind; + + /** + * The kind of the memory copy, stored as a byte to reduce record + * size. \see TApti_ActivityMemcpyKind + */ + uint8_t copyKind; + + /** + * The source memory kind read by the memory copy, stored as a byte + * to reduce record size. \see TApti_ActivityMemoryKind + */ + uint8_t srcKind; + + /** + * The destination memory kind read by the memory copy, stored as a + * byte to reduce record size. \see TApti_ActivityMemoryKind + */ + uint8_t dstKind; + + /** + * The flags associated with the memory copy. \see TApti_ActivityFlag + */ + uint8_t flags; + + /** + * The number of bytes transferred by the memory copy. + */ + uint64_t bytes; + + /** + * The start timestamp for the memory copy, in ns. A value of 0 for + * both the start and end timestamps indicates that timestamp + * information could not be collected for the memory copy. + */ + uint64_t start; + + /** + * The end timestamp for the memory copy, in ns. A value of 0 for + * both the start and end timestamps indicates that timestamp + * information could not be collected for the memory copy. + */ + uint64_t end; + + /** + * The ID of the device where the memory copy is occurring. + */ + uint32_t deviceId; + + /** + * The ID of the context where the memory copy is occurring. + */ + uint32_t contextId; + + /** + * The ID of the stream where the memory copy is occurring. + */ + uint32_t streamId; + + /** + * The correlation ID of the memory copy. Each memory copy is + * assigned a unique correlation ID that is identical to the + * correlation ID in the driver API activity record that launched + * the memory copy. + */ + uint32_t correlationId; + + /** + * The runtime correlation ID of the memory copy. Each memory copy + * is assigned a unique runtime correlation ID that is identical to + * the correlation ID in the runtime API activity record that + * launched the memory copy. + */ + uint32_t runtimeCorrelationId; + +#ifdef TAptiLP64 + /** + * Undefined. Reserved for internal use. + */ + uint32_t pad; +#endif + + /** + * Undefined. Reserved for internal use. + */ + void *reserved0; +} TApti_ActivityMemcpy; + +/** + * \brief The activity record for peer-to-peer memory copies. + * + * This activity record represents a peer-to-peer memory copy + * (TAPTI_ACTIVITY_KIND_MEMCPY2) but is no longer generated + * by TAPTI. Peer-to-peer memory copy activities are now reported using the + * TApti_ActivityMemcpyPtoP2 activity record.. + */ + typedef struct PACKED_ALIGNMENT { + /** + * The activity record kind, must be TAPTI_ACTIVITY_KIND_MEMCPY2. + */ + TApti_ActivityKind kind; + + /** + * The kind of the memory copy, stored as a byte to reduce record + * size. \see TApti_ActivityMemcpyKind + */ + uint8_t copyKind; + + /** + * The source memory kind read by the memory copy, stored as a byte + * to reduce record size. \see TApti_ActivityMemoryKind + */ + uint8_t srcKind; + + /** + * The destination memory kind read by the memory copy, stored as a + * byte to reduce record size. \see TApti_ActivityMemoryKind + */ + uint8_t dstKind; + + /** + * The flags associated with the memory copy. \see + * TApti_ActivityFlag + */ + uint8_t flags; + + /** + * The number of bytes transferred by the memory copy. + */ + uint64_t bytes; + + /** + * The start timestamp for the memory copy, in ns. A value of 0 for + * both the start and end timestamps indicates that timestamp + * information could not be collected for the memory copy. + */ + uint64_t start; + + /** + * The end timestamp for the memory copy, in ns. A value of 0 for + * both the start and end timestamps indicates that timestamp + * information could not be collected for the memory copy. + */ + uint64_t end; + + /** + * The ID of the device where the memory copy is occurring. + */ + uint32_t deviceId; + + /** + * The ID of the context where the memory copy is occurring. + */ + uint32_t contextId; + + /** + * The ID of the stream where the memory copy is occurring. + */ + uint32_t streamId; + + /** + * The ID of the device where memory is being copied from. + */ + uint32_t srcDeviceId; + + /** + * The ID of the context owning the memory being copied from. + */ + uint32_t srcContextId; + + /** + * The ID of the device where memory is being copied to. + */ + uint32_t dstDeviceId; + + /** + * The ID of the context owning the memory being copied to. + */ + uint32_t dstContextId; + + /** + * The correlation ID of the memory copy. Each memory copy is + * assigned a unique correlation ID that is identical to the + * correlation ID in the driver and runtime API activity record that + * launched the memory copy. + */ + uint32_t correlationId; + +#ifndef TAptiLP64 + /** + * Undefined. Reserved for internal use. + */ + uint32_t pad; +#endif + + /** + * Undefined. Reserved for internal use. + */ + void *reserved0; +} TApti_ActivityMemcpyPtoP; + +/** + * \brief The activity record for memset. (deprecated) + * + * This activity record represents a memory set operation + * (TAPTI_ACTIVITY_KIND_MEMSET). + */ + typedef struct PACKED_ALIGNMENT { + /** + * The activity record kind, must be TAPTI_ACTIVITY_KIND_MEMSET. + */ + TApti_ActivityKind kind; + + /** + * The value being assigned to memory by the memory set. + */ + uint32_t value; + + /** + * The number of bytes being set by the memory set. + */ + uint64_t bytes; + + /** + * The start timestamp for the memory set, in ns. A value of 0 for + * both the start and end timestamps indicates that timestamp + * information could not be collected for the memory set. + */ + uint64_t start; + + /** + * The end timestamp for the memory set, in ns. A value of 0 for + * both the start and end timestamps indicates that timestamp + * information could not be collected for the memory set. + */ + uint64_t end; + + /** + * The ID of the device where the memory set is occurring. + */ + uint32_t deviceId; + + /** + * The ID of the context where the memory set is occurring. + */ + uint32_t contextId; + + /** + * The ID of the stream where the memory set is occurring. + */ + uint32_t streamId; + + /** + * The correlation ID of the memory set. Each memory set is assigned + * a unique correlation ID that is identical to the correlation ID + * in the driver API activity record that launched the memory set. + */ + uint32_t correlationId; + + /** + * The flags associated with the memset. \see TApti_ActivityFlag + */ + uint16_t flags; + + /** + * The memory kind of the memory set \see TApti_ActivityMemoryKind + */ + uint16_t memoryKind; + + +} TApti_ActivityMemset; + +#ifdef __cplusplus +extern "C" { +#endif //! __cplusplus + +#if defined(_MSC_VER) +#define TAPTI_DEPRECATED __declspec(deprecated) +#define TAPTI_API_EXPORT __declspec(dllexport) +#define TAPTI_API_IMPORT __declspec(dllimport) +#elif defined(__GNUC__) || defined(__clang__) +#define TAPTI_DEPRECATED __attribute__((deprecated)) +#define TAPTI_API_EXPORT __attribute__((visibility("default"))) +#define TAPTI_API_IMPORT __attribute__((visibility("default"))) +#else +#define TAPTI_DEPRECATED +#define TAPTI_API_EXPORT +#define TAPTI_API_IMPORT +#endif //! UNKNOWN COMPILER + +#if defined(tapti_shared_EXPORTS) +#define TAPTI_API TAPTI_API_EXPORT +#else +#define TAPTI_API TAPTI_API_IMPORT +#endif //! For user + +/** + * \brief Enable collection of a specific kind of activity record. + * + * Enable collection of a specific kind of activity record. Multiple + * kinds can be enabled by calling this function multiple times. By + * default all activity kinds are disabled for collection. + * + * \param kind The kind of activity record to collect + * + * \retval TAPTI_SUCCESS + * \retval TAPTI_ERROR_NOT_INITIALIZED + * \retval TAPTI_ERROR_NOT_COMPATIBLE if the activity kind cannot be enabled + * \retval TAPTI_ERROR_INVALID_KIND if the activity kind is not supported + */ +TAptiResult TAPTI_API taptiActivityEnable(TApti_ActivityKind kind); + +/** + * \brief Disable collection of a specific kind of activity record. + * + * Disable collection of a specific kind of activity record. Multiple + * kinds can be disabled by calling this function multiple times. By + * default all activity kinds are disabled for collection. + * + * \param kind The kind of activity record to stop collecting + * + * \retval TAPTI_SUCCESS + * \retval TAPTI_ERROR_NOT_INITIALIZED + * \retval TAPTI_ERROR_INVALID_KIND if the activity kind is not supported + */ +TAptiResult TAPTI_API taptiActivityDisable(TApti_ActivityKind kind); + +/** + * \brief Iterate over the activity records in a buffer. + * + * This is a helper function to iterate over the activity records in a + * buffer. A buffer of activity records is typically obtained by + * receiving a TApti_BuffersCallbackCompleteFunc callback. + * + * An example of typical usage: + * \code + * TApti_Activity *record = NULL; + * TAptiResult status = TAPTI_SUCCESS; + * do { + * status = taptiActivityGetNextRecord(buffer, validSize, &record); + * if(status == TAPTI_SUCCESS) { + * // Use record here... + * } + * else if (status == TAPTI_ERROR_MAX_LIMIT_REACHED) + * break; + * else { + * goto Error; + * } + * } while (1); + * \endcode + * + * \param buffer The buffer containing activity records + * \param record Inputs the previous record returned by + * taptiActivityGetNextRecord and returns the next activity record + * from the buffer. If input value is NULL, returns the first activity + * record in the buffer. Records of kind TAPTI_ACTIVITY_KIND_CONCURRENT_KERNEL + * may contain invalid (0) timestamps, indicating that no timing information could + * be collected for lack of device memory. + * \param validBufferSizeBytes The number of valid bytes in the buffer. + * + * \retval TAPTI_SUCCESS + * \retval TAPTI_ERROR_NOT_INITIALIZED + * \retval TAPTI_ERROR_MAX_LIMIT_REACHED if no more records in the buffer + * \retval TAPTI_ERROR_INVALID_PARAMETER if \p buffer is NULL. + */ +TAptiResult TAPTI_API taptiActivityGetNextRecord(uint8_t* buffer, size_t validBufferSizeBytes, TApti_Activity **record); + +/** + * \brief Push an external correlation id for the calling thread + * + * This function notifies TAPTI that the calling thread is entering an external API region. + * When a TAPTI activity API record is created while within an external API region and + * TAPTI_ACTIVITY_KIND_EXTERNAL_CORRELATION is enabled, the activity API record will + * be preceeded by a TApti_ActivityExternalCorrelation record for each \ref TApti_ExternalCorrelationKind. + * + * \param kind The kind of external API activities should be correlated with. + * \param id External correlation id. + * + * \retval TAPTI_SUCCESS + * \retval TAPTI_ERROR_INVALID_PARAMETER The external API kind is invalid + */ +TAptiResult TAPTI_API taptiActivityPushExternalCorrelationId(TApti_ExternalCorrelationKind kind, uint64_t id); + +/** + * \brief Pop an external correlation id for the calling thread + * + * This function notifies TAPTI that the calling thread is leaving an external API region. + * + * \param kind The kind of external API activities should be correlated with. + * \param lastId If the function returns successful, contains the last external correlation id for this \p kind, can be NULL. + * + * \retval TAPTI_SUCCESS + * \retval TAPTI_ERROR_INVALID_PARAMETER The external API kind is invalid. + * \retval TAPTI_ERROR_QUEUE_EMPTY No external id is currently associated with \p kind. + */ +TAptiResult TAPTI_API taptiActivityPopExternalCorrelationId(TApti_ExternalCorrelationKind kind, uint64_t *lastId); + +/** + * \brief Request to deliver activity records via the buffer completion callback. + * + * This function returns the activity records associated with all contexts/streams + * (and the global buffers not associated with any stream) to the TAPTI client + * using the callback registered in taptiActivityRegisterCallbacks. + * + * This is a blocking call but it doesn't issue any TANG synchronization calls + * implicitly thus it's not guaranteed that all activities are completed on the + * underlying devices. Activity record is considered as completed if it has all + * the information filled up including the timestamps if any. It is the client's + * responsibility to issue necessary TANG synchronization calls before calling + * this function if all activity records with complete information are expected + * to be delivered. + * + * Behavior of the function based on the input flag: + * - ::For default flush i.e. when flag is set as 0, it returns all the + * activity buffers which have all the activity records completed, buffers need not + * to be full though. It doesn't return buffers which have one or more incomplete + * records. Default flush can be done at a regular interval in a separate thread. + * - ::For forced flush i.e. when flag TAPTI_ACTIVITY_FLAG_FLUSH_FORCED is passed + * to the function, it returns all the activity buffers including the ones which have + * one or more incomplete activity records. It's suggested for clients to do the + * force flush before the termination of the profiling session to allow remaining + * buffers to be delivered. In general, it can be done in the at-exit handler. + * + * Before calling this function, the buffer handling callback api must be activated + * by calling taptiActivityRegisterCallbacks. + * + * \param flag The flag can be set to indicate a forced flush. See TApti_ActivityFlag + * + * \retval TAPTI_SUCCESS + * \retval TAPTI_ERROR_NOT_INITIALIZED + * \retval TAPTI_ERROR_INVALID_OPERATION if not preceeded by a + * successful call to taptiActivityRegisterCallbacks + * \retval TAPTI_ERROR_UNKNOWN an internal error occurred + * + * \see taptiActivityFlushPeriod + */ +TAptiResult TAPTI_API taptiActivityFlushAll(void); + +/** + * \brief Function type for callback used by TAPTI to request an empty + * buffer for storing activity records. + * + * This callback function signals the TAPTI client that an activity + * buffer is needed by TAPTI. The activity buffer is used by TAPTI to + * store activity records. The callback function can decline the + * request by setting \p *buffer to NULL. In this case TAPTI may drop + * activity records. + * + * \param buffer Returns the new buffer. If set to NULL then no buffer + * is returned. + * \param size Returns the size of the returned buffer. + * \param maxNumRecords Returns the maximum number of records that + * should be placed in the buffer. If 0 then the buffer is filled with + * as many records as possible. If > 0 the buffer is filled with at + * most that many records before it is returned. + */ +typedef void (*TApti_BuffersCallbackRequestFunc)(uint8_t **buffer, size_t *size, size_t *maxNumRecords); + +/** + * \brief Function type for callback used by TAPTI to return a buffer + * of activity records. + * + * This callback function returns to the TAPTI client a buffer + * containing activity records. The buffer contains \p validSize + * bytes of activity records which should be read using + * taptiActivityGetNextRecord. The number of dropped records can be + * read using taptiActivityGetNumDroppedRecords. After this call TAPTI + * relinquished ownership of the buffer and will not use it + * anymore. The client may return the buffer to TAPTI using the + * TApti_BuffersCallbackRequestFunc callback. + + * \param buffer The activity record buffer. + * \param size The total size of the buffer in bytes as set in + * TApti_BuffersCallbackRequestFunc. + * \param validSize The number of valid bytes in the buffer. + */ +typedef void (*TApti_BuffersCallbackCompleteFunc)(uint8_t* buffer, size_t size, size_t validSize); + +/** + * \brief Registers callback functions with TAPTI for activity buffer + * handling. + * + * This function registers two callback functions to be used in asynchronous + * buffer handling. If registered, activity record buffers are handled using + * asynchronous requested/completed callbacks from TAPTI. + * + * Registering these callbacks prevents the client from using TAPTI's + * blocking enqueue/dequeue functions. + * + * \param funcBufferRequested callback which is invoked when an empty + * buffer is requested by TAPTI + * \param funcBufferCompleted callback which is invoked when a buffer + * containing activity records is available from TAPTI + * + * \retval TAPTI_SUCCESS + * \retval TAPTI_ERROR_INVALID_PARAMETER if either \p + * funcBufferRequested or \p funcBufferCompleted is NULL + */ +TAptiResult TAPTI_API taptiActivityRegisterCallbacks(TApti_BuffersCallbackRequestFunc funcBufferRequested, + TApti_BuffersCallbackCompleteFunc funcBufferCompleted); + +TAptiResult TAPTI_API taptiActivityPostProcess(void); + +#ifdef __cplusplus +} +#endif //! __cplusplus + +#endif // _TAPTI_ACTIVITY_HPP_ + diff --git a/third_party/sunrise/backend/include/tapti/tapti_callbacks.h b/third_party/sunrise/backend/include/tapti/tapti_callbacks.h new file mode 100755 index 0000000000..bf9e7061d6 --- /dev/null +++ b/third_party/sunrise/backend/include/tapti/tapti_callbacks.h @@ -0,0 +1,109 @@ +#ifndef __TAPTI_CALLBACKS_HPP__ +#define __TAPTI_CALLBACKS_HPP__ + +#include +#include "tapti_result.h" + +/** + * \brief Callback domains. + * + * Callback domains. Each domain represents callback points for a + * group of related API functions or TANG driver activity. + */ +typedef enum { + /** + * Invalid domain. + */ + TAPTI_CB_DOMAIN_INVALID = 0, + /** + * Domain containing callback points for all driver API functions. + */ + TAPTI_CB_DOMAIN_DRIVER_API = 1, + /** + * Domain containing callback points for all runtime API + * functions. + */ + TAPTI_CB_DOMAIN_RUNTIME_API = 2, + TAPTI_CB_DOMAIN_SIZE, + + TAPTI_CB_DOMAIN_FORCE_INT = 0x7fffffff +}TApti_CallbackDomain; + +/** + * \brief An ID for a driver API, runtime API, resource or + * synchronization callback. + * + * An ID for a driver API, runtime API, resource or synchronization + * callback. Within a driver API callback this should be interpreted + * as a tapti_driver_api_trace_cbid value. + * Within a runtime API callback this should be interpreted as a + * TAPTI_runtime_api_trace_cbid value. + * Within a resource API callback this should be interpreted as a + * ref TAPTI_CallbackIdResource value. + * Within a synchronize API callback this should be interpreted as a + * ref TAPTI_CallbackIdSync value. + */ +typedef uint32_t TApti_CallbackId; + +#ifdef __cplusplus +extern "C" { +#endif //! __cplusplus + +#if defined(_MSC_VER) +#define TAPTI_DEPRECATED __declspec(deprecated) +#define TAPTI_API_EXPORT __declspec(dllexport) +#define TAPTI_API_IMPORT __declspec(dllimport) +#elif defined(__GNUC__) || defined(__clang__) +#define TAPTI_DEPRECATED __attribute__((deprecated)) +#define TAPTI_API_EXPORT __attribute__((visibility("default"))) +#define TAPTI_API_IMPORT __attribute__((visibility("default"))) +#else +#define TAPTI_DEPRECATED +#define TAPTI_API_EXPORT +#define TAPTI_API_IMPORT +#endif //! UNKNOWN COMPILER + +#if defined(tapti_shared_EXPORTS) +#define TAPTI_API TAPTI_API_EXPORT +#else +#define TAPTI_API TAPTI_API_IMPORT +#endif //! For user + +/** + * \brief Get the name of a callback for a specific domain and callback ID. + * + * Returns a pointer to the name c_string in \p **name. + * + * \note \b Names are available only for the DRIVER and RUNTIME domains. + * + * \param domain The domain of the callback + * \param cbid The ID of the callback + * \param name Returns pointer to the name string on success, NULL otherwise + * + * \retval TAPTI_SUCCESS on success + * \retval TAPTI_ERROR_INVALID_PARAMETER if \p name is NULL, or if + * \p domain or \p cbid is invalid. + */ +TAptiResult TAPTI_API taptiGetCallbackName(TApti_CallbackDomain domain, + uint32_t cbid, + const char **name); +/** + * \brief Get the TAPTI timestamp. + * + * Returns a timestamp normalized to correspond with the start and end + * timestamps reported in the TAPTI activity records. The timestamp is + * reported in nanoseconds. + * + * \param timestamp Returns the TAPTI timestamp + * + * \retval TAPTI_SUCCESS + * \retval TAPTI_ERROR_INVALID_PARAMETER if \p timestamp is NULL + */ +TAptiResult TAPTI_API taptiGetTimestamp(uint64_t *timestamp); + +#ifdef __cplusplus +} +#endif //! __cplusplus + +#endif // __TAPTI_CALLBACKS_HPP__ + diff --git a/third_party/sunrise/backend/include/tapti/tapti_driver_cbid.h b/third_party/sunrise/backend/include/tapti/tapti_driver_cbid.h new file mode 100755 index 0000000000..075daf5c92 --- /dev/null +++ b/third_party/sunrise/backend/include/tapti/tapti_driver_cbid.h @@ -0,0 +1,203 @@ +// ************************************************************************* +// Definitions of indices for API functions, unique across entire API +// ************************************************************************* + +// This file is generated. + +typedef enum tapti_driver_api_trace_cbid_enum { + TAPTI_DRIVER_TRACE_CBID_INVALID = 0, + TAPTI_DRIVER_TRACE_CBID_taRuntimeGetVersion = 1, + TAPTI_DRIVER_TRACE_CBID_taGetExportTable = 2, + TAPTI_DRIVER_TRACE_CBID_taDriverGetVersion = 3, + TAPTI_DRIVER_TRACE_CBID_taKernelDriverGetVersion = 4, + TAPTI_DRIVER_TRACE_CBID_taGetErrorString = 5, + TAPTI_DRIVER_TRACE_CBID_taGetErrorName = 6, + TAPTI_DRIVER_TRACE_CBID_taInit = 7, + TAPTI_DRIVER_TRACE_CBID_taDeviceGet = 8, + TAPTI_DRIVER_TRACE_CBID_taDeviceReset = 9, + TAPTI_DRIVER_TRACE_CBID_taDeviceSynchronize = 10, + TAPTI_DRIVER_TRACE_CBID___taDeviceSynchronizeCurrent = 11, + TAPTI_DRIVER_TRACE_CBID_taDeviceGetByPCIBusId = 12, + TAPTI_DRIVER_TRACE_CBID_taDeviceGetPCIBusId = 13, + TAPTI_DRIVER_TRACE_CBID_taDeviceGetCount = 14, + TAPTI_DRIVER_TRACE_CBID_taDeviceGetAttribute = 15, + TAPTI_DRIVER_TRACE_CBID_taDeviceGetName = 16, + TAPTI_DRIVER_TRACE_CBID_taDeviceGetUuid = 17, + TAPTI_DRIVER_TRACE_CBID_taDeviceTotalMem = 18, + TAPTI_DRIVER_TRACE_CBID_taMemGetInfo = 19, + TAPTI_DRIVER_TRACE_CBID_taDeviceMemGetInfo = 20, + TAPTI_DRIVER_TRACE_CBID_taCtxCreate = 21, + TAPTI_DRIVER_TRACE_CBID_taCtxDestroy = 22, + TAPTI_DRIVER_TRACE_CBID_taCtxSetCurrent = 23, + TAPTI_DRIVER_TRACE_CBID_taCtxQueryCurrent = 24, + TAPTI_DRIVER_TRACE_CBID_taCtxGetCurrent = 25, + TAPTI_DRIVER_TRACE_CBID_taCtxGetCurrentDevice = 26, + TAPTI_DRIVER_TRACE_CBID_taCtxGetDevice = 27, + TAPTI_DRIVER_TRACE_CBID_taCtxGetOrdinal = 28, + TAPTI_DRIVER_TRACE_CBID_taCtxPushCurrent = 29, + TAPTI_DRIVER_TRACE_CBID_taCtxPopCurrent = 30, + TAPTI_DRIVER_TRACE_CBID_taCtxSynchronize = 31, + TAPTI_DRIVER_TRACE_CBID_taDevicePrimaryCtxRetain = 32, + TAPTI_DRIVER_TRACE_CBID_taDevicePrimaryCtxRelease = 33, + TAPTI_DRIVER_TRACE_CBID_taDevicePrimaryCtxReset = 34, + TAPTI_DRIVER_TRACE_CBID_taDevicePrimaryCtxSetFlags = 35, + TAPTI_DRIVER_TRACE_CBID_taDevicePrimaryCtxGetState = 36, + TAPTI_DRIVER_TRACE_CBID_taDeviceGetOrdinal = 37, + TAPTI_DRIVER_TRACE_CBID_taCtxGetFunction = 38, + TAPTI_DRIVER_TRACE_CBID_taCtxRegisterFunction = 39, + TAPTI_DRIVER_TRACE_CBID_taCtxGetVariable = 40, + TAPTI_DRIVER_TRACE_CBID_taCtxRegisterVariable = 41, + TAPTI_DRIVER_TRACE_CBID_taCtxGetLimit = 42, + TAPTI_DRIVER_TRACE_CBID_taCtxSetLimit = 43, + TAPTI_DRIVER_TRACE_CBID___taCtxQueryLimit = 44, + TAPTI_DRIVER_TRACE_CBID_taCtxGetBuiltInFunction = 45, + TAPTI_DRIVER_TRACE_CBID_taCtxRegisterBuiltInFunction = 46, + TAPTI_DRIVER_TRACE_CBID_taMemAlloc = 47, + TAPTI_DRIVER_TRACE_CBID_taMemFree = 48, + TAPTI_DRIVER_TRACE_CBID_taMemFreeAsync = 49, + TAPTI_DRIVER_TRACE_CBID_taMemFreeAsync_ptsz = 50, + TAPTI_DRIVER_TRACE_CBID_taMemAllocHost = 51, + TAPTI_DRIVER_TRACE_CBID_taMemHostAlloc = 52, + TAPTI_DRIVER_TRACE_CBID_taMemFreeHost = 53, + TAPTI_DRIVER_TRACE_CBID_taMemHostGetDevicePointer = 54, + TAPTI_DRIVER_TRACE_CBID_taMemHostGetFlags = 55, + TAPTI_DRIVER_TRACE_CBID_taMemHostRegister = 56, + TAPTI_DRIVER_TRACE_CBID_taMemHostUnregister = 57, + TAPTI_DRIVER_TRACE_CBID_taPointerGetAttribute = 58, + TAPTI_DRIVER_TRACE_CBID_taMemset = 59, + TAPTI_DRIVER_TRACE_CBID_taMemset_ptds = 60, + TAPTI_DRIVER_TRACE_CBID_taMemsetAsync = 61, + TAPTI_DRIVER_TRACE_CBID_taMemsetAsync_ptsz = 62, + TAPTI_DRIVER_TRACE_CBID_taMemcpyH2HAsync = 63, + TAPTI_DRIVER_TRACE_CBID_taMemcpyH2HAsync_ptsz = 64, + TAPTI_DRIVER_TRACE_CBID_taMemcpyH2DAsync = 65, + TAPTI_DRIVER_TRACE_CBID_taMemcpyH2DAsync_ptsz = 66, + TAPTI_DRIVER_TRACE_CBID_taMemcpyD2HAsync = 67, + TAPTI_DRIVER_TRACE_CBID_taMemcpyD2HAsync_ptsz = 68, + TAPTI_DRIVER_TRACE_CBID_taMemcpyD2DAsync = 69, + TAPTI_DRIVER_TRACE_CBID_taMemcpyD2DAsync_ptsz = 70, + TAPTI_DRIVER_TRACE_CBID_taMemcpyH2H = 71, + TAPTI_DRIVER_TRACE_CBID_taMemcpyH2H_ptds = 72, + TAPTI_DRIVER_TRACE_CBID_taMemcpyH2D = 73, + TAPTI_DRIVER_TRACE_CBID_taMemcpyH2D_ptds = 74, + TAPTI_DRIVER_TRACE_CBID_taMemcpyD2H = 75, + TAPTI_DRIVER_TRACE_CBID_taMemcpyD2H_ptds = 76, + TAPTI_DRIVER_TRACE_CBID_taMemcpyD2D = 77, + TAPTI_DRIVER_TRACE_CBID_taMemcpyD2D_ptds = 78, + TAPTI_DRIVER_TRACE_CBID_taStreamCreate = 79, + TAPTI_DRIVER_TRACE_CBID_taStreamCreateWithPriority = 80, + TAPTI_DRIVER_TRACE_CBID_taStreamGetPriority = 81, + TAPTI_DRIVER_TRACE_CBID_taStreamGetPriority_ptsz = 82, + TAPTI_DRIVER_TRACE_CBID_taStreamGetFlags = 83, + TAPTI_DRIVER_TRACE_CBID_taStreamGetFlags_ptsz = 84, + TAPTI_DRIVER_TRACE_CBID_taStreamGetId = 85, + TAPTI_DRIVER_TRACE_CBID_taStreamGetId_ptsz = 86, + TAPTI_DRIVER_TRACE_CBID_taStreamDestroy = 87, + TAPTI_DRIVER_TRACE_CBID_taStreamC2Ctransfers = 88, + TAPTI_DRIVER_TRACE_CBID_taDeviceGetP2PAttribute = 89, + TAPTI_DRIVER_TRACE_CBID_taDeviceGetPeerPointer = 90, + TAPTI_DRIVER_TRACE_CBID_taDeviceCanAccessPeer = 91, + TAPTI_DRIVER_TRACE_CBID_taDeviceEnablePeerAccess = 92, + TAPTI_DRIVER_TRACE_CBID_taDeviceDisablePeerAccess = 93, + TAPTI_DRIVER_TRACE_CBID_taMemcpyPeer = 94, + TAPTI_DRIVER_TRACE_CBID_taMemcpyPeerAsync = 95, + TAPTI_DRIVER_TRACE_CBID_taMemcpyPeer_v2 = 96, + TAPTI_DRIVER_TRACE_CBID_taMemcpyPeer_v2_ptds = 97, + TAPTI_DRIVER_TRACE_CBID_taMemcpyPeerAsync_v2 = 98, + TAPTI_DRIVER_TRACE_CBID_taMemcpyPeerAsync_v2_ptsz = 99, + TAPTI_DRIVER_TRACE_CBID_taMemcpyFromPeerAsync = 100, + TAPTI_DRIVER_TRACE_CBID_taMemcpyFromPeerAsync_ptsz = 101, + TAPTI_DRIVER_TRACE_CBID_taMemcpyToPeerAsync = 102, + TAPTI_DRIVER_TRACE_CBID_taMemcpyToPeerAsync_ptsz = 103, + TAPTI_DRIVER_TRACE_CBID_taStreamWaitEvent = 104, + TAPTI_DRIVER_TRACE_CBID_taStreamWaitEvent_ptsz = 105, + TAPTI_DRIVER_TRACE_CBID_taStreamSynchronize = 106, + TAPTI_DRIVER_TRACE_CBID_taStreamSynchronize_ptsz = 107, + TAPTI_DRIVER_TRACE_CBID_taStreamQuery = 108, + TAPTI_DRIVER_TRACE_CBID_taStreamQuery_ptsz = 109, + TAPTI_DRIVER_TRACE_CBID_taEventCreate = 110, + TAPTI_DRIVER_TRACE_CBID_taEventDestroy = 111, + TAPTI_DRIVER_TRACE_CBID_taEventRecord = 112, + TAPTI_DRIVER_TRACE_CBID_taEventRecord_ptsz = 113, + TAPTI_DRIVER_TRACE_CBID_taEventRecordWithFlags = 114, + TAPTI_DRIVER_TRACE_CBID_taEventRecordWithFlags_ptsz = 115, + TAPTI_DRIVER_TRACE_CBID_taEventSynchronize = 116, + TAPTI_DRIVER_TRACE_CBID_taEventSynchronizeWithFlags = 117, + TAPTI_DRIVER_TRACE_CBID_taEventElapsedTime = 118, + TAPTI_DRIVER_TRACE_CBID_taEventQuery = 119, + TAPTI_DRIVER_TRACE_CBID_taEventQueryTimestamp = 120, + TAPTI_DRIVER_TRACE_CBID_taGetBuiltinModule = 121, + TAPTI_DRIVER_TRACE_CBID_taModuleLoad = 122, + TAPTI_DRIVER_TRACE_CBID_taModuleLoadData = 123, + TAPTI_DRIVER_TRACE_CBID_taModuleUnload = 124, + TAPTI_DRIVER_TRACE_CBID_taModuleLoadFatBinaryManaged = 125, + TAPTI_DRIVER_TRACE_CBID_taModuleUnloadManaged = 126, + TAPTI_DRIVER_TRACE_CBID_taModuleSymbolTypeGetName = 127, + TAPTI_DRIVER_TRACE_CBID_taModuleIterateSymbols = 128, + TAPTI_DRIVER_TRACE_CBID_taModuleGetFunction = 129, + TAPTI_DRIVER_TRACE_CBID_taVariableGetInfo = 130, + TAPTI_DRIVER_TRACE_CBID_taFuncGetAttribute = 131, + TAPTI_DRIVER_TRACE_CBID_taFuncGetModule = 132, + TAPTI_DRIVER_TRACE_CBID_taFunctionGetAddress = 133, + TAPTI_DRIVER_TRACE_CBID_taFunctionGetNumArgs = 134, + TAPTI_DRIVER_TRACE_CBID_taFunctionGetInfo = 135, + TAPTI_DRIVER_TRACE_CBID_taModuleGetVariable = 136, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyFromDevice = 137, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyFromDevice_ptds = 138, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyFromDeviceAsync = 139, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyFromDeviceAsync_ptsz = 140, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyFromHost = 141, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyFromHost_ptds = 142, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyFromHostAsync = 143, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyFromHostAsync_ptsz = 144, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyToDevice = 145, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyToDevice_ptds = 146, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyToDeviceAsync = 147, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyToDeviceAsync_ptsz = 148, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyToHost = 149, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyToHost_ptds = 150, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyToHostAsync = 151, + TAPTI_DRIVER_TRACE_CBID_taVariableCopyToHostAsync_ptsz = 152, + TAPTI_DRIVER_TRACE_CBID_taEnqueueCommand = 153, + TAPTI_DRIVER_TRACE_CBID_taEnqueueCommand_ptsz = 154, + TAPTI_DRIVER_TRACE_CBID_taLaunchFunction = 155, + TAPTI_DRIVER_TRACE_CBID_taLaunchFunction_ptsz = 156, + TAPTI_DRIVER_TRACE_CBID_taLaunchKernel = 157, + TAPTI_DRIVER_TRACE_CBID_taLaunchKernel_ptsz = 158, + TAPTI_DRIVER_TRACE_CBID_taLaunchHostFunc = 159, + TAPTI_DRIVER_TRACE_CBID_taLaunchHostFunc_ptsz = 160, + TAPTI_DRIVER_TRACE_CBID_taLaunchHostFuncProxy = 161, + TAPTI_DRIVER_TRACE_CBID_taLaunchHostFuncProxy_ptsz = 162, + TAPTI_DRIVER_TRACE_CBID_taStreamAddCallback = 163, + TAPTI_DRIVER_TRACE_CBID_taStreamAddCallback_ptsz = 164, + TAPTI_DRIVER_TRACE_CBID_taOccupancyMaxActiveBlocksPerMultiprocessor = 165, + TAPTI_DRIVER_TRACE_CBID_taStreamBeginCapture = 166, + TAPTI_DRIVER_TRACE_CBID_taStreamBeginCapture_ptsz = 167, + TAPTI_DRIVER_TRACE_CBID_taThreadExchangeStreamCaptureMode = 168, + TAPTI_DRIVER_TRACE_CBID_taStreamEndCapture = 169, + TAPTI_DRIVER_TRACE_CBID_taStreamEndCapture_ptsz = 170, + TAPTI_DRIVER_TRACE_CBID_taStreamIsCapturing = 171, + TAPTI_DRIVER_TRACE_CBID_taStreamIsCapturing_ptsz = 172, + TAPTI_DRIVER_TRACE_CBID_taStreamGetCaptureInfo = 173, + TAPTI_DRIVER_TRACE_CBID_taStreamGetCaptureInfo_ptsz = 174, + TAPTI_DRIVER_TRACE_CBID_taGraphInstantiateWithFlags = 175, + TAPTI_DRIVER_TRACE_CBID_taGraphLaunch = 176, + TAPTI_DRIVER_TRACE_CBID_taGraphLaunch_ptsz = 177, + TAPTI_DRIVER_TRACE_CBID_taGraphDestroy = 178, + TAPTI_DRIVER_TRACE_CBID_taGraphExecDestroy = 179, + TAPTI_DRIVER_TRACE_CBID_taGraphGetInfo = 180, + TAPTI_DRIVER_TRACE_CBID_taGraphCreate = 181, + TAPTI_DRIVER_TRACE_CBID_taGraphAddHostNode = 182, + TAPTI_DRIVER_TRACE_CBID_taGraphAddKernelNode = 183, + TAPTI_DRIVER_TRACE_CBID_taProfilerStart = 184, + TAPTI_DRIVER_TRACE_CBID_taProfilerStop = 185, + TAPTI_DRIVER_TRACE_CBID_taIpcGetMemHandle = 186, + TAPTI_DRIVER_TRACE_CBID_taIpcOpenMemHandle = 187, + TAPTI_DRIVER_TRACE_CBID_taIpcCloseMemHandle = 188, + TAPTI_DRIVER_TRACE_CBID_taIpcGetEventHandle = 189, + TAPTI_DRIVER_TRACE_CBID_taIpcOpenEventHandle = 190, + TAPTI_DRIVER_TRACE_CBID_taMemAllocAsync = 191, + TAPTI_DRIVER_TRACE_CBID_taMemAllocAsync_ptsz = 192, + TAPTI_DRIVER_TRACE_CBID_FORCE_INT = 0x7fffffff, +} tapti_driver_api_trace_cbid; + diff --git a/third_party/sunrise/backend/include/tapti/tapti_result.h b/third_party/sunrise/backend/include/tapti/tapti_result.h new file mode 100755 index 0000000000..645afddeaa --- /dev/null +++ b/third_party/sunrise/backend/include/tapti/tapti_result.h @@ -0,0 +1,248 @@ +#ifndef _TAPTI_RESULT_HPP_ +#define _TAPTI_RESULT_HPP_ + +/** + * \brief TAPTI result codes. + * + * Error and result codes returned by TAPTI functions. + */ +typedef enum { + /** + * No error. + */ + TAPTI_SUCCESS = 0, + /** + * One or more of the parameters is invalid. + */ + TAPTI_ERROR_INVALID_PARAMETER = 1, + /** + * The device does not correspond to a valid TANG device. + */ + TAPTI_ERROR_INVALID_DEVICE = 2, + /** + * The context is NULL or not valid. + */ + TAPTI_ERROR_INVALID_CONTEXT = 3, + /** + * The event domain id is invalid. + */ + TAPTI_ERROR_INVALID_EVENT_DOMAIN_ID = 4, + /** + * The event id is invalid. + */ + TAPTI_ERROR_INVALID_EVENT_ID = 5, + /** + * The event name is invalid. + */ + TAPTI_ERROR_INVALID_EVENT_NAME = 6, + /** + * The current operation cannot be performed due to dependency on + * other factors. + */ + TAPTI_ERROR_INVALID_OPERATION = 7, + /** + * Unable to allocate enough memory to perform the requested + * operation. + */ + TAPTI_ERROR_OUT_OF_MEMORY = 8, + /** + * An error occurred on the performance monitoring hardware. + */ + TAPTI_ERROR_HARDWARE = 9, + /** + * The output buffer size is not sufficient to return all + * requested data. + */ + TAPTI_ERROR_PARAMETER_SIZE_NOT_SUFFICIENT = 10, + /** + * API is not implemented. + */ + TAPTI_ERROR_API_NOT_IMPLEMENTED = 11, + /** + * The maximum limit is reached. + */ + TAPTI_ERROR_MAX_LIMIT_REACHED = 12, + /** + * The object is not yet ready to perform the requested operation. + */ + TAPTI_ERROR_NOT_READY = 13, + /** + * The current operation is not compatible with the current state + * of the object + */ + TAPTI_ERROR_NOT_COMPATIBLE = 14, + /** + * TAPTI is unable to initialize its connection to the TANG + * driver. + */ + TAPTI_ERROR_NOT_INITIALIZED = 15, + /** + * The metric id is invalid. + */ + TAPTI_ERROR_INVALID_METRIC_ID = 16, + /** + * The metric name is invalid. + */ + TAPTI_ERROR_INVALID_METRIC_NAME = 17, + /** + * The queue is empty. + */ + TAPTI_ERROR_QUEUE_EMPTY = 18, + /** + * Invalid handle (internal?). + */ + TAPTI_ERROR_INVALID_HANDLE = 19, + /** + * Invalid stream. + */ + TAPTI_ERROR_INVALID_STREAM = 20, + /** + * Invalid kind. + */ + TAPTI_ERROR_INVALID_KIND = 21, + /** + * Invalid event value. + */ + TAPTI_ERROR_INVALID_EVENT_VALUE = 22, + /** + * TAPTI is disabled due to conflicts with other enabled profilers + */ + TAPTI_ERROR_DISABLED = 23, + /** + * Invalid module. + */ + TAPTI_ERROR_INVALID_MODULE = 24, + /** + * Invalid metric value. + */ + TAPTI_ERROR_INVALID_METRIC_VALUE = 25, + /** + * The performance monitoring hardware is in use by other client. + */ + TAPTI_ERROR_HARDWARE_BUSY = 26, + /** + * The attempted operation is not supported on the current + * system or device. + */ + TAPTI_ERROR_NOT_SUPPORTED = 27, + /** + * Unified memory profiling is not supported on the system. + * Potential reason could be unsupported OS or architecture. + */ + TAPTI_ERROR_UM_PROFILING_NOT_SUPPORTED = 28, + /** + * Unified memory profiling is not supported on the device + */ + TAPTI_ERROR_UM_PROFILING_NOT_SUPPORTED_ON_DEVICE = 29, + /** + * Unified memory profiling is not supported on a multi-GPU + * configuration without P2P support between any pair of devices + */ + TAPTI_ERROR_UM_PROFILING_NOT_SUPPORTED_ON_NON_P2P_DEVICES = 30, + /** + * Profiling on virtualized GPU is not supported. + */ + TAPTI_ERROR_VIRTUALIZED_DEVICE_NOT_SUPPORTED = 33, + /** + * User doesn't have sufficient privileges which are required to + * start the profiling session. + * One possible reason for this may be that the NVIDIA driver or your system + * administrator may have restricted access to the NVIDIA GPU performance counters. + * To learn how to resolve this issue and find more information, please visit + * https://developer.nvidia.com/TAPTI_ERROR_INSUFFICIENT_PRIVILEGES + */ + TAPTI_ERROR_INSUFFICIENT_PRIVILEGES = 35, + /** + * Legacy TAPTI Profiling API i.e. event API from the header TAPTI_events.h and + * metric API from the header TAPTI_metrics.h are not compatible with the + * Profiling API in the header TAPTI_profiler_target.h and Perfworks metrics API + * in the headers nvperf_host.h and nvperf_target.h. + */ + TAPTI_ERROR_OLD_PROFILER_API_INITIALIZED = 36, + /** + * Missing definition of the OpenACC API routine in the linked OpenACC library. + * + * One possible reason is that OpenACC library is linked statically in the + * user application, which might not have the definition of all the OpenACC + * API routines needed for the OpenACC profiling, as compiler might ignore + * definitions for the functions not used in the application. This issue + * can be mitigated by linking the OpenACC library dynamically. + */ + TAPTI_ERROR_OPENACC_UNDEFINED_ROUTINE = 37, + /** + * Legacy TAPTI Profiling API i.e. event API from the header TAPTI_events.h and + * metric API from the header TAPTI_metrics.h are not supported on devices with + * compute capability 7.5 and higher (i.e. Turing and later GPU architectures). + * These API will be deprecated in a future TANG release. These are replaced by + * Profiling API in the header TAPTI_profiler_target.h and Perfworks metrics API + * in the headers nvperf_host.h and nvperf_target.h. + */ + TAPTI_ERROR_LEGACY_PROFILER_NOT_SUPPORTED = 38, + /** + * TAPTI doesn't allow multiple callback subscribers. Only a single subscriber + * can be registered at a time. + * Same error code is used when application is launched using NVIDIA tools + * like nvprof, Visual Profiler, Nsight Systems, Nsight Compute, cuda-gdb and + * cuda-memcheck. + */ + TAPTI_ERROR_MULTIPLE_SUBSCRIBERS_NOT_SUPPORTED = 39, + /** + * Profiling on virtualized GPU is not allowed by hypervisor. + */ + TAPTI_ERROR_VIRTUALIZED_DEVICE_INSUFFICIENT_PRIVILEGES = 40, + /** + * Profiling and tracing are not allowed when confidential computing mode + * is enabled. + */ + TAPTI_ERROR_CONFIDENTIAL_COMPUTING_NOT_SUPPORTED = 41, + /** + * An unknown internal error has occurred. + */ + TAPTI_ERROR_UNKNOWN = 999, + TAPTI_ERROR_FORCE_INT = 0x7fffffff +} TAptiResult; + +#ifdef __cplusplus +extern "C" { +#endif //! __cplusplus + +#if defined(_MSC_VER) +#define TAPTI_DEPRECATED __declspec(deprecated) +#define TAPTI_API_EXPORT __declspec(dllexport) +#define TAPTI_API_IMPORT __declspec(dllimport) +#elif defined(__GNUC__) || defined(__clang__) +#define TAPTI_DEPRECATED __attribute__((deprecated)) +#define TAPTI_API_EXPORT __attribute__((visibility("default"))) +#define TAPTI_API_IMPORT __attribute__((visibility("default"))) +#else +#define TAPTI_DEPRECATED +#define TAPTI_API_EXPORT +#define TAPTI_API_IMPORT +#endif //! UNKNOWN COMPILER + +#if defined(tapti_shared_EXPORTS) +#define TAPTI_API TAPTI_API_EXPORT +#else +#define TAPTI_API TAPTI_API_IMPORT +#endif //! For user + +/** + * \brief Get the descriptive string for a TAptiResult. + * + * Return the descriptive string for a TAptiResult in \p *str. + * \note \b Thread-safety: this function is thread safe. + * + * \param result The result to get the string for + * \param str Returns the string + * + * \retval TAPTI_SUCCESS on success + * \retval TAPTI_ERROR_INVALID_PARAMETER if \p str is NULL or \p + * result is not a valid TAptiResult + */ +TAptiResult TAPTI_API taptiGetResultString(TAptiResult result, const char **str); + +#ifdef __cplusplus +} +#endif //! __cplusplus + +#endif // _TAPTI_RESULT_HPP_ diff --git a/third_party/sunrise/backend/include/tapti/tapti_runtime_cbid.h b/third_party/sunrise/backend/include/tapti/tapti_runtime_cbid.h new file mode 100755 index 0000000000..75cb5fa3b6 --- /dev/null +++ b/third_party/sunrise/backend/include/tapti/tapti_runtime_cbid.h @@ -0,0 +1,147 @@ +// ************************************************************************* +// Definitions of indices for API functions, unique across entire API +// ************************************************************************* + +// This file is generated. + +typedef enum tapti_runtime_api_trace_cbid_enum { + TAPTI_RUNTIME_TRACE_CBID_INVALID = 0, + TAPTI_RUNTIME_TRACE_CBID_tangRuntimeGetVersion = 1, + TAPTI_RUNTIME_TRACE_CBID_tangDriverGetVersion = 2, + TAPTI_RUNTIME_TRACE_CBID_tangGetDeviceCount = 3, + TAPTI_RUNTIME_TRACE_CBID_tangGetDevice = 4, + TAPTI_RUNTIME_TRACE_CBID_tangSetDevice = 5, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceSynchronize = 6, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceReset = 7, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceGetByPCIBusId = 8, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceGetPCIBusId = 9, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceSetCacheConfig = 10, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceGetCacheConfig = 11, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceSetSharedMemConfig = 12, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceGetSharedMemConfig = 13, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceGetStreamPriorityRange = 14, + TAPTI_RUNTIME_TRACE_CBID_tangSetValidDevices = 15, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceGetLimit = 16, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceSetLimit = 17, + TAPTI_RUNTIME_TRACE_CBID_tangMalloc = 18, + TAPTI_RUNTIME_TRACE_CBID_tangMallocAsync = 19, + TAPTI_RUNTIME_TRACE_CBID_tangMallocAsync_ptsz = 20, + TAPTI_RUNTIME_TRACE_CBID_tangFree = 21, + TAPTI_RUNTIME_TRACE_CBID_tangFreeAsync = 22, + TAPTI_RUNTIME_TRACE_CBID_tangFreeAsync_ptsz = 23, + TAPTI_RUNTIME_TRACE_CBID_tangMallocHost = 24, + TAPTI_RUNTIME_TRACE_CBID_tangHostAlloc = 25, + TAPTI_RUNTIME_TRACE_CBID_tangHostGetDevicePointer = 26, + TAPTI_RUNTIME_TRACE_CBID_tangHostGetFlags = 27, + TAPTI_RUNTIME_TRACE_CBID_tangFreeHost = 28, + TAPTI_RUNTIME_TRACE_CBID_tangHostRegister = 29, + TAPTI_RUNTIME_TRACE_CBID_tangHostUnregister = 30, + TAPTI_RUNTIME_TRACE_CBID_tangMemGetInfo = 31, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpy = 32, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpy_ptds = 33, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyAsync = 34, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyAsync_ptsz = 35, + TAPTI_RUNTIME_TRACE_CBID_tangMemset = 36, + TAPTI_RUNTIME_TRACE_CBID_tangMemset_ptds = 37, + TAPTI_RUNTIME_TRACE_CBID_tangMemsetAsync = 38, + TAPTI_RUNTIME_TRACE_CBID_tangMemsetAsync_ptsz = 39, + TAPTI_RUNTIME_TRACE_CBID_tangGetSymbolAddress = 40, + TAPTI_RUNTIME_TRACE_CBID_tangGetSymbolSize = 41, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyFromSymbol = 42, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyFromSymbol_ptds = 43, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyFromSymbolAsync = 44, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyFromSymbolAsync_ptsz = 45, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyToSymbol = 46, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyToSymbol_ptds = 47, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyToSymbolAsync = 48, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyToSymbolAsync_ptsz = 49, + TAPTI_RUNTIME_TRACE_CBID_tangStreamCreate = 50, + TAPTI_RUNTIME_TRACE_CBID_tangStreamCreateWithFlags = 51, + TAPTI_RUNTIME_TRACE_CBID_tangStreamCreateWithPriority = 52, + TAPTI_RUNTIME_TRACE_CBID_tangStreamGetPriority = 53, + TAPTI_RUNTIME_TRACE_CBID_tangStreamGetPriority_ptsz = 54, + TAPTI_RUNTIME_TRACE_CBID_tangStreamAddCallback = 55, + TAPTI_RUNTIME_TRACE_CBID_tangStreamAddCallback_ptsz = 56, + TAPTI_RUNTIME_TRACE_CBID_tangLaunchHostFunc = 57, + TAPTI_RUNTIME_TRACE_CBID_tangLaunchHostFunc_ptsz = 58, + TAPTI_RUNTIME_TRACE_CBID_tangStreamDestroy = 59, + TAPTI_RUNTIME_TRACE_CBID_tangStreamSynchronize = 60, + TAPTI_RUNTIME_TRACE_CBID_tangStreamSynchronize_ptsz = 61, + TAPTI_RUNTIME_TRACE_CBID_tangStreamC2Ctransfers = 62, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceGetPeerPointer = 63, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceGetP2PAttribute = 64, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceCanAccessPeer = 65, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceEnablePeerAccess = 66, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceDisablePeerAccess = 67, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyPeer = 68, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyPeerAsync = 69, + TAPTI_RUNTIME_TRACE_CBID_tangEngineCollAssign = 70, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyPeer_v2 = 71, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyPeer_v2_ptds = 72, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyPeerAsync_v2 = 73, + TAPTI_RUNTIME_TRACE_CBID_tangMemcpyPeerAsync_v2_ptsz = 74, + TAPTI_RUNTIME_TRACE_CBID_tangStreamQuery = 75, + TAPTI_RUNTIME_TRACE_CBID_tangStreamQuery_ptsz = 76, + TAPTI_RUNTIME_TRACE_CBID_tangStreamWaitEvent = 77, + TAPTI_RUNTIME_TRACE_CBID_tangStreamWaitEvent_ptsz = 78, + TAPTI_RUNTIME_TRACE_CBID_tangStreamGetFlags = 79, + TAPTI_RUNTIME_TRACE_CBID_tangStreamGetFlags_ptsz = 80, + TAPTI_RUNTIME_TRACE_CBID_tangStreamGetId = 81, + TAPTI_RUNTIME_TRACE_CBID_tangStreamGetId_ptsz = 82, + TAPTI_RUNTIME_TRACE_CBID_tangEventCreateWithFlags = 83, + TAPTI_RUNTIME_TRACE_CBID_tangEventCreate = 84, + TAPTI_RUNTIME_TRACE_CBID_tangEventRecord = 85, + TAPTI_RUNTIME_TRACE_CBID_tangEventRecord_ptsz = 86, + TAPTI_RUNTIME_TRACE_CBID_tangEventRecordWithFlags = 87, + TAPTI_RUNTIME_TRACE_CBID_tangEventRecordWithFlags_ptsz = 88, + TAPTI_RUNTIME_TRACE_CBID_tangEventDestroy = 89, + TAPTI_RUNTIME_TRACE_CBID_tangEventSynchronize = 90, + TAPTI_RUNTIME_TRACE_CBID_tangEventSynchronizeWithFlags = 91, + TAPTI_RUNTIME_TRACE_CBID_tangEventElapsedTime = 92, + TAPTI_RUNTIME_TRACE_CBID_tangEventQuery = 93, + TAPTI_RUNTIME_TRACE_CBID_tangEventQueryTimestamp = 94, + TAPTI_RUNTIME_TRACE_CBID_tangDeviceGetAttribute = 95, + TAPTI_RUNTIME_TRACE_CBID_tangGetDeviceProperties = 96, + TAPTI_RUNTIME_TRACE_CBID_tangChooseDevice = 97, + TAPTI_RUNTIME_TRACE_CBID_tangFuncGetAttributes = 98, + TAPTI_RUNTIME_TRACE_CBID_tangFuncSetAttribute = 99, + TAPTI_RUNTIME_TRACE_CBID_tangFuncSetCacheConfig = 100, + TAPTI_RUNTIME_TRACE_CBID_tangFuncSetSharedMemConfig = 101, + TAPTI_RUNTIME_TRACE_CBID_tangSetDoubleForDevice = 102, + TAPTI_RUNTIME_TRACE_CBID_tangSetDoubleForHost = 103, + TAPTI_RUNTIME_TRACE_CBID_tangOccupancyMaxActiveBlocksPerMultiprocessor = 104, + TAPTI_RUNTIME_TRACE_CBID_tangOccupancyMaxActiveBlocksPerMultiprocessorWithFlags = 105, + TAPTI_RUNTIME_TRACE_CBID_tangPointerGetAttributes = 106, + TAPTI_RUNTIME_TRACE_CBID_tangStreamBeginCapture = 107, + TAPTI_RUNTIME_TRACE_CBID_tangStreamBeginCapture_ptsz = 108, + TAPTI_RUNTIME_TRACE_CBID_tangStreamEndCapture = 109, + TAPTI_RUNTIME_TRACE_CBID_tangStreamEndCapture_ptsz = 110, + TAPTI_RUNTIME_TRACE_CBID_tangStreamIsCapturing = 111, + TAPTI_RUNTIME_TRACE_CBID_tangStreamIsCapturing_ptsz = 112, + TAPTI_RUNTIME_TRACE_CBID_tangStreamGetCaptureInfo = 113, + TAPTI_RUNTIME_TRACE_CBID_tangStreamGetCaptureInfo_ptsz = 114, + TAPTI_RUNTIME_TRACE_CBID_tangGraphInstantiate = 115, + TAPTI_RUNTIME_TRACE_CBID_tangGraphLaunch = 116, + TAPTI_RUNTIME_TRACE_CBID_tangGraphLaunch_ptsz = 117, + TAPTI_RUNTIME_TRACE_CBID_tangGraphInstantiateWithFlags = 118, + TAPTI_RUNTIME_TRACE_CBID_tangGraphDestroy = 119, + TAPTI_RUNTIME_TRACE_CBID_tangGraphExecDestroy = 120, + TAPTI_RUNTIME_TRACE_CBID_tangGraphGetInfo = 121, + TAPTI_RUNTIME_TRACE_CBID_tangGraphCreate = 122, + TAPTI_RUNTIME_TRACE_CBID_tangGraphAddHostNode = 123, + TAPTI_RUNTIME_TRACE_CBID_tangGraphAddKernelNode = 124, + TAPTI_RUNTIME_TRACE_CBID_tangProfilerStart = 125, + TAPTI_RUNTIME_TRACE_CBID_tangProfilerStop = 126, + TAPTI_RUNTIME_TRACE_CBID_tangIpcGetMemHandle = 127, + TAPTI_RUNTIME_TRACE_CBID_tangIpcOpenMemHandle = 128, + TAPTI_RUNTIME_TRACE_CBID_tangIpcCloseMemHandle = 129, + TAPTI_RUNTIME_TRACE_CBID_tangIpcGetEventHandle = 130, + TAPTI_RUNTIME_TRACE_CBID_tangIpcOpenEventHandle = 131, + TAPTI_RUNTIME_TRACE_CBID_tangGetFuncBySymbol = 132, + TAPTI_RUNTIME_TRACE_CBID_tangGetExportTable = 133, + TAPTI_RUNTIME_TRACE_CBID_tangLaunchKernel = 134, + TAPTI_RUNTIME_TRACE_CBID_tangLaunchKernel_ptsz = 135, + TAPTI_RUNTIME_TRACE_CBID_tangThreadExchangeStreamCaptureMode = 136, + TAPTI_RUNTIME_TRACE_CBID_FORCE_INT = 0x7fffffff, +} tapti_runtime_api_trace_cbid; + diff --git a/third_party/sunrise/backend/include/tapti/tapti_version.h b/third_party/sunrise/backend/include/tapti/tapti_version.h new file mode 100755 index 0000000000..f548f07ecf --- /dev/null +++ b/third_party/sunrise/backend/include/tapti/tapti_version.h @@ -0,0 +1,39 @@ +#ifndef _TAPTI_VERSION_ +#define _TAPTI_VERSION_ + +#include +#include "tapti_result.h" + +#define TAPTI_API_VERSION 1 + +#ifdef __cplusplus +extern "C" { +#endif //! __cplusplus + +#if defined(_MSC_VER) +#define TAPTI_DEPRECATED __declspec(deprecated) +#define TAPTI_API_EXPORT __declspec(dllexport) +#define TAPTI_API_IMPORT __declspec(dllimport) +#elif defined(__GNUC__) || defined(__clang__) +#define TAPTI_DEPRECATED __attribute__((deprecated)) +#define TAPTI_API_EXPORT __attribute__((visibility("default"))) +#define TAPTI_API_IMPORT __attribute__((visibility("default"))) +#else +#define TAPTI_DEPRECATED +#define TAPTI_API_EXPORT +#define TAPTI_API_IMPORT +#endif //! UNKNOWN COMPILER + +#if defined(tapti_shared_EXPORTS) +#define TAPTI_API TAPTI_API_EXPORT +#else +#define TAPTI_API TAPTI_API_IMPORT +#endif //! For user + +TAptiResult TAPTI_API taptiGetVersion(uint32_t *version); + +#ifdef __cplusplus +} +#endif //! __cplusplus + +#endif // _TAPTI_VERSION_ diff --git a/third_party/sunrise/backend/lib/.empty b/third_party/sunrise/backend/lib/.empty new file mode 100644 index 0000000000..e69de29bb2 diff --git a/third_party/sunrise/backend/spec/__init__.py b/third_party/sunrise/backend/spec/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/third_party/sunrise/backend/spec/include/flagtree_spec.h b/third_party/sunrise/backend/spec/include/flagtree_spec.h new file mode 100644 index 0000000000..2546be760b --- /dev/null +++ b/third_party/sunrise/backend/spec/include/flagtree_spec.h @@ -0,0 +1,12 @@ +#ifndef SUNRISE_FLAGTREE_SPEC_H_ +#define SUNRISE_FLAGTREE_SPEC_H_ + +#include "triton/Conversion/TritonToTritonGPU/sunrise_TritonToTritonGPUPass.h" +#include "triton/Dialect/Triton/Transforms/sunrise_RewriteTensorPointer.h" +#include "triton/Dialect/TritonGPU/IR/sunrise_Dialect.h" +#include "triton/Dialect/TritonGPU/IR/sunrise_LinearLayoutConversions.h" +#include "triton/Dialect/TritonGPU/Transforms/sunrise_Prefetch.h" +#include "triton/Dialect/TritonGPU/Transforms/sunrise_RemoveLayoutConversions.h" +#include "triton/Dialect/TritonGPU/Transforms/sunrise_Utility.h" + +#endif // SUNRISE_FLAGTREE_SPEC_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Conversion/TritonToTritonGPU/sunrise_TritonToTritonGPUPass.h b/third_party/sunrise/backend/spec/include/triton/Conversion/TritonToTritonGPU/sunrise_TritonToTritonGPUPass.h new file mode 100644 index 0000000000..ba3dfb3467 --- /dev/null +++ b/third_party/sunrise/backend/spec/include/triton/Conversion/TritonToTritonGPU/sunrise_TritonToTritonGPUPass.h @@ -0,0 +1,6 @@ +#ifndef SUNRISE_TRITON_CONVERSION_TRITONTOTRITONGPU_TRITONTOTRITONGPUPASS_H_ +#define SUNRISE_TRITON_CONVERSION_TRITONTOTRITONGPU_TRITONTOTRITONGPUPASS_H_ + +#define FLAGTREE_SPEC_triton_Conversion_TritonToTritonGPU_sunrise_TritonToTritonGPUPass + +#endif // SUNRISE_TRITON_CONVERSION_TRITONTOTRITONGPU_TRITONTOTRITONGPUPASS_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/Triton/Transforms/sunrise_RewriteTensorPointer.h b/third_party/sunrise/backend/spec/include/triton/Dialect/Triton/Transforms/sunrise_RewriteTensorPointer.h new file mode 100644 index 0000000000..3e960483ae --- /dev/null +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/Triton/Transforms/sunrise_RewriteTensorPointer.h @@ -0,0 +1,6 @@ +#ifndef SUNRISE_TRITON_DIALECT_TRITON_TRANSFORMS_REWRITETENSORPOINTER_H_ +#define SUNRISE_TRITON_DIALECT_TRITON_TRANSFORMS_REWRITETENSORPOINTER_H_ + +#define FLAGTREE_SPEC_triton_Dialect_Triton_Transforms_sunrise_RewriteTensorPointer + +#endif // SUNRISE_TRITON_DIALECT_TRITON_TRANSFORMS_REWRITETENSORPOINTER_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td new file mode 100644 index 0000000000..731be21840 --- /dev/null +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td @@ -0,0 +1,1533 @@ +#ifndef TRITONGPU_ATTRDEFS +#define TRITONGPU_ATTRDEFS + +include "triton/Dialect/TritonGPU/IR/TritonGPUAttrBase.td" + +//===----------------------------------------------------------------------===// +// Traits, Interfaces and shared Parameters +//===----------------------------------------------------------------------===// + +def LayoutEncodingTrait : AttrInterface<"LayoutEncodingTrait"> { + let cppNamespace = "::mlir::triton::gpu"; + let description = [{ + Common trait for all TTGIR layouts. + }]; + let methods = [ + InterfaceMethod<"Get the CTA layout backing this encoding.", + "CTAEncodingAttr", "getCTALayout">, + InterfaceMethod<"Get the rank of the layout.", "unsigned", "getRank", + (ins), [{}], [{ + return $_attr.getCTALayout().getRank(); + }]> + ]; +} +def DeclareLayoutEncodingMethods : DeclareAttrInterfaceMethods< + LayoutEncodingTrait, ["getCTALayout"]>; + +def SharedEncodingTrait : AttrInterface<"SharedEncodingTrait"> { + let cppNamespace = "::mlir::triton::gpu"; + + let description = [{ + Common trait describing shared memory. + }]; + let methods = [ + InterfaceMethod<"Return the default alignment for the layout.", + "int32_t", "getAlignment", (ins), [{}], [{ return 16; }]>, + ]; +} +def DeclareSharedEncodingMethods : DeclareAttrInterfaceMethods< + SharedEncodingTrait, ["getAlignment"]>; + +//===----------------------------------------------------------------------===// +// Shared Layout Encoding +//===----------------------------------------------------------------------===// + +def SwizzledSharedEncodingAttr + : TritonGPU_Attr<"SwizzledSharedEncoding", "swizzled_shared_encoding", + [SharedEncodingTrait, LayoutEncodingTrait, + DeclareLayoutEncodingMethods]> { + let mnemonic = "swizzled_shared"; + + let description = [{ +An encoding for tensors whose elements may be simultaneously accessed by +different GPU threads in the programs, via shared memory. In other words, +for all indices i \in Z^d, \mathcal{L}(i) = {0, 1, ..., 32*num_warps - 1}. + +In order to avoid shared memory bank conflicts, elements may be swizzled. +Here are some examples. In all cases, the input tensor is [0, 1, ..., n-1]. + +1. Basic swizzling + + #ttg.swizzled_shared<{vec=1, perPhase=1, maxPhase=4, order=[1,0]}> + [ 0, 1, 2, 3], // xor with 0 + [ 5, 4, 7, 6], // xor with 1 + [10, 11, 8, 9], // xor with 2 + [15, 14, 13, 12] // xor with 3 + +Here elements of row r are xor'ed with r (or more properly, in[r][c] -> +out[r][c^r]). + +2. Multiple rows per phase + + #ttg.swizzled_shared<{vec=1, perPhase=2, maxPhase=4, order=[1,0]}> + [ 0, 1, 2, 3], // phase 0 (xor with 0) + [ 4, 5, 6, 7], + [ 9, 8, 11, 10], // phase 1 (xor with 1) + [13, 12, 15, 14] + +Elements of row r are xor'ed with r/2. In other words, perPhase=2 +means that pairs of 2 rows get the same swizzling. + +3. Max-phase applied + + #ttg.swizzled_shared<{vec=1, perPhase=1, maxPhase=2, order=[1,0]}> + [ 0, 1, 2, 3], // phase 0 (xor with 0) + [ 5, 4, 7, 6], // phase 1 (xor with 1) + [ 8, 9, 10, 11], // phase 0 + [13, 12, 15, 14], // phase 1 + [16, 17, 18, 19], // ... + [21, 20, 23, 22], + [24, 25, 26, 27], + [29, 28, 31, 30] + +Elements of row r are xor'ed with (r/2) % 2. In other words, maxPhase=m has the +effect of limiting the maximum value of the xor to m-1. + +4. Max-phase and per-phase + + #ttg.swizzled_shared<{vec=1, perPhase=2, maxPhase=2, order=[1,0]}> + [ 0, 1, 2, 3], // phase 0 (xor with 0) + [ 4, 5, 6, 7], // phase 0 + [ 9, 8, 11, 10], // phase 1 (xor with 1) + [13, 12, 15, 14], // phase 1 + [16, 17, 18, 19], // phase 0 + [20, 21, 22, 23], // phase 0 + [25, 24, 27, 26], // phase 1 + [29, 28, 31, 30]] // phase 1 + +Here the xor value (the "phase", I guess?) changes every perPhase rows, up to a +maximum value of maxPhase-1. In other words, elements of row r are xor'ed with +(r/2) % 2. + +5. Adding vec + + #ttg.swizzled_shared<{vec=2, perPhase=1, maxPhase=4, order=[1,0]}> + [ 0, 1, 2, 3, 4, 5, 6, 7], + [10, 11, 8, 9, 14, 15, 12, 13], + [20, 21, 22, 23, 16, 17, 18, 19], + [30, 31, 28, 29, 26, 27, 24, 25] + +When vec=2, elements are swizzled in pairs of 2. In other words, the element at +(r,c) has value + + ((c / 2) ^ r) * 2 + (c % 2). + }]; + + // swizzle info: vec, perPhase, maxPhase + // order: the fastest-changing axis first + let parameters = ( + ins + "unsigned":$vec, + "unsigned":$perPhase, + "unsigned":$maxPhase, + ArrayRefParameter<"unsigned">:$order, + "CTAEncodingAttr":$CTALayout + ); + + let builders = [ + AttrBuilder<(ins "DotOperandEncodingAttr":$dotOpEnc, + "ArrayRef":$shape, + "ArrayRef":$order, + "CTAEncodingAttr":$CTALayout, + "unsigned":$typeWidthInBit), [{ + bool needTrans = false; // default value + return get(context, dotOpEnc, shape, order, CTALayout, typeWidthInBit, needTrans); + }]>, + + // TODO(jlebar): This should not be an overload of + // SwizzledSharedEncodingAttr::get(). It's misleading, because it does a bunch of + // nontrivial work based on the given dotOpEnc. + AttrBuilder<(ins "DotOperandEncodingAttr":$dotOpEnc, + "ArrayRef":$shape, + "ArrayRef":$order, + "CTAEncodingAttr":$CTALayout, + "unsigned":$typeWidthInBit, + "bool":$needTrans), [{ + + // ---- begin MFMA ---- + if (auto mfmaEnc = mlir::dyn_cast(dotOpEnc.getParent())) { + return mfmaEnc.composeSharedLayoutForOperand( + CTALayout, dotOpEnc.getOpIdx(), shape, order, dotOpEnc.getKWidth(), + typeWidthInBit, needTrans); + } + + // ---- begin WMMA ---- + if (auto wmmaEnc = mlir::dyn_cast(dotOpEnc.getParent())) { + return wmmaEnc.composeSharedLayoutForOperand( + CTALayout, dotOpEnc.getOpIdx(), shape, order, dotOpEnc.getKWidth(), + typeWidthInBit, needTrans); + } + + // ---- begin SunriseMMA ---- + auto sunriseMmaEnc = mlir::dyn_cast(dotOpEnc.getParent()); + if(sunriseMmaEnc) { + if (dotOpEnc.getOpIdx() == 0) { + constexpr int numBanks = 32; + constexpr int bankBitWidth = 32; + constexpr int totalBits = numBanks * bankBitWidth; + + // number of inner dimension rows per one pattern repeat + int innerDimLength = shape[order[0]]; + int elemsPerOneBanksRow = totalBits / typeWidthInBit; + + int perPhase = std::max(1, elemsPerOneBanksRow / innerDimLength); + int vecSize = 32 / typeWidthInBit; + int maxPhase = 16 / perPhase; + + return get(context, vecSize, perPhase, maxPhase, order, CTALayout); + } else { + // Do not swizzle in case k dimension is not innermost. + // In this case accesses will go in different banks even without swizzling. + return get(context, 2, 1, 1, order, CTALayout); + } + } + + auto blockEnc = mlir::dyn_cast(dotOpEnc.getParent()); + if (blockEnc) { + // 针对fp16类形使用FMA的情况，可生成fp16x2类形的fma指令 + auto sizePerThread = blockEnc.getSizePerThread(); + unsigned vec = sizePerThread[sizePerThread.size() - 1]; + return get(context, vec, 1, 1, order, CTALayout); + } + + auto mmaEnc = mlir::dyn_cast(dotOpEnc.getParent()); + + if(!mmaEnc) + return get(context, 1, 1, 1, order, CTALayout); + + // ---- begin Ampere & Hopper ---- + if (mmaEnc.isAmpere() || mmaEnc.isHopper()) { + return get(context, dotOpEnc.getOpIdx(), dotOpEnc.getKWidth(), shape, order, CTALayout, typeWidthInBit, needTrans); + } + + // ---- not implemented ---- + llvm_unreachable("unsupported swizzling for provided MMA version"); + }]>, + + // NVIDIA constructor! + // TODO(lezcano): We should totally get rid of all these constructors... + AttrBuilder<(ins "int":$opIdx, + "unsigned":$kWidth, + "ArrayRef":$shape, + "ArrayRef":$order, + "CTAEncodingAttr":$CTALayout, + "unsigned":$bitwidth, + "bool":$needTrans), [{ + int K = getShapePerCTA(CTALayout.getCTASplitNum(), shape)[order[0]]; + // Elems necessary to cover all the banks divided by the inner dimension + // This packs a few rows together for small K + int perPhase = std::max(1024 / (bitwidth * K), 1); + + int mmaStride = 8; + int vec = 4 * kWidth; + // needsTrans is equiv. to flipping the opIdx + if (needTrans) + std::swap(vec, mmaStride); + assert(opIdx == 0 || opIdx == 1); + int rank = order.size(); + int kDim = opIdx == 0 ? rank-1 : rank-2; + if (order[0] != kDim) + std::swap(vec, mmaStride); + // Count how many vec elements are needed to cover all the banks + int maxPhase = std::max(std::min(mmaStride, 1024 / (vec * bitwidth)), 1); + // Account for the row packing from perPhase: mmaStride / perPhase + maxPhase = std::max(maxPhase / perPhase, 1); + return get(context, vec, perPhase, maxPhase, order, CTALayout); + }]>, + + AttrBuilder<(ins "DotOperandEncodingAttr":$dotOpEnc, + "ArrayRef":$shape, + "ArrayRef":$order, + "CTAEncodingAttr":$CTALayout, + "Type":$eltTy), [{ + unsigned bitwidth = eltTy.getIntOrFloatBitWidth(); + return get(context, dotOpEnc, shape, order, CTALayout, bitwidth); + }]>, + + AttrBuilder<(ins "DotOperandEncodingAttr":$dotOpEnc, + "ArrayRef":$shape, + "ArrayRef":$order, + "CTAEncodingAttr":$CTALayout, + "Type":$eltTy, + "bool":$needTrans), [{ + unsigned bitwidth = eltTy.getIntOrFloatBitWidth(); + return get(context, dotOpEnc, shape, order, CTALayout, bitwidth, needTrans); + }]>, + ]; + + let hasCustomAssemblyFormat = 1; + let genVerifyDecl = 1; +} + +def PaddedSharedEncodingAttr + : TritonGPU_Attr<"PaddedSharedEncoding", "padded_shared_encoding", + [SharedEncodingTrait, DeclareLayoutEncodingMethods]> { + let mnemonic = "padded_shared"; + + let description = [{ +An encoding for tensors whose elements may be simultaneously accessed by +different GPU threads in the programs, via shared memory. In other words, +for all indices i \in Z^d, \mathcal{L}(i) = {0, 1, ..., 32*num_warps - 1}. +Compared to SwizzledSharedEncodingAttr, this encoding combines padding with +element reordering via linear transformation (e.g. row permutation) to avoid +shared memory bank conflicts. + +Formally, given a layout: + padded_shared<[:+, :+, ...]> +We insert a padding of `` elements after every `` elements. +Multi interval-padding pairs are supported for flexibility of multi tiered +padding schemes; they compose in an additive manner. So for a 1-D tensor element +at index i, the corresponding shared memory location index is + i + \sum_{k} (i / interval_k) * pad_k = 1 +`` and `` all need to be power of two. + +Some concrete examples ignoring the linear component, using `eM` to mean tensor +elements and `pN` to mean padding: + +1. Single interval-padding pair: + + #ttg.padded_shared<[2:+2], {...}> + [e0, e1, p0, p1, + e2, e3, p2, p3, + ...] + +2. Double interval-padding pairs: + + #ttg.padded_shared<[2:+1, 4:+2], {...}> + [e0, e1, p0, + e2, e3, p1, p2, p3, + e4, e5, p4, + e6, e7, p5, p6, p7, + ...] + +Furthermore this encoding allows for a linear remapping from the 1-D shared +memory offset to logical n-D tensor elements. The remapping is given in the form +of linear bases mapping from offset to [dim0, dim1...dimN-1]. +See LinearLayout.h for more details how linear layouts are applied to remap +elements. +Some concrete examples using `xN` and `yN` to mean the logical n-D tensor elements +and `pN` to mean padding: + +1. 1D Single interval-padding with strided elements + + #ttg.padded_shared<[2:+2] {offset = [[2], [1]], block = []}> + [x0, x2, p0 p1, + x1, x3, p2, p3 + ...] + +2. 2D single interval-padding with rearranged rows. + + #ttg.padded_shared<[16:+1] {offset = [[0, 1], [0, 2], /*gap, stride by 2 rows*/[2, 0], [4, 0], [1, 0]]], block = []}> + [ + x0y0, x0y1, x0y2, x0y3, + x2y0, x2y1, x2y2, x2y3, + x4y0, x4y1, x4y2, x4y3, + x6y0, x6y1, x6y2, x6y3, + p0, + x1y0, x1y1, x1y2, x1y3, + x3y0, x3y1, x3y2, x3y3, + x5y0, x5y1, x5y2, x5y3, + x7y0, x7y1, x7y2, x7y3, + p1, + ] + +For identity mappings a short form based on order and shape is used to increase readability. The following two encodings are the same: + + #ttg.padded_shared<[2:+2] {order = [1, 0], shape = [16, 32]}> + #ttg.padded_shared<[2:+2] {offset = [[0, 1], [0, 2], [0, 4], [0, 8], [0, 16], [1, 0], [2, 0], [4, 0], [8, 0]], block = []}> + + + }]; + + let parameters = (ins + ArrayRefParameter<"unsigned">:$intervals, + ArrayRefParameter<"unsigned">:$paddings, + LinearLayoutParam:$linearComponent + ); + + let builders = [ + AttrBuilder<(ins "ArrayRef>":$intervalPads, + "LinearLayout":$linearComponent)>, + + // Builder to create an identity mapping as the linear component + AttrBuilder<(ins "ArrayRef>":$intervalPads, + "ArrayRef":$order, "ArrayRef":$shape, + "CTAEncodingAttr":$ctaLayout)>, + ]; + + let extraClassDeclaration = extraBaseClassDeclaration # [{ + // Returns the order of the dimensions `dimName` of the layout. + // If more than dimension is of size one, it uses defaultOrder to determine + // the order of the dimensions of size one. + SmallVector orderPerDim(StringAttr dimName, + ArrayRef defaultOrder) const; + SmallVector getOrder() const; + + // Returns the bases of the dimensions `dimName` of the linear_component. + // If skipBroadcast is false, we count a base zero + SmallVector basesPerDim(StringAttr dimName, + bool skipBroadcast = true) const; + + unsigned getMinInterval() const { + return *llvm::min_element(getIntervals()); + } + + // Returns the total number of elements including padding given the input + // tensor shape. + int64_t getPaddedSize(ArrayRef shape) const; + }]; + let hasCustomAssemblyFormat = 1; + let genVerifyDecl = 1; +} + +def SharedLinearEncodingAttr + : TritonGPU_Attr<"SharedLinearEncoding", "shared_linear_encoding", + [SharedEncodingTrait, LayoutEncodingTrait, + DeclareLayoutEncodingMethods]> { + let mnemonic = "shared_linear"; + + let description = [{ + Linear shared encodings mirror LinearEncodingAttr but operate on shared + memory layouts. The LinearLayout parameter captures how shared memory + offsets (and optionally blocks) map to logical tensor indices. + }]; + + let parameters = (ins LinearLayoutParam:$linearLayout, "unsigned":$layoutAlignment); + + let extraClassDeclaration = [{ + SmallVector basesPerDim(StringAttr dimName, + bool skipBroadcast = true) const; + SmallVector orderPerDim(StringAttr dimName, + ArrayRef defaultOrder) const; + + SmallVector getOrder() const; + + LinearLayout toLinearLayout(ArrayRef shape) const; + + int32_t getAlignment() const { return static_cast(getLayoutAlignment()); } + }]; + + let genVerifyDecl = 1; + let hasCustomAssemblyFormat = 1; +} + +def NVMMASharedEncodingAttr : TritonGPU_Attr<"NVMMASharedEncoding", "nvmma_shared_encoding", + [DeclareSharedEncodingMethods, LayoutEncodingTrait, + DeclareLayoutEncodingMethods]> { + let mnemonic = "nvmma_shared"; + + let description = [{ + Represent blocked shared memory matching MMAv3/MMAv5 shared memory input. + This is meant to represent 2d tiled blocked layout. + The full layout representation is described here: + https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-shared-memory-layout + When the memdesc has more than 2 dimensions the tiling is applied to 8 rows even if the first outer dimension is smaller than 8. + In this case `transposed` means that the contiguous dimension is the most outer dimension of the memdesc. + }]; + + + // fp4Padded: Indicates that this encoding represents a mixed-precision fp4 operand in MMAv5 scaled dot, which needs + // to be in the special padded layout as described in https://docs.nvidia.com/cuda/parallel-thread-execution/#packing-format-used-for-matrix-a-and-b-by-kind-mxf8f6f4-in-shared-memory + let parameters = ( + ins + "unsigned":$swizzlingByteWidth, + "bool":$transposed, + "unsigned":$elementBitWidth, + "bool":$fp4Padded, + "CTAEncodingAttr":$CTALayout + ); + + let builders = [ + AttrBuilder<(ins "ArrayRef":$shape, + "ArrayRef":$order, + "CTAEncodingAttr":$CTALayout, + "Type":$eltTy, + "bool": $fp4Padded), [{ + auto shapePerCTA = getShapePerCTA(CTALayout.getCTASplitNum(), shape); + int32_t swizzlingByteWidth = 0; + unsigned eleBitWidth = eltTy.getIntOrFloatBitWidth(); + int packingFactor = fp4Padded ? 2 : 1; + + // get proper shared memory swizzling mode from the contiguous dimension + // size of the origin blocked layout. + auto contigDimSizeInByte = shapePerCTA[order[0]] * packingFactor * eleBitWidth / 8; + if (contigDimSizeInByte >= 128 && contigDimSizeInByte % 128 == 0) { + swizzlingByteWidth = 128; + } else if (contigDimSizeInByte >= 64 && contigDimSizeInByte % 64 == 0) { + swizzlingByteWidth = 64; + } else if (contigDimSizeInByte >= 32 && contigDimSizeInByte % 32 == 0) { + swizzlingByteWidth = 32; + } else { + swizzlingByteWidth = 0; + } + int flattenOutterDim = 1; + for (int i = 1; i < shapePerCTA.size(); i++) { + flattenOutterDim *= shapePerCTA[order[i]]; + } + if (shapePerCTA.size() < 2 || flattenOutterDim < 8) { + swizzlingByteWidth = 0; + } + bool transposed = order[0] == 0; + return $_get(context, swizzlingByteWidth, transposed, eleBitWidth, fp4Padded, CTALayout); + }]> + ]; + + let extraClassDeclaration = extraBaseClassDeclaration # [{ + int getPerPhase() const; + int getMaxPhase() const; + int getVec() const; + }]; + let hasCustomAssemblyFormat = 1; +} + +def AMDRotatingSharedEncodingAttr : + TritonGPU_Attr<"AMDRotatingSharedEncoding", "amd_rotating_shared_encoding", + [SharedEncodingTrait, LayoutEncodingTrait, + DeclareLayoutEncodingMethods]> { + let mnemonic = "amd_rotating_shared"; + + let description = [{ +This shared encoding is similar to SwizzledSharedEncodingAttr, but instead of +repeating swizzling pattern every `maxPhase*perPhase` rows of the memory object, +called a block, this layout changes swizzling pattern `maxPhase` times, then +repeats the pattern. The name "rotating" comes from the fact that first tensor +element of each block is swizzled with different phase, which is equal to +current block number: 0, 1, 2.. maxPhase-1, 0, 1, 2 ... + +This layout is used to reduce bank conflicts in cases where shared memory writes +and reads are performed on layouts with different order. It's meant for hardware +without native shared memory tranpose support. + +Swizzling pattern affects only 2 fastest dimensions of a tensor. +In the following text these two dimensions are called row and column: +- row is a fastest dimension +- column is a second fastest dimension + +Elements in a row dimension are stored in memory contiguously. + +If a matrix of size [128x64] is stored in this shared layout with order [1, 0], +dim 1 (64) will be stored contiguously and called row, dim 0 (128) is will be +called column. If order of shared layout is [0, 1], dim 0 (128) is stored +contiguously becomes a row, dim 1 (64) becomes a column. + +Swizzling pattern is following: + +Let's consider an element with logical coordinates = (inRowId, inColId). +For simplicity, we do not vectorize memory in examples, +i.e. vec == 1 and layout swizzles inidividual elements. +For vec != 1 example, take a look at SwizzledSharedEncodingAttr documentation. + +Swizzled coordinates within memory object are (outRowId, outColId): + + outRowId = inRowId + phase = (inRowId / perPhase) % maxPhase + blockNo = (inRowId / (perPhase * maxPhase)) % maxPhase + combinedPhase = phase ^ blockNo + outColId = inColId ^ combinedPhase + +Actual offset in memory could be computed with following function: + +memmory_offset = (outColId + outRowId * num_of_element_in_row) * sizeof(element) + + +Swizzling examples (matrix is filled with numbers 0, 1, 2, .. columns*rows-1): + + #shared<{vec=1, perPhase=1, maxPhase=2, order=[1,0]}> + row elements + 0 [ 0, 1, 2, 3], // phase = 0 blockNo = 0 (xor with 0) + 1 [ 5, 4, 7, 6], // phase = 1 blockNo = 0 (xor with 1) + 2 [ 9, 8, 11, 10], // phase = 0 blockNo = 1 (xor with 1) + 3 [12, 13, 14, 15] // phase = 1 blockNo = 1 (xor with 0) + 4 [16, 17, 18, 19], // phase = 0 blockNo = 0 (xor with 0) + 5 [21, 20, 23, 22], // phase = 1 blockNo = 0 (xor with 1) + 6 [25, 24, 27, 26], // phase = 0 blockNo = 1 (xor with 1) + 7 [28, 29, 30, 31] // phase = 1 blockNo = 1 (xor with 0) + + #shared<{vec=1, perPhase=2, maxPhase=2, order=[1,0]}> + row elements + 0 [ 0, 1, 2, 3], // phase = 0 blockNo = 0 (xor with 0) + 1 [ 4, 5, 6, 7], // phase = 0 blockNo = 0 (xor with 0) + 2 [ 9, 8, 11, 10], // phase = 1 blockNo = 0 (xor with 1) + 3 [13, 12, 15, 14] // phase = 1 blockNo = 0 (xor with 1) + 4 [17, 16, 19, 18], // phase = 0 blockNo = 1 (xor with 1) + 5 [21, 20, 23, 22], // phase = 0 blockNo = 1 (xor with 1) + 6 [24, 25, 26, 27], // phase = 1 blockNo = 1 (xor with 0) + 7 [28, 29, 30, 31] // phase = 1 blockNo = 1 (xor with 0) + + #shared<{vec=1, perPhase=1, maxPhase=4, order=[1,0]}> + row elements + 0 [ 0, 1, 2, 3], // phase = 0 blockNo = 0 (xor with 0) + 1 [ 5, 4, 7, 6], // phase = 1 blockNo = 0 (xor with 1) + 2 [10, 11, 8, 9], // phase = 2 blockNo = 0 (xor with 2) + 3 [15, 14, 13, 12] // phase = 3 blockNo = 0 (xor with 3) + 4 [17, 16, 19, 18], // phase = 0 blockNo = 1 (xor with 1) + 5 [20, 21, 22, 23], // phase = 1 blockNo = 1 (xor with 0) + 6 [27, 26, 25, 24], // phase = 2 blockNo = 1 (xor with 3) + 7 [30, 31, 28, 29] // phase = 3 blockNo = 1 (xor with 2) + }]; + + let parameters = ( + ins + "unsigned":$vec, + "unsigned":$perPhase, + "unsigned":$maxPhase, + ArrayRefParameter<"unsigned">:$order, + "CTAEncodingAttr":$CTALayout + ); + + let hasCustomAssemblyFormat = 1; +} + + +//===----------------------------------------------------------------------===// +// Distributed Layout Encoding +//===----------------------------------------------------------------------===// + +def DistributedEncodingTrait : AttrInterface<"DistributedEncodingTrait"> { + let cppNamespace = "::mlir::triton::gpu"; + + let description = [{ +The Distributed encoding describes the layout L with the 4-level compute hierarchy on GPU. +It is abstracted from the top to the bottom as CTAs Per CGA->Warps Per CTA->Threads Per Warp->Values Per Thread. + +For CTAs Per CGA and Warps Per CTA level, the linear id is distributed contiguously with the shape and order. +For example, for a shape/order pair defines a distribution layout +shape = [4, 4] +order = [0, 1] // The fastest-changing axis first +-> +layout = [0 4 8 12] + [1 5 9 13] + [2 6 10 14] + [3 7 11 15] + +For the Threads Per Warp and Values Per Thread level, the linear id distribution is variant for each sub-class encoding. + +If the layout does not completely cover the tensor, we tile it until we cover the entire tensor. +We call each individual tile "rep". + }]; + + let methods = [ + InterfaceMethod<"Get the order of reps (tiles of this layout that tile the whole tensor). The fastest-changing axis first", + "SmallVector", + "getRepOrder">, + InterfaceMethod<"Return total element size per thread.", + "unsigned", + "getTotalElemsPerThread", + (ins "ArrayRef":$shape), + /*defaultImplementation=*/[{ + return toLinearEncoding($_self, shape).getTotalElemsPerThread(shape); + }]>, + InterfaceMethod<"Return element size per thread in each dimension.", + "SmallVector", + "getElemsPerThread", + (ins "ArrayRef":$shape), + /*defaultImplementation=*/[{ + return toLinearEncoding($_self, shape).getElemsPerThread(shape); + }]>, + InterfaceMethod<"Convert to LinearLayout.", + "LinearLayout", + "toLinearLayout", + (ins "ArrayRef":$shape)>, + ]; +} + +class DistributedEncoding traits = []> + : TritonGPU_Attr { + + let description = [{ +Distributed encodings have a layout function L that is entirely characterized +by a d-dimensional tensor T. Note that L doesn't need to have the same shape +(or even the same rank) as the tensor it is encoding. + +The layout function \mathcal{L} of this layout is then defined, for an +index `i` \in Z^d, as follows: + +\mathcal{L}(T)[i_d] = L[(i_d + k_d*T.shape[d]) % L.shape[d]] \forall k_d such as i_d + k_d*T.shape[d] < L.shape[d] + +Intuitively, when the tensor dim size T.shape[d] is larger than the layout +dim size L.shape[d], on that particular dim, we distribute values from the +tensor to threads mapped in the layout in a "wrapped around" manner, with +each thread owning multiple values. + +OTOH, when the tensor dim size T.shape[d] is smaller than the layout +dim size L.shape[d], on that particular dim, we distribute values from the +tensor to threads mapped in the layout in a "broadcasted" manner, with +each value owned by multiple threads. + +For example, for a tensor/layout pair +T = [x x x x x x x x] + [x x x x x x x x] +L = [0 1 2 3 ] + [4 5 6 7 ] + [8 9 10 11] + [12 13 14 15] + +Then the data of T would be distributed as follow between the 16 CUDA threads: +L(T) = [ {0,8} , {1,9} , {2,10}, {3,11}, {0,8} , {1, 9} , {2, 10}, {3, 11}, + {4,12}, {5,13}, {6,14}, {7,15}, {4,12}, {5, 13}, {6, 14}, {7, 15} ] + }]; + + code extraDistributedDeclaration = extraBaseClassDeclaration # [{ + // Implemented in subclasses + SmallVector getRepOrder() const; + + LinearLayout toLinearLayout(ArrayRef shape) const; + }]; +} + +//===----------------------------------------------------------------------===// +// Linear Layout Encoding +//===----------------------------------------------------------------------===// + +def LinearEncodingAttr : DistributedEncoding<"LinearEncoding", "linear_encoding"> { + let mnemonic = "linear"; + + let description = [{ + See the docs in LinearLayout.h for the definition of linear layouts. + }]; + + let parameters = (ins LinearLayoutParam:$linearLayout); + + let extraClassDeclaration = extraDistributedDeclaration # [{ + // Generic distributed encoding methods + unsigned getTotalElemsPerThread(ArrayRef shape) const; + SmallVector getElemsPerThread(ArrayRef shape) const; + + SmallVector getContig(const char *, SmallVector) const; + SmallVector getContigPerThread() const; + SmallVector getContigPerWarp() const; + SmallVector getOrder() const; + SmallVector getWarpOrder() const; + SmallVector getThreadOrder() const; + + + // Generalizes get{Warp,Thread,CTA}Order to linear layouts. + // Returns the order of the dimensions `dimName` of the layout. + // If more than dimension is of size one, it uses defaultOrder to determine + // the order of the dimensions of size one. + SmallVector orderPerDim(StringAttr dimName, + ArrayRef defaultOrder) const; + + // Generalizes getThreadsPerWarp, getWarpsPerCTA, getCTAsPerCGA to linear layouts. + // Returns the bases of the dimensions `dimName` of the layout. + // If skipBroadcast is false, we count a base zero + SmallVector basesPerDim(StringAttr dimName, + bool skipBroadcast = true) const; + SmallVector getThreadsPerWarp() const; + SmallVector getWarpsPerCTA() const; + + // [FIXME LL] Supports legacy behaviour. We should remove these functions + SmallVector getSizePerThread() const; + }]; + + let genVerifyDecl = 1; + // Example of assembly format: + // <{register = [[0, 1], [8, 0], [0, 8], [64, 0]], + // lane = [[0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], + // warp = [[16, 0], [32, 0]], + // block = []}> + let hasCustomAssemblyFormat = 1; +} + + +//===----------------------------------------------------------------------===// +// Blocked Layout Encoding +//===----------------------------------------------------------------------===// + +def BlockedEncodingAttr : DistributedEncoding<"BlockedEncoding", "blocked_encoding"> { + let mnemonic = "blocked"; + + let description = [{ +An encoding where each warp owns a contiguous portion of the target tensor. This is typically the kind of data layout +used to promote memory coalescing in LoadInst and StoreInst. +It is characterized by three tuples -- thread tile size, warp tile size, and block tile size -- which +specify the amount of elements owned by each CUDA thread, warp and CTA respectively. + +Example 1, a row-major coalesced layout may partition a 16x16 tensor over 2 warps (i.e. 64 threads) as follows: + +[ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] +[ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] +[ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] +[ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] +... +[ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] +[ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] + +for + +#ttg.blocked_layout<{ + sizePerThread = {2, 2} + threadsPerWarp = {8, 4} + blocked = {{0, 1}} +}> + +Example 2, a row-major coalesced layout may partition a 32x32 tensor over 2 warps (i.e. 64 threads) as follows: + +[ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] +[ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] +[ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] +[ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] +... ... +[ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] +[ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] +[ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] +[ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] +[ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] +[ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] +... ... +[ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] +[ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] +for + +#ttg.blocked_layout<{ + sizePerThread = {2, 2} + threadsPerWarp = {8, 4} + blocked = {{0, 1}} +}> + +Example 3, A row-major coalesced layout may partition a 32x32 tensor over 2 warps (i.e. 64 threads) and +4 CTAs (taking 2x2 for example) as follows: + +CTA [0,0] CTA [0,1] +[ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] [ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] +[ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] [ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] +[ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] [ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] +[ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] [ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] +... ... +[ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] [ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] +[ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] [ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] + +CTA [1,0] CTA [1,1] +[ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] [ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] +[ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] [ 0 0 1 1 2 2 3 3 ; 32 32 33 33 34 34 35 35 ] +[ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] [ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] +[ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] [ 4 4 5 5 6 6 7 7 ; 36 36 37 37 38 38 39 39 ] +... ... +[ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] [ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] +[ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] [ 28 28 29 29 30 30 31 31 ; 60 60 61 61 62 62 63 63 ] +for + +#ttg.blocked_layout<{ + sizePerThread = {2, 2} + threadsPerWarp = {8, 4} + blocked = {{0, 1}, {1, 0}} +}> +}]; + + let parameters = ( + ins + ArrayRefParameter<"unsigned">:$sizePerThread, + ArrayRefParameter<"unsigned">:$threadsPerWarp, + ArrayRefParameter<"unsigned">:$warpsPerCTA, + ArrayRefParameter<"unsigned">:$order, // the fastest-changing axis first + + // CTALayout is optional in the textual IR. If omitted, we infer it to be a + // single CTA (i.e. the trivial map onto dim0..dimn-1) + "CTAEncodingAttr":$CTALayout + ); + let genVerifyDecl = 1; + + let builders = [ + AttrBuilder<(ins "ArrayRef":$shape, + "ArrayRef":$sizePerThread, + "ArrayRef":$order, + "unsigned":$numWarps, + "unsigned":$numThreadsPerWarp, + "CTAEncodingAttr":$CTALayout), [{ + unsigned rank = sizePerThread.size(); + SmallVector threadsPerWarp(rank); + SmallVector warpsPerCTA(rank); + SmallVector shapePerCTA = getShapePerCTA(CTALayout.getCTASplitNum(), shape); + + unsigned remainingLanes = numThreadsPerWarp; + unsigned remainingThreads = numWarps * numThreadsPerWarp; + unsigned remainingWarps = numWarps; + unsigned prevLanes = 1; + unsigned prevWarps = 1; + + // starting from the contiguous dimension + for (unsigned d = 0; d < rank - 1; ++d) { + unsigned i = order[d]; + unsigned threadsPerCTA = std::clamp(remainingThreads, 1, std::max(1, shapePerCTA[i] / sizePerThread[i])); + threadsPerWarp[i] = std::clamp(threadsPerCTA, 1, remainingLanes); + warpsPerCTA[i] = std::clamp(threadsPerCTA / threadsPerWarp[i], 1, remainingWarps); + remainingWarps /= warpsPerCTA[i]; + remainingLanes /= threadsPerWarp[i]; + remainingThreads /= threadsPerCTA; + prevLanes *= threadsPerWarp[i]; + prevWarps *= warpsPerCTA[i]; + } + + // Expand the last dimension to fill the remaining lanes and warps + threadsPerWarp[order[rank - 1]] = numThreadsPerWarp / prevLanes; + warpsPerCTA[order[rank - 1]] = numWarps / prevWarps; + + return $_get(context, sizePerThread, threadsPerWarp, warpsPerCTA, order, CTALayout); + }]>, + + AttrBuilder<(ins "ArrayRef":$shape, + "ArrayRef":$sizePerThread, + "ArrayRef":$order, + "unsigned":$numWarps, + "unsigned":$numThreadsPerWarp, + "unsigned":$numCTAs), [{ + unsigned rank = sizePerThread.size(); + SmallVector CTAsPerCGA(rank); + SmallVector CTASplitNum(rank); + ArrayRef CTAOrder = order; + + unsigned remainingCTAs = numCTAs; + + // starting from the most strided dimension + for (int d = rank - 1; d >= 0; --d) { + unsigned i = order[d]; + CTAsPerCGA[i] = std::clamp(remainingCTAs, 1, std::max(1, shape[i] / sizePerThread[i])); + CTASplitNum[i] = CTAsPerCGA[i]; + remainingCTAs /= CTAsPerCGA[i]; + } + + CTAsPerCGA[rank - 1] *= remainingCTAs; // wrap at CTA level + + CTAEncodingAttr CTALayout = CTAEncodingAttr::fromSplitParams(context, CTAsPerCGA, CTASplitNum, CTAOrder); + return get(context, shape, sizePerThread, order, numWarps, numThreadsPerWarp, CTALayout); + }]> + ]; + + let extraClassDeclaration = extraDistributedDeclaration; + + let hasCustomAssemblyFormat = 1; +} + +//===----------------------------------------------------------------------===// +// MMA Layout Encoding +//===----------------------------------------------------------------------===// + +def MmaEncodingTrait : AttrInterface<"MmaEncodingTrait"> { + let cppNamespace = "::mlir::triton::gpu"; + let methods = [ + InterfaceMethod<"Get the order of reps (tiles of this layout that tile the whole tensor). The fastest-changing axis first", + "SmallVector", + "getRepOrderForOperand", + (ins "int":$opIdx)>, + ]; +} + +def AMDMfmaEncodingAttr : DistributedEncoding<"AMDMfmaEncoding", "amd_mfma_encoding", [MmaEncodingTrait]> { + let mnemonic = "amd_mfma"; + + let description = [{ +An encoding for tensors that have been produced by MFMA matrix core instructions, +available on AMD Instinct GPUs of CDNA architectures. + +It is characterized by the following parameters: +- `version`: The GPU architecture: + - 1: gfx908: CDNA1 + - 2: gfx90a: CDNA2 + - 3: gfx942: CDNA3 + - 4: gfx950: CDNA4 +- `warpsPerCTA`: The warp layout in the block. +- `instrShape`: The shape in the form of (M, N, K) of the matrix. +- `isTransposed`: Indicates the result tensor is transposed so that it can be converted to dotOperand layout +without going to shared memory. This is used in the case of chained dot (E.g. Flash-Attention kernel). +- `tilesPerWarp`: The tile layout within a warp. Defaults to unit tile layout, i.e., single tile on all dimensions. +- `elementBitWidth`: Bit width of the output element type. Supported values are 32 and 64. Defaults to 32. + +Example 1: +Suppose we have a tensor with a shape of [32, 64], warpsPerCTA set to [1, 2] and MDim=NDim=32. +The data will be distributed between threads as follows: + + warp 0 warp 1 +-----------------/\-------------- -----------------/\-------------- +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 0 1 2 3 ...... 30 31 ] [ 64 65 66 67 ...... 94 95 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] +[ 32 33 34 35 ...... 62 63 ] [ 96 97 98 99 ...... 126 127 ] + +Example 2: +Suppose we have a tensor with a shape of [16, 32], warpsPerCTA set to [1, 2] and MDim=NDim=16. +The data will be distributed between threads as follows: + + warp 0 warp 1 +-----------------/\------------- ------------------/\--------------- +[ 0 1 2 3 ...... 14 15 ] [ 64 65 66 67 ...... 78 79 ] +[ 0 1 2 3 ...... 14 15 ] [ 64 65 66 67 ...... 78 79 ] +[ 0 1 2 3 ...... 14 15 ] [ 64 65 66 67 ...... 78 79 ] +[ 0 1 2 3 ...... 14 15 ] [ 64 65 66 67 ...... 78 79 ] +[ 16 17 18 19 ...... 30 31 ] [ 80 81 82 83 ...... 94 95 ] +[ 16 17 18 19 ...... 30 31 ] [ 80 81 82 83 ...... 94 95 ] +[ 16 17 18 19 ...... 30 31 ] [ 80 81 82 83 ...... 94 95 ] +[ 16 17 18 19 ...... 30 31 ] [ 80 81 82 83 ...... 94 95 ] +[ 32 33 34 35 ...... 46 47 ] [ 96 97 98 99 ...... 110 111 ] +[ 32 33 34 35 ...... 46 47 ] [ 96 97 98 99 ...... 110 111 ] +[ 32 33 34 35 ...... 46 47 ] [ 96 97 98 99 ...... 110 111 ] +[ 32 33 34 35 ...... 46 47 ] [ 96 97 98 99 ...... 110 111 ] +[ 48 49 50 51 ...... 62 63 ] [ 112 113 114 115 ...... 126 127 ] +[ 48 49 50 51 ...... 62 63 ] [ 112 113 114 115 ...... 126 127 ] +[ 48 49 50 51 ...... 62 63 ] [ 112 113 114 115 ...... 126 127 ] +[ 48 49 50 51 ...... 62 63 ] [ 112 113 114 115 ...... 126 127 ] + +Example 3: +Suppose we have a tensor with a shape of [8, 8], warpsPerCTA set to [2, 2] and nonKDim set to 4. +The data will be distributed between threads as follows(note that each element is duplicated in 16 threads): +Suppose we have a tensor with a shape of [8, 8], warpsPerCTA set to [2, 2] and MDim=NDim=4. +The data will be distributed between threads as follows(note that each element is duplicated in 16 threads): + +M N -> warp 0 warp 2 +| --------------------------/\-------------------------- ------------------------------/\------------------------------ +V [ 0,4,8...60 1,5...61 2,6...62 3,7...63 ] [ 128,132...188 129,133...189 130,134...190 131,135...191 ] + [ 0,4,8...60 1,5...61 2,6...62 3,7...63 ] [ 128,132...188 129,133...189 130,134...190 131,135...191 ] + [ 0,4,8...60 1,5...61 2,6...62 3,7...63 ] [ 128,132...188 129,133...189 130,134...190 131,135...191 ] + [ 0,4,8...60 1,5...61 2,6...62 3,7...63 ] [ 128,132...188 129,133...189 130,134...190 131,135...191 ] + warp 1 warp 3 + --------------------------/\-------------------------- ------------------------------/\------------------------------ + [ 64,68...124 65,69...125 66,70...126 67,71...127 ] [ 192,196...252 193,197...253 194,198...254 195,199...255 ] + [ 64,68...124 65,69...125 66,70...126 67,71...127 ] [ 192,196...252 193,197...253 194,198...254 195,199...255 ] + [ 64,68...124 65,69...125 66,70...126 67,71...127 ] [ 192,196...252 193,197...253 194,198...254 195,199...255 ] + [ 64,68...124 65,69...125 66,70...126 67,71...127 ] [ 192,196...252 193,197...253 194,198...254 195,199...255 ] + +Example 4: +This example demonstrates semantics of tilesPerWarp parameter. The MFMA layout (with tilesPerWarp=[1,1]) +assumes that each warp within a CTA tile computes a single MFMA tile. When the tensor is larger than +a single CTA tile, these tiles are repeated across the tensor. In this setup, the output tiles computed +by each warp were strided by the number of warps per CTA tile in both row and column dimensions. + +For instance, with 16 MFMA tiles and warpsPerCTA = [2, 2], the distribution of warps across the MFMA +tiles looked like: + +w0 w1 w0 w1 +w2 w3 w2 w3 +w0 w1 w0 w1 +w2 w3 w2 w3 + +tilesPerWarp parameter allows each warp to compute contiguous MFMA tiles in the row and/or column dimensions. +Using the same example with tilesPerWarp = [2, 2], the layout becomes: + +w0 w0 w1 w1 +w0 w0 w1 w1 +w2 w2 w3 w3 +w2 w2 w3 w3 +}]; + + let parameters = ( + ins + "unsigned": $version, + ArrayRefParameter<"unsigned">:$warpsPerCTA, + ArrayRefParameter<"unsigned">:$instrShape, + "bool":$isTransposed, + "CTAEncodingAttr":$CTALayout, + ArrayRefParameter<"unsigned">:$tilesPerWarp, + "unsigned":$elementBitWidth + ); + + let builders = [ + AttrBuilder<(ins "unsigned":$version, + "ArrayRef":$warpsPerCTA, + "ArrayRef":$instrShape, + "bool":$isTransposed, + "CTAEncodingAttr":$CTALayout, + CArg<"ArrayRef", "{}">:$tpw, + CArg<"unsigned", "0">:$elementBitWidth), [{ + SmallVector tilesPerWarp(tpw); + if (tilesPerWarp.empty()) + tilesPerWarp = SmallVector(warpsPerCTA.size(), 1); + if (elementBitWidth == 0) + elementBitWidth = 32; + return $_get($_ctxt, version, warpsPerCTA, instrShape, isTransposed, CTALayout, tilesPerWarp, elementBitWidth); + }]> + ]; + + let extraClassDeclaration = extraDistributedDeclaration # [{ + SmallVector getInstrShapeForOperand(int kWidth, int opIdx) const; + SmallVector getRepForOperand(ArrayRef operandShape, int kWidth, int opIdx) const; + SmallVector getRepOrderForOperand(int opIdx) const; + + // Check if tilesPerWarp is 1 in every dimension. + bool hasUnitTilesPerWarp() const; + + // Returns a swizzled shared layout matching this MFMA layout for the + // dot operand at the given |operandIdx| with |operandShape|. + SwizzledSharedEncodingAttr composeSharedLayoutForOperand( + CTAEncodingAttr ctaLayout, int operandIdx, ArrayRef operandShape, + ArrayRef sharedOrder, unsigned vectorSize, + unsigned elemBitWidth, bool needTrans) const; + }]; + + let genVerifyDecl = 1; + let hasCustomAssemblyFormat = 1; + let skipDefaultBuilders = 1; +} + +def AMDWmmaEncodingAttr : DistributedEncoding<"AMDWmmaEncoding", "amd_wmma_encoding", [MmaEncodingTrait]> { + let mnemonic = "amd_wmma"; + + let description = [{ +An encoding for tensors that have been produced by WMMA matrix core instructions, +available on AMD Radeon GPUs of RDNA architectures. + +It is characterized by the following parameters: +- `version` indicates the GPU architecture: + - 1: RDNA3; e.g., gfx1100, gfx1101 + - 2: RDNA4; e.g., gfx1200, gfx1201 + - 3: gfx1250 +- `warpsPerCTA` indicates the warp layout in the block. +- `tilesPerWarp` The tile layout within a warp. Defaults to unit tile layout, i.e., single tile on all dimensions. +- `instrShape` indicates the shape in the form of (M, N, K) of the matrix + operation performed by a single WMMA instruction. Defaults to (16, 16, 16). +- `isTransposed` indicates the layout of the result tensor is transposed. + +Example 1: +Suppose we have a tensor with shape [32, 64], `warpsPerCTA` set to [2, 2]. +Matrix elements represent which lane owns the element. Currently only wave32 mode +is supported. + +// ----------------------------------- version = 1 ----------------------------------- // + +Row | warp 0 warp 1 + |/-------------------^-------------------\ /-------------------^-------------------\ +0 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] +1 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] +2 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] +3 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] + | ... ... ... ... +14 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] +15 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] + + | warp 2 warp 3 +16 |/-------------------^-------------------\ /-------------------^-------------------\ +17 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] +18 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] +19 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] +20 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] + | ... ... ... ... +30 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] [0 1 2 ... 14 15] +31 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] [16 17 18 ... 30 31] + +// ------------------------ version = 2/3, isTransposed = false ------------------------ // + +Row | warp 0 warp 1 + |/--------^---------\ /---------^--------\ +0 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] +1 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] +.. | ... ... +6 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] +7 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] +8 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] +9 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] +.. | ... ... +14 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] +15 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] + | + | warp 2 warp 3 + |/--------^---------\ /---------^--------\ +16 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] +17 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] +.. | ... ... +22 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] +23 |[0 1 2 ... 14 15] [0 1 2 ... 14 15] +24 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] +25 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] +.. | ... ... +30 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] +31 |[16 17 18 ... 30 31] [16 17 18 ... 30 31] + +// ------------------------ version = 2/3, isTransposed = true ------------------------ // + + | warp 0 warp 1 + |/----------------^----------------\ /-------^-------\ +Col>| 0 1 2 3 4 5 6 7 8 ... 15 16 17 18 ... 32 +Row | +0 |[0 0 0 0 0 0 0 0 16 ... 16] [0 0 0 ... 16] +1 |[1 1 1 1 1 1 1 1 17 ... 17] [1 1 1 ... 17] +.. | ... ... +14 |[14 14 14 14 14 14 14 14 30 ... 30] [14 14 14 ... 30] +15 |[15 15 15 15 15 15 15 15 31 ... 31] [15 15 15 ... 31] + | + | warp 2 warp 3 + |/----------------^----------------\ /-------^-------\ +16 |[0 0 0 0 0 0 0 0 16 ... 16] [0 0 0 ... 16] +17 |[1 1 1 1 1 1 1 1 17 ... 17] [1 1 1 ... 17] +.. | ... ... +30 |[14 14 14 14 14 14 14 14 30 ... 30] [14 14 14 ... 30] +31 |[15 15 15 15 15 15 15 15 31 ... 31] [15 15 15 ... 31] + +Example 2: +This example demonstrates the tilesPerWarp parameter, which shares the same sematics with +AMDMfmaEncodingAttr. + +By default, WMMA layout assumes that each warp within a CTA tile computes a single WMMA tile. +When the tensor is larger than a single CTA tile, these tiles are repeated across the tensor. +In this setup, the output tiles computed by each warp are strided by the number of warps per CTA +tile in both row and column dimensions. + +For instance, with 16 WMMA tiles and warpsPerCTA = [2, 2], the default(tilesPerWarp = [1, 1]) +distribution of warps across the WMMA tiles looked like: + +w0 w1 w0 w1 +w2 w3 w2 w3 +w0 w1 w0 w1 +w2 w3 w2 w3 + +* Each unit reprsents a WMMA tile. w* shows which warp occupies that WMMA tile. + +tilesPerWarp parameter allows each warp to compute contiguous WMMA tiles in the row and/or column dimensions. +Using the same example with tilesPerWarp = [2, 2], the layout becomes: + +w0 w0 w1 w1 +w0 w0 w1 w1 +w2 w2 w3 w3 +w2 w2 w3 w3 + }]; + + let parameters = ( + ins + "unsigned": $version, + "bool":$isTransposed, + ArrayRefParameter<"unsigned">:$warpsPerCTA, + ArrayRefParameter<"unsigned">:$tilesPerWarp, + "CTAEncodingAttr":$CTALayout, + ArrayRefParameter<"unsigned">:$instrShape + ); + + let genVerifyDecl = 1; + let hasCustomAssemblyFormat = 1; + + let builders = [ + AttrBuilder<(ins "unsigned":$version, + "bool":$isTransposed, + "ArrayRef":$warpsPerCTA, + "CTAEncodingAttr":$CTALayout, + "ArrayRef":$instrShape), [{ + SmallVector tilesPerWarp(warpsPerCTA.size(), 1); + return $_get(context, version, isTransposed, warpsPerCTA, tilesPerWarp, CTALayout, instrShape); + }]> + ]; + + let extraClassDeclaration = extraDistributedDeclaration # [{ + SmallVector getRepForOperand(ArrayRef operandShape, int kDim, int opIdx) const; + SmallVector getRepOrderForOperand(int opIdx) const; + + static SmallVector getDefaultInstrShape() { + return {16, 16, 16}; + } + + // Check if tilesPerWarp is 1 in every dimension. + bool hasUnitTilesPerWarp() const; + + // Returns a swizzled shared layout matching this WMMA layout for the + // dot operand at the given |operandIdx| with |operandShape|. + SwizzledSharedEncodingAttr composeSharedLayoutForOperand( + CTAEncodingAttr ctaLayout, int operandIdx, ArrayRef operandShape, + ArrayRef sharedOrder, unsigned kWidth, + unsigned elemBitWidth, bool needTrans) const; + }]; +} + +def SunriseMmaEncodingAttr : DistributedEncoding<"SunriseMmaEncoding", "sunrise_mma_encoding", [MmaEncodingTrait]> { + let mnemonic = "sunrise_mma"; + + let description = [{ + TODO ... + }]; + + let parameters = ( + ins + "unsigned":$versionMajor, + "unsigned":$versionMinor, + ArrayRefParameter<"unsigned">:$warpsPerCTA, + // "int":$numWarps, + "CTAEncodingAttr":$CTALayout, + "SunriseMmaEncodingAttr::TMMAOutLayout":$outLayout, + "unsigned":$inputElemBitWidth, + "unsigned":$outputElemBitWidth, + "bool":$isACol, + "bool":$isBCol + + ); + + let hasCustomAssemblyFormat = 1; + + let extraClassDeclaration = extraDistributedDeclaration # [{ + enum class TMMAOutLayout : unsigned { // 只用于mma中c, d的布局 + NotAvailable, + Row_2B, Col_2B, + ARow_4B_8x4, BRow_4B_4x8 + }; + + SmallVector getRepForOperand(ArrayRef operandShape, Type elemType, int opIdx) const; + SmallVector getRepOrderForOperand(int opIdx) const; + SmallVector getInstrShapeForOperand(unsigned opIdx) const; + SmallVector getInstrShapeForOperand(unsigned opIdx, unsigned elemBitWdith) const; + SmallVector getShapePerCTATileForOperand(unsigned opIdx) const; + SmallVector getShapePerCTATile() const; + }]; +} + +def NvidiaMmaEncodingAttr : DistributedEncoding<"NvidiaMmaEncoding", "nvidia_mma_encoding", [MmaEncodingTrait]> { + let mnemonic = "nvidia_mma"; + + let description = [{ +An encoding for tensors that have been produced by tensor cores. + +It is characterized by two parameters: +- A 'versionMajor' which specifies the generation the tensor cores + whose output is being partitioned: + - 1 for first-gen tensor cores (Volta), and + - 2 for second-gen tensor cores (Turing/Ampere). +- A 'versionMinor' which indicates the specific layout of a tensor core + generation, e.g. for Volta, there might be multiple kinds of layouts + annotated by 0,1,2 and so on. +- A `blockTileSize` to indicate how data should be partitioned between warps. + +// -------------------------------- version = 1 --------------------------- // + +For first-gen tensor cores, the implicit warpTileSize is [16, 16]. +Note: the layout is different from the recommended in PTX ISA +https://docs.nvidia.com/cuda/parallel-thread-execution/index.html +(mma.884 section, FP32 accumulator). + +For example, when versionMinor=1, the matrix L corresponding to +blockTileSize=[32,16] is: + + warp 0 +--------------------------------/\------------------------------- +[ 0 0 2 2 8 8 10 10 0 0 2 2 8 8 10 10 ] +[ 1 1 3 3 9 9 11 11 1 1 3 3 9 9 11 11 ] +[ 0 0 2 2 8 8 10 10 0 0 2 2 8 8 10 10 ] +[ 1 1 3 3 9 9 11 11 1 1 3 3 9 9 11 11 ] +[ 4 4 6 6 12 12 14 14 4 4 6 6 12 12 14 14 ] +[ 5 5 7 7 13 13 15 15 5 5 7 7 13 13 15 15 ] +[ 4 4 6 6 12 12 14 14 4 4 6 6 12 12 14 14 ] +[ 5 5 7 7 13 13 15 15 5 5 7 7 13 13 15 15 ] +[ 16 16 18 18 20 20 22 22 16 16 18 18 20 20 22 22 ] +[ 17 17 19 19 21 21 23 23 17 17 19 19 21 21 23 23 ] +[ 16 16 18 18 20 20 22 22 16 16 18 18 20 20 22 22 ] +[ 17 17 19 19 21 21 23 23 17 17 19 19 21 21 23 23 ] +[ 24 24 26 26 28 28 30 30 24 24 26 26 28 28 30 30 ] +[ 25 25 27 27 29 29 31 31 25 25 27 27 29 29 31 31 ] +[ 24 24 26 26 28 28 30 30 24 24 26 26 28 28 30 30 ] +[ 25 25 27 27 29 29 31 31 25 25 27 27 29 29 31 31 ] + + warp 1 = warp0 + 32 +--------------------------------/\------------------------------- +[ 32 32 34 34 40 40 42 42 32 32 34 34 40 40 42 42 ] +[ 33 33 35 35 41 41 43 43 33 33 35 35 41 41 43 43 ] +[ ............................................................... ] + + +// -------------------------------- version = 2 --------------------------- // + +For second-gen tensor cores, the implicit warpTileSize is [16, 8]. +Information about this layout can be found in the official PTX documentation +https://docs.nvidia.com/cuda/parallel-thread-execution/index.html +(mma.16816 section, FP32 accumulator). + +For example, the matrix L corresponding to blockTileSize=[32,16] is: + warp 0 warp 2 +-----------------/\------------- ----------------/\------------- +[ 0 0 1 1 2 2 3 3 32 32 33 33 34 34 35 35 +[ 4 4 5 5 6 6 7 7 36 36 37 37 38 38 39 39 +[ .............................. .............................. +[ 28 28 29 29 30 30 31 31 60 60 61 61 62 62 63 63 +[ 0 0 1 1 2 2 3 3 32 32 33 33 34 34 35 35 +[ 4 4 5 5 6 6 7 7 36 36 37 37 38 38 39 39 +[ .............................. .............................. +[ 28 28 29 29 30 30 31 31 60 60 61 61 62 62 63 63 + + warp 1 warp 3 +----------------/\------------- ----------------/\------------- +[ 64 64 65 65 66 66 67 67 96 96 97 97 98 98 99 99 +[ 68 68 69 69 70 70 71 71 100 100 101 101 102 102 103 103 +[ .............................. ............................... +[ 92 92 93 93 94 94 95 95 124 124 125 125 126 126 127 127 +[ 64 64 65 65 66 66 67 67 96 96 97 97 98 98 99 99 +[ 68 68 69 69 70 70 71 71 100 100 101 101 102 102 103 103 +[ .............................. ............................... +[ 92 92 93 93 94 94 95 95 124 124 125 125 126 126 127 127 + +}]; + + let parameters = ( + ins + "unsigned":$versionMajor, + "unsigned":$versionMinor, + ArrayRefParameter<"unsigned">:$warpsPerCTA, + "CTAEncodingAttr":$CTALayout, + ArrayRefParameter<"unsigned">:$instrShape + ); + + + let extraClassDeclaration = extraDistributedDeclaration # [{ + bool isVolta() const; + bool isTuring() const; + bool isAmpere() const; + bool isHopper() const; + + SmallVector getRepForOperand(ArrayRef shape, + int bitwidth, int kWidth, + int opIdx) const; + SmallVector getRepOrderForOperand(int opIdx) const; + }]; + + let hasCustomAssemblyFormat = 1; +} + +def SliceEncodingAttr : DistributedEncoding<"SliceEncoding", "slice_encoding"> { + let mnemonic = "slice"; + + let description = [{ + Given a `parent` layout and a `dim`, squeezes the given `dim` in the `parent` + layout and distributes values in a tensor T according to the new layout. + + For example, given + + T = [x x x x x x x x] + L_parent = [0 1 2 3 ] + [4 5 6 7 ] + [8 9 10 11] + [12 13 14 15] (with 16 CUDA threads) + + With dim = 0, squeezing out dim 0, we have + L = [{0,4,8,12}, {1,5,9,13}, {2,6,10,14}, {3,7,11,15} ] + + Then the data of T would be distributed as follow between the 16 CUDA threads: + L(T) = [ {0,4,8,12} , {1,5,9,13} , ... {3,7,11,15}, {0,4,8,12} , ..., {3,7,11,15} ] + + With dim = 1, squeezing out dim 1, we have + L = [ {0,1,2,3}, {4,5,6,7}, {8,9,10,11}, {12,13,14,15} ] + + Then the data of T would be distributed as follow between the 16 CUDA threads: + L = [ {0,1,2,3}, {4,5,6,7}, ..., {12,13,14,15}, {0,1,2,3}, ..., {12,13,14,15} ] + + This is useful for constructing the inverse layout of an expand_dims operation + during some optimization passes. + }]; + + let parameters = ( + ins + "unsigned":$dim, + "DistributedEncodingTrait":$parent + ); + + let extraClassDeclaration = extraDistributedDeclaration # [{ + template + SmallVector paddedShape(ArrayRef shape) const; + }]; + + let hasCustomAssemblyFormat = 1; + let genVerifyDecl = 1; +} + +def DotOperandEncodingAttr : DistributedEncoding<"DotOperandEncoding", "dot_operand_encoding"> { + let mnemonic = "dot_op"; + + let description = [{ +In the TritonGPU dialect, given `d = tt.dot a, b, c` tt.dot's operands a and b +must be of DotOperandEncodingAttr layout, if the dot is MMA v1 or v2 (i.e. +pre-Hopper). For MMA v3, the operands are *almost always* in a regular shared +encoding, but sometimes the LHS is also a dot-operand encoding. + +a's opIdx is 0, b's opIdx is 1. + +The parent field is the layout of d. + +kWidth defines number of consecutive elements stored by one thread along k dimension. +Some layouts do not use this parameter, either because they have a fixed number of +elements along the K dim, or they use all elements of the tensor along the K dim. + +# WGMMA Notes +We require kWidth to be provided for Hopper because the dtype at loading might be +different from the dtype at WGMMA, due to casting. The kWidth is determined by the +dtype at WGMMA. + +The encoded tensor consists of operand A for possibly multiple wgmma instructions. +For each wgmma, each warp in a warp group feeds a single "warp matrix" +Each warp matrix consists of 2x2 "quads". +Each thread holds several elements in each quad. Right before a wgmma, +the sum of bitwidth of +the elements in each quad should add up to 32. + +These values are stored unrolled in `elements`. +The ordering of dimensions is as follows by convention: +batch (only 1 batch for Hopper currently) +matM (m-index of the "warp matrix") +matK (k-index of the "warp matrix") +quadK (k-index of the "quad" in the core matrix) +quadM (m-index of the "quad" in the core matrix) +vecIdx (index of the element in the quad; this is always along the k-dim) + }]; + + let parameters = ( + ins + "unsigned":$opIdx, + "Attribute":$parent, + DefaultValuedParameter<"unsigned", "0">:$kWidth + ); + + let builders = [ + AttrBuilder<(ins "unsigned":$opIdx, + "Attribute":$parent, + "Type":$eltTy), [{ + NvidiaMmaEncodingAttr parentAttr = mlir::dyn_cast(parent); + if (!parentAttr || (!parentAttr.isAmpere() && !parentAttr.isHopper())) + return $_get(context, opIdx, parent, 0); + // For MMAV2 and V3 + unsigned bitwidth = eltTy.getIntOrFloatBitWidth(); + unsigned kWidth = std::max(32 / bitwidth, 1u); + return $_get(context, opIdx, parent, kWidth); + }]> + ]; + + let assemblyFormat = "`<` `{` struct(params) `}` `>`"; + let genVerifyDecl = 1; + let extraClassDeclaration = extraDistributedDeclaration; +} + +def TTG_SharedMemorySpace : AttrDef { + let mnemonic = "shared_memory"; + let description = [{ + Attribute to indicate that the memory descriptor points to shared memory. + }]; +} + +#endif diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_Dialect.h b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_Dialect.h new file mode 100644 index 0000000000..72408db721 --- /dev/null +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_Dialect.h @@ -0,0 +1,6 @@ +#ifndef SUNRISE_TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_ +#define SUNRISE_TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_ + +#define FLAGTREE_SPEC_triton_Dialect_TritonGPU_IR_sunrise_Dialect + +#endif // SUNRISE_TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_LinearLayoutConversions.h b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_LinearLayoutConversions.h new file mode 100644 index 0000000000..38b20f334a --- /dev/null +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_LinearLayoutConversions.h @@ -0,0 +1,6 @@ +#ifndef SUNRISE_TRITON_DIALECT_TRITONGPU_IR_LINEARLAYOUTCONVERSIONS_H_ +#define SUNRISE_TRITON_DIALECT_TRITONGPU_IR_LINEARLAYOUTCONVERSIONS_H_ + +#define FLAGTREE_SPEC_triton_Dialect_TritonGPU_IR_sunrise_LinearLayoutConversion + +#endif // SUNRISE_TRITON_DIALECT_TRITONGPU_IR_LINEARLAYOUTCONVERSIONS_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Prefetch.h b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Prefetch.h new file mode 100644 index 0000000000..f85de4efa1 --- /dev/null +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Prefetch.h @@ -0,0 +1,6 @@ +#ifndef SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_PREFETCH_H_ +#define SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_PREFETCH_H_ + +#define FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_Prefetch + +#endif // SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_PREFETCH_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_RemoveLayoutConversions.h b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_RemoveLayoutConversions.h new file mode 100644 index 0000000000..35a4d9f665 --- /dev/null +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_RemoveLayoutConversions.h @@ -0,0 +1,6 @@ +#ifndef SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_REMOVELAYOUTCONVERSIONS_H_ +#define SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_REMOVELAYOUTCONVERSIONS_H_ + +#define FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_RemoveLayoutConversion + +#endif // SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_REMOVELAYOUTCONVERSIONS_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Utility.h b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Utility.h new file mode 100644 index 0000000000..43b33be56b --- /dev/null +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Utility.h @@ -0,0 +1,6 @@ +#ifndef SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_UTILITY_H_ +#define SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_UTILITY_H_ + +#define FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_Utility + +#endif // SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_UTILITY_H_ diff --git a/third_party/sunrise/backend/spec/lib/CMakeLists.txt b/third_party/sunrise/backend/spec/lib/CMakeLists.txt new file mode 100644 index 0000000000..629c08af67 --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/CMakeLists.txt @@ -0,0 +1,2 @@ +add_subdirectory(Conversion) +add_subdirectory(Dialect) diff --git a/third_party/sunrise/backend/spec/lib/Conversion/CMakeLists.txt b/third_party/sunrise/backend/spec/lib/Conversion/CMakeLists.txt new file mode 100644 index 0000000000..5cbcea5da2 --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Conversion/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(TritonToTritonGPU) diff --git a/third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/CMakeLists.txt b/third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/CMakeLists.txt new file mode 100644 index 0000000000..2ec9d7e00c --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/CMakeLists.txt @@ -0,0 +1,14 @@ +add_triton_library(FlagTree_sunrise_TritonToTritonGPU + TritonToTritonGPUPass.cpp + + DEPENDS + TritonConversionPassIncGen + + LINK_LIBS PUBLIC + MLIRIR + MLIRPass + MLIRTransforms + TritonIR + ProtonIR + TritonGPUIR +) diff --git a/third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp new file mode 100644 index 0000000000..5c208d8d1a --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp @@ -0,0 +1,1043 @@ +#include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/UB/IR/UBOps.h" +#include "mlir/IR/Value.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Transforms/DialectConversion.h" +#ifdef __TLE__ +#include "tle/dialect/include/IR/Dialect.h" +#endif +#include "triton/Conversion/TritonToTritonGPU/Passes.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h" +#include "triton/Dialect/TritonGPU/Transforms/Utility.h" +#include "triton/Tools/LayoutUtils.h" + +namespace mlir::triton { +#define GEN_PASS_DEF_CONVERTTRITONTOTRITONGPU +#include "triton/Conversion/TritonToTritonGPU/Passes.h.inc" +} // namespace mlir::triton + +namespace { + +using namespace mlir; +using namespace mlir::triton; +using namespace mlir::triton::gpu; + +// [Sunrise Optim] 即使前端没手工写 num_warps=1/2, +// 也会把这类 small rank-1 reduce (size <= 128) 自动从坏的 4/8-warp 配置拉回到 2-warp。 +static bool shouldClampSunriseNumWarpsForSmallReduce(ModuleOp mod, + StringRef target, + mlir::Pass::Option &numWarps) { + int32_t clampedNumWarps = numWarps.getValue(); + if (!target.starts_with("tang:") || clampedNumWarps <= 2) + return false; + + bool hasDot = false; + bool hasSmallRank1Reduce = false; + mod.walk([&](Operation *op) { + if (isa(op)) { + hasDot = true; + return WalkResult::interrupt(); + } + auto reduce = dyn_cast(op); + if (!reduce) + return WalkResult::advance(); + + auto operandTy = dyn_cast(reduce.getOperands()[0].getType()); + if (!operandTy || operandTy.getRank() != 1 || reduce.getAxis() != 0) + return WalkResult::advance(); + + int64_t reduceWidth = operandTy.getShape()[0]; + if (reduceWidth > 0 && reduceWidth <= 128) { + hasSmallRank1Reduce = true; + // 1 warp for reduceWidth <= 64, 2 warps for reduceWidth > 64 and <= 128 + clampedNumWarps = reduceWidth <= 64 ? 1 : 2; + } + return WalkResult::advance(); + }); + if (hasSmallRank1Reduce && !hasDot) + numWarps = clampedNumWarps; + return hasSmallRank1Reduce && !hasDot; +} + +// pass named attrs (e.g., tt.contiguity) from Triton to Triton +static void addNamedAttrs(Operation *op, DictionaryAttr dictAttrs) { + for (const NamedAttribute attr : dictAttrs.getValue()) + if (!op->hasAttr(attr.getName())) + op->setAttr(attr.getName(), attr.getValue()); +} + +template struct GenericOpPattern : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(Op op, typename Op::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + SmallVector retTypes; +#ifdef __TLE__ + if (failed( + this->getTypeConverter()->convertTypes(op->getResults(), retTypes))) +#else + if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), + retTypes))) +#endif + return failure(); + rewriter.replaceOpWithNewOp(op, retTypes, adaptor.getOperands(), + op->getAttrs()); + + return success(); + } +}; + +class ArithConstantPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(arith::ConstantOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Type retType = +#ifdef __TLE__ + getTypeConverter()->convertType(op.getResult()); +#else + getTypeConverter()->convertType(op.getType()); +#endif + auto retShapedType = cast(retType); + auto value = dyn_cast(adaptor.getValue()); + if (isa(retShapedType)) { + assert(value && "expected a dense elements attribute"); + // This is a hack. We just want to add encoding. + value = value.reshape(retShapedType); + } + addNamedAttrs(rewriter.replaceOpWithNewOp( + op, retShapedType, value), + adaptor.getAttributes()); + return success(); + } +}; + +void populateArithPatternsAndLegality(TritonGPUTypeConverter &typeConverter, + RewritePatternSet &patterns, + TritonGPUConversionTarget &target) { + // -------------- + // Add legality and rewrite pattern rules for operations + // from the Arith dialect. The basic premise is that + // Arith operations require both inputs to have the same + // non-null encoding + // -------------- + MLIRContext *context = patterns.getContext(); + // TODO: there's probably a better way to avoid adding all ops one-by-one + patterns.add< + ArithConstantPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, // NegFOp + // Floating point + GenericOpPattern, GenericOpPattern, + // MaxMin + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + // Floating point + GenericOpPattern, GenericOpPattern, + GenericOpPattern, + // Cmp + GenericOpPattern, GenericOpPattern, + // Select + GenericOpPattern, + // Cast Ops + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern>(typeConverter, context); +} + +void populateMathPatternsAndLegality(TritonGPUTypeConverter &typeConverter, + RewritePatternSet &patterns, + TritonGPUConversionTarget &target) { + MLIRContext *context = patterns.getContext(); + // Rewrite rule + patterns.add, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern, + GenericOpPattern, GenericOpPattern>( + typeConverter, context); +} + +// +// Triton patterns +// +struct TritonExpandDimsPattern + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::ExpandDimsOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + // Type retType = op.getType()); + RankedTensorType argType = + cast(adaptor.getSrc().getType()); + Attribute _argEncoding = argType.getEncoding(); + if (!_argEncoding) + return failure(); + auto argEncoding = cast(_argEncoding); + // return shape + auto retShape = argType.getShape().vec(); + retShape.insert(retShape.begin() + op.getAxis(), 1); + auto newRank = retShape.size(); + // return encoding + auto retSizePerThread = llvm::to_vector(argEncoding.getSizePerThread()); + retSizePerThread.insert(retSizePerThread.begin() + op.getAxis(), 1); + auto retThreadsPerWarp = to_vector(argEncoding.getThreadsPerWarp()); + retThreadsPerWarp.insert(retThreadsPerWarp.begin() + op.getAxis(), 1); + auto retWarpsPerCTA = to_vector(argEncoding.getWarpsPerCTA()); + retWarpsPerCTA.insert(retWarpsPerCTA.begin() + op.getAxis(), 1); + SmallVector retOrder(retShape.size()); + std::iota(retOrder.begin(), retOrder.end(), 0); + + auto ctaLl = argEncoding.getCTALayout().getLinearLayout(); + auto kBlock = *ctaLl.getInDimNames().begin(); + auto *ctx = kBlock.getContext(); + auto newDim = standardOutDimNames(ctx, newRank)[newRank - 1]; + ctaLl *= LinearLayout::identity1D(1, kBlock, newDim); + // Move last dim to op.getAxis(). nb is this a std::rotate? + auto newOrder = to_vector(llvm::seq(newRank)); + for (int i = newRank - 1; i >= op.getAxis() + 1; --i) { + std::swap(newOrder[i], newOrder[i - 1]); + } + ctaLl = transposeLinearLayout(ctaLl, newOrder); + auto retCTALayout = CTAEncodingAttr::get(ctx, std::move(ctaLl)); + triton::gpu::BlockedEncodingAttr retEncoding = + triton::gpu::BlockedEncodingAttr::get(getContext(), retSizePerThread, + retThreadsPerWarp, retWarpsPerCTA, + retOrder, retCTALayout); + // convert operand to slice of return type + Attribute newArgEncoding = triton::gpu::SliceEncodingAttr::get( + getContext(), op.getAxis(), retEncoding); + RankedTensorType newArgType = argType.cloneWithEncoding(newArgEncoding); + // construct new op + auto newSrc = triton::gpu::ConvertLayoutOp::create( + rewriter, op.getLoc(), newArgType, adaptor.getSrc()); + addNamedAttrs(rewriter.replaceOpWithNewOp( + op, newSrc, adaptor.getAxis()), + adaptor.getAttributes()); + return success(); + } + +private: + template + SmallVector insertOne(ArrayRef vec, unsigned axis) const { + SmallVector res(vec.begin(), vec.end()); + res.insert(res.begin() + axis, 1); + return res; + } + + // Example: order = [ 0, 2, 1, 3], dim = 2 + // resOrder = [2, 0, 3, 1, 4] + SmallVector insertOrder(ArrayRef order, + unsigned axis) const { + SmallVector resOrder(order.begin(), order.end()); + for (unsigned i = 0; i < resOrder.size(); ++i) + if (resOrder[i] >= axis) + ++resOrder[i]; + resOrder.insert(resOrder.begin(), axis); + return resOrder; + } +}; + +struct TritonDotPattern : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::DotOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + RankedTensorType origType = op.getType(); + auto origShape = origType.getShape(); + auto typeConverter = getTypeConverter(); + int numWarps = typeConverter->getNumWarps(); + int threadsPerWarp = typeConverter->getThreadsPerWarp(); + int numCTAs = typeConverter->getNumCTAs(); + auto rank = origShape.size(); + SmallVector retSizePerThread(rank, 1); + auto numElements = product(origShape); + if (numElements / (numWarps * threadsPerWarp) >= 4) { + retSizePerThread[rank - 1] = 2; + retSizePerThread[rank - 2] = 2; + } + if (numElements / (numWarps * threadsPerWarp) >= 16) { + retSizePerThread[rank - 1] = 4; + retSizePerThread[rank - 2] = 4; + } + retSizePerThread[rank - 1] = std::min( + retSizePerThread[rank - 1], static_cast(origShape[rank - 1])); + retSizePerThread[rank - 2] = std::min( + retSizePerThread[rank - 2], static_cast(origShape[rank - 2])); + + SmallVector retOrder(rank); + for (unsigned i = 0; i < rank; ++i) + retOrder[i] = rank - 1 - i; + Attribute dEncoding = triton::gpu::BlockedEncodingAttr::get( + getContext(), origShape, retSizePerThread, retOrder, numWarps, + threadsPerWarp, numCTAs); + RankedTensorType retType = origType.cloneWithEncoding(dEncoding); + // a & b must be of smem layout + auto aType = cast(adaptor.getA().getType()); + auto bType = cast(adaptor.getB().getType()); + Type aEltType = aType.getElementType(); + Type bEltType = bType.getElementType(); + Attribute aEncoding = aType.getEncoding(); + Attribute bEncoding = bType.getEncoding(); + if (!aEncoding || !bEncoding) + return failure(); + Value a = adaptor.getA(); + Value b = adaptor.getB(); + Value c = adaptor.getC(); + if (!mlir::isa(aEncoding)) { + Attribute encoding = triton::gpu::DotOperandEncodingAttr::get( + getContext(), 0, dEncoding, aEltType); + auto dstType = aType.cloneWithEncoding(encoding); + a = triton::gpu::ConvertLayoutOp::create(rewriter, a.getLoc(), dstType, + a); + } + if (!mlir::isa(bEncoding)) { + Attribute encoding = triton::gpu::DotOperandEncodingAttr::get( + getContext(), 1, dEncoding, bEltType); + auto dstType = bType.cloneWithEncoding(encoding); + b = triton::gpu::ConvertLayoutOp::create(rewriter, b.getLoc(), dstType, + b); + } + c = triton::gpu::ConvertLayoutOp::create(rewriter, c.getLoc(), retType, c); + + addNamedAttrs(rewriter.replaceOpWithNewOp( + op, retType, a, b, c, adaptor.getInputPrecision(), + adaptor.getMaxNumImpreciseAcc()), + adaptor.getAttributes()); + return success(); + } +}; + +struct TritonCatPattern : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::CatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + // The cat op satisfy two conditions: + // 1. output.numel = lhs.numel + rhs.numel + // 2. output.total_elems_per_thread = + // next_power_of_2(lhs.total_elems_per_thread + rhs.total_elems_per_thread) + // For now, this behaves like generic, but this + // will evolve when we add support for `can_reorder=False`. + auto retType = cast( +#ifdef __TLE__ + this->getTypeConverter()->convertType(op.getResult())); +#else + this->getTypeConverter()->convertType(op.getType())); +#endif + auto retEncoding = + cast(retType.getEncoding()); + auto lhsType = adaptor.getLhs().getType(); + auto rhsType = adaptor.getRhs().getType(); + auto lhsTotalElemsPerThread = triton::gpu::getTotalElemsPerThread(lhsType); + auto rhsTotalElemsPerThread = triton::gpu::getTotalElemsPerThread(rhsType); + auto retTotalElemsPerThread = triton::gpu::getTotalElemsPerThread(retType); + auto retShape = retType.getShape(); + auto retOrder = retEncoding.getOrder(); + auto retThreadsPerWarp = retEncoding.getThreadsPerWarp(); + auto retWarpsPerCTA = retEncoding.getWarpsPerCTA(); + // Get new retSizePerThread if ret elems per thread is not enough. + // We have to round it up to the next power of 2 due to triton's tensor size + // constraint. + auto newRetTotalElemsPerThread = + nextPowOf2(lhsTotalElemsPerThread + rhsTotalElemsPerThread); + auto newRetSizePerThread = llvm::to_vector(retEncoding.getSizePerThread()); + newRetSizePerThread[retOrder[0]] *= + newRetTotalElemsPerThread / retTotalElemsPerThread; + triton::gpu::BlockedEncodingAttr newRetEncoding = + triton::gpu::BlockedEncodingAttr::get( + getContext(), newRetSizePerThread, retThreadsPerWarp, + retWarpsPerCTA, retOrder, retEncoding.getCTALayout()); + auto newRetType = retType.cloneWithEncoding(newRetEncoding); + addNamedAttrs(rewriter.replaceOpWithNewOp( + op, newRetType, adaptor.getOperands()), + adaptor.getAttributes()); + return success(); + } +}; + +struct TritonJoinOpPattern : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(JoinOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + // Simply rely on type inference for this op. (Notably, GenericOpPattern + // does not do this, instead it assigns the default layout to the ins and + // outs.) + addNamedAttrs(rewriter.replaceOpWithNewOp( + op, adaptor.getLhs(), adaptor.getRhs()), + adaptor.getAttributes()); + return success(); + } +}; + +struct TritonSplitOpPattern : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult matchAndRewrite(SplitOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const { + auto src = adaptor.getSrc(); + auto srcTy = cast(src.getType()); + auto srcEnc = dyn_cast(srcTy.getEncoding()); + int rank = srcEnc.getOrder().size(); + auto typeConverter = getTypeConverter(); + + // The operand to split must have: + // - a blocked layout, with + // - sizePerThread = 2 in the last dimension, + // - threadsPerWarp, warpsPerCTA, and CTAsPerCGA = 1 in the last dim, and + // - the last dimension minor. + // If that's not the case, add a convert before the split. + if (!srcEnc || srcEnc.getSizePerThread().back() != 2 || + srcEnc.getOrder().front() != rank - 1) { + // If we take the default encoding for the op's result (i.e. post-split) + // and add 1 to the end of each dim, that gives us what we want. Other + // than making a legal src encoding, our choice of layout doesn't matter; + // it'll get fixed by RemoveLayoutConversions. + auto defaultEnc = getDefaultBlockedEncoding( + getContext(), + cast(op.getResult(0).getType()).getShape(), + typeConverter->getNumWarps(), typeConverter->getThreadsPerWarp(), + typeConverter->getNumCTAs()); + + auto append = [&](ArrayRef vals, unsigned val) { + SmallVector res(vals); + res.push_back(val); + return res; + }; + auto prepend = [&](ArrayRef vals, unsigned val) { + SmallVector res; + res.push_back(val); + res.append(vals.begin(), vals.end()); + return res; + }; + + auto layout = defaultEnc.getCTALayout().getLinearLayout(); + auto kBlock = StringAttr::get(getContext(), "block"); + auto newDim = standardOutDimNames(getContext(), rank)[rank - 1]; + layout *= LinearLayout::identity1D(1, kBlock, newDim); + srcEnc = BlockedEncodingAttr::get( + getContext(), append(defaultEnc.getSizePerThread(), 2), + append(defaultEnc.getThreadsPerWarp(), 1), + append(defaultEnc.getWarpsPerCTA(), 1), + prepend(defaultEnc.getOrder(), rank - 1), + CTAEncodingAttr::get(getContext(), layout)); + srcTy = srcTy.cloneWithEncoding(srcEnc); + src = ConvertLayoutOp::create(rewriter, op.getLoc(), srcTy, src); + } + + addNamedAttrs(rewriter.replaceOpWithNewOp(op, src), + adaptor.getAttributes()); + return success(); + } +}; + +struct TritonTransPattern : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(TransOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + Value src = adaptor.getSrc(); + auto srcTy = cast(src.getType()); + auto srcEnc = srcTy.getEncoding(); + if (!srcEnc) + return failure(); + addNamedAttrs(rewriter.replaceOpWithNewOp(op, src, op.getOrder()), + adaptor.getAttributes()); + return success(); + } +}; + +struct TritonBroadcastPattern + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + // This creates a tensor with the new shape but the argument's layout + LogicalResult + matchAndRewrite(BroadcastOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto srcType = cast(adaptor.getSrc().getType()); + auto srcEncoding = srcType.getEncoding(); + if (!srcEncoding) + return failure(); + Type retType = op.getType().cloneWithEncoding(srcEncoding); + // Type retType = this->getTypeConverter()->convertType(op.getType()); + addNamedAttrs(rewriter.replaceOpWithNewOp( + op, retType, adaptor.getOperands()), + adaptor.getAttributes()); + return success(); + } +}; + +struct TritonReducePattern : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::ReduceOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto newReduce = triton::ReduceOp::create( + rewriter, op.getLoc(), adaptor.getOperands(), adaptor.getAxis()); + addNamedAttrs(newReduce, adaptor.getAttributes()); + + auto &newCombineOp = newReduce.getCombineOp(); + rewriter.cloneRegionBefore(op.getCombineOp(), newCombineOp, + newCombineOp.end()); + rewriter.replaceOp(op, newReduce.getResult()); + return success(); + } +}; + +struct TritonScanPattern : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::ScanOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto newScan = + triton::ScanOp::create(rewriter, op.getLoc(), adaptor.getOperands(), + adaptor.getAxis(), op.getReverse()); + addNamedAttrs(newScan, adaptor.getAttributes()); + + auto &newCombineOp = newScan.getCombineOp(); + rewriter.cloneRegionBefore(op.getCombineOp(), newCombineOp, + newCombineOp.end()); + rewriter.replaceOp(op, newScan.getResult()); + return success(); + } +}; + +struct TritonMapElementwisePattern + : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::MapElementwiseOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto converter = getTypeConverter(); + SmallVector resultTys; +#ifdef __TLE__ + auto err = converter->convertTypes(op.getResults(), resultTys); +#else + auto err = converter->convertTypes(op.getResults().getType(), resultTys); +#endif + if (failed(err)) { + return err; + } + + auto newMapOp = triton::MapElementwiseOp::create( + rewriter, op.getLoc(), resultTys, adaptor.getOperands(), op.getPack()); + addNamedAttrs(newMapOp, adaptor.getAttributes()); + + auto &newScalarOp = newMapOp.getScalarOp(); + rewriter.cloneRegionBefore(op.getScalarOp(), newScalarOp, + newScalarOp.end()); + rewriter.replaceOp(op, newMapOp.getResult()); + return success(); + } +}; + +class TritonFuncOpPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::FuncOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto converter = getTypeConverter(); + TypeConverter::SignatureConversion result(op.getNumArguments()); + auto newOp = rewriter.replaceOpWithNewOp( + op, op.getName(), op.getFunctionType()); + addNamedAttrs(newOp, adaptor.getAttributes()); + rewriter.inlineRegionBefore(op.getBody(), newOp.getBody(), + newOp.getBody().end()); + // Convert just the entry block. The remaining unstructured control flow is + // converted by br patterns. + if (!newOp.getBody().empty()) + rewriter.applySignatureConversion(&newOp.getBody().front(), result, + converter); + return success(); + } +}; + +class TritonCallOpPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::CallOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto newOp = rewriter.replaceOpWithNewOp( + op, op.getCallee(), op.getResultTypes(), adaptor.getOperands()); + addNamedAttrs(newOp, adaptor.getAttributes()); + return success(); + } +}; + +class TritonReturnOpPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(ReturnOp op, ReturnOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + rewriter.replaceOpWithNewOp(op, adaptor.getOperands()); + return success(); + } +}; + +void populateTritonPatterns(TritonGPUTypeConverter &typeConverter, + RewritePatternSet &patterns, unsigned numCTAs) { + MLIRContext *context = patterns.getContext(); + patterns.insert< // TODO: view should have custom pattern that views the + // layout + // clang-format off + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + TritonBroadcastPattern, + TritonCatPattern, + TritonJoinOpPattern, + TritonSplitOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + TritonReducePattern, + GenericOpPattern, + TritonScanPattern, + GenericOpPattern, + GenericOpPattern, +#ifdef __TLE__ + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, +#endif + TritonExpandDimsPattern, + TritonTransPattern, + TritonDotPattern, + TritonMapElementwisePattern, + GatherScatterOpPattern, + GatherScatterOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + // this assumes the right layout will be set later for dot scaled. + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + TritonFuncOpPattern + // clang-format on + >(typeConverter, context); +} +// +// SCF patterns +// +// This is borrowed from ConvertForOpTypes in +// SCF/Transforms/StructuralTypeConversions.cpp +struct SCFForPattern : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + // Ref: ConvertForOpTypes + LogicalResult + matchAndRewrite(scf::ForOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto newOp = + cast(rewriter.cloneWithoutRegions(*op.getOperation())); + rewriter.inlineRegionBefore(op.getRegion(), newOp.getRegion(), + newOp.getRegion().end()); + + // Now, update all the types. + + // Convert the types of block arguments within the given region. This + // replaces each block with a new block containing the updated signature. + // The entry block may have a special conversion if `entryConversion` is + // provided. On success, the new entry block to the region is returned for + // convenience. Otherwise, failure is returned. + if (failed(rewriter.convertRegionTypes(&newOp.getRegion(), + *getTypeConverter()))) { + return rewriter.notifyMatchFailure(op, "could not convert body types"); + } + // Change the clone to use the updated operands. We could have cloned with + // a IRMapping, but this seems a bit more direct. + newOp->setOperands(adaptor.getOperands()); + // Update the result types to the new converted types. + SmallVector newResultTypes; +#ifdef __TLE__ + for (Value result : op.getResults()) { + Type newType = typeConverter->convertType(result); +#else + for (Type type : op.getResultTypes()) { + Type newType = typeConverter->convertType(type); +#endif + if (!newType) + return rewriter.notifyMatchFailure(op, "not a 1:1 type conversion"); + newResultTypes.push_back(newType); + } + for (auto t : llvm::zip(newOp.getResults(), newResultTypes)) + std::get<0>(t).setType(std::get<1>(t)); + + rewriter.replaceOp(op, newOp.getResults()); + + return success(); + } +}; + +// This is borrowed from ConvertFIfOpTypes in +// SCF/Transforms/StructuralTypeConversions.cpp +class SCFIfPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(scf::IfOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + // TODO: Generalize this to any type conversion, not just 1:1. + // + // We need to implement something more sophisticated here that tracks which + // types convert to which other types and does the appropriate + // materialization logic. + // For example, it's possible that one result type converts to 0 types and + // another to 2 types, so newResultTypes would at least be the right size to + // not crash in the llvm::zip call below, but then we would set the the + // wrong type on the SSA values! These edge cases are also why we cannot + // safely use the TypeConverter::convertTypes helper here. + SmallVector newResultTypes; +#ifdef __TLE__ + for (Value result : op.getResults()) { + Type newType = typeConverter->convertType(result); +#else + for (auto type : op.getResultTypes()) { + Type newType = typeConverter->convertType(type); +#endif + if (!newType) + return rewriter.notifyMatchFailure(op, "not a 1:1 type conversion"); + newResultTypes.push_back(newType); + } + + // See comments in the ForOp pattern for why we clone without regions and + // then inline. + scf::IfOp newOp = + cast(rewriter.cloneWithoutRegions(*op.getOperation())); + rewriter.inlineRegionBefore(op.getThenRegion(), newOp.getThenRegion(), + newOp.getThenRegion().end()); + rewriter.inlineRegionBefore(op.getElseRegion(), newOp.getElseRegion(), + newOp.getElseRegion().end()); + + // Update the operands and types. + newOp->setOperands(adaptor.getOperands()); + for (auto t : llvm::zip(newOp.getResults(), newResultTypes)) + std::get<0>(t).setType(std::get<1>(t)); + rewriter.replaceOp(op, newOp.getResults()); + return success(); + } +}; + +// This is borrowed from ConvertFIfOpTypes in +// SCF/Transforms/StructuralTypeConversions.cpp +class SCFWhilePattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(scf::WhileOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto *converter = getTypeConverter(); + assert(converter); + SmallVector newResultTypes; +#ifdef __TLE__ + if (failed(converter->convertTypes(op.getResults(), newResultTypes))) +#else + if (failed(converter->convertTypes(op.getResultTypes(), newResultTypes))) +#endif + return failure(); + + auto newOp = scf::WhileOp::create(rewriter, op.getLoc(), newResultTypes, + adaptor.getOperands()); + for (auto i : {0u, 1u}) { + auto &dstRegion = newOp.getRegion(i); + rewriter.inlineRegionBefore(op.getRegion(i), dstRegion, dstRegion.end()); + if (failed(rewriter.convertRegionTypes(&dstRegion, *converter))) + return rewriter.notifyMatchFailure(op, "could not convert body types"); + } + rewriter.replaceOp(op, newOp.getResults()); + return success(); + } +}; + +class SCFConditionPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + LogicalResult + matchAndRewrite(scf::ConditionOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + rewriter.modifyOpInPlace(op, + [&]() { op->setOperands(adaptor.getOperands()); }); + return success(); + } +}; + +void populateSCFPatterns(TritonGPUTypeConverter &typeConverter, + RewritePatternSet &patterns) { + MLIRContext *context = patterns.getContext(); + patterns.add, SCFForPattern, SCFIfPattern, + SCFWhilePattern, SCFConditionPattern>(typeConverter, context); +} + +// CF + +class CFBranchPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(cf::BranchOp op, cf::BranchOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto converter = getTypeConverter(); + auto newOp = rewriter.replaceOpWithNewOp( + op, op.getSuccessor(), adaptor.getOperands()); + if (failed(rewriter.convertRegionTypes(newOp.getSuccessor()->getParent(), + *converter))) + return failure(); + return success(); + } +}; + +class CFCondBranchPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(cf::CondBranchOp op, cf::CondBranchOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto converter = getTypeConverter(); + auto newOp = rewriter.replaceOpWithNewOp( + op, adaptor.getCondition(), op.getTrueDest(), + adaptor.getTrueDestOperands(), op.getFalseDest(), + adaptor.getFalseDestOperands()); + addNamedAttrs(newOp, adaptor.getAttributes()); + + if (failed(rewriter.convertRegionTypes(newOp.getTrueDest()->getParent(), + *converter))) + return failure(); + if (failed(rewriter.convertRegionTypes(newOp.getFalseDest()->getParent(), + *converter))) + return failure(); + return success(); + } +}; + +void populateCFPatterns(TritonGPUTypeConverter &typeConverter, + RewritePatternSet &patterns) { + MLIRContext *context = patterns.getContext(); + patterns.add(typeConverter, context); +} + +#ifdef __TLE__ +class TleDSLRegionOpPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(tle::DSLRegionOp op, tle::DSLRegionOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto newOp = rewriter.cloneWithoutRegions(op); + Region &body = op.getBody(), &newBody = newOp.getBody(); + rewriter.inlineRegionBefore(body, newBody, newBody.end()); + + if (failed(rewriter.convertRegionTypes(&newBody, *getTypeConverter()))) { + return rewriter.notifyMatchFailure(op, "could not convert body types"); + } + newOp->setOperands(adaptor.getOperands()); + for (OpResult result : newOp->getResults()) { +#ifdef __TLE__ + result.setType(getTypeConverter()->convertType(result)); +#else + result.setType(getTypeConverter()->convertType(result.getType())); +#endif + } + + rewriter.replaceOp(op, newOp->getResults()); + return success(); + } +}; + +class TleExtractTileOpPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(tle::ExtractTileOp op, tle::ExtractTileOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + + auto srcType = dyn_cast(adaptor.getSrc().getType()); + if (!srcType) { + return op.emitError("source must be a ranked tensor"); + } + auto srcEnc = srcType.getEncoding(); + if (!srcEnc) { + return op.emitError("source tensor must have encoding attribute"); + } + + Type retType = op.getType().cloneWithEncoding(srcEnc); + + auto newOp = rewriter.replaceOpWithNewOp( + op, retType, adaptor.getSrc(), adaptor.getIndex()); + + if (auto tileShapeAttr = op->getAttr("tile_shape")) + newOp->setAttr("tile_shape", tileShapeAttr); + + addNamedAttrs(newOp, adaptor.getAttributes()); + + return success(); + } +}; + +// insert_tile op pattern +class TleInsertTileOpPattern : public OpConversionPattern { +public: + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(tle::InsertTileOp op, tle::InsertTileOp::Adaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + + auto srcType = dyn_cast(adaptor.getSrc().getType()); + if (!srcType) { + return op.emitError("source must be a ranked tensor"); + } + + auto srcEnc = srcType.getEncoding(); + if (!srcEnc) { + return op.emitError("source tensor must have encoding attribute"); + } + + auto tileType = dyn_cast(adaptor.getTile().getType()); + if (!tileType) { + return op.emitError("tile must be a ranked tensor"); + } + + auto tileEnc = tileType.getEncoding(); + if (!tileEnc) { + return op.emitError("tile tensor must have encoding attribute"); + } + + Type retType = op.getType().cloneWithEncoding(srcEnc); + + auto newOp = rewriter.replaceOpWithNewOp( + op, retType, adaptor.getSrc(), adaptor.getTile(), adaptor.getIndex()); + + addNamedAttrs(newOp, adaptor.getAttributes()); + + return success(); + } +}; + +// flagtree tle raw +void populateTleRawPatterns(TritonGPUTypeConverter &typeConverter, + RewritePatternSet &patterns) { + MLIRContext *context = patterns.getContext(); + patterns + .add, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, + GenericOpPattern, GenericOpPattern>( + typeConverter, context); +} +#endif + +class ConvertTritonToTritonGPU + : public triton::impl::ConvertTritonToTritonGPUBase< + ConvertTritonToTritonGPU> { +public: + using ConvertTritonToTritonGPUBase::ConvertTritonToTritonGPUBase; + + void runOnOperation() override { + if (target.getValue().empty()) { + mlir::emitError( + getOperation().getLoc(), + "'convert-triton-to-tritongpu' requires 'target' option to be set"); + return signalPassFailure(); + } + + MLIRContext *context = &getContext(); + ModuleOp mod = getOperation(); + shouldClampSunriseNumWarpsForSmallReduce(mod, this->target.getValue(), numWarps); + // type converter + TritonGPUTypeConverter typeConverter(context, numWarps, threadsPerWarp, + numCTAs, enableSourceRemat); + TritonGPUConversionTarget target(*context, typeConverter); + // rewrite patterns + RewritePatternSet patterns(context); + // add rules + populateArithPatternsAndLegality(typeConverter, patterns, target); + populateMathPatternsAndLegality(typeConverter, patterns, target); + populateTritonPatterns(typeConverter, patterns, numCTAs); + // TODO: can we use + // mlir::scf::populateSCFStructurealTypeConversionsAndLegality(...) here? + populateSCFPatterns(typeConverter, patterns); + populateCFPatterns(typeConverter, patterns); +#ifdef __TLE__ + populateTleRawPatterns(typeConverter, patterns); +#endif + patterns.insert>(typeConverter, context); + + Builder b(&getContext()); + mod->setAttr(AttrNumWarpsName, b.getI32IntegerAttr(numWarps)); + mod->setAttr("ttg.total-num-warps", b.getI32IntegerAttr(numWarps)); + mod->setAttr(AttrNumThreadsPerWarp, b.getI32IntegerAttr(threadsPerWarp)); + mod->setAttr(AttrNumCTAsName, b.getI32IntegerAttr(numCTAs)); + mod->setAttr(AttrTargetName, b.getStringAttr(this->target.getValue())); + + if (failed(applyPartialConversion(mod, target, std::move(patterns)))) + return signalPassFailure(); + } +}; + +} // namespace diff --git a/third_party/sunrise/backend/spec/lib/Dialect/CMakeLists.txt b/third_party/sunrise/backend/spec/lib/Dialect/CMakeLists.txt new file mode 100644 index 0000000000..27cb65ce51 --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/CMakeLists.txt @@ -0,0 +1,2 @@ +add_subdirectory(Triton) +add_subdirectory(TritonGPU) diff --git a/third_party/sunrise/backend/spec/lib/Dialect/Triton/CMakeLists.txt b/third_party/sunrise/backend/spec/lib/Dialect/Triton/CMakeLists.txt new file mode 100644 index 0000000000..e31af32661 --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/Triton/CMakeLists.txt @@ -0,0 +1 @@ +add_subdirectory(Transforms) diff --git a/third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/CMakeLists.txt b/third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/CMakeLists.txt new file mode 100644 index 0000000000..937ec444db --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/CMakeLists.txt @@ -0,0 +1,13 @@ +add_triton_library(FlagTree_sunrise_TritonTransforms + RewriteTensorPointer.cpp + + DEPENDS + TritonTransformsIncGen + + LINK_LIBS PUBLIC + MLIRPass + MLIRTransformUtils + MLIRTransforms + MLIRSCFToControlFlow + TritonIR +) diff --git a/third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp b/third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp new file mode 100644 index 0000000000..29bc998b65 --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp @@ -0,0 +1,619 @@ +#include + +#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Support/LLVM.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/Triton/Transforms/Passes.h" + +namespace mlir::triton { + +#define GEN_PASS_DEF_TRITONREWRITETENSORPOINTER +#include "triton/Dialect/Triton/Transforms/Passes.h.inc" + +namespace { + +/// An additional struct to record the meta information of operations +/// with tensor pointers +struct RewritedInfo { +private: + Value base; + SmallVector shape; + SmallVector strides; + SmallVector offsets; + ArrayRef tensorShape; + bool hasRowStride; + unsigned rowStrideParamIndex; + + // A cache to avoid generating the same offset with range + DenseMap cachedOffsetWithRange; + +public: + RewritedInfo() = default; + + RewritedInfo(const RewritedInfo &other) = default; + + RewritedInfo &operator=(const RewritedInfo &other) = default; + + RewritedInfo(Value base, const SmallVector &shape, + const SmallVector &strides, + const SmallVector &offsets, + const ArrayRef &tensorShape, + bool hasRowStride, unsigned rowStrideParamIndex) + : base(base), shape(shape), strides(strides), offsets(offsets), + tensorShape(tensorShape), hasRowStride(hasRowStride), rowStrideParamIndex(rowStrideParamIndex) { + assert(shape.size() == strides.size() && shape.size() == offsets.size() && + shape.size() == tensorShape.size()); + } + + unsigned int length() const { return shape.size(); } + + Value getOffset(unsigned i) { return offsets[i]; } + + SmallVector getOffsets() { return offsets; } + + void setOffset(unsigned i, Value newOffset) { + offsets[i] = newOffset; + cachedOffsetWithRange.clear(); + } + + void setOffsets(const SmallVector &newOffsets) { + offsets = newOffsets; + cachedOffsetWithRange.clear(); + } + + unsigned getRowStrideParamIndex() const { return rowStrideParamIndex; } + + bool getHasRowStride() const { return hasRowStride; } + + Value getExpandedOffsetWithRange(OpBuilder &builder, const Location &loc, + unsigned i) { + if (cachedOffsetWithRange.count(i)) + return cachedOffsetWithRange[i]; + + // Add range + auto indexI32RowType = + RankedTensorType::get({tensorShape[i]}, builder.getI32Type()); + auto indexRowType = + RankedTensorType::get({tensorShape[i]}, builder.getI64Type()); + Value splatOffset = + triton::SplatOp::create(builder, loc, indexRowType, offsets[i]); + Value range = triton::MakeRangeOp::create(builder, loc, indexI32RowType, 0, + tensorShape[i]); + Value i64Range = arith::ExtSIOp::create(builder, loc, indexRowType, range); + + // Expand dimensions + Value expandedResult = + arith::AddIOp::create(builder, loc, splatOffset, i64Range); + for (size_t j = 0; j < tensorShape.size(); ++j) { + if (j == i) + continue; + expandedResult = + triton::ExpandDimsOp::create(builder, loc, expandedResult, j); + } + + return cachedOffsetWithRange[i] = expandedResult; + } + + Value generatePtr(OpBuilder &builder, const Location &loc) { + assert(tensorShape.size() == offsets.size() && + tensorShape.size() == strides.size()); + auto indexTensorType = + RankedTensorType::get(tensorShape, builder.getI64Type()); + auto ptrType = cast(base.getType()); + auto ptrTensorType = RankedTensorType::get(tensorShape, ptrType); + + // Generate offsets per dimension + Value ptr = triton::SplatOp::create(builder, loc, ptrTensorType, base); + for (unsigned i = 0; i < tensorShape.size(); ++i) { + auto offsetWithRange = getExpandedOffsetWithRange(builder, loc, i); + + // We must splat strides into the expanded shape not a row for retaining + // the divisibility information given by strides + Value splatStride = triton::SplatOp::create( + builder, loc, offsetWithRange.getType(), strides[i]); + Value offsetWithStride = + arith::MulIOp::create(builder, loc, offsetWithRange, splatStride); + Value broadcasted = triton::BroadcastOp::create( + builder, loc, indexTensorType, offsetWithStride); + + // Add to the pointer + ptr = triton::AddPtrOp::create(builder, loc, ptrTensorType, ptr, + broadcasted); + } + + return ptr; + } + + Value generateMask(OpBuilder &builder, const Location &loc, + const std::optional> &boundaryCheck) { + if (!boundaryCheck.has_value()) + return {}; + + // Generate mask per dimension + auto maskTensorType = + RankedTensorType::get(tensorShape, builder.getI1Type()); + Value mask; + for (auto i : boundaryCheck.value()) { + auto offsetWithRange = getExpandedOffsetWithRange(builder, loc, i); + + // Compare with lower bound + Value lowerBound = mlir::arith::ConstantIntOp::create( + builder, loc, builder.getI64Type(), 0); + Value splatLowerBound = triton::SplatOp::create( + builder, loc, offsetWithRange.getType(), lowerBound); + Value cmpLower = + arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::sge, + offsetWithRange, splatLowerBound); + + // Compare with upper bound + Value splatUpperBound = triton::SplatOp::create( + builder, loc, offsetWithRange.getType(), shape[i]); + Value cmpUpper = + arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::slt, + offsetWithRange, splatUpperBound); + + // And and broadcast + Value andResult = arith::AndIOp::create(builder, loc, cmpLower, cmpUpper); + Value broadcasted = + triton::BroadcastOp::create(builder, loc, maskTensorType, andResult); + + // And up all results + if (!mask) { + mask = broadcasted; + } else { + mask = arith::AndIOp::create(builder, loc, mask, broadcasted); + } + } + + return mask; + } + + Value generateOther(OpBuilder &builder, const Location &loc, + const std::optional &padding) { + if (!padding.has_value()) + return Value(); + + // Create element attribute + auto elementType = + cast(base.getType()).getPointeeType(); + auto otherTensorType = RankedTensorType::get(tensorShape, elementType); + + // Set zero padding value + TypedAttr attr = builder.getZeroAttr(elementType); + + // Float NaN padding case + if (padding.value() == triton::PaddingOption::PAD_NAN) { + assert(!elementType.isIntOrIndex()); + auto apNaN = llvm::APFloat::getNaN( + cast(attr).getValue().getSemantics()); + attr = builder.getFloatAttr(elementType, apNaN); + } + + // Create tensor + Value constant = arith::ConstantOp::create(builder, loc, attr); + return triton::SplatOp::create(builder, loc, otherTensorType, constant); + } +}; + +// 成功返回true +bool findRowStrideParamIndexFromValue(Value v, unsigned* paramIdx) { + Operation* op = v.getDefiningOp(); + while(op != nullptr) { + if(op->getOperands().size() < 1) { + return false; + } + v = op->getOperands()[0]; + op = v.getDefiningOp(); + } + if(auto blockArg = dyn_cast(v)) { + *paramIdx = blockArg.getArgNumber(); + Block* parentBlock = blockArg.getOwner(); + if(parentBlock->isEntryBlock()) { + return true; + } + } + return false; +} + +} // namespace + +// TODO: this pass relies on assumptions of how block pointers are created and +// on pattern matches that walks the SSA links to find the base/strides. This is +// very fragile and to solve we should expose convert Ptr of tensor to a +// structure containins all values and not only offsets. +class RewriteTensorPointerPass + : public impl::TritonRewriteTensorPointerBase { +private: + DenseMap rewritedInfo; + +public: + static bool needRewrite(Operation *op) { + return std::any_of(op->getOperands().begin(), op->getOperands().end(), + [](Value operand) { + return triton::isTensorPointerType(operand.getType()); + }); + } + + static void generateNewOperands(SmallVector &oldOperands, + unsigned index, ArrayRef newValues) { + size_t size = oldOperands.size(); + assert(index < size); + SmallVector operands = oldOperands; + oldOperands.reserve(size - 1 + newValues.size()); + oldOperands.clear(); + if (index != 0) { + oldOperands.append(operands.begin(), operands.begin() + index); + } + oldOperands.append(newValues.begin(), newValues.end()); + if (index != size - 1) { + oldOperands.append(operands.begin() + index + 1, operands.end()); + } + } + + Operation *rewriteMakeTensorPtrOp(OpBuilder &builder, + triton::MakeTensorPtrOp op, + std::stack &eraser) { + // Save info for later use + auto ptrType = cast(op.getType()); + auto tensorType = cast(ptrType.getPointeeType()); + + // Cast I32 offsets into I64 + SmallVector i64Offsets; + for (auto offset : op.getOffsets()) { + auto i64Offset = arith::ExtSIOp::create(builder, op.getLoc(), + builder.getI64Type(), offset); + i64Offsets.push_back(i64Offset); + } + + unsigned rowStrideParamIdx = 0; + bool hasRowStride = false; + auto strides = op.getStrides(); + if(strides.size() == 2) { + hasRowStride = findRowStrideParamIndexFromValue(strides[0], &rowStrideParamIdx); + } + + // Save information + rewritedInfo[op.getResult()] = + RewritedInfo(op.getBase(), op.getShape(), op.getStrides(), i64Offsets, + tensorType.getShape(), hasRowStride, rowStrideParamIdx); + + // Erase the original operation + eraser.push(op); + return nullptr; + } + + Operation *rewriteAdvanceOp(OpBuilder &builder, triton::AdvanceOp op, + std::stack &eraser) { + // Get info from previous results + assert(rewritedInfo.count(op.getPtr())); + auto info = rewritedInfo[op.getPtr()]; + + // Calculate new offsets + assert(info.length() == op.getOffsets().size()); + SmallVector newOffsets; + for (size_t i = 0; i < info.length(); ++i) { + Value i64Offset = arith::ExtSIOp::create( + builder, op.getLoc(), builder.getI64Type(), op.getOffsets()[i]); + Value newOffset = arith::AddIOp::create(builder, op.getLoc(), + info.getOffset(i), i64Offset); + newOffsets.push_back(newOffset); + } + + // Save info for later use + info.setOffsets(newOffsets); + rewritedInfo[op.getResult()] = info; + + // Erase the original operation + eraser.push(op); + return nullptr; + } + + Operation *rewriteLoadStoreOp(OpBuilder &builder, Operation *op, + std::stack &eraser) { + assert(isa(op) || isa(op)); + + // We only have to rewrite load/stores with tensor pointers + auto ptr = op->getOperand(0); + if (!triton::isTensorPointerType(ptr.getType())) + return nullptr; + + // Get info from previous results + assert(rewritedInfo.count(ptr)); + auto info = rewritedInfo[ptr]; + + // Load/store with tensor pointers implicitly will check the bound while + // accessing memory, so we should set `mask` and `other` (according to the + // padding). Also note that load with tensor pointers do not have `mask` and + // `other` while building IR from Python AST + std::optional> boundaryCheck; + if (auto loadOp = dyn_cast(op)) { + assert(!loadOp.getMask() && !loadOp.getOther()); + boundaryCheck = loadOp.getBoundaryCheck(); + } else if (auto storeOp = dyn_cast(op)) { + assert(!storeOp.getMask()); + boundaryCheck = storeOp.getBoundaryCheck(); + } + + // Generate new `ptr`, `mask` and `other` + auto newPtr = info.generatePtr(builder, op->getLoc()); + auto newMask = info.generateMask(builder, op->getLoc(), boundaryCheck); + Value newOther; + if (auto loadOp = dyn_cast(op)) + newOther = info.generateOther(builder, op->getLoc(), loadOp.getPadding()); + + // Create a new operation + if (auto loadOp = dyn_cast(op)) { + auto newResult = triton::LoadOp::create( + builder, loadOp.getLoc(), newPtr, newMask, newOther, + loadOp.getCache(), loadOp.getEvict(), loadOp.getIsVolatile(), + loadOp.getFlagtreeHintsAttr()); // flagtree hints + MLIRContext *ctx = builder.getContext(); + if(info.getHasRowStride() == true) { + newResult->setAttr("sunrise.rowStrideParamIdx", IntegerAttr::get(IntegerType::get(ctx, 32), info.getRowStrideParamIndex())); + } + op->getResult(0).replaceAllUsesWith(newResult); + } else if (auto storeOp = dyn_cast(op)) { + triton::StoreOp::create(builder, storeOp.getLoc(), newPtr, + storeOp.getValue(), newMask, storeOp.getCache(), + storeOp.getEvict()); + } + + // Erase the original operation + eraser.push(op); + return nullptr; + } + + Operation *rewriteIfOp(OpBuilder &builder, scf::IfOp op, + std::stack &eraser) { + auto thenYieldOp = op.thenYield(); + assert(op.getNumResults() == thenYieldOp.getNumOperands()); + SmallVector results = thenYieldOp.getOperands(); + + // get new result types + SmallVector newRetTypes; + bool needRewrite = false; + for (unsigned i = 0; i < results.size(); ++i) { + if (!triton::isTensorPointerType(results[i].getType())) { + newRetTypes.push_back(results[i].getType()); + continue; + } + needRewrite = true; + auto makeTensorPtrOp = triton::getMakeTensorPtrOp(results[i]); + assert(rewritedInfo.count(makeTensorPtrOp.getResult())); + const auto &info = rewritedInfo[makeTensorPtrOp.getResult()]; + for (unsigned j = 0; j < info.length(); ++j) { + newRetTypes.push_back(builder.getI64Type()); + } + } + if (!needRewrite) + return op; + // create and clone new IfOp + bool hasElse = !op.getElseRegion().empty(); + scf::IfOp newOp = scf::IfOp::create(builder, op.getLoc(), newRetTypes, + op.getCondition(), hasElse); + IRMapping mapping; + for (unsigned i = 0; i < op->getNumOperands(); ++i) { + mapping.map(op->getOperand(i), newOp->getOperand(i)); + } + auto rematerialize = [&](Block *block) { + for (Operation &opInIf : block->getOperations()) { + builder.clone(opInIf, mapping); + } + }; + builder.setInsertionPointToStart(newOp.thenBlock()); + rematerialize(op.thenBlock()); + if (hasElse) { + builder.setInsertionPointToStart(newOp.elseBlock()); + rematerialize(op.elseBlock()); + } + + // update rewritedInfo + auto opResults = op.getResults(); + unsigned oldResIdx = 0, newResIdx = 0; + while (oldResIdx < results.size()) { + if (!triton::isTensorPointerType(results[oldResIdx].getType())) { + opResults[oldResIdx].replaceAllUsesWith(newOp.getResult(newResIdx)); + oldResIdx++; + newResIdx++; + } else { + auto makeTensorPtrOp = triton::getMakeTensorPtrOp(results[oldResIdx]); + assert(rewritedInfo.count(makeTensorPtrOp.getResult())); + auto info = rewritedInfo[makeTensorPtrOp.getResult()]; + for (unsigned j = 0; j < info.length(); ++j) { + info.setOffset(j, newOp->getResult(newResIdx++)); + } + rewritedInfo[op.getResult(oldResIdx)] = info; + oldResIdx++; + } + } + + eraser.push(op); + return newOp; + } + + Operation *rewriteForOp(OpBuilder &builder, scf::ForOp op, + std::stack &eraser) { + // Generate new iteration operands and set rewritten information + SmallVector oldIterOperands = llvm::to_vector(op.getInitArgs()); + SmallVector newIterOperands = llvm::to_vector(op.getInitArgs()); + for (unsigned i = 0, oldI = 0, size = op.getInitArgs().size(); i < size; + ++i, ++oldI) { + if (!triton::isTensorPointerType(newIterOperands[i].getType())) + continue; + + // Expand the tensor pointer into offsets + assert(rewritedInfo.count(newIterOperands[i])); + auto info = rewritedInfo[newIterOperands[i]]; + generateNewOperands(newIterOperands, i, info.getOffsets()); + i += info.length() - 1; + size += info.length() - 1; + } + + // Rebuild the loop type + auto newForOp = + scf::ForOp::create(builder, op.getLoc(), op.getLowerBound(), + op.getUpperBound(), op.getStep(), newIterOperands); + newForOp->setAttrs(op->getAttrs()); + + // Create value mapping. Note that for tensor pointers, we use identity + // mapping. It may refer to a value in the old loop, but we will rewrite it + // later + IRMapping mapping; + for (unsigned i = 0, oldI = 0, sz = op.getInitArgs().size(); oldI < sz; + ++i, ++oldI) { + auto oldRegionIterArg = op.getRegionIterArg(oldI); + if (triton::isTensorPointerType(oldRegionIterArg.getType())) { + // Pass rewritten info inside + assert(rewritedInfo.count(oldIterOperands[oldI])); + auto info = rewritedInfo[oldIterOperands[oldI]]; + mapping.map(oldRegionIterArg, oldRegionIterArg); + for (unsigned j = 0; j < info.length(); ++j) + info.setOffset(j, newForOp.getRegionIterArg(i + j)); + rewritedInfo[oldRegionIterArg] = info; + i += info.length() - 1; + } else { + mapping.map(oldRegionIterArg, newForOp.getRegionIterArg(i)); + } + } + mapping.map(op.getInductionVar(), newForOp.getInductionVar()); + + // Clone body + builder.setInsertionPointToStart(newForOp.getBody()); + for (auto &opInFor : *op.getBody()) { + builder.clone(opInFor, mapping); + } + + // Replace later usages + assert(op.getNumResults() == op.getInitArgs().size()); + for (unsigned i = 0, oldI = 0; oldI < op.getNumResults(); ++i, ++oldI) { + auto oldResult = op.getResult(oldI); + if (triton::isTensorPointerType(oldResult.getType())) { + // Pack new offsets into rewritten info + assert(rewritedInfo.count(oldIterOperands[oldI])); + auto info = rewritedInfo[oldIterOperands[oldI]]; + for (unsigned j = 0; j < info.length(); ++j) + info.setOffset(j, newForOp.getResult(i + j)); + i += info.length() - 1; + rewritedInfo[oldResult] = info; + } else { + oldResult.replaceAllUsesWith(newForOp.getResult(i)); + } + } + + // Erase later + eraser.push(op); + return newForOp; + } + + Operation *rewriteYieldOp(OpBuilder &builder, scf::YieldOp op, + std::stack &eraser) { + // Replace tensor pointers with offsets + SmallVector newOperands = op->getOperands(); + for (unsigned i = 0, size = op.getNumOperands(); i < size; ++i) { + if (!triton::isTensorPointerType(newOperands[i].getType())) + continue; + + assert(rewritedInfo.count(newOperands[i])); + auto info = rewritedInfo[newOperands[i]]; + generateNewOperands(newOperands, i, info.getOffsets()); + i += info.length() - 1; + size += info.length() - 1; + } + op->setOperands(newOperands); + + // No need to erase + return nullptr; + } + + Operation *rewriteOp(Operation *op, std::stack &eraser) { + OpBuilder builder(op); + + // Rewrite `make_tensor_ptr` and `advance` and make a tensor of pointers + // Rewriting functions return the next operation to visit, if there is no + // next one, simply return `nullptr` + if (auto makeTensorPtrOp = dyn_cast(op)) { + return rewriteMakeTensorPtrOp(builder, makeTensorPtrOp, eraser); + } else if (auto advanceOp = dyn_cast(op)) { + return rewriteAdvanceOp(builder, advanceOp, eraser); + } else if (isa(op) || isa(op)) { + return rewriteLoadStoreOp(builder, op, eraser); + } else if (isa(op->getDialect())) { + if (auto ifOp = dyn_cast(op)) { + return rewriteIfOp(builder, ifOp, eraser); + } + if (!needRewrite(op)) + return op; + + if (auto forOp = dyn_cast(op)) { + return rewriteForOp(builder, forOp, eraser); + } else if (auto yieldOp = dyn_cast(op)) { + return rewriteYieldOp(builder, yieldOp, eraser); + } else { + llvm_unreachable("Currently we only support tensor pointer usages " + "inside a `scf::ForOp` or `scf::IfOp`, others such as " + "`scf::WhileOp`, `cf::BranchOp` or `cf::CondBranchOp` " + "are not supported yet"); + } + } + + // Otherwise return the original one + return op; + } + + void visitOperation(Operation *op, std::stack &eraser) { + for (Region ®ion : op->getRegions()) { + for (Block &block : region) { + for (Operation &nestedOp : llvm::make_early_inc_range(block)) { + if (auto newOp = rewriteOp(&nestedOp, eraser)) { + visitOperation(newOp, eraser); + } + } + } + } + } + + void runOnOperation() override { + // NOTES(Chenggang): we don't use `ConversionPatternRewriter`, because + // MLIR does not support one-multiple value mapping. For example, if we use + // `ConversionPatternRewriter`, we can not make a type converter, which + // converts `ptr` into multiple types `ptr<>, int64, int64, ...` + // (containing the base/offsets/strides...). What we can do is to convert + // `ptr` into a single type `Tuple, int64, int64, ...>`. But + // in this way, we also have to define `PackTuple` and `UnpackTuple` + // operations and make a canonicalization pass to optimize, which is much + // So here we recursively build the IR, to be specific, we have to rewrite + // `tt.make_tensor_ptr`, `tt.advance`, `tt.load`, `tt.store`, + // `scf.for` (tensor pointer usages may be in a loop fashion) + std::stack eraser; + visitOperation(getOperation(), eraser); + + // The operation could not be erased during visit, because they may have + // later usages, so we erase after visit + rewritedInfo.clear(); + while (!eraser.empty()) { + auto op = eraser.top(); + eraser.pop(); + op->erase(); + } + + // 此时loadOp不再有boundaryCheck，根据参数个数判断是否有mask + getOperation()->walk([](Operation* op){ + if(isa(*op) == false) { + return; + } + OpBuilder builder(op); + MLIRContext *ctx = builder.getContext(); + int hasOriMask = 0; + if(op->getNumOperands() > 1) { + hasOriMask = 1; + } + op->setAttr("sunrise.hasOriMask", IntegerAttr::get(IntegerType::get(ctx, 32), hasOriMask)); + }); + } +}; + +} // namespace mlir::triton diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/CMakeLists.txt b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/CMakeLists.txt new file mode 100644 index 0000000000..9f57627c32 --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/CMakeLists.txt @@ -0,0 +1,2 @@ +add_subdirectory(IR) +add_subdirectory(Transforms) diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/CMakeLists.txt b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/CMakeLists.txt new file mode 100644 index 0000000000..e462f09336 --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/CMakeLists.txt @@ -0,0 +1,14 @@ +add_triton_library(FlagTree_sunrise_TritonGPUIR + Dialect.cpp + LinearLayoutConversions.cpp + + DEPENDS + TritonGPUTableGen + TritonGPUAttrDefsIncGen + TritonGPUTypeInterfacesIncGen + + LINK_LIBS PUBLIC + MLIRGPUDialect + TritonIR + TritonTools +) diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp new file mode 100644 index 0000000000..3a43eebb05 --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp @@ -0,0 +1,4182 @@ +#include "triton/Dialect/Triton/IR/Dialect.h" + +#include +#include +#include + +#include "mlir/Dialect/UB/IR/UBOps.h" +#include "mlir/IR/DialectImplementation.h" +#include "mlir/IR/OpImplementation.h" +#include "mlir/IR/OperationSupport.h" +#include "mlir/Support/LLVM.h" +#include "triton/Analysis/Utility.h" +#include "triton/Dialect/Triton/IR/Interfaces.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Attributes.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h" +#include "triton/Dialect/TritonGPU/IR/TritonGPUInterfaces.h" +#include "triton/Dialect/TritonGPU/IR/Types.h" +#include "triton/Dialect/TritonGPU/Transforms/Utility.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +#include "triton/Tools/LayoutUtils.h" +#include "triton/Tools/LinearLayout.h" +#include "triton/Tools/StrUtil.h" +#include "llvm/ADT/SmallSet.h" +#include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/MathExtras.h" + +// Include TableGen'erated code +#include "triton/Dialect/TritonGPU/IR/Dialect.cpp.inc" +#include "triton/Dialect/TritonGPU/IR/OpInterfaces.cpp.inc" +#include "triton/Dialect/TritonGPU/IR/TypeInterfaces.cpp.inc" + +using namespace mlir; +using namespace mlir::triton; +using namespace mlir::triton::gpu; + +static SmallVector +basesPerDimImpl(const LinearLayout::BasesT &namedBases, StringAttr dimName, + size_t rank, bool skipBroadcast = true); + +// Utility +namespace mlir { +namespace triton { +namespace gpu { + +LinearEncodingAttr TritonGPUDialect::toLinearEncoding(ArrayRef shape, + Attribute layout) { + // LinearEncoding is a DistributedLayout + std::vector allocationShape; + CacheKey key{std::vector(shape.begin(), shape.end()), layout}; + if (auto result = leCache.get(key)) { + return *result; + } + auto linearLayout = toLinearLayout(shape, layout); + auto linearEncoding = + LinearEncodingAttr::get(layout.getContext(), std::move(linearLayout)); + leCache.set(key, linearEncoding); + return linearEncoding; +} + +LinearEncodingAttr toLinearEncoding(DistributedEncodingTrait layout, + ArrayRef shape) { + auto *ctx = layout.getContext(); + return ctx->getLoadedDialect()->toLinearEncoding(shape, + layout); +} + +LinearEncodingAttr toLinearEncoding(RankedTensorType type) { + auto *ctx = type.getContext(); + return ctx->getLoadedDialect()->toLinearEncoding( + type.getShape(), type.getEncoding()); +} + +unsigned getTotalElemsPerThread(Attribute layout, ArrayRef shape) { + return toLinearEncoding(cast(layout), shape) + .getTotalElemsPerThread(shape); +} + +SmallVector getElemsPerThread(Attribute layout, + ArrayRef shape) { + return toLinearEncoding(cast(layout), shape) + .getElemsPerThread(shape); +} + +SmallVector getElemsPerThread(Type type) { + if (type.isIntOrIndexOrFloat() || isa(type)) + return SmallVector(1, 1); + auto tensorType = cast(type); + return getElemsPerThread(tensorType.getEncoding(), tensorType.getShape()); +} + +unsigned getTotalElemsPerThread(Type type) { + if (type.isIntOrIndexOrFloat() || isa(type)) + return 1; + auto tensorType = cast(type); + return getTotalElemsPerThread(tensorType.getEncoding(), + tensorType.getShape()); +} + +SmallVector getThreadsPerWarp(Attribute layout, + ArrayRef shape) { + return toLinearEncoding(cast(layout), shape) + .getThreadsPerWarp(); +} + +SmallVector getWarpsPerCTA(Attribute layout, + ArrayRef shape) { + return toLinearEncoding(cast(layout), shape) + .getWarpsPerCTA(); +} + +SmallVector getContigPerThread(RankedTensorType type) { + return toLinearEncoding(type).getContigPerThread(); +} + +bool isExpensiveView(Type srcType, Type dstType) { + auto tensorSrcType = cast(srcType); + auto tensorDstType = cast(dstType); + auto llSrc = toLinearLayout(tensorSrcType); + auto llDst = toLinearLayout(tensorDstType); + // In case there are replicated value we need to make sure the new and old + // layout have matching masks. + for (auto [srcMask, dstMask] : + llvm::zip(llSrc.getFreeVariableMasks(), llDst.getFreeVariableMasks())) { + assert(srcMask.first == dstMask.first); + if (srcMask.second != dstMask.second) + return true; + } + return getTotalElemsPerThread(srcType) != getTotalElemsPerThread(dstType); +} + +/* Utility function used by get.*Order methods of SliceEncodingAttr. + * Erase dim and decrease all values larger than dim by 1. + * Example: order = [0, 2, 4, 3, 1], dim = 2 + * resOrder = [0, 3, 2, 1] + */ +static SmallVector eraseOrder(ArrayRef order, + unsigned dim) { + unsigned rank = order.size(); + assert(dim < rank && "Invalid dim to erase"); + SmallVector resOrder; + for (unsigned i : order) + if (i < dim) + resOrder.push_back(i); + else if (i > dim) + resOrder.push_back(i - 1); + return resOrder; +} + +SmallVector getMatrixOrder(unsigned rank, bool rowMajor) { + // Return the order that represents that the batch is in row-major or + // column-major order for a batch of matrices of shape [*, m, n] with + // len(shape) == rank. + SmallVector order(rank); + if (rank < 2) { + return order; + } + std::iota(order.rbegin(), order.rend(), 0); + if (!rowMajor) { + std::swap(order[0], order[1]); + } + return order; +} + +SmallVector getOrderForDotOperand(unsigned opIdx, unsigned rank, + bool kContig) { + // kContig: if true, the matrix is fastest-running on k, + // otherwise it is on m (resp. n) + // opIdx=0: [*batch, m, k] + // opIdx=1: [*batch, k, n] + assert(opIdx == 0 || opIdx == 1); + auto rowMajor = bool(opIdx) != kContig; + return getMatrixOrder(rank, rowMajor); +} + +SmallVector getRepOrder(RankedTensorType type) { + auto layout = type.getEncoding(); + if (auto distributedLayout = mlir::dyn_cast(layout)) + return distributedLayout.getRepOrder(); + else + llvm::report_fatal_error("Unimplemented usage of getRepOrder"); + return {}; +} + +// Legacy impl for now +// This one's not terribly bad as we don't broadcast ShareEncodings +SmallVector getOrder(SharedEncodingTrait layout, + ArrayRef shape) { + if (auto swizzledLayout = dyn_cast(layout)) { + return llvm::to_vector(swizzledLayout.getOrder()); + } + if (auto paddedEnc = dyn_cast(layout)) { + return paddedEnc.getOrder(); + } + if (auto linearEnc = dyn_cast(layout)) { + return linearEnc.getOrder(); + } + if (auto sharedLayout = dyn_cast(layout)) { + if (shape.size() == 1) { + return {0}; + } + return getMatrixOrder(shape.size(), !sharedLayout.getTransposed()); + } + if (auto sharedLayout = dyn_cast(layout)) { + return llvm::to_vector(sharedLayout.getOrder()); + } + llvm::report_fatal_error("Unimplemented usage of getOrder for MemDescType"); + return {}; +} + +SmallVector getOrder(DistributedEncodingTrait layout, + ArrayRef shape) { + return toLinearEncoding(layout, shape).getOrder(); +} + +SmallVector getOrderForMemory(DistributedEncodingTrait layout, + ArrayRef shape) { + auto linear = toLinearEncoding(layout, shape); + auto order = linear.getOrder(); + auto threadOrder = linear.getThreadOrder(); + if (order == threadOrder) { + return order; + } + // Heuristic: + // If the element contiguity does not align with the thread order + // because the thread order dimension has contiguity of 1---meaning that + // the order position of this dimension is irrelevant---we prefer + // to use the thread order for the memory layout + auto contig = linear.getElemsPerThread(shape); + if (contig[threadOrder[0]] == 1) { + return threadOrder; + } + return order; +} + +SmallVector getThreadOrder(DistributedEncodingTrait layout, + ArrayRef shape) { + return toLinearEncoding(layout, shape).getThreadOrder(); +} + +SmallVector getWarpOrder(DistributedEncodingTrait layout, + ArrayRef shape) { + return toLinearEncoding(layout, shape).getWarpOrder(); +} + +CTAEncodingAttr getCTALayout(Attribute layout) { + if (auto ttgLayout = mlir::dyn_cast(layout)) + return ttgLayout.getCTALayout(); + llvm::report_fatal_error("Unimplemented usage of getCTALayout"); + return {}; +} + +SmallVector getCTAsPerCGA(Attribute layout) { + if (auto ttgLayout = mlir::dyn_cast(layout)) + return ttgLayout.getCTALayout().getCTAsPerCGA(); + llvm::report_fatal_error("Unimplemented usage of getCTAsPerCGA"); +} + +SmallVector getCTASplitNum(Attribute layout) { + SmallVector res; + if (auto ttgLayout = mlir::dyn_cast(layout)) { + return ttgLayout.getCTALayout().getCTASplitNum(); + } else if (auto tmemLayout = + mlir::dyn_cast( + layout)) { + res.resize(2); + res[0] = tmemLayout.getCTASplitM(); + res[1] = tmemLayout.getCTASplitN(); + } else if (auto tmemScaleLayout = mlir::dyn_cast< + triton::nvidia_gpu::TensorMemoryScalesEncodingAttr>(layout)) { + res.resize(2); + res[0] = tmemScaleLayout.getCTASplitM(); + res[1] = tmemScaleLayout.getCTASplitN(); + } else { + assert(false && "Unimplemented usage of getCTASplitNum"); + } + return res; +} + +SmallVector getCTAOrder(Attribute layout) { + SmallVector res; + if (auto ttgLayout = mlir::dyn_cast(layout)) { + res = ttgLayout.getCTALayout().getCTAOrder(); + } else { + llvm::report_fatal_error("Unimplemented usage of getCTAOrder"); + } + return res; +} + +SmallVector getShapePerCTA(ArrayRef CTASplitNum, + ArrayRef shape) { + unsigned rank = shape.size(); + auto splitNum = llvm::to_vector(CTASplitNum); + if (splitNum.size() <= rank) { // pipelining + splitNum.insert(splitNum.begin(), rank - splitNum.size(), 1); + } else { // memory slicing + splitNum = + llvm::to_vector(llvm::drop_begin(splitNum, splitNum.size() - rank)); + } + SmallVector shapePerCTA(rank); + for (unsigned i = 0; i < rank; ++i) { + shapePerCTA[i] = shape[i] / std::min(shape[i], splitNum[i]); + } + return shapePerCTA; +} + +SmallVector getShapePerCTA(Attribute layout, ArrayRef shape) { + return getShapePerCTA(getCTASplitNum(layout), shape); +} + +SmallVector getAllocationShapePerCTA(Attribute layout, + ArrayRef shapeLogical) { + SmallVector shape(shapeLogical); + if (auto sharedMMALayout = dyn_cast(layout)) { + if (sharedMMALayout.getFp4Padded()) { + auto packedAxis = getOrder(sharedMMALayout, shapeLogical)[0]; + shape[packedAxis] *= 2; + } + } + return getShapePerCTA(layout, shape); +} + +SmallVector getShapePerCTA(Type type) { + auto tensorType = cast(type); + return getShapePerCTA(tensorType.getEncoding(), tensorType.getShape()); +} + +SmallVector getAllocationShapePerCTA(Type type) { + auto tensorType = cast(type); + return getAllocationShapePerCTA(tensorType.getEncoding(), + tensorType.getShape()); +} + +unsigned getNumCTAs(Attribute layout) { + return product(getCTAsPerCGA(layout)); +} + +SmallVector orderPerDimImpl(const LinearLayout &ll, + StringAttr dimName, + ArrayRef defaultOrder) { + assert(ll.getBases().contains(dimName)); + const auto &bases = ll.getBases().find(dimName)->second; + llvm::SetVector order; + auto nonZero = [](auto val) { return val != 0; }; + for (const auto &basis : bases) { + // Bases can have one or zero non-zero elements + // Skip a basis if it's broadcasting (all zeros) + // e.g. warps for DotOperandEncodingAttr (see ampereDotToLinearLayout) + auto it = std::find_if(basis.begin(), basis.end(), nonZero); + if (it != basis.end()) { + auto i = it - basis.begin(); + order.insert(i); + } + } + // If any dim is missing, we add them in the defaultOrder + for (auto i : defaultOrder) { + order.insert(i); + } + return order.takeVector(); +} + +bool isExpensiveCat(CatOp cat, Attribute targetEncoding) { + // If the new elements per thread is less than the old one, we will need to + // do convert encoding that goes through shared memory anyway. So we + // consider it as expensive. + RankedTensorType tensorTy = cat.getType(); + auto totalElemsPerThread = gpu::getTotalElemsPerThread(tensorTy); + auto shape = tensorTy.getShape(); + auto newTotalElemsPerThread = + gpu::getTotalElemsPerThread(targetEncoding, shape); + return newTotalElemsPerThread < totalElemsPerThread; +} + +static LogicalResult +verifyLayoutOrder(function_ref emitError, + ArrayRef order) { + if (!isPermutationOfIota(order)) { + return emitError() + << "order must be a permutation of 0..(rank-1), but was [" << order + << "]"; + } + return success(); +} + +LogicalResult +CTAEncodingAttr::verify(function_ref emitError, + LinearLayout linearLayout) { + if (linearLayout.getNumInDims() != 1) { + return emitError() << "CTA encoding must have exactly one input dimension " + "named 'block'."; + } + auto dim = *linearLayout.getInDimNames().begin(); + auto ctx = dim.getContext(); + if (dim != StringAttr::get(ctx, "block")) { + return emitError() << "CTA encoding must have exactly one input dimension " + "named 'block'."; + } + + auto outDimNames = linearLayout.getOutDimNames(); + auto expected = standardOutDimNames(ctx, linearLayout.getNumOutDims()); + if (!llvm::equal(outDimNames, expected)) { + return emitError() << "CTA encoding output dims must be [dim0, dim1, ...], " + "but got [" + << outDimNames << "]."; + } + + return success(); +} + +CTAEncodingAttr CTAEncodingAttr::getDefault(MLIRContext *ctx, int rank) { + auto kBlock = StringAttr::get(ctx, "block"); + LinearLayout::BasesT bases; + bases[kBlock] = {}; + auto dims = standardOutDimNames(ctx, rank); + return get(ctx, LinearLayout(bases, dims)); +} + +CTAEncodingAttr CTAEncodingAttr::fromSplitParams(MLIRContext *ctx, + ArrayRef CTAsPerCGA, + ArrayRef CTASplitNum, + ArrayRef CTAOrder) { + int rank = CTAOrder.size(); + auto outDimNames = standardOutDimNames(ctx, rank); + StringAttr kBlock = StringAttr::get(ctx, "block"); + + LinearLayout layout = LinearLayout::empty(); + SmallVector splitNums(CTASplitNum.begin(), CTASplitNum.end()); + SmallVector ctas(CTAsPerCGA.begin(), CTAsPerCGA.end()); + + for (int i = 0; i < rank; ++i) { + int dim = CTAOrder[i]; + unsigned split = splitNums[dim]; + unsigned total = ctas[dim]; + assert(total % split == 0 && "invalid CTA encoding parameters"); + layout *= LinearLayout::identity1D(split, kBlock, outDimNames[dim]) * + LinearLayout::zeros1D(total / split, kBlock, outDimNames[dim]); + } + + layout = layout.transposeOuts(outDimNames); + return CTAEncodingAttr::get(ctx, layout); +} + +SmallVector CTAEncodingAttr::getCTAsPerCGA() const { + auto ll = getLinearLayout(); + auto rank = ll.getNumOutDims(); + return basesPerDimImpl(ll.getBases(), StringAttr::get(getContext(), "block"), + rank, /*skipBroadcast=*/false); +} + +SmallVector CTAEncodingAttr::getCTASplitNum() const { + auto ll = getLinearLayout(); + auto rank = ll.getNumOutDims(); + return basesPerDimImpl(ll.getBases(), StringAttr::get(getContext(), "block"), + rank); +} + +SmallVector CTAEncodingAttr::getCTAOrder() const { + auto rank = getRank(); + SmallVector defaultOrder(rank); + std::iota(defaultOrder.begin(), defaultOrder.end(), 0); + return orderPerDimImpl(getLinearLayout(), + StringAttr::get(getContext(), "block"), defaultOrder); +} + +LogicalResult BlockedEncodingAttr::verify( + function_ref emitError, + ArrayRef sizePerThread, ArrayRef threadsPerWarp, + ArrayRef warpsPerCTA, ArrayRef order, + CTAEncodingAttr CTALayout) { + if (!llvm::all_equal({sizePerThread.size(), threadsPerWarp.size(), + warpsPerCTA.size(), order.size()})) { + return emitError() << "sizePerThread, threadsPerWarp, warpsPerCTA, and " + "order must all have the same rank."; + } + if (llvm::any_of(sizePerThread, + [](unsigned x) { return !llvm::isPowerOf2_64(x); })) { + return emitError() + << "Every element in sizePerThread must be a power of two."; + } + if (llvm::any_of(threadsPerWarp, + [](unsigned x) { return !llvm::isPowerOf2_64(x); })) { + return emitError() + << "Every element in threadsPerWarp must be a power of two."; + } + if (llvm::any_of(warpsPerCTA, + [](unsigned x) { return !llvm::isPowerOf2_64(x); })) { + return emitError() + << "Every element in warpsPerCTA must be a power of two."; + } + + // Empty CTALayout is allowed, but if it's present its rank must match the + // BlockedEncodingAttr's rank. + if (order.size() != CTALayout.getRank()) { + return emitError() << "BlockedEncodingAttr and CTALayout's fields must " + "have the same rank."; + } + return verifyLayoutOrder(emitError, order); +} + +// 1 element per thread +// order = reverse(arange(rank)) +triton::gpu::BlockedEncodingAttr +getDefaultBlockedEncoding(MLIRContext *context, ArrayRef shape, + int numWarps, int threadsPerWarp, int numCTAs) { + int rank = shape.size(); + llvm::SmallVector order(rank); + std::iota(order.begin(), order.end(), 0); + std::reverse(order.begin(), order.end()); + llvm::SmallVector sizePerThread(rank, 1); + triton::gpu::BlockedEncodingAttr encoding = + triton::gpu::BlockedEncodingAttr::get(context, shape, sizePerThread, + order, numWarps, threadsPerWarp, + numCTAs); + return encoding; +} + +LogicalResult tryJoinOnAxis(MLIRContext *ctx, const LinearLayout &inLl, + LinearLayout &outLl, bool fwdInference, int axis, + std::optional loc) { + auto kRegister = StringAttr::get(ctx, "register"); + auto outDims = llvm::to_vector(inLl.getOutDimNames()); + if (fwdInference) { + auto split = LinearLayout::identity1D(2, kRegister, outDims[axis]); + outLl = split * inLl; + } else { + // Assert that there is a dimension with size 2 in the axis + // that has contiguous elements + // Note that this is more general than the fwdInference case in that + // - It allows the dimension not to be the fastest running + // - It allows broadcasting + // In general, this allows us to split along any axis as long as + // the basis (0, 0, ..., 0, 1, 0, ..., 0) is in the registers. + bool found = false; + LinearLayout::BasesT newBases; + for (const auto &basesDim : inLl.getBases()) { + std::vector> newBasesDim; + for (auto base : basesDim.second) { + if (base[axis] == 1 && basesDim.first == kRegister) { + found = true; + continue; + } + base[axis] /= 2; + newBasesDim.push_back(std::move(base)); + } + newBases.insert({basesDim.first, std::move(newBasesDim)}); + } + if (!found) + return emitOptionalError(loc, + "Fp4ToFpOp/SplitOp requires at least 2 elements " + "per thread in the axis/last dimension"); + outLl = LinearLayout(std::move(newBases), std::move(outDims)); + } + return success(); +} + +} // namespace gpu +} // namespace triton +} // namespace mlir + +static LogicalResult parseIntAttrValue(AsmParser &parser, Attribute attr, + unsigned &value, StringRef desc) { + auto intAttr = mlir::dyn_cast(attr); + if (!intAttr) { + parser.emitError(parser.getNameLoc(), "expected an integer type in ") + << desc; + return failure(); + } + if (intAttr.getType().isSignedInteger()) { + int64_t attrVal = intAttr.getSInt(); + if (attrVal < 0) { + parser.emitError(parser.getNameLoc(), + "expected an unsigned integer value in ") + << desc; + return failure(); + } + value = attrVal; + } else if (intAttr.getType().isSignlessInteger()) { + int64_t attrVal = intAttr.getInt(); + if (attrVal < 0) { + parser.emitError(parser.getNameLoc(), + "expected an unsigned integer value in ") + << desc; + return failure(); + } + value = attrVal; + } else { + value = intAttr.getUInt(); + } + return success(); +} + +static LogicalResult parseBoolAttrValue(AsmParser &parser, Attribute attr, + bool &value, StringRef desc) { + auto boolAttr = mlir::dyn_cast(attr); + if (!boolAttr) { + parser.emitError(parser.getNameLoc(), "expected a bool type in ") << desc; + return failure(); + } + value = boolAttr.getValue(); + return success(); +} + +// parse an array of integers +static LogicalResult parseIntArrayAttr(AsmParser &parser, + const NamedAttribute &attr, + SmallVector &res, + StringRef desc) { + auto arrayAttr = mlir::dyn_cast(attr.getValue()); + if (!arrayAttr) { + parser.emitError(parser.getNameLoc(), "expected an array for ") << desc; + return failure(); + } + for (Attribute i : arrayAttr) { + unsigned value; + if (parseIntAttrValue(parser, i, value, desc).failed()) + return failure(); + res.push_back(value); + } + return success(); +}; + +static LogicalResult parseUInt(AsmParser &parser, const NamedAttribute &attr, + unsigned &value, StringRef desc) { + return parseIntAttrValue(parser, attr.getValue(), value, desc); +}; + +static LogicalResult parseBool(AsmParser &parser, const NamedAttribute &attr, + bool &value, StringRef desc) { + return parseBoolAttrValue(parser, attr.getValue(), value, desc); +}; + +static LogicalResult parseType(AsmParser &parser, const NamedAttribute &attr, + Type &value, StringRef desc) { + auto typeAttr = mlir::dyn_cast(attr.getValue()); + if (!typeAttr) { + parser.emitError(parser.getNameLoc(), "expected a Type in ") << desc; + return failure(); + } + value = typeAttr.getValue(); + return success(); +} + +std::optional +parseLinearLayout(const DictionaryAttr &dict, AsmParser &parser, + ArrayRef inDimNames) { + LinearLayout::BasesT bases; + + // Parse the basis names in order (the order is relevant) + for (const auto &inDimNameStr : inDimNames) { + auto inDimName = StringAttr::get(parser.getContext(), inDimNameStr); + Attribute value = dict.get(inDimName); + if (!value) { + parser.emitError(parser.getCurrentLocation(), "Expected basis of '") + << inDimName.getValue() << "' not found"; + return {}; + } + // Expecting an array of arrays + auto arrayOfArraysAttr = mlir::dyn_cast(value); + if (!arrayOfArraysAttr) { + parser.emitError(parser.getCurrentLocation(), + "Expected array of arrays for basis of '") + << inDimName.getValue() << "'"; + return {}; + } + + std::vector> inDimBases; + for (Attribute arrayAttr : arrayOfArraysAttr) { + auto intArrayAttr = mlir::dyn_cast(arrayAttr); + if (!intArrayAttr) { + parser.emitError(parser.getCurrentLocation(), + "Expected array of integers in basis for '") + << inDimName.getValue() << "'"; + return {}; + } + std::vector basis; + for (Attribute intAttr : intArrayAttr) { + auto intValueAttr = mlir::dyn_cast(intAttr); + if (!intValueAttr) { + parser.emitError(parser.getCurrentLocation(), + "Expected integer in basis for '") + << inDimName.getValue() << "'"; + return {}; + } + basis.push_back(intValueAttr.getInt()); + } + inDimBases.push_back(std::move(basis)); + } + bases[inDimName] = std::move(inDimBases); + } + size_t rank = 0; + for (const auto &basesDim : llvm::make_second_range(bases)) { + if (!basesDim.empty()) { + rank = basesDim[0].size(); + break; + } + } + + // To implement this we'd need to serialise the rank as well. + // We can do this if we ever need it + if (rank == 0) { + parser.emitError(parser.getCurrentLocation(), "Empty Layout not supported"); + return {}; + } + + // Generate standared outDimNames (dim0, dim1, ...) + SmallVector outDimNames; + for (int i = 0; i < rank; ++i) { + outDimNames.push_back( + StringAttr::get(parser.getContext(), "dim" + llvm::Twine(i))); + } + + // Create LinearLayout + return LinearLayout(std::move(bases), std::move(outDimNames)); +} + +// We don't use the default implementation as it's a bit too verbose +// This prints in the following format that is shape agnostic, in the sense +// that we don't print explicitly the outShape of the LL +// We always assume LLs to be surjective +// <{register = [[0, 1], [8, 0], [0, 8], [64, 0]], +// lane = [[0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], +// warp = [[16, 0], [32, 0]], +// block = []}> +static void printLinearLayout(AsmPrinter &printer, const LinearLayout &ll) { + printer << join(ll.getBases(), ", ", [](const auto &base) { + return base.first.str() + " = " + "[" + + join(base.second, ", ", + [](const std::vector &vec) { + return "[" + join(vec, ", ") + "]"; + }) + + "]"; + }); +} + +// Print the CTA encoding as `CGALayout = [[...]]` when the layout is +// non-trivial. +static void maybePrintCTALayout(mlir::MLIRContext *context, + mlir::AsmPrinter &printer, + CTAEncodingAttr layout, unsigned rank) { + if (layout == CTAEncodingAttr::getDefault(context, rank)) + return; + + auto kBlock = StringAttr::get(context, "block"); + const auto &basesMap = layout.getLinearLayout().getBases(); + auto it = basesMap.find(kBlock); + assert(it != basesMap.end()); + const auto &bases = it->second; + // This is the default layout + assert(!bases.empty()); + + printer << ", CGALayout = ["; + llvm::interleaveComma(bases, printer, [&](const std::vector &vec) { + printer << "["; + llvm::interleaveComma(vec, printer); + printer << "]"; + }); + printer << "]"; +} + +//===----------------------------------------------------------------------===// +// Attribute methods +//===----------------------------------------------------------------------===// + +#include "triton/Dialect/TritonGPU/IR/AttrInterfaces.cpp.inc" + +#define GET_ATTRDEF_CLASSES +#include "triton/Dialect/TritonGPU/IR/AttrDefs.cpp.inc" +#undef GET_ATTRDEF_CLASSES + +//===----------------------------------------------------------------------===// +// Blocked Encoding +//===----------------------------------------------------------------------===// + +std::optional parseCTAAttr(AsmParser &parser, Attribute attr, + unsigned rank) { + if (!attr) + return CTAEncodingAttr::getDefault(parser.getContext(), rank); + + auto array = llvm::dyn_cast(attr); + if (!array) { + parser.emitError(parser.getNameLoc(), + "expected array value for 'CGALayout'"); + return {}; + } + + auto ctx = parser.getContext(); + auto cgaName = StringAttr::get(ctx, "CGALayout"); + std::vector> bases; + bases.reserve(array.size()); + for (Attribute vecAttr : array) { + SmallVector basisValues; + NamedAttribute basisAttr(cgaName, vecAttr); + if (parseIntArrayAttr(parser, basisAttr, basisValues, "CGALayout entry") + .failed()) + return {}; + if (basisValues.size() != rank) { + parser.emitError(parser.getNameLoc()) + << "'CGALayout' entry length does not match rank " << rank; + return {}; + } + std::vector basis; + basis.reserve(basisValues.size()); + for (unsigned value : basisValues) + basis.push_back(static_cast(value)); + bases.push_back(std::move(basis)); + } + + LinearLayout::BasesT namedBases; + namedBases.insert( + std::make_pair(StringAttr::get(ctx, "block"), std::move(bases))); + LinearLayout ll(namedBases, standardOutDimNames(ctx, rank)); + return CTAEncodingAttr::get(ctx, std::move(ll)); +} + +Attribute BlockedEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + // Parse the data as a dictionary + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + SmallVector sizePerThread; + SmallVector threadsPerWarp; + SmallVector warpsPerCTA; + SmallVector order; + Attribute ctaAttr = nullptr; + + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "sizePerThread") { + if (parseIntArrayAttr(parser, attr, sizePerThread, + "number of elements per thread") + .failed()) + return {}; + } else if (attr.getName() == "threadsPerWarp") { + if (parseIntArrayAttr(parser, attr, threadsPerWarp, + "number of threads per warp") + .failed()) + return {}; + } else if (attr.getName() == "warpsPerCTA") { + if (parseIntArrayAttr(parser, attr, warpsPerCTA, + "number of warps per CTA") + .failed()) + return {}; + } else if (attr.getName() == "order") { + if (parseIntArrayAttr(parser, attr, order, "order").failed()) + return {}; + } else if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + } else { + parser.emitError(parser.getNameLoc(), "unexpected key: ") + << attr.getName().strref(); + return {}; + } + } + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, /*rank=*/sizePerThread.size()); + if (!CTALayout.has_value()) + return {}; + + return parser.getChecked(parser.getContext(), + sizePerThread, threadsPerWarp, + warpsPerCTA, order, *CTALayout); +} + +void BlockedEncodingAttr::print(mlir::AsmPrinter &printer) const { + printer << "<{" + << "sizePerThread = [" << ArrayRef(getSizePerThread()) << "]" + << ", threadsPerWarp = [" << ArrayRef(getThreadsPerWarp()) << "]" + << ", warpsPerCTA = [" << ArrayRef(getWarpsPerCTA()) << "]" + << ", order = [" << getOrder() << "]"; + + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getSizePerThread().size()); + + printer << "}>"; +} + +// FIXME Can we take the LinearLayout by const&? +LogicalResult +LinearEncodingAttr::verify(function_ref emitError, + LinearLayout linearLayout) { + // Example of LinearEncodingAttr + // <{register = [[0, 1], [8, 0], [0, 8], [64, 0]], + // lane = [[0, 2], [0, 4], [1, 0], [2, 0], [4, 0]], + // warp = [[16, 0], [32, 0]], + // block = []}> + // The input dims must be {register, lane, warp, block} + // The output dims of the linear layout should be dim0..dim[rank-1] + + static const auto expectedInDims = + SmallVector({"register", "lane", "warp", "block"}); + for (const auto &[i, dims] : llvm::enumerate( + llvm::zip(linearLayout.getInDimNames(), expectedInDims))) { + const auto &[dim, expectedDimStr] = dims; + if (dim.str() != expectedDimStr) { + return emitError() << "Expected input dimension " << i << " to be '" + << expectedDimStr << "'. Got " << dim; + } + } + + // outDims are ['dim0', 'dim1', ...] + for (auto [i, dim] : llvm::enumerate(linearLayout.getOutDimNames())) { + if (dim.str() != ("dim" + llvm::Twine(i)).str()) { + return emitError() + << "Expected output dimensions to be ['dim0', 'dim1', ...]. Got " + << dim << " at position " << i; + } + } + + const auto &bases = linearLayout.getBases(); + auto nonZero = [](auto val) { return val != 0; }; + for (const auto &dimBases : llvm::make_second_range(bases)) { + if (!llvm::all_of(dimBases, [&](const auto &basis) { + return std::count_if(basis.begin(), basis.end(), nonZero) <= 1; + })) { + return emitError() + << "In a distributed layout, each base must move in at most one " + "dimension."; + } + } + + return success(); +} + +// If we only had BlockedEncodingAttr, we could simply return ArrayRefs here. +// But we need to have a consistent interface with e.g. SliceEncodingAttr, which +// computes some of these fields. +SmallVector BlockedEncodingAttr::getRepOrder() const { + return SmallVector(getOrder()); +} + +//===----------------------------------------------------------------------===// +// Linear Encoding +//===----------------------------------------------------------------------===// + +void LinearEncodingAttr::print(mlir::AsmPrinter &printer) const { + printer << "<{"; + printLinearLayout(printer, getLinearLayout()); + printer << "}>"; +} + +Attribute LinearEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + + if (parser.parseGreater().failed()) + return {}; + + std::vector inDimNames = {"register", "lane", "warp", "block"}; + auto maybeLL = parseLinearLayout(dict, parser, inDimNames); + if (!maybeLL.has_value()) + return {}; + + // Create and return the LinearEncodingAttr + return parser.getChecked(parser.getContext(), + std::move(*maybeLL)); +} + +static SmallVector +basesPerDimImpl(const LinearLayout::BasesT &namedBases, StringAttr dimName, + size_t rank, bool skipBroadcast) { + const auto &bases = namedBases.find(dimName)->second; + + if (bases.empty()) { + return SmallVector(rank, 1); + } + + SmallVector ret(rank, 1); + auto nonZero = [](auto val) { return val != 0; }; + int nonZeroIdx = 0; + for (const auto &basis : bases) { + auto it = std::find_if(basis.begin(), basis.end(), nonZero); + // Bases can have one or zero non-zero elements + // Skip a basis if it's broadcasting (all zeros) + // e.g. warps for DotOperandEncodingAttr (see ampereDotToLinearLayout) + if (it != basis.end()) { + nonZeroIdx = it - basis.begin(); + ret[nonZeroIdx] *= 2; + } else if (!skipBroadcast) { + // If we've seen a non-zero basis, we double the size of the previous dim + // This is just needed to count the CTAsPerCGA + ret[nonZeroIdx] *= 2; + } + } + return ret; +} + +SmallVector +LinearEncodingAttr::basesPerDim(StringAttr dimName, bool skipBroadcast) const { + auto ll = getLinearLayout(); + auto rank = ll.getNumOutDims(); + return basesPerDimImpl(ll.getBases(), dimName, rank, skipBroadcast); +} + +CTAEncodingAttr linearToCTAEncodingAttr(const LinearLayout &ll, + ArrayRef cgaLogicalShape) { + // Compute the shapePerCTA + auto shape = ll.getOutDims(); + for (int i = 0; i < shape.size(); ++i) { + shape[i].second /= cgaLogicalShape[i]; + } + auto inDims = to_vector(ll.getInDimNames()); + auto kBlock = inDims.back(); + assert(kBlock.str() == "block"); + inDims.pop_back(); + auto outDims = to_vector(ll.getOutDimNames()); + auto subLl = ll.sublayout(inDims, outDims); + // sublayout returns the same output size. We trim it to the + // real size + subLl = LinearLayout(subLl.getBases(), shape, false); + // The ctaLayout is what we get after dividing on the left by + // the layout in a single CTA + auto maybeCtaLayout = divideLeft(ll, subLl); + assert(maybeCtaLayout.has_value()); + auto *ctx = inDims[0].getContext(); + auto ctaLayout = maybeCtaLayout->sublayout({kBlock}, outDims); + return CTAEncodingAttr::get(ctx, std::move(ctaLayout)); +} + +SmallVector +LinearEncodingAttr::orderPerDim(StringAttr dimName, + ArrayRef defaultOrder) const { + return orderPerDimImpl(getLinearLayout(), dimName, defaultOrder); +} + +// [Note. Divergence of methods wrt. legacy layouts] +// For smaller shapes where the CTATile is larger than the output +// tensor, some methods return different values than the legacy layouts. I think +// this is benign tho. An example: what is the vector of `warpsPerCTA` if +// all the warps hold the same data? I think it should be [1, 1], even if we +// have 4 warps. But perhaps for this we have to add some masking in some +// places... We'll see +SmallVector LinearEncodingAttr::getRepOrder() const { + // This is not correct, but: + // - It happens to agree in most places with the legacy layout + // - getRepOrder does not make sense for LinearEncodingAttr as it already has + // the same shape as the tensor that uses it + return getOrder(); +} + +CTAEncodingAttr LinearEncodingAttr::getCTALayout() const { + auto splitNum = basesPerDim(StringAttr::get(getContext(), "block")); + return linearToCTAEncodingAttr(getLinearLayout(), splitNum); +} +SmallVector LinearEncodingAttr::getWarpsPerCTA() const { + return basesPerDim(StringAttr::get(getContext(), "warp")); +} +SmallVector LinearEncodingAttr::getWarpOrder() const { + return orderPerDim(StringAttr::get(getContext(), "warp"), getOrder()); +} +SmallVector LinearEncodingAttr::getThreadsPerWarp() const { + return basesPerDim(StringAttr::get(getContext(), "lane")); +} +SmallVector LinearEncodingAttr::getThreadOrder() const { + return orderPerDim(StringAttr::get(getContext(), "lane"), getOrder()); +} + +SmallVector LinearEncodingAttr::getSizePerThread() const { + auto rank = getOrder().size(); + auto ll = getLinearLayout(); + auto ctx = getContext(); + auto kRegister = StringAttr::get(ctx, "register"); + auto splitNum = getCTALayout().getCTASplitNum(); + + // We canonicalize on the spot, as if we use CGAs the regs are not in + // canonical form The order is [reg, lane, warp, rep, block], so we first + // remove the blocks + llvm::SmallVector ctaShape; + for (auto [shape, cgaNum] : llvm::zip(ll.getOutDimSizes(), splitNum)) { + ctaShape.push_back(shape / cgaNum); + } + LinearLayout::BasesT bases = ll.getBases(); + + llvm::SetVector reverseRepOrder; + auto nonZero = [](auto val) { return val != 0; }; + auto ®isters = bases[kRegister]; + while (!registers.empty()) { + auto &basis = registers.back(); + auto it = std::find_if(basis.begin(), basis.end(), nonZero); + // If there's broadcasting (base == zeros) there are no more reps + if (it == basis.end()) { + break; + } + auto dim = it - basis.begin(); + reverseRepOrder.insert(dim); + // As soon as we stop finding reps, we stop + if (dim != reverseRepOrder.back() || 2 * basis[dim] != ctaShape[dim]) { + break; + } + ctaShape[dim] /= 2; + registers.pop_back(); + } + return basesPerDimImpl(bases, kRegister, rank); +} + +SmallVector LinearEncodingAttr::getOrder() const { + auto rank = getLinearLayout().getNumOutDims(); + SmallVector order(rank); + // Choose [rank-1, rank-2, ... 0] as the default order in case + // there are dims that do not move in the register + // This order is as good as any really + std::iota(order.rbegin(), order.rend(), 0); + + return orderPerDim(StringAttr::get(getContext(), "register"), order); +} + +LinearLayout LinearEncodingAttr::toLinearLayout(ArrayRef shape) const { + auto ll = getLinearLayout(); + auto canonicalDims = llvm::to_vector(ll.getOutDimNames()); + llvm::SmallDenseMap namedShape; + llvm::SmallVector permutedDims; + for (auto dim : getRepOrder()) { + permutedDims.push_back(canonicalDims[dim]); + namedShape[canonicalDims[dim]] = shape[dim]; + } + ll = ll.transposeOuts(permutedDims); + ll = ensureLayoutNotSmallerThan(ll, namedShape); + ll = ensureLayoutNotLargerThan(ll, namedShape, /*broadcastRegisters=*/false); + ll = ll.transposeOuts(canonicalDims); + return ll; +} + +SmallVector +LinearEncodingAttr::getElemsPerThread(ArrayRef shape) const { + // When broadcasting the layout the shape changes, otherwise the shape is + // the same as the shape of the tensor + // We can either have BroadcastOp with SameOperandsAndResultEncoding, or keep + // the invariant that the shape of the LL is that of the tensor + // We choose the former for BC + auto scaledLayout = get(getContext(), toLinearLayout(shape)); + auto kRegister = StringAttr::get(getContext(), "register"); + return scaledLayout.basesPerDim(kRegister, /*skipBroadcast=*/false); +} + +SmallVector +LinearEncodingAttr::getContig(const char *inDim, + SmallVector lowerContig) const { + auto ll = getLinearLayout(); + const auto &bases = + ll.getBases().find(StringAttr::get(getContext(), inDim))->second; + auto order = getOrder(); + auto rank = order.size(); + + SmallVector contig(lowerContig); + auto basisIt = bases.begin(); + for (unsigned dim : order) { + std::vector basis(rank, 0); + basis[dim] = contig[dim]; + + while (basisIt != bases.end() && *basisIt == basis) { + contig[dim] *= 2; + basis[dim] *= 2; + ++basisIt; + } + } + return contig; +} + +SmallVector LinearEncodingAttr::getContigPerThread() const { + SmallVector contig(getOrder().size(), 1); + return getContig("register", contig); +} + +SmallVector LinearEncodingAttr::getContigPerWarp() const { + return getContig("lane", getContigPerThread()); +} + +unsigned +LinearEncodingAttr::getTotalElemsPerThread(ArrayRef shape) const { + return product(getElemsPerThread(shape)); +} + +//===----------------------------------------------------------------------===// +// MMA encoding +//===----------------------------------------------------------------------===// + +Attribute NvidiaMmaEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + unsigned versionMajor = 0; + unsigned versionMinor = 0; + SmallVector warpsPerCTA; + SmallVector instrShape; + Attribute ctaAttr = nullptr; + + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "versionMajor") { + if (parseUInt(parser, attr, versionMajor, "versionMajor").failed()) + return {}; + } + if (attr.getName() == "versionMinor") { + if (parseUInt(parser, attr, versionMinor, "versionMinor").failed()) + return {}; + } + if (attr.getName() == "warpsPerCTA") { + if (parseIntArrayAttr(parser, attr, warpsPerCTA, "warpsPerCTA").failed()) + return {}; + } + if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + continue; + } + if (attr.getName() == "instrShape") { + if (parseIntArrayAttr(parser, attr, instrShape, "instrShape").failed()) { + return {}; + } + } + } + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, /*rank=*/warpsPerCTA.size()); + if (!CTALayout.has_value()) + return {}; + + return parser.getChecked( + parser.getContext(), versionMajor, versionMinor, warpsPerCTA, *CTALayout, + instrShape); +} + +void NvidiaMmaEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "versionMajor = " << getVersionMajor() + << ", versionMinor = " << getVersionMinor() // + << ", warpsPerCTA = [" << ArrayRef(getWarpsPerCTA()) << "]"; + + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getRank()); + + printer << ", instrShape = [" << getInstrShape() << "]}>"; +} + +//===----------------------------------------------------------------------===// +// MFMA encoding +//===----------------------------------------------------------------------===// + +Attribute AMDMfmaEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + unsigned version = 0; + SmallVector warpsPerCTA; + SmallVector instrShape; + bool isTransposed; + SmallVector tilesPerWarp = {}; + unsigned elementBitWidth = 32; + Attribute ctaAttr = nullptr; + + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "version") { + if (parseUInt(parser, attr, version, "version").failed()) + return {}; + } + if (attr.getName() == "warpsPerCTA") { + if (parseIntArrayAttr(parser, attr, warpsPerCTA, "warpsPerCTA").failed()) + return {}; + } + if (attr.getName() == "instrShape") { + if (parseIntArrayAttr(parser, attr, instrShape, "instrShape").failed()) + return {}; + } + if (attr.getName() == "isTransposed") { + if (parseBool(parser, attr, isTransposed, "isTransposed").failed()) + return {}; + } + if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + continue; + } + if (attr.getName() == "tilesPerWarp") { + if (parseIntArrayAttr(parser, attr, tilesPerWarp, "tilesPerWarp") + .failed()) + return {}; + } + if (attr.getName() == "elementBitWidth") { + if (parseUInt(parser, attr, elementBitWidth, "elementBitWidth").failed()) + return {}; + } + } + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, /*rank=*/warpsPerCTA.size()); + if (!CTALayout.has_value()) + return {}; + + if (tilesPerWarp.empty()) + tilesPerWarp = SmallVector(instrShape.size(), 1); + + return parser.getChecked( + parser.getContext(), version, warpsPerCTA, instrShape, isTransposed, + *CTALayout, tilesPerWarp, elementBitWidth); +} + +void AMDMfmaEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "version = " << getVersion() // + << ", warpsPerCTA = [" << getWarpsPerCTA() << "]" // + << ", instrShape = [" << getInstrShape() << "]"; + + printer << ", isTransposed = " << getIsTransposed(); + + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getRank()); + + auto tilesPerWarp = getTilesPerWarp(); + if (!hasUnitTilesPerWarp()) + printer << ", tilesPerWarp = [" << getTilesPerWarp() << "]"; + + auto elementBitWidth = getElementBitWidth(); + if (elementBitWidth != 32) + printer << ", elementBitWidth = " << elementBitWidth; + + printer << "}>"; +} + +LogicalResult AMDMfmaEncodingAttr::verify( + function_ref emitError, unsigned version, + llvm::ArrayRef warpsPerCTA, + llvm::ArrayRef instrShape, bool isTransposed, + mlir::triton::gpu::CTAEncodingAttr, + llvm::ArrayRef tilesPerWarp, unsigned elementBitWidth) { + if (!(version >= 0 && version <= 4)) { + return emitError() << "version must be in the [0, 4] range"; + } + + auto mDim = instrShape[0]; + auto nDim = instrShape[1]; + const std::array, 4> validDims = { + {{32, 32}, {16, 16}, {64, 4}, {4, 64}}}; + if (!llvm::is_contained(validDims, std::make_pair(mDim, nDim))) { + return emitError() << "invalid (mDim, nDim) combination: (" << mDim << ", " + << nDim << ")"; + } + + if (!(elementBitWidth == 32 || elementBitWidth == 64)) + return emitError() << "elementBitWidth must be 32 or 64"; + + return success(); +} + +//===----------------------------------------------------------------------===// +// WMMA encoding +//===----------------------------------------------------------------------===// +bool AMDWmmaEncodingAttr::hasUnitTilesPerWarp() const { + return llvm::all_of(getTilesPerWarp(), [](int x) { return x == 1; }); +} + +Attribute AMDWmmaEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + unsigned version = 0; + bool isTransposed = false; + SmallVector warpsPerCTA; + SmallVector tilesPerWarp = {}; + SmallVector instrShape = getDefaultInstrShape(); + Attribute ctaAttr = nullptr; + + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "version") { + if (parseUInt(parser, attr, version, "version").failed()) + return {}; + } + if (attr.getName() == "isTranspose") { + if (parseBool(parser, attr, isTransposed, "isTranspose").failed()) + return {}; + } + if (attr.getName() == "warpsPerCTA") { + if (parseIntArrayAttr(parser, attr, warpsPerCTA, "warpsPerCTA").failed()) + return {}; + } + if (attr.getName() == "tilesPerWarp") { + if (parseIntArrayAttr(parser, attr, tilesPerWarp, "tilesPerWarp") + .failed()) + return {}; + } + if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + continue; + } + if (attr.getName() == "instrShape") { + instrShape.clear(); + if (parseIntArrayAttr(parser, attr, instrShape, "instrShape").failed()) { + return {}; + } + } + } + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, /*rank=*/warpsPerCTA.size()); + if (!CTALayout.has_value()) + return {}; + + if (tilesPerWarp.empty()) + tilesPerWarp = SmallVector(instrShape.size(), 1); + + return parser.getChecked( + parser.getContext(), version, isTransposed, warpsPerCTA, tilesPerWarp, + *CTALayout, instrShape); +} + +void AMDWmmaEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "version = " << getVersion() + << ", isTranspose = " << getIsTransposed() // + << ", warpsPerCTA = [" << ArrayRef(getWarpsPerCTA()) << "]"; + + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getWarpsPerCTA().size()); + + auto tilesPerWarp = getTilesPerWarp(); + if (!hasUnitTilesPerWarp()) + printer << ", tilesPerWarp = [" << getTilesPerWarp() << "]"; + + if (getInstrShape() != ArrayRef(getDefaultInstrShape())) { + printer << ", instrShape = [" << getInstrShape() << "]"; + } + printer << "}>"; +} + +LogicalResult AMDWmmaEncodingAttr::verify( + function_ref emitError, unsigned version, + bool isTransposed, llvm::ArrayRef warpsPerCTA, + llvm::ArrayRef tilesPerWarp, CTAEncodingAttr ctaLayout, + llvm::ArrayRef instrShape) { + if (!(version >= 1 && version <= 3)) + return emitError() << "WMMA version must be in the [1, 3] range"; + + auto shape = SmallVector(instrShape); + auto validShapesV1 = std::vector>{{16, 16, 16}}; + if (version == 1 && !llvm::is_contained(validShapesV1, shape)) + return emitError() << "invalid WMMA v1 instruction shape"; + + auto validShapesV2 = + std::vector>{{16, 16, 16}, {16, 16, 32}}; + if (version == 2 && !llvm::is_contained(validShapesV2, shape)) + return emitError() << "invalid WMMA v2 instruction shape"; + + auto validShapesV3 = std::vector>{ + {16, 16, 4}, {16, 16, 32}, {16, 16, 64}, {16, 16, 128}}; + if (version == 3 && !llvm::is_contained(validShapesV3, shape)) + return emitError() << "invalid WMMA v3 instruction shape"; + + return success(); +} + +//===----------------------------------------------------------------------===// +// Sunrise_MMA Encoding +//===----------------------------------------------------------------------===// + +Attribute SunriseMmaEncodingAttr::parse(AsmParser &parser, Type type) { + DictionaryAttr dict; + if (parser.parseLess().failed()) { return {}; } + if (parser.parseAttribute(dict).failed()) { return {};} + if (parser.parseGreater().failed()) {return {};} + + unsigned versionMajor = 0; + unsigned versionMinor = 0; + + SmallVector warpsPerCTA; + SunriseMmaEncodingAttr::TMMAOutLayout outLayout; + unsigned outLayoutUint = static_cast(SunriseMmaEncodingAttr::TMMAOutLayout::NotAvailable); + unsigned inputElemBitWidth = 0; + unsigned outputElemBitWidth = 0; + bool isACol = false, isBCol = false; + Attribute ctaAttr = nullptr; + + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "versionMajor") { if (parseUInt(parser, attr, versionMajor, "versionMajor").failed()) {return {};} } + if (attr.getName() == "versionMinor") { if (parseUInt(parser, attr, versionMinor, "versionMinor").failed()) {return {};} } + if (attr.getName() == "warpsPerCTA") { if (parseIntArrayAttr(parser, attr, warpsPerCTA, "warpsPerCTA").failed()) {return {};} } + if (attr.getName() == "CGALayout") { ctaAttr = attr.getValue(); continue; } + if (attr.getName() == "outLayout") { if(parseUInt(parser, attr, outLayoutUint, "outLayout").failed()) {return {}; } } + if (attr.getName() == "inputElemBitWidth") { if(parseUInt(parser, attr, inputElemBitWidth, "inputElemBitWidth").failed()) {return {}; } } + if (attr.getName() == "outputElemBitWidth") { if(parseUInt(parser, attr, outputElemBitWidth, "outputElemBitWidth").failed()) {return {}; } } + if (attr.getName() == "isACol") { if(parseBool(parser, attr, isACol, "isACol").failed()) {return {}; } } + if (attr.getName() == "isBCol") { if(parseBool(parser, attr, isBCol, "isBCol").failed()) {return {}; } } + } + outLayout = static_cast(outLayoutUint); + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, /*rank=*/warpsPerCTA.size()); + if (!CTALayout.has_value()) + return {}; + + return parser.getChecked( + parser.getContext(), versionMajor, versionMinor, warpsPerCTA, *CTALayout, outLayout, inputElemBitWidth, outputElemBitWidth, isACol, isBCol); +} + +void SunriseMmaEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "versionMajor = " << getVersionMajor() + << ", versionMinor = " << getVersionMinor() + << ", warpsPerCTA = [" << ArrayRef(getWarpsPerCTA()) << "]" + << ", isACol = " << getIsACol() + << ", isBCol = " << getIsBCol() + << ", inputElemBitWidth = " << getInputElemBitWidth() + << ", outputElemBitWidth = " << getOutputElemBitWidth(); + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getWarpsPerCTA().size()); + + // printer << ", order = " << sv2str(getOrder(*this)) + // << ", threadsPerWarp = " << sv2str(getThreadsPerWarp()) + // << ", threadOrder = " << sv2str(getThreadOrder()) + // << ", sizePerThread = " << sv2str(getSizePerThread()) + // << ", shapePerCTATile = " << sv2str(getShapePerCTATile()) + // << ", outLayout = " << static_cast(getOutLayout()) + + printer << "}>"; +} + + +//===----------------------------------------------------------------------===// +// Sliced Encoding +//===----------------------------------------------------------------------===// + +Attribute SliceEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + NamedAttrList attrs; + if (parser.parseOptionalAttrDict(attrs).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + unsigned dim = mlir::cast(attrs.get("dim")).getInt(); + auto parent = mlir::dyn_cast(attrs.get("parent")); + if (!parent) { + parser.emitError(parser.getNameLoc(), + "expected a distributed encoding trait"); + return {}; + } + return parser.getChecked(parser.getContext(), dim, parent); +} + +void SliceEncodingAttr::print(mlir::AsmPrinter &printer) const { + printer << "<{" + << "dim = " << getDim() << ", " + << "parent = " << getParent() << "}>"; +} + +LogicalResult +SliceEncodingAttr::verify(function_ref emitError, + unsigned dim, DistributedEncodingTrait parent) { + unsigned rank = ::getCTALayout(parent).getRank(); + if (rank <= 1) + return emitError() << "parent layout must have at least rank >= 2"; + if (dim >= rank) { + return emitError() << "slice dim=" << dim + << " must be less than the parent rank=" << rank; + } + return success(); +} + +SmallVector SliceEncodingAttr::getRepOrder() const { + auto parentRepOrder = getParent().getRepOrder(); + return eraseOrder(parentRepOrder, getDim()); +} + +CTAEncodingAttr SliceEncodingAttr::getCTALayout() const { + auto layout = ::getCTALayout(getParent()).getLinearLayout(); + layout = removeStandardDim(layout, getDim()); + return CTAEncodingAttr::get(getContext(), layout); +} + +template +SmallVector SliceEncodingAttr::paddedShape(ArrayRef shape) const { + size_t rank = shape.size(); + unsigned dim = getDim(); + SmallVector retShape(rank + 1); + for (unsigned d = 0; d < rank + 1; ++d) { + if (d < dim) + retShape[d] = shape[d]; + else if (d == dim) + retShape[d] = 1; + else + retShape[d] = shape[d - 1]; + } + return retShape; +} +template SmallVector +SliceEncodingAttr::paddedShape(ArrayRef shape) const; +template SmallVector +SliceEncodingAttr::paddedShape(ArrayRef shape) const; + +template +Attribute parseSwizzledEncoding(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + // Parse the data as a dictionary + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + unsigned vec = 0; + unsigned perPhase = 0; + unsigned maxPhase = 0; + SmallVector order; + Attribute ctaAttr = nullptr; + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "vec") { + if (parseUInt(parser, attr, vec, "vec").failed()) + return {}; + } else if (attr.getName() == "perPhase") { + if (parseUInt(parser, attr, perPhase, "perPhase").failed()) + return {}; + } else if (attr.getName() == "maxPhase") { + if (parseUInt(parser, attr, maxPhase, "maxPhase").failed()) + return {}; + } else if (attr.getName() == "order") { + if (parseIntArrayAttr(parser, attr, order, "order").failed()) + return {}; + } else { + if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + } else { + parser.emitError(parser.getNameLoc(), "unexpected key: ") + << attr.getName().strref(); + return {}; + } + } + } + + if (auto CTALayout = parseCTAAttr(parser, ctaAttr, order.size())) + return parser.getChecked( + parser.getContext(), vec, perPhase, maxPhase, order, *CTALayout); + return {}; +} + +//===----------------------------------------------------------------------===// +// SwizzledShared encoding +//===----------------------------------------------------------------------===// + +LogicalResult +SwizzledSharedEncodingAttr::verify(function_ref emitError, + unsigned vec, unsigned perPhase, + unsigned maxPhase, ArrayRef order, + CTAEncodingAttr ctaLayout) { + if (order.size() != ctaLayout.getRank()) { + return emitError() << "order size (" << order.size() + << ") must match CTALayout rank (" << ctaLayout.getRank() + << ")"; + } + return verifyLayoutOrder(emitError, order); +} + +Attribute SwizzledSharedEncodingAttr::parse(AsmParser &parser, Type type) { + return parseSwizzledEncoding(parser, type); +} + +void SwizzledSharedEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "vec = " << getVec() // + << ", perPhase = " << getPerPhase() + << ", maxPhase = " << getMaxPhase() // + << ", order = [" << getOrder() << "]"; + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getOrder().size()); + printer << "}>"; +} + +//===----------------------------------------------------------------------===// +// SharedLinear encoding +//===----------------------------------------------------------------------===// + +LogicalResult +SharedLinearEncodingAttr::verify(function_ref emitError, + LinearLayout linearLayout, + unsigned layoutAlignment) { + if (layoutAlignment == 0 || !llvm::isPowerOf2_32(layoutAlignment)) { + return emitError() << "alignment must be a positive power of two"; + } + static const auto expectedInDims = + SmallVector({"offset", "block"}); + for (const auto &[index, dims] : llvm::enumerate( + llvm::zip(linearLayout.getInDimNames(), expectedInDims))) { + const auto &[dim, expected] = dims; + if (dim.str() != expected) { + return emitError() << "Expected input dimension " << index << " to be '" + << expected << "'. Got " << dim; + } + } + + for (auto [i, dim] : llvm::enumerate(linearLayout.getOutDimNames())) { + if (dim.str() != ("dim" + llvm::Twine(i)).str()) { + return emitError() + << "Expected output dimensions to be ['dim0', 'dim1', ...]. Got " + << dim << " at position " << i; + } + } + + SmallVector outDimNames = + llvm::to_vector(linearLayout.getOutDimNames()); + if (outDimNames.empty()) { + return emitError() + << "SharedLinearEncodingAttr requires at least one output" + " dimension."; + } + + auto *ctx = outDimNames.front().getContext(); + auto kOffset = StringAttr::get(ctx, "offset"); + auto kBlock = StringAttr::get(ctx, "block"); + + if (!linearLayout.isSurjective()) { + return emitError() << "The layout must be surjective"; + } + + LinearLayout withoutBroadcast = + linearLayout.removeZeroBasesAlongDim(kOffset).removeZeroBasesAlongDim( + kBlock); + if (!withoutBroadcast.isInvertible()) { + return emitError() + << "After removing the zero bases the layout must be bijective"; + } + + return success(); +} + +void SharedLinearEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{"; + auto layout = getLinearLayout(); + auto kBlock = StringAttr::get(getContext(), "block"); + auto kOffset = StringAttr::get(getContext(), "offset"); + if (layout.getBases().lookup(kBlock).empty()) { + layout = + layout.sublayout({kOffset}, llvm::to_vector(layout.getOutDimNames())); + } + printLinearLayout(printer, layout); + printer << "}, alignment = " << getAlignment() << ">"; +} + +Attribute SharedLinearEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + + DictionaryAttr layoutDictRaw; + if (parser.parseAttribute(layoutDictRaw).failed()) + return {}; + + if (layoutDictRaw.get("alignment")) { + parser.emitError(parser.getCurrentLocation()) + << "alignment must be specified outside of the linear layout braces"; + return {}; + } + + NamedAttrList layoutAttrList(layoutDictRaw.getValue()); + auto *ctx = parser.getContext(); + auto kBlock = StringAttr::get(ctx, "block"); + if (!layoutAttrList.get(kBlock)) { + layoutAttrList.push_back({kBlock, ArrayAttr::get(ctx, {})}); + } + + DictionaryAttr layoutDict = layoutAttrList.getDictionary(ctx); + + // Parse alignment + unsigned layoutAlignment; + if (parser.parseComma().failed()) + return {}; + if (parser.parseKeyword("alignment").failed() || parser.parseEqual().failed()) + return {}; + if (parser.parseInteger(layoutAlignment).failed()) + return {}; + + if (parser.parseGreater().failed()) + return {}; + + std::vector inDimNames = {"offset", "block"}; + auto maybeLL = parseLinearLayout(layoutDict, parser, inDimNames); + if (!maybeLL.has_value()) + return {}; + + // Special case for cleaner errors + if (layoutDict.get("alignment")) { + parser.emitError(parser.getCurrentLocation()) + << "alignment must be specified outside of the linear layout braces"; + return {}; + } + + if (layoutDict.size() != 2) { + parser.emitError(parser.getCurrentLocation()) + << "SharedLinearEncodingAttr must have exactly two attributes: offset " + "and block"; + return {}; + } + + return parser.getChecked( + parser.getContext(), std::move(*maybeLL), layoutAlignment); +} + +SmallVector +SharedLinearEncodingAttr::basesPerDim(StringAttr dimName, + bool skipBroadcast) const { + auto ll = getLinearLayout(); + auto rank = ll.getNumOutDims(); + return basesPerDimImpl(ll.getBases(), dimName, rank, skipBroadcast); +} + +SmallVector +SharedLinearEncodingAttr::orderPerDim(StringAttr dimName, + ArrayRef defaultOrder) const { + return orderPerDimImpl(getLinearLayout(), dimName, defaultOrder); +} + +SmallVector SharedLinearEncodingAttr::getOrder() const { + auto ll = getLinearLayout(); + auto rank = ll.getNumOutDims(); + SmallVector defaultOrder(rank); + std::iota(defaultOrder.rbegin(), defaultOrder.rend(), 0); + return orderPerDim(StringAttr::get(getContext(), "offset"), defaultOrder); +} + +CTAEncodingAttr SharedLinearEncodingAttr::getCTALayout() const { + auto splitNum = basesPerDim(StringAttr::get(getContext(), "block")); + return linearToCTAEncodingAttr(getLinearLayout(), splitNum); +} +LinearLayout +SharedLinearEncodingAttr::toLinearLayout(ArrayRef shape) const { + auto ll = getLinearLayout(); + auto outDimNames = llvm::to_vector(ll.getOutDimNames()); + assert(shape.size() == outDimNames.size()); + // We don't support automatic broadcasting for shared linear layouts + for (auto [size, llSize] : llvm::zip(shape, ll.getOutDimSizes())) { + assert(size == llSize); + } + return ll; +} + +//===----------------------------------------------------------------------===// +// PaddedShared encoding +//===----------------------------------------------------------------------===// + +Attribute PaddedSharedEncodingAttr::parse(AsmParser &parser, Type type) { + // <[ + if (failed(parser.parseLess()) || failed(parser.parseLSquare())) + return {}; + + // :+ + SmallVector intervals, paddings; + auto parseIntervalPaddingPair = [&]() { + unsigned interval = 0, padding = 0; + if (failed(parser.parseInteger(interval)) || failed(parser.parseColon()) || + failed(parser.parsePlus()) || failed(parser.parseInteger(padding))) + return failure(); + intervals.push_back(interval); + paddings.push_back(padding); + return success(); + }; + // ] + if (failed(parser.parseCommaSeparatedList(parseIntervalPaddingPair)) || + failed(parser.parseRSquare())) + return {}; + + // {} + auto attrList = DictionaryAttr::get(parser.getContext()); + if (failed(parser.parseAttribute(attrList))) + return {}; + + // We have 2 possible formats for the attr-dict: + // 1) offset=[..], block=[..] handled by parseLinearLayout + // 2) order=[..], shape=[..] which creates an identity mapping + + std::optional maybeLL; + // Assume it's the first variant if offset or block is defined + if (attrList.contains("offset") || attrList.contains("block")) { + std::vector inDimNames = {"offset", "block"}; + // Error out on additional attribute names + for (const NamedAttribute &attr : attrList) { + if (!llvm::is_contained(inDimNames, attr.getName())) { + parser.emitError(parser.getCurrentLocation(), "Unexpected attribute ") + << attr.getName() << " found"; + } + } + maybeLL = parseLinearLayout(attrList, parser, inDimNames); + } else { + // Parse the second form + SmallVector order; + SmallVector shape; + for (const NamedAttribute &attr : attrList) { + if (attr.getName() == "order") { + if (parseIntArrayAttr(parser, attr, order, "order").failed()) + return {}; + } else if (attr.getName() == "shape") { + if (parseIntArrayAttr(parser, attr, shape, "shape").failed()) + return {}; + } else { + parser.emitError(parser.getCurrentLocation(), "Unexpected attribute ") + << attr.getName() << " found"; + return {}; + } + } + + if (order.size() != shape.size()) { + parser.emitError(parser.getCurrentLocation(), + "Mismatch of shape and order ranks in padded layout"); + return {}; + } + + // Create identity mapping based on shape and order + auto kOffset = StringAttr::get(parser.getContext(), "offset"); + maybeLL = identityStandardND(kOffset, shape, order); + maybeLL = combineCtaCgaWithShape( + *maybeLL, + CTAEncodingAttr::getDefault(parser.getContext(), shape.size()), + SmallVector(ArrayRef(shape))); + } + + if (!maybeLL.has_value()) + return {}; + + // > + if (parser.parseGreater().failed()) + return {}; + + return parser.getChecked( + parser.getContext(), intervals, paddings, *maybeLL); +} + +void PaddedSharedEncodingAttr::print(AsmPrinter &printer) const { + + auto *ctx = getContext(); + const auto &ll = getLinearComponent(); + + printer << "<["; + llvm::interleaveComma(llvm::zip(getIntervals(), getPaddings()), printer, + [&](std::tuple intervalPad) { + printer << std::get<0>(intervalPad) << ":+" + << std::get<1>(intervalPad); + }); + printer << "] {"; + + // We have a short hand form if linearComponent: + // 1) does have an empty CTA layout (empty block dim) + // 2) offsets are an identity mapping + auto kOffset = StringAttr::get(ctx, "offset"); + auto kBlock = StringAttr::get(ctx, "block"); + auto shape = SmallVector(ll.getOutDimSizes()); + + bool hasEmptyBlock = ll.getInDimSizeLog2(kBlock) == 0; + + LinearLayout identity = identityStandardND(kOffset, shape, getOrder()) + .transposeOuts(to_vector(ll.getOutDimNames())); + auto offsetLayout = ll.sublayout({kOffset}, to_vector(ll.getOutDimNames())); + + if (hasEmptyBlock && offsetLayout == identity) { + printer << "order = [" << ArrayRef(getOrder()) << "], shape = [" + << ArrayRef(shape) << "]"; + } else { + printLinearLayout(printer, getLinearComponent()); + } + + printer << "}>"; +} + +LogicalResult PaddedSharedEncodingAttr::verify( + function_ref emitError, ArrayRef intervals, + ArrayRef paddings, LinearLayout linearComponent) { + if (intervals.size() != paddings.size()) + return emitError() << "intervals size (" << intervals.size() + << ") must match paddings size (" << paddings.size() + << ")"; + + if (intervals.empty()) + return emitError() << "must have at least one interval-padding pair"; + + if (!llvm::all_of(intervals, llvm::isPowerOf2_32)) + return emitError() << "interval values must all be power of two"; + if (!llvm::all_of(paddings, llvm::isPowerOf2_32)) + return emitError() << "padding values must all be power of two"; + + llvm::SmallSet intervalValues(intervals.begin(), + intervals.end()); + if (intervalValues.size() != intervals.size()) + return emitError() << "interval values cannot have duplicates"; + + const auto &ll = linearComponent; + // The linear layout should map from [offset, block] to [dim0..dimN). All + // bases should be 0 or power of twos and move in a single direction without + // broadcasting + + if (ll == LinearLayout::empty()) + return emitError() << "linearComponent cannot be empty"; + + assert(!ll.getInDimNames().empty()); + auto *ctx = ll.getInDimNames().begin()->getContext(); + + if (!llvm::equal(ll.getInDimNames(), + std::array{StringAttr::get(ctx, "offset"), + StringAttr::get(ctx, "block")})) { + return emitError() + << "linearComponent must have [offset, block] as input dims"; + } + + if (!llvm::equal(ll.getOutDimNames(), + standardOutDimNames(ctx, ll.getNumOutDims()))) { + return emitError() + << "Expected output dimensions to be ['dim0', 'dim1', ...]."; + } + + const auto &bases = ll.getBases(); + + // Check that we are not broadcasting or having repeated bases + if (!ll.isInvertible()) { + return emitError() << "Broadcasting is not supported."; + } + + auto nonZero = [](auto val) { return val != 0; }; + for (const auto &dimBases : llvm::make_second_range(bases)) { + if (!llvm::all_of(dimBases, [&](const auto &basis) { + return llvm::count_if(basis, nonZero) <= 1; + })) { + return emitError() + << "Each offset basis must move in at most one dimension."; + } + // Ensure all non zero elements are a power of 2. Combined with the + // broadcast check above this prevents per element swizzling. The intent of + // the linear component is to rearrange whole rows or cache-line sized + // chunks of rows. + if (!llvm::all_of(dimBases, [&](const auto &basis) { + return llvm::all_of( + basis, [](auto v) { return v == 0 || llvm::isPowerOf2_32(v); }); + })) { + return emitError() << "Each offset basis must be 0 or a power of two."; + } + } + + return success(); +} + +PaddedSharedEncodingAttr PaddedSharedEncodingAttr::get( + MLIRContext *context, ArrayRef> intervalPads, + ArrayRef order, ArrayRef shape, + CTAEncodingAttr ctaLayout) { + auto outDimNames = standardOutDimNames(context, shape.size()); + StringAttr kOffset = StringAttr::get(context, "offset"); + + // Create identity mapping based on shape and order + LinearLayout linearComponent = + identityStandardND(kOffset, SmallVector(shape), order); + linearComponent = combineCtaCgaWithShape(linearComponent, ctaLayout, shape); + + return get(context, intervalPads, linearComponent); +} + +PaddedSharedEncodingAttr PaddedSharedEncodingAttr::get( + MLIRContext *context, ArrayRef> intervalPads, + LinearLayout linearComponent) { + SmallVector intervals, paddings; + intervals.reserve(intervalPads.size()); + paddings.reserve(intervalPads.size()); + for (auto [interval, padding] : intervalPads) { + intervals.push_back(interval); + paddings.push_back(padding); + } + return get(context, intervals, paddings, linearComponent); +} + +SmallVector +PaddedSharedEncodingAttr::basesPerDim(StringAttr dimName, + bool skipBroadcast) const { + const auto &ll = getLinearComponent(); + auto rank = ll.getNumOutDims(); + return basesPerDimImpl(ll.getBases(), dimName, rank, skipBroadcast); +} + +int64_t PaddedSharedEncodingAttr::getPaddedSize(ArrayRef shape) const { + int64_t unpaddedSize = product(shape); + int64_t paddingSize = 0; + for (auto [interval, padding] : + llvm::zip_equal(getIntervals(), getPaddings())) { + paddingSize += (unpaddedSize >> llvm::Log2_32(interval)) + << llvm::Log2_32(padding); + // There is no need for padding after the last element + if (unpaddedSize % interval == 0) + paddingSize -= padding; + } + return unpaddedSize + paddingSize; +} + +SmallVector +PaddedSharedEncodingAttr::orderPerDim(StringAttr dimName, + ArrayRef defaultOrder) const { + return orderPerDimImpl(getLinearComponent(), dimName, defaultOrder); +} + +SmallVector PaddedSharedEncodingAttr::getOrder() const { + auto rank = getLinearComponent().getNumOutDims(); + SmallVector order(rank); + // Choose [rank-1, rank-2, ... 0] as the default order in case + // there are dims that do not move in the offsets + std::iota(order.rbegin(), order.rend(), 0); + + return orderPerDim(StringAttr::get(getContext(), "offset"), order); +} + +CTAEncodingAttr PaddedSharedEncodingAttr::getCTALayout() const { + auto splitNum = basesPerDim(StringAttr::get(getContext(), "block")); + return linearToCTAEncodingAttr(getLinearComponent(), splitNum); +} +//===----------------------------------------------------------------------===// +// NVMMAShared encoding +//===----------------------------------------------------------------------===// + +Attribute NVMMASharedEncodingAttr::parse(AsmParser &parser, Type type) { + if (parser.parseLess().failed()) + return {}; + // Parse the data as a dictionary + DictionaryAttr dict; + if (parser.parseAttribute(dict).failed()) + return {}; + if (parser.parseGreater().failed()) + return {}; + + unsigned swizzlingByteWidth; + bool transposed = false; + bool fp4Padded = false; + unsigned elementBitWidth; + unsigned layoutRank = 2; + Attribute ctaAttr = nullptr; + for (const NamedAttribute &attr : dict) { + if (attr.getName() == "swizzlingByteWidth") { + if (parseUInt(parser, attr, swizzlingByteWidth, "swizzlingByteWidth") + .failed()) + return {}; + } else if (attr.getName() == "transposed") { + if (parseBool(parser, attr, transposed, "transposed").failed()) + return {}; + } else if (attr.getName() == "elementBitWidth") { + if (parseUInt(parser, attr, elementBitWidth, "elementBitWidth").failed()) + return {}; + } else if (attr.getName() == "fp4Padded") { + if (parseBool(parser, attr, fp4Padded, "fp4Padded").failed()) + return {}; + } else if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + } else if (attr.getName() == "rank") { + if (parseUInt(parser, attr, layoutRank, "rank").failed()) + return {}; + } else { + parser.emitError(parser.getNameLoc(), "unexpected key: ") + << attr.getName().strref(); + return {}; + } + } + + std::optional CTALayout = + parseCTAAttr(parser, ctaAttr, layoutRank); + if (!CTALayout.has_value()) + return {}; + + return parser.getChecked( + parser.getContext(), swizzlingByteWidth, transposed, elementBitWidth, + fp4Padded, *CTALayout); +} + +void NVMMASharedEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "swizzlingByteWidth = " << getSwizzlingByteWidth() // + << ", transposed = " << getTransposed() // + << ", elementBitWidth = " << getElementBitWidth(); + if (getFp4Padded()) { + // Print only in this case to reduce the noise for the more common case. + printer << ", fp4Padded = true"; + } + unsigned rank = getCTALayout().getCTAOrder().size(); + auto *ctx = getContext(); + auto defaultLayout = CTAEncodingAttr::getDefault(ctx, rank); + if (getCTALayout() == defaultLayout && rank != 2) { + printer << ", rank = " << rank; + } else { + maybePrintCTALayout(ctx, printer, getCTALayout(), rank); + } + printer << "}>"; +} + +int NVMMASharedEncodingAttr::getVec() const { + if (getSwizzlingByteWidth() == 0) + return 1; + return 128 / getElementBitWidth(); +} + +int NVMMASharedEncodingAttr::getPerPhase() const { + if (getSwizzlingByteWidth() == 0) + return 1; + return 128 / getSwizzlingByteWidth(); +} + +int NVMMASharedEncodingAttr::getMaxPhase() const { + if (getSwizzlingByteWidth() == 0) + return 1; + return getSwizzlingByteWidth() / 16; +} + +int32_t NVMMASharedEncodingAttr::getAlignment() const { + return 128 * getMaxPhase(); +} + +//===----------------------------------------------------------------------===// +// AMDRotatingShared encoding +//===----------------------------------------------------------------------===// + +Attribute AMDRotatingSharedEncodingAttr::parse(AsmParser &parser, Type type) { + return parseSwizzledEncoding(parser, type); +} + +void AMDRotatingSharedEncodingAttr::print(AsmPrinter &printer) const { + printer << "<{" + << "vec = " << getVec() // + << ", perPhase = " << getPerPhase() + << ", maxPhase = " << getMaxPhase() // + << ", order = [" << getOrder() << "]"; + maybePrintCTALayout(getContext(), printer, getCTALayout(), + /*rank=*/getOrder().size()); + printer << "}>"; +} + +//===----------------------------------------------------------------------===// +// Mfma encoding +//===----------------------------------------------------------------------===// +// TODO: there is a lot of common code with MmaEncoding here + +bool AMDMfmaEncodingAttr::hasUnitTilesPerWarp() const { + return llvm::all_of(getTilesPerWarp(), [](int x) { return x == 1; }); +} + +SmallVector +AMDMfmaEncodingAttr::getInstrShapeForOperand(int kWidth, int opIdx) const { + auto mnkDim = getInstrShape(); + unsigned mDim = mnkDim[0]; + unsigned nDim = mnkDim[1]; + assert((mDim == nDim) && (mDim == 32 || mDim == 16 || mDim == 4) || + (mDim == 64 && nDim == 4) || (mDim == 4 && nDim == 64)); + + constexpr int warpSize = 64; // MFMA is always based on the 64-wide warps. + int kGroups = warpSize / std::min(mDim, nDim); // for 64x4 and 4x64, + // kGroups = 16 + int64_t kDim = kWidth * kGroups; + + if (opIdx == 0) + return {mDim, kDim}; + else + assert(opIdx == 1); + return {kDim, nDim}; +} + +SmallVector AMDMfmaEncodingAttr::getRepOrder() const { + return getMatrixOrder(getRank(), /*rowMajor*/ true); +} + +SmallVector +AMDMfmaEncodingAttr::getRepOrderForOperand(int opIdx) const { + return getOrderForDotOperand(opIdx, getRank(), /*kContig*/ true); +} + +SmallVector +AMDMfmaEncodingAttr::getRepForOperand(ArrayRef operandShape, + int kWidth, int opIdx) const { + auto operandTileShape = getInstrShapeForOperand(kWidth, opIdx); + auto rank = operandShape.size(); + auto warpsPerCTA = getWarpsPerCTA(); + auto tilesPerWarp = getTilesPerWarp(); + + int numRepBatch = + rank == 3 ? std::max(1, operandShape[0] / warpsPerCTA[0]) : 1; + if (opIdx == 0) + return { + numRepBatch, + std::max(1, operandShape[rank - 2] / + (operandTileShape[0] * tilesPerWarp[rank - 2] * + warpsPerCTA[rank - 2])) * + tilesPerWarp[rank - 2], + std::max(1, operandShape[rank - 1] / operandTileShape[1])}; + else { + assert(opIdx == 1); + return { + numRepBatch, + std::max(1, operandShape[rank - 2] / operandTileShape[0]), + std::max(1, operandShape[rank - 1] / + (operandTileShape[1] * tilesPerWarp[rank - 1] * + warpsPerCTA[rank - 1])) * + tilesPerWarp[rank - 1]}; + } +} + +SwizzledSharedEncodingAttr AMDMfmaEncodingAttr::composeSharedLayoutForOperand( + CTAEncodingAttr ctaLayout, int operandIdx, ArrayRef operandShape, + ArrayRef sharedOrder, unsigned vectorSize, unsigned elemBitWidth, + bool needTrans) const { + int kDimIndex = operandIdx == 0 ? 1 : 0; + + // Disable swizzling for scales + if (operandIdx >= 2) { + return SwizzledSharedEncodingAttr::get(getContext(), 1, 1, 1, sharedOrder, + ctaLayout); + } + + if (needTrans) + kDimIndex = 1 - kDimIndex; + + bool isKContig = sharedOrder[0] == kDimIndex; + // GFX950 supports LDS transpose load instructions, so we need swizzling even + // when K dimension is not the contiguous dimension. + bool isGFX950 = getVersion() == 4; + bool swizzleNonKContig = + isGFX950 && (elemBitWidth == 8 || elemBitWidth == 16); + + if (!isKContig && !swizzleNonKContig) { + // Do not swizzle. In this case accesses will go in different banks even + // without swizzling. + return SwizzledSharedEncodingAttr::get(getContext(), 1, 1, 1, sharedOrder, + ctaLayout); + } + + const unsigned numBanks = isGFX950 ? 64 : 32; + const unsigned bankBitWidth = 32; + const unsigned simdWidth = 16; + + // Number of inner dimension rows per one pattern repeat + int innerDimLength = operandShape[sharedOrder[0]]; + int elemsPerOneBanksRow = (numBanks * bankBitWidth) / elemBitWidth; + + int perPhase = std::max(1, elemsPerOneBanksRow / innerDimLength); + int maxPhase = + std::max(std::min(simdWidth / perPhase, innerDimLength / vectorSize), 1u); + + // TODO (zhanglx): figure out better parameters for mfma4 + if (getInstrShape()[0] == 4) + maxPhase = 4; + + return SwizzledSharedEncodingAttr::get(getContext(), vectorSize, perPhase, + maxPhase, sharedOrder, ctaLayout); +} + +//===----------------------------------------------------------------------===// +// Wmma encoding +//===----------------------------------------------------------------------===// + +SmallVector AMDWmmaEncodingAttr::getRepOrder() const { + return getMatrixOrder(getRank(), /*rowMajor*/ true); +} + +SmallVector +AMDWmmaEncodingAttr::getRepOrderForOperand(int opIdx) const { + return getOrderForDotOperand(opIdx, getRank(), /*kContig*/ true); +} + +SmallVector +AMDWmmaEncodingAttr::getRepForOperand(ArrayRef operandShape, int kDim, + int opIdx) const { + auto mnkDim = getInstrShape(); + SmallVector operandTileShape{opIdx == 0 ? mnkDim[0] : kDim, + opIdx == 0 ? kDim : mnkDim[1]}; + + assert(operandTileShape.size() == 2); + auto warpsPerCTA = getWarpsPerCTA(); + auto tilesPerWarp = getTilesPerWarp(); + + auto rank = operandShape.size(); + assert(rank == 2 || rank == 3); + int numRepBatch = + rank == 3 ? std::max(1, operandShape[0] / warpsPerCTA[0]) : 1; + if (opIdx == 0) + return { + numRepBatch, + std::max(1, operandShape[rank - 2] / + (operandTileShape[0] * tilesPerWarp[rank - 2] * + warpsPerCTA[rank - 2])) * + tilesPerWarp[rank - 2], + std::max(1, operandShape[rank - 1] / operandTileShape[1])}; + else { + assert(opIdx == 1); + return { + numRepBatch, + std::max(1, operandShape[rank - 2] / operandTileShape[0]), + std::max(1, operandShape[rank - 1] / + (operandTileShape[1] * tilesPerWarp[rank - 1] * + warpsPerCTA[rank - 1])) * + tilesPerWarp[rank - 1]}; + } +} + +SwizzledSharedEncodingAttr AMDWmmaEncodingAttr::composeSharedLayoutForOperand( + CTAEncodingAttr ctaLayout, int operandIdx, ArrayRef operandShape, + ArrayRef sharedOrder, unsigned kWidth, unsigned elemBitWidth, + bool needTrans) const { + int kDimIndex = operandIdx == 0 ? 1 : 0; + bool isKContig = sharedOrder[0] == kDimIndex; + + if (!isKContig) { + // Do not swizzle. In this case accesses will go in different banks even + // without swizzling. + return SwizzledSharedEncodingAttr::get(getContext(), 1, 1, 1, sharedOrder, + ctaLayout); + } + + // max vectorization size for ds_load is 128 bits + int vectorSize = std::min(kWidth * elemBitWidth, 128u) / elemBitWidth; + + const int numBanks = 32; + const int bankBitWidth = 32; + + // Number of inner dimension rows per one pattern repeat + int innerDimLength = operandShape[sharedOrder[0]]; + int elemsPerOneBanksRow = (numBanks * bankBitWidth) / elemBitWidth; + + int perPhase = std::max(1, elemsPerOneBanksRow / innerDimLength); + // for both RDNA3 and RDNA4, the M/N dimension of wmma is 16 + // This represents the max number of rows that can be accessed + // at the same time + int mDim = getInstrShape()[0]; + int maxPhase = + std::max(std::min(mDim / perPhase, innerDimLength / vectorSize), 1); + + return SwizzledSharedEncodingAttr::get(getContext(), vectorSize, perPhase, + maxPhase, sharedOrder, ctaLayout); +} + +//===----------------------------------------------------------------------===// +// sunrise_mma encoding +//===----------------------------------------------------------------------===// + +SmallVector SunriseMmaEncodingAttr::getRepOrder() const { + return getMatrixOrder(getRank(), /*rowMajor*/ true); +} + +SmallVector +SunriseMmaEncodingAttr::getRepOrderForOperand(int opIdx) const { + return getOrderForDotOperand(opIdx, getRank(), /*kContig*/ true); +} + +SmallVector +SunriseMmaEncodingAttr::getInstrShapeForOperand(unsigned opIdx, unsigned elemBitWdith) const { // MK, KN + // A, B矩阵每个warp加载的尺寸 + int packPer32bit = 32 / this->getInputElemBitWidth(); + if(opIdx == 0) { + switch(elemBitWdith) { + case 32: return {8, 4}; + case 16: return {8, 8}; + case 8: return {8, 16}; + case 4: return {8, 32}; + default: llvm_unreachable("unsupported packPer32bit"); + } + } + else { + switch(elemBitWdith) { + case 32: return {4, 8}; + case 16: return {8, 8}; + case 8: return {16, 8}; + case 4: return {32, 8}; + default: llvm_unreachable("unsupported packPer32bit"); + } + } return {0, 0}; +} + +SmallVector +SunriseMmaEncodingAttr::getShapePerCTATileForOperand(unsigned opIdx) const { + // CTA内所有warp执行一次mma操作对应输入tensor的尺寸 + unsigned k = 0; // mma指令k维度的元素个数 + int elemBitWidth = this->getInputElemBitWidth(); + switch (elemBitWidth) { + case 4: k = 32; break; + case 8: k = 16; break; + case 16: k = 8; break; + case 32: k = 4; break; + default: + llvm::report_fatal_error("SunriseMmaEncodingAttr::getShapePerCTATileForOperand " + "unsuppored inputElemBitWidth"); + } + auto shapePerCTATile = getShapePerCTATile(); + SmallVector ret; + if (opIdx == 0) { + ret = {shapePerCTATile[0], k}; + } else { + ret = {k, shapePerCTATile[1]}; + } + return ret; +} + +SmallVector SunriseMmaEncodingAttr::getShapePerCTATile() const { + // CTA内所有warp执行一次mma操作的结果的尺寸 + auto warpsPerCTA = getWarpsPerCTA(); + assert(warpsPerCTA.size() == 2); + SmallVector shapePerCTATile(2); + shapePerCTATile[0] = warpsPerCTA[0] * 8; + shapePerCTATile[1] = warpsPerCTA[1] * 8; + return shapePerCTATile; +} + +// 获取mma的两个输入a、b的每个维度的切割数，即每个维度有几个CTATile +// 【注意】两个a和b的元素类型应该一致 +// 返回值：[repM, repK]或[repK, repN] +SmallVector SunriseMmaEncodingAttr::getRepForOperand(ArrayRef operandShape, + Type elemType, int opIdx) const { + int instrM = 8, instrN = 8, instrK = 0; + if(elemType.isF32() || elemType.isInteger(32)) { instrK = 4; } + else if(elemType.isF16() || elemType.isBF16()) { instrK = 8; } + else if(elemType.isInteger(8)) { instrK = 16; } + else if(elemType.isInteger(4)) { instrK = 32; } + else { + llvm::report_fatal_error("unsupported tensor data type for tmma!"); + return {}; + } + + auto warpsPerCTA = getWarpsPerCTA(); + SmallVector ret; + if (opIdx == 0) + ret = {std::max(1, operandShape[0] / (instrM * warpsPerCTA[0])), + std::max(1, operandShape[1] / instrK)}; + else { + assert(opIdx == 1); + ret = {std::max(1, operandShape[0] / instrK), + std::max(1, operandShape[1] / (instrN * warpsPerCTA[1]))}; + } + return ret; +} + +//===----------------------------------------------------------------------===// +// Mma encoding +//===----------------------------------------------------------------------===// + +bool NvidiaMmaEncodingAttr::isVolta() const { return getVersionMajor() == 1; } + +bool NvidiaMmaEncodingAttr::isTuring() const { + return getVersionMajor() == 2 && getVersionMinor() == 1; +} + +bool NvidiaMmaEncodingAttr::isAmpere() const { return getVersionMajor() == 2; } + +bool NvidiaMmaEncodingAttr::isHopper() const { return getVersionMajor() == 3; } + +SmallVector NvidiaMmaEncodingAttr::getRepOrder() const { + return getMatrixOrder(getRank(), /*rowMajor*/ true); +} + +SmallVector +NvidiaMmaEncodingAttr::getRepOrderForOperand(int opIdx) const { + return getOrderForDotOperand(opIdx, getRank(), /*kContig*/ true); +} + +SmallVector +NvidiaMmaEncodingAttr::getRepForOperand(ArrayRef shape, int bitwidth, + int kWidth, int opIdx) const { + assert(kWidth >= std::max(32 / bitwidth, 1) && + "kWidth must be >= max(32 / bitwidth, 1) for this function to be " + "well-defined"); + auto rank = shape.size(); + // Broadcast long K + auto warpsPerCTA = to_vector(getWarpsPerCTA()); + auto kDim = opIdx == 0 ? rank - 1 : rank - 2; + warpsPerCTA[kDim] = 1; + + SmallVector tileSize; + if (rank == 3) { + tileSize.push_back(1); + } + // warpSizeK * (warpRepK * VecBitWidth) + auto tileBitWidthK = (isAmpere() && bitwidth == 64) ? (4 * 256) : (4 * 64); + if (opIdx == 0) { + // m x k + tileSize.push_back(16); + tileSize.push_back(tileBitWidthK / bitwidth); + } else { + // k x n + // Hopper path never uses the n value, since this method is only invoked + // for in-RF (dotOpEnc) operands, but WGMMA only supports in A to be in RF + // so it's fine if the n is incorrect here + tileSize.push_back(tileBitWidthK / bitwidth); + tileSize.push_back(8); + } + + SmallVector numRep; + // Lezcano: This is odd. Why do we always return a vector of size 3? + if (rank != 3) { + numRep.push_back(1); + } + for (auto [s, size, warp] : llvm::zip(shape, tileSize, warpsPerCTA)) { + numRep.push_back(std::max(1, s / (size * warp))); + } + return numRep; +} + +//===----------------------------------------------------------------------===// +// DotOperand Encoding +//===----------------------------------------------------------------------===// + +SmallVector DotOperandEncodingAttr::getRepOrder() const { + if (auto mma = mlir::dyn_cast(getParent())) { + return mma.getRepOrderForOperand(getOpIdx()); + } else if (auto blocked = mlir::dyn_cast(getParent())) { + return to_vector(blocked.getOrder()); + } + llvm::report_fatal_error( + "getRepOrder not implemented for DotOperandEncodingAttr"); + return {}; +} + +CTAEncodingAttr DotOperandEncodingAttr::getCTALayout() const { + auto layout = ::getCTALayout(getParent()).getLinearLayout(); + auto bases = layout.getBases(); + auto kBlock = StringAttr::get(getContext(), "block"); + auto &blockBases = bases[kBlock]; + auto rank = layout.getNumOutDims(); + auto kDim = getOpIdx() == 0 ? rank - 1 : rank - 2; + for (auto &basis : blockBases) { + basis[kDim] = 0; + } + auto dims = layout.getOutDims(); + dims[kDim].second = 1; + return CTAEncodingAttr::get(getContext(), LinearLayout(bases, dims, true)); +} +LogicalResult DotOperandEncodingAttr::verify( + ::llvm::function_ref<::mlir::InFlightDiagnostic()> emitError, + unsigned opIdx, Attribute parent, unsigned kWidth) { + if (opIdx != 0 && opIdx != 1) { + return emitError() << "ttg.dot_op opIdx parameter can be 0 or 1, got: " + << opIdx; + } + if (!parent) { + return emitError() << "ttg.dot_op parent parameter cannot be null"; + } + if (auto parentAttr = mlir::dyn_cast(parent)) { + if (kWidth != 0 && !(parentAttr.isAmpere() || parentAttr.isHopper())) + return emitError() << "ttg.dot_op kWidth parameter can only be " + "non-zero for Ampere or Hopper MMA parent"; + if (kWidth == 0 && (parentAttr.isAmpere() || parentAttr.isHopper())) + return emitError() << "ttg.dot_op kWidth parameter is mandatory for " + "Ampere or Hopper MMA parent"; + if (opIdx != 0 && parentAttr.isHopper()) + return emitError() + << "ttg.dot_op opIdx parameter must be 0 for " + "Hopper MMA parent, since Hopper WGMMA only allows first " + "operand to be in registers"; + return success(); + } + + if (auto parentAttr = mlir::dyn_cast(parent)) { + if (parentAttr.getVersion() == 1 && (kWidth != 8 && kWidth != 16)) + return emitError() + << "ttg.dot_op kWidth parameter must be 8/16 for WMMA v1 " + "(including packed cases for `scaled_dot`)"; + if (parentAttr.getVersion() == 2 && !llvm::is_contained({4, 8, 16}, kWidth)) + return emitError() + << "ttg.dot_op kWidth parameter must be 4/8/16 for WMMA v2 " + "(including packed cases for `scaled_dot`)"; + if (parentAttr.getVersion() == 3 && !llvm::is_contained({2, 8, 16}, kWidth)) + return emitError() + << "ttg.dot_op kWidth parameter must be 2/8/16 for WMMA v3"; + return success(); + } + + if (auto parentAttr = mlir::dyn_cast(parent)) { + if (kWidth == 0) + return emitError() << "ttg.dot_op kWidth parameter is mandatory for " + "MFMA parent"; + return success(); + } + + if (auto parentAttr = mlir::dyn_cast(parent)) { + return success(); + } + + if (auto parentAttr = mlir::dyn_cast(parent)) { + if (kWidth != 0) + return emitError() << "ttg.dot_op kWidth parameter is not supported " + "when the parent is a blocked layout"; + return success(); + } + + return emitError() << "ttg.dot_op unexpected parent layout: " << parent; +} + +//===----------------------------------------------------------------------===// +// ASM Interface (i.e.: alias) +//===----------------------------------------------------------------------===// + +class TritonGPUOpAsmInterface : public OpAsmDialectInterface { +public: + using OpAsmDialectInterface::OpAsmDialectInterface; + + AliasResult getAlias(Attribute attr, raw_ostream &os) const override { + // Encoding attributes + if (auto mmaAttr = mlir::dyn_cast(attr)) { + os << "mma"; + return AliasResult::FinalAlias; + } else if (auto sharedAttr = mlir::dyn_cast(attr)) { + os << "shared"; + return AliasResult::FinalAlias; + } else if (auto blockedAttr = mlir::dyn_cast(attr)) { + os << "blocked"; + return AliasResult::FinalAlias; + } else if (auto linearAttr = mlir::dyn_cast(attr)) { + os << "linear"; + return AliasResult::FinalAlias; + } /* else if (auto sliceAttr = dyn_cast(attr)) { + os << "slice"; + return AliasResult::FinalAlias; + } */ + // Memory space attributes + if (auto smem = mlir::dyn_cast(attr)) { + os << "smem"; + return AliasResult::FinalAlias; + } + return OpAsmDialectInterface::getAlias(attr, os); + } +}; + +struct TritonGPUInferLayoutInterface + : public triton::DialectInferLayoutInterface { + using DialectInferLayoutInterface::DialectInferLayoutInterface; + + LogicalResult + inferReduceOpEncoding(Attribute operandEncoding, unsigned axis, + Attribute &resultEncoding, + std::optional loc) const override { + resultEncoding = + SliceEncodingAttr::get(getDialect()->getContext(), axis, + cast(operandEncoding)); + return success(); + } + + // Infer the encoding of a tt.trans(x) given the encoding of x. + // + // Our goal is to choose an encoding so that the trans is a "nop". For + // example, in a blocked encoding, the same GPU threads hold the same + // elements, they're just "renamed" -- what was element [i,j] of the tensor is + // now element [j,i], but that element is held by the same GPU thread. + // + // For most properties of the encoding, we let + // outputEnc.prop = inputEnc.prop * trans.order, + // where `x * y` means we apply permutation y to x. + // + // This works because prop[i] tells you something about the i'th dimension of + // the tensor. (For example, sizePerThread[2] == 4 means that one GPU thread + // contains 4 elements along dim 2 of the tensor.) The transpose reorders the + // dimensions according to the perm trans.order, so we achieve our goal of + // having a "nop" transpose by reordering the values in the prop the same way. + // + // The big exception to this is the encoding's `order`. + // + // An encoding's order is a list of dimensions, from fastest moving (most + // minor) to slowest moving. Thus enc.order[i] does not tell you something + // about the i'th dimension of the tensor, and it would be disasterously + // incorrect to do enc.order * trans.order. + // + // But! If we invert enc.order, it *does* meet this criterion. For example, + // if enc.order = [2,0,1], inverse(enc.order) = [1,2,0]. If you stare at it, + // you'll see that inverse(enc.order)[i] == j means that dimension i is the + // j'th most minor. Therefore we can safely permute *this* by trans.order. + // + // Thus we have + // + // outputEnc.order = inverse(inverse(inputEnc.order) * trans.order) + // = inverse(trans.order) * inputEnc.order. + // + LogicalResult + inferTransOpEncoding(Attribute operandEncoding, ArrayRef shape, + ArrayRef order, Attribute &resultEncoding, + std::optional loc) const override { + // Note: inferFooOpEncoding should not crash if given invalid inputs, which + // happens when someone creates invalid IR. If we return failure() on + // error, then MLIR will generate a helpful error message. + if (isIota(order)) { + resultEncoding = operandEncoding; + return success(); + } + if (shape.size() != order.size()) { + return emitOptionalError(loc, "shape and order rank do not match: ", + shape.size(), " vs ", order.size()); + } + auto checkRank = [&](unsigned rank) { + if (rank != order.size()) { + return emitOptionalError(loc, "rank of encoding does not match order: ", + rank, " vs ", order.size()); + } + return success(); + }; + auto *ctx = getDialect()->getContext(); + + auto permuteCTALayout = [ctx](CTAEncodingAttr layout, + ArrayRef order) { + auto ll = transposeLinearLayout(layout.getLinearLayout(), order); + return CTAEncodingAttr::get(ctx, std::move(ll)); + }; + + auto invOrder = inversePermutation(order); + SmallVector invOrderUnsigned(invOrder.begin(), invOrder.end()); + + if (auto enc = dyn_cast(operandEncoding)) { + if (failed(checkRank(enc.getCTALayout().getRank()))) + return failure(); + + CTAEncodingAttr ctaLayout = permuteCTALayout(enc.getCTALayout(), order); + resultEncoding = SwizzledSharedEncodingAttr::get( + ctx, enc.getVec(), enc.getPerPhase(), enc.getMaxPhase(), + applyPermutation(invOrderUnsigned, enc.getOrder()), ctaLayout); + return success(); + } + + if (auto enc = dyn_cast(operandEncoding)) { + if (order == ArrayRef({1, 0})) { + if (failed(checkRank(enc.getCTALayout().getRank()))) + return failure(); + + CTAEncodingAttr ctaLayout = permuteCTALayout(enc.getCTALayout(), order); + resultEncoding = NVMMASharedEncodingAttr::get( + ctx, enc.getSwizzlingByteWidth(), !enc.getTransposed(), + enc.getElementBitWidth(), enc.getFp4Padded(), ctaLayout); + return success(); + } + } + + if (auto enc = dyn_cast(operandEncoding)) { + if (failed(checkRank(enc.getCTALayout().getRank()))) + return failure(); + + CTAEncodingAttr ctaLayout = permuteCTALayout(enc.getCTALayout(), order); + resultEncoding = BlockedEncodingAttr::get( + ctx, applyPermutation(enc.getSizePerThread(), order), + applyPermutation(enc.getThreadsPerWarp(), order), + applyPermutation(enc.getWarpsPerCTA(), order), + applyPermutation(invOrderUnsigned, enc.getOrder()), ctaLayout); + return success(); + } + // Generic case + auto padded = dyn_cast(operandEncoding); + + auto ll = padded ? padded.getLinearComponent() + : toLinearLayout(shape, operandEncoding); + if (failed(checkRank(ll.getNumOutDims()))) + return failure(); + auto transposedLl = transposeLinearLayout(ll, order); + if (isa(operandEncoding)) { + resultEncoding = LinearEncodingAttr::get(ctx, std::move(transposedLl)); + } else if (padded) { + resultEncoding = PaddedSharedEncodingAttr::get(ctx, padded.getIntervals(), + padded.getPaddings(), + std::move(transposedLl)); + } else { + auto shared = cast(operandEncoding); + resultEncoding = SharedLinearEncodingAttr::get( + ctx, std::move(transposedLl), shared.getAlignment()); + } + return success(); + } + + LogicalResult + inferExpandDimsOpEncoding(Attribute operandEncoding, unsigned axis, + Attribute &resultEncoding, + std::optional location) const override { + auto sliceEncoding = mlir::dyn_cast(operandEncoding); + if (!sliceEncoding) + return emitOptionalError( + location, "ExpandDimsOp operand encoding must be SliceEncodingAttr"); + if (sliceEncoding.getDim() != axis) + return emitOptionalError( + location, "Incompatible slice dimension for ExpandDimsOp operand"); + resultEncoding = sliceEncoding.getParent(); + return success(); + } + + LogicalResult + inferDotOpEncoding(Attribute operandEncoding, unsigned opIdx, + Attribute retEncoding, + std::optional location) const override { + auto mmaRetEncoding = mlir::dyn_cast(retEncoding); + if (mmaRetEncoding && mmaRetEncoding.isHopper()) { + auto dotOpEnc = mlir::dyn_cast(operandEncoding); + if (!mlir::isa( + operandEncoding) && + !(opIdx == 0 && dotOpEnc && dotOpEnc.getOpIdx() == 0 && + mlir::isa(dotOpEnc.getParent()))) { + return emitOptionalError( + location, "unexpected operand layout for NvidiaMmaEncodingAttr v3"); + } + } else if (auto dotOpEnc = + mlir::dyn_cast(operandEncoding)) { + if (opIdx != dotOpEnc.getOpIdx()) + return emitOptionalError(location, "Wrong opIdx"); + if (retEncoding != dotOpEnc.getParent()) + return emitOptionalError(location, "Incompatible parent encoding"); + } else + return emitOptionalError( + location, "Dot's a/b's encoding should be of DotOperandEncodingAttr"); + return success(); + } + + LogicalResult + verifyDotOpEncodingCompatibility(Operation *op, Attribute operandEncodingA, + Attribute operandEncodingB) const override { + auto aEncoding = + mlir::dyn_cast(operandEncodingA); + auto bEncoding = + mlir::dyn_cast(operandEncodingB); + if (!aEncoding && !bEncoding) + return mlir::success(); + if (!aEncoding || !bEncoding) + return op->emitError("mismatching encoding between A and B operands"); + // Verify that the encodings are valid. + if (aEncoding.getKWidth() != bEncoding.getKWidth()) + return op->emitError("mismatching kWidth between A and B operands"); + + // Check if we have already selected an MMA version for Nvidia. If so, + // validate that the encodings are correct and compatible. + auto mmaAEncoding = + dyn_cast_or_null(aEncoding.getParent()); + auto mmaBEncoding = + dyn_cast_or_null(bEncoding.getParent()); + auto dotOp = cast(op); + auto resEnc = dotOp.getResult().getType().getEncoding(); + auto mmaResEncoding = dyn_cast(resEnc); + if (mmaAEncoding || mmaBEncoding || mmaResEncoding) { + // Check that they are all set and have the same version. + if (!mmaAEncoding || !mmaBEncoding || !mmaResEncoding) + return op->emitError("mismatching MMA encoding"); + auto mmaBEncoding = cast(bEncoding.getParent()); + if (mmaAEncoding.getVersionMajor() != mmaBEncoding.getVersionMajor() || + mmaAEncoding.getVersionMajor() != mmaResEncoding.getVersionMajor()) { + return op->emitError("mismatched MMA version."); + } + // Verify that the operands are supported on the selected MMA version. + if (!supportMMA(dotOp, mmaResEncoding.getVersionMajor())) + return op->emitError("unsupported MMA version"); + } + return success(); + } + + // Given a src shape + encoding and a dst shape, our goal is to compute a dst + // encoding that makes the reshape a "nop". That is, if GPU thread [x,y,z] + // contains elements [a,b,c,d] before the reshape, it contains those same + // elements after the reshape, they're just "renamed". + // + // Using legacy layouts, a dst encoding that satisfies this property may not + // exist. Here are some positive and negative examples. + // + // - NOT OK: 4x4 order=[0,1] -> 16. Reshape merges elements so + // dim 1 is the fastest-changing in the dst, but the src has the opposite + // order. + // - OK: 2x2x32 order=[1,0,2] -> 4x32. We choose dst order [0,1]. + // What's important is that the 2x2 dimensions appear in major-to-minor + // order. + // - NOT OK: 32x32 sizePerThread=[2,2] -> 1024. Thread 0 in the src + // contains elements [(0,0), (0,1), (1,0), and (1,1)]. We cannot express + // this with an encoding based on the dst shape. + // - OK: 32x4 sizePerThread=[4,4] -> 128. dst with sizePerThread=[16] will + // contain the same elements as before. + // + // With linear layouts, we can always find a dst encoding that satisfies + // this property. See inferReshapeOpEncoding. + // + // Users of this function require that it is symmetrical: if + // (srcShape,srcEnc,dstShape) => dstEnc, then (dstShape,dstEnc,srcShape) => + // srcEnc. + LogicalResult inferReshapeOpLegacyEncoding(ArrayRef srcShape, + Attribute srcEnc, + ArrayRef dstShape, + Attribute &dstEnc) const { + auto src = mlir::dyn_cast(srcEnc); + if (!src) { + return failure(); + } + + // Nop reshape; we can always infer an encoding. + if (srcShape == dstShape) { + dstEnc = srcEnc; + return success(); + } + + // default -> default encoding is always a nop. + auto context = srcEnc.getContext(); + int32_t numWarps = product(src.getWarpsPerCTA()); + int32_t threadsPerWarp = product(src.getThreadsPerWarp()); + int32_t numCTAs = product(src.getCTALayout().getCTAsPerCGA()); + if (srcEnc == getDefaultBlockedEncoding(context, srcShape, numWarps, + threadsPerWarp, numCTAs)) { + dstEnc = getDefaultBlockedEncoding(context, dstShape, numWarps, + threadsPerWarp, numCTAs); + return success(); + } + + // Cowardly refuse to handle encodings with multiple CTAs. CTAsPerCGA + // should be like the other fields in blocked encoding, but I'm not sure how + // to handle CTASplitNum. + auto srcCTALayout = src.getCTALayout(); + if (!all_of(srcCTALayout.getCTAsPerCGA(), + [](int32_t x) { return x == 1; }) || + !all_of(srcCTALayout.getCTASplitNum(), + [](int32_t x) { return x == 1; })) { + return failure(); + } + + // Cowardly refuse to handle encodings where shape[dim] is not divisible by + // sizePerThread[dim], threadsPerWarp[dim], and warpsPerCTA[dim]. (We make + // an exception if the block is larger than the shape.) + auto checkDivisibility = [&](StringRef name, ArrayRef subblock) { + for (int dim = 0; dim < srcShape.size(); dim++) { + if (srcShape[dim] >= subblock[dim] && + srcShape[dim] % subblock[dim] != 0) { + return failure(); + } + } + return success(); + }; + if (!succeeded( + checkDivisibility("sizePerThread", src.getSizePerThread())) || + !succeeded( + checkDivisibility("threadsPerWarp", src.getThreadsPerWarp())) || + !succeeded(checkDivisibility("warpsPerCTA", src.getWarpsPerCTA()))) { + return failure(); + } + + SmallVector, SmallVector>> decomp = + getReshapeDecomposition(srcShape, dstShape); + + // enc.order[i] == j means that dimension j is the enc.order[i]'th most + // minor. But what we usually want is the inverse: inverse(enc.order)[i] = j + // means that dimension i is the j'th most minor (larger means more major). + auto srcInvOrder = inversePermutation(src.getOrder()); + + // If src dims [a,b,c] are to be merged, then they must be consecutive in + // physical order, with `a` being the most major. + for (const auto &[srcDims, dstDims] : decomp) { + if (!isConsecutive(to_vector(reverse(gather(srcInvOrder, srcDims))))) { + return failure(); + } + } + + // If src dims [a,b,c] are to be merged, then `c` must fill up sizePerThread + // / threadsPerWarp / blocksPerCTA before `b` can have any non-1 values. + // Examples: + // + // - NOT OK: shape=[4,4,4], sizePerThread=[1,2,2]. + // The total sizePerThread for dim 2 is 2, which is less than dim 2's + // size of 4. Therefore dim 1 cannot have non-1 sizePerThread. + // + // - OK: shape=[4,4,4], sizePerThread=[1,2,4]. + // Dim 2's sizePerThread covers its whole size, so dim 1 is allowed to + // have non-1 sizePerThread. + // + // - NOT OK: shape=[4,4,4], sizePerThread=[2,1,4]. + // Dim 1's sizePerThread does not cover its whole size, so dim 0 is not + // allowed to have non-1 sizePerThread. + // + // - NOT OK: shape=[4,4,4], sizePerThread=[1,1,2], + // threadsPerWarp=[1,2,1]. + // Dim 2 has 2 elems per thread and 1 thread per warp. 2*1 is less than + // dim 2's size. Therefore dim 1 must have threadsPerWarp=1. + // + // In addition, the encoding's block can be larger than the shape, but only + // in the most-major dimension of each decomposed chunk, and only after + // we've "used up" the more minor dims. Examples: + // + // - OK: shape=[4,4,4], sizePerThread=[1,2,4], threadsPerWarp=[16,2,1], + // warpsPerCTA=[4,1,1]. + // The whole size of dims 0 and 1 are covered by sizePerThread * + // threadsPerWarp. Therefore dim 2 is allowed to have threadsPerWarp and + // warpsPerCTA larger than its size. + for (const auto &[srcDims, dstDims] : decomp) { + auto shapeRemaining = gather(srcShape, srcDims); + auto checkSubblock = [&, srcDims = srcDims](ArrayRef subblock) { + // Iterate minor-to-major (i==0 is most major). + for (int i = srcDims.size() - 1; i >= 0; i--) { + int dim = srcDims[i]; + if (subblock[dim] == 1) { + continue; + } + + // Check that more-minor dims all have 1 in shapeRemaining. + for (int j = i + 1; j < srcDims.size(); j++) { + if (shapeRemaining[j] != 1) { + return failure(); + } + } + + if (shapeRemaining[i] >= subblock[dim]) { + assert(shapeRemaining[i] % subblock[dim] == 0); // checked earlier + shapeRemaining[i] /= subblock[dim]; + } else { + shapeRemaining[i] = 0; + } + + // Is the block larger than the shape in this dimension? This is OK + // only if we're the most-major dimension of the chunk and in all + // future chunks, only this most-major dim has a non-1 size. + if (shapeRemaining[i] == 0 && i != 0) { + return failure(); + } + } + return success(); + }; + if (!succeeded(checkSubblock(src.getSizePerThread())) || + !succeeded(checkSubblock(src.getThreadsPerWarp())) || + !succeeded(checkSubblock(src.getWarpsPerCTA()))) { + return failure(); + } + } + + // Given e.g. src.getSizePerThread(), computeSubblockSize computes e.g. + // dst.getSizePerThread(). This should be called for each of sizePerThread, + // threadsPerWarp, and warpsPerCTA, in that order. + SmallVector dstShapeRemaining(dstShape); + auto computeSubblockSize = [&](ArrayRef srcSubblock, + SmallVector &dstSubblock, + StringRef fieldName) -> LogicalResult { + // The dst subblock is "filled up" greedily starting with the most minor + // dim. When we're done, we are left with a smaller shape, of size + // dstShape / dstSubblock, which we store in dstShapeRemaining and use for + // the next call to computeSubblockSize. + dstSubblock.resize(dstShape.size()); + for (const auto &[srcDims, dstDims] : decomp) { + int64_t subblockRemaining = product(gather(srcSubblock, srcDims)); + for (int i = dstDims.size() - 1; i >= 0; i--) { + auto &val = dstSubblock[dstDims[i]]; + auto &shapeRemaining = dstShapeRemaining[dstDims[i]]; + val = std::min(subblockRemaining, shapeRemaining); + + assert(shapeRemaining % val == 0); // Checked earlier. + subblockRemaining /= val; + shapeRemaining /= val; + } + + // If there are any elems remaining in the subblock, it must be because + // the block is larger than the shape. This excess goes into the + // most-major dim of the subblock. + dstSubblock[dstDims[0]] *= subblockRemaining; + } + return success(); + }; + + SmallVector dstSizePerThread; + SmallVector dstThreadsPerWarp; + SmallVector dstWarpsPerCTA; + if (!succeeded(computeSubblockSize(src.getSizePerThread(), dstSizePerThread, + "sizePerThread")) || + !succeeded(computeSubblockSize(src.getThreadsPerWarp(), + dstThreadsPerWarp, "threadsPerWarp")) || + !succeeded(computeSubblockSize(src.getWarpsPerCTA(), dstWarpsPerCTA, + "warpsPerCTA"))) { + return failure(); + } + + // Since we know that each set of srcDims is consecutive, we can + // meaningfully sort decomp by the physical order of the src dimensions, + // major-to-minor. This will also be the order of the dst dimensions. + llvm::sort(decomp, [&](const auto &a, const auto &b) { + const auto &[srcDimsA, dstDimsA] = a; + const auto &[srcDimsB, dstDimsB] = b; + return srcInvOrder[srcDimsA.front()] < srcInvOrder[srcDimsB.front()]; + }); + + // Compute the dst order. Make the dimensions appear in the same order as + // their corresponding src dimensions. + SmallVector dstInvOrder(dstShape.size()); + int i = 0; + for (const auto &[srcDims, dstDims] : decomp) { + for (auto dim : reverse(dstDims)) { + dstInvOrder[dim] = i++; + } + } + auto dstOrder = inversePermutation(dstInvOrder); + + // CTALayout can be all 1's because we bailed on multi-CTA layouts above. + auto CTALayout = + CTAEncodingAttr::getDefault(src.getContext(), dstShape.size()); + + dstEnc = BlockedEncodingAttr::get(src.getContext(), dstSizePerThread, + dstThreadsPerWarp, dstWarpsPerCTA, + dstOrder, CTALayout); + + return success(); + } + + LogicalResult + verifyLayoutsAreEqual(ArrayRef shape, Attribute expected, + Attribute got, + std::optional loc) const override { + if (expected == got) { + return success(); + } + if (!expected || !got) + return failure(); + + // Check whether the encodings are structurally the same. + if (!areLayoutsEquivalent(shape, cast(expected), + cast(got))) { + return emitOptionalError(loc, "Expected result encoding ", expected, + " but was ", got); + } + return success(); + } + + LogicalResult + inferReshapeOpEncoding(ArrayRef srcShape, Attribute srcEnc, + ArrayRef dstShape, Attribute &dstEnc, + std::optional loc) const override { + if (product(srcShape) != product(dstShape)) { + return emitOptionalError(loc, "numel of dst shape does not match " + "numel of src shape"); + } + auto result = + inferReshapeOpLegacyEncoding(srcShape, srcEnc, dstShape, dstEnc); + if (succeeded(result)) { + return result; + } + if (!isa(srcEnc)) { + return emitOptionalError(loc, + "Failed MemDescReshapeOp encoding inference"); + } + // If the legacy encoding failed use LinearLayouts. + // Once LinearLayouts are more widely used, we can remove + // inferReshapeOpLegacyEncoding and simply use LLs. + + // HACK: We create a dummy tensor type to pass to inferReshapeLinearLayout. + auto ctx = srcEnc.getContext(); + auto fp32Type = IntegerType::get(ctx, 32, IntegerType::Unsigned); + auto srcTy = RankedTensorType::get(srcShape, fp32Type, srcEnc); + LinearLayout ll = + inferReshapeLinearLayout(cast(srcTy), dstShape); + + dstEnc = LinearEncodingAttr::get(srcEnc.getContext(), ll); + return success(); + } + + LogicalResult + inferDefaultJoinOpEncoding(Attribute srcEnc, Attribute &dstEnc, + ArrayRef shape, + std::optional loc) const override { + auto ctx = getContext(); + if (auto enc = mlir::dyn_cast(srcEnc); + enc && enc.getDim() == shape.size()) { + SmallVector joinedShape(shape); + joinedShape.push_back(2); + auto parent = enc.getParent(); + auto parentLL = toLinearLayout(joinedShape, parent); + + Attribute splitEnc; + auto result = inferSplitOpEncoding(parent, splitEnc, joinedShape, loc); + if (succeeded(result) && + areLayoutsEquivalent(shape, cast(splitEnc), + cast(srcEnc))) { + dstEnc = parent; + return success(); + } + } else if (auto enc = mlir::dyn_cast(srcEnc)) { + // JoinOp takes two tensors of shape AxBxC and generates a tensor of shape + // AxBxCx2. The encoding is the same as the input, but with 2 elems per + // thread in the new dimension. The new dimension is the fastest running + // dimension. + auto append = [](ArrayRef vals, int val) { + SmallVector ret(vals); + ret.push_back(val); + return ret; + }; + auto appendMajorDim = [](ArrayRef order) { + SmallVector ret(order); + ret.insert(ret.begin(), ret.size()); + return ret; + }; + auto ctall = enc.getCTALayout().getLinearLayout(); + auto kBlock = StringAttr::get(enc.getContext(), "block"); + auto newDim = standardOutDimNames( + enc.getContext(), ctall.getNumOutDims() + 1)[ctall.getNumOutDims()]; + ctall *= LinearLayout::identity1D(1, kBlock, newDim); + dstEnc = BlockedEncodingAttr::get( + enc.getContext(), append(enc.getSizePerThread(), 2), + append(enc.getThreadsPerWarp(), 1), append(enc.getWarpsPerCTA(), 1), + appendMajorDim(enc.getOrder()), + CTAEncodingAttr::get(enc.getContext(), ctall)); + return success(); + } + + // Append dim to shape + auto ll = toLinearLayout(shape, srcEnc); + SmallVector dstShape(shape.begin(), shape.end()); + dstShape.push_back(1); + ll = ll.reshapeOuts(standardOutDimPairs(ctx, dstShape)); + + // Try join on last dim + auto axis = dstShape.size() - 1; + auto newLl = LinearLayout::empty(); + auto result = + tryJoinOnAxis(ctx, ll, newLl, /*fwdInference=*/true, axis, loc); + + assert(result.succeeded()); + dstEnc = LinearEncodingAttr::get(ctx, newLl); + return success(); + } + + LogicalResult + inferSplitOpEncoding(Attribute srcEnc, Attribute &dstEnc, + ArrayRef shape, + std::optional loc) const override { + // SplitOp takes a tensor of shape AxBxCx2 and generates two tensors of + // shape AxBxC. The input must have 2 elements per thread in the last + // dimension, which must be the fastest running dimension. The result + // encoding is the same as the input, but with the last dimension removed. + auto enc = mlir::dyn_cast(srcEnc); + bool isSimpleSplit = (enc && (enc.getSizePerThread().back() == 2) && + (enc.getThreadsPerWarp().back() == 1) && + (enc.getWarpsPerCTA().back() == 1) && + (enc.getCTALayout().getCTAsPerCGA().back() == 1)); + if (isSimpleSplit) { + SmallVector newOrder(enc.getOrder()); + auto ctall = enc.getCTALayout().getLinearLayout(); + int splitDim = newOrder.size() - 1; + // Remove splitDim from order. + newOrder.erase(std::remove(newOrder.begin(), newOrder.end(), splitDim), + newOrder.end()); + // Remove last dimension from ctall. + ctall = ctall.unsqueezeOut(to_vector(ctall.getOutDimNames()).back()); + dstEnc = BlockedEncodingAttr::get( + enc.getContext(), // + ArrayRef(enc.getSizePerThread()).drop_back(1), + ArrayRef(enc.getThreadsPerWarp()).drop_back(1), + ArrayRef(enc.getWarpsPerCTA()).drop_back(1), ArrayRef(newOrder), + CTAEncodingAttr::get(enc.getContext(), ctall)); + return success(); + } + + auto axis = shape.size() - 1; + if (shape[axis] != 2) { + return emitOptionalError( + loc, "SplitOp input shape should have 2 in the last dim"); + } + + auto ctx = getContext(); + + // Split on last dim + auto ll = toLinearLayout(shape, srcEnc); + auto newLl = LinearLayout::empty(); + auto result = + tryJoinOnAxis(ctx, ll, newLl, /*fwdInference=*/false, axis, loc); + if (!result.succeeded()) { + return failure(); + } + // Remove last dim from newLl (which should be 1) + SmallVector dstShape(shape.begin(), shape.end()); + dstShape.pop_back(); + newLl = newLl.reshapeOuts(standardOutDimPairs(ctx, dstShape)); + dstEnc = LinearEncodingAttr::get(ctx, newLl); + return success(); + } + + LogicalResult + inferFp4ToFpOpEncoding(ArrayRef shape, int axis, Attribute inEnc, + Attribute &outEnc, bool fwdInference, + std::optional loc) const override { + // We implement two legacy layout propagations + // Once we fully migrate to LinearLayouts, we can remove these. + auto *ctx = getContext(); + // The output encoding will only be a legacy encoding if the axis is the + // fastest running dimension. + // FIXME: We should make sure that there are enough elements along the axis + // axis whenever fwdInference is false + if (getOrder(cast(inEnc), shape)[axis] == 0) { + // Dot operand: double kWidth if kDim == axis. + if (auto dotEnc = mlir::dyn_cast(inEnc)) { + auto kWidth = dotEnc.getKWidth(); + if (fwdInference) { + kWidth *= 2; + } else { + if (kWidth > 1) { + // bwd inference + kWidth /= 2; + } else { + return emitOptionalError(loc, + "Fp4ToFpOp requires at least 2 elements " + "per thread in the axis dimension"); + } + } + outEnc = DotOperandEncodingAttr::get(ctx, dotEnc.getOpIdx(), + dotEnc.getParent(), kWidth); + return success(); + } + + // Blocked layout: double elemsPerThread[axis]. + if (auto blockedEnc = mlir::dyn_cast(inEnc)) { + auto sizePerThread = llvm::to_vector(blockedEnc.getSizePerThread()); + if (fwdInference) { + sizePerThread[axis] *= 2; + } else { + if (sizePerThread[axis] > 1) { + sizePerThread[axis] /= 2; + } else { + return emitOptionalError( + loc, "Fp4ToFpOp requires at least 2 elements per " + "thread in the axis dimension"); + } + } + outEnc = BlockedEncodingAttr::get( + ctx, sizePerThread, blockedEnc.getThreadsPerWarp(), + blockedEnc.getWarpsPerCTA(), blockedEnc.getOrder(), + blockedEnc.getCTALayout()); + return success(); + } + } + + auto ll = toLinearLayout(shape, inEnc); + auto newLl = LinearLayout::empty(); + auto result = tryJoinOnAxis(ctx, ll, newLl, fwdInference, axis, loc); + if (!result.succeeded()) + return result; + outEnc = LinearEncodingAttr::get(ctx, newLl); + return success(); + } +}; + +struct TritonGPUVerifyTensorLayoutInterface + : public triton::DialectVerifyTensorLayoutInterface { + using DialectVerifyTensorLayoutInterface::DialectVerifyTensorLayoutInterface; + + LogicalResult verifyTensorLayout( + Attribute layout, RankedTensorType rankedTy, Operation *op, + function_ref makeErr) const override { + auto distr = dyn_cast(layout); + if (!distr) + return makeErr() + << "Non-distributed layout is not allowed in tensor type."; + auto rank = distr.getRepOrder().size(); + if (rank != rankedTy.getRank()) + return makeErr() << "Layout has rank " << rank + << ", but the tensor it's attached to has rank " + << rankedTy.getRank() << "."; + if (llvm::any_of(rankedTy.getShape(), + [](int64_t i) { return !llvm::isPowerOf2_64(i); })) { + return makeErr() << "Layout has shape " << rankedTy.getShape() + << ", but the tensor it's attached to has shape " + << rankedTy.getShape() + << " which is not a power of two."; + } + auto ll = toLinearLayout(rankedTy); + ModuleOp module = op->getParentOfType(); + + // Number of threads per warp. + auto kLane = StringAttr::get(module.getContext(), "lane"); + int moduleThreadsPerWarp = TritonGPUDialect::getThreadsPerWarp(module); + if (ll.getInDimSize(kLane) != moduleThreadsPerWarp) { + return makeErr() << layout << ".\nLayout has " << ll.getInDimSize(kLane) + << " threads per warp, but the module specifies " + << moduleThreadsPerWarp << " threads per warp."; + } + + // Number of warps per CTA. + std::optional moduleWarpsPerCTA = maybeLookupNumWarps(op); + if (!moduleWarpsPerCTA) { + return makeErr() + << "Could not determine the number of warps per CTA. Operation " + "is not in a context with `ttg.num-warps`."; + } + auto kWarp = StringAttr::get(module.getContext(), "warp"); + if (ll.getInDimSize(kWarp) != *moduleWarpsPerCTA) { + return makeErr() << layout << ".\nLayout has " << ll.getInDimSize(kWarp) + << " warps per CTA, but the context requires " + << *moduleWarpsPerCTA << " warps per CTA."; + } + + // Number of CTAs per CGA. + auto kBlock = StringAttr::get(module.getContext(), "block"); + int moduleCTAsPerCGA = TritonGPUDialect::getNumCTAs(module); + if (ll.getInDimSize(kBlock) != moduleCTAsPerCGA) { + return makeErr() << layout << ".\nLayout has " << ll.getInDimSize(kBlock) + << " CTAs per CGA, but the context requires " + << moduleCTAsPerCGA << " CTAs per CGA."; + } + return success(); + } +}; + +//===----------------------------------------------------------------------===// +// Layout debug printing +//===----------------------------------------------------------------------===// + +// Return N-D delinearized indices from a linear index. +static SmallVector delinearizeIndex(int64_t idx, + ArrayRef shape) { + SmallVector ret(shape.size()); + for (int i = shape.size() - 1; i >= 0; i--) { + ret[i] = idx % shape[i]; + idx /= shape[i]; + } + return ret; +} + +// Returns how many padding characters are needed for the string representation +// of value to be the same as max. +static int numCharacterPadding(int value, int max) { + return std::to_string(max).size() - std::to_string(value).size(); +} + +// return the string padded to have the same length as max. +static std::string paddedString(int value, int max) { + int nbChar = numCharacterPadding(value, max); + std::string str; + for (int i = 0; i < nbChar; i++) + str += " "; + str += std::to_string(value); + return str; +} + +std::string mlir::triton::gpu::getSharedLayoutStr(LinearLayout &ll, + bool useHWPointOfView) { + // This RankedTensorType is a MemDescType (?!) + auto outDimNames = llvm::to_vector(ll.getOutDimNames()); + auto shape = convertType(llvm::to_vector(ll.getOutDimSizes())); + auto *ctx = outDimNames[0].getContext(); + + StringAttr kOffset = StringAttr::get(ctx, "offset"); + StringAttr kBlock = StringAttr::get(ctx, "block"); + int64_t tensorSize = product(shape); + unsigned numBlocks = ll.getInDimSize(kBlock); + int32_t blockSize = tensorSize / numBlocks; + + // elementMapping is for the non-hw layout, offsetMapping for hw-layout + std::vector elementMapping(tensorSize); + std::vector offsetMapping; + + // Shared layouts are a mapping of (block, offset) --> (...) + + // We can just use a single int to index into elementMapping because + // the 'swizzle' operation rearranges the indices---and we want to keep it + // that way + int32_t idx = 0; + // Enumerate all the offsets for each block + for (int32_t block = 0; block < numBlocks; block++) { + for (int32_t offset = 0; offset < blockSize; offset++) { + SmallVector> inputs = { + {kBlock, block}, + {kOffset, offset}, + }; + + SmallVector> outputs = ll.apply(inputs); + + std::string sharedInfo = "("; + std::string &value = elementMapping[idx]; + + if (!value.empty()) + value += "|"; + + value += "("; + // We can build up both strings (for hw/non-hw layouts) concurrently + for (int i = 0; i < outputs.size(); i++) { + // Based on the formatting from LinearLayout::toString, the format for + // the hw layout is slightly different. HW layouts use "," vs ":". + if (i > 0) { + sharedInfo += ","; + value += ":"; + } + auto index = paddedString(outputs[i].second, shape[i]); + sharedInfo += index; + value += index; + } + value += ")"; + sharedInfo += ")"; + + offsetMapping.push_back(sharedInfo); + + idx++; + } + } + + std::string layoutStr; + + if (!useHWPointOfView) { + int rank = shape.size(); + bool newLine = true; + for (int i = 0; i < tensorSize; i++) { + auto indices = delinearizeIndex(i, shape); + int numOpenBracket = 0; + for (int j = rank - 1; j >= 0; j--) { + if (indices[j] % shape[j] != 0) + break; + layoutStr += "["; + numOpenBracket++; + } + if (newLine) { + for (int j = 0; j < rank - numOpenBracket; j++) + layoutStr += " "; + newLine = false; + } + + layoutStr += elementMapping[i]; + auto nextIndices = delinearizeIndex(i + 1, shape); + for (int j = rank - 1; j >= 0; j--) { + if (nextIndices[j] % shape[j] != 0) + break; + layoutStr += "]"; + } + if (nextIndices.back() % shape.back() == 0) { + layoutStr += "\n"; + newLine = true; + } else { + layoutStr += ","; + } + } + } else { + // For the HW view here, print the (block, offset) --> (r,c) mapping + uint32_t idx = 0; + for (int32_t block = 0; block < numBlocks; block++) { + layoutStr += "Block: " + std::to_string(block) + ":\n"; + for (int32_t offset = 0; offset < (tensorSize / numBlocks); offset++) { + layoutStr += "Offset: " + std::to_string(offset) + " -> "; + layoutStr += offsetMapping[idx]; + layoutStr += "\n"; + idx++; + } + } + } + + return layoutStr; +} + +std::string mlir::triton::gpu::getDistributedLayoutStr(LinearLayout &ll, + bool useHWPointOfView) { + auto inDimNames = llvm::to_vector(ll.getInDimNames()); + auto *ctx = inDimNames[0].getContext(); + StringAttr kRegister = StringAttr::get(ctx, "register"); + StringAttr kLane = StringAttr::get(ctx, "lane"); + StringAttr kWarp = StringAttr::get(ctx, "warp"); + StringAttr kBlock = StringAttr::get(ctx, "block"); + + int64_t tensorSize = ll.getTotalOutDimSize(); + std::vector elementMapping(tensorSize); + std::vector threadMapping; + auto shape = convertType(llvm::to_vector(ll.getOutDimSizes())); + unsigned threadsPerWarp = ll.getInDimSize(kLane); + unsigned numWarpsPerCTA = ll.getInDimSize(kWarp); + unsigned numBlocks = ll.getInDimSize(kBlock); + int numElementsPerThreads = ll.getInDimSize(kRegister); + for (int blockId = 0; blockId < numBlocks; ++blockId) { + for (int warpId = 0; warpId < numWarpsPerCTA; warpId++) { + for (int tid = 0; tid < threadsPerWarp; ++tid) { + for (int idx = 0; idx < numElementsPerThreads; ++idx) { + SmallVector> inputs = { + {kBlock, blockId}, + {kWarp, warpId}, + {kLane, tid}, + {kRegister, idx}}; + SmallVector> outputs = + ll.apply(inputs); + int32_t linearizedIdx = 0; + int stride = 1; + for (int i = outputs.size() - 1; i >= 0; i--) { + linearizedIdx += outputs[i].second * stride; + stride *= shape[i]; + } + std::string &value = elementMapping[linearizedIdx]; + if (!value.empty()) + value += "|"; + int padding = numCharacterPadding(blockId, numBlocks) + + numCharacterPadding(tid + warpId * threadsPerWarp, + numWarpsPerCTA * threadsPerWarp) + + numCharacterPadding(idx, numElementsPerThreads); + for (int i = 0; i < padding; i++) + value += " "; + if (numBlocks > 1) + value += "B" + std::to_string(blockId) + ":"; + value += "T" + std::to_string(tid + warpId * threadsPerWarp) + ":" + + std::to_string(idx); + // Now also compute the thread mapping. + std::string threadInfo = "("; + for (int i = 0; i < outputs.size(); i++) { + if (i > 0) + threadInfo += ","; + threadInfo += paddedString(outputs[i].second, shape[i]); + } + threadInfo += ")"; + threadMapping.push_back(threadInfo); + } + } + } + } + std::string layoutStr; + if (!useHWPointOfView) { + // Printing the threads containing each elements of the tensor. + int rank = ll.getNumOutDims(); + bool newLine = true; + for (int i = 0; i < tensorSize; i++) { + auto indices = delinearizeIndex(i, shape); + int numOpenBracket = 0; + for (int j = rank - 1; j >= 0; j--) { + if (indices[j] % shape[j] != 0) + break; + layoutStr += "["; + numOpenBracket++; + } + if (newLine) { + for (int j = 0; j < rank - numOpenBracket; j++) + layoutStr += " "; + newLine = false; + } + + layoutStr += elementMapping[i]; + auto nextIndices = delinearizeIndex(i + 1, shape); + for (int j = rank - 1; j >= 0; j--) { + if (nextIndices[j] % shape[j] != 0) + break; + layoutStr += "]"; + } + if (nextIndices.back() % shape.back() == 0) { + layoutStr += "\n"; + newLine = true; + } else { + layoutStr += ", "; + } + } + } else { + // Printing the elements in each physical reg/warps/threads. + for (int blockId = 0; blockId < numBlocks; blockId++) { + if (numBlocks > 1) + layoutStr += "Block" + std::to_string(blockId) + ":\n"; + for (int warpId = 0; warpId < numWarpsPerCTA; warpId++) { + layoutStr += "Warp" + std::to_string(warpId) + ":\n"; + for (int idx = 0; idx < numElementsPerThreads; ++idx) { + for (int tid = 0; tid < threadsPerWarp; ++tid) { + int linearizedIdx = + blockId * numWarpsPerCTA * threadsPerWarp * + numElementsPerThreads + + warpId * threadsPerWarp * numElementsPerThreads + + tid * numElementsPerThreads + idx; + layoutStr += threadMapping[linearizedIdx]; + if (tid < threadsPerWarp - 1) + layoutStr += ", "; + } + layoutStr += "\n"; + } + } + } + } + return layoutStr; +} + +template +llvm::SmallVector +mlir::triton::gpu::expandMatrixShapeWithBatch(llvm::ArrayRef s) { + auto rank = s.size(); + assert(rank == 2 || rank == 3); + if (rank == 3) + return llvm::SmallVector(s); + return {1, s[0], s[1]}; +} + +template llvm::SmallVector +mlir::triton::gpu::expandMatrixShapeWithBatch( + llvm::ArrayRef s); + +template llvm::SmallVector +mlir::triton::gpu::expandMatrixShapeWithBatch( + llvm::ArrayRef s); + +llvm::SmallVector +mlir::triton::gpu::expandMatrixOrderWithBatch(llvm::ArrayRef o) { + int rank = o.size(); + assert(rank == 2 || rank == 3); + if (rank == 3) + return llvm::SmallVector(o); + llvm::SmallVector expanded(3, 0); + for (int i = 0; i < rank; ++i) + expanded[i] += o[i] + 1; + return expanded; +} + +std::string mlir::triton::gpu::getLayoutStr(RankedTensorType tensorType, + bool useHWPointOfView) { + auto layout = tensorType.getEncoding(); + LinearLayout ll = triton::gpu::toLinearLayout(tensorType.getShape(), layout); + + // tensorType is needed later on (e.g., getDimSize(j)), so we still have to + // pass it as a param + // TODO: Pass TensorOrMemDesc instead of RankedTensorType in + // triton-tensor-layout.cpp + if (mlir::isa(layout)) { + return getSharedLayoutStr(ll, useHWPointOfView); + } else if (mlir::isa(layout)) { + return getDistributedLayoutStr(ll, useHWPointOfView); + } + + // else unimplemented, return error + llvm::report_fatal_error("Unimplemented usage of getLayoutStr"); + return ""; +} + +void mlir::triton::gpu::dumpLayout(RankedTensorType tensorType) { + llvm::errs() << getLayoutStr(tensorType, /*useHWPointOfView=*/false); +} + +void mlir::triton::gpu::dumpHWLayout(RankedTensorType tensorType) { + llvm::errs() << getLayoutStr(tensorType, /*useHWPointOfView=*/true); +} + +namespace { +struct TensorModel + : public triton::gpu::TensorOrMemDesc::ExternalModel { + Type getElementType(Type pointer) const { + return cast(pointer).getElementType(); + } + Attribute getEncoding(Type pointer) const { + return cast(pointer).getEncoding(); + } + ArrayRef getShape(Type pointer) const { + return cast(pointer).getShape(); + } + int64_t getRank(Type pointer) const { + return cast(pointer).getRank(); + } + int64_t getElementTypeBitWidth(Type pointer) const { + return cast(pointer).getElementTypeBitWidth(); + } +}; + +struct MemDescModel + : public triton::gpu::TensorOrMemDesc::ExternalModel { + Type getElementType(Type pointer) const { + return cast(pointer).getElementType(); + } + Attribute getEncoding(Type pointer) const { + return cast(pointer).getEncoding(); + } + ArrayRef getShape(Type pointer) const { + return cast(pointer).getShape(); + } + int64_t getRank(Type pointer) const { + return cast(pointer).getShape().size(); + } + int64_t getElementTypeBitWidth(Type pointer) const { + return cast(pointer).getElementType().getIntOrFloatBitWidth(); + } +}; +} // namespace + +void TritonGPUDialect::initialize() { + registerTypes(); + + addAttributes< +#define GET_ATTRDEF_LIST +#include "triton/Dialect/TritonGPU/IR/AttrDefs.cpp.inc" + >(); + addOperations< +#define GET_OP_LIST +#include "triton/Dialect/TritonGPU/IR/Ops.cpp.inc" +#include "triton/Dialect/TritonGPU/IR/OpsEnums.cpp.inc" + >(); + addInterfaces(); + addInterfaces(); + addInterfaces(); + addInterfaces(); + + RankedTensorType::attachInterface(*getContext()); + MemDescType::attachInterface(*getContext()); +} + +LogicalResult TritonGPUDialect::verifyOperationAttribute(Operation *op, + NamedAttribute attr) { + // Verify that dialect attributes are attached to the right ops. + if (llvm::is_contained( + {AttrNumCTAsName, AttrTargetName, AttrNumThreadsPerWarp}, + attr.getName()) && + !isa(op)) { + return op->emitOpError("has unexpected attribute ") + << attr.getName() << " which is expected only on `module` ops"; + } + if (attr.getName() == AttrNumWarpsName && !isa(op)) { + return op->emitOpError("has unexpected attribute ") + << attr.getName() + << " which is expected only on `module` or `tt.func` ops"; + } + + // Verify that all ops in a tt.warp_specialize op have partition ids + if (attr.getName() == "tt.warp_specialize") { + if (!isa(op)) { + return op->emitOpError("has unexpected attribute ") + << attr.getName() << " which is expected only on `scf.for` ops"; + } + Operation *failedOp = nullptr; + op->walk([&](Operation *childOp) { + if (!childOp->hasAttr(kPartitionAttrName)) { + failedOp = childOp; + WalkResult::interrupt(); + } + }); + if (failedOp) { + return failedOp->emitOpError("does not have expected attribute ") + << kPartitionAttrName + << " which is expected on all child ops of an op with " + "attribute `tt.warp_specialize`"; + } + } + + // Verify that partition id lists are non-empty, sorted and have no duplicates + auto verifyPartitionIds = + [&](const ArrayRef &partitionIds) -> LogicalResult { + SetVector idSet; + for (auto id : partitionIds) { + if (idSet.contains(id)) + return op->emitOpError("has duplicated partition ids in attribute ") + << attr.getName(); + idSet.insert(id); + } + if (idSet.empty()) + return op->emitOpError("has no partition ids in attribute ") + << attr.getName(); + auto ids = idSet.takeVector(); + SmallVector sortedIds(ids.begin(), ids.end()); + std::sort(sortedIds.begin(), sortedIds.end()); + if (ids != sortedIds) + return op->emitOpError("partition ids not in sorted order in attribute ") + << attr.getName(); + return success(); + }; + + if (attr.getName() == kPartitionAttrName) { + auto result = verifyPartitionIds( + cast(attr.getValue()).asArrayRef()); + if (failed(result)) + return result; + } + if (attr.getName() == kPartitionOutputsAttrName) { + auto arrayAttr = cast(attr.getValue()); + for (auto idx = 0; idx < arrayAttr.size(); idx++) { + auto result = verifyPartitionIds( + cast(arrayAttr[idx]).asArrayRef()); + if (failed(result)) + return result; + } + } + + // Verify that op partitions include partitions of all child ops + if (attr.getName() == kPartitionAttrName && op->getNumRegions() != 0) { + SetVector expectedIds; + for (auto ®ion : op->getRegions()) { + for (auto &block : region.getBlocks()) { + for (auto &childOp : block.getOperations()) { + if (isa(childOp)) { + // yield ops and ub.poison do not need partition ids + continue; + } + if (!childOp.hasAttr(kPartitionAttrName)) + return childOp.emitOpError("does not have expected attribute ") + << kPartitionAttrName + << " which is expected for ops whose parent has partitions"; + auto ids = getPartitionIds(&childOp); + expectedIds.insert(ids.begin(), ids.end()); + } + } + } + auto partitionIds = getPartitionIds(op); + for (auto id : expectedIds) { + if (!partitionIds.contains(id)) { + return op->emitOpError("partition ids in attr ") + << attr.getName() + << " does not contain partition ids of all child ops"; + } + } + } + + if (attr.getName() == kPartitionOutputsAttrName) { + if (!isa(op)) + return op->emitOpError("has unexpected attribute ") << attr.getName(); + + // Verify that number of output partitions matches number of For/If results + size_t numResults = 0; + if (isa(op)) { + numResults = cast(op).getResults().size(); + } else if (isa(op)) { + numResults = cast(op).getResults().size(); + } else { + numResults = cast(op).getResults().size(); + } + + if (cast(attr.getValue()).size() != numResults) { + return op->emitOpError("does not have expected number of output " + "partition sets in attr ") + << attr.getName() << "; should match number of results"; + } + + // Verify that union of op output partitions is a subset of op partitions + if (!op->hasAttr(kPartitionAttrName)) + return op->emitOpError("does not have expected attribute ") + << kPartitionAttrName << " which is expected for ops with attr " + << kPartitionOutputsAttrName; + auto partitionIds = getPartitionIds(op); + + SetVector outputPartitionIdsUnion; + for (auto outputPartitionIds : getPartitionOutputs(op)) { + outputPartitionIdsUnion.insert(outputPartitionIds.begin(), + outputPartitionIds.end()); + } + if (!std::all_of(outputPartitionIdsUnion.begin(), + outputPartitionIdsUnion.end(), + [&](int id) { return partitionIds.contains(id); })) { + return op->emitOpError("partition ids in attr ") + << kPartitionAttrName + << " must be the union of all partition ids in " << attr.getName(); + } + } + + return success(); +} + +int TritonGPUDialect::getNumCTAs(ModuleOp module) { + if (auto attr = module->getAttrOfType(AttrNumCTAsName)) + return attr.getInt(); + return 1; +} + +int TritonGPUDialect::getThreadsPerWarp(ModuleOp module) { + if (auto attr = module->getAttrOfType(AttrNumThreadsPerWarp)) + return attr.getInt(); + return 32; +} + +std::optional triton::gpu::maybeLookupNumWarps(Operation *op) { + if (isa(op)) { + if (auto attr = op->getAttrOfType(AttrNumWarpsName)) + return attr.getInt(); + } else if (auto partitions = + dyn_cast(op->getParentOp())) { + unsigned idx = op->getParentRegion()->getRegionNumber(); + return partitions.getParentOp().getPartitionNumWarps()[idx]; + } + if (Operation *parent = op->getParentOp()) + return maybeLookupNumWarps(parent); + return {}; +} + +int triton::gpu::lookupNumWarps(Operation *op) { + std::optional numWarps = maybeLookupNumWarps(op); + if (!numWarps) { + op->emitOpError( + "is not contained within a context that specifies the number of warps"); + llvm::report_fatal_error("failed to lookup the number of warps, the " + "surrounding module should contain a " + + Twine(AttrNumWarpsName) + " attribute"); + } + return *numWarps; +} + +int triton::gpu::lookupNumWarps(Region *region) { + if (auto partitions = + dyn_cast(region->getParentOp())) { + unsigned idx = region->getRegionNumber(); + return partitions.getParentOp().getPartitionNumWarps()[idx]; + } + return lookupNumWarps(region->getParentOp()); +} + +int triton::gpu::lookupThreadsPerWarp(OpBuilder &rewriter) { + assert(rewriter.getInsertionBlock() && "expected an insertion point"); + Operation *op = + rewriter.getInsertionBlock()->getParentOp()->getParentOfType(); + assert(op && "cannot check threads per warp outside of module"); + return triton::gpu::TritonGPUDialect::getThreadsPerWarp(cast(op)); +} + +int triton::gpu::lookupNumCTAs(Operation *op) { + auto mod = op->getParentOfType(); + if (!mod) { + op->emitOpError( + "is not contained within a module, cannot lookup number of CTAs"); + llvm::report_fatal_error( + "failed to lookup the number of CTAs, the surrounding module should " + "contain a ModuleOp"); + } + return triton::gpu::TritonGPUDialect::getNumCTAs(mod); +} + +int triton::gpu::lookupNumCTAs(OpBuilder &rewriter) { + assert(rewriter.getInsertionBlock() && "expected an insertion point"); + Operation *op = + rewriter.getInsertionBlock()->getParentOp()->getParentOfType(); + assert(op && "cannot check number of CTAs outside of module"); + return triton::gpu::TritonGPUDialect::getNumCTAs(cast(op)); +} + +bool triton::gpu::areLayoutsEquivalent(ArrayRef shape, + LayoutEncodingTrait lhs, + LayoutEncodingTrait rhs) { + auto lhsLL = triton::gpu::toLinearLayout(shape, lhs); + auto rhsLL = triton::gpu::toLinearLayout(shape, rhs); + return lhsLL == rhsLL; +} + +bool triton::gpu::isInnermostContiguous(MemDescType type, unsigned numElems) { + ArrayRef shape = type.getShape(); + Attribute enc = type.getEncoding(); + MLIRContext *ctx = enc.getContext(); + + LinearLayout actual = toLinearLayout(type); + StringAttr fastestIn = *actual.getInDimNames().begin(); + + // Flatten actual outs in reverse order to produce a row-major flattening + // of the layout + auto outNames = actual.getOutDimNames(); + SmallVector revOut(outNames.begin(), outNames.end()); + std::reverse(revOut.begin(), revOut.end()); + actual = actual.transposeOuts(revOut).flattenOuts(); + + return actual.getNumConsecutiveInOut() >= numElems; +} + +LinearLayout triton::gpu::inferReshapeLinearLayout(TensorOrMemDesc srcTy, + ArrayRef dstShape) { + auto *ctx = srcTy.getContext(); + auto src = toLinearLayout(srcTy); + assert(product(srcTy.getShape()) == product(dstShape)); + auto dst = reshapeLayout(ctx, src, dstShape); + return dst; +} + +SetVector triton::gpu::getPartitionIds(Operation *op) { + auto attrs = op->getAttr(kPartitionAttrName); + SmallVector partitionIds; + for (auto id : cast(attrs).asArrayRef()) { + partitionIds.push_back(id); + } + std::sort(partitionIds.begin(), partitionIds.end()); + return SetVector(partitionIds.begin(), partitionIds.end()); +} + +SmallVector, 4> triton::gpu::getPartitionOutputs(Operation *op) { + SmallVector, 4> partitionOutputsIds; + if (op->getNumResults() == 0) { + return partitionOutputsIds; + } + auto arrayAttr = cast(op->getAttr(kPartitionOutputsAttrName)); + for (auto attr : arrayAttr) { + auto ids = cast(attr).asArrayRef(); + partitionOutputsIds.push_back(SetVector(ids.begin(), ids.end())); + } + return partitionOutputsIds; +} + +SetVector triton::gpu::getPartitionIds(OpOperand *use) { + auto owner = use->getOwner(); + if (isa(owner)) { + return getPartitionOutputs(owner->getParentOp())[use->getOperandNumber()]; + } else if (scf::ForOp forOp = dyn_cast(owner)) { + int idx = use->getOperandNumber() - forOp.getNumControlOperands(); + return idx >= 0 ? getPartitionOutputs(owner)[idx] : getPartitionIds(forOp); + } else { + return getPartitionIds(owner); + } +} + +bool triton::gpu::hasPartition(Operation *op) { + return op && op->hasAttr(kPartitionAttrName); +} + +bool triton::gpu::hasWarpSpecializeTag(Operation *op) { + return op && op->hasAttr(kWarpSpecializeTagAttrName); +} + +std::optional triton::gpu::getWarpSpecializeTag(Operation *op) { + if (hasWarpSpecializeTag(op)) { + return cast(op->getAttr(kWarpSpecializeTagAttrName)).getInt(); + } + return std::nullopt; +} diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp new file mode 100644 index 0000000000..457196235e --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp @@ -0,0 +1,1827 @@ +#include + +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Attributes.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h" +#include "triton/Dialect/TritonGPU/IR/TritonGPUInterfaces.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +#include "triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h" +#include "triton/Tools/LayoutUtils.h" +#include "triton/Tools/LinearLayout.h" +#include "triton/Tools/StrUtil.h" +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/Twine.h" +#include "llvm/Support/ErrorHandling.h" +#include "llvm/Support/MathExtras.h" + +using mlir::triton::nvidia_gpu::TensorMemoryEncodingAttr; +using mlir::triton::nvidia_gpu::TensorMemoryScalesEncodingAttr; + +namespace mlir::triton::gpu { +namespace { + +// We use the following nomenclature in this file. +// +// - ctaLayout: A layout for one block, i.e. input dims [register, lane, warp] +// for register layouts, and input dims [offset] for shared layouts. +// - cgaLayout: Arrangement of multiple blocks, i.e. input dims [block]. +// +// Note that this is inconsistent with the type name CTAEncodingAttr. That type +// is equivalent to our cgaLayout. +// +// IMO the name CTAEncodingAttr is wrong. If we tried to be consistent anyway, +// then we'd have to rename ctaLayout to "warpLayout". I think that's more +// confusing than being inconsistent about "cgaLayout", especially when we have +// to consider the size of the warpLayout (surely that's not the "warpSize"). + +#define S(v) StringAttr::get(ctx, (v)) + +SmallVector getDefaultMmaOrder(MmaEncodingTrait layout) { + auto rank = layout.getRepOrderForOperand(0).size(); + return getMatrixOrder(rank, /*rowMajor*/ true); +} + +// TODO Have order be a mandatory argument of standardOutDimNames. +SmallVector permuteDimNames(const SmallVector &names, + const SmallVector &order) { + assert(names.size() == order.size()); + SmallVector ret; + for (unsigned i : order) { + ret.push_back(names[i]); + } + return ret; +} + +LinearLayout swizzledSharedToLinearLayout(ArrayRef shape, + SwizzledSharedEncodingAttr shared) { + MLIRContext *ctx = shared.getContext(); + + auto shapePerCTA = getShapePerCTA(shared, shape); + + int rank = shape.size(); + if (rank == 1) { + return combineCtaCgaWithShape( + LinearLayout::identity1D(shapePerCTA[0], S("offset"), S("dim0")), + shared.getCTALayout(), shape); + } + + auto outDimNames = standardOutDimNames(ctx, rank); + + // Construct bases for the 2 most minor dimensions of the layout. These are + // the dims that get swizzled. + assert(shape.size() >= 2); + int colDim = shared.getOrder()[0]; + int rowDim = shared.getOrder()[1]; + int numCols = shapePerCTA[colDim]; + int numRows = shapePerCTA[rowDim]; + StringAttr colDimName = outDimNames[colDim]; + StringAttr rowDimName = outDimNames[rowDim]; + + std::vector> bases2D; + for (int col = 1; col < numCols; col *= 2) { + bases2D.push_back({0, col}); + } + for (int row = 1; row < numRows; row *= 2) { + int vec = shared.getVec(); + int perPhase = shared.getPerPhase(); + int maxPhase = shared.getMaxPhase(); + bases2D.push_back({row, (vec * ((row / perPhase) % maxPhase)) % numCols}); + } + LinearLayout ctaLayout = + LinearLayout({{S("offset"), bases2D}}, {rowDimName, colDimName}); + + // Add the remaining dimensions. + for (int i = 2; i < rank; i++) { + int dim = shared.getOrder()[i]; + ctaLayout *= LinearLayout::identity1D(shapePerCTA[dim], S("offset"), + outDimNames[dim]); + } + + return combineCtaCgaWithShape(ctaLayout, shared.getCTALayout(), shape); +} + +LinearLayout +sharedToLinearLayoutAMDRotating(ArrayRef shape, + AMDRotatingSharedEncodingAttr shared) { + MLIRContext *ctx = shared.getContext(); + + auto shapePerCTA = getShapePerCTA(shared, shape); + + int rank = shape.size(); + if (rank == 1) { + return combineCtaCgaWithShape( + LinearLayout::identity1D(shapePerCTA[0], S("offset"), S("dim0")), + shared.getCTALayout(), shape); + } + + auto outDimNames = standardOutDimNames(ctx, rank); + + // Construct bases for the 2 most minor dimensions of the layout. These are + // the dims that get swizzled. + assert(shape.size() >= 2); + int colDim = shared.getOrder()[0]; + int rowDim = shared.getOrder()[1]; + int numCols = shape[colDim]; + int numRows = shape[rowDim]; + StringAttr colDimName = outDimNames[colDim]; + StringAttr rowDimName = outDimNames[rowDim]; + + std::vector> bases2D; + for (int col = 1; col < numCols; col *= 2) { + bases2D.push_back({0, col}); + } + for (int row = 1; row < numRows; row *= 2) { + int vec = shared.getVec(); + int perPhase = shared.getPerPhase(); + int maxPhase = shared.getMaxPhase(); + + int phase = (row / perPhase) % maxPhase; + int blockNo = row / maxPhase / perPhase % maxPhase; + int combinedPhase = phase ^ blockNo; + bases2D.push_back({row, (vec * combinedPhase) % numCols}); + } + LinearLayout ctaLayout = + LinearLayout({{S("offset"), bases2D}}, {rowDimName, colDimName}); + + // Add the remaining dimensions. + for (int i = 2; i < rank; i++) { + int dim = shared.getOrder()[i]; + ctaLayout *= + LinearLayout::identity1D(shape[dim], S("offset"), outDimNames[dim]); + } + + return combineCtaCgaWithShape(ctaLayout, shared.getCTALayout(), shape); +} + +} // namespace + +// Returns the layout of a single core matrix which tiles the nvmma layout +LinearLayout getCoreMatrixLinearLayout(NVMMASharedEncodingAttr shared, + bool disableSwizzle) { + auto *ctx = shared.getContext(); + + int elemBitWidth = shared.getElementBitWidth(); + int tileWidthBytes = shared.getSwizzlingByteWidth(); + int vec = shared.getVec(); + int perPhase = shared.getPerPhase(); + int maxPhase = shared.getMaxPhase(); + + int tileRows = 8; + int tileCols = 8 * std::max(16, tileWidthBytes) / elemBitWidth; + bool isFp4Padded = shared.getFp4Padded(); + + std::vector> bases2D; + for (int col = 1; col < tileCols; col *= 2) { + if (isFp4Padded) { + // Each group of 16 offsets consists of 8 "real" and 8 "padded" offsets. + // We represent the padded layout by mapping 8 padded offsets to the same + // coordinates as the real ones. When computing the inverse of this LL, + // the offsets correspoding to the real ones are picked in the image by + // invertAndCompose. + int colPacked = col / 16 * 8 + col % 8; + bases2D.push_back({0, colPacked}); + } else { + bases2D.push_back({0, col}); + } + } + for (int row = 1; row < tileRows; row *= 2) { + if (disableSwizzle) { + bases2D.push_back({row, 0}); + } else if (isFp4Padded) { + int colPadded = vec * ((row / perPhase) % maxPhase); + int colPacked = colPadded / 16 * 8 + colPadded % 8; + bases2D.push_back({row, colPacked}); + } else { + bases2D.push_back({row, vec * ((row / perPhase) % maxPhase)}); + } + } + auto outDimNames = standardOutDimNames(ctx, 2); + return LinearLayout({{S("offset"), bases2D}}, outDimNames); +} + +LinearLayout nvmmaSharedToLinearLayout(ArrayRef shape, + NVMMASharedEncodingAttr shared, + bool disableSwizzle) { + MLIRContext *ctx = shared.getContext(); + int rank = shape.size(); + auto shapePerCTA = getShapePerCTA(shared, shape); + auto kOffset = S("offset"); + auto tmaShape = triton::nvidia_gpu::getTMABlockShape(shared, shapePerCTA, + /*packedSize=*/true); + if (shared.getSwizzlingByteWidth() == 0) { + auto outDimNames = standardOutDimNames(ctx, rank); + LinearLayout layout = LinearLayout::identity1D(tmaShape[rank - 1], kOffset, + outDimNames[rank - 1]); + for (int i = rank - 2; i >= 0; --i) { + layout *= LinearLayout::identity1D(tmaShape[i], kOffset, outDimNames[i]); + } + layout = ensureLayoutNotSmallerThan(layout, outDimNames, shapePerCTA); + return combineCtaCgaWithShape(layout, shared.getCTALayout(), shape); + } + assert(rank >= 2); + + // Collapse all the outer dim into one. We will then create a layout for this + // shape and reshape it to the original shape. + std::array collapsedTmaShape{1, tmaShape.back()}; + for (int i = 0; i + 1 < rank; i++) + collapsedTmaShape[0] *= tmaShape[i]; + if (shared.getTransposed()) { + std::swap(collapsedTmaShape[0], collapsedTmaShape[1]); + } + + auto tileLayout = getCoreMatrixLinearLayout(shared, disableSwizzle); + auto outDimNames = standardOutDimNames(ctx, 2); + auto kRow = outDimNames[0]; + auto kCol = outDimNames[1]; + auto tileRows = tileLayout.getOutDimSize(kRow); + auto tileCols = tileLayout.getOutDimSize(kCol); + + int packingFactor = shared.getFp4Padded() ? 2 : 1; + if (collapsedTmaShape[1] * packingFactor < tileCols || + collapsedTmaShape[0] < tileRows) { + llvm::errs() << "Illegal shared layout; expected collapsed shapePerCTA to " + "be at least [" + << tileRows << ", " << (tileCols / packingFactor) + << "], collapsedTmaShape: [" << collapsedTmaShape[0] << ", " + << collapsedTmaShape[1] << "]\n"; + llvm::report_fatal_error("Illegal shared layout"); + } + + // Distribute the remaining rows and cols. + auto layout = + ensureLayoutNotSmallerThan(tileLayout, outDimNames, collapsedTmaShape); + + // Reshape the layout to the N-D pre-transposed shape per CTA. + SmallVector maybeTransposedTmaShape = tmaShape; + if (shared.getTransposed()) { + // Move the outer dim to the inner position. + // TODO: we should move back to using `order` instead of transposed to make + // the order more explicit. + std::rotate(maybeTransposedTmaShape.begin(), + maybeTransposedTmaShape.begin() + 1, + maybeTransposedTmaShape.end()); + } + auto reshapedLayout = reshapeLayout(ctx, layout, maybeTransposedTmaShape); + + if (shared.getTransposed()) { + SmallVector order = {rank - 1}; + for (int i = 0; i < rank - 1; i++) { + order.push_back(i); + } + reshapedLayout = transposeLinearLayout(reshapedLayout, order); + } + + reshapedLayout = ensureLayoutNotSmallerThan( + reshapedLayout, standardOutDimNames(ctx, shapePerCTA.size()), + shapePerCTA); + return combineCtaCgaWithShape(reshapedLayout, shared.getCTALayout(), shape); +} + +/// Function to generate lane and warp layout for dot operands. +static LinearLayout broadcastedDotOperandLayout(MLIRContext *ctx, + ArrayRef shape, + ArrayRef order, + unsigned kDim, + StringAttr inDimName) { + // Let warpsPerCTAMma = {2, 2}, then + // warpsPerCTA = {2, 1} for opA and warpsPerCTA = {1, 2} for opB + // assume warpOrder = {1, 0} + // Assume that C is tiled by 2x2 tiles. Since warpOrder={1, 0}, we have that + // the C is owned as per the following layout: + // C: 0 | 1 + // - | - + // 2 | 3 + // In order to be able to compute C, we need the following warp tiling of + // A and B: + // A: 0 1 | 0 1 B: 0 2 | 1 3 + // - - | - - - - | - - + // 2 3 | 2 3 0 2 | 1 3 + // In other words, we need to broadcast along K + auto rank = shape.size(); + auto dimNames = standardOutDimNames(ctx, rank); + LinearLayout layout = LinearLayout::empty(); + + // We have to broadcast along the inner dimension + // For A, when moving along M we go from 0 to 2. + // For B, when moving along N we go from 0 to 1. + // As such, choosing the order of A {1, 0}, gives us the correct broadcasting + // Same happens if the warpOrder is {0, 1}, like in Hopper + for (auto d : order) { + if (d == kDim) { + layout *= LinearLayout::zeros1D(shape[d], inDimName, dimNames[d]); + } else { + layout *= LinearLayout::identity1D(shape[d], inDimName, dimNames[d]); + } + } + return layout; +} + +LinearLayout +AMDMfmaEncodingAttr::toLinearLayout(ArrayRef shape) const { + int rank = shape.size(); + assert(rank == getRank()); + + bool hasBatchDim = rank == 3; + int mIndex = 0 + hasBatchDim; + int nIndex = 1 + hasBatchDim; + (void)mIndex, (void)nIndex; + + MLIRContext *ctx = getContext(); + SmallVector outDimNames = standardOutDimNames(ctx, rank); + + StringAttr kRegister = S("register"); + StringAttr kLane = S("lane"); + StringAttr kWarp = S("warp"); + + // https://github.com/ROCm/amd_matrix_instruction_calculator can print the + // register and lane layout for mfma instructions. + + // We use the order from fastest varying to slowest varying. So each base + // vector is a tuple of values mapping to matrix C's (N, M[, B]) indices, + // which will be [1, 0] / [2, 1, 0]. + SmallVector order = getDefaultMmaOrder(*this); + auto dimM = outDimNames[order[1]]; + auto dimN = outDimNames[order[0]]; + + auto mDim = getInstrShape()[0]; + auto nDim = getInstrShape()[1]; + auto elementBitWidth = getElementBitWidth(); + int height = elementBitWidth == 64 ? 1 : 4; + constexpr int warpSize = 64; + + bool isTransposed = getIsTransposed(); + // Special case for 64x4 mfma: we always transpose the output to turn + // the 64x4 mfma into a equalvalent 4x64 mfma and swap operand A and B, so + // that we can use the mfma broadcast. + if (mDim == 64 && nDim == 4) + assert(isTransposed && "64x4 mfma must be transposed"); + + int tiles = (mDim * nDim) / (warpSize * height); + + LinearLayout tileLayout = LinearLayout::empty(); + if (!isTransposed) { + // Each lane holds 'height' elements along the M dimension. + LinearLayout regs = LinearLayout::identity1D(height, kRegister, dimM); + // First, distribute the lanes along the N dimension. + // Then, distribute the lanes along the M dimension. If the #elements + // exceeds the mDim, duplicate elements across lanes - this can happen for + // 4x4 output. + LinearLayout lanes = LinearLayout::identity1D(nDim, kLane, dimN) * + LinearLayout::identity1D(warpSize / nDim, kLane, dimM); + tileLayout = (regs * lanes); + + // Repeat the above distribution along the M dimension to fits the tile. + if (tiles > 0) + tileLayout *= LinearLayout::identity1D(tiles, kRegister, dimM); + } else { + // For the transposed output, we will use the same method for layout but + // swap the order of the M and N dimensions. + LinearLayout regs = LinearLayout::identity1D(height, kRegister, dimN); + LinearLayout lanes = LinearLayout::identity1D(mDim, kLane, dimM) * + LinearLayout::identity1D(warpSize / mDim, kLane, dimN); + tileLayout = (regs * lanes); + + if (tiles > 0) + tileLayout *= LinearLayout::identity1D(tiles, kRegister, dimN); + } + + tileLayout = tileLayout.transposeOuts({dimN, dimM}); + + // Instead of defining the layout on a CTA tile and using the + // combineCtaCgaWithShape function to extend it to the whole tensor, we take a + // different approach. Suppose tilesPerWarp is 2x2—meaning a warp computes a + // 2x2 block of MFMA tiles. If we define the layout only on the CTA tile and + // extend it across the tensor, the resulting tile order won’t be N-contiguous + // (i.e., row-major). Due to the 2x2 shape, the third tile would fall in the M + // dimension. While defining the layout per CTA tile might seem more + // intuitive, the current dot op lowering assumes an N-contiguous ordering of + // MFMA tiles across the entire tensor. In other words, the lowering logic + // isn't layout-aware, it only supports a fixed N-contiguous MFMA tile + // ordering. Supporting other orderings would require extending the dot + // lowering implementation. For now, we conform to the current lowering + // algorithm by defining the MFMA linear layout globally, with N-contiguous + // tiles across the tensor and across CTA tile boundaries. + auto tilesPerWarp = getTilesPerWarp(); + auto warpsPerCTA = getWarpsPerCTA(); + + const unsigned tilesPerWarpM = tilesPerWarp[mIndex]; + const unsigned tilesPerWarpN = tilesPerWarp[nIndex]; + const unsigned warpsPerCTAM = warpsPerCTA[mIndex]; + const unsigned warpsPerCTAN = warpsPerCTA[nIndex]; + + // First, extend the layout along the N dimension: + // - registers are distributed across tilesPerWarpN + // - then across warpsPerCTAN in the N dimension. + tileLayout *= LinearLayout::identity1D(tilesPerWarpN, kRegister, dimN); + tileLayout *= LinearLayout::identity1D(warpsPerCTAN, kWarp, dimN); + + // At this point, the layout is defined across the N dimension within a CTA + // tile. Instead of switching to the M dimension now, we continue extending + // the layout along the remaining N dimension, and only then proceed along M, + // following the tilesPerWarp configuration. + // If the N dimension is not large enough to span multiple CTA tiles (i.e., + // the first argument is 0), an empty layout is created, so this identity + // layout will not introduce any new registers. + tileLayout *= LinearLayout::identity1D( + shape[nIndex] / (nDim * warpsPerCTAN * tilesPerWarpN), kRegister, dimN); + tileLayout *= LinearLayout::identity1D(tilesPerWarpM, kRegister, dimM); + + // Finally, extend the layout across warps in the M dimension. + // After this step, the layout covers a sub-tensor of size ctaTileM × N, + // i.e., the full N dimension and a CTA tile's extent in M. + // The rest of the layout will be defined by combineCtaCgaWithShape. + tileLayout *= LinearLayout::identity1D(warpsPerCTAM, kWarp, dimM); + + // Adjust spatial ordering if batch dimension is present + if (hasBatchDim) { + assert(order[2] == 0); + // Extend the base vector with one value to accommodate for the batch + // dimension, which appears at the last. + tileLayout *= LinearLayout::identity1D(1, kRegister, outDimNames[order[2]]); + tileLayout *= LinearLayout::identity1D(1, kLane, outDimNames[order[2]]); + tileLayout *= + LinearLayout::identity1D(warpsPerCTA[0], kWarp, outDimNames[order[2]]); + } + + return combineCtaCgaWithShape(tileLayout, getCTALayout(), shape); +} + +std::optional +chooseDotDsReadTrLayout(DotOperandEncodingAttr dotMfmaLayout, + ArrayRef shape, int32_t elemBitWidth, + unsigned instBitWidth, + unsigned numLanesInShuffleGroup) { + if (instBitWidth != 64 || numLanesInShuffleGroup != 16) + return std::nullopt; + auto mfmaLayout = llvm::cast(dotMfmaLayout.getParent()); + auto mDim = mfmaLayout.getInstrShape()[0]; + assert(mDim == 16 || mDim == 32); + + assert(elemBitWidth == 4); + // When doing ds_read_tr4 we actually write the LL as if it were on i8 + // elements this is becasue LL needs to be described for the i8 tensor + // elements. + elemBitWidth = 8; + + auto rank = shape.size(); + bool hasBatchDim = rank == 3; + int32_t kWidthDot = dotMfmaLayout.getKWidth(); + auto kDim = dotMfmaLayout.getOpIdx() == 0 ? rank - 1 : rank - 2; + + int32_t kSize = shape[kDim]; + auto warpsPerCTA = mfmaLayout.getWarpsPerCTA(); + + MLIRContext *ctx = dotMfmaLayout.getContext(); + SmallVector outDimNames = standardOutDimNames(ctx, rank); + + StringAttr kRegister = S("register"); + StringAttr kLane = S("lane"); + StringAttr kWarp = S("warp"); + + // register order + // operand A: [1, 0] / [2, 1, 0] + // operand B: [0, 1] / [1, 2, 0] + // Regular dot mfma order for both cases is [k, nonk]/[k, nonk, batch] + // For LDS transpose layout swap order to [nonk, k]/[nonk, k, batch] + SmallVector order = + getOrderForDotOperand(dotMfmaLayout.getOpIdx(), rank, /*kContig*/ false); + + std::vector> registerBase; + std::vector> laneBase; + + const bool isMfma32 = (mDim == 32); + // ds_read_b64_tr4 operates on FP4 values swapping the packing of them. Look + // at i8 values for the ownership of register/lane since it's the data type + // of the tensor. Register dimension: what i8 in the tile are held by thread + // 0? Lane dimension: what i8 in the tile are held in register 0 of each + // thread? + registerBase.push_back({1, 0}); + registerBase.push_back({2, 0}); + registerBase.push_back({4, 0}); + registerBase.push_back({0, 16}); + + // If more than one tile needs to be loaded, populate registerBase + // dimension for the other tiles + const int kTileSize = isMfma32 ? 64 : 128; + for (int reg = kTileSize; reg < kSize; reg *= 2) { + registerBase.push_back({0, reg}); + } + + // When mDim == 16 we have 16x128 mfma, otherwise it's 16x64 + // The LL for the two is different + laneBase.push_back({0, 1}); + laneBase.push_back({0, 2}); + laneBase.push_back({0, 4}); + laneBase.push_back({0, 8}); + if (mDim == 16) { + laneBase.push_back({0, 32}); + laneBase.push_back({0, 64}); + } else { + assert(mDim == 32); + laneBase.push_back({8, 0}); + laneBase.push_back({0, 32}); + } + + // Base vectors above are defined in a fixed order [non-k-dim, k-dim]. + // To assign them to actual matrix dimensions we associate with register + // `order` which is also [nonk, k] given we set kContig to false. + LinearLayout tileLayout({{kRegister, registerBase}, {kLane, laneBase}}, + {outDimNames[order[0]], outDimNames[order[1]]}); + if (hasBatchDim) { + assert(order[2] == 0); + // Extend the base vector with one value to accommodate for the batch + // dimension, which appears at the last. + tileLayout *= LinearLayout::identity1D(1, kRegister, outDimNames[order[2]]); + tileLayout *= LinearLayout::identity1D(1, kLane, outDimNames[order[2]]); + } + + // warp order + // common for both operand A and B: [0, 1] / [0, 1, 2] + // in both cases it is [M dim, N dim]/[batch, M dim, N dim] + auto warpOrder = getDefaultMmaOrder(mfmaLayout); + LinearLayout warpLayout = identityStandardND(kWarp, warpsPerCTA, warpOrder); + + LinearLayout ctaLayout = tileLayout.transposeOuts(outDimNames) * + warpLayout.transposeOuts(outDimNames); + return combineCtaCgaWithShape(ctaLayout, mfmaLayout.getCTALayout(), shape); +} + +LinearLayout mfmaDotToLinearLayout(DotOperandEncodingAttr dotMfmaLayout, + ArrayRef shape) { + auto mfmaLayout = llvm::cast(dotMfmaLayout.getParent()); + + auto rank = shape.size(); + bool hasBatchDim = rank == 3; + int mIndex = 0 + hasBatchDim; + + int32_t kWidth = dotMfmaLayout.getKWidth(); + auto nonKDimIndex = dotMfmaLayout.getOpIdx() == 0 ? rank - 2 : rank - 1; + + auto warpsPerCTA = mfmaLayout.getWarpsPerCTA(); + auto tilesPerWarp = mfmaLayout.getTilesPerWarp(); + auto tilePerWarpNonK = tilesPerWarp[nonKDimIndex]; + + auto mDim = mfmaLayout.getInstrShape()[0]; + auto nDim = mfmaLayout.getInstrShape()[1]; + auto opIdx = dotMfmaLayout.getOpIdx(); + auto nonKDim = opIdx == 0 ? mDim : nDim; + constexpr int warpSize = 64; + + auto kDimIndex = dotMfmaLayout.getOpIdx() == 0 ? rank - 1 : rank - 2; + int32_t kSize = shape[kDimIndex]; + + MLIRContext *ctx = dotMfmaLayout.getContext(); + SmallVector outDimNames = standardOutDimNames(ctx, rank); + + StringAttr kRegister = S("register"); + StringAttr kLane = S("lane"); + StringAttr kWarp = S("warp"); + + // register order + // operand A: [1, 0] / [2, 1, 0] + // operand B: [0, 1] / [1, 2, 0] + // for both cases it is [k, nonk]/[k, nonk, batch] + auto order = + getOrderForDotOperand(dotMfmaLayout.getOpIdx(), rank, /*kContig*/ true); + auto dimK = outDimNames[order[0]]; + auto dimNonK = outDimNames[order[1]]; + + // warp order + // common for both operand A and B: [0, 1] / [0, 1, 2] + // in both cases it is [M dim, N dim]/[batch, M dim, N dim] + auto warpOrder = getDefaultMmaOrder(mfmaLayout); + + // Each lane holds kWidth elements along the K dimension + LinearLayout regs = LinearLayout::identity1D(kWidth, kRegister, dimK); + // First distribute nonKDim elements along the non-K dimension, + // then distribute remaining elements along the K dimension + LinearLayout lanes = + LinearLayout::identity1D(nonKDim, kLane, dimNonK) * + LinearLayout::identity1D(warpSize / nonKDim, kLane, dimK); + LinearLayout tileLayout = regs * lanes; + + int kTileSize = warpSize / nonKDim * kWidth; + // Special case for 4x64 and 64x4 mfma: for the 64x64 operand, + // we need to repeat the layout 16 times along the K dimension + if ((mDim == 64 && nDim == 4 && opIdx == 0) || + (mDim == 4 && nDim == 64 && opIdx == 1)) { + tileLayout *= LinearLayout::identity1D(16, kRegister, dimK); + kTileSize *= 16; + } + + // If shape K is larger than the tile size, repeat the tile + // along the K dimension. + if (kSize > kTileSize) { + tileLayout *= LinearLayout::identity1D(kSize / kTileSize, kRegister, dimK); + } + + // Follow the tiles per warp property, repeat the tile layout + // along the non-K dimension. + tileLayout *= LinearLayout::identity1D(tilePerWarpNonK, kRegister, dimNonK); + + tileLayout = tileLayout.transposeOuts({dimK, dimNonK}); + if (hasBatchDim) { + assert(order[2] == 0); + // Extend the base vector with one value to accommodate for the batch + // dimension, which appears at the last. + tileLayout *= LinearLayout::identity1D(1, kRegister, outDimNames[order[2]]); + tileLayout *= LinearLayout::identity1D(1, kLane, outDimNames[order[2]]); + } + + LinearLayout warpLayout = identityStandardND(kWarp, warpsPerCTA, warpOrder); + LinearLayout ctaLayout = tileLayout * warpLayout; + + // Note the current the output order is [k, nonk]/[k, nonk, batch]. If the + // layout's out-size is smaller than the shape, we follow this order to + // extend each dimension to match the shape. After that, we can transpose + // to match the standard output order. + return combineCtaCgaWithShape(ctaLayout, mfmaLayout.getCTALayout(), shape) + .transposeOuts(outDimNames); +} + +LinearLayout +AMDWmmaEncodingAttr::toLinearLayout(ArrayRef shape) const { + int rank = shape.size(); + assert(rank == getRank()); + + bool hasBatchDim = rank == 3; + int mIndex = 0 + hasBatchDim; + int nIndex = 1 + hasBatchDim; + (void)mIndex, (void)nIndex; + + auto mnkDim = getInstrShape(); + unsigned mDim = mnkDim[0], nDim = mnkDim[1]; + (void)mDim, (void)nDim; + + assert(((shape[mIndex] == 1 || shape[mIndex] >= mDim) && + (shape[nIndex] == 1 || shape[nIndex] >= nDim)) && + "Unsupported tensor shape for given wmma layout"); + + MLIRContext *ctx = getContext(); + SmallVector outDimNames = standardOutDimNames(ctx, rank); + + StringAttr kRegister = S("register"); + StringAttr kLane = S("lane"); + StringAttr kWarp = S("warp"); + + // https://github.com/ROCm/amd_matrix_instruction_calculator can print the + // register and lane layout for mfma instructions. + + // We use the order from fastest varying to slowest varying. So each base + // vector is a tuple of values mapping to matrix C's (N, M[, B]) indices. + auto threadOrder = getMatrixOrder(rank, /*rowMajor*/ !getIsTransposed()); + assert(threadOrder[0] == mIndex || threadOrder[0] == nIndex); + assert(threadOrder[1] == mIndex || threadOrder[1] == nIndex); + + // For wmma with 16x16 output, each of the 32 threads holds 8 elements. + // + // The first version of WMMA layout has following specific: + // for the register (i.e., element) dimension, these 8 elements are + // along the matrix C's M dimension, with 1 consecutive elements + // spanning 1 row and then the next 1 row being a gap. + // + // For the lane (i.e., thread) dimension, these threads are along the + // matrix C's N dimension, with 16 consecutive threads covering a whole + // row and the next 16 threads start at the next row. + // + // The second version of wmma layout is less tricky: + // for the register dimension 8 elements are along the matrix C's M + // dimension. First 16 lanes take 0-8 elems along M, second 16 take 8-15. + // We have 16 pair of threads in each warp, one pair covers the whole + // column. + // + // Please also check explaining comments in TritonGPUAttrDefs.td at the + // AMDWmmaEncodingAttr section. + unsigned version = getVersion(); + assert(version >= 1 && version <= 3 && "unexpected wmma version"); + LinearLayout tileLayout = + version == 1 + ? LinearLayout( + {{kRegister, {/*gap*/ {0, 2}, {0, 4}, {0, 8}}}, + {kLane, {{1, 0}, {2, 0}, {4, 0}, {8, 0}, /*gap*/ {0, 1}}}}, + {outDimNames[threadOrder[0]], outDimNames[threadOrder[1]]}) + : LinearLayout( + {{kRegister, {{0, 1}, {0, 2}, {0, 4}}}, + {kLane, {{1, 0}, {2, 0}, {4, 0}, {8, 0}, /*gap*/ {0, 8}}}}, + {outDimNames[threadOrder[0]], outDimNames[threadOrder[1]]}); + + auto tilesPerWarp = getTilesPerWarp(); + auto warpsPerCTA = getWarpsPerCTA(); + + const unsigned tilesPerWarpM = tilesPerWarp[mIndex]; + const unsigned tilesPerWarpN = tilesPerWarp[nIndex]; + const unsigned warpsPerCTAM = warpsPerCTA[mIndex]; + const unsigned warpsPerCTAN = warpsPerCTA[nIndex]; + + auto warpOrder = getDefaultMmaOrder(*this); + auto dimM = outDimNames[warpOrder[1]]; + auto dimN = outDimNames[warpOrder[0]]; + tileLayout = tileLayout.transposeOuts({dimN, dimM}); + + // First, extend the layout along the N dimension: + // - registers are distributed across tilesPerWarpN + // - then across warpsPerCTAN in the N dimension. + tileLayout *= LinearLayout::identity1D(tilesPerWarpN, kRegister, dimN); + tileLayout *= LinearLayout::identity1D(warpsPerCTAN, kWarp, dimN); + + // At this point, the layout is defined across the N dimension within a CTA + // tile. Instead of switching to the M dimension now, we continue extending + // the layout along the remaining N dimension, and only then proceed along M, + // following the tilesPerWarp configuration. + // If the N dimension is not large enough to span multiple CTA tiles (i.e., + // the first argument is 0), an empty layout is created, so this identity + // layout will not introduce any new registers. + tileLayout *= LinearLayout::identity1D( + shape[nIndex] / (nDim * warpsPerCTAN * tilesPerWarpN), kRegister, dimN); + tileLayout *= LinearLayout::identity1D(tilesPerWarpM, kRegister, dimM); + + // Finally, extend the layout across warps in the M dimension. + // After this step, the layout covers a sub-tensor of size ctaTileM × N, + // i.e., the full N dimension and a CTA tile's extent in M. + // The rest of the layout will be defined by combineCtaCgaWithShape. + tileLayout *= LinearLayout::identity1D(warpsPerCTAM, kWarp, dimM); + + if (hasBatchDim) { + int batchIndex = 0; + // Extend the base vector with one value to accommodate for the batch + // dimension, which appears at the last. + tileLayout *= + LinearLayout::identity1D(1, kRegister, outDimNames[batchIndex]); + tileLayout *= LinearLayout::identity1D(1, kLane, outDimNames[batchIndex]); + tileLayout *= LinearLayout::identity1D(warpsPerCTA[0], kWarp, + outDimNames[batchIndex]); + } + + return combineCtaCgaWithShape(tileLayout, getCTALayout(), shape); +} + +LinearLayout wmmaDotOperandToLinearLayout(DotOperandEncodingAttr dotWmmaLayout, + ArrayRef shape) { + auto wmmaLayout = llvm::cast(dotWmmaLayout.getParent()); + unsigned version = wmmaLayout.getVersion(); + assert(version >= 1 && version <= 3 && "unexpected wmma version"); + + auto rank = shape.size(); + bool hasBatchDim = rank == 3; + + MLIRContext *ctx = dotWmmaLayout.getContext(); + SmallVector outDimNames = standardOutDimNames(ctx, rank); + StringAttr kRegister = S("register"); + StringAttr kLane = S("lane"); + StringAttr kWarp = S("warp"); + + // lane order + // operand A: [1, 0] / [2, 1, 0] + // operand B: [0, 1] / [1, 2, 0] + // for both cases it is [k, nonk]/[k, nonk, batch] + auto order = + getOrderForDotOperand(dotWmmaLayout.getOpIdx(), rank, /*kContig*/ true); + auto dimK = outDimNames[order[0]]; + auto dimNonK = outDimNames[order[1]]; + + auto mnkDim = wmmaLayout.getInstrShape(); + auto kDim = mnkDim[2]; + auto warpsPerCTA = wmmaLayout.getWarpsPerCTA(); + auto tilesPerWarp = wmmaLayout.getTilesPerWarp(); + auto nonKDimIndex = dotWmmaLayout.getOpIdx() == 0 ? rank - 2 : rank - 1; + auto tilePerWarpNonK = tilesPerWarp[nonKDimIndex]; + auto kDimIndex = dotWmmaLayout.getOpIdx() == 0 ? rank - 1 : rank - 2; + unsigned kSize = shape[kDimIndex]; + + auto nonKDim = dotWmmaLayout.getOpIdx() == 0 ? mnkDim[0] : mnkDim[1]; + auto kWidth = dotWmmaLayout.getKWidth(); + constexpr int warpSize = 32; + + // The relative order of registers and lanes is given by: + // - k dim: kWidth registers + // - non-k dim: nonKDim lanes + // - k dim: depth = warpSize / nonKDim lanes + // version 1 duplicates these values across k dim + // version 2/3 offsets these values across k dim + // - k dim: repeat kDim / (kWidth * depth) times to fit k dim + LinearLayout tileLayout; + int depth = warpSize / nonKDim; + tileLayout = LinearLayout::identity1D(kWidth, kRegister, dimK) * + LinearLayout::identity1D(nonKDim, kLane, dimNonK); + tileLayout *= version == 1 ? LinearLayout::zeros1D(depth, kLane, dimK) + : LinearLayout::identity1D(depth, kLane, dimK); + + // When tilePerWarpNonK > 1, we can't rely on the traditional way to fill the + // block along K. Instead, we need to manually fill the whole kSize, then + // apply tilePerWarpNonK along nonK direction. + int kTileSize = depth * kWidth; + if (tilePerWarpNonK > 1) { + tileLayout *= LinearLayout::identity1D(std::max(kSize, kDim) / kTileSize, + kRegister, dimK); + tileLayout *= LinearLayout::identity1D(tilePerWarpNonK, kRegister, dimNonK); + } else { + tileLayout *= LinearLayout::identity1D(kDim / kTileSize, kRegister, dimK); + } + + if (hasBatchDim) { + assert(order[2] == 0); + // Extend the base vector with one value to accommodate for the batch + // dimension, which appears at the last. + tileLayout *= LinearLayout::identity1D(1, kRegister, outDimNames[order[2]]); + tileLayout *= LinearLayout::identity1D(1, kLane, outDimNames[order[2]]); + } + + // Generate warp layout + auto warpOrder = getDefaultMmaOrder(wmmaLayout); + LinearLayout warpLayout = broadcastedDotOperandLayout( + ctx, warpsPerCTA, warpOrder, kDimIndex, S("warp")); + + // reorder dim names in rep order, so combineCtaCgaWithShape generate proper + // extension of layout + auto repOrder = wmmaLayout.getRepOrderForOperand(dotWmmaLayout.getOpIdx()); + SmallVector repDimNames; + for (auto dim : repOrder) + repDimNames.push_back(outDimNames[dim]); + + // join instruction layout and warps using repetition order of dimensions + LinearLayout ctaLayout = tileLayout.transposeOuts(repDimNames) * + warpLayout.transposeOuts(repDimNames); + + return combineCtaCgaWithShape(ctaLayout, wmmaLayout.getCTALayout(), shape); +} + +LinearLayout sunrisemmaDotOperandToLinearLayout(DotOperandEncodingAttr dotEncAttr, ArrayRef shape) { + MLIRContext *ctx = dotEncAttr.getContext(); + auto rank = shape.size(); + SmallVector outDimNames = standardOutDimNames(ctx, rank); + unsigned dotOpIdx = dotEncAttr.getOpIdx(); + SunriseMmaEncodingAttr sunriseMmaAttr = cast(dotEncAttr.getParent()); + bool kContig = (dotOpIdx == 0 ? sunriseMmaAttr.getIsACol() == false : sunriseMmaAttr.getIsBCol() == true); + auto order = getOrderForDotOperand(dotOpIdx, rank, kContig); + unsigned elemBitWidth = sunriseMmaAttr.getInputElemBitWidth(); + auto tileLayout = LinearLayout::empty(); + switch(elemBitWidth) { + case 32: + if( (dotOpIdx == 0 && sunriseMmaAttr.getIsACol() == false) + || (dotOpIdx == 1 && sunriseMmaAttr.getIsBCol() == true) ) { + tileLayout = LinearLayout( + {{S("register"), {}}, + {S("lane"), {{1,0}, {2,0}, {0,1}, {0,2}, {0,4}}} + }, {outDimNames[order[0]], outDimNames[order[1]]} + ); + } else { + tileLayout = LinearLayout( + {{S("register"), {}}, + {S("lane"), {{1,0}, {2,0}, {4,0}, {0,1}, {0,2}}} + }, {outDimNames[order[0]], outDimNames[order[1]]} + ); + } + break; + case 16: + if(dotOpIdx == 0) { + tileLayout = LinearLayout( + //{{S("register"), {{1,0}, {8,0}}}, // 有么有8,0好像都行？ + {{S("register"), {{1,0}}}, + {S("lane"), {{2,0}, {4,0}, {0,1}, {0,2}, {0,4}}} + }, {outDimNames[order[0]], outDimNames[order[1]]} + ); + } else { + tileLayout = LinearLayout( + {{S("register"), {{1,0}}}, + {S("lane"), {{2,0}, {4,0}, {0,1}, {0,2}, {0,4}}} + }, {outDimNames[order[0]], outDimNames[order[1]]} + ); + } + break; + case 8: + if( (dotOpIdx == 0 && sunriseMmaAttr.getIsACol() == false) + || (dotOpIdx == 1 && sunriseMmaAttr.getIsBCol() == true) ) { + tileLayout = LinearLayout( + {{S("register"), {{1,0}, {2,0}}}, + {S("lane"), {{4,0}, {8,0}, {0,1}, {0,2}, {0,4}}} + }, {outDimNames[order[0]], outDimNames[order[1]]} + ); + } else { + tileLayout = LinearLayout( + {{S("register"), {{1,0}, {2,0}}}, + {S("lane"), {{4,0}, {0,1}, {0,2}, {0,4}, {0,8}}} + }, {outDimNames[order[0]], outDimNames[order[1]]} + ); + } + break; + case 4: + default: + llvm::report_fatal_error((std::string("linearlayout not implemented! elemBitWidth:") + std::to_string(elemBitWidth)).c_str()); + break; + } + + auto kDim = dotOpIdx == 0 ? rank - 1 : rank - 2; + auto warpOrder = getDefaultMmaOrder(sunriseMmaAttr); + // SmallVector warpOrder = dotOpIdx == 0 ? SmallVector({1,0}) : SmallVector({0,1}); + + // LinearLayout warpLayout = identityStandardND(S("warp"), sunriseMmaAttr.getWarpsPerCTA(), warpOrder); + LinearLayout warpLayout = broadcastedDotOperandLayout(ctx, sunriseMmaAttr.getWarpsPerCTA(), warpOrder, kDim, S("warp")); + + // reorder dim names in rep order, so combineCtaCgaWithShape generate proper + // extension of layout + auto repOrder = sunriseMmaAttr.getRepOrderForOperand(dotOpIdx); + SmallVector repDimNames; + for (auto dim : repOrder) + repDimNames.push_back(outDimNames[dim]); + + // join instruction layout and warps using repetition order of dimensions + LinearLayout ctaLayout = tileLayout.transposeOuts(repDimNames) * + warpLayout.transposeOuts(repDimNames); + return combineCtaCgaWithShape(ctaLayout, sunriseMmaAttr.getCTALayout(), shape); +} + +LinearLayout SunriseMmaEncodingAttr::toLinearLayout(ArrayRef shape) const { + int rank = shape.size(); + assert(rank == 2); + + MLIRContext *ctx = getContext(); + SmallVector outDimNames = standardOutDimNames(ctx, rank); + // SmallVector order = getDefaultMmaOrder(*this); + // SmallVector order = getOrderForDotOperand(dotWmmaLayout.getOpIdx(), rank, /*kContig*/ true); + SmallVector order = SmallVector({1,0}); + + SunriseMmaEncodingAttr::TMMAOutLayout outLayout = getOutLayout(); + if ((outLayout != SunriseMmaEncodingAttr::TMMAOutLayout::Row_2B) && (outLayout != SunriseMmaEncodingAttr::TMMAOutLayout::ARow_4B_8x4)) { + assert(0 && "Unsupport SunriseMmaEncodingAttr::TMMAOutLayout yet"); + } + auto tileLayout = LinearLayout::empty(); + if(outLayout == SunriseMmaEncodingAttr::TMMAOutLayout::Row_2B) { + // fp16 row + tileLayout = LinearLayout( + {{S("register"), {{1, 0}}}, + //{{S("register"), {}}, + {S("lane"), {{2,0}, {4,0}, {0,1}, {0,2}, {0,4}}} + }, + {outDimNames[order[0]], outDimNames[order[1]]} + ); + } else if(outLayout == SunriseMmaEncodingAttr::TMMAOutLayout::ARow_4B_8x4) { + // fp32 8x4 + tileLayout = LinearLayout( + {{S("register"), {{4, 0}}}, + {S("lane"), {{1,0}, {2,0}, {0,1}, {0,2}, {0,4}}} + }, + {outDimNames[order[0]], outDimNames[order[1]]} + ); + } else { + llvm::report_fatal_error("Unsupport sunrisemma outLayout"); + } + LinearLayout warpLayout = identityStandardND(S("warp"), getWarpsPerCTA(), order); + // reorder dim names in rep order, so combineCtaCgaWithShape generate proper + // extension of layout + auto repOrder = getRepOrder(); + SmallVector repDimNames; + for (auto dim : repOrder) + repDimNames.push_back(outDimNames[dim]); + + // join instruction layout and warps using repetition order of dimensions + LinearLayout ctaLayout = tileLayout.transposeOuts(repDimNames) * + warpLayout.transposeOuts(repDimNames); + return combineCtaCgaWithShape(ctaLayout, getCTALayout(), shape); +} + +LinearLayout +BlockedEncodingAttr::toLinearLayout(ArrayRef shape) const { + MLIRContext *ctx = getContext(); + auto order = getOrder(); + LinearLayout ctaLayout = + identityStandardND(S("register"), getSizePerThread(), order) * + identityStandardND(S("lane"), getThreadsPerWarp(), order) * + identityStandardND(S("warp"), getWarpsPerCTA(), order); + + return combineCtaCgaWithShape(ctaLayout, getCTALayout(), shape); +} + +LinearLayout fmaDotToLinearLayout(DotOperandEncodingAttr operandLayout, + ArrayRef shape) { + int rank = shape.size(); + auto blocked = cast(operandLayout.getParent()); + MLIRContext *ctx = operandLayout.getContext(); + + // TODO: introduce registerOrder or use getDefaultOrder(operandLayout) + // Currently this order is used in legacy converter, because we do not + // have access to full dot operand layout, only parent part. + auto regOrder = blocked.getOrder(); + auto threadOrder = blocked.getOrder(); + auto warpOrder = blocked.getOrder(); + auto repOrder = blocked.getRepOrder(); + + StringAttr kReg = S("register"); + StringAttr kLane = S("lane"); + StringAttr kWarp = S("warp"); + + auto threadSize = llvm::to_vector(blocked.getSizePerThread()); + auto kDimIdx = operandLayout.getOpIdx() == 0 ? rank - 1 : rank - 2; + threadSize[kDimIdx] = shape[kDimIdx]; + auto threadShape = blocked.getThreadsPerWarp(); + auto warpShape = blocked.getWarpsPerCTA(); + + SmallVector repDimNames = + permuteDimNames(standardOutDimNames(ctx, rank), repOrder); + + auto registersLayout = identityStandardND(kReg, threadSize, regOrder); + auto lanesLayout = broadcastedDotOperandLayout(ctx, threadShape, threadOrder, + kDimIdx, kLane); + auto warpsLayout = + broadcastedDotOperandLayout(ctx, warpShape, warpOrder, kDimIdx, kWarp); + + LinearLayout ctaLayout = registersLayout.transposeOuts(repDimNames) * + lanesLayout.transposeOuts(repDimNames) * + warpsLayout.transposeOuts(repDimNames); + + return combineCtaCgaWithShape(ctaLayout, getCTALayout(operandLayout), shape); +} + +LinearLayout nvidiaMmaTile(MLIRContext *ctx, ArrayRef tileShape, + unsigned kWidth, ArrayRef order, + ArrayRef repOrder) { + // Trivial layout mapping 0 -> (0, 0), but we set the order to repOrder + // Like LinearLayout::empty() but with a rank and an order + int rank = repOrder.size(); + auto dimNames = standardOutDimNames(ctx, rank); + auto trivialShape = SmallVector(rank, 1); + LinearLayout ctaLayout = + identityStandardND(S("register"), trivialShape, repOrder); + + assert(rank >= 2); + auto inner = order[0]; + auto outer = order[1]; + + assert(tileShape.size() == rank); + int m = tileShape[outer]; + int n = tileShape[inner]; + + // The relative order of registers and lanes is given by: + // - Inner dim: kWidth registers + // - Inner dim: 4 lanes + // - Outer dim: 8 lanes + // - Outer dim: repeat m / 8 times + // - Inner dim: repeat n / (kWidth * 4) times + assert(m % 8 == 0); + assert(n % (kWidth * 4) == 0); + // There is at least one subtile on the inner-most dimension + // FIXME. We should implement operator* in terms of operator*= + // and chain *= instead of using * + auto outDimNames = llvm::to_vector(ctaLayout.getOutDimNames()); + ctaLayout = ctaLayout * + LinearLayout::identity1D(kWidth, S("register"), dimNames[inner]) * + LinearLayout::identity1D(4, S("lane"), dimNames[inner]) * + LinearLayout::identity1D(8, S("lane"), dimNames[outer]) * + LinearLayout::identity1D(m / 8, S("register"), dimNames[outer]) * + LinearLayout::identity1D(n / (kWidth * 4), S("register"), + dimNames[inner]); + return ctaLayout; +} + +LinearLayout +NvidiaMmaEncodingAttr::toLinearLayout(ArrayRef shape) const { + auto ctx = getContext(); + int rank = shape.size(); + assert(rank == getRank()); + + SmallVector tileShape; + if (isAmpere()) { + // Ampere.getInstrShape() returns the tile shape + tileShape = SmallVector(getInstrShape()); + } else { + assert(isHopper()); + auto instrShapeMNK = getInstrShape(); + tileShape = SmallVector({instrShapeMNK[0], instrShapeMNK[1]}); + } + // nvidiamma layout always assumes kWidth = 2 + constexpr auto kWidth = 2; + auto order = getDefaultMmaOrder(*this); + auto ctaLayout = nvidiaMmaTile(ctx, tileShape, kWidth, order, getRepOrder()); + + auto warpOrder = getMatrixOrder(rank, /*rowMajor*/ !isHopper()); + ctaLayout *= identityStandardND(S("warp"), getWarpsPerCTA(), warpOrder) + .transposeOuts(llvm::to_vector(ctaLayout.getOutDimNames())); + + return combineCtaCgaWithShape(ctaLayout, getCTALayout(), shape); +} + +LinearLayout nvidiaDotToLinearLayout(ArrayRef shape, + DotOperandEncodingAttr dot) { + int rank = shape.size(); + auto mma = cast(dot.getParent()); + int kWidth = dot.getKWidth(); + bool isA = dot.getOpIdx() == 0; + MLIRContext *ctx = mma.getContext(); + + SmallVector tileShape(rank, 1); + if (isA) { + tileShape[rank - 2] = 16; + tileShape[rank - 1] = kWidth * 8; + } else { + // Hopper takes the rhs via shared memory + assert(mma.isAmpere()); + tileShape[rank - 2] = kWidth * 8; + tileShape[rank - 1] = 8; + } + auto order = getOrderForDotOperand(dot.getOpIdx(), rank, /*kContig*/ true); + auto ctaLayout = + nvidiaMmaTile(ctx, tileShape, kWidth, order, dot.getRepOrder()); + auto kDim = isA ? rank - 1 : rank - 2; + auto warpOrder = getMatrixOrder(rank, /*rowMajor*/ !mma.isHopper()); + ctaLayout *= broadcastedDotOperandLayout(ctx, mma.getWarpsPerCTA(), warpOrder, + kDim, S("warp")) + .transposeOuts(llvm::to_vector(ctaLayout.getOutDimNames())); + + return combineCtaCgaWithShape(ctaLayout, getCTALayout(dot), shape); +} + +LinearLayout +DotOperandEncodingAttr::toLinearLayout(ArrayRef shape) const { + auto parent = getParent(); + if (auto blockedLayout = mlir::dyn_cast(parent)) { + return fmaDotToLinearLayout(*this, shape); + } else if (auto mfmaLayout = mlir::dyn_cast(parent)) { + return mfmaDotToLinearLayout(*this, shape); + } else if (auto wmmaLayout = mlir::dyn_cast(parent)) { + return wmmaDotOperandToLinearLayout(*this, shape); + } else if (auto sunrisemmaLayout = mlir::dyn_cast(parent)){ + return sunrisemmaDotOperandToLinearLayout(*this, shape); + } else { + auto mma = mlir::cast(parent); + return nvidiaDotToLinearLayout(shape, *this); + } +} + +LinearLayout SliceEncodingAttr::toLinearLayout(ArrayRef shape) const { + MLIRContext *ctx = getContext(); + + // First compute the linear layout for this layout's parent. + SmallVector parentShape(shape); + parentShape.insert(parentShape.begin() + getDim(), 1); + LinearLayout parentLL = triton::gpu::toLinearLayout(parentShape, getParent()); + + auto sliceLL = removeStandardDim(parentLL, getDim()); + + // Step 3: Along the "register" dim, remove any all-zero bases. + auto bases = sliceLL.getBases(); + std::vector> newRegBases; + for (const auto &basis : bases[S("register")]) { + if (llvm::any_of(basis, [](int b) { return b != 0; })) { + newRegBases.push_back(basis); + } + } + bases[S("register")] = newRegBases; + + return LinearLayout(std::move(bases), + llvm::to_vector(sliceLL.getOutDimNames())); +} + +LinearLayout tensorMemoryToLinearLayout(ArrayRef shape, + TensorMemoryEncodingAttr encoding) { + // [Zeros in TMEM LinearLayouts] + // If there is a zero in bases rows=32,64 this means that there is + // broadcasting, i.e. the same tensor element is duplicated in different + // addressable blocks If the zero is in any other row/col (i.e. within a given + // warp-addressable tmem space) it means it is not defined + + // We model packed layouts as having the rows/cols dimensions of bitWidth=16 + // This means that a layout with unpacked=True is the same as one with + // unpacked=False + assert(shape.size() == 2); + auto *ctx = encoding.getContext(); + auto kRow = S("row"); + auto kCol = S("col"); + auto dims = standardOutDimNames(ctx, 2); + // The CTAOrder = [0, 1] so se start by N so that it ends up as + // ((tile * splitM) * splitN) + if (encoding.getCTASplitN() > 1) { + auto split = + LinearLayout::identity1D(encoding.getCTASplitN(), kCol, dims[1]); + auto newEncoding = TensorMemoryEncodingAttr::get( + ctx, encoding.getBlockM(), encoding.getBlockN(), + encoding.getColStride(), encoding.getCTASplitM(), 1, + encoding.getTwoCTAs()); + return tensorMemoryToLinearLayout( + {shape[0], shape[1] / encoding.getCTASplitN()}, newEncoding) * + split; + } + if (encoding.getCTASplitM() > 1) { + auto splitM = encoding.getCTASplitM(); + auto blockM = encoding.getBlockM(); + bool isM64TwoCTA = blockM == 64 && encoding.getTwoCTAs(); + if (isM64TwoCTA) { + // blockM == 64 and twoCTAs is laid out as the transpose of 128xblockN + // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-b + blockM *= 2; + splitM /= 2; + } + auto split = LinearLayout::identity1D(splitM, kCol, dims[0]); + auto newEncoding = TensorMemoryEncodingAttr::get( + ctx, blockM, encoding.getBlockN(), encoding.getColStride(), 1, + encoding.getCTASplitN(), encoding.getTwoCTAs()); + auto ret = + tensorMemoryToLinearLayout({shape[0] / splitM, shape[1]}, newEncoding) * + split; + // In this case, we swap the basis of the last row and last column as per + // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-bny + if (isM64TwoCTA) { + auto bases = ret.getBases(); + auto &rowBases = bases[kRow]; + auto &colBases = bases[kCol]; + std::swap(rowBases[rowBases.size() - 1], colBases[colBases.size() - 1]); + ret = LinearLayout(bases, ret.getOutDims(), ret.isSurjective()); + } + return ret; + } + assert(encoding.getCTASplitM() == 1 && encoding.getCTASplitN() == 1); + + auto blockM = encoding.getBlockM(); + auto blockN = std::min(encoding.getBlockN(), shape[1]); + assert(blockM == 64 || blockM == 128); + LinearLayout tile = + LinearLayout::zeros1D(encoding.getColStride(), kCol, dims[1]); + if (blockM == 64) { + tile *= LinearLayout::identity1D(16, kRow, dims[0]) * + LinearLayout::identity1D(blockN, kCol, dims[1]); + auto bases = tile.getBases(); + if (shape[0] > blockM) { + bases[kRow].push_back({64, 0}); + } else if (shape[1] > blockN) { + bases[kRow].push_back({0, blockN}); + } else { + // Empty, meaning the element is not defined + bases[kRow].push_back({0, 0}); + } + bases[kRow].push_back({16, 0}); + bases[kRow].push_back({32, 0}); + tile = LinearLayout(bases, dims); + } else { + tile *= LinearLayout::identity1D(blockM, kRow, dims[0]) * + LinearLayout::identity1D(blockN, kCol, dims[1]); + } + auto repsM = shape[0] / tile.getOutDimSize(dims[0]); + auto repsN = shape[1] / tile.getOutDimSize(dims[1]); + assert(repsM >= 1 && repsN >= 1); + // Broadcast the remaining dimensions in order [0, 1] + tile = tile * LinearLayout::identity1D(repsM, kCol, dims[0]) * + LinearLayout::identity1D(repsN, kCol, dims[1]); + return tile; +} + +LinearLayout +tensorMemoryScalesToLinearLayout(ArrayRef shape, + TensorMemoryScalesEncodingAttr encoding) { + assert(shape.size() == 2); + auto *ctx = encoding.getContext(); + auto kRow = S("row"); + auto kCol = S("col"); + auto dims = standardOutDimNames(ctx, 2); + + // The CTAOrder = [0, 1] so se start by N so that it ends up as + // ((tile * splitM) * splitN) + if (encoding.getCTASplitN() > 1) { + auto split = + LinearLayout::identity1D(encoding.getCTASplitN(), kCol, dims[1]); + auto newEncoding = + TensorMemoryScalesEncodingAttr::get(ctx, encoding.getCTASplitM(), 1); + return tensorMemoryScalesToLinearLayout( + {shape[0], shape[1] / encoding.getCTASplitN()}, newEncoding) * + split; + } + if (encoding.getCTASplitM() > 1) { + auto split = + LinearLayout::identity1D(encoding.getCTASplitM(), kCol, dims[0]); + auto newEncoding = + TensorMemoryScalesEncodingAttr::get(ctx, 1, encoding.getCTASplitN()); + return tensorMemoryScalesToLinearLayout( + {shape[0] / encoding.getCTASplitM(), shape[1]}, newEncoding) * + split; + } + assert(encoding.getCTASplitM() == 1 && encoding.getCTASplitN() == 1); + + // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-a-layout-1x + auto tile = LinearLayout::identity1D(32, kRow, dims[0]) * + // Broadcasting along 'warps' + LinearLayout::zeros1D(4, kRow, dims[0]) * + LinearLayout::identity1D(4, kCol, dims[1]) * + LinearLayout::identity1D(2, kCol, dims[0]); + // We choose repOrder = [0, 1] + tile *= LinearLayout::identity1D( + llvm::divideCeil(shape[0], tile.getOutDimSize(dims[0])), kCol, + dims[0]) * + LinearLayout::identity1D( + llvm::divideCeil(shape[1], tile.getOutDimSize(dims[1])), kCol, + dims[1]); + // See [Zeros in TMEM LinearLayouts] + // Set some rows/cols to 0 if shape is smaller than 64 x 4 + llvm::SmallDenseMap shapeMap; + for (auto [dim, size] : llvm::zip(dims, shape)) { + shapeMap[dim] = size; + } + return ensureLayoutNotLargerThan(tile, shapeMap); +} + +LinearLayout TritonGPUDialect::toLinearLayout(ArrayRef shape, + Attribute layout) { + CacheKey key{std::vector(shape.begin(), shape.end()), layout}; + if (auto result = llCache.get(key)) { + return *result; + } + + // Layouts are distributed or shared in triton core + // To add a new layout add an else-if clause + LinearLayout result = LinearLayout::empty(); + if (auto distributed = dyn_cast(layout)) { + result = distributed.toLinearLayout(shape); + } else { + assert(llvm::all_of(shape, + [](int64_t dim) { + return llvm::isPowerOf2_32(dim) && dim >= 1; + }) && + "shape must be a postive power of 2"); + if (auto shared = dyn_cast(layout)) { + result = swizzledSharedToLinearLayout(shape, shared); + } else if (auto shared = dyn_cast(layout)) { + result = shared.toLinearLayout(shape); + } else if (auto shared = dyn_cast(layout)) { + result = nvmmaSharedToLinearLayout(shape, shared); + } else if (auto sbl = dyn_cast(layout)) { + result = sharedToLinearLayoutAMDRotating(shape, sbl); + } else if (auto tensorMemoryEncoding = + dyn_cast(layout)) { + result = tensorMemoryToLinearLayout(shape, tensorMemoryEncoding); + } else if (auto tensorMemoryScalesEncoding = + dyn_cast(layout)) { + result = + tensorMemoryScalesToLinearLayout(shape, tensorMemoryScalesEncoding); + } else { + assert(0 && "unknown layout"); + } + } + + llCache.set(std::move(key), result); + return result; +} + +LinearLayout toLinearLayout(RankedTensorType type) { + return toLinearLayout(type.getShape(), type.getEncoding()); +} + +LinearLayout toLinearLayout(MemDescType type) { + // Pass in the allocation shape. Then when using invertAndCompose it will + // trim the allocationShape to the shape if they are different. + // We also remove the first dimension of the allocationShape if there was a + // call to memdesc_index + auto shape = type.getAllocShape().take_back(type.getRank()); + return toLinearLayout(shape, type.getEncoding()); +} + +LinearLayout toLinearLayout(TensorOrMemDesc type) { + if (auto ranked = dyn_cast(type)) { + return toLinearLayout(ranked); + } else { + auto memDesc = cast(type); + return toLinearLayout(memDesc); + } +} + +// UNSAFE OVERLOAD! +// If you call this with a SharedMemoryEncodingAttr, you should call it +// with the allocShape as the shape, otherwise the layout will be incorrect! +LinearLayout toLinearLayout(ArrayRef shape, Attribute layout) { + auto *ctx = layout.getContext(); + return ctx->getLoadedDialect()->toLinearLayout(shape, + layout); +} + +LinearLayout getLayoutWithinBlock(const LinearLayout &layout) { + assert(!layout.getInDimNames().empty()); + MLIRContext *ctx = layout.getInDimNames().begin()->getContext(); + + StringAttr kBlock = S("block"); + assert(layout.hasInDim(kBlock)); + auto bases = layout.getBases(); + bases[kBlock] = {}; + return LinearLayout(bases, llvm::to_vector<4>(layout.getOutDimNames())); +} + +LinearLayout combineCtaCgaWithShape(LinearLayout ctaLayout, + CTAEncodingAttr cgaLayoutAttr, + ArrayRef shape) { + int rank = shape.size(); + assert(ctaLayout.getNumOutDims() == rank); + assert(cgaLayoutAttr.getCTAOrder().size() == rank); + MLIRContext *ctx = cgaLayoutAttr.getContext(); + + SmallVector outDimNames = standardOutDimNames(ctx, rank); + + llvm::SmallDenseMap labeledShape; + for (auto [dim, size] : llvm::zip(outDimNames, shape)) { + labeledShape[dim] = size; + } + + LinearLayout cgaLayout = + ensureLayoutNotLargerThan(cgaLayoutAttr.getLinearLayout(), labeledShape) + .transposeOuts(llvm::to_vector(ctaLayout.getOutDimNames())); + + // Calculate the shape of the ctaLayout, which is `shape` divided by the + // cgaLayout's size. + llvm::SmallDenseMap ctaShape; + assert(llvm::to_vector(ctaLayout.getOutDimNames()) == + llvm::to_vector(cgaLayout.getOutDimNames())); + for (auto dim : ctaLayout.getOutDimNames()) { + ctaShape[dim] = + std::max(int64_t{1}, labeledShape[dim] / cgaLayout.getOutDimSize(dim)); + } + + ctaLayout = ensureLayoutNotSmallerThan(ctaLayout, ctaShape); + ctaLayout = ensureLayoutNotLargerThan(ctaLayout, ctaShape); + + LinearLayout ret = (ctaLayout * cgaLayout).transposeOuts(outDimNames); + for (auto dim : ret.getOutDimNames()) { + assert(ret.getOutDimSize(dim) == labeledShape[dim]); + } + return ret; +} + +LinearLayout chooseShemLayoutForRegToRegConversion( + MLIRContext *ctx, ArrayRef tensorShape, + ArrayRef repShape, ArrayRef order) { + auto outDimNames = standardOutDimNames(ctx, tensorShape.size()); + LinearLayout layout = LinearLayout::empty(); + SmallVector kRepDims; + SmallVector kOffsetDims; + auto totalIters = 1; + auto totalOffsets = 1; + for (int i = 0; i < tensorShape.size(); i++) { + int dim = order[i]; + StringAttr kIteration = S("iteration" + std::to_string(dim)); + StringAttr kOffset = S("offset" + std::to_string(dim)); + kRepDims.push_back(kIteration); + kOffsetDims.push_back(kOffset); + assert(llvm::isPowerOf2_32(repShape[dim])); + assert(llvm::isPowerOf2_32(tensorShape[dim])); + auto numIters = tensorShape[dim] / repShape[dim]; + layout *= + LinearLayout::identity1D(repShape[dim], kOffset, outDimNames[dim]); + layout *= LinearLayout::identity1D(numIters, kIteration, outDimNames[dim]); + totalIters *= numIters; + totalOffsets *= repShape[dim]; + } + StringAttr kOffset = S("offset"); + StringAttr kIteration = S("iteration"); + StringAttr kBlock = S("block"); + SmallVector newDims; + newDims.append(kOffsetDims.begin(), kOffsetDims.end()); + newDims.append(kRepDims.begin(), kRepDims.end()); + // Transpose layout from [offset0, rep0, offset1, rep1, ...] to + // [offset0, offset1, ..., rep0, rep1, ...] + auto ret = layout.transposeIns(newDims); + // Reshape layout from [offset0, offset1, ..., rep0, rep1, ...] to + // [offset, rep, block] + return ret.reshapeIns( + {{kOffset, totalOffsets}, {kIteration, totalIters}, {kBlock, 1}}); +} + +std::optional +chooseDsReadTrLayout(Attribute enc, ArrayRef shape, + int32_t elemBitWidth, unsigned instBitWidth, + unsigned numLanesInShuffleGroup) { + assert(elemBitWidth == 4); + auto dot = cast(enc); + return chooseDotDsReadTrLayout(dot, shape, elemBitWidth, instBitWidth, + numLanesInShuffleGroup); +} + +LinearLayout chooseScaledWmmaScaleLayout(MLIRContext *ctx, int dotOperandIdx, + ArrayRef dotOperandShape, + unsigned wmmaMDim, + ArrayRef tilesPerWarp, + ArrayRef warpsPerCTA) { + using basisT = std::vector>; + unsigned rank = dotOperandShape.size(); + auto order = mlir::triton::gpu::getMatrixOrder(rank, /*rowMajor=*/true); + auto outDimNames = standardOutDimNames(ctx, rank); + + StringAttr kRegister = StringAttr::get(ctx, "register"); + StringAttr kLane = StringAttr::get(ctx, "lane"); + StringAttr kWarp = StringAttr::get(ctx, "warp"); + StringAttr kBlock = StringAttr::get(ctx, "block"); + + // In scaled dot, the shapes of operands(without batch dimension) are, + // respectively: + // - A: [M, K] + // - B: [K, N] + // - aScale: [M, K / 32 or 16] + // - bScale: [N, K / 32 or 16] + auto dimK = outDimNames[order[0]]; + auto dimNonK = outDimNames[order[1]]; + + // Each lane holds kWidth=4 consecutive values along the K dim. + // The first 16 lanes are distributed along the nonK dim. + unsigned scaleKWidth = 4; + auto kSize = dotOperandShape[1]; + LinearLayout tileLayout = + LinearLayout::identity1D(scaleKWidth, kRegister, dimK) * + LinearLayout::identity1D(16, kLane, dimNonK); + + // If there's 1 tile per warp, we are not using the remaining 16 lanes, so + // just let them duplicate values of the first 16 lanes. + // Otherwise, we put consecutive values along the nonK dim in the remaining + // 16 lanes. + unsigned mnDim = dotOperandIdx == 0 ? rank - 2 : rank - 1; + unsigned tilePerWarpMN = tilesPerWarp[mnDim]; + if (tilePerWarpMN > 1) { + assert(tilePerWarpMN == 2 && "TilesPerWarp > 2 is not supported."); + tileLayout *= LinearLayout::identity1D(tilePerWarpMN, kLane, dimNonK); + } else { + tileLayout *= LinearLayout::zeros1D(2, kLane, dimNonK); + } + + // If the shape along the K dim is larger than kWidth, repeat this + // pattern to fill the K dim. + tileLayout *= LinearLayout::identity1D(kSize / scaleKWidth, kRegister, dimK); + + auto warpsPerCTANew = (dotOperandIdx == 1) + ? SmallVector{warpsPerCTA[1], warpsPerCTA[0]} + : SmallVector{warpsPerCTA[0], warpsPerCTA[1]}; + + auto warpOrder = (dotOperandIdx == 1) ? SmallVector{0, 1} + : SmallVector{1, 0}; + LinearLayout warpLayout = + identityStandardND(kWarp, warpsPerCTANew, warpOrder); + LinearLayout ctaLayout = tileLayout.transposeOuts(outDimNames) * + warpLayout.transposeOuts(outDimNames); + + return combineCtaCgaWithShape( + ctaLayout, CTAEncodingAttr::getDefault(ctx, /*rank=*/2), dotOperandShape); +} + +// PTX ISA - Warp-level MMA Block Scaling +// https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-block-scaling +// This function generates layouts for scale tensors used in scaled dot +// operations. +// Implementation notes: +// - We choose a fixed provider for A (thread-id-a = 0) and B (thread-id-b = +// 0) +// - We choose a fixed byte selector for A (byte-id-a = 0) and B (byte-id-b = +// 0) +// - Each lane in a quad has the same scale factor. +LinearLayout getSM120DotScaledScaleLayout(MLIRContext *ctx, + ArrayRef shape, int opIdx, + ArrayRef warpsPerCTA, + CTAEncodingAttr ctaLayout) { + unsigned rank = shape.size(); + auto outDims = standardOutDimNames(ctx, rank); + StringAttr kRegister = StringAttr::get(ctx, "register"); + StringAttr kLane = StringAttr::get(ctx, "lane"); + StringAttr kWarp = StringAttr::get(ctx, "warp"); + // - A: [M, K] + // - B: [K, N] + // - aScale: [M, K / K_GROUP_SIZE] + // - bScale: [N, K / K_GROUP_SIZE] + const unsigned kIdx = 1; + const unsigned mnIdx = 0; + + std::vector> laneBase; + SmallVector order; + SmallVector mmaWarpsPerCTA; + if (opIdx == 0) { + laneBase = {{8, 0}, {0, 0}, {1, 0}, {2, 0}, {4, 0}}; + order = SmallVector{1u, 0u}; + mmaWarpsPerCTA = SmallVector{warpsPerCTA[0], warpsPerCTA[1]}; + } else { + laneBase = {{0, 0}, {0, 0}, {1, 0}, {2, 0}, {4, 0}}; + order = SmallVector{0u, 1u}; + mmaWarpsPerCTA = SmallVector{warpsPerCTA[1], warpsPerCTA[0]}; + } + LinearLayout LL = + LinearLayout::identity1D(shape[1], kRegister, outDims[kIdx]) * + LinearLayout({{kLane, laneBase}}, {outDims[mnIdx], outDims[kIdx]}) * + broadcastedDotOperandLayout(ctx, mmaWarpsPerCTA, order, 1u, kWarp); + return combineCtaCgaWithShape(LL, ctaLayout, shape); +} + +LinearLayout chooseScaledMfmaScaleLayout(MLIRContext *ctx, int dotOperandIdx, + ArrayRef dotOperandShape, + unsigned mfmaMDim, + ArrayRef tilesPerWarp, + ArrayRef warpsPerCTA) { + using basisT = std::vector>; + unsigned rank = dotOperandShape.size(); + auto order = mlir::triton::gpu::getMatrixOrder(rank, /*rowMajor=*/true); + auto standardOutDims = standardOutDimNames(ctx, rank); + StringAttr kRegister = StringAttr::get(ctx, "register"); + StringAttr kLane = StringAttr::get(ctx, "lane"); + StringAttr kWarp = StringAttr::get(ctx, "warp"); + StringAttr kBlock = StringAttr::get(ctx, "block"); + + // Fetch the tilesPerWarp value in the M dimension for operand A, or in the N + // dimension for operand B. + unsigned mnDim = dotOperandIdx == 0 ? rank - 2 : rank - 1; + unsigned tilePerWarpMN = tilesPerWarp[mnDim]; + + // In scaled dot, the shapes of operands(without batch dimension) are, + // respectively: + // - A: [M, K] + // - B: [K, N] + // - aScale: [M, K / 32] + // - bScale: [N, K / 32] + // + // In general, for both 32x32 and 16x16 scaled mfma, and no matter what + // data type the A/B operand is, each lane takes 32 elements from A/B + // alone K dim, and 1 or 2 elements from scale accordingly. The number of + // scale's elements in a lane varies because the 32 elements from A/B may + // not be consecutive. + // + // For mxfp4, these 32 elements are consecutive, so only 1 scale element + // is required. But for mxfp6/mxfp8, there are 2 16-consecutive elements + // blocks, so 2 scale elements are required. + int32_t kSize = dotOperandShape[1]; + + std::vector> registerBase; + std::vector> laneBase; + + auto threadsInKDim = mfmaMDim == 32 ? 2 : 4; + for (int32_t elem = threadsInKDim; elem < kSize; elem *= 2) + registerBase.emplace_back(std::vector{elem, 0}); + + for (int32_t elem = mfmaMDim; elem < tilePerWarpMN * mfmaMDim; elem *= 2) + registerBase.emplace_back(std::vector{0, elem}); + + if (mfmaMDim == 32) { + // For ROCDL::mfma_scale_f32_32x32x64_f8f6f4 with fp4 input, each lane + // takes 32 consecutive elements from A alone K dimension. The first + // 32 lanes collectively handle A[0:32][0:32], and the other 32 lanes + // collectively handle A[0:32][32:64]. Each lane take 1 scale element + // accordingly. Similar to B and bScale. + laneBase = {{0, 1}, {0, 2}, {0, 4}, {0, 8}, {0, 16}, {1, 0}}; + } else { + assert(mfmaMDim == 16); + // For ROCDL::mfma_scale_f32_16x16x128_f8f6f4 with fp4 input, each lane + // takes 32 consecutive elements from A alone K dimension. The first + // 16 lanes collectively handle A[0:16][0:32], and another 16 lanes + // collectively handle A[0:16][32:64] and so on. Each lane take 1 scale + // element accordingly. Similar to B and bScale. + laneBase = {{0, 1}, {0, 2}, {0, 4}, {0, 8}, {1, 0}, {2, 0}}; + } + + SmallVector outDimNames = standardOutDimNames(ctx, rank); + LinearLayout tileLayout({{kRegister, registerBase}, {kLane, laneBase}}, + {outDimNames[order[0]], outDimNames[order[1]]}); + + SmallVector warpsPerCTANew = + (dotOperandIdx == 1) + ? SmallVector{warpsPerCTA[1], warpsPerCTA[0]} + : SmallVector{warpsPerCTA[0], warpsPerCTA[1]}; + + SmallVector warpOrder = (dotOperandIdx == 1) + ? SmallVector{0, 1} + : SmallVector{1, 0}; + + LinearLayout warpLayout = + identityStandardND(kWarp, warpsPerCTANew, warpOrder); + LinearLayout ctaLayout = tileLayout.transposeOuts(outDimNames) * + warpLayout.transposeOuts(outDimNames); + + auto ctaLay = CTAEncodingAttr::getDefault(ctx, 2); + auto finalLay = combineCtaCgaWithShape(ctaLayout, ctaLay, dotOperandShape); + return finalLay; +} + +std::optional +chooseMfmaLikeStoreLayout(RankedTensorType valType) { + // TODO: WMMA Support on RDNA + if (!isa(valType.getEncoding())) + return {}; + auto mfmaLayout = cast(valType.getEncoding()); + + // We currently only support transposed [B]F16 MFMA32x32 and MFMA16x16 on + // CDNA4. + auto mnkDim = mfmaLayout.getInstrShape(); + bool isMfma32 = mnkDim[0] == 32 && mnkDim[1] == 32; + bool isMfma16 = mnkDim[0] == 16 && mnkDim[1] == 16; + + auto valShape = valType.getShape(); + // For mfma16x16, to use in-wavefront swap, we need to make sure the tiles + // used are in one wavefront if there are multiple tiles, which means + // warpsPerCTA = [numWarps, 1] and at least two tiles along the N dim. For + // now, it is only possible for FA-like kernels since during mfma generation, + // the WarpsPerCTA of the head dot in the chain will be reshaped to [numWaprs, + // 1]. + // TODO: For gemm-like kernel, the transformation here cannot be applied for + // now and will support it. + bool validForMfma16 = isMfma16 && valShape.back() >= 16 * 2 && + mfmaLayout.getWarpsPerCTA().back() == 1; + + Type elemType = valType.getElementType(); + if (!(valType.getRank() == 2 && (elemType.isF16() || elemType.isBF16()) && + mfmaLayout.getVersion() == 4 && mfmaLayout.getIsTransposed() && + (isMfma32 || validForMfma16))) + return {}; + + LinearLayout mfmaLL = mfmaLayout.toLinearLayout(valShape); + auto mfmaOutDims = llvm::to_vector(mfmaLL.getOutDimNames()); + StringAttr dimM = mfmaOutDims[0]; + StringAttr dimN = mfmaOutDims[1]; + auto swapLL = LinearLayout::empty(); + // The rows are kept as is with an identity linear layout. + swapLL *= LinearLayout::identity1D(valShape[0], dimM, dimM); + /* + clang-format off + In transposed mfma32 layout, Each thread holds 4 consecutive values along N + dim. We want to exchange column 4-7 (owned by thread 32-63, BLK0) and column + 8-11 (owned by thread 0-31, BLK1) every 16 columns to make each thread holds 8 + elements. This would mean exchange the 2nd and 3rd basis vector from an + identity linear layout on tensor elements. + + Correspondingly, the transposed mfma16 layout, the output of + transposed of mfma16x16 is: + + N/register + M/Lane v0 v1 v2 v3 v4 v5 v6 v7 + ------------------------------------------------------------------------- + row0: 0-15 | tile-0 | tile-0 | tile-0 | tile-0 | tile-1 | tile-1 | tile-1 | tile-1 | + ------------------------------------------------------------------------- + row1: 16-31 | tile-0 | tile-0 | tile-0 | tile-0 | tile-1 | tile-1 | tile-1 | tile-1 | + ------------------------------------------------------------------------- + row2: 32-47 | tile-0 | tile-0 | tile-0 | tile-0 | tile-1 | tile-1 | tile-1 | tile-1 | + ------------------------------------------------------------------------- + row3: 48-63 | tile-0 | tile-0 | tile-0 | tile-0 | tile-1 | tile-1 | tile-1 | tile-1 | + ------------------------------------------------------------------------- + which means: + The columns from v0 to v3 are in the one output of mfma16x16 and + the columns from v4 to v7 are in the one output of mfma16x16, + + The following graph is the same as the one above, execept the tile number is replaced with coordinates in the tenor, + N/register + ----------------------------------------------- + M/lane |(0, 0) ... (0, 3) | (0, 16) ... (0, 19) | + |.... | sub-tensor-0 | + |(15, 0) ... (15, 3) | (15, 16) ... (15, 19) | + ----------------------------------------------- + |(0, 4) ... (0, 7) | (0, 20) ... (0, 23) | + |sub-tensor-1 | .... | + |(15, 0) ... (15, 3) | (15, 20) ... (15, 23) | + ----------------------------------------------- + |(0, 8) ... (0, 11)| (0, 24) ... (0, 27) | + |.... | sub-tensor-2 | + |(15, 8) ... (15, 11)| (15, 24) ... (15, 27) | + ----------------------------------------------- + |(0, 12) ... (0, 15)| (0, 28) ... (0, 31) | + |sub-tensor-3 | .... | + |(15, 12) ... (15, 15)| (15, 28) ... (15, 31) | + ----------------------------------------------- + The basis vector for lane and register are: + Register = {{0, 1}, {0, 2}} + Lane = {{1, 0}, {2, 0}, {4, 0}, {8, 0}, {0, 4}, {0, 8}} + With this layout, only 4xfp16 can be packed in the final global store. + + To use 128-bits global store, we need to pack 8 elements, which means the layout looks like: + N/register + M/Lane v0 v1 v2 v3 v4 v5 v6 v7 + ------------------------------------------------------------------------- + row0: 0-15 | tile-0 | tile-0 | tile-0 | tile-0 | tile-0 | tile-0 | tile-0 | tile-0 | + ------------------------------------------------------------------------- + row1: 16-31 | tile-1 | tile-1 | tile-1 | tile-1 | tile-1 | tile-1 | tile-1 | tile-1 | + ------------------------------------------------------------------------- + row2: 32-47 | tile-0 | tile-0 | tile-0 | tile-0 | tile-0 | tile-0 | tile-0 | tile-0 | + ------------------------------------------------------------------------- + row3: 48-63 | tile-1 | tile-1 | tile-1 | tile-1 | tile-1 | tile-1 | tile-1 | tile-1 | + ------------------------------------------------------------------------- + + The following graph is the same as the one above, execept the tile number is replaced with coordinates in the tenor: + N/register + ----------------------------------------------- + |(0, 0) ... (0, 3) | (0, 4) ... (0, 7) | + |.... | sub-tensor-1 | + |(15, 0) ... (15, 3) | (15, 16) ... (15, 19) | + ----------------------------------------------- + |(0, 16) ... (0, 19) | (0, 20) ... (0, 23) | + |sub-tensor-0 | .... | + |(15, 16) ... (15, 19)| (15, 20) ... (15, 23) | + ----------------------------------------------- + |(0, 8) ... (0, 11)| (0, 12) ... (0, 15) | + |.... | sub-tensor-3 | + |(15, 8) ... (15, 11)| (15, 12) ... (15, 15) | + ----------------------------------------------- + |(0, 24) ... (0, 27)| (0, 28) ... (0, 31) | + |sub-tensor-2 | .... | + |(15, 24) ... (15, 27)| (15, 28) ... (15, 31) | + ----------------------------------------------- + which means we need to exchange sub-tensor-0 with sub-tensor-1 and sub-tensor-2 and sub-tensor-3. + And basis vector for lane and register are: + Register = {{0, 1}, {0, 2}, {0, 4}} + Lane = {{1, 0}, {2, 0, [4, 0}, {8, 0}, {0, 16}, {0, 8}} + + The steps to get this layout are, firstly we check the last dim of WarpsPerCTA is 1, so we can use v_permlane16. + Then, we exchange the 2nd and 4th elements in the basis vector of an identity linear and then it will be composed with + the original mfma16 LL. + clang-format on + */ + auto destIdxInBases = isMfma32 ? 3 : 4; + std::vector> dimNBases(mfmaLL.getOutDimSizeLog2(dimN)); + std::generate(dimNBases.begin(), dimNBases.end(), + [i = 0]() mutable { return std::vector{1 << i++}; }); + std::swap(dimNBases[2], dimNBases[destIdxInBases]); + swapLL *= LinearLayout({{dimN, dimNBases}}, {dimN}); + + return mfmaLL.compose(swapLL); +} + +} // namespace mlir::triton::gpu diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt new file mode 100644 index 0000000000..318e86333a --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/CMakeLists.txt @@ -0,0 +1,9 @@ +add_triton_library(FlagTree_sunrise_TritonGPUTransforms + Prefetch.cpp + RemoveLayoutConversions.cpp + Utility.cpp + + DEPENDS + TritonTableGen + TritonGPUAttrDefsIncGen +) diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp new file mode 100644 index 0000000000..f673708d1e --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp @@ -0,0 +1,459 @@ +//===----------------------------------------------------------------------===// +// +// This pass tries to prefetch operands (a and b) of tt.dot. +// Those ConvertLayoutOps will be lowered to shared memory loads. +// +// For example: +// %a: tensor<128x32xf16, #enc> +// scf.for %iv = ... iter_args(%a_arg = %a, ...) { +// %d = tt.dot %a_arg, %b, %c +// ... +// scf.yield %a_next, ... +// } +// +// will be translated to +// +// %a: tensor<128x32xf16, #enc> +// %a_tmp = tensor.subview %a[0, 0] [128, 16] +// %a_prefetch = ttg.local_load %a_tmp +// scf.for %iv = ... iter_args(%a_buf = %a, ..., %a_prefetch_arg = %a_prefetch) +// { +// %x = tt.dot %a_prefetch_arg, %b, %c +// %a_tmp_rem = tensor.subview %a_buf[0, 16] [128, 16] +// %a_prefetch_next = ttg.local_load %a_tmp_rem +// ... +// scf.yield %next_a, ..., %a_prefetch_next +// } +//===----------------------------------------------------------------------===// + +#include "mlir/IR/IRMapping.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/Transforms/Passes.h" +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "tritongpu-prefetch" +#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") +#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n") + +namespace mlir { +namespace triton { +namespace gpu { + +#define GEN_PASS_DEF_TRITONGPUPREFETCH +#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc" + +namespace { + +class Prefetcher { + /// cache the ForOp we are working on + scf::ForOp forOp; + /// cache the YieldOp of this ForOp + scf::YieldOp yieldOp; + /// + // TODO: add a hook to infer prefetchWidth + unsigned prefetchWidth = 32; + + /// dots to be prefetched + SetVector dots; + /// dot => dot operand + DenseMap dot2aLoopArg; + DenseMap dot2aHeaderDef; + DenseMap dot2bLoopArg; + DenseMap dot2bHeaderDef; + DenseMap dot2aYield; + DenseMap dot2bYield; + DenseMap> dot2aVals; + DenseMap> dot2bVals; + /// operand => defining + DenseMap operand2headPrefetch; + + LogicalResult isForOpOperand(Value v); + + Value generatePrefetch(Value v, unsigned opIdx, bool isPrologue, + Attribute dotEncoding, OpBuilder &builder, + std::optional offsetK = std::nullopt, + std::optional shapeK = std::nullopt); + + void cloneElementwiseOps(Value &bRem, const SmallVector &vals, + OpBuilder &builder); + +public: + Prefetcher() = delete; + + Prefetcher(scf::ForOp forOp) : forOp(forOp) { + yieldOp = cast(forOp.getBody()->getTerminator()); + } + + LogicalResult initialize(); + + void emitPrologue(); + + scf::ForOp createNewForOp(); +}; + +void Prefetcher::cloneElementwiseOps(Value &ret, const SmallVector &vals, + OpBuilder &builder) { + IRMapping mapping; + mapping.map(vals[1], ret); + for (int i = 2; i < vals.size(); i++) { + Value v = vals[i]; + Value curr = builder.clone(*v.getDefiningOp(), mapping)->getResult(0); + if (isa(curr.getType())) { + auto retType = RankedTensorType::get( + cast(ret.getType()).getShape(), + cast(curr.getType()).getElementType(), + cast(curr.getDefiningOp()->getOperand(0).getType()) + .getEncoding()); + curr.setType(retType); + } + mapping.map(v, curr); + } + if (vals.size() > 1) + ret = mapping.lookup(vals.back()); +} + +Value Prefetcher::generatePrefetch(Value v, unsigned opIdx, bool isPrologue, + Attribute dotEncoding, OpBuilder &builder, + std::optional offsetK, + std::optional shapeK) { + // opIdx: 0 => a, 1 => b + auto type = cast(v.getType()); + SmallVector shape{type.getShape().begin(), type.getShape().end()}; + auto rank = shape.size(); + SmallVector offset(rank, 0); + Type elementType = type.getElementType(); + + // k => (prefetchWidth, k - prefetchWidth) + int64_t kIdx = opIdx == 0 ? rank - 1 : rank - 2; + + offset[kIdx] = isPrologue ? 0 : prefetchWidth; + shape[kIdx] = isPrologue ? prefetchWidth : (shape[kIdx] - prefetchWidth); + + if (shapeK) + shape[kIdx] = *shapeK; + if (offsetK) + offset[kIdx] = *offsetK; + + Value newSmem = triton::gpu::MemDescSubsliceOp::create( + builder, v.getLoc(), + triton::gpu::MemDescType::get( + shape, elementType, type.getEncoding(), type.getMemorySpace(), + type.getMutableMemory(), type.getAllocShape()), + v, offset); + + auto dotOperandEnc = triton::gpu::DotOperandEncodingAttr::get( + builder.getContext(), opIdx, dotEncoding, prefetchWidth / 8); + Value prefetchSlice = triton::gpu::LocalLoadOp::create( + builder, v.getLoc(), + RankedTensorType::get(shape, elementType, dotOperandEnc), newSmem); + + return prefetchSlice; +} + +LogicalResult Prefetcher::initialize() { + Block *loop = forOp.getBody(); + + auto getEncoding = [](Value v) { + return cast(v.getType()).getEncoding(); + }; + + SmallVector dotsInFor; + for (Operation &op : *loop) + if (auto dotOp = dyn_cast(op)) { + // Only accepts dotOps encoded as Nvidia MMA v2 or AMD MFMA + auto dstMmaEnc = + dyn_cast(getEncoding(dotOp.getResult())); + auto dstMfmaEnc = + dyn_cast(getEncoding(dotOp.getResult())); + auto dstTmmaEnc = + dyn_cast(getEncoding(dotOp.getResult())); + if (!dstTmmaEnc && !dstMfmaEnc && (!dstMmaEnc || dstMmaEnc.getVersionMajor() != 2)) + // Don't rewrite if any other type is found. + return failure(); + dotsInFor.push_back(dotOp); + } + + if (dotsInFor.empty()) + return failure(); + + // TODO: segfault (original for still has uses) + // when used in flash attention that has 2 dots in the loop + if (dotsInFor.size() > 1) + return failure(); + + // returns source of cvt + auto getPrefetchSrc = [](Value v) -> SmallVector { + // walk back to conversion + Operation *op = v.getDefiningOp(); + bool foundConvertFromShared = false; + SmallVector rets; + rets.push_back(op->getResult(0)); + LDBG("Prefetch src: " << *op); + while (op) { + if (op->getNumOperands() != 1) + break; + if (!op->getResult(0).hasOneUse()) + break; + rets.push_back(op->getOperand(0)); + if (auto cvt = dyn_cast(op)) { + // NYI for other encodings, for example if we have transpose + // in the chain + if (isa(cvt.getType().getEncoding())) + foundConvertFromShared = true; + break; + } + op = op->getOperand(0).getDefiningOp(); + if (op) + LDBG("op: " << *op); + } + std::reverse(rets.begin(), rets.end()); + + if (foundConvertFromShared) + return rets; + return {}; + }; + + auto getIncomingOp = [this](Value v) -> Value { + if (auto arg = mlir::dyn_cast(v)) + if (arg.getOwner()->getParentOp() == forOp.getOperation()) + return forOp.getTiedLoopInit(arg)->get(); + return Value(); + }; + + auto getYieldOperand = [this](Value v) -> Value { + auto arg = mlir::cast(v); + unsigned yieldIdx = arg.getArgNumber() - forOp.getNumInductionVars(); + return yieldOp.getOperand(yieldIdx); + }; + + for (triton::DotOp dot : dotsInFor) { + auto aType = dot.getA().getType(); + auto bType = dot.getB().getType(); + auto aEnc = + mlir::cast(aType.getEncoding()); + auto bEnc = + mlir::cast(bType.getEncoding()); + int aKWidth = aEnc.getKWidth(); + int bKWidth = bEnc.getKWidth(); + assert(aKWidth == bKWidth); + + auto kSize = aType.getShape().back(); + + // works better with nvidia tensor cores + unsigned elementWidth = aType.getElementTypeBitWidth(); + if (aKWidth == 0) + prefetchWidth = 256 / elementWidth; + else + prefetchWidth = 8 * aKWidth; + + // Skip prefetching if kSize is less than prefetchWidth + if (kSize < prefetchWidth) + continue; + auto aVals = getPrefetchSrc(dot.getA()); + auto bVals = getPrefetchSrc(dot.getB()); + + if (aVals.size() && bVals.size()) { + Value aSmem = aVals.front(); + Value bSmem = bVals.front(); + Value aHeaderDef = getIncomingOp(aSmem); + Value bHeaderDef = getIncomingOp(bSmem); + // Only prefetch loop arg + if (aHeaderDef && bHeaderDef) { + dots.insert(dot); + dot2aVals[dot] = aVals; + dot2bVals[dot] = bVals; + dot2aHeaderDef[dot] = aHeaderDef; + dot2bHeaderDef[dot] = bHeaderDef; + dot2aLoopArg[dot] = aSmem; + dot2bLoopArg[dot] = bSmem; + dot2aYield[dot] = getYieldOperand(aSmem); + dot2bYield[dot] = getYieldOperand(bSmem); + } + } + } + + return success(); +} + +void Prefetcher::emitPrologue() { + OpBuilder builder(forOp); + + for (triton::DotOp dot : dots) { + Attribute dotEncoding = dot.getType().getEncoding(); + Value aPrefetched = + generatePrefetch(dot2aHeaderDef[dot], 0, true, dotEncoding, builder); + cloneElementwiseOps(aPrefetched, dot2aVals[dot], builder); + Value bPrefetched = + generatePrefetch(dot2bHeaderDef[dot], 1, true, dotEncoding, builder); + cloneElementwiseOps(bPrefetched, dot2bVals[dot], builder); + + operand2headPrefetch[dot.getA()] = aPrefetched; + operand2headPrefetch[dot.getB()] = bPrefetched; + } +} + +scf::ForOp Prefetcher::createNewForOp() { + OpBuilder builder(forOp); + + SmallVector loopArgs; + for (auto v : forOp.getInitArgs()) + loopArgs.push_back(v); + for (triton::DotOp dot : dots) { + loopArgs.push_back(operand2headPrefetch[dot.getA()]); + loopArgs.push_back(operand2headPrefetch[dot.getB()]); + } + + auto newForOp = + scf::ForOp::create(builder, forOp.getLoc(), forOp.getLowerBound(), + forOp.getUpperBound(), forOp.getStep(), loopArgs); + + builder.setInsertionPointToStart(newForOp.getBody()); + IRMapping mapping; + for (const auto &arg : llvm::enumerate(forOp.getRegionIterArgs())) + mapping.map(arg.value(), newForOp.getRegionIterArgs()[arg.index()]); + mapping.map(forOp.getInductionVar(), newForOp.getInductionVar()); + + // The insertion point should be placed before the yield op + auto setInsertionPointBeforeYield = [](OpBuilder &builder, + scf::ForOp newForOp) { + if (newForOp.getBody()->mightHaveTerminator()) { + builder.setInsertionPoint(newForOp.getBody()->getTerminator()); + } else { + builder.setInsertionPointToEnd(newForOp.getBody()); + } + }; + + for (Operation &op : forOp.getBody()->without_terminator()) { + // If we're currently trying to sink a prefetched dot, we need to stop + // sinking it (by resetting the insertion point to the end) if we find + // control flow, or anything that depends on the dot op. + if (op.getNumRegions() > 0) { + setInsertionPointBeforeYield(builder, newForOp); + } + for (auto operand : op.getOperands()) { + if (auto def = operand.getDefiningOp()) { + auto dot = dyn_cast(def); + if (dot && dots.contains(dot)) { + setInsertionPointBeforeYield(builder, newForOp); + } + } + } + Operation *newOp = builder.clone(op, mapping); + auto dot = dyn_cast(&op); + if (dot && dots.contains(dot)) { + Attribute dotEncoding = dot.getType().getEncoding(); + // prefetched dot + Operation *firstDot = builder.clone(*dot, mapping); + if (Value a = operand2headPrefetch.lookup(dot.getA())) + firstDot->setOperand( + 0, newForOp.getTiedLoopRegionIterArg(&*a.use_begin())); + if (Value b = operand2headPrefetch.lookup(dot.getB())) + firstDot->setOperand( + 1, newForOp.getTiedLoopRegionIterArg(&*b.use_begin())); + + // remaining part + int64_t kOff = prefetchWidth; + int64_t kRem = dot.getA().getType().getShape().back() - prefetchWidth; + Operation *prevDot = firstDot; + if (kRem == 0) { + // There is only one dot while prefetchWidth == kSize so delay issuing + // it. Meanwhile, newOp should be set to firstDot to make sure the dot + // result is updated to yield. + builder.setInsertionPoint(prevDot); + newOp = firstDot; + } + + while (kRem != 0) { + // int64_t kShape = largestPow2(kRem); + int64_t kShape = prefetchWidth; + auto insertionPoint = builder.saveInsertionPoint(); + builder.setInsertionPoint(prevDot); + Value aRem = + generatePrefetch(mapping.lookup(dot2aLoopArg[dot]), 0, false, + dotEncoding, builder, kOff, kShape); + cloneElementwiseOps(aRem, dot2aVals[dot], builder); + Value bRem = + generatePrefetch(mapping.lookup(dot2bLoopArg[dot]), 1, false, + dotEncoding, builder, kOff, kShape); + cloneElementwiseOps(bRem, dot2bVals[dot], builder); + builder.restoreInsertionPoint(insertionPoint); + newOp = builder.clone(*dot, mapping); + newOp->setOperand(0, aRem); + newOp->setOperand(1, bRem); + newOp->setOperand(2, prevDot->getResult(0)); + prevDot = newOp; + kOff += kShape; + kRem -= kShape; + if (kRem == 0) { + // We want to delay issuing the last dot as long as possible, ideally + // until after the prefetch. To accomplish this, set the insertion + // point above the dot. If we find anything dependent on the dot (at + // the top of this loop), we resume inserting after it. + builder.setInsertionPoint(prevDot); + } + } + } + // update mapping of results + for (unsigned dstIdx : llvm::seq(unsigned(0), op.getNumResults())) + mapping.map(op.getResult(dstIdx), newOp->getResult(dstIdx)); + } + + // prefetch next iteration + SmallVector yieldValues; + for (Value v : forOp.getBody()->getTerminator()->getOperands()) + yieldValues.push_back(mapping.lookupOrDefault(v)); + for (triton::DotOp dot : dots) { + Attribute dotEncoding = dot.getType().getEncoding(); + Value aToYield = generatePrefetch(mapping.lookup(dot2aYield[dot]), 0, true, + dotEncoding, builder); + cloneElementwiseOps(aToYield, dot2aVals[dot], builder); + yieldValues.push_back(aToYield); + // bToYield + Value bToYield = generatePrefetch(mapping.lookup(dot2bYield[dot]), 1, true, + dotEncoding, builder); + cloneElementwiseOps(bToYield, dot2bVals[dot], builder); + yieldValues.push_back(bToYield); + } + // Update ops of yield + builder.setInsertionPointToEnd(newForOp.getBody()); + if (!yieldValues.empty()) + scf::YieldOp::create(builder, yieldOp.getLoc(), yieldValues); + return newForOp; +} + +} // anonymous namespace + +struct PrefetchPass : public impl::TritonGPUPrefetchBase { + void runOnOperation() override { + + // Canonicalize convert ops to make the pattern matching easier. + RewritePatternSet cleanUpPatterns(&getContext()); + triton::gpu::ConvertLayoutOp::getCanonicalizationPatterns(cleanUpPatterns, + &getContext()); + if (mlir::applyPatternsGreedily(getOperation(), std::move(cleanUpPatterns)) + .failed()) { + signalPassFailure(); + } + getOperation()->walk([&](scf::ForOp forOp) { + Prefetcher prefetcher(forOp); + + if (prefetcher.initialize().failed()) + return; + + prefetcher.emitPrologue(); + + scf::ForOp newForOp = prefetcher.createNewForOp(); + + // replace the original loop + for (unsigned i = 0; i < forOp->getNumResults(); ++i) + forOp->getResult(i).replaceAllUsesWith(newForOp->getResult(i)); + forOp->erase(); + }); + } +}; + +} // namespace gpu +} // namespace triton +} // namespace mlir diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp new file mode 100644 index 0000000000..8a88bc0a6e --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp @@ -0,0 +1,1894 @@ +#include "mlir/Analysis/SliceAnalysis.h" +#include "mlir/Analysis/TopologicalSortUtils.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/Dominance.h" +#include "mlir/IR/IRMapping.h" +#include "mlir/IR/Matchers.h" +#include "mlir/IR/PatternMatch.h" +#include "mlir/IR/Verifier.h" +#include "mlir/Interfaces/InferTypeOpInterface.h" +#include "mlir/Interfaces/SideEffectInterfaces.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Support/LogicalResult.h" +#include "mlir/Transforms/GreedyPatternRewriteDriver.h" +#include "mlir/Transforms/Passes.h" +#include "mlir/Transforms/RegionUtils.h" +#include "triton/Analysis/Utility.h" +#ifdef __TLE__ +#include "tle/dialect/include/Transforms/TransformAttrs.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#endif +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/TritonGPUInterfaces.h" +#ifdef __TLE__ +#include "tle/dialect/include/Transforms/EncodingRematerialization.h" +#endif +#include "triton/Dialect/TritonGPU/Transforms/Passes.h" +#include "triton/Dialect/TritonGPU/Transforms/TritonGPUConversion.h" +#include "triton/Dialect/TritonGPU/Transforms/Utility.h" +#include + +namespace mlir::triton::gpu { + +#define GEN_PASS_DEF_TRITONGPUREMOVELAYOUTCONVERSIONS +#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc" + +#define DEBUG_TYPE "tritongpu-remove-layout-conversions" +#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") +#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n") + +namespace { + +#ifdef __TLE__ +static bool touchesTleRemotePointerPath(Value value, DenseSet &visited) { + if (!visited.insert(value).second) + return false; + Operation *def = value.getDefiningOp(); + if (!def) + return false; + StringRef opName = def->getName().getStringRef(); + if (opName == "tle.remote_pointers") + return true; + if (auto ifOp = dyn_cast(def)) { + auto result = dyn_cast(value); + if (!result) + return false; + unsigned idx = result.getResultNumber(); + return touchesTleRemotePointerPath(ifOp.thenYield().getOperand(idx), + visited) || + touchesTleRemotePointerPath(ifOp.elseYield().getOperand(idx), + visited); + } + for (Value operand : def->getOperands()) { + if (touchesTleRemotePointerPath(operand, visited)) + return true; + } + return false; +} +#endif + +// ----------------------------------------------------------------------------- +// +// ----------------------------------------------------------------------------- + +// The current algorithm works by analyzing the IR and doing a one-shot rewrite +// based on the analysis. The algorithm is as follows. +// +// 1. Find all the anchor ops. These are ops that have a layout we want to +// preserve. +// +// 2. For each anchor, propagate its layout to all its descendants. +// An op can have multiple ancestors that are anchors, so at this stage an op +// may have multiple layouts associated with it. +// +// 3. Resolve conflicts by deciding which of the multiple layouts the op should +// keep, inserting convert-layout ops to resolve conflicts. After this +// stage, each value has only one layout associated with it. +// +// 4. Rewrite the IR by walking the function in dominance order. Since we +// assume the IR is structured we just need to process the regions in the +// correct order. For each op, rewrite it using the layout decided by the +// analysis phase. +class LayoutPropagation { +public: + // Structure to keep track of the layout associated to a value. + struct LayoutInfo { + LayoutInfo(Attribute encoding) { encodings.insert(encoding); } + LayoutInfo() {} + llvm::SmallSetVector encodings; + }; + LayoutPropagation(FuncOp F) : funcOp(F) {} + // Find the anchor ops and set their layout in the data structure. + void initAnchorLayout(); + // Recursively Propagate the layout to all the users of the anchor ops until + // we reach a fix point. + void propagateLayout(); + // Add layouts given in `Info` to the uses of `value`. + SmallVector propagateToUsers(Value value, LayoutInfo &info); + // Set the encoding to all the values and fill out the values with new layout + // in `changed`. + void setEncoding(ValueRange values, LayoutInfo &info, + SmallVector &changed, Operation *op); + // Resolve cases where a value has multiple layouts associated to it. + void resolveConflicts(); + // Rewrite the IR for the full module. + void rewrite(); + // Rewrite the IR for a region. + void rewriteRegion(Region &R); + // Rewrite an op based on the layout picked by the analysis. + Operation *rewriteOp(Operation *op); + // Rewrite a for op based on the layout picked by the analysis. + Operation *rewriteForOp(scf::ForOp forOp); + Operation *rewriteWhileOp(scf::WhileOp whileOp); + Operation *rewriteIfOp(scf::IfOp ifOp); + void rewriteYieldOp(scf::YieldOp yieldOp); + void rewriteConditionOp(scf::ConditionOp conditionOp); + void rewriteReduceToScalar(Operation *reduceOp); + void rewriteAssertOp(AssertOp assertOp); + Operation *cloneElementwise(OpBuilder &rewriter, Operation *op, + Attribute encoding); + // Map the original value to the rewritten one. + void map(Value old, Value newV); + // Return the mapped value in the given encoding. This will insert a convert + // if the encoding is different than the encoding decided at resolve time. + Value getValueAs(Value value, Attribute encoding); + // Return the original value mapped to the new desired encoding. + Value getRewrittenValue(Value value); + // Dump the current stage of layout information. + void dump(); + +private: + // map from value to layout information. + llvm::MapVector layouts; + // map of the values rewrite based on their encoding. + DenseMap, Value> rewriteMapping; + SetVector opToDelete; + FuncOp funcOp; +}; + +class LayoutRematerialization { +public: + LayoutRematerialization(FuncOp F) : funcOp(F) {} + + // Map the original value to the remat'ed one. + void addRematValue(Value old, Attribute encoding, Value newV); + // Get the remat'ed value in the given encoding, if one already exists and + // is different then the layout conversion root. + Value getRematValue(Value value, Attribute encoding) const { + return rematMapping.lookup({value, encoding}); + } + + void cleanup(); + bool backwardRematerialization(); + void backwardRematerialization(ConvertLayoutOp convertOp); + // TODO: Merge the three hoistConvert*(); functions as they are duplicate code + void hoistConvertDotOperand(); + void hoistConvertDotOperand(ConvertLayoutOp convertOp); + void hoistConvertOnTopOfExtOrBroadcast(); + void hoistConvertOnTopOfExtOrBroadcast(ConvertLayoutOp convertOp); + void hoistConvertIntoConditionals(); + void hoistConvertIntoConditionals(ConvertLayoutOp convertOp); + void rewriteSlice(SetVector &slice, DenseMap &layout, + ConvertLayoutOp convertOp, IRMapping &mapping); + void rewriteSlice(SetVector &slice, DenseMap &layout, + ConvertLayoutOp convertOp); + + LogicalResult + getConvertBackwardSlice(OpOperand &root, Attribute rootEncoding, + SetVector &slice, + DenseMap &layout, + std::function stopPropagation); + + LogicalResult getRematerializableSlice( + OpOperand &root, Attribute rootEncoding, SetVector &slice, + DenseMap &layout, + std::function stopPropagation = nullptr); + +private: + void updateRematMapping(SmallVector> &values); + // Existing tuples of (value, layout) that needs to be updated when recreating + // scf ops. This prevents keeping track of Values that have been delete when + // rewriting slices. + DenseMap mappedValues; + // map of the values remat based on encoding. + DenseMap, Value> rematMapping; + // DenseMap, Operation*> + SetVector opToDelete; + FuncOp funcOp; + DominanceInfo domInfo; + PostDominanceInfo postDomInfo; +}; + +void LayoutRematerialization::addRematValue(Value old, Attribute encoding, + Value newV) { + LDBG("addRematValue " << old << " encoding " << encoding << " " << newV); + rematMapping[{old, encoding}] = newV; + mappedValues[old] = encoding; +} + +// Remove unneeded values now that we are done with the rematMapping. +void LayoutRematerialization::cleanup() { + for (Operation *op : llvm::reverse(opToDelete)) + op->erase(); +} + +// Return true if the op is an op with a layout we don't want to change. We will +// propagate the layout starting from anchor ops. +bool isLayoutAnchor(Operation *op) { + if (isa(op)) + return true; + if (isa(op)) + return isExpensiveLoadOrStore(op); + if (isa(op)) + return true; + if (auto gatherOp = dyn_cast(op)) + return gatherOp.getEfficientLayout(); + + // Heuristic: Mark permuting reshape as a layout anchor. Its dst can be + // anything, so it stops forward-propagation of layouts. We rely on the + // backwards pass to fix it up if necessary. (If we didn't do this, then + // anything following the reshape won't be covered by the forward pass at + // all.) + if (auto reshape = dyn_cast(op)) + return reshape.getAllowReorder(); + + return false; +} + +void LayoutPropagation::initAnchorLayout() { + auto addAnchor = [&](Value v) { + if (auto tensorType = dyn_cast(v.getType())) { + layouts.insert({v, LayoutInfo(tensorType.getEncoding())}); + } + }; + + // Consider function args as anchors. This makes it easier to write tests -- + // you can pass a tensor with an encoding as an arg, instead of explicitly + // calling tt.load. + for (auto arg : funcOp.getArguments()) { + addAnchor(arg); + } + + funcOp.walk([&](Operation *op) { + if (isLayoutAnchor(op)) { + for (auto result : op->getResults()) { + addAnchor(result); + } + } + }); +} + +void LayoutPropagation::setEncoding(ValueRange values, LayoutInfo &info, + SmallVector &changed, + Operation *op) { + for (Value value : values) { + if (!isa(value.getType())) + continue; + bool hasChanged = false; + for (auto encoding : info.encodings) { + Attribute dstEncoding; + if (isa(op)) { + // Try to remove the convert by making the dst encoding match the source + // encoding. + dstEncoding = encoding; + } else { + dstEncoding = inferDstEncoding(op, encoding); + } + if (dstEncoding) + hasChanged |= layouts[value].encodings.insert(dstEncoding); + } + if (hasChanged) + changed.push_back(value); + } +} + +SmallVector LayoutPropagation::propagateToUsers(Value value, + LayoutInfo &info) { + SmallVector changed; + for (OpOperand &use : value.getUses()) { + Operation *user = use.getOwner(); + if (auto forOp = dyn_cast(user)) { + Value arg = forOp.getTiedLoopRegionIterArg(&use); + Value result = forOp.getTiedLoopResult(&use); + setEncoding({arg, result}, info, changed, user); + continue; + } + if (auto whileOp = dyn_cast(user)) { + Value arg = whileOp.getBeforeArguments()[use.getOperandNumber()]; + setEncoding({arg}, info, changed, user); + continue; + } + if (auto yieldOp = dyn_cast(user)) { + auto parent = yieldOp->getParentOp(); + SmallVector valuesToPropagate; + if (isa(parent)) + valuesToPropagate.push_back(parent->getResult(use.getOperandNumber())); + if (auto forOp = dyn_cast(parent)) + valuesToPropagate.push_back( + forOp.getRegionIterArg(use.getOperandNumber())); + if (auto whileOp = dyn_cast(parent)) + valuesToPropagate.push_back( + whileOp.getBeforeArguments()[use.getOperandNumber()]); + if (isa(parent)) + setEncoding(valuesToPropagate, info, changed, user); + continue; + } + if (auto conditionOp = dyn_cast(user)) { + auto whileOp = cast(conditionOp->getParentOp()); + // Skip arg 0 as it is the condition. + unsigned argIndex = use.getOperandNumber() - 1; + Value afterArg = whileOp.getAfterArguments()[argIndex]; + Value result = whileOp->getResult(argIndex); + setEncoding({afterArg, result}, info, changed, user); + continue; + } + if (auto dotWaitOp = dyn_cast(user)) { + unsigned opIndex = use.getOperandNumber(); + Value result = dotWaitOp->getResult(opIndex); + setEncoding(result, info, changed, user); + continue; + } + if (auto gatherOp = dyn_cast(user)) { + // Propagate the layout through the indices only, and if the layout does + // not have an efficient layout set. + if (!gatherOp.getEfficientLayout() && + &use == &gatherOp.getIndicesMutable()) { + setEncoding(gatherOp.getResult(), info, changed, user); + continue; + } + } + if (user->hasTrait() || + user->hasTrait() || + isa(user)) { + // sunrise fix:S2的mma进行32bit转16bit时布局不一样，不能直接转换 + if(isa(user)) { + auto srcType = dyn_cast(user->getOperand(0).getType()); + auto dstType = dyn_cast(user->getResult(0).getType()); + if(srcType != nullptr && dstType != nullptr) { + auto srcElemType = srcType.getElementType(); + auto dstElemType = dstType.getElementType(); + if((srcElemType.getIntOrFloatBitWidth() == 32 && dstElemType.getIntOrFloatBitWidth() == 16) + || (dstElemType.getIntOrFloatBitWidth() == 8)){ + continue; + } + } + } + setEncoding(user->getResults(), info, changed, user); + continue; + } + } + return changed; +} + +void LayoutPropagation::propagateLayout() { + SmallVector queue; + for (auto it : layouts) { + queue.push_back(it.first); + } + while (!queue.empty()) { + Value currentValue = queue.back(); + LayoutInfo info = layouts[currentValue]; + queue.pop_back(); + SmallVector changed = propagateToUsers(currentValue, info); + + LLVM_DEBUG({ + DBGS() << "propagateLayout considering " << currentValue << ", which has " + << info.encodings.size() << " candidate encoding(s):\n"; + for (Attribute encoding : info.encodings) + DBGS() << " " << encoding << "\n"; + DBGS() << "changed: " << changed.size() << "\n"; + }); + + queue.insert(queue.end(), changed.begin(), changed.end()); + } +} + +void LayoutPropagation::resolveConflicts() { + for (auto &it : layouts) { + Operation *op = it.first.getDefiningOp(); + LayoutInfo &info = it.second; + if (info.encodings.size() <= 1) + continue; + // Hacky resolve, prefer block encoding. + // TODO: add a proper heuristic. + Attribute encoding = *info.encodings.begin(); + bool isLoadOrStore = + op && isa(op); + for (Attribute e : info.encodings) { + if ((isLoadOrStore && isa(e)) || + (!isLoadOrStore && isa(e))) { + encoding = e; + break; + } + } + info.encodings.clear(); + info.encodings.insert(encoding); + } +} + +void LayoutPropagation::dump() { + for (auto it : layouts) { + llvm::errs() << "Value: "; + OpPrintingFlags flags; + flags.skipRegions(); + it.first.print(llvm::errs(), flags); + llvm::errs() << " \n encoding:\n"; + for (auto encoding : it.second.encodings) { + encoding.print(llvm::errs()); + llvm::errs() << "\n"; + } + llvm::errs() << "--\n"; + } +} + +void LayoutPropagation::rewrite() { rewriteRegion(funcOp->getRegion(0)); } + +bool reduceToScalar(Operation *op) { + // For reductions returning a scalar we can change the src encoding without + // affecting the output. + return isa(op) && !isa(op->getResultTypes()[0]); +} + +void LayoutPropagation::rewriteRegion(Region ®ion) { + std::deque queue = {®ion}; + while (!queue.empty()) { + Region *currentRegion = queue.front(); + queue.pop_front(); + for (Operation &op : currentRegion->getOps()) { + bool needRewrite = false; + SmallVector results = op.getResults(); + for (Value result : results) { + auto it = layouts.find(result); + // If we haven't mapped this value skip. + if (it == layouts.end()) + continue; + LayoutInfo &info = it->second; + assert(info.encodings.size() == 1 && + "we should have resolved to a single encoding"); + auto encoding = cast(result.getType()).getEncoding(); + // If the encoding is already what we want skip. + if (encoding == *info.encodings.begin()) + continue; + needRewrite = true; + } + if (needRewrite) { + Operation *newOp = rewriteOp(&op); + for (Region &R : newOp->getRegions()) + queue.push_back(&R); + } else if (auto yieldOp = dyn_cast(&op)) { + rewriteYieldOp(yieldOp); + } else if (auto conditionOp = dyn_cast(&op)) { + rewriteConditionOp(conditionOp); + } else if (reduceToScalar(&op)) { + rewriteReduceToScalar(&op); + } else if (auto assertOp = dyn_cast(&op)) { + rewriteAssertOp(assertOp); + } else { + // If we don't need to rewrite the op we still need to remap the + // operands. + for (OpOperand &operand : op.getOpOperands()) { + auto it = layouts.find(operand.get()); + if (it == layouts.end()) + continue; + Attribute encoding = + cast(operand.get().getType()).getEncoding(); + Value newOperand = getValueAs(operand.get(), encoding); + op.setOperand(operand.getOperandNumber(), newOperand); + } + for (Region &R : op.getRegions()) + queue.push_back(&R); + } + } + } + for (Operation *op : llvm::reverse(opToDelete)) + op->erase(); +} + +void LayoutPropagation::map(Value old, Value newV) { + rewriteMapping[{old, cast(newV.getType()).getEncoding()}] = + newV; +} + +Value LayoutPropagation::getRewrittenValue(Value value) { + auto tensorType = dyn_cast(value.getType()); + if (!tensorType) + return value; + auto layoutIt = layouts.find(value); + if (layoutIt == layouts.end()) { + return value; + } + assert(layoutIt->second.encodings.size() == 1 && + "we should have resolved to a single encoding"); + Attribute encodingPicked = *(layoutIt->second.encodings.begin()); + if (encodingPicked == tensorType.getEncoding()) + return value; + return rewriteMapping.at({value, encodingPicked}); +} + +Value LayoutPropagation::getValueAs(Value value, Attribute encoding) { + if (auto tensorType = dyn_cast(value.getType())) { + Value rewrittenValue = getRewrittenValue(value); + if (cast(rewrittenValue.getType()).getEncoding() == + encoding) + return rewrittenValue; + OpBuilder rewriter(value.getContext()); + rewriter.setInsertionPointAfterValue(rewrittenValue); + auto tmpType = tensorType.cloneWithEncoding(encoding); + Value converted = ConvertLayoutOp::create(rewriter, value.getLoc(), tmpType, + rewrittenValue); + // TODO: we could cache the conversion. + return converted; + } + return value; +} + +Operation *LayoutPropagation::cloneElementwise(OpBuilder &rewriter, + Operation *op, + Attribute encoding) { + Operation *newOp = rewriter.clone(*op); + + Attribute operandEnc; + if (op->getNumOperands() > 0) { + for (auto operand : op->getOperands()) { + auto ty = + dyn_cast(getRewrittenValue(operand).getType()); + if (!ty) + continue; + auto enc = ty.getEncoding(); + if (inferDstEncoding(op, enc) == encoding) { + operandEnc = enc; + break; + } + } + if (!operandEnc) + operandEnc = inferSrcEncoding(op, encoding); + assert(operandEnc); + } + + for (OpOperand &operand : op->getOpOperands()) { + newOp->setOperand(operand.getOperandNumber(), + getValueAs(operand.get(), operandEnc)); + } + + for (unsigned i = 0, e = op->getNumResults(); i < e; ++i) { + auto origType = dyn_cast(op->getResult(i).getType()); + if (!origType) + continue; + auto newType = origType.cloneWithEncoding(encoding); + newOp->getResult(i).setType(newType); + } + return newOp; +} + +Operation *LayoutPropagation::rewriteForOp(scf::ForOp forOp) { + SmallVector operands; + OpBuilder rewriter(forOp); + for (auto [operand, result] : + llvm::zip(forOp.getInitArgs(), forOp.getResults())) { + Value convertedOperand = operand; + if (layouts.count(result)) + convertedOperand = + getValueAs(operand, *layouts[result].encodings.begin()); + operands.push_back(convertedOperand); + } + auto newForOp = + scf::ForOp::create(rewriter, forOp.getLoc(), forOp.getLowerBound(), + forOp.getUpperBound(), forOp.getStep(), operands); + newForOp->setAttrs(forOp->getAttrs()); + newForOp.getBody()->getOperations().splice( + newForOp.getBody()->getOperations().begin(), + forOp.getBody()->getOperations()); + + for (auto [oldResult, newResult] : + llvm::zip(forOp.getResults(), newForOp.getResults())) { + if (oldResult.getType() == newResult.getType()) { + oldResult.replaceAllUsesWith(newResult); + continue; + } + map(oldResult, newResult); + } + + for (auto [oldArg, newArg] : llvm::zip(forOp.getBody()->getArguments(), + newForOp.getBody()->getArguments())) { + if (oldArg.getType() == newArg.getType()) { + oldArg.replaceAllUsesWith(newArg); + continue; + } + map(oldArg, newArg); + } + return newForOp.getOperation(); +} + +Operation *LayoutPropagation::rewriteWhileOp(scf::WhileOp whileOp) { + SmallVector operands; + SmallVector returnTypes; + OpBuilder rewriter(whileOp); + for (auto [operand, arg] : + llvm::zip(whileOp->getOperands(), whileOp.getBeforeArguments())) { + Value convertedOperand = operand; + if (layouts.count(arg)) + convertedOperand = getValueAs(operand, *layouts[arg].encodings.begin()); + operands.push_back(convertedOperand); + } + for (Value ret : whileOp.getResults()) { + auto it = layouts.find(ret); + if (it == layouts.end()) { + returnTypes.push_back(ret.getType()); + continue; + } + auto origType = dyn_cast(ret.getType()); + auto newType = origType.cloneWithEncoding(it->second.encodings[0]); + returnTypes.push_back(newType); + } + + auto newWhileOp = + scf::WhileOp::create(rewriter, whileOp.getLoc(), returnTypes, operands); + SmallVector argsTypesBefore; + for (Value operand : operands) + argsTypesBefore.push_back(operand.getType()); + SmallVector bbArgLocsBefore(argsTypesBefore.size(), + whileOp.getLoc()); + SmallVector bbArgLocsAfter(returnTypes.size(), whileOp.getLoc()); + rewriter.createBlock(&newWhileOp.getBefore(), {}, argsTypesBefore, + bbArgLocsBefore); + rewriter.createBlock(&newWhileOp.getAfter(), {}, returnTypes, bbArgLocsAfter); + + for (int i = 0; i < whileOp.getNumRegions(); ++i) { + newWhileOp->getRegion(i).front().getOperations().splice( + newWhileOp->getRegion(i).front().getOperations().begin(), + whileOp->getRegion(i).front().getOperations()); + } + + auto remapArg = [&](Value oldVal, Value newVal) { + if (oldVal.getType() == newVal.getType()) + oldVal.replaceAllUsesWith(newVal); + else + map(oldVal, newVal); + }; + for (auto [oldResult, newResult] : + llvm::zip(whileOp.getResults(), newWhileOp.getResults())) + remapArg(oldResult, newResult); + for (auto [oldArg, newArg] : + llvm::zip(whileOp.getBeforeArguments(), newWhileOp.getBeforeArguments())) + remapArg(oldArg, newArg); + for (auto [oldArg, newArg] : + llvm::zip(whileOp.getAfterArguments(), newWhileOp.getAfterArguments())) + remapArg(oldArg, newArg); + return newWhileOp.getOperation(); +} + +Operation *LayoutPropagation::rewriteIfOp(scf::IfOp ifOp) { + SmallVector operands; + OpBuilder rewriter(ifOp); + SmallVector newResultTypes(ifOp->getResultTypes()); + for (unsigned i = 0, e = ifOp->getNumResults(); i < e; ++i) { + auto it = layouts.find(ifOp->getResult(i)); + if (it == layouts.end()) + continue; + auto origType = cast(ifOp->getResult(i).getType()); + Attribute encoding = *(it->second.encodings.begin()); + newResultTypes[i] = origType.cloneWithEncoding(encoding); + } + auto newIfOp = scf::IfOp::create(rewriter, ifOp.getLoc(), newResultTypes, + ifOp.getCondition(), true, true); + newIfOp.getThenRegion().takeBody(ifOp.getThenRegion()); + newIfOp.getElseRegion().takeBody(ifOp.getElseRegion()); + for (auto [oldResult, newResult] : + llvm::zip(ifOp.getResults(), newIfOp.getResults())) { + if (oldResult.getType() == newResult.getType()) { + oldResult.replaceAllUsesWith(newResult); + continue; + } + map(oldResult, newResult); + } + return newIfOp.getOperation(); +} + +void LayoutPropagation::rewriteYieldOp(scf::YieldOp yieldOp) { + Operation *parentOp = yieldOp->getParentOp(); + for (OpOperand &operand : yieldOp->getOpOperands()) { + Type yieldType = operand.get().getType(); + if (isa(parentOp)) + yieldType = parentOp->getResult(operand.getOperandNumber()).getType(); + if (auto whileOp = dyn_cast(parentOp)) + yieldType = + whileOp.getBeforeArguments()[operand.getOperandNumber()].getType(); + auto tensorType = dyn_cast(yieldType); + if (!tensorType) + continue; + Value newOperand = getValueAs(operand.get(), tensorType.getEncoding()); + yieldOp->setOperand(operand.getOperandNumber(), newOperand); + } +} + +void LayoutPropagation::rewriteConditionOp(scf::ConditionOp conditionOp) { + scf::WhileOp whileOp = cast(conditionOp->getParentOp()); + for (unsigned i = 1; i < conditionOp->getNumOperands(); ++i) { + OpOperand &operand = conditionOp->getOpOperand(i); + Type argType = whileOp->getResult(operand.getOperandNumber() - 1).getType(); + auto tensorType = dyn_cast(argType); + if (!tensorType) + continue; + Value newOperand = getValueAs(operand.get(), tensorType.getEncoding()); + conditionOp->setOperand(operand.getOperandNumber(), newOperand); + } +} + +void LayoutPropagation::rewriteReduceToScalar(Operation *reduceOp) { + OpBuilder rewriter(reduceOp); + Attribute srcEncoding; + // Since all the operands need to have the same encoding pick the first one + // and use it for all the operands. + for (Value operand : reduceOp->getOperands()) { + auto it = layouts.find(operand); + if (it != layouts.end()) { + srcEncoding = it->second.encodings[0]; + break; + } + } + if (!srcEncoding) + return; + for (OpOperand &operand : reduceOp->getOpOperands()) { + Value newOperand = getValueAs(operand.get(), srcEncoding); + reduceOp->setOperand(operand.getOperandNumber(), newOperand); + } +} + +void LayoutPropagation::rewriteAssertOp(AssertOp assertOp) { + Attribute srcEncoding; + // Only need to deal with the first operand which is the condition tensor. + Value operand = assertOp->getOperand(0); + auto it = layouts.find(operand); + if (it == layouts.end()) + return; + srcEncoding = it->second.encodings[0]; + Value newOperand = getValueAs(operand, srcEncoding); + assertOp->setOperand(0, newOperand); +} + +Operation *LayoutPropagation::rewriteOp(Operation *op) { + opToDelete.insert(op); + if (auto forOp = dyn_cast(op)) + return rewriteForOp(forOp); + if (auto whileOp = dyn_cast(op)) + return rewriteWhileOp(whileOp); + if (auto ifOp = dyn_cast(op)) + return rewriteIfOp(ifOp); + OpBuilder rewriter(op); + Attribute encoding = *layouts[op->getResult(0)].encodings.begin(); + if (auto convertOp = dyn_cast(op)) { + Attribute srcEncoding = convertOp.getSrc().getType().getEncoding(); + auto it = layouts.find(convertOp.getSrc()); + if (it != layouts.end()) + srcEncoding = *(it->second.encodings.begin()); + Value src = getValueAs(convertOp.getSrc(), srcEncoding); + auto tensorType = cast(op->getResult(0).getType()); + auto newType = tensorType.cloneWithEncoding(encoding); + auto cvt = ConvertLayoutOp::create(rewriter, op->getLoc(), newType, src); + map(op->getResult(0), cvt.getResult()); + return cvt.getOperation(); + } + if (canFoldIntoConversion(op, encoding)) { + Operation *newOp = rewriter.clone(*op); + auto tensorType = cast(op->getResult(0).getType()); + auto newType = tensorType.cloneWithEncoding(encoding); + auto cvt = ConvertLayoutOp::create(rewriter, op->getLoc(), newType, + newOp->getResult(0)); + map(op->getResult(0), cvt.getResult()); + return cvt.getOperation(); + } + if (op->hasTrait() || + op->hasTrait() || + isa(op)) { + Operation *newOp = cloneElementwise(rewriter, op, encoding); + for (auto [oldResult, newResult] : + llvm::zip(op->getResults(), newOp->getResults())) { + if (oldResult.getType() == newResult.getType()) { + oldResult.replaceAllUsesWith(newResult); + continue; + } + map(oldResult, newResult); + } + return newOp; + } + llvm::report_fatal_error("unexpected op in rewrite"); + return nullptr; +} + +bool canBeRemat(Operation *op) { + if (isa(op)) + return !isExpensiveLoadOrStore(op); + if (isa(op)) + return false; + if (auto gather = dyn_cast(op)) + return !gather.getEfficientLayout(); + + if (isa(op)) + return false; + + return true; +} + +void LayoutRematerialization::updateRematMapping( + SmallVector> &values) { + for (auto [old, newV] : values) { + auto it = mappedValues.find(old); + if (it != mappedValues.end()) { + Attribute encoding = it->second; + auto rematIt = rematMapping.find({old, it->second}); + assert(rematIt != rematMapping.end()); + Value replacedValue = rematIt->second; + rematMapping.erase(rematIt); + mappedValues.erase(it); + // Loop through the replacement value to find the new version of remat + // value. This should be okay as the number of values should be small. + for (auto [before, after] : values) { + if (before == replacedValue) { + replacedValue = after; + break; + } + } + rematMapping[{newV, encoding}] = replacedValue; + mappedValues[newV] = encoding; + } + } +} + +void LayoutRematerialization::rewriteSlice(SetVector &slice, + DenseMap &layout, + ConvertLayoutOp convertOp, + IRMapping &mapping) { + SetVector opsToRewrite; + // Keep track of yield operands that need to be duplicated. + DenseMap> yieldOperandsMap; + // Keep these around to remove them from the slice after our collection pass + // This ensures we don't duplicate them during an for rewrite or causing the + // for/yield to fall out of sync + SetVector valuesWithExistingRemat; + for (Value v : slice) { + auto layoutIt = layout.find(v); + assert(layoutIt != layout.end()); + // If we already have a remat value for this value, use it. + if (Value remat = getRematValue(v, layoutIt->second)) { + mapping.map(v, remat); + valuesWithExistingRemat.insert(v); + continue; + } + if (v.getDefiningOp()) { + opsToRewrite.insert(v.getDefiningOp()); + if (auto ifOp = v.getDefiningOp()) { + unsigned operandIdx = cast(v).getResultNumber(); + opsToRewrite.insert(ifOp.thenYield().getOperation()); + yieldOperandsMap[ifOp.thenYield()].push_back(operandIdx); + opsToRewrite.insert(ifOp.elseYield().getOperation()); + yieldOperandsMap[ifOp.elseYield()].push_back(operandIdx); + } + } else { + BlockArgument blockArg = cast(v); + Operation *parentOp = blockArg.getOwner()->getParentOp(); + if (auto loopOp = cast(parentOp)) { + opsToRewrite.insert(loopOp.getOperation()); + OpOperand *operand = loopOp.getTiedLoopYieldedValue(blockArg); + auto yieldOp = blockArg.getOwner()->getTerminator(); + yieldOperandsMap[yieldOp].push_back(operand->getOperandNumber()); + opsToRewrite.insert(yieldOp); + } + } + } + slice.set_subtract(valuesWithExistingRemat); + opsToRewrite = mlir::topologicalSort(opsToRewrite); + + // replaceAllUsesWith calls delayed until after initial rewrite. + // This is required for slice.count(value) to work mid rewrite. + SmallVector> replacements; + + SmallVector deadOps; + IRRewriter builder(slice.begin()->getContext()); + for (Operation *op : opsToRewrite) { + if (auto forOp = dyn_cast(op)) { + // Keep a mapping of the operands index to the new operands index. + SmallVector> argMapping; + SmallVector newOperands; + for (auto arg : forOp.getRegionIterArgs()) { + if (slice.count(arg)) { + OpOperand &initVal = *forOp.getTiedLoopInit(arg); + argMapping.push_back(std::make_pair( + forOp.getTiedLoopResult(&initVal).getResultNumber(), + forOp.getInitArgs().size() + newOperands.size())); + newOperands.push_back(mapping.lookup(initVal.get())); + } + } + // Create a new for loop with the new operands. + scf::ForOp newForOp = replaceForOpWithNewSignature( + builder, forOp, newOperands, replacements); + deadOps.push_back(forOp.getOperation()); + Block &loopBody = *newForOp.getBody(); + for (auto m : argMapping) { + mapping.map(forOp.getResult(m.first), newForOp.getResult(m.second)); + int numIndVars = newForOp.getNumInductionVars(); + mapping.map(loopBody.getArgument(m.first + numIndVars), + loopBody.getArgument(m.second + numIndVars)); + LLVM_DEBUG({ + DBGS() << "mapping forOp " + << loopBody.getArgument(m.first + numIndVars) << " to " + << loopBody.getArgument(m.second + numIndVars) << '\n'; + }); + // The result is not in the layout/slice, the argument is. + Value oldArg = loopBody.getArgument(m.first + numIndVars); + addRematValue(newForOp.getResult(m.first), layout[oldArg], + newForOp.getResult(m.second)); + addRematValue(oldArg, layout[oldArg], + loopBody.getArgument(m.second + numIndVars)); + } + continue; + } + if (auto ifOp = dyn_cast(op)) { + SmallVector newTypes; + for (auto res : ifOp.getResults()) { + if (slice.count(res)) { + auto it = layout.find(res); + assert(it != layout.end()); + + auto oldType = cast(res.getType()); + auto newType = oldType.cloneWithEncoding(it->second); + newTypes.push_back(newType); + } + } + scf::IfOp newIfOp = + replaceIfOpWithNewSignature(builder, ifOp, newTypes, replacements); + unsigned oldIdx = 0; + unsigned newIdx = ifOp.getNumResults(); + for (auto res : ifOp.getResults()) { + if (slice.count(res)) { + // Why can't we use res instead of ifOp.getResult(oldIdx)? + mapping.map(ifOp.getResult(oldIdx), newIfOp.getResult(newIdx)); + addRematValue(ifOp.getResult(oldIdx), layout[res], + newIfOp.getResult(newIdx)); + ++newIdx; + } + ++oldIdx; + } + deadOps.push_back(ifOp.getOperation()); + continue; + } + builder.setInsertionPoint(op); + if (auto yieldOp = dyn_cast(op)) { + auto yieldOperands = llvm::to_vector(yieldOp.getOperands()); + SmallVector operandsToRewrite = yieldOperandsMap[op]; + // Sort so that operands are added in the same order as the new scf + // results/arguments. + std::sort(operandsToRewrite.begin(), operandsToRewrite.end()); + for (int operandIdx : operandsToRewrite) { + yieldOperands.push_back(mapping.lookup(yieldOp.getOperand(operandIdx))); + } + scf::YieldOp::create(builder, op->getLoc(), yieldOperands); + op->erase(); + continue; + } + if (isa(op)) { + Operation *newOp = builder.clone(*op); + auto tensorType = cast(op->getResult(0).getType()); + auto newType = tensorType.cloneWithEncoding(layout[op->getResult(0)]); + auto cvt = ConvertLayoutOp::create(builder, op->getLoc(), newType, + newOp->getResult(0)); + mapping.map(op->getResult(0), cvt.getResult()); + addRematValue(op->getResult(0), layout[op->getResult(0)], + cvt.getResult()); + continue; + } + Operation *newOp = builder.clone(*op, mapping); + for (auto [old, newV] : llvm::zip(op->getResults(), newOp->getResults())) { + auto it = layout.find(old); + if (it == layout.end()) + continue; + auto newType = + cast(old.getType()).cloneWithEncoding(it->second); + newV.setType(newType); + addRematValue(old, it->second, newV); + } + } + // Check mapping and see if there are existing convertOps on the old Argument + convertOp.replaceAllUsesWith(mapping.lookup(convertOp.getSrc())); + opToDelete.insert(convertOp); + + updateRematMapping(replacements); + for (auto &kv : replacements) { + builder.replaceAllUsesWith(std::get<0>(kv), std::get<1>(kv)); + } + + for (Operation *op : deadOps) + opToDelete.insert(op); +} + +void LayoutRematerialization::rewriteSlice(SetVector &slice, + DenseMap &layout, + ConvertLayoutOp convertOp) { + IRMapping mapping; + rewriteSlice(slice, layout, convertOp, mapping); +} + +LogicalResult LayoutRematerialization::getConvertBackwardSlice( + OpOperand &root, Attribute rootEncoding, SetVector &slice, + DenseMap &layout, + std::function stopPropagation) { + // Allow re-using existing conversions for a value. Check dominance of any + // reusable materializations against the root value. This is sufficient + // because the conversions are processed in post-order. + auto getExistingConversion = [&](OpOperand &value, Attribute encoding) { + Value remat = getRematValue(value.get(), encoding); + if (!remat) + return Value(); + // `value` can be replaced with an existing rematerialization if it + // dominates the current use of value. + Operation *user = value.getOwner(); + if (domInfo.properlyDominates(remat, user)) { + return remat; + } + // FIXME: If the current user is a conversion, then we know it will become + // a no-op when its operand is replaced with `remat`, but we need to check + // that its users are all dominated by `remat` so the IR is valid. + // if (isa(user) && remat.getDefiningOp() && + // domInfo.properlyDominates(user, remat.getDefiningOp())) { + // for (Operation *op : user->getUsers()) { + // if (!domInfo.dominates(remat, op)) + // return Value(); + // } + // return remat; + // } + return Value(); + }; + + return mlir::getConvertBackwardSlice(root, slice, rootEncoding, layout, + stopPropagation, getExistingConversion); +} + +LogicalResult LayoutRematerialization::getRematerializableSlice( + OpOperand &root, Attribute rootEncoding, SetVector &slice, + DenseMap &layout, + std::function stopPropagation) { + LogicalResult result = getConvertBackwardSlice(root, rootEncoding, slice, + layout, stopPropagation); + if (result.failed() || slice.empty()) + return failure(); + + // Check if all the operations in the slice can be rematerialized. + for (Value v : slice) { + if (Operation *op = v.getDefiningOp()) { + if (!canBeRemat(op)) + return failure(); + } + } + return success(); +} + +bool LayoutRematerialization::backwardRematerialization() { + bool changed = false; + // Go through each ConvertLayoutOp. + SmallVector convertOps; + funcOp.walk( + [&](ConvertLayoutOp convertOp) { convertOps.push_back(convertOp); }); + for (ConvertLayoutOp convertOp : convertOps) { + backwardRematerialization(convertOp); + if (!opToDelete.contains(convertOp)) { + // If the conversion didn't get removed, consider it for reuse in future + // backward slices. + addRematValue(convertOp.getSrc(), convertOp.getType().getEncoding(), + convertOp.getResult()); + } else { + changed = true; + } + } + return changed; +} + +void LayoutRematerialization::hoistConvertOnTopOfExtOrBroadcast() { + // Go through each ConvertLayoutOp. + SmallVector convertOps; + funcOp.walk( + [&](ConvertLayoutOp convertOp) { convertOps.push_back(convertOp); }); + for (ConvertLayoutOp convertOp : convertOps) { + hoistConvertOnTopOfExtOrBroadcast(convertOp); + if (!opToDelete.contains(convertOp)) { + // If the conversion didn't get removed, consider it for reuse in future + // backward slices. + addRematValue(convertOp.getSrc(), convertOp.getType().getEncoding(), + convertOp.getResult()); + } + } +} + +void LayoutRematerialization::hoistConvertIntoConditionals() { + // Go through each ConvertLayoutOp. + SmallVector convertOps; + funcOp.walk( + [&](ConvertLayoutOp convertOp) { convertOps.push_back(convertOp); }); + for (ConvertLayoutOp convertOp : convertOps) { + hoistConvertIntoConditionals(convertOp); + if (!opToDelete.contains(convertOp)) { + // If the conversion didn't get removed, consider it for reuse in future + // backward slices. + addRematValue(convertOp.getSrc(), convertOp.getType().getEncoding(), + convertOp.getResult()); + } + } +} + +static bool isExpensiveMathOp(Operation *op) { + // These operations are either multiple instructions or have throughput + // lower than 16 according to the arithmetic instructions table in: + // https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#arithmetic-instructions + return isa(op); +} + +static int64_t getByteCount(Value result, int64_t minElementCount = 0, + int64_t minBitWidth = 0) { + int64_t elementCount = 0; + int64_t dtypeBitWidth = 0; + if (auto tensorTy = dyn_cast(result.getType())) { + elementCount = tensorTy.getNumElements(); + auto elemType = tensorTy.getElementType(); + if (elemType.isIntOrFloat()) { + dtypeBitWidth = elemType.getIntOrFloatBitWidth(); + } + } + if (elementCount < minElementCount) { + elementCount = minElementCount; + } + if (dtypeBitWidth < minBitWidth) { + dtypeBitWidth = minBitWidth; + } + return (elementCount * dtypeBitWidth) >> 3; +} + +void LayoutRematerialization::backwardRematerialization( + ConvertLayoutOp convertOp) { + // DotOperand is hoisted by hoistDotOperand + RankedTensorType targetType = convertOp.getType(); + if (isa(targetType.getEncoding())) + return; + Value oldV = convertOp.getSrc(); + LDBG("check backward remat with source " << oldV << " encoding " + << targetType.getEncoding()); + // Check to see if there are existing remat'ed values for the pair of oldValue + // and encoding. Make sure it dominates the current conversion. + Value newV = getRematValue(oldV, targetType.getEncoding()); + if (newV && domInfo.properlyDominates(newV, convertOp)) { + // Replace it with the remat'ed value. + convertOp.replaceAllUsesWith(newV); + opToDelete.insert(convertOp); + LDBG("found remat'ed value" << newV); + return; + } + + // 1. Take a backward slice of all the tensor dependencies that can be + // rematerialized. + SetVector slice; + DenseMap layout; + LogicalResult result = getRematerializableSlice( + convertOp.getSrcMutable(), targetType.getEncoding(), slice, layout); + if (result.failed()) { + LDBG(" getRematerializableSlice failed"); + return; + } + + // 2. Determine whether rematerialisation is beneficial. + + // Identify all operations in the slice + SetVector sliceOps; + for (Value v : slice) { + if (Operation *op = v.getDefiningOp()) { + sliceOps.insert(op); + } + } + + // Compute single-use operations + DenseMap isSingleUse; + std::function isOpSingleUse; + isOpSingleUse = [&](Operation *op) -> bool { + // lookup in memoization array: + auto it = isSingleUse.find(op); + if (it != isSingleUse.end()) { + return it->second; + } + + bool singleUse = true; + + for (Value result : op->getResults()) { + for (Operation *user : result.getUsers()) { + if (user == convertOp) { + continue; + } + if (sliceOps.contains(user)) { + if (!isOpSingleUse(user)) { + singleUse = false; + break; + } + } else { + singleUse = false; + break; + } + } + if (!singleUse) { + break; + } + } + + // insert into memoization array: + isSingleUse[op] = singleUse; + return singleUse; + }; + + // Measure the number of bytes that we're manipulating with the + // ConvertLayoutOp. We pessimistically assume that we round-trip + // through shared memory and that we cannot vectorise sub-register + // loads/stores, so we set a minimum element count of 32 (the warp + // size and number of shared memory banks) and minimum bitwidth of + // 32 (the width per bank of the shared memory load/store unit). + int64_t convertLayoutBytes = getByteCount(convertOp.getSrc(), 32, 32); + + // We measure costs in standardised milli-SM-cycles. The smem load + // and store each cost 8 * convertLayoutBytes, and then we double + // it to account for extra cost due to synchronisation. + int64_t convertLayoutCost = 32 * convertLayoutBytes; + int64_t rematerialisationCost = 0; + + // Evaluate single-use status for every operation in slice + for (Operation *op : sliceOps) { + auto dialect = op->getDialect(); + if (isOpSingleUse(op)) { + // when we rematerialise, this operation does not get duplicated + // so it does not contribute to our cost model: + continue; + } else if (isa(op)) { + // special-case: arith.constant has zero cost + continue; + } else if (isa(op) || isa(op)) { + // optimistically assume L1-cached: + for (Value result : op->getResults()) { + rematerialisationCost += 8 * getByteCount(result); + } + } else if (isa(dialect)) { + // this is an arithmetic operation; we distinguish between cheap + // operations (such as floating point add/mul which can be fused + // as halves of a single-cycle FMA instruction) and expensive + // operations which use the special function unit and/or involve + // multiple instructions. + int64_t multiplier = isExpensiveMathOp(op) ? 8 : 1; + for (Value result : op->getResults()) { + rematerialisationCost += multiplier * getByteCount(result); + } + } else if (isa(op)) { + // Reduce op introduce much cost. + auto reduceOp = dyn_cast(op); + ReduceOpHelper helper(reduceOp); + if (!helper.isAssociative()) { + // We shouldn't rematerize a no associative reduce op if it has multiple + // use chain. + LDBG(" skipped rematerialization due to non-associative reduce in the " + "slice"); + return; + } + rematerialisationCost += helper.getIntraWarpSizeWithUniqueData(); + rematerialisationCost += 8 * helper.getInterWarpSizeWithUniqueData(); + } + } + + LLVM_DEBUG({ + DBGS() << " convert layout cost: " << convertLayoutCost << "\n"; + DBGS() << " rematerialisation cost: " << rematerialisationCost << "\n"; + }); + + if (rematerialisationCost > convertLayoutCost) { + LDBG(" skipped rematerialization due to higher cost"); + return; + } + + LLVM_DEBUG({ + DBGS() << " remat convert op " << convertOp << '\n'; + for (Value v : slice) + DBGS() << " " << v << '\n'; + }); + + // 3. Rewrite the slice. + rewriteSlice(slice, layout, convertOp); +} + +void LayoutRematerialization::hoistConvertDotOperand() { + // Go through each ConvertLayoutOp. + SmallVector convertOps; + funcOp.walk( + [&](ConvertLayoutOp convertOp) { convertOps.push_back(convertOp); }); + for (ConvertLayoutOp convertOp : convertOps) { + hoistConvertDotOperand(convertOp); + if (!opToDelete.contains(convertOp)) { + // If the conversion didn't get removed, consider it for reuse in future + // backward slices. + addRematValue(convertOp.getSrc(), convertOp.getType().getEncoding(), + convertOp.getResult()); + } + } +} + +void LayoutRematerialization::hoistConvertDotOperand( + ConvertLayoutOp convertOp) { + auto targetType = convertOp.getType(); + // The pass is targeted to MMA dot operands + +#ifdef __TLE__ + { + DenseSet visited; + if (touchesTleRemotePointerPath(convertOp.getSrc(), visited)) + return; + } +#endif + + auto canBePipelined = [&](ConvertLayoutOp convertOp) { + // FIXME: Check that the parent is a for loop + auto parent = convertOp->getParentOp(); + if (!parent) + return false; + + // Find all the dot-like ops in the for loop that have a dot operand + // encoding on the lhs and check if any of them post-dominates the load + + // cvt + SmallVector dotLikeOps; + parent->walk([&](Operation *op) { + if (!isa(op)) + return; + auto opType = dyn_cast(op->getOperand(0).getType()); + if (!opType) + return; + auto dotEnc = dyn_cast(opType.getEncoding()); + if (!dotEnc) + return; + if (isa(dotEnc.getParent())) + dotLikeOps.push_back(op); + }); + if (dotLikeOps.empty()) + return false; + return llvm::any_of(dotLikeOps, [&](Operation *dot) { + return postDomInfo.postDominates(dot, convertOp); + }); + }; + + // We move convert #dot_operand next to their loads. This is done + // so that it's then easy to pipeline these loads + if (!canBePipelined(convertOp)) + return; + + // We hoist over any operation that can be done without data movement between + // threads We do views and elementwise pure ops for now + auto noDataMovement = [](Operation *op) { + return (op->hasTrait() && isMemoryEffectFree(op)) || + isa( + op) || + isView(op); + }; + // Stop the slice as soon as we find an operation that cannot be done without + // data movement between threads + auto stop = std::not_fn(noDataMovement); + + SetVector slice; + DenseMap layout; + // Set-up the conversion "cache" + LogicalResult result = getConvertBackwardSlice( + convertOp.getSrcMutable(), targetType.getEncoding(), slice, layout, stop); + if (result.failed()) + return; + + IRMapping mapping; + OpBuilder builder(convertOp.getContext()); + SetVector innerSlice; + for (Value v : slice) { + if (!v.getDefiningOp()) { + LLVM_DEBUG( + { DBGS() << " Block arguments not supported. Got " << v << "\n"; }); + return; + } + + // We expect the leaves of the slice to be Load, DescriptorLoad or + // arith::Constant This could be generalised if necessary + if (!isa(v.getDefiningOp())) { + auto op = v.getDefiningOp(); + if (isa(op) || noDataMovement(op)) { + innerSlice.insert(v); + continue; + } else { + LLVM_DEBUG({ + DBGS() << " Leaves must be Load, DescriptorLoad or Constant. Got " + << v << "\n"; + }); + return; + } + } + Operation *loadOp = v.getDefiningOp(); + builder.setInsertionPointAfter(loadOp); + auto type = dyn_cast(loadOp->getResult(0).getType()); + if (!type) + continue; + auto newType = type.cloneWithEncoding(layout[loadOp->getResult(0)]); + auto newConvertOp = ConvertLayoutOp::create(builder, convertOp.getLoc(), + newType, loadOp->getResult(0)); + mapping.map(loadOp->getResult(0), newConvertOp.getResult()); + } + + if (innerSlice.empty()) { + return; + } + + LLVM_DEBUG({ + DBGS() << " Hoisting " << convertOp << '\n'; + for (Value v : innerSlice) + DBGS() << " " << v << '\n'; + }); + + rewriteSlice(innerSlice, layout, convertOp, mapping); +} + +// For convert left we try to hoist them above type extension to reduce the cost +// of the convert. +void LayoutRematerialization::hoistConvertOnTopOfExtOrBroadcast( + ConvertLayoutOp convertOp) { + // DotOperand is hoisted by hoistDotOperand + RankedTensorType targetType = convertOp.getType(); + if (isa(targetType.getEncoding())) + return; + + auto isExtOrBroadcastOp = [](Operation *op) { + if (isa(op)) { + return true; + } + if (auto fpToFpOp = dyn_cast(op)) { + auto srcType = cast(fpToFpOp.getOperand().getType()); + return getElementBitWidth(srcType) < + getElementBitWidth(cast(fpToFpOp.getType())); + } + return false; + }; + // 1. Take a backward slice of all the tensor dependencies. + SetVector slice; + DenseMap layout; + LogicalResult result = getRematerializableSlice( + convertOp.getSrcMutable(), targetType.getEncoding(), slice, layout, + isExtOrBroadcastOp); + if (result.failed()) + return; + + Operation *extOrBroadcastOp = nullptr; + unsigned sliceSize = slice.size(); + for (unsigned i = 0; i < sliceSize; i++) { + Value v = slice[i]; + Operation *op = v.getDefiningOp(); + if (!op) + continue; + if (isExtOrBroadcastOp(op)) { + SetVector tempSlice; + DenseMap tempLayout; + Attribute srcEncoding = inferSrcEncoding(op, layout[v]); + if (!srcEncoding) + return; + LogicalResult result = getRematerializableSlice( + op->getOpOperand(0), srcEncoding, tempSlice, tempLayout); + + // If a value is already assigned to a _different_ layout, + // we cannot propagate past this op (as it would conflict with + // an already-assigned layout). + for (auto [val, enc] : tempLayout) { + auto preexistingLayout = layout.find(val); + if (preexistingLayout != layout.end() && + preexistingLayout->second != enc) { + result = failure(); + break; + } + } + + // If we can rematerialize the rest of the ext slice we can ignore this + // ext as it won't need a convert. + if (result.succeeded()) { + slice.insert(tempSlice.begin(), tempSlice.end()); + layout.insert(tempLayout.begin(), tempLayout.end()); + continue; + } + // Only apply it if there is a single ext op otherwise we would have to + // duplicate the convert. + if (extOrBroadcastOp != nullptr) + return; + extOrBroadcastOp = op; + } + } + + if (extOrBroadcastOp == nullptr) + return; + Attribute dstEncoding = layout[extOrBroadcastOp->getResult(0)]; + Attribute srcEncoding = inferSrcEncoding(extOrBroadcastOp, dstEncoding); + if (!srcEncoding) + return; + // Move the convert before the ext op and rewrite the slice. + OpBuilder builder(extOrBroadcastOp); + auto tensorType = + cast(extOrBroadcastOp->getOperand(0).getType()); + auto newType = tensorType.cloneWithEncoding(srcEncoding); + auto newConvertOp = ConvertLayoutOp::create( + builder, convertOp.getLoc(), newType, extOrBroadcastOp->getOperand(0)); + Operation *newExtOrBroadcast = builder.clone(*extOrBroadcastOp); + newExtOrBroadcast->setOperand(0, newConvertOp.getResult()); + auto oldExtOrBroadcastType = + cast(extOrBroadcastOp->getResult(0).getType()); + Type newExtOrBroadcastType = + oldExtOrBroadcastType.cloneWithEncoding(dstEncoding); + newExtOrBroadcast->getResult(0).setType(newExtOrBroadcastType); + IRMapping mapping; + mapping.map(extOrBroadcastOp->getResult(0), newExtOrBroadcast->getResult(0)); + slice.remove(extOrBroadcastOp->getResult(0)); + // 3. Rewrite the slice. + rewriteSlice(slice, layout, convertOp, mapping); +} + +void LayoutRematerialization::hoistConvertIntoConditionals( + ConvertLayoutOp convertOp) { + // Take the backward slice of tensor dependencies rooted at the conversion, + // stopping at conditionals. This subslice is used to initialize the analysis. + SetVector slice; + DenseMap layout; + auto isIfOp = [](Operation *op) { return isa(op); }; + if (failed(getRematerializableSlice(convertOp.getSrcMutable(), + convertOp.getType().getEncoding(), slice, + layout, isIfOp))) + return; + + // These are the conditional edges above which conversions should be hoisted. + // The value represents the `scf.if` op result and the operand represents the + // edge into one of the branches. + SmallVector> hoistAbove; + + // The list of `scf.if` op results in the slice that are not rematerializable. + // Hoisting is terminated at these values. + SmallVector terminals; + + // This loop recurses through the subslices of the backwards dependencies, so + // re-query the size of `slice`. + for (unsigned i = 0; i != slice.size(); ++i) { + Value v = slice[i]; + auto ifOp = v.getDefiningOp(); + if (!ifOp) + continue; + + Attribute rootLayout = layout.at(v); + unsigned resIdx = cast(v).getResultNumber(); + + // Take the backward slice along each branch. + auto thenYield = + cast(ifOp.getThenRegion().front().getTerminator()); + auto elseYield = + cast(ifOp.getElseRegion().front().getTerminator()); + + OpOperand &thenRes = thenYield.getResultsMutable()[resIdx]; + OpOperand &elseRes = elseYield.getResultsMutable()[resIdx]; + + SetVector thenSlice, elseSlice; + DenseMap thenLayout, elseLayout; + + LogicalResult thenResult = getRematerializableSlice( + thenRes, rootLayout, thenSlice, thenLayout, isIfOp); + LogicalResult elseResult = getRematerializableSlice( + elseRes, rootLayout, elseSlice, elseLayout, isIfOp); + + // If propagation across both edges of this conditional succeeded, then we + // don't need to hoist across it. Merge into the current slice. + if (succeeded(thenResult) && succeeded(elseResult)) { + slice.insert(thenSlice.begin(), thenSlice.end()); + slice.insert(elseSlice.begin(), elseSlice.end()); + layout.insert(thenLayout.begin(), thenLayout.end()); + layout.insert(elseLayout.begin(), elseLayout.end()); + continue; + } + + // If propagation across both edges failed, then this conditional + // terminates backwards rematerialization. + if (failed(thenResult) && failed(elseResult)) { + terminals.push_back(cast(v)); + continue; + } + + // Only hoist into conditionals inside loops. The assumption is that an if + // inside a loop executes fewer than the total number of loop iterations, + // making this hoist profitable. + if (!isa(ifOp->getParentOp())) { + terminals.push_back(cast(v)); + continue; + } + + // The layout conversion can be rematerialized along one edge but not the + // other. We can hoist the conversion into the other branch. Push this + // into the subslice list for analysis. + if (succeeded(thenResult)) { + hoistAbove.emplace_back(v, &elseRes); + slice.insert(thenSlice.begin(), thenSlice.end()); + layout.insert(thenLayout.begin(), thenLayout.end()); + } else { + hoistAbove.emplace_back(v, &thenRes); + slice.insert(elseSlice.begin(), elseSlice.end()); + layout.insert(elseLayout.begin(), elseLayout.end()); + } + } + + // Exit early if there is nothing to do. + if (hoistAbove.empty()) + return; + + // Rematerialize failed hoists right before the condtional, and hoist those + // that succeeded into the branch and then rewrite the slice. + IRMapping mapping; + auto hoistRemat = [&](OpBuilder &b, Value v, Attribute encoding) { + auto tensorType = cast(v.getType()); + auto newType = tensorType.cloneWithEncoding(encoding); + Value newCvt = ConvertLayoutOp::create(b, convertOp.getLoc(), newType, v); + + mapping.map(v, newCvt); + slice.remove(v); + }; + for (Value v : terminals) { + OpBuilder b(v.getContext()); + b.setInsertionPointAfter(v.getDefiningOp()); + hoistRemat(b, v, layout.at(v)); + } + for (auto [result, edge] : hoistAbove) { + OpBuilder b(edge->getOwner()); + hoistRemat(b, edge->get(), layout.at(result)); + } + rewriteSlice(slice, layout, convertOp, mapping); +} + +bool backwardRematerialization(ModuleOp module) { + bool changed = false; + module.walk([&](FuncOp funcOp) { + LayoutRematerialization layoutRemat(funcOp); + changed |= layoutRemat.backwardRematerialization(); + layoutRemat.cleanup(); + }); + return changed; +} + +#ifdef __TLE__ +static void eraseOpsCreatedAfter(Operation *previous, Operation *insertBefore) { + Operation *cur = + previous ? previous->getNextNode() : &insertBefore->getBlock()->front(); + while (cur && cur != insertBefore) { + Operation *next = cur->getNextNode(); + cur->erase(); + cur = next; + } +} + +bool rematerializeStoreLayoutDemands(ModuleOp module) { + bool changed = false; + EncodingRematerializationPolicy policy; + IRRewriter rewriter(module.getContext()); + DominanceInfo dominance(module); + + SmallVector stores; + module.walk([&](StoreOp store) { stores.push_back(store); }); + + for (StoreOp store : stores) { + if (!store || !store->getBlock()) + continue; + auto valueConvert = store.getValue().getDefiningOp(); + if (!valueConvert) + continue; + auto targetTy = dyn_cast(valueConvert.getSrc().getType()); + if (!targetTy || !targetTy.getEncoding()) + continue; + if (!isa(targetTy.getEncoding())) + continue; + + Operation *rollbackPrevious = store->getPrevNode(); + EncodingRematerializationCache cache; + auto ptr = rematerializeWithEncoding(rewriter, store, store.getPtr(), + targetTy.getEncoding(), cache, + dominance, policy); + if (failed(ptr)) + continue; + + Value mask; + if (Value oldMask = store.getMask()) { + auto rematMask = rematerializeWithEncoding(rewriter, store, oldMask, + targetTy.getEncoding(), cache, + dominance, policy); + if (failed(rematMask)) { + eraseOpsCreatedAfter(rollbackPrevious, store); + continue; + } + mask = *rematMask; + } + + store.getPtrMutable().assign(*ptr); + store.getValueMutable().assign(valueConvert.getSrc()); + if (mask) + store.getMaskMutable().assign(mask); + if (valueConvert->use_empty()) + valueConvert.erase(); + changed = true; + } + + return changed; +} + +FailureOr +inferLocalStoreSourceTargetEncoding(LocalStoreOp store, + DominanceInfo &dominance, + EncodingRematerializationPolicy &policy) { + auto sourceTy = dyn_cast(store.getSrc().getType()); + if (!sourceTy || !sourceTy.getEncoding()) + return failure(); + + if (auto convert = store.getSrc().getDefiningOp()) { + auto convertSrcTy = dyn_cast(convert.getSrc().getType()); + if (convertSrcTy && convertSrcTy.getEncoding() && + isa(convertSrcTy.getEncoding()) && + convertSrcTy.getEncoding() != sourceTy.getEncoding()) + return convertSrcTy.getEncoding(); + } + + SmallVector equivalentMmaEncodings; + collectAvailableEquivalentNvidiaMmaEncodings(store.getSrc(), + store.getOperation(), dominance, + policy, equivalentMmaEncodings); + if (equivalentMmaEncodings.empty()) + return failure(); + if (equivalentMmaEncodings.front() == sourceTy.getEncoding()) + return failure(); + return equivalentMmaEncodings.front(); +} + +bool rematerializeLocalStoreLayoutDemands(ModuleOp module) { + bool changed = false; + EncodingRematerializationPolicy policy; + IRRewriter rewriter(module.getContext()); + DominanceInfo dominance(module); + + SmallVector stores; + module.walk([&](LocalStoreOp store) { stores.push_back(store); }); + + for (LocalStoreOp store : stores) { + if (!store || !store->getBlock()) + continue; + + auto targetEncoding = + inferLocalStoreSourceTargetEncoding(store, dominance, policy); + if (failed(targetEncoding)) + continue; + + EncodingRematerializationCache cache; + auto source = + rematerializeWithEncoding(rewriter, store, store.getSrc(), + *targetEncoding, cache, dominance, policy); + if (failed(source)) + continue; + + store.getSrcMutable().assign(*source); + changed = true; + } + + return changed; +} +#endif + +void hoistConvert(ModuleOp module) { + SmallVector convertOps; + module.walk([](FuncOp funcOp) { + LayoutRematerialization layoutRemat(funcOp); + layoutRemat.hoistConvertOnTopOfExtOrBroadcast(); + layoutRemat.cleanup(); + + layoutRemat = LayoutRematerialization(funcOp); + layoutRemat.hoistConvertIntoConditionals(); + layoutRemat.cleanup(); + + layoutRemat = LayoutRematerialization(funcOp); + layoutRemat.hoistConvertDotOperand(); + layoutRemat.cleanup(); + }); +} +} // namespace + +class TritonGPURemoveLayoutConversionsPass + : public impl::TritonGPURemoveLayoutConversionsBase< + TritonGPURemoveLayoutConversionsPass> { +public: + // Cleanup convert ops. + void cleanupConvertOps() { + MLIRContext *context = &getContext(); + ModuleOp m = getOperation(); + RewritePatternSet cleanUpPatterns(context); + ConvertLayoutOp::getCanonicalizationPatterns(cleanUpPatterns, context); + if (applyPatternsGreedily(m, std::move(cleanUpPatterns)).failed()) { + signalPassFailure(); + } + + LLVM_DEBUG({ + DBGS() << "Module after canonicalizing:\n"; + m.dump(); + }); + } + + void runOnOperation() override { + MLIRContext *context = &getContext(); + ModuleOp m = getOperation(); + + // 1. Propagate layout forward starting from "anchor" ops. + m.walk([](FuncOp funcOp) { + LayoutPropagation layoutPropagation(funcOp); + layoutPropagation.initAnchorLayout(); + layoutPropagation.propagateLayout(); + layoutPropagation.resolveConflicts(); + layoutPropagation.rewrite(); + }); + + LLVM_DEBUG({ + DBGS() << "Module after propagating layouts forward:\n"; + m.dump(); + }); + + cleanupConvertOps(); + + bool changed = false; + do { + changed = false; + // 2. For remaining convert ops, try to rematerialize the slice of + // producer operation to avoid having to convert. + changed = backwardRematerialization(m); + LLVM_DEBUG({ + DBGS() << "Module after backward remat:\n"; + m.dump(); + }); + + // Cleanup dummy converts created during backward remat. + cleanupConvertOps(); +#ifdef __TLE__ + if (m->hasAttr( + ::mlir::triton::tle::kTleEnableEncodingRematerializationAttr)) { + changed |= rematerializeStoreLayoutDemands(m); + changed |= rematerializeLocalStoreLayoutDemands(m); + cleanupConvertOps(); + } +#endif + } while (changed); + // 3. For remaining converts, try to hoist them above cast generating larger + // size types in order to reduce the cost of the convert op. + hoistConvert(m); + LLVM_DEBUG({ + DBGS() << "Module after hoisting converts:\n"; + m.dump(); + }); + + // 4. Apply clean up patterns to remove remove dead convert and dead code + // generated by the previous transformations. + RewritePatternSet cleanUpPatterns2(context); + populateForOpDeadArgumentElimination(cleanUpPatterns2); + scf::ForOp::getCanonicalizationPatterns(cleanUpPatterns2, context); + scf::IfOp::getCanonicalizationPatterns(cleanUpPatterns2, context); + ConvertLayoutOp::getCanonicalizationPatterns(cleanUpPatterns2, context); + if (applyPatternsGreedily(m, std::move(cleanUpPatterns2)).failed()) { + signalPassFailure(); + } + LLVM_DEBUG({ + DBGS() << "Module after final cleanups:\n"; + m.dump(); + }); + } +}; + +} // namespace mlir::triton::gpu diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Utility.cpp b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Utility.cpp new file mode 100644 index 0000000000..ca1b93eab3 --- /dev/null +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Utility.cpp @@ -0,0 +1,1777 @@ +#include "triton/Analysis/Utility.h" + +#include + +#include "mlir/Analysis/SliceAnalysis.h" +#include "mlir/Dialect/SCF/IR/SCF.h" +#include "mlir/IR/Dominance.h" +#include "mlir/IR/IRMapping.h" +#include "triton/Analysis/AxisInfo.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h" +#include "triton/Dialect/TritonGPU/Transforms/Utility.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +#ifdef __TLE__ +#include "tle/dialect/include/IR/Dialect.h" +#endif +#include "llvm/Support/Debug.h" + +#define DEBUG_TYPE "ttg-utility" +#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ") +#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n") + +namespace tt = mlir::triton; +namespace ttg = mlir::triton::gpu; +namespace ttng = mlir::triton::nvidia_gpu; +namespace mlir { + +using namespace triton; + +SmallVector mmaVersionToInstrShape(int version, + const ArrayRef &shape, + Type eltType, int numWarps) { + if (version == 1) + return {16, 16}; + else if (version == 2) { + auto rank = shape.size(); + SmallVector ret(rank, 1); + ret[rank - 1] = 8; + ret[rank - 2] = 16; + return ret; + } else if (version == 3) { + unsigned k = 256 / eltType.getIntOrFloatBitWidth(); + if (shape[0] % 64 != 0 || shape[1] % 8 != 0) { + assert(false && "type not supported"); + return {0, 0, 0}; + } + SmallVector validN; + + // MMAv3 with larger instruction shape is preferred. + if (llvm::isa( + eltType) || + eltType.isF16() || eltType.isBF16() || eltType.isF32()) { + validN.assign({256, 248, 240, 232, 224, 216, 208, 200, 192, 184, 176, + 168, 160, 152, 144, 136, 128, 120, 112, 104, 96, 88, + 80, 72, 64, 56, 48, 40, 32, 24, 16, 8}); + } + + if (eltType.isInteger(8)) { + validN.assign({224, 208, 192, 176, 160, 144, 128, 112, 96, 80, 64, 48, 32, + 24, 16, 8}); + } + + unsigned m = 16; + unsigned mWarps = std::max(shape[0] / m, 1); + unsigned nWarps = std::max(numWarps / mWarps, 1); + unsigned maxN = std::max(shape[1] / nWarps, 8); + for (auto n : validN) { + if (shape[1] % n == 0 && n <= maxN) { + return {m, n, k}; + } + } + + assert(false && "type not supported"); + return {0, 0, 0}; + } else if (version == 5) { + unsigned m = shape[0] >= 128 ? 128 : 64; + // Right now default to distributing along N. TODO: For cases where we have + // dot followed by reduction we need to be able to distribute along M. + // if (numWarps > 4) + // m = 64; + unsigned n = shape[1] >= 256 ? 256 : shape[1]; + unsigned k = 256 / eltType.getIntOrFloatBitWidth(); + return {m, n, k}; + } else { + assert(false && "version not supported"); + return {0, 0}; + } +} + +bool isLoadFromTensorPtr(triton::LoadOp op) { + return mlir::triton::isTensorPointerType(op.getPtr().getType()); +} + +SmallVector argSort(const SmallVector &arr) { + SmallVector ret(arr.size()); + std::iota(ret.begin(), ret.end(), 0); + std::stable_sort(ret.begin(), ret.end(), + [&](unsigned x, unsigned y) { return arr[x] > arr[y]; }); + return ret; +} + +Value getMemAccessPtr(Operation *op) { + if (auto ld = dyn_cast(op)) + return ld.getPtr(); + if (auto atomic = dyn_cast(op)) + return atomic.getPtr(); + if (auto atomic = dyn_cast(op)) + return atomic.getPtr(); + if (auto copy = dyn_cast(op)) + return copy.getSrc(); + if (auto store = dyn_cast(op)) + return store.getPtr(); + return nullptr; +} + +unsigned getElementBitWidth(RankedTensorType type) { + auto typeForMem = + isa(type.getElementType()) + ? cast(type.getElementType()).getPointeeType() + : type.getElementType(); + return typeForMem.getIntOrFloatBitWidth(); +} + +unsigned getNumElementsPerThread(Operation *op, SmallVector order, + ModuleAxisInfoAnalysis &axisInfoAnalysis, + SmallVector &shapePerCTA) { + Value val = getMemAccessPtr(op); + auto ty = cast(val.getType()); + AxisInfo &valInfo = *axisInfoAnalysis.getAxisInfo(val); + unsigned elemNumBits = getElementBitWidth(ty); + unsigned elemNumBytes = std::max(elemNumBits / 8, 1u); + unsigned maxMultipleBytes = valInfo.getDivisibility(order[0]); + unsigned maxMultiple = std::max(maxMultipleBytes / elemNumBytes, 1u); + unsigned maxContig = + std::min(valInfo.getContiguity(order[0]), shapePerCTA[order[0]]); + unsigned alignment = std::min(maxMultiple, maxContig); + // sunrise specific: 128->32 + // unsigned currPerThread = std::min(alignment, 128 / elemNumBits); + unsigned currPerThread = std::min(alignment, std::max(32 / elemNumBits, 1u)); + LDBG("elemNumBytes: " << elemNumBytes + << ", divisibility: " << maxMultipleBytes + << ", contig: " << valInfo.getContiguity(order[0]) + << ", alignment: " << alignment); + return currPerThread; +} + +bool isView(Operation *op) { + return isa(op); +} + +bool isNoop(Operation *op) { + if (isa(op)) + return true; + if (auto cvt = dyn_cast(op)) { + // The conversion op is a noop if the conversion layout is trivial + return minimalCvtLayout(cvt.getSrc().getType(), + cvt.getResult().getType()) == LinearLayout::empty(); + } + return false; +} + +//===----------------------------------------------------------------------===// +// GraphDumper +//===----------------------------------------------------------------------===// + +GraphDumper::NodeInfo GraphDumper::onValue(Value value) const { + return {{"shape", "box"}, {"style", "filled"}, {"fillcolor", "white"}}; +} + +GraphDumper::NodeInfo GraphDumper::onOperation(Operation *op) const { + return {{"shape", "ellipse"}, {"style", "filled"}, {"fillcolor", "white"}}; +} + +std::string GraphDumper::dump(triton::FuncOp func) const { + llvm::SetVector values; + llvm::SetVector operations; + + func.walk([&](Operation *op) { + operations.insert(op); + for (Value operand : op->getOperands()) + values.insert(operand); + for (Value result : op->getResults()) + values.insert(result); + }); + + std::ostringstream oss; + oss << "// Generated by Triton GraphDumper\n" + << "\n" + << "digraph {\n"; + + oss << " // Value Nodes\n"; + for (Value value : values) + oss << " " << emitValueNode(value) << "\n"; + oss << "\n"; + + oss << " // Operation Nodes\n"; + for (Operation *op : operations) + oss << " " << emitOperationNode(op) << "\n"; + oss << "\n"; + + oss << " // Edges\n"; + for (Operation *op : operations) { + for (Value operand : op->getOperands()) + oss << " " << emitEdge(getUniqueId(operand), getUniqueId(op)) << "\n"; + for (Value result : op->getResults()) + oss << " " << emitEdge(getUniqueId(op), getUniqueId(result)) << "\n"; + } + + oss << "}\n"; + return oss.str(); +} + +void GraphDumper::dumpToFile(triton::FuncOp func, + const std::string &filename) const { + std::ofstream ofs(filename); + ofs << dump(func); +} + +std::string GraphDumper::getShapeStr(const Type &type) const { + std::ostringstream oss; + oss << "["; + if (auto tensorTy = dyn_cast(type)) { + auto shape = tensorTy.getShape(); + for (unsigned i = 0; i < shape.size(); ++i) { + if (i > 0) + oss << ", "; + oss << shape[i]; + } + } + oss << "]"; + return oss.str(); +} + +std::string GraphDumper::getUniqueId(Value value) const { + std::ostringstream oss; + oss << value.getImpl(); + return oss.str(); +} + +std::string GraphDumper::getUniqueId(Operation *op) const { + std::ostringstream oss; + oss << op; + return oss.str(); +} + +std::string GraphDumper::emitNode(const std::string &id, + const GraphDumper::NodeInfo info) const { + std::ostringstream oss; + oss << "\"" << id << "\" ["; + for (auto it = info.begin(); it != info.end(); ++it) { + if (it != info.begin()) + oss << ", "; + oss << it->first << " = \"" << it->second << "\""; + } + oss << "];"; + return oss.str(); +} + +std::string GraphDumper::emitEdge(const std::string &srcId, + const std::string &destId) const { + std::ostringstream oss; + oss << "\"" << srcId << "\" -> \"" << destId << "\";"; + return oss.str(); +} + +std::string GraphDumper::emitValueNode(Value value) const { + NodeInfo info = onValue(value); + if (info.find("label") == info.end()) { + std::string shapeStr = getShapeStr(value.getType()); + if (auto arg = mlir::dyn_cast(value)) + info["label"] = + "BlockArg" + std::to_string(arg.getArgNumber()) + " " + shapeStr; + else + info["label"] = shapeStr; + } + return emitNode(getUniqueId(value), info); +} + +std::string GraphDumper::emitOperationNode(Operation *op) const { + NodeInfo info = onOperation(op); + if (info.find("label") == info.end()) + info["label"] = op->getName().getStringRef().str(); + return emitNode(getUniqueId(op), info); +} + +//===----------------------------------------------------------------------===// +// GraphLayoutMarker +//===----------------------------------------------------------------------===// + +GraphDumper::NodeInfo GraphLayoutMarker::onValue(Value value) const { + std::string color = getColor(value.getType()); + return {{"shape", "box"}, {"style", "filled"}, {"fillcolor", color}}; +} + +std::string GraphLayoutMarker::getColor(const Type &type) const { + if (auto tensorTy = dyn_cast(type)) { + auto layout = tensorTy.getEncoding(); + if (isa(layout)) + return "green"; + else if (isa(layout)) + return "yellow"; + else if (isa(layout)) + return "lightslateblue"; + else if (isa(layout)) + return "orange"; + else if (isa(layout)) + return "orangered"; + else { + llvm::report_fatal_error("Unrecognized layout"); + return "unknown"; + } + } else { + return "white"; + } +} +// -------------------------------------------------------------------------- // + +static Attribute inferDstEncoding(triton::ReduceOp op, Attribute encoding) { + return triton::gpu::SliceEncodingAttr::get( + op->getContext(), op.getAxis(), + cast(encoding)); +} + +static Attribute inferDstEncoding(triton::ExpandDimsOp op, Attribute encoding) { + auto sliceEncoding = mlir::dyn_cast(encoding); + if (!sliceEncoding) + return {}; + if (op.getAxis() != sliceEncoding.getDim()) + return {}; + return sliceEncoding.getParent(); +} + +static Attribute inferDstEncoding(JoinOp op, Attribute srcEnc) { + Attribute dstEnc; + auto shape = op.getLhs().getType().getShape(); + if (srcEnc.getDialect() + .getRegisteredInterface() + ->inferDefaultJoinOpEncoding(srcEnc, dstEnc, shape, + /*loc=*/std::nullopt) + .succeeded()) { + return dstEnc; + } + return {}; +} + +static Attribute inferDstEncoding(SplitOp op, Attribute srcEnc) { + Attribute dstEnc; + auto shape = op.getSrc().getType().getShape(); + if (srcEnc.getDialect() + .getRegisteredInterface() + ->inferSplitOpEncoding(srcEnc, dstEnc, shape, + /*loc=*/std::nullopt) + .succeeded()) { + return dstEnc; + } + return {}; +} + +static Attribute inferSrcEncoding(triton::ReduceOp op, Attribute encoding) { + auto sliceEncoding = mlir::dyn_cast(encoding); + if (!sliceEncoding) + return {}; + if (op.getAxis() != sliceEncoding.getDim()) + return {}; + return sliceEncoding.getParent(); +} + +static Attribute inferSrcEncoding(triton::ExpandDimsOp op, Attribute encoding) { + return triton::gpu::SliceEncodingAttr::get( + op->getContext(), op.getAxis(), + cast(encoding)); +} + +static Attribute inferSrcEncoding(JoinOp op, Attribute dstEnc) { + // Split is the inverse of join. + auto shape = op.getResult().getType().getShape(); + Attribute srcEnc; + if (dstEnc.getDialect() + .getRegisteredInterface() + ->inferSplitOpEncoding(dstEnc, srcEnc, shape, /*loc=*/std::nullopt) + .succeeded()) { + return srcEnc; + } + return {}; +} + +static Attribute inferSrcEncoding(SplitOp op, Attribute dstEnc) { + // Join is the inverse of split. + Attribute srcEnc; + auto shape = op.getOutLHS().getType().getShape(); + if (dstEnc.getDialect() + .getRegisteredInterface() + ->inferDefaultJoinOpEncoding(dstEnc, srcEnc, shape, + /*loc=*/std::nullopt) + .succeeded()) { + return srcEnc; + } + return {}; +} + +static Attribute inferSrcEncoding(GatherOp op, Attribute dstEnc) { + // The index encoding is the same as the output encoding. + return dstEnc; +} + +static Attribute inferTransOpDstEncoding(Attribute srcEnc, + ArrayRef shape, + ArrayRef order) { + // Simply forward to the existing inferTransOpEncoding function. + Attribute retEncoding; + if (succeeded( + srcEnc.getDialect() + .getRegisteredInterface() + ->inferTransOpEncoding(srcEnc, shape, order, retEncoding, + /*loc=*/{}))) { + return retEncoding; + } + return {}; +} + +static Attribute inferDstEncoding(triton::gpu::Fp4ToFpOp op, Attribute srcEnc) { + Attribute dstEnc; + auto shape = op.getSrc().getType().getShape(); + auto result = + srcEnc.getDialect() + .getRegisteredInterface() + ->inferFp4ToFpOpEncoding(shape, op.getAxis(), srcEnc, dstEnc, + /*fwdInference*/ true, std::nullopt); + assert(succeeded(result)); + return dstEnc; +} + +static Attribute inferSrcEncoding(triton::gpu::Fp4ToFpOp op, Attribute dstEnc) { + Attribute srcEnc; + auto shape = op.getType().getShape(); + if (succeeded( + dstEnc.getDialect() + .getRegisteredInterface() + ->inferFp4ToFpOpEncoding(shape, op.getAxis(), dstEnc, srcEnc, + /*fwdInference*/ false, std::nullopt))) { + return srcEnc; + } + return {}; +} + +static Attribute inferDstEncoding(triton::TransposeOpInterface op, + Attribute encoding) { + return inferTransOpDstEncoding( + encoding, cast(op.getSrc().getType()).getShape(), + op.getOrder()); +} + +static Attribute inferSrcEncoding(triton::TransposeOpInterface op, + Attribute encoding) { + // We want to solve for srcEnc in + // transpose(srcEnc, order) -> dstEnc. + // Given the identity + // transpose(transpose(x, order), inverse(order)) == x, + // we can see this is equivalent to + // transpose(dstEnc, inverse(order)) -> srcEnc. + auto shape = cast(op->getResult(0).getType()).getShape(); + return inferTransOpDstEncoding(encoding, shape, + triton::inversePermutation(op.getOrder())); +} + +static Attribute inferReshapeOpDstEncoding(ArrayRef srcShape, + Attribute srcEnc, + ArrayRef dstShape, + bool allowReorder) { + // We don't do anything smart to allow-reorder reshapes here. They are + // handled in OptimizeThreadLocality. + if (allowReorder) + return {}; + + Attribute dstEnc; + auto result = + srcEnc.getDialect() + .getRegisteredInterface() + ->inferReshapeOpEncoding(srcShape, srcEnc, dstShape, dstEnc, + /*loc=*/std::nullopt); + assert(succeeded(result)); + return dstEnc; +} + +static Attribute inferDstEncoding(triton::ReshapeOp op, Attribute encoding) { + return inferReshapeOpDstEncoding(op.getSrc().getType().getShape(), encoding, + op.getType().getShape(), + op.getAllowReorder()); +} + +static Attribute inferDstEncoding(GatherOp op, Attribute encoding) { + // The output encoding is the same as the index encoding. + // FIXME: This assumes `encoding` is the index encoding, which can be + // different than the source encoding. + return encoding; +} + +static Attribute inferSrcEncoding(triton::ReshapeOp op, Attribute encoding) { + // The encoding of x given the encoding of y in `reshape(x) -> y` is the same + // as the encoding of x given the encoding of y in `reshape(y) -> x`. It's an + // invariant of inferReshapeOpNoReorderEncoding that it's symmetric in this + // way. + return inferReshapeOpDstEncoding(op.getType().getShape(), encoding, + op.getSrc().getType().getShape(), + op.getAllowReorder()); +} + +static bool isSingleValue(Value value) { + // Don't consider load as expensive if it is loading a scalar. + if (auto tensorTy = dyn_cast(value.getType())) + return tensorTy.getNumElements() == 1; + // TODO: Handle other cases. + // For example, when ptr is a tensor of single value. + // It means that ptr is a resultant of broadcast or generated through + // a chain of broadcast and other operations. + // Rematerialize it without considering contiguous memory access pattern is + // fine. + return true; +} + +Attribute inferSrcEncoding(Operation *op, Attribute encoding) { + if (isa(op)) { + // Scan only supports blocked encoding at the moment. + if (!isa(encoding)) + return {}; + } + + if (isa(op)) + return {}; + + if (op->hasTrait() || + op->hasTrait() || + op->hasTrait() || + isa(op)) { + return encoding; + } + + if (auto reduceOp = dyn_cast(op)) + return inferSrcEncoding(reduceOp, encoding); + if (auto expand = dyn_cast(op)) + return inferSrcEncoding(expand, encoding); + if (auto join = dyn_cast(op)) + return inferSrcEncoding(join, encoding); + if (auto split = dyn_cast(op)) + return inferSrcEncoding(split, encoding); + if (auto trans = dyn_cast(op)) + return inferSrcEncoding(trans, encoding); + if (auto reshape = dyn_cast(op)) + return inferSrcEncoding(reshape, encoding); + if (auto gather = dyn_cast(op)) + return inferSrcEncoding(gather, encoding); + if (auto fp4ToFp = dyn_cast(op)) + return inferSrcEncoding(fp4ToFp, encoding); + + return {}; +} + +Attribute inferDstEncoding(Operation *op, Attribute encoding) { + if (isa(op)) { + if (!isa(encoding)) + return {}; + } + if (isa(op)) + return {}; + + if (op->hasTrait() || + op->hasTrait() || + op->hasTrait() || + isa(op)) + return encoding; + if (auto reduceOp = dyn_cast(op)) + return inferDstEncoding(reduceOp, encoding); + if (auto expand = dyn_cast(op)) + return inferDstEncoding(expand, encoding); + if (auto join = dyn_cast(op)) + return inferDstEncoding(join, encoding); + if (auto split = dyn_cast(op)) + return inferDstEncoding(split, encoding); + if (auto trans = dyn_cast(op)) + return inferDstEncoding(trans, encoding); + if (auto reshape = dyn_cast(op)) + return inferDstEncoding(reshape, encoding); + if (auto gather = dyn_cast(op)) + return inferDstEncoding(gather, encoding); + if (auto fp4ToFp = dyn_cast(op)) + return inferDstEncoding(fp4ToFp, encoding); + + return {}; +} + +bool isExpensiveLoadOrStore(Operation *op) { + // Case 1: Pointer of tensor is always expensive + auto operandType = op->getOperand(0).getType(); + if (triton::isTensorPointerType(operandType)) + return true; + // Case 2a: A size 1 tensor is not expensive since all threads will load the + // same + if (isSingleValue(op->getOperand(0))) + return false; + // Case 2b: Tensor of pointers has more threads than elements + // we can presume a high hit-rate that makes it cheap to load + auto ptrType = cast(op->getOperand(0).getType()); + auto mod = op->getParentOfType(); + int numWarps = triton::gpu::lookupNumWarps(op); + int threadsPerWarp = triton::gpu::TritonGPUDialect::getThreadsPerWarp(mod); + if (ptrType.getNumElements() < numWarps * threadsPerWarp) + return false; + return true; +} + +bool isExpensiveToRemat(Operation *op, Attribute &targetEncoding) { + if (!op) + return true; + if (isa(op)) + return isExpensiveLoadOrStore(op); + if (isa(op)) + return triton::gpu::isExpensiveCat(cast(op), targetEncoding); + if (isa(op)) + return true; + if (isa( + op)) + return true; + return false; +} + +bool canFoldIntoConversion(Operation *op, Attribute targetEncoding) { + if (isa(op)) + return !triton::gpu::isExpensiveCat(cast(op), + targetEncoding); + if (auto convert = dyn_cast(op)) { + if (mlir::isa(targetEncoding)) { + auto srcEncoding = convert.getSrc().getType().getEncoding(); + if (targetEncoding != srcEncoding) + return false; + } + return true; + } + + if (auto reshape = dyn_cast(op)) { + auto reshapeDstType = reshape.getType(); + RankedTensorType newDstType = + reshapeDstType.cloneWithEncoding(targetEncoding); + return reshape.getAllowReorder() && !reshape.getEfficientLayout() && + !triton::gpu::isExpensiveView(reshape.getSrc().getType(), + newDstType); + } + return isa(op); +} + +scf::ForOp replaceForOpWithNewSignature( + OpBuilder &rewriter, scf::ForOp loop, ValueRange newIterOperands, + SmallVectorImpl> &replacements) { + OpBuilder::InsertionGuard g(rewriter); + rewriter.setInsertionPoint(loop); + + // Create a new loop before the existing one, with the extra operands. + auto operands = llvm::to_vector<4>(loop.getInitArgs()); + operands.append(newIterOperands.begin(), newIterOperands.end()); + scf::ForOp newLoop = + scf::ForOp::create(rewriter, loop.getLoc(), loop.getLowerBound(), + loop.getUpperBound(), loop.getStep(), operands); + newLoop->setAttrs(loop->getAttrs()); + newLoop.getBody()->erase(); + newLoop.getRegion().getBlocks().splice( + newLoop.getRegion().getBlocks().begin(), loop.getRegion().getBlocks()); + for (Value operand : newIterOperands) + newLoop.getBody()->addArgument(operand.getType(), operand.getLoc()); + + for (auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front( + loop.getNumResults()))) + replacements.push_back(it); + return newLoop; +} + +scf::ForOp replaceForOpWithNewSignature(OpBuilder &rewriter, scf::ForOp loop, + ValueRange newIterOperands) { + SmallVector> replacements; + auto newForOp = replaceForOpWithNewSignature(rewriter, loop, newIterOperands, + replacements); + for (auto [result, value] : replacements) { + result.replaceAllUsesWith(value); + } + return newForOp; +} + +scf::ForOp addIterArgsToLoop(OpBuilder &rewriter, scf::ForOp loop, + ValueRange newIterOperands) { + scf::ForOp newLoop = + replaceForOpWithNewSignature(rewriter, loop, newIterOperands); + // Save the caller from insertion point invalidation. + if (rewriter.getInsertionPoint() == loop->getIterator()) + rewriter.setInsertionPoint(newLoop); + loop.erase(); + return newLoop; +} + +scf::WhileOp replaceWhileOpWithNewSignature( + OpBuilder &rewriter, scf::WhileOp loop, ValueRange newIterOperands, + TypeRange newResultTypes, + SmallVectorImpl> &replacements) { + OpBuilder::InsertionGuard g(rewriter); + rewriter.setInsertionPoint(loop); + + // Create a new loop before the existing one, with the extra operands. + auto operands = llvm::to_vector<4>(loop.getInits()); + operands.append(newIterOperands.begin(), newIterOperands.end()); + + // Result and operand types + SmallVector resultTypes; + SmallVector argsTypesBefore; + for (auto res : loop.getResults()) + resultTypes.push_back(res.getType()); + for (auto type : newResultTypes) + resultTypes.push_back(type); + for (Value operand : operands) + argsTypesBefore.push_back(operand.getType()); + scf::WhileOp newLoop = + scf::WhileOp::create(rewriter, loop.getLoc(), resultTypes, operands); + newLoop->setAttrs(loop->getAttrs()); + + SmallVector bbArgLocsBefore(argsTypesBefore.size(), loop.getLoc()); + SmallVector bbArgLocsAfter(resultTypes.size(), loop.getLoc()); + rewriter.createBlock(&newLoop.getBefore(), {}, argsTypesBefore, + bbArgLocsBefore); + rewriter.createBlock(&newLoop.getAfter(), {}, resultTypes, bbArgLocsAfter); + + // Copy regions + for (int i = 0; i < loop.getNumRegions(); ++i) + newLoop->getRegion(i).front().getOperations().splice( + newLoop->getRegion(i).front().getOperations().begin(), + loop->getRegion(i).front().getOperations()); + + // Remap arguments + for (auto [oldArg, newArg] : llvm::zip( + loop.getBeforeArguments(), newLoop.getBeforeArguments().take_front( + loop.getBeforeArguments().size()))) + oldArg.replaceAllUsesWith(newArg); + for (auto [oldArg, newArg] : llvm::zip(loop.getAfterArguments(), + newLoop.getAfterArguments().take_front( + loop.getAfterArguments().size()))) + oldArg.replaceAllUsesWith(newArg); + + // Stack the new results + for (auto it : llvm::zip(loop.getResults(), newLoop.getResults().take_front( + loop.getNumResults()))) + replacements.push_back(it); + + return newLoop; +} + +scf::WhileOp replaceWhileOpWithNewSignature(OpBuilder &rewriter, + scf::WhileOp loop, + ValueRange newIterOperands, + TypeRange newResultTypes) { + SmallVector> replacements; + auto newWhileOp = replaceWhileOpWithNewSignature( + rewriter, loop, newIterOperands, newResultTypes, replacements); + for (auto &kv : replacements) { + std::get<0>(kv).replaceAllUsesWith(std::get<1>(kv)); + } + return newWhileOp; +} + +scf::IfOp replaceIfOpWithNewSignature( + OpBuilder &rewriter, scf::IfOp ifOp, TypeRange newResultTypes, + SmallVectorImpl> &replacements) { + OpBuilder::InsertionGuard g(rewriter); + rewriter.setInsertionPoint(ifOp); + + // Create a new loop before the existing one, with the extra operands. + auto resultTypes = llvm::to_vector<4>(ifOp.getResults().getTypes()); + resultTypes.append(newResultTypes.begin(), newResultTypes.end()); + scf::IfOp newIf = scf::IfOp::create(rewriter, ifOp.getLoc(), resultTypes, + ifOp.getCondition()); + newIf->setAttrs(ifOp->getAttrs()); + + newIf.getThenRegion().takeBody(ifOp.getThenRegion()); + newIf.getElseRegion().takeBody(ifOp.getElseRegion()); + scf::IfOp::ensureTerminator(newIf.getThenRegion(), rewriter, ifOp.getLoc()); + scf::IfOp::ensureTerminator(newIf.getElseRegion(), rewriter, ifOp.getLoc()); + + for (auto it : llvm::zip(ifOp.getResults(), + newIf.getResults().take_front(ifOp.getNumResults()))) + replacements.push_back(it); + return newIf; +} + +void appendToForOpYield(scf::ForOp forOp, ArrayRef newOperands) { + Operation *yieldOp = forOp.getBody()->getTerminator(); + SmallVector operands(yieldOp->getOperands()); + operands.append(newOperands.begin(), newOperands.end()); + + OpBuilder builder(yieldOp); + scf::YieldOp::create(builder, yieldOp->getLoc(), operands); + yieldOp->erase(); +} + +scf::IfOp replaceIfOpWithNewSignature(OpBuilder &rewriter, scf::IfOp ifOp, + TypeRange newResultTypes) { + SmallVector> replacements; + auto newIfOp = + replaceIfOpWithNewSignature(rewriter, ifOp, newResultTypes, replacements); + for (auto &kv : replacements) + std::get<0>(kv).replaceAllUsesWith(std::get<1>(kv)); + return newIfOp; +} + +Operation *cloneWithInferType(mlir::OpBuilder &rewriter, Operation *op, + IRMapping &mapping) { + Operation *newOp = rewriter.clone(*op, mapping); + // if input types haven't changed, we're done + bool preserveTypes = + std::all_of(op->operand_begin(), op->operand_end(), [&](Value v) { + return !mapping.contains(v) || + v.getType() == mapping.lookup(v).getType(); + }); + if (preserveTypes) + return newOp; + + if (newOp->getNumResults() == 0) + return newOp; + auto origType = dyn_cast(op->getResult(0).getType()); + auto argType = dyn_cast(newOp->getOperand(0).getType()); + if (!origType || !argType) + return newOp; + auto newType = origType.cloneWithEncoding(argType.getEncoding()); + newOp->getResult(0).setType(newType); + auto typeInfer = dyn_cast(newOp); + if (typeInfer) { + SmallVector newTypes; + auto success = typeInfer.inferReturnTypes( + newOp->getContext(), newOp->getLoc(), newOp->getOperands(), + newOp->getAttrDictionary(), newOp->getPropertiesStorage(), + newOp->getRegions(), newTypes); + if (succeeded(success)) { + for (size_t i = 0; i < newTypes.size(); i++) + newOp->getResult(i).setType(newTypes[i]); + } + } + return newOp; +} + +// Check if the convert will be performed by reordering registers. +static bool isFreeConvert(Operation *op) { + auto convertOp = dyn_cast(op); + if (!convertOp) + return false; + return cvtReordersRegisters(convertOp.getSrc().getType(), + convertOp.getType()); +} + +LogicalResult getConvertBackwardSlice( + OpOperand &root, SetVector &slice, Attribute rootEncoding, + DenseMap &layout, + std::function stopPropagation, + std::function getExistingConversion) { + DenseSet> seen; + SmallVector> queue; + + auto enqueue = [&](OpOperand &operand, Attribute encoding) { + auto x = std::make_pair(&operand, encoding); + if (!seen.insert(x).second) { + return; // Already enqueued, skip + } + queue.push_back(x); + }; + enqueue(root, rootEncoding); + + auto updateLayout = [&](Value value, Attribute encoding) { + assert((isa(value.getType()))); + slice.insert(value); + Attribute &existing = layout[value]; + if (existing && existing != encoding) + return failure(); + existing = encoding; + return success(); + }; + + while (!queue.empty()) { + auto [currentValueUse, encoding] = queue.back(); + Value currentValue = currentValueUse->get(); + queue.pop_back(); + if (!isa(currentValue.getType())) + continue; + // Skip propagating through for op/while op results for now. + // TODO: enable this based on needs. + if (currentValue.getDefiningOp() || + currentValue.getDefiningOp()) + return failure(); + if (failed(updateLayout(currentValue, encoding))) + return failure(); + + Value existing; + if (getExistingConversion && + (existing = getExistingConversion(*currentValueUse, encoding))) { + if (failed(updateLayout(existing, encoding))) + return failure(); + currentValue = existing; + } + + if (auto ifOp = currentValue.getDefiningOp()) { + if (stopPropagation && stopPropagation(ifOp)) + continue; + unsigned argIdx = mlir::cast(currentValue).getResultNumber(); + + OpOperand &thenValue = ifOp.thenYield()->getOpOperand(argIdx); + OpOperand &elseValue = ifOp.elseYield()->getOpOperand(argIdx); + + enqueue(thenValue, encoding); + enqueue(elseValue, encoding); + + continue; + } + if (auto *definingOp = currentValue.getDefiningOp()) { + // If the op has multiple results we need to update all results layout. + for (Value result : definingOp->getResults()) { + if (result == currentValue || !isa(result.getType())) + continue; + if (failed(updateLayout(result, encoding))) + return failure(); + } + if (isFreeConvert(definingOp)) { + enqueue(definingOp->getOpOperand(0), encoding); + continue; + } + if (canFoldIntoConversion(definingOp, encoding)) + continue; + if (stopPropagation && stopPropagation(definingOp)) + continue; + if (isa(definingOp)) + return failure(); + if (auto gather = dyn_cast(definingOp)) { + // Specially handle gather since its transfer function only applies + // between its index operand and result. + auto srcEncoding = inferSrcEncoding(gather, encoding); + if (!srcEncoding) + return failure(); + enqueue(gather.getIndicesMutable(), srcEncoding); + continue; + } + for (auto [i, operand] : llvm::enumerate(definingOp->getOpOperands())) { + Attribute srcEncoding; + if (auto upcast = + dyn_cast(definingOp)) { + srcEncoding = upcast.inferSrcEncoding(i, encoding); + } else { + srcEncoding = inferSrcEncoding(definingOp, encoding); + } + if (!srcEncoding) + return failure(); + // If the infered layout matches the original one we don't need to keep + // propagating. + if (auto operandType = + dyn_cast(operand.get().getType())) { + if (srcEncoding == operandType.getEncoding()) + continue; + } + enqueue(operand, srcEncoding); + } + continue; + } + auto blockArg = cast(currentValue); + Block *block = blockArg.getOwner(); + Operation *parentOp = block->getParentOp(); + if (auto forOp = dyn_cast(parentOp)) { + OpOperand *initOperand = forOp.getTiedLoopInit(blockArg); + OpOperand &yieldOperand = forOp.getBody()->getTerminator()->getOpOperand( + blockArg.getArgNumber() - forOp.getNumInductionVars()); + enqueue(*initOperand, encoding); + enqueue(yieldOperand, encoding); + continue; + } + // TODO: add support for WhileOp and other region types. + return failure(); + } + return success(); +} + +// TODO(thomas): this is duplicated with what is in GPUToLLVM +// Convert an \param index to a multi-dim coordinate given \param shape and +// \param order. +SmallVector delinearize(OpBuilder &b, Location loc, Value linear, + ArrayRef shape, + ArrayRef order) { + unsigned rank = shape.size(); + assert(rank == order.size()); + auto reordered = triton::applyPermutation(shape, order); + auto reorderedMultiDim = delinearize(b, loc, linear, reordered); + SmallVector multiDim(rank); + for (unsigned i = 0; i < rank; ++i) { + multiDim[order[i]] = reorderedMultiDim[i]; + } + return multiDim; +} + +SmallVector delinearize(OpBuilder &b, Location loc, Value linear, + ArrayRef shape) { + unsigned rank = shape.size(); + assert(rank > 0); + SmallVector multiDim(rank); + if (rank == 1) { + multiDim[0] = linear; + } else { + Value remained = linear; + for (auto &&en : llvm::enumerate(shape.drop_back())) { + auto dimSize = arith::ConstantIntOp::create(b, loc, en.value(), 32); + multiDim[en.index()] = arith::RemSIOp::create(b, loc, remained, dimSize); + remained = arith::DivSIOp::create(b, loc, remained, dimSize); + } + multiDim[rank - 1] = remained; + } + return multiDim; +} + +Value linearize(OpBuilder &b, Location loc, ArrayRef multiDim, + ArrayRef shape, ArrayRef order) { + return linearize(b, loc, triton::applyPermutation(multiDim, order), + triton::applyPermutation(shape, order)); +} + +Value linearize(OpBuilder &b, Location loc, ArrayRef multiDim, + ArrayRef shape) { + auto rank = multiDim.size(); + Value linear = arith::ConstantIntOp::create(b, loc, 0, 32); + if (rank > 0) { + linear = multiDim.back(); + for (auto [dim, dimShape] : + llvm::reverse(llvm::zip(multiDim.drop_back(), shape.drop_back()))) { + Value dimSize = arith::ConstantIntOp::create(b, loc, dimShape, 32); + linear = arith::AddIOp::create( + b, loc, arith::MulIOp::create(b, loc, linear, dimSize), dim); + } + } + return linear; +} + +bool isPureUnaryInlineAsm(Operation *op) { + auto inlineAsmOp = dyn_cast(op); + if (!inlineAsmOp) + return false; + return op->getNumOperands() == 1 && op->getNumResults() == 1 && + inlineAsmOp.getPure(); +} + +int getNVIDIAComputeCapability(Operation *module) { + StringAttr targetAttr = + module->getAttrOfType(triton::gpu::AttrTargetName); + assert(targetAttr && "Expected a target attribute on the module operation"); + + StringRef ref = targetAttr.strref(); + assert(ref.starts_with("cuda:") && + "expected target attribute to be prefixed with \"cuda:\""); + + StringRef capabilityStr = ref.drop_front(5); // drop the "cuda:" + int computeCapability; + bool parseError = capabilityStr.getAsInteger(10, computeCapability); + assert(!parseError && + "invalid compute capability string in target attribute"); + + return computeCapability; +} + +std::optional getAMDArch(Operation *module) { + StringAttr targetAttr = + module->getAttrOfType(triton::gpu::AttrTargetName); + if (!targetAttr) { + LDBG("Expected a target attribute on the module operation"); + return {}; + } + + StringRef ref = targetAttr.strref(); + if (!ref.starts_with("hip:")) { + LDBG("expected target attribute to be prefixed with \"hip:\""); + return {}; + } + + return ref.drop_front(4); // drop the "hip:" +} + +inline ttg::SwizzledSharedEncodingAttr +swizzleDotOperandLike(RankedTensorType type, ttg::CTAEncodingAttr ctaLayout) { + // We want to see if the linear layout has the same order as an mma microtile + // of shape (8, 4*kWidth) or (4*kWidth, 8). If so, we return a + // DotOperandEncodingAttr with a tile of this shape This works because + // SwizzledSharedEncodingAttr::get just looks at the microtile to determine + // the swizzling + + auto *ctx = type.getContext(); + auto layout = ttg::toLinearEncoding(type); + auto order = layout.getThreadOrder(); + auto rank = order.size(); + if (rank < 2) { + return {}; + } + int opIdx; + if (ttg::getOrderForDotOperand(0, rank, /*kContig=*/true) == order) { + opIdx = 0; + } else if (ttg::getOrderForDotOperand(1, rank, /*kContig=*/true) == order) { + opIdx = 1; + } else { + return {}; + } + auto kWidth = layout.getContigPerThread()[order[0]]; + SmallVector microtileShape(rank, 1); + microtileShape[order[0]] = 4 * kWidth; + microtileShape[order[1]] = 8; + // All the LinearLayouts contained within LinearEncoidngAttr have order [0, 1, + // 2, ...] + auto repOrder = to_vector(llvm::seq(rank)); + auto tile = ttg::nvidiaMmaTile(ctx, microtileShape, kWidth, order, repOrder); + if (!divideLeft(layout.getLinearLayout(), tile).has_value()) { + return {}; + } + return ttg::SwizzledSharedEncodingAttr::get( + ctx, opIdx, kWidth, type.getShape(), order, ctaLayout, + type.getElementTypeBitWidth(), false); +} + +// If all the transitive uses of the given value have are used by a convert to +// the same dot operand encoding, return the shared encoding that needs to be +// used to be compatible with users' layouts. If there are incompatible shared +// encodings, set incompatible to true. +std::optional +getSharedEncIfAllUsersAreDotEnc(Value val, bool &incompatible) { + ttg::SwizzledSharedEncodingAttr attr; + incompatible = false; + for (Operation *user : val.getUsers()) { + ttg::SwizzledSharedEncodingAttr tempAttr; + if (user->getNumResults() != 1) + return std::nullopt; + if (auto memDesc = + dyn_cast(user->getResult(0).getType())) { + // First time we find a shared encoding in the chain, save it and try to + // use it if it is compatible with the other users. + tempAttr = + dyn_cast(memDesc.getEncoding()); + if (!tempAttr) + return std::nullopt; + if (!getSharedEncIfAllUsersAreDotEnc(user->getResult(0), incompatible) + .has_value()) + return std::nullopt; + } else { + if (!isa(user)) + return std::nullopt; + auto srcTy = cast(val.getType()); + auto dstTy = cast(user->getResult(0).getType()); + + // FIXME This may not be correct for multiple CTA, but getCTALayout is NYI + // for LinearEncodingAttr + auto CTALayout = isa(dstTy.getEncoding()) + ? ttg::getCTALayout(srcTy.getEncoding()) + : ttg::getCTALayout(dstTy.getEncoding()); + + if (auto dot = + dyn_cast(dstTy.getEncoding())) { + auto order = getOrderForMemory(srcTy); + unsigned bitWidth = srcTy.getElementTypeBitWidth(); + tempAttr = ttg::SwizzledSharedEncodingAttr::get( + val.getContext(), dot, srcTy.getShape(), order, CTALayout, bitWidth, + /*needTrans=*/false); + } else { + // Try to see if the layout is like an mma microtile + tempAttr = swizzleDotOperandLike(dstTy, CTALayout); + } + if (!tempAttr) + return std::nullopt; + } + // Check that the shared encodings needed by the users are compatible. + if (attr != nullptr && attr != tempAttr) { + incompatible = true; + return std::nullopt; + } + attr = tempAttr; + } + return attr; +} + +static Type getNewType(Type type, Attribute encoding) { + RankedTensorType tensorType = cast(type); + return RankedTensorType::get(tensorType.getShape(), + tensorType.getElementType(), encoding); +} + +static bool skipOperand(Operation *op, unsigned operandNumber) { + if (auto gather = dyn_cast(op)) { + return operandNumber == gather.getXOffsetsMutable().getOperandNumber(); + } + if (auto scatter = dyn_cast(op)) { + return operandNumber == scatter.getXOffsetsMutable().getOperandNumber(); + } + return false; +} + +Operation *convertDistributedOpEncoding(Attribute encoding, Operation *op) { + OpBuilder builder(op); + // Convert operands + // For load/store with tensor pointers, we don't have to change the + // operands' type, we do this by changing the outputs' type of + // `make_tensor_ptr` + SmallVector newArgs; + for (auto &opOperand : op->getOpOperands()) { + Value operand = opOperand.get(); + auto tensorType = dyn_cast(operand.getType()); + bool skip = skipOperand(op, opOperand.getOperandNumber()); + if (tensorType && !skip) { + Type newType = getNewType(tensorType, encoding); + newArgs.push_back(triton::gpu::ConvertLayoutOp::create( + builder, op->getLoc(), newType, operand)); + } else { + newArgs.push_back(operand); + } + } + + // Convert output types + SmallVector newTypes; + for (auto t : op->getResultTypes()) { + bool isAsync = isa(op); + newTypes.push_back(isAsync ? t : getNewType(t, encoding)); + } + + // Construct new op with the new encoding + Operation *newOp = builder.create(op->getLoc(), op->getName().getIdentifier(), + newArgs, newTypes, op->getAttrs()); + + // Cast the results back to the original layout + for (size_t i = 0; i < op->getNumResults(); i++) { + Value newResult = newOp->getResult(i); + if (newTypes[i] != op->getResultTypes()[i]) { + newResult = triton::gpu::ConvertLayoutOp::create( + builder, op->getLoc(), op->getResult(i).getType(), newResult); + } + op->getResult(i).replaceAllUsesWith(newResult); + } + op->erase(); + return newOp; +} + +namespace { + +/// Detect dead arguments in scf.for op by assuming all the values are dead and +/// propagate liveness property. +struct ForOpDeadArgElimination : public OpRewritePattern { + using OpRewritePattern::OpRewritePattern; + + LogicalResult matchAndRewrite(scf::ForOp forOp, + PatternRewriter &rewriter) const final { + Block &block = *forOp.getBody(); + auto yieldOp = cast(block.getTerminator()); + // Assume that nothing is live at the beginning and mark values as live + // based on uses. + DenseSet aliveValues; + SmallVector queue; + // Helper to mark values as live and add them to the queue of value to + // propagate if it is the first time we detect the value as live. + auto markLive = [&](Value val) { + if (!forOp->isAncestor(val.getParentRegion()->getParentOp())) + return; + if (aliveValues.insert(val).second) + queue.push_back(val); + }; + // Mark all yield operands as live if the associated forOp result has any + // use. + for (auto result : llvm::enumerate(forOp.getResults())) { + if (!result.value().use_empty()) + markLive(yieldOp.getOperand(result.index())); + } + if (aliveValues.size() == forOp.getNumResults()) + return failure(); + // Operations with side-effects are always live. Mark all theirs operands as + // live. + block.walk([&](Operation *op) { + if (!isa(op) && !wouldOpBeTriviallyDead(op)) { + for (Value operand : op->getOperands()) + markLive(operand); + } + }); + // Propagate live property until reaching a fixed point. + while (!queue.empty()) { + Value value = queue.pop_back_val(); + if (auto nestedFor = value.getDefiningOp()) { + auto result = mlir::cast(value); + OpOperand &forOperand = *nestedFor.getTiedLoopInit(result); + markLive(forOperand.get()); + auto nestedYieldOp = + cast(nestedFor.getBody()->getTerminator()); + Value nestedYieldOperand = + nestedYieldOp.getOperand(result.getResultNumber()); + markLive(nestedYieldOperand); + continue; + } + if (auto nestedIf = value.getDefiningOp()) { + auto result = mlir::cast(value); + // mark condition as live. + markLive(nestedIf.getCondition()); + for (scf::YieldOp nestedYieldOp : + {nestedIf.thenYield(), nestedIf.elseYield()}) { + Value nestedYieldOperand = + nestedYieldOp.getOperand(result.getResultNumber()); + markLive(nestedYieldOperand); + } + continue; + } + if (Operation *def = value.getDefiningOp()) { + // TODO: support while ops. + if (isa(def)) + return failure(); + for (Value operand : def->getOperands()) + markLive(operand); + continue; + } + // If an argument block is live then the associated yield operand and + // forOp operand are live. + auto arg = mlir::cast(value); + if (auto forOwner = dyn_cast(arg.getOwner()->getParentOp())) { + if (arg.getArgNumber() < forOwner.getNumInductionVars()) + continue; + unsigned iterIdx = arg.getArgNumber() - forOwner.getNumInductionVars(); + Value yieldOperand = + forOwner.getBody()->getTerminator()->getOperand(iterIdx); + markLive(yieldOperand); + markLive(forOwner.getInitArgs()[iterIdx]); + } + } + SmallVector deadArg; + for (auto yieldOperand : llvm::enumerate(yieldOp->getOperands())) { + if (aliveValues.contains(yieldOperand.value())) + continue; + if (yieldOperand.value() == block.getArgument(yieldOperand.index() + 1)) + continue; + + // The yield operand might live outside the loop, e.g. + // %init = ... + // %x = ... + // %y = for iter_args(%unused = %init) { + // yield %x + // } + // + // In this case, the loop returns %x if it runs 1 or more times, and + // otherwise it returns %init. We cowardly refuse to remove this operand + // from the yield. (We could, but we'd need to prove that the loop runs 0 + // or >=1 times.) + // + // As a special case, if it doesn't matter whether the loop runs 0 or >=1 + // times (because the loop returns the same value in both cases) then we + // can still mark the operand as dead. This occurs in the above example + // when %init is the same as %x. + if (!forOp->isAncestor( + yieldOperand.value().getParentRegion()->getParentOp()) && + yieldOperand.value() != forOp.getInitArgs()[yieldOperand.index()]) + continue; + + deadArg.push_back(yieldOperand.index()); + } + if (deadArg.empty()) + return failure(); + rewriter.modifyOpInPlace(forOp, [&]() { + // For simplicity we just change the dead yield operand to use the + // associated argument and leave the operations and argument removal to + // dead code elimination. + for (unsigned deadArgIdx : deadArg) { + BlockArgument arg = block.getArgument(deadArgIdx + 1); + yieldOp.setOperand(deadArgIdx, arg); + } + }); + return success(); + } +}; + +} // namespace + +void populateForOpDeadArgumentElimination(RewritePatternSet &patterns) { + patterns.add(patterns.getContext()); +} + +ttg::LocalAllocOp findShmemAlloc(Value operand) { + // If it's a shmem operand, it must either be defined outside the loop, or + // come from an MemDescIndex op. Only ConvertLayout and MemdescView ops are + // allowed in between. + Value transitiveOperand = operand; + while (isa_and_nonnull( + transitiveOperand.getDefiningOp()) || + isa(transitiveOperand)) { + if (auto blockArg = dyn_cast(transitiveOperand)) { + assert(isa(blockArg.getOwner()->getParentOp()) && + "Block argument must come from a for loop"); + transitiveOperand = + cast(blockArg.getOwner()->getTerminator()) + .getOperand(blockArg.getArgNumber() - 1); + } else { + transitiveOperand = transitiveOperand.getDefiningOp()->getOperand(0); + } + } + if (auto subView = dyn_cast_or_null( + transitiveOperand.getDefiningOp())) { + // Multi-buffered operand + return dyn_cast_or_null( + subView.getSrc().getDefiningOp()); + } else { + // Single bufferred operand that does not require a subview (not loaded in + // the loop) + return dyn_cast_or_null( + transitiveOperand.getDefiningOp()); + } + return nullptr; +} + +SmallVector +getMMAsWithMultiBufferredOperands(scf::ForOp forOp, + SmallVector &mmaOps) { + // The A and B operands of the mmaOp should be multi-buffered + SmallVector eligible; + for (auto mmaOp : mmaOps) { + auto a = findShmemAlloc(mmaOp->getOperand(0)); + auto b = findShmemAlloc(mmaOp->getOperand(1)); + if (a && forOp.isDefinedOutsideOfLoop(a) && b && + forOp.isDefinedOutsideOfLoop(b)) { + eligible.push_back(mmaOp); + } + } + + return eligible; +} + +template +static Operation *findNearestCommonDominatorImpl( + ArrayRef ops, DomInfoT &domInfo, + function_ref isBefore) { + if (ops.size() == 0) { + return nullptr; + } + if (ops.size() == 1) { + return ops[0]; + } + llvm::SmallPtrSet blocks; + for (auto op : ops) { + blocks.insert(op->getBlock()); + } + Block *domBlock = domInfo.findNearestCommonDominator(blocks); + if (domBlock == nullptr) { + return nullptr; + } + SmallVector ancestorOps; + for (auto op : ops) { + ancestorOps.push_back(domBlock->findAncestorOpInBlock(*op)); + } + Operation *dom = ancestorOps[0]; + for (unsigned i = 1; i < ops.size(); i++) { + if (isBefore(ancestorOps[i], dom)) { + dom = ancestorOps[i]; + } + } + return dom; +} + +Operation *findNearestCommonDominator(ArrayRef ops, + DominanceInfo &domInfo) { + return findNearestCommonDominatorImpl( + ops, domInfo, + [](Operation *a, Operation *b) { return a->isBeforeInBlock(b); }); +} + +Operation *findNearestCommonPostDominator(ArrayRef ops, + PostDominanceInfo &domInfo) { + return findNearestCommonDominatorImpl( + ops, domInfo, + [](Operation *a, Operation *b) { return b->isBeforeInBlock(a); }); +} + +void visitNestedOperands(Operation *op, + function_ref visitor) { + op->walk([&](Operation *nestedOp) { + for (OpOperand &operand : nestedOp->getOpOperands()) { + if (operand.get().getParentBlock()->getParentOp()->isProperAncestor(op)) + visitor(operand); + } + }); +} + +void visitNestedOperands(Operation *op, function_ref visitor) { + visitNestedOperands(op, [&](OpOperand &operand) { visitor(operand.get()); }); +} + +SetVector getNestedOperands(Operation *op) { + SetVector result; + visitNestedOperands(op, [&](Value operand) { result.insert(operand); }); + return result; +} + +void eraseLoopCarriedValues(scf::ForOp &loop, llvm::BitVector indices) { + // Pad the indices in case new arguments were added. + while (indices.size() != loop.getInitArgs().size()) + indices.push_back(false); + + loop.getBody()->getTerminator()->eraseOperands(indices); + loop.getBody()->eraseArguments([&](BlockArgument arg) { + int idx = arg.getArgNumber(); + return idx != 0 && indices.test(idx - 1); + }); + + llvm::BitVector loopOperandIndices(loop->getNumOperands()); + for (auto [i, operand] : llvm::enumerate(loop.getInitArgsMutable())) { + if (indices.test(i)) + loopOperandIndices.set(operand.getOperandNumber()); + } + loop->eraseOperands(loopOperandIndices); + + // Rewrite the loop to erase results. + OperationState state(loop.getLoc(), loop->getName(), loop->getOperands(), + loop.getInitArgs().getTypes(), loop->getAttrs()); + state.addRegion()->takeBody(loop.getBodyRegion()); + + OpBuilder b(loop); + auto newLoop = cast(b.create(state)); + + // Replace uses of the old loop with the new loop. + unsigned newResultIdx = 0; + for (auto [i, result] : llvm::enumerate(loop.getResults())) { + if (indices.test(i)) { + assert(result.use_empty() && "loop carried value still has uses"); + continue; + } + result.replaceAllUsesWith(newLoop.getResult(newResultIdx++)); + } + + loop.erase(); + loop = newLoop; +} + +} // namespace mlir + +namespace mlir::triton { + +void replaceUsesAndPropagateType( + OpBuilder &builder, Operation *oldUse, Value val, + std::function callback) { + OpBuilder::InsertionGuard guard(builder); + SmallVector opsToDelete; + SmallVector operandsToReplace; + + // Save the operand to replace / delete later (avoid iterator invalidation). + // TODO: can we use an early_inc iterator? + for (OpOperand &use : oldUse->getUses()) { + // Propagate through `ttg.warp_specialize`. + if (auto wsOp = dyn_cast(use.getOwner())) { + for (Region *region : wsOp.getPartitionRegions()) + region->getArgument(use.getOperandNumber()).setType(val.getType()); + } + + // Non-subview/trans ops will be replaced by `val`. + if (!use.getOwner()->hasTrait()) { + operandsToReplace.push_back(&use); + continue; + } + + Operation *user = use.getOwner(); + // `subview(old_op)` is replaced by a new `subview(val)`. + builder.setInsertionPoint(user); + Value newVal; + if (auto subview = dyn_cast(user)) { + ttg::MemDescType oldType = subview.getType(); + bool isMutable = cast(val.getType()).getMutableMemory(); + Type newDstType = ttg::MemDescType::get( + oldType.getShape(), oldType.getElementType(), oldType.getEncoding(), + oldType.getMemorySpace(), isMutable); + newVal = ttg::MemDescIndexOp::create(builder, subview.getLoc(), + newDstType, val, subview.getIndex()); + } else if (auto subslice = dyn_cast(user)) { + ttg::MemDescType oldType = subslice.getType(); + bool isMutable = cast(val.getType()).getMutableMemory(); + Type newDstType = ttg::MemDescType::get( + oldType.getShape(), oldType.getElementType(), oldType.getEncoding(), + oldType.getMemorySpace(), isMutable, oldType.getAllocShape()); + newVal = ttg::MemDescSubsliceOp::create( + builder, subslice.getLoc(), newDstType, val, subslice.getOffsets()); + } else if (auto trans = dyn_cast(user)) { + newVal = ttg::MemDescTransOp::create(builder, trans.getLoc(), val, + trans.getOrder()); +#ifdef __TLE__ + } else if (auto view = dyn_cast(user)) { + ttg::MemDescType oldType = view.getType(); + bool isMutable = cast(val.getType()).getMutableMemory(); + Type newDstType = ttg::MemDescType::get( + oldType.getShape(), oldType.getElementType(), oldType.getEncoding(), + oldType.getMemorySpace(), isMutable, oldType.getAllocShape()); + builder.getContext()->getOrLoadDialect(); + newVal = triton::tle::MemDescWGMMAViewOp::create( + builder, view.getLoc(), newDstType, val, view.getOrder()); +#endif + } else if (auto reshape = dyn_cast(user)) { + auto shape = reshape.getType().getShape(); + newVal = + ttg::MemDescReshapeOp::create(builder, reshape.getLoc(), val, shape); + } + assert(newVal && "unhandled memdesc view"); + newVal.getDefiningOp()->setAttrs(user->getAttrs()); + replaceUsesAndPropagateType(builder, user, newVal); + opsToDelete.push_back(user); + if (callback) { + callback(user, newVal.getDefiningOp()); + } + } + + // Perform late replacement. +#ifdef __TLE__ + SmallVector nonWaitOperandsToReplace; + SmallVector>> + waitOperandsToReplace; + for (OpOperand *operand : operandsToReplace) { + if (auto wait = dyn_cast(operand->getOwner())) { + auto existing = llvm::find_if(waitOperandsToReplace, [&](auto &entry) { + return entry.first == wait; + }); + if (existing == waitOperandsToReplace.end()) { + waitOperandsToReplace.push_back({wait, {}}); + waitOperandsToReplace.back().second.push_back( + operand->getOperandNumber()); + continue; + } + existing->second.push_back(operand->getOperandNumber()); + } else { + nonWaitOperandsToReplace.push_back(operand); + } + } + for (OpOperand *operand : nonWaitOperandsToReplace) + operand->set(val); + for (auto &[wait, operandNumbers] : waitOperandsToReplace) { + // Need to update the return type on the wait op as well. + builder.setInsertionPointAfter(wait); + auto operands = llvm::to_vector(wait.getOperands()); + for (unsigned operandNumber : operandNumbers) + operands[operandNumber] = val; + auto newWait = ttng::WarpGroupDotWaitOp::create( + builder, wait.getLoc(), operands, wait.getPendings()); + wait.replaceAllUsesWith(newWait.getResults()); + wait.erase(); + } +#else + for (OpOperand *operand : operandsToReplace) { + if (auto wait = dyn_cast(operand->getOwner())) { + // Need to update the return type on the wait op as well + builder.setInsertionPointAfter(wait); + auto operands = llvm::to_vector(wait.getOperands()); + operands[operand->getOperandNumber()] = val; + auto newWait = ttng::WarpGroupDotWaitOp::create( + builder, wait.getLoc(), operands, wait.getPendings()); + wait.replaceAllUsesWith(newWait.getResults()); + wait.erase(); + } else { + operand->set(val); + } + } +#endif + + // Perform late op erasure. + for (Operation *op : opsToDelete) + op->erase(); +} + +ttg::LocalLoadOp +replaceUsesWithLocalLoad(OpBuilder &builder, OpResult old, + TypedValue alloc, + TypedValue token) { + // Remove redundant local_load -> local_alloc + auto allocTy = alloc.getType(); + SmallVector allocsToErase; + for (Operation *user : old.getUsers()) { + if (auto userAlloc = dyn_cast(user)) { + if (allocTy.getEncoding() == userAlloc.getType().getEncoding()) { + replaceUsesAndPropagateType(builder, userAlloc, alloc); + allocsToErase.push_back(userAlloc); + } + } + } + + // If there are some uses that were not local_allocs, we need to create a + // local_load for them. + ttg::LocalLoadOp maybeLocalLoad; + if (std::distance(old.getUsers().begin(), old.getUsers().end()) > + allocsToErase.size()) { + auto loc = old.getOwner()->getLoc(); + maybeLocalLoad = + ttg::LocalLoadOp::create(builder, loc, old.getType(), alloc, token); + old.replaceAllUsesWith(maybeLocalLoad); + } + for (auto alloc : allocsToErase) { + alloc.erase(); + } + return maybeLocalLoad; +} + +bool comesFromLoadOrBlockArg(Value v) { + // Peel out the original cvt dot_op<..., #blocked> + // and any other potential cvt/trans ops + while (true) { + Operation *def = v.getDefiningOp(); + if (!def) + break; + if (auto cvtOp = dyn_cast(def)) { + v = cvtOp.getSrc(); + continue; + } + if (auto transOp = dyn_cast(def)) { + v = transOp.getSrc(); + continue; + } + if (def->hasTrait()) { + v = def->getOperand(0); + continue; + } + break; + } + // We also accept block arguments as they appear in many MLIR tests + // If this is problematic we can totally drop them + Operation *def = v.getDefiningOp(); + bool fromLoad = def && isa(def); +#ifdef __TLE__ + bool fromSharedMemory = + def && def->hasAttrOfType("tt.memory_space") && + def->getAttrOfType("tt.memory_space").getValue() == + "shared_memory"; + return isa(v) || fromLoad || fromSharedMemory; +#else + return isa(v) || fromLoad; +#endif +} + +SmallVector getTiedArgs(Operation *op, int resultIdx) { + if (auto forOp = dyn_cast(op)) { + auto iterArg = forOp.getRegionIterArg(resultIdx); + auto result = forOp.getResult(resultIdx); + auto yieldVal = forOp.getBody()->getTerminator()->getOperand(resultIdx); + auto initVal = forOp.getInitArgs()[resultIdx]; + return {iterArg, result, yieldVal, initVal}; + } else if (auto whileOp = dyn_cast(op)) { + auto iterArg = whileOp.getBeforeArguments()[resultIdx]; + auto result = whileOp.getResults()[resultIdx]; + auto yieldVal = whileOp.getConditionOp().getArgs()[resultIdx]; + auto initVal = whileOp.getOperands()[resultIdx]; + auto bodyArg = whileOp.getAfterArguments()[resultIdx]; + return {iterArg, result, yieldVal, initVal, bodyArg}; + } else if (auto ifOp = dyn_cast(op)) { + SmallVector values; + for (auto &block : ifOp.getThenRegion().getBlocks()) { + auto terminator = block.getTerminator(); + if (isa(terminator)) + values.push_back(terminator->getOperands()[resultIdx]); + } + for (auto &block : ifOp.getElseRegion().getBlocks()) { + auto terminator = block.getTerminator(); + if (isa(terminator)) + values.push_back(terminator->getOperands()[resultIdx]); + } + values.push_back(ifOp->getResults()[resultIdx]); + return values; + } + return {}; +} + +LogicalResult verifyBarrierType(Operation *op, + mlir::triton::gpu::MemDescType barrierType) { + if (!barrierType.getElementType().isInteger(64) || + barrierType.getShape() != ArrayRef({1})) + return op->emitOpError( + "barrier allocation must be a descriptor of 1xi64 type"); + return success(); +} + +} // namespace mlir::triton diff --git a/third_party/sunrise/backend/sunrise_hint_handler.py b/third_party/sunrise/backend/sunrise_hint_handler.py new file mode 100644 index 0000000000..300b34e44c --- /dev/null +++ b/third_party/sunrise/backend/sunrise_hint_handler.py @@ -0,0 +1,54 @@ +# should store at third_party/sunrise/backend/ +# flagtree tle +# Sunrise hint handler: parses `#@hint: shared_memory` source comments and +# injects them as the `flagtree_hints` kwarg on tl.load, so the +# process-shared-memory-hint TTGIR pass can rewrite those loads into the async +# copy -> shared -> local_load chain. Mirrors nvidia_hint_handler.py. +import functools +import tokenize +from io import StringIO + +from triton.compiler.hint_manager import BaseHintHandler + + +class SunriseHintHandler(BaseHintHandler): + + @staticmethod + def get_node_hints(code_generator, node): + line_flagtree_hints = getattr(code_generator, 'flagtree_line_hints', {}) + return line_flagtree_hints.get(node.lineno) + + @staticmethod + def inject_kwargs_with_hints(fn, flagtree_hints, line_num, kws): + if fn.__name__ == "load" and flagtree_hints is not None: + print(f"[FLAGTREE] tl.load at line {line_num} has annotation {flagtree_hints}") + if 'flagtree_hints' not in kws: + kws['flagtree_hints'] = "" + if flagtree_hints not in kws['flagtree_hints']: + kws['flagtree_hints'] = flagtree_hints + + @staticmethod + def maps_line_numbers_to_comment_hints(jit_fn): + return dict(SunriseHintHandler._maps_line_numbers_to_comment_hints_from_source(jit_fn.src)) + + @staticmethod + @functools.lru_cache(maxsize=256) + def _maps_line_numbers_to_comment_hints_from_source(code_str): + # Maps line numbers to comment hints + line_flagtree_hints = {} + g = tokenize.generate_tokens(StringIO(code_str).readline) + for tok_type, tok_text, start, end, _ in g: + if tok_type == tokenize.COMMENT: + comment = tok_text.replace(" ", "").strip() + if comment.startswith('#@hint:'): + flagtree_hints = comment[len('#@hint:'):].strip() + # Record the line number of the comment + line_num = start[0] + line_flagtree_hints[line_num] = flagtree_hints + + return line_flagtree_hints + + @staticmethod + def attach_line_number_to_comment_mapping(tree, line_flagtree_hints): + if tree.body: + tree.body[0].line_flagtree_hints = line_flagtree_hints diff --git a/third_party/sunrise/language/tang/__init__.py b/third_party/sunrise/language/tang/__init__.py new file mode 100644 index 0000000000..2ce93570d3 --- /dev/null +++ b/third_party/sunrise/language/tang/__init__.py @@ -0,0 +1,3 @@ +from . import libdevice + +__all__ = ["libdevice"] \ No newline at end of file diff --git a/third_party/sunrise/language/tang/libdevice.py b/third_party/sunrise/language/tang/libdevice.py new file mode 100644 index 0000000000..b0cc4fb21e --- /dev/null +++ b/third_party/sunrise/language/tang/libdevice.py @@ -0,0 +1,203 @@ +from triton.language import core + +@core.extern +def erf(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_erf_f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("__ocml_erf_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def pow(arg0, arg1, _semantic=None): + return core.extern_elementwise( + "", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("int32")): ("__ocml_pown_f32", core.dtype("fp32")), + (core.dtype("fp64"), core.dtype("int32")): ("__ocml_pown_f32", core.dtype("fp64")), + (core.dtype("fp16"), core.dtype("fp16")): ("__ocml_pow_f16", core.dtype("fp16")), + (core.dtype("fp32"), core.dtype("fp32")): ("__ocml_pow_f32", core.dtype("fp32")), + (core.dtype("fp64"), core.dtype("fp64")): ("__ocml_pow_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def tanh(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_tanh_f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("__ocml_tanh_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def atan2(arg0, arg1, _semantic=None): + return core.extern_elementwise( + "", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("fp32")): ("__ocml_atan2_f32", core.dtype("fp32")), + (core.dtype("fp64"), core.dtype("fp64")): ("__ocml_atan2_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +# atanh, tanh2 +@core.extern +def atan(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_atan_f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("__ocml_atan_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def asin(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_asin_f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("__ocml_asin_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def acos(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_acos_f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("__ocml_acos_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def div_rd(arg0, arg1, _semantic=None): + return core.extern_elementwise( + "", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("fp32")): ("llvm.stvm.div.rm.f32", core.dtype("fp32")), + (core.dtype("fp64"), core.dtype("fp64")): ("llvm.stvm.div.rm.f32", core.dtype("fp64")), + }, 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")): ("llvm.stvm.div.rz.f32", core.dtype("fp32")), + (core.dtype("fp64"), core.dtype("fp64")): ("llvm.stvm.div.rz.f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def rsqrt(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("llvm.stvm.rsqrt.f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("llvm.stvm.rsqrt.f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def isinf(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("llvm.stvm.testp.f32.inf", core.dtype("int32")), + (core.dtype("fp64"), ): ("llvm.stvm.testp.f32.inf", core.dtype("int32")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def isnan(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [ + arg0, + ], { + (core.dtype("fp32"), ): ("llvm.stvm.testp.f32.not", core.dtype("int32")), + (core.dtype("fp64"), ): ("llvm.stvm.testp.f32.not", core.dtype("int32")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def sin(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_sin_f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("__ocml_sin_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def cos(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_cos_f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("__ocml_cos_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def tan(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_tan_f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("__ocml_tan_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def erf(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_erf_f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("__ocml_erf_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + + +@core.extern +def exp(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_exp_f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("__ocml_exp_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def exp2(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("llvm.stvm.exp2.f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("llvm.stvm.exp2.f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def div_rn(arg0, arg1, _semantic=None): + return core.extern_elementwise( + "", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("fp32")): ("llvm.stvm.div.rn.f32", core.dtype("fp32")), + (core.dtype("fp64"), core.dtype("fp64")): ("llvm.stvm.div.rn.f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def trunc(arg0, _semantic=None): + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("llvm.stvm.trunc.f", core.dtype("fp32")), + (core.dtype("fp64"), ): ("llvm.stvm.trunc.f", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def fmod(arg0, arg1, _semantic=None): + return core.extern_elementwise( + "", "", [arg0, arg1], { + (core.dtype("fp32"), core.dtype("fp32")): ("__ocml_fmod_f32", core.dtype("fp32")), + (core.dtype("fp64"), core.dtype("fp64")): ("__ocml_fmod_f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def isfinited(arg0, _semantic=None): + return core.extern_elementwise("", "", [arg0], { + (core.dtype("fp64"), ): ("__ocml_isfinite_f64", core.dtype("int32")), + }, is_pure=True, _semantic=_semantic).to(core.int1, _semantic=_semantic) + +@core.extern +def finitef(arg0, _semantic=None): + return core.extern_elementwise("", "", [arg0], { + (core.dtype("fp32"), ): ("__ocml_isfinite_f32", core.dtype("int32")), + }, is_pure=True, _semantic=_semantic).to(core.int1, _semantic=_semantic) + +@core.extern +def rint(arg0, _semantic=None): + return core.extern_elementwise("", "", [arg0], { + (core.dtype("fp32"), ): ("llvm.rint.f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("llvm.rint.f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + +@core.extern +def ffs(arg0, _semantic=None): + return core.extern_elementwise("", "", [arg0], { + (core.dtype("int32"), ): ("_Z5__ffsi", core.dtype("int32")), + (core.dtype("int64"), ): ("_Z7__ffsllx", core.dtype("int32")), + }, is_pure=True, _semantic=_semantic) diff --git a/third_party/sunrise/plugin/.empty b/third_party/sunrise/plugin/.empty new file mode 100644 index 0000000000..e69de29bb2 diff --git a/third_party/sunrise/python/src/gluon_ir.cc b/third_party/sunrise/python/src/gluon_ir.cc new file mode 100644 index 0000000000..c4009bb424 --- /dev/null +++ b/third_party/sunrise/python/src/gluon_ir.cc @@ -0,0 +1,984 @@ +#include "ir.h" +#include "pybind11/pybind11.h" +#include + +#include +#include + +#include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/DialectRegistry.h" +#include "mlir/IR/Types.h" +#include "third_party/amd/include/Dialect/TritonAMDGPU/IR/Dialect.h" +#include "triton/Analysis/Utility.h" +#include "triton/Dialect/Gluon/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/Attributes.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonGPU/IR/LinearLayoutConversions.h" +#include "triton/Dialect/TritonGPU/IR/Types.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +#include "triton/Tools/GenericSwizzling.h" +#include "triton/Tools/LayoutUtils.h" +#include "triton/Tools/LinearLayout.h" +#include "llvm/ADT/StringSwitch.h" +#include "llvm/Support/MathExtras.h" + +using namespace mlir; +namespace py = pybind11; +namespace tt = triton; +namespace ttg = triton::gpu; +namespace ttng = triton::nvidia_gpu; +namespace gluon = mlir::triton::gluon; +namespace ttag = mlir::triton::amdgpu; + +static ttg::CTAEncodingAttr +buildCtaLayoutAttr(MLIRContext *ctx, + const std::vector> &layout, + unsigned rank) { + auto kBlock = StringAttr::get(ctx, "block"); + tt::LinearLayout::BasesT bases; + bases[kBlock] = layout; + auto outDims = tt::standardOutDimNames(ctx, rank); + tt::LinearLayout ll(std::move(bases), outDims); + return ttg::CTAEncodingAttr::get(ctx, std::move(ll)); +} + +static std::vector> +getCgaLayoutBases(ttg::CTAEncodingAttr layout) { + std::vector> result; + auto ctx = layout.getContext(); + auto block = StringAttr::get(ctx, "block"); + const auto &basesMap = layout.getLinearLayout().getBases(); + auto it = basesMap.find(block); + assert(it != basesMap.end()); + return it->second; +} + +// Helper to check if an MLIR type or attribute has a verifier method. +template +static constexpr auto hasVerifier(AttrOrType t) + -> decltype(t.verifyInvariants, true) { + return true; +} +static constexpr auto hasVerifier(...) { return false; } + +// Print a diagnostic without its location. The frontend will attach the AST +// location to the error message. +static void printDiagStr(llvm::raw_ostream &os, const Diagnostic &diag) { + for (const DiagnosticArgument &arg : diag.getArguments()) + arg.print(os); + os << "\n"; + for (const Diagnostic ¬e : diag.getNotes()) + printDiagStr(os, note); +} + +struct GluonOpBuilder : public TritonOpBuilder { + using TritonOpBuilder::TritonOpBuilder; + // Construct an attribute or type while calling its verifier. Error messages + // are intercepted and sent back to Python via a C++ exception. + template + std::enable_if_t + getChecked(ArgTs &&...args) { + // Set up a scoped handler to intercept errors. + std::string msg; + llvm::raw_string_ostream os(msg); + ScopedDiagnosticHandler handler( + getContext(), [&](Diagnostic &diag) { printDiagStr(os, diag); }); + + auto result = + AttrOrType::getChecked([&] { return mlir::emitError(getLastLoc()); }, + std::forward(args)...); + if (!result) + throw std::runtime_error(os.str()); + return result; + } + + // A variant of the above due to issues with C++ overload resolution and how + // MLIR sets up the default `getChecked` implementation. + template + std::enable_if_t + getChecked(MLIRContext *ctx, ArgTs &&...args) { + // Set up a scoped handler to intercept errors. + std::string msg; + llvm::raw_string_ostream os(msg); + ScopedDiagnosticHandler handler( + getContext(), [&](Diagnostic &diag) { printDiagStr(os, diag); }); + + if (failed(AttrOrType::verifyInvariants( + [&] { return mlir::emitError(getLastLoc()); }, + std::forward(args)...))) + throw std::runtime_error(os.str()); + + return AttrOrType::get(ctx, std::forward(args)...); + } + + // Fallback method for types or attributes that do not have a verifier. + template + std::enable_if_t + getChecked(ArgTs &&...args) { + return AttrOrType::get(std::forward(args)...); + } +}; + +struct GluonLayouts { + py::handle AutoLayout; + py::handle CoalescedLayout; + py::handle BlockedLayout; + py::handle SliceLayout; + py::handle DistributedLinearLayout; + py::handle DotOperandLayout; + py::handle NVMMADistributedLayout; + py::handle TensorMemoryScalesLayout; + py::handle TensorMemoryLayout; + py::handle NVMMASharedLayout; + py::handle SwizzledSharedLayout; + py::handle SharedLinearLayout; + py::handle AMDMFMALayout; + py::handle AMDWMMALayout; + py::handle PaddedSharedLayout; + + GluonLayouts() { + auto layouts = + py::module::import("triton.experimental.gluon.language._layouts"); + auto amdLayouts = + py::module::import("triton.experimental.gluon.language.amd._layouts"); + auto blackwellLayouts = py::module::import( + "triton.experimental.gluon.language.nvidia.blackwell"); + AutoLayout = py::object(layouts.attr("AutoLayout")).release(); + CoalescedLayout = py::object(layouts.attr("CoalescedLayout")).release(); + BlockedLayout = py::object(layouts.attr("BlockedLayout")).release(); + SliceLayout = py::object(layouts.attr("SliceLayout")).release(); + DistributedLinearLayout = + py::object(layouts.attr("DistributedLinearLayout")).release(); + DotOperandLayout = py::object(layouts.attr("DotOperandLayout")).release(); + NVMMADistributedLayout = + py::object(layouts.attr("NVMMADistributedLayout")).release(); + TensorMemoryScalesLayout = + py::object(blackwellLayouts.attr("TensorMemoryScalesLayout")).release(); + TensorMemoryLayout = + py::object(blackwellLayouts.attr("TensorMemoryLayout")).release(); + NVMMASharedLayout = py::object(layouts.attr("NVMMASharedLayout")).release(); + SwizzledSharedLayout = + py::object(layouts.attr("SwizzledSharedLayout")).release(); + SharedLinearLayout = + py::object(layouts.attr("SharedLinearLayout")).release(); + AMDMFMALayout = py::object(amdLayouts.attr("AMDMFMALayout")).release(); + AMDWMMALayout = py::object(amdLayouts.attr("AMDWMMALayout")).release(); + PaddedSharedLayout = + py::object(layouts.attr("PaddedSharedLayout")).release(); + + auto core = py::module::import("triton.language.core"); + } +}; + +static bool isConvertLayoutTrivial(RankedTensorType dstTy, Value value) { + auto srcTy = cast(value.getType()); + if (srcTy.getEncoding() == dstTy.getEncoding()) + return true; + // Fail safe on unresolved layouts. + if (isa(srcTy.getEncoding())) + return false; + if (isa(dstTy.getEncoding())) + return false; + + // Check concrete layouts. + triton::LinearLayout cvt = minimalCvtLayout(srcTy, dstTy); + auto dims = llvm::to_vector(cvt.getInDimNames()); + return dims.empty() || (dims.size() == 1 && dims.front() == "register"); +} + +template +std::vector> toStdVector(R &&range) { + return {range.begin(), range.end()}; +} + +py::object layoutToGluon(Attribute layout) { + static GluonLayouts layouts; + if (auto blocked = dyn_cast(layout)) { + auto cgaBases = getCgaLayoutBases(blocked.getCTALayout()); + return layouts.BlockedLayout(toStdVector(blocked.getSizePerThread()), + toStdVector(blocked.getThreadsPerWarp()), + toStdVector(blocked.getWarpsPerCTA()), + toStdVector(blocked.getOrder()), cgaBases); + } else if (auto sliced = dyn_cast(layout)) { + return layouts.SliceLayout(sliced.getDim(), + layoutToGluon(sliced.getParent())); + } else if (auto linear = dyn_cast(layout)) { + const auto &ll = linear.getLinearLayout(); + auto ctx = layout.getContext(); + auto kReg = mlir::StringAttr::get(ctx, "register"); + auto kLane = mlir::StringAttr::get(ctx, "lane"); + auto kWarp = mlir::StringAttr::get(ctx, "warp"); + auto kBlock = mlir::StringAttr::get(ctx, "block"); + return layouts.DistributedLinearLayout( + ll.getBases().lookup(kReg), ll.getBases().lookup(kLane), + ll.getBases().lookup(kWarp), ll.getBases().lookup(kBlock), + toStdVector(ll.getOutDimSizes())); + } else if (auto dotOp = dyn_cast(layout)) { + return layouts.DotOperandLayout( + dotOp.getOpIdx(), layoutToGluon(dotOp.getParent()), dotOp.getKWidth()); + } else if (auto mma = dyn_cast(layout)) { + auto cgaBases = getCgaLayoutBases(mma.getCTALayout()); + return layouts.NVMMADistributedLayout( + std::vector{mma.getVersionMajor(), mma.getVersionMinor()}, + toStdVector(mma.getWarpsPerCTA()), toStdVector(mma.getInstrShape()), + cgaBases); + } else if (auto nvmma = dyn_cast(layout)) { + auto ctaLayout = nvmma.getCTALayout(); + auto cgaBases = getCgaLayoutBases(ctaLayout); + return layouts.NVMMASharedLayout(nvmma.getSwizzlingByteWidth(), + nvmma.getElementBitWidth(), + ctaLayout.getRank(), nvmma.getTransposed(), + nvmma.getFp4Padded(), cgaBases); + } else if (auto swizzled = + dyn_cast(layout)) { + auto cgaBases = getCgaLayoutBases(swizzled.getCTALayout()); + return layouts.SwizzledSharedLayout( + swizzled.getVec(), swizzled.getPerPhase(), swizzled.getMaxPhase(), + toStdVector(swizzled.getOrder()), cgaBases); + } else if (auto sharedLl = dyn_cast(layout)) { + const auto &ll = sharedLl.getLinearLayout(); + auto ctx = layout.getContext(); + auto kOffset = mlir::StringAttr::get(ctx, "offset"); + auto kBlock = mlir::StringAttr::get(ctx, "block"); + return layouts.SharedLinearLayout( + toStdVector(ll.getBases().lookup(kOffset)), + toStdVector(ll.getBases().lookup(kBlock)), sharedLl.getAlignment()); + } else if (auto autoEnc = dyn_cast(layout)) { + return layouts.AutoLayout(); + } else if (auto autoEnc = dyn_cast(layout)) { + return layouts.CoalescedLayout(); + } else if (auto amdMfma = dyn_cast(layout)) { + auto cgaBases = getCgaLayoutBases(amdMfma.getCTALayout()); + return layouts.AMDMFMALayout( + amdMfma.getVersion(), toStdVector(amdMfma.getInstrShape()), + amdMfma.getIsTransposed(), toStdVector(amdMfma.getWarpsPerCTA()), + amdMfma.getElementBitWidth(), toStdVector(amdMfma.getTilesPerWarp()), + cgaBases); + } else if (auto amdWmma = dyn_cast(layout)) { + auto cgaBases = getCgaLayoutBases(amdWmma.getCTALayout()); + return layouts.AMDWMMALayout( + amdWmma.getVersion(), amdWmma.getIsTransposed(), + toStdVector(amdWmma.getWarpsPerCTA()), + toStdVector(amdWmma.getInstrShape()), + toStdVector(amdWmma.getTilesPerWarp()), cgaBases); + } else if (auto paddedShared = + dyn_cast(layout)) { + auto *ctx = paddedShared.getContext(); + std::vector> intervalPaddingPairs; + for (auto [interval, padding] : + llvm::zip(paddedShared.getIntervals(), paddedShared.getPaddings())) { + intervalPaddingPairs.push_back({interval, padding}); + } + auto kOffset = mlir::StringAttr::get(ctx, "offset"); + auto kBlock = mlir::StringAttr::get(ctx, "block"); + const auto &ll = paddedShared.getLinearComponent(); + auto shape = toStdVector(ll.getOutDimSizes()); + return layouts.PaddedSharedLayout(intervalPaddingPairs, + ll.getBases().lookup(kOffset), + ll.getBases().lookup(kBlock), shape); + } else if (auto tmemScales = + dyn_cast(layout)) { + return layouts.TensorMemoryScalesLayout(std::vector{ + tmemScales.getCTASplitM(), tmemScales.getCTASplitN()}); + } else if (auto tmem = dyn_cast(layout)) { + return layouts.TensorMemoryLayout( + std::vector{tmem.getBlockM(), tmem.getBlockN()}, + tmem.getColStride(), + std::vector{tmem.getCTASplitM(), tmem.getCTASplitN()}); + } + + throw py::value_error("Unhandled encoding encountered"); +} + +template static void check(CondT &&cond, const char *msg) { + if (!std::forward(cond)) + throw py::value_error(msg); +} + +void init_gluon_ir(py::module &&m) { + using ret = py::return_value_policy; + + py::class_( + m, "GluonOpBuilder", py::module_local(), py::dynamic_attr()) + .def(py::init()) + .def("get_op_builder", &GluonOpBuilder::getBuilder, ret::reference) + .def("get_distributed_ty", + [](GluonOpBuilder &self, Type &elementType, + std::vector &shape, Attribute layout) -> Type { + return self.getChecked(shape, elementType, + layout); + }) + .def("get_shared_mem_desc_ty", + [](GluonOpBuilder &self, Type &elementType, + std::vector &shape, Attribute layout, + std::vector &allocShape) -> Type { + auto ctx = self.getContext(); + return self.getChecked( + shape, elementType, layout, + ttg::SharedMemorySpaceAttr::get(ctx), + /*mutableMemory=*/true, + /*allocShape=*/allocShape); + }) + .def("get_tensor_mem_desc_ty", + [](GluonOpBuilder &self, Type &elementType, + std::vector &shape, Attribute layout, + std::vector &allocShape) -> Type { + auto ctx = self.getContext(); + return self.getChecked( + shape, elementType, layout, + ttng::TensorMemorySpaceAttr::get(ctx), + /*mutableMemory=*/true, + /*allocShape=*/allocShape); + }) + .def("get_blocked_layout", + [](GluonOpBuilder &self, std::vector &sizePerThread, + std::vector &threadsPerWarp, + std::vector &warpsPerCta, std::vector &order, + std::vector> &cgaBases) -> Attribute { + auto ctx = self.getContext(); + unsigned rank = order.size(); + auto ctaLayout = buildCtaLayoutAttr(ctx, cgaBases, rank); + return self.getChecked( + ctx, sizePerThread, threadsPerWarp, warpsPerCta, order, + ctaLayout); + }) + .def("get_slice_layout", + [](GluonOpBuilder &self, unsigned dim, + Attribute parent) -> Attribute { + auto ctx = self.getContext(); + auto dist = cast(parent); + return self.getChecked(ctx, dim, dist); + }) + .def("get_distributed_linear_layout", + [](GluonOpBuilder &self, std::vector> regBases, + std::vector> laneBases, + std::vector> warpBases, + std::vector> blockBases, + std::vector shape) -> Attribute { + auto ctx = self.getContext(); + auto kReg = mlir::StringAttr::get(ctx, "register"); + auto kLane = mlir::StringAttr::get(ctx, "lane"); + auto kWarp = mlir::StringAttr::get(ctx, "warp"); + auto kBlock = mlir::StringAttr::get(ctx, "block"); + auto outDims = tt::standardOutDimPairs(ctx, shape); + auto ll = tt::LinearLayout({{kReg, regBases}, + {kLane, laneBases}, + {kWarp, warpBases}, + {kBlock, blockBases}}, + outDims, + /*requiresSurjective=*/true); + return ttg::LinearEncodingAttr::get(ctx, ll); + }) + .def("to_linear_layout", + [](GluonOpBuilder &self, Attribute layout, + std::vector &shape) -> py::object { + auto ctx = self.getContext(); + auto linearLayout = ttg::toLinearLayout(shape, layout); + auto attr = ttg::LinearEncodingAttr::get(ctx, linearLayout); + return layoutToGluon(attr); + }) + .def("get_dot_operand_layout", + [](GluonOpBuilder &self, unsigned opIdx, Attribute parent, + unsigned kWidth) -> Attribute { + return self.getChecked( + self.getContext(), opIdx, parent, kWidth); + }) + .def("get_mma_layout", + [](GluonOpBuilder &self, std::vector &version, + std::vector &warpsPerCta, + std::vector> &cgaBases, + std::vector &instrShape) -> Attribute { + auto ctx = self.getContext(); + unsigned rank = warpsPerCta.size(); + auto ctaLayout = buildCtaLayoutAttr(ctx, cgaBases, rank); + return self.getChecked( + ctx, version[0], version[1], warpsPerCta, ctaLayout, + instrShape); + }) + .def("get_amd_mfma_layout", + [](GluonOpBuilder &self, unsigned version, + std::vector &warpsPerCta, + std::vector &instrShape, bool transposed, + std::vector> &cgaBases, + std::vector &tilesPerWarp, + unsigned elementBitWidth) -> Attribute { + auto ctx = self.getContext(); + unsigned rank = warpsPerCta.size(); + auto ctaLayout = buildCtaLayoutAttr(ctx, cgaBases, rank); + return ttg::AMDMfmaEncodingAttr::get( + ctx, version, warpsPerCta, instrShape, transposed, ctaLayout, + tilesPerWarp, elementBitWidth); + }) + .def("get_amd_wmma_layout", + [](GluonOpBuilder &self, unsigned version, bool transposed, + std::vector &warpsPerCta, + std::vector &tilesPerWarp, + std::vector> &cgaBases, + std::vector &instrShape) -> Attribute { + auto ctx = self.getContext(); + unsigned rank = warpsPerCta.size(); + auto ctaLayout = buildCtaLayoutAttr(ctx, cgaBases, rank); + return ttg::AMDWmmaEncodingAttr::get(ctx, version, transposed, + warpsPerCta, tilesPerWarp, + ctaLayout, instrShape); + }) + .def("get_padded_shared_layout", + [](GluonOpBuilder &self, std::vector &intervals, + std::vector &paddings, + std::vector> &offsetBases, + std::vector> &blockBases, + std::vector &shape) -> Attribute { + auto ctx = self.getContext(); + auto rank = shape.size(); + auto kOffset = mlir::StringAttr::get(ctx, "offset"); + auto kBlock = mlir::StringAttr::get(ctx, "block"); + auto ll = tt::LinearLayout( + {{kOffset, offsetBases}, {kBlock, blockBases}}, + tt::standardOutDimNames(ctx, rank)); + return ttg::PaddedSharedEncodingAttr::get(ctx, intervals, paddings, + ll); + }) + .def("get_shared_linear_layout", + [](GluonOpBuilder &self, std::vector> &offsetBases, + std::vector> &blockBases, + unsigned alignment) -> Attribute { + auto ctx = self.getContext(); + auto kOffset = mlir::StringAttr::get(ctx, "offset"); + auto kBlock = mlir::StringAttr::get(ctx, "block"); + auto outDims = tt::standardOutDimNames(ctx, offsetBases[0].size()); + auto ll = tt::LinearLayout( + {{kOffset, offsetBases}, {kBlock, blockBases}}, outDims); + return self.getChecked(ctx, ll, + alignment); + }) + .def("get_nvmma_shared_layout", + [](GluonOpBuilder &self, unsigned swizzleByteWidth, + unsigned elementBitwidth, bool transposed, bool fp4Padded, + std::vector> &cgaBases, + unsigned rank) -> Attribute { + auto ctx = self.getContext(); + auto ctaLayout = buildCtaLayoutAttr(ctx, cgaBases, rank); + return self.getChecked( + ctx, swizzleByteWidth, transposed, elementBitwidth, fp4Padded, + ctaLayout); + }) + .def("get_auto_layout", + [](GluonOpBuilder &self) -> Attribute { + return self.getChecked(self.getContext()); + }) + .def("get_coalesced_layout", + [](GluonOpBuilder &self) -> Attribute { + return self.getChecked( + self.getContext()); + }) + .def("get_swizzled_shared_layout", + [](GluonOpBuilder &self, int vec, int perPhase, int maxPhase, + std::vector &order, + std::vector> &cgaBases) -> Attribute { + auto ctx = self.getContext(); + unsigned rank = order.size(); + auto ctaLayout = buildCtaLayoutAttr(ctx, cgaBases, rank); + return self.getChecked( + ctx, vec, perPhase, maxPhase, order, ctaLayout); + }) + .def("get_tensor_memory_layout", + [](GluonOpBuilder &self, std::vector &block, + unsigned colStride, std::vector &ctaSplitNum, + bool twoCTAs) -> Attribute { + auto ctx = self.getContext(); + check(block.size() == 2, "expected a 2D block"); + check(ctaSplitNum.size() == 2, "expected 2D CTA dimensions"); + return self.getChecked( + ctx, block[0], block[1], colStride, ctaSplitNum[0], + ctaSplitNum[1], twoCTAs); + }) + .def("get_tensor_memory_scales_layout", + [](GluonOpBuilder &self, + std::vector &ctaSplitNum) -> Attribute { + auto ctx = self.getContext(); + check(ctaSplitNum.size() == 2, "expected 2D CTA dimensions"); + return self.getChecked( + ctx, ctaSplitNum[0], ctaSplitNum[1]); + }) + .def("get_gluon_layout_from_tensor", + [](GluonOpBuilder &self, Value tensor) -> py::object { + auto ty = dyn_cast(tensor.getType()); + check(ty.getEncoding(), "expected a tensor with an encoding"); + return layoutToGluon(ty.getEncoding()); + }) + .def("get_gluon_layout_from_memdesc", + [](GluonOpBuilder &self, Value memdesc) -> py::object { + auto ty = dyn_cast(memdesc.getType()); + check(ty.getEncoding(), "expected a memdesc with an encoding"); + return layoutToGluon(ty.getEncoding()); + }) + .def("get_tensor_descriptor_layout_type", + [](GluonOpBuilder &self, Type blockType, bool isSigned, + Attribute layout) -> Type { + auto ctx = self.getContext(); + auto blockTy = cast(blockType); + auto blockTyLayout = blockTy.cloneWithEncoding(layout); + return triton::TensorDescType::get(ctx, blockTyLayout, isSigned); + }) + .def("is_convert_layout_trivial", + [](GluonOpBuilder &self, Type resultTy, Value value) -> bool { + auto dstTy = cast(resultTy); + return isConvertLayoutTrivial(dstTy, value); + }) + .def("create_histogram", + [](GluonOpBuilder &self, Value operand, int numBins, + std::optional mask, Attribute layout) -> Value { + auto *ctx = self.getContext(); + auto resultTy = + RankedTensorType::get({static_cast(numBins)}, + IntegerType::get(ctx, 32), layout); + if (!mask) { + return self.create(resultTy, operand); + } else { + return self.create(resultTy, operand, + *mask); + } + }) + .def("create_cat", + [](GluonOpBuilder &self, Value &lhs, Value &rhs, + Type retType) -> Value { + return self.create(retType, lhs, rhs); + }) + .def("create_fp4_to_fp", + [](GluonOpBuilder &self, Value src, Type elemType, + int axis) -> Value { + return self.create( + cast>(src), elemType, axis); + }) + .def("create_async_copy_global_to_local", + [](GluonOpBuilder &self, Value smem, Value pointer, Value mask, + Value other, tt::CacheModifier cacheModifier, + tt::EvictionPolicy evictionPolicy, bool isVolatile) { + self.create( + pointer, smem, mask, other, cacheModifier, evictionPolicy, + isVolatile); + }) + .def("create_async_copy_mbarrier_arrive", + [](GluonOpBuilder &self, Value mbarrier, bool incrementCount) { + self.create(mbarrier, + !incrementCount); + }) + .def("create_async_commit_group", + [](GluonOpBuilder &self) { + ValueRange tokens; + self.create(tokens); + }) + .def("create_async_wait_group", + [](GluonOpBuilder &self, int num) { + ValueRange tokens; + self.create(tokens, num); + }) + .def("create_convert_layout", + [](GluonOpBuilder &self, Type resultTy, Value value) -> Value { + return self.create(resultTy, value); + }) + .def("create_local_alloc", + [](GluonOpBuilder &self, Type resultTy) -> Value { + return self.create(resultTy); + }) + .def("create_local_alloc", + [](GluonOpBuilder &self, Type resultTy, Value value) -> Value { + return self.create(resultTy, value); + }) + .def("create_local_store", + [](GluonOpBuilder &self, Value memDesc, Value value) { + self.create(value, memDesc); + }) + .def("create_local_load", + [](GluonOpBuilder &self, Type resultTy, Value memDesc) -> Value { + return self.create(resultTy, memDesc); + }) + .def("get_shared_bank_conflicts", + [](GluonOpBuilder &self, Attribute regLayoutAttr, + Attribute sharedLayoutAttr, std::vector &shape, + int bitwidth) -> int { + auto regLayout = ttg::toLinearLayout(shape, regLayoutAttr); + auto smemLayout = ttg::toLinearLayout(shape, sharedLayoutAttr); + return ttg::bankConflictsMemDesc(regLayout, smemLayout, bitwidth); + }) + .def("create_local_dealloc", + [](GluonOpBuilder &self, Value memDesc) -> Operation * { + return self.create(memDesc); + }) + + .def("create_memdesc_index", + [](GluonOpBuilder &self, Type resultType, Value src, + Value index) -> Value { + return self.create(resultType, src, index); + }) + .def("create_memdesc_subslice", + [](GluonOpBuilder &self, Type resultType, Value src, + std::vector &offsets) -> Value { + return self.create(resultType, src, + offsets); + }) + .def("create_memdesc_trans", + [](GluonOpBuilder &self, Value src, + std::vector &order) -> Value { + return self.create(src, order); + }) + .def("create_memdesc_reshape", + [](GluonOpBuilder &self, Value src, + std::vector &shape) -> Value { + return self.create(src, shape); + }) + .def("create_memdesc_reinterpret", + [](GluonOpBuilder &self, Type resultType, Value src) -> Value { + return self.create(resultType, src); + }) + .def("create_set_auto_layout", + [](GluonOpBuilder &self, Attribute layout, Value value) -> Value { + return self.create(layout, value); + }) + .def("create_split", + [](GluonOpBuilder &self, Value &a) -> py::tuple { + auto argTy = cast(a.getType()); + auto ctx = argTy.getContext(); + auto enc = ttg::SliceEncodingAttr::get( + ctx, argTy.getRank() - 1, + cast(argTy.getEncoding())); + auto resTy = + RankedTensorType::get(ArrayRef(argTy.getShape()).drop_back(), + argTy.getElementType(), enc); + auto op = self.create(TypeRange{resTy, resTy}, a); + return py::make_tuple(op->getResult(0), op->getResult(1)); + }) + .def("create_warpgroup_mma", + [](GluonOpBuilder &self, Value a, Value b, Value acc, Value useAcc, + triton::InputPrecision precision = triton::InputPrecision::IEEE, + int maxNumImpreciseAcc = 0, bool isAsync = false) -> Value { + return self.create( + a, b, acc, useAcc, precision, maxNumImpreciseAcc, isAsync); + }) + .def("create_warpgroup_mma_wait", + [](GluonOpBuilder &self, std::vector &deps, int pendings) { + std::vector results; + auto wait = self.create(deps, pendings); + llvm::append_range(results, wait.getResults()); + return results; + }) + .def("create_tmem_alloc", + [](GluonOpBuilder &self, Type resultTy, Value value) -> Value { + return self.create(resultTy, value); + }) + .def("create_tmem_alloc", + [](GluonOpBuilder &self, Type resultTy, py::none value) -> Value { + return self.create(resultTy, Value{}); + }) + .def("create_tmem_store", + [](GluonOpBuilder &self, Value memDesc, Value value, Value pred) { + self.create(memDesc, value, pred); + }) + .def("create_tmem_load", + [](GluonOpBuilder &self, Type resultTy, Value memDesc) -> Value { + return self.create(resultTy, memDesc); + }) + .def("create_tmem_copy", + [](GluonOpBuilder &self, Value src, Value dst) { + self.create(src, dst, /*barrier=*/Value()); + }) + .def("create_tmem_subslice", + [](GluonOpBuilder &self, Type resultTy, Value memDesc, + int N) -> Value { + return self.create(resultTy, memDesc, N); + }) + .def("create_mbarrier_init", + [](GluonOpBuilder &self, Value memDesc, int count) { + self.create(memDesc, count); + }) + .def("create_mbarrier_inval", + [](GluonOpBuilder &self, Value memDesc) { + self.create(memDesc); + }) + .def("create_mbarrier_expect", + [](GluonOpBuilder &self, Value memDesc, int bytes, Value pred) { + self.create(memDesc, bytes, pred); + }) + .def("create_mbarrier_wait", + [](GluonOpBuilder &self, Value memDesc, Value phase, Value pred, + std::vector &deps) { + self.create(memDesc, phase, pred, deps); + }) + .def("create_mbarrier_arrive", + [](GluonOpBuilder &self, Value memDesc, int count, Value pred) { + self.create(memDesc, count, pred); + }) + .def("create_tcgen05_mma", + [](GluonOpBuilder &self, Value a, Value b, Value acc, Value useAcc, + Value pred, std::vector &mbarriers, + std::vector &mbarrier_preds, bool two_ctas) { + Value accDep; + auto tokType = self.getBuilder().getType(); + self.create(tokType, a, b, acc, accDep, useAcc, + pred, two_ctas, mbarriers, + mbarrier_preds); + }) + .def("create_tcgen05_mma_scaled", + [](GluonOpBuilder &self, Value a, Value b, Value acc, Value aScale, + Value bScale, tt::ScaleDotElemType aType, + tt::ScaleDotElemType bType, Value useAcc, Value pred, + std::vector &mbarriers, + std::vector &mbarrier_preds) { + Value accDep; + auto tokType = self.getBuilder().getType(); + self.create( + tokType, a, b, acc, accDep, aScale, bScale, aType, bType, + useAcc, pred, mbarriers, mbarrier_preds); + }) + .def("create_tcgen05_commit", + [](GluonOpBuilder &self, Value &barrier) { + self.create(barrier); + }) + + .def("create_async_tma_copy_global_to_local", + [](GluonOpBuilder &self, Value descPtr, std::vector &coord, + Value barrier, Value result, Value pred) { + self.create( + descPtr, coord, barrier, result, pred); + }) + .def("create_async_tma_copy_local_to_global", + [](GluonOpBuilder &self, Value descPtr, std::vector &coord, + Value src) { + self.create(descPtr, coord, + src); + }) + .def("create_async_tma_reduce", + [](GluonOpBuilder &self, triton::DescriptorReduceKind kind, + Value descPtr, std::vector &coord, Value src) { + self.create(kind, descPtr, coord, src); + }) + .def("create_async_tma_store_wait", + [](GluonOpBuilder &self, int pendings) { + self.create(pendings); + }) + .def("create_async_tma_gather", + [](GluonOpBuilder &self, Value descPtr, Value xOffsets, + Value yOffset, Value barrier, Value result, Value pred) { + self.create(descPtr, xOffsets, yOffset, + barrier, result, pred); + }) + .def("create_async_tma_scatter", + [](GluonOpBuilder &self, Value descPtr, Value xOffsets, + Value yOffset, Value src) { + self.create(descPtr, xOffsets, yOffset, + src); + }) + .def("create_fence_async_shared", + [](GluonOpBuilder &self, bool bCluster) -> OpState { + return self.create(bCluster); + }) + + .def("create_broadcast", + [](TritonOpBuilder &self, Value &arg, Type retTy) -> Value { + return self.create(retTy, arg); + }) + .def("create_warp_return", + [](GluonOpBuilder &self) -> Operation * { + return self.create(); + }) + .def("create_warp_yield", + [](GluonOpBuilder &self, std::vector &values) -> Operation * { + return self.create(values); + }) + .def("create_warp_specialize_partitions", + [](GluonOpBuilder &self, int numPartitions) -> Operation * { + return self.create(numPartitions); + }) + .def("create_warp_specialize", + [](GluonOpBuilder &self, std::vector &resultTypes, + std::vector &explicitCaptures, + std::vector &partitionNumWarps) { + return self.create( + resultTypes, explicitCaptures, partitionNumWarps); + }) + .def("create_buffer_load", + [](GluonOpBuilder &self, Type resultType, Value ptr, Value offsets, + Value mask, Value other, tt::CacheModifier cache) -> Value { + return self.create(resultType, ptr, offsets, + Value() /*stride*/, cache, + mask, other); + }) + .def("create_buffer_store", + [](GluonOpBuilder &self, Value storedValue, Value ptr, Value offsets, + Value mask, tt::CacheModifier cache) { + self.create(storedValue, ptr, offsets, + Value() /*stride*/, cache, mask); + }) + .def("create_buffer_atomic_rmw", + [](GluonOpBuilder &self, tt::RMWOp op, Value ptr, Value offsets, + Value value, tt::MemSemantic sem, tt::MemSyncScope scope, + Value mask) -> Value { + return self.create( + value.getType(), op, ptr, offsets, value, Value() /*stride*/, + sem, scope, mask); + }) + .def("create_buffer_load_to_local", + [](GluonOpBuilder &self, Value dest, Value ptr, Value offsets, + Value mask, Value other, Value stride, + tt::CacheModifier cacheModifier) { + self.create( + dest, ptr, offsets, mask, other, stride, cacheModifier); + }) + .def("create_make_tensor_descriptor", + [](TritonOpBuilder &self, Type resultTy, Value &base, + std::vector &shape, std::vector &strides, + tt::PaddingOption paddingOption) -> Value { + return self.create(resultTy, base, shape, + strides, paddingOption); + }) + .def("create_async_tdm_copy_global_to_local", + [](GluonOpBuilder &self, Value descPtr, std::vector &indices, + Value result, Value pred, Value barrier) { + self.create( + descPtr, indices, result, pred, barrier); + }) + .def("create_async_tdm_copy_local_to_global", + [](GluonOpBuilder &self, Value descPtr, std::vector &indices, + Value src) { + self.create(descPtr, indices, + src); + }) + .def("create_async_tdm_wait", + [](GluonOpBuilder &self, int num) { + ValueRange tokens; + self.create(tokens, num); + }) + .def("create_async_copy_lds_barrier_arrive", + [](GluonOpBuilder &self, Value mbarrier) { + self.create(mbarrier); + }) + .def("create_lds_barrier_init", + [](GluonOpBuilder &self, Value memDesc, int count) { + self.create(memDesc, count); + }) + .def("create_lds_barrier_wait", + [](GluonOpBuilder &self, Value memDesc, Value phase) { + self.create(memDesc, phase); + }) + .def("create_lds_barrier_arrive", + [](GluonOpBuilder &self, Value memDesc, int count) -> Value { + return self.create(memDesc, count); + }); + + m.def( + "compute_tmem_reg_layout", + [](py::object elementTyObj, std::vector shape, + py::object layoutObj, unsigned numWarps, const std::string &atomName, + std::vector> cgaBases) -> py::object { + DialectRegistry registry; + registry.insert(); + MLIRContext context(MLIRContext::Threading::DISABLED); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + + GluonOpBuilder builder(&context); + auto builderObj = + py::cast(&builder, py::return_value_policy::reference); + + auto elementType = elementTyObj.attr("to_ir")(builderObj).cast(); + auto layoutAttr = + layoutObj.attr("_to_ir")(builderObj).cast(); + auto allocShape = shape; + + auto ctx = builder.getContext(); + unsigned rank = shape.size(); + auto memDescTy = builder.getChecked( + shape, elementType, layoutAttr, + ttng::TensorMemorySpaceAttr::get(ctx), + /*mutableMemory=*/true, allocShape); + auto ctaLayoutAttr = buildCtaLayoutAttr(ctx, cgaBases, rank); + + auto maybeAtom = + llvm::StringSwitch>(atomName) + .Case("32x32b", ttng::TMemAccessAtom::I32x32b) + .Case("16x64b", ttng::TMemAccessAtom::I16x64b) + .Case("16x128b", ttng::TMemAccessAtom::I16x128b) + .Case("16x256b", ttng::TMemAccessAtom::I16x256b) + .Case("16x32bx2", ttng::TMemAccessAtom::I16x32bx2) + .Default(std::nullopt); + if (!maybeAtom) + throw std::invalid_argument("unknown TMEM access atom: " + atomName); + auto atom = *maybeAtom; + if (atom == ttng::TMemAccessAtom::I16x32bx2) + throw std::invalid_argument( + "Atom 16x32bx2 is inferred implicitly and cannot be requested " + "explicitly"); + if (numWarps < 4 || !llvm::isPowerOf2_32(numWarps)) + throw std::invalid_argument( + "numWarps must be a power of two and >= 4"); + + auto layout = ttng::getDistributedLayoutForTmemLdSt( + memDescTy, atom, numWarps, ctaLayoutAttr); + if (!layout) + return py::none(); + + auto attr = ttg::LinearEncodingAttr::get(ctx, *layout); + return layoutToGluon(attr); + }); + + m.def( + "make_cga_layout", + [](std::vector ctasPerCga, std::vector ctaSplitNum, + std::vector ctaOrder) -> std::vector> { + DialectRegistry registry; + registry.insert(); + MLIRContext ctx(MLIRContext::Threading::DISABLED); + ctx.appendDialectRegistry(registry); + ctx.loadAllAvailableDialects(); + auto attr = ttg::CTAEncodingAttr::fromSplitParams( + &ctx, ctasPerCga, ctaSplitNum, ctaOrder); + return getCgaLayoutBases(attr); + }); + + m.def("get_amd_mfma_scale_layout", + [](unsigned opIdx, std::vector &shape, unsigned mfmaMDim, + std::vector &tilesPerWarp, + std::vector &warpsPerCTA) -> py::object { + DialectRegistry registry; + registry.insert(); + MLIRContext ctx(MLIRContext::Threading::DISABLED); + ctx.appendDialectRegistry(registry); + ctx.loadAllAvailableDialects(); + + auto ll = ttg::chooseScaledMfmaScaleLayout( + &ctx, opIdx, shape, mfmaMDim, tilesPerWarp, warpsPerCTA); + auto attr = ttg::LinearEncodingAttr::get(&ctx, ll); + return layoutToGluon(attr); + }); + + m.def("get_amd_wmma_scale_layout", + [](unsigned opIdx, std::vector &shape, unsigned wmmaMDim, + std::vector &tilesPerWarp, + std::vector &warpsPerCTA) -> py::object { + DialectRegistry registry; + registry.insert(); + MLIRContext ctx(MLIRContext::Threading::DISABLED); + ctx.appendDialectRegistry(registry); + ctx.loadAllAvailableDialects(); + + auto ll = ttg::chooseScaledWmmaScaleLayout( + &ctx, opIdx, shape, wmmaMDim, tilesPerWarp, warpsPerCTA); + auto attr = ttg::LinearEncodingAttr::get(&ctx, ll); + return layoutToGluon(attr); + }); + + py::class_(m, "WarpSpecializeOp", + py::module_local()) + .def("get_default_region", &ttg::WarpSpecializeOp::getDefaultRegion, + ret::reference) + .def("get_partition_op_holder", + &ttg::WarpSpecializeOp::getPartitionOpHolder, ret::reference) + .def("set_requested_registers", [](ttg::WarpSpecializeOp &self, + std::vector &requestedRegisters) { + self.setRequestedRegisters(requestedRegisters); + }); +} diff --git a/third_party/sunrise/python/src/interpreter.cc b/third_party/sunrise/python/src/interpreter.cc new file mode 100644 index 0000000000..747a0cc171 --- /dev/null +++ b/third_party/sunrise/python/src/interpreter.cc @@ -0,0 +1,740 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace py = pybind11; + +namespace { + +struct npy_half { + uint16_t value; +}; + +enum class MemSemantic { ACQUIRE_RELEASE, ACQUIRE, RELEASE, RELAXED }; + +std::mutex atomic_op_guard; + +template +constexpr bool is_reinterpret_cast_to_atomic_safe = + std::is_trivially_copyable_v && + std::is_trivially_copyable_v> && + std::is_standard_layout_v && std::is_standard_layout_v> && + sizeof(T) == sizeof(std::atomic) && + alignof(T) == alignof(std::atomic); + +enum class RMWOp { ADD, FADD, AND, OR, XOR, XCHG, MAX, MIN, UMIN, UMAX }; + +std::map mem_semantic_map = { + {MemSemantic::ACQUIRE_RELEASE, std::memory_order_acq_rel}, + {MemSemantic::ACQUIRE, std::memory_order_acquire}, + {MemSemantic::RELEASE, std::memory_order_release}, + {MemSemantic::RELAXED, std::memory_order_relaxed}, +}; + +template +T atomic_cmp(T *ptr, T val, std::memory_order order) { + auto cmp = [](T old, T val) { + if constexpr (is_min) { + return old > val; + } else { + return old < val; + } + }; + + T old_val; + if constexpr (is_reinterpret_cast_to_atomic_safe) { + std::atomic *atomic_ptr = reinterpret_cast *>(ptr); + old_val = atomic_ptr->load(order); + while (cmp(old_val, val)) { + if (atomic_ptr->compare_exchange_weak(old_val, val, order, order)) { + break; + } + } + } else { + const std::lock_guard lock(atomic_op_guard); + old_val = *ptr; + if (cmp(old_val, val)) { + *ptr = val; + } + } + return old_val; +} + +template T atomic_fadd(T *loc, T value, std::memory_order order) { + static_assert(std::is_floating_point::value, + "T must be a floating-point type"); + T old_value; + + if constexpr (is_reinterpret_cast_to_atomic_safe) { + T new_value; + std::atomic *atomic_loc = reinterpret_cast *>(loc); + old_value = atomic_loc->load(order); + do { + new_value = old_value + value; + } while ( + !atomic_loc->compare_exchange_weak(old_value, new_value, order, order)); + } else { + const std::lock_guard lock(atomic_op_guard); + old_value = *loc; + *loc = old_value + value; + } + + return old_value; +} + +/** Create a value of type `To` from the bits of `from`. + * + * similar to `std::bit_cast` but compatible with C++17, + * should perform similar to `*reinterpret_cast(&from)` + * or through punning without expecting any undefined behaviors. + * + * Note: taken from + * https://github.com/numpy/numpy/blob/70fde29fdd4d8fcc6098df7ef8a34c84844e347f/numpy/_core/src/common/utils.hpp#L32 + * with simplification. + */ +template +inline To BitCast(const From &from) noexcept { + static_assert(sizeof(To) == sizeof(From), + "both data types must have the same size"); + + static_assert(std::is_trivially_copyable_v && + std::is_trivially_copyable_v, + "both data types must be trivially copyable"); + + To to; + memcpy(&to, &from, sizeof(from)); + return to; +} + +// Taken from +// https://github.com/numpy/numpy/blob/70fde29fdd4d8fcc6098df7ef8a34c84844e347f/numpy/_core/src/common/half_private.hpp#L14 +template +inline uint16_t FromFloatBits(uint32_t f) { + uint32_t f_exp, f_sig; + uint16_t h_sgn, h_exp, h_sig; + + h_sgn = (uint16_t)((f & 0x80000000u) >> 16); + f_exp = (f & 0x7f800000u); + + /* Exponent overflow/NaN converts to signed inf/NaN */ + if (f_exp >= 0x47800000u) { + if (f_exp == 0x7f800000u) { + /* Inf or NaN */ + f_sig = (f & 0x007fffffu); + if (f_sig != 0) { + /* NaN - propagate the flag in the significand... */ + uint16_t ret = (uint16_t)(0x7c00u + (f_sig >> 13)); + /* ...but make sure it stays a NaN */ + if (ret == 0x7c00u) { + ret++; + } + return h_sgn + ret; + } else { + /* signed inf */ + return (uint16_t)(h_sgn + 0x7c00u); + } + } else { + if constexpr (gen_overflow) { + // FloatStatus::RaiseOverflow(); + throw std::overflow_error("overflow to signed inf"); + } + return (uint16_t)(h_sgn + 0x7c00u); + } + } + + /* Exponent underflow converts to a subnormal half or signed zero */ + if (f_exp <= 0x38000000u) { + /* + * Signed zeros, subnormal floats, and floats with small + * exponents all convert to signed zero half-floats. + */ + if (f_exp < 0x33000000u) { + if constexpr (gen_underflow) { + /* If f != 0, it underflowed to 0 */ + if ((f & 0x7fffffff) != 0) { + // FloatStatus::RaiseUnderflow(); + throw std::underflow_error(""); + } + } + return h_sgn; + } + /* Make the subnormal significand */ + f_exp >>= 23; + f_sig = (0x00800000u + (f & 0x007fffffu)); + if constexpr (gen_underflow) { + /* If it's not exactly represented, it underflowed */ + if ((f_sig & (((uint32_t)1 << (126 - f_exp)) - 1)) != 0) { + // FloatStatus::RaiseUnderflow(); + throw std::underflow_error(""); + } + } + /* + * Usually the significand is shifted by 13. For subnormals an + * additional shift needs to occur. This shift is one for the largest + * exponent giving a subnormal `f_exp = 0x38000000 >> 23 = 112`, which + * offsets the new first bit. At most the shift can be 1+10 bits. + */ + f_sig >>= (113 - f_exp); + /* Handle rounding by adding 1 to the bit beyond half precision */ + if constexpr (round_even) { + /* + * If the last bit in the half significand is 0 (already even), and + * the remaining bit pattern is 1000...0, then we do not add one + * to the bit after the half significand. However, the (113 - f_exp) + * shift can lose up to 11 bits, so the || checks them in the original. + * In all other cases, we can just add one. + */ + if (((f_sig & 0x00003fffu) != 0x00001000u) || (f & 0x000007ffu)) { + f_sig += 0x00001000u; + } + } else { + f_sig += 0x00001000u; + } + h_sig = (uint16_t)(f_sig >> 13); + /* + * If the rounding causes a bit to spill into h_exp, it will + * increment h_exp from zero to one and h_sig will be zero. + * This is the correct result. + */ + return (uint16_t)(h_sgn + h_sig); + } + + /* Regular case with no overflow or underflow */ + h_exp = (uint16_t)((f_exp - 0x38000000u) >> 13); + /* Handle rounding by adding 1 to the bit beyond half precision */ + f_sig = (f & 0x007fffffu); + if constexpr (round_even) { + /* + * If the last bit in the half significand is 0 (already even), and + * the remaining bit pattern is 1000...0, then we do not add one + * to the bit after the half significand. In all other cases, we do. + */ + if ((f_sig & 0x00003fffu) != 0x00001000u) { + f_sig += 0x00001000u; + } + } else { + f_sig += 0x00001000u; + } + h_sig = (uint16_t)(f_sig >> 13); + /* + * If the rounding causes a bit to spill into h_exp, it will + * increment h_exp by one and h_sig will be zero. This is the + * correct result. h_exp may increment to 15, at greatest, in + * which case the result overflows to a signed inf. + */ + if constexpr (gen_overflow) { + h_sig += h_exp; + if (h_sig == 0x7c00u) { + // FloatStatus::RaiseOverflow(); + throw std::overflow_error(""); + } + return h_sgn + h_sig; + } else { + return h_sgn + h_exp + h_sig; + } +} + +// Taken from +// https://github.com/numpy/numpy/blob/70fde29fdd4d8fcc6098df7ef8a34c84844e347f/numpy/_core/src/common/half_private.hpp#L269 +constexpr uint32_t ToFloatBits(uint16_t h) { + uint16_t h_exp = (h & 0x7c00u); + uint32_t f_sgn = ((uint32_t)h & 0x8000u) << 16; + switch (h_exp) { + case 0x0000u: { // 0 or subnormal + uint16_t h_sig = (h & 0x03ffu); + // Signed zero + if (h_sig == 0) { + return f_sgn; + } + // Subnormal + h_sig <<= 1; + while ((h_sig & 0x0400u) == 0) { + h_sig <<= 1; + h_exp++; + } + uint32_t f_exp = ((uint32_t)(127 - 15 - h_exp)) << 23; + uint32_t f_sig = ((uint32_t)(h_sig & 0x03ffu)) << 13; + return f_sgn + f_exp + f_sig; + } + case 0x7c00u: // inf or NaN + // All-ones exponent and a copy of the significand + return f_sgn + 0x7f800000u + (((uint32_t)(h & 0x03ffu)) << 13); + default: // normalized + // Just need to adjust the exponent and shift + return f_sgn + (((uint32_t)(h & 0x7fffu) + 0x1c000u) << 13); + } +} + +npy_half npy_float_to_half(float f) { + return {FromFloatBits(BitCast(f))}; +} + +float npy_half_to_float(npy_half h) { + return BitCast(ToFloatBits(h.value)); +} + +template <> +npy_half atomic_fadd(npy_half *loc, npy_half value, + std::memory_order order) { + npy_half old_value; + + const std::lock_guard lock(atomic_op_guard); + old_value = *loc; + *loc = npy_float_to_half(npy_half_to_float(old_value) + + npy_half_to_float(value)); + + return old_value; +} + +class AtomicOp { +public: + AtomicOp(const uint64_t *ptr, size_t numel, std::memory_order order) + : ptr(ptr), numel(numel), order(order) {} + + void apply() { + for (size_t i = 0; i < numel; ++i) { + applyAt(reinterpret_cast(ptr[i]), i); + } + } + + virtual ~AtomicOp() = default; + +protected: + virtual void applyAt(void *, size_t i) = 0; + + const uint64_t *ptr; + size_t numel; + std::memory_order order; +}; + +template class AtomicRMWOpBase : public AtomicOp { +public: + AtomicRMWOpBase(const uint64_t *ptr, const void *val, void *ret, + const bool *mask, size_t numel, std::memory_order order) + : AtomicOp(ptr, numel, order), val(val), ret(ret), mask(mask) {} + +protected: + void applyAt(void *loc, size_t i) override final { + if (mask[i]) { + DType *ptr = static_cast(loc); + *(static_cast(ret) + i) = + applyAtMasked(ptr, *(static_cast(val) + i), order); + } + } + + virtual DType applyAtMasked(DType *loc, const DType value, + std::memory_order order) = 0; + + const void *val; + void *ret; + const bool *mask; +}; + +template +class AtomicRMWOp : public AtomicRMWOpBase { +public: + using AtomicRMWOpBase::AtomicRMWOpBase; +}; + +template +class AtomicRMWOp> + : public AtomicRMWOpBase { +public: + using AtomicRMWOpBase::AtomicRMWOpBase; + +protected: + DType applyAtMasked(DType *loc, const DType value, + std::memory_order order) override { + DType old_val; + if constexpr (is_reinterpret_cast_to_atomic_safe) { + std::atomic *atomic_loc = + reinterpret_cast *>(loc); + old_val = std::atomic_fetch_add_explicit(atomic_loc, value, order); + } else { + const std::lock_guard lock(atomic_op_guard); + old_val = *loc; + *loc = *loc + value; + } + return old_val; + } +}; + +template +class AtomicRMWOp> + : public AtomicRMWOpBase { +public: + using AtomicRMWOpBase::AtomicRMWOpBase; + +protected: + DType applyAtMasked(DType *loc, const DType value, + std::memory_order order) override { + return atomic_fadd(loc, value, order); + } +}; + +template +class AtomicRMWOp> + : public AtomicRMWOpBase { +public: + using AtomicRMWOpBase::AtomicRMWOpBase; + +protected: + DType applyAtMasked(DType *loc, const DType value, + std::memory_order order) override { + DType old_val; + if constexpr (is_reinterpret_cast_to_atomic_safe) { + std::atomic *atomic_loc = + reinterpret_cast *>(loc); + old_val = std::atomic_fetch_and_explicit(atomic_loc, value, order); + } else { + const std::lock_guard lock(atomic_op_guard); + old_val = *loc; + *loc = *loc & value; + } + return old_val; + } +}; + +template +class AtomicRMWOp> + : public AtomicRMWOpBase { +public: + using AtomicRMWOpBase::AtomicRMWOpBase; + +protected: + DType applyAtMasked(DType *loc, const DType value, + std::memory_order order) override { + DType old_val; + if constexpr (is_reinterpret_cast_to_atomic_safe) { + std::atomic *atomic_loc = + reinterpret_cast *>(loc); + old_val = std::atomic_fetch_or_explicit(atomic_loc, value, order); + } else { + const std::lock_guard lock(atomic_op_guard); + old_val = *loc; + *loc = *loc | value; + } + return old_val; + } +}; + +template +class AtomicRMWOp> + : public AtomicRMWOpBase { +public: + using AtomicRMWOpBase::AtomicRMWOpBase; + +protected: + DType applyAtMasked(DType *loc, const DType value, + std::memory_order order) override { + DType old_val; + if constexpr (is_reinterpret_cast_to_atomic_safe) { + std::atomic *atomic_loc = + reinterpret_cast *>(loc); + old_val = std::atomic_fetch_xor_explicit(atomic_loc, value, order); + } else { + const std::lock_guard lock(atomic_op_guard); + old_val = *loc; + *loc = *loc ^ value; + } + return old_val; + } +}; + +template +class AtomicRMWOp> + : public AtomicRMWOpBase { +public: + using AtomicRMWOpBase::AtomicRMWOpBase; + +protected: + DType applyAtMasked(DType *loc, const DType value, + std::memory_order order) override { + return atomic_cmp(loc, value, order); + } +}; + +template +class AtomicRMWOp> + : public AtomicRMWOpBase { +public: + using AtomicRMWOpBase::AtomicRMWOpBase; + +protected: + DType applyAtMasked(DType *loc, const DType value, + std::memory_order order) override { + return atomic_cmp(loc, value, order); + } +}; + +template +class AtomicRMWOp> + : public AtomicRMWOpBase { +public: + using AtomicRMWOpBase::AtomicRMWOpBase; + +protected: + DType applyAtMasked(DType *loc, const DType value, + std::memory_order order) override { + DType old_val; + if constexpr (is_reinterpret_cast_to_atomic_safe) { + std::atomic *atomic_loc = + reinterpret_cast *>(loc); + old_val = atomic_loc->exchange(value, order); + } else { + const std::lock_guard lock(atomic_op_guard); + old_val = *loc; + *loc = value; + } + return old_val; + } +}; + +template +void atomic_compare_exchange_strong(void *loc, void *expected, + const void *desired, size_t i, + std::memory_order order) { + T desired_val = *(static_cast(desired) + i); + T *expected_uint = static_cast(expected) + i; + + if constexpr (is_reinterpret_cast_to_atomic_safe) { + std::atomic *atomic_loc = reinterpret_cast *>(loc); + atomic_loc->compare_exchange_strong(*expected_uint, desired_val, order, + order); + } else { + const std::lock_guard lock(atomic_op_guard); + T *atomic_loc = static_cast(loc); + if (*atomic_loc == *expected_uint) { + *atomic_loc = desired_val; + } else { + *expected_uint = *atomic_loc; + } + } +} + +class AtomicCASOp : public AtomicOp { +public: + AtomicCASOp(const uint64_t *ptr, void *expected, const void *desired, + size_t itemsize, size_t numel, std::memory_order order) + : AtomicOp(ptr, numel, order), expected(expected), desired(desired), + itemsize(itemsize) {} + +protected: + void applyAt(void *loc, size_t i) override { + // Atomic operations perform bitwise comparison, so it's safe to + // use number of bytes (itemsize) to determine the type of pointers + if (itemsize == 1) { + atomic_compare_exchange_strong(loc, expected, desired, i, order); + } else if (itemsize == 2) { + atomic_compare_exchange_strong(loc, expected, desired, i, + order); + } else if (itemsize == 4) { + atomic_compare_exchange_strong(loc, expected, desired, i, + order); + } else if (itemsize == 8) { + atomic_compare_exchange_strong(loc, expected, desired, i, + order); + } else { + throw std::invalid_argument("Invalid byte size"); + } + } + +private: + void *expected; + const void *desired; + size_t itemsize; +}; + +// This is a workaround because explicit template parameter list for lambdas is +// a C++20 extension: +// auto try_make_op = [&]() { +// if (dtype.is(pybind11::dtype::of())) { +// atomic_op = std::make_unique>(ptr, val, ret, mask, +// numel, order); +// } +// }; +template struct OpCreator { + pybind11::dtype dtype; + const uint64_t *ptr; + const void *val; + void *ret; + const bool *mask; + size_t numel; + std::memory_order order; + std::unique_ptr &atomic_op; + + template void create() { + if (!atomic_op && dtype.is(pybind11::dtype::of())) { + atomic_op = std::make_unique>(ptr, val, ret, mask, + numel, order); + } + } +}; + +template <> template <> void OpCreator::create() { + if (!atomic_op && dtype.char_() == 'e') { // float16 + // workaround until https://github.com/pybind/pybind11/issues/4061 is + // implemented + atomic_op = std::make_unique>( + ptr, val, ret, mask, numel, order); + } +}; + +template +std::unique_ptr +makeAtomicRMWOp(pybind11::dtype dtype, const uint64_t *ptr, const void *val, + void *ret, const bool *mask, size_t numel, + std::memory_order order) { + // Iterate over all supported data types, make one that matches, and return + std::unique_ptr atomic_op; + OpCreator try_make_op{dtype, ptr, val, ret, + mask, numel, order, atomic_op}; + + (try_make_op.template create(), ...); + if (!atomic_op) { + throw std::invalid_argument("Unsupported data type"); + } + // Make it a unique_ptr + return atomic_op; +} + +} // namespace + +void init_triton_interpreter(py::module &&m) { + using ret = py::return_value_policy; + + py::enum_(m, "MEM_SEMANTIC", py::module_local()) + .value("ACQUIRE_RELEASE", MemSemantic::ACQUIRE_RELEASE) + .value("ACQUIRE", MemSemantic::ACQUIRE) + .value("RELEASE", MemSemantic::RELEASE) + .value("RELAXED", MemSemantic::RELAXED) + .export_values(); + + py::enum_(m, "RMW_OP", py::module_local()) + .value("ADD", RMWOp::ADD) + .value("FADD", RMWOp::FADD) + .value("AND", RMWOp::AND) + .value("OR", RMWOp::OR) + .value("XOR", RMWOp::XOR) + .value("XCHG", RMWOp::XCHG) + .value("MAX", RMWOp::MAX) + .value("MIN", RMWOp::MIN) + .value("UMIN", RMWOp::UMIN) + .value("UMAX", RMWOp::UMAX) + .export_values(); + + m.def("load", + [](py::array_t ptr, py::array_t mask, py::array other, + py::dtype ret_dtype) -> py::array { + int numel = ptr.size(); + auto shape = + std::vector(ptr.shape(), ptr.shape() + ptr.ndim()); + py::array ret(ret_dtype, py::array::ShapeContainer{numel}); + py::array_t reshaped_ptr = ptr.reshape({numel}); + py::array_t reshaped_mask = mask.reshape({numel}); + py::array reshaped_others = other.reshape({numel}); + for (size_t i = 0; i < ptr.size(); ++i) { + if (reshaped_mask.at(i)) + memcpy(ret.mutable_data(i), + reinterpret_cast(reshaped_ptr.at(i)), + ret_dtype.itemsize()); + else + memcpy(ret.mutable_data(i), reshaped_others.data(i), + ret_dtype.itemsize()); + } + return ret.reshape(shape); + }); + + m.def("store", + [](py::array_t ptr, py::array value, py::array_t mask) { + int numel = ptr.size(); + py::array_t reshaped_ptr = ptr.reshape({numel}); + py::array_t reshaped_mask = mask.reshape({numel}); + py::array reshaped_value = value.reshape({numel}); + for (size_t i = 0; i < ptr.size(); ++i) { + if (reshaped_mask.at(i)) { + memcpy(reinterpret_cast(reshaped_ptr.mutable_at(i)), + reshaped_value.data(i), value.dtype().itemsize()); + } + } + }); + + m.def("atomic_rmw", + [](RMWOp rmw_op, py::array_t ptr, py::array val, + py::array_t mask, MemSemantic sem) -> py::array { + std::memory_order order = mem_semantic_map[sem]; + int numel = ptr.size(); + auto shape = + std::vector(ptr.shape(), ptr.shape() + ptr.ndim()); + auto ret_dtype = val.dtype(); + py::array ret(ret_dtype, py::array::ShapeContainer{numel}); + py::array_t reshaped_ptr = ptr.reshape({numel}); + py::array_t reshaped_mask = mask.reshape({numel}); + py::array reshaped_val = val.reshape({numel}); + auto *ptr_data = reshaped_ptr.data(); + auto *mask_data = reshaped_mask.data(); + auto *val_data = static_cast(reshaped_val.data()); + auto *ret_data = static_cast(ret.mutable_data()); + + std::unique_ptr atomic_op; + +#define MAKE_ATOMIC_RMW_OP(OP_NAME, ...) \ + case OP_NAME: \ + atomic_op = makeAtomicRMWOp( \ + ret_dtype, ptr_data, val_data, ret_data, mask_data, numel, order); \ + break; + + switch (rmw_op) { + MAKE_ATOMIC_RMW_OP(RMWOp::ADD, int32_t, uint32_t, int64_t, uint64_t) + MAKE_ATOMIC_RMW_OP(RMWOp::FADD, npy_half, float, double) + MAKE_ATOMIC_RMW_OP(RMWOp::AND, int32_t, uint32_t, int64_t, uint64_t) + MAKE_ATOMIC_RMW_OP(RMWOp::OR, int32_t, uint32_t, int64_t, uint64_t) + MAKE_ATOMIC_RMW_OP(RMWOp::XOR, int32_t, uint32_t, int64_t, uint64_t) + MAKE_ATOMIC_RMW_OP(RMWOp::MAX, int32_t, int64_t) + MAKE_ATOMIC_RMW_OP(RMWOp::UMAX, uint32_t, uint64_t) + MAKE_ATOMIC_RMW_OP(RMWOp::MIN, int32_t, int64_t) + MAKE_ATOMIC_RMW_OP(RMWOp::UMIN, uint32_t, uint64_t) + MAKE_ATOMIC_RMW_OP(RMWOp::XCHG, int32_t, uint32_t, int64_t, + uint64_t) + default: + throw std::invalid_argument("Unsupported RMW operation"); + } + +#undef MAKE_ATOMIC_RMW_OP + + atomic_op->apply(); + return ret.reshape(shape); + }); + + m.def("atomic_cas", + [](py::array_t ptr, py::array &cmp, py::array &val, + MemSemantic sem) -> py::array { + std::memory_order order = mem_semantic_map[sem]; + int numel = ptr.size(); + auto shape = + std::vector(ptr.shape(), ptr.shape() + ptr.ndim()); + auto ret_dtype = cmp.dtype(); + py::array ret(ret_dtype, py::array::ShapeContainer{numel}); + py::array_t reshaped_ptr = ptr.reshape({numel}); + py::array reshaped_cmp = cmp.reshape({numel}); + py::array reshaped_val = val.reshape({numel}); + auto itemsize = cmp.itemsize(); + memcpy(static_cast(ret.mutable_data()), + static_cast(reshaped_cmp.data()), + itemsize * numel); + AtomicCASOp(reshaped_ptr.data(), ret.mutable_data(), + static_cast(reshaped_val.data()), itemsize, + numel, order) + .apply(); + return ret.reshape(shape); + }); +} diff --git a/third_party/sunrise/python/src/ir.cc b/third_party/sunrise/python/src/ir.cc new file mode 100644 index 0000000000..800a61e9fa --- /dev/null +++ b/third_party/sunrise/python/src/ir.cc @@ -0,0 +1,2110 @@ +#include "ir.h" + +#include +#include +#include +#include +#include + +#include "mlir/Bytecode/BytecodeWriter.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h" +#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" +#include "mlir/Dialect/DLTI/DLTI.h" +#include "mlir/Dialect/LLVMIR/LLVMAttrs.h" +#include "mlir/Dialect/LLVMIR/LLVMDialect.h" +#include "mlir/Dialect/LLVMIR/Transforms/InlinerInterfaceImpl.h" +#include "mlir/Dialect/UB/IR/UBOps.h" +#include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Diagnostics.h" +#include "mlir/IR/MLIRContext.h" +#include "mlir/IR/Verifier.h" +#include "mlir/Parser/Parser.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "mlir/Support/FileUtilities.h" +#include "mlir/Support/LLVM.h" +#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h" +#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h" +#include "mlir/Transforms/LocationSnapshot.h" + +#ifdef __TLE__ +#include "tle/dialect/include/IR/Dialect.h" // flagtree tle raw +#endif +#include "triton/Conversion/TritonGPUToLLVM/Utility.h" +#include "triton/Dialect/Gluon/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Dialect.h" +#include "triton/Dialect/Triton/IR/Types.h" +#include "triton/Dialect/Triton/IR/Utility.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Dialect/TritonInstrument/IR/Dialect.h" +#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h" +#include "triton/Dialect/TritonNvidiaGPU/Transforms/TMAUtilities.h" +#include "triton/Tools/Sys/GetEnv.hpp" +#include "llvm/ADT/SmallVector.h" +#include "llvm/Support/FileSystem.h" +#include "llvm/Support/SourceMgr.h" + +#ifdef __TLE__ +// flagtree tle +// Pointer to the TritonOpBuilder class, used to register IR ops for third-party +// dialects. +static py::class_ *builderClassPtr = nullptr; +namespace ir { +py::class_ *getBuilderClass() { return builderClassPtr; } +} // namespace ir +#endif + +namespace { +namespace py = pybind11; +using namespace mlir; +using namespace triton; +namespace tt = triton; +namespace ttg = triton::gpu; +namespace ttng = triton::nvidia_gpu; + +llvm::raw_fd_ostream &mlir_dumps() { + std::error_code EC; + static llvm::raw_fd_ostream S(::triton::tools::getStrEnv("MLIR_DUMP_PATH"), + EC, llvm::sys::fs::CD_CreateAlways); + assert(!EC); + return S; +} + +llvm::raw_ostream &mlir_dumps_or_dbgs() { + if (!::triton::tools::getStrEnv("MLIR_DUMP_PATH").empty()) { + return mlir_dumps(); + } else { + return llvm::dbgs(); + } +} + +// Function to parse a comma-separated string into a vector of C-style strings +llvm::SmallVector +parseCommaSeparatedValues(const std::string &input, + llvm::SmallVector &storage) { + llvm::SmallVector split; + llvm::SmallVector result; + StringRef(input.c_str()).split(split, ','); + llvm::transform(split, std::back_inserter(result), [&storage](StringRef str) { + // StringRefs are not always null-terminated. + // The purpose for this storage pattern is to + // produce a collection of C-strings that are. + storage.push_back(str.str()); + return storage.back().c_str(); + }); + return result; +} + +// Run the pass manager under a source manager diagnostic handler, which +// enables emitted MLIR diagnostics to directly reference Python source +// code. This diagnostic handler supports filtering diagnostic info by +// severity levels. +struct TritonSourceMgrDiagnosticHandler : public SourceMgrDiagnosticHandler { + TritonSourceMgrDiagnosticHandler(MLIRContext *ctx, + DiagnosticSeverity minSeverity) + : SourceMgrDiagnosticHandler(sourceMgr, ctx, llvm::errs()) { + setHandler([this, minSeverity](Diagnostic &diag) { + auto severity = diag.getSeverity(); + switch (severity) { + case DiagnosticSeverity::Error: + break; + case DiagnosticSeverity::Warning: + if (minSeverity == DiagnosticSeverity::Error) + return success(); + break; + case DiagnosticSeverity::Remark: + if (minSeverity == DiagnosticSeverity::Error || + minSeverity == DiagnosticSeverity::Warning) + return success(); + break; + case DiagnosticSeverity::Note: + // notes are handled somewhere else. + return failure(); + default: + llvm_unreachable("Unknown diagnostic severity"); + } + emitDiagnostic(diag); + return success(); + }); + } + + llvm::SourceMgr sourceMgr; +}; + +TritonSourceMgrDiagnosticHandler +setupTritonDiagnosticHandler(MLIRContext *context) { + bool showOperations = false, showStacktraces = false, showRemarks = false, + showWarnings = false; + + if (auto enableDiagnostics = + triton::tools::getStrEnv("MLIR_ENABLE_DIAGNOSTICS"); + !enableDiagnostics.empty()) { + llvm::SmallVector storage; + parseCommaSeparatedValues(enableDiagnostics, storage); + for (auto &str : storage) { + if (str == "warnings") { + showWarnings = true; + } else if (str == "remarks") { + showRemarks = true; + } else if (str == "stacktraces") { + showStacktraces = true; + } else if (str == "operations") { + showOperations = true; + } + // we show errors by default, so no need to set it + } + } + + DiagnosticSeverity minSeverity = + showWarnings ? DiagnosticSeverity::Warning : DiagnosticSeverity::Error; + minSeverity = showRemarks ? DiagnosticSeverity::Remark : minSeverity; + + context->printOpOnDiagnostic(showOperations); + context->printStackTraceOnDiagnostic(showStacktraces); + if (showStacktraces) { + context->disableMultithreading(); + } + + return TritonSourceMgrDiagnosticHandler(context, minSeverity); +} + +std::string locationToString(Location loc) { + std::string str; + llvm::raw_string_ostream os(str); + loc.print(os); + os.flush(); // Make sure all the content is dumped into the 'str' string + return str; +} + +void outputWarning(Location loc, const std::string &msg) { + std::string locStr = locationToString(loc); + + PyErr_WarnEx(PyExc_UserWarning, (locStr + ": " + msg).c_str(), + /*stack_level=*/2); +} + +// Allow dump a reproducer in the console on crash. +struct ConsoleReproducerStream : public mlir::ReproducerStream { + ~ConsoleReproducerStream() override {} + + StringRef description() override { + return "std::errs, please share the reproducer above with Triton project."; + } + raw_ostream &os() override { return llvm::errs(); } +}; + +ReproducerStreamFactory makeConsoleReproducer() { + return [](std::string &error) -> std::unique_ptr { + return std::make_unique(); + }; +} + +OpPrintingFlags getOpPrintingFlags() { + auto printingFlags = OpPrintingFlags(); + printingFlags.enableDebugInfo(); + printingFlags.printNameLocAsPrefix(true); + return printingFlags; +} + +py::list getTensorDescMetadata(ModuleOp &mod) { + py::list result; + triton::FuncOp kernelFunc; + mod.walk([&](triton::FuncOp func) { + if (triton::isKernel(func)) { + kernelFunc = func; + return WalkResult::interrupt(); + } + return WalkResult::skip(); + }); + assert(kernelFunc); + + for (auto [i, argTy] : llvm::enumerate(kernelFunc.getArgumentTypes())) { + auto descTy = dyn_cast(argTy); + if (!descTy) + continue; + + auto blockType = descTy.getBlockType(); + auto encoding = blockType.getEncoding(); + + py::dict metadata; + if (isa(encoding)) { + auto mmaEncoding = dyn_cast(encoding); + auto swizzle = ttng::getTMASwizzleMode(nullptr, descTy); + auto elemType = ttng::getTMAElementType(nullptr, descTy); + assert(swizzle.has_value()); + assert(elemType.has_value()); + auto blockSize = ttng::getTMABlockShape(blockType, /*packedSize=*/false); + metadata["swizzle"] = *swizzle; + metadata["elem_size"] = + descTy.getBlockType().getElementTypeBitWidth() / 8; + metadata["elem_type"] = *elemType; + metadata["block_size"] = + std::vector(blockSize.begin(), blockSize.end()); + metadata["fp4_padded"] = mmaEncoding && mmaEncoding.getFp4Padded(); + } else { + auto blockShape = blockType.getShape(); + metadata["block_size"] = + std::vector(blockShape.begin(), blockShape.end()); + metadata["elem_bits"] = blockType.getElementTypeBitWidth(); + + if (auto paddedEnc = dyn_cast(encoding)) { + py::list intervalPaddingPairs; + for (auto [interval, padding] : llvm::zip_equal( + paddedEnc.getIntervals(), paddedEnc.getPaddings())) { + py::list pair; + pair.append(interval); + pair.append(padding); + intervalPaddingPairs.append(pair); + } + metadata["interval_padding_pairs"] = intervalPaddingPairs; + + auto blockShape = blockType.getShape(); + } + } + result.append(std::move(metadata)); + } + return result; +} + +} // anonymous namespace + +/*****************************************************************************/ +/* Python bindings for ir */ +/*****************************************************************************/ + +void init_triton_ir(py::module &&m) { + using ret = py::return_value_policy; + using namespace pybind11::literals; + + py::enum_(m, "PADDING_OPTION", py::module_local()) + .value("PAD_ZERO", PaddingOption::PAD_ZERO) + .value("PAD_NAN", PaddingOption::PAD_NAN) + .export_values(); + + py::enum_(m, "CACHE_MODIFIER", py::module_local()) + .value("NONE", CacheModifier::NONE) + .value("CA", CacheModifier::CA) + .value("CG", CacheModifier::CG) + .value("WB", CacheModifier::WB) + .value("CS", CacheModifier::CS) + .value("WT", CacheModifier::WT) + .value("CV", CacheModifier::CV) + .export_values(); + + py::enum_(m, "MEM_SEMANTIC", py::module_local()) + .value("ACQUIRE_RELEASE", MemSemantic::ACQUIRE_RELEASE) + .value("ACQUIRE", MemSemantic::ACQUIRE) + .value("RELEASE", MemSemantic::RELEASE) + .value("RELAXED", MemSemantic::RELAXED) + .export_values(); + + py::enum_(m, "MEM_SYNC_SCOPE", py::module_local()) + .value("GPU", MemSyncScope::GPU) + .value("CTA", MemSyncScope::CTA) + .value("SYSTEM", MemSyncScope::SYSTEM) + .export_values(); + + py::enum_(m, "EVICTION_POLICY", py::module_local()) + .value("NORMAL", EvictionPolicy::NORMAL) + .value("EVICT_FIRST", EvictionPolicy::EVICT_FIRST) + .value("EVICT_LAST", EvictionPolicy::EVICT_LAST) + .export_values(); + + py::enum_(m, "ATOMIC_OP", py::module_local()) + .value("ADD", RMWOp::ADD) + .value("FADD", RMWOp::FADD) + .value("AND", RMWOp::AND) + .value("OR", RMWOp::OR) + .value("XOR", RMWOp::XOR) + .value("XCHG", RMWOp::XCHG) + .value("MAX", RMWOp::MAX) + .value("MIN", RMWOp::MIN) + .value("UMIN", RMWOp::UMIN) + .value("UMAX", RMWOp::UMAX); + + py::enum_(m, "DESCRIPTOR_REDUCE_KIND", + py::module_local()) + .value("ADD", DescriptorReduceKind::ADD) + .value("AND", DescriptorReduceKind::AND) + .value("OR", DescriptorReduceKind::OR) + .value("XOR", DescriptorReduceKind::XOR) + .value("MAX", DescriptorReduceKind::MAX) + .value("MIN", DescriptorReduceKind::MIN) + .value("INC", DescriptorReduceKind::INC) + .value("DEC", DescriptorReduceKind::DEC); + + py::enum_(m, "ROUNDING_MODE", py::module_local()) + .value("RTZ", RoundingMode::RTZ) + .value("RTNE", RoundingMode::RTNE); + + py::enum_(m, "PROPAGATE_NAN", py::module_local()) + .value("NONE", PropagateNan::NONE) + .value("ALL", PropagateNan::ALL); + + py::enum_(m, "INPUT_PRECISION", py::module_local()) + .value("TF32", InputPrecision::TF32) + .value("TF32x3", InputPrecision::TF32x3) + .value("IEEE", InputPrecision::IEEE) + .value("BF16x3", InputPrecision::BF16x3) + .value("BF16x6", InputPrecision::BF16x6) + .export_values(); + + py::enum_(m, "ScaleDotElemTypeTY", py::module_local()) + .value("E4M3", ScaleDotElemType::E4M3) + .value("E5M2", ScaleDotElemType::E5M2) + .value("E2M3", ScaleDotElemType::E2M3) + .value("E3M2", ScaleDotElemType::E3M2) + .value("E2M1", ScaleDotElemType::E2M1) + .value("BF16", ScaleDotElemType::BF16) + .value("FP16", ScaleDotElemType::FP16) + .export_values(); + + py::class_(m, "context", py::module_local()) + .def(py::init<>([]() { + return std::make_unique(MLIRContext::Threading::DISABLED); + })) + .def("printOpOnDiagnostic", + [](MLIRContext &self, bool v) { self.printOpOnDiagnostic(v); }) + .def("printStackTraceOnDiagnostic", [](MLIRContext &self, bool v) { + self.printStackTraceOnDiagnostic(v); + }); + + py::class_(m, "source_mgr_diag", + py::module_local()) + .def(py::init()); + + m.def("load_dialects", [](MLIRContext &context) { + DialectRegistry registry; + registry.insert(); + mlir::LLVM::registerInlinerInterface(registry); + registerBuiltinDialectTranslation(registry); + registerLLVMDialectTranslation(registry); + mlir::LLVM::registerInlinerInterface(registry); + context.appendDialectRegistry(registry); + context.loadAllAvailableDialects(); + }); + + py::class_(m, "type", py::module_local()) + .def("is_integer", + [](Type &self, unsigned width) { return self.isInteger(width); }) + .def("is_fp16", &Type::isF16) + .def("__eq__", + [](Type &self, py::object &other) { + Type *other_ty = py::cast(other); + return (other_ty != nullptr) && (*other_ty == self); + }) + .def("__ne__", + [](Type &self, py::object &other) { + Type *other_ty = py::cast(other); + return (other_ty == nullptr) || (*other_ty != self); + }) + .def("__str__", [](Type &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return os.str(); + }); + + py::class_(m, "function_type", py::module_local()) + .def("param_types", [](FunctionType &self) { + return std::vector(self.getInputs().begin(), + self.getInputs().end()); + }); + + py::class_(m, "location", py::module_local()) + .def("__str__", + [](Location &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return os.str(); + }) + .def("set_name", [](Location &self, std::string &name) { + mlir::StringAttr nameAttr = + mlir::StringAttr::get(self.getContext(), name); + mlir::NameLoc nameLoc = mlir::NameLoc::get(nameAttr, self); + self = dyn_cast(nameLoc); + }); + + py::class_(m, "value", py::module_local()) + .def(py::init<>()) + .def("set_attr", + [](Value &self, std::string &name, Attribute &attr) -> void { + if (Operation *definingOp = self.getDefiningOp()) + definingOp->setAttr(name, attr); + else { + auto arg = mlir::cast(self); + int id = arg.getArgNumber(); + std::string attrName = name + "_arg" + std::to_string(id); + Block *owner = arg.getOwner(); + if (owner->isEntryBlock() && + !isa(owner->getParentOp())) { + owner->getParentOp()->setAttr(attrName, attr); + } + } + }) + .def("get_context", &Value::getContext) + .def("get_loc", &Value::getLoc) + .def("set_loc", &Value::setLoc) + .def("replace_all_uses_with", + [](Value &self, Value &newValue) { + self.replaceAllUsesWith(newValue); + }) + .def("get_type", &Value::getType) + .def("id", + [](Value &self) { + // The Value is identified by and compared with + // other Values via the underlying ValueImpl + return (uint64_t)self.getImpl(); + }) + .def("set_loc", + [](Value &self, Location loc) { return self.setLoc(loc); }) + .def("get_loc", [](Value &self) { return self.getLoc(); }); + + py::class_(m, "op_result", py::module_local()); + + py::class_(m, "block_argument", py::module_local()) + .def("get_loc", &BlockArgument::getLoc) + .def("set_loc", &BlockArgument::setLoc); + + py::class_(m, "region", py::module_local()) + .def("get_parent_region", &Region::getParentRegion, ret::reference) + .def("size", [](Region &self) { return self.getBlocks().size(); }) + .def("empty", &Region::empty) + .def("id", [](Region &self) { return (uint64_t)&self; }) + .def("push_back", + [](Region &self, Block *block) { self.push_back(block); }) + .def("push_front", + [](Region &self, Block *block) { self.push_front(block); }); + + py::class_(m, "block", py::module_local()) + .def("arg", + [](Block &self, int index) -> BlockArgument { + if (index >= self.getNumArguments()) + throw pybind11::index_error("Block argument index out of range"); + return self.getArgument(index); + }) + .def("add_argument", + [](Block &self, Type ty) { + auto loc = UnknownLoc::get(ty.getContext()); + self.addArgument(ty, loc); + }) + .def("add_argument_at", [](Block &self, Type ty, + Location loc) { self.addArgument(ty, loc); }) + .def("get_num_arguments", &Block::getNumArguments) + .def("get_argument", &Block::getArgument) + .def("dump", &Block::dump) + .def("move_before", + [](Block &self, Block &dst) { self.moveBefore(&dst); }) + .def("insert_before", &Block::insertBefore) + .def("get_parent", &Block::getParent, ret::reference) + .def("merge_block_before", + [](Block &self, Block &dst) { + // ref: RewriterBase::mergeBlocks() + if (self.getNumArguments() != 0) + throw std::runtime_error( + "This block has arguments, don't merge"); + dst.getOperations().splice(dst.begin(), self.getOperations()); + self.dropAllUses(); + self.erase(); + }) + .def("replace_use_in_block_with", + [](Block &self, Value &v, Value &newVal) { + v.replaceUsesWithIf(newVal, [&](OpOperand &operand) { + Operation *user = operand.getOwner(); + Block *currentBlock = user->getBlock(); + while (currentBlock) { + if (currentBlock == &self) + return true; + // Move up one level + currentBlock = + currentBlock->getParent()->getParentOp()->getBlock(); + } + return false; + }); + }) + .def("__str__", + [](Block &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.print(os); + return str; + }) + .def("has_terminator", + [](Block &self) { + return !self.empty() && + self.back().hasTrait(); + }) + .def("has_return", + [](Block &self) { + return !self.empty() && + self.back().hasTrait(); + }) + .def("erase", [](Block &self) { self.erase(); }) + .def("id", [](Block &self) { return (uint64_t)&self; }); + + py::class_(m, "attribute", py::module_local()); + py::class_(m, "integer_attr", py::module_local()); + py::class_(m, "bool_attr", py::module_local()); + py::class_(m, "unit_attr", py::module_local()); + + // Ops + py::class_(m, "OpState", py::module_local()) + .def("set_attr", + [](OpState &self, std::string &name, Attribute &attr) -> void { + self->setAttr(name, attr); + }) + .def("get_num_results", + [](OpState &self) -> unsigned { return self->getNumResults(); }) + .def("get_result", + [](OpState &self, unsigned idx) -> Value { + if (idx >= self->getNumResults()) + throw pybind11::index_error("Op result index out of range"); + return self->getResult(idx); + }) + .def( + "get_region", + [](OpState &self, unsigned idx) -> Region & { + if (idx >= self->getNumRegions()) + throw pybind11::index_error("Op region index out of range"); + return self->getRegion(idx); + }, + ret::reference) + .def( + "get_body", + [](scf::ForOp &self, unsigned idx) -> Block * { + if (idx >= self->getNumRegions()) + throw pybind11::index_error("Op region index out of range"); + return self.getBody(idx); + }, + ret::reference) + .def("dump", [](OpState &self) { self->dump(); }) + .def("__str__", + [](OpState &self) -> std::string { + std::string str; + llvm::raw_string_ostream os(str); + auto printingFlags = getOpPrintingFlags(); + self->print(os, printingFlags); + return str; + }) + .def("str_nodebug", + [](OpState &self) -> std::string { + std::string str; + llvm::raw_string_ostream os(str); + self->print(os); + return str; + }) + .def("append_operand", + [](OpState &self, Value &val) { + self->insertOperands(self->getNumOperands(), val); + }) + .def("verify", + [](OpState &self) -> bool { + TritonSourceMgrDiagnosticHandler handler = + setupTritonDiagnosticHandler(self.getContext()); + return succeeded(verify(self.getOperation())); + }) + .def("get_operation", [](OpState &self) { return self.getOperation(); }); + + // scf Ops + py::class_(m, "ForOp", py::module_local()) + .def("get_induction_var", &scf::ForOp::getInductionVar); + + py::class_(m, "IfOp", py::module_local()) + .def("get_then_block", &scf::IfOp::thenBlock, ret::reference) + .def("get_else_block", &scf::IfOp::elseBlock, ret::reference) + .def("get_then_yield", &scf::IfOp::thenYield) + .def("get_else_yield", &scf::IfOp::elseYield); + py::class_(m, "YieldOp", py::module_local()); + py::class_(m, "WhileOp", py::module_local()) + .def("get_before", &scf::WhileOp::getBefore, ret::reference) + .def("get_after", &scf::WhileOp::getAfter, ret::reference); + py::class_(m, "ConditionOp", py::module_local()); + + py::class_>( + m, "operation", py::module_local()) + .def("get_name", + [](Operation &self) { + llvm::StringRef opName = self.getName().getStringRef(); + return opName.str(); + }) + .def("get_num_operands", &Operation::getNumOperands) + .def("get_operand", &Operation::getOperand) + .def("get_num_results", &Operation::getNumResults) + .def("get_result", &Operation::getResult) + .def("get_num_regions", &Operation::getNumRegions) + .def("get_region", &Operation::getRegion, ret::reference) + .def("get_block", &Operation::getBlock, ret::reference) + .def("get_str_attr", + [](Operation &self, const std::string &name) -> py::object { + auto ret = self.getAttrOfType(name); + if (!ret) + return py::none(); + return py::str(ret.getValue().str()); + }) + .def("get_bool_attr", + [](Operation &self, const std::string &name) -> py::object { + auto ret = self.getAttrOfType(name); + if (!ret) + return py::none(); + return py::bool_(ret.getValue()); + }) + .def("get_flat_symbol_ref_attr", + [](Operation &self, const std::string &name) -> py::object { + auto ret = self.getAttrOfType(name); + if (!ret) + return py::none(); + return py::str(ret.getValue().str()); + }); + + // dynamic_attr is used to transfer ownership of the MLIR context to the + // module + py::class_(m, "module", py::module_local(), + py::dynamic_attr()) + .def("dump", &ModuleOp::dump) + .def("str", + [](ModuleOp &self) -> std::string { + std::string str; + llvm::raw_string_ostream os(str); + auto printingFlags = getOpPrintingFlags(); + self.print(os, printingFlags); + return str; + }) + .def("push_back", + [](ModuleOp &self, FuncOp &funcOp) -> void { + self.push_back(funcOp); + }) + .def("get_entry_func_name", + [](ModuleOp &self) -> std::string { + for (auto &op : self.getOps()) { + if (auto func = dyn_cast(op)) { + if (triton::isKernel(func)) + return func.getName().str(); + } + } + return ""; + }) + .def("has_function", + [](ModuleOp &self, std::string &funcName) -> bool { + if (self.lookupSymbol(funcName)) + return true; + return false; + }) + .def("get_function", + [](ModuleOp &self, std::string &funcName) -> FuncOp { + return self.lookupSymbol(funcName); + }) + /* + * def ty_to_cpp(ty) is the consumer of this function. + * If the type is a ptr it expects ty[0] == '*', else the type itself. + */ + + .def("get_function_signature", + [](ModuleOp &self, FuncOp &func) -> std::vector { + std::vector strVec; + + auto type = func.getFunctionType(); + unsigned numArgs = type.getNumInputs(); + for (unsigned i = 0; i != numArgs; ++i) { + std::string tempType; + llvm::raw_string_ostream os(tempType); + + auto ty = type.getInput(i); + if (auto attributes = func.getCallableArgAttrs()) { + Attribute attr = attributes[i]; + // Check for tt.nv_tma_desc = 1 + if (auto dAttr = dyn_cast(attr)) { + if (dAttr.contains("tt.nv_tma_desc")) { + strVec.push_back("nvTmaDesc"); + continue; + } + } + } + if (auto ptrType = dyn_cast(ty)) { + auto pType = ptrType.getPointeeType(); + os << "*"; + pType.print(os); + } else { + ty.print(os); + } + strVec.push_back(tempType); + } + return strVec; + }) + .def("get_int_attr", + [](ModuleOp &self, std::string name) -> py::object { + auto ret = self->getAttrOfType(name); + if (!ret) + return py::none(); + return py::int_(ret.getInt()); + }) + .def("get_tensordesc_metadata", getTensorDescMetadata) + .def("create_location_snapshot", + [](ModuleOp &self, const std::string &fileName) -> void { + auto printingFlags = getOpPrintingFlags(); + if (failed(generateLocationsFromIR(fileName, self, printingFlags))) + throw std::runtime_error("Failed to create location snapshot"); + }) + .def("walk", + [](ModuleOp &self, const std::function &fn) { + self.walk(fn); + }); + + m.def("make_attr", [](const std::vector &values, MLIRContext &context) { + return mlir::cast(DenseIntElementsAttr::get( + RankedTensorType::get({static_cast(values.size())}, + IntegerType::get(&context, 32)), + values)); + }); + + m.def( + "parse_mlir_module", + [](const std::string &inputFilename, MLIRContext &context) { + // parse module + OwningOpRef module = + parseSourceFile(inputFilename, &context); + if (!module) + throw std::runtime_error("Parse MLIR file failed."); + return module->clone(); + }, + ret::take_ownership); + + py::class_(m, "function", py::module_local()) + // .def_property_readonly("attrs", &ir::function::attrs) + // .def("add_attr", &ir::function::add_attr); + .def("args", + [](FuncOp &self, unsigned idx) -> BlockArgument { + if (idx >= self.getNumArguments()) + throw pybind11::index_error( + "Function argument index out of range"); + return self.getArgument(idx); + }) + .def("get_num_args", &FuncOp::getNumArguments) + .def( + "add_entry_block", + [](FuncOp &self) -> Block * { return self.addEntryBlock(); }, + ret::reference) + .def( + "set_arg_attr", + [](FuncOp &self, int arg_no, const std::string &name, int val) { + if (arg_no >= self.getNumArguments()) + throw pybind11::index_error( + "Function argument index out of range"); + // set arg attributes "name" to value "val" + auto attrTy = IntegerType::get(self.getContext(), 32); + self.setArgAttr(arg_no, name, IntegerAttr::get(attrTy, val)); + }, + ret::reference) + // .def("has_attr", &::FuncOp::hasAttr) + .def("finalize", [](FuncOp &self) -> void {}) + .def_property_readonly("type", &FuncOp::getFunctionType) + .def("reset_type", &FuncOp::setType); + + py::class_(m, "op_builder", py::module_local(), + py::dynamic_attr()) + .def(py::init()); + + py::class_(m, "InsertPoint", py::module_local()); + + static py::class_ builderClass( + m, "builder", py::module_local(), py::dynamic_attr()); + builderClassPtr = &builderClass; + builderClass.def(py::init()) + .def("get_op_builder", &TritonOpBuilder::getBuilder, ret::reference) + // getters + .def("create_module", + [](TritonOpBuilder &self) -> ModuleOp { + return self.create(); + }) + // insertion block/point + .def("set_insertion_point_to_start", + [](TritonOpBuilder &self, Block &block) -> void { + self.setInsertionPointToStart(block); + }) + .def("set_insertion_point_to_end", + [](TritonOpBuilder &self, Block &block) { + self.setInsertionPointToEnd(block); + }) + .def("set_insertion_point_after", + [](TritonOpBuilder &self, Operation &op) { + self.setInsertionPointAfter(op); + }) + .def( + "get_insertion_block", + [](TritonOpBuilder &self) -> Block * { + return self.getBuilder().getInsertionBlock(); + }, + ret::reference) + .def("get_insertion_point", + [](TritonOpBuilder &self) { + return self.getBuilder().saveInsertionPoint(); + }) + .def("restore_insertion_point", + [](TritonOpBuilder &self, OpBuilder::InsertPoint pt) { + self.restoreInsertionPoint(pt); + }) + // Attr + .def( + "get_unit_attr", + [](TritonOpBuilder &self) { return self.getBuilder().getUnitAttr(); }) + .def("get_bool_attr", + [](TritonOpBuilder &self, bool value) { + return self.getBuilder().getBoolAttr(value); + }) + .def("get_int32_attr", + [](TritonOpBuilder &self, int32_t value) { + return self.getBuilder().getI32IntegerAttr(value); + }) + .def("get_string_attr", + [](TritonOpBuilder &self, std::string value) -> Attribute { + return self.getBuilder().getStringAttr(value); + }) + .def("get_disable_loop_licm_attr", + [](TritonOpBuilder &self) -> Attribute { + auto licmAttr = + LLVM::LoopLICMAttr::get(self.getBuilder().getContext(), + self.getBuilder().getBoolAttr(true), + self.getBuilder().getBoolAttr(true)); + mlir::LLVM::LoopAnnotationAttr la = + mlir::LLVM::LoopAnnotationAttr::get( + self.getBuilder().getContext(), {}, {}, {}, {}, {}, + licmAttr, {}, {}, {}, {}, {}, {}, {}, {}, {}); + return la; + }) + // Use arith.ConstantOp to create constants + // Constants + .def("get_int1", + [](TritonOpBuilder &self, bool v) -> Value { + return Value(self.create( + self.getBuilder().getI1Type(), v)); + }) + .def("get_int8", + [](TritonOpBuilder &self, int64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI8Type(), v)); + }) + .def("get_int16", + [](TritonOpBuilder &self, int64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI16Type(), v)); + }) + .def("get_int32", + [](TritonOpBuilder &self, int64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI32Type(), v)); + }) + .def("get_int64", + [](TritonOpBuilder &self, int64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI64Type(), v)); + }) + .def("get_uint8", + [](TritonOpBuilder &self, uint64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI8Type(), v)); + }) + .def("get_uint16", + [](TritonOpBuilder &self, uint64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI16Type(), v)); + }) + .def("get_uint32", + [](TritonOpBuilder &self, uint64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI32Type(), v)); + }) + .def("get_uint64", + [](TritonOpBuilder &self, uint64_t v) -> Value { + return Value(self.create( + self.getBuilder().getI64Type(), v)); + }) + .def("get_bf16", + [](TritonOpBuilder &self, float v) -> Value { + auto type = self.getBuilder().getBF16Type(); + return self.create( + type, APFloat(type.getFloatSemantics(), std::to_string(v))); + }) + .def("get_fp16", + [](TritonOpBuilder &self, float v) -> Value { + return self.create( + self.getBuilder().getF16FloatAttr(v)); + }) + .def("get_fp32", + [](TritonOpBuilder &self, float v) -> Value { + return self.create( + self.getBuilder().getF32FloatAttr(v)); + }) + .def("get_fp64", + [](TritonOpBuilder &self, double v) -> Value { + return self.create( + self.getBuilder().getF64FloatAttr(v)); + }) + .def("get_null_value", + [](TritonOpBuilder &self, Type type) -> Value { + if (auto floatTy = dyn_cast(type)) + return self.create( + floatTy, APFloat(floatTy.getFloatSemantics(), 0)); + else if (auto intTy = dyn_cast(type)) + return self.create(intTy, 0); + else + throw std::runtime_error("Not implemented"); + }) + .def("get_all_ones_value", + [](TritonOpBuilder &self, Type type) -> Value { + uint64_t val = 0xFFFFFFFFFFFFFFFF; + if (auto intTy = dyn_cast(type)) + return self.create(intTy, val); + else + throw std::runtime_error("Not implemented"); + }) + + // Types + .def("get_void_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getNoneType(); + }) + .def("get_int1_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getI1Type(); + }) // or ret::copy? + .def("get_int8_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getI8Type(); + }) + .def("get_int16_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getType(16); + }) + .def("get_int32_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getI32Type(); + }) + .def("get_int64_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getI64Type(); + }) + .def("get_fp8e4nv_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getType(); + }) + .def("get_fp8e4b8_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getType(); + }) + .def("get_fp8e4b15_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getI8Type(); + }) + .def("get_fp8e5_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getType(); + }) + .def("get_fp8e5b16_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getType(); + }) + .def("get_half_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getF16Type(); + }) + .def("get_bf16_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getBF16Type(); + }) + .def("get_float_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getF32Type(); + }) + .def("get_double_ty", + [](TritonOpBuilder &self) -> Type { + return self.getBuilder().getF64Type(); + }) + .def("get_ptr_ty", + [](TritonOpBuilder &self, Type &type, int addrSpace) -> Type { + return PointerType::get(type, addrSpace); + }) + .def("get_block_ty", + [](TritonOpBuilder &self, Type &elementType, + std::vector &shape) -> Type { + return RankedTensorType::get(shape, elementType); + }) + .def("get_function_ty", + [](TritonOpBuilder &self, std::vector inTypes, + std::vector outTypes) -> Type { + return self.getBuilder().getFunctionType(inTypes, outTypes); + }) + // locs + .def("set_loc", + [](TritonOpBuilder &self, Location loc) { self.setLastLoc(loc); }) + .def("set_loc", + [](TritonOpBuilder &self, std::string name) { + auto nameAttr = StringAttr::get(self.getContext(), name); + auto loc = NameLoc::get(nameAttr); + self.setLastLoc(loc); + }) + .def("create_loc", + [](TritonOpBuilder &self, const std::string &fileName, int line, + int column) -> Location { + return mlir::FileLineColLoc::get(self.getContext(), fileName, line, + column); + }) + .def( + "create_name_loc", + [](TritonOpBuilder &self, std::string name, + std::optional childLoc) -> Location { + auto nameAttr = StringAttr::get(self.getContext(), name); + if (childLoc) + return NameLoc::get(nameAttr, *childLoc); + return NameLoc::get(nameAttr); + }, + py::arg("name"), py::arg("child_loc") = py::none()) + .def("set_loc", + [](TritonOpBuilder &self, const std::string &fileName, int line, + int column) { self.setLastLoc(fileName, line, column); }) + .def("get_loc", + [](TritonOpBuilder &self) -> Location { return self.getLastLoc(); }) + + // Ops + .def("get_or_insert_function", + [](TritonOpBuilder &self, ModuleOp &module, std::string &funcName, + Type &funcType, std::string &visibility, + bool noinline) -> FuncOp { + if (Operation *funcOperation = module.lookupSymbol(funcName)) + return llvm::dyn_cast(funcOperation); + if (auto funcTy = dyn_cast(funcType)) { + llvm::SmallVector attrs = { + NamedAttribute( + self.getBuilder().getStringAttr("sym_visibility"), + self.getBuilder().getStringAttr(visibility)), + NamedAttribute(self.getBuilder().getStringAttr("noinline"), + self.getBuilder().getBoolAttr(noinline))}; + return self.create(funcName, funcTy, attrs); + } + throw std::invalid_argument("invalid function type"); + }) + .def( + "create_block", + [](TritonOpBuilder &self) -> Block * { + Region *parent = self.getBuilder().getBlock()->getParent(); + return self.getBuilder().createBlock(parent); + }, + ret::reference) + .def( + "create_block_with_parent", + [](TritonOpBuilder &self, Region &parent, + std::vector &argTypes) -> Block * { + // TODO: update arg loc + auto loc = self.getBuilder().getUnknownLoc(); + llvm::SmallVector argLocs(argTypes.size(), loc); + return self.getBuilder().createBlock(&parent, {}, argTypes, + argLocs); + }, + ret::reference) + .def( + "new_block", + [](TritonOpBuilder &self) -> Block * { return new Block(); }, + ret::reference) + // Function + .def("ret", + [](TritonOpBuilder &self, std::vector &vals) -> OpState { + return self.create(vals); + }) + .def("call", + [](TritonOpBuilder &self, FuncOp &func, std::vector &args) + -> OpState { return self.create(func, args); }) + // Unstructured control flow + .def("create_cond_branch", + [](TritonOpBuilder &self, Value condition, Block *trueDest, + Block *falseDest) -> OpState { + return self.create(condition, trueDest, + falseDest); + }) + .def("create_branch", + [](TritonOpBuilder &self, Block *dest, std::vector &args) + -> OpState { return self.create(dest, args); }) + // Structured control flow + .def("create_for_op", + [](TritonOpBuilder &self, Value &lb, Value &ub, Value &step, + std::vector &initArgs) -> scf::ForOp { + return self.create(lb, ub, step, initArgs); + }) + .def("create_if_op", + [](TritonOpBuilder &self, std::vector &retTypes, + Value &condition, bool withElse) -> scf::IfOp { + return self.create(retTypes, condition, withElse); + }) + .def("create_yield_op", + [](TritonOpBuilder &self, std::vector &yields) + -> scf::YieldOp { return self.create(yields); }) + .def("create_while_op", + [](TritonOpBuilder &self, std::vector &retTypes, + std::vector &initArgs) -> scf::WhileOp { + return self.create(retTypes, initArgs); + }) + .def("create_condition_op", + [](TritonOpBuilder &self, Value &cond, + std::vector &args) -> scf::ConditionOp { + return self.create(cond, args); + }) + + // miscellaneous + .def("create_make_range", + [](TritonOpBuilder &self, Type retTy, int start, int end) -> Value { + return self.create(retTy, start, end); + }) + + // Cast instructions + // Conversions for custom FP types (FP8 and non-standard rounding modes) + .def("create_fp_to_fp", + [](TritonOpBuilder &self, Value &src, Type &dstType, + std::optional roundingMode) -> Value { + if (roundingMode.has_value()) + return self.create( + dstType, src, + RoundingModeAttr::get(self.getBuilder().getContext(), + roundingMode.value())); + else + return self.create(dstType, src); + }) + // Conversions for standard LLVM builtin types + .def("create_bitcast", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_si_to_fp", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_ui_to_fp", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_fp_to_si", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_fp_to_ui", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_fp_ext", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_fp_trunc", + [](TritonOpBuilder &self, Value &src, Type &dstType) -> Value { + return self.create(dstType, src); + }) + .def("create_int_cast", + [](TritonOpBuilder &self, Value &src, Type &dstType, + bool isSigned) -> Value { + // get element type if necessary + Type srcType = src.getType(); + auto srcTensorType = dyn_cast(srcType); + auto dstTensorType = dyn_cast(dstType); + Type srcEltType = srcType; + Type dstEltType = dstType; + if (dstTensorType && srcTensorType) { + dstEltType = dstTensorType.getElementType(); + srcEltType = srcTensorType.getElementType(); + } + unsigned srcWidth = srcEltType.getIntOrFloatBitWidth(); + unsigned dstWidth = dstEltType.getIntOrFloatBitWidth(); + if (srcWidth == dstWidth) + return self.create(dstType, src); + else if (srcWidth > dstWidth) + return self.create(dstType, src); + else if (isSigned) + return self.create(dstType, src); + else + return self.create(dstType, src); + }) + .def("create_fmul", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_fdiv", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_frem", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_fadd", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_fsub", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_mul", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_umulhi", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_sdiv", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_udiv", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_srem", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_urem", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_add", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_sub", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_fma", + [](TritonOpBuilder &self, Value &a, Value &b, Value &c) -> Value { + return Value(self.create(a, b, c)); + }) + .def("create_shl", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_lshr", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_ashr", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_minsi", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_minui", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + // minimumf follows the torch.minimum convention and returns NaN if either + // operand is NaN + .def("create_minimumf", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + // minnumf follows the torch.fmin convention and returns the non-NaN + // operand + .def("create_minnumf", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_maxsi", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_maxui", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + // maximumf follows the torch.maximum convention and returns NaN if either + // operand is NaN + .def("create_maximumf", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + // maxnumf follows the torch.fmax convention and returns the non-NaN + // operand + .def("create_maxnumf", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + .def("create_clampf", + [](TritonOpBuilder &self, Value &input, Value &min, Value &max, + PropagateNan propagateNan) -> Value { + return Value(self.create(input, min, max, propagateNan)); + }) + .def("create_precise_sqrt", + [](TritonOpBuilder &self, Value &input) -> Value { + return Value(self.create(input)); + }) + .def("create_precise_divf", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return Value(self.create(lhs, rhs)); + }) + // AddPtr (similar to GEP) + .def("create_addptr", + [](TritonOpBuilder &self, Value &ptr, Value &offset) -> Value { + return self.create(ptr.getType(), ptr, offset); + }) + // Comparison (int) + .def("create_icmpSLE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::sle, lhs, + rhs); + }) + .def("create_icmpSLT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::slt, lhs, + rhs); + }) + .def("create_icmpSGE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::sge, lhs, + rhs); + }) + .def("create_icmpSGT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::sgt, lhs, + rhs); + }) + .def("create_icmpULE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::ule, lhs, + rhs); + }) + .def("create_icmpULT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::ult, lhs, + rhs); + }) + .def("create_icmpUGE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::uge, lhs, + rhs); + }) + .def("create_icmpUGT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::ugt, lhs, + rhs); + }) + .def("create_icmpEQ", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::eq, lhs, + rhs); + }) + .def("create_icmpNE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpIPredicate::ne, lhs, + rhs); + }) + // Comparison (float) + .def("create_fcmpOLT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::OLT, lhs, + rhs); + }) + .def("create_fcmpOGT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::OGT, lhs, + rhs); + }) + .def("create_fcmpOLE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::OLE, lhs, + rhs); + }) + .def("create_fcmpOGE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::OGE, lhs, + rhs); + }) + .def("create_fcmpOEQ", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::OEQ, lhs, + rhs); + }) + .def("create_fcmpONE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::ONE, lhs, + rhs); + }) + .def("create_fcmpULT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::ULT, lhs, + rhs); + }) + .def("create_fcmpUGT", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::UGT, lhs, + rhs); + }) + .def("create_fcmpULE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::ULE, lhs, + rhs); + }) + .def("create_fcmpUGE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::UGE, lhs, + rhs); + }) + .def("create_fcmpUEQ", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::UEQ, lhs, + rhs); + }) + .def("create_fcmpUNE", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(arith::CmpFPredicate::UNE, lhs, + rhs); + }) + // // Logical + .def("create_and", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_xor", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + .def("create_or", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + return self.create(lhs, rhs); + }) + // Input/Output + .def("create_load", + [](TritonOpBuilder &self, Value &ptrs, CacheModifier cacheModifier, + EvictionPolicy evictionPolicy, bool isVolatile, + std::optional flagtree_hints) -> Value { + // flagtree hints + auto hintsAttr = + flagtree_hints + ? mlir::StringAttr::get(self.getContext(), *flagtree_hints) + : mlir::StringAttr::get(self.getContext(), ""); + return self.create(ptrs, cacheModifier, evictionPolicy, + isVolatile, hintsAttr); + }) + .def("create_store", + [](TritonOpBuilder &self, Value &ptrs, Value &value, + CacheModifier cacheModifier, + EvictionPolicy evictionPolicy) -> void { + self.create(ptrs, value, cacheModifier, evictionPolicy); + }) + .def("create_tensor_pointer_load", + [](TritonOpBuilder &self, Value &ptr, + std::vector &boundaryCheck, + std::optional paddingOption, + CacheModifier cacheModifier, EvictionPolicy evictionPolicy, + bool isVolatile, + std::optional flagtree_hints) -> Value { + // flagtree hints + auto hintsAttr = + flagtree_hints + ? mlir::StringAttr::get(self.getContext(), *flagtree_hints) + : mlir::StringAttr::get(self.getContext(), ""); + return self.create(ptr, boundaryCheck, paddingOption, + cacheModifier, evictionPolicy, + isVolatile, hintsAttr); + }) + .def("create_tensor_pointer_store", + [](TritonOpBuilder &self, Value &ptr, Value &val, + std::vector &boundaryCheck, CacheModifier cacheModifier, + EvictionPolicy evictionPolicy) -> void { + self.create(ptr, val, boundaryCheck, cacheModifier, + evictionPolicy); + }) + .def("create_masked_load", + [](TritonOpBuilder &self, Value &ptrs, Value &mask, + std::optional &other, CacheModifier cacheModifier, + EvictionPolicy evictionPolicy, bool isVolatile, + std::optional flagtree_hints) -> Value { + // flagtree hints + auto hintsAttr = + flagtree_hints + ? mlir::StringAttr::get(self.getContext(), *flagtree_hints) + : mlir::StringAttr::get(self.getContext(), ""); + return self.create(ptrs, mask, other.value_or(Value()), + cacheModifier, evictionPolicy, + isVolatile, hintsAttr); + }) + .def("create_masked_store", + [](TritonOpBuilder &self, Value &ptrs, Value &val, Value &mask, + CacheModifier cacheModifier, + EvictionPolicy evictionPolicy) -> void { + self.create(ptrs, val, mask, cacheModifier, + evictionPolicy); + }) + .def("create_tensor_descriptor_type", + [](TritonOpBuilder &self, Type blockTy, bool isSigned) -> Type { + auto ctx = self.getContext(); + return triton::TensorDescType::get( + ctx, cast(blockTy), isSigned); + }) + .def("create_descriptor_load", + [](TritonOpBuilder &self, Value desc, std::vector &indices, + CacheModifier cacheModifier, + EvictionPolicy evictionPolicy) -> Value { + auto descTy = cast(desc.getType()); + auto resTy = descTy.getSignlessBlockType(); + return self.create( + resTy, desc, indices, cacheModifier, evictionPolicy); + }) + .def("create_descriptor_gather", + [](TritonOpBuilder &self, Value desc, Value x_indices, Value y_index, + Type type) -> Value { + return self.create(type, desc, x_indices, + y_index); + }) + .def("create_descriptor_store", + [](TritonOpBuilder &self, Value desc, Value value, + std::vector &indices) -> void { + self.create(desc, value, indices); + }) + .def("create_descriptor_reduce", + [](TritonOpBuilder &self, DescriptorReduceKind kind, Value desc, + Value value, std::vector &indices) -> void { + self.create(kind, desc, value, indices); + }) + .def("create_descriptor_scatter", + [](TritonOpBuilder &self, Value desc, Value value, Value x_indices, + Value y_index) -> void { + self.create(desc, x_indices, y_index, value); + }) + .def("create_reshape", + [](TritonOpBuilder &self, Value &arg, std::vector &shape, + bool allowReorder) -> Value { + return self.create(shape, arg, allowReorder); + }) + .def("create_expand_dims", + [](TritonOpBuilder &self, Value &arg, int axis) -> Value { + return self.create(arg, axis); + }) + .def("create_cat", + [](TritonOpBuilder &self, Value &lhs, Value &rhs) -> Value { + auto lhsType = dyn_cast(lhs.getType()); + auto rhsType = dyn_cast(rhs.getType()); + if (!(lhsType.getShape().size() == 1 && + rhsType.getShape().size() == 1)) + throw std::invalid_argument( + "shape not supported by cat. Expecting rank-1 inputs"); + std::vector shape{lhsType.getShape()[0] + + rhsType.getShape()[0]}; + return self.create(lhsType.clone(shape), lhs, rhs); + }) + .def("create_join", + [](TritonOpBuilder &self, Value &a, Value &b) -> Value { + return self.create(a, b); + }) + .def("create_split", + [](TritonOpBuilder &self, Value &a) -> std::vector { + auto op = self.create(a); + return std::vector(op->result_begin(), op->result_end()); + }) + // Implements tl.trans and tl.permute. + .def("create_trans", + [](TritonOpBuilder &self, Value &arg, std::vector &order) + -> Value { return self.create(arg, order); }) + .def("create_broadcast", + [](TritonOpBuilder &self, Value &arg, + std::vector &shape) -> Value { + if (auto argType = dyn_cast(arg.getType())) + return self.createOrFold(argType.clone(shape), arg); + throw std::invalid_argument( + "arg is not of RankedTensorType, use create_splat"); + }) + .def("create_splat", + [](TritonOpBuilder &self, Type &retTy, Value &arg) -> Value { + return self.createOrFold(retTy, arg); + }) + .def("create_unsplat", + [](TritonOpBuilder &self, Value &arg) -> Value { + return self.createOrFold(arg); + }) + // // atomic + .def("create_atomic_cas", + [](TritonOpBuilder &self, Value &ptr, Value &cmp, Value &val, + MemSemantic sem, MemSyncScope scope) -> Value { + Type dstType; + if (auto srcTensorType = + dyn_cast(ptr.getType())) { + Type dstElemType = + cast(srcTensorType.getElementType()) + .getPointeeType(); + dstType = srcTensorType.clone(dstElemType); + } else { + auto ptrType = cast(getElementTypeOrSelf(ptr)); + dstType = ptrType.getPointeeType(); + } + return self.create(dstType, ptr, cmp, val, sem, + scope); + }) + .def("create_atomic_rmw", + [](TritonOpBuilder &self, RMWOp rmwOp, Value &ptr, Value &val, + Value &mask, MemSemantic sem, MemSyncScope scope) -> Value { + Type dstType; + if (auto srcTensorType = + dyn_cast(ptr.getType())) { + Type dstElemType = + cast(srcTensorType.getElementType()) + .getPointeeType(); + dstType = srcTensorType.clone(dstElemType); + } else { + auto ptrType = cast(getElementTypeOrSelf(ptr)); + dstType = ptrType.getPointeeType(); + } + return self.create(dstType, rmwOp, ptr, val, mask, + sem, scope); + }) + // External + .def("create_extern_elementwise", + [](TritonOpBuilder &self, const std::string &libName, + const std::string &libPath, const std::string &symbol, + std::vector &argList, Type retType, bool isPure) -> Value { + return self.create(retType, argList, libName, + libPath, symbol, isPure); + }) + // Built-in instruction + .def("create_get_program_id", + [](TritonOpBuilder &self, int axis) -> Value { + if (axis < 0 || axis > 3) + throw pybind11::index_error("program_id must be in [0,3]"); + return self.create(axis); + }) + .def("create_get_num_programs", + [](TritonOpBuilder &self, int axis) -> Value { + if (axis < 0 || axis > 3) + throw pybind11::index_error("program_id must be in [0,3]"); + return self.create(axis); + }) + .def("create_dot", + [](TritonOpBuilder &self, mlir::Value &a, mlir::Value &b, + mlir::Value &c, InputPrecision inputPrecision, + int maxNumImpreciseAcc) -> mlir::Value { + return self.create(c.getType(), a, b, c, inputPrecision, + maxNumImpreciseAcc); + }) + .def("create_dot_scaled", + [](TritonOpBuilder &self, mlir::Value &lhs, + std::optional &lhs_scale, + ScaleDotElemType lhs_format, mlir::Value &rhs, + std::optional &rhs_scale, + ScaleDotElemType rhs_format, bool fast_math, bool lhs_k_pack, + bool rhs_k_pack, mlir::Value &c) -> mlir::Value { + return self.create( + c.getType(), lhs, rhs, c, lhs_scale.value_or(Value()), + rhs_scale.value_or(Value()), lhs_format, rhs_format, fast_math, + lhs_k_pack, rhs_k_pack); + }) + .def("create_floor", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_ceil", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_exp", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_exp2", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_cos", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_sin", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_log", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_log2", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_erf", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_sqrt", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_rsqrt", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_fabs", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_iabs", + [](TritonOpBuilder &self, Value &val) -> Value { + return self.create(val); + }) + .def("create_reduce", + [](TritonOpBuilder &self, std::vector operands, int axis) + -> OpState { return self.create(operands, axis); }) + .def("create_reduce_ret", + [](TritonOpBuilder &self, py::args args) -> OpState { + llvm::SmallVector return_values; + for (const auto &arg : args) { + return_values.push_back(py::cast(arg)); + } + return self.create(return_values); + }) + .def("create_scan", + [](TritonOpBuilder &self, std::vector operands, int axis, + bool reverse) -> OpState { + return self.create(operands, axis, reverse); + }) + .def("create_scan_ret", + [](TritonOpBuilder &self, py::args args) -> OpState { + llvm::SmallVector return_values; + for (const auto &arg : args) { + return_values.push_back(py::cast(arg)); + } + return self.create(return_values); + }) + .def("create_map_elementwise", + [](TritonOpBuilder &self, std::vector inputs, + std::vector returnTys, int pack) -> OpState { + return self.create(returnTys, inputs, pack); + }) + .def("create_map_elementwise_ret", + [](TritonOpBuilder &self, std::vector returnVals) -> OpState { + return self.create(returnVals); + }) + .def("create_ptr_to_int", + [](TritonOpBuilder &self, Value &val, Type &type) -> Value { + return self.create(type, val); + }) + .def("create_int_to_ptr", + [](TritonOpBuilder &self, Value &val, Type &type) -> Value { + return self.create(type, val); + }) + .def("create_select", + [](TritonOpBuilder &self, Value &condition, Value &trueValue, + Value &falseValue) -> Value { + return self.create(condition, trueValue, + falseValue); + }) + .def("create_inline_asm", + [](TritonOpBuilder &self, const std::string &inlineAsm, + const std::string &constraints, const std::vector &values, + const std::vector &types, bool isPure, + int pack) -> OpState { + return self.create( + types, inlineAsm, constraints, isPure, pack, values); + }) + .def("create_print", + [](TritonOpBuilder &self, const std::string &prefix, bool hex, + const std::vector &values, + const std::vector &isSigned) -> void { + auto prefixAttr = StringAttr::get(self.getBuilder().getContext(), + llvm::StringRef(prefix)); + self.create(prefixAttr, hex, values, isSigned); + }) + .def("create_assert", + [](TritonOpBuilder &self, Value &condition, + const std::string &message) -> void { + auto messageAttr = StringAttr::get(self.getBuilder().getContext(), + llvm::StringRef(message)); + self.create(condition, messageAttr); + }) + .def("create_assume", + [](TritonOpBuilder &self, Value &condition) { + self.create(condition); + }) + .def("create_poison", + [](TritonOpBuilder &self, Type &type) -> Value { + return self.create(type); + }) + .def("create_histogram", + [](TritonOpBuilder &self, Value operand, int numBins, + std::optional mask) -> Value { + if (!mask) { + return self.create( + RankedTensorType::get( + {static_cast(numBins)}, + IntegerType::get(operand.getContext(), 32)), + operand); + } else { + return self.create( + RankedTensorType::get( + {static_cast(numBins)}, + IntegerType::get(operand.getContext(), 32)), + operand, *mask); + } + }) + .def("create_gather", + [](TritonOpBuilder &self, Value src, Value indices, int axis) + -> Value { return self.create(src, indices, axis); }) + // Force GPU barrier + .def("create_barrier", + [](TritonOpBuilder &self) { self.create(); }) + // Make a block pointer (tensor pointer in Triton IR) + .def("create_make_block_ptr", + [](TritonOpBuilder &self, Value &base, std::vector &shape, + std::vector &strides, std::vector &offsets, + std::vector &tensorShape, + std::vector &order) -> Value { + return self.create(base, shape, strides, offsets, + tensorShape, order); + }) + // Advance a block pointer + .def("create_advance", + [](TritonOpBuilder &self, Value &ptr, + std::vector &offsets) -> Value { + return self.create(ptr.getType(), ptr, offsets); + }) + // Make a tensor descriptor + .def("create_make_tensor_descriptor", + [](TritonOpBuilder &self, Value &base, std::vector &shape, + std::vector &strides, std::vector &tensorShape, + bool isSignedInteger, PaddingOption paddingOption) -> Value { + return self.create(base, shape, strides, + tensorShape, isSignedInteger, + paddingOption); + }) + + ; + + static unsigned g_passGroupNumber=1; + py::class_(m, "pass_manager", py::module_local()) + .def(py::init()) + .def( "get_dump_dir_name", [](PassManager &self) -> std::string { + return std::to_string(g_passGroupNumber++); + }) + .def("enable_debug", + [](PassManager &self) -> bool { + auto *context = self.getContext(); + bool haveDump = ::triton::tools::getBoolEnv("MLIR_ENABLE_DUMP"); + std::string funcToDump; + if (!haveDump) { + funcToDump = triton::tools::getStrEnv("MLIR_ENABLE_DUMP"); + bool isEnvValueBool = + triton::tools::isEnvValueBool(funcToDump).has_value(); + if (!funcToDump.empty() && !isEnvValueBool) + haveDump = true; + } + if (haveDump) { + context->disableMultithreading(); + auto printingFlags = getOpPrintingFlags(); + auto printAlways = [funcToDump](Pass *, Operation *op) -> bool { + if (funcToDump.empty()) + return true; + if (auto mod = dyn_cast(op)) { + return mod.lookupSymbol(funcToDump); + } + if (auto func = dyn_cast(op)) { + return SymbolTable::getSymbolName(func).getValue() == + funcToDump; + } + + return false; + }; + if(std::string("1") == getenv("MLIR_DUMP_FILE_TREE")) { + std::string dumpDirName = std::to_string(g_passGroupNumber++); + auto returnTrue = [](Pass *, Operation *) { return true; }; + self.enableIRPrintingToFileTree( + /*shouldPrintBeforePass=*/returnTrue, + /*shouldPrintAfterPass=*/nullptr, + /*printModuleScope=*/true, + /*printAfterOnlyOnChange=*/true, + /*printAfterOnlyOnFailure*/ false, + /*printTreeDir*/ dumpDirName, + printingFlags); + } + else { + self.enableIRPrinting( + /*shouldPrintBeforePass=*/printAlways, + /*shouldPrintAfterPass=*/printAlways, + /*printModuleScope=*/true, + /*printAfterOnlyOnChange=*/ true, + /*printAfterOnlyOnFailure*/ true, mlir_dumps_or_dbgs(), + printingFlags); + } + } + return haveDump; + }) + .def("get_pipeline_str", + [](PassManager &self) { + std::string str; + llvm::raw_string_ostream os(str); + self.printAsTextualPipeline(os); + return str; + }) + .def( + "run", + [](PassManager &self, ModuleOp &mod, std::string repro_pipeline_tag) { + // TODO: maybe dump module to file and print error for better + // diagnostics + + auto *context = mod.getContext(); + if (::triton::tools::getBoolEnv("MLIR_DISABLE_MULTITHREADING")) + context->disableMultithreading(); + + auto reproducerPath = + triton::tools::getStrEnv("TRITON_REPRODUCER_PATH"); + if (!reproducerPath.empty()) { + if (reproducerPath != "-") { + std::string repro_suffix = + "." + repro_pipeline_tag + ".repro.mlir"; + reproducerPath += repro_suffix; + } + auto anchorName = self.getOpAnchorName(); + auto passes = self.getPasses(); + Operation *op = mod.getOperation(); + // Save a reproducer for the current pass manager invocation + // immediately. + makeReproducer(anchorName, passes, op, reproducerPath); + // But if the pass manager crashes, attempt to generate a local + // reproducer instead. + context->disableMultithreading(); + self.enableCrashReproducerGeneration(reproducerPath, + /*genLocalReproducer=*/true); + } else { + self.enableCrashReproducerGeneration(makeConsoleReproducer()); + } + + if (triton::tools::getBoolEnv("TRITON_ENABLE_LLVM_DEBUG")) { + ::llvm::DebugFlag = true; + } + + if (auto debugOnly = + triton::tools::getStrEnv("TRITON_LLVM_DEBUG_ONLY"); + !debugOnly.empty()) { + llvm::SmallVector storage; + llvm::SmallVector debugTypes = + parseCommaSeparatedValues(debugOnly, storage); + ::llvm::DebugFlag = true; + using namespace llvm; + setCurrentDebugTypes(debugTypes.data(), debugTypes.size()); + } + + bool haveTiming = ::triton::tools::getBoolEnv("MLIR_ENABLE_TIMING"); + if (haveTiming) { + self.enableTiming(); + } + + TritonSourceMgrDiagnosticHandler diagHandler = + setupTritonDiagnosticHandler(context); + if (failed(self.run(mod.getOperation()))) + throw std::runtime_error("PassManager::run failed"); + }, + py::call_guard()); +} + +bool str_eq_ignore_case(const char *s1, const char *s2, int n) { + for (int i = 0; i < n; ++i) { + if (tolower(s1[i]) != s2[i]) + return false; + } + return true; +} + +int strlen_max(const char *str, int max) { + for (int i = 0; i <= max; ++i) { + if (str[i] == '\0') { + return i; + } + } + return 0; +} + +bool is_truthy(char *str) { + int len = strlen_max(str, 4); + switch (len) { + case 1: + return str[0] == '1' || tolower(str[0]) == 'y'; + case 2: + return str_eq_ignore_case(str, "on", len); + case 3: + return str_eq_ignore_case(str, "yes", len); + case 4: + return str_eq_ignore_case(str, "true", len); + default: + return false; + } +} + +PyObject *py_getenv(PyObject *self, PyObject *const *args, Py_ssize_t nargs) { + if (!(nargs == 1 || nargs == 2)) { + PyErr_SetString(PyExc_TypeError, "getenv expected 1 or 2 arguments"); + return NULL; + } + PyObject *name = args[0]; + PyObject *default_val = nargs == 2 ? args[1] : Py_None; + if (!PyUnicode_CheckExact(name)) { + PyErr_SetString(PyExc_TypeError, "name must be a string"); + return NULL; + } + char *env_val = getenv(PyUnicode_AsUTF8(name)); + if (!env_val) { + Py_INCREF(default_val); + return default_val; + } + return PyUnicode_FromString(env_val); +} + +PyObject *py_getenv_bool(PyObject *self, PyObject *const *args, + Py_ssize_t nargs) { + if (nargs != 2) { + PyErr_SetString(PyExc_TypeError, "getenv_bool expected 2 arguments"); + return NULL; + } + PyObject *name = args[0]; + PyObject *default_val = args[1]; + if (!PyUnicode_CheckExact(name)) { + PyErr_SetString(PyExc_TypeError, "name must be a string"); + return NULL; + } + char *env_val = getenv(PyUnicode_AsUTF8(name)); + PyObject *res = default_val; + if (env_val) { + res = is_truthy(env_val) ? Py_True : Py_False; + } + Py_INCREF(res); + return res; +} + +static PyMethodDef ModuleMethods[] = { + {"getenv", (PyCFunction)py_getenv, METH_FASTCALL, NULL}, + {"getenv_bool", (PyCFunction)py_getenv_bool, METH_FASTCALL, NULL}, + {NULL, NULL, 0, NULL} // sentinel +}; + +void init_triton_env_vars(py::module &m) { + m.def("get_cache_invalidating_env_vars", + []() -> std::map { + std::map ret; + for (const auto &envVar : CACHE_INVALIDATING_ENV_VARS) { + auto strVal = triton::tools::getStrEnv(envVar); + if (strVal.empty()) + continue; + auto boolV = triton::tools::isEnvValueBool(strVal); + if (boolV.has_value()) + ret[envVar] = boolV.value() ? "true" : "false"; + else + ret[envVar] = strVal; + } + return ret; + }); + PyModule_AddFunctions(m.ptr(), ModuleMethods); +} diff --git a/third_party/sunrise/python/src/ir.h b/third_party/sunrise/python/src/ir.h new file mode 100644 index 0000000000..b36c3cd5d8 --- /dev/null +++ b/third_party/sunrise/python/src/ir.h @@ -0,0 +1,113 @@ +#pragma once +#include "mlir/IR/Builders.h" +#include "triton/Tools/Sys/GetEnv.hpp" +#include + +#ifdef __TLE__ +#include +#include +#include +namespace py = pybind11; +#endif + +// A custom op builder that keeps track of the last location +class TritonOpBuilder { +public: + TritonOpBuilder(mlir::MLIRContext *context) { + builder = std::make_unique(context); + lastLoc = std::make_unique(builder->getUnknownLoc()); + } + + mlir::OpBuilder &getBuilder() { return *builder; } + mlir::MLIRContext *getContext() { return builder->getContext(); } + + bool isLineInfoEnabled() { return lineInfoEnabled; } + + void setLastLoc(mlir::Location loc) { + if (lineInfoEnabled) + lastLoc = std::make_unique(loc); + } + + void setLastLoc(const std::string &fileName, int line, int column) { + auto context = builder->getContext(); + setLastLoc(mlir::FileLineColLoc::get(context, fileName, line, column)); + } + + mlir::Location getLastLoc() { + assert(lastLoc); + return *lastLoc; + } + + void setInsertionPointToStart(mlir::Block &block) { + if (!block.empty()) + setLastLoc(block.begin()->getLoc()); + else + setLastLoc(getLocForBlock(&block)); + builder->setInsertionPointToStart(&block); + } + + void setInsertionPointToEnd(mlir::Block &block) { + if (!block.empty()) + setLastLoc(block.back().getLoc()); + else + setLastLoc(getLocForBlock(&block)); + builder->setInsertionPointToEnd(&block); + } + + void setInsertionPointAfter(mlir::Operation &op) { + setLastLoc(op.getLoc()); + builder->setInsertionPointAfter(&op); + } + + void restoreInsertionPoint(mlir::OpBuilder::InsertPoint pt) { + setLastLoc(builder->getUnknownLoc()); + if (pt.isSet()) { + if (pt.getPoint() != pt.getBlock()->end()) + setLastLoc(pt.getPoint()->getLoc()); + else + setLastLoc(getLocForBlock(pt.getBlock())); + } + + builder->restoreInsertionPoint(pt); + } + + template OpTy create(Args &&...args) { + auto loc = getLastLoc(); + return OpTy::create(*builder, loc, std::forward(args)...); + } + + // Overload to create or fold a single result operation. + template + std::enable_if_t(), + mlir::Value> + createOrFold(Args &&...args) { + auto loc = getLastLoc(); + return builder->createOrFold(loc, std::forward(args)...); + } + + // Overload to create or fold a zero result operation. + template + std::enable_if_t(), OpTy> + createOrFold(Args &&...args) { + auto loc = getLastLoc(); + return builder->createOrFold(loc, std::forward(args)...); + } + +private: + std::unique_ptr builder; + std::unique_ptr lastLoc; + bool lineInfoEnabled = + !mlir::triton::tools::getBoolEnv("TRITON_DISABLE_LINE_INFO"); + + mlir::Location getLocForBlock(mlir::Block *block) { + if (auto parentOp = block->getParentOp()) + return parentOp->getLoc(); + return builder->getUnknownLoc(); + } +}; + +#ifdef __TLE__ +namespace ir { +extern py::class_ *getBuilderClass(); +} // namespace ir +#endif diff --git a/third_party/sunrise/python/src/linear_layout.cc b/third_party/sunrise/python/src/linear_layout.cc new file mode 100644 index 0000000000..21ad101257 --- /dev/null +++ b/third_party/sunrise/python/src/linear_layout.cc @@ -0,0 +1,223 @@ +#include "pybind11/numpy.h" +#include "pybind11/pybind11.h" +#include "pybind11/stl.h" + +#include "mlir/IR/Attributes.h" +#include "mlir/IR/MLIRContext.h" +#include "triton/Dialect/TritonGPU/IR/Dialect.h" +#include "triton/Tools/LinearLayout.h" +#include "llvm/ADT/STLExtras.h" +#include +#include +#include + +namespace py = pybind11; +using LinearLayout = mlir::triton::LinearLayout; + +namespace { + +mlir::MLIRContext *getLinearLayoutContext() { + static PyObject *ctxObject = []() { + py::module irMod = py::module::import("triton._C.libtriton.ir"); + // Keep the Python object alive for the life of the process without running + // its destructor during interpreter shutdown (avoids segfaults). + py::object ctx = irMod.attr("context")(); + return ctx.release().ptr(); + }(); + return py::cast(py::handle(ctxObject)); +} + +} // namespace + +void init_linear_layout(py::module &&m) { + py::class_(m, "LinearLayout", py::module_local(false)) + .def(py::init<>()) + .def_static( + "identity_1d", + [](int32_t size, std::string inDim, std::string outDim) { + auto *ctx = getLinearLayoutContext(); + return LinearLayout::identity1D(size, + mlir::StringAttr::get(ctx, inDim), + mlir::StringAttr::get(ctx, outDim)); + }, + py::arg("size"), py::arg("inDim"), py::arg("outDim")) + .def_static( + "strided_1d", + [](int32_t size, int32_t stride, std::string inDim, + std::string outDim) { + auto *ctx = getLinearLayoutContext(); + return LinearLayout::strided1D(size, stride, + mlir::StringAttr::get(ctx, inDim), + mlir::StringAttr::get(ctx, outDim)); + }, + py::arg("size"), py::arg("stride"), py::arg("inDim"), + py::arg("outDim")) + .def_static( + "zeros_1d", + [](int32_t size, std::string inDim, std::string outDim, + int32_t outDimSize) { + auto *ctx = getLinearLayoutContext(); + return LinearLayout::zeros1D( + size, mlir::StringAttr::get(ctx, inDim), + mlir::StringAttr::get(ctx, outDim), outDimSize); + }, + py::arg("size"), py::arg("inDim"), py::arg("outDim"), + py::arg("outDimSize") = 1) + .def_static( + "from_bases", + [](const std::vector>>> &bases, + const std::vector &outDimNames, + std::optional> outDimSizes, + bool requireSurjective) { + auto *ctx = getLinearLayoutContext(); + + std::vector< + std::pair>>> + convertedBases; + convertedBases.reserve(bases.size()); + for (const auto &entry : bases) { + std::vector> converted; + converted.reserve(entry.second.size()); + for (const auto &vec : entry.second) + converted.emplace_back(vec.begin(), vec.end()); + convertedBases.emplace_back( + mlir::StringAttr::get(ctx, entry.first), + std::move(converted)); + } + + if (outDimSizes) { + if (outDimSizes->size() != outDimNames.size()) + throw std::invalid_argument("out_dim_names and out_dim_sizes " + "must have the same length"); + std::vector> outDims; + outDims.reserve(outDimNames.size()); + for (auto it : llvm::enumerate(outDimNames)) + outDims.emplace_back(mlir::StringAttr::get(ctx, it.value()), + (*outDimSizes)[it.index()]); + return LinearLayout(convertedBases, outDims, requireSurjective); + } + + if (!requireSurjective) + throw std::invalid_argument("out_dim_sizes must be provided when " + "require_surjective is false"); + + std::vector convertedNames; + convertedNames.reserve(outDimNames.size()); + for (const auto &name : outDimNames) + convertedNames.push_back(mlir::StringAttr::get(ctx, name)); + return LinearLayout(convertedBases, convertedNames); + }, + py::arg("bases"), py::arg("out_dim_names"), + py::arg("out_dim_sizes") = py::none(), + py::arg("require_surjective") = true) + .def("compose", &LinearLayout::compose) + .def("invert_and_compose", &LinearLayout::invertAndCompose) + .def("invert", &LinearLayout::invert) + .def("pseudoinvert", &LinearLayout::pseudoinvert) + .def("is_surjective", &LinearLayout::isSurjective) + .def("is_injective", &LinearLayout::isInjective) + .def("is_invertible", &LinearLayout::isInvertible) + .def("get_in_dim_names", + [](const LinearLayout &self) { + std::vector dims; + dims.reserve(self.getNumInDims()); + for (mlir::StringAttr dim : self.getInDimNames()) + dims.push_back(dim.str()); + return dims; + }) + .def("get_out_dim_names", + [](const LinearLayout &self) { + std::vector dims; + dims.reserve(self.getNumOutDims()); + for (mlir::StringAttr dim : self.getOutDimNames()) + dims.push_back(dim.str()); + return dims; + }) + .def_property_readonly( + "bases", + [](const LinearLayout &self) { + auto bases = self.getBases(); + pybind11::list result; + for (const auto &it : bases) { + pybind11::list dimBases; + for (const auto &vec : it.second) + dimBases.append(pybind11::cast( + std::vector(vec.begin(), vec.end()))); + result.append(pybind11::make_tuple(it.first.str(), dimBases)); + } + return result; + }) + .def_property_readonly( + "out_dims", + [](const LinearLayout &self) { + pybind11::list result; + for (const auto &it : self.getOutDims()) { + result.append(pybind11::make_tuple(it.first.str(), it.second)); + } + return result; + }) + .def_property_readonly("num_in_dims", &LinearLayout::getNumInDims) + .def_property_readonly("num_out_dims", &LinearLayout::getNumOutDims) + .def("__mul__", [](const LinearLayout &lhs, + const LinearLayout &rhs) { return lhs * rhs; }) + .def( + "__imul__", + [](LinearLayout &lhs, const LinearLayout &rhs) -> LinearLayout & { + lhs *= rhs; + return lhs; + }, + py::return_value_policy::reference_internal) + .def("__eq__", [](const LinearLayout &lhs, + const LinearLayout &rhs) { return lhs == rhs; }) + .def("__ne__", [](const LinearLayout &lhs, + const LinearLayout &rhs) { return lhs != rhs; }) + .def("__repr__", [](const LinearLayout &self) { return self.toString(); }) + .def("__str__", [](const LinearLayout &self) { return self.toString(); }) + .def("get_shared_view", + [](const LinearLayout &self, bool useHWPointOfView) { + return mlir::triton::gpu::getSharedLayoutStr( + const_cast(self), useHWPointOfView); + }) + .def("get_distributed_view", + [](const LinearLayout &self, bool useHWPointOfView) { + return mlir::triton::gpu::getDistributedLayoutStr( + const_cast(self), useHWPointOfView); + }) + .def( + "apply", + [](const LinearLayout &self, py::dict inputsDict) { + std::vector> inputs; + inputs.reserve(inputsDict.size()); + for (auto item : inputsDict) { + inputs.emplace_back(py::cast(item.first), + py::cast(item.second)); + } + auto *ctx = getLinearLayoutContext(); + std::vector> converted; + converted.reserve(inputs.size()); + for (const auto &it : inputs) { + converted.emplace_back(mlir::StringAttr::get(ctx, it.first), + it.second); + } + auto outputs = self.apply(converted); + py::dict result; + for (const auto &out : outputs) { + result[py::str(out.first.str())] = out.second; + } + return result; + }, + py::arg("inputs")) + .def("get_matrix_view", [](const LinearLayout &self) { + std::unique_ptr matrix = mlir::triton::getMatrix(self); + auto nRows = self.getTotalOutDimSizeLog2(); + auto nCols = self.getTotalInDimSizeLog2(); + std::vector> result(nRows, std::vector(nCols)); + for (size_t i = 0; i < nRows; ++i) { + for (size_t j = 0; j < nCols; ++j) { + result[i][j] = (matrix[i] >> j) & 1; + } + } + return result; + }); +} diff --git a/third_party/sunrise/python/src/llvm.cc b/third_party/sunrise/python/src/llvm.cc new file mode 100644 index 0000000000..7f3d5b5514 --- /dev/null +++ b/third_party/sunrise/python/src/llvm.cc @@ -0,0 +1,826 @@ +#include "mlir/IR/BuiltinOps.h" // mlir::ModuleOp +#include "mlir/Target/LLVMIR/LLVMTranslationInterface.h" +#include "mlir/Target/LLVMIR/ModuleTranslation.h" +#include "triton/Tools/Sys/GetEnv.hpp" +#include "llvm/ADT/SmallVector.h" +#include "llvm/CodeGen/MIRParser/MIRParser.h" +#include "llvm/CodeGen/MachineModuleInfo.h" +#include "llvm/CodeGen/MachineRegisterInfo.h" +#include "llvm/IR/LLVMContext.h" +#include "llvm/IR/LegacyPassManager.h" +#include "llvm/IR/Module.h" +#include "llvm/IR/PassManager.h" +#include "llvm/IR/Verifier.h" +#include "llvm/IRReader/IRReader.h" +#include "llvm/Linker/Linker.h" +#include "llvm/MC/TargetRegistry.h" +#include "llvm/Pass.h" +#include "llvm/Passes/OptimizationLevel.h" +#include "llvm/Passes/PassBuilder.h" +#include "llvm/Passes/PassPlugin.h" +#include "llvm/Passes/StandardInstrumentations.h" +#include "llvm/Support/CodeGen.h" +#include "llvm/Support/Signals.h" +#include "llvm/Support/SourceMgr.h" +#include "llvm/Support/TargetSelect.h" +#include "llvm/Target/TargetMachine.h" +#include "llvm/Transforms/IPO/AlwaysInliner.h" +#include "llvm/Transforms/InstCombine/InstCombine.h" +#include "llvm/Transforms/Instrumentation/AddressSanitizer.h" +#include "llvm/Transforms/Instrumentation/AddressSanitizerOptions.h" +#include +#include +#include +#include +#include +#include +#include + +namespace py = pybind11; + +namespace llvm { +struct BreakStructPhiNodesPass : PassInfoMixin { + PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM); + static StringRef name() { return "BreakStructPhiNodesPass"; } +}; +} // namespace llvm + +using namespace llvm; + +std::unique_ptr +createTargetMachine(llvm::Module *module, std::string proc, + bool enable_fp_fusion, const std::string &features) { + std::string error; + auto target = + llvm::TargetRegistry::lookupTarget(module->getTargetTriple(), error); + llvm::TargetOptions opt; + bool disableLLVMOpt = mlir::triton::tools::getBoolEnv("DISABLE_LLVM_OPT"); + if (enable_fp_fusion) + opt.AllowFPOpFusion = llvm::FPOpFusion::Fast; + opt.UnsafeFPMath = false; + opt.NoInfsFPMath = false; + opt.NoNaNsFPMath = true; + opt.TrapUnreachable = true; + opt.MCOptions.AsmVerbose = true; + opt.MCOptions.PreserveAsmComments = true; + std::unique_ptr machine{target->createTargetMachine( + module->getTargetTriple(), proc, features, opt, llvm::Reloc::PIC_, + std::nullopt, + disableLLVMOpt ? llvm::CodeGenOptLevel::None + : llvm::CodeGenOptLevel::Aggressive)}; + return machine; +} + +void dumpSchedulingDAG(llvm::Module &module, const std::string &triple, + const std::string &proc, const std::string &features, + const std::vector &flags, + bool enable_fp_fusion, const std::string &dumpFileId) { + using namespace mlir; + + // Check if we should dump sched DAG + std::string dumpMirBase = triton::tools::getStrEnv("TRITON_DUMP_MIR"); + bool dumpMir = !dumpMirBase.empty(); + if (!dumpMir) { + return; + } + + // options + auto options = llvm::cl::getRegisteredOptions(); + for (std::string flag : flags) { + auto *shortPtr = static_cast *>(options[flag]); + assert(shortPtr); + shortPtr->setValue(true); + } + bool disableLLVMOpt = triton::tools::getBoolEnv("DISABLE_LLVM_OPT"); + if (!disableLLVMOpt) { + // Check to see if we are passing a list of flags to disable optimizations. + auto flagList = triton::tools::getStrEnv("DISABLE_LLVM_OPT"); + if (!flagList.empty()) { + llvm::SmallVector split; + StringRef(flagList.c_str()).split(split, ','); + for (auto flag : split) { + auto optIt = options.find(flag); + if (optIt != options.end()) { + auto optPtr = static_cast *>(optIt->second); + *optPtr = true; + } + } + } + } + + // inline everything + for (llvm::Function &f : module.functions()) + if (!f.hasFnAttribute(llvm::Attribute::NoInline)) + f.addFnAttr(llvm::Attribute::AlwaysInline); + // verify and store llvm + llvm::legacy::PassManager pm; + pm.add(llvm::createAlwaysInlinerLegacyPass()); + pm.add(llvm::createVerifierPass()); + + pm.run(module); + + // create machine + module.setTargetTriple(Triple(triple)); + auto machine = createTargetMachine(&module, proc, enable_fp_fusion, features); + // set data layout + module.setDataLayout(machine->createDataLayout()); + + int saved_stderr_fd = -1; + std::string dumpFilename = dumpMirBase + "/" + dumpFileId + ".txt"; + + // Save and set stop-after + std::string originalStopAfter; + auto stopAfterOpt = options.find("stop-after"); + if (stopAfterOpt != options.end()) { + auto *optPtr = + static_cast *>(stopAfterOpt->second); + originalStopAfter = optPtr->getValue(); + optPtr->setValue("machine-scheduler"); + } + + // Enable misched-print-dags for DAG + auto mischedPrintOpt = options.find("misched-print-dags"); + if (mischedPrintOpt != options.end()) { + auto *optPtr = static_cast *>(mischedPrintOpt->second); + optPtr->setValue(true); + } + + // Save original stderr file descriptor + saved_stderr_fd = dup(fileno(stderr)); + + // Redirect stderr to append to dump file + FILE *redirected = freopen(dumpFilename.c_str(), "a", stderr); + if (!redirected) { + llvm::errs() << "Warning: Failed to redirect stderr to " << dumpFilename + << "\n"; + } + + // emit machine code + std::string result; + { + llvm::raw_string_ostream stream(result); + llvm::buffer_ostream pstream(stream); + llvm::legacy::PassManager pass; + // emit + machine->addPassesToEmitFile(pass, pstream, nullptr, + llvm::CodeGenFileType::AssemblyFile); + pass.run(module); + } + + // Restore stderr and reset options + fflush(stderr); + if (saved_stderr_fd != -1) { + dup2(saved_stderr_fd, fileno(stderr)); + close(saved_stderr_fd); + clearerr(stderr); + } + + if (stopAfterOpt != options.end()) { + auto *optPtr = + static_cast *>(stopAfterOpt->second); + optPtr->setValue(originalStopAfter); + } + + if (mischedPrintOpt != options.end()) { + auto *optPtr = static_cast *>(mischedPrintOpt->second); + optPtr->setValue(false); + } + + llvm::errs() << "MIR and DAG dumped to: " << dumpFilename << "\n"; +} + +std::string +translateLLVMIRToMIR(llvm::Module &module, const std::string &triple, + const std::string &proc, const std::string &features, + const std::vector &flags, + bool enable_fp_fusion, const std::string &dumpFileId) { + using namespace mlir; + + // Check if we should dump MIR + std::string dumpMirBase = triton::tools::getStrEnv("TRITON_DUMP_MIR"); + bool dumpMir = !dumpMirBase.empty(); + if (!dumpMir) { + return ""; + } + + // options + auto options = llvm::cl::getRegisteredOptions(); + for (std::string flag : flags) { + auto *shortPtr = static_cast *>(options[flag]); + assert(shortPtr); + shortPtr->setValue(true); + } + bool disableLLVMOpt = triton::tools::getBoolEnv("DISABLE_LLVM_OPT"); + if (!disableLLVMOpt) { + // Check to see if we are passing a list of flags to disable optimizations. + auto flagList = triton::tools::getStrEnv("DISABLE_LLVM_OPT"); + if (!flagList.empty()) { + llvm::SmallVector split; + StringRef(flagList.c_str()).split(split, ','); + for (auto flag : split) { + auto optIt = options.find(flag); + if (optIt != options.end()) { + auto optPtr = static_cast *>(optIt->second); + *optPtr = true; + } + } + } + } + + // Save and set stop-before if needed (for MIR output or custom stop point) + std::string originalStopBefore; + auto stopBeforeOpt = options.find("stop-before"); + if (stopBeforeOpt != options.end()) { + auto *optPtr = + static_cast *>(stopBeforeOpt->second); + originalStopBefore = optPtr->getValue(); + optPtr->setValue("machine-scheduler"); + } + + // inline everything + for (llvm::Function &f : module.functions()) + if (!f.hasFnAttribute(llvm::Attribute::NoInline)) + f.addFnAttr(llvm::Attribute::AlwaysInline); + // verify and store llvm + llvm::legacy::PassManager pm; + pm.add(llvm::createAlwaysInlinerLegacyPass()); + pm.add(llvm::createVerifierPass()); + + pm.run(module); + + // create machine + module.setTargetTriple(Triple(triple)); + auto machine = createTargetMachine(&module, proc, enable_fp_fusion, features); + // set data layout + module.setDataLayout(machine->createDataLayout()); + + // emit machine code + std::string result; + { + llvm::raw_string_ostream stream(result); + llvm::buffer_ostream pstream(stream); + llvm::legacy::PassManager pass; + // emit + machine->addPassesToEmitFile(pass, pstream, nullptr, + llvm::CodeGenFileType::AssemblyFile); + pass.run(module); + } + + if (stopBeforeOpt != options.end()) { + auto *optPtr = + static_cast *>(stopBeforeOpt->second); + optPtr->setValue(originalStopBefore); + } + + std::string dumpFilename = dumpMirBase + "/" + dumpFileId + ".txt"; + { + std::error_code EC; + llvm::raw_fd_ostream outFile(dumpFilename, EC, llvm::sys::fs::OF_None); + if (EC) { + llvm::errs() << "Error opening file " << dumpFilename << ": " + << EC.message() << "\n"; + } else { + outFile << result; + outFile << "---"; + outFile << "\n========== SCHEDULING DAG ==========\n"; + } + } + + return result; +} + +std::string translateLLVMIRToASM(llvm::Module &module, + const std::string &triple, + const std::string &proc, + const std::string &features, + const std::vector &flags, + bool enable_fp_fusion, bool isObject) { + using namespace mlir; + auto isInteger=[&](const std::string& str) { + return !str.empty() && std::all_of(str.begin(), str.end(), ::isdigit); + }; + // options + auto options = llvm::cl::getRegisteredOptions(); + for (std::string flag : flags) { + auto pos = flag.find("="); // sunrise fix + if(pos == std::string::npos){ + auto *shortPtr = static_cast *>(options[flag]); + assert(shortPtr); + shortPtr->setValue(true); + }else{ + auto optIt = options.find(flag.substr(0,pos)); + assert(optIt != options.end() && "Option not found in cl::opt!"); + + if (isInteger(flag.substr(pos+1))) { + if(auto *unsignedOpt = static_cast*>(options[flag.substr(0,pos)])) { + unsignedOpt->setValue(std::stoi(flag.substr(pos+1))); + } + } + else { + if (auto *stringOpt = static_cast*>(options[flag.substr(0,pos)])) { + stringOpt->setValue(flag.substr(pos+1)); + } + } + } + } + if (triton::tools::getBoolEnv("LLVM_IR_ENABLE_DUMP")) { + auto optIt = options.find("print-after-all"); + if (optIt != options.end()) { + auto optPtr = static_cast *>(optIt->second); + *optPtr = true; + } + } + bool disableLLVMOpt = triton::tools::getBoolEnv("DISABLE_LLVM_OPT"); + if (!disableLLVMOpt) { + // Check to see if we are passing a list of flags to disable optimizations. + auto flagList = triton::tools::getStrEnv("DISABLE_LLVM_OPT"); + if (!flagList.empty()) { + llvm::SmallVector split; + StringRef(flagList.c_str()).split(split, ','); + for (auto flag : split) { + auto optIt = options.find(flag); + if (optIt != options.end()) { + auto optPtr = static_cast *>(optIt->second); + *optPtr = true; + } + } + } + } + + // inline everything + for (llvm::Function &f : module.functions()) + if (!f.hasFnAttribute(llvm::Attribute::NoInline)) + f.addFnAttr(llvm::Attribute::AlwaysInline); + // verify and store llvm + llvm::legacy::PassManager pm; + pm.add(llvm::createAlwaysInlinerLegacyPass()); + pm.add(llvm::createVerifierPass()); + + const bool enabledTiming = triton::tools::getBoolEnv("LLVM_ENABLE_TIMING"); + if (enabledTiming) { + llvm::TimePassesIsEnabled = true; + llvm::TimePassesPerRun = true; + } + + pm.run(module); + + SmallString<0> timePassesStr; + raw_svector_ostream reportStream(timePassesStr); + + if (enabledTiming) { + reportAndResetTimings(&reportStream); + llvm::dbgs() << reportStream.str(); + timePassesStr.clear(); + } + + // create machine + module.setTargetTriple(Triple(triple)); + auto machine = createTargetMachine(&module, proc, enable_fp_fusion, features); + // set data layout + module.setDataLayout(machine->createDataLayout()); + // emit machine code + std::string result; + { + llvm::raw_string_ostream stream(result); + llvm::buffer_ostream pstream(stream); + llvm::legacy::PassManager pass; + // emit + auto fileType = isObject ? llvm::CodeGenFileType::ObjectFile + : llvm::CodeGenFileType::AssemblyFile; + machine->addPassesToEmitFile(pass, pstream, nullptr, fileType); + pass.run(module); + + if (enabledTiming) { + reportAndResetTimings(&reportStream); + llvm::dbgs() << reportStream.str(); + timePassesStr.clear(); + } + } + return result; +} + +using ret = py::return_value_policy; + +void init_triton_llvm(py::module &&m) { + + py::class_(m, "context", py::module_local()) + .def(py::init<>()); + py::class_(m, "source_mgr", py::module_local()) + .def(py::init<>()); + + py::class_(m, "function_list") + .def( + "__iter__", + [](llvm::Module::FunctionListType &s) { + return py::make_iterator(s.begin(), s.end()); + }, + py::keep_alive<0, 1>()); + + // Module Flag behavior. See + // https://llvm.org/doxygen/classllvm_1_1Module.html#a0a5c55e12c97b80021330fe82b642293 + // for details. + py::class_(m, "module_flag_behavior", + py::module_local()); + m.attr("MODULE_FLAG_BEHAVIOR_ERROR") = llvm::Module::Error; + m.attr("MODULE_FLAG_BEHAVIOR_WARNING") = llvm::Module::Warning; + m.attr("MODULE_FLAG_BEHAVIOR_REQUIRE") = llvm::Module::Require; + m.attr("MODULE_FLAG_BEHAVIOR_OVERRIDE") = llvm::Module::Override; + m.attr("MODULE_FLAG_BEHAVIOR_APPEND") = llvm::Module::Append; + m.attr("MODULE_FLAG_BEHAVIOR_APPEND_UNIQUE") = llvm::Module::AppendUnique; + m.attr("MODULE_FLAG_BEHAVIOR_MAX") = llvm::Module::Max; + m.attr("MODULE_FLAG_BEHAVIOR_MIN") = llvm::Module::Min; + + py::class_(m, "module", py::module_local()) + .def( + "__str__", + [](llvm::Module *self) { + std::string str; + llvm::raw_string_ostream os(str); + os << *self; + return os.str(); + }, + ret::take_ownership) + .def( + "get_functions", + [](llvm::Module *mod) -> llvm::Module::FunctionListType & { + // Note: Backends assume that we are compiling exactly one kernel + // (i.e. one function that's that's called by the CPU) and that it's + // the first function in this list. + return mod->getFunctionList(); + }, + ret::reference_internal) + .def("add_flag", + [](llvm::Module *mod, llvm::Module::ModFlagBehavior behavior, + std::string &key, uint32_t value) { + return mod->addModuleFlag(behavior, key, value); + }); + + py::class_(m, "function", py::module_local()) + .def_property_readonly( + "name", [](llvm::Function *fn) { return fn->getName().str(); }) + .def("set_calling_conv", &llvm::Function::setCallingConv) + .def("add_fn_attr", [](llvm::Function *fn, std::string &name, + std::string &val) { fn->addFnAttr(name, val); }) + .def("remove_fn_attr", [](llvm::Function *fn, + std::string &name) { fn->removeFnAttr(name); }) + .def("add_fn_asan_attr", + [](llvm::Function *fn) { + fn->addFnAttr(llvm::Attribute::SanitizeAddress); + }) + .def("add_fn_target_feature", + [](llvm::Function *fn, std::string &val) { + fn->addFnAttr("target-features", val); + }) + // Sets the nvvm.maxreg property on the given function. + .def("set_nvvm_maxnreg", + [](llvm::Function *fn, int maxnreg) { + auto op = MDNode::get( + fn->getContext(), + { + ValueAsMetadata::get(fn), + MDString::get(fn->getContext(), "maxnreg"), + ConstantAsMetadata::get(ConstantInt::get( + Type::getInt32Ty(fn->getContext()), maxnreg)), + }); + fn->getParent() + ->getOrInsertNamedMetadata("nvvm.annotations") + ->addOperand(op); + }) + // External functions that are definitions (i.e. not declarations) are + // kernel functions. + .def("is_declaration", &llvm::Function::isDeclaration) + .def("is_external_linkage", [](llvm::Function *fn) { + return fn->getLinkage() == llvm::GlobalValue::ExternalLinkage; + }); + + // optimization levels + py::class_(m, "optimization_level", + py::module_local()); + m.attr("OPTIMIZE_O0") = llvm::OptimizationLevel::O0; + m.attr("OPTIMIZE_O1") = llvm::OptimizationLevel::O1; + m.attr("OPTIMIZE_O2") = llvm::OptimizationLevel::O2; + m.attr("OPTIMIZE_O3") = llvm::OptimizationLevel::O3; + m.attr("OPTIMIZE_Os") = llvm::OptimizationLevel::Os; + m.attr("OPTIMIZE_Oz") = llvm::OptimizationLevel::Oz; + + m.def( + "to_module", + [](mlir::ModuleOp &mod, llvm::LLVMContext &ctx) { + std::unique_ptr llvmMod = + mlir::translateModuleToLLVMIR(mod, ctx); + if (!llvmMod) { + throw std::runtime_error("failed to translate module to LLVM IR"); + } + return llvmMod; + }, + py::keep_alive<0, 2>(), py::call_guard()); + + m.def("attach_datalayout", [](llvm::Module *mod, const std::string triple, + const std::string proc, + const std::string features) { + std::string error; + llvm::Triple targetTriple(triple); + auto target = llvm::TargetRegistry::lookupTarget(targetTriple, error); + if (!target) { + throw std::runtime_error("target lookup error: " + error); + } + llvm::TargetOptions opt; + // Target machine is only used to create the data layout. + std::unique_ptr machine{target->createTargetMachine( + targetTriple, proc, features, opt, llvm::Reloc::PIC_, std::nullopt, + llvm::CodeGenOptLevel::None)}; + // set data layout + mod->setDataLayout(machine->createDataLayout()); + }); + + m.def( + "optimize_module", + [](llvm::Module *mod, const llvm::OptimizationLevel &opt, + std::string arch, std::string features, std::vector flags, + bool enable_fp_fusion) { + if (mlir::triton::tools::getBoolEnv("DISABLE_LLVM_OPT")) + return; + auto options = llvm::cl::getRegisteredOptions(); + // Hack for the 3.6 release only. Vectorization of copyable elements + // exposed a bug in ptxas. Manually disable it by modifying the command + // line option for it. Note that we can abuse DISABLE_LLVM_OPT to + // override this, since setting it to slp-copyable-elements will set the + // flag back to true. + auto it = options.find("slp-copyable-elements"); + if (it != options.end()) + *static_cast *>(it->second) = false; + // Check to see if we are passing a list of flags to disable + // optimizations. + auto flagList = mlir::triton::tools::getStrEnv("DISABLE_LLVM_OPT"); + if (!flagList.empty()) { + llvm::SmallVector split; + StringRef(flagList.c_str()).split(split, ','); + for (auto flag : split) { + auto optIt = options.find(flag); + if (optIt != options.end()) { + auto optPtr = static_cast *>(optIt->second); + *optPtr = true; + } + } + } + using namespace llvm; + LoopAnalysisManager lam; + FunctionAnalysisManager fam; + CGSCCAnalysisManager cgam; + ModuleAnalysisManager mam; + + if (arch.empty()) { + llvm::TargetLibraryInfoImpl TLII(mod->getTargetTriple()); + TLII.disableAllFunctions(); + fam.registerPass([TLII = std::move(TLII)] { + return llvm::TargetLibraryAnalysis(TLII); + }); + } + + PassInstrumentationCallbacks *instrCbPtr = nullptr; + PassInstrumentationCallbacks passInstrCb; + StandardInstrumentations standardInstr(mod->getContext(), + /*DebugLogging*/ true); + if (mlir::triton::tools::getBoolEnv("LLVM_IR_ENABLE_DUMP")) { + auto optMap = llvm::cl::getRegisteredOptions(); + auto optIt = optMap.find("print-after-all"); + if (optIt != optMap.end()) { + auto optPtr = static_cast *>(optIt->second); + *optPtr = true; + } + standardInstr.registerCallbacks(passInstrCb, &mam); + instrCbPtr = &passInstrCb; + } + + PipelineTuningOptions tuningOptions; + tuningOptions.LoopUnrolling = true; + tuningOptions.LoopInterleaving = true; + tuningOptions.LoopVectorization = true; + // TODO: currently we run SLP vectorizer with an empty target machine. + // This cause the vectorizer to create larger vector which could be bad. + // Disabling it would currently cause regressions as this pass also + // applies some scheduling that helps performance in some cases. We + // should work on using NVPTX target instead and address the performance + // regressions with some scheduling solution. + tuningOptions.SLPVectorization = true; + + std::string pluginFile = + mlir::triton::tools::getStrEnv("LLVM_PASS_PLUGIN_PATH"); + + // We don't pass the targetMachine to the LLVM-IR pass builder, unless + // `arch` is specified. + // + // Don't set target machine in LLVM pass builder when using LLVM IR + // level plugins. LLVM IR level plugin passes typically want to insert + // calls to externally generated code (i.e. precompile a Cuda/Hip kernel + // with Clang and then insert a call to it within an instrumentation + // pass) setting the targetMachine value here can can cause a mismatch + // in the target machine between the MLIR and Clang generated kernels + // and break the lowering of some target specific intrinsics. + std::unique_ptr targetMachine = nullptr; + if (!arch.empty() && pluginFile.empty()) + targetMachine = + createTargetMachine(mod, arch, enable_fp_fusion, features); + PassBuilder pb(/*targetMachine=*/targetMachine.get(), tuningOptions, + std::nullopt, instrCbPtr); + + if (!pluginFile.empty()) { + // TODO: Add some logging here that we inserted a pass into the LLVM + // pass pipeline + auto passPlugin = llvm::PassPlugin::Load(pluginFile); + if (!passPlugin) { + llvm::Error Err = passPlugin.takeError(); + std::string ErrMsg = + "Pass Plugin Error: " + llvm::toString(std::move(Err)); + throw std::runtime_error(ErrMsg); + } + passPlugin->registerPassBuilderCallbacks(pb); + } + + pb.registerModuleAnalyses(mam); + pb.registerCGSCCAnalyses(cgam); + pb.registerFunctionAnalyses(fam); + pb.registerLoopAnalyses(lam); + pb.crossRegisterProxies(lam, fam, cgam, mam); + + ModulePassManager mpm; + pb.registerVectorizerStartEPCallback( + [&](llvm::FunctionPassManager &fpm, llvm::OptimizationLevel level) { + // Triton generates large structure of scalars which may pessimise + // optimizations, we run a pass to break up phi of struct to make + // sure all the struct are removed for the following passes. + fpm.addPass(BreakStructPhiNodesPass()); + fpm.addPass(InstCombinePass()); + }); + bool enableAddressSanitizer = + mlir::triton::tools::getBoolEnv("TRITON_ENABLE_ASAN"); + if (enableAddressSanitizer) { + AddressSanitizerOptions Opts; + mpm.addPass(AddressSanitizerPass(Opts)); + } + mpm.addPass(pb.buildPerModuleDefaultPipeline(opt)); + mpm.run(*mod, mam); + }, + // Mandatory parameters + py::arg("mod"), py::arg("opt"), + // If we want to specify the target machine, we require additional + // (optional) parameters + py::arg("arch") = "", py::arg("features") = "", + py::arg("flags") = std::vector{}, + py::arg("enable_fp_fusion") = false, + py::call_guard()); + + m.def( + "translate_to_asm", + [](std::string llvmIR, std::string triple, std::string proc, + std::string features, std::vector flags, + bool enable_fp_fusion, bool isObject) -> py::object { + std::string obj; + { + // when allow_threads goes out of scope, gil will be released + py::gil_scoped_release allow_threads; + // create LLVM module from C++ + llvm::LLVMContext context; + std::unique_ptr buffer = + llvm::MemoryBuffer::getMemBuffer(llvmIR.c_str()); + llvm::SMDiagnostic error; + std::unique_ptr module = + llvm::parseIR(buffer->getMemBufferRef(), error, context); + if (!module) { + llvm::report_fatal_error( + "failed to parse IR: " + error.getMessage() + + "lineno: " + std::to_string(error.getLineNo())); + } + obj = translateLLVMIRToASM(*module, triple, proc, features, flags, + enable_fp_fusion, isObject); + } + if (isObject) + return py::bytes(obj); + else + return py::str(obj); + }, + ret::take_ownership); + + m.def("dump_sched_dag", [](std::string llvmIR, std::string triple, + std::string proc, std::string features, + std::vector flags, + bool enable_fp_fusion, std::string dumpFileId) { + // when allow_threads goes out of scope, gil will be released + py::gil_scoped_release allow_threads; + // create LLVM module from C++ + llvm::LLVMContext context; + std::unique_ptr buffer = + llvm::MemoryBuffer::getMemBuffer(llvmIR.c_str()); + llvm::SMDiagnostic error; + std::unique_ptr module = + llvm::parseIR(buffer->getMemBufferRef(), error, context); + if (!module) { + llvm::report_fatal_error("failed to parse IR: " + error.getMessage() + + "lineno: " + std::to_string(error.getLineNo())); + } + dumpSchedulingDAG(*module, triple, proc, features, flags, enable_fp_fusion, + dumpFileId); + }); + + m.def( + "translate_to_mir", + [](std::string llvmIR, std::string triple, std::string proc, + std::string features, std::vector flags, + bool enable_fp_fusion, std::string dumpFileId) -> py::object { + std::string obj; + { + // when allow_threads goes out of scope, gil will be released + py::gil_scoped_release allow_threads; + // create LLVM module from C++ + llvm::LLVMContext context; + std::unique_ptr buffer = + llvm::MemoryBuffer::getMemBuffer(llvmIR.c_str()); + llvm::SMDiagnostic error; + std::unique_ptr module = + llvm::parseIR(buffer->getMemBufferRef(), error, context); + if (!module) { + llvm::report_fatal_error( + "failed to parse IR: " + error.getMessage() + + "lineno: " + std::to_string(error.getLineNo())); + } + obj = translateLLVMIRToMIR(*module, triple, proc, features, flags, + enable_fp_fusion, dumpFileId); + } + return py::str(obj); + }, + ret::take_ownership); + + m.def("init_targets", []() { + static std::once_flag init_flag; + std::call_once(init_flag, []() { + llvm::InitializeAllTargetInfos(); + llvm::InitializeAllTargets(); + llvm::InitializeAllTargetMCs(); + llvm::InitializeAllAsmParsers(); + llvm::InitializeAllAsmPrinters(); + }); + }); + + m.def("link_extern_libs", [](llvm::Module *dstMod, + const std::vector &paths) { + if (paths.empty()) + return; + + auto shouldOverrideFromSrc = [](const std::string &path) { + return path.find("builtin_math_double_S2.bc") != std::string::npos || + path.find("builtin_math_double_S3.bc") != std::string::npos || + path.find("builtin_math_fdiv_S3.bc") != std::string::npos; + }; + + LLVMContext &ctx = dstMod->getContext(); + llvm::Linker linker(*dstMod); + for (const std::string &path : paths) { + llvm::SMDiagnostic err; + std::unique_ptr libMod = llvm::parseIRFile(path, err, ctx); + if (!libMod) { + std::string message = "Failed to parse library at " + path; + throw std::invalid_argument(message); + } + libMod->setTargetTriple(Triple(dstMod->getTargetTriple())); + libMod->setDataLayout(dstMod->getDataLayout()); + + bool keepExternalLinkage = shouldOverrideFromSrc(path); + unsigned linkFlags = keepExternalLinkage + ? llvm::Linker::Flags::OverrideFromSrc + : llvm::Linker::Flags::LinkOnlyNeeded; + std::unordered_set externalFns; + if (!keepExternalLinkage) { + for (llvm::Function &fn : libMod->functions()) { + if (!fn.isDeclaration()) + externalFns.insert(fn.getName().str()); + } + } + + if (linker.linkInModule(std::move(libMod), linkFlags)) { + std::string message = "Failed to link library at " + path; + throw std::invalid_argument(message); + } + + // Mark linked-in functions as internal because backends use external + // linkage as a signifier of kernel functions. + if (!keepExternalLinkage) { + for (llvm::Function &fn : dstMod->functions()) { + if (externalFns.count(fn.getName().str())) { + fn.setLinkage(llvm::GlobalValue::InternalLinkage); + } + } + } + } + }); +} + +void triton_stacktrace_signal_handler(void *) { + llvm::sys::PrintStackTrace(llvm::errs()); + raise(SIGABRT); +} + +void init_triton_stacktrace_hook(pybind11::module &m) { + if (mlir::triton::tools::getBoolEnv("TRITON_ENABLE_PYTHON_STACKTRACE")) { + llvm::sys::AddSignalHandler(triton_stacktrace_signal_handler, nullptr); + } +} diff --git a/third_party/sunrise/python/src/main.cc b/third_party/sunrise/python/src/main.cc new file mode 100644 index 0000000000..40b19f77ea --- /dev/null +++ b/third_party/sunrise/python/src/main.cc @@ -0,0 +1,62 @@ +#include "llvm/ADT/StringRef.h" +#include "llvm/Support/Signals.h" +#include + +namespace py = pybind11; + +#define FOR_EACH_1(MACRO, X) MACRO(X) +#define FOR_EACH_2(MACRO, X, ...) MACRO(X) FOR_EACH_1(MACRO, __VA_ARGS__) +#define FOR_EACH_3(MACRO, X, ...) MACRO(X) FOR_EACH_2(MACRO, __VA_ARGS__) +#define FOR_EACH_4(MACRO, X, ...) MACRO(X) FOR_EACH_3(MACRO, __VA_ARGS__) +#define FOR_EACH_5(MACRO, X, ...) MACRO(X) FOR_EACH_4(MACRO, __VA_ARGS__) + +#define FOR_EACH_NARG(...) FOR_EACH_NARG_(__VA_ARGS__, FOR_EACH_RSEQ_N()) +#define FOR_EACH_NARG_(...) FOR_EACH_ARG_N(__VA_ARGS__) +#define FOR_EACH_ARG_N(_1, _2, _3, _4, _5, N, ...) N +#define FOR_EACH_RSEQ_N() 5, 4, 3, 2, 1, 0 + +#define CONCATENATE(x, y) CONCATENATE1(x, y) +#define CONCATENATE1(x, y) x##y + +#define FOR_EACH(MACRO, ...) \ + CONCATENATE(FOR_EACH_, FOR_EACH_NARG_HELPER(__VA_ARGS__))(MACRO, __VA_ARGS__) +#define FOR_EACH_NARG_HELPER(...) FOR_EACH_NARG(__VA_ARGS__) + +// New macro to remove parentheses +#define REMOVE_PARENS(...) __VA_ARGS__ + +// Intermediate macro to ensure correct expansion +#define FOR_EACH_P_INTERMEDIATE(MACRO, ...) FOR_EACH(MACRO, __VA_ARGS__) + +// Modified FOR_EACH to handle parentheses +#define FOR_EACH_P(MACRO, ARGS_WITH_PARENS) \ + FOR_EACH_P_INTERMEDIATE(MACRO, REMOVE_PARENS ARGS_WITH_PARENS) + +#define DECLARE_BACKEND(name) void init_triton_##name(pybind11::module &&m); + +#define INIT_BACKEND(name) init_triton_##name(m.def_submodule(#name)); + +void init_triton_env_vars(pybind11::module &m); +void init_triton_ir(pybind11::module &&m); +void init_triton_llvm(pybind11::module &&m); +void init_triton_interpreter(pybind11::module &&m); +void init_triton_passes(pybind11::module &&m); +void init_triton_stacktrace_hook(pybind11::module &m); +void init_gluon_ir(pybind11::module &&m); +void init_linear_layout(pybind11::module &&m); +void init_native_specialize(pybind11::module &m); +FOR_EACH_P(DECLARE_BACKEND, TRITON_BACKENDS_TUPLE) + +PYBIND11_MODULE(libtriton, m) { + m.doc() = "Python bindings to the C++ Triton API"; + init_triton_stacktrace_hook(m); + init_triton_env_vars(m); + init_native_specialize(m); + init_triton_ir(m.def_submodule("ir")); + init_triton_passes(m.def_submodule("passes")); + init_triton_interpreter(m.def_submodule("interpreter")); + init_triton_llvm(m.def_submodule("llvm")); + init_linear_layout(m.def_submodule("linear_layout")); + init_gluon_ir(m.def_submodule("gluon_ir")); + FOR_EACH_P(INIT_BACKEND, TRITON_BACKENDS_TUPLE) +} diff --git a/third_party/sunrise/python/src/passes.cc b/third_party/sunrise/python/src/passes.cc new file mode 100644 index 0000000000..f264f4005b --- /dev/null +++ b/third_party/sunrise/python/src/passes.cc @@ -0,0 +1,135 @@ +#include "mlir/Transforms/Passes.h" +#include "mlir/Conversion/Passes.h" +#include "mlir/Pass/Pass.h" +#include "mlir/Pass/PassManager.h" +#include "passes.h" +#include "triton/Analysis/Allocation.h" +#include "triton/Analysis/Membar.h" +#include "triton/Conversion/TritonGPUToLLVM/Passes.h" +#include "triton/Conversion/TritonToTritonGPU/Passes.h" +#include "triton/Dialect/Gluon/Transforms/Passes.h" +#include "triton/Dialect/Triton/Transforms/Passes.h" +#include "triton/Dialect/TritonGPU/Transforms/Passes.h" +#include "triton/Dialect/TritonInstrument/Transforms/Passes.h" +#include "triton/Target/LLVMIR/Passes.h" +#include +#include + +namespace py = pybind11; + +void init_triton_analysis(py::module &&m) { + py::class_(m, "allocation", py::module_local()) + .def(py::init()); + py::class_(m, "membar", py::module_local()) + .def(py::init()) + .def("run", &mlir::ModuleMembarAnalysis::run); +} + +void init_triton_passes_common(py::module &&m) { + using namespace mlir; + ADD_PASS_WRAPPER_0("add_sccp", createSCCPPass); + ADD_PASS_WRAPPER_0("add_symbol_dce", createSymbolDCEPass); + ADD_PASS_WRAPPER_0("add_inliner", createInlinerPass); + ADD_PASS_WRAPPER_0("add_canonicalizer", createCanonicalizerPass); + ADD_PASS_WRAPPER_0("add_cse", createCSEPass); + ADD_PASS_WRAPPER_0("add_licm", createLoopInvariantCodeMotionPass); + ADD_PASS_WRAPPER_0("print_ir", createPrintIRPass); +} + +void init_triton_passes_ttir(py::module &&m) { + using namespace mlir::triton; + ADD_PASS_WRAPPER_0("add_combine", createTritonCombineOps); + ADD_PASS_WRAPPER_0("add_reorder_broadcast", createTritonReorderBroadcast); + ADD_PASS_WRAPPER_0("add_rewrite_tensor_pointer", + createTritonRewriteTensorPointer); + ADD_PASS_WRAPPER_0("add_rewrite_tensor_descriptor_to_pointer", + createTritonRewriteTensorDescriptorToPointer); + ADD_PASS_WRAPPER_0("add_loop_unroll", createTritonLoopUnroll); + ADD_PASS_WRAPPER_0("add_triton_licm", createTritonLoopInvariantCodeMotion); + ADD_PASS_WRAPPER_0("add_loop_aware_cse", createTritonLoopAwareCSE); + ADD_PASS_OPTION_WRAPPER_4("add_convert_to_ttgpuir", + createConvertTritonToTritonGPU, const std::string &, + int, int, int); +} + +void init_triton_passes_ttgpuir(py::module &&m) { + using namespace mlir; + using namespace mlir::triton::gpu; + using namespace mlir::triton::instrument; + ADD_PASS_WRAPPER_0("add_process_shared_memory_hint", + createTritonGPUProcessSharedMemoryHint); // flagtree hints + ADD_PASS_WRAPPER_0("add_coalesce", createTritonGPUCoalesce); + ADD_PASS_WRAPPER_0("add_optimize_thread_locality", + createTritonGPUOptimizeThreadLocality); + ADD_PASS_OPTION_WRAPPER_1("add_hoist_tmem_alloc", + createTritonGPUHoistTMEMAlloc, bool); + ADD_PASS_OPTION_WRAPPER_1("add_assign_latencies", + createTritonGPUAssignLatencies, int); + ADD_PASS_WRAPPER_0("add_schedule_loops", createTritonGPUScheduleLoops); + ADD_PASS_OPTION_WRAPPER_2("add_pipeline", createTritonGPUPipeline, int, bool); + ADD_PASS_OPTION_WRAPPER_1("add_warp_specialize", + createTritonGPUAutomaticWarpSpecialization, int); + ADD_PASS_WRAPPER_0("add_prefetch", createTritonGPUPrefetch); + ADD_PASS_WRAPPER_0("add_accelerate_matmul", createTritonGPUAccelerateMatmul); + ADD_PASS_WRAPPER_0("add_reorder_instructions", + createTritonGPUReorderInstructions); + ADD_PASS_OPTION_WRAPPER_1("add_f32_dot_tc", createTritonGPUF32DotTC, bool); + ADD_PASS_OPTION_WRAPPER_1("add_optimize_dot_operands", + createTritonGPUOptimizeDotOperands, bool); + ADD_PASS_WRAPPER_0("add_remove_layout_conversions", + createTritonGPURemoveLayoutConversions); + ADD_PASS_WRAPPER_0("add_reduce_data_duplication", + createTritonGPUReduceDataDuplication); + ADD_PASS_WRAPPER_0("add_allocate_warp_groups", + createTritonGPUAllocateWarpGroups); + ADD_PASS_WRAPPER_0("add_allocate_shared_memory", createAllocateSharedMemory); + ADD_PASS_WRAPPER_0("add_allocate_global_scratch_memory", + createTritonGPUGlobalScratchAllocationPass); + ADD_PASS_WRAPPER_0("add_combine_tensor_select_and_if", + createTritonGPUCombineTensorSelectAndIf); + ADD_PASS_WRAPPER_0("add_optimize_accumulator_init", + createTritonGPUOptimizeAccumulatorInit); + ADD_PASS_WRAPPER_0("add_fuse_nested_loops", createTritonGPUFuseNestedLoops); + ADD_PASS_WRAPPER_0("add_coalesce_async_copy", + createTritonGPUCoalesceAsyncCopy); + ADD_PASS_WRAPPER_0("add_concurrency_sanitizer", + createTritonInstrumentConcurrencySanitizer); + ADD_PASS_WRAPPER_0("add_optimize_partition_warps", + createTritonGPUOptimizePartitionWarps); +} + +void init_triton_passes_convert(py::module &&m) { + using namespace mlir; + ADD_PASS_WRAPPER_0("add_scf_to_cf", createSCFToControlFlowPass); + ADD_PASS_WRAPPER_0("add_cf_to_llvmir", createConvertControlFlowToLLVMPass); + ADD_PASS_WRAPPER_0("add_index_to_llvmir", createConvertIndexToLLVMPass); + ADD_PASS_WRAPPER_0("add_arith_to_llvmir", createArithToLLVMConversionPass); + ADD_PASS_WRAPPER_0("add_nvvm_to_llvm", createConvertNVVMToLLVMPass); +} + +void init_triton_passes_llvmir(py::module &&m) { + using namespace mlir; + ADD_PASS_WRAPPER_0("add_di_scope", mlir::createLLVMDIScope); + ADD_PASS_WRAPPER_0("add_di_local_variable", mlir::createLLVMDILocalVariable); +} + +void init_gluon_passes(py::module &&m) { + using namespace mlir; + namespace gluon = mlir::triton::gluon; + ADD_PASS_WRAPPER_0("add_resolve_auto_encodings", + gluon::createGluonResolveAutoEncodingsPass); + ADD_PASS_WRAPPER_0("add_canonicalizer", gluon::createGluonCanonicalize); + ADD_PASS_WRAPPER_0("add_inliner", gluon::createGluonInline); + ADD_PASS_WRAPPER_0("add_infer_coalesced_encodings", + gluon::createGluonInferCoalescedEncodingsPass); +} + +void init_triton_passes(py::module &&m) { + init_triton_analysis(m.def_submodule("analysis")); + init_triton_passes_common(m.def_submodule("common")); + init_triton_passes_convert(m.def_submodule("convert")); + init_triton_passes_ttir(m.def_submodule("ttir")); + init_triton_passes_ttgpuir(m.def_submodule("ttgpuir")); + init_triton_passes_llvmir(m.def_submodule("llvmir")); + init_gluon_passes(m.def_submodule("gluon")); +} diff --git a/third_party/sunrise/python/src/passes.h b/third_party/sunrise/python/src/passes.h new file mode 100644 index 0000000000..62f5986a07 --- /dev/null +++ b/third_party/sunrise/python/src/passes.h @@ -0,0 +1,43 @@ +#define ADD_PASS_WRAPPER_0(name, builder) \ + m.def(name, [](mlir::PassManager &pm) { pm.addPass(builder()); }) + +#define ADD_PASS_WRAPPER_1(name, builder, ty0) \ + m.def(name, \ + [](mlir::PassManager &pm, ty0 val0) { pm.addPass(builder(val0)); }) + +#define ADD_PASS_WRAPPER_2(name, builder, ty0, ty1) \ + m.def(name, [](mlir::PassManager &pm, ty0 val0, ty1 val1) { \ + pm.addPass(builder(val0, val1)); \ + }) + +#define ADD_PASS_WRAPPER_3(name, builder, ty0, ty1, ty2) \ + m.def(name, [](mlir::PassManager &pm, ty0 val0, ty1 val1, ty2 val2) { \ + pm.addPass(builder(val0, val1, val2)); \ + }) + +#define ADD_PASS_WRAPPER_4(name, builder, ty0, ty1, ty2, ty3) \ + m.def(name, [](mlir::PassManager &pm, ty0 val0, ty1 val1, ty2 val2, \ + ty3 val3) { pm.addPass(builder(val0, val1, val2, val3)); }) + +#define ADD_PASS_OPTION_WRAPPER_1(name, builder, ty0) \ + m.def(name, \ + [](mlir::PassManager &pm, ty0 val0) { pm.addPass(builder({val0})); }) + +#define ADD_PASS_OPTION_WRAPPER_2(name, builder, ty0, ty1) \ + m.def(name, [](mlir::PassManager &pm, ty0 val0, ty1 val1) { \ + pm.addPass(builder({val0, val1})); \ + }) + +#define ADD_PASS_OPTION_WRAPPER_3(name, builder, ty0, ty1, ty2) \ + m.def(name, [](mlir::PassManager &pm, ty0 val0, ty1 val1, ty2 val2) { \ + pm.addPass(builder({val0, val1, val2})); \ + }) + +#define ADD_PASS_OPTION_WRAPPER_4(name, builder, ty0, ty1, ty2, ty3) \ + m.def(name, [](mlir::PassManager &pm, ty0 val0, ty1 val1, ty2 val2, \ + ty3 val3) { pm.addPass(builder({val0, val1, val2, val3})); }) + +#define ADD_PASS_OPTION_WRAPPER_5(name, builder, ty0, ty1, ty2, ty3, ty4) \ + m.def(name, \ + [](mlir::PassManager &pm, ty0 val0, ty1 val1, ty2 val2, ty3 val3, \ + ty4 val4) { pm.addPass(builder({val0, val1, val2, val3, val4})); }) diff --git a/third_party/sunrise/python/src/specialize.cc b/third_party/sunrise/python/src/specialize.cc new file mode 100644 index 0000000000..3449e3c900 --- /dev/null +++ b/third_party/sunrise/python/src/specialize.cc @@ -0,0 +1,584 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +namespace py = pybind11; + +using DTypePtrKey = std::pair; +using DTypeKey = Py_hash_t; + +struct DTypePtrKeyHash { + std::size_t operator()(const DTypePtrKey &k) const { + return std::hash()(k.first) ^ (std::hash()(k.second) << 1); + } +}; + +using DtypePtr2Str = + std::unordered_map; +using Dtype2Str = std::unordered_map; + +using TypeHandler = std::pair (*)(PyObject *, + PyObject *, bool, + bool, bool); +using TypeHandlerCache = std::unordered_map; + +static std::pair +specialize_arg(PyObject *backend, PyObject *arg, bool is_const, + bool specialize_value, bool align); + +static bool init_called = false; + +static PyObject *constexpr_cls = nullptr; +static PyObject *jit_callable_cls = nullptr; +static PyObject *tensor_descriptor_cls = nullptr; +static PyObject *nvidia_tensor_descriptor_cls = nullptr; +static PyObject *amd_tensor_descriptor_cls = nullptr; +static PyObject *canonicalize_dtype_fn = nullptr; +static PyObject *canonicalize_ptr_dtype_fn = nullptr; +static PyObject *torch_tensor_cls = nullptr; + +static PyObject *i32_str = nullptr; +static PyObject *i64_str = nullptr; +static PyObject *u64_str = nullptr; +static PyObject *fp32_str = nullptr; +static PyObject *u1_str = nullptr; +static PyObject *D_str = nullptr; +static PyObject *constexpr_str = nullptr; +static PyObject *empty_str = nullptr; +static PyObject *nvTmaDesc_str = nullptr; + +static PyObject *base_attr = nullptr; +static PyObject *data_ptr_attr = nullptr; +static PyObject *dtype_attr = nullptr; +static PyObject *cache_key_attr = nullptr; +static PyObject *_fields_attr = nullptr; +static PyObject *block_shape_attr = nullptr; +static PyObject *layout_attr = nullptr; +static PyObject *has_native_tensor_spec_attr = nullptr; +static PyObject *get_tensor_spec_attr = nullptr; +static PyObject *align_kwarg = nullptr; + +static DtypePtr2Str dtype_ptr2str; +static Dtype2Str dtype2str; +static TypeHandlerCache type_handler_cache; + +// Wrappers to make steal and borrow slightly simpler. We use raw CPython API +// with py::object to handle decref, as using the pybind11 APIs adds exception +// handling overhead which is quite significant here. +py::object from_new_ref(py::handle val) { + return py::reinterpret_steal(val); +} +py::object from_borrowed_ref(py::handle val) { + return py::reinterpret_borrow(val); +} + +PyObject *intern_from_string(const char *str) { + PyObject *obj = PyUnicode_InternFromString(str); + if (!obj) + throw py::error_already_set(); + return obj; +} + +PyObject *import_from(const char *module_name, const char *var_name) { + py::object var = py::module_::import(module_name).attr(var_name); + return var.release().ptr(); +} + +void init_interned_strings() { + i32_str = intern_from_string("i32"); + i64_str = intern_from_string("i64"); + u64_str = intern_from_string("u64"); + fp32_str = intern_from_string("fp32"); + u1_str = intern_from_string("u1"); + D_str = intern_from_string("D"); + constexpr_str = intern_from_string("constexpr"); + empty_str = intern_from_string(""); + nvTmaDesc_str = intern_from_string("nvTmaDesc"); + + base_attr = intern_from_string("base"); + data_ptr_attr = intern_from_string("data_ptr"); + dtype_attr = intern_from_string("dtype"); + cache_key_attr = intern_from_string("cache_key"); + _fields_attr = intern_from_string("_fields"); + block_shape_attr = intern_from_string("block_shape"); + layout_attr = intern_from_string("layout"); + has_native_tensor_spec_attr = + intern_from_string("supports_native_tensor_specialization"); + get_tensor_spec_attr = intern_from_string("get_tensor_specialization"); + + align_kwarg = py::make_tuple("align").release().ptr(); +} + +void init_type_handler_cache(); + +bool init_globals() noexcept try { + // Import releavant symbols + jit_callable_cls = import_from("triton.runtime.jit", "JITCallable"); + tensor_descriptor_cls = + import_from("triton.tools.tensor_descriptor", "TensorDescriptor"); + nvidia_tensor_descriptor_cls = import_from( + "triton.experimental.gluon.nvidia.hopper", "TensorDescriptor"); + amd_tensor_descriptor_cls = + import_from("triton.experimental.gluon.amd.gfx1250", "TensorDescriptor"); + + auto m_canonicalize = py::module_::import("triton._utils"); + canonicalize_dtype_fn = import_from("triton._utils", "canonicalize_dtype"); + canonicalize_ptr_dtype_fn = + import_from("triton._utils", "canonicalize_ptr_dtype"); + constexpr_cls = import_from("triton.language", "constexpr"); + + try { + torch_tensor_cls = import_from("torch", "Tensor"); + } catch (py::error_already_set &e) { + } + + init_interned_strings(); + init_type_handler_cache(); + + init_called = true; + return true; +} catch (py::error_already_set &e) { + e.restore(); + return false; +} + +std::pair specialize_tensordesc(PyObject *arg, + bool has_layout) { + auto base = from_new_ref(PyObject_GetAttr(arg, base_attr)); + if (!base) + return {}; + + auto dtype = from_new_ref(PyObject_GetAttr(base.ptr(), dtype_attr)); + if (!dtype) + return {}; + + PyObject *type_str; + Py_hash_t dtype_hash = PyObject_Hash(dtype.ptr()); + if (dtype_hash == -1) + return {}; + DTypeKey dsk{dtype_hash}; + auto it = dtype2str.find(dsk); + if (it != dtype2str.end()) { + type_str = it->second; + } else { + auto res = from_new_ref(PyObject_CallFunctionObjArgs(canonicalize_dtype_fn, + dtype.ptr(), nullptr)); + if (!res) + return {}; + dtype2str[dsk] = res.ptr(); + type_str = res.release().ptr(); + } + + std::string desc_cstr; + desc_cstr.reserve(128); + desc_cstr = "tensordesc<"; + auto dtype_str = from_new_ref(PyObject_Str(type_str)); + if (!dtype_str) + return {}; + + const char *dtype_cstr = PyUnicode_AsUTF8(dtype_str.ptr()); + if (!dtype_cstr) + return {}; + desc_cstr += dtype_cstr; + + auto block_shape_obj = from_new_ref(PyObject_GetAttr(arg, block_shape_attr)); + if (!block_shape_obj) + return {}; + auto block_shape_list = from_new_ref(PySequence_List(block_shape_obj.ptr())); + if (!block_shape_list) + return {}; + auto block_shape_str = from_new_ref(PyObject_Str(block_shape_list.ptr())); + if (!block_shape_str) + return {}; + const char *block_shape_cstr = PyUnicode_AsUTF8(block_shape_str.ptr()); + if (!block_shape_cstr) + return {}; + desc_cstr += block_shape_cstr; + + if (has_layout) { + auto layout_obj = from_new_ref(PyObject_GetAttr(arg, layout_attr)); + if (!layout_obj) + return {}; + auto layout_repr = from_new_ref(PyObject_Repr(layout_obj.ptr())); + if (!layout_repr) + return {}; + desc_cstr += ","; + const char *layout_cstr = PyUnicode_AsUTF8(layout_repr.ptr()); + if (!layout_cstr) + return {}; + desc_cstr += layout_cstr; + } + + desc_cstr += ">"; + auto type_str_result = from_new_ref(PyUnicode_FromString(desc_cstr.c_str())); + if (!type_str_result) + return {}; + + return {std::move(type_str_result), py::none()}; +} + +std::pair handle_long_type(PyObject *backend, + PyObject *arg, bool is_const, + bool specialize_value, + bool align) { + int overflow; + long long val = PyLong_AsLongLongAndOverflow(arg, &overflow); + if (PyErr_Occurred()) { + return {}; + } + + if (specialize_value && (val == 1)) { + return {from_borrowed_ref(constexpr_str), from_borrowed_ref(arg)}; + } + + py::handle type_str; + py::handle key_obj; + if (overflow == 0) { + type_str = (val >= INT32_MIN && val <= INT32_MAX) ? i32_str : i64_str; + if (specialize_value) { + key_obj = (align && ((val & 15) == 0)) ? D_str : empty_str; + } + } else { + unsigned long long val_64 = PyLong_AsUnsignedLongLong(arg); + if (PyErr_Occurred()) { + // this runs into an edge-case where the Python reference + // returns i64 as type and alignment of the value despite + // not being representable as such which at kernel launch later + // will throw an OverflowError nevertheless, here we throw + // OverflowError immediately + PyErr_SetString(PyExc_OverflowError, + "integer to be specialized too large to represent"); + return {}; + } + type_str = u64_str; + if (specialize_value) { + key_obj = (align && ((val_64 & 15) == 0)) ? D_str : empty_str; + } + } + if (!key_obj) { + return {from_borrowed_ref(type_str), py::none()}; + } + return {from_borrowed_ref(type_str), from_borrowed_ref(key_obj)}; +} + +std::pair handle_tensor(PyObject *backend, + PyObject *arg, bool is_const, + bool specialize_value, + bool align) { + // handle type_str specialization of a tensor + auto dtype = from_new_ref(PyObject_GetAttr(arg, dtype_attr)); + if (!dtype) + return {}; + + Py_hash_t dtype_hash = PyObject_Hash(dtype.ptr()); + if (dtype_hash == -1) + return {}; + + DTypePtrKey dsk{dtype_hash, is_const}; + auto it = dtype_ptr2str.find(dsk); + + py::handle type_str; + if (it != dtype_ptr2str.end()) { + type_str = it->second; + } else { + auto canon_res = + PyObject_CallFunctionObjArgs(canonicalize_ptr_dtype_fn, dtype.ptr(), + is_const ? Py_True : Py_False, nullptr); + if (!canon_res) + return {}; + dtype_ptr2str[dsk] = canon_res; + type_str = canon_res; + } + + // handle alignment specialization of a tensor + if (!specialize_value) { + return {from_borrowed_ref(type_str), py::none()}; + } + + bool native_impl_available = false; + auto native_spec_obj = + from_new_ref(PyObject_GetAttr(backend, has_native_tensor_spec_attr)); + if (native_spec_obj) { + native_impl_available = PyObject_IsTrue(native_spec_obj.ptr()); + } else { + PyErr_Clear(); + // on error we fall back to native_impl_available = false gracefully + } + + py::object key; + if (native_impl_available) { + auto data_ptr_result = + from_new_ref(PyObject_CallMethodNoArgs(arg, data_ptr_attr)); + if (!data_ptr_result) + return {}; + + auto data_ptr = PyLong_AsUnsignedLongLong(data_ptr_result.ptr()); + if (PyErr_Occurred()) + return {}; + + auto key_obj = (align && ((data_ptr & 15) == 0)) ? D_str : empty_str; + key = from_borrowed_ref(key_obj); + } else { + PyObject *args[3] = {backend, arg, align ? Py_True : Py_False}; + PyObject *kwnames = align_kwarg; + key = from_new_ref( + PyObject_VectorcallMethod(get_tensor_spec_attr, args, 2, kwnames)); + if (!key) + return {}; + } + + return {from_borrowed_ref(type_str), std::move(key)}; +} + +std::pair handle_bool_type(PyObject *backend, + PyObject *arg, bool is_const, + bool specialize_value, + bool align) { + return {from_borrowed_ref(u1_str), py::none()}; +} + +std::pair +handle_float_type(PyObject *backend, PyObject *arg, bool is_const, + bool specialize_value, bool align) { + return {from_borrowed_ref(fp32_str), py::none()}; +} + +std::pair +handle_tensor_descriptor(PyObject *backend, PyObject *arg, bool is_const, + bool specialize_value, bool align) { + return specialize_tensordesc(arg, false); +} + +std::pair +handle_gluon_tensor_descriptor(PyObject *backend, PyObject *arg, bool is_const, + bool specialize_value, bool align) { + return specialize_tensordesc(arg, true); +} + +std::pair +handle_constexpr_type(PyObject *backend, PyObject *arg, bool is_const, + bool specialize_value, bool align) { + return {from_borrowed_ref(constexpr_str), from_borrowed_ref(arg)}; +} + +std::pair +handle_jit_callable(PyObject *backend, PyObject *arg, bool is_const, + bool specialize_value, bool align) { + auto cache_key = from_new_ref(PyObject_GetAttr(arg, cache_key_attr)); + if (!cache_key) + return {}; + return {from_borrowed_ref(constexpr_str), std::move(cache_key)}; +} + +std::pair handle_tuple(PyObject *backend, PyObject *arg, + bool is_const, + bool specialize_value, + bool align) { + Py_ssize_t size = PyTuple_GET_SIZE(arg); + if (size == 0) { + // return tuple of empty tuples as in python reference + return {from_borrowed_ref(arg), from_borrowed_ref(arg)}; + } + + bool is_namedtuple = PyObject_HasAttr(arg, _fields_attr); + auto tuple_type = Py_TYPE(arg); + + // Create tuples directly instead of lists + auto tys_tuple = from_new_ref(PyTuple_New(size)); + if (!tys_tuple) + return {}; + + auto keys_tuple = from_new_ref(PyTuple_New(size)); + if (!keys_tuple) + return {}; + + for (Py_ssize_t i = 0; i < size; ++i) { + PyObject *item = PyTuple_GET_ITEM(arg, i); // Borrowed reference + // python reference calls specialize recursively with default arguments set + // currently this is is_const=False, specialize_value=True, align=True + auto [type, key] = specialize_arg(backend, item, false, true, true); + if (!type || !key) + return {}; + // Steals reference + PyTuple_SET_ITEM(tys_tuple.ptr(), i, type.release().ptr()); + PyTuple_SET_ITEM(keys_tuple.ptr(), i, key.release().ptr()); + } + + if (is_namedtuple) { + tys_tuple = from_new_ref( + PyObject_CallObject((PyObject *)tuple_type, tys_tuple.ptr())); + if (!tys_tuple) + return {}; + keys_tuple = from_new_ref( + PyObject_CallObject((PyObject *)tuple_type, keys_tuple.ptr())); + if (!keys_tuple) + return {}; + } + + return {std::move(tys_tuple), std::move(keys_tuple)}; +} + +// initialize type handler which returns specialize impelemntations based on +// type(arg) +void init_type_handler_cache() { + // Python Types (int, bool, float, tuple) + type_handler_cache[&PyLong_Type] = handle_long_type; + type_handler_cache[&PyBool_Type] = handle_bool_type; + type_handler_cache[&PyFloat_Type] = handle_float_type; + type_handler_cache[&PyTuple_Type] = handle_tuple; + + // torch.Tensor + if (torch_tensor_cls && PyType_Check(torch_tensor_cls)) { + type_handler_cache[(PyTypeObject *)torch_tensor_cls] = handle_tensor; + } + // TensorDescriptor + if (tensor_descriptor_cls && PyType_Check(tensor_descriptor_cls)) { + type_handler_cache[(PyTypeObject *)tensor_descriptor_cls] = + handle_tensor_descriptor; + } + // GluonTensorDescriptor + if (nvidia_tensor_descriptor_cls && + PyType_Check(nvidia_tensor_descriptor_cls)) { + type_handler_cache[(PyTypeObject *)nvidia_tensor_descriptor_cls] = + handle_gluon_tensor_descriptor; + } + if (amd_tensor_descriptor_cls && PyType_Check(amd_tensor_descriptor_cls)) { + type_handler_cache[(PyTypeObject *)amd_tensor_descriptor_cls] = + handle_gluon_tensor_descriptor; + } + // constexpr + if (constexpr_cls && PyType_Check(constexpr_cls)) { + type_handler_cache[(PyTypeObject *)constexpr_cls] = handle_constexpr_type; + } + // JITCallable + if (jit_callable_cls && PyType_Check(jit_callable_cls)) { + type_handler_cache[(PyTypeObject *)jit_callable_cls] = handle_jit_callable; + } +} + +// specialization logic without passing of objects from Python (to be called in +// specialize_impl only) +std::pair specialize_arg(PyObject *backend, + PyObject *arg, bool is_const, + bool specialize_value, + bool align) { + // fast-path for default types + PyTypeObject *arg_type = Py_TYPE(arg); + auto it = type_handler_cache.find(arg_type); + if (it != type_handler_cache.end()) { + return it->second(backend, arg, is_const, specialize_value, align); + } + + // separate handling of None + if (Py_IsNone(arg)) { + return {from_borrowed_ref(constexpr_str), py::none()}; + } + + // handling of sublcasses of tuples + if (PyTuple_Check(arg)) { + return handle_tuple(backend, arg, is_const, specialize_value, align); + } + + // fallback paths checking full inheritance + if (PyObject_IsInstance(arg, constexpr_cls)) { + return handle_constexpr_type(backend, arg, is_const, specialize_value, + align); + } + + if (PyObject_IsInstance(arg, tensor_descriptor_cls)) { + return handle_tensor_descriptor(backend, arg, is_const, specialize_value, + align); + } + + if (PyObject_IsInstance(arg, nvidia_tensor_descriptor_cls)) { + return handle_gluon_tensor_descriptor(backend, arg, is_const, + specialize_value, align); + } + + if (PyObject_IsInstance(arg, amd_tensor_descriptor_cls)) { + return handle_gluon_tensor_descriptor(backend, arg, is_const, + specialize_value, align); + } + + if (PyObject_IsInstance(arg, jit_callable_cls)) { + return handle_jit_callable(backend, arg, is_const, specialize_value, align); + } + + // fallback paths checking attributes directly + if (PyObject_HasAttr(arg, data_ptr_attr)) { + return handle_tensor(backend, arg, is_const, specialize_value, align); + } + + // fallback for default types + if (PyLong_Check(arg)) { + return handle_long_type(backend, arg, is_const, specialize_value, align); + } + if (PyFloat_Check(arg)) { + return handle_float_type(backend, arg, is_const, specialize_value, align); + } + + return {}; +} + +// main entry-point from Python implementing specialization logic natively +PyObject *specialize_impl(PyObject *self, PyObject *const *args, + Py_ssize_t nargs) { + if (!init_called) { + if (!init_globals()) { + return nullptr; + } + } + + if (nargs != 5) { + PyErr_SetString(PyExc_TypeError, + "native_specialize_impl expected 5 arguments"); + return nullptr; + } + + PyObject *backend = args[0]; + PyObject *arg = args[1]; + int is_const = PyObject_IsTrue(args[2]); + int specialize_value = PyObject_IsTrue(args[3]); + int align = PyObject_IsTrue(args[4]); + + if (is_const == -1 || specialize_value == -1 || align == -1) { + PyErr_SetString(PyExc_TypeError, "native_specialize_impl expected boolean " + "arguments for args2, args3, args4"); + return nullptr; + } + + auto [type, key] = + specialize_arg(backend, arg, is_const, specialize_value, align); + + // check if specialization failed + if (!type || !key) { + if (!PyErr_Occurred()) { + PyErr_Format(PyExc_TypeError, "failed to specialize argument of type: %s", + Py_TYPE(arg)->tp_name); + } + return nullptr; + } + + return PyTuple_Pack(2, type.ptr(), key.ptr()); +} + +static PyMethodDef module_methods[] = { + {"native_specialize_impl", (PyCFunction)specialize_impl, METH_FASTCALL, + nullptr}, + {nullptr, nullptr, 0, nullptr} // sentinel +}; + +} // anonymous namespace + +void init_native_specialize(pybind11::module &m) { + // add functions to module + PyModule_AddFunctions(m.ptr(), module_methods); +} diff --git a/third_party/sunrise/python/test/01-vector-add.py b/third_party/sunrise/python/test/01-vector-add.py new file mode 100644 index 0000000000..e62b49795a --- /dev/null +++ b/third_party/sunrise/python/test/01-vector-add.py @@ -0,0 +1,50 @@ +""" +Vector Addition (sunrise) +========================= + +Basic Triton vector-add smoke test on the sunrise backend. Rewritten to pytest +form to match the other sunrise test cases (02-06). +""" + +import torch +import triton +import triton.language as tl +import pytest + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@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(x, y, BLOCK_SIZE=1024): + output = torch.empty_like(x) + assert x.device == y.device == output.device + n_elements = output.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), ) + add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=BLOCK_SIZE) + return output + + +@pytest.mark.parametrize("size", [1024, 98432, 100000]) +def test_vector_add(size): + torch.manual_seed(0) + x = torch.rand(size, device=DEVICE) + y = torch.rand(size, device=DEVICE) + + out = add(x, y) + + torch.testing.assert_close(out.cpu(), (x + y).cpu(), atol=1e-6, rtol=1e-5) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/third_party/sunrise/python/test/02-tle-local_ptr.py b/third_party/sunrise/python/test/02-tle-local_ptr.py new file mode 100644 index 0000000000..ff982a968a --- /dev/null +++ b/third_party/sunrise/python/test/02-tle-local_ptr.py @@ -0,0 +1,110 @@ +# flagtree tle +""" +TLE Local Pointer (no-TMA) End-to-End Test +========================================== + +Exercises the TLE local-pointer path that is portable across backends that +lower ``tle.local_pointers`` to ``ttg.local_*`` (NVIDIA, sunrise, ...): + + tle.gpu.alloc(nv_mma_shared_layout=False) # SwizzledShared smem buffer + tle.gpu.local_ptr(buffer, indices) # materialize shared pointers + tl.store(local_ptr, value) # stage into shared memory + tl.load(local_ptr) # read back from shared memory + +It deliberately AVOIDS ``tle.gpu.copy`` because that lowers to a TMA copy +(``create_tma_copy``), which is NVIDIA Hopper specific and unsupported on +sunrise. + +What this covers +---------------- +* The local-pointer TTGIR pass chain wired into the backend pipeline + (early_assign_memory_space / select_encodings / insert_local_pointer_barriers + / optimize_local_pointer_loads / optimize_local_pointer_stores). +* The store->load visibility barrier inserted by + insert_local_pointer_barriers (correctness: the value read back must equal + the value staged). +* The shared-store/-load alignment guard in the backend LoadStore lowering -- + forced to matter by sweeping inner block widths, including non-power-of-two + widths (48 / 96) that the guard must clamp the vector width for. +""" + +import pytest +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@triton.jit +def shared_roundtrip_kernel( + src_ptr, + dst_ptr, + xstride, + ynumel, + XBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, +): + """Per program (one X-block), for each Y-block: + + gmem -> (tl.store) shared -> (tl.load) shared -> (scale) -> gmem + + The shared store and load both go through ``tle.gpu.local_ptr`` pointers. + """ + pid = tl.program_id(0) + + xoffs = pid * XBLOCK + tl.arange(0, XBLOCK) + src_rows = src_ptr + xstride * xoffs[:, None] + dst_rows = dst_ptr + xstride * xoffs[:, None] + + smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, + scope=tle.gpu.smem, nv_mma_shared_layout=False) + + row_ids = tl.arange(0, XBLOCK)[:, None] + col_ids = tl.arange(0, YBLOCK)[None, :] + row_ids = tl.broadcast_to(row_ids, (XBLOCK, YBLOCK)) + col_ids = tl.broadcast_to(col_ids, (XBLOCK, YBLOCK)) + smem_ptrs = tle.gpu.local_ptr(smem, (row_ids, col_ids)) + + for yoff in range(0, ynumel, YBLOCK): + yoffs = tl.arange(0, YBLOCK) + yoff + gval = tl.load(src_rows + yoffs[None, :]) + + tl.store(smem_ptrs, gval) # vectorized shared store (guarded) + sval = tl.load(smem_ptrs) # vectorized shared load (guarded) + + tl.store(dst_rows + yoffs[None, :], sval * 2.0) + + +def shared_roundtrip(src, dst, XBLOCK, YBLOCK): + assert src.shape == dst.shape + xnumel, ynumel = src.shape + assert ynumel % YBLOCK == 0, "kernel assumes ynumel divisible by YBLOCK" + grid = (triton.cdiv(xnumel, XBLOCK), ) + return shared_roundtrip_kernel[grid](src, dst, src.stride(0), ynumel, XBLOCK, YBLOCK) + + +class TestTLELocalPtrRoundtrip: + """Local-pointer staging correctness across vectorizable / non-vectorizable widths.""" + + @pytest.mark.parametrize("XBLOCK,YBLOCK", [ + (32, 64), + (32, 128), + (16, 64), + (64, 64), + ]) + def test_roundtrip(self, XBLOCK, YBLOCK): + torch.manual_seed(0) + xnumel, ynumel = 256, YBLOCK * 4 + + src = torch.randn(xnumel, ynumel, device=DEVICE, dtype=torch.float32) + dst = torch.empty_like(src) + + shared_roundtrip(src, dst, XBLOCK, YBLOCK) + + torch.testing.assert_close(dst.cpu(), src.cpu() * 2.0, atol=1e-5, rtol=1e-5) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/third_party/sunrise/python/test/03-tle-extract_insert_tile.py b/third_party/sunrise/python/test/03-tle-extract_insert_tile.py new file mode 100644 index 0000000000..dd9b007ef0 --- /dev/null +++ b/third_party/sunrise/python/test/03-tle-extract_insert_tile.py @@ -0,0 +1,261 @@ +# flagtree tle +""" +TLE extract_tile / insert_tile end-to-end test (sunrise) +======================================================== + +Exercises ``tle.extract_tile`` and ``tle.insert_tile`` on the sunrise backend, +covering BOTH lowering paths now that sunrise has its own tile-op lowering +(third_party/sunrise/plugin/lib/TritonSunriseGPUToLLVM/TileOpsToLLVM.cpp): + + * STATIC + CTA-tile aligned index -> register-permutation path (no SMEM). + * dynamic index OR non-aligned tile -> SMEM relay path. The sunrise lowering + routes the barrier through targetInfo.barrier() (-> STVM::Barrier0Op) and + the thread id through getThreadId() (-> STVM), so unlike the common TLE + lowering it does not require NVVM and works on sunrise. + +SMEM relay is reached via two distinct entries, both covered here: + - dynamic (runtime) index (test_*_dynamic_index), and + - static index with a CTA-tile non-aligned (thin) tile shape + (test_*_static_nonaligned). + +The (M, N, tile, index) values mirror python/test/tle/unit/test_extract_tile_static_index.py. +""" + +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle +import pytest + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +# --------------------------------------------------------------------------- +# Static, CTA-tile-aligned index -> register path +# --------------------------------------------------------------------------- +@triton.jit +def extract_tile_kernel(x_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, TN: tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :]) + + tile = tle.extract_tile(x, index=[1, 1], tile_shape=[TM, TN]) + + out_m = tl.arange(0, TM) + out_n = tl.arange(0, TN) + tl.store(out_ptr + out_m[:, None] * TN + out_n[None, :], tile) + + +@triton.jit +def insert_tile_kernel(x_ptr, y_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, + TN: tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :]) + + tile_m = tl.arange(0, TM) + tile_n = tl.arange(0, TN) + y = tl.load(y_ptr + tile_m[:, None] * TN + tile_n[None, :]) + + z = tle.insert_tile(x, y, index=[1, 1]) + + tl.store(out_ptr + offs_m[:, None] * N + offs_n[None, :], z) + + +# --------------------------------------------------------------------------- +# Dynamic (runtime) index -> SMEM relay path +# --------------------------------------------------------------------------- +@triton.jit +def extract_tile_dyn_kernel(x_ptr, out_ptr, idx, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, + TN: tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :]) + + # Runtime linear tile index -> always uses the SMEM relay path. + lin = tl.load(idx) + tile = tle.extract_tile(x, index=lin, tile_shape=[TM, TN]) + + out_m = tl.arange(0, TM) + out_n = tl.arange(0, TN) + tl.store(out_ptr + out_m[:, None] * TN + out_n[None, :], tile) + + +@triton.jit +def insert_tile_dyn_kernel(x_ptr, y_ptr, out_ptr, idx, M: tl.constexpr, N: tl.constexpr, + TM: tl.constexpr, TN: tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :]) + + tile_m = tl.arange(0, TM) + tile_n = tl.arange(0, TN) + y = tl.load(y_ptr + tile_m[:, None] * TN + tile_n[None, :]) + + lin = tl.load(idx) + z = tle.insert_tile(x, y, index=lin) + + tl.store(out_ptr + offs_m[:, None] * N + offs_n[None, :], z) + + +def _lin_index(grid_n, ti, tj, device): + """row-major linear tile index for tile coord (ti, tj).""" + return torch.tensor([ti * grid_n + tj], device=device, dtype=torch.int32) + + +# --------------------------------------------------------------------------- +# Register-path tests (static aligned) +# --------------------------------------------------------------------------- +def test_extract_tile_static_aligned(): + M, N = 512, 512 + TM, TN = 128, 128 + + x = torch.arange(M * N, device=DEVICE, dtype=torch.float32).reshape(M, N) + out = torch.zeros((TM, TN), device=DEVICE, dtype=torch.float32) + + extract_tile_kernel[(1, )](x, out, M, N, TM, TN) + + expected = x[TM:2 * TM, TN:2 * TN] + torch.testing.assert_close(out.cpu(), expected.cpu(), atol=0, rtol=0) + + +def test_insert_tile_static_aligned(): + M, N = 128, 128 + TM, TN = 32, 32 + + x = torch.arange(M * N, device=DEVICE, dtype=torch.float32).reshape(M, N) + y = (100000 + torch.arange(TM * TN, device=DEVICE, dtype=torch.float32)).reshape(TM, TN) + out = torch.empty_like(x) + + insert_tile_kernel[(1, )](x, y, out, M, N, TM, TN) + + expected = x.clone() + expected[TM:2 * TM, TN:2 * TN] = y + torch.testing.assert_close(out.cpu(), expected.cpu(), atol=0, rtol=0) + + +def test_extract_then_insert_roundtrip(): + """extract a tile then insert it back at the same index -> identity.""" + M, N = 128, 128 + TM, TN = 32, 32 + + x = torch.randn(M, N, device=DEVICE, dtype=torch.float32) + tile = torch.empty((TM, TN), device=DEVICE, dtype=torch.float32) + out = torch.empty_like(x) + + extract_tile_kernel[(1, )](x, tile, M, N, TM, TN) + insert_tile_kernel[(1, )](x, tile, out, M, N, TM, TN) + + torch.testing.assert_close(out.cpu(), x.cpu(), atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# SMEM-path tests (dynamic index) +# --------------------------------------------------------------------------- +@pytest.mark.parametrize("ti,tj", [(0, 0), (1, 1), (2, 3)]) +def test_extract_tile_dynamic_index(ti, tj): + M, N = 128, 128 + TM, TN = 32, 32 + grid_n = N // TN + + x = torch.arange(M * N, device=DEVICE, dtype=torch.float32).reshape(M, N) + out = torch.zeros((TM, TN), device=DEVICE, dtype=torch.float32) + idx = _lin_index(grid_n, ti, tj, DEVICE) + + extract_tile_dyn_kernel[(1, )](x, out, idx, M, N, TM, TN) + + expected = x[ti * TM:(ti + 1) * TM, tj * TN:(tj + 1) * TN] + torch.testing.assert_close(out.cpu(), expected.cpu(), atol=0, rtol=0) + + +@pytest.mark.parametrize("ti,tj", [(0, 0), (1, 1), (2, 3)]) +def test_insert_tile_dynamic_index(ti, tj): + M, N = 128, 128 + TM, TN = 32, 32 + grid_n = N // TN + + x = torch.arange(M * N, device=DEVICE, dtype=torch.float32).reshape(M, N) + y = (100000 + torch.arange(TM * TN, device=DEVICE, dtype=torch.float32)).reshape(TM, TN) + out = torch.empty_like(x) + idx = _lin_index(grid_n, ti, tj, DEVICE) + + insert_tile_dyn_kernel[(1, )](x, y, out, idx, M, N, TM, TN) + + expected = x.clone() + expected[ti * TM:(ti + 1) * TM, tj * TN:(tj + 1) * TN] = y + torch.testing.assert_close(out.cpu(), expected.cpu(), atol=0, rtol=0) + + +# --------------------------------------------------------------------------- +# Static but CTA-tile NON-aligned index -> SMEM relay path +# +# isCTATileAligned() returns false when tileShape[i] % ctaTile[i] != 0 (or the +# tile offset is not a multiple of ctaTile[i]); ctaTile = getShapePerCTATile is +# derived from the source blocked encoding. A very "thin" tile dimension (e.g. +# 2) is almost never a multiple of the per-CTA tile, so even with a static index +# the conversion takes the SMEM relay path (lowerExtractTileViaSMEM), exercising +# the targetInfo.barrier()/getThreadId() lowering rather than the register path. +# This is a different SMEM-path entry than the dynamic-index tests above. +# --------------------------------------------------------------------------- +@triton.jit +def extract_tile_thin_kernel(x_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, + TN: tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :]) + + # static index, but thin tile -> CTA-tile non-aligned -> SMEM relay path + tile = tle.extract_tile(x, index=[1, 0], tile_shape=[TM, TN]) + + out_m = tl.arange(0, TM) + out_n = tl.arange(0, TN) + tl.store(out_ptr + out_m[:, None] * TN + out_n[None, :], tile) + + +@triton.jit +def insert_tile_thin_kernel(x_ptr, y_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, + TN: tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :]) + + tile_m = tl.arange(0, TM) + tile_n = tl.arange(0, TN) + y = tl.load(y_ptr + tile_m[:, None] * TN + tile_n[None, :]) + + z = tle.insert_tile(x, y, index=[1, 0]) + + tl.store(out_ptr + offs_m[:, None] * N + offs_n[None, :], z) + + +@pytest.mark.parametrize("M,N,TM,TN", [(32, 32, 2, 32), (32, 32, 4, 32)]) +def test_extract_tile_static_nonaligned(M, N, TM, TN): + """Static index + thin tile -> SMEM relay path (non-aligned).""" + torch.manual_seed(3) + x = torch.arange(M * N, device=DEVICE, dtype=torch.float32).reshape(M, N) + out = torch.zeros((TM, TN), device=DEVICE, dtype=torch.float32) + + extract_tile_thin_kernel[(1, )](x, out, M, N, TM, TN) + + # index=[1,0] -> rows [TM:2*TM], cols [0:TN] + expected = x[TM:2 * TM, 0:TN] + torch.testing.assert_close(out.cpu(), expected.cpu(), atol=0, rtol=0) + + +@pytest.mark.parametrize("M,N,TM,TN", [(32, 32, 2, 32), (32, 32, 4, 32)]) +def test_insert_tile_static_nonaligned(M, N, TM, TN): + """Static index + thin tile insert -> SMEM relay path (non-aligned).""" + torch.manual_seed(4) + x = torch.arange(M * N, device=DEVICE, dtype=torch.float32).reshape(M, N) + y = (100000 + torch.arange(TM * TN, device=DEVICE, dtype=torch.float32)).reshape(TM, TN) + out = torch.empty_like(x) + + insert_tile_thin_kernel[(1, )](x, y, out, M, N, TM, TN) + + expected = x.clone() + expected[TM:2 * TM, 0:TN] = y + torch.testing.assert_close(out.cpu(), expected.cpu(), atol=0, rtol=0) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/third_party/sunrise/python/test/04-tle-hint-shared-memory.py b/third_party/sunrise/python/test/04-tle-hint-shared-memory.py new file mode 100644 index 0000000000..e605deae1b --- /dev/null +++ b/third_party/sunrise/python/test/04-tle-hint-shared-memory.py @@ -0,0 +1,55 @@ +# flagtree tle +""" +TLE #@hint: shared_memory end-to-end test (sunrise) +=================================================== + +Verifies the source-comment hint path on sunrise: + + x = tl.load(x_ptr + offs) #@hint: shared_memory + +The SunriseHintHandler parses the `#@hint: shared_memory` comment and injects it +as the load's `flagtree_hints` attr; the `add_process_shared_memory_hint` TTGIR +pass then rewrites that load into the async copy -> shared -> local_load chain +(local_alloc + async_copy_global_to_local + async_commit_group + async_wait + +local_load). On sunrise these lower to async_load.* + STVM::Barrier0Op, so the +result must be numerically identical to a plain load. + +Requires the sunrise hint handler to be registered in +triton/compiler/hint_manager.py (backend 'sunrise', aliases 'ptpu'/'tang'). +""" + +import torch +import triton +import triton.language as tl +import pytest + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@triton.jit +def hinted_add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + # The hint comment must sit on the load line so the handler maps it by lineno. + x = tl.load(x_ptr + offs, mask=mask) #@hint: shared_memory + y = tl.load(y_ptr + offs, mask=mask) + tl.store(out_ptr + offs, x + y, mask=mask) + + +@pytest.mark.parametrize("n_elements", [128, 256, 512]) +def test_hint_shared_memory(n_elements): + torch.manual_seed(0) + BLOCK = 128 + x = torch.randn(n_elements, device=DEVICE, dtype=torch.float32) + y = torch.randn(n_elements, device=DEVICE, dtype=torch.float32) + out = torch.empty_like(x) + + grid = (triton.cdiv(n_elements, BLOCK), ) + hinted_add_kernel[grid](x, y, out, n_elements, BLOCK) + + torch.testing.assert_close(out.cpu(), (x + y).cpu(), atol=1e-5, rtol=1e-5) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/third_party/sunrise/python/test/05-tle-copy-normcopy.py b/third_party/sunrise/python/test/05-tle-copy-normcopy.py new file mode 100644 index 0000000000..5b1b06c44a --- /dev/null +++ b/third_party/sunrise/python/test/05-tle-copy-normcopy.py @@ -0,0 +1,144 @@ +# flagtree tle +""" +TLE tle.gpu.copy (normcopy / no-TMA) end-to-end test (sunrise) +============================================================= + +tle.gpu.copy has two paths, chosen by operand type: + + * normcopy -- src/dst are tl.tensor <-> tle.buffered_tensor (plain pointers). + Lowers to tl.load + local_ptr + tl.store. NO TMA. Portable. + * tmacopy -- src/dst involve tl.tensor_descriptor. Lowers to create_tma_copy + (NVIDIA Hopper TMA). NOT supported on sunrise. + +This test exercises ONLY the normcopy path (plain pointer tensors), in both +directions: + + * GM_TO_LOCAL : copy(global_ptrs, smem_buffer, shape) + * LOCAL_TO_GM : copy(smem_buffer, global_ptrs, shape) + +It relies on the local_ptr support wired into the sunrise backend. Because +normcopy is just load/store through local pointers, no TMA is involved and it +should run on sunrise. +""" + +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle +import pytest + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@triton.jit +def copy_roundtrip_kernel( + a_ptr, + c_ptr, + xnumel, + ynumel, + xstride, + ystride, + XBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, +): + """For each X-block, for each Y-block: + + gmem --copy(GM_TO_LOCAL)--> smem --(scale via local_ptr load/store)--> + smem --copy(LOCAL_TO_GM)--> gmem + + Exercises both copy directions plus local_ptr load/store on the staged tile. + """ + pid = tl.program_id(0) + xoffs = pid * XBLOCK + tl.arange(0, XBLOCK) + a_rows = a_ptr + xstride * xoffs[:, None] + c_rows = c_ptr + xstride * xoffs[:, None] + + smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, + scope=tle.gpu.smem, nv_mma_shared_layout=False) + + row_ids = tl.arange(0, XBLOCK)[:, None] + col_ids = tl.arange(0, YBLOCK)[None, :] + row_ids = tl.broadcast_to(row_ids, (XBLOCK, YBLOCK)) + col_ids = tl.broadcast_to(col_ids, (XBLOCK, YBLOCK)) + smem_ptrs = tle.gpu.local_ptr(smem, (row_ids, col_ids)) + + for yoff in range(0, ynumel, YBLOCK): + yoffs = tl.arange(0, YBLOCK) + yoff + + # GM_TO_LOCAL: plain pointer tensor src -> smem buffer dst (normcopy) + tle.gpu.copy(a_rows + ystride * yoffs[None, :], smem, [XBLOCK, YBLOCK]) + + # scale in shared memory through local pointers + val = tl.load(smem_ptrs) + tl.store(smem_ptrs, val * 2.0) + + # LOCAL_TO_GM: smem buffer src -> plain pointer tensor dst (normcopy) + tle.gpu.copy(smem, c_rows + ystride * yoffs[None, :], [XBLOCK, YBLOCK]) + + +def copy_roundtrip(a, c, XBLOCK, YBLOCK): + assert a.shape == c.shape + xnumel, ynumel = a.shape + assert ynumel % YBLOCK == 0 + grid = (triton.cdiv(xnumel, XBLOCK), ) + return copy_roundtrip_kernel[grid](a, c, xnumel, ynumel, a.stride(0), a.stride(1), XBLOCK, YBLOCK) + + +@pytest.mark.parametrize("XBLOCK,YBLOCK", [(32, 64), (64, 64), (32, 128)]) +def test_copy_roundtrip(XBLOCK, YBLOCK): + torch.manual_seed(0) + xnumel, ynumel = 256, YBLOCK * 4 + + a = torch.randn(xnumel, ynumel, device=DEVICE, dtype=torch.float32) + c = torch.empty_like(a) + + copy_roundtrip(a, c, XBLOCK, YBLOCK) + + torch.testing.assert_close(c.cpu(), (a * 2.0).cpu(), atol=1e-5, rtol=1e-5) + + +@triton.jit +def copy_gm_to_local_kernel( + a_ptr, + c_ptr, + ynumel, + xstride, + ystride, + XBLOCK: tl.constexpr, + YBLOCK: tl.constexpr, +): + """Pure GM_TO_LOCAL copy then read back via local_ptr and store to gmem.""" + pid = tl.program_id(0) + xoffs = pid * XBLOCK + tl.arange(0, XBLOCK) + a_rows = a_ptr + xstride * xoffs[:, None] + c_rows = c_ptr + xstride * xoffs[:, None] + + smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, + scope=tle.gpu.smem, nv_mma_shared_layout=False) + row_ids = tl.broadcast_to(tl.arange(0, XBLOCK)[:, None], (XBLOCK, YBLOCK)) + col_ids = tl.broadcast_to(tl.arange(0, YBLOCK)[None, :], (XBLOCK, YBLOCK)) + smem_ptrs = tle.gpu.local_ptr(smem, (row_ids, col_ids)) + + for yoff in range(0, ynumel, YBLOCK): + yoffs = tl.arange(0, YBLOCK) + yoff + tle.gpu.copy(a_rows + ystride * yoffs[None, :], smem, [XBLOCK, YBLOCK]) + val = tl.load(smem_ptrs) + tl.store(c_rows + ystride * yoffs[None, :], val) + + +def test_copy_gm_to_local_identity(): + torch.manual_seed(1) + XBLOCK, YBLOCK = 32, 64 + xnumel, ynumel = 256, YBLOCK * 4 + + a = torch.randn(xnumel, ynumel, device=DEVICE, dtype=torch.float32) + c = torch.empty_like(a) + + grid = (triton.cdiv(xnumel, XBLOCK), ) + copy_gm_to_local_kernel[grid](a, c, ynumel, a.stride(0), a.stride(1), XBLOCK, YBLOCK) + + torch.testing.assert_close(c.cpu(), a.cpu(), atol=0, rtol=0) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/third_party/sunrise/python/test/06-tle-cumsum.py b/third_party/sunrise/python/test/06-tle-cumsum.py new file mode 100644 index 0000000000..ef501f3fb1 --- /dev/null +++ b/third_party/sunrise/python/test/06-tle-cumsum.py @@ -0,0 +1,82 @@ +# flagtree tle +""" +TLE tle.cumsum end-to-end test (sunrise) +======================================== + +tle.cumsum computes an EXCLUSIVE CTA-wide prefix sum and returns +(exclusive_sum, total_sum). It is a single TLE op (ExclusiveCumsumOp), distinct +from upstream tl.cumsum (which is inclusive and tensor-local). + +On sunrise the lowering uses a sunrise-local conversion +(third_party/sunrise/plugin/lib/TritonSunriseGPUToLLVM/ExclusiveCumsumOpToLLVM.cpp) +that drops the NVVM `shfl.sync.up.b32` i32 fast path and instead always uses the +generic targetInfo.shuffleUp / shuffleIdx / shared-memory cross-warp scan, which +lowers to STVM. Scratch for the scan is reserved by the sunrise TLE-aware +allocation pass. + +reference: exclusive[i] = sum(x[:i]); total = sum(x). +""" + +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle +import pytest + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@triton.jit +def cumsum_kernel(x_ptr, exclusive_ptr, total_ptr, n, BLOCK: tl.constexpr, + REVERSE: tl.constexpr): + offs = tl.arange(0, BLOCK) + mask = offs < n + x = tl.load(x_ptr + offs, mask=mask, other=0) + exclusive, total = tle.cumsum(x, axis=0, reverse=REVERSE) + tl.store(exclusive_ptr + offs, exclusive, mask=mask) + tl.store(total_ptr, total) + + +def _ref_exclusive(x, reverse): + if reverse: + # reverse-exclusive: exclusive[i] = sum(x[i+1:]) + rev = torch.flip(x, dims=[0]) + excl_rev = torch.cumsum(rev, dim=0) - rev + return torch.flip(excl_rev, dims=[0]) + return torch.cumsum(x, dim=0) - x + + +@pytest.mark.parametrize("reverse", [False, True]) +@pytest.mark.parametrize("n,BLOCK", [(128, 128), (256, 256), (100, 128)]) +def test_cumsum_int32(n, BLOCK, reverse): + torch.manual_seed(0) + x = torch.randint(0, 10, (BLOCK, ), device=DEVICE, dtype=torch.int32) + x[n:] = 0 # masked-out tail + exclusive = torch.empty(BLOCK, device=DEVICE, dtype=torch.int32) + total = torch.zeros(1, device=DEVICE, dtype=torch.int32) + + cumsum_kernel[(1, )](x, exclusive, total, n, BLOCK, reverse) + + xv = x[:n].to(torch.int64) + ref_excl = _ref_exclusive(xv, reverse) + torch.testing.assert_close(exclusive[:n].to(torch.int64).cpu(), ref_excl.cpu(), atol=0, rtol=0) + torch.testing.assert_close(int(total.item()), int(xv.sum().item())) + + +@pytest.mark.parametrize("n,BLOCK", [(128, 128), (256, 256)]) +def test_cumsum_float32(n, BLOCK): + torch.manual_seed(1) + x = torch.randn(BLOCK, device=DEVICE, dtype=torch.float32) + x[n:] = 0.0 + exclusive = torch.empty(BLOCK, device=DEVICE, dtype=torch.float32) + total = torch.zeros(1, device=DEVICE, dtype=torch.float32) + + cumsum_kernel[(1, )](x, exclusive, total, n, BLOCK, False) + + ref_excl = torch.cumsum(x[:n], dim=0) - x[:n] + torch.testing.assert_close(exclusive[:n].cpu(), ref_excl.cpu(), atol=1e-4, rtol=1e-4) + torch.testing.assert_close(total.item(), x[:n].sum().item(), atol=1e-4, rtol=1e-4) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/third_party/sunrise/python/test/07-tle-load.py b/third_party/sunrise/python/test/07-tle-load.py new file mode 100644 index 0000000000..a75a5dd9f9 --- /dev/null +++ b/third_party/sunrise/python/test/07-tle-load.py @@ -0,0 +1,78 @@ +# flagtree tle +""" +TLE tle.load end-to-end test (sunrise) +====================================== + +tle.load is tl.load plus a `tt.load.async` bool attribute (see +python/triton/experimental/tle/language/core.py). When is_async=True, the +add_lower_async_load TTGIR pass (wired into the sunrise pipeline) rewrites the +load into the async copy chain: local_alloc + async_copy_global_to_local + +async_commit_group + async_wait + local_load. On sunrise these lower to +async_load.* + STVM::Barrier0Op (the async_wait barrier), so the result must be +numerically identical to a plain tl.load. is_async=False is exactly tl.load. + +This test loads through tle.load (both async modes), adds, and stores; result +must match a torch reference. +""" + +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle +import pytest + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@triton.jit +def tle_load_add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK: tl.constexpr, + IS_ASYNC: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n_elements + # tle.load: tl.load + tt.load.async attr. IS_ASYNC=True -> async copy chain. + x = tle.load(x_ptr + offs, mask=mask, other=0.0, is_async=IS_ASYNC) + y = tle.load(y_ptr + offs, mask=mask, other=0.0, is_async=IS_ASYNC) + tl.store(out_ptr + offs, x + y, mask=mask) + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("n_elements", [128, 256, 1000]) +def test_tle_load_add(n_elements, is_async): + torch.manual_seed(0) + BLOCK = 128 + x = torch.randn(n_elements, device=DEVICE, dtype=torch.float32) + y = torch.randn(n_elements, device=DEVICE, dtype=torch.float32) + out = torch.empty_like(x) + + grid = (triton.cdiv(n_elements, BLOCK), ) + tle_load_add_kernel[grid](x, y, out, n_elements, BLOCK, is_async) + + torch.testing.assert_close(out.cpu(), (x + y).cpu(), atol=1e-5, rtol=1e-5) + + +@triton.jit +def tle_load_2d_kernel(x_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, + IS_ASYNC: tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + ptrs = x_ptr + offs_m[:, None] * N + offs_n[None, :] + # 2D async load -> async copy of a tile into shared, then local_load. + x = tle.load(ptrs, is_async=IS_ASYNC) + tl.store(out_ptr + offs_m[:, None] * N + offs_n[None, :], x * 2.0) + + +@pytest.mark.parametrize("is_async", [False, True]) +@pytest.mark.parametrize("M,N", [(32, 64), (64, 64)]) +def test_tle_load_2d(M, N, is_async): + torch.manual_seed(1) + x = torch.randn(M, N, device=DEVICE, dtype=torch.float32) + out = torch.empty_like(x) + + tle_load_2d_kernel[(1, )](x, out, M, N, is_async) + + torch.testing.assert_close(out.cpu(), (x * 2.0).cpu(), atol=1e-5, rtol=1e-5) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/third_party/sunrise/python/test/08-tle-dot-local-ptr.py b/third_party/sunrise/python/test/08-tle-dot-local-ptr.py new file mode 100644 index 0000000000..0aeb44e160 --- /dev/null +++ b/third_party/sunrise/python/test/08-tle-dot-local-ptr.py @@ -0,0 +1,94 @@ +# flagtree tle +""" +TLE GEMM via local_ptr staging + tl.dot (sunrise) +================================================= + +Exercises the most common TLE pattern that was previously untested on sunrise: +shared-memory staging through tle.gpu.local_ptr feeding tl.dot (a GEMM tile). + +This also indirectly covers: + * the local_ptr TTGIR pass chain (early_assign_memory_space / select_encodings + / insert_local_pointer_barriers / optimize_local_pointer_loads/_stores), + * the gpu.barrier -> STVM::Barrier0Op lowering between local store and load, + * the AABS dot-dimension floor for sunrise (M>=8, N>=8, K>=16) when run under + autotune with FLAGTREE_AABS=1. + +Single-block GEMM (one program computes the whole MxN tile from MxK * KxN) keeps +the kernel simple and backend-portable; dims are multiples of the sunrise +min_dot_size (M=8,N=8,K=16/4). +""" + +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle +import pytest + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + + +@triton.jit +def gemm_local_ptr_kernel(a_ptr, b_ptr, c_ptr, M: tl.constexpr, N: tl.constexpr, + K: tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + offs_k = tl.arange(0, K) + + # Load A (MxK) and B (KxN) from global memory. + a = tl.load(a_ptr + offs_m[:, None] * K + offs_k[None, :]) + b = tl.load(b_ptr + offs_k[:, None] * N + offs_n[None, :]) + + # Stage A through a shared-memory buffer via local_ptr (the TLE path), then + # read it back and feed tl.dot. This forces local store -> barrier -> load. + a_smem = tle.gpu.alloc([M, K], dtype=tl.float32, layout=None, scope=tle.gpu.smem, + nv_mma_shared_layout=False) + row = tl.broadcast_to(tl.arange(0, M)[:, None], (M, K)) + col = tl.broadcast_to(tl.arange(0, K)[None, :], (M, K)) + a_ptrs = tle.gpu.local_ptr(a_smem, (row, col)) + tl.store(a_ptrs, a) + a_staged = tl.load(a_ptrs) + + acc = tl.dot(a_staged, b) + tl.store(c_ptr + offs_m[:, None] * N + offs_n[None, :], acc) + + +@pytest.mark.parametrize("M,N,K", [(16, 16, 16), (32, 32, 32), (8, 8, 16)]) +def test_gemm_local_ptr(M, N, K): + torch.manual_seed(0) + a = torch.randn(M, K, device=DEVICE, dtype=torch.float32) + b = torch.randn(K, N, device=DEVICE, dtype=torch.float32) + c = torch.empty(M, N, device=DEVICE, dtype=torch.float32) + + gemm_local_ptr_kernel[(1, )](a, b, c, M, N, K) + + ref = a @ b + torch.testing.assert_close(c.cpu(), ref.cpu(), atol=1e-2, rtol=1e-2) + + +@triton.jit +def gemm_plain_kernel(a_ptr, b_ptr, c_ptr, M: tl.constexpr, N: tl.constexpr, + K: tl.constexpr): + offs_m = tl.arange(0, M) + offs_n = tl.arange(0, N) + offs_k = tl.arange(0, K) + a = tl.load(a_ptr + offs_m[:, None] * K + offs_k[None, :]) + b = tl.load(b_ptr + offs_k[:, None] * N + offs_n[None, :]) + acc = tl.dot(a, b) + tl.store(c_ptr + offs_m[:, None] * N + offs_n[None, :], acc) + + +@pytest.mark.parametrize("M,N,K", [(16, 16, 16), (32, 64, 32)]) +def test_gemm_plain(M, N, K): + """Plain tl.dot (no local_ptr) baseline; confirms dot path itself works.""" + torch.manual_seed(1) + a = torch.randn(M, K, device=DEVICE, dtype=torch.float32) + b = torch.randn(K, N, device=DEVICE, dtype=torch.float32) + c = torch.empty(M, N, device=DEVICE, dtype=torch.float32) + + gemm_plain_kernel[(1, )](a, b, c, M, N, K) + + torch.testing.assert_close(c.cpu(), (a @ b).cpu(), atol=1e-2, rtol=1e-2) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/third_party/sunrise/python/test/09-tle-negative-contract.py b/third_party/sunrise/python/test/09-tle-negative-contract.py new file mode 100644 index 0000000000..318faf52d3 --- /dev/null +++ b/third_party/sunrise/python/test/09-tle-negative-contract.py @@ -0,0 +1,113 @@ +# flagtree tle +""" +TLE negative / contract tests (sunrise) +======================================= + +SKILL requires every semantic surface to have a negative (contract) case, not +just positive ones. These assert that misuse raises at compile time rather than +miscompiling or crashing, and that Hopper-only TLE features fail with a clear +error on sunrise instead of reaching an NVVM lowering that cannot legalize. + +All checks fire during JIT tracing, so the frontend ValueError is wrapped in a +triton CompilationError (original is in __cause__). We accept either. +""" + +import torch +import triton +import triton.language as tl +import triton.experimental.tle.language as tle +import pytest + +DEVICE = triton.runtime.driver.active.get_active_torch_device() + +# JIT tracing wraps frontend ValueError in CompilationError; accept both. +try: + from triton.compiler.errors import CompilationError + _RAISES = (CompilationError, ValueError) +except Exception: + _RAISES = (Exception, ) + + +# --- extract_tile: tile_shape rank must match source rank --- +@triton.jit +def _extract_tile_rank_mismatch(x_ptr): + offs = tl.arange(0, 64) + x = tl.load(x_ptr + offs[:, None] * 64 + tl.arange(0, 64)[None, :]) + # 2D source but 1D tile_shape -> rank mismatch -> ValueError + t = tle.extract_tile(x, index=0, tile_shape=[32]) + tl.store(x_ptr + offs, tl.sum(t)) + + +def test_extract_tile_rank_mismatch_raises(): + x = torch.zeros(64 * 64, device=DEVICE, dtype=torch.float32) + with pytest.raises(_RAISES): + _extract_tile_rank_mismatch[(1, )](x) + + +# --- insert_tile: source rank must match tile rank --- +@triton.jit +def _insert_tile_rank_mismatch(x_ptr, y_ptr): + offs = tl.arange(0, 64) + x = tl.load(x_ptr + offs[:, None] * 64 + tl.arange(0, 64)[None, :]) + y = tl.load(y_ptr + tl.arange(0, 32)) # 1D tile vs 2D source + z = tle.insert_tile(x, y, index=0) + tl.store(x_ptr + offs[:, None] * 64 + tl.arange(0, 64)[None, :], z) + + +def test_insert_tile_rank_mismatch_raises(): + x = torch.zeros(64 * 64, device=DEVICE, dtype=torch.float32) + y = torch.zeros(32, device=DEVICE, dtype=torch.float32) + with pytest.raises(_RAISES): + _insert_tile_rank_mismatch[(1, )](x, y) + + +# --- tle.gpu.alloc: source dim not divisible by tile dim (extract) --- +@triton.jit +def _extract_tile_not_divisible(x_ptr): + offs = tl.arange(0, 64) + x = tl.load(x_ptr + offs[:, None] * 64 + tl.arange(0, 64)[None, :]) + # 64 not divisible by tile dim 48 -> ValueError + t = tle.extract_tile(x, index=[0, 0], tile_shape=[48, 48]) + tl.store(x_ptr + offs, tl.sum(t)) + + +def test_extract_tile_not_divisible_raises(): + x = torch.zeros(64 * 64, device=DEVICE, dtype=torch.float32) + with pytest.raises(_RAISES): + _extract_tile_not_divisible[(1, )](x) + + +# --- tle.pipe: capacity must be positive (Hopper-only feature; also a contract +# check that fires before any NVWS lowering) --- +@triton.jit +def _pipe_bad_capacity(out_ptr): + buf = tle.gpu.alloc([1, 16, 16], dtype=tl.float32, layout=None, + scope=tle.gpu.smem, nv_mma_shared_layout=False) + # capacity must be > 0 -> ValueError + _ = tle.pipe(capacity=0, data=buf) + tl.store(out_ptr + tl.arange(0, 16), tl.zeros([16], tl.float32)) + + +def test_pipe_bad_capacity_raises(): + out = torch.zeros(16, device=DEVICE, dtype=torch.float32) + with pytest.raises(_RAISES): + _pipe_bad_capacity[(1, )](out) + + +# --- tle.pipe: scope must be 'cta' in the MVP --- +@triton.jit +def _pipe_bad_scope(out_ptr): + buf = tle.gpu.alloc([2, 16, 16], dtype=tl.float32, layout=None, + scope=tle.gpu.smem, nv_mma_shared_layout=False) + _ = tle.pipe(capacity=2, scope="gpu", data=buf) # only 'cta' supported + tl.store(out_ptr + tl.arange(0, 16), tl.zeros([16], tl.float32)) + + +def test_pipe_bad_scope_raises(): + out = torch.zeros(16, device=DEVICE, dtype=torch.float32) + with pytest.raises(_RAISES): + _pipe_bad_scope[(1, )](out) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"]) diff --git a/third_party/tle/dialect/lib/Transforms/TlePromoteLocalStoreStaging.cpp b/third_party/tle/dialect/lib/Transforms/TlePromoteLocalStoreStaging.cpp index 6c128af17b..9d1fa7e440 100644 --- a/third_party/tle/dialect/lib/Transforms/TlePromoteLocalStoreStaging.cpp +++ b/third_party/tle/dialect/lib/Transforms/TlePromoteLocalStoreStaging.cpp @@ -166,13 +166,14 @@ class PromoteLocalStoreStagingPass auto stage = builder.create( store.getLoc(), store.getDst().getType(), store.getSrc()); + Operation *storeOp = store.getOperation(); dst.replaceUsesWithIf(stage.getResult(), [&](OpOperand &use) { Operation *owner = use.getOwner(); - if (owner == store.getOperation()) + if (owner == storeOp) return false; if (!forOp->isAncestor(owner)) return false; - return domInfo.properlyDominates(store.getOperation(), owner); + return domInfo.properlyDominates(storeOp, owner); }); store.erase(); From 121db06f8f0c6b48f05cca87f9e4f20c1b80f9e6 Mon Sep 17 00:00:00 2001 From: lixingda Date: Fri, 26 Jun 2026 12:29:09 +0800 Subject: [PATCH 2/6] [TODO] verify the url path. --- python/setup_tools/setup_helper.py | 4 ++-- third_party/sunrise/CMakeLists.txt | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/setup_tools/setup_helper.py b/python/setup_tools/setup_helper.py index 12a579bbac..142ed249b7 100644 --- a/python/setup_tools/setup_helper.py +++ b/python/setup_tools/setup_helper.py @@ -622,7 +622,7 @@ def sunrise_pre_llvm_env(): cache.store( file="sunrise_llvm22_dev_release", condition=("sunrise" == flagtree_backend), - url= None, + url= None, # TODO pre_hook=sunrise_pre_llvm_env, post_hook=sunrise_set_llvm_env, ) @@ -630,5 +630,5 @@ def sunrise_pre_llvm_env(): cache.store( file="sunriseTritonPlugin.so", condition=("sunrise" == flagtree_backend) and (not configs.flagtree_plugin), - url=None, + url=None, # TODO copy_dst_path=f"third_party/{flagtree_backend}") diff --git a/third_party/sunrise/CMakeLists.txt b/third_party/sunrise/CMakeLists.txt index 0aea96bd52..6c870525f9 100644 --- a/third_party/sunrise/CMakeLists.txt +++ b/third_party/sunrise/CMakeLists.txt @@ -2,7 +2,7 @@ add_compile_options("-Wno-deprecated-declarations") add_compile_options("-Wno-error=deprecated-declarations") option(EDITABLE_MODE "Build in developer (editable) mode" OFF) -set(SUNRISE_PLUGIN_DIR ${CMAKE_LIBRARY_OUTPUT_DIRECTORY}) +set(SUNRISE_PLUGIN_DIR "$ENV{FLAGTREE_CACHE_DIR}/sunrise") if(TRITON_BUILD_PYTHON_MODULE) if(FLAGTREE_PLUGIN) From 552a1510c8e432e75a52ef8b5f5dad7de5e9df7b Mon Sep 17 00:00:00 2001 From: lixingda Date: Mon, 29 Jun 2026 16:08:53 +0800 Subject: [PATCH 3/6] [BACKEND] fix comments. --- .github/workflows/sunrise-build-and-test.yml | 4 +- .../TritonToTritonGPUPass.cpp | 2 +- .../Transforms/RewriteTensorPointer.cpp | 2 +- lib/Dialect/TritonGPU/IR/Dialect.cpp | 2 +- .../TritonGPU/IR/LinearLayoutConversions.cpp | 2 +- lib/Dialect/TritonGPU/Transforms/Prefetch.cpp | 2 +- .../Transforms/RemoveLayoutConversions.cpp | 2 +- lib/Dialect/TritonGPU/Transforms/Utility.cpp | 2 +- python/setup_tools/setup_helper.py | 33 +++------------ python/setup_tools/utils/sunrise.py | 13 ++++++ third_party/sunrise/CMakeLists.txt | 40 +++++++------------ .../sunrise_TritonToTritonGPUPass.h | 2 +- .../Transforms/sunrise_RewriteTensorPointer.h | 2 +- .../Dialect/TritonGPU/IR/sunrise_Dialect.h | 2 +- .../IR/sunrise_LinearLayoutConversions.h | 2 +- .../TritonGPU/Transforms/sunrise_Prefetch.h | 2 +- .../sunrise_RemoveLayoutConversions.h | 2 +- .../TritonGPU/Transforms/sunrise_Utility.h | 2 +- 18 files changed, 50 insertions(+), 68 deletions(-) create mode 100644 python/setup_tools/utils/sunrise.py diff --git a/.github/workflows/sunrise-build-and-test.yml b/.github/workflows/sunrise-build-and-test.yml index 6aded30788..080fa005a0 100644 --- a/.github/workflows/sunrise-build-and-test.yml +++ b/.github/workflows/sunrise-build-and-test.yml @@ -12,7 +12,7 @@ concurrency: jobs: sunrise-build-and-test: - runs-on: flagtree-sunrise + runs-on: sunrise3.6 if: ${{ github.repository == 'FlagTree/flagtree' || github.repository == 'flagos-ai/flagtree' }} steps: - name: Setup environment @@ -58,4 +58,4 @@ jobs: python3 -m pytest third_party/sunrise/python/test/06-tle-cumsum.py python3 -m pytest third_party/sunrise/python/test/07-tle-load.py python3 -m pytest third_party/sunrise/python/test/08-tle-dot-local-ptr.py - python3 -m pytest third_party/sunrise/python/test/09-tle-negative-contract.py \ No newline at end of file + python3 -m pytest third_party/sunrise/python/test/09-tle-negative-contract.py diff --git a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp index defd11073a..cd81ac070f 100644 --- a/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp +++ b/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp @@ -2,7 +2,7 @@ #include "flagtree_spec.h" #endif -#ifndef FLAGTREE_SPEC_triton_Conversion_TritonToTritonGPU_sunrise_TritonToTritonGPUPass +#ifndef FLAGTREE_SPEC_Conversion_TritonToTritonGPU_TritonToTritonGPUPass #include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h" diff --git a/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp b/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp index 3a745727ed..a0437199ee 100644 --- a/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp +++ b/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp @@ -2,7 +2,7 @@ #include "flagtree_spec.h" #endif -#ifndef FLAGTREE_SPEC_triton_Dialect_Triton_Transforms_sunrise_RewriteTensorPointer +#ifndef FLAGTREE_SPEC_Dialect_Triton_Transforms_RewriteTensorPointer #include diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp index b8c6e2716c..893bc6dab7 100644 --- a/lib/Dialect/TritonGPU/IR/Dialect.cpp +++ b/lib/Dialect/TritonGPU/IR/Dialect.cpp @@ -2,7 +2,7 @@ #include "flagtree_spec.h" #endif -#ifndef FLAGTREE_SPEC_triton_Dialect_TritonGPU_IR_sunrise_Dialect +#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_IR_Dialect #include "triton/Dialect/Triton/IR/Dialect.h" diff --git a/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp b/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp index 7e45fe9125..4bc9b6d9df 100644 --- a/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp +++ b/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp @@ -2,7 +2,7 @@ #include "flagtree_spec.h" #endif -#ifndef FLAGTREE_SPEC_triton_Dialect_TritonGPU_IR_sunrise_LinearLayoutConversion +#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_IR_LinearLayoutConversion #include diff --git a/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp b/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp index cff02296f6..f38960f6b1 100644 --- a/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp +++ b/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp @@ -29,7 +29,7 @@ #include "flagtree_spec.h" #endif -#ifndef FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_Prefetch +#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_Transforms_Prefetch #include "mlir/IR/IRMapping.h" #include "mlir/Support/LLVM.h" diff --git a/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp b/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp index fd5766b625..2b9fb15cec 100644 --- a/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp +++ b/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp @@ -2,7 +2,7 @@ #include "flagtree_spec.h" #endif -#ifndef FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_RemoveLayoutConversion +#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_Transforms_RemoveLayoutConversion #include "mlir/Analysis/SliceAnalysis.h" #include "mlir/Analysis/TopologicalSortUtils.h" diff --git a/lib/Dialect/TritonGPU/Transforms/Utility.cpp b/lib/Dialect/TritonGPU/Transforms/Utility.cpp index 50ee6955ee..ca06f838ed 100644 --- a/lib/Dialect/TritonGPU/Transforms/Utility.cpp +++ b/lib/Dialect/TritonGPU/Transforms/Utility.cpp @@ -2,7 +2,7 @@ #include "flagtree_spec.h" #endif -#ifndef FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_Utility +#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_Transforms_Utility #include "triton/Analysis/Utility.h" diff --git a/python/setup_tools/setup_helper.py b/python/setup_tools/setup_helper.py index 142ed249b7..855cc28187 100644 --- a/python/setup_tools/setup_helper.py +++ b/python/setup_tools/setup_helper.py @@ -595,40 +595,19 @@ def uninstall_triton(): post_hook=set_llvm_env, ) -# sunrise -def sunrise_cp_bc_files(path): - # mkdir -p third_party/sunrise/backend/lib - lib_dir = Path("third_party/sunrise/backend/lib") - os.makedirs(lib_dir, exist_ok=True) - # cp ${LLVM_SYSPATH}/stpu/bitcode/*.bc third_party/sunrise/backend/lib - bc_dir = Path(path) / "stpu" / "bitcode" - for bc_file in bc_dir.glob("*.bc"): - shutil.copy(bc_file, lib_dir) - - -def sunrise_set_llvm_env(path): - set_llvm_env(path) - sunrise_cp_bc_files(path) - - -def sunrise_pre_llvm_env(): - llvm_path = os.environ.get('LLVM_SYSPATH', '') - ret = llvm_path != '' - if ret: - sunrise_cp_bc_files(llvm_path) - return ret cache.store( file="sunrise_llvm22_dev_release", condition=("sunrise" == flagtree_backend), - url= None, # TODO - pre_hook=sunrise_pre_llvm_env, - post_hook=sunrise_set_llvm_env, + url= "https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/llvm-34b694004c-triton-v3.6.x.tar.gz", + pre_hook=lambda: check_env('LLVM_SYSPATH'), + post_hook=lambda path: [f(path) for f in (set_llvm_env, utils.activate("sunrise").sunrise_cp_bc_files)], ) cache.store( file="sunriseTritonPlugin.so", condition=("sunrise" == flagtree_backend) and (not configs.flagtree_plugin), - url=None, # TODO - copy_dst_path=f"third_party/{flagtree_backend}") + url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/sunrise-plugin-triton-v3.6.x.tar.gz", + md5_digest="3526d699", +) diff --git a/python/setup_tools/utils/sunrise.py b/python/setup_tools/utils/sunrise.py new file mode 100644 index 0000000000..d1efaa88dd --- /dev/null +++ b/python/setup_tools/utils/sunrise.py @@ -0,0 +1,13 @@ +import os +import shutil +from pathlib import Path + +# sunrise +def sunrise_cp_bc_files(path): + # mkdir -p third_party/sunrise/backend/lib + lib_dir = Path("third_party/sunrise/backend/lib") + os.makedirs(lib_dir, exist_ok=True) + # cp ${LLVM_SYSPATH}/stpu/bitcode/*.bc third_party/sunrise/backend/lib + bc_dir = Path(path) / "stpu" / "bitcode" + for bc_file in bc_dir.glob("*.bc"): + shutil.copy(bc_file, lib_dir) \ No newline at end of file diff --git a/third_party/sunrise/CMakeLists.txt b/third_party/sunrise/CMakeLists.txt index 6c870525f9..b71051996d 100644 --- a/third_party/sunrise/CMakeLists.txt +++ b/third_party/sunrise/CMakeLists.txt @@ -1,8 +1,11 @@ add_compile_options("-Wno-deprecated-declarations") add_compile_options("-Wno-error=deprecated-declarations") -option(EDITABLE_MODE "Build in developer (editable) mode" OFF) -set(SUNRISE_PLUGIN_DIR "$ENV{FLAGTREE_CACHE_DIR}/sunrise") +if(DEFINED ENV{FLAGTREE_CACHE_DIR} AND NOT "$ENV{FLAGTREE_CACHE_DIR}" STREQUAL "") + set(SUNRISE_PLUGIN_DIR "$ENV{FLAGTREE_CACHE_DIR}/sunrise") +else() + set(SUNRISE_PLUGIN_DIR "$ENV{HOME}/.flagtree/sunrise") +endif() if(TRITON_BUILD_PYTHON_MODULE) if(FLAGTREE_PLUGIN) @@ -11,29 +14,16 @@ if(TRITON_BUILD_PYTHON_MODULE) SHARED_LIB sunriseTritonPlugin ) else() - if(EDITABLE_MODE) - find_library(sunriseTritonPluginLib - NAMES - sunriseTritonPlugin.so - PATHS - ${SUNRISE_PLUGIN_DIR} - REQUIRED - ) - add_triton_plugin(TritonSunrise - SHARED_LIB ${sunriseTritonPluginLib} - ) - else() - find_library(sunriseTritonPluginLib - NAMES - sunriseTritonPlugin.so - PATHS - ${SUNRISE_PLUGIN_DIR} - REQUIRED - ) - add_triton_plugin(TritonSunrise - SHARED_LIB ${sunriseTritonPluginLib} - ) - endif() + find_library(sunriseTritonPluginLib + NAMES + sunriseTritonPlugin.so + PATHS + ${SUNRISE_PLUGIN_DIR} + REQUIRED + ) + add_triton_plugin(TritonSunrise + SHARED_LIB ${sunriseTritonPluginLib} + ) endif() endif() diff --git a/third_party/sunrise/backend/spec/include/triton/Conversion/TritonToTritonGPU/sunrise_TritonToTritonGPUPass.h b/third_party/sunrise/backend/spec/include/triton/Conversion/TritonToTritonGPU/sunrise_TritonToTritonGPUPass.h index ba3dfb3467..8d90f7a20c 100644 --- a/third_party/sunrise/backend/spec/include/triton/Conversion/TritonToTritonGPU/sunrise_TritonToTritonGPUPass.h +++ b/third_party/sunrise/backend/spec/include/triton/Conversion/TritonToTritonGPU/sunrise_TritonToTritonGPUPass.h @@ -1,6 +1,6 @@ #ifndef SUNRISE_TRITON_CONVERSION_TRITONTOTRITONGPU_TRITONTOTRITONGPUPASS_H_ #define SUNRISE_TRITON_CONVERSION_TRITONTOTRITONGPU_TRITONTOTRITONGPUPASS_H_ -#define FLAGTREE_SPEC_triton_Conversion_TritonToTritonGPU_sunrise_TritonToTritonGPUPass +#define FLAGTREE_SPEC_Conversion_TritonToTritonGPU_TritonToTritonGPUPass #endif // SUNRISE_TRITON_CONVERSION_TRITONTOTRITONGPU_TRITONTOTRITONGPUPASS_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/Triton/Transforms/sunrise_RewriteTensorPointer.h b/third_party/sunrise/backend/spec/include/triton/Dialect/Triton/Transforms/sunrise_RewriteTensorPointer.h index 3e960483ae..171b79b11f 100644 --- a/third_party/sunrise/backend/spec/include/triton/Dialect/Triton/Transforms/sunrise_RewriteTensorPointer.h +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/Triton/Transforms/sunrise_RewriteTensorPointer.h @@ -1,6 +1,6 @@ #ifndef SUNRISE_TRITON_DIALECT_TRITON_TRANSFORMS_REWRITETENSORPOINTER_H_ #define SUNRISE_TRITON_DIALECT_TRITON_TRANSFORMS_REWRITETENSORPOINTER_H_ -#define FLAGTREE_SPEC_triton_Dialect_Triton_Transforms_sunrise_RewriteTensorPointer +#define FLAGTREE_SPEC_Dialect_Triton_Transforms_RewriteTensorPointer #endif // SUNRISE_TRITON_DIALECT_TRITON_TRANSFORMS_REWRITETENSORPOINTER_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_Dialect.h b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_Dialect.h index 72408db721..db1d7cc7e3 100644 --- a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_Dialect.h +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_Dialect.h @@ -1,6 +1,6 @@ #ifndef SUNRISE_TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_ #define SUNRISE_TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_ -#define FLAGTREE_SPEC_triton_Dialect_TritonGPU_IR_sunrise_Dialect +#define FLAGTREE_SPEC_Dialect_TritonGPU_IR_Dialect #endif // SUNRISE_TRITON_DIALECT_TRITONGPU_IR_DIALECT_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_LinearLayoutConversions.h b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_LinearLayoutConversions.h index 38b20f334a..0118f0ab7a 100644 --- a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_LinearLayoutConversions.h +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/IR/sunrise_LinearLayoutConversions.h @@ -1,6 +1,6 @@ #ifndef SUNRISE_TRITON_DIALECT_TRITONGPU_IR_LINEARLAYOUTCONVERSIONS_H_ #define SUNRISE_TRITON_DIALECT_TRITONGPU_IR_LINEARLAYOUTCONVERSIONS_H_ -#define FLAGTREE_SPEC_triton_Dialect_TritonGPU_IR_sunrise_LinearLayoutConversion +#define FLAGTREE_SPEC_Dialect_TritonGPU_IR_LinearLayoutConversion #endif // SUNRISE_TRITON_DIALECT_TRITONGPU_IR_LINEARLAYOUTCONVERSIONS_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Prefetch.h b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Prefetch.h index f85de4efa1..da77c5706c 100644 --- a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Prefetch.h +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Prefetch.h @@ -1,6 +1,6 @@ #ifndef SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_PREFETCH_H_ #define SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_PREFETCH_H_ -#define FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_Prefetch +#define FLAGTREE_SPEC_Dialect_TritonGPU_Transforms_Prefetch #endif // SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_PREFETCH_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_RemoveLayoutConversions.h b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_RemoveLayoutConversions.h index 35a4d9f665..74bb0c312a 100644 --- a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_RemoveLayoutConversions.h +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_RemoveLayoutConversions.h @@ -1,6 +1,6 @@ #ifndef SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_REMOVELAYOUTCONVERSIONS_H_ #define SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_REMOVELAYOUTCONVERSIONS_H_ -#define FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_RemoveLayoutConversion +#define FLAGTREE_SPEC_Dialect_TritonGPU_Transforms_RemoveLayoutConversion #endif // SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_REMOVELAYOUTCONVERSIONS_H_ diff --git a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Utility.h b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Utility.h index 43b33be56b..bf5e6de66f 100644 --- a/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Utility.h +++ b/third_party/sunrise/backend/spec/include/triton/Dialect/TritonGPU/Transforms/sunrise_Utility.h @@ -1,6 +1,6 @@ #ifndef SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_UTILITY_H_ #define SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_UTILITY_H_ -#define FLAGTREE_SPEC_triton_Dialect_TritonGPU_Transforms_sunrise_Utility +#define FLAGTREE_SPEC_Dialect_TritonGPU_Transforms_Utility #endif // SUNRISE_TRITON_DIALECT_TRITONGPU_TRANSFORMS_UTILITY_H_ From 27f20b3f41b6bc513f610331f0669f5d605da525 Mon Sep 17 00:00:00 2001 From: lixingda Date: Mon, 29 Jun 2026 19:45:02 +0800 Subject: [PATCH 4/6] [BACKEND] fix comments --- CMakeLists.txt | 36 ++++++++++------------------- python/setup_tools/utils/sunrise.py | 2 +- setup.py | 4 ++-- third_party/sunrise/CMakeLists.txt | 12 +++++++--- 4 files changed, 24 insertions(+), 30 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 7fc372dfdc..8ff15d0b39 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -15,18 +15,6 @@ include(CTest) list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") -# Options -option(TRITON_BUILD_PYTHON_MODULE "Build Python Triton bindings" OFF) -if(FLAGTREE_BACKEND) - option(TRITON_BUILD_PROTON "Build the Triton Proton profiler" OFF) - option(TRITON_BUILD_UT "Build C++ Triton Unit Tests" OFF) -else() - option(TRITON_BUILD_PROTON "Build the Triton Proton profiler" ON) - option(TRITON_BUILD_UT "Build C++ Triton Unit Tests" ON) -endif() -option(TRITON_BUILD_WITH_CCACHE "Build with ccache (if available)" ON) -set(TRITON_CODEGEN_BACKENDS "" CACHE STRING "Enable different codegen backends") - # FLAGTREE Options set(FLAGTREE_BACKEND "$ENV{FLAGTREE_BACKEND}") if(FLAGTREE_BACKEND) @@ -84,6 +72,18 @@ function(get_flagtree_backend_lib lib_name output_lib) set(${output_lib} ${ret} PARENT_SCOPE) endfunction() +# FLAGTREE SPEC TD FILE GET FUNC +function(set_flagtree_backend_td output_td td_filename) + set(ret ${td_filename}) + file(RELATIVE_PATH relative_path "${PROJECT_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}") + get_filename_component(BACKEND_SPEC_ROOT "${BACKEND_SPEC_INCLUDE_DIR}" DIRECTORY) + set(BACKEND_SPEC_TD ${BACKEND_SPEC_ROOT}/${relative_path}/${td_filename}) + if(EXISTS ${BACKEND_SPEC_TD}) + set(ret ${BACKEND_SPEC_TD}) + endif() + set(${output_td} ${ret} PARENT_SCOPE) +endfunction() + list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") # Options @@ -102,18 +102,6 @@ if(FLAGTREE_BACKEND AND FLAGTREE_BACKEND STREQUAL "metax") list(APPEND TRITON_CODEGEN_BACKENDS "amd") endif() -# FLAGTREE SPEC TD FILE GET FUNC -function(set_flagtree_backend_td output_td td_filename) - set(ret ${td_filename}) - file(RELATIVE_PATH relative_path "${PROJECT_SOURCE_DIR}" "${CMAKE_CURRENT_SOURCE_DIR}") - get_filename_component(BACKEND_SPEC_ROOT "${BACKEND_SPEC_INCLUDE_DIR}" DIRECTORY) - set(BACKEND_SPEC_TD ${BACKEND_SPEC_ROOT}/${relative_path}/${td_filename}) - if(EXISTS ${BACKEND_SPEC_TD}) - set(ret ${BACKEND_SPEC_TD}) - endif() - set(${output_td} ${ret} PARENT_SCOPE) -endfunction() - if(TRITON_BUILD_WITH_CCACHE) find_program(CCACHE_PROGRAM ccache) if(CCACHE_PROGRAM) diff --git a/python/setup_tools/utils/sunrise.py b/python/setup_tools/utils/sunrise.py index d1efaa88dd..83fb9e4405 100644 --- a/python/setup_tools/utils/sunrise.py +++ b/python/setup_tools/utils/sunrise.py @@ -10,4 +10,4 @@ def sunrise_cp_bc_files(path): # cp ${LLVM_SYSPATH}/stpu/bitcode/*.bc third_party/sunrise/backend/lib bc_dir = Path(path) / "stpu" / "bitcode" for bc_file in bc_dir.glob("*.bc"): - shutil.copy(bc_file, lib_dir) \ No newline at end of file + shutil.copy(bc_file, lib_dir) diff --git a/setup.py b/setup.py index 8cb4b2777e..fc0d0fd470 100644 --- a/setup.py +++ b/setup.py @@ -16,7 +16,7 @@ from pathlib import Path from typing import Optional -from setuptools import Extension, find_namespace_packages, setup +from setuptools import Extension, find_packages, setup from setuptools.command.build_ext import build_ext from setuptools.command.build_py import build_py from setuptools.command.develop import develop @@ -686,7 +686,7 @@ def get_package_dirs(): def get_packages(): - yield from find_namespace_packages(where="python", include=["triton", "triton.*"]) + yield from find_packages(where="python", include=["triton", "triton.*"]) for backend in backends: yield f"triton.backends.{backend.name}" diff --git a/third_party/sunrise/CMakeLists.txt b/third_party/sunrise/CMakeLists.txt index b71051996d..b5240cec79 100644 --- a/third_party/sunrise/CMakeLists.txt +++ b/third_party/sunrise/CMakeLists.txt @@ -1,10 +1,16 @@ add_compile_options("-Wno-deprecated-declarations") add_compile_options("-Wno-error=deprecated-declarations") -if(DEFINED ENV{FLAGTREE_CACHE_DIR} AND NOT "$ENV{FLAGTREE_CACHE_DIR}" STREQUAL "") - set(SUNRISE_PLUGIN_DIR "$ENV{FLAGTREE_CACHE_DIR}/sunrise") + +option(EDITABLE_MODE "Build in developer (editable) mode" OFF) +if(EDITABLE_MODE) + if(DEFINED ENV{FLAGTREE_CACHE_DIR} AND NOT "$ENV{FLAGTREE_CACHE_DIR}" STREQUAL "") + set(SUNRISE_PLUGIN_DIR "$ENV{FLAGTREE_CACHE_DIR}/sunrise") + else() + set(SUNRISE_PLUGIN_DIR "$ENV{HOME}/.flagtree/sunrise") + endif() else() - set(SUNRISE_PLUGIN_DIR "$ENV{HOME}/.flagtree/sunrise") + set(SUNRISE_PLUGIN_DIR "${Python3_SITELIB}/triton/_C") endif() if(TRITON_BUILD_PYTHON_MODULE) From 67517a8f9f9ea14c0d66a0765954b7c0e5bb3e55 Mon Sep 17 00:00:00 2001 From: lixingda Date: Mon, 29 Jun 2026 19:52:53 +0800 Subject: [PATCH 5/6] format cmake. --- CMakeLists.txt | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 8ff15d0b39..002db738a6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -10,11 +10,6 @@ set(CMAKE_CXX_STANDARD 17) set(CMAKE_INCLUDE_CURRENT_DIR ON) -project(triton CXX C) -include(CTest) - -list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") - # FLAGTREE Options set(FLAGTREE_BACKEND "$ENV{FLAGTREE_BACKEND}") if(FLAGTREE_BACKEND) @@ -84,6 +79,9 @@ function(set_flagtree_backend_td output_td td_filename) set(${output_td} ${ret} PARENT_SCOPE) endfunction() +project(triton CXX C) +include(CTest) + list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/cmake") # Options From 47e066a6f7bf5bcb570e73b14a9f779a2fa88a22 Mon Sep 17 00:00:00 2001 From: lixingda Date: Tue, 30 Jun 2026 16:38:14 +0800 Subject: [PATCH 6/6] format code. --- .pre-commit-config.yaml | 3 +- python/setup_tools/setup_helper.py | 7 +- python/setup_tools/utils/sunrise.py | 1 + third_party/sunrise/backend/compiler.py | 25 +-- third_party/sunrise/backend/driver.c | 28 +-- third_party/sunrise/backend/driver.py | 10 +- .../TritonToTritonGPUPass.cpp | 14 +- .../Transforms/RewriteTensorPointer.cpp | 41 +++-- .../spec/lib/Dialect/TritonGPU/IR/Dialect.cpp | 174 ++++++++++++------ .../TritonGPU/IR/LinearLayoutConversions.cpp | 136 +++++++------- .../Dialect/TritonGPU/Transforms/Prefetch.cpp | 3 +- .../Transforms/RemoveLayoutConversions.cpp | 12 +- third_party/sunrise/language/tang/__init__.py | 2 +- .../sunrise/language/tang/libdevice.py | 41 ++++- third_party/sunrise/python/src/ir.cc | 49 +++-- third_party/sunrise/python/src/llvm.cc | 25 +-- .../sunrise/python/test/02-tle-local_ptr.py | 8 +- .../python/test/03-tle-extract_insert_tile.py | 13 +- .../python/test/05-tle-copy-normcopy.py | 8 +- .../sunrise/python/test/06-tle-cumsum.py | 3 +- .../sunrise/python/test/07-tle-load.py | 6 +- .../python/test/08-tle-dot-local-ptr.py | 9 +- .../python/test/09-tle-negative-contract.py | 6 +- 23 files changed, 366 insertions(+), 258 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 382afc8da8..9e455d43ce 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -55,5 +55,6 @@ exclude: | ^third_party/amd/backend/include/roctracer/| ^third_party/amd/backend/lib/| ^third_party/nvidia/backend/include/cuda.h| - ^third_party/f2reduce + ^third_party/f2reduce/| + ^third_party/sunrise/backend/include/ ) diff --git a/python/setup_tools/setup_helper.py b/python/setup_tools/setup_helper.py index 855cc28187..5fcf447b1c 100644 --- a/python/setup_tools/setup_helper.py +++ b/python/setup_tools/setup_helper.py @@ -243,6 +243,7 @@ def store(self, file=None, condition=None, url=None, copy_src_path=None, copy_ds def get(self, file_name) -> Path: return self.cache_files[file_name] + cache = FlagTreeCache() # -----flagtree-tle-raw-----flagtree-mlir--- @@ -396,7 +397,7 @@ def handle_plugin_backend(editable): os.makedirs(dst_build_plugin_dir) dst_build_plugin_path = dst_build_plugin_dir / flagtree_plugin_so shutil.copy(src_build_plugin_path, dst_build_plugin_path) - if flagtree_backend in ("mthreads",): + if flagtree_backend in ("mthreads", ): dst_install_plugin_dir = Path( __file__).resolve().parent.parent.parent / "third_party" / flagtree_backend / "python" / "triton" / "_C" else: @@ -595,12 +596,10 @@ def uninstall_triton(): post_hook=set_llvm_env, ) - - cache.store( file="sunrise_llvm22_dev_release", condition=("sunrise" == flagtree_backend), - url= "https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/llvm-34b694004c-triton-v3.6.x.tar.gz", + url="https://baai-cp-web.ks3-cn-beijing.ksyuncs.com/trans/llvm-34b694004c-triton-v3.6.x.tar.gz", pre_hook=lambda: check_env('LLVM_SYSPATH'), post_hook=lambda path: [f(path) for f in (set_llvm_env, utils.activate("sunrise").sunrise_cp_bc_files)], ) diff --git a/python/setup_tools/utils/sunrise.py b/python/setup_tools/utils/sunrise.py index 83fb9e4405..abb6de6027 100644 --- a/python/setup_tools/utils/sunrise.py +++ b/python/setup_tools/utils/sunrise.py @@ -2,6 +2,7 @@ import shutil from pathlib import Path + # sunrise def sunrise_cp_bc_files(path): # mkdir -p third_party/sunrise/backend/lib diff --git a/third_party/sunrise/backend/compiler.py b/third_party/sunrise/backend/compiler.py index a80f9f8ecf..a0b0b9f316 100644 --- a/third_party/sunrise/backend/compiler.py +++ b/third_party/sunrise/backend/compiler.py @@ -30,11 +30,13 @@ class sunrise_knobs(base_knobs): def min_dot_size(target: GPUTarget): return lambda lhsType, rhsType: (8, 8, 16) if lhsType.is_int8() else (8, 8, 4) + @functools.lru_cache(None) def file_hash(path): with open(path, "rb") as f: return hashlib.sha256(f.read()).hexdigest() + @dataclass(frozen=True) class SunriseOptions: num_warps: int = 4 @@ -59,7 +61,7 @@ def __post_init__(self): warp_size = 32 object.__setattr__(self, 'warp_size', warp_size) default_libdir = Path(__file__).parent / 'lib' - extern_libs ={} if self.extern_libs is None else dict(self.extern_libs) + extern_libs = {} if self.extern_libs is None else dict(self.extern_libs) lib_ver = 'S2' if knobs_sunrise.triple and ('stcuv2' in knobs_sunrise.triple): lib_ver = 'S3' @@ -89,7 +91,7 @@ def hash(self): class SunriseBackend(BaseBackend): instrumentation = None supports_native_tensor_specialization = False - stage_stcu_dump_dir_name = '0' # 用于保存 .asm 等文件 + stage_stcu_dump_dir_name = '0' # 用于保存 .asm 等文件 @staticmethod def supports_target(target: GPUTarget): @@ -106,7 +108,7 @@ def parse_options(self, opts) -> Any: args = {'arch': knobs.runtime.override_arch or self.target.arch} if "enable_fp_fusion" not in opts: args["enable_fp_fusion"] = knobs.language.default_fp_fusion - args["max_num_imprecise_acc_default"] = 0 # TODO + args["max_num_imprecise_acc_default"] = 0 # TODO args.update({k: opts[k] for k in SunriseOptions.__dataclass_fields__.keys() \ if k in opts and opts[k] is not None}) return SunriseOptions(**args) @@ -119,9 +121,7 @@ def pack_metadata(self, metadata): ) def get_codegen_implementation(self, options): - codegen_fns = { - "min_dot_size": min_dot_size(self.target) - } + codegen_fns = {"min_dot_size": min_dot_size(self.target)} return codegen_fns def get_module_map(self) -> Dict[str, ModuleType]: @@ -163,7 +163,7 @@ def get_flag(metadata, opt): flag.append('thread-regfile-size=64') for name, path in opt.extern_libs: if name == "ockl": - flag.append('ocklPath='+path) + flag.append('ocklPath=' + path) return flag @staticmethod @@ -236,7 +236,7 @@ def make_ttgir(mod, metadata, opt): if os.getenv('OFF_MMA', '0') == '1': print('not run accelerate_matmul pass') else: - sunrise.passes.ttgpuir.add_accelerate_matmul(pm, 1, 0) # 版本：1.0 + sunrise.passes.ttgpuir.add_accelerate_matmul(pm, 1, 0) # 版本：1.0 # sunrise.passes.ttgpuir.add_mma_direct_store(pm) passes.ttgpuir.add_remove_layout_conversions(pm) passes.ttgpuir.add_optimize_dot_operands(pm, True) @@ -246,7 +246,7 @@ def make_ttgir(mod, metadata, opt): if os.getenv('OFF_ASYNC', '0') == '0': passes.ttgpuir.add_assign_latencies(pm, num_stages) passes.ttgpuir.add_schedule_loops(pm) - passes.ttgpuir.add_pipeline(pm, num_stages, True ) + passes.ttgpuir.add_pipeline(pm, num_stages, True) if os.getenv('OFF_PREF', '0') == '0': passes.ttir.add_loop_aware_cse(pm) passes.common.add_canonicalizer(pm) @@ -256,7 +256,7 @@ def make_ttgir(mod, metadata, opt): if os.getenv('OFF_ASYNC', '0') == '0': passes.ttgpuir.add_assign_latencies(pm, num_stages) passes.ttgpuir.add_schedule_loops(pm) - sunrise.passes.ttgpuir.add_pipeline(pm, num_stages, 1, 0) # 版本：1.0 + sunrise.passes.ttgpuir.add_pipeline(pm, num_stages, 1, 0) # 版本：1.0 if os.getenv('OFF_PREF', '0') == '0': passes.ttir.add_loop_aware_cse(pm) passes.common.add_canonicalizer(pm) @@ -265,7 +265,7 @@ def make_ttgir(mod, metadata, opt): if os.getenv('OFF_MMA', '0') == '1': print('not run accelerate_matmul pass') else: - sunrise.passes.ttgpuir.add_accelerate_matmul(pm, 1, 0) # 版本：1.0 + sunrise.passes.ttgpuir.add_accelerate_matmul(pm, 1, 0) # 版本：1.0 sunrise.passes.ttgpuir.add_mma_direct_store(pm) passes.ttgpuir.add_remove_layout_conversions(pm) # sunrise.passes.ttgpuir.add_optimize_dot_operands(pm) @@ -441,5 +441,6 @@ def add_stages(self, stages, options, language): @functools.lru_cache() def hash(self): - version = subprocess.check_output([SunriseBackend.path_to_clang_offload_bundler(), "--version"], encoding='utf-8') + version = subprocess.check_output([SunriseBackend.path_to_clang_offload_bundler(), "--version"], + encoding='utf-8') return f'{version}-{self.target.arch}' diff --git a/third_party/sunrise/backend/driver.c b/third_party/sunrise/backend/driver.c index 336dfd0102..1c155a531f 100644 --- a/third_party/sunrise/backend/driver.c +++ b/third_party/sunrise/backend/driver.c @@ -57,16 +57,15 @@ static PyObject *getDeviceProperties(PyObject *self, PyObject *args) { int mem_clock_rate; int mem_bus_width; TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( - &max_shared_mem, TA_DEV_ATTR_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, - device)); - TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( - &max_num_regs, TA_DEV_ATTR_REGS_PER_BLOCK, device)); + &max_shared_mem, TA_DEV_ATTR_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, device)); + TANG_CHECK_AND_RETURN_NULL( + taDeviceGetAttribute(&max_num_regs, TA_DEV_ATTR_REGS_PER_BLOCK, device)); TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( &multiprocessor_count, TA_DEV_ATTR_MULTIPROCESSOR_COUNT, device)); - TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( - &warp_size, TA_DEV_ATTR_WARP_SIZE, device)); - TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( - &sm_clock_rate, TA_DEV_ATTR_CLOCK_RATE, device)); + TANG_CHECK_AND_RETURN_NULL( + taDeviceGetAttribute(&warp_size, TA_DEV_ATTR_WARP_SIZE, device)); + TANG_CHECK_AND_RETURN_NULL( + taDeviceGetAttribute(&sm_clock_rate, TA_DEV_ATTR_CLOCK_RATE, device)); TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( &mem_clock_rate, TA_DEV_ATTR_MEMORY_CLOCK_RATE, device)); TANG_CHECK_AND_RETURN_NULL(taDeviceGetAttribute( @@ -108,7 +107,8 @@ static PyObject *loadBinary(PyObject *self, PyObject *args) { TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(taCtxSetCurrent(pctx)); } - TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(taModuleLoadData(&mod, data, (size_t)data_size)); + TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( + taModuleLoadData(&mod, data, (size_t)data_size)); TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( taModuleGetFunction(&fun, mod, name)); // get allocated registers and spilled registers from the function @@ -116,7 +116,8 @@ static PyObject *loadBinary(PyObject *self, PyObject *args) { // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( // taFuncGetAttribute(&n_regs, TA_FUNC_ATTRIBUTE_NUM_REGS, fun)); // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( - // taFuncGetAttribute(&n_spills, CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, fun)); + // taFuncGetAttribute(&n_spills, CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES, + // fun)); // n_spills /= 4; TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(taFuncGetAttribute( &n_max_threads, TA_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK, fun)); @@ -130,12 +131,13 @@ static PyObject *loadBinary(PyObject *self, PyObject *args) { // cuFuncSetCacheConfig(fun, CU_FUNC_CACHE_PREFER_SHARED)); // int shared_total, shared_static; // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(cuDeviceGetAttribute( - // &shared_total, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_MULTIPROCESSOR, - // device)); + // &shared_total, + // CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_MULTIPROCESSOR, device)); // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS(cuFuncGetAttribute( // &shared_static, CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES, fun)); // TANG_CHECK_AND_RETURN_NULL_ALLOW_THREADS( - // cuFuncSetAttribute(fun, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, + // cuFuncSetAttribute(fun, + // CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, // shared_optin - shared_static)); // } Py_END_ALLOW_THREADS; diff --git a/third_party/sunrise/backend/driver.py b/third_party/sunrise/backend/driver.py index 583fd61e90..38586d5d76 100644 --- a/third_party/sunrise/backend/driver.py +++ b/third_party/sunrise/backend/driver.py @@ -17,6 +17,7 @@ libraries = ['tang', 'tangrt_shared'] arch = platform.machine() + @functools.lru_cache() def libtang_dirs(): if env_libtang_path := knobs.env_opt_str("TRITON_LIBTANG_PATH"): @@ -98,6 +99,7 @@ def ty_to_cpp(ty): "fp64": "double", }[ty] + FLOAT_STORAGE_TYPE = { "fp16": "uint16_t", "bf16": "uint16_t", @@ -115,7 +117,9 @@ def ty_to_cpp(ty): _BASE_ARGS_FORMAT = "iiiKKOOOOO" + def make_launcher(constants, signature, warp_size, tensordesc_meta): + def _expand_signature(signature): output = [] # Expand tensor descriptor arguments into base pointer, shape, and @@ -446,6 +450,7 @@ def format_of(ty): """ return src + def wrap_handle_tensordesc(launcher, signature, tensordesc_meta): has_tensor_desc_arg = any(isinstance(sig, str) and sig.startswith("tensordesc") for sig in signature.values()) if not has_tensor_desc_arg: @@ -472,6 +477,7 @@ def inner(*args): return inner + class SunriseLauncher(object): def __init__(self, src, metadata): @@ -494,6 +500,7 @@ def __init__(self, src, metadata): self.profile_scratch_align = metadata.profile_scratch_align def __call__(self, gridX, gridY, gridZ, stream, function, *args): + def allocate_scratch(size, align, allocator): if size > 0: grid_size = gridX * gridY * gridZ @@ -504,7 +511,7 @@ def allocate_scratch(size, align, allocator): profile_scratch = allocate_scratch(self.profile_scratch_size, self.profile_scratch_align, _allocation._profile_allocator) - self.launch(gridX, gridY, gridZ, stream, function, profile_scratch,*args) + self.launch(gridX, gridY, gridZ, stream, function, profile_scratch, *args) class SunriseDriver(GPUDriver): @@ -520,7 +527,6 @@ def __init__(self): self.get_current_device = torch.ptpu.current_device self.set_current_device = torch.ptpu.set_device - def get_current_target(self): arch = knobs.runtime.override_arch if not arch: diff --git a/third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp b/third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp index 5c208d8d1a..85eacde2dc 100644 --- a/third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp +++ b/third_party/sunrise/backend/spec/lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp @@ -27,10 +27,10 @@ using namespace mlir::triton; using namespace mlir::triton::gpu; // [Sunrise Optim] 即使前端没手工写 num_warps=1/2, -// 也会把这类 small rank-1 reduce (size <= 128) 自动从坏的 4/8-warp 配置拉回到 2-warp。 -static bool shouldClampSunriseNumWarpsForSmallReduce(ModuleOp mod, - StringRef target, - mlir::Pass::Option &numWarps) { +// 也会把这类 small rank-1 reduce (size <= 128) 自动从坏的 4/8-warp 配置拉回到 +// 2-warp。 +static bool shouldClampSunriseNumWarpsForSmallReduce( + ModuleOp mod, StringRef target, mlir::Pass::Option &numWarps) { int32_t clampedNumWarps = numWarps.getValue(); if (!target.starts_with("tang:") || clampedNumWarps <= 2) return false; @@ -46,7 +46,8 @@ static bool shouldClampSunriseNumWarpsForSmallReduce(ModuleOp mod, if (!reduce) return WalkResult::advance(); - auto operandTy = dyn_cast(reduce.getOperands()[0].getType()); + auto operandTy = + dyn_cast(reduce.getOperands()[0].getType()); if (!operandTy || operandTy.getRank() != 1 || reduce.getAxis() != 0) return WalkResult::advance(); @@ -1008,7 +1009,8 @@ class ConvertTritonToTritonGPU MLIRContext *context = &getContext(); ModuleOp mod = getOperation(); - shouldClampSunriseNumWarpsForSmallReduce(mod, this->target.getValue(), numWarps); + shouldClampSunriseNumWarpsForSmallReduce(mod, this->target.getValue(), + numWarps); // type converter TritonGPUTypeConverter typeConverter(context, numWarps, threadsPerWarp, numCTAs, enableSourceRemat); diff --git a/third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp b/third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp index 29bc998b65..ffe066c0c9 100644 --- a/third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp +++ b/third_party/sunrise/backend/spec/lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp @@ -41,10 +41,11 @@ struct RewritedInfo { RewritedInfo(Value base, const SmallVector &shape, const SmallVector &strides, const SmallVector &offsets, - const ArrayRef &tensorShape, - bool hasRowStride, unsigned rowStrideParamIndex) + const ArrayRef &tensorShape, bool hasRowStride, + unsigned rowStrideParamIndex) : base(base), shape(shape), strides(strides), offsets(offsets), - tensorShape(tensorShape), hasRowStride(hasRowStride), rowStrideParamIndex(rowStrideParamIndex) { + tensorShape(tensorShape), hasRowStride(hasRowStride), + rowStrideParamIndex(rowStrideParamIndex) { assert(shape.size() == strides.size() && shape.size() == offsets.size() && shape.size() == tensorShape.size()); } @@ -200,19 +201,19 @@ struct RewritedInfo { }; // 成功返回true -bool findRowStrideParamIndexFromValue(Value v, unsigned* paramIdx) { - Operation* op = v.getDefiningOp(); - while(op != nullptr) { - if(op->getOperands().size() < 1) { +bool findRowStrideParamIndexFromValue(Value v, unsigned *paramIdx) { + Operation *op = v.getDefiningOp(); + while (op != nullptr) { + if (op->getOperands().size() < 1) { return false; } v = op->getOperands()[0]; op = v.getDefiningOp(); } - if(auto blockArg = dyn_cast(v)) { + if (auto blockArg = dyn_cast(v)) { *paramIdx = blockArg.getArgNumber(); - Block* parentBlock = blockArg.getOwner(); - if(parentBlock->isEntryBlock()) { + Block *parentBlock = blockArg.getOwner(); + if (parentBlock->isEntryBlock()) { return true; } } @@ -272,8 +273,9 @@ class RewriteTensorPointerPass unsigned rowStrideParamIdx = 0; bool hasRowStride = false; auto strides = op.getStrides(); - if(strides.size() == 2) { - hasRowStride = findRowStrideParamIndexFromValue(strides[0], &rowStrideParamIdx); + if (strides.size() == 2) { + hasRowStride = + findRowStrideParamIndexFromValue(strides[0], &rowStrideParamIdx); } // Save information @@ -352,8 +354,10 @@ class RewriteTensorPointerPass loadOp.getCache(), loadOp.getEvict(), loadOp.getIsVolatile(), loadOp.getFlagtreeHintsAttr()); // flagtree hints MLIRContext *ctx = builder.getContext(); - if(info.getHasRowStride() == true) { - newResult->setAttr("sunrise.rowStrideParamIdx", IntegerAttr::get(IntegerType::get(ctx, 32), info.getRowStrideParamIndex())); + if (info.getHasRowStride() == true) { + newResult->setAttr("sunrise.rowStrideParamIdx", + IntegerAttr::get(IntegerType::get(ctx, 32), + info.getRowStrideParamIndex())); } op->getResult(0).replaceAllUsesWith(newResult); } else if (auto storeOp = dyn_cast(op)) { @@ -601,17 +605,18 @@ class RewriteTensorPointerPass } // 此时loadOp不再有boundaryCheck，根据参数个数判断是否有mask - getOperation()->walk([](Operation* op){ - if(isa(*op) == false) { + getOperation()->walk([](Operation *op) { + if (isa(*op) == false) { return; } OpBuilder builder(op); MLIRContext *ctx = builder.getContext(); int hasOriMask = 0; - if(op->getNumOperands() > 1) { + if (op->getNumOperands() > 1) { hasOriMask = 1; } - op->setAttr("sunrise.hasOriMask", IntegerAttr::get(IntegerType::get(ctx, 32), hasOriMask)); + op->setAttr("sunrise.hasOriMask", + IntegerAttr::get(IntegerType::get(ctx, 32), hasOriMask)); }); } }; diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp index 3a43eebb05..8658d3dd53 100644 --- a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/Dialect.cpp @@ -1474,50 +1474,95 @@ LogicalResult AMDWmmaEncodingAttr::verify( Attribute SunriseMmaEncodingAttr::parse(AsmParser &parser, Type type) { DictionaryAttr dict; - if (parser.parseLess().failed()) { return {}; } - if (parser.parseAttribute(dict).failed()) { return {};} - if (parser.parseGreater().failed()) {return {};} + if (parser.parseLess().failed()) { + return {}; + } + if (parser.parseAttribute(dict).failed()) { + return {}; + } + if (parser.parseGreater().failed()) { + return {}; + } unsigned versionMajor = 0; unsigned versionMinor = 0; SmallVector warpsPerCTA; SunriseMmaEncodingAttr::TMMAOutLayout outLayout; - unsigned outLayoutUint = static_cast(SunriseMmaEncodingAttr::TMMAOutLayout::NotAvailable); + unsigned outLayoutUint = static_cast( + SunriseMmaEncodingAttr::TMMAOutLayout::NotAvailable); unsigned inputElemBitWidth = 0; unsigned outputElemBitWidth = 0; bool isACol = false, isBCol = false; Attribute ctaAttr = nullptr; for (const NamedAttribute &attr : dict) { - if (attr.getName() == "versionMajor") { if (parseUInt(parser, attr, versionMajor, "versionMajor").failed()) {return {};} } - if (attr.getName() == "versionMinor") { if (parseUInt(parser, attr, versionMinor, "versionMinor").failed()) {return {};} } - if (attr.getName() == "warpsPerCTA") { if (parseIntArrayAttr(parser, attr, warpsPerCTA, "warpsPerCTA").failed()) {return {};} } - if (attr.getName() == "CGALayout") { ctaAttr = attr.getValue(); continue; } - if (attr.getName() == "outLayout") { if(parseUInt(parser, attr, outLayoutUint, "outLayout").failed()) {return {}; } } - if (attr.getName() == "inputElemBitWidth") { if(parseUInt(parser, attr, inputElemBitWidth, "inputElemBitWidth").failed()) {return {}; } } - if (attr.getName() == "outputElemBitWidth") { if(parseUInt(parser, attr, outputElemBitWidth, "outputElemBitWidth").failed()) {return {}; } } - if (attr.getName() == "isACol") { if(parseBool(parser, attr, isACol, "isACol").failed()) {return {}; } } - if (attr.getName() == "isBCol") { if(parseBool(parser, attr, isBCol, "isBCol").failed()) {return {}; } } + if (attr.getName() == "versionMajor") { + if (parseUInt(parser, attr, versionMajor, "versionMajor").failed()) { + return {}; + } + } + if (attr.getName() == "versionMinor") { + if (parseUInt(parser, attr, versionMinor, "versionMinor").failed()) { + return {}; + } + } + if (attr.getName() == "warpsPerCTA") { + if (parseIntArrayAttr(parser, attr, warpsPerCTA, "warpsPerCTA") + .failed()) { + return {}; + } + } + if (attr.getName() == "CGALayout") { + ctaAttr = attr.getValue(); + continue; + } + if (attr.getName() == "outLayout") { + if (parseUInt(parser, attr, outLayoutUint, "outLayout").failed()) { + return {}; + } + } + if (attr.getName() == "inputElemBitWidth") { + if (parseUInt(parser, attr, inputElemBitWidth, "inputElemBitWidth") + .failed()) { + return {}; + } + } + if (attr.getName() == "outputElemBitWidth") { + if (parseUInt(parser, attr, outputElemBitWidth, "outputElemBitWidth") + .failed()) { + return {}; + } + } + if (attr.getName() == "isACol") { + if (parseBool(parser, attr, isACol, "isACol").failed()) { + return {}; + } + } + if (attr.getName() == "isBCol") { + if (parseBool(parser, attr, isBCol, "isBCol").failed()) { + return {}; + } + } } outLayout = static_cast(outLayoutUint); std::optional CTALayout = - parseCTAAttr(parser, ctaAttr, /*rank=*/warpsPerCTA.size()); + parseCTAAttr(parser, ctaAttr, /*rank=*/warpsPerCTA.size()); if (!CTALayout.has_value()) return {}; return parser.getChecked( - parser.getContext(), versionMajor, versionMinor, warpsPerCTA, *CTALayout, outLayout, inputElemBitWidth, outputElemBitWidth, isACol, isBCol); + parser.getContext(), versionMajor, versionMinor, warpsPerCTA, *CTALayout, + outLayout, inputElemBitWidth, outputElemBitWidth, isACol, isBCol); } void SunriseMmaEncodingAttr::print(AsmPrinter &printer) const { printer << "<{" << "versionMajor = " << getVersionMajor() - << ", versionMinor = " << getVersionMinor() - << ", warpsPerCTA = [" << ArrayRef(getWarpsPerCTA()) << "]" - << ", isACol = " << getIsACol() - << ", isBCol = " << getIsBCol() + << ", versionMinor = " << getVersionMinor() << ", warpsPerCTA = [" + << ArrayRef(getWarpsPerCTA()) << "]" + << ", isACol = " << getIsACol() << ", isBCol = " << getIsBCol() << ", inputElemBitWidth = " << getInputElemBitWidth() << ", outputElemBitWidth = " << getOutputElemBitWidth(); maybePrintCTALayout(getContext(), printer, getCTALayout(), @@ -1533,7 +1578,6 @@ void SunriseMmaEncodingAttr::print(AsmPrinter &printer) const { printer << "}>"; } - //===----------------------------------------------------------------------===// // Sliced Encoding //===----------------------------------------------------------------------===// @@ -2452,28 +2496,38 @@ SunriseMmaEncodingAttr::getRepOrderForOperand(int opIdx) const { return getOrderForDotOperand(opIdx, getRank(), /*kContig*/ true); } -SmallVector -SunriseMmaEncodingAttr::getInstrShapeForOperand(unsigned opIdx, unsigned elemBitWdith) const { // MK, KN +SmallVector SunriseMmaEncodingAttr::getInstrShapeForOperand( + unsigned opIdx, unsigned elemBitWdith) const { // MK, KN // A, B矩阵每个warp加载的尺寸 int packPer32bit = 32 / this->getInputElemBitWidth(); - if(opIdx == 0) { - switch(elemBitWdith) { - case 32: return {8, 4}; - case 16: return {8, 8}; - case 8: return {8, 16}; - case 4: return {8, 32}; - default: llvm_unreachable("unsupported packPer32bit"); + if (opIdx == 0) { + switch (elemBitWdith) { + case 32: + return {8, 4}; + case 16: + return {8, 8}; + case 8: + return {8, 16}; + case 4: + return {8, 32}; + default: + llvm_unreachable("unsupported packPer32bit"); } - } - else { - switch(elemBitWdith) { - case 32: return {4, 8}; - case 16: return {8, 8}; - case 8: return {16, 8}; - case 4: return {32, 8}; - default: llvm_unreachable("unsupported packPer32bit"); + } else { + switch (elemBitWdith) { + case 32: + return {4, 8}; + case 16: + return {8, 8}; + case 8: + return {16, 8}; + case 4: + return {32, 8}; + default: + llvm_unreachable("unsupported packPer32bit"); } - } return {0, 0}; + } + return {0, 0}; } SmallVector @@ -2482,13 +2536,22 @@ SunriseMmaEncodingAttr::getShapePerCTATileForOperand(unsigned opIdx) const { unsigned k = 0; // mma指令k维度的元素个数 int elemBitWidth = this->getInputElemBitWidth(); switch (elemBitWidth) { - case 4: k = 32; break; - case 8: k = 16; break; - case 16: k = 8; break; - case 32: k = 4; break; + case 4: + k = 32; + break; + case 8: + k = 16; + break; + case 16: + k = 8; + break; + case 32: + k = 4; + break; default: - llvm::report_fatal_error("SunriseMmaEncodingAttr::getShapePerCTATileForOperand " - "unsuppored inputElemBitWidth"); + llvm::report_fatal_error( + "SunriseMmaEncodingAttr::getShapePerCTATileForOperand " + "unsuppored inputElemBitWidth"); } auto shapePerCTATile = getShapePerCTATile(); SmallVector ret; @@ -2513,14 +2576,19 @@ SmallVector SunriseMmaEncodingAttr::getShapePerCTATile() const { // 获取mma的两个输入a、b的每个维度的切割数，即每个维度有几个CTATile // 【注意】两个a和b的元素类型应该一致 // 返回值：[repM, repK]或[repK, repN] -SmallVector SunriseMmaEncodingAttr::getRepForOperand(ArrayRef operandShape, - Type elemType, int opIdx) const { +SmallVector +SunriseMmaEncodingAttr::getRepForOperand(ArrayRef operandShape, + Type elemType, int opIdx) const { int instrM = 8, instrN = 8, instrK = 0; - if(elemType.isF32() || elemType.isInteger(32)) { instrK = 4; } - else if(elemType.isF16() || elemType.isBF16()) { instrK = 8; } - else if(elemType.isInteger(8)) { instrK = 16; } - else if(elemType.isInteger(4)) { instrK = 32; } - else { + if (elemType.isF32() || elemType.isInteger(32)) { + instrK = 4; + } else if (elemType.isF16() || elemType.isBF16()) { + instrK = 8; + } else if (elemType.isInteger(8)) { + instrK = 16; + } else if (elemType.isInteger(4)) { + instrK = 32; + } else { llvm::report_fatal_error("unsupported tensor data type for tmma!"); return {}; } @@ -2529,11 +2597,11 @@ SmallVector SunriseMmaEncodingAttr::getRepForOperand(ArrayRef SmallVector ret; if (opIdx == 0) ret = {std::max(1, operandShape[0] / (instrM * warpsPerCTA[0])), - std::max(1, operandShape[1] / instrK)}; + std::max(1, operandShape[1] / instrK)}; else { assert(opIdx == 1); ret = {std::max(1, operandShape[0] / instrK), - std::max(1, operandShape[1] / (instrN * warpsPerCTA[1]))}; + std::max(1, operandShape[1] / (instrN * warpsPerCTA[1]))}; } return ret; } diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp index 457196235e..2be6183195 100644 --- a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp @@ -846,77 +846,81 @@ LinearLayout wmmaDotOperandToLinearLayout(DotOperandEncodingAttr dotWmmaLayout, return combineCtaCgaWithShape(ctaLayout, wmmaLayout.getCTALayout(), shape); } -LinearLayout sunrisemmaDotOperandToLinearLayout(DotOperandEncodingAttr dotEncAttr, ArrayRef shape) { +LinearLayout +sunrisemmaDotOperandToLinearLayout(DotOperandEncodingAttr dotEncAttr, + ArrayRef shape) { MLIRContext *ctx = dotEncAttr.getContext(); auto rank = shape.size(); SmallVector outDimNames = standardOutDimNames(ctx, rank); unsigned dotOpIdx = dotEncAttr.getOpIdx(); - SunriseMmaEncodingAttr sunriseMmaAttr = cast(dotEncAttr.getParent()); - bool kContig = (dotOpIdx == 0 ? sunriseMmaAttr.getIsACol() == false : sunriseMmaAttr.getIsBCol() == true); + SunriseMmaEncodingAttr sunriseMmaAttr = + cast(dotEncAttr.getParent()); + bool kContig = (dotOpIdx == 0 ? sunriseMmaAttr.getIsACol() == false + : sunriseMmaAttr.getIsBCol() == true); auto order = getOrderForDotOperand(dotOpIdx, rank, kContig); unsigned elemBitWidth = sunriseMmaAttr.getInputElemBitWidth(); auto tileLayout = LinearLayout::empty(); - switch(elemBitWidth) { + switch (elemBitWidth) { case 32: - if( (dotOpIdx == 0 && sunriseMmaAttr.getIsACol() == false) - || (dotOpIdx == 1 && sunriseMmaAttr.getIsBCol() == true) ) { - tileLayout = LinearLayout( - {{S("register"), {}}, - {S("lane"), {{1,0}, {2,0}, {0,1}, {0,2}, {0,4}}} - }, {outDimNames[order[0]], outDimNames[order[1]]} - ); + if ((dotOpIdx == 0 && sunriseMmaAttr.getIsACol() == false) || + (dotOpIdx == 1 && sunriseMmaAttr.getIsBCol() == true)) { + tileLayout = + LinearLayout({{S("register"), {}}, + {S("lane"), {{1, 0}, {2, 0}, {0, 1}, {0, 2}, {0, 4}}}}, + {outDimNames[order[0]], outDimNames[order[1]]}); } else { - tileLayout = LinearLayout( - {{S("register"), {}}, - {S("lane"), {{1,0}, {2,0}, {4,0}, {0,1}, {0,2}}} - }, {outDimNames[order[0]], outDimNames[order[1]]} - ); + tileLayout = + LinearLayout({{S("register"), {}}, + {S("lane"), {{1, 0}, {2, 0}, {4, 0}, {0, 1}, {0, 2}}}}, + {outDimNames[order[0]], outDimNames[order[1]]}); } break; case 16: - if(dotOpIdx == 0) { + if (dotOpIdx == 0) { tileLayout = LinearLayout( - //{{S("register"), {{1,0}, {8,0}}}, // 有么有8,0好像都行？ - {{S("register"), {{1,0}}}, - {S("lane"), {{2,0}, {4,0}, {0,1}, {0,2}, {0,4}}} - }, {outDimNames[order[0]], outDimNames[order[1]]} - ); + //{{S("register"), {{1,0}, {8,0}}}, // 有么有8,0好像都行？ + {{S("register"), {{1, 0}}}, + {S("lane"), {{2, 0}, {4, 0}, {0, 1}, {0, 2}, {0, 4}}}}, + {outDimNames[order[0]], outDimNames[order[1]]}); } else { - tileLayout = LinearLayout( - {{S("register"), {{1,0}}}, - {S("lane"), {{2,0}, {4,0}, {0,1}, {0,2}, {0,4}}} - }, {outDimNames[order[0]], outDimNames[order[1]]} - ); + tileLayout = + LinearLayout({{S("register"), {{1, 0}}}, + {S("lane"), {{2, 0}, {4, 0}, {0, 1}, {0, 2}, {0, 4}}}}, + {outDimNames[order[0]], outDimNames[order[1]]}); } break; case 8: - if( (dotOpIdx == 0 && sunriseMmaAttr.getIsACol() == false) - || (dotOpIdx == 1 && sunriseMmaAttr.getIsBCol() == true) ) { - tileLayout = LinearLayout( - {{S("register"), {{1,0}, {2,0}}}, - {S("lane"), {{4,0}, {8,0}, {0,1}, {0,2}, {0,4}}} - }, {outDimNames[order[0]], outDimNames[order[1]]} - ); + if ((dotOpIdx == 0 && sunriseMmaAttr.getIsACol() == false) || + (dotOpIdx == 1 && sunriseMmaAttr.getIsBCol() == true)) { + tileLayout = + LinearLayout({{S("register"), {{1, 0}, {2, 0}}}, + {S("lane"), {{4, 0}, {8, 0}, {0, 1}, {0, 2}, {0, 4}}}}, + {outDimNames[order[0]], outDimNames[order[1]]}); } else { - tileLayout = LinearLayout( - {{S("register"), {{1,0}, {2,0}}}, - {S("lane"), {{4,0}, {0,1}, {0,2}, {0,4}, {0,8}}} - }, {outDimNames[order[0]], outDimNames[order[1]]} - ); + tileLayout = + LinearLayout({{S("register"), {{1, 0}, {2, 0}}}, + {S("lane"), {{4, 0}, {0, 1}, {0, 2}, {0, 4}, {0, 8}}}}, + {outDimNames[order[0]], outDimNames[order[1]]}); } break; case 4: default: - llvm::report_fatal_error((std::string("linearlayout not implemented! elemBitWidth:") + std::to_string(elemBitWidth)).c_str()); + llvm::report_fatal_error( + (std::string("linearlayout not implemented! elemBitWidth:") + + std::to_string(elemBitWidth)) + .c_str()); break; } auto kDim = dotOpIdx == 0 ? rank - 1 : rank - 2; auto warpOrder = getDefaultMmaOrder(sunriseMmaAttr); - // SmallVector warpOrder = dotOpIdx == 0 ? SmallVector({1,0}) : SmallVector({0,1}); + // SmallVector warpOrder = dotOpIdx == 0 ? + // SmallVector({1,0}) : SmallVector({0,1}); - // LinearLayout warpLayout = identityStandardND(S("warp"), sunriseMmaAttr.getWarpsPerCTA(), warpOrder); - LinearLayout warpLayout = broadcastedDotOperandLayout(ctx, sunriseMmaAttr.getWarpsPerCTA(), warpOrder, kDim, S("warp")); + // LinearLayout warpLayout = identityStandardND(S("warp"), + // sunriseMmaAttr.getWarpsPerCTA(), warpOrder); + LinearLayout warpLayout = broadcastedDotOperandLayout( + ctx, sunriseMmaAttr.getWarpsPerCTA(), warpOrder, kDim, S("warp")); // reorder dim names in rep order, so combineCtaCgaWithShape generate proper // extension of layout @@ -928,45 +932,46 @@ LinearLayout sunrisemmaDotOperandToLinearLayout(DotOperandEncodingAttr dotEncAtt // join instruction layout and warps using repetition order of dimensions LinearLayout ctaLayout = tileLayout.transposeOuts(repDimNames) * warpLayout.transposeOuts(repDimNames); - return combineCtaCgaWithShape(ctaLayout, sunriseMmaAttr.getCTALayout(), shape); + return combineCtaCgaWithShape(ctaLayout, sunriseMmaAttr.getCTALayout(), + shape); } -LinearLayout SunriseMmaEncodingAttr::toLinearLayout(ArrayRef shape) const { +LinearLayout +SunriseMmaEncodingAttr::toLinearLayout(ArrayRef shape) const { int rank = shape.size(); assert(rank == 2); MLIRContext *ctx = getContext(); SmallVector outDimNames = standardOutDimNames(ctx, rank); // SmallVector order = getDefaultMmaOrder(*this); - // SmallVector order = getOrderForDotOperand(dotWmmaLayout.getOpIdx(), rank, /*kContig*/ true); - SmallVector order = SmallVector({1,0}); + // SmallVector order = + // getOrderForDotOperand(dotWmmaLayout.getOpIdx(), rank, /*kContig*/ true); + SmallVector order = SmallVector({1, 0}); SunriseMmaEncodingAttr::TMMAOutLayout outLayout = getOutLayout(); - if ((outLayout != SunriseMmaEncodingAttr::TMMAOutLayout::Row_2B) && (outLayout != SunriseMmaEncodingAttr::TMMAOutLayout::ARow_4B_8x4)) { + if ((outLayout != SunriseMmaEncodingAttr::TMMAOutLayout::Row_2B) && + (outLayout != SunriseMmaEncodingAttr::TMMAOutLayout::ARow_4B_8x4)) { assert(0 && "Unsupport SunriseMmaEncodingAttr::TMMAOutLayout yet"); } auto tileLayout = LinearLayout::empty(); - if(outLayout == SunriseMmaEncodingAttr::TMMAOutLayout::Row_2B) { + if (outLayout == SunriseMmaEncodingAttr::TMMAOutLayout::Row_2B) { // fp16 row - tileLayout = LinearLayout( - {{S("register"), {{1, 0}}}, - //{{S("register"), {}}, - {S("lane"), {{2,0}, {4,0}, {0,1}, {0,2}, {0,4}}} - }, - {outDimNames[order[0]], outDimNames[order[1]]} - ); - } else if(outLayout == SunriseMmaEncodingAttr::TMMAOutLayout::ARow_4B_8x4) { + tileLayout = + LinearLayout({{S("register"), {{1, 0}}}, + //{{S("register"), {}}, + {S("lane"), {{2, 0}, {4, 0}, {0, 1}, {0, 2}, {0, 4}}}}, + {outDimNames[order[0]], outDimNames[order[1]]}); + } else if (outLayout == SunriseMmaEncodingAttr::TMMAOutLayout::ARow_4B_8x4) { // fp32 8x4 - tileLayout = LinearLayout( - {{S("register"), {{4, 0}}}, - {S("lane"), {{1,0}, {2,0}, {0,1}, {0,2}, {0,4}}} - }, - {outDimNames[order[0]], outDimNames[order[1]]} - ); + tileLayout = + LinearLayout({{S("register"), {{4, 0}}}, + {S("lane"), {{1, 0}, {2, 0}, {0, 1}, {0, 2}, {0, 4}}}}, + {outDimNames[order[0]], outDimNames[order[1]]}); } else { llvm::report_fatal_error("Unsupport sunrisemma outLayout"); } - LinearLayout warpLayout = identityStandardND(S("warp"), getWarpsPerCTA(), order); + LinearLayout warpLayout = + identityStandardND(S("warp"), getWarpsPerCTA(), order); // reorder dim names in rep order, so combineCtaCgaWithShape generate proper // extension of layout auto repOrder = getRepOrder(); @@ -1139,7 +1144,8 @@ DotOperandEncodingAttr::toLinearLayout(ArrayRef shape) const { return mfmaDotToLinearLayout(*this, shape); } else if (auto wmmaLayout = mlir::dyn_cast(parent)) { return wmmaDotOperandToLinearLayout(*this, shape); - } else if (auto sunrisemmaLayout = mlir::dyn_cast(parent)){ + } else if (auto sunrisemmaLayout = + mlir::dyn_cast(parent)) { return sunrisemmaDotOperandToLinearLayout(*this, shape); } else { auto mma = mlir::cast(parent); diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp index f673708d1e..f8288aba85 100644 --- a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/Prefetch.cpp @@ -169,7 +169,8 @@ LogicalResult Prefetcher::initialize() { dyn_cast(getEncoding(dotOp.getResult())); auto dstTmmaEnc = dyn_cast(getEncoding(dotOp.getResult())); - if (!dstTmmaEnc && !dstMfmaEnc && (!dstMmaEnc || dstMmaEnc.getVersionMajor() != 2)) + if (!dstTmmaEnc && !dstMfmaEnc && + (!dstMmaEnc || dstMmaEnc.getVersionMajor() != 2)) // Don't rewrite if any other type is found. return failure(); dotsInFor.push_back(dotOp); diff --git a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp index 8a88bc0a6e..00e6e8e40d 100644 --- a/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp +++ b/third_party/sunrise/backend/spec/lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp @@ -345,14 +345,16 @@ SmallVector LayoutPropagation::propagateToUsers(Value value, isa(user)) { // sunrise fix:S2的mma进行32bit转16bit时布局不一样，不能直接转换 - if(isa(user)) { - auto srcType = dyn_cast(user->getOperand(0).getType()); + if (isa(user)) { + auto srcType = + dyn_cast(user->getOperand(0).getType()); auto dstType = dyn_cast(user->getResult(0).getType()); - if(srcType != nullptr && dstType != nullptr) { + if (srcType != nullptr && dstType != nullptr) { auto srcElemType = srcType.getElementType(); auto dstElemType = dstType.getElementType(); - if((srcElemType.getIntOrFloatBitWidth() == 32 && dstElemType.getIntOrFloatBitWidth() == 16) - || (dstElemType.getIntOrFloatBitWidth() == 8)){ + if ((srcElemType.getIntOrFloatBitWidth() == 32 && + dstElemType.getIntOrFloatBitWidth() == 16) || + (dstElemType.getIntOrFloatBitWidth() == 8)) { continue; } } diff --git a/third_party/sunrise/language/tang/__init__.py b/third_party/sunrise/language/tang/__init__.py index 2ce93570d3..229b57d87d 100644 --- a/third_party/sunrise/language/tang/__init__.py +++ b/third_party/sunrise/language/tang/__init__.py @@ -1,3 +1,3 @@ from . import libdevice -__all__ = ["libdevice"] \ No newline at end of file +__all__ = ["libdevice"] diff --git a/third_party/sunrise/language/tang/libdevice.py b/third_party/sunrise/language/tang/libdevice.py index b0cc4fb21e..e7c00aef51 100644 --- a/third_party/sunrise/language/tang/libdevice.py +++ b/third_party/sunrise/language/tang/libdevice.py @@ -1,5 +1,6 @@ from triton.language import core + @core.extern def erf(arg0, _semantic=None): return core.extern_elementwise( @@ -8,6 +9,7 @@ def erf(arg0, _semantic=None): (core.dtype("fp64"), ): ("__ocml_erf_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def pow(arg0, arg1, _semantic=None): return core.extern_elementwise( @@ -19,6 +21,7 @@ def pow(arg0, arg1, _semantic=None): (core.dtype("fp64"), core.dtype("fp64")): ("__ocml_pow_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def tanh(arg0, _semantic=None): return core.extern_elementwise( @@ -27,6 +30,7 @@ def tanh(arg0, _semantic=None): (core.dtype("fp64"), ): ("__ocml_tanh_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def atan2(arg0, arg1, _semantic=None): return core.extern_elementwise( @@ -35,6 +39,7 @@ def atan2(arg0, arg1, _semantic=None): (core.dtype("fp64"), core.dtype("fp64")): ("__ocml_atan2_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + # atanh, tanh2 @core.extern def atan(arg0, _semantic=None): @@ -44,6 +49,7 @@ def atan(arg0, _semantic=None): (core.dtype("fp64"), ): ("__ocml_atan_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def asin(arg0, _semantic=None): return core.extern_elementwise( @@ -52,6 +58,7 @@ def asin(arg0, _semantic=None): (core.dtype("fp64"), ): ("__ocml_asin_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def acos(arg0, _semantic=None): return core.extern_elementwise( @@ -60,6 +67,7 @@ def acos(arg0, _semantic=None): (core.dtype("fp64"), ): ("__ocml_acos_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def div_rd(arg0, arg1, _semantic=None): return core.extern_elementwise( @@ -68,6 +76,7 @@ def div_rd(arg0, arg1, _semantic=None): (core.dtype("fp64"), core.dtype("fp64")): ("llvm.stvm.div.rm.f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def div_rz(arg0, arg1, _semantic=None): return core.extern_elementwise( @@ -76,6 +85,7 @@ def div_rz(arg0, arg1, _semantic=None): (core.dtype("fp64"), core.dtype("fp64")): ("llvm.stvm.div.rz.f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def rsqrt(arg0, _semantic=None): return core.extern_elementwise( @@ -84,6 +94,7 @@ def rsqrt(arg0, _semantic=None): (core.dtype("fp64"), ): ("llvm.stvm.rsqrt.f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def isinf(arg0, _semantic=None): return core.extern_elementwise( @@ -92,6 +103,7 @@ def isinf(arg0, _semantic=None): (core.dtype("fp64"), ): ("llvm.stvm.testp.f32.inf", core.dtype("int32")), }, is_pure=True, _semantic=_semantic) + @core.extern def isnan(arg0, _semantic=None): return core.extern_elementwise( @@ -102,6 +114,7 @@ def isnan(arg0, _semantic=None): (core.dtype("fp64"), ): ("llvm.stvm.testp.f32.not", core.dtype("int32")), }, is_pure=True, _semantic=_semantic) + @core.extern def sin(arg0, _semantic=None): return core.extern_elementwise( @@ -119,6 +132,7 @@ def cos(arg0, _semantic=None): (core.dtype("fp64"), ): ("__ocml_cos_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def tan(arg0, _semantic=None): return core.extern_elementwise( @@ -127,6 +141,7 @@ def tan(arg0, _semantic=None): (core.dtype("fp64"), ): ("__ocml_tan_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def erf(arg0, _semantic=None): return core.extern_elementwise( @@ -144,6 +159,7 @@ def exp(arg0, _semantic=None): (core.dtype("fp64"), ): ("__ocml_exp_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def exp2(arg0, _semantic=None): return core.extern_elementwise( @@ -152,6 +168,7 @@ def exp2(arg0, _semantic=None): (core.dtype("fp64"), ): ("llvm.stvm.exp2.f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def div_rn(arg0, arg1, _semantic=None): return core.extern_elementwise( @@ -160,6 +177,7 @@ def div_rn(arg0, arg1, _semantic=None): (core.dtype("fp64"), core.dtype("fp64")): ("llvm.stvm.div.rn.f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def trunc(arg0, _semantic=None): return core.extern_elementwise( @@ -168,6 +186,7 @@ def trunc(arg0, _semantic=None): (core.dtype("fp64"), ): ("llvm.stvm.trunc.f", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def fmod(arg0, arg1, _semantic=None): return core.extern_elementwise( @@ -176,28 +195,34 @@ def fmod(arg0, arg1, _semantic=None): (core.dtype("fp64"), core.dtype("fp64")): ("__ocml_fmod_f32", core.dtype("fp64")), }, is_pure=True, _semantic=_semantic) + @core.extern def isfinited(arg0, _semantic=None): return core.extern_elementwise("", "", [arg0], { (core.dtype("fp64"), ): ("__ocml_isfinite_f64", core.dtype("int32")), }, is_pure=True, _semantic=_semantic).to(core.int1, _semantic=_semantic) + @core.extern def finitef(arg0, _semantic=None): return core.extern_elementwise("", "", [arg0], { (core.dtype("fp32"), ): ("__ocml_isfinite_f32", core.dtype("int32")), }, is_pure=True, _semantic=_semantic).to(core.int1, _semantic=_semantic) + @core.extern def rint(arg0, _semantic=None): - return core.extern_elementwise("", "", [arg0], { - (core.dtype("fp32"), ): ("llvm.rint.f32", core.dtype("fp32")), - (core.dtype("fp64"), ): ("llvm.rint.f32", core.dtype("fp64")), - }, is_pure=True, _semantic=_semantic) + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("fp32"), ): ("llvm.rint.f32", core.dtype("fp32")), + (core.dtype("fp64"), ): ("llvm.rint.f32", core.dtype("fp64")), + }, is_pure=True, _semantic=_semantic) + @core.extern def ffs(arg0, _semantic=None): - return core.extern_elementwise("", "", [arg0], { - (core.dtype("int32"), ): ("_Z5__ffsi", core.dtype("int32")), - (core.dtype("int64"), ): ("_Z7__ffsllx", core.dtype("int32")), - }, is_pure=True, _semantic=_semantic) + return core.extern_elementwise( + "", "", [arg0], { + (core.dtype("int32"), ): ("_Z5__ffsi", core.dtype("int32")), + (core.dtype("int64"), ): ("_Z7__ffsllx", core.dtype("int32")), + }, is_pure=True, _semantic=_semantic) diff --git a/third_party/sunrise/python/src/ir.cc b/third_party/sunrise/python/src/ir.cc index 800a61e9fa..2c729fdac9 100644 --- a/third_party/sunrise/python/src/ir.cc +++ b/third_party/sunrise/python/src/ir.cc @@ -1885,12 +1885,13 @@ void init_triton_ir(py::module &&m) { ; - static unsigned g_passGroupNumber=1; + static unsigned g_passGroupNumber = 1; py::class_(m, "pass_manager", py::module_local()) .def(py::init()) - .def( "get_dump_dir_name", [](PassManager &self) -> std::string { - return std::to_string(g_passGroupNumber++); - }) + .def("get_dump_dir_name", + [](PassManager &self) -> std::string { + return std::to_string(g_passGroupNumber++); + }) .def("enable_debug", [](PassManager &self) -> bool { auto *context = self.getContext(); @@ -1919,27 +1920,25 @@ void init_triton_ir(py::module &&m) { return false; }; - if(std::string("1") == getenv("MLIR_DUMP_FILE_TREE")) { - std::string dumpDirName = std::to_string(g_passGroupNumber++); - auto returnTrue = [](Pass *, Operation *) { return true; }; - self.enableIRPrintingToFileTree( - /*shouldPrintBeforePass=*/returnTrue, - /*shouldPrintAfterPass=*/nullptr, - /*printModuleScope=*/true, - /*printAfterOnlyOnChange=*/true, - /*printAfterOnlyOnFailure*/ false, - /*printTreeDir*/ dumpDirName, - printingFlags); - } - else { - self.enableIRPrinting( - /*shouldPrintBeforePass=*/printAlways, - /*shouldPrintAfterPass=*/printAlways, - /*printModuleScope=*/true, - /*printAfterOnlyOnChange=*/ true, - /*printAfterOnlyOnFailure*/ true, mlir_dumps_or_dbgs(), - printingFlags); - } + if (std::string("1") == getenv("MLIR_DUMP_FILE_TREE")) { + std::string dumpDirName = std::to_string(g_passGroupNumber++); + auto returnTrue = [](Pass *, Operation *) { return true; }; + self.enableIRPrintingToFileTree( + /*shouldPrintBeforePass=*/returnTrue, + /*shouldPrintAfterPass=*/nullptr, + /*printModuleScope=*/true, + /*printAfterOnlyOnChange=*/true, + /*printAfterOnlyOnFailure*/ false, + /*printTreeDir*/ dumpDirName, printingFlags); + } else { + self.enableIRPrinting( + /*shouldPrintBeforePass=*/printAlways, + /*shouldPrintAfterPass=*/printAlways, + /*printModuleScope=*/true, + /*printAfterOnlyOnChange=*/true, + /*printAfterOnlyOnFailure*/ true, mlir_dumps_or_dbgs(), + printingFlags); + } } return haveDump; }) diff --git a/third_party/sunrise/python/src/llvm.cc b/third_party/sunrise/python/src/llvm.cc index 7f3d5b5514..ab945b0826 100644 --- a/third_party/sunrise/python/src/llvm.cc +++ b/third_party/sunrise/python/src/llvm.cc @@ -296,29 +296,30 @@ std::string translateLLVMIRToASM(llvm::Module &module, const std::vector &flags, bool enable_fp_fusion, bool isObject) { using namespace mlir; - auto isInteger=[&](const std::string& str) { + auto isInteger = [&](const std::string &str) { return !str.empty() && std::all_of(str.begin(), str.end(), ::isdigit); }; // options auto options = llvm::cl::getRegisteredOptions(); for (std::string flag : flags) { - auto pos = flag.find("="); // sunrise fix - if(pos == std::string::npos){ + auto pos = flag.find("="); // sunrise fix + if (pos == std::string::npos) { auto *shortPtr = static_cast *>(options[flag]); assert(shortPtr); shortPtr->setValue(true); - }else{ - auto optIt = options.find(flag.substr(0,pos)); + } else { + auto optIt = options.find(flag.substr(0, pos)); assert(optIt != options.end() && "Option not found in cl::opt!"); - if (isInteger(flag.substr(pos+1))) { - if(auto *unsignedOpt = static_cast*>(options[flag.substr(0,pos)])) { - unsignedOpt->setValue(std::stoi(flag.substr(pos+1))); + if (isInteger(flag.substr(pos + 1))) { + if (auto *unsignedOpt = static_cast *>( + options[flag.substr(0, pos)])) { + unsignedOpt->setValue(std::stoi(flag.substr(pos + 1))); } - } - else { - if (auto *stringOpt = static_cast*>(options[flag.substr(0,pos)])) { - stringOpt->setValue(flag.substr(pos+1)); + } else { + if (auto *stringOpt = static_cast *>( + options[flag.substr(0, pos)])) { + stringOpt->setValue(flag.substr(pos + 1)); } } } diff --git a/third_party/sunrise/python/test/02-tle-local_ptr.py b/third_party/sunrise/python/test/02-tle-local_ptr.py index ff982a968a..0d6b02838a 100644 --- a/third_party/sunrise/python/test/02-tle-local_ptr.py +++ b/third_party/sunrise/python/test/02-tle-local_ptr.py @@ -58,8 +58,8 @@ def shared_roundtrip_kernel( src_rows = src_ptr + xstride * xoffs[:, None] dst_rows = dst_ptr + xstride * xoffs[:, None] - smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, - scope=tle.gpu.smem, nv_mma_shared_layout=False) + smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, scope=tle.gpu.smem, + nv_mma_shared_layout=False) row_ids = tl.arange(0, XBLOCK)[:, None] col_ids = tl.arange(0, YBLOCK)[None, :] @@ -71,8 +71,8 @@ def shared_roundtrip_kernel( yoffs = tl.arange(0, YBLOCK) + yoff gval = tl.load(src_rows + yoffs[None, :]) - tl.store(smem_ptrs, gval) # vectorized shared store (guarded) - sval = tl.load(smem_ptrs) # vectorized shared load (guarded) + tl.store(smem_ptrs, gval) # vectorized shared store (guarded) + sval = tl.load(smem_ptrs) # vectorized shared load (guarded) tl.store(dst_rows + yoffs[None, :], sval * 2.0) diff --git a/third_party/sunrise/python/test/03-tle-extract_insert_tile.py b/third_party/sunrise/python/test/03-tle-extract_insert_tile.py index dd9b007ef0..58fba2be34 100644 --- a/third_party/sunrise/python/test/03-tle-extract_insert_tile.py +++ b/third_party/sunrise/python/test/03-tle-extract_insert_tile.py @@ -47,8 +47,7 @@ def extract_tile_kernel(x_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, TM: tl @triton.jit -def insert_tile_kernel(x_ptr, y_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, - TN: tl.constexpr): +def insert_tile_kernel(x_ptr, y_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, TN: tl.constexpr): offs_m = tl.arange(0, M) offs_n = tl.arange(0, N) x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :]) @@ -66,8 +65,7 @@ def insert_tile_kernel(x_ptr, y_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, # Dynamic (runtime) index -> SMEM relay path # --------------------------------------------------------------------------- @triton.jit -def extract_tile_dyn_kernel(x_ptr, out_ptr, idx, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, - TN: tl.constexpr): +def extract_tile_dyn_kernel(x_ptr, out_ptr, idx, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, TN: tl.constexpr): offs_m = tl.arange(0, M) offs_n = tl.arange(0, N) x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :]) @@ -82,8 +80,8 @@ def extract_tile_dyn_kernel(x_ptr, out_ptr, idx, M: tl.constexpr, N: tl.constexp @triton.jit -def insert_tile_dyn_kernel(x_ptr, y_ptr, out_ptr, idx, M: tl.constexpr, N: tl.constexpr, - TM: tl.constexpr, TN: tl.constexpr): +def insert_tile_dyn_kernel(x_ptr, y_ptr, out_ptr, idx, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, + TN: tl.constexpr): offs_m = tl.arange(0, M) offs_n = tl.arange(0, N) x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :]) @@ -198,8 +196,7 @@ def test_insert_tile_dynamic_index(ti, tj): # This is a different SMEM-path entry than the dynamic-index tests above. # --------------------------------------------------------------------------- @triton.jit -def extract_tile_thin_kernel(x_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, - TN: tl.constexpr): +def extract_tile_thin_kernel(x_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, TM: tl.constexpr, TN: tl.constexpr): offs_m = tl.arange(0, M) offs_n = tl.arange(0, N) x = tl.load(x_ptr + offs_m[:, None] * N + offs_n[None, :]) diff --git a/third_party/sunrise/python/test/05-tle-copy-normcopy.py b/third_party/sunrise/python/test/05-tle-copy-normcopy.py index 5b1b06c44a..6008b59068 100644 --- a/third_party/sunrise/python/test/05-tle-copy-normcopy.py +++ b/third_party/sunrise/python/test/05-tle-copy-normcopy.py @@ -53,8 +53,8 @@ def copy_roundtrip_kernel( a_rows = a_ptr + xstride * xoffs[:, None] c_rows = c_ptr + xstride * xoffs[:, None] - smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, - scope=tle.gpu.smem, nv_mma_shared_layout=False) + smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, scope=tle.gpu.smem, + nv_mma_shared_layout=False) row_ids = tl.arange(0, XBLOCK)[:, None] col_ids = tl.arange(0, YBLOCK)[None, :] @@ -113,8 +113,8 @@ def copy_gm_to_local_kernel( a_rows = a_ptr + xstride * xoffs[:, None] c_rows = c_ptr + xstride * xoffs[:, None] - smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, - scope=tle.gpu.smem, nv_mma_shared_layout=False) + smem = tle.gpu.alloc([XBLOCK, YBLOCK], dtype=tl.float32, layout=None, scope=tle.gpu.smem, + nv_mma_shared_layout=False) row_ids = tl.broadcast_to(tl.arange(0, XBLOCK)[:, None], (XBLOCK, YBLOCK)) col_ids = tl.broadcast_to(tl.arange(0, YBLOCK)[None, :], (XBLOCK, YBLOCK)) smem_ptrs = tle.gpu.local_ptr(smem, (row_ids, col_ids)) diff --git a/third_party/sunrise/python/test/06-tle-cumsum.py b/third_party/sunrise/python/test/06-tle-cumsum.py index ef501f3fb1..cf3c381e6e 100644 --- a/third_party/sunrise/python/test/06-tle-cumsum.py +++ b/third_party/sunrise/python/test/06-tle-cumsum.py @@ -27,8 +27,7 @@ @triton.jit -def cumsum_kernel(x_ptr, exclusive_ptr, total_ptr, n, BLOCK: tl.constexpr, - REVERSE: tl.constexpr): +def cumsum_kernel(x_ptr, exclusive_ptr, total_ptr, n, BLOCK: tl.constexpr, REVERSE: tl.constexpr): offs = tl.arange(0, BLOCK) mask = offs < n x = tl.load(x_ptr + offs, mask=mask, other=0) diff --git a/third_party/sunrise/python/test/07-tle-load.py b/third_party/sunrise/python/test/07-tle-load.py index a75a5dd9f9..524f4e00b3 100644 --- a/third_party/sunrise/python/test/07-tle-load.py +++ b/third_party/sunrise/python/test/07-tle-load.py @@ -25,8 +25,7 @@ @triton.jit -def tle_load_add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK: tl.constexpr, - IS_ASYNC: tl.constexpr): +def tle_load_add_kernel(x_ptr, y_ptr, out_ptr, n_elements, BLOCK: tl.constexpr, IS_ASYNC: tl.constexpr): pid = tl.program_id(0) offs = pid * BLOCK + tl.arange(0, BLOCK) mask = offs < n_elements @@ -52,8 +51,7 @@ def test_tle_load_add(n_elements, is_async): @triton.jit -def tle_load_2d_kernel(x_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, - IS_ASYNC: tl.constexpr): +def tle_load_2d_kernel(x_ptr, out_ptr, M: tl.constexpr, N: tl.constexpr, IS_ASYNC: tl.constexpr): offs_m = tl.arange(0, M) offs_n = tl.arange(0, N) ptrs = x_ptr + offs_m[:, None] * N + offs_n[None, :] diff --git a/third_party/sunrise/python/test/08-tle-dot-local-ptr.py b/third_party/sunrise/python/test/08-tle-dot-local-ptr.py index 0aeb44e160..a6918ae3d8 100644 --- a/third_party/sunrise/python/test/08-tle-dot-local-ptr.py +++ b/third_party/sunrise/python/test/08-tle-dot-local-ptr.py @@ -28,8 +28,7 @@ @triton.jit -def gemm_local_ptr_kernel(a_ptr, b_ptr, c_ptr, M: tl.constexpr, N: tl.constexpr, - K: tl.constexpr): +def gemm_local_ptr_kernel(a_ptr, b_ptr, c_ptr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr): offs_m = tl.arange(0, M) offs_n = tl.arange(0, N) offs_k = tl.arange(0, K) @@ -40,8 +39,7 @@ def gemm_local_ptr_kernel(a_ptr, b_ptr, c_ptr, M: tl.constexpr, N: tl.constexpr, # Stage A through a shared-memory buffer via local_ptr (the TLE path), then # read it back and feed tl.dot. This forces local store -> barrier -> load. - a_smem = tle.gpu.alloc([M, K], dtype=tl.float32, layout=None, scope=tle.gpu.smem, - nv_mma_shared_layout=False) + a_smem = tle.gpu.alloc([M, K], dtype=tl.float32, layout=None, scope=tle.gpu.smem, nv_mma_shared_layout=False) row = tl.broadcast_to(tl.arange(0, M)[:, None], (M, K)) col = tl.broadcast_to(tl.arange(0, K)[None, :], (M, K)) a_ptrs = tle.gpu.local_ptr(a_smem, (row, col)) @@ -66,8 +64,7 @@ def test_gemm_local_ptr(M, N, K): @triton.jit -def gemm_plain_kernel(a_ptr, b_ptr, c_ptr, M: tl.constexpr, N: tl.constexpr, - K: tl.constexpr): +def gemm_plain_kernel(a_ptr, b_ptr, c_ptr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr): offs_m = tl.arange(0, M) offs_n = tl.arange(0, N) offs_k = tl.arange(0, K) diff --git a/third_party/sunrise/python/test/09-tle-negative-contract.py b/third_party/sunrise/python/test/09-tle-negative-contract.py index 318faf52d3..2d91af6283 100644 --- a/third_party/sunrise/python/test/09-tle-negative-contract.py +++ b/third_party/sunrise/python/test/09-tle-negative-contract.py @@ -81,8 +81,7 @@ def test_extract_tile_not_divisible_raises(): # check that fires before any NVWS lowering) --- @triton.jit def _pipe_bad_capacity(out_ptr): - buf = tle.gpu.alloc([1, 16, 16], dtype=tl.float32, layout=None, - scope=tle.gpu.smem, nv_mma_shared_layout=False) + buf = tle.gpu.alloc([1, 16, 16], dtype=tl.float32, layout=None, scope=tle.gpu.smem, nv_mma_shared_layout=False) # capacity must be > 0 -> ValueError _ = tle.pipe(capacity=0, data=buf) tl.store(out_ptr + tl.arange(0, 16), tl.zeros([16], tl.float32)) @@ -97,8 +96,7 @@ def test_pipe_bad_capacity_raises(): # --- tle.pipe: scope must be 'cta' in the MVP --- @triton.jit def _pipe_bad_scope(out_ptr): - buf = tle.gpu.alloc([2, 16, 16], dtype=tl.float32, layout=None, - scope=tle.gpu.smem, nv_mma_shared_layout=False) + buf = tle.gpu.alloc([2, 16, 16], dtype=tl.float32, layout=None, scope=tle.gpu.smem, nv_mma_shared_layout=False) _ = tle.pipe(capacity=2, scope="gpu", data=buf) # only 'cta' supported tl.store(out_ptr + tl.arange(0, 16), tl.zeros([16], tl.float32))