build_thrust.sh

#!/bin/bash
# Thrust 并行编译和测试脚本 (使用 Ninja)
#
# 所有参数均可通过环境变量设置，命令行参数优先级更高
# 环境变量:
#   THRUST_JOBS          编译并行数 (默认: nproc)
#   THRUST_TEST_JOBS     测试并行数 (默认: 1)
#   THRUST_MUSA_ARCH     MUSA 目标架构 (默认: mp_31)
#   THRUST_MUSA_DEVICES  设置 MUSA_VISIBLE_DEVICES
#   THRUST_NO_CLEAN      设置为 1 不删除 build 目录
#   THRUST_EXCLUDE_TESTS 排除匹配正则表达式的测试
#   THRUST_BUILD_ONLY    设置为 1 仅编译不测试

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
THRUST_DIR="${SCRIPT_DIR}"
SOURCE_DIR="${THRUST_DIR}"

# 默认值 - 可通过环境变量覆盖
JOBS="${THRUST_JOBS:-$(nproc)}"
TEST_JOBS="${THRUST_TEST_JOBS:-1}"
RUN_TEST=true
TEST_VERBOSE="-V"
MUSA_DEVICES="${THRUST_MUSA_DEVICES:-}"  # 默认所有GPU可见
SKIP_CLEAN="${THRUST_NO_CLEAN:-false}"
BUILD_ONLY="${THRUST_BUILD_ONLY:-false}"
EXCLUDE_TESTS="${THRUST_EXCLUDE_TESTS:-async\\.(exclusive_scan|inclusive_scan)\\.large_indices}"  # 默认仅排除已知失败的测试用例
MUSA_ARCH="${THRUST_MUSA_ARCH:-mp_31}"  # 默认 MUSA 架构
BUILD_DIR=""  # 将在参数解析后设置
LOG_FILE=""
REPORT_FILE=""

# CUB 相关 - Thrust CUDA/MUSA 后端依赖 CUB
CUB_CMAKE_ARG=""
if [ -n "${CUB_DIR:-}" ]; then
    CUB_CMAKE_ARG="-DCUB_DIR=${CUB_DIR}"
fi
THRUST_CMAKE_ARG=""
if [ -n "${Thrust_DIR:-}" ]; then
    THRUST_CMAKE_ARG="-DThrust_DIR=${Thrust_DIR}"
fi

show_help() {
    cat << EOF
用法: $0 [选项] [命令]

命令:
  clean         仅清理 build 目录
  build         仅编译，不运行测试

选项:
  -j, --jobs N      编译并行数 (默认: $(nproc))
  -T, --test-jobs N 测试并行数 (默认: 1)
  -g, --gpus DEVICES 设置 MUSA_VISIBLE_DEVICES (如: 0,1,2,3)
  -a, --arch ARCH   MUSA 目标架构 (默认: mp_31, 支持: mp_21, mp_22, mp_31)
  -b, --build-dir DIR 指定构建目录名 (默认: build_<arch>, 如 build_mp_31)
  -n, --no-clean    不删除 build 目录 (增量编译)
  -E, --exclude RE  排除匹配正则表达式的测试 (默认: ${EXCLUDE_TESTS})
                    传空字符串 "" 可取消默认排除
  -h, --help        显示帮助

默认行为:
  1. 检查并安装 cub (如果缺失)
  2. 删除 build 目录
  3. CMake 配置
  4. 编译
  5. 运行测试 (ctest -V)
  6. 生成 markdown 报告

架构说明:
  mp_21  - S3000 系列 (PTX 210, Warp 128 线程)
  mp_22  - S4000 系列 (PTX 220, Warp 128 线程)
  mp_31  - S5000 系列 (PTX 310, Warp 32 线程)

示例:
  $0                          # 完整流程：清理、编译、测试、生成报告 (使用 build_mp_31)
  $0 -a mp_22                 # 使用 build_mp_22 目录
  $0 build                    # 仅编译
  $0 -n                       # 增量编译并测试
  $0 -T 4 -g 0,1,2,3          # 用4个并行测试，只用GPU 0-3
  $0 --build-dir custom       # 使用自定义构建目录 build_custom

环境变量:
  THRUST_JOBS          编译并行数
  THRUST_TEST_JOBS     测试并行数
  THRUST_MUSA_ARCH     MUSA 目标架构
  THRUST_MUSA_DEVICES  设置 MUSA_VISIBLE_DEVICES
  THRUST_NO_CLEAN      设置为 1 不删除 build 目录
  THRUST_EXCLUDE_TESTS 排除匹配正则表达式的测试
  THRUST_BUILD_ONLY    设置为 1 仅编译不测试
EOF
}

