# JAX-Flagtree-HCU

**Repository Path**: onescience-ai/jax-fl

## Basic Information

- **Project Name**: JAX-Flagtree-HCU
- **Description**: 基于flagtree实现的jax 0.6.0/xla适配框架，运行在海光HCU(gfx936&gfx938)。
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-09-02
- **Last Updated**: 2026-11-10

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# JAX–FlagTree DCU

面向 DCU 的 JAX 0.6.0 编译与运行系统，以 **TTIR → FlagTree → PJRT** 为主要数值执行路线，保留计算库、控制流、通信和其他必要的原生执行路径。

本项目复用 JAX 的 Python API、JAXPR、自动微分和程序变换，通过仓库内修改后的 XLA 完成融合、数值发射、编译及设备执行。用户安装兼容的 FlagTree 和本项目的四个 JAX wheel 后，可以直接运行自己的 JAX 程序，不需要开发者的 POC 目录或构建目录。

> **当前源码快照：2026-11-03，候选构建标识 `1021045`，四个JAX包的版本号仍为 `0.6.0`。** 本次同步 `jax/`、`xla/` 和本文档，不重新构建、不安装、不替换 `wheels/`，也不提交或推送Git。仓库现有 `wheels/` 是历史交付集合，不能作为当前源码的配套包；应以各自的SHA256和源码清单识别版本。

当前候选源码快照SHA256为 `7dd46d92c9e5878c3f650761dce6aaa0e1f5342f05eae9b4bea5836cf1216b46`，已安装四包的680份有效载荷清单SHA256为 `171b95d71fd5d080f53e0efec35a032a0a65e0b8ac2bf14ec547a505e7d63ba6`。完整清单用于开发站点溯源，不是安装运行依赖。版本说明和验证边界见[当前候选记录](docs/CANDIDATE_SNAPSHOT_20261103.md)。

本次保留既有别名/输出生命周期、显存预算、BF16舍入及有限性保护，并同步以下通用改造：

- 按图依赖划分输出与归约区域，减少重复计算和中间张量读写；保留合法原执行方案参与比较。
- 基于实际代码生成、资源反馈和有预算的同步实测选择布局与执行方案，缓存结果；结构化只读索引和数据坐标变换保留精确输出比较。
- 在HLO数值权限内选择矩阵计算库方案，保留合法库路径和有界调优；融合偏置、残差等尾部计算时保留应有的中间舍入。
- 按真实数据类型和已证明的坐标域处理索引位宽，并改进编译缓存复用和完成事件调度。
- `--xla_gpu_exclude_nondeterministic_ops` 开启时使用确定性Scatter等专门路径；这一改造不把默认未开启状态改成确定性执行，也不承诺跨平台逐位一致。

这些是通用编译/执行机制，不以模型名、测试case名或固定rank作为放行条件。有限的验收集合不能证明所有形状都更快。

## 1. 架构与执行边界

```text
Python / JAX API
  → JAXPR / autodiff / vmap / pmap 等变换
  → StableHLO → XLA HLO 优化与融合
      ├─ TTIR 数值路线 → TTIR 契约转换 → FlagTree → HSACO → kernel thunk ──┐
      ├─ 计算库路线 → BLAS / MIOpen / FFT 等库调用 thunk ────────────────┤
      ├─ 控制流、通信、传输及其他功能性调度 → 对应 thunk ────────────────┤
      └─ 模式允许的原生数值路线 → LLVM IR → DTK GCVM → HSACO ────────────┤
                                                                       → PJRT / DCU
```

- 分流依据是 HLO 语义、融合结构、索引、布局和发射能力，不是 Python API 或模型名称。
- TTIR 主线的融合候选阶段不执行 tile 搜索；融合图稳定后才选择最终 tile，采用有预算的搜索，而不是在每个融合候选中反复展开 tile 组合。
- GEMM、卷积及既定库算子保留合法的计算库路线。控制流、collective、custom call 等也保留其合法执行契约。这些不是“严格模式偷偷 fallback”。
- 严格模式拒绝未经允许的数值回退；它不要求调度、库调用或所有辅助工作都变成 TTIR。
- 兼容模式允许回退，但“程序运行成功”不等于所有计算都由 FlagTree 执行。路线验收应使用 `fl_only` 并保留测试产物。

