[AMD] Support ds3.2 on gfx942 platform (#17504)
Co-authored-by: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
This commit is contained in:
wufann
+76
-56
@@ -69,6 +69,10 @@ ARG TILELANG_REPO="https://github.com/HaiShaw/tilelang.git"
|
||||
ARG TILELANG_BRANCH="dsv32-mi35x"
|
||||
ARG TILELANG_COMMIT="ae938cf885743f165a19656d1122ad42bb0e30b8"
|
||||
|
||||
ARG TILELANG_GFX942_REPO="https://github.com/tile-ai/tilelang.git"
|
||||
ARG TILELANG_GFX942_BRANCH="main"
|
||||
ARG TILELANG_GFX942_COMMIT="2d8d3676eda18bd3d8e6fa783399ff96d3cd4ded"
|
||||
|
||||
ARG FHT_REPO="https://github.com/jeffdaily/fast-hadamard-transform.git"
|
||||
ARG FHT_BRANCH="rocm"
|
||||
ARG FHT_COMMIT="46efb7d776d38638fc39f3c803eaee3dd7016bd1"
|
||||
@@ -215,66 +219,82 @@ ENV LIBGL_ALWAYS_INDIRECT=1
|
||||
RUN echo "LC_ALL=en_US.UTF-8" >> /etc/environment
|
||||
|
||||
RUN /bin/bash -lc 'set -euo pipefail; \
|
||||
# Build TileLang only for gfx950
|
||||
if [ "${GPU_ARCH:-}" != "gfx950" ]; then \
|
||||
# Build TileLang for gfx950 and gfx942-rocm700
|
||||
if [ "${GPU_ARCH:-}" != "gfx950" ] && [ "${GPU_ARCH:-}" != "gfx942-rocm700" ]; then \
|
||||
echo "[TileLang] Skipping (GPU_ARCH=${GPU_ARCH:-unset})"; \
|
||||
exit 0; \
|
||||
fi; \
|
||||
echo "[TileLang] Building TileLang for ${GPU_ARCH}"; \
|
||||
\
|
||||
# System dependencies (NO llvm-dev to avoid llvm-config-16 shadowing)
|
||||
apt-get update && apt-get install -y --no-install-recommends \
|
||||
build-essential git wget curl ca-certificates gnupg \
|
||||
libgtest-dev libgmock-dev \
|
||||
libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev \
|
||||
python3 python3-dev python3-setuptools python3-pip \
|
||||
gcc libtinfo-dev zlib1g-dev libedit-dev libxml2-dev \
|
||||
cmake ninja-build pkg-config libstdc++6 \
|
||||
&& rm -rf /var/lib/apt/lists/*; \
|
||||
\
|
||||
# Build GoogleTest static libs (Ubuntu package ships sources only)
|
||||
cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \
|
||||
cmake --build /tmp/build-gtest -j"$(nproc)" && \
|
||||
cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \
|
||||
rm -rf /tmp/build-gtest; \
|
||||
\
|
||||
# Keep setuptools < 80 (compat with base image)
|
||||
python3 -m pip install --upgrade "setuptools>=77.0.3,<80" wheel cmake ninja && \
|
||||
python3 -m pip cache purge || true; \
|
||||
\
|
||||
# Locate ROCm llvm-config; fallback to installing LLVM 18 if missing
|
||||
LLVM_CONFIG_PATH=""; \
|
||||
for p in /opt/rocm/llvm/bin/llvm-config /opt/rocm/llvm-*/bin/llvm-config /opt/rocm-*/llvm*/bin/llvm-config; do \
|
||||
if [ -x "$p" ]; then LLVM_CONFIG_PATH="$p"; break; fi; \
|
||||
done; \
|
||||
if [ -z "$LLVM_CONFIG_PATH" ]; then \
|
||||
echo "[TileLang] ROCm llvm-config not found; installing LLVM 18..."; \
|
||||
curl -fsSL https://apt.llvm.org/llvm.sh -o /tmp/llvm.sh; \
|
||||
chmod +x /tmp/llvm.sh; \
|
||||
/tmp/llvm.sh 18; \
|
||||
LLVM_CONFIG_PATH="$(command -v llvm-config-18)"; \
|
||||
if [ -z "$LLVM_CONFIG_PATH" ]; then echo "ERROR: llvm-config-18 not found after install"; exit 1; fi; \
|
||||
fi; \
|
||||
echo "[TileLang] Using LLVM_CONFIG at: $LLVM_CONFIG_PATH"; \
|
||||
export PATH="$(dirname "$LLVM_CONFIG_PATH"):/usr/local/bin:${PATH}"; \
|
||||
export LLVM_CONFIG="$LLVM_CONFIG_PATH"; \
|
||||
\
|
||||
# Optional shim for tools that expect llvm-config-16
|
||||
mkdir -p /usr/local/bin && \
|
||||
printf "#!/usr/bin/env bash\nexec \"%s\" \"\$@\"\n" "$LLVM_CONFIG_PATH" > /usr/local/bin/llvm-config-16 && \
|
||||
chmod +x /usr/local/bin/llvm-config-16; \
|
||||
\
|
||||
# TVM Python bits need Cython
|
||||
python3 -m pip install --no-cache-dir "cython>=0.29.36,<3.0"; \
|
||||
\
|
||||
# Clone + pin TileLang (bundled TVM), then build
|
||||
git clone --recursive --branch "${TILELANG_BRANCH}" "${TILELANG_REPO}" /opt/tilelang && \
|
||||
cd /opt/tilelang && \
|
||||
git fetch --depth=1 origin "${TILELANG_COMMIT}" || true && \
|
||||
git checkout -f "${TILELANG_COMMIT}" && \
|
||||
git submodule update --init --recursive && \
|
||||
export CMAKE_ARGS="-DLLVM_CONFIG=${LLVM_CONFIG} ${CMAKE_ARGS:-}" && \
|
||||
bash ./install_rocm.sh'
|
||||
if [ "$GPU_ARCH" = "gfx950" ]; then \
|
||||
\
|
||||
# System dependencies (NO llvm-dev to avoid llvm-config-16 shadowing)
|
||||
apt-get update && apt-get install -y --no-install-recommends \
|
||||
build-essential git wget curl ca-certificates gnupg \
|
||||
libgtest-dev libgmock-dev \
|
||||
libprotobuf-dev protobuf-compiler libgflags-dev libsqlite3-dev \
|
||||
python3 python3-dev python3-setuptools python3-pip \
|
||||
gcc libtinfo-dev zlib1g-dev libedit-dev libxml2-dev \
|
||||
cmake ninja-build pkg-config libstdc++6 \
|
||||
&& rm -rf /var/lib/apt/lists/*; \
|
||||
\
|
||||
# Build GoogleTest static libs (Ubuntu package ships sources only)
|
||||
cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \
|
||||
cmake --build /tmp/build-gtest -j"$(nproc)" && \
|
||||
cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \
|
||||
rm -rf /tmp/build-gtest; \
|
||||
\
|
||||
# Keep setuptools < 80 (compat with base image)
|
||||
python3 -m pip install --upgrade "setuptools>=77.0.3,<80" wheel cmake ninja && \
|
||||
python3 -m pip cache purge || true; \
|
||||
\
|
||||
# Locate ROCm llvm-config; fallback to installing LLVM 18 if missing
|
||||
LLVM_CONFIG_PATH=""; \
|
||||
for p in /opt/rocm/llvm/bin/llvm-config /opt/rocm/llvm-*/bin/llvm-config /opt/rocm-*/llvm*/bin/llvm-config; do \
|
||||
if [ -x "$p" ]; then LLVM_CONFIG_PATH="$p"; break; fi; \
|
||||
done; \
|
||||
if [ -z "$LLVM_CONFIG_PATH" ]; then \
|
||||
echo "[TileLang] ROCm llvm-config not found; installing LLVM 18..."; \
|
||||
curl -fsSL https://apt.llvm.org/llvm.sh -o /tmp/llvm.sh; \
|
||||
chmod +x /tmp/llvm.sh; \
|
||||
/tmp/llvm.sh 18; \
|
||||
LLVM_CONFIG_PATH="$(command -v llvm-config-18)"; \
|
||||
if [ -z "$LLVM_CONFIG_PATH" ]; then echo "ERROR: llvm-config-18 not found after install"; exit 1; fi; \
|
||||
fi; \
|
||||
echo "[TileLang] Using LLVM_CONFIG at: $LLVM_CONFIG_PATH"; \
|
||||
export PATH="$(dirname "$LLVM_CONFIG_PATH"):/usr/local/bin:${PATH}"; \
|
||||
export LLVM_CONFIG="$LLVM_CONFIG_PATH"; \
|
||||
\
|
||||
# Optional shim for tools that expect llvm-config-16
|
||||
mkdir -p /usr/local/bin && \
|
||||
printf "#!/usr/bin/env bash\nexec \"%s\" \"\$@\"\n" "$LLVM_CONFIG_PATH" > /usr/local/bin/llvm-config-16 && \
|
||||
chmod +x /usr/local/bin/llvm-config-16; \
|
||||
\
|
||||
# TVM Python bits need Cython
|
||||
python3 -m pip install --no-cache-dir "cython>=0.29.36,<3.0"; \
|
||||
\
|
||||
# Clone + pin TileLang (bundled TVM), then build
|
||||
git clone --recursive --branch "${TILELANG_BRANCH}" "${TILELANG_REPO}" /opt/tilelang && \
|
||||
cd /opt/tilelang && \
|
||||
git fetch --depth=1 origin "${TILELANG_COMMIT}" || true && \
|
||||
git checkout -f "${TILELANG_COMMIT}" && \
|
||||
git submodule update --init --recursive && \
|
||||
export CMAKE_ARGS="-DLLVM_CONFIG=${LLVM_CONFIG} ${CMAKE_ARGS:-}" && \
|
||||
bash ./install_rocm.sh; \
|
||||
else \
|
||||
# Build GoogleTest static libs (Ubuntu package ships sources only)
|
||||
apt-get install -y libgtest-dev libgmock-dev && \
|
||||
cmake -S /usr/src/googletest -B /tmp/build-gtest -DBUILD_GTEST=ON -DBUILD_GMOCK=ON -DCMAKE_BUILD_TYPE=Release && \
|
||||
cmake --build /tmp/build-gtest -j && \
|
||||
cp -v /tmp/build-gtest/lib/*.a /usr/lib/x86_64-linux-gnu/ && \
|
||||
rm -rf /tmp/build-gtest; \
|
||||
# Build TileLang for gfx942-rocm700
|
||||
git clone --branch "${TILELANG_GFX942_BRANCH}" "${TILELANG_GFX942_REPO}" /opt/tilelang && \
|
||||
cd /opt/tilelang && \
|
||||
git checkout -f "${TILELANG_GFX942_COMMIT}" && \
|
||||
git submodule update --init --recursive && \
|
||||
sed -i "/^[[:space:]]*\"torch/d" pyproject.toml && \
|
||||
USE_ROCM=1 USE_CUDA=0 pip install -e . -v ; \
|
||||
fi'
|
||||
|
||||
# -----------------------
|
||||
# Hadamard-transform (HIP build)
|
||||
|
||||
Reference in New Issue
Block a user