Addressing review comments

Change-Id: I032eedad3f7fdd0d24b8a2e0e6951c6081fed124

.github/workflows/ci.yml

@@ -15,7 +15,7 @@ name: CI
 # https://github.community/t/support-for-yaml-anchors/16128/92
 # https://github.com/actions/runner/issues/1182
 # Neither does it have any contexts that are available everywhere. The
-# top-level `env` fbuild_test_runtime_arm64:ield is available in many places, but not all. We already
+# top-level `env` field is available in many places, but not all. We already
 # have a "should-run" job that every other job depends on, so we leverage that
 # for variables that every other job can use, since that *is* available in all
 # sub-fields of the job.
@@ -366,7 +366,7 @@ jobs:
     if: contains(fromJson(needs.setup.outputs.enabled-jobs), 'build_test_runtime_arm64')
     runs-on: <TODO>
     env:
-      BUILD_DIR: build-runtime
+      BUILD_DIR: build-runtime-arm64
     steps:
       - name: "Checking out repository"
         uses: actions/checkout@ac593985615ec2ede58e132d2e21d2b1cbd6127c # v3.3.0

build_tools/docker/context/fetch_cuda_deps.sh

@@ -25,10 +25,10 @@
 #   ./fetch_cuda_deps.sh /usr/local/iree_cuda_deps
 set -e
 
-machine=$(uname -m)
-if [[ $machine == aarch64 ]]; then
-  echo "CUDA for ${machine} is not supported yet."
-  exit 0
+ARCH="$(uname -m)"
+if [[ "${ARCH}" == "aarch64" ]]; then
+  echo "ERROR: Script does not support ${ARCH}."
+  exit 1
 fi
 
 TARGET_DIR="$1"

build_tools/docker/context/install_basics.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+set -euo pipefail
+
+# Useful utilities for building child images. Best practices would tell us to
+# use multi-stage builds
+# (https://docs.docker.com/develop/develop-images/multistage-build/) but it
+# turns out that Dockerfile is a thoroughly non-composable awful format and that
+# doesn't actually work that well. These deps are pretty small.
+apt-get update && apt-get install -y git unzip wget curl gnupg2 lsb-release

build_tools/docker/context/install_bazel.sh

@@ -14,16 +14,15 @@ fi
 
 BAZEL_VERSION="$(cat .bazelversion)"
 
-# We could do the whole apt install dance, but this technique works across a
-# range of platforms, allowing us to use a single script. See
-# https://bazel.build/install/ubuntu#binary-installer
-machine=$(uname -m)
-if [ "$machine" == "aarch64" ]; then
-  machine="arm64"
+ARCH="$(uname -m)"
+if [[ "${ARCH}" == "aarch64" ]]; then
+  ARCH="arm64"
 fi
 
+# We could do the whole apt install dance, but this technique works across a
+# range of platforms.
 curl --silent --fail --show-error --location \
-  "https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION?}/bazel-${BAZEL_VERSION?}-linux-${machine}" \
+  "https://github.com/bazelbuild/bazel/releases/download/${BAZEL_VERSION?}/bazel-${BAZEL_VERSION?}-linux-${ARCH}" \
   --output /usr/bin/bazel
 chmod +x /usr/bin/bazel
 

build_tools/docker/context/install_ccache.sh

@@ -1,14 +1,15 @@
 #!/bin/bash
-
 # Copyright 2022 The IREE Authors
 #
 # Licensed under the Apache License v2.0 with LLVM Exceptions.
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+set -euo pipefail
+
 CCACHE_VERSION=4.7.4
-machine=$(uname -m)
-if [[ $machine == x86_64 ]]; then
+ARCH="$(uname -m)"
+if [[ "${ARCH}" == "x86_64" ]]; then
   curl --silent --show-error --fail --location \
       "https://github.com/ccache/ccache/releases/download/v${CCACHE_VERSION}/ccache-${CCACHE_VERSION}-linux-x86_64.tar.xz" \
       --output ccache.tar.xz \
@@ -20,12 +21,8 @@ if [[ $machine == x86_64 ]]; then
   gpg --verify ccache.tar.xz.asc
   tar -xvf ccache.tar.xz --strip-components=1
   cp ccache /usr/bin/
-elif  [[ $machine == aarch64 ]]; then
-  # Apt install for Ubuntu 20.04 gives ccache version 3.7.7, I assume you want
-  # something newer. Two alternatives from what I can find
-  #   1) archlinuxarm.org has a 4.8.2 version
-  #   2) build from source
-  # but I may be missing something...
+elif [[ "${ARCH}" == "aarch64" ]]; then
+  # Latest version of ccache is not released for arm64, use distro version
   apt-get install -y ccache
 fi
 

build_tools/docker/context/install_cmake.sh

@@ -9,10 +9,10 @@ set -euo pipefail
 
 CMAKE_VERSION="$1"
 
-machine=$(uname -m)
+ARCH="$(uname -m)"
 
 curl --silent --fail --show-error --location \
-    "https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-linux-${machine}.sh" \
+    "https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-linux-${ARCH}.sh" \
     --output cmake-installer.sh
 
 chmod +x cmake-installer.sh

build_tools/docker/context/install_rocm.sh

