Skip to content

Commit c911c9d

Browse files
committed
feat: add LLVM setup script for modular installation
Add scripts/setup_llvm.sh to support LLVM installation and configuration within containers. This script: - Downloads LLVM source from GitHub when not mounted as a volume - Installs build dependencies for LLVM compilation - Provides a clean interface for LLVM integration The script supports the INSTALL_LLVM environment variable with 'source' or 'skip' options, enabling flexible LLVM setup for different container configurations. This is part of the modular script architecture introduced in PR #115. Signed-off-by: Craig Magina <cmagina@redhat.com>
1 parent f8c375c commit c911c9d

File tree

8 files changed

+129
-14
lines changed

8 files changed

+129
-14
lines changed

.github/workflows/llvm-image.yml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,12 @@ on: # yamllint disable-line rule:truthy
66
- main
77
paths:
88
- .github/workflows/llvm-image.yml
9+
- scripts/setup_llvm.sh
910
- dockerfiles/Dockerfile.llvm
1011
pull_request:
1112
paths:
1213
- .github/workflows/llvm-image.yml
14+
- scripts/setup_llvm.sh
1315
- dockerfiles/Dockerfile.llvm
1416
schedule:
1517
- cron: '0 0 * * 0' # Runs every Sunday at midnight UTC

Makefile

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,9 @@ USERNAME ?=triton
4343
create_user ?=true
4444
# NOTE: Requires host build system to have a valid Red Hat Subscription if true
4545
INSTALL_NSIGHT ?=false
46+
llvm_path ?=
4647
user_path ?=
48+
INSTALL_LLVM ?= skip # Options: source, skip
4749
INSTALL_TRITON ?= source # Options: release, source, skip
4850
INSTALL_JUPYTER ?= true
4951
USE_CCACHE ?= 0
@@ -114,6 +116,9 @@ define run_container
114116
else \
115117
volume_arg=""; \
116118
fi; \
119+
if [ -n "$(llvm_path)" ]; then \
120+
volume_arg+=" -v $(llvm_path):/workspace/llvm-project$(SELINUXFLAG)"; \
121+
fi; \
117122
if [ -n "$(user_path)" ]; then \
118123
volume_arg+=" -v $(user_path):/workspace/user$(SELINUXFLAG)"; \
119124
fi; \
@@ -152,16 +157,19 @@ define run_container
152157
else \
153158
port_arg=""; \
154159
fi; \
160+
if [ "$(CUSTOM_LLVM)" = "false" ]; then \
161+
install_llvm="-e INSTALL_LLVM=$(INSTALL_LLVM)"; \
162+
fi; \
155163
env_vars="-e USERNAME=$(USER) -e TORCH_VERSION=$(torch_version) -e CUSTOM_LLVM=$(CUSTOM_LLVM) -e INSTALL_TOOLS=$(DEMO_TOOLS) -e INSTALL_JUPYTER=$(INSTALL_JUPYTER) -e NOTEBOOK_PORT=$(NOTEBOOK_PORT) -e INSTALL_TRITON=$(INSTALL_TRITON) -e USE_CCACHE=$(USE_CCACHE) -e MAX_JOBS=$(MAX_JOBS)"; \
156164
if [ "$(create_user)" = "true" ]; then \
157-
$(CTR_CMD) run -e CREATE_USER=$(create_user) $$env_vars $$port_arg \
165+
$(CTR_CMD) run -e CREATE_USER=$(create_user) $$env_vars $$install_llvm $$port_arg \
158166
-e USER_UID=`id -u $(USER)` -e USER_GID=`id -g $(USER)` $$gpu_args $$profiling_args $$keep_ns_arg \
159167
-ti $$volume_arg $$gitconfig_arg $(IMAGE_REPO)/$(strip $(1)):$(TRITON_TAG) bash; \
160168
elif [ "$(STRIPPED_CMD)" = "docker" ]; then \
161-
$(CTR_CMD) run --user $(shell id -u):$(shell id -g) $$env_vars $$gpu_args $$profiling_args $$port_arg \
169+
$(CTR_CMD) run --user $(shell id -u):$(shell id -g) $$env_vars $$install_llvm $$gpu_args $$profiling_args $$port_arg \
162170
-ti $$volume_arg $$gitconfig_arg $(IMAGE_REPO)/$(strip $(1)):$(TRITON_TAG) bash; \
163171
elif [ "$(STRIPPED_CMD)" = "podman" ]; then \
164-
$(CTR_CMD) run --user $(USER) $$env_vars $$keep_ns_arg $$gpu_args $$profiling_args $$port_arg \
172+
$(CTR_CMD) run --user $(USER) $$env_vars $$install_llvm $$keep_ns_arg $$gpu_args $$profiling_args $$port_arg \
165173
-ti $$volume_arg $$gitconfig_arg $(IMAGE_REPO)/$(strip $(1)):$(TRITON_TAG) bash; \
166174
fi
167175
endef