## 2. 仓库与四个安装包

```text
jax/       JAX Python、jaxlib、ROCm 插件及 wheel 构建入口
xla/       本项目使用的 XLA、TTIR 发射与后端改造
docs/      覆盖审计、设计与历史验收记录
wheels/    历史交付集合及SHA256SUMS；本次源码更新不替换此目录
```

| 安装包 | 作用 | 目标版本 |
| --- | --- | --- |
| `jax` | Python API、变换和前端 | `0.6.0` |
| `jaxlib` | JAX/XLA Python 原生扩展 | `0.6.0` |
| `jax-rocm60-plugin` | ROCm 原生扩展与算子注册 | `0.6.0` |
| `jax-rocm60-pjrt` | ROCm PJRT、XLA 后端和打包的 FlagTree 接入组件 | `0.6.0` |

四包必须来自同一已验证交付集合，不能按文件名相同就混用不同构建。当前源码同步没有更新仓库内历史wheel；下面的构建和安装命令针对你选定并核对SHA256的交付集合。版本均不带 `dev`、日期或其他后缀；源码修订身份需另外用 Git 状态、补丁记录和 wheel SHA256 追溯。

FlagTree 编译器本身单独安装；PJRT wheel 包含适配器、TTIR 转换组件、worker 和 `jax-flagtree-compile` 入口。DTK、驱动、系统动态库、模型数据和编译缓存不包含在四个 wheel 中。

## 3. 环境要求

以下是当前验证环境，不代表其他平台已获得同等覆盖：

| 项目 | 当前基线 |
| --- | --- |
| 系统 / Python | Linux x86_64 / Python 3.10 |
| DCU目标 | 当前SCNet验证为`gfx936`；源码包含`gfx938`适配，具体wheel的目标覆盖以构建记录为准 |
| DTK | `26.04`，须有完整运行库和编译组件 |
| 主机编译器 | GCC `12.2.0`，以及匹配的 `libstdc++` |
| Bazel | `7.4.1`，只在源码构建时需要 |
| ROCm wheel 包名 | `rocm60`；不是对任意 ROCm 6.x 安装的兼容承诺 |
| FlagTree | 插件接受`>=0.6.0,<0.8.0`，并校验实际编译器文件与接口；当前SCNet实测版本为`0.6.0+hcu3.6.glibc228.compat1.git8a1dda94` |

当前 FlagTree 接入使用 DTK 中的 DCC、链接器及设备 bitcode；原生兼容路线需要 DTK 的 GCVM。仅有 HIP 运行库而没有这些组件的精简安装不满足完整要求。安装位置可不同，但必须与本项目及所安装 FlagTree 的 ABI/目录契约一致。

## 4. 从源码构建四个 wheel

本仓库的 `jax/` 和 `xla/` 必须一起使用。BW1100 的入口是 `scripts/build_bw1100.sh`；脚本通过 `--local_xla_path` 使用当前仓库的 XLA，生成同一源码版本的 `jax`、`jaxlib`、`jax-rocm60-plugin` 和 `jax-rocm60-pjrt` 四个 `0.6.0` wheel。它只构建和核验，不安装到现有 JAX 运行环境。

### 4.1 一键构建与指定输出目录

在 BW1100 上进入**包含 `jax/`、`xla/` 和完整 `scripts/` 的同一份仓库**，运行：

```bash
cd /home/testuser/liuyx19/jax-fl
bash scripts/build_bw1100.sh
```

默认输出到 `仓库根目录/wheelhouse/bw1100/<运行时间>/`。若需指定 wheel 输出位置，给 `WHEEL_ROOT` 一个可写的**绝对路径**；每次使用新的空目录，便于四包和构建回执保持独立：

```bash
cd /home/testuser/liuyx19/jax-fl
WHEEL_ROOT=/data_add/testuser/liuyx19/wheels/jax-fl-$(date +%Y%m%d_%H%M%S) \
  bash scripts/build_bw1100.sh
```

