DragonArrayLoongson optimization lab

DragonArray docs

Build TensorFlow 2.21.0 for LoongArch from source

DragonArray's TensorFlow patch adds the LoongArch Bazel platform and wheel tag, CPython 3.13 toolchain support, GCC 15 compatibility fixes, old-world build shims, sparse BF16 fixes, and the Eigen convolution changes used by the LASX path.

The same patch supports three CPU profiles, but the compiler flags are not interchangeable. Choose one profile and use it consistently for TensorFlow, the static C++ runtime, NumPy, and every other native dependency.

Source revision and patch

Build this exact TensorFlow revision:

tag:     v2.21.0
commit:  a481b10260dfdf833a1b16007eead49c1d7febf3

Download the tensorflow-2.21.0-loongarch64.patch.gz patch. The published file has these checksums:

compressed .gz: b6804b551a01693c517eab7f26f7bd571e558b52df94c3b87425f9a45a8ad9af
uncompressed:   15de4da62c2d1309a39362a7a44cf2a9f948cab7aaf9abc772a969b4e47134d4

Start with a clean checkout and apply the consolidated patch once:

git clone --branch v2.21.0 --depth 1 \
  https://github.com/tensorflow/tensorflow.git tensorflow-2.21.0
cd tensorflow-2.21.0

test "$(git rev-parse HEAD)" = \
  a481b10260dfdf833a1b16007eead49c1d7febf3

curl --fail --location --output tensorflow.patch.gz \
  https://dragon-array.org/download/tensorflow-2.21.0-loongarch64.patch.gz
printf '%s  %s\n' \
  b6804b551a01693c517eab7f26f7bd571e558b52df94c3b87425f9a45a8ad9af \
  tensorflow.patch.gz | sha256sum -c -

gzip -dc tensorflow.patch.gz > tensorflow.patch
printf '%s  %s\n' \
  15de4da62c2d1309a39362a7a44cf2a9f948cab7aaf9abc772a969b4e47134d4 \
  tensorflow.patch | sha256sum -c -

git apply --check tensorflow.patch
git apply tensorflow.patch
git diff --check
git apply --reverse --check tensorflow.patch

The final command is read-only: it proves that the complete patch is present. Do not stack older incremental TensorFlow or XLA patches on top of this file; their changes are already included. Patches for h5py, grpcio, and prepared Bazel override repositories remain separate because they modify other source trees.

Choose the CPU profile

Use lscpu on the target system if the processor model is not known:

LC_ALL=C lscpu | grep -E 'Architecture|Model name|Flags'
Target systems ISA level Required compiler flags Example Bazel output base
Loongson 3A5000 and 3B5000 LA464, LSX only -march=la464 -mtune=la464 -mlsx -mno-lasx ~/.cache/bazel/tf-2.21.0-la464-lsx
Loongson 3C5000 LA464, LASX -march=la464 -mtune=la464 -mlasx ~/.cache/bazel/tf-2.21.0-la464-lasx
Loongson 3A6000, 3B6000, and 3C6000 series LA664, LASX -march=la664 -mtune=la664 -mlasx ~/.cache/bazel/tf-2.21.0-la664-lasx

Every profile also uses -O2 -g0 -fPIC -mcmodel=medium for native release builds.

The differences matter:

Do not use -march=native for a redistributable wheel. Apply the selected ISA flags to target and host C/C++ actions, the static libstdc++.a, TensorFlow, NumPy, HDF5/h5py, grpcio, ml_dtypes, and optree. Keep a separate build tree, virtual environment, wheel directory, and Bazel output base for every profile.

Tested build environment

The DragonArray wheels were produced with:

GCC                 15.2.0  /opt/loongson-gcc-15.2.0
Bazel                8.4.2
CMake                4.3.2
Ninja                 1.10
OpenJDK              21.0.10
pyenv                 2.8.1
build Python         3.13.13
validation Python    3.13.14

