From bde7986d08ab6f58de02ad02059a849fcae513d9 Mon Sep 17 00:00:00 2001
From: Xu Zhao <i@xuzhao.net>
Date: Thu, 21 Nov 2024 12:22:13 -0800
Subject: [PATCH] Install patch in the docker (#60)

Summary:
Now we have to patch third-party submodules like hstu and xformers, install the patch utility to the CI docker.

Pull Request resolved: https://github.com/pytorch-labs/tritonbench/pull/60

Test Plan:
CI

https://github.com/pytorch-labs/tritonbench/actions/runs/11946317137

Reviewed By: FindHao

Differential Revision: D66273502

Pulled By: xuzhao9

fbshipit-source-id: 78552c53b7d0de008081c4fc94851308bff4350a
---
 .ci/tritonbench/test-gpu.sh           | 7 -------
 docker/tritonbench-nightly.dockerfile | 2 +-
 tools/cuda_utils.py                   | 6 ++++++
 3 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/.ci/tritonbench/test-gpu.sh b/.ci/tritonbench/test-gpu.sh
index b41331aa..f05a965f 100644
--- a/.ci/tritonbench/test-gpu.sh
+++ b/.ci/tritonbench/test-gpu.sh
@@ -8,11 +8,4 @@ fi
 
 . "${SETUP_SCRIPT}"
 
-# FIXME: patch and install hstu
-sudo apt-get install  -y patch
-python install.py --hstu
-
-# FIXME: install colfax
-python install.py --colfax
-
 python -m unittest test.test_gpu.main -v
diff --git a/docker/tritonbench-nightly.dockerfile b/docker/tritonbench-nightly.dockerfile
index 148df0cc..a2d09fde 100644
--- a/docker/tritonbench-nightly.dockerfile
+++ b/docker/tritonbench-nightly.dockerfile
@@ -44,7 +44,7 @@ RUN cd /workspace/tritonbench && \
 
 # Tritonbench library build and test require libcuda.so.1
 # which is from NVIDIA driver
-RUN sudo apt update && sudo apt-get install -y libnvidia-compute-550 patchelf
+RUN sudo apt update && sudo apt-get install -y libnvidia-compute-550 patchelf patch
 
 # Install Tritonbench
 RUN cd /workspace/tritonbench && \
diff --git a/tools/cuda_utils.py b/tools/cuda_utils.py
index 91bcbbbe..ba46be49 100644
--- a/tools/cuda_utils.py
+++ b/tools/cuda_utils.py
@@ -129,6 +129,12 @@ def install_torch_deps(cuda_version: str):
     ]
     cmd = ["conda", "install", "-y"] + torch_deps
     subprocess.check_call(cmd)
+    # conda forge deps
+    # ubuntu 22.04 comes with libstdcxx6 12.3.0
+    # we need to install the same library version in conda
+    conda_deps = ["libstdcxx-ng=12.3.0"]
+    cmd = ["conda", "install", "-y", "-c", "conda-forge"] + conda_deps
+    subprocess.check_call(cmd)
 
 
 def install_torch_build_deps(cuda_version: str):