Merge remote-tracking branch 'upstream/main' into unity-staging

[MERGE] Merge main into unity 2023-08-01

Author: tqchen

.gitmodules

@@ -25,3 +25,6 @@
 [submodule "3rdparty/cutlass_fpA_intB_gemm"]
 	path = 3rdparty/cutlass_fpA_intB_gemm
 	url = https://github.com/tlc-pack/cutlass_fpA_intB_gemm
+[submodule "3rdparty/libflash_attn"]
+	path = 3rdparty/libflash_attn
+	url = https://github.com/tlc-pack/libflash_attn

3rdparty/libflash_attn

@@ -214,7 +214,7 @@ Apache Software Foundation License 2.0
 3rdparty/OpenCL-Headers
 3rdparty/mlperftiny
 3rdparty/nvbench (with LLVM exception)
-
+3rdparty/cutlass_fpA_intB_gemm
 
 BSD 2-clause License
 --------------------
@@ -247,3 +247,4 @@ BSD 3-Clause "New" or "Revised" License
 
 3rdparty/cutlass
 3rdparty/libbacktrace
+3rdparty/libflash_attn

LICENSE

@@ -19,7 +19,7 @@
 
 {% call m.invoke_build(
   name='BUILD: arm',
-  node='ARM-SMALL',
+  node='ARM-GRAVITON3',
   condition='!skip_ci && is_docs_only_build != 1',
   ws='tvm/build-arm',
   docker_image='ci_arm',
@@ -40,7 +40,7 @@
 {% call(shard_index, num_shards) m.sharded_test_step(
   name="integration: aarch64",
   num_shards=4,
-  node="ARM-SMALL",
+  node="ARM-GRAVITON3",
   ws="tvm/ut-python-arm",
   platform="arm",
   docker_image="ci_arm",
@@ -57,7 +57,7 @@
 
 {% call(shard_index, num_shards) m.sharded_test_step(
   name="topi: aarch64",
-  node="ARM-SMALL",
+  node="ARM-GRAVITON3",
   ws="tvm/ut-python-arm",
   platform="arm",
   docker_image="ci_arm",
@@ -82,7 +82,7 @@
 
 {% call(shard_index, num_shards) m.sharded_test_step(
   name="frontend: aarch64",
-  node="ARM-SMALL",
+  node="ARM-GRAVITON3",
   ws="tvm/frontend-python-arm",
   platform="arm",
   docker_image="ci_arm",

ci/jenkins/generated/arm_jenkinsfile.groovy

@@ -32,5 +32,7 @@ if(USE_NNPACK)
   list(APPEND TVM_RUNTIME_LINKER_LIBS ${NNPACK_CONTRIB_LIB})
   list(APPEND TVM_RUNTIME_LINKER_LIBS ${NNPACK_PTHREAD_CONTRIB_LIB})
   list(APPEND TVM_RUNTIME_LINKER_LIBS ${NNPACK_CPUINFO_CONTRIB_LIB})
-  list(APPEND TVM_RUNTIME_LINKER_LIBS ${NNPACK_CLOG_CONTRIB_LIB})
+  if(NNPACK_CLOG_CONTRIB_LIB)
+    list(APPEND TVM_RUNTIME_LINKER_LIBS ${NNPACK_CLOG_CONTRIB_LIB})
+  endif(NNPACK_CLOG_CONTRIB_LIB)
 endif(USE_NNPACK)

ci/jenkins/templates/arm_jenkinsfile.groovy.j2

@@ -43,4 +43,5 @@ pip3 install --upgrade \
     junitparser==2.4.2 \
     six \
     tornado \
-    pytest-lazy-fixture
+    pytest-lazy-fixture \
+    git+https://github.com/jax-ml/ml_dtypes.git@v0.2.0

cmake/modules/contrib/NNPack.cmake

@@ -28,7 +28,7 @@ if [ "$architecture_type" != "aarch64" ]; then
     gcc-aarch64-linux-gnu
 fi
 
-compute_lib_version="v23.02.1"
+compute_lib_version="v23.05"
 compute_lib_variant="arm64-v8a-neon"
 compute_lib_full_name="arm_compute-${compute_lib_version}-bin-linux-${compute_lib_variant}"
 compute_lib_base_url="https://github.com/ARM-software/ComputeLibrary/releases/download/${compute_lib_version}"

docker/install/ubuntu2004_install_python_package.sh

@@ -23,13 +23,13 @@ set -o pipefail
 # Install jax and jaxlib
 if [ "$1" == "cuda" ]; then
     pip3 install --upgrade \
-        jaxlib==0.4.7 \
-        "jax[cuda11_pip]==0.4.7" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+        jaxlib~=0.4.9 \
+        "jax[cuda11_pip]~=0.4.9" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 else
     pip3 install --upgrade \
-        jaxlib==0.4.7 \
-        "jax[cpu]==0.4.7"
+        jaxlib~=0.4.9 \
+        "jax[cpu]~=0.4.9"
 fi
 
 # Install flax
-pip3 install flax==0.6.8
+pip3 install flax~=0.6.9

docker/install/ubuntu_download_arm_compute_lib_binaries.sh

@@ -25,7 +25,7 @@ pip3 install --upgrade \
     "Pygments>=2.4.0" \
     attrs \
     cloudpickle \
-    cython \
+    cython==0.29.34 \
     decorator \
     mypy \
     numpy==1.21.* \