Support additional architectures beyond x86_64 in ubuntu_install_java (#3546)

* Support additional architectures beyond x86_64 in ubuntu_install_java

While attempting to get a development environment going for TVM
on my AArch64 desktop I ran into some hardcoding of relevant architectures.

diff --git a/docker/install/ubuntu_install_java.sh b/docker/install/ubuntu_install_java.sh
index aefe2e2..7e8ee39 100755 (executable)
--- a/docker/install/ubuntu_install_java.sh
+++ b/docker/install/ubuntu_install_java.sh
@@ -20,5 +20,23 @@ set -o errexit -o nounset
 set -o pipefail
 
 apt-get update && apt-get install -y openjdk-8-jdk maven
-test -d "/usr/lib/jvm/java-8-openjdk-amd64/jre"
-echo "export JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64/jre" >> /etc/profile
+arch=$(uname -m)
+jre_arch="unknown"
+case $arch in
+    'x86_64')
+       jre_arch="amd64"
+       ;;
+    'aarch64')
+       jre_arch="arm64"
+       ;;
+    default)
+       echo "Unknown architecture $arch" >&2
+       exit 1
+        ;;
+esac
+
+if [ ! -d "/usr/lib/jvm/java-8-openjdk-$jre_arch/jre" ]; then
+  echo "error: missing openjdk for $jre_arch" >&2
+  exit 1
+fi
+echo "export JAVA_HOME=/usr/lib/jvm/java-8-openjdk-$jre_arch/jre" >> /etc/profile