[BACKEND] Support metax plugin (#6)

* [BACKEND] Support metax plugin

* [BACKEND] [TEST] Add metax python unit test

---------

Co-authored-by: ArinaJJH <jiangjh0930@163.com>

Author: zhzhcookie

.github/workflows/metax-build-and-test.yml

@@ -1,7 +1,10 @@
 name: Metax-Build-And-Test
 
 on:
-  workflow_call:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
@@ -12,18 +15,46 @@ jobs:
     runs-on: metax
     if: ${{ github.repository == 'FlagTree/flagtree' }}
     steps:
-      - name: Checkout code
+      - name: Checkout code (attempt 1)
+        id: checkout1
         uses: actions/checkout@v4
+        continue-on-error: true
+
+      - name: Sleep before checkout2
+        if: steps.checkout1.outcome == 'failure'
+        run: |
+          echo "First checkout attempt failed. Sleeping for 120 seconds before retry..."
+          sleep 120
+
+      - name: Checkout code (attempt 2)
+        id: checkout2
+        if: steps.checkout1.outcome == 'failure'
+        uses: actions/checkout@v4
+        continue-on-error: true
+
+      - name: Sleep before final checkout
+        if: steps.checkout1.outcome == 'failure' && steps.checkout2.outcome == 'failure'
+        run: |
+          echo "Second checkout attempt failed. Sleeping for 180 seconds before final retry..."
+          sleep 180
+
+      - name: Checkout code (final attempt)
+        if: steps.checkout1.outcome == 'failure' && steps.checkout2.outcome == 'failure'
+        uses: actions/checkout@v4
+
+      - name: Verify checkout success
+        if: success()
+        run: echo "Checkout completed successfully"
 
       - name: FlagTree Build on Metax
         shell: bash
         run: |
           source ~/env.sh
           export FLAGTREE_BACKEND=metax
           cd python
-          MAX_JOBS=20 pip3 install . --no-build-isolation
+          MAX_JOBS=32 python3.10 -m pip install . --no-build-isolation
 
       - name: FlagTree Test on Metax
         shell: bash
         run: |
-          pytest -s python/test/unit
+          python3.10 -m pytest -s third_party/metax/python/test/unit

CMakeLists.txt

@@ -79,6 +79,9 @@ else()
 endif()
 set(CMAKE_CXX_FLAGS "${CMAKE_C_FLAGS} -D__STDC_FORMAT_MACROS  -fPIC -std=gnu++17")
 
+if(FLAGTREE_BACKEND STREQUAL "metax")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DUSE_MACA -DUSE_MACA_OPAQUE_PTR -DUSE_BUILTIN -Wno-unused-result -Wno-attributes")
+endif()
 
 # #########
 # LLVM
@@ -322,6 +325,33 @@ if(TRITON_BUILD_PYTHON_MODULE)
       LLVMXCNCodeGen
       LLVMXCNAsmParser
     )
+  elseif(FLAGTREE_BACKEND STREQUAL "metax")
+    set(TRITON_LIBRARIES
+      ${triton_libs}
+      ${triton_plugins}
+
+      # mlir
+      MLIRMACADialect
+      MLIRGPUToMACATransforms
+      MLIRGPUToGPURuntimeTransforms
+      MLIRGPUTransforms
+      MLIRIR
+      MLIRControlFlowToLLVM
+      MLIRBytecodeWriter
+      MLIRPass
+      MLIRTransforms
+      MLIRLLVMDialect
+      MLIRSupport
+      MLIRTargetLLVMIRExport
+      MLIRMathToLLVM
+      MLIRGPUDialect
+      MLIRSCFToControlFlow
+      MLIRIndexToLLVM
+
+      # LLVM
+      LLVMPasses
+      LLVMNVPTXCodeGen
+    )
   endif()
 
   if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64" OR # Linux arm64
@@ -352,6 +382,10 @@ if(TRITON_BUILD_PYTHON_MODULE)
   add_compile_definitions(TRITON_BACKENDS_TUPLE=${TRITON_BACKENDS_TUPLE})
   if(FLAGTREE_BACKEND STREQUAL "cambricon")
     add_library(triton SHARED)
+  elseif(FLAGTREE_BACKEND STREQUAL "metax")
+    add_library(triton SHARED ${PYTHON_SRC_PATH}/main.cc
+                ${PYTHON_SRC_PATH}/interpreter.cc
+                ${PYTHON_SRC_PATH}/llvm.cc)
   else()
     add_library(triton SHARED ${PYTHON_SRC_PATH}/main.cc
                 ${PYTHON_SRC_PATH}/ir.cc

third_party/metax/CMakeLists.txt

@@ -0,0 +1,24 @@
+add_subdirectory(include)
+add_subdirectory(lib)
+
+if(TRITON_BUILD_PYTHON_MODULE)
+  if(FLAGTREE_PLUGIN)
+    add_subdirectory(plugin)
+    add_triton_plugin(TritonMetax
+      SHARED_LIB metaxTritonPlugin
+    )
+  else()
+    find_library(metaxTritonPluginLib
+      NAMES
+        metaxTritonPlugin.so
+      PATHS
+        ${CMAKE_CURRENT_SOURCE_DIR}
+      REQUIRED
+    )
+    add_triton_plugin(TritonMetax
+      SHARED_LIB ${metaxTritonPluginLib}
+    )
+  endif()
+endif()
+
+add_subdirectory(bin)

third_party/metax/LICENSE

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