# Dockerfile.build-paddle-cpu - Build PaddlePaddle CPU wheel for ARM64
#
# Required because PyPI wheels don't work on ARM64 (x86 SSE instructions).
#
# Build time: ~1-2 hours
# Output: /output/paddlepaddle-*.whl
#
# Usage:
#   docker build -t paddle-builder:cpu-arm64 -f Dockerfile.build-paddle-cpu .
#   docker run --rm -v ./wheels:/wheels paddle-builder:cpu-arm64

# syntax=docker/dockerfile:1.4
FROM ubuntu:22.04

LABEL maintainer="Sergio Jimenez"
LABEL description="PaddlePaddle CPU wheel builder for ARM64"

ARG PADDLE_VERSION=v3.0.0
ARG PYTHON_VERSION=3.11

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1
ENV CCACHE_DIR=/ccache
ENV PATH="/usr/lib/ccache:${PATH}"

# Install build dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    python${PYTHON_VERSION} \
    python${PYTHON_VERSION}-dev \
    python${PYTHON_VERSION}-venv \
    python3-pip \
    build-essential \
    cmake \
    ninja-build \
    git \
    wget \
    curl \
    pkg-config \
    ccache \
    libssl-dev \
    libffi-dev \
    zlib1g-dev \
    libbz2-dev \
    libreadline-dev \
    libsqlite3-dev \
    liblzma-dev \
    libncurses5-dev \
    libncursesw5-dev \
    libgflags-dev \
    libgoogle-glog-dev \
    libprotobuf-dev \
    protobuf-compiler \
    patchelf \
    libopenblas-dev \
    liblapack-dev \
    swig \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sf /usr/bin/python${PYTHON_VERSION} /usr/bin/python \
    && ln -sf /usr/bin/python${PYTHON_VERSION} /usr/bin/python3

# Setup ccache
RUN mkdir -p /usr/lib/ccache && \
    ln -sf /usr/bin/ccache /usr/lib/ccache/gcc && \
    ln -sf /usr/bin/ccache /usr/lib/ccache/g++ && \
    ln -sf /usr/bin/ccache /usr/lib/ccache/cc && \
    ln -sf /usr/bin/ccache /usr/lib/ccache/c++

RUN python -m pip install --upgrade pip setuptools wheel && \
    python -m pip install numpy protobuf pyyaml requests packaging astor decorator paddle-bfloat opt-einsum

WORKDIR /build
RUN git clone --depth 1 --branch ${PADDLE_VERSION} https://github.com/PaddlePaddle/Paddle.git

WORKDIR /build/Paddle

# Patch -m64 flag (x86_64 specific)
RUN sed -i 's/-m64//g' cmake/flags.cmake && \
    sed -i 's/-m64//g' CMakeLists.txt 2>/dev/null || true && \
    find . -name "*.cmake" -exec sed -i 's/-m64//g' {} \; 2>/dev/null || true

# Install sse2neon for x86 SSE -> ARM NEON translation
RUN git clone --depth 1 https://github.com/DLTcollab/sse2neon.git /tmp/sse2neon && \
    mkdir -p /usr/local/include/sse2neon && \
    cp /tmp/sse2neon/sse2neon.h /usr/local/include/sse2neon/ && \
    rm -rf /tmp/sse2neon

# Create x86 intrinsic wrapper headers
RUN mkdir -p /usr/local/include/x86_stubs && \
    for h in immintrin xmmintrin emmintrin pmmintrin smmintrin; do \
        echo "#ifndef __x86_64__" > /usr/local/include/x86_stubs/${h}.h && \
        echo "#include <sse2neon/sse2neon.h>" >> /usr/local/include/x86_stubs/${h}.h && \
        echo "#else" >> /usr/local/include/x86_stubs/${h}.h && \
        echo "#include_next <${h}.h>" >> /usr/local/include/x86_stubs/${h}.h && \
        echo "#endif" >> /usr/local/include/x86_stubs/${h}.h; \
    done

RUN pip install -r python/requirements.txt || true

RUN mkdir -p build
WORKDIR /build/Paddle/build

# Configure for CPU-only ARM64 build
# WITH_ARM=ON enables ARM NEON optimizations and disables x86-specific code (XBYAK, MKL)
RUN cmake .. \
    -GNinja \
    -DCMAKE_BUILD_TYPE=Release \
    -DPY_VERSION=${PYTHON_VERSION} \
    -DWITH_GPU=OFF \
    -DWITH_ARM=ON \
    -DWITH_TESTING=OFF \
    -DWITH_DISTRIBUTE=OFF \
    -DWITH_NCCL=OFF \
    -DWITH_MKL=OFF \
    -DWITH_MKLDNN=OFF \
    -DWITH_XBYAK=OFF \
    -DON_INFER=OFF \
    -DWITH_PYTHON=ON \
    -DWITH_AVX=OFF \
    -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
    -DCMAKE_C_COMPILER_LAUNCHER=ccache \
    -DCMAKE_CXX_FLAGS="-Wno-class-memaccess -Wno-error=class-memaccess -I/usr/local/include/x86_stubs"

# Build external dependencies
RUN --mount=type=cache,target=/ccache \
    ninja extern_gflags extern_glog extern_protobuf extern_zlib extern_eigen3

# Note: extern_xbyak excluded - it's x86-only and disabled with WITH_ARM=ON
RUN --mount=type=cache,target=/ccache \
    ninja extern_openblas extern_pybind extern_utf8proc extern_xxhash extern_yaml extern_cryptopp extern_warpctc extern_warprnnt extern_gloo

# Build PaddlePaddle
RUN --mount=type=cache,target=/ccache \
    ninja -j$(nproc) || ninja -j$(($(nproc)/2)) || ninja -j4

RUN ninja paddle_python || true

RUN mkdir -p /output

WORKDIR /build/Paddle
RUN if [ -f build/python/setup.py ]; then \
        cd build/python && python setup.py bdist_wheel; \
    elif [ -f python/setup.py ]; then \
        cd python && python setup.py bdist_wheel; \
    fi

RUN find /build -name "paddlepaddle*.whl" -type f -exec cp {} /output/ \; && \
    ls -la /output/

CMD ["sh", "-c", "cp /output/*.whl /wheels/ && ls -la /wheels/"]