默认 `GPU_TARGETS=gfx938`、`JOBS=64`、`DTK_ROOT=/home/testuser/dtk-26.04`、`BUILD_ROOT=/data_add/testuser/liuyx19/jax-fl-build`。例如，可在保持脚本和源码不变的情况下设置 `JOBS=32` 或另行设置 `BUILD_ROOT`。构建日志和阶段结果写入 `$BUILD_ROOT/logs/<运行时间>/`；输出目录包含四个 wheel、`SHA256SUMS` 和 `build_manifest.json`。构建盘需有 **超过 40 GiB** 可用空间；首次构建通常还需更多空间。脚本不要求占用 GPU，但原生构建耗用较多 CPU、内存和磁盘，应在获准的构建节点执行。

### 4.2 当前 BW1100 已备好的工具、库与离线缓存

以下是**当前这台机器**的路径，不是 wheel 的运行时内容。使用默认 `BUILD_ROOT` 时，脚本读取或复用这些位置：

| 用途 | 路径 |
| --- | --- |
| DTK 26.04 环境及编译器 | `/home/testuser/dtk-26.04/env.sh`、`/home/testuser/dtk-26.04/llvm/bin/clang` |
| 独立 Python 3.10、GCC/G++ 12 构建环境 | `/data_add/testuser/liuyx19/jax-fl-build/env/` |
| Bazel 7.4.1 可执行文件 | `/data_add/testuser/liuyx19/jax-fl-build/tools/bazel-7.4.1` |
| Bazel 外部源码下载缓存 | `/data_add/testuser/liuyx19/jax-fl-build/repository-cache/`、`/data_add/testuser/liuyx19/jax-fl-build/downloads/` |
| 固定的 Triton 源码包 | `/data_add/testuser/liuyx19/jax-fl-build/downloads/triton_integrate_branch-1.3.tar.gz` |
| Bazel 已展开的外部仓库和增量构建状态 | `/data_add/testuser/liuyx19/jax-fl-build/bazel-user-root/` |
| 构建用 Python 依赖 wheel 与哈希清单 | `/data_add/testuser/liuyx19/jax-fl-build/python-wheels/`、其中的 `manifest.json` |
| Conda 和 pip 下载缓存 | `/data_add/testuser/liuyx19/jax-fl-build/conda-pkgs/`、`/data_add/testuser/liuyx19/jax-fl-build/pip-cache/` |

脚本还需要系统已有的 `libnuma` 开发链接库；只会在隔离构建环境中补建指向它的链接。Bazel、DTK、固定源码依赖和构建用 Python wheel 都不是四个 JAX 成品包。`python-wheels/` 存在时，脚本按其中的哈希清单检查文件，并只为 Bazel 生成的构建 Python 配置离线 wheel 索引；不修改用户全局 pip 设置。`downloads/` 和 `repository-cache/` 用于加速或支持外部依赖获取，`bazel-user-root/` 用于复用增量构建；脚本不执行 `bazel clean`。

### 4.3 在一台没有上述缓存的新 BW1100 上运行

先取得含 `jax/`、`xla/` 和 `scripts/build_bw1100.sh` 及其辅助脚本的同一份源码，并准备 DTK 26.04、Conda、系统 `libnuma` 开发库、可写构建盘及所需的依赖下载源。**当前这台机器的 `scripts/` 尚是 Git 未跟踪目录；仅克隆现有 Git 版本不能保证取得这些脚本，须先确认脚本已随源码一并发布或复制。** 直接按 4.1 运行即可，无须手动预建构建环境或先加载 DTK。

脚本会创建 `$BUILD_ROOT` 及独立 Conda Python 3.10 / GCC 12 环境，加载 `$DTK_ROOT/env.sh`，下载并校验缺失的 Bazel 7.4.1，配置隔离主机编译器和 DTK 库链接，准备 Bazel 锁定的外部依赖与构建 Python，然后依次构建 `jaxlib`、ROCm plugin/PJRT、`jax`。最后核对四包的版本、ZIP 完整性、PJRT 接入组件、源码身份和 SHA256，并在输出目录生成清单。它不会安装四包、修改运行中的容器，也不会运行模型测试。