Allow at least 15 GiB of RAM plus swap and 60 GiB of free disk space for an isolated Bazel output tree. The reference builds use four jobs and a 10 GB Bazel memory budget. They use the old-world LoongArch LP64 ABI and glibc 2.28.

A LASX-capable 3A6000 can build all three profiles. The 3A5000 profile is safe only when -mno-lasx is applied everywhere. Code produced for LA664 cannot be executed for validation on an LA464 machine.

Prepare Python and source inputs

TensorFlow 2.21.0 supports CPython 3.13. The reference build uses CPython 3.13.13 and validates the finished wheel in a clean CPython 3.13.14 environment:

export PYENV_ROOT="$HOME/.pyenv"
export PATH="$PYENV_ROOT/bin:$PATH"
eval "$(pyenv init -)"

pyenv install 3.13.13
pyenv install 3.13.14

The Bazel build uses a hermetic CPython 3.13.13 archive with a top-level python/ directory:

python-3.13.13-loongarch64.tar.gz
SHA-256 ca540b5ab4804d27fb8f6ddae1e4792fcf2b4050dfc88f22ed227fd5fc673b84

Pass it through these repository settings:

--repo_env=HERMETIC_PYTHON_VERSION=3.13
--repo_env=HERMETIC_PYTHON_URL=file:///path/to/python-3.13.13-loongarch64.tar.gz
--repo_env=HERMETIC_PYTHON_SHA256=ca540b5ab4804d27fb8f6ddae1e4792fcf2b4050dfc88f22ed227fd5fc673b84
--repo_env=HERMETIC_PYTHON_PREFIX=python

The complete wheel set also requires source builds of:

The experimental NumPy variant is built from adamdeprince/numpy, branch loongson-experimental, pinned by the reference build to commit 79b033101a38a0f1fa2c6b94380f31041f74da55.

Prepare the Bazel overrides

The TensorFlow patch enables the LoongArch build, but TensorFlow's external repository graph still needs LoongArch-capable source snapshots. The reference build supplies prepared overrides for:

platforms
rules_python
eigen_archive
ducc
tf_runtime
io_bazel_rules_closure
build_bazel_rules_android
local_execution_config_python
local_config_python
local_execution_config_platform
system_python
cpuinfo
skcms
png
pybind11_protobuf
pybind11_abseil
boringssl
org_brotli
hwloc

Pass each directory with --override_repository=name=/absolute/source/path. Use empty local repository stubs for unsupported prebuilt x86/AArch64 PyPI repositories such as grpcio, optree, h5py, SciPy, libclang, and TensorBoard's data server. This keeps foreign wheels out of the LoongArch build graph.

The consolidated TensorFlow patch does not embed these third-party source trees. Preserve their licenses and source revision records in any distributed build kit.

Build a target-matched static C++ runtime

TensorFlow's roughly 2 GiB shared library exceeds the LoongArch normal code model. Rebuild a PIC, medium-model libstdc++.a from GCC 15.2 source using the same CPU flags chosen above:

gcc_root=/opt/loongson-gcc-15.2.0
source_root=/path/to/gcc-15.2.0
build_root=/path/to/libstdcxx-build
prefix=/path/to/target-libstdcxx
cpu_flags='-march=la464 -mtune=la464 -mlsx -mno-lasx'  # choose from the table
common_flags="-O2 -g0 -fPIC -mcmodel=medium $cpu_flags"
bits_include="$gcc_root/include/c++/15.2.0/loongarch64-linux-gnu/bits"

mkdir -p "$build_root"
cd "$build_root"
CC="$gcc_root/bin/gcc" \
CXX="$gcc_root/bin/g++" \
CFLAGS="$common_flags" \
CXXFLAGS="$common_flags -I$bits_include" \
  "$source_root/libstdc++-v3/configure" \
    --prefix="$prefix" \
    --build=loongarch64-linux-gnu \
    --host=loongarch64-linux-gnu \
    --disable-multilib \
    --disable-shared \
    --enable-static \
    --disable-libstdcxx-pch \
    --enable-libstdcxx-threads