dockerfiles/Dockerfile.llvm

Lines changed: 6 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,20 +8,17 @@ ARG TRITON_CPU_BACKEND=0
88
ARG LLVM_TAG
99
USER 0
1010

11+
COPY scripts/setup_llvm.sh setup_llvm.sh
12+
1113
# Conditionally execute the build based on CUSTOM_LLVM
1214
RUN if [ "$CUSTOM_LLVM" = "true" ]; then \
1315
dnf update -y && \
1416
dnf -y install clang rpm-build git ninja-build cmake lld && \
1517
dnf clean all && rm -rf /var/cache/dnf && \
16-
python3 -m pip install --upgrade pip && \
17-
python3 -m pip install --upgrade cmake ninja sccache pybind11 && \
18-
if [ ! -d "/llvm-project" ]; then \
19-
git clone https://github.com/llvm/llvm-project /llvm-project; \
20-
else \
21-
cd /llvm-project && git fetch origin; \
22-
fi && \
23-
cd /llvm-project && \
24-
python3 -m pip install -r mlir/python/requirements.txt && \
18+
python3 -m pip install --upgrade uv && \
19+
uv pip install --upgrade cmake ninja sccache pybind11 && \
20+
./setup_llvm.sh source && \
21+
cd "${HOME}/llvm-project" && \
2522
REPO="triton"; \
2623
PROJECTS="mlir;llvm;lld"; \
2724
if [ "$TRITON_CPU_BACKEND" = "1" ]; then REPO="triton-cpu"; fi; \
@@ -63,8 +60,6 @@ RUN if [ "$CUSTOM_LLVM" = "true" ]; then \
6360
-DLLVM_INCLUDE_BENCHMARKS=OFF \
6461
-DMLIR_ENABLE_EXECUTION_ENGINE:bool=OFF && \
6562
ninja install; \
66-
else \
67-
echo "LLVM is already up-to-date. Skipping rebuild."; \
6863
fi \
6964
else \
7065
echo "Skipping LLVM build because CUSTOM_LLVM is not true"; \

dockerfiles/Dockerfile.triton

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ ENV BASH_ENV=/workspace/bin/activate \
6565
UV_HTTP_TIMEOUT=60
6666

6767
COPY --from=quay.io/triton-dev-containers/gosu /usr/local/bin/gosu /usr/local/bin/gosu
68+
COPY scripts/setup_llvm.sh setup_llvm.sh
6869
COPY scripts/setup_triton.sh setup_triton.sh
6970
COPY scripts/setup_user.sh setup_user.sh
7071
COPY scripts/setup.sh setup.sh

dockerfiles/Dockerfile.triton-amd

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ ENV BASH_ENV=/workspace/bin/activate \
8383
WORKDIR /workspace
8484

8585
COPY --from=quay.io/triton-dev-containers/gosu /usr/local/bin/gosu /usr/local/bin/gosu
86+
COPY scripts/setup_llvm.sh setup_llvm.sh
8687
COPY scripts/setup_triton.sh setup_triton.sh
8788
COPY scripts/setup_user.sh setup_user.sh
8889
COPY scripts/setup.sh setup.sh

dockerfiles/Dockerfile.triton-cpu

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ ENV BASH_ENV=/workspace/bin/activate \
6262
WORKDIR /workspace
6363

6464
COPY --from=quay.io/triton-dev-containers/gosu /usr/local/bin/gosu /usr/local/bin/gosu
65+
COPY scripts/setup_llvm.sh setup_llvm.sh
6566
COPY scripts/setup_triton.sh setup_triton.sh
6667
COPY scripts/setup_user.sh setup_user.sh
6768
COPY scripts/setup.sh setup.sh