新机器若**没有网络**，须事先准备可用的 Conda 环境/包缓存、Bazel 可执行文件、构建用 Python wheel 与 `manifest.json`，以及覆盖本次源码依赖的 `downloads/`、`repository-cache/` 或等效 Bazel 外部仓库缓存；仅有空目录不会完成离线构建。若依赖站点代理，可设置 `BUILD_NETWORK_ENV=/绝对路径/network.env`，由脚本加载网络环境变量；脚本不会关闭 TLS 或哈希校验。新机器首次构建会完整编译，只有同一构建根目录内的后续构建才可复用增量缓存。

### 5 最小设备与数值检查

请在自己的程序目录运行，避免工作目录是 `jax/` 源码根目录而意外导入源码包。

```bash
python -u - <<'PY'
import jax
import jax.numpy as jnp
import numpy as np

print("JAX:", jax.__version__, jax.__file__, flush=True)
print("Devices:", jax.devices(), flush=True)
assert jax.default_backend() == "gpu", "没有使用 DCU/GPU，不能记为通过"

x = np.linspace(-1, 1, 256, dtype=np.float32)
y = jax.jit(lambda a: jnp.tanh(a * 1.25 + 0.5))(jnp.asarray(x))
y.block_until_ready()
np.testing.assert_allclose(np.asarray(y), np.tanh(x * 1.25 + 0.5),
                           rtol=1e-5, atol=1e-6)
print("DCU smoke test: PASS", flush=True)
PY
```

此检查不等价于完整算子覆盖。若需证明没有数值 fallback，使用下一节的严格模式运行测试。

## 6. 模式、日志与调优

在启动 Python **之前**设置模式，使用小写的完整名称：

| `XLA_MODE` | 行为 |
| --- | --- |
| 不设置 / `back_unsupported` | 默认。TTIR–FlagTree 优先；对明确未实现的状态允许原生回退，不把所有编译错误都当作不支持。 |
| `fl_only` | 严格。禁止未经允许的数值 fallback，保留既定库与功能性路线；开发和路线验收使用此模式。 |
| `back_fail` | 对接入层允许回退的编译失败扩大原生回退范围；不保证从 Python 异常、设备运行错误或错误数值中恢复。 |
| `native` | 关闭本项目 TTIR–FlagTree 及原生 Triton GPU 选择，采用原生非 Triton 路线；需生成设备代码的 LLVM IR 由 DTK GCVM 编译，库和功能性算子仍走既定路径。 |

```bash
XLA_MODE=fl_only python your_program.py
XLA_MODE=back_unsupported python your_program.py
XLA_MODE=back_fail python your_program.py
XLA_MODE=native python your_program.py

# 仅排查问题时打开本项目诊断。
XLA_MODE=fl_only XLA_FLAGTREE_LOG=1 python your_program.py
```

模式开关不能弥补缺少编译器、DTK 或动态库的环境错误。不要使用历史 `XLA_FLAGTREE_BACKEND` / `XLA_FLAGTREE_FALLBACK_POLICY` 配置代替 `XLA_MODE`，也不需要照搬早期试验的内部开关组合。

当前源码将 `xla_gpu_autotune_level` 默认设为 `0`，仍保留该参数供高级用户显式修改。它不是开启 TTIR 路线的必需环境变量，也不等于关闭最终融合图的 tile 选择。下面是需要显式固定时的示例；已有其他 `XLA_FLAGS` 时应合并并消除同名重复项：

```bash
XLA_FLAGS='--xla_gpu_autotune_level=0' python your_program.py
```

## 7. 测试与发布判定

`tests/jax/tests/flagtree_regression/cases.json` 固定了 JAX 0.6.0 的 **11328 个不同 nodeid**，SHA256 为 `260dd3152aa0ae2c057463cb51df1fecd5d4bd5e5cf83bde9ae8a85592d4bba7`。配套的 137 个官方测试源码文件、原始运行器、离线依赖 wheel、PJRT 测试配置及 `docs/autodidax.py` 均位于 `tests/jax/`；迁移文件的逐项哈希在 `tests/source_manifest.json`。这是一组历史通过用例的固定回归集合，**不是上游 JAX 全部测试的覆盖承诺**。

在 BW1100 宿主机上，先用不占 GPU 的命令核对清单和容器入口：

```bash
cd /home/testuser/liuyx19/jax-fl
bash tests/run_bw1100_official.sh --list
```