@@ -5,13 +5,13 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-set -e
+set -euo pipefail
 
 ROCM_VERSION=$1
 AMDGPU_VERSION=$2
 
-machine=$(uname -m)
-if [[ $machine == x86_64 ]]; then
+ARCH="$(uname -m)"
+if [[ "${ARCH}" == "x86_64" ]]; then
   apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends ca-certificates curl libnuma-dev gnupg \
     && curl -sL https://repo.radeon.com/rocm/rocm.gpg.key | apt-key add - \
     && printf "deb [arch=amd64] https://repo.radeon.com/rocm/apt/${ROCM_VERSION}/ ubuntu main" | tee /etc/apt/sources.list.d/rocm.list \
@@ -24,5 +24,5 @@ if [[ $machine == x86_64 ]]; then
     rocm-dev \
     build-essential
 else
-  echo "Installing ROCM for ${machine} is not supported yet."
+  echo "Installing ROCM for ${ARCH} is not supported yet."
 fi

build_tools/docker/context/install_vulkan.sh

@@ -5,16 +5,16 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-set -e
+set -euo pipefail
 
 VULKAN_SDK_VERSION=$1
 
-machine=$(uname -m)
-if [[ $machine == x86_64 ]]; then
+ARCH="$(uname -m)"
+if [[ "${ARCH}" == "x86_64" ]]; then
   wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | apt-key add - \
     && wget -qO /etc/apt/sources.list.d/lunarg-vulkan-${VULKAN_SDK_VERSION}-focal.list https://packages.lunarg.com/vulkan/${VULKAN_SDK_VERSION}/lunarg-vulkan-${VULKAN_SDK_VERSION}-focal.list \
     && apt-get update \
     && apt-get install -y --no-install-recommends vulkan-sdk
 else
-  echo "Installing Vulkan for ${machine} is not supported yet."
+  echo "Installing Vulkan for ${ARCH} is not supported yet."
 fi

build_tools/docker/dockerfiles/base-arm64.Dockerfile

@@ -5,7 +5,7 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # 20.04
-FROM ubuntu@sha256:7bdccf116db125b3e6e39eb67ca9e2ae890386acf95a13a4e8b69466b6eba5e2
+FROM --platform=linux/arm64 ubuntu@sha256:7bdccf116db125b3e6e39eb67ca9e2ae890386acf95a13a4e8b69466b6eba5e2
 
 SHELL ["/bin/bash", "-e", "-u", "-o", "pipefail", "-c"]
 
@@ -16,19 +16,8 @@ ARG APT_KEY_DONT_WARN_ON_DANGEROUS_USAGE=1
 
 ######## Basic stuff ########
 WORKDIR /install-basics
-# Useful utilities for building child images. Best practices would tell us to
-# use multi-stage builds
-# (https://docs.docker.com/develop/develop-images/multistage-build/) but it
-# turns out that Dockerfile is a thoroughly non-composable awful format and that
-# doesn't actually work that well. These deps are pretty small.
-RUN apt-get update \
-  && apt-get install -y \
-    git \
-    unzip \
-    wget \
-    curl \
-    gnupg2 \
-    lsb-release
+COPY build_tools/docker/context/install_basics.sh ./
+RUN ./install_basics.sh && rm -rf /install-basics
 
 # Install the oldest supported compiler tools
 ARG LLVM_VERSION=9

build_tools/docker/dockerfiles/base.Dockerfile

@@ -16,19 +16,8 @@ ARG APT_KEY_DONT_WARN_ON_DANGEROUS_USAGE=1
 
 ######## Basic stuff ########
 WORKDIR /install-basics
-# Useful utilities for building child images. Best practices would tell us to
-# use multi-stage builds
-# (https://docs.docker.com/develop/develop-images/multistage-build/) but it
-# turns out that Dockerfile is a thoroughly non-composable awful format and that
-# doesn't actually work that well. These deps are pretty small.
-RUN apt-get update \
-  && apt-get install -y \
-    git \
-    unzip \
-    wget \
-    curl \
-    gnupg2 \
-    lsb-release
+COPY build_tools/docker/context/install_basics.sh ./
+RUN ./install_basics.sh && rm -rf /install-basics
 
 # Install the oldest supported compiler tools
 ARG LLVM_VERSION=9

build_tools/docker/manage_images.py

@@ -247,11 +247,21 @@ def parse_prod_digests() -> Dict[str, str]:
             # access in builds.
             # We're assuming this is being run from the root of the repository.
             # FIXME: make this more robust to where it is run from.
-            cmd = ["docker", "buildx", "build", "--file", image_path, "--load"]
-            if image == "base-arm64":
-                cmd.extend(["--platform", "linux/arm64"])
-            cmd.extend(["--tag", tagged_image_url, "."])
-            utils.run_command(cmd, dry_run=args.dry_run)
+            utils.run_command(
+                [
+                    "docker",
+                    "buildx",
+                    "build",
+                    "--file",
+                    image_path,
+                    "--load",
+                    "--tag",
+                    tagged_image_url,
+                    ".",
+                ],
+                dry_run=args.dry_run,
+            )
+
 
             utils.run_command(
                 ["docker", "push", tagged_image_url], dry_run=args.dry_run