scripts/setup.sh

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ declare -a SAVE_VARS=(
2525
"DISPLAY"
2626
"HOME"
2727
"INSTALL_JUPYTER"
28+
"INSTALL_LLVM"
2829
"INSTALL_TOOLS"
2930
"INSTALL_TRITON"
3031
"MAX_JOBS"
@@ -62,6 +63,10 @@ fi
6263

6364
run_as_user ./install_software.sh
6465

66+
if [ "${INSTALL_LLVM:-skip}" != "skip" ]; then
67+
run_as_user ./setup_llvm.sh "$INSTALL_LLVM"
68+
fi
69+
6570
if [ "${INSTALL_TRITON:-skip}" != "skip" ]; then
6671
run_as_user ./setup_triton.sh "$INSTALL_TRITON"
6772
fi

scripts/setup_llvm.sh

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
#! /bin/bash -e
2+
3+
trap "echo -e '\nScript interrupted. Exiting gracefully.'; exit 1" SIGINT
4+
5+
# Copyright (C) 2024-2025 Red Hat, Inc.
6+
#
7+
# Licensed under the Apache License, Version 2.0 (the "License");
8+
# you may not use this file except in compliance with the License.
9+
# You may obtain a copy of the License at
10+
#
11+
# http://www.apache.org/licenses/LICENSE-2.0
12+
#
13+
# Unless required by applicable law or agreed to in writing, software
14+
# distributed under the License is distributed on an "AS IS" BASIS,
15+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
16+
# See the License for the specific language governing permissions and
17+
# limitations under the License.
18+
#
19+
# SPDX-License-Identifier: Apache-2.0
20+
set -euo pipefail
21+
22+
WORKSPACE=${WORKSPACE:-${HOME}}
23+
24+
LLVM_DIR=${WORKSPACE}/llvm-project
25+
LLVM_REPO=https://github.com/llvm/llvm-project.git
26+
LLVM_BUILD_PATH=$LLVM_DIR/build
27+
28+
pip_install() {
29+
if command -v uv &>/dev/null; then
30+
uv pip install "$@"
31+
else
32+
pip install "$@"
33+
fi
34+
}
35+
36+
setup_src() {
37+
echo "Downloading LLVM source code and setting up the environment for building from source..."
38+
39+
if [ ! -d "$LLVM_DIR" ]; then
40+
echo "Cloning the LLVM Project repo $LLVM_REPO to $LLVM_DIR ..."
41+
git clone "$LLVM_REPO" "$LLVM_DIR"
42+
if [ ! -d "$LLVM_DIR" ]; then
43+
echo "$LLVM_DIR not found. ERROR Cloning repository..."
44+
exit 1
45+
else
46+
pushd "$LLVM_DIR" 1>/dev/null || exit 1
47+
git fetch origin
48+
49+
# shellcheck source=/dev/null
50+
[ -f "${HOME}"/.bashrc ] && source "${HOME}"/.bashrc
51+
52+
if [ -n "${LLVM_GITREF:-}" ]; then
53+
git checkout "$LLVM_GITREF"
54+
fi
55+
popd 1>/dev/null
56+
fi
57+
else
58+
echo "LLVM repo already present, not cloning ..."
59+
fi
60+
61+
echo "Adding LLVM_BUILD_PATH to ${HOME}/.bashrc ..."
62+
echo "export LLVM_BUILD_PATH=$LLVM_BUILD_PATH" >>"${HOME}/.bashrc"
63+
echo "Run 'source ${HOME}/.bashrc' to update the current shell"
64+
}
65+
66+
install_build_deps() {
67+
pushd "$LLVM_DIR" 1>/dev/null || exit 1
68+
if [ -f mlir/python/requirements.txt ]; then
69+
echo "Installing LLVM build dependencies ..."
70+
pip_install -r mlir/python/requirements.txt
71+
fi
72+
popd 1>/dev/null
73+
}
74+
75+
usage() {
76+
cat >&2 <<EOF
77+
Usage: $(basename "$0") [COMMAND]
78+
source Download LLVM's source (if needed) and install the build deps
79+
EOF
80+
}
81+
82+
##
83+
## Main
84+
##
85+
86+
if [ $# -ne 1 ]; then
87+
usage
88+
exit 1
89+
fi
90+
91+
COMMAND=${1,,}
92+
93+
case $COMMAND in
94+
source)
95+
setup_src
96+
install_build_deps
97+
;;
98+
*)
99+
usage
100+
exit 1
101+
;;
102+
esac

0 commit comments

Comments
 (0)