make -j4
make install
test -f "$prefix/lib/libstdc++.a"

Use this archive for target and host links with -static-libstdc++. Keep libgcc_s.so.1 dynamic; forcing the installed libgcc_eh.a into the large TensorFlow link can reintroduce normal-model relocation overflows.

For the LA464/LSX profile, disassemble the archive and reject any LASX instruction before continuing. For the LASX profiles, verify that representative LASX instructions are present and execute correctly on a compatible host.

Build the native dependency wheels

Create a clean CPython 3.13.13 virtual environment for the selected profile. Export the chosen flags through CFLAGS, CXXFLAGS, and the build-system specific options for every native project:

export CFLAGS="-O2 -g0 -fPIC -mcmodel=medium $cpu_flags"
export CXXFLAGS="$CFLAGS"

Build the standard NumPy, HDF5/h5py, ml_dtypes, optree, and grpcio wheels before TensorFlow. Build the experimental NumPy separately if desired.

NumPy's Meson option remains -Dcpu-baseline=LSX for all three profiles because NumPy exposes LSX, not LASX, as its named LoongArch baseline. The global compiler flags still control whether LASX code may be emitted. The 3A5000/3B5000 build must retain -mno-lasx.

When statically linking the private C++ runtime into ml_dtypes, use:

-Wl,--exclude-libs,libstdc++.a

This prevents private static C++ symbols from interposing on TensorFlow, grpcio, optree, or another extension when import order changes.

Configure and build TensorFlow

Use GCC 15.2 for both target and host actions. The common Bazel configuration used by the reference build includes:

--enable_workspace
--noenable_bzlmod
--config=opt
--config=nogcp
--config=nonccl
--config=clang_local
--define=tflite_with_ruy=true
--define=tensorflow_mkldnn_contraction_kernel=0
--platforms=//:linux_loongarch64
--host_platform=//:linux_loongarch64
--repo_env=TF_DOWNLOAD_CLANG=0
--repo_env=USE_PYWRAP_RULES=1
--repo_env=WHEEL_NAME=tensorflow
--local_resources=cpu=4
--local_resources=memory=10000
--jobs=4

For every flag in the selected CPU profile, pass all four forms:

--copt=FLAG
--cxxopt=FLAG
--host_copt=FLAG
--host_cxxopt=FLAG

Also pass -mcmodel=medium in all four forms. Put the target-matched static runtime first in the search path and link it for both target and host:

--linkopt=-L/path/to/target-libstdcxx/lib
--host_linkopt=-L/path/to/target-libstdcxx/lib
--linkopt=/path/to/target-libstdcxx/lib/libstdc++.a
--host_linkopt=/path/to/target-libstdcxx/lib/libstdc++.a
--linkopt=-static-libstdc++
--host_linkopt=-static-libstdc++

Run an analysis-only pass first by adding --nobuild, then compile the wheel:

bazel --output_base="$HOME/.cache/bazel/tf-2.21.0-TARGET" \
  build //tensorflow/tools/pip_package:wheel \
  "${bazel_args[@]}" --nobuild

bazel --output_base="$HOME/.cache/bazel/tf-2.21.0-TARGET" \
  build //tensorflow/tools/pip_package:wheel \
  "${bazel_args[@]}"

Here bazel_args contains the common settings, hermetic Python settings, repository overrides, compiler environment, selected CPU flags, and link flags described above. Replace TARGET with la464-lsx, la464-lasx, or la664-lasx; never share an output base between profiles.

Validate and package the result

Test the wheel in a new CPython 3.13.14 environment. At minimum, verify:

Package TensorFlow together with the dependency wheels built under the same CPU profile. Do not combine a TensorFlow wheel from one profile with native dependency wheels from another. Record SHA-256 checksums, source revisions, third-party licenses, and the selected compiler flags with the release.

Common failure modes