# 解析参数
CUSTOM_BUILD_DIR=""
DO_CLEAN=false
while [[ $# -gt 0 ]]; do
    case $1 in
        -j|--jobs)
            JOBS="$2"
            shift 2
            ;;
        -T|--test-jobs)
            TEST_JOBS="$2"
            shift 2
            ;;
        -g|--gpus)
            MUSA_DEVICES="$2"
            shift 2
            ;;
        -a|--arch)
            MUSA_ARCH="$2"
            shift 2
            ;;
        -b|--build-dir)
            CUSTOM_BUILD_DIR="$2"
            shift 2
            ;;
        -E|--exclude)
            EXCLUDE_TESTS="$2"
            shift 2
            ;;
        -n|--no-clean)
            SKIP_CLEAN=true
            shift
            ;;
        clean)
            DO_CLEAN=true
            shift
            ;;
        build)
            BUILD_ONLY=true
            RUN_TEST=false
            shift
            ;;
        -h|--help)
            show_help
            exit 0
            ;;
        *)
            echo "未知选项: $1"
            show_help
            exit 1
            ;;
    esac
done

# 设置构建目录 (基于架构自动命名，除非指定了自定义目录)
if [ -n "$CUSTOM_BUILD_DIR" ]; then
    BUILD_DIR="${THRUST_DIR}/build_${CUSTOM_BUILD_DIR}"
else
    BUILD_DIR="${THRUST_DIR}/build_${MUSA_ARCH}"
fi
LOG_FILE="${BUILD_DIR}/test_verbose.log"
REPORT_FILE="${THRUST_DIR}/test_report_${MUSA_ARCH}.md"

# 处理 clean 命令 (需要在设置 BUILD_DIR 后)
if [ "$DO_CLEAN" = true ]; then
    echo "清理构建目录: ${BUILD_DIR}"
    rm -rf "${BUILD_DIR}"
    exit 0
fi

# 1. 清理 build 目录
if [ "$SKIP_CLEAN" = false ] && [ -d "${BUILD_DIR}" ]; then
    echo "删除 build 目录..."
    rm -rf "${BUILD_DIR}"
fi

# 3. CMake 配置
echo ""
echo "=========================================="
echo "CMake 配置 (Ninja)..."
echo "MUSA 架构: ${MUSA_ARCH}"
echo "=========================================="

# 从 MUSA_ARCH 提取架构数字 (mp_22 -> 22)
ARCH_NUM="${MUSA_ARCH#mp_}"

cmake -G Ninja \
    -DCMAKE_CXX_COMPILER=mcc \
    ${THRUST_CMAKE_ARG} \
    ${CUB_CMAKE_ARG} \
    -DMUSA_64_BIT_DEVICE_CODE=ON \
    -DMUSA_ARCH=${MUSA_ARCH} \
    -DTHRUST_DISABLE_ARCH_BY_DEFAULT=ON \
    -DTHRUST_ENABLE_COMPUTE_${ARCH_NUM}=ON \
    -DTHRUST_ENABLE_TESTING=ON \
    -DTHRUST_ENABLE_EXAMPLES=OFF \
    -DTHRUST_ENABLE_HEADER_TESTING=OFF \
    -DTHRUST_INSTALL_CUB_HEADERS=ON \
    -S "${SOURCE_DIR}" -B "${BUILD_DIR}"

# 4. 并行编译
echo ""
echo "=========================================="
echo "编译 (并行数: ${JOBS})..."
echo "=========================================="
cmake --build "${BUILD_DIR}" -j "${JOBS}"

if [ "$RUN_TEST" = true ]; then
    echo ""
    echo "=========================================="
    echo "运行测试 (并行数: ${TEST_JOBS})"
    echo "=========================================="

    # 设置环境变量
    if [ -n "$MUSA_DEVICES" ]; then
        export MUSA_VISIBLE_DEVICES="$MUSA_DEVICES"
        echo "MUSA_VISIBLE_DEVICES=${MUSA_DEVICES}"
    fi

    # 运行 ctest 并保存输出
    echo "测试输出保存到: ${LOG_FILE}"
    EXCLUDE_ARG=""
    if [ -n "$EXCLUDE_TESTS" ]; then
        EXCLUDE_ARG="-E ${EXCLUDE_TESTS}"
        echo "排除测试: ${EXCLUDE_TESTS}"
    fi
    ctest --test-dir "${BUILD_DIR}" -j "${TEST_JOBS}" ${TEST_VERBOSE} ${EXCLUDE_ARG} 2>&1 | tee "${LOG_FILE}"

    # 生成 markdown 报告
    echo ""
    echo "=========================================="
    echo "生成测试报告"
    echo "=========================================="
    python3 "${THRUST_DIR}/parse_ctest_log.py" "${LOG_FILE}" "${REPORT_FILE}"
    echo "测试报告: ${REPORT_FILE}"
fi

echo ""
echo "完成!"