Update llvm commit ID to 1d01fc1 (#292)

* Fix for LLVM revision D85495 * Fix for LLVM revision DD86121 * Fix for LLVM revision D85622 (f9dc2b7) TODO: Change preloadDialectsInContext to false Memo for previous fixes: D86121 (250f43d), D85495 (575b22b) * clang-format * Update llvm commit ID of README and clone-mlir.sh * Updated llvm commit ID of README.md * Fix for passing backend tests * Removed the commented code * Empty commit for triggering rebuild * Test multi-stage travis build * Specify stage order. * Empty commit for triggering rebuild * Update prereq.s390x.Dockerfile Make it possible to execute s390x prereq docker multiple times. * Build prereq for each arch * Fix multi-arch prereq build. * timeout at 40m * Update .travis.yml * add ppc64le prereq builder * Run ppc docker prereq build multiple times * Do not test branch update unless it's mater. * Fix dockerfile. * Fix typo in travis.yml. * Fix ppc64 docker file * Update .travis.yml * turn off metacopy on ppc64le * Update .travis.yml * Turn off metacopy. * Turn off metacopy inside Dockerfile in ppc64. * No sudo in Docker. * Remove metacopy config from Dockerfile. * Change base image to be bionic. * Using newer linux distro for ppc64. * Turn off metacopy in before_install. * Fix sudo permission issue. * Run docker info. * Allow amd64 docker file to be built multiple times * Support building amd64 prereq. * Fix amd64 docker file typo. * fix ppc64le dockerfile typo. * timeout from 40m -> 30m * 40m->30m * 40m->30m * fix bug preventing incremental build. * fix bug preventing incremental build. * Bump CircleCI cache version. * Push to production prereq container repository and condition prereq docker rebuild on commit message. * Rebuild prereq docker. * Move default script to top-level. * Python not properly installed. * amd64 -> x86 * Rebuild prereq docker. * Rebuild prereq docker. * Rebuild prereq docker. * Restart all CI. * Disallow cache on Jenkins docker build. * Restart zJenkins. * Restart zJenkins. Co-authored-by: Haruki Imai <imaihal@jp.ibm.com> Co-authored-by: Alexandre Eichenberger <alexe@us.ibm.com>
2020-09-09 23:12:01 +08:00 · 2020-09-09 23:12:01 +08:00 · dbc41d2330
parent c738e0fb07
commit dbc41d2330
21 changed files with 440 additions and 275 deletions
--- 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:
--- 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
--- a/MLIR.cmake
+++ b/MLIR.cmake
@ -253,8 +253,8 @@ set(MLIRLibs
        ${MLIRSideEffectInterfaces}
        ${MLIRStandardOps}
        ${MLIRStandardToLLVM}
-        ${MLIRSupport}
        ${MLIRTranslation}
+        ${MLIRSupport}
        ${MLIRLinalgOps}
        ${MLIRLinalgEDSC}
        ${MLIRLinalgAnalysis}
--- 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)
--- 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
--- 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
--- 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
--- 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)
--- 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);
 }
--- 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"
--- 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"
--- 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;
-    }
+    });
  }
 };

--- 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
--- 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() {
--- 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);

--- 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) {
--- 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);

--- 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));
 }
--- 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);

--- 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);
--- 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 ..