确认有 **8 张独占空闲 HCU** 后，一键启动整套测试：

```bash
cd /home/testuser/liuyx19/jax-fl
bash tests/run_bw1100_official.sh
```

脚本默认从 16 张卡中选出 8 张当前无计算进程的卡；也可指定已确认空闲的物理卡，例如 `GPU_DEVICES=0,1,2,3,4,5,6,7 bash tests/run_bw1100_official.sh`。若不足 8 张或指定卡已被占用，脚本会拒绝启动，不会取消其他任务。空闲检查不能代替独占资源安排，应避免其他任务在检查后抢占这些卡。

宿主机入口会先核对 11328 条清单及全部迁移文件的哈希，再将原始测试文件的副本放到现有 `flagos-jax-test` 容器的 `/work` 映射工作区，使用**该容器当前已安装的** JAX 环境运行原始 `flagtree_regression/run.py`。运行模式固定为 `fl_only`，默认 8 个 worker、每批 128 条、每批超时 1800 秒；包括单卡、8 卡、x64、可选依赖和独立进程 profile。不会在这一步重新构建或安装 wheel。

结果保存在宿主机 `/home/testuser/liuyx19/flagos-test-20261027/official_11328_runs/<运行时间>/`，包括 `inventory.json`、`provenance.json`、逐批日志与结果、`summary.json` 和 `nonpassing_cases.json`。只有同一轮固定清单中的 11328 条全部报告 `PASS`，且无新增 `SKIP`、`FAIL`、`ERROR`、`TIMEOUT`、`MISSING` 或进程/收集失败，才能将这项固定回归判为通过。`--list` 成功或脚本启动成功均不等于测试通过；正式发布还须结合相应版本的模型精度、性能及源码与 wheel 身份核验。

## 8. 常见问题

| 现象 | 优先检查 |
| --- | --- |
| 插件找不到 FlagTree、版本不匹配或 `jax-flagtree-compile` 缺失 | 当前Python环境、FlagTree版本范围及实际接口/ABI、四包是否属于同一已验证交付集合，以及是否被外部 `PYTHONPATH` 覆盖。 |
| `librccl`、`libgalaxyhip` 等动态库找不到 | 是否在计算节点加载了匹配 DTK，是否保留 DTK 设置的动态库搜索路径。运行时不要照抄构建时仅保留 GCC 的 `LD_LIBRARY_PATH`。 |
| `libhipfftMp` / `ompi_mpi_int` 错误 | 当前源码已收紧普通 hipFFT 的链接选择，避免把多节点 hipFFT-Mp 误链接成普通 FFT 依赖。先核对实际加载的 PJRT wheel 和动态依赖，不要默认使用旧包或盲目预加载任意 MPI。四包不捆绑站点 MPI。 |
| `CMAKE_PREFIX_PATH: unbound variable` | 在严格 shell 脚本中，将 `set -u` 放在两份 DTK 环境脚本加载之后。 |
| 默认设备只有 CPU | 检查作业分配、设备可见性和 PJRT 初始化错误；不能通过忽略 GPU 错误宣称 DCU 测试通过。 |
| wheel 名称相同但结果不同 | 比较四包 SHA256、源码快照、DTK、FlagTree 与测试配置；`0.6.0` 版本号本身不能唯一标识本项目构建。 |

## 9. 许可证与源码基线

保留 JAX、XLA 及第三方组件的原有许可声明，参见 [JAX LICENSE](jax/LICENSE) 和 [XLA LICENSE](xla/LICENSE)。FlagTree 与 DTK 按各自的分发和许可条件使用。

当前构建树的上游起点分别为 JAX `127aa7621868cb77e552b5d1f90e4a42b09c13fa`、XLA `0d1b60216ea13b0d261d59552a0f7ef20c4f76c5`。本项目改造位于这些起点之上；仅检出这两个上游提交不会得到本项目实现。

本次仅同步当前候选JAX/XLA源码及文档，不替换历史wheel，不重新构建或安装，也不提交、推送Git。必要的站点外部工具链适配仍按第4.1节处理，本次不是全新克隆的干净构建验收。正式支持声明应绑定源码快照、实际wheel的SHA256和独立验收报告；历史版本结果不能拼接为当前构建全通过。