[BACKEND][TLE] add sunrise backend with tle

.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: sunrise3.6
+    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

.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/

.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/
   )

CMakeLists.txt

@@ -50,12 +50,35 @@ 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()
 
+# 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()
+
+# 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()
+
 project(triton CXX C)
 include(CTest)
 
@@ -131,7 +154,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 +494,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 +581,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 +613,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 +647,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()

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)

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}
 )

lib/Conversion/TritonToTritonGPU/TritonToTritonGPUPass.cpp

@@ -1,3 +1,9 @@
+#if __has_include("flagtree_spec.h")
+#include "flagtree_spec.h"
+#endif
+
+#ifndef FLAGTREE_SPEC_Conversion_TritonToTritonGPU_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

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}
 )

lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp

@@ -1,3 +1,9 @@
+#if __has_include("flagtree_spec.h")
+#include "flagtree_spec.h"
+#endif
+
+#ifndef FLAGTREE_SPEC_Dialect_Triton_Transforms_RewriteTensorPointer
+
 #include <stack>
 
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
@@ -565,3 +571,5 @@ class RewriteTensorPointerPass
 };
 
 } // namespace mlir::triton
+
+#endif

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}
 )

lib/Dialect/TritonGPU/IR/Dialect.cpp

@@ -1,3 +1,9 @@
+#if __has_include("flagtree_spec.h")
+#include "flagtree_spec.h"
+#endif
+
+#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_IR_Dialect
+
 #include "triton/Dialect/Triton/IR/Dialect.h"
 
 #include <cstdint>
@@ -4011,3 +4017,5 @@ std::optional<int> triton::gpu::getWarpSpecializeTag(Operation *op) {
   }
   return std::nullopt;
 }
+
+#endif

lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp

@@ -1,3 +1,9 @@
+#if __has_include("flagtree_spec.h")
+#include "flagtree_spec.h"
+#endif
+
+#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_IR_LinearLayoutConversion
+
 #include <vector>
 
 #include "triton/Dialect/Triton/IR/Utility.h"
@@ -1689,3 +1695,5 @@ chooseMfmaLikeStoreLayout(RankedTensorType valType) {
 }
 
 } // namespace mlir::triton::gpu
+
+#endif

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}
 )

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_Dialect_TritonGPU_Transforms_Prefetch
 
 #include "mlir/IR/IRMapping.h"
 #include "mlir/Support/LLVM.h"
@@ -455,3 +460,5 @@ struct PrefetchPass : public impl::TritonGPUPrefetchBase<PrefetchPass> {
 } // namespace gpu
 } // namespace triton
 } // namespace mlir
+
+#endif

lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp

@@ -1,3 +1,9 @@
+#if __has_include("flagtree_spec.h")
+#include "flagtree_spec.h"
+#endif
+
+#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_Transforms_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

lib/Dialect/TritonGPU/Transforms/Utility.cpp

@@ -1,3 +1,9 @@
+#if __has_include("flagtree_spec.h")
+#include "flagtree_spec.h"
+#endif
+
+#ifndef FLAGTREE_SPEC_Dialect_TritonGPU_Transforms_Utility
+
 #include "triton/Analysis/Utility.h"
 
 #include <fstream>
@@ -1773,3 +1779,5 @@ LogicalResult verifyBarrierType(Operation *op,
 }
 
 } // namespace mlir::triton
+
+#endif