Install all tools required to run MLIR code

setup-mlir.sh

#!/usr/bin/env bash
# Run this script using the command "bash setup-mlir.sh"
# One-shot setup for the MLIR getting-started environment on Debian/Ubuntu-like systems.

set -euo pipefail

# -------------
# Configuration
# -------------
GETTING_STARTED_ROOT="${HOME}/mlir-getting-started"
LLVM_ZIP_URL="https://github.com/llvm/llvm-project/archive/refs/tags/llvmorg-19.1.3.zip"
LLVM_ZIP_NAME="llvmorg-19.1.3.zip"
LLVM_SRC_DIR="${GETTING_STARTED_ROOT}/llvm-project-llvmorg-19.1.3"
LLVM_BUILD_DIR="${LLVM_SRC_DIR}/build"

# Frequently used executables symlinks
MLIR_OPT_LINK="${GETTING_STARTED_ROOT}/mlir-opt"
MLIR_CPU_RUNNER_LINK="${GETTING_STARTED_ROOT}/mlir-cpu-runner"

# Packages from instructions
PKGS=("cmake" "ninja-build" "python3.12-dev" "python3-numpy" "pybind11-dev")

# -----------------------------
# Helper functions
# -----------------------------

info()  { printf "\033[1;34m[INFO]\033[0m %s\n" "$*"; }
warn()  { printf "\033[1;33m[WARN]\033[0m %s\n" "$*"; }
error() { printf "\033[1;31m[ERROR]\033[0m %s\n" "$*"; }

# -----------------------------
# Create getting-started directory
# -----------------------------
info "Creating getting-started directory at: ${GETTING_STARTED_ROOT}"
mkdir -p "${GETTING_STARTED_ROOT}"
cd "${GETTING_STARTED_ROOT}"

# -----------------------------
# Install required packages
# -----------------------------
info "Updating apt package lists..."
sudo apt-get update -y

info "Installing required packages: ${PKGS[*]}"
# We install exactly the packages listed in the instructions
sudo apt-get install -y "${PKGS[@]}"

# -----------------------------
# Download and unpack LLVM 19.1.3
# -----------------------------
if [[ -d "${LLVM_SRC_DIR}" ]]; then
  info "LLVM source directory already exists: ${LLVM_SRC_DIR} (skipping download/unpack)"
else
  info "Downloading LLVM ${LLVM_ZIP_NAME} from ${LLVM_ZIP_URL}"
  # Remove any previous partial zip to avoid unzip confusion
  [[ -f "${LLVM_ZIP_NAME}" ]] && rm -f "${LLVM_ZIP_NAME}"
  wget "${LLVM_ZIP_URL}" -O "${LLVM_ZIP_NAME}"

  info "Unpacking ${LLVM_ZIP_NAME}"
  unzip "${LLVM_ZIP_NAME}"

  info "Unpacked to: ${LLVM_SRC_DIR}"
fi

# -------------------------------------
# Build LLVM/MLIR with Python bindings
# -------------------------------------
info "Configuring LLVM build directory: ${LLVM_BUILD_DIR}"
mkdir -p "${LLVM_BUILD_DIR}"
cd "${LLVM_BUILD_DIR}"

# Configure with CMake (Ninja) per instructions
info "Running CMake configure for MLIR (Release, assertions ON, Python bindings ON)"
cmake -G Ninja ../llvm \
  -DLLVM_ENABLE_PROJECTS="mlir" \
  -DLLVM_BUILD_EXAMPLES=OFF \
  -DLLVM_TARGETS_TO_BUILD="Native;NVPTX;AMDGPU" \
  -DCMAKE_BUILD_TYPE=Release \
  -DLLVM_ENABLE_ASSERTIONS=ON \
  -DMLIR_ENABLE_BINDINGS_PYTHON=ON \
  -DCMAKE_C_FLAGS="-O0" \
  -DCMAKE_CXX_FLAGS="-O0"

info "Building LLVM/MLIR (this may take a while)..."
cmake --build .

# -----------------------------
# Create symbolic links to frequently used executables
# -----------------------------
cd "${GETTING_STARTED_ROOT}"

MLIR_OPT_BIN="${LLVM_BUILD_DIR}/bin/mlir-opt"
MLIR_CPU_RUNNER_BIN="${LLVM_BUILD_DIR}/bin/mlir-cpu-runner"

if [[ -x "${MLIR_OPT_BIN}" ]]; then
  info "Creating symlink: ${MLIR_OPT_LINK} -> ${MLIR_OPT_BIN}"
  ln -sf "${MLIR_OPT_BIN}" "${MLIR_OPT_LINK}"
else
  warn "mlir-opt not found at ${MLIR_OPT_BIN}"
fi

if [[ -x "${MLIR_CPU_RUNNER_BIN}" ]]; then
  info "Creating symlink: ${MLIR_CPU_RUNNER_LINK} -> ${MLIR_CPU_RUNNER_BIN}"
  ln -sf "${MLIR_CPU_RUNNER_BIN}" "${MLIR_CPU_RUNNER_LINK}"
else
  warn "mlir-cpu-runner not found at ${MLIR_CPU_RUNNER_BIN}"
fi

# -----------------------------
# Create run-mlir helper script
# -----------------------------
RUN_MLIR_SCRIPT="${GETTING_STARTED_ROOT}/run-mlir"
info "Creating helper script: ${RUN_MLIR_SCRIPT}"

cat > "${RUN_MLIR_SCRIPT}" <<'EOF'
#!/usr/bin/env bash
set -euo pipefail

GETTING_STARTED_ROOT="$HOME/mlir-getting-started"
MLIR_OPT="$GETTING_STARTED_ROOT/mlir-opt"
MLIR_CPU_RUNNER="$GETTING_STARTED_ROOT/mlir-cpu-runner"
LLVM_ROOT="$GETTING_STARTED_ROOT/llvm-project-llvmorg-19.1.3"

# Pipeline from instructions:
# - convert elementwise to linalg
# - one-shot bufferize with function boundaries
# - convert linalg to loops
# - convert scf to cf
# - memref expand
# - expand strided metadata
# - lower affine
# - convert to LLVM
# - reconcile unrealized casts
# Then run with mlir-cpu-runner and the runner utils libs.

"$MLIR_OPT" \
  --convert-elementwise-to-linalg \
  --one-shot-bufferize="bufferize-function-boundaries" \
  --convert-linalg-to-loops \
  --convert-scf-to-cf \
  --memref-expand \
  --expand-strided-metadata \
  --lower-affine \
  --convert-to-llvm \
  --reconcile-unrealized-casts \
  "$@" \
| "$MLIR_CPU_RUNNER" \
  --shared-libs="$LLVM_ROOT/build/lib/libmlir_runner_utils.so,$LLVM_ROOT/build/lib/libmlir_c_runner_utils.so"
EOF

chmod +x "${RUN_MLIR_SCRIPT}"

# -----------------------------
# Create MLIR test file
# -----------------------------
TEST_FILE="${GETTING_STARTED_ROOT}/constant.mlir"
info "Creating simple MLIR test file: ${TEST_FILE}"

cat > "${TEST_FILE}" <<'EOF'
func.func @main() -> f32 {
  %0 = arith.constant 5.0 : f32
  return %0 : f32
}
EOF

# -----------------------------
# Run the MLIR test
# -----------------------------
info "Running MLIR test: ./run-mlir constant.mlir"
cd "${GETTING_STARTED_ROOT}"
./run-mlir constant.mlir || {
  error "MLIR test run failed."
  exit 1
}

info "Setup complete. Expected output includes: 5.000000e+00"