Fix workflow and scripts for environment settings

.github/workflows/backend-test.yaml

@@ -88,4 +88,5 @@ jobs:
           CHANGED_FILES: ${{ inputs.changed_files }}
         run: |
           source .venv/bin/activate
+          source tools/set-env.sh ${{ inputs.vendor }}
           tools/test-op.sh ${{ inputs.pr_id }}

tools/set-env.sh

@@ -0,0 +1,47 @@
+VENDOR=$1
+echo "Setting up environment variable for vendor $VENDOR"
+
+case $VENDOR in
+  ascend)
+    if [ -f /usr/local/Ascend/ascend-toolkit/set_env.sh ]; then
+      source /usr/local/Ascend/ascend-toolkit/set_env.sh
+    fi
+    ;;
+  hygon)
+    source /opt/dtk-26.04/env.sh
+    ;;
+  iluvatar)
+    export LD_LIBRARY_PATH=/usr/local/corex/lib:$LD_LIBRARY_PATH
+    ;;
+  kunlunxin)
+    export LD_LIBRARY_PATH=/xcudart/lib:/usr/local/cuda/lib64
+    ;;
+  metax)
+    export MACA_PATH=/opt/maca
+    export LD_LIBRARY_PATH=$MACA_PATH/lib:$LD_LIBRARY_PATH
+    export LD_LIBRARY_PATH=$MACA_PATH/mxgpu_llvm/lib:$LD_LIBRARY_PATH
+    if [ -z "${USE_TRITON}" ]; then
+      SITE_PACKAGES=$VIRTUAL_ENV/lib/python3.12/site-packages
+      export LD_LIBRARY_PATH=${SITE_PACKAGES}/triton/backends/metax/lib:$LD_LIBRARY_PATH
+    fi
+    ;;
+  mthreads)
+    export MUSA_HOME=/usr/local/musa
+    export PATH=$MUSA_HOME/bin:$PATH
+    export LD_LIBRARY_PATH=$MUSA_HOME/lib:$LD_LIBRARY_PATH
+    export LD_LIBRARY_PATH=$VIRTUAL_ENV/lib:$LD_LIBRARY_PATH
+    if [ -z "${USE_TRITON}" ]; then
+      SITE_PACKAGES=$VIRTUAL_ENV/lib/python3.10/site-packages
+      export LD_LIBRARY_PATH=${SITE_PACKAGES}/triton/_C:$LD_LIBRARY_PATH
+    fi
+  tsingmicro)
+    export TX8_DEPS_ROOT=/opt/tsingmicro/tx8_deps
+    export LLVM_SYSPATH=/opt/tsingmicro/llvm21
+    export LLVM_BINARY_DIR=${LLVM_SYSPATH}/bin
+    export PYTHONPATH=${LLVM_SYSPATH}/python_packages/mlir_core:$PYTHONPATH
+    export LD_LIBRARY_PATH=$TX8_DEPS_ROOT/lib:/usr/local/kuiper/lib:$LD_LIBRARY_PATH
+    if [ -n "${USE_TRITON}" ]; then
+      SITE_PACKAGES=$VIRTUAL_ENV/lib/python3.10/site-packages
+      export PYTHONPATH=$SITE_PACKAGES/triton/backends/tsingmicro/llvm/python_packages/mlir_core
+    fi
+esac

tools/setup_iluvatar.sh

@@ -13,5 +13,5 @@ uv pip install --index ${FLAGOS_PYPI} \
 if [ -n "${USE_TRITON}" ]; then
   uv pip uninstall flagtree
   uv pip install --index $FLAGOS_PYPI \
-    "flagtree==0.5.0+iluvatar3.1"
+    "triton==3.1.0+corex.4.4.0"
 fi

tools/setup_metax.sh

@@ -11,14 +11,12 @@ uv pip install --index ${FLAGOS_PYPI} \
     "torchaudio==2.4.1+metax3.5.3.9" \
     "torchvision==0.15.1+metax3.5.3.9" \
 
-if [ -n "${USE_FLAGTREE}" ]; then
-  uv pip uninstall triton
+if [ -n "${USE_TRITON}" ]; then
+  uv pip uninstall flagtree
   uv pip install --index ${FLAGOS_PYPI} \
-    "flagtree==3.1.0+metax3.5.3.9"
-
+    "triton==3.0.0+metax3.5.3.9"
+else
   SITE_PACKAGES=$VIRTUAL_ENV/lib/python3.12/site-packages
   export LD_LIBRARY_PATH=${SITE_PACKAGES}/triton/backends/metax/lib:$LD_LIBRARY_PATH
 else
-  uv pip install --index ${FLAGOS_PYPI} \
-    "triton==3.0.0+metax3.5.3.9"
 fi

tools/setup_tsingmicro.sh

@@ -19,7 +19,8 @@ uv pip install --index ${FLAGOS_PYPI} \
     "txops==0.1.0+20260225.5cc33e4e"
 
 # Replace flagtree with Triton if requested
-if [ -n "${USE_FLAGTREE}" ]; then
+if [ -n "${USE_TRITON}" ]; then
+  uv pip uninstall flagtree
   uv pip install --index ${FLAGOS_PYPI} \
     "triton==3.3.0++gitfe2a28fa"
   # The following is needed when using `triton==3.3.0+gitfe2a28fa` rather than `flagtree`