From 01a26544a3547c5d454581f1890d7a9805f388a9 Mon Sep 17 00:00:00 2001
From: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Date: Tue, 30 Sep 2025 20:00:37 -0700
Subject: [PATCH] [AMD] Add Tilelang and Fast Hadamard Transform builds to
 Dockerfile.rocm (#11114)

---
 docker/Dockerfile.rocm | 93 ++++++++++++++++++++++++++++++++++++++++--
 1 file changed, 90 insertions(+), 3 deletions(-)

diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm
index d4378236a..001c2ccd3 100644
--- a/docker/Dockerfile.rocm
+++ b/docker/Dockerfile.rocm
@@ -69,6 +69,13 @@ ARG LLVM_COMMIT="6520ace8227ffe2728148d5f3b9872a870b0a560"
 ARG MOONCAKE_REPO="https://github.com/kvcache-ai/Mooncake.git"
 ARG MOONCAKE_COMMIT="dcdf1c784b40aa6975a8ed89fe26321b028e40e8"
 
+ARG TILELANG_REPO="https://github.com/HaiShaw/tilelang.git"
+ARG TILELANG_BRANCH="dsv32-mi35x"
+ARG TILELANG_COMMIT="ae938cf885743f165a19656d1122ad42bb0e30b8"
+
+ARG FHT_REPO="https://github.com/jeffdaily/fast-hadamard-transform.git"
+ARG FHT_BRANCH="rocm"
+ARG FHT_COMMIT="46efb7d776d38638fc39f3c803eaee3dd7016bd1"
 USER root
 
 # Install some basic utilities
@@ -90,8 +97,6 @@ RUN if [ "$BUILD_LLVM" = "1" ]; then \
      && make -j$(nproc); \
     fi
 
-# -----------------------
-
 # -----------------------
 # AITER
 RUN pip uninstall -y aiter
@@ -155,7 +160,6 @@ RUN if [ "$BUILD_MOONCAKE" = "1" ]; then \
      make -j "$(nproc)" && make install; \
     fi
 
-
 # -----------------------
 # Build SGLang
 ARG BUILD_TYPE=all
@@ -207,6 +211,89 @@ RUN python3 -m pip install --no-cache-dir setuptools-rust \
     && python3 -m pip install --no-cache-dir . \
     && rm -rf /root/.cache
 
+# -----------------------
+# TileLang
+ENV DEBIAN_FRONTEND=noninteractive
+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 \
+    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'
+
+# -----------------------
+# Hadamard-transform (HIP build)
+RUN /bin/bash -lc 'set -euo pipefail; \
+    git clone --branch "${FHT_BRANCH}" "${FHT_REPO}" fast-hadamard-transform; \
+    cd fast-hadamard-transform; \
+    git checkout -f "${FHT_COMMIT}"; \
+    python setup.py install'
+
+# -----------------------
+# Python tools
+RUN python3 -m pip install --no-cache-dir \
+    py-spy \
+    pre-commit
+
+# -----------------------
 # Performance environment variable.
 ENV HIP_FORCE_DEV_KERNARG=1
 ENV HSA_NO_SCRATCH_RECLAIM=1