diff --git a/.circleci/config.yml b/.circleci/config.yml
index 9ccba31..b7291ff 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -18,7 +18,7 @@ jobs:
             git submodule update --init --recursive
       # Use cached mlir installation if possible.
       - restore_cache:
-          key: V17-LLVM-PROJECT-{{ arch }}
+          key: V18-LLVM-PROJECT-{{ arch }}
       - run:
           name: Install MLIR
           command: |
@@ -29,7 +29,7 @@ jobs:
               source onnx-mlir/utils/install-mlir.sh
             fi
       - save_cache:
-          key: V17-LLVM-PROJECT-{{ arch }}
+          key: V18-LLVM-PROJECT-{{ arch }}
           paths:
             - llvm-project
       - run:
diff --git a/.travis.yml b/.travis.yml
index 525852a..2d2166e 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -11,9 +11,17 @@ arch:
 env:
   - NIGHTLY=true
   - NIGHTLY=false
-jobs:
-  allow_failures:
-    - env: NIGHTLY=true
+# Do not test branch update unless it's mater.
+branches:
+  only:
+  - master
+stages:
+  - name: prereq-init
+  - name: prereq-0
+  - name: prereq-1
+  - name: prereq-2
+  - name: prereq-final
+  - name: test
 script:
   - export CPU_ARCH=$TRAVIS_CPU_ARCH
   - if [ "$CPU_ARCH" == "amd64" ] ; then
@@ -42,11 +50,153 @@ script:
   - docker exec build compile-onnx-mlir.sh
   - docker exec build test-onnx-mlir.sh
   - docker commit build onnxmlirczar/onnx-mlir-build:$CPU_ARCH
-after_success: 
+jobs:
+  allow_failures:
+    - env: NIGHTLY=true
+  include:
+    ##########################################################################################
+    #
+    # Incrementally build s390 prereq docker images through 5 time-constrained steps.
+    #
+    ##########################################################################################
+    - stage: prereq-init
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - s390x
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:s390x -f docker/prereq.s390x.Dockerfile --build-arg BASE_IMAGE="ubuntu:focal" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:s390x
+    - stage: prereq-0
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - s390x
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:s390x -f docker/prereq.s390x.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:s390x" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:s390x
+    - stage: prereq-1
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - s390x
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:s390x -f docker/prereq.s390x.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:s390x" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:s390x
+    - stage: prereq-2
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - s390x
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:s390x -f docker/prereq.s390x.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:s390x" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:s390x
+    - stage: prereq-final
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - s390x
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage:s390x -f docker/prereq.s390x.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:s390x" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage:s390x
+
+    ##########################################################################################
+    #
+    # Incrementally build ppc64le prereq docker images through 5 time-constrained steps.
+    #
+    ##########################################################################################
+    - stage: prereq-init
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - ppc64le
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le -f docker/prereq.ppc64le.Dockerfile --build-arg BASE_IMAGE="ubuntu:bionic" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le
+    - stage: prereq-0
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - ppc64le
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le -f docker/prereq.ppc64le.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le
+    - stage: prereq-1
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - ppc64le
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le -f docker/prereq.ppc64le.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le
+    - stage: prereq-2
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - ppc64le
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le -f docker/prereq.ppc64le.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le
+    - stage: prereq-final
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - ppc64le
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage:ppc64le -f docker/prereq.ppc64le.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:ppc64le" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage:ppc64le
+
+    ##########################################################################################
+    #
+    # Incrementally build amd64 prereq docker images through 5 time-constrained steps.
+    #
+    ##########################################################################################
+    - stage: prereq-init
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - amd64
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:x86 -f docker/prereq.amd64.Dockerfile --build-arg BASE_IMAGE="ubuntu:focal" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:x86
+    - stage: prereq-0
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - amd64
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:x86 -f docker/prereq.amd64.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:x86" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:x86
+    - stage: prereq-1
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - amd64
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:x86 -f docker/prereq.amd64.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:x86" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:x86
+    - stage: prereq-2
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - amd64
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage-partial:x86 -f docker/prereq.amd64.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:x86" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage-partial:x86
+    - stage: prereq-final
+      if: commit_message = "Rebuild prereq docker."
+      arch:
+        - amd64
+      script:
+        - docker build -t onnxmlirczar/onnx-mlir-llvmimage:x86 -f docker/prereq.amd64.Dockerfile --build-arg BASE_IMAGE="onnxmlirczar/onnx-mlir-llvmimage-partial:x86" ./utils
+        - docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN"
+        - docker push onnxmlirczar/onnx-mlir-llvmimage:x86
+
+after_success:
   - if [ $TRAVIS_PULL_REQUEST == false ] && [ $TRAVIS_BRANCH == "master" ] && [ $NIGHTLY == false ]; then
             echo "Pushing new build to docker hub";
             docker login -u onnxmlirczar -p "$DOCKER_HUB_TOKEN";
             docker push onnxmlirczar/onnx-mlir-build:$CPU_ARCH;
-        else 
-            echo "Skip pushing build"; 
+        else
+            echo "Skip pushing build";
         fi
