Repository: alibaba/BladeDISC (language: python)

Issue:
[bug][AArch64] cpu demo crash
We encounter following error on AArch64 platform when running cpu demo.
```
*** stack smashing detected ***: terminated
``` 

call stack:
![image](https://user-images.githubusercontent.com/15364516/163111293-af6c0018-a41a-4e6f-89b4-22c41ea50471.png)



Discussion / hints:
to #227 

Fix patch:
diff --git a/docker/dev/Dockerfile.aarch64 b/docker/dev/Dockerfile.aarch64
index b7a6e34a914..bba03ca9b20 100644
--- a/docker/dev/Dockerfile.aarch64
+++ b/docker/dev/Dockerfile.aarch64
@@ -98,8 +98,8 @@ RUN mkdir -p ~/.pip && \
 
 RUN pip3 install --upgrade pip
 
-COPY docker/scripts/install-python.sh /install/scripts/install-python.sh
-RUN bash /install/scripts/install-python.sh
+COPY docker/scripts/install-python-aarch64.sh /install/scripts/install-python-aarch64.sh
+RUN bash /install/scripts/install-python-aarch64.sh
 
 RUN chmod a+x /usr/local/bin/bazel
 
diff --git a/docker/runtime/Dockerfile.pytorch.aarch64 b/docker/runtime/Dockerfile.pytorch.aarch64
index dc40ff4a98e..da2a9d67fa3 100644
--- a/docker/runtime/Dockerfile.pytorch.aarch64
+++ b/docker/runtime/Dockerfile.pytorch.aarch64
@@ -5,4 +5,5 @@ ARG WHEEL_FILE=torch_blade*.whl
 
 ADD ./build/${WHEEL_FILE}  /install/python/
 
-RUN pip install /install/python/${WHEEL_FILE}
+RUN source /opt/venv_disc/bin/activate && \
+    pip install /install/python/${WHEEL_FILE}
diff --git a/docker/runtime/Dockerfile.tf.aarch64 b/docker/runtime/Dockerfile.tf.aarch64
index 9ca547ff7fc..515a2dd6d1a 100644
--- a/docker/runtime/Dockerfile.tf.aarch64
+++ b/docker/runtime/Dockerfile.tf.aarch64
@@ -6,4 +6,5 @@ ARG WHEEL_FILE=blade_disc*.whl
 ADD ./build/${WHEEL_FILE}  /install/python/
 ADD ./build/disc-replay-main /usr/bin/disc-replay-main
 
-RUN pip install /install/python/${WHEEL_FILE}
+RUN source /opt/venv_disc/bin/activate && \
+    pip install /install/python/${WHEEL_FILE}
diff --git a/docker/scripts/install-python-aarch64.sh b/docker/scripts/install-python-aarch64.sh
new file mode 100644
index 00000000000..a3ac948bbcb
--- /dev/null
+++ b/docker/scripts/install-python-aarch64.sh
@@ -0,0 +1,35 @@
+#!/bin/bash
+# Copyright 2021 The BladeDISC Authors. All rights reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+set -ex
+python3 -m pip install virtualenv numpy
+DISC_VENV=/opt/venv_disc
+
+function install_venv() {
+  python3 -m virtualenv ${DISC_VENV}
+  source ${DISC_VENV}/bin/activate
+  # TODO(disc): tensorflow-aarch64 wheel is not compatible with `libtao_ops.so`.
+  # Thus we use the wheel package provided by the ARM:
+  #   https://github.com/ARM-software/Tool-Solutions/tree/master/docker/tensorflow-aarch64
+  # See #244 for more details.
+  # if [[ ! -z "${DISC_HOST_TF_VERSION}" ]]; then
+  #   echo "install TensorFlow: "${DISC_HOST_TF_VERSION} "..."
+  #   pip install -q ${DISC_HOST_TF_VERSION}
+  # fi
+  pip install https://hlomodule.oss-cn-zhangjiakou.aliyuncs.com/bladedisc/aarch64/tensorflow-2.8.0-cp38-cp38-linux_aarch64.whl
+
+  pip install pytest pytest-forked
+  deactivate
+}
+
+install_venv
\ No newline at end of file
diff --git a/tao_compiler/mlir/xla/ral/context/mkldnn/ideep/ideep/operators/matmul.hpp b/tao_compiler/mlir/xla/ral/context/mkldnn/ideep/ideep/operators/matmul.hpp
index 876adb95190..752451ff0cb 100644
--- a/tao_compiler/mlir/xla/ral/context/mkldnn/ideep/ideep/operators/matmul.hpp
+++ b/tao_compiler/mlir/xla/ral/context/mkldnn/ideep/ideep/operators/matmul.hpp
@@ -53,10 +53,10 @@ struct matmul_forward : public dnnl::matmul,
       data_type x_dtype = data_type::f32,
       const engine& aengine = engine::cpu_engine()) {
     auto ndims = weights_dims.size();
-    auto x_dims = weights_dims;
+    dims x_dims = weights_dims;
     x_dims[ndims - 2] = 1;
     x_dims[ndims - 1] = weights_dims[ndims - 2];
-    auto y_dims = {x_dims[0], weights_dims[1]};
+    dims y_dims = {x_dims[0], weights_dims[1]};
     if (ndims == 3) y_dims = {x_dims[0], x_dims[1], weights_dims[2]};
     auto y_dtype = (dtype != data_type::s8) ? dtype : data_type::s32;
 


Test patch:
diff --git a/scripts/ci/build_and_test.sh b/scripts/ci/build_and_test.sh
index 64c47abb45e..bde801f2037 100644
--- a/scripts/ci/build_and_test.sh
+++ b/scripts/ci/build_and_test.sh
@@ -52,8 +52,4 @@ cp tao/dist/blade_disc*.whl ./build
 cp tf_community/bazel-bin/tensorflow/compiler/mlir/disc/tools/disc-replay/disc-replay-main ./build/
 
 # test example models
-arch=`uname -p`
-if [[ $arch == "x86_64" ]]; then
-  # TODO(disc): figure out the root cause of failure on aarch64
-  source ${SCRIPT_DIR}/test_cpu_examples.sh
-fi
+source ${SCRIPT_DIR}/test_cpu_examples.sh
diff --git a/scripts/ci/test_cpu_examples.sh b/scripts/ci/test_cpu_examples.sh
index 7d32c838c2b..012240fcfc1 100644
--- a/scripts/ci/test_cpu_examples.sh
+++ b/scripts/ci/test_cpu_examples.sh
@@ -13,6 +13,11 @@ set -ex
 if [ ! -z "$CPU_ONLY" ]; then
   # install disc python wheel
   python -m virtualenv venv && source venv/bin/activate
+  arch=`uname -p`
+  if [[ $arch == "aarch64" ]]; then
+    # TODO(disc): a workaround for issue #224
+    pip install https://hlomodule.oss-cn-zhangjiakou.aliyuncs.com/bladedisc/aarch64/tensorflow-2.8.0-cp38-cp38-linux_aarch64.whl
+  fi
   python -m pip install ./build/blade_disc*.whl
 
   pushd examples/TensorFlow/Inference/X86/BERT