diff --git a/flang/unittests/Lower/OpenMPLoweringTest.cpp b/flang/unittests/Lower/OpenMPLoweringTest.cpp
index ad6fe739d16b..dc002cef1a52 100644
--- a/flang/unittests/Lower/OpenMPLoweringTest.cpp
+++ b/flang/unittests/Lower/OpenMPLoweringTest.cpp
@@ -9,14 +9,15 @@
 #include "gtest/gtest.h"
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/IR/Builders.h"
+#include "mlir/InitAllDialects.h"
 #include "flang/Parser/parse-tree.h"
 #include "llvm/Frontend/OpenMP/OMPConstants.h"
 
 class OpenMPLoweringTest : public testing::Test {
 protected:
   void SetUp() override {
-    mlir::registerDialect<mlir::omp::OpenMPDialect>();
-    mlir::registerAllDialects(&ctx);
+    ctx.loadDialect<mlir::omp::OpenMPDialect>();
+    mlir::registerAllDialects(ctx.getDialectRegistry());
     mlirOpBuilder.reset(new mlir::OpBuilder(&ctx));
   }