diff --git a/MLIR.cmake b/MLIR.cmake
index c8a7761..96882e9 100644
--- a/MLIR.cmake
+++ b/MLIR.cmake
@@ -253,8 +253,8 @@ set(MLIRLibs
         ${MLIRSideEffectInterfaces}
         ${MLIRStandardOps}
         ${MLIRStandardToLLVM}
-        ${MLIRSupport}
         ${MLIRTranslation}
+        ${MLIRSupport}
         ${MLIRLinalgOps}
         ${MLIRLinalgEDSC}
         ${MLIRLinalgAnalysis}
diff --git a/README.md b/README.md
index 151f55e..d7641d7 100644
--- a/README.md
+++ b/README.md
@@ -62,7 +62,7 @@ Firstly, install MLIR (as a part of LLVM-Project):
 ``` bash
 git clone https://github.com/llvm/llvm-project.git
 # Check out a specific branch that is known to work with ONNX MLIR.
-cd llvm-project && git checkout 9c94908320549a1a2328c758d6bbb694466021e7 && cd ..
+cd llvm-project && git checkout 1d01fc100bb5bef5f5eaf92520b2e52f64ee1d6e && cd ..
 ```
 
 [same-as-file]: <> (utils/build-mlir.sh)
@@ -152,7 +152,7 @@ Install MLIR (as a part of LLVM-Project):
 ```shell
 git clone https://github.com/llvm/llvm-project.git
 # Check out a specific branch that is known to work with ONNX MLIR.
-cd llvm-project && git checkout 9c94908320549a1a2328c758d6bbb694466021e7 && cd ..
+cd llvm-project && git checkout 1d01fc100bb5bef5f5eaf92520b2e52f64ee1d6e && cd ..
 ```
 
 [same-as-file]: <> (utils/build-mlir.cmd)
diff --git a/docker/prereq.amd64.Dockerfile b/docker/prereq.amd64.Dockerfile
index 5b184ad..94b48d0 100644
--- a/docker/prereq.amd64.Dockerfile
+++ b/docker/prereq.amd64.Dockerfile
@@ -1,43 +1,56 @@
-FROM ubuntu:focal
+ARG BASE_IMAGE
+
+# By default, use ubuntu:focal;
+FROM $BASE_IMAGE
 
 WORKDIR /build
 
 # install stuff that is needed for compiling LLVM, MLIR and ONNX
 RUN apt-get update
-RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git cmake ninja-build libprotobuf-dev protobuf-compiler 
+RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git cmake ninja-build libprotobuf-dev protobuf-compiler
 RUN DEBIAN_FRONTEND=noninteractive \
     apt-get update && \
     apt-get -y install \
         make build-essential libssl-dev zlib1g-dev libbz2-dev libreadline-dev \
         libsqlite3-dev wget curl llvm libncurses5-dev libncursesw5-dev xz-utils \
         libffi-dev liblzma-dev
-RUN git clone git://github.com/yyuu/pyenv.git .pyenv
-RUN git clone https://github.com/yyuu/pyenv-virtualenv.git ~/.pyenv/plugins/pyenv-virtualenv
+
+RUN if [ ! -d .pyenv ]; then \
+      git clone git://github.com/yyuu/pyenv.git .pyenv && \
+      git clone https://github.com/yyuu/pyenv-virtualenv.git ~/.pyenv/plugins/pyenv-virtualenv; \
+    fi
 
 ENV HOME=/build
 ENV PYENV_ROOT=$HOME/.pyenv
 ENV PATH=$PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH
 
-RUN pyenv install 3.7.0
-
-RUN pyenv global 3.7.0
-RUN pyenv rehash
+RUN if [ ! -f "./pyenv-installed" ]; then \
+      pyenv install 3.7.0 && \
+      pyenv global 3.7.0 && \
+      pyenv rehash && \
+      touch ./pyenv-installed; \
+    fi
 
 # first install MLIR in llvm-project
-RUN mkdir bin
+RUN mkdir -p bin
 ENV PATH=$PATH:/build/bin
 COPY clone-mlir.sh bin/clone-mlir.sh
 RUN chmod a+x bin/clone-mlir.sh
-RUN clone-mlir.sh
+RUN if [ ! -d /build/llvm-project ]; then \
+      clone-mlir.sh; \
+    fi
 
 WORKDIR /build/llvm-project/build
-RUN cmake -G Ninja ../llvm \
-   -DLLVM_ENABLE_PROJECTS=mlir \
-   -DLLVM_BUILD_EXAMPLES=ON \
-   -DLLVM_TARGETS_TO_BUILD="host" \
-   -DCMAKE_BUILD_TYPE=Release \
-   -DLLVM_ENABLE_ASSERTIONS=ON \
-   -DLLVM_ENABLE_RTTI=ON
+RUN if [ ! -f "/build/llvm-project/build/CMakeCache.txt" ]; then \
+      cmake -G Ninja ../llvm \
+       -DLLVM_ENABLE_PROJECTS=mlir \
+       -DLLVM_BUILD_EXAMPLES=ON \
+       -DLLVM_TARGETS_TO_BUILD="host" \
+       -DCMAKE_BUILD_TYPE=Release \
+       -DLLVM_ENABLE_ASSERTIONS=ON \
+       -DLLVM_ENABLE_RTTI=ON; \
+    fi
 
-RUN cmake --build . --target -- ${MAKEFLAGS}
-RUN cmake --build . --target check-mlir
+RUN if timeout 30m cmake --build . --target -- ${MAKEFLAGS} ; then \
+      cmake --build . --target check-mlir; \
+    fi
diff --git a/docker/prereq.ppc64le.Dockerfile b/docker/prereq.ppc64le.Dockerfile
index 5b184ad..94b48d0 100644
--- a/docker/prereq.ppc64le.Dockerfile
+++ b/docker/prereq.ppc64le.Dockerfile
@@ -1,43 +1,56 @@
-FROM ubuntu:focal
+ARG BASE_IMAGE
+
+# By default, use ubuntu:focal;
+FROM $BASE_IMAGE
 
 WORKDIR /build
 
 # install stuff that is needed for compiling LLVM, MLIR and ONNX
 RUN apt-get update
-RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git cmake ninja-build libprotobuf-dev protobuf-compiler 
+RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git cmake ninja-build libprotobuf-dev protobuf-compiler
 RUN DEBIAN_FRONTEND=noninteractive \
     apt-get update && \
     apt-get -y install \
         make build-essential libssl-dev zlib1g-dev libbz2-dev libreadline-dev \
         libsqlite3-dev wget curl llvm libncurses5-dev libncursesw5-dev xz-utils \
         libffi-dev liblzma-dev
-RUN git clone git://github.com/yyuu/pyenv.git .pyenv
-RUN git clone https://github.com/yyuu/pyenv-virtualenv.git ~/.pyenv/plugins/pyenv-virtualenv
+
+RUN if [ ! -d .pyenv ]; then \
+      git clone git://github.com/yyuu/pyenv.git .pyenv && \
+      git clone https://github.com/yyuu/pyenv-virtualenv.git ~/.pyenv/plugins/pyenv-virtualenv; \
+    fi
 
 ENV HOME=/build
 ENV PYENV_ROOT=$HOME/.pyenv
 ENV PATH=$PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH
 
-RUN pyenv install 3.7.0
-
-RUN pyenv global 3.7.0
-RUN pyenv rehash
+RUN if [ ! -f "./pyenv-installed" ]; then \
+      pyenv install 3.7.0 && \
+      pyenv global 3.7.0 && \
+      pyenv rehash && \
+      touch ./pyenv-installed; \
+    fi
 
 # first install MLIR in llvm-project
-RUN mkdir bin
+RUN mkdir -p bin
 ENV PATH=$PATH:/build/bin
 COPY clone-mlir.sh bin/clone-mlir.sh
 RUN chmod a+x bin/clone-mlir.sh
-RUN clone-mlir.sh
+RUN if [ ! -d /build/llvm-project ]; then \
+      clone-mlir.sh; \
+    fi
 
 WORKDIR /build/llvm-project/build
-RUN cmake -G Ninja ../llvm \
-   -DLLVM_ENABLE_PROJECTS=mlir \
-   -DLLVM_BUILD_EXAMPLES=ON \
-   -DLLVM_TARGETS_TO_BUILD="host" \
-   -DCMAKE_BUILD_TYPE=Release \
-   -DLLVM_ENABLE_ASSERTIONS=ON \
-   -DLLVM_ENABLE_RTTI=ON
+RUN if [ ! -f "/build/llvm-project/build/CMakeCache.txt" ]; then \
+      cmake -G Ninja ../llvm \
+       -DLLVM_ENABLE_PROJECTS=mlir \
+       -DLLVM_BUILD_EXAMPLES=ON \
+       -DLLVM_TARGETS_TO_BUILD="host" \
+       -DCMAKE_BUILD_TYPE=Release \
+       -DLLVM_ENABLE_ASSERTIONS=ON \
+       -DLLVM_ENABLE_RTTI=ON; \
+    fi
 
-RUN cmake --build . --target -- ${MAKEFLAGS}
-RUN cmake --build . --target check-mlir
+RUN if timeout 30m cmake --build . --target -- ${MAKEFLAGS} ; then \
+      cmake --build . --target check-mlir; \
+    fi
diff --git a/docker/prereq.s390x.Dockerfile b/docker/prereq.s390x.Dockerfile
index 2047b24..a408716 100644
--- a/docker/prereq.s390x.Dockerfile
+++ b/docker/prereq.s390x.Dockerfile
@@ -1,43 +1,57 @@
-FROM ubuntu:focal
+ARG BASE_IMAGE
+
+# By default, use ubuntu:focal;
+FROM $BASE_IMAGE
 
 WORKDIR /build
 
 # install stuff that is needed for compiling LLVM, MLIR and ONNX
 RUN apt-get update
-RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git cmake ninja-build libprotobuf-dev protobuf-compiler 
+RUN DEBIAN_FRONTEND=noninteractive apt-get install -y git cmake ninja-build libprotobuf-dev protobuf-compiler
 RUN DEBIAN_FRONTEND=noninteractive \
     apt-get update && \
     apt-get -y install \
         make build-essential libssl-dev zlib1g-dev libbz2-dev libreadline-dev \
         libsqlite3-dev wget curl llvm libncurses5-dev libncursesw5-dev xz-utils \
         libffi-dev liblzma-dev
-RUN git clone git://github.com/yyuu/pyenv.git .pyenv
-RUN git clone https://github.com/yyuu/pyenv-virtualenv.git ~/.pyenv/plugins/pyenv-virtualenv
+
+RUN if [ ! -d .pyenv ]; then \
+      git clone git://github.com/yyuu/pyenv.git .pyenv && \
+      git clone https://github.com/yyuu/pyenv-virtualenv.git ~/.pyenv/plugins/pyenv-virtualenv; \
+    fi
 
 ENV HOME=/build
 ENV PYENV_ROOT=$HOME/.pyenv
 ENV PATH=$PYENV_ROOT/shims:$PYENV_ROOT/bin:$PATH
 
-RUN pyenv install 3.7.0
-
-RUN pyenv global 3.7.0
-RUN pyenv rehash
+RUN if [ ! -f "./pyenv-installed" ]; then \
+      pyenv install 3.7.0 && \
+      pyenv global 3.7.0 && \
+      pyenv rehash && \
+      touch ./pyenv-installed; \
+    fi
 
 # first install MLIR in llvm-project
-RUN mkdir bin
+RUN mkdir -p bin
 ENV PATH=$PATH:/build/bin
 COPY clone-mlir.sh bin/clone-mlir.sh
 RUN chmod a+x bin/clone-mlir.sh
-RUN clone-mlir.sh
+RUN if [ ! -d /build/llvm-project ]; then \
+      clone-mlir.sh; \
+    fi
 
 WORKDIR /build/llvm-project/build
-RUN cmake -G Ninja ../llvm \
-   -DLLVM_ENABLE_PROJECTS=mlir \
-   -DLLVM_BUILD_EXAMPLES=ON \
-   -DLLVM_TARGETS_TO_BUILD="host" \
-   -DCMAKE_BUILD_TYPE=Release \
-   -DLLVM_ENABLE_ASSERTIONS=ON \
-   -DLLVM_ENABLE_RTTI=ON
+RUN if [ ! -f "/build/llvm-project/build/CMakeCache.txt" ]; then \
+      cmake -G Ninja ../llvm \
+       -DLLVM_ENABLE_PROJECTS=mlir \
+       -DLLVM_BUILD_EXAMPLES=ON \
+       -DLLVM_TARGETS_TO_BUILD="host" \
+       -DCMAKE_BUILD_TYPE=Release \
+       -DLLVM_ENABLE_ASSERTIONS=ON \
+       -DLLVM_ENABLE_RTTI=ON; \
+    fi
+
+# Build for 30 minutes:
+RUN timeout 30m cmake --build . --target -- ${MAKEFLAGS} || true
 
-RUN cmake --build . --target -- ${MAKEFLAGS}
 # RUN cmake --build . --target check-mlir
diff --git a/docs/README.md b/docs/README.md
index e9cf1f4..8885378 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -20,7 +20,7 @@ Firstly, install MLIR (as a part of LLVM-Project):
 ``` bash
 git clone https://github.com/llvm/llvm-project.git
 # Check out a specific branch that is known to work with ONNX MLIR.
-cd llvm-project && git checkout 9c94908320549a1a2328c758d6bbb694466021e7 && cd ..
+cd llvm-project && git checkout 1d01fc100bb5bef5f5eaf92520b2e52f64ee1d6e && cd ..
 ```
 
 [same-as-file]: <> (utils/build-mlir.sh)
@@ -110,7 +110,7 @@ Install MLIR (as a part of LLVM-Project):
 ```shell
 git clone https://github.com/llvm/llvm-project.git
 # Check out a specific branch that is known to work with ONNX MLIR.
-cd llvm-project && git checkout 9c94908320549a1a2328c758d6bbb694466021e7 && cd ..
+cd llvm-project && git checkout 1d01fc100bb5bef5f5eaf92520b2e52f64ee1d6e && cd ..
 ```
 
 [same-as-file]: <> (utils/build-mlir.cmd)
diff --git a/src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.cpp b/src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.cpp
index 76bddc8..1d68774 100644
--- a/src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.cpp
+++ b/src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.cpp
@@ -309,162 +309,167 @@ std::vector<Value> getLoopIVsForBroadcasting(Location loc,
 Value emitConstantOp(
     PatternRewriter &rewriter, Location loc, Type type, double value) {
   Attribute constantAttr;
-  auto typeKind = type.getKind();
-  if (typeKind == StandardTypes::F16) {
-    constantAttr = rewriter.getF16FloatAttr((float)value);
-  } else if (typeKind == StandardTypes::F32) {
-    constantAttr = rewriter.getF32FloatAttr((float)value);
-  } else if (typeKind == StandardTypes::F64) {
-    constantAttr = rewriter.getF64FloatAttr(value);
-  } else if (typeKind == StandardTypes::Integer) {
-    auto width = type.cast<IntegerType>().getWidth();
-    if (width == 1) {
-      constantAttr = rewriter.getBoolAttr(false);
-    } else {
-      constantAttr =
-          rewriter.getIntegerAttr(type, APInt(width, (int64_t)value));
-    }
-  } else if (typeKind == StandardTypes::Index) {
-    constantAttr = rewriter.getIntegerAttr(type, (int64_t)value);
-  } else {
-    llvm_unreachable("unsupported element type");
-  }
 
+  TypeSwitch<Type>(type)
+      .Case<Float16Type>(
+          [&](Type) { constantAttr = rewriter.getF16FloatAttr((float)value); })
+      .Case<Float32Type>(
+          [&](Type) { constantAttr = rewriter.getF32FloatAttr((float)value); })
+      .Case<Float64Type>(
+          [&](Type) { constantAttr = rewriter.getF64FloatAttr((float)value); })
+      .Case<IntegerType>([&](Type) {
+        auto width = type.cast<IntegerType>().getWidth();
+        if (width == 1) {
+          constantAttr = rewriter.getBoolAttr(false);
+        } else {
+          constantAttr =
+              rewriter.getIntegerAttr(type, APInt(width, (int64_t)value));
+        }
+      })
+      .Case<IndexType>([&](Type) {
+        constantAttr = rewriter.getIntegerAttr(type, (int64_t)value);
+      })
+      .Default([](Type) { llvm_unreachable("unsupported element type"); });
   return rewriter.create<ConstantOp>(loc, constantAttr);
 }
 
 Value emitPositiveInfinityConstantOp(
     ConversionPatternRewriter &rewriter, Location loc, Type type) {
   Attribute constantAttr;
-  auto typeKind = type.getKind();
-  if (typeKind == StandardTypes::F16) {
-    // 0x7C00
-    float value = std::numeric_limits<float>::infinity();
-    constantAttr = rewriter.getF16FloatAttr(value);
-  } else if (typeKind == StandardTypes::F32) {
-    // 0x7F800000
-    float value = std::numeric_limits<float>::infinity();
-    constantAttr = rewriter.getF32FloatAttr(value);
-  } else if (typeKind == StandardTypes::F64) {
-    // 0x7FF0000000000000
-    double value = std::numeric_limits<double>::infinity();
-    constantAttr = rewriter.getF64FloatAttr(value);
-  } else if (typeKind == StandardTypes::Integer) {
-    auto width = type.cast<IntegerType>().getWidth();
-    // The latest llvm-project includes a patch which allows getting the sign of
-    // IntegerType:
-    // https://github.com/llvm/llvm-project/commit/35b685270b410f6a1351c2a527021f22330c25b9
-    // as follows:
-    //   auto isSigned = type.cast<IntegerType>().isSigned();
-    // TODO (tungld): update the following statement once our llvm-project is
-    // upgraded to include the patch.
-    auto isSigned = true;
-    if (width == 8) {
-      if (isSigned) {
-        int8_t value = std::numeric_limits<int8_t>::max();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      } else {
-        uint8_t value = std::numeric_limits<uint8_t>::max();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      }
-    } else if (width == 16) {
-      if (isSigned) {
-        int16_t value = std::numeric_limits<int16_t>::max();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      } else {
-        uint16_t value = std::numeric_limits<uint16_t>::max();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      }
-    } else if (width == 32) {
-      if (isSigned) {
-        int32_t value = std::numeric_limits<int32_t>::max();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      } else {
-        uint32_t value = std::numeric_limits<uint32_t>::max();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      }
-    } else if (width == 64) {
-      if (isSigned) {
-        int64_t value = std::numeric_limits<int64_t>::max();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      } else {
-        uint64_t value = std::numeric_limits<uint64_t>::max();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      }
-    } else {
-      llvm_unreachable("unsupported element type");
-    }
-  } else {
-    llvm_unreachable("unsupported element type");
-  }
 
+  TypeSwitch<Type>(type)
+      .Case<Float16Type>([&](Type) {
+        // 0x7C00
+        float value = std::numeric_limits<float>::infinity();
+        constantAttr = rewriter.getF16FloatAttr(value);
+      })
+      .Case<Float32Type>([&](Type) {
+        // 0x7F800000
+        float value = std::numeric_limits<float>::infinity();
+        constantAttr = rewriter.getF32FloatAttr(value);
+      })
+      .Case<Float64Type>([&](Type) {
+        // 0x7FF0000000000000
+        double value = std::numeric_limits<double>::infinity();
+        constantAttr = rewriter.getF64FloatAttr(value);
+      })
+      .Case<IntegerType>([&](Type) {
+        auto width = type.cast<IntegerType>().getWidth();
+        // The latest llvm-project includes a patch which allows getting the
+        // sign of IntegerType:
+        // https://github.com/llvm/llvm-project/commit/35b685270b410f6a1351c2a527021f22330c25b9
+        // as follows:
+        //   auto isSigned = type.cast<IntegerType>().isSigned();
+        // TODO (tungld): update the following statement once our llvm-project
+        // is upgraded to include the patch.
+        auto isSigned = true;
+        if (width == 8) {
+          if (isSigned) {
+            int8_t value = std::numeric_limits<int8_t>::max();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          } else {
+            uint8_t value = std::numeric_limits<uint8_t>::max();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          }
+        } else if (width == 16) {
+          if (isSigned) {
+            int16_t value = std::numeric_limits<int16_t>::max();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          } else {
+            uint16_t value = std::numeric_limits<uint16_t>::max();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          }
+        } else if (width == 32) {
+          if (isSigned) {
+            int32_t value = std::numeric_limits<int32_t>::max();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          } else {
+            uint32_t value = std::numeric_limits<uint32_t>::max();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          }
+        } else if (width == 64) {
+          if (isSigned) {
+            int64_t value = std::numeric_limits<int64_t>::max();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          } else {
+            uint64_t value = std::numeric_limits<uint64_t>::max();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          }
+        } else {
+          llvm_unreachable("unsupported element type");
+        }
+      })
+      .Default([](Type) { llvm_unreachable("unsupported element type"); });
   return rewriter.create<ConstantOp>(loc, constantAttr);
 }
 
 Value emitNegativeInfinityConstantOp(
     ConversionPatternRewriter &rewriter, Location loc, Type type) {
   Attribute constantAttr;
-  auto typeKind = type.getKind();
-  if (typeKind == StandardTypes::F16) {
-    // 0xFC00
-    float value = -std::numeric_limits<float>::infinity();
-    constantAttr = rewriter.getF16FloatAttr(value);
-  } else if (typeKind == StandardTypes::F32) {
-    // 0xFF800000
-    float value = -std::numeric_limits<float>::infinity();
-    constantAttr = rewriter.getF32FloatAttr(value);
-  } else if (typeKind == StandardTypes::F64) {
-    // 0xFFF0000000000000
-    double value = -std::numeric_limits<double>::infinity();
-    constantAttr = rewriter.getF64FloatAttr(value);
-  } else if (typeKind == StandardTypes::Integer) {
-    auto width = type.cast<IntegerType>().getWidth();
-    // The latest llvm-project includes a patch which allows getting the sign of
-    // IntegerType:
-    // https://github.com/llvm/llvm-project/commit/35b685270b410f6a1351c2a527021f22330c25b9
-    // as follows:
-    //   auto isSigned = type.cast<IntegerType>().isSigned();
-    // TODO (tungld): update the following statement once our llvm-project is
-    // upgraded to include the patch.
-    auto isSigned = true;
-    if (width == 8) {
-      if (isSigned) {
-        int8_t value = std::numeric_limits<int8_t>::min();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      } else {
-        uint8_t value = std::numeric_limits<uint8_t>::min();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      }
-    } else if (width == 16) {
-      if (isSigned) {
-        int16_t value = std::numeric_limits<int16_t>::min();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      } else {
-        uint16_t value = std::numeric_limits<uint16_t>::min();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      }
-    } else if (width == 32) {
-      if (isSigned) {
-        int32_t value = std::numeric_limits<int32_t>::min();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      } else {
-        uint32_t value = std::numeric_limits<uint32_t>::min();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      }
-    } else if (width == 64) {
-      if (isSigned) {
-        int64_t value = std::numeric_limits<int64_t>::min();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      } else {
-        uint64_t value = std::numeric_limits<uint64_t>::min();
-        constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
-      }
-    } else {
-      llvm_unreachable("unsupported element type");
-    }
-  } else {
-    llvm_unreachable("unsupported element type");
-  }
+
+  TypeSwitch<Type>(type)
+      .Case<Float16Type>([&](Type) {
+        // 0xFC00
+        float value = -std::numeric_limits<float>::infinity();
+        constantAttr = rewriter.getF16FloatAttr(value);
+      })
+      .Case<Float32Type>([&](Type) {
+        // 0xFF800000
+        float value = -std::numeric_limits<float>::infinity();
+        constantAttr = rewriter.getF32FloatAttr(value);
+      })
+      .Case<Float64Type>([&](Type) {
+        // 0xFFF0000000000000
+        double value = -std::numeric_limits<double>::infinity();
+        constantAttr = rewriter.getF64FloatAttr(value);
+      })
+      .Case<IntegerType>([&](Type) {
+        auto width = type.cast<IntegerType>().getWidth();
+        // The latest llvm-project includes a patch which allows getting the
+        // sign of IntegerType:
+        // https://github.com/llvm/llvm-project/commit/35b685270b410f6a1351c2a527021f22330c25b9
+        // as follows:
+        //   auto isSigned = type.cast<IntegerType>().isSigned();
+        // TODO (tungld): update the following statement once our llvm-project
+        // is upgraded to include the patch.
+        auto isSigned = true;
+        if (width == 8) {
+          if (isSigned) {
+            int8_t value = std::numeric_limits<int8_t>::min();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          } else {
+            uint8_t value = std::numeric_limits<uint8_t>::min();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          }
+        } else if (width == 16) {
+          if (isSigned) {
+            int16_t value = std::numeric_limits<int16_t>::min();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          } else {
+            uint16_t value = std::numeric_limits<uint16_t>::min();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          }
+        } else if (width == 32) {
+          if (isSigned) {
+            int32_t value = std::numeric_limits<int32_t>::min();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          } else {
+            uint32_t value = std::numeric_limits<uint32_t>::min();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          }
+        } else if (width == 64) {
+          if (isSigned) {
+            int64_t value = std::numeric_limits<int64_t>::min();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          } else {
+            uint64_t value = std::numeric_limits<uint64_t>::min();
+            constantAttr = rewriter.getIntegerAttr(type, APInt(width, value));
+          }
+        } else {
+          llvm_unreachable("unsupported element type");
+        }
+      })
+      .Default([](Type) { llvm_unreachable("unsupported element type"); });
 
   return rewriter.create<ConstantOp>(loc, constantAttr);
 }
diff --git a/src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp b/src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp
index fe1c74c..e9e9416 100644
--- a/src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp
+++ b/src/Conversion/ONNXToKrnl/ONNXToKrnlCommon.hpp
@@ -20,6 +20,7 @@
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Sequence.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 #include "src/Dialect/Krnl/KrnlHelper.hpp"
 #include "src/Dialect/Krnl/KrnlOps.hpp"
diff --git a/src/Dialect/Krnl/KrnlOps.cpp b/src/Dialect/Krnl/KrnlOps.cpp
index f15c21e..05fc53b 100644
--- a/src/Dialect/Krnl/KrnlOps.cpp
+++ b/src/Dialect/Krnl/KrnlOps.cpp
@@ -34,7 +34,7 @@ using namespace mlir;
 
 namespace mlir {
 KrnlOpsDialect::KrnlOpsDialect(MLIRContext *context)
-    : Dialect(getDialectNamespace(), context) {
+    : Dialect(getDialectNamespace(), context, TypeID::get<KrnlOpsDialect>()) {
   addOperations<
 #define GET_OP_LIST
 #include "src/Dialect/Krnl/KrnlOps.cpp.inc"
diff --git a/src/Dialect/Krnl/KrnlOps.hpp b/src/Dialect/Krnl/KrnlOps.hpp
index 9460f00..04f0c5d 100644
--- a/src/Dialect/Krnl/KrnlOps.hpp
+++ b/src/Dialect/Krnl/KrnlOps.hpp
@@ -17,6 +17,7 @@
 #include "mlir/IR/Function.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/StandardTypes.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 #include "KrnlHelper.hpp"
 #include "KrnlTypes.hpp"
@@ -38,11 +39,10 @@ public:
 
   /// Print a type registered to this dialect.
   void printType(Type type, DialectAsmPrinter &os) const override {
-    switch (type.getKind()) {
-    case KrnlTypes::Loop:
+    TypeSwitch<Type>(type).Case<LoopType>([&](Type) {
       os << "loop";
       return;
-    }
+    });
   }
 };
 
diff --git a/src/Dialect/Krnl/KrnlTypes.hpp b/src/Dialect/Krnl/KrnlTypes.hpp
index 4adc1c6..2cc9130 100644
--- a/src/Dialect/Krnl/KrnlTypes.hpp
+++ b/src/Dialect/Krnl/KrnlTypes.hpp
@@ -13,16 +13,6 @@
 #include <mlir/IR/Types.h>
 
 namespace mlir {
-
-namespace KrnlTypes {
-enum Kinds {
-  // A krnl.loop is simply a reference to a for loop and will be used to:
-  // - Indicate the presence of a for loop in krnl.iterate.
-  // - Identify the loop in optimization intrinsics.
-  Loop = mlir::Type::Kind::FIRST_PRIVATE_EXPERIMENTAL_0_TYPE,
-};
-}
-
 class LoopType
     : public mlir::Type::TypeBase<LoopType, mlir::Type, mlir::TypeStorage> {
 
@@ -30,11 +20,8 @@ public:
   using Base::Base;
 
   // Support type inquiry through isa, cast and dyn_cast.
-  static bool kindof(unsigned kind) { return kind == KrnlTypes::Loop; }
 
   // Get a unique instance of Loop type.
-  static LoopType get(mlir::MLIRContext *context) {
-    return Base::get(context, KrnlTypes::Loop);
-  }
+  static LoopType get(mlir::MLIRContext *context) { return Base::get(context); }
 };
 } // namespace mlir
diff --git a/src/Dialect/ONNX/ONNXOps.cpp b/src/Dialect/ONNX/ONNXOps.cpp
index c18cf6a..ac680b4 100644
--- a/src/Dialect/ONNX/ONNXOps.cpp
+++ b/src/Dialect/ONNX/ONNXOps.cpp
@@ -478,7 +478,7 @@ static void insertConvTransposeSpatialDim(SmallVectorImpl<int64_t> &outputDims,
 /// Dialect creation, the instance will be owned by the context. This is the
 /// point of registration of custom types and operations for the dialect.
 ONNXOpsDialect::ONNXOpsDialect(mlir::MLIRContext *ctx)
-    : mlir::Dialect(getDialectNamespace(), ctx) {
+    : mlir::Dialect(getDialectNamespace(), ctx, TypeID::get<ONNXOpsDialect>()) {
   addOperations<
 #define GET_OP_LIST
 #include "src/Dialect/ONNX/ONNXOps.cpp.inc"
@@ -2885,7 +2885,7 @@ struct SeqTypeStorage : public mlir::TypeStorage {
 SeqType SeqType::get(llvm::ArrayRef<mlir::Type> elementTypes) {
   assert(!elementTypes.empty() && "expected non-empty seq");
   mlir::MLIRContext *ctx = elementTypes.front().getContext();
-  return Base::get(ctx, ONNXTypes::SEQ, elementTypes);
+  return Base::get(ctx, elementTypes);
 }
 
 llvm::ArrayRef<mlir::Type> SeqType::getElementTypes() {
diff --git a/src/Dialect/ONNX/ONNXOps.hpp b/src/Dialect/ONNX/ONNXOps.hpp
index f54a92d..f70a83c 100644
--- a/src/Dialect/ONNX/ONNXOps.hpp
+++ b/src/Dialect/ONNX/ONNXOps.hpp
@@ -54,30 +54,12 @@ public:
 // Or we need two namespace?
 // Will put all the ONNXOps into this namespace
 namespace onnxmlir {
-
-namespace ONNXTypes {
-
-enum Kind {
-  FIRST_USED_ONNX_TYPE = Type::FIRST_PRIVATE_EXPERIMENTAL_1_TYPE,
-  //#define HANDLE_TF_TYPE(tftype, enumerant, name) enumerant,
-  //#include "src/Dialect/ONNX/ONXTypes.def"
-  STRING,
-  SEQ,
-  LAST_USED_ONNX_TYPE,
-};
-} // namespace ONNXTypes
-
 class StringType
     : public mlir::Type::TypeBase<StringType, mlir::Type, mlir::TypeStorage> {
 public:
   using Base::Base;
-  static bool kindof(unsigned kind) { return kind == ONNXTypes::STRING; }
 
-  static unsigned getTypeKind() { return ONNXTypes::STRING; }
-
-  static StringType get(MLIRContext *ctx) {
-    return Base::get(ctx, ONNXTypes::STRING);
-  }
+  static StringType get(MLIRContext *ctx) { return Base::get(ctx); }
 };
 
 namespace detail {
@@ -88,9 +70,6 @@ class SeqType
     : public mlir::Type::TypeBase<SeqType, mlir::Type, detail::SeqTypeStorage> {
 public:
   using Base::Base;
-  static bool kindof(unsigned kind) { return kind == ONNXTypes::SEQ; }
-
-  static unsigned getTypeKind() { return ONNXTypes::SEQ; }
 
   static SeqType get(llvm::ArrayRef<mlir::Type> elementTypes);
 
diff --git a/src/MainUtils.cpp b/src/MainUtils.cpp
index 4dfe1de..35e22ac 100644
--- a/src/MainUtils.cpp
+++ b/src/MainUtils.cpp
@@ -376,14 +376,15 @@ void compileModuleToJniJar(
   genJniJar(module, modelSharedLibPath, modelJniJarPath);
 }
 
-void registerDialects() {
-  mlir::registerDialect<mlir::AffineDialect>();
-  mlir::registerDialect<mlir::LLVM::LLVMDialect>();
-  mlir::registerDialect<mlir::scf::SCFDialect>();
-  mlir::registerDialect<mlir::StandardOpsDialect>();
-  mlir::registerDialect<mlir::shape::ShapeDialect>();
-  mlir::registerDialect<mlir::ONNXOpsDialect>();
-  mlir::registerDialect<mlir::KrnlOpsDialect>();
+void registerDialects(mlir::MLIRContext &context) {
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::AffineDialect>();
+  context.getOrLoadDialect<mlir::LLVM::LLVMDialect>();
+  context.getOrLoadDialect<mlir::scf::SCFDialect>();
+  context.getOrLoadDialect<mlir::StandardOpsDialect>();
+  context.getOrLoadDialect<mlir::shape::ShapeDialect>();
+  context.getOrLoadDialect<mlir::ONNXOpsDialect>();
+  context.getOrLoadDialect<mlir::KrnlOpsDialect>();
 }
 
 void addONNXToMLIRPasses(mlir::PassManager &pm) {
diff --git a/src/MainUtils.hpp b/src/MainUtils.hpp
index 1155d08..346a20a 100644
--- a/src/MainUtils.hpp
+++ b/src/MainUtils.hpp
@@ -59,7 +59,7 @@ void compileModuleToSharedLibrary(
 void compileModuleToJniJar(
     const mlir::OwningModuleRef &module, std::string outputBaseName);
 
-void registerDialects();
+void registerDialects(mlir::MLIRContext &context);
 
 void addONNXToMLIRPasses(mlir::PassManager &pm);
 
diff --git a/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp b/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp
index 2f78bfd..f35f262 100644
--- a/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp
+++ b/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp
@@ -60,13 +60,14 @@ static llvm::cl::opt<bool> allowUnregisteredDialects(
     llvm::cl::init(false));
 
 int main(int argc, char **argv) {
-  mlir::registerDialect<mlir::linalg::LinalgDialect>();
-  mlir::registerDialect<mlir::AffineDialect>();
-  mlir::registerDialect<mlir::LLVM::LLVMDialect>();
-  mlir::registerDialect<mlir::scf::SCFDialect>();
-  mlir::registerDialect<mlir::StandardOpsDialect>();
-  mlir::registerDialect<mlir::vector::VectorDialect>();
-  mlir::registerDialect<mlir::shape::ShapeDialect>();
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::linalg::LinalgDialect>();
+  registry.insert<mlir::AffineDialect>();
+  registry.insert<mlir::LLVM::LLVMDialect>();
+  registry.insert<mlir::scf::SCFDialect>();
+  registry.insert<mlir::StandardOpsDialect>();
+  registry.insert<mlir::vector::VectorDialect>();
+  registry.insert<mlir::shape::ShapeDialect>();
 
   registerTransformsPasses();
   registerAffinePasses();
@@ -97,7 +98,8 @@ int main(int argc, char **argv) {
   auto output = mlir::openOutputFile(output_filename, &error_message);
   assert(output);
 
+  // TODO(imaihal): Change preloadDialectsInContext to false.
   return failed(mlir::MlirOptMain(output->os(), std::move(file), passPipeline,
-      split_input_file, verify_diagnostics, verify_passes,
-      allowUnregisteredDialects));
+      registry, split_input_file, verify_diagnostics, verify_passes,
+      allowUnregisteredDialects, /*preloadDialectsInContext*/ true));
 }
diff --git a/src/main.cpp b/src/main.cpp
index 5585c07..d273af1 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -13,7 +13,8 @@ using namespace onnx_mlir;
 
 int main(int argc, char *argv[]) {
   setExecPath(argv[0], (void *)main);
-  registerDialects();
+  mlir::MLIRContext context;
+  registerDialects(context);
 
   llvm::cl::OptionCategory OnnxMlirOptions(
       "ONNX MLIR Options", "These are frontend options.");
@@ -43,7 +44,6 @@ int main(int argc, char *argv[]) {
   llvm::cl::ParseCommandLineOptions(
       argc, argv, "ONNX MLIR modular optimizer driver\n");
 
-  mlir::MLIRContext context;
   mlir::OwningModuleRef module;
   processInputFile(inputFilename, emissionTarget, context, module);
 
diff --git a/test/numerical/TestConv.cpp b/test/numerical/TestConv.cpp
index c880afa..6b77026 100644
--- a/test/numerical/TestConv.cpp
+++ b/test/numerical/TestConv.cpp
@@ -23,8 +23,8 @@ using namespace std;
 bool isOMConvTheSameAsNaiveImplFor(const int N, const int C, const int H,
     const int W, const int kH, const int kW, const int pHBegin, const int pHEnd,
     const int pWBegin, const int pWEnd) {
-  registerDialects();
   MLIRContext ctx;
+  registerDialects(ctx);
 
   auto module = ModuleOp::create(UnknownLoc::get(&ctx));
   OpBuilder builder(&ctx);
diff --git a/utils/clone-mlir.sh b/utils/clone-mlir.sh
index 7c1e059..b7db56e 100644
--- a/utils/clone-mlir.sh
+++ b/utils/clone-mlir.sh
@@ -1,3 +1,3 @@
 git clone https://github.com/llvm/llvm-project.git
 # Check out a specific branch that is known to work with ONNX MLIR.
-cd llvm-project && git checkout 9c94908320549a1a2328c758d6bbb694466021e7 && cd ..
+cd llvm-project && git checkout 1d01fc100bb5bef5f5eaf92520b2e52f64ee1d6